flashmd 0.1.2__tar.gz → 0.2.0__tar.gz

{flashmd-0.1.2/src/flashmd.egg-info → flashmd-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.1.2
+Version: 0.2.0
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -22,8 +22,8 @@ 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: ase
-Requires-Dist: pet-mad>=1.2.0
 Requires-Dist: huggingface_hub
 Dynamic: license-file
 
@@ -35,6 +35,8 @@ designed to learn and predict molecular dynamics trajectories using long strides
 very large time steps. Before 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
 ----------
 
@@ -51,35 +53,41 @@ import ase.build
 import ase.units
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 import torch
-from pet_mad.calculator import PETMADCalculator
+from metatomic.torch.ase_calculator import MetatomicCalculator
 
-from flashmd import get_universal_model
+from flashmd import get_pretrained
 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 = 16  # 16 fs; also available: 1, 2, 4, 8, 32, 64, 128 fs
+
 # Create a structure and initialize velocities
 atoms = ase.build.bulk("Al", "fcc", cubic=True)
 MaxwellBoltzmannDistribution(atoms, temperature_K=300)
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"
-calculator = PETMADCalculator("1.0.1", device=device)
+energy_model, flashmd_model = get_pretrained("pet-omatpes", time_step)  
+
+calculator = MetatomicCalculator(energy_model, device=device)
 atoms.calc = calculator
-model = get_universal_model(16)  # 16 fs model; also available: 1, 4, 8, 32, 64 fs
-model = model.to(device)
 
 # Run MD
 dyn = Langevin(
     atoms=atoms,
-    timestep=16*ase.units.fs,
+    timestep=time_step*ase.units.fs,
     temperature_K=300,
     time_constant=100*ase.units.fs,
-    model=model,
+    model=flashmd_model,
     device=device
 )
 dyn.run(1000)
 ```
 
+[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
@@ -91,11 +99,35 @@ 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 to perform constant
-pressure, NpT molecular dynamics. You can see
+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. 
+
+
+Publication
+-----------
+
+If you found FlashMD useful, you can cite the corresponding article:
+
+```
+@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}
+}
+```
+
+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.

{flashmd-0.1.2 → flashmd-0.2.0}/README.md

@@ -6,6 +6,8 @@ designed to learn and predict molecular dynamics trajectories using long strides
 very large time steps. Before 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
 ----------
 
@@ -22,35 +24,41 @@ import ase.build
 import ase.units
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 import torch
-from pet_mad.calculator import PETMADCalculator
+from metatomic.torch.ase_calculator import MetatomicCalculator
 
-from flashmd import get_universal_model
+from flashmd import get_pretrained
 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 = 16  # 16 fs; also available: 1, 2, 4, 8, 32, 64, 128 fs
+
 # Create a structure and initialize velocities
 atoms = ase.build.bulk("Al", "fcc", cubic=True)
 MaxwellBoltzmannDistribution(atoms, temperature_K=300)
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"
-calculator = PETMADCalculator("1.0.1", device=device)
+energy_model, flashmd_model = get_pretrained("pet-omatpes", time_step)  
+
+calculator = MetatomicCalculator(energy_model, device=device)
 atoms.calc = calculator
-model = get_universal_model(16)  # 16 fs model; also available: 1, 4, 8, 32, 64 fs
-model = model.to(device)
 
 # Run MD
 dyn = Langevin(
     atoms=atoms,
-    timestep=16*ase.units.fs,
+    timestep=time_step*ase.units.fs,
     temperature_K=300,
     time_constant=100*ase.units.fs,
-    model=model,
+    model=flashmd_model,
     device=device
 )
 dyn.run(1000)
 ```
 
