Add MACE force field support for OpenMM simulations

pyproject.toml

@@ -65,7 +65,7 @@ ase = ["ase>=3.23.0"]
 # tblite py3.12 support tracked in https://github.com/tblite/tblite/issues/198
 ase-ext = ["tblite>=0.3.0; python_version < '3.12'"]
 openmm = [
-    "mdanalysis>=2.7.0",
+    "mdanalysis>=2.8.0",
     "openmm-mdanalysis-reporter>=0.1.0",
     "openmm>=8.1.0",
 ]
@@ -153,14 +153,17 @@ ignore_missing_imports = true
 no_strict_optional = true
 
 [tool.pytest.ini_options]
-addopts = "-p no:warnings --import-mode=importlib --cov-config=pyproject.toml"
+addopts = "-p no:warnings --import-mode=importlib --cov-config=pyproject.toml -m 'not openmm_mace'"
 filterwarnings = [
     "ignore:.*POTCAR.*:UserWarning",
     "ignore:.*input structure.*:UserWarning",
     "ignore:.*is not gzipped.*:UserWarning",
     "ignore:.*magmom.*:UserWarning",
     "ignore::DeprecationWarning",
 ]
+markers = [
+    "openmm_mace: tests marked openmm_mace are skipped by default because they are very slow (unskip with pytest -m openmm_mace)",
+]
 
 [tool.coverage.run]
 include = ["src/*"]

src/atomate2/openmm/jobs/mace.py

@@ -0,0 +1,89 @@
+"""Run MACE on randomly packed benchmarking structures."""
+
+import io
+import json
+from pathlib import Path
+
+import numpy as np
+import openmm
+import openmm.unit as omm_unit
+from emmet.core.openmm import OpenMMInterchange, OpenMMTaskDocument
+from emmet.core.vasp.task_valid import TaskState
+from jobflow import Response
+from mace.calculators.foundations_models import download_mace_mp_checkpoint
+from monty.json import MontyEncoder
+from openmm import Context, XmlSerializer
+from openmm.app.pdbfile import PDBFile
+from pymatgen.core import Structure
+
+from atomate2.openmm.jobs.base import openmm_job
+from atomate2.openmm.mace_utils import MacePotential
+from atomate2.openmm.utils import structure_to_topology
+
+
+@openmm_job
+def generate_mace_interchange(
+    structure: Structure,
+    ff_path: str | Path | None = None,
+    tags: list[str] | None = None,
+) -> Response:
+    """Generate an OpenMMInterchange object with the MACE force-field.
+
+    Parameters
+    ----------
+    structure : Structure
+        The structure to simulate.
+    ff_path : str | Path, optional
+        The path to the MACE force-field. Must be accessible where the job is run.
+        Defaults to None.
+    tags : list[str], optional
+        Tags to add to the task document. Defaults to None.
+
+    Returns
+    -------
+    Response
+        The response containing the OpenMMTaskDocument.
+    """
+    if not ff_path:
+        ff_path = Path(download_mace_mp_checkpoint())
+
+    potential = MacePotential(model_path=ff_path)
+
+    topology = structure_to_topology(structure)
+    topology.setPeriodicBoxVectors(structure.lattice.matrix / 10)
+    system = potential.create_system(topology)
+    integrator = openmm.LangevinIntegrator(
+        300 * omm_unit.kelvin, 10.0 / omm_unit.picoseconds, 1.0 * omm_unit.femtosecond
+    )
+    context = Context(system, integrator)
+    context.setPositions(structure.cart_coords / 10)
+    state = context.getState(getPositions=True)
+    with io.StringIO() as buffer:
+        PDBFile.writeFile(topology, np.zeros(shape=(len(structure), 3)), file=buffer)
+        buffer.seek(0)
+        pdb = buffer.read()
+
+    interchange = OpenMMInterchange(
+        system=XmlSerializer.serialize(system),
+        state=XmlSerializer.serialize(state),
+        topology=pdb,
+    )
+
+    interchange_json = interchange.model_dump_json()
+
+    dir_name = Path.cwd()
+
+    task_doc = OpenMMTaskDocument(
+        dir_name=str(dir_name),
+        state=TaskState.SUCCESS,
+        interchange=interchange_json,
+        structure=structure,
+        force_field=Path(ff_path).stem,
+        tags=tags,
+    )
+
+    # write out task_doc json to output dir
+    with open(dir_name / "taskdoc.json", "w") as file:
+        json.dump(task_doc.model_dump(), file, cls=MontyEncoder)
+
+    return Response(output=task_doc)

