Merge pull request #168 from Ethan-Norch/devel

Update classical.py and xmlio.py
deepmodeling · Jul 4, 2024 · 64c5fbc · 64c5fbc
2 parents 8c49937 + f7cccc9
commit 64c5fbc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/dmff/api/xmlio.py b/dmff/api/xmlio.py
@@ -123,6 +123,8 @@ def parseForce(self, nodes):
                 inner = {}
                 inner["name"] = child.tag
                 inner["attrib"] = child.attrib
+                if child.text is not None:
+                    inner["formula"] = child.text
                 ret["node"].append(inner)
         return ret
 
@@ -178,6 +180,8 @@ def writeXML(
                 for node in force_info["node"]:
                     subnode = ET.SubElement(fnode, node["name"])
                     subnode.attrib = genStrDict(node["attrib"])
+                    if "formula" in node:
+                        subnode.text = node["formula"]
 
         tree = ET.ElementTree(root)
         xmlstr = minidom.parseString(ET.tostring(root)).toprettyxml(indent="   ")

diff --git a/dmff/generators/classical.py b/dmff/generators/classical.py
@@ -897,7 +897,7 @@ def overwrite(self, paramset):
         proper_phase = paramset[self.name]["proper_phase"]
         proper_k = paramset[self.name]["proper_k"]
         proper_shift = paramset[self.name]["proper_shift"]
-        proper_msks = paramset.mask[self.name]["proper"]
+        proper_msks = paramset.mask[self.name]["proper_phase"]
         for nnode, key in enumerate(self.proper_keys):
             self.ffinfo["Forces"][self.name]["node"][proper_node_indices[nnode]]["attrib"] = {
             }
@@ -929,7 +929,7 @@ def overwrite(self, paramset):
         improper_phase = paramset[self.name]["improper_phase"]
         improper_k = paramset[self.name]["improper_k"]
         improper_shift = paramset[self.name]["improper_shift"]
-        improper_msks = paramset.mask[self.name]["improper"]
+        improper_msks = paramset.mask[self.name]["improper_phase"]
         for nnode, key in enumerate(self.imp_keys):
             self.ffinfo["Forces"][self.name]["node"][improper_node_indices[nnode]]["attrib"] = {
             }
@@ -2028,9 +2028,9 @@ def overwrite(self, paramset: ParamSet) -> None:
         """
         radius = paramset[self.name]["radius"]
         scale = paramset[self.name]["scale"]
-        for i in self.ffinfo.perParticleParamIndices:
-            self.ffinfo["Forces"][self.name]["node"][i]["attrib"]["radius"] = str(radius[i])
-            self.ffinfo["Forces"][self.name]["node"][i]["attrib"]["scale"] = str(scale[i])
+        for j, i in enumerate(self.perParticleParamIndices):
+            self.ffinfo["Forces"][self.name]["node"][i]["attrib"]["radius"] = str(radius[j])
+            self.ffinfo["Forces"][self.name]["node"][i]["attrib"]["scale"] = str(scale[j])
 
     def _find_key_index(self, key: Tuple[str]) -> int:
         """
@@ -2109,4 +2109,4 @@ def potential_fn(positions: jnp.ndarray, box: jnp.ndarray, pairs: jnp.ndarray, p
         return potential_fn
 
 
-_DMFFGenerators["CustomGBForce"] = CustomGBGenerator
+_DMFFGenerators["CustomGBForce"] = CustomGBGenerator