flashmd 0.2.8__tar.gz → 0.3.0__tar.gz

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

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: flashmd
-Version: 0.2.8
+Version: 0.3.0
 Summary: Accelerated molecular dynamics with large-time-step predictions
 Author: flashmd developers
-License: Apache-2.0
+License-Expression: Apache-2.0
 Keywords: machine learning,molecular modeling,molecular dynamics
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: POSIX
 Classifier: Operating System :: MacOS :: MacOS X
 Classifier: Operating System :: Microsoft :: Windows
@@ -19,12 +18,15 @@ 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.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: metatrain==2026.2
-Requires-Dist: ase
+Requires-Dist: metatrain==2026.3.1
+Requires-Dist: metatomic-ase
+Requires-Dist: vesin>=0.5.8
+Requires-Dist: ase>=3.29.0
 Requires-Dist: huggingface_hub
+Requires-Dist: scipy
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
@@ -56,7 +58,7 @@ After installation, you can run accelerated molecular dynamics as follows:
 import ase.build
 import ase.units
 import torch
-from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase.md.velocitydistribution import thermalize_momenta
 
 from flashmd import get_pretrained
 from flashmd.ase import EnergyCalculator
@@ -68,7 +70,7 @@ time_step = 64  # 64 fs; also available: 1, 2, 4, 8, 16, 32, 128 fs
 
 # Create a structure and initialize velocities
 atoms = ase.build.bulk("Al", "fcc", cubic=True)
-MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+thermalize_momenta(atoms, temperature_K=300)
 atoms.set_velocities(  # it is generally a good idea to remove any net velocity
     atoms.get_velocities() - atoms.get_momenta().sum(axis=0) / atoms.get_masses().sum()
 )
@@ -128,7 +130,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.
 