src/atomate2/openmm/jobs/random_structure.json

@@ -0,0 +1 @@
+{"@module": "pymatgen.core.structure", "@class": "Structure", "charge": 0, "lattice": {"matrix": [[8.945722384641158, 0.0, 0.0], [0.0, 8.945722384641158, 0.0], [0.0, 0.0, 8.945722384641158]], "pbc": [true, true, true], "a": 8.945722384641158, "b": 8.945722384641158, "c": 8.945722384641158, "alpha": 90.0, "beta": 90.0, "gamma": 90.0, "volume": 715.8899231699995}, "properties": {}, "sites": [{"species": [{"element": "Al", "occu": 1}], "abc": [0.64659998950236, 0.38190599407216497, 0.8952527985610246], "properties": {}, "label": "Al", "xyz": [5.784304, 3.4164249999999994, 8.008683]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8834840452425946, 0.31815720157899446, 0.8839606976376343], "properties": {}, "label": "Al", "xyz": [7.903403, 2.846146, 7.907666999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.3468634355708847, 0.757030534686289, 0.817740668161081], "properties": {}, "label": "Al", "xyz": [3.1029439999999995, 6.772185, 7.315281]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.11223308267689708, 0.8869685039212494, 0.1123682310693931], "properties": {}, "label": "Al", "xyz": [1.004006, 7.934574, 1.005215]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.16107788035921722, 0.3912898086363331, 0.6561033025211006], "properties": {}, "label": "Al", "xyz": [1.4409579999999997, 3.50037, 5.869317999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.68638146099213, 0.5142043093017935, 0.491245626797559], "properties": {}, "label": "Al", "xyz": [6.140178, 4.5999289999999995, 4.394547]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5147756438212701, 0.11389152895569865, 0.3368813462369576], "properties": {}, "label": "Al", "xyz": [4.60504, 1.018842, 3.0136469999999997]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5836308992735897, 0.3558866308511883, 0.1110339620761486], "properties": {}, "label": "Al", "xyz": [5.221, 3.183663, 0.9932789999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.7667138219911492, 0.8885189656283793, 0.19472759438531093], "properties": {}, "label": "Al", "xyz": [6.858809, 7.948444, 1.7419789999999997]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.4653749379868551, 0.11650641001216443, 0.675629730068182], "properties": {}, "label": "Al", "xyz": [4.163115, 1.042234, 6.043995999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8434003063803603, 0.6355647711314559, 0.3462858410762388], "properties": {}, "label": "Al", "xyz": [7.544825, 5.685586, 3.097777]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.10461547539266058, 0.268665167178269, 0.45120213063284215], "properties": {}, "label": "Al", "xyz": [0.935861, 2.4034039999999997, 4.036329]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.10656545765792173, 0.48920040348150656, 0.34230941542043336], "properties": {}, "label": "Al", "xyz": [0.953305, 4.376251, 3.062205]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.16209926237926373, 0.532217276066407, 0.10737713050979408], "properties": {}, "label": "Al", "xyz": [1.450095, 4.761068, 0.9605659999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.10987653738143902, 0.6517177427738681, 0.531342779892299], "properties": {}, "label": "Al", "xyz": [0.9829249999999999, 5.830086, 4.753245]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6220034292091662, 0.11033888126180591, 0.8866595294325322], "properties": {}, "label": "Al", "xyz": [5.56427, 0.987061, 7.93181]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.837565227025604, 0.33900448388681453, 0.11178527088174481], "properties": {}, "label": "Al", "xyz": [7.492625999999999, 3.0326399999999993, 1.0]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.12504557506955002, 0.1168739599828236, 0.6426506158821079], "properties": {}, "label": "Al", "xyz": [1.118623, 1.045522, 5.748974]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6789277309150062, 0.2765485998366623, 0.41951491882234837], "properties": {}, "label": "Al", "xyz": [6.073499, 2.473927, 3.752864]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8902798072153085, 0.6174102842139062, 0.1067077603077542], "properties": {}, "label": "Al", "xyz": [7.964196, 5.523181, 0.9545779999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5625750256503042, 0.780675019827383, 0.11049593956739476], "properties": {}, "label": "Al", "xyz": [5.03264, 6.983702, 0.988466]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5418058812468304, 0.5988211761631681, 0.8458476213157743], "properties": {}, "label": "Al", "xyz": [4.846844999999999, 5.356887999999999, 7.566718]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6963827773926489, 0.10453331321856249, 0.5926859533560945], "properties": {}, "label": "Al", "xyz": [6.229647, 0.9351259999999999, 5.302004]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8508692392543241, 0.7985831320079093, 0.8003875698504825], "properties": {}, "label": "Al", "xyz": [7.61164, 7.143902999999999, 7.160045]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.419600322769302, 0.24957928538373256, 0.8982807261933984], "properties": {}, "label": "Al", "xyz": [3.7536279999999995, 2.232667, 8.03577]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.3342419842061687, 0.18765795849893677, 0.4822187426033163], "properties": {}, "label": "Al", "xyz": [2.9900359999999995, 1.678736, 4.313795]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.7144616974670827, 0.8272436458240091, 0.5737339903160743], "properties": {}, "label": "Al", "xyz": [6.391376, 7.400292, 5.132464999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6282204788345254, 0.852304226776633, 0.8892708333603299], "properties": {}, "label": "Al", "xyz": [5.619886, 7.624477, 7.95517]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8881081547578888, 0.40276344883963533, 0.4077231377381166], "properties": {}, "label": "Al", "xyz": [7.944769, 3.6030099999999994, 3.647378]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.3606916089347692, 0.1239122959933509, 0.10891205406427132], "properties": {}, "label": "Al", "xyz": [3.2266469999999994, 1.108485, 0.974297]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.13044402115624212, 0.5731628793671395, 0.8127437547673961], "properties": {}, "label": "Al", "xyz": [1.1669159999999998, 5.127356, 7.27058]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8207171745688493, 0.11312177558040695, 0.22061840454477363], "properties": {}, "label": "Al", "xyz": [7.341908, 1.011956, 1.973591]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.3506572040940694, 0.8885147177880858, 0.17155763771910992], "properties": {}, "label": "Al", "xyz": [3.136882, 7.948406, 1.5347069999999998]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5031197936264306, 0.8860720978340487, 0.685552461423271], "properties": {}, "label": "Al", "xyz": [4.500769999999999, 7.926555, 6.132761999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.38536317714472473, 0.5908122086455747, 0.19223835997331626], "properties": {}, "label": "Al", "xyz": [3.447352, 5.285242, 1.719711]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6037762818655432, 0.11123148664979664, 0.1103842660817839], "properties": {}, "label": "Al", "xyz": [5.401215, 0.9950459999999999, 0.9874669999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.5199556614886593, 0.6769940692992937, 0.5641595818650529], "properties": {}, "label": "Al", "xyz": [4.651379, 6.056201, 5.046815]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.4035626017411689, 0.3583061112541528, 0.6792217261774252], "properties": {}, "label": "Al", "xyz": [3.6101589999999995, 3.205307, 6.076128999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.12481473848517921, 0.8587373573306065, 0.8211087583727481], "properties": {}, "label": "Al", "xyz": [1.116558, 7.682026, 7.345411]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.14827310115025513, 0.3429417846978114, 0.8948110231705], "properties": {}, "label": "Al", "xyz": [1.3264099999999999, 3.067862, 8.004731]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6645119023820963, 0.5724788652946141, 0.19265699581276838], "properties": {}, "label": "Al", "xyz": [5.944538999999999, 5.121237, 1.7234559999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8900891015431842, 0.15954295680474, 0.4518175085490461], "properties": {}, "label": "Al", "xyz": [7.96249, 1.427227, 4.041834]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.88777995320258, 0.868444119542352, 0.4076669097468631], "properties": {}, "label": "Al", "xyz": [7.941832999999999, 7.76886, 3.6468749999999996]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6454857139222108, 0.33587248416724985, 0.6551271935357612], "properties": {}, "label": "Al", "xyz": [5.774336, 3.004622, 5.860586]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.6042424264451202, 0.7678240733015467, 0.3522360592600032], "properties": {}, "label": "Al", "xyz": [5.405385, 6.868740999999999, 3.1510059999999998]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.21920409729757778, 0.10988827483488159, 0.8689416757830466], "properties": {}, "label": "Al", "xyz": [1.960939, 0.9830299999999998, 7.773311]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8854965154742787, 0.6515903075650629, 0.5878952837871864], "properties": {}, "label": "Al", "xyz": [7.921406, 5.828945999999999, 5.259147999999999]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.17644310119772572, 0.8866629947759297, 0.5117170870246747], "properties": {}, "label": "Al", "xyz": [1.578411, 7.931841, 4.577679]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.8902248088620346, 0.3336060378001225, 0.6406499948891374], "properties": {}, "label": "Al", "xyz": [7.963703999999999, 2.984347, 5.731077]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.3383039255941987, 0.3644458054720618, 0.11053484084166161], "properties": {}, "label": "Al", "xyz": [3.026373, 3.260231, 0.988814]}, {"species": [{"element": "Al", "occu": 1}], "abc": [0.109196882934478, 0.2875773346618552, 0.15432973891189894], "properties": {}, "label": "Al", "xyz": [0.976845, 2.5725869999999995, 1.380591]}, {"species": [{"element": "Fe", "occu": 1}], "abc": [0.5150752283472332, 0.43701378512617667, 0.33287015536190795], "properties": {}, "label": "Fe", "xyz": [4.60772, 3.909404, 2.977764]}, {"species": [{"element": "Fe", "occu": 1}], "abc": [0.3231218090293945, 0.6687096628742476, 0.4165378534882257], "properties": {}, "label": "Fe", "xyz": [2.890558, 5.982090999999999, 3.7262319999999995]}, {"species": [{"element": "Fe", "occu": 1}], "abc": [0.8868078684869924, 0.11344483501325518, 0.7480456817538979], "properties": {}, "label": "Fe", "xyz": [7.933137, 1.014846, 6.691809]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.4182014418894879, 0.8897721903002345, 0.4561953551125879], "properties": {}, "label": "Ni", "xyz": [3.741114, 7.959655, 4.080997]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.3127311445303946, 0.5842903205865211, 0.6508213366866673], "properties": {}, "label": "Ni", "xyz": [2.797606, 5.226899, 5.822067]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.11719713120094273, 0.7311663294212961, 0.3001219895454755], "properties": {}, "label": "Ni", "xyz": [1.048413, 6.540811, 2.6848079999999994]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.19971142890230265, 0.10653147493557369, 0.29398352496554625], "properties": {}, "label": "Ni", "xyz": [1.786563, 0.953001, 2.629895]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.77572270875678, 0.5601726483937846, 0.7856451047560562], "properties": {}, "label": "Ni", "xyz": [6.9394, 5.011149, 7.028162999999999]}, {"species": [{"element": "Ni", "occu": 1}], "abc": [0.30994746771489734, 0.42839927679621675, 0.4650978223003392], "properties": {}, "label": "Ni", "xyz": [2.7727039999999996, 3.8323409999999996, 4.160636]}]}