Support more MACE foundation models (#109)

* Fixed errors in documentation

* Support more MACE foundation models

* Tests for MACE foundation models

* Removed documentation on MACE-OMOL-0

* Don't register MACE-OMOL-0

Author: peastman

doc/userguide.md

@@ -78,7 +78,9 @@ are supported.
 
 | Name | Model |
 | --- | --- |
-| `mace-off23-small`<br>`mace-off23-medium`<br>`mace-off23-large` | Pretrained [MACE-OFF23](https://arxiv.org/abs/2312.15211) models |
+| `mace-off23-small`<br>`mace-off23-medium`<br>`mace-off23-large`<br>`mace-off24-medium` | Pretrained [MACE-OFF](https://pubs.acs.org/doi/10.1021/jacs.4c07099) models |
+| `mace-mpa-0-medium` | Pretrained [MACE-MPA-0](https://github.com/ACEsuit/mace-foundations) model |
+| `mace-omat-0-small`<br>`mace-omat-0-medium` | Pretrained [MACE-OMAT-0](https://github.com/ACEsuit/mace-foundations) models |
 | `mace` | Custom MACE models specified with the `modelPath` argument |
 
 When creating MACE models, the following keyword arguments to the `MLPotential` constructor are supported.
@@ -92,7 +94,7 @@ When using MACE models, the following extra keyword arguments to `createSystem()
 | Argument | Description |
 | --- | --- |
 | `precision` | The numerical precision of the model. Supported options are `'single'` and `'double'`.  If `None`, the default precision of the model is used. |
- | `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy.  The default is `'interaction_energy'`. Supported options are `'interaction_energy'` and `'energy'`. |
+| `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy.  The default is `'interaction_energy'`. Supported options are `'interaction_energy'` and `'energy'`. |
 
 ### AIMNet2
 
@@ -106,9 +108,9 @@ are supported.
 
 When using AIMNet2 models, the following extra keyword arguments to `createSystem()` and `createMixedSystem()` are supported.
 
-| Argument       | Description |
+| Argument | Description |
 | --- | --- |
-| `charge`       | The total charge of the system.  If omitted, it is assumed to be 0. |
+| `charge` | The total charge of the system.  If omitted, it is assumed to be 0. |
 | `multiplicity` | The spin multiplicity of the system.  If omitted, it is assumed to be 1. |
 
 ### NequIP

openmmml/models/macepotential.py

@@ -47,13 +47,13 @@ class MACEPotentialImpl(MLPotentialImpl):
 
     The MACE potential is constructed using MACE to build a PyTorch model,
     and then integrated into the OpenMM System using a TorchForce.
-    This implementation supports both MACE-OFF23 and locally trained MACE models.
+    This implementation supports both foundation models and locally trained MACE models.
 
-    To use one of the pre-trained MACE-OFF23 models, specify the model name. For example:
+    To use one of the pre-trained MACE foundation models, specify the model name. For example:
 
     >>> potential = MLPotential('mace-off23-small')
 
-    Other available MACE-OFF23 models include 'mace-off23-medium' and 'mace-off23-large'.
+    Other available models include 'mace-off23-medium', 'mace-off23-large', and 'mace-off24-medium'.
 
     To use a locally trained MACE model, provide the path to the model file. For example:
 
@@ -93,7 +93,8 @@ def __init__(self, name: str, modelPath) -> None:
         ----------
         name : str
             The name of the MACE model.
-            Options include 'mace-off23-small', 'mace-off23-medium', 'mace-off23-large', and 'mace'.
+            Options include 'mace-off23-small', 'mace-off23-medium', 'mace-off23-large',
+            'mace-off24-medium', and 'mace'.
         modelPath : str, optional
             The path to the locally trained MACE model if ``name`` is 'mace'.
         """
@@ -135,7 +136,7 @@ def addForces(
 
         try:
             from mace.tools import utils, to_one_hot, atomic_numbers_to_indices
-            from mace.calculators.foundations_models import mace_off
+            from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol
         except ImportError as e:
             raise ImportError(
                 f"Failed to import mace with error: {e}. "
@@ -161,13 +162,24 @@ def addForces(
             "energy",
         ], f"Unsupported returnEnergyType: '{returnEnergyType}'. Supported options are 'interaction_energy' or 'energy'."
 
+        models = {
+            'mace-off23-small': (mace_off, 'small', True),
+            'mace-off23-medium': (mace_off, 'medium', True),
+            'mace-off23-large': (mace_off, 'large', True),
+            'mace-off24-medium': (mace_off, 'https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true', True),
+            'mace-mpa-0-medium': (mace_mp, 'medium-mpa-0', False),
+            'mace-omat-0-small': (mace_mp, 'small-omat-0', True),
+            'mace-omat-0-medium': (mace_mp, 'medium-omat-0', True),
+            'mace-omol-0-extra-large': (mace_omol, 'extra_large', True)
+        }
+
         # Load the model to the CPU (OpenMM-Torch takes care of loading to the right devices)
-        if self.name.startswith("mace-off23"):
-            size = self.name.split("-")[-1]
-            assert (
-                size in ["small", "medium", "large"]
-            ), f"Unsupported MACE model: '{self.name}'. Available MACE-OFF23 models are 'mace-off23-small', 'mace-off23-medium', 'mace-off23-large'"
-            model = mace_off(model=size, device="cpu", return_raw_model=True)
+        if self.name in models:
+            fn, name, warn = models[self.name]
+            model = fn(model=name, device="cpu", return_raw_model=True)
+            if warn:
+                import logging
+                logging.warning(f'The model {self.name} is distributed under the restrictive ASL license.  Commercial use is not permitted.')
         elif self.name == "mace":
             if self.modelPath is not None:
                 model = torch.load(self.modelPath, map_location="cpu")
@@ -231,6 +243,10 @@ class MACEForce(torch.nn.Module):
                 Conversion factor for the length, viz. nm to Angstrom.
             indices : torch.Tensor
                 The indices of the atoms to calculate the energy for.
+            charge : float
+                Total charge of the system
+            multiplicity : float
+                Spin multiplicity of the system
             returnEnergyType : str
                 Whether to return the interaction energy or the energy including the self-energy.
             inputDict : dict
@@ -242,6 +258,8 @@ def __init__(
                 model: torch.jit._script.RecursiveScriptModule,
                 nodeAttrs: torch.Tensor,
                 atoms: Optional[Iterable[int]],
+                charge: float,
+                multiplicity: float,
                 periodic: bool,
                 dtype: torch.dtype,
                 returnEnergyType: str,
@@ -282,6 +300,8 @@ def __init__(
                 self.register_buffer("node_attrs", nodeAttrs.to(self.dtype))
                 self.register_buffer("batch", torch.zeros(nodeAttrs.shape[0], dtype=torch.long, requires_grad=False))
                 self.register_buffer("pbc", torch.tensor([periodic, periodic, periodic], dtype=torch.bool, requires_grad=False))
+                self.register_buffer("charge", torch.tensor([charge], dtype=dtype, requires_grad=False))
+                self.register_buffer("multiplicity", torch.tensor([multiplicity], dtype=dtype, requires_grad=False))
 
             def _getNeighborPairs(
                 self, positions: torch.Tensor, cell: Optional[torch.Tensor]
@@ -371,7 +391,9 @@ def forward(
                     "edge_index": edgeIndex,
                     "shifts": shifts,
                     "cell": cell if cell is not None else torch.zeros(3, 3, dtype=self.dtype),
-                }                    
+                    "total_charge": self.charge,
+                    "total_spin": self.multiplicity
+                }
 
                 # Predict the energy.
                 energy = self.model(inputDict, compute_force=False)[
@@ -392,6 +414,8 @@ def forward(
             model,
             nodeAttrs,
             atoms,
+            float(args.get('charge', 0)),
+            float(args.get('multiplicity', 1)),
             isPeriodic,
             dtype,
             returnEnergyType,

setup.py

@@ -48,6 +48,10 @@
             'mace-off23-small = openmmml.models.macepotential:MACEPotentialImplFactory',
             'mace-off23-medium = openmmml.models.macepotential:MACEPotentialImplFactory',
             'mace-off23-large = openmmml.models.macepotential:MACEPotentialImplFactory',
+            'mace-off24-medium = openmmml.models.macepotential:MACEPotentialImplFactory',
+            'mace-mpa-0-medium = openmmml.models.macepotential:MACEPotentialImplFactory',
+            'mace-omat-0-small = openmmml.models.macepotential:MACEPotentialImplFactory',
+            'mace-omat-0-medium = openmmml.models.macepotential:MACEPotentialImplFactory',
             'nequip = openmmml.models.nequippotential:NequIPPotentialImplFactory',
             'deepmd = openmmml.models.deepmdpotential:DeepmdPotentialImplFactory',
             'torchmdnet = openmmml.models.torchmdnetpotential:TorchMDNetPotentialImplFactory',
@@ -56,4 +60,3 @@
         ]
     }
 )
-

test/TestMACEPotential.py

@@ -0,0 +1,37 @@
+import os
+
+import numpy as np
+import openmm as mm
+import openmm.app as app
+import openmm.unit as unit
+import pytest
+
+from openmmml import MLPotential
+
+mace = pytest.importorskip("mace", reason="mace is not installed")
+platform_ints = range(mm.Platform.getNumPlatforms())
+# Get the path to the test data
+test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+
+@pytest.mark.parametrize("platform_int", list(platform_ints))
+@pytest.mark.parametrize("model", ['mace-off23-small', 'mace-off23-medium', 'mace-off23-large', 'mace-off24-medium',
+                                   'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium'])#, 'mace-omol-0-extra-large'])
+class TestMACE:
+    def testCreatePureMLSystem(self, platform_int, model):
+        pdb = app.PDBFile(os.path.join(test_data_dir, "toluene", "toluene.pdb"))
+        potential = MLPotential(model)
+        system = potential.createSystem(pdb.topology, returnEnergyType='energy')
+        platform = mm.Platform.getPlatform(platform_int)
+        context = mm.Context(system, mm.VerletIntegrator(0.001), platform)
+        context.setPositions(pdb.getPositions(asNumpy=True))
+        energyML = context.getState(energy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+        # Reference energies are calculated with MACECalculator
+        refEnergy = {'mace-off23-small': -713468.6327560507,
+                     'mace-off23-medium': -713468.0563706581,
+                     'mace-off23-large': -713467.7476380612,
+                     'mace-off24-medium': -713467.9394350434,
+                     'mace-mpa-0-medium': -8839.299589829867,
+                     'mace-omat-0-small': -8726.63865431241,
+                     'mace-omat-0-medium': -8679.026847088873,
+                     'mace-omol-0-extra-large': -712903.4934289698}
+        assert np.isclose(refEnergy[model], energyML, rtol=1e-6)

test/data/toluene/mace_energies.py

@@ -0,0 +1,26 @@
+# This script computes reference energies for the MACE foundation models.
+
+import ase.io
+from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol
+from openmm.unit import kilojoules_per_mole, ev, item
+
+atoms = ase.io.read('toluene.pdb')
+results = {}
+atoms.calc = mace_off('small')
+results['mace-off23-small'] = atoms.get_potential_energy()
+atoms.calc = mace_off('medium')
+results['mace-off23-medium'] = atoms.get_potential_energy()
+atoms.calc = mace_off('large')
+results['mace-off23-large'] = atoms.get_potential_energy()
+atoms.calc = mace_off('https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true')
+results['mace-off24-medium'] = atoms.get_potential_energy()
+atoms.calc = mace_mp('medium-mpa-0')
+results['mace-mpa-0-medium'] = atoms.get_potential_energy()
+atoms.calc = mace_mp('small-omat-0')
+results['mace-omat-0-small'] = atoms.get_potential_energy()
+atoms.calc = mace_mp('medium-omat-0')
+results['mace-omat-0-medium'] = atoms.get_potential_energy()
+atoms.calc = mace_omol('extra_large')
+results['mace-omol-0-extra-large'] = atoms.get_potential_energy()
+for key in results:
+    print(f'{key}: {(results[key]*ev/item).value_in_unit(kilojoules_per_mole)}')