flashmd 0.2.5__tar.gz → 0.2.7__tar.gz

{flashmd-0.2.5 → flashmd-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.5
+Version: 0.2.7
 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.12
+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,16 +50,16 @@ 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
 import ase.units
-from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 import torch
-from metatomic.torch.ase_calculator import MetatomicCalculator
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 
 from flashmd import get_pretrained
+from flashmd.ase import EnergyCalculator
 from flashmd.ase.langevin import Langevin
 
 
@@ -70,11 +74,11 @@ 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)  
+device = "cuda" if torch.cuda.is_available() else "cpu"
+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)
+# Set the energy model (optional, see below for more precise usage)
+calculator = EnergyCalculator(energy_model, device=device)
 atoms.calc = calculator
 
 # Run MD
@@ -84,9 +88,9 @@ dyn = Langevin(
     temperature_K=300,
     time_constant=100*ase.units.fs,
     model=flashmd_model,
-    device=device
+    device=device,
 )
-dyn.run(1000)
+dyn.run(1000)  # this is 64 ps!
 ```
 
 [The first time you use this code and call the `get_pretrained` function, the
@@ -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
------------------------------------------------------
+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.
 
-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...
+Using FlashMD in i-PI
+---------------------
 
-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
-```
+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 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.
+Models
+------
 
-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:
+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.
 
-```
-from metatomic.torch.ase_calculator import MetatomicCalculator
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
 
-...  # 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 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.7/README.md

@@ -0,0 +1,152 @@
+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. 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.
+
+ASE Quickstart (see below for LAMMPS)
+-------------------------------------
+
+You can install the package with
+
+```bash
+  pip install flashmd
+```
+
+After installation, you can run accelerated molecular dynamics as follows:
+
+```py
+import ase.build
+import ase.units
+import torch
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+
+from flashmd import get_pretrained
+from flashmd.ase import EnergyCalculator
+from flashmd.ase.langevin import Langevin
+
+
+# Choose your time step (go for 10-30x what you would use in normal MD for your system)
+time_step = 64  # 64 fs; also available: 1, 2, 4, 8, 16, 32, 128 fs
+
+# Create a structure and initialize velocities
+atoms = ase.build.bulk("Al", "fcc", cubic=True)
+MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+atoms.set_velocities(  # it is generally a good idea to remove any net velocity
+    atoms.get_velocities() - atoms.get_momenta().sum(axis=0) / atoms.get_masses().sum()
+)
+
+# Load models
+device = "cuda" if torch.cuda.is_available() else "cpu"
+energy_model, flashmd_model = get_pretrained("pet-omatpes-v2", time_step)
+
+# Set the energy model (optional, see below for more precise usage)
+calculator = EnergyCalculator(energy_model, device=device)
+atoms.calc = calculator
+
+# Run MD
+dyn = Langevin(
+    atoms=atoms,
+    timestep=time_step*ase.units.fs,
+    temperature_K=300,
+    time_constant=100*ase.units.fs,
+    model=flashmd_model,
+    device=device,
+)
+dyn.run(1000)  # this is 64 ps!
+```
+
+[The first time you use this code and call the `get_pretrained` function, the
+pre-trained models will be downloaded. This might take a bit.]
+
+Other available integrators:
+
+```py
+  from flashmd.ase.velocity_verlet import VelocityVerlet
+  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
+---------------
+
+Stick to 10-30x what you would use in normal MD for your system! The 64 fs example
+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
+-----------------------
+
+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.
+
+Using FlashMD in i-PI
+---------------------
+
+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.
+
+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.
+
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
+
+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 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](https://arxiv.org/abs/2505.19350):
+
+```
+@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}
+}
+```
+
+Reproducing the results in the article is supported with FlashMD v0.1.2:
+
+```bash
+pip install flashmd==0.1.2 ase==3.24.0 pet-mad==1.4.3
+```
+
+and using the "PET-MAD" models (PBEsol) from https://huggingface.co/lab-cosmo/flashmd.
+Note that the results were obtained through the i-PI interface.
+
+Instructions and material to reproduce the results in the paper are available on
+Materials Cloud at https://doi.org/10.24435/materialscloud:b7-xq.

{flashmd-0.2.5 → flashmd-0.2.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flashmd"
-version = "0.2.5"
+version = "0.2.7"
 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.12",
+    "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.5 → flashmd-0.2.7}/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.7/src/flashmd/ase/__init__.py

@@ -0,0 +1,26 @@
+import tempfile
+
+from metatomic.torch import AtomisticModel
+from metatomic.torch.ase_calculator import MetatomicCalculator
+
+
+class EnergyCalculator(MetatomicCalculator):
+    """
+    ASE calculator for energy predictions using a metatomic AtomisticModel.
+
+    Slightly modified to save the model to a temporary file to ensure compatibility
+    with ase.io.Trajectory.
+    """
+
+    def __init__(self, model, *args, **kwargs):
+        # save the model to a path otherwise it won't work with ase.io.Trajectory
+        # which calls todict on the calculator
+
+        if isinstance(model, AtomisticModel):
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f:
+                path = f.name
+            model.save(path)
+        else:
+            path = model
+
+        super().__init__(path, *args, **kwargs)

{flashmd-0.2.5 → flashmd-0.2.7}/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.5 → flashmd-0.2.7}/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.5 → flashmd-0.2.7}/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.5 → flashmd-0.2.7}/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)