@@ -191,3 +193,10 @@ 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 [@johannes-spies](https://github.com/johannes-spies) and [@zyxwwxyz](https://github.com/zyxwwxyz), 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.
+
+It was previously maintained by [@frostedoyster](https://github.com/frostedoyster) (2025 to 2026), many thanks to them for creating FlashMD!

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

@@ -23,7 +23,7 @@ After installation, you can run accelerated molecular dynamics as follows:
 import ase.build
 import ase.units
 import torch
-from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase.md.velocitydistribution import thermalize_momenta
 
 from flashmd import get_pretrained
 from flashmd.ase import EnergyCalculator
@@ -35,7 +35,7 @@ time_step = 64  # 64 fs; also available: 1, 2, 4, 8, 16, 32, 128 fs
 
 # Create a structure and initialize velocities
 atoms = ase.build.bulk("Al", "fcc", cubic=True)
-MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+thermalize_momenta(atoms, temperature_K=300)
 atoms.set_velocities(  # it is generally a good idea to remove any net velocity
     atoms.get_velocities() - atoms.get_momenta().sum(axis=0) / atoms.get_masses().sum()
 )
@@ -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.
 
@@ -158,3 +158,10 @@ 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 [@johannes-spies](https://github.com/johannes-spies) and [@zyxwwxyz](https://github.com/zyxwwxyz), 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.
+
+It was previously maintained by [@frostedoyster](https://github.com/frostedoyster) (2025 to 2026), many thanks to them for creating FlashMD!

flashmd-0.3.0/pyproject.toml

@@ -0,0 +1,105 @@
+[project]
+name = "flashmd"
+version = "0.3.0"
+requires-python = ">=3.11"
+
+readme = "README.md"
+license = "Apache-2.0"
+description = "Accelerated molecular dynamics with large-time-step predictions"
+authors = [{name = "flashmd developers"}]
+
+dependencies = [
+    "metatrain==2026.3.1",
+    "metatomic-ase",
+    "vesin>=0.5.8",
+    "ase>=3.29.0",  # thermalize_momenta replaces MaxwellBoltzmannDistribution
+    "huggingface_hub",
+    "scipy",
+]
+
+keywords = ["machine learning", "molecular modeling", "molecular dynamics"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Operating System :: POSIX",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: Microsoft :: Windows",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Topic :: Scientific/Engineering :: Chemistry",
+    "Topic :: Scientific/Engineering :: Physics",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+# TODO: add project URLs
+
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff",
+    "mypy",
+]
+
+[tool.ruff]
+line-length = 88
+
+[tool.ruff.lint]
+select = ["E", "F", "B", "I"]
+ignore = ["B018", "B904", "E501"]
+
+[tool.ruff.lint.per-file-ignores]
+"**/{tests}/*" = ["D"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+known-first-party = ["flashmd"]
+
+[tool.ruff.format]
+docstring-code-format = true
+
+[tool.mypy]
+follow_imports = 'skip'
+ignore_missing_imports = true
+disable_error_code = ["union-attr"]
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "error",
+    "ignore:custom data:UserWarning",
+    "ignore:`torch.jit.script` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.load` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.save` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.save` is not supported: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 '.*' Structure will use memory layout compatible with MSVC:DeprecationWarning",
+    "ignore:warp.config.quiet is deprecated: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",
+    "ignore:`torch.jit.interface` is deprecated:DeprecationWarning",
+    "ignore:`torch.jit.set_fusion_strategy` is deprecated:DeprecationWarning",
+    "ignore:the '.*' output name is deprecated:UserWarning",
+    "ignore:the '.*' quantity is deprecated:UserWarning",
+    "ignore:the '.*' quantity is deprecated:DeprecationWarning",
+    "ignore:ModelOutput.quantity is deprecated:DeprecationWarning",
+    "ignore:Setting the shape on a NumPy array has been deprecated in NumPy 2.5:DeprecationWarning",
+    "ignore:currentThread\\(\\) is deprecated:DeprecationWarning",
+    "ignore:unclosed file:ResourceWarning",  # i-PI's initializer leaks its handle
+]

{flashmd-0.2.8 → flashmd-0.3.0}/src/flashmd/ase/__init__.py

@@ -1,4 +1,4 @@
-from metatomic.torch.ase_calculator import MetatomicCalculator
+from metatomic_ase import MetatomicCalculator
 
 
 class EnergyCalculator(MetatomicCalculator):

{flashmd-0.2.8 → flashmd-0.3.0}/src/flashmd/ase/bussi.py

@@ -18,15 +18,29 @@ class Bussi(VelocityVerlet):
         device: str | torch.device = "auto",
         rescale_energy: bool = False,
         random_rotation: bool = False,
+        rng: np.random.Generator | None = None,
         **kwargs,
     ):
         super().__init__(
-            atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs
+            atoms,
+            timestep,
+            model,
+            device,
+            rescale_energy,
+            random_rotation,
+            rng,
+            **kwargs,
         )
 
         self.temperature_K = temperature_K
         self.time_constant = time_constant
 
+        if np.isclose(self.atoms.get_kinetic_energy(), 0.0, rtol=0, atol=1e-12):
+            raise ValueError(
+                "Initial kinetic energy is zero. "
+                "Please set the initial velocities before running Bussi NVT."
+            )
+
     def step(self):
         self.apply_bussi_half_step()
         super().step()
@@ -34,7 +48,7 @@ class Bussi(VelocityVerlet):
 
     def apply_bussi_half_step(self):
         old_kinetic_energy = self.atoms.get_kinetic_energy()
-        n_degrees_of_freedom = 3 * len(self.atoms)
+        n_degrees_of_freedom = self.atoms.get_number_of_degrees_of_freedom()
         target_kinetic_energy = (
             0.5 * ase.units.kB * self.temperature_K * n_degrees_of_freedom
         )
@@ -46,7 +60,7 @@ class Bussi(VelocityVerlet):
             / old_kinetic_energy
             / n_degrees_of_freedom
         )
-        r = np.random.randn(n_degrees_of_freedom)
+        r = self.rng.standard_normal(n_degrees_of_freedom)
         alpha_sq = (
             exp_term
             + energy_scaling_term * np.sum(r**2)

{flashmd-0.2.8 → flashmd-0.3.0}/src/flashmd/ase/langevin.py

@@ -19,10 +19,18 @@ class Langevin(VelocityVerlet):
         device: str | torch.device = "auto",
         rescale_energy: bool = False,
         random_rotation: bool = False,
+        rng: np.random.Generator | None = None,
         **kwargs,
     ):
         super().__init__(
-            atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs
+            atoms,
+            timestep,
+            model,
+            device,
+            rescale_energy,
+            random_rotation,
+            rng,
+            **kwargs,
         )
 
         self.temperature_K = temperature_K
@@ -45,7 +53,7 @@ class Langevin(VelocityVerlet):
             1.0 - np.exp(-self.friction * self.dt)
         ) * np.sqrt(
             ase.units.kB * self.temperature_K * self.atoms.get_masses()[:, None]
-        ) * np.random.randn(*old_momenta.shape)
+        ) * self.rng.standard_normal(old_momenta.shape)
         self.atoms.set_momenta(new_momenta)
         if self.fixcm:
             self.atoms.set_velocities(

{flashmd-0.2.8 → flashmd-0.3.0}/src/flashmd/ase/velocity_verlet.py

@@ -1,3 +1,5 @@
+from types import ModuleType
+
 import ase
 import ase.units
 import numpy as np
@@ -5,10 +7,11 @@ import torch
 from ase.md.md import MolecularDynamics
 from metatensor.torch import Labels, TensorBlock, TensorMap
 from metatomic.torch import AtomisticModel, System
-from metatomic.torch.ase_calculator import _ase_to_torch_data
+from metatomic_ase._calculator import _ase_to_torch_data
 from scipy.spatial.transform import Rotation
 
 from ..stepper import FlashMDStepper
+from ..symplectic_stepper import SymplecticStepper
 
 
 class VelocityVerlet(MolecularDynamics):
@@ -20,13 +23,25 @@ class VelocityVerlet(MolecularDynamics):
         device: str | torch.device = "auto",
         rescale_energy: bool = True,
         random_rotation: bool = False,
+        rng: np.random.Generator | None = None,
         **kwargs,
     ):
         super().__init__(atoms, timestep, **kwargs)
 
