Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for template_molecule to allow charged molecules in classical MD workflows #1023

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add support for template_molecule to allow charged molecules in classical MD workflows #1023

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/atomate2/openmm/jobs/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."
)

Expand Down
14 changes: 6 additions & 8 deletions src/atomate2/openmm/jobs/generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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"):
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 12 additions & 1 deletion tests/openmm_md/flows/test_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand Down
17 changes: 12 additions & 5 deletions tests/openmm_md/jobs/test_generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
41 changes: 41 additions & 0 deletions tests/test_data/openmm/opls_xml_files/ClO4.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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>
Loading