flashmd 0.2.7__tar.gz → 0.2.9__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.7
+Version: 0.2.9
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -19,10 +19,13 @@ 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.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2026.1
+Requires-Dist: metatrain==2026.2
+Requires-Dist: metatomic-ase
+Requires-Dist: vesin>=0.5.4
+Requires-Dist: nvalchemi-toolkit-ops==0.3.0
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -128,7 +131,7 @@ Using FlashMD in i-PI
 ---------------------
 
 You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+[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.
 
@@ -141,10 +144,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 +162,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
 -----------
@@ -183,3 +194,8 @@ 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.
+
+Maintainers
+-----------
+
+This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@frostedoyster](https://github.com/frostedoyster) and [@johannes-spies](https://github.com/johannes-spies), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you haven't received an answer after a couple of days.

{flashmd-0.2.7 → flashmd-0.2.9}/README.md

@@ -95,7 +95,7 @@ Using FlashMD in i-PI
 ---------------------
 
 You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+[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.
 
@@ -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
 -----------
@@ -150,3 +158,8 @@ 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.
+
+Maintainers
+-----------
+
+This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@frostedoyster](https://github.com/frostedoyster) and [@johannes-spies](https://github.com/johannes-spies), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you haven't received an answer after a couple of days.

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

@@ -1,7 +1,7 @@
 [project]
 name = "flashmd"
-version = "0.2.7"
-requires-python = ">=3.9"
+version = "0.2.9"
+requires-python = ">=3.11"
 
 readme = "README.md"
 license = {text = "Apache-2.0"}
@@ -9,7 +9,10 @@ description = "Accelerated molecular dynamics with large-time-step predictions"
 authors = [{name = "flashmd developers"}]
 
 dependencies = [
-    "metatrain==2026.1",
+    "metatrain==2026.2",
+    "metatomic-ase",
+    "vesin>=0.5.4",
+    "nvalchemi-toolkit-ops==0.3.0",  # tricks metatomic-ase into avoiding buggy vesin
     "ase",
     "huggingface_hub",
 ]
@@ -74,10 +77,18 @@ 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",
+    "ignore:`compute_requested_neighbors_from_options` is deprecated and will be removed in a future version:UserWarning",
+    "ignore:hf_xet\\.download_files\\(\\) is deprecated:DeprecationWarning",
+    "ignore:Due to '_pack_', the 'APICLaunchParamRecord' Structure will use memory layout compatible with MSVC:DeprecationWarning",
+    "ignore:`torch.jit.script` is not supported:DeprecationWarning",
+    "ignore:`torch.jit.script_method` is not supported:DeprecationWarning",
+    "ignore:`torch.jit.script_method` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.load` is not supported:DeprecationWarning",
+    "ignore:The 'nvalchemiops.neighborlist' module has been renamed to 'nvalchemiops.neighbors':DeprecationWarning",
+    "ignore:From version 0.3.0 onwards, PyTorch is now an optional dependency:DeprecationWarning",
 ]

flashmd-0.2.9/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.9/src/flashmd/ase/__init__.py

@@ -0,0 +1,13 @@
+from metatomic_ase 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.7 → flashmd-0.2.9}/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.7 → flashmd-0.2.9}/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
 
@@ -28,6 +28,15 @@ class FlashMDStepper:
 
         self.dtype = getattr(torch, self.model.capabilities().dtype)
         self.device = device
+        self.neighbor_list_calculators = vesin.metatomic.neighbor_lists_for_model(
+            "angstrom", self.model
+        )
+
+        # vesin's CUDA brute_force algorithm is broken for triclinic cells,
+        # so we use cell_list for all CUDA calculations to avoid potential issues. See
+        # https://github.com/Luthaf/vesin/issues/157
+        for neighbor_list in self.neighbor_list_calculators:
+            neighbor_list._nl.algorithm = "cell_list"
 
     def step(self, system: System):
         if system.device.type != self.device.type:
@@ -35,9 +44,8 @@ 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()
-        )
+        for calculator in self.neighbor_list_calculators:
+            calculator.add_neighbor_list(system)
 
         masses = system.get_data("masses").block().values
         model_outputs = self.model(

{flashmd-0.2.7 → flashmd-0.2.9}/src/flashmd.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.7
+Version: 0.2.9
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -19,10 +19,13 @@ 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.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2026.1
+Requires-Dist: metatrain==2026.2
+Requires-Dist: metatomic-ase
+Requires-Dist: vesin>=0.5.4
+Requires-Dist: nvalchemi-toolkit-ops==0.3.0
 Requires-Dist: ase
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -128,7 +131,7 @@ Using FlashMD in i-PI
 ---------------------
 
 You can see
-[this cookbook recipe](https://atomistic-cookbook.org/examples/flashmd/flashmd-demo.html) 
+[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.
 
@@ -141,10 +144,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 +162,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
 -----------
@@ -183,3 +194,8 @@ 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.
+
+Maintainers
+-----------
+
+This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@frostedoyster](https://github.com/frostedoyster) and [@johannes-spies](https://github.com/johannes-spies), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you haven't received an answer after a couple of days.

{flashmd-0.2.7 → flashmd-0.2.9}/src/flashmd.egg-info/SOURCES.txt

@@ -15,6 +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.9/src/flashmd.egg-info/requires.txt

@@ -0,0 +1,11 @@
+metatrain==2026.2
+metatomic-ase
+vesin>=0.5.4
+nvalchemi-toolkit-ops==0.3.0
+ase
+huggingface_hub
+
+[dev]
+pytest
+ruff
+mypy

flashmd-0.2.9/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.7/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")

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

@@ -1,26 +0,0 @@
-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.7/src/flashmd.egg-info/requires.txt

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