Merge pull request #178 from Jeremydream/devel

[Enhancement] Improved EANN Model with Hard Cutoff for Efficient Inte…
deepmodeling · Apr 7, 2024 · 90cfca1 · 90cfca1
2 parents 8c4c767 + f73eddb
commit 90cfca1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dmff/eann/eann.py b/dmff/eann/eann.py
@@ -307,6 +307,9 @@ def get_energy(positions, box, pairs, params):
             dr = pbc_shift(dr, box, box_inv)
 
             dr_norm = jnp.linalg.norm(dr, axis=1)
+            buffer_scales2 = jnp.piecewise(buffer_scales, (dr_norm <= self.rc, dr_norm > self.rc),
+                              (lambda x: jnp.array(1), lambda x: jnp.array(0)))
+            buffer_scales = buffer_scales2 * buffer_scales
             self.rs = params['density.rs']
             self.inta = params['density.inta']