Add pet-omatpes-v2 models (#27)

Author: frostedoyster

README.md

@@ -42,7 +42,7 @@ atoms.set_velocities(  # it is generally a good idea to remove any net velocity
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"
-energy_model, flashmd_model = get_pretrained("pet-omatpes", time_step)  
+energy_model, flashmd_model = get_pretrained("pet-omatpes-v2", time_step)  
 
 # Set the energy model (see below for more precise usage)
 calculator = MetatomicCalculator(energy_model, device=device)
@@ -70,6 +70,11 @@ Other available integrators:
   from flashmd.ase.bussi import Bussi
 ```
 
+Along with all FlashMD models, we also provide the potential energy model whose
+dynamics they are trained to reproduce. See this short [guide](docs/energy.md) on how to
+best use the energy models if you want to enforce exact energy conservation during
+FlashMD runs, run traditional MD with the energy model, and more.
+
 Common pitfalls
 ---------------
 
@@ -78,52 +83,6 @@ above is good for metals. However,
 - for most materials: try 32 fs (aggressive) or 16 fs (conservative)
 - for aqueous and/or organic systems: try 16 fs (aggressive) or 8 fs (conservative)
 
-
-Companion energy models and exact energy conservation
------------------------------------------------------
-
-You might have noticed that ``get_pretrained()`` does not only return a FlashMD model,
-but also an energy model, which is itself just a machine-learned interatomic potential.
-This is the energy model that the FlashMD model was trained on. You might want to use it
-if...
-
-**Case 1**: you want to run FlashMD with exact energy conservation, available through the
-integrator's (``dyn`` above) parameter ``rescale_energy=True`` (this is enabled by
-default only when targeting the NVE ensemble with ``VelocityVerlet``). In that case,
-besides setting this flag, you should attach the energy calculator to the atoms before
-running FlashMD, exactly as shown above (and below with the more precise
-``do_gradients_with_energy=False`` which will save you memory and computation):
-
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
-
-...  # setting up atoms
-calculator = MetatomicCalculator(energy_model, device=device, do_gradients_with_energy=False)
-atoms.calc = calculator
-...  # running FlashMD
-```
-
-**Case 2**: you want to compute energies after running FlashMD for your own analysis. In
-this case, you can create the calculator just like in case 1, but possibly after running
-FlashMD and/or in a different script.
-
-**Case 3**: you found something interesting during a FlashMD run and you want to confirm it
-with traditional MD. Then, you can just use ASE's MD modules as usual after attaching
-the energy calculator:
-
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
-
-...  # setting up atoms
-calculator = MetatomicCalculator(energy_model, device=device)
-atoms.calc = calculator
-...  # running MD
-```
-
-In general, the energy models are slower and have a larger memory footprint compared to
-the FlashMD models. As summarized above, you should use `do_gradients_with_energy=False`
-to save computation and memory when you do not need forces.
-
 Using FlashMD in LAMMPS
 -----------------------
 
@@ -140,6 +99,12 @@ You can see
 for usage examples. i-PI is our most mature interface, and the one that was used to
 generate all our published results.
 
+Models
+------
+
+See [here](docs/models.md) for the complete list of the models we provide. If you are
+new to FlashMD, we recommend starting with the ``pet-omatpes-v2`` models.
+
 Disclaimer
 ----------
 

docs/energy.md

@@ -0,0 +1,44 @@
+Companion energy models and exact energy conservation
+-----------------------------------------------------
+
+You might have noticed that ``get_pretrained()`` does not only return a FlashMD model,
+but also an energy model, which is itself just a machine-learned interatomic potential.
+This is the energy model that the FlashMD model was trained on. You might want to use it
+if...
+
+**Case 1**: you want to run FlashMD with exact energy conservation, available through the
+integrator's (often named ``dyn``) parameter ``rescale_energy=True`` (this is enabled by
+default only when targeting the NVE ensemble with ``VelocityVerlet``). In that case,
+besides setting this flag, you should attach the energy calculator to the atoms before
+running FlashMD, exactly as shown in the opening example (and below with the more precise
+``do_gradients_with_energy=False`` which will save you memory and computation):
+
+```
+from metatomic.torch.ase_calculator import MetatomicCalculator
+
+...  # setting up atoms
+calculator = MetatomicCalculator(energy_model, device=device, do_gradients_with_energy=False)
+atoms.calc = calculator
+...  # running FlashMD
+```
+
+**Case 2**: you want to compute energies after running FlashMD for your own analysis. In
+this case, you can create the calculator just like in case 1, but possibly after running
+FlashMD and/or in a different script.
+
+**Case 3**: you found something interesting during a FlashMD run and you want to confirm it
+with traditional MD. Then, you can just use ASE's MD modules as usual after attaching
+the energy calculator:
+
+```
+from metatomic.torch.ase_calculator import MetatomicCalculator
+
+...  # setting up atoms
+calculator = MetatomicCalculator(energy_model, device=device)
+atoms.calc = calculator
+...  # running MD
+```
+
+In general, the energy models are slower and have a larger memory footprint compared to
+the FlashMD models. As summarized above, you should use `do_gradients_with_energy=False`
+to save computation and memory when you do not need forces.

docs/lammps.md

@@ -28,7 +28,7 @@ number you get and match it to the number from the following list:
 where you should replace ``ADA89`` with the string you found above. If you do not have
 conda yet, we recommend the [Miniforge](https://github.com/conda-forge/miniforge) conda provider.
 On HPC systems, it is safer to execute this command on a GPU-enabled compute (or debug)
-node, as sometimes nvidia drivers are not present on login nodes and this can prevent
+node, as sometimes Nvidia drivers are not present on login nodes and this can prevent
 conda from installing the correct GPU libraries.
 
 You are now ready to run FlashMD in LAMMPS!
@@ -50,7 +50,7 @@ trained on. Here is how you can get them in the current directory from Python:
     from flashmd import get_pretrained
 
     time_step = 16  # in fs, also available: 1, 2, 4, 8, 32, 64, 128
-    energy_model, flashmd_model = get_pretrained("pet-omatpes", time_step)
+    energy_model, flashmd_model = get_pretrained("pet-omatpes-v2", time_step)
     energy_model.save("mlip.pt")
     flashmd_model.save(f"flashmd-{time_step}.pt")
 ```
@@ -61,8 +61,7 @@ Here below you will see how to run different types of molecular dynamics. In all
 you should launch LAMMPS as ``lmp -in in.flashmd -k on g 1 -pk kokkos newton on neigh half -sf kk``
 (assuming your input file is named ``in.flashmd``), or ``lmp -in in.flashmd`` if you
 want to run without Kokkos acceleration. The following sections will present some input
-files that you can take inspiration from. Trivially, we advise you to consult the [LAMMPS
-documentation](https://docs.lammps.org) for further explanation of all the parameters.
+files that you can take inspiration from. 
 
 # NVT (Langevin thermostat)
 

docs/models.md

@@ -0,0 +1,6 @@
+This is the complete list of FlashMD we make available through this library.
+
+| Name                | Available timesteps (fs)        | Description                                                                 |
+|---------------------|---------------------------------|-----------------------------------------------------------------------------|
+| ``pet-omatpes-v2``  | ``1, 2, 4, 8, 16, 32, 64, 128`` | New version of the ``pet-omatpes`` models, more robust at high temperatures | 
+| ``pet-omatpes``     | ``1, 2, 4, 8, 16, 32, 64, 128`` | Trained on the MATPES-r2SCAN dataset                                        |

src/flashmd/models.py

@@ -6,11 +6,14 @@
 from metatomic.torch import AtomisticModel, load_atomistic_model
 
 
-AVAILABLE_MLIPS = ["pet-omatpes"]
-AVAILABLE_TIME_STEPS = {"pet-omatpes": [1, 2, 4, 8, 16, 32, 64, 128]}
+AVAILABLE_MLIPS = ["pet-omatpes", "pet-omatpes-v2"]
+AVAILABLE_TIME_STEPS = {
+    "pet-omatpes": [1, 2, 4, 8, 16, 32, 64, 128],
+    "pet-omatpes-v2": [1, 2, 4, 8, 16, 32, 64, 128],
+}
 
 
-def get_pretrained(mlip: str = "pet-omatpes", time_step: int = 16) -> AtomisticModel:
+def get_pretrained(mlip: str = "pet-omatpes-v2", time_step: int = 16) -> AtomisticModel:
     if mlip not in AVAILABLE_MLIPS:
         raise ValueError(
             f"MLIP '{mlip}' is not available. "

tests/test_models.py

@@ -8,6 +8,7 @@ def test_available_mlips():
     assert isinstance(AVAILABLE_MLIPS, list)
     assert len(AVAILABLE_MLIPS) > 0
     assert "pet-omatpes" in AVAILABLE_MLIPS
+    assert "pet-omatpes-v2" in AVAILABLE_MLIPS
 
 
 def test_available_time_steps():