-        capabilities = model.capabilities()
+        self.rng = np.random if rng is None else rng
+
+        if isinstance(model, tuple):
+            flashmd_model, symplectic_part = model
+            if isinstance(symplectic_part, tuple):
+                symplectic_model, symplectic_config = symplectic_part
+            else:
+                symplectic_model, symplectic_config = symplectic_part, None
+        else:
+            flashmd_model, symplectic_model, symplectic_config = model, None, None
+
+        capabilities = flashmd_model.capabilities()
 
-        model_timestep = float(model.module.timestep)
+        model_timestep = float(flashmd_model.module.timestep)
         if not np.allclose(model_timestep, self.dt / ase.units.fs):
             raise ValueError(
                 f"Mismatch between timestep ({self.dt / ase.units.fs} fs) "
@@ -40,7 +55,14 @@ class VelocityVerlet(MolecularDynamics):
         self.device = torch.device(device)
         self.dtype = getattr(torch, capabilities.dtype)
 
-        self.stepper = FlashMDStepper(model, self.device)
+        flashmd_stepper = FlashMDStepper(flashmd_model, self.device)
+        self.stepper: FlashMDStepper | SymplecticStepper
+        if symplectic_model is not None:
+            self.stepper = SymplecticStepper(
+                flashmd_stepper, symplectic_model, symplectic_config
+            )
+        else:
+            self.stepper = flashmd_stepper
         self.rescale_energy = rescale_energy
         self.random_rotation = random_rotation
 
@@ -55,7 +77,7 @@ class VelocityVerlet(MolecularDynamics):
         if self.random_rotation:
             # generate a random rotation matrix with SciPy
             R = torch.tensor(
-                _get_random_rotation(),
+                _get_random_rotation(self.rng),
                 device=system.positions.device,
                 dtype=system.positions.dtype,
             )
@@ -132,6 +154,14 @@ def _convert_atoms_to_system(
 ) -> System:
     system_data = _ase_to_torch_data(atoms, dtype=dtype, device=device)
     system = System(*system_data)
+    n = len(atoms)
+    atom_samples = Labels(
+        names=["system", "atom"],
+        values=torch.vstack(
+            [torch.full((n,), 0, device=device), torch.arange(n, device=device)]
+        ).T,
+        assume_unique=True,
+    )
     system.add_data(
         "momenta",
         TensorMap(
@@ -141,12 +171,7 @@ def _convert_atoms_to_system(
                     values=torch.tensor(
                         atoms.get_momenta(), dtype=dtype, device=device
                     ).unsqueeze(-1),
-                    samples=Labels(
-                        names=["system", "atom"],
-                        values=torch.tensor(
-                            [[0, j] for j in range(len(atoms))], device=device
-                        ),
-                    ),
+                    samples=atom_samples,
                     components=[
                         Labels(
                             names="xyz",
@@ -167,12 +192,7 @@ def _convert_atoms_to_system(
                     values=torch.tensor(
                         atoms.get_masses(), dtype=dtype, device=device
                     ).unsqueeze(-1),
-                    samples=Labels(
-                        names=["system", "atom"],
-                        values=torch.tensor(
-                            [[0, j] for j in range(len(atoms))], device=device
-                        ),
-                    ),
+                    samples=atom_samples,
                     components=[],
                     properties=Labels.single().to(device),
                 )
@@ -182,8 +202,13 @@ def _convert_atoms_to_system(
     return system
 
 
-def _get_random_rotation():
-    R = Rotation.random().as_matrix()
-    if np.random.rand() < 0.5:
+def _get_random_rotation(rng: np.random.Generator | ModuleType = np.random):
+    # Haar-uniform rotation from a normalised Gaussian quaternion. This only
+    # needs standard_normal/random, so it works with any RNG following the
+    # ASE convention (the np.random module, a Generator, ...), unlike
+    # Rotation.random(rng=...), which requires an actual Generator.
+    quaternion = rng.standard_normal(4)
+    R = Rotation.from_quat(quaternion / np.linalg.norm(quaternion)).as_matrix()
+    if rng.random() < 0.5:
         R *= -1  # allow improper rotations
     return R

flashmd-0.3.0/src/flashmd/fpi.py

@@ -0,0 +1,86 @@
+from typing import Callable
+
+import torch
+
+
+def anderson_solver(
+    f: Callable[[torch.Tensor], torch.Tensor],
+    x0: torch.Tensor,
+    m: int = 5,
+    max_iter: int = 50,
+    tol: float = 1e-5,
+    beta: float = 1.0,
+    lambda_reg: float = 1e-4,
+    return_residual_norms: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, list[float]]:
+    """
+    Solve fixed-point problem x = f(x) using Anderson acceleration.
+
+    Args:
+        f: Fixed-point mapping.
+        x0: Initial guess.
+        m: Number of previous iterates to use for acceleration.
+        max_iter: Maximum number of iterations.
+        tol: Convergence tolerance based on residual norm.
+        beta: Mixing parameter for the fixed-point step.
+        lambda_reg: Regularization parameter for least-squares solve.
+        return_residual_norms: If True, also return list of residual norms.
+
+    Returns:
+        Approximate solution x, and optionally list of residual norms.
+    """
+    # history buffers
+    delta_xs: list[torch.Tensor] = []
+    delta_gs: list[torch.Tensor] = []
+    residual_norms = []
+
+    # run fixed-pointer iteration
+    x = x0
+    fx = f(x)
+    g = fx - x
+    x_prev, g_prev = None, None
+    for k in range(max_iter):
+        # evaluate residual and compute convergence
+        res_norm = torch.norm(g).item()
+        residual_norms.append(res_norm)
+        if res_norm < tol:
+            break
+
+        # update history
+        if k > 0:
+            assert x_prev is not None and g_prev is not None
+            delta_xs.append(x - x_prev)
+            delta_gs.append(g - g_prev)
+
+            # truncate history to hold at most m elements
+            if len(delta_xs) > m:
+                delta_xs.pop(0)
+                delta_gs.pop(0)
+        x_prev, g_prev = x, g
+
+        # compute Anderson acceleration step
+        if len(delta_xs) > 0:
+            # create matrices from history of shape (features, history_length)
+            X = torch.stack(delta_xs, dim=1)  # (n, k)
+            G = torch.stack(delta_gs, dim=1)  # (n, k)
+
+            # solve regularized least-squares problem
+            A = G.T @ G + lambda_reg * torch.eye(G.shape[1], device=G.device)
+            b = G.T @ g
+            try:
+                coeffs = torch.linalg.solve(A, b)
+                # update iterate with momentum + Anderson step
+                x = x + beta * g - (X + beta * G) @ coeffs
+            except RuntimeError:
+                x = x + beta * g  # fallback to fixed-point step if matrix is singular
+        else:
+            x = x + beta * g  # fixed-point step if there is no history
+
+        # update iterate and residual
+        fx = f(x)
+        g = fx - x
+
+    if return_residual_norms:
+        return x, residual_norms
+    else:
+        return x