flashmd 0.2.6__tar.gz → 0.2.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.6
+Version: 0.2.7
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -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

{flashmd-0.2.6 → flashmd-0.2.7}/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

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

@@ -1,6 +1,6 @@
 [project]
 name = "flashmd"
-version = "0.2.6"
+version = "0.2.7"
 requires-python = ">=3.9"
 
 readme = "README.md"

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.6 → flashmd-0.2.7}/src/flashmd.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.6
+Version: 0.2.7
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -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

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

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

flashmd-0.2.7/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()