diff --git a/src/atomate2/openmm/jobs/base.py b/src/atomate2/openmm/jobs/base.py index 759acdd197..c5f0471ba2 100644 --- 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 index a155af20f2..c806ee4106 100644 --- 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 index f30a23ce85..34237cd7a8 100644 --- 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 index 6c934cfa58..79dbe876ec 100644 --- 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 new file mode 100644 index 0000000000..7dde254ed6 --- /dev/null +++ b/tests/test_data/openmm/opls_xml_files/ClO4.xml @@ -0,0 +1,41 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +