Several changes to support using charged anions in classical MD workf…

…lows
materialsproject · Nov 16, 2024 · 5f103ed · 5f103ed
1 parent 5f0bc75
commit 5f103ed
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 18 deletions.
diff --git a/src/atomate2/openmm/jobs/base.py b/src/atomate2/openmm/jobs/base.py
@@ -39,10 +39,12 @@
 
 
 try:
-    # so we can load OpenMM Interchange created by openmmml
-    import openmmml
+    # so we can load OpenMM Interchange created with nnpops
+    import nnpops
+    import openmmtorch
 except ImportError:
-    openmmml = None
+    nnpops = None
+    openmmtorch = None
 
 try:
     from openff.interchange import Interchange
@@ -54,7 +56,7 @@ class Interchange:  # type: ignore[no-redef]
         def model_validate(self, _: str) -> None:
             """Parse raw is the first method called on the Interchange object."""
             raise ImportError(
-                "openff-interchange must be installed for OpenMM makers to"
+                "openff-i nterchange must be installed for OpenMM makers to"
                 "to support OpenFF Interchange objects."
             )
 

diff --git a/src/atomate2/openmm/jobs/generate.py b/src/atomate2/openmm/jobs/generate.py
@@ -5,7 +5,6 @@
 import copy
 import io
 import re
-import warnings
 import xml.etree.ElementTree as ET
 from pathlib import Path
 from xml.etree.ElementTree import tostring
@@ -20,7 +19,7 @@
 from openmm.app.pdbfile import PDBFile
 from openmm.unit import kelvin, picoseconds
 from pymatgen.core import Element
-from pymatgen.io.openff import get_atom_map
+from pymatgen.io.openff import coerce_formal_charges, get_atom_map
 
 from atomate2.openff.utils import create_mol_spec, merge_specs_by_name_and_smiles
 from atomate2.openmm.jobs.base import openmm_job
@@ -85,12 +84,8 @@ def increment_types(self, increment: str) -> None:
                     if type_stub in key:
                         element.attrib[key] += increment
 
-    def to_openff_molecule(self) -> tk.Molecule:
+    def to_openff_molecule(self, template_molecule: tk.Molecule = None) -> tk.Molecule:
         """Convert the XMLMoleculeFF to an openff_toolkit Molecule."""
-        if sum(self.partial_charges) > 1e-3:
-            # TODO: update message
-            warnings.warn("Formal charges not considered.", stacklevel=1)
-
         p_table = {e.symbol: e.number for e in Element}
         openff_mol = tk.Molecule()
         for atom in self.tree.getroot().findall(".//Residues/Residue/Atom"):
@@ -108,6 +103,9 @@ def to_openff_molecule(self) -> tk.Molecule:
 
         openff_mol.partial_charges = self.partial_charges * unit.elementary_charge
 
+        if template_molecule:
+            openff_mol = coerce_formal_charges(openff_mol, template_molecule)
+
         return openff_mol
 
     @property
