flashmd 0.2.1__tar.gz → 0.2.2__tar.gz

{flashmd-0.2.1/src/flashmd.egg-info → flashmd-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.1
+Version: 0.2.2
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -65,6 +65,7 @@ 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)
+atoms.set_velocities(atoms.get_velocities() - atoms.get_velocities().mean(axis=0))
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"

{flashmd-0.2.1 → flashmd-0.2.2}/README.md

@@ -36,6 +36,7 @@ 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)
+atoms.set_velocities(atoms.get_velocities() - atoms.get_velocities().mean(axis=0))
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"

{flashmd-0.2.1 → flashmd-0.2.2}/pyproject.toml

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

{flashmd-0.2.1 → flashmd-0.2.2}/src/flashmd/ase/bussi.py

@@ -16,10 +16,11 @@ class Bussi(VelocityVerlet):
         model: AtomisticModel | List[AtomisticModel],
         time_constant: float = 10.0 * ase.units.fs,
         device: str | torch.device = "auto",
-        rescale_energy: bool = True,
+        rescale_energy: bool = False,
+        random_rotation: bool = False,
         **kwargs,
     ):
-        super().__init__(atoms, timestep, model, device, rescale_energy, **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.1 → flashmd-0.2.2}/src/flashmd/ase/langevin.py

@@ -1,8 +1,6 @@
 from .velocity_verlet import VelocityVerlet
 import ase.units
 from typing import List
-
-# from ..utils.pretrained import load_pretrained_models
 from metatomic.torch import AtomisticModel
 import torch
 import ase
@@ -17,14 +15,19 @@ class Langevin(VelocityVerlet):
         temperature_K: float,
         model: AtomisticModel | List[AtomisticModel],
         time_constant: float = 100.0 * ase.units.fs,
+        fixcm: bool = True,
         device: str | torch.device = "auto",
-        rescale_energy: bool = True,
+        rescale_energy: bool = False,
+        random_rotation: bool = False,
         **kwargs,
     ):
-        super().__init__(atoms, timestep, model, device, rescale_energy, **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_velocities().mean(axis=0))
 
     def step(self):
         self.apply_langevin_half_step()
@@ -39,3 +42,5 @@ class Langevin(VelocityVerlet):
             ase.units.kB * self.temperature_K * self.atoms.get_masses()[:, None]
         ) * 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_velocities().mean(axis=0))

{flashmd-0.2.1 → flashmd-0.2.2}/src/flashmd/ase/velocity_verlet.py

@@ -9,6 +9,7 @@ from metatomic.torch import System
 import ase
 from ..stepper import FlashMDStepper
 import numpy as np
+from scipy.spatial.transform import Rotation
 
 
 class VelocityVerlet(MolecularDynamics):
@@ -19,6 +20,7 @@ class VelocityVerlet(MolecularDynamics):
         model: AtomisticModel | List[AtomisticModel],
         device: str | torch.device = "auto",
         rescale_energy: bool = True,
+        random_rotation: bool = False,
         **kwargs,
     ):
         super().__init__(atoms, timestep, **kwargs)
@@ -41,6 +43,7 @@ class VelocityVerlet(MolecularDynamics):
 
         self.stepper = FlashMDStepper(model, self.device)
         self.rescale_energy = rescale_energy
+        self.random_rotation = random_rotation
 
     def step(self):
         if self.rescale_energy:
@@ -49,7 +52,33 @@ class VelocityVerlet(MolecularDynamics):
         system = _convert_atoms_to_system(
             self.atoms, device=self.device, dtype=self.dtype
         )
+
+        if self.random_rotation:
+            # generate a random rotation matrix with SciPy
+            R = torch.tensor(
+                _get_random_rotation(),
+                device=system.positions.device,
+                dtype=system.positions.dtype,
+            )
+            # apply the random rotation
+            old_cell = system.cell
+            system.cell = system.cell @ R.T
+            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)
+            )
+
         new_system = self.stepper.step(system)
+
+        if self.random_rotation:
+            # revert q, p to the original reference frame, load old cell
+            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)
+            )
+
         self.atoms.set_positions(new_system.positions.detach().cpu().numpy())
         self.atoms.set_momenta(
             new_system.get_data("momenta")
@@ -120,3 +149,10 @@ def _convert_atoms_to_system(
         ),
     )
     return system
+
+
+def _get_random_rotation():
+    R = Rotation.random().as_matrix()
+    if np.random.rand() < 0.5:
+        R *= -1  # allow improper rotations
+    return R

flashmd-0.2.2/src/flashmd/constraints.py

