flashmd 0.2.6__tar.gz → 0.2.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.6
+Version: 0.2.8
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -19,10 +19,10 @@ Classifier: Topic :: Scientific/Engineering :: Chemistry
 Classifier: Topic :: Scientific/Engineering :: Physics
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2026.1
+Requires-Dist: metatrain==2026.2
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -55,11 +55,11 @@ 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
 
 
@@ -74,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-v2", 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
@@ -88,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
@@ -141,10 +141,13 @@ 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
+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.
 
+The trajectory datasets generated with the PET-MAD baseline MLIPs used in the paper can be found in [this
+repository](https://zenodo.org/records/17904449).
+
 Disclaimer
 ----------
 
@@ -156,6 +159,11 @@ parts of the simulated system. For the time being, we also recommend checking al
 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.
 
+#### More energy models
+
+Do you want to try out different energy models (MLIPs),
+trained on different levels of theory and/or with different accuracy-speed tradeoffs? Check
+out the **[UPET](https://github.com/lab-cosmo/upet) repository.**
 
 Publication
 -----------

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

@@ -22,11 +22,11 @@ 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
 
 
@@ -41,11 +41,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-v2", 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
@@ -55,9 +55,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
@@ -108,10 +108,13 @@ 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
+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.
 
+The trajectory datasets generated with the PET-MAD baseline MLIPs used in the paper can be found in [this
+repository](https://zenodo.org/records/17904449).
+
 Disclaimer
 ----------
 
@@ -123,6 +126,11 @@ parts of the simulated system. For the time being, we also recommend checking al
 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.
 
+#### More energy models
+
+Do you want to try out different energy models (MLIPs),
+trained on different levels of theory and/or with different accuracy-speed tradeoffs? Check
+out the **[UPET](https://github.com/lab-cosmo/upet) repository.**
 
 Publication
 -----------

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

@@ -1,7 +1,7 @@
 [project]
 name = "flashmd"
-version = "0.2.6"
-requires-python = ">=3.9"
+version = "0.2.8"
+requires-python = ">=3.10"
 
 readme = "README.md"
 license = {text = "Apache-2.0"}
@@ -9,7 +9,7 @@ description = "Accelerated molecular dynamics with large-time-step predictions"
 authors = [{name = "flashmd developers"}]
 
 dependencies = [
-    "metatrain==2026.1",
+    "metatrain==2026.2",
     "ase",
     "huggingface_hub",
 ]
@@ -74,10 +74,9 @@ 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",
+    "ignore:`torch.jit.script` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.load` is deprecated:DeprecationWarning",
 ]

flashmd-0.2.8/src/flashmd/__init__.py

@@ -0,0 +1,21 @@
+import warnings
+
+import torch
+
+from .models import get_pretrained as get_pretrained
+from .models import save_checkpoint as save_checkpoint
+
+
+warnings.filterwarnings("ignore", category=UserWarning, message="custom data")
+warnings.filterwarnings(
+    "ignore", category=DeprecationWarning, message="`torch.jit.script`"
+)
+warnings.filterwarnings(
+    "ignore", category=DeprecationWarning, message="`torch.jit.load`"
+)
+
+
+# Disable static fusion. Besides the fact that atomistic batches have variable
+# sizes, statically fused CUDA kernels cannot allocate new tensors at runtime,
+# causing "Global alloc not supported yet" errors (cuda 13+) at the time of writing
+torch.jit.set_fusion_strategy([("DYNAMIC", 10)])

flashmd-0.2.8/src/flashmd/ase/__init__.py

@@ -0,0 +1,13 @@
+from metatomic.torch.ase_calculator import MetatomicCalculator
+
+
+class EnergyCalculator(MetatomicCalculator):
+    """
+    ASE calculator for energy predictions using a metatomic AtomisticModel.
+
+    Slightly modified to ensure compatibility with ase.io.Trajectory, otherwise
+    completely equivalent to `metatomic.torch.ase_calculator.MetatomicCalculator`.
+    """
+
+    def todict(self):
+        return {"name": "flashmd.ase.EnergyCalculator"}

{flashmd-0.2.6 → flashmd-0.2.8}/src/flashmd/constraints.py

@@ -27,20 +27,29 @@ def enforce_physical_constraints(
             total_masses = [m.sum() for m in masses]
             momenta_before = [s.get_data("momenta").block().values for s in systems]
             momenta_now = torch.split(prediction_tmap.block().values, system_sizes)
-            velocities_now = [p / m[:, None] for p, m in zip(momenta_now, masses)]
+            velocities_now = [
+                p / m[:, None] for p, m in zip(momenta_now, masses, strict=True)
+            ]
             velocities_com_before = [
-                torch.sum(p, dim=0) / M for p, M in zip(momenta_before, total_masses)
+                torch.sum(p, dim=0) / M
+                for p, M in zip(momenta_before, total_masses, strict=True)
             ]
             velocities_com_now = [
-                torch.sum(p, dim=0) / M for p, M in zip(momenta_now, total_masses)
+                torch.sum(p, dim=0) / M
+                for p, M in zip(momenta_now, total_masses, strict=True)
             ]
             velocities_now = [
                 v - v_com_now_i + v_com_before_i
                 for v, v_com_before_i, v_com_now_i in zip(
-                    velocities_now, velocities_com_before, velocities_com_now
+                    velocities_now,
+                    velocities_com_before,
+                    velocities_com_now,
+                    strict=True,
                 )
             ]
-            momenta_now = [v * m[:, None] for v, m in zip(velocities_now, masses)]
+            momenta_now = [
+                v * m[:, None] for v, m in zip(velocities_now, masses, strict=True)
+            ]
             new_predictions[key] = TensorMap(
                 prediction_tmap.keys,
                 [
@@ -61,15 +70,16 @@ def enforce_physical_constraints(
             momenta = [s.get_data("momenta").block().values for s in systems]
             positions_now = torch.split(prediction_tmap.block().values, system_sizes)
             velocities_com = [
-                torch.sum(p, dim=0) / M for p, M in zip(momenta, total_masses)
+                torch.sum(p, dim=0) / M
+                for p, M in zip(momenta, total_masses, strict=True)
             ]
             positions_com_before = [
                 torch.sum(q * m[:, None], dim=0) / M
-                for q, m, M in zip(positions_before, masses, total_masses)
+                for q, m, M in zip(positions_before, masses, total_masses, strict=True)
             ]
             positions_com_now = [
                 torch.sum(q * m[:, None], dim=0) / M
-                for q, m, M in zip(positions_now, masses, total_masses)
+                for q, m, M in zip(positions_now, masses, total_masses, strict=True)
             ]
             positions_now = [
                 q - q_com_now_i + q_com_before_i + v_com_i * timestep
@@ -78,6 +88,7 @@ def enforce_physical_constraints(
                     positions_com_now,
                     positions_com_before,
                     velocities_com,
+                    strict=True,
                 )
             ]
             new_predictions[key] = TensorMap(

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

@@ -1,9 +1,9 @@
 # from ..utils.pretrained import load_pretrained_models
 import ase.units
 import torch
+import vesin.metatomic
 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 .constraints import enforce_physical_constraints
 
@@ -35,9 +35,7 @@ class FlashMDStepper:
         if system.positions.dtype != self.dtype:
             raise ValueError("System dtype does not match stepper dtype.")
 
-        system = get_system_with_neighbor_lists(
-            system, self.model.requested_neighbor_lists()
-        )
+        vesin.metatomic.compute_requested_neighbors([system], "angstrom", self.model)
 
         masses = system.get_data("masses").block().values
         model_outputs = self.model(

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.6
+Version: 0.2.8
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -19,10 +19,10 @@ Classifier: Topic :: Scientific/Engineering :: Chemistry
 Classifier: Topic :: Scientific/Engineering :: Physics
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2026.1
+Requires-Dist: metatrain==2026.2
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -55,11 +55,11 @@ 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
 
 
@@ -74,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-v2", 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
@@ -88,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
@@ -141,10 +141,13 @@ 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
+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.
 
+The trajectory datasets generated with the PET-MAD baseline MLIPs used in the paper can be found in [this
+repository](https://zenodo.org/records/17904449).
+
 Disclaimer
 ----------
 
@@ -156,6 +159,11 @@ parts of the simulated system. For the time being, we also recommend checking al
 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.
 
+#### More energy models
+
+Do you want to try out different energy models (MLIPs),
+trained on different levels of theory and/or with different accuracy-speed tradeoffs? Check
+out the **[UPET](https://github.com/lab-cosmo/upet) repository.**
 
 Publication
 -----------

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

@@ -15,5 +15,7 @@ src/flashmd/ase/__init__.py
 src/flashmd/ase/bussi.py
 src/flashmd/ase/langevin.py
 src/flashmd/ase/velocity_verlet.py
+tests/test_edge_cases.py
+tests/test_energy.py
 tests/test_models.py
 tests/test_readme.py

{flashmd-0.2.6 → flashmd-0.2.8}/src/flashmd.egg-info/requires.txt

@@ -1,4 +1,4 @@
-metatrain==2026.1
+metatrain==2026.2
 ase
 huggingface_hub
 

flashmd-0.2.8/tests/test_edge_cases.py

@@ -0,0 +1,55 @@
+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.velocity_verlet import VelocityVerlet
+
+
+def test_isolated_atom(monkeypatch, tmp_path):
+    """Test that a short MD run completes without errors on an isolated atom."""
+    monkeypatch.chdir(tmp_path)
+
+    atoms = ase.Atoms("O", positions=[[0, 0, 0]])
+    MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+
+    time_step = 8
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    energy_model, flashmd_model = get_pretrained("pet-omatpes-v2", time_step)
+    calculator = EnergyCalculator(energy_model, device=device)
+    atoms.calc = calculator
+
+    dyn = VelocityVerlet(
+        atoms=atoms,
+        timestep=time_step * ase.units.fs,
+        model=flashmd_model,
+        device=device,
+    )
+    dyn.run(10)
+
+
+def test_slab_plus_isolated_atom(monkeypatch, tmp_path):
+    """Test that a short MD run completes without errors on a slab plus an isolated atom."""
+    monkeypatch.chdir(tmp_path)
+
+    # Create a slab and an isolated atom
+    slab = ase.build.fcc111("Al", size=(2, 2, 3), vacuum=10)
+    isolated_atom = ase.Atoms("O", positions=[[0, 0, 24]])
+    atoms = slab + isolated_atom
+    MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+
+    time_step = 8
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    energy_model, flashmd_model = get_pretrained("pet-omatpes-v2", time_step)
+    calculator = EnergyCalculator(energy_model, device=device)
+    atoms.calc = calculator
+
+    dyn = VelocityVerlet(
+        atoms=atoms,
+        timestep=time_step * ase.units.fs,
+        model=flashmd_model,
+        device=device,
+    )
+    dyn.run(10)

flashmd-0.2.8/tests/test_energy.py

@@ -0,0 +1,29 @@
+import ase.build
+import ase.io
+import ase.units
+import torch
+from ase.md import VelocityVerlet
+
+from flashmd import get_pretrained
+from flashmd.ase import EnergyCalculator
+
+
+def test_md(monkeypatch, tmp_path):
+    """Test that a short MD run completes without errors with a Trajectory file."""
+    monkeypatch.chdir(tmp_path)
+
+    atoms = ase.build.bulk("Al", "fcc", cubic=True)
+
+    time_step = 64
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    energy_model, _ = get_pretrained("pet-omatpes-v2", time_step)
+    calculator = EnergyCalculator(
+        energy_model, device=device, do_gradients_with_energy=False
+    )
+    atoms.calc = calculator
+
+    dyn = VelocityVerlet(atoms=atoms, timestep=time_step * ase.units.fs)
+    traj = ase.io.Trajectory("test_md.traj", "w", atoms)
+    dyn.attach(traj.write)
+    dyn.run(10)
+    traj.close()

flashmd-0.2.6/src/flashmd/__init__.py

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