@@ -253,7 +251,7 @@ def generate_openmm_interchange(
 
     for mol_spec, xml_mol in zip(mol_specs, xml_mols, strict=True):
         openff_mol = tk.Molecule.from_json(mol_spec.openff_mol)
-        xml_openff_mol = xml_mol.to_openff_molecule()
+        xml_openff_mol = xml_mol.to_openff_molecule(template_molecule=openff_mol)
         is_isomorphic, _atom_map = get_atom_map(openff_mol, xml_openff_mol)
         if not is_isomorphic:
             raise ValueError(

diff --git a/tests/openmm_md/flows/test_core.py b/tests/openmm_md/flows/test_core.py
@@ -63,7 +63,7 @@ def test_hdf5_writing(interchange, run_job):
     import MDAnalysis
     from packaging.version import Version
 
-    if Version(MDAnalysis.__version__) < Version("2.8.0"):
+    if Version(MDAnalysis.__version__) < Version("2.8.0.dev0"):
         return
 
     anneal_maker = OpenMMFlowMaker.anneal_flow(
@@ -90,6 +90,17 @@ def test_hdf5_writing(interchange, run_job):
         "trajectory.h5md",
     }
 
+    with io.StringIO(interchange.topology) as s:
+        pdb = PDBFile(s)
+        openmm_topology = pdb.getTopology()
+
+    traj_file = Path(task_doc.calcs_reversed[0].output.dir_name) / "trajectory3.h5md"
+    Universe(
+        openmm_topology,
+        str(traj_file),
+        format="h5md",
+    )
+
 
 def test_collect_outputs(interchange, run_job):
     # Create an instance of ProductionMaker with custom parameters

diff --git a/tests/openmm_md/jobs/test_generate.py b/tests/openmm_md/jobs/test_generate.py
@@ -52,11 +52,18 @@ def test_xml_molecule_from_file(openmm_data):
 
 
 def test_to_openff_molecule(openmm_data):
-    xml_mol = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
-
-    mol = xml_mol.to_openff_molecule()
-    assert len(mol.atoms) == 6
-    assert len(mol.bonds) == 5
+    co_xml = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
+
+    co = co_xml.to_openff_molecule()
+    assert len(co.atoms) == 6
+    assert len(co.bonds) == 5
+
+    clo4_template = tk.Molecule.from_smiles("[O-]Cl(=O)(=O)=O")
+    clo4_xml = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "ClO4.xml")
+    clo4 = clo4_xml.to_openff_molecule(clo4_template)
+    assert len(clo4.atoms) == 5
+    assert len(clo4.bonds) == 4
+    assert clo4.to_smiles()
 
 
 def test_assign_partial_charges_w_mol(openmm_data):

diff --git a/tests/test_data/openmm/opls_xml_files/ClO4.xml b/tests/test_data/openmm/opls_xml_files/ClO4.xml
@@ -0,0 +1,41 @@
+<ForceField>
+	<AtomTypes>
+		<Type name="Clc" class="Clc" element="C" mass="35.45" />
+		<Type name="Oc" class="Oc" element="O" mass="15.99" />
+	</AtomTypes>
+	<Residues>
+		<Residue name="PF6">
+			<Atom name="Cl1" type="Clc" />
+			<Atom name="O2" type="Oc" />
+			<Atom name="O3" type="Oc" />
+			<Atom name="O4" type="Oc" />
+			<Atom name="O5" type="Oc" />
+			<Bond from="0" to="1" />
+			<Bond from="0" to="2" />
+			<Bond from="0" to="3" />
+			<Bond from="0" to="4" />
+		</Residue>
+	</Residues>
+	<HarmonicBondForce>
+		<Bond class1="Clc" class2="Oc" length="0.1506" k="633959.68" />
+		<Bond class1="Clc" class2="Oc" length="0.1506" k="633959.68" />
+		<Bond class1="Clc" class2="Oc" length="0.1506" k="633959.68" />
+		<Bond class1="Clc" class2="Oc" length="0.1506" k="633959.68" />
+	</HarmonicBondForce>
+	<HarmonicAngleForce>
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+		<Angle class1="Oc" class2="Clc" class3="Oc" angle="1.911135530933791" k="1739.707" />
+	</HarmonicAngleForce>
+	<PeriodicTorsionForce />
+	<NonbondedForce coulomb14scale="0.5" lj14scale="0.5">
+		<Atom type="Clc" charge="1.176" sigma="0.350000" epsilon="0.492830" />
+		<Atom type="Oc" charge="-0.544" sigma="0.290000" epsilon="0.878640" />
+		<Atom type="Oc" charge="-0.544" sigma="0.290000" epsilon="0.878640" />
+		<Atom type="Oc" charge="-0.544" sigma="0.290000" epsilon="0.878640" />
+		<Atom type="Oc" charge="-0.544" sigma="0.290000" epsilon="0.878640" />
+	</NonbondedForce>
+</ForceField>