@@ -0,0 +1,97 @@
+from typing import Dict, List
+
+import torch
+from metatensor.torch import TensorBlock, TensorMap
+from metatomic.torch import System
+
+
+@torch.jit.script
+def enforce_physical_constraints(
+    systems: List[System],
+    predictions: Dict[str, TensorMap],
+    timestep: float,
+) -> Dict[str, TensorMap]:
+    """
+    Enforces physical constraints in the predictions of a FlashMD model, namely
+    conservation of momentum of the center of mass and uniform linear motion of the
+    center of mass.
+    """
+
+    new_predictions: Dict[str, TensorMap] = {}
+
+    for key, prediction_tmap in predictions.items():
+        if key == "momenta":
+            # conservation of momentum of the center of mass
+            system_sizes = [len(s) for s in systems]
+            masses = [s.get_data("masses").block().values for s in systems]
+            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_com_before = [
+                torch.sum(p, dim=0) / M for p, M in zip(momenta_before, total_masses)
+            ]
+            velocities_com_now = [
+                torch.sum(p, dim=0) / M for p, M in zip(momenta_now, total_masses)
+            ]
+            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
+                )
+            ]
+            momenta_now = [v * m[:, None] for v, m in zip(velocities_now, masses)]
+            new_predictions[key] = TensorMap(
+                prediction_tmap.keys,
+                [
+                    TensorBlock(
+                        values=torch.concatenate(momenta_now),
+                        samples=prediction_tmap.block().samples,
+                        components=prediction_tmap.block().components,
+                        properties=prediction_tmap.block().properties,
+                    )
+                ],
+            )
+        elif key == "positions":
+            # uniform linear motion of the center of mass
+            system_sizes = [len(s) for s in systems]
+            masses = [s.get_data("masses").block().values for s in systems]
+            total_masses = [m.sum() for m in masses]
+            positions_before = [s.positions.unsqueeze(-1) for s in systems]
+            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)
+            ]
+            positions_com_before = [
+                torch.sum(q * m[:, None], dim=0) / M
+                for q, m, M in zip(positions_before, masses, total_masses)
+            ]
+            positions_com_now = [
+                torch.sum(q * m[:, None], dim=0) / M
+                for q, m, M in zip(positions_now, masses, total_masses)
+            ]
+            positions_now = [
+                q - q_com_now_i + q_com_before_i + v_com_i * timestep
+                for q, q_com_now_i, q_com_before_i, v_com_i in zip(
+                    positions_now,
+                    positions_com_now,
+                    positions_com_before,
+                    velocities_com,
+                )
+            ]
+            new_predictions[key] = TensorMap(
+                prediction_tmap.keys,
+                [
+                    TensorBlock(
+                        values=torch.concatenate(positions_now),
+                        samples=prediction_tmap.block().samples,
+                        components=prediction_tmap.block().components,
+                        properties=prediction_tmap.block().properties,
+                    )
+                ],
+            )
+        else:
+            new_predictions[key] = prediction_tmap
+
+    return new_predictions

{flashmd-0.2.1 → flashmd-0.2.2}/src/flashmd/ipi.py

@@ -15,7 +15,7 @@ from metatensor.torch import Labels, TensorBlock, TensorMap
 
 
 def get_standard_vv_step(
-    sim, model=None, device=None, rescale_energy=True, random_rotation=False
+    sim, model=None, device=None, rescale_energy=False, random_rotation=False
 ):
     """
     Returns a velocity Verlet stepper function for i-PI simulations.
@@ -56,7 +56,7 @@ def get_standard_vv_step(
     return vv_step
 
 
-def get_flashmd_vv_step(sim, model, device, rescale_energy=True, 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)
@@ -113,7 +113,7 @@ def get_flashmd_vv_step(sim, model, device, rescale_energy=True, random_rotation
             kinetic_energy = sim.properties("kinetic_md")
             alpha = np.sqrt(1.0 - (new_energy - old_energy) / kinetic_energy)
             motion.beads.p[:] = alpha * dstrip(motion.beads.p)
-            motion.integrator.pconstraints()  # just to make sure; should be cheap
+            motion.integrator.pconstraints()  # just to be sure
 
         info("@flashmd: End of VV step", verbosity.debug)
 
@@ -156,7 +156,7 @@ def get_nvt_stepper(
     sim,
     model,
     device,
-    rescale_energy=True,
+    rescale_energy=False,
     random_rotation=False,
     use_standard_vv=False,
 ):
@@ -224,7 +224,7 @@ def get_npt_stepper(
     sim,
     model,
     device,
-    rescale_energy=True,
+    rescale_energy=False,
     random_rotation=False,
     use_standard_vv=False,
 ):

{flashmd-0.2.1 → flashmd-0.2.2}/src/flashmd/stepper.py

@@ -5,6 +5,8 @@ import torch
 from metatomic.torch import System
 from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists
 from metatomic.torch import AtomisticModel
+from .constraints import enforce_physical_constraints
+import ase.units
 
 
 class FlashMDStepper:
@@ -14,6 +16,7 @@ class FlashMDStepper:
         device: torch.device,
     ):
         self.model = model.to(device)
+        self.time_step = float(model.module.timestep) * ase.units.fs
 
         # one of these for each model:
         self.evaluation_options = ModelEvaluationOptions(
@@ -41,6 +44,10 @@ class FlashMDStepper:
         model_outputs = self.model(
             [system], self.evaluation_options, check_consistency=False
         )
+        model_outputs = enforce_physical_constraints(
+            [system], model_outputs, timestep=self.time_step
+        )
+
         new_q = model_outputs["positions"].block().values.squeeze(-1)
         new_p = model_outputs["momenta"].block().values.squeeze(-1)
 

{flashmd-0.2.1 → flashmd-0.2.2/src/flashmd.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.1
+Version: 0.2.2
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
 License: Apache-2.0
@@ -65,6 +65,7 @@ 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)
+atoms.set_velocities(atoms.get_velocities() - atoms.get_velocities().mean(axis=0))
 
 # Load models
 device="cuda" if torch.cuda.is_available() else "cpu"

{flashmd-0.2.1 → flashmd-0.2.2}/src/flashmd.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ LICENSE
 README.md
 pyproject.toml
 src/flashmd/__init__.py
+src/flashmd/constraints.py
 src/flashmd/ipi.py
 src/flashmd/models.py
 src/flashmd/stepper.py