From 2f22a4da3281e6c78acd1d592497821afcf3d92e Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:40:05 +0200 Subject: [PATCH 1/6] Support MTK ('flexible') barostats in get_npt_stepper `_qbaro` and `_pbaro` were transcriptions of `BaroBZP.qcstep`/`pstep` with the scalar algebra baked in, so they only worked for the isotropic (BZP) barostat. An MTK barostat, whose `p` is a 3x3 upper-triangular matrix, passed the integrator type check (i-PI restricts the `nst` ensemble to the RGB barostat, so `flexible` runs with `NPTIntegrator`) and then propagated silently wrong dynamics through broadcasting. Both helpers now take a `mode` argument using i-PI's barostat mode names, with the flexible branch following `BaroMTK` -- matrix exponentials for the cell scaling, the full `np.triu` stress for the cell momentum, and the `hmask`/`vol_constraint` projection with its `ethermo` bookkeeping. As in the isotropic branch, the terms that the plain velocity Verlet splitting needs but `OAbBbBABbAbPO` does not (the `sinh` drift term and the `pc.fc` force corrections) are left out. The mode is derived from `type(motion.barostat)`, so it cannot disagree with the input file, and unsupported barostats (SC-BZP, RGB, dummy) now raise instead of being silently mishandled. Co-Authored-By: Claude Opus 5 (1M context) --- src/flashmd/ipi.py | 73 +++++++++++++++++++++++++++++++++------------- 1 file changed, 53 insertions(+), 20 deletions(-) diff --git a/src/flashmd/ipi.py b/src/flashmd/ipi.py index 1cbdbaf..81eee78 100644 --- a/src/flashmd/ipi.py +++ b/src/flashmd/ipi.py @@ -2,8 +2,10 @@ import ase.units import numpy as np import torch +from ipi.engine.barostats import BaroBZP, BaroMTK from ipi.engine.motion.dynamics import NPTIntegrator, NVEIntegrator, NVTIntegrator from ipi.utils.depend import dstrip +from ipi.utils.mathtools import matrix_exp from ipi.utils.mathtools import random_rotation as random_rotation_matrix from ipi.utils.messages import info, verbosity from ipi.utils.units import Constants @@ -206,35 +208,56 @@ def nvt_stepper(motion, *_, **__): return nvt_stepper -def _qbaro(baro): +def _qbaro(baro, mode): """Propagation step for the cell volume (adjusting atomic positions and momenta).""" - v = baro.p[0] / baro.m[0] halfdt = ( baro.qdt ) # this is set to half the inner loop in all integrators that use a barostat - expq, expp = (np.exp(v * halfdt), np.exp(-v * halfdt)) - baro.nm.qnm[0, :] *= expq - baro.nm.pnm[0, :] *= expp - baro.cell.h *= expq + if mode == "isotropic": + v = baro.p[0] / baro.m[0] + expq, expp = (np.exp(v * halfdt), np.exp(-v * halfdt)) + baro.nm.qnm[0, :] *= expq + baro.nm.pnm[0, :] *= expp + baro.cell.h *= expq + else: + v = baro.p / baro.m[0] + expq, expp = (matrix_exp(v * halfdt), matrix_exp(-v * halfdt)) + + baro.nm.qnm[0] = (dstrip(baro.nm.qnm)[0].reshape(-1, 3) @ expq.T).reshape(-1) + baro.nm.pnm[0] = (dstrip(baro.nm.pnm)[0].reshape(-1, 3) @ expp.T).reshape(-1) + baro.cell.h = expq @ dstrip(baro.cell.h) -def _pbaro(baro): + +def _pbaro(baro, mode): """Propagation step for the cell momentum (adjusting atomic positions and momenta).""" # we are assuming then that p the coupling between p^2 and dp/dt only involves the fast force dt = baro.pdt[0] - - # computes the pressure associated with the forces at the outer level MTS level. - press = np.trace(baro.stress_mts(0)) / 3.0 - # integerates the kinetic part of the pressure with the force at the inner-most level. nbeads = baro.beads.nbeads - baro.p += ( - 3.0 - * dt - * (baro.cell.V * (press - nbeads * baro.pext) + Constants.kb * baro.temp) - ) + + if mode == "isotropic": + # computes the pressure associated with the forces at the outer level MTS level. + press = np.trace(baro.stress_mts(0)) / 3.0 + # integerates the kinetic part of the pressure with the force at the inner-most level. + baro.p += ( + 3.0 + * dt + * (baro.cell.V * (press - nbeads * baro.pext) + Constants.kb * baro.temp) + ) + else: + stress = np.triu(dstrip(baro.stress_mts(0)) - nbeads * np.eye(3) * baro.pext) + baro.p += dt * (baro.cell.V * stress + Constants.kb * baro.temp * baro.L) + + # zero out the fixed cell components, tracking the change in the conserved + # quantity through the barostat thermostat (as in i-PI's BaroMTK.pstep) + baro.thermostat.ethermo += baro.kin + baro.p *= baro.hmask + if baro.vol_constraint: + baro.p -= np.eye(3) * np.trace(baro.p) / 3.0 + baro.thermostat.ethermo -= baro.kin def get_npt_stepper( @@ -251,6 +274,16 @@ def get_npt_stepper( f"Base i-PI integrator is of type {motion.integrator.__class__.__name__}, use a NPT setup." ) + if type(motion.barostat) is BaroBZP: + baro_mode = "isotropic" + elif type(motion.barostat) is BaroMTK: + baro_mode = "flexible" + else: + raise TypeError( + f"Barostat is of type {motion.barostat.__class__.__name__}, use an " + "'isotropic' (BZP) or 'flexible' (MTK) barostat." + ) + if use_standard_vv: # use the standard velocity Verlet integrator vv_step = get_standard_vv_step( @@ -273,15 +306,15 @@ def npt_stepper(motion, *_, **__): info("@flashmd: Barostat thermo", verbosity.debug) motion.barostat.thermostat.step() info("@flashmd: Barostat q", verbosity.debug) - _qbaro(motion.barostat) + _qbaro(motion.barostat, baro_mode) info("@flashmd: Barostat p", verbosity.debug) - _pbaro(motion.barostat) + _pbaro(motion.barostat, baro_mode) info("@flashmd: FlashVV", verbosity.debug) vv_step(motion) info("@flashmd: Barostat p", verbosity.debug) - _pbaro(motion.barostat) + _pbaro(motion.barostat, baro_mode) info("@flashmd: Barostat q", verbosity.debug) - _qbaro(motion.barostat) + _qbaro(motion.barostat, baro_mode) info("@flashmd: Barostat thermo", verbosity.debug) motion.barostat.thermostat.step() info("@flashmd: Particle thermo", verbosity.debug) From 8080b8b16636fe0fd4cfc7d761fc6beb3a0f791a Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:14:42 +0200 Subject: [PATCH 2/6] Add regression tests for vol_constraint and hmask failures --- pyproject.toml | 3 ++ tests/test_ipi.py | 111 ++++++++++++++++++++++++++++++++++++++++++++++ tox.ini | 1 + 3 files changed, 115 insertions(+) create mode 100644 tests/test_ipi.py diff --git a/pyproject.toml b/pyproject.toml index 379b60c..8ec56e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,7 @@ filterwarnings = [ "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:`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", @@ -98,4 +99,6 @@ filterwarnings = [ "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 ] diff --git a/tests/test_ipi.py b/tests/test_ipi.py new file mode 100644 index 0000000..b57bab5 --- /dev/null +++ b/tests/test_ipi.py @@ -0,0 +1,111 @@ +import ase.build +import ase.io +import numpy as np +import pytest +import torch +from ipi.scripting import InteractiveSimulation +from ipi.utils.depend import dstrip + +from flashmd import get_pretrained +from flashmd.ipi import get_npt_stepper + + +# same models as the other tests, so nothing extra is downloaded +TIME_STEP = 64 + +# input file for i-PI to run a simulation in the isotheral-isobaric ensemble with +# flexible cell vectors. `run_npt` below replaces the {variables} with concrete values +# for each run. `barostat_extra` is used to test different barostat options. +INPUT = """ + 1 + + 32123 + + metatomic + {{model: {model}, template: ./structure.xyz, device: {device}}} + + + + + ./structure.xyz + 300 + + + 300 + 0 + + + + {time_step} + 100 + + 2000 + 100 + {barostat_extra} + + + + + +""" + + +@pytest.fixture(scope="module") +def models(tmp_path_factory): + """The MLIP (as a file, for i-PI's metatomic driver) and the FlashMD model.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + mlip, flashmd_model = get_pretrained("pet-omatpes-v2", TIME_STEP) + mlip_path = tmp_path_factory.mktemp("models") / "mlip.pt" + mlip.save(str(mlip_path)) + return mlip_path, flashmd_model, device + + +@pytest.fixture(autouse=True) +def in_tmp_path(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + + +def run_npt(models, barostat_extra, n_steps=10): + """Run the NPT stepper and return the initial and the final cell.""" + mlip_path, flashmd_model, device = models + + # define a simple test system: + # fcc Al in a diagonal cell, so `h[0, 1]` starts at zero up to round-off + ase.io.write("structure.xyz", ase.build.bulk("Al", "fcc", cubic=True)) + + # create a simulation with a FlashMD NPT stepper + simulation = InteractiveSimulation( + # replace {variables} in INPUT with concrete values for this run + INPUT.format( + model=mlip_path, + device=device, + time_step=TIME_STEP, + barostat_extra=barostat_extra, + ) + ) + motion = simulation.syslist[0].motion + step = get_npt_stepper(simulation, flashmd_model, device) + + # run a few steps of NPT and track the cell + cell = initial_cell = dstrip(motion.cell.h).copy() # type: ignore + for _ in range(n_steps): + step(motion) + cell = dstrip(motion.cell.h).copy() # type: ignore + + # an unstable piston explodes the cell, and the neighbor list with it, so stop + # here instead of leaving the model to churn on a nonsensical structure + ratio = np.linalg.det(cell) / np.linalg.det(initial_cell) + assert 0.5 < ratio < 2.0, f"the barostat went unstable: V / V0 = {ratio}" + return initial_cell, cell + + +def test_hfix_freezes_cell_component(models): + """`hfix` keeps a cell component that starts at zero at zero.""" + initial_cell, cell = run_npt(models, " [ xy ] ") + assert cell[0, 1] == pytest.approx(initial_cell[0, 1], abs=1e-10) + + +def test_vol_constraint_conserves_volume(models): + """`vol_constraint` keeps the cell volume constant.""" + initial_cell, cell = run_npt(models, " True ") + assert np.linalg.det(cell) == pytest.approx(np.linalg.det(initial_cell), rel=1e-10) diff --git a/tox.ini b/tox.ini index 7883750..2970287 100644 --- a/tox.ini +++ b/tox.ini @@ -45,6 +45,7 @@ description = Run package tests with pytest passenv = * deps = pytest + ipi changedir = tests commands = pytest \ From 0f05d6e41b2369875df1b2d0b49bc6b8f03a1a31 Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:16:42 +0200 Subject: [PATCH 3/6] Remove trailing space --- tests/test_ipi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_ipi.py b/tests/test_ipi.py index b57bab5..8907818 100644 --- a/tests/test_ipi.py +++ b/tests/test_ipi.py @@ -13,7 +13,7 @@ # same models as the other tests, so nothing extra is downloaded TIME_STEP = 64 -# input file for i-PI to run a simulation in the isotheral-isobaric ensemble with +# input file for i-PI to run a simulation in the isotheral-isobaric ensemble with # flexible cell vectors. `run_npt` below replaces the {variables} with concrete values # for each run. `barostat_extra` is used to test different barostat options. INPUT = """ From cb79efcafed4c8dc6b209560f27f0ed5a144a0a1 Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:43:36 +0200 Subject: [PATCH 4/6] Ignore the Python 3.14 wording of the torch.jit.save deprecation --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 8ec56e5..9505e1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ filterwarnings = [ "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", From 4d540df1ea1551fbc58de266be63bacc5b85aed7 Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:38:56 +0200 Subject: [PATCH 5/6] Pull projection out --- src/flashmd/ipi.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/src/flashmd/ipi.py b/src/flashmd/ipi.py index 81eee78..7c655be 100644 --- a/src/flashmd/ipi.py +++ b/src/flashmd/ipi.py @@ -208,6 +208,23 @@ def nvt_stepper(motion, *_, **__): return nvt_stepper +def _project_baro(baro): + """Zero the cell momentum components held fixed by `hfix`/`vol_constraint`. + + The removed kinetic energy is credited to the barostat thermostat's `ethermo`, as + in i-PI's `BaroMTK.pstep`, so that the conserved quantity stays correct. + """ + + if not baro.vol_constraint and np.all(baro.hmask == 1.0): + return + + baro.thermostat.ethermo += baro.kin + baro.p *= baro.hmask + if baro.vol_constraint: + baro.p -= np.eye(3) * np.trace(baro.p) / 3.0 + baro.thermostat.ethermo -= baro.kin + + def _qbaro(baro, mode): """Propagation step for the cell volume (adjusting atomic positions and momenta).""" @@ -223,6 +240,10 @@ def _qbaro(baro, mode): baro.nm.pnm[0, :] *= expp baro.cell.h *= expq else: + # the barostat thermostat kicks all six components, so `p` cannot be assumed + # to be projected here + _project_baro(baro) + v = baro.p / baro.m[0] expq, expp = (matrix_exp(v * halfdt), matrix_exp(-v * halfdt)) @@ -251,13 +272,7 @@ def _pbaro(baro, mode): stress = np.triu(dstrip(baro.stress_mts(0)) - nbeads * np.eye(3) * baro.pext) baro.p += dt * (baro.cell.V * stress + Constants.kb * baro.temp * baro.L) - # zero out the fixed cell components, tracking the change in the conserved - # quantity through the barostat thermostat (as in i-PI's BaroMTK.pstep) - baro.thermostat.ethermo += baro.kin - baro.p *= baro.hmask - if baro.vol_constraint: - baro.p -= np.eye(3) * np.trace(baro.p) / 3.0 - baro.thermostat.ethermo -= baro.kin + _project_baro(baro) def get_npt_stepper( From 661811fb68fa89b711803b43826e9c1ff5613a05 Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:34:49 +0200 Subject: [PATCH 6/6] Raise on invalid mode in step functions --- src/flashmd/ipi.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/flashmd/ipi.py b/src/flashmd/ipi.py index 7c655be..f429c7e 100644 --- a/src/flashmd/ipi.py +++ b/src/flashmd/ipi.py @@ -239,7 +239,7 @@ def _qbaro(baro, mode): baro.nm.qnm[0, :] *= expq baro.nm.pnm[0, :] *= expp baro.cell.h *= expq - else: + elif mode == "flexible": # the barostat thermostat kicks all six components, so `p` cannot be assumed # to be projected here _project_baro(baro) @@ -250,6 +250,8 @@ def _qbaro(baro, mode): baro.nm.qnm[0] = (dstrip(baro.nm.qnm)[0].reshape(-1, 3) @ expq.T).reshape(-1) baro.nm.pnm[0] = (dstrip(baro.nm.pnm)[0].reshape(-1, 3) @ expp.T).reshape(-1) baro.cell.h = expq @ dstrip(baro.cell.h) + else: + raise TypeError(f"Unknown barostat mode '{mode}'.") def _pbaro(baro, mode): @@ -268,11 +270,13 @@ def _pbaro(baro, mode): * dt * (baro.cell.V * (press - nbeads * baro.pext) + Constants.kb * baro.temp) ) - else: + elif mode == "flexible": stress = np.triu(dstrip(baro.stress_mts(0)) - nbeads * np.eye(3) * baro.pext) baro.p += dt * (baro.cell.V * stress + Constants.kb * baro.temp * baro.L) _project_baro(baro) + else: + raise TypeError(f"Unknown barostat mode '{mode}'.") def get_npt_stepper(