+[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
@@ -62,11 +70,35 @@ 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 to perform constant
-pressure, NpT molecular dynamics. You can see
+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. 
+
+
+Publication
+-----------
+
+If you found FlashMD useful, you can cite the corresponding article:
+
+```
+@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}
+}
+```
+
+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.

{flashmd-0.1.2 → flashmd-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flashmd"
-version = "0.1.2"
+version = "0.2.0"
 requires-python = ">=3.9"
 
 readme = "README.md"
@@ -9,8 +9,8 @@ description = "Accelerated molecular dynamics with large-time-step predictions"
 authors = [{name = "flashmd developers"}]
 
 dependencies = [
+    "metatrain==2025.11",
     "ase",
-    "pet-mad>=1.2.0",
     "huggingface_hub",
 ]
 

{flashmd-0.1.2 → flashmd-0.2.0}/src/flashmd/__init__.py

@@ -1,4 +1,4 @@
-from .models import get_universal_model as get_universal_model
+from .models import get_pretrained as get_pretrained
 
 import warnings
 

{flashmd-0.1.2 → flashmd-0.2.0}/src/flashmd/ase/velocity_verlet.py

@@ -23,19 +23,13 @@ class VelocityVerlet(MolecularDynamics):
     ):
         super().__init__(atoms, timestep, **kwargs)
 
-        models = model if isinstance(model, list) else [model]
-        capabilities = models[0].capabilities()
+        capabilities = model.capabilities()
 
-        base_timestep = float(models[0].module.base_time_step) * ase.units.fs
-
-        n_time_steps = int(
-            [k for k in capabilities.outputs.keys() if "mtt::delta_" in k][0].split(
-                "_"
-            )[1]
-        )
-        if n_time_steps != self.dt / base_timestep:
+        model_timestep = float(model.module.timestep)
+        if not np.allclose(model_timestep, self.dt / ase.units.fs):
             raise ValueError(
-                f"Mismatch between timestep ({self.dt}) and model timestep ({base_timestep})."
+                f"Mismatch between timestep ({self.dt / ase.units.fs} fs) "
+                f"and model timestep ({model_timestep} fs)."
             )
 
         if device == "auto":
@@ -45,7 +39,7 @@ class VelocityVerlet(MolecularDynamics):
         self.device = torch.device(device)
         self.dtype = getattr(torch, capabilities.dtype)
 
-        self.stepper = FlashMDStepper(models, n_time_steps, self.device)
+        self.stepper = FlashMDStepper(model, self.device)
         self.rescale_energy = rescale_energy
 
     def step(self):

{flashmd-0.1.2 → flashmd-0.2.0}/src/flashmd/ipi.py

@@ -39,7 +39,6 @@ def get_standard_vv_step(
             info("@flashmd: Old energy", verbosity.debug)
             old_energy = sim.properties("potential") + sim.properties("kinetic_md")
 
-        print(motion.integrator.pdt, motion.integrator.qdt)
         motion.integrator.pstep(level=0)
         motion.integrator.pconstraints()
         motion.integrator.qcstep()  # does two steps because qdt is halved in the i-PI integrator
@@ -60,21 +59,18 @@ def get_standard_vv_step(
 def get_flashmd_vv_step(sim, model, device, rescale_energy=True, random_rotation=False):
     capabilities = model.capabilities()
 
-    base_timestep = float(model.module.base_time_step) * ase.units.fs
+    model_timestep = float(model.module.timestep)
 
-    dt = sim.syslist[0].motion.dt * 2.4188843e-17 * ase.units.s
+    dt = sim.syslist[0].motion.dt * 2.4188843e-17 * ase.units.s / ase.units.fs
 
-    n_time_steps = int(
-        [k for k in capabilities.outputs.keys() if "mtt::delta_" in k][0].split("_")[1]
-    )
-    if not np.allclose(dt, n_time_steps * base_timestep):
+    if not np.allclose(dt, model_timestep):
         raise ValueError(
-            f"Mismatch between timestep ({dt}) and model timestep ({base_timestep})."
+            f"Mismatch between timestep ({dt} fs) and model timestep ({model_timestep} fs)."
         )
 
     device = torch.device(device)
     dtype = getattr(torch, capabilities.dtype)
-    stepper = FlashMDStepper([model], n_time_steps, device)
+    stepper = FlashMDStepper(model, device)
 
     def flashmd_vv(motion):
         info("@flashmd: Starting VV", verbosity.debug)

flashmd-0.2.0/src/flashmd/models.py

@@ -0,0 +1,61 @@
+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.1.2 → flashmd-0.2.0}/src/flashmd/stepper.py

@@ -4,42 +4,27 @@ from metatensor.torch import Labels, TensorBlock, TensorMap
 import torch
 from metatomic.torch import System
 from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists
-from typing import List
 from metatomic.torch import AtomisticModel
 
 
 class FlashMDStepper:
     def __init__(
         self,
-        models: List[AtomisticModel],
-        n_time_steps: int,
+        model: AtomisticModel,
         device: torch.device,
-        # q_error_threshold: float = 0.1,
-        # p_error_threshold: float = 0.1,
-        # energy_error_threshold: float = 0.1,
     ):
-        self.n_time_steps = n_time_steps
-
-        # internally, turn list of models into a dict and send to device
-        self.models = {}
-        for model in models:
-            n_time_steps_model = int(
-                [k for k in model.capabilities().outputs.keys() if "mtt::delta_" in k][
-                    0
-                ].split("_")[1]
-            )
-            self.models[n_time_steps_model] = model.to(device)
+        self.model = model.to(device)
 
         # one of these for each model:
         self.evaluation_options = ModelEvaluationOptions(
             length_unit="Angstrom",
             outputs={
-                f"mtt::delta_{self.n_time_steps}_q": ModelOutput(per_atom=True),
-                f"mtt::p_{self.n_time_steps}": ModelOutput(per_atom=True),
+                "positions": ModelOutput(per_atom=True),
+                "momenta": ModelOutput(per_atom=True),
             },
         )
 
-        self.dtype = getattr(torch, self.models[n_time_steps].capabilities().dtype)
+        self.dtype = getattr(torch, self.model.capabilities().dtype)
         self.device = device
 
     def step(self, system: System):
@@ -49,27 +34,18 @@ class FlashMDStepper:
             raise ValueError("System dtype does not match stepper dtype.")
 
         system = get_system_with_neighbor_lists(
-            system, self.models[self.n_time_steps].requested_neighbor_lists()
+            system, self.model.requested_neighbor_lists()
         )
 
         masses = system.get_data("masses").block().values
-        model_outputs = self.models[self.n_time_steps](
+        model_outputs = self.model(
             [system], self.evaluation_options, check_consistency=False
         )
-        delta_q_scaled = (
-            model_outputs[f"mtt::delta_{self.n_time_steps}_q"]
-            .block()
-            .values.squeeze(-1)
-        )
-        p_scaled = (
-            model_outputs[f"mtt::p_{self.n_time_steps}"].block().values.squeeze(-1)
-        )
-        sqrt_masses = torch.sqrt(masses)
-        delta_q = delta_q_scaled / sqrt_masses
-        p = p_scaled * sqrt_masses
+        new_q = model_outputs["positions"].block().values.squeeze(-1)
+        new_p = model_outputs["momenta"].block().values.squeeze(-1)
 
         new_system = System(
-            positions=system.positions + delta_q,
+            positions=new_q,
             types=system.types,
             cell=system.cell,
             pbc=system.pbc,
@@ -80,7 +56,7 @@ class FlashMDStepper:
                 keys=Labels.single().to(self.device),
                 blocks=[
                     TensorBlock(
-                        values=p.unsqueeze(-1),
+                        values=new_p.unsqueeze(-1),
                         samples=Labels(
                             names=["system", "atom"],
                             values=torch.tensor(

{flashmd-0.1.2 → flashmd-0.2.0/src/flashmd.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.1.2
+Version: 0.2.0
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -22,8 +22,8 @@ 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: ase
-Requires-Dist: pet-mad>=1.2.0
 Requires-Dist: huggingface_hub
 Dynamic: license-file
 
@@ -35,6 +35,8 @@ designed to learn and predict molecular dynamics trajectories using long strides
 very large time steps. Before 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
 ----------
 
@@ -51,35 +53,41 @@ import ase.build
 import ase.units
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 import torch
-from pet_mad.calculator import PETMADCalculator
+from metatomic.torch.ase_calculator import MetatomicCalculator
 
-from flashmd import get_universal_model
+from flashmd import get_pretrained
 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 = 16  # 16 fs; also available: 1, 2, 4, 8, 32, 64, 128 fs
+
 # Create a structure and initialize velocities
 atoms = ase.build.bulk("Al", "fcc", cubic=True)
 MaxwellBoltzmannDistribution(atoms, temperature_K=300)
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"
-calculator = PETMADCalculator("1.0.1", device=device)
+energy_model, flashmd_model = get_pretrained("pet-omatpes", time_step)  
+
+calculator = MetatomicCalculator(energy_model, device=device)
 atoms.calc = calculator
-model = get_universal_model(16)  # 16 fs model; also available: 1, 4, 8, 32, 64 fs
-model = model.to(device)
 
 # Run MD
 dyn = Langevin(
     atoms=atoms,
-    timestep=16*ase.units.fs,
+    timestep=time_step*ase.units.fs,
     temperature_K=300,
     time_constant=100*ase.units.fs,
-    model=model,
+    model=flashmd_model,
     device=device
 )
 dyn.run(1000)
 ```
 
+[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
@@ -91,11 +99,35 @@ 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 to perform constant
-pressure, NpT molecular dynamics. You can see
+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. 
+
+
+Publication
+-----------
+
+If you found FlashMD useful, you can cite the corresponding article:
+
+```
+@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}
+}
+```
+
+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.

{flashmd-0.1.2 → flashmd-0.2.0}/src/flashmd.egg-info/requires.txt

@@ -1,3 +1,3 @@
+metatrain==2025.11
 ase
-pet-mad>=1.2.0
 huggingface_hub

flashmd-0.1.2/src/flashmd/models.py

@@ -1,18 +0,0 @@
-from huggingface_hub import hf_hub_download
-from metatomic.torch import AtomisticModel, load_atomistic_model
-
-
-def get_universal_model(time_step: int = 16) -> AtomisticModel:
-    if time_step not in [1, 4, 8, 16, 32, 64]:
-        raise ValueError(
-            "Universal FlashMD models are only available for"
-            " time steps of 1, 4, 8, 16, 32, 64 fs."
-        )
-
-    model_path = hf_hub_download(
-        repo_id="lab-cosmo/flashmd",
-        filename=f"flashmd_{time_step}fs.pt",
-        cache_dir=None,
-        revision="main",
-    )
-    return load_atomistic_model(model_path)