Integration of OpenMM

.github/workflows/pytest.yml

@@ -38,7 +38,7 @@ jobs:
       run: ./install_gap.sh
 
     - name: Test basic install
-      run: pytest --cov
+      run: pytest --cov -k "not test_openmm"
 
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v4

README.md

@@ -33,14 +33,19 @@ Each model is installed into individual conda environment:
 
 - Units are: distance (Å), energy (eV), force (eV Å$`^{-1}`$), time (fs)
 
-## Using with OpenMM (Experimental!)
+## Using with OpenMM
 
 The OpenMM backend only works with MACE at the moment. The necessary dependencies are installed automatically via conda:
 
-```console
+```
 ./install_mace.sh
 ```
 
+Depending on your machine, you might need to prefix the command above with something like `CONDA_OVERRIDE_CUDA="11.2"` in two scenarios:
+
+- To ensure an environment that is compatible with your CUDA driver.
+- To force CUDA builds to be installed, even if the installation is being done from a CPU-only machine. This is typical in a situation where you are installing from a head node without GPUs but intend to run on GPUs and want to install the CUDA builds.
+
 You should now be able to run `water_openmm.py` in `./examples` or run the jupyter notebook on Google Colab [`water_openmm_colab.ipynb`](./examples/water_openmm_colab.ipynb).
 
 You can use OpenMM during active learning by passing the keyword argument `md_program="OpenMM"` to the `al_train` method.

environment_mace.yml

@@ -24,17 +24,7 @@ dependencies:
   - openmm
   - openmm-torch
   - nnpops
-  # MACE dependencies
-  - pytorch=2.0
-  - torchvision
-  - torchaudio
-  - torch-ema
-  # TODO: You might also need CUDA-specific libraries,
-  # but that depends on CUDA version
-  # https://pytorch.org/get-started/locally/
-  # - pytorch-cuda=11.8
-  # - pytorch-cuda=12.1
   - pip:
-    - mace-torch
     - openmmml@git+https://github.com/openmm/openmm-ml.git@main
     - ase@git+https://gitlab.com/ase/ase.git@f2615a6e9a # For PLUMED
+    - mace-torch

mlptrain/potentials/mace/mace.py

@@ -689,6 +689,9 @@ def opt_param_options(self) -> Dict:
             ],
             lr=self.args.lr,
             amsgrad=Config.mace_params['amsgrad'],
+            foreach=False
+            if torch.get_default_dtype() == torch.float64
+            else True,
         )
 
         return param_options

mlptrain/sampling/md.py

@@ -557,17 +557,24 @@ def add_momenta(idx, vector, energy):
 
 
 def _initialise_traj(
-    ase_atoms: 'ase.atoms.Atoms', restart: bool, traj_name: str
+    ase_atoms: 'ase.atoms.Atoms',
+    restart: bool,
+    traj_name: str,
+    remove_last: bool = True,
 ) -> 'ase.io.trajectory.Trajectory':
     """Initialise ASE trajectory object"""
 
     if not restart:
         traj = ASETrajectory(traj_name, 'w', ase_atoms)
 
     else:
-        # Remove the last frame to avoid duplicate frames
         previous_traj = ASETrajectory(traj_name, 'r', ase_atoms)
-        previous_atoms = previous_traj[:-1]
+
+        if remove_last:
+            # Remove the last frame to avoid duplicate frames
+            previous_atoms = previous_traj[:-1]
+        else:
+            previous_atoms = previous_traj
 
         os.remove(traj_name)