From 533981f4d94c47e14e17e62a1541c5c4a99d7770 Mon Sep 17 00:00:00 2001 From: WangXinyan940 Date: Tue, 15 Nov 2022 16:03:43 +0800 Subject: [PATCH 1/4] Hook jax force to generator for intra potentials --- dmff/generators/classical.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dmff/generators/classical.py b/dmff/generators/classical.py index b4584ba9a..d13e244cb 100644 --- a/dmff/generators/classical.py +++ b/dmff/generators/classical.py @@ -97,6 +97,7 @@ def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args): map_param = np.array(map_param, dtype=int) bforce = HarmonicBondJaxForce(map_atom1, map_atom2, map_param) + self._force_latest = bforce def potential_fn(positions, box, pairs, params): return bforce.get_energy(positions, box, pairs, @@ -163,6 +164,7 @@ def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args): aforce = HarmonicAngleJaxForce(map_atom1, map_atom2, map_atom3, map_param) + self._force_latest = aforce def potential_fn(positions, box, pairs, params): return aforce.get_energy(positions, box, pairs, @@ -384,6 +386,8 @@ def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args): jnp.array(map_impr_param[p], dtype=int), p) for p in range(1, self.max_pred_impr + 1) ] + self._props_latest = props + self._imprs_latest = imprs def potential_fn(positions, box, pairs, params): prop_sum = sum([ From 9fcd2fcbcebff3fcdcf24c1ccf7ac893eee336d5 Mon Sep 17 00:00:00 2001 From: WangXinyan940 Date: Tue, 22 Nov 2022 12:29:07 +0800 Subject: [PATCH 2/4] Let NeighborList.update return (N, 3) pairs with colv_map information --- dmff/common/nblist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmff/common/nblist.py b/dmff/common/nblist.py index c57b481be..be4411d30 100644 --- a/dmff/common/nblist.py +++ b/dmff/common/nblist.py @@ -44,7 +44,7 @@ def allocate(self, positions: jnp.ndarray, box: Optional[jnp.ndarray] = None): self.nblist = self.neighborlist_fn.allocate(positions) else: self.update(positions, box) - return self.nblist + return self.pairs() def update(self, positions: jnp.ndarray, box: Optional[jnp.ndarray] = None): """ A function to update a neighbor list given a new set of positions and a previously allocated neighbor list. From d0cd7a3ac4c73319b7b86a4feb9e5e33d1d7c3f0 Mon Sep 17 00:00:00 2001 From: WangXinyan940 Date: Tue, 22 Nov 2022 12:31:42 +0800 Subject: [PATCH 3/4] Update nblist.py --- dmff/common/nblist.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dmff/common/nblist.py b/dmff/common/nblist.py index be4411d30..1b06631ec 100644 --- a/dmff/common/nblist.py +++ b/dmff/common/nblist.py @@ -44,7 +44,7 @@ def allocate(self, positions: jnp.ndarray, box: Optional[jnp.ndarray] = None): self.nblist = self.neighborlist_fn.allocate(positions) else: self.update(positions, box) - return self.pairs() + return self.pairs def update(self, positions: jnp.ndarray, box: Optional[jnp.ndarray] = None): """ A function to update a neighbor list given a new set of positions and a previously allocated neighbor list. @@ -59,7 +59,7 @@ def update(self, positions: jnp.ndarray, box: Optional[jnp.ndarray] = None): self.nblist = self.nblist.update(positions) else: self.nblist = self.nblist.update(positions, box) - return self.nblist + return self.pairs @property def pairs(self): From 92b35ef8d947cf8623414fcae0ab642648cadf08 Mon Sep 17 00:00:00 2001 From: WangXinyan940 Date: Tue, 22 Nov 2022 13:47:43 +0800 Subject: [PATCH 4/4] bugfix: correctly recognize if the LJForce card uses type or class --- dmff/generators/classical.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dmff/generators/classical.py b/dmff/generators/classical.py index d13e244cb..e88c909d2 100644 --- a/dmff/generators/classical.py +++ b/dmff/generators/classical.py @@ -826,8 +826,8 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, map_lj = jnp.array(maps["sigma"]) - ifType = len(self.fftree.get_attribs("LennardJonesForce/Atom", - "type")) != 0 + ifType = len([i for i in self.fftree.get_attribs("LennardJonesForce/Atom", + "type") if i is not None]) != 0 if ifType: atom_labels = self.fftree.get_attribs("LennardJonesForce/Atom", "type")