flashmd 0.2.4__tar.gz → 0.2.6__tar.gz

{flashmd-0.2.4 → flashmd-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.4
+Version: 0.2.6
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -22,9 +22,13 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2025.11
+Requires-Dist: metatrain==2026.1
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
 Dynamic: license-file
 
 FlashMD: universal long-stride molecular dynamics
@@ -32,13 +36,13 @@ FlashMD: universal long-stride molecular dynamics
 
 This repository contains custom integrators to run MD trajectories with FlashMD models. These models are
 designed to learn and predict molecular dynamics trajectories using long strides, therefore allowing
-very large time steps. Before using this method, make sure you are aware of its limitations, which are
+very large time steps. When using this method, make sure you are aware of its limitations, which are
 discussed in [this preprint](http://arxiv.org/abs/2505.19350).
 
 The pre-trained models we make available are trained to reproduce ab-initio MD at the r2SCAN level of theory.
 
-Quickstart
-----------
+ASE Quickstart (see below for LAMMPS)
+-------------------------------------
 
 You can install the package with
 
@@ -46,7 +50,7 @@ You can install the package with
   pip install flashmd
 ```
 
-After installation, you can run accelerated molecular dynamics with ASE as follows:
+After installation, you can run accelerated molecular dynamics as follows:
 
 ```py
 import ase.build
@@ -71,7 +75,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)
@@ -99,6 +103,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
 ---------------
 
@@ -107,77 +116,59 @@ 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)
 
+Using FlashMD in LAMMPS
+-----------------------
 
-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):
+LAMMPS can allow you to run FlashMD with better computational and memory efficiency.
+Furthermore, it will give you access to more sophisticated types of simulations, such as
+simulations in the NPT ensemble and metadynamics. See [here](docs/lammps.md) for a guide
+on using FlashMD in LAMMPS.
 
-```
-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
-```
+Using FlashMD in i-PI
+---------------------
 
-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.
+You can see
+[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+for usage examples. i-PI is our most mature interface, and the one that was used to
+generate all our published results.
 
-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:
+Models
+------
 
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
+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.
 
-...  # setting up atoms
-calculator = MetatomicCalculator(energy_model, device=device)
-atoms.calc = calculator
-...  # running MD
-```
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
 
-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 don't need forces.
+FlashMD models can be trained from the metatrain library. This
+[tutorial](https://docs.metatensor.org/metatrain/latest/generated_examples/1-advanced/04-flashmd.html)
+shows how to train your own FlashMD model, either from scratch or by fine-tuning of one of our universal models.
 
 Disclaimer
 ----------
 
-This is experimental software and should only be used if you know what you're doing.
-We recommend using the i-PI integrators for any serious work and/or if you need to perform
-constant-pressure (NPT) molecular dynamics. You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
-for a usage example.
-Given that the main issue we observe in direct MD trajectories is loss of equipartition
-of energy between different degrees of freedom, we recommend using a local Langevin
-thermostat, and to monitor the temperature of different atomic types or different
-parts of the simulated system. 
+This is experimental software and should only be used if you know what you are doing.
+Given that the main issue we observe in FlashMD is loss of equipartition
+of energy between different degrees of freedom, we recommend using a Langevin
+thermostat, possibly monitoring the temperature of different atomic types or different
+parts of the simulated system. For the time being, we also recommend checking all
+FlashMD-powered findings with traditional MD. The energy models that were used to train
+FlashMD, and that we make available in this repository, can be used for this purpose.
 
 
 Publication
 -----------
 
-If you found FlashMD useful, you can cite the corresponding article:
+If you found FlashMD useful, you can cite the corresponding [article](https://arxiv.org/abs/2505.19350):
 
 ```
-@article{FlashMD,
-  title={FlashMD: long-stride, universal prediction of molecular dynamics},
-  author={Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
-  journal={arXiv preprint arXiv:2505.19350},
-  year={2025}
+@inproceedings{FlashMD,
+  title     = {FlashMD: long-stride, universal prediction of molecular dynamics},
+  author    = {Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
+  booktitle = {Advances in Neural Information Processing Systems},
+  year      = {2025},
+  url       = {https://arxiv.org/abs/2505.19350}
 }
 ```
 

{flashmd-0.2.4 → flashmd-0.2.6}/README.md

@@ -3,13 +3,13 @@ FlashMD: universal long-stride molecular dynamics
 
 This repository contains custom integrators to run MD trajectories with FlashMD models. These models are
 designed to learn and predict molecular dynamics trajectories using long strides, therefore allowing
-very large time steps. Before using this method, make sure you are aware of its limitations, which are
+very large time steps. When using this method, make sure you are aware of its limitations, which are
 discussed in [this preprint](http://arxiv.org/abs/2505.19350).
 
 The pre-trained models we make available are trained to reproduce ab-initio MD at the r2SCAN level of theory.
 
-Quickstart
-----------
+ASE Quickstart (see below for LAMMPS)
+-------------------------------------
 
 You can install the package with
 
@@ -17,7 +17,7 @@ You can install the package with
   pip install flashmd
 ```
 
-After installation, you can run accelerated molecular dynamics with ASE as follows:
+After installation, you can run accelerated molecular dynamics as follows:
 
 ```py
 import ase.build
@@ -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,77 +83,59 @@ 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)
 
+Using FlashMD in LAMMPS
+-----------------------
 
-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):
+LAMMPS can allow you to run FlashMD with better computational and memory efficiency.
+Furthermore, it will give you access to more sophisticated types of simulations, such as
+simulations in the NPT ensemble and metadynamics. See [here](docs/lammps.md) for a guide
+on using FlashMD in LAMMPS.
 
-```
-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
-```
+Using FlashMD in i-PI
+---------------------
 
-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.
+You can see
+[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+for usage examples. i-PI is our most mature interface, and the one that was used to
+generate all our published results.
 
-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:
+Models
+------
 
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
+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.
 
-...  # setting up atoms
-calculator = MetatomicCalculator(energy_model, device=device)
-atoms.calc = calculator
-...  # running MD
-```
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
 
-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 don't need forces.
+FlashMD models can be trained from the metatrain library. This
+[tutorial](https://docs.metatensor.org/metatrain/latest/generated_examples/1-advanced/04-flashmd.html)
+shows how to train your own FlashMD model, either from scratch or by fine-tuning of one of our universal models.
 
 Disclaimer
 ----------
 
-This is experimental software and should only be used if you know what you're doing.
-We recommend using the i-PI integrators for any serious work and/or if you need to perform
-constant-pressure (NPT) molecular dynamics. You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
-for a usage example.
-Given that the main issue we observe in direct MD trajectories is loss of equipartition
-of energy between different degrees of freedom, we recommend using a local Langevin
-thermostat, and to monitor the temperature of different atomic types or different
-parts of the simulated system. 
+This is experimental software and should only be used if you know what you are doing.
+Given that the main issue we observe in FlashMD is loss of equipartition
+of energy between different degrees of freedom, we recommend using a Langevin
+thermostat, possibly monitoring the temperature of different atomic types or different
+parts of the simulated system. For the time being, we also recommend checking all
+FlashMD-powered findings with traditional MD. The energy models that were used to train
+FlashMD, and that we make available in this repository, can be used for this purpose.
 
 
 Publication
 -----------
 
-If you found FlashMD useful, you can cite the corresponding article:
+If you found FlashMD useful, you can cite the corresponding [article](https://arxiv.org/abs/2505.19350):
 
 ```
-@article{FlashMD,
-  title={FlashMD: long-stride, universal prediction of molecular dynamics},
-  author={Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
-  journal={arXiv preprint arXiv:2505.19350},
-  year={2025}
+@inproceedings{FlashMD,
+  title     = {FlashMD: long-stride, universal prediction of molecular dynamics},
+  author    = {Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
+  booktitle = {Advances in Neural Information Processing Systems},
+  year      = {2025},
+  url       = {https://arxiv.org/abs/2505.19350}
 }
 ```
 

{flashmd-0.2.4 → flashmd-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flashmd"
-version = "0.2.4"
+version = "0.2.6"
 requires-python = ">=3.9"
 
 readme = "README.md"
@@ -9,7 +9,7 @@ description = "Accelerated molecular dynamics with large-time-step predictions"
 authors = [{name = "flashmd developers"}]
 
 dependencies = [
-    "metatrain==2025.11",
+    "metatrain==2026.1",
     "ase",
     "huggingface_hub",
 ]
@@ -44,5 +44,40 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff",
+    "mypy",
+]
+
 [tool.ruff]
 line-length = 88
+
+[tool.ruff.lint]
+select = ["E", "F", "B", "I"]
+ignore = ["B018", "B904", "E501"]
+
+[tool.ruff.lint.per-file-ignores]
+"**/{tests}/*" = ["D"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+known-first-party = ["flashmd"]
+
+[tool.ruff.format]
+docstring-code-format = true
+
+[tool.mypy]
+follow_imports = 'skip'
+ignore_missing_imports = true
+disable_error_code = ["union-attr"]
+
+[tool.pytest.ini_options]
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+]
+filterwarnings = [
+    "error",
+    "ignore:custom data:UserWarning",
+]

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/__init__.py

@@ -1,5 +1,7 @@
+import warnings
+
 from .models import get_pretrained as get_pretrained
+from .models import save_checkpoint as save_checkpoint
 
-import warnings
 
 warnings.filterwarnings("ignore", category=UserWarning, message="custom data")

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/ase/bussi.py

@@ -1,10 +1,10 @@
-from .velocity_verlet import VelocityVerlet
-import ase.units
-from typing import List
-from metatomic.torch import AtomisticModel
-import torch
 import ase
+import ase.units
 import numpy as np
+import torch
+from metatomic.torch import AtomisticModel
+
+from .velocity_verlet import VelocityVerlet
 
 
 class Bussi(VelocityVerlet):
@@ -13,14 +13,16 @@ class Bussi(VelocityVerlet):
         atoms: ase.Atoms,
         timestep: float,
         temperature_K: float,
-        model: AtomisticModel | List[AtomisticModel],
+        model: AtomisticModel,
         time_constant: float = 10.0 * ase.units.fs,
         device: str | torch.device = "auto",
         rescale_energy: bool = False,
         random_rotation: bool = False,
         **kwargs,
     ):
-        super().__init__(atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs)
+        super().__init__(
+            atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs
+        )
 
         self.temperature_K = temperature_K
         self.time_constant = time_constant

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/ase/langevin.py

@@ -1,10 +1,10 @@
-from .velocity_verlet import VelocityVerlet
-import ase.units
-from typing import List
-from metatomic.torch import AtomisticModel
-import torch
 import ase
+import ase.units
 import numpy as np
+import torch
+from metatomic.torch import AtomisticModel
+
+from .velocity_verlet import VelocityVerlet
 
 
 class Langevin(VelocityVerlet):
@@ -13,7 +13,7 @@ class Langevin(VelocityVerlet):
         atoms: ase.Atoms,
         timestep: float,
         temperature_K: float,
-        model: AtomisticModel | List[AtomisticModel],
+        model: AtomisticModel,
         time_constant: float = 100.0 * ase.units.fs,
         fixcm: bool = True,
         device: str | torch.device = "auto",
@@ -21,13 +21,18 @@ class Langevin(VelocityVerlet):
         random_rotation: bool = False,
         **kwargs,
     ):
-        super().__init__(atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs)
+        super().__init__(
+            atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs
+        )
 
         self.temperature_K = temperature_K
         self.friction = 1.0 / time_constant
         self.fixcm = fixcm
         if self.fixcm:
-            self.atoms.set_velocities(self.atoms.get_velocities() - self.atoms.get_momenta().sum(axis=0) / self.atoms.get_masses().sum())
+            self.atoms.set_velocities(
+                self.atoms.get_velocities()
+                - self.atoms.get_momenta().sum(axis=0) / self.atoms.get_masses().sum()
+            )
 
     def step(self):
         self.apply_langevin_half_step()
@@ -43,4 +48,7 @@ class Langevin(VelocityVerlet):
         ) * np.random.randn(*old_momenta.shape)
         self.atoms.set_momenta(new_momenta)
         if self.fixcm:
-            self.atoms.set_velocities(self.atoms.get_velocities() - self.atoms.get_momenta().sum(axis=0) / self.atoms.get_masses().sum())
+            self.atoms.set_velocities(
+                self.atoms.get_velocities()
+                - self.atoms.get_momenta().sum(axis=0) / self.atoms.get_masses().sum()
+            )

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/ase/velocity_verlet.py

@@ -1,23 +1,22 @@
-from ase.md.md import MolecularDynamics
-from typing import List
-from metatomic.torch import AtomisticModel
-from metatensor.torch import Labels, TensorBlock, TensorMap
+import ase
 import ase.units
+import numpy as np
 import torch
+from ase.md.md import MolecularDynamics
+from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatomic.torch import AtomisticModel, System
 from metatomic.torch.ase_calculator import _ase_to_torch_data
-from metatomic.torch import System
-import ase
-from ..stepper import FlashMDStepper
-import numpy as np
 from scipy.spatial.transform import Rotation
 
+from ..stepper import FlashMDStepper
+
 
 class VelocityVerlet(MolecularDynamics):
     def __init__(
         self,
         atoms: ase.Atoms,
         timestep: float,
-        model: AtomisticModel | List[AtomisticModel],
+        model: AtomisticModel,
         device: str | torch.device = "auto",
         rescale_energy: bool = True,
         random_rotation: bool = False,
@@ -66,8 +65,8 @@ class VelocityVerlet(MolecularDynamics):
             system.positions = system.positions @ R.T
             # change momentum TensorMap in place
             system.get_data("momenta").block().values[:] = (
-                (system.get_data("momenta").block().values.squeeze(-1) @ R.T).unsqueeze(-1)
-            )
+                system.get_data("momenta").block().values.squeeze(-1) @ R.T
+            ).unsqueeze(-1)
 
         new_system = self.stepper.step(system)
 
@@ -76,8 +75,8 @@ class VelocityVerlet(MolecularDynamics):
             new_system.cell = old_cell
             new_system.positions = new_system.positions @ R
             new_system.get_data("momenta").block().values[:] = (
-                (new_system.get_data("momenta").block().values.squeeze(-1) @ R).unsqueeze(-1)
-            )
+                new_system.get_data("momenta").block().values.squeeze(-1) @ R
+            ).unsqueeze(-1)
 
         self.atoms.set_positions(new_system.positions.detach().cpu().numpy())
         self.atoms.set_momenta(
@@ -127,6 +126,7 @@ class VelocityVerlet(MolecularDynamics):
         for _ in self.irun(steps=steps):
             pass
 
+
 def _convert_atoms_to_system(
     atoms: ase.Atoms, dtype: str, device: str | torch.device
 ) -> System:

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/ipi.py

@@ -1,17 +1,16 @@
+import ase.data
+import ase.units
+import numpy as np
+import torch
+from ipi.engine.motion.dynamics import NPTIntegrator, NVEIntegrator, NVTIntegrator
 from ipi.utils.depend import dstrip
-from ipi.utils.units import Constants
-from ipi.utils.messages import verbosity, info
 from ipi.utils.mathtools import random_rotation as random_rotation_matrix
-from ipi.engine.motion.dynamics import NVEIntegrator, NVTIntegrator, NPTIntegrator
+from ipi.utils.messages import info, verbosity
+from ipi.utils.units import Constants
+from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatomic.torch import System
 
 from flashmd.stepper import FlashMDStepper
-import ase.units
-import torch
-import numpy as np
-import ase.data
-
-from metatomic.torch import System
-from metatensor.torch import Labels, TensorBlock, TensorMap
 
 
 def get_standard_vv_step(
@@ -56,7 +55,9 @@ def get_standard_vv_step(
     return vv_step
 
 
-def get_flashmd_vv_step(sim, model, device, rescale_energy=False, random_rotation=False):
+def get_flashmd_vv_step(
+    sim, model, device, rescale_energy=False, random_rotation=False
+):
     capabilities = model.capabilities()
 
     model_timestep = float(model.module.timestep)

flashmd-0.2.6/src/flashmd/models.py

@@ -0,0 +1,102 @@
+import os
+import shutil
+import subprocess
+import time
+
+from huggingface_hub import hf_hub_download
+from metatomic.torch import AtomisticModel, load_atomistic_model
+
+
+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-v2", time_step: int = 16) -> AtomisticModel:
+    if mlip not in AVAILABLE_MLIPS:
+        raise ValueError(
+            f"MLIP '{mlip}' is not available. "
+            f"Available MLIPs are: {', '.join(AVAILABLE_MLIPS)}."
+        )
+
+    if time_step not in AVAILABLE_TIME_STEPS[mlip]:
+        raise ValueError(
+            f"Pre-trained FlashMD models based on the {mlip} MLIP are only available "
+            f"for time steps of {', '.join(map(str, AVAILABLE_TIME_STEPS[mlip]))} fs."
+        )
+
+    # Get checkpoints corresponding to the selected MLIP and FlashMD models
+    mlip_path = hf_hub_download(
+        repo_id="lab-cosmo/flashmd",
+        filename=f"mlip_{mlip}.ckpt",
+        cache_dir=None,
+        revision="main",
+    )
+    flashmd_path = hf_hub_download(
+        repo_id="lab-cosmo/flashmd",
+        filename=f"flashmd_{mlip}_{time_step}fs.ckpt",
+        cache_dir=None,
+        revision="main",
+    )
+
+    # Now we need to export both using metatrain. However, we don't want to do it if
+    # HuggingFace hasn't downloaded a new version of the files, so we only re-export
+    # if the files above have changed in the last 10 seconds.
+    reexport = False
+    exported_mlip_path = mlip_path.replace(".ckpt", ".pt")
+    exported_flashmd_path = flashmd_path.replace(".ckpt", ".pt")
+    if not os.path.exists(exported_mlip_path) or not os.path.exists(
+        exported_flashmd_path
+    ):
+        reexport = True
+    mlip_mtime = os.path.getmtime(mlip_path)
+    flashmd_mtime = os.path.getmtime(flashmd_path)
+    if (time.time() - mlip_mtime < 10) or (time.time() - flashmd_mtime < 10):
+        reexport = True
+    if reexport:
+        subprocess.run(
+            ["mtt", "export", mlip_path, "-o", exported_mlip_path], capture_output=True
+        )
+        subprocess.run(
+            ["mtt", "export", flashmd_path, "-o", exported_flashmd_path],
+            capture_output=True,
+        )
+
+    # Load as AtomisticModel instances.
+    # If it doesn't work, try to re-export once  and load again (this will, among
+    # others, catch upgrades in metatomic that break compatibility)
+    try:
+        mlip_model = load_atomistic_model(exported_mlip_path)
+        flashmd_model = load_atomistic_model(exported_flashmd_path)
+    except Exception:
+        subprocess.run(
+            ["mtt", "export", mlip_path, "-o", exported_mlip_path], capture_output=True
+        )
+        subprocess.run(
+            ["mtt", "export", flashmd_path, "-o", exported_flashmd_path],
+            capture_output=True,
+        )
+        mlip_model = load_atomistic_model(exported_mlip_path)
+        flashmd_model = load_atomistic_model(exported_flashmd_path)
+
+    return mlip_model, flashmd_model
+
+
+def save_checkpoint(mlip: str = "pet-omatpes-v2", time_step: int = 16):
+    if time_step not in AVAILABLE_TIME_STEPS[mlip]:
+        raise ValueError(
+            f"Pre-trained FlashMD models based on the {mlip} MLIP are only available "
+            f"for time steps of {', '.join(map(str, AVAILABLE_TIME_STEPS[mlip]))} fs."
+        )
+
+    checkpoint_path = hf_hub_download(
+        repo_id="lab-cosmo/flashmd",
+        filename=f"flashmd_{mlip}_{time_step}fs.ckpt",
+        cache_dir=None,
+        revision="main",
+    )
+
+    # Copy it to the current directory
+    shutil.copyfile(checkpoint_path, f"flashmd_{mlip}_{time_step}fs.ckpt")

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd/stepper.py

@@ -1,12 +1,11 @@
 # from ..utils.pretrained import load_pretrained_models
-from metatomic.torch import ModelEvaluationOptions, ModelOutput
-from metatensor.torch import Labels, TensorBlock, TensorMap
+import ase.units
 import torch
-from metatomic.torch import System
+from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatomic.torch import AtomisticModel, ModelEvaluationOptions, ModelOutput, System
 from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists
-from metatomic.torch import AtomisticModel
+
 from .constraints import enforce_physical_constraints
-import ase.units
 
 
 class FlashMDStepper:

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.4
+Version: 0.2.6
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -22,9 +22,13 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2025.11
+Requires-Dist: metatrain==2026.1
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
 Dynamic: license-file
 
 FlashMD: universal long-stride molecular dynamics
@@ -32,13 +36,13 @@ FlashMD: universal long-stride molecular dynamics
 
 This repository contains custom integrators to run MD trajectories with FlashMD models. These models are
 designed to learn and predict molecular dynamics trajectories using long strides, therefore allowing
-very large time steps. Before using this method, make sure you are aware of its limitations, which are
+very large time steps. When using this method, make sure you are aware of its limitations, which are
 discussed in [this preprint](http://arxiv.org/abs/2505.19350).
 
 The pre-trained models we make available are trained to reproduce ab-initio MD at the r2SCAN level of theory.
 
-Quickstart
-----------
+ASE Quickstart (see below for LAMMPS)
+-------------------------------------
 
 You can install the package with
 
@@ -46,7 +50,7 @@ You can install the package with
   pip install flashmd
 ```
 
-After installation, you can run accelerated molecular dynamics with ASE as follows:
+After installation, you can run accelerated molecular dynamics as follows:
 
 ```py
 import ase.build
@@ -71,7 +75,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)
@@ -99,6 +103,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
 ---------------
 
@@ -107,77 +116,59 @@ 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)
 
+Using FlashMD in LAMMPS
+-----------------------
 
-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):
+LAMMPS can allow you to run FlashMD with better computational and memory efficiency.
+Furthermore, it will give you access to more sophisticated types of simulations, such as
+simulations in the NPT ensemble and metadynamics. See [here](docs/lammps.md) for a guide
+on using FlashMD in LAMMPS.
 
-```
-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
-```
+Using FlashMD in i-PI
+---------------------
 
-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.
+You can see
+[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+for usage examples. i-PI is our most mature interface, and the one that was used to
+generate all our published results.
 
-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:
+Models
+------
 
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
+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.
 
-...  # setting up atoms
-calculator = MetatomicCalculator(energy_model, device=device)
-atoms.calc = calculator
-...  # running MD
-```
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
 
-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 don't need forces.
+FlashMD models can be trained from the metatrain library. This
+[tutorial](https://docs.metatensor.org/metatrain/latest/generated_examples/1-advanced/04-flashmd.html)
+shows how to train your own FlashMD model, either from scratch or by fine-tuning of one of our universal models.
 
 Disclaimer
 ----------
 
-This is experimental software and should only be used if you know what you're doing.
-We recommend using the i-PI integrators for any serious work and/or if you need to perform
-constant-pressure (NPT) molecular dynamics. You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
-for a usage example.
-Given that the main issue we observe in direct MD trajectories is loss of equipartition
-of energy between different degrees of freedom, we recommend using a local Langevin
-thermostat, and to monitor the temperature of different atomic types or different
-parts of the simulated system. 
+This is experimental software and should only be used if you know what you are doing.
+Given that the main issue we observe in FlashMD is loss of equipartition
+of energy between different degrees of freedom, we recommend using a Langevin
+thermostat, possibly monitoring the temperature of different atomic types or different
+parts of the simulated system. For the time being, we also recommend checking all
+FlashMD-powered findings with traditional MD. The energy models that were used to train
+FlashMD, and that we make available in this repository, can be used for this purpose.
 
 
 Publication
 -----------
 
-If you found FlashMD useful, you can cite the corresponding article:
+If you found FlashMD useful, you can cite the corresponding [article](https://arxiv.org/abs/2505.19350):
 
 ```
-@article{FlashMD,
-  title={FlashMD: long-stride, universal prediction of molecular dynamics},
-  author={Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
-  journal={arXiv preprint arXiv:2505.19350},
-  year={2025}
+@inproceedings{FlashMD,
+  title     = {FlashMD: long-stride, universal prediction of molecular dynamics},
+  author    = {Bigi, Filippo and Chong, Sanggyu and Kristiadi, Agustinus and Ceriotti, Michele},
+  booktitle = {Advances in Neural Information Processing Systems},
+  year      = {2025},
+  url       = {https://arxiv.org/abs/2505.19350}
 }
 ```
 

{flashmd-0.2.4 → flashmd-0.2.6}/src/flashmd.egg-info/SOURCES.txt

@@ -14,5 +14,6 @@ src/flashmd.egg-info/top_level.txt
 src/flashmd/ase/__init__.py
 src/flashmd/ase/bussi.py
 src/flashmd/ase/langevin.py
-src/flashmd/ase/npt.py
-src/flashmd/ase/velocity_verlet.py
+src/flashmd/ase/velocity_verlet.py
+tests/test_models.py
+tests/test_readme.py

flashmd-0.2.6/src/flashmd.egg-info/requires.txt

@@ -0,0 +1,8 @@
+metatrain==2026.1
+ase
+huggingface_hub
+
+[dev]
+pytest
+ruff
+mypy

flashmd-0.2.6/tests/test_models.py

@@ -0,0 +1,65 @@
+import pytest
+
+from flashmd.models import AVAILABLE_MLIPS, AVAILABLE_TIME_STEPS
+
+
+def test_available_mlips():
+    """Test that AVAILABLE_MLIPS is defined and contains expected values."""
+    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():
+    """Test that AVAILABLE_TIME_STEPS is defined and contains expected values."""
+    assert isinstance(AVAILABLE_TIME_STEPS, dict)
+    assert len(AVAILABLE_TIME_STEPS) > 0
+
+    # Check that each MLIP has time steps defined
+    for mlip in AVAILABLE_MLIPS:
+        assert mlip in AVAILABLE_TIME_STEPS
+        assert isinstance(AVAILABLE_TIME_STEPS[mlip], list)
+        assert len(AVAILABLE_TIME_STEPS[mlip]) > 0
+
+
+def test_get_pretrained_invalid_mlip():
+    """Test that get_pretrained raises ValueError for invalid MLIP."""
+    from flashmd.models import get_pretrained
+
+    with pytest.raises(ValueError, match="MLIP 'invalid_mlip' is not available"):
+        get_pretrained(mlip="invalid_mlip", time_step=16)
+
+
+def test_get_pretrained_invalid_time_step():
+    """Test that get_pretrained raises ValueError for invalid time step."""
+    from flashmd.models import get_pretrained
+
+    with pytest.raises(ValueError, match="Pre-trained FlashMD models"):
+        get_pretrained(mlip="pet-omatpes", time_step=999)
+
+
+def test_save_checkpoint_invalid_time_step():
+    """Test that save_checkpoint raises ValueError for invalid time step."""
+    from flashmd.models import save_checkpoint
+
+    with pytest.raises(ValueError, match="Pre-trained FlashMD models"):
+        save_checkpoint(mlip="pet-omatpes-v2", time_step=999)
+
+
+def test_save_checkpoint(monkeypatch):
+    """Test that save_checkpoint saves the checkpoint file."""
+    from flashmd.models import save_checkpoint
+
+    # Mock hf_hub_download and shutil.copyfile
+    def mock_hf_hub_download(repo_id, filename, cache_dir, revision):
+        return f"/path/to/{filename}"
+
+    def mock_copyfile(src, dst):
+        assert src == "/path/to/flashmd_pet-omatpes-v2_16fs.ckpt"
+        assert dst == "flashmd_pet-omatpes-v2_16fs.ckpt"
+
+    monkeypatch.setattr("flashmd.models.hf_hub_download", mock_hf_hub_download)
+    monkeypatch.setattr("flashmd.models.shutil.copyfile", mock_copyfile)
+
+    save_checkpoint(mlip="pet-omatpes-v2", time_step=16)

flashmd-0.2.6/tests/test_readme.py

@@ -0,0 +1,37 @@
+"""Test that the README example works correctly."""
+
+import re
+from pathlib import Path
+
+
+def extract_python_code_from_readme():
+    """Extract the first Python code block from README.md."""
+    readme_path = Path(__file__).parent.parent / "README.md"
+    content = readme_path.read_text()
+
+    # Find the first ```py or ```python code block
+    pattern = r"```(?:py|python)\n(.*?)\n```"
+    matches = re.findall(pattern, content, re.DOTALL)
+
+    if not matches:
+        raise ValueError("No Python code block found in README.md")
+
+    # Return the first code block (the ASE usage example)
+    return matches[0]
+
+
+def test_readme_example():
+    """Test that the README example runs successfully with 10 steps instead of 1000."""
+    # Extract the code from README
+    code = extract_python_code_from_readme()
+
+    # Modify the code to run for only 10 steps instead of 1000
+    modified_code = code.replace("dyn.run(1000)", "dyn.run(10)")
+
+    # Execute the code
+    exec_globals = {}
+    exec(modified_code, exec_globals)
+
+    # Basic validation: check that the dynamics object was created
+    assert "dyn" in exec_globals
+    assert exec_globals["dyn"] is not None

flashmd-0.2.4/src/flashmd/models.py

@@ -1,61 +0,0 @@
-from huggingface_hub import hf_hub_download
-from metatomic.torch import AtomisticModel, load_atomistic_model
-import subprocess
-import os
-import time
-
-
-AVAILABLE_MLIPS = ["pet-omatpes"]
-AVAILABLE_TIME_STEPS = {
-    "pet-omatpes": [1, 2, 4, 8, 16, 32, 64, 128]
-}
-
-
-def get_pretrained(mlip: str = "pet-omatpes", time_step: int = 16) -> AtomisticModel:
-    if mlip not in AVAILABLE_MLIPS:
-        raise ValueError(
-            f"MLIP '{mlip}' is not available. "
-            f"Available MLIPs are: {', '.join(AVAILABLE_MLIPS)}."
-        )
-
-    if time_step not in AVAILABLE_TIME_STEPS[mlip]:
-        raise ValueError(
-            f"Pretrained FlashMD models based on the {mlip} MLIP are only available "
-            f"for time steps of {', '.join(map(str, AVAILABLE_TIME_STEPS[mlip]))} fs."
-        )
-
-    # Get checkpoints corresponding to the selected MLIP and FlashMD models
-    mlip_path = hf_hub_download(
-        repo_id="lab-cosmo/flashmd",
-        filename=f"mlip_{mlip}.ckpt",
-        cache_dir=None,
-        revision="main",
-    )
-    flashmd_path = hf_hub_download(
-        repo_id="lab-cosmo/flashmd",
-        filename=f"flashmd_{mlip}_{time_step}fs.ckpt",
-        cache_dir=None,
-        revision="main",
-    )
-
-    # Now we need to export both using metatrain. However, we don't want to do it if
-    # HuggingFace hasn't downloaded a new version of the files, so we only re-export
-    # if the files above have changed in the last 10 seconds.
-    reexport = False
-    exported_mlip_path = mlip_path.replace(".ckpt", ".pt")
-    exported_flashmd_path = flashmd_path.replace(".ckpt", ".pt")
-    if not os.path.exists(exported_mlip_path) or not os.path.exists(exported_flashmd_path):
-        reexport = True
-    mlip_mtime = os.path.getmtime(mlip_path)
-    flashmd_mtime = os.path.getmtime(flashmd_path)
-    if (time.time() - mlip_mtime < 10) or (time.time() - flashmd_mtime < 10):
-        reexport = True
-    if reexport:
-        subprocess.run(["mtt", "export", mlip_path, "-o", exported_mlip_path], capture_output=True)
-        subprocess.run(["mtt", "export", flashmd_path, "-o", exported_flashmd_path], capture_output=True)
-
-    # Load as AtomisticModel instances
-    mlip_model = load_atomistic_model(exported_mlip_path)
-    flashmd_model = load_atomistic_model(exported_flashmd_path)
-
-    return mlip_model, flashmd_model

flashmd-0.2.4/src/flashmd.egg-info/requires.txt

@@ -1,3 +0,0 @@
-metatrain==2025.11
-ase
-huggingface_hub