From f73eddb2f0accace3b6238bf04692e0c3d6376e5 Mon Sep 17 00:00:00 2001 From: Jeremydream Date: Sun, 7 Apr 2024 13:01:14 +0800 Subject: [PATCH] [Enhancement] Improved EANN Model with Hard Cutoff for Efficient Interaction Handling Enabled hard cutoff in EANN model, aligning with cutoff rules in XML file. Allows processing of larger atom pairs, minimizing unnecessary interaction counts. --- dmff/eann/eann.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dmff/eann/eann.py b/dmff/eann/eann.py index dc1edd5ca..1849aeaa5 100644 --- 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']