From c56612965cfe1e6c63f118dae033977d0973adcc Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 13 Feb 2026 16:14:32 +0000 Subject: [PATCH 01/15] vibecoded (need to check) --- pyproject.toml | 6 +- tests/test_physical_validation.py | 345 ++++++++++++++++++++++++++++++ 2 files changed, 350 insertions(+), 1 deletion(-) create mode 100644 tests/test_physical_validation.py diff --git a/pyproject.toml b/pyproject.toml index 8093d764a..09c5e5c06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ test = [ "ase>=3.26", "phonopy>=2.37.0", + "physical-validation>=1.0.5", "platformdirs>=4.0.0", "psutil>=7.0.0", "pymatgen>=2025.6.14", @@ -132,8 +133,11 @@ check-filenames = true ignore-words-list = ["convertor"] [tool.pytest.ini_options] -addopts = ["-p no:warnings"] +addopts = ["-p no:warnings", "-m not physical_validation"] testpaths = ["tests"] +markers = [ + "physical_validation: long-running physical validation tests (run with: pytest -m physical_validation)", +] [tool.uv] # make these dependencies mutually exclusive since they use incompatible e3nn versions diff --git a/tests/test_physical_validation.py b/tests/test_physical_validation.py new file mode 100644 index 000000000..e57f85c8b --- /dev/null +++ b/tests/test_physical_validation.py @@ -0,0 +1,345 @@ +"""Physical validation tests for torch-sim MD integrators. + +Uses the physical_validation library (https://github.com/shirtsgroup/physical_validation) +to verify that integrators produce physically correct results. These tests are +long-running (~5 min total) and excluded by default. Run with: + + pytest -m physical_validation -v +""" + +import numpy as np +import pytest +import torch +from ase.build import bulk + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.units import MetalUnits + +physical_validation = pytest.importorskip("physical_validation") + +DEVICE = torch.device("cpu") +DTYPE = torch.float64 + +# LJ Argon parameters +SIGMA = 3.405 +EPSILON = 0.0104 +CUTOFF = 2.5 * SIGMA + + +def _make_unit_data(): + """Create UnitData for torch-sim's MetalUnits system.""" + return physical_validation.data.UnitData( + kb=float(MetalUnits.temperature), # k_B in eV/K = 8.617e-5 + energy_str="eV", + energy_conversion=1.0, + length_str="Ang", + length_conversion=1.0, + volume_str="Ang^3", + volume_conversion=1.0, + temperature_str="K", + temperature_conversion=1.0, + pressure_str="eV/Ang^3", + pressure_conversion=1.0, + time_str="internal", + time_conversion=1.0, + ) + + +def _make_lj_model(): + """Create a Lennard-Jones model for Argon.""" + return LennardJonesModel( + use_neighbor_list=False, + sigma=SIGMA, + epsilon=EPSILON, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=False, + cutoff=CUTOFF, + ) + + +def _make_ar_supercell(repeat=(2, 2, 2)): + """Create an FCC Argon supercell SimState.""" + atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) + return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) + + +def _run_nvt_langevin( + sim_state, + model, + temperature, + timestep_ps, + n_steps, + n_equilibration, + seed=42, +): + """Run NVT Langevin simulation and collect per-step observables.""" + kT = temperature * float(MetalUnits.temperature) + dt_internal = timestep_ps * float(MetalUnits.time) + natoms = int(sim_state.positions.shape[0]) + + state = ts.nvt_langevin_init(sim_state, model, kT=kT, seed=seed) + + # Equilibration + for _ in range(n_equilibration): + state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) + + # Production - collect observables + ke_list = [] + pe_list = [] + total_e_list = [] + position_list = [] + velocity_list = [] + + for _ in range(n_steps): + state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) + + ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + pe = state.energy.sum() + ke_list.append(float(ke)) + pe_list.append(float(pe)) + total_e_list.append(float(ke + pe)) + position_list.append(state.positions.detach().cpu().numpy().copy()) + velocity_list.append(state.velocities.detach().cpu().numpy().copy()) + + # Compute volume from cell + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "positions": np.array(position_list), + "velocities": np.array(velocity_list), + "volume": volume, + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": dt_internal, + "natoms": natoms, + } + + +def _run_nve(sim_state, model, kT_init, timestep_ps, n_steps, seed=42): + """Run NVE simulation and collect constant of motion.""" + dt_internal = timestep_ps * float(MetalUnits.time) + + state = ts.nve_init(sim_state, model, kT=kT_init, seed=seed) + + com_list = [] + for _ in range(n_steps): + state = ts.nve_step(state, model, dt=dt_internal) + ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + pe = state.energy.sum() + com_list.append(float(ke + pe)) + + return { + "constant_of_motion": np.array(com_list), + "dt_internal": dt_internal, + } + + +def _build_nvt_simulation_data(run_data, temperature): + """Build a physical_validation SimulationData from NVT run results.""" + units = _make_unit_data() + + system = physical_validation.data.SystemData( + natoms=run_data["natoms"], + nconstraints=0, + ndof_reduction_tra=3, + ndof_reduction_rot=0, + mass=run_data["masses"], + ) + + ensemble_data = physical_validation.data.EnsembleData( + ensemble="NVT", + natoms=run_data["natoms"], + volume=run_data["volume"], + temperature=temperature, + ) + + observables = physical_validation.data.ObservableData( + kinetic_energy=run_data["kinetic_energy"], + potential_energy=run_data["potential_energy"], + total_energy=run_data["total_energy"], + ) + + trajectory = physical_validation.data.TrajectoryData( + position=run_data["positions"], + velocity=run_data["velocities"], + ) + + return physical_validation.data.SimulationData( + units=units, + dt=run_data["dt_internal"], + system=system, + ensemble=ensemble_data, + observables=observables, + trajectory=trajectory, + ) + + +@pytest.mark.physical_validation +def test_ke_distribution(): + """Test that kinetic energy follows the Maxwell-Boltzmann distribution. + + Runs NVT Langevin at 100K on a 2x2x2 Ar supercell (32 atoms) and checks + that the KE distribution matches the analytical Maxwell-Boltzmann prediction. + """ + sim_state = _make_ar_supercell(repeat=(2, 2, 2)) + model = _make_lj_model() + temperature = 100.0 # K + + run_data = _run_nvt_langevin( + sim_state, + model, + temperature=temperature, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=42, + ) + + data = _build_nvt_simulation_data(run_data, temperature) + + result = physical_validation.kinetic_energy.distribution( + data, + strict=False, + verbosity=0, + ) + # strict=False returns (d_mean, d_width) in sigma units + d_mean, d_width = result + + assert abs(d_mean) < 3, ( + f"KE mean deviation {d_mean:.2f} sigma exceeds threshold" + ) + assert abs(d_width) < 3, ( + f"KE width deviation {d_width:.2f} sigma exceeds threshold" + ) + + +@pytest.mark.physical_validation +def test_integrator_convergence(): + """Test that NVE energy error scales as dt^2 (velocity Verlet). + + Runs NVE at 3 different timesteps from identical initial conditions on a + 4-atom Ar unit cell at low temperature (5K). Low temperature minimizes + thermal fluctuations so the integration error dominates the RMSD of the + conserved quantity, allowing the dt^2 convergence to be observed. + """ + sim_state = _make_ar_supercell(repeat=(1, 1, 1)) # 4 atoms + model = _make_lj_model() + temperature = 5.0 # K, low T so integration error dominates + kT_init = temperature * float(MetalUnits.temperature) + + # Timesteps chosen so integration error >> thermal fluctuations at all dt. + # Factor of ~sqrt(2) spacing gives dt^2 ratio of ~2.0 per step. + timesteps = [0.008, 0.00566, 0.004] # ps + n_steps = 5_000 + seed = 42 + + natoms = int(sim_state.positions.shape[0]) + masses = sim_state.masses.detach().cpu().numpy() + volume = float( + np.abs(np.linalg.det(sim_state.cell[0].detach().cpu().numpy())) + ) + units = _make_unit_data() + + simulations = [] + for dt_ps in timesteps: + run_data = _run_nve( + sim_state, + model, + kT_init=kT_init, + timestep_ps=dt_ps, + n_steps=n_steps, + seed=seed, + ) + + system = physical_validation.data.SystemData( + natoms=natoms, + nconstraints=0, + ndof_reduction_tra=3, + ndof_reduction_rot=0, + mass=masses, + ) + + ensemble_data = physical_validation.data.EnsembleData( + ensemble="NVE", + natoms=natoms, + volume=volume, + ) + + observables = physical_validation.data.ObservableData( + constant_of_motion=run_data["constant_of_motion"], + ) + + sim_data = physical_validation.data.SimulationData( + units=units, + dt=run_data["dt_internal"], + system=system, + ensemble=ensemble_data, + observables=observables, + ) + simulations.append(sim_data) + + result = physical_validation.integrator.convergence( + simulations, + verbose=False, + ) + + assert result < 0.5, ( + f"Integrator convergence deviation {result:.3f} exceeds threshold 0.5" + ) + + +@pytest.mark.physical_validation +def test_ensemble_check(): + """Test NVT ensemble validity via Boltzmann weight ratio at two temperatures. + + Runs NVT Langevin at 80K and 100K on a 2x2x2 Ar supercell (32 atoms), + then checks that the total energy distributions satisfy the expected + Boltzmann weight relationship. Uses total_energy=True for the ensemble + check which includes both kinetic and potential energy contributions. + """ + sim_state = _make_ar_supercell(repeat=(2, 2, 2)) + model = _make_lj_model() + + temp_low = 80.0 + temp_high = 100.0 + + run_low = _run_nvt_langevin( + sim_state, + model, + temperature=temp_low, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=42, + ) + run_high = _run_nvt_langevin( + sim_state, + model, + temperature=temp_high, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=123, + ) + + data_low = _build_nvt_simulation_data(run_low, temp_low) + data_high = _build_nvt_simulation_data(run_high, temp_high) + + quantiles = physical_validation.ensemble.check( + data_low, + data_high, + total_energy=True, + data_is_uncorrelated=True, + verbosity=0, + ) + + for i, q in enumerate(quantiles): + assert abs(q) < 3, ( + f"Ensemble quantile {i} = {q:.2f} sigma exceeds threshold" + ) From 4e6dccc02fb947824d34c54b73ebfcf9739428e7 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Tue, 24 Mar 2026 17:20:28 +0000 Subject: [PATCH 02/15] code for physical_validation --- fast_integrator_tests/README.md | 70 +++ fast_integrator_tests/analyze.py | 871 +++++++++++++++++++++++++++++++ fast_integrator_tests/common.py | 48 ++ fast_integrator_tests/run_npt.py | 174 ++++++ fast_integrator_tests/run_nve.py | 91 ++++ fast_integrator_tests/run_nvt.py | 116 ++++ 6 files changed, 1370 insertions(+) create mode 100644 fast_integrator_tests/README.md create mode 100644 fast_integrator_tests/analyze.py create mode 100644 fast_integrator_tests/common.py create mode 100644 fast_integrator_tests/run_npt.py create mode 100644 fast_integrator_tests/run_nve.py create mode 100644 fast_integrator_tests/run_nvt.py diff --git a/fast_integrator_tests/README.md b/fast_integrator_tests/README.md new file mode 100644 index 000000000..a60ee9446 --- /dev/null +++ b/fast_integrator_tests/README.md @@ -0,0 +1,70 @@ +# Integrator Physical Validation Tests + +Physical validation of all torch-sim MD integrators using the +[physical_validation](https://physical-validation.readthedocs.io/) library +and LJ Argon as a test system. + +## What's here + +| File | Purpose | +|---|---| +| `common.py` | Shared constants, model/structure builders | +| `run_nvt.py` | Run NVT simulations (Langevin, Nose-Hoover, V-Rescale) at 80K and 100K | +| `run_npt.py` | Run NPT simulations (Langevin, Nose-Hoover, iso/aniso C-Rescale) at 80K and 100K | +| `run_nve.py` | Run NVE simulations at 8 timesteps for convergence analysis | +| `analyze.ipynb` | Jupyter notebook with all diagnostic plots and `physical_validation` checks | +| `data/` | Output directory for `.npz` trajectory data (gitignored) | + +## Quick start + +```bash +# 1. Generate trajectory data (from this directory) +python run_nve.py +python run_nvt.py +python run_npt.py + +# Run a single integrator if needed +python run_nvt.py --integrator nvt_langevin +python run_npt.py --integrator npt_nose_hoover + +# NPT runs both a pressure and temperature sweep, so you can specify which to do: +python run_npt.py --mode temperature # Vary T at P=0 +python run_npt.py --mode pressure # Vary P at fixed T + +# 2. Open the notebook +jupyter notebook analyze.ipynb +``` + +## What the notebook shows + +### Custom plots +- **Time series**: Temperature, total energy, volume (NPT) vs step +- **KE distribution**: Observed histogram vs theoretical Gamma(Nf/2, k_BT) distribution +- **Ensemble check**: Overlapping energy distributions at two temperatures with log-ratio inset +- **NVE convergence**: Log-log RMSD vs dt plot, energy drift traces, convergence ratio table + +### physical_validation native plots +The notebook also calls `physical_validation` functions with `screen=True` to produce +their built-in diagnostic figures: +- `kinetic_energy.distribution()` — KS test or mean/width comparison plot +- `ensemble.check()` — Forward/reverse work distributions and linear fit +- `integrator.convergence()` — RMSD vs dt with reference line + +### Summary table +Final cell runs all checks programmatically and prints a PASS/FAIL table for every integrator. + +## Validation checks + +| Check | Integrators | What it tests | +|---|---|---| +| KE distribution | All NVT + NPT | Kinetic energy follows Maxwell-Boltzmann (gamma) distribution | +| Ensemble check | All NVT + NPT | Energy distributions at T=80K and T=100K satisfy Boltzmann weight ratio | +| NVE convergence | NVE (velocity Verlet) | Energy drift RMSD scales as dt^2 | + +## System details + +- **Structure**: FCC Argon (a=5.26 A), 2x2x2 supercell (32 atoms) for NVT/NPT, unit cell (4 atoms) for NVE +- **Model**: Lennard-Jones (sigma=3.405, epsilon=0.0104 eV, cutoff=8.5125 A), no neighbor list +- **Timestep**: 0.004 ps for NVT/NPT; sweep 0.002-0.010 ps for NVE +- **Production**: 10,000 steps after 2,000 (NVT) or 3,000 (NPT) equilibration steps +- **Threshold**: |deviation| < 3 sigma for all statistical checks diff --git a/fast_integrator_tests/analyze.py b/fast_integrator_tests/analyze.py new file mode 100644 index 000000000..444948d65 --- /dev/null +++ b/fast_integrator_tests/analyze.py @@ -0,0 +1,871 @@ +# %% [markdown] +# # Physical Validation of torch-sim Integrators +# +# This notebook analyzes MD trajectory data produced by `run_nvt.py`, `run_npt.py`, and `run_nve.py`. +# +# **Tests performed:** +# 1. **KE Distribution** — Does kinetic energy follow the Maxwell-Boltzmann (gamma) distribution? +# 2. **Ensemble Check** — Do energy distributions at two temperatures satisfy the Boltzmann weight relationship? +# 3. **Pressure Ensemble Check** — Do volume distributions at different pressures satisfy the expected Boltzmann weight relationship? +# 4. **NVE Convergence** — Does the energy drift RMSD scale as dt² (velocity Verlet)? +# 5. **Time Series** — Visual inspection of temperature, energy, and volume stability. + +# %% +# Create a folder to save plot +import os +os.makedirs("plots", exist_ok=True) + +# %% +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +from scipy import stats + +plt.rcParams.update({ + "figure.dpi": 120, + "axes.grid": True, + "grid.alpha": 0.3, + "figure.facecolor": "white", +}) + +DATA_DIR = Path("fast_integrator_tests/data") + +def load(name): + """Load an npz file from the data directory.""" + return dict(np.load(DATA_DIR / f"{name}.npz", allow_pickle=True)) + +# Check what data is available +available = sorted(p.stem for p in DATA_DIR.glob("*.npz")) +print("Available datasets:") +for a in available: + print(f" {a}") + +# %% +import physical_validation +from torch_sim.units import MetalUnits + +k_B_eV = float(MetalUnits.temperature) # 8.617333e-5 eV/K +THRESHOLD = 3.0 + +# def make_unit_data(): +# return physical_validation.data.UnitData( +# kb=k_B_eV, +# energy_str="eV", energy_conversion=1.0, +# length_str="Ang", length_conversion=1.0, +# volume_str="Ang^3", volume_conversion=1.0, +# temperature_str="K", temperature_conversion=1.0, +# pressure_str="eV/Ang^3", pressure_conversion=1.0, +# time_str="internal", time_conversion=1.0, +# ) +def make_unit_data(): + return physical_validation.data.UnitData( + kb=k_B_eV, + energy_str="eV", energy_conversion=96.485, # Convert to kJ/mol + length_str="Ang", length_conversion=1e-1, # Convert to nm + volume_str="Ang^3", volume_conversion=1e-3, # Convert to nm^3 + temperature_str="K", temperature_conversion=1.0, + # pressure_str="eV/Ang^3", pressure_conversion=1.6e6, # Convert to bar + pressure_str="bar", pressure_conversion=1, + time_str="fs", time_conversion=1e-3, # Convert to ps + ) + +def build_sim_data(d, temperature, ensemble="NVT", pressure=None): + """Build physical_validation.SimulationData from a loaded npz dict.""" + units = make_unit_data() + natoms = int(d["natoms"]) + system = physical_validation.data.SystemData( + natoms=natoms, nconstraints=0, + ndof_reduction_tra=3, ndof_reduction_rot=0, + mass=d["masses"], + ) + ens_kw = dict(natoms=natoms, temperature=temperature) + if ensemble == "NVT": + ens_kw["ensemble"] = "NVT" + ens_kw["volume"] = float(d.get("volume", np.mean(d.get("volumes", [0])))) + else: + ens_kw["ensemble"] = "NPT" + ens_kw["pressure"] = pressure if pressure is not None else 0.0 + + obs_kw = dict( + kinetic_energy=d["kinetic_energy"], + potential_energy=d["potential_energy"], + total_energy=d["total_energy"], + ) + if "volumes" in d: + obs_kw["volume"] = d["volumes"] + + return physical_validation.data.SimulationData( + units=units, + dt=float(d["dt_internal"]), + system=system, + ensemble=physical_validation.data.EnsembleData(**ens_kw), + observables=physical_validation.data.ObservableData(**obs_kw), + ) + +print("physical_validation helpers loaded") + +# %% [markdown] +# ## 1. NVT Time Series +# +# Temperature and energy vs step for each NVT integrator at both temperatures. + +# %% +NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] +TEMPS = [88.0, 100.0] + +fig, axes = plt.subplots(len(NVT_INTEGRATORS), 2, figsize=(14, 3.5 * len(NVT_INTEGRATORS)), + squeeze=False, sharex=True) +fig.suptitle("NVT Integrators — Time Series", fontsize=14, y=1.01) + +for row, name in enumerate(NVT_INTEGRATORS): + for temp in TEMPS: + label = f"{name}_T{temp:.0f}K" + try: + d = load(label) + except FileNotFoundError: + continue + + steps = np.arange(len(d["temperature"])) + target_T = float(d["target_temperature"]) + + # Temperature + ax = axes[row, 0] + ax.plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") + ax.axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) + ax.set_ylabel("Temperature (K)") + ax.set_title(name) + ax.legend(fontsize=8) + + # Total energy + ax = axes[row, 1] + ax.plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") + ax.set_ylabel("Total Energy (eV)") + ax.set_title(name) + ax.legend(fontsize=8) + +axes[-1, 0].set_xlabel("Step") +axes[-1, 1].set_xlabel("Step") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ## 2. KE Distribution — NVT Integrators +# +# For each integrator at 100K, compare the observed KE distribution against the theoretical gamma distribution. +# The KE of an ideal gas with $N_f$ degrees of freedom at temperature $T$ follows $\text{Gamma}(N_f/2, k_BT)$. + +# %% +fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NVT, T=100K)", fontsize=13) + +for ax, name in zip(axes, NVT_INTEGRATORS): + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + ke = d["kinetic_energy"] + natoms = int(d["natoms"]) + target_T = float(d["target_temperature"]) + + # Degrees of freedom: 3*N - 3 (COM removed) + if name == "nvt_langevin": + ndof = 3 * natoms + else: + ndof = 3 * natoms - 3 + # Theoretical gamma: shape=ndof/2, scale=k_B*T + shape = ndof / 2 + scale = k_B_eV * target_T + + # Histogram of observed KE + ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") + + # Theoretical curve + x = np.linspace(ke.min(), ke.max(), 300) + pdf = stats.gamma.pdf(x, a=shape, scale=scale) + ax.plot(x, pdf, "r-", lw=2, label="Theory") + + # Stats + d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) + d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) + ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=10) + ax.set_xlabel("Kinetic Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in KE distribution plots (NVT) +# +# Uses `physical_validation.kinetic_energy.distribution(..., screen=True)` which overlays the observed and theoretical distributions with a KS-test or mean/width comparison. + +# %% +for name in NVT_INTEGRATORS: + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — KE distribution (physical_validation)") + print(f"{'='*60}") + sd = build_sim_data(d, 100.0, ensemble="NVT") + result = physical_validation.kinetic_energy.distribution( + sd, strict=False, screen=True, verbosity=2, filename=f"plots/{name}_ke_distribution.png" + ) + print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") + +plt.show() + +# %% [markdown] +# ## 3. Ensemble Check — NVT Integrators +# +# For each integrator, compare total energy distributions at T_low=88K and T_high=100K. +# The log ratio of the energy histograms should be linear with slope $\Delta\beta = 1/k_BT_\text{low} - 1/k_BT_\text{high}$. + +# %% +fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("NVT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) + +temp_low, temp_high = 88.0, 100.0 +delta_beta = 1 / (k_B_eV * temp_low) - 1 / (k_B_eV * temp_high) + +for ax, name in zip(axes, NVT_INTEGRATORS): + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + e_lo = d_lo["total_energy"] + e_hi = d_hi["total_energy"] + + # Overlapping histogram bins + e_min = min(e_lo.min(), e_hi.min()) + e_max = max(e_lo.max(), e_hi.max()) + bins = np.linspace(e_min, e_max, 50) + bin_centers = (bins[:-1] + bins[1:]) / 2 + + h_lo, _ = np.histogram(e_lo, bins=bins, density=True) + h_hi, _ = np.histogram(e_hi, bins=bins, density=True) + + # Plot overlapping distributions + ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") + ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") + + # Inset: log ratio + mask = (h_lo > 0) & (h_hi > 0) + if mask.sum() > 2: + log_ratio = np.log(h_lo[mask] / h_hi[mask]) + bc = bin_centers[mask] + # Linear fit + slope, intercept = np.polyfit(bc, log_ratio, 1) + inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) + inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) + inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, + label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") + inset.set_xlabel("E (eV)", fontsize=7) + inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) + inset.legend(fontsize=6) + inset.tick_params(labelsize=6) + + ax.set_title(name, fontsize=10) + ax.set_xlabel("Total Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in ensemble check plots (NVT) +# +# Uses `physical_validation.ensemble.check(..., screen=True)` which shows the forward/reverse work distributions and the linear fit of the log-likelihood ratio. + +# %% +for name in NVT_INTEGRATORS: + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — Ensemble check (physical_validation)") + print(f"{'='*60}") + sd_lo = build_sim_data(d_lo, temp_low, ensemble="NVT") + sd_hi = build_sim_data(d_hi, temp_high, ensemble="NVT") + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=True, data_is_uncorrelated=True, + screen=True, verbosity=2, filename=f"plots/{name}_ensemble_check.png" + ) + print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") + +plt.show() + +# %% [markdown] +# ## 4. NPT Time Series +# +# Temperature, energy, and volume vs step for each NPT integrator. + +# %% +NPT_INTEGRATORS = ["npt_langevin", "npt_nose_hoover", "npt_isotropic_crescale", "npt_anisotropic_crescale"] + +fig, axes = plt.subplots(len(NPT_INTEGRATORS), 3, figsize=(16, 3.5 * len(NPT_INTEGRATORS)), + squeeze=False, sharex=True) +fig.suptitle("NPT Integrators — Time Series", fontsize=14, y=1.01) + +for row, name in enumerate(NPT_INTEGRATORS): + for temp in TEMPS: + label = f"{name}_T{temp:.0f}K" + try: + d = load(label) + except FileNotFoundError: + continue + + steps = np.arange(len(d["temperature"])) + target_T = float(d["target_temperature"]) + tag = f"T={target_T:.0f}K" + + # Temperature + axes[row, 0].plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=tag) + axes[row, 0].axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) + axes[row, 0].set_ylabel("Temperature (K)") + axes[row, 0].set_title(name) + axes[row, 0].legend(fontsize=7) + + # Total energy + axes[row, 1].plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=tag) + axes[row, 1].set_ylabel("Total Energy (eV)") + axes[row, 1].set_title(name) + axes[row, 1].legend(fontsize=7) + + # Volume + axes[row, 2].plot(steps, d["volumes"], alpha=0.5, lw=0.5, label=tag) + axes[row, 2].set_ylabel("Volume (ų)") + axes[row, 2].set_title(name) + axes[row, 2].legend(fontsize=7) + +for j in range(3): + axes[-1, j].set_xlabel("Step") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ## 5. KE Distribution — NPT Integrators +# +# Same gamma distribution check, but for NPT integrators at 100K. + +# %% +fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NPT, T=100K)", fontsize=13) + +for ax, name in zip(axes, NPT_INTEGRATORS): + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + ke = d["kinetic_energy"] + natoms = int(d["natoms"]) + target_T = float(d["target_temperature"]) + + ndof = 3 * natoms - 3 + shape = ndof / 2 + scale = k_B_eV * target_T + + ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") + + x = np.linspace(ke.min(), ke.max(), 300) + pdf = stats.gamma.pdf(x, a=shape, scale=scale) + ax.plot(x, pdf, "r-", lw=2, label="Theory") + + d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) + d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) + ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=9) + ax.set_xlabel("Kinetic Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in KE distribution plots (NPT) + +# %% +for name in NPT_INTEGRATORS: + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — KE distribution (physical_validation)") + print(f"{'='*60}") + sd = build_sim_data(d, 100.0, ensemble="NPT") + result = physical_validation.kinetic_energy.distribution( + sd, strict=False, screen=True, verbosity=2, + ) + print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") + +plt.show() + +# %% [markdown] +# ## 6. Ensemble Check — NPT Integrators +# +# Same Boltzmann weight ratio check at T=88K vs T=100K for NPT integrators. + +# %% +fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("NPT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) + +for ax, name in zip(axes, NPT_INTEGRATORS): + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + e_lo = d_lo["total_energy"] + e_hi = d_hi["total_energy"] + + e_min = min(e_lo.min(), e_hi.min()) + e_max = max(e_lo.max(), e_hi.max()) + bins = np.linspace(e_min, e_max, 50) + bin_centers = (bins[:-1] + bins[1:]) / 2 + + h_lo, _ = np.histogram(e_lo, bins=bins, density=True) + h_hi, _ = np.histogram(e_hi, bins=bins, density=True) + + ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") + ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") + + mask = (h_lo > 0) & (h_hi > 0) + if mask.sum() > 2: + log_ratio = np.log(h_lo[mask] / h_hi[mask]) + bc = bin_centers[mask] + slope, intercept = np.polyfit(bc, log_ratio, 1) + inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) + inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) + inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, + label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") + inset.set_xlabel("E (eV)", fontsize=7) + inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) + inset.legend(fontsize=6) + inset.tick_params(labelsize=6) + + ax.set_title(name, fontsize=9) + ax.set_xlabel("Total Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in ensemble check plots (NPT) + +# %% +for name in NPT_INTEGRATORS: + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — Ensemble check (physical_validation)") + print(f"{'='*60}") + sd_lo = build_sim_data(d_lo, temp_low, ensemble="NPT") + sd_hi = build_sim_data(d_hi, temp_high, ensemble="NPT") + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, # data_is_uncorrelated=True, + screen=True, verbosity=2, + ) + print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: #ConvergenceError + print(f" ConvergenceError: {e}") + +plt.show() + +# %% [markdown] +# ## 6b. NPT Pressure Ensemble Check +# +# Compare volume distributions at two pressures (same temperature). +# For correct NPT sampling, `physical_validation` checks that volume distributions +# at different pressures satisfy the expected Boltzmann weight relationship. +# This is a 1D check on volumes only — no energy data needed. + +# %% +from torch_sim.units import MetalUnits + +P_BAR_CONVERSION = float(MetalUnits.pressure) # 1 bar in eV/Ang^3 + +# Discover available pressure-sweep data files +pressure_files = sorted(DATA_DIR.glob("*_P*bar.npz")) +print("Pressure-sweep data files:") +for f in pressure_files: + print(f" {f.name}") + +# Parse integrator -> {pressure_bar: filename} mapping +pressure_data = {} # integrator_name -> [(p_bar, filename), ...] +for f in pressure_files: + stem = f.stem # e.g. npt_langevin_T100K_P0bar + parts = stem.rsplit("_P", 1) + if len(parts) != 2: + continue + prefix = parts[0] # e.g. npt_langevin_T100K + p_bar = float(parts[1].replace("bar", "")) + int_prefix = prefix.rsplit("_T", 1)[0] # e.g. npt_langevin + pressure_data.setdefault(int_prefix, []).append((p_bar, f.stem)) + +for name, entries in pressure_data.items(): + entries.sort(key=lambda x: x[0]) + print(f"\n{name}: {[(p, fn) for p, fn in entries]}") + +# %% +# Volume distributions at different pressures for each NPT integrator +integrators_with_pressure = [n for n in NPT_INTEGRATORS if n in pressure_data] + +if integrators_with_pressure: + fig, axes = plt.subplots(1, len(integrators_with_pressure), + figsize=(5 * len(integrators_with_pressure), 4), squeeze=False) + axes = axes[0] + fig.suptitle("NPT Pressure Check \u2014 Volume Distributions at Same T, Different P", fontsize=13) + + for ax, name in zip(axes, integrators_with_pressure): + entries = pressure_data[name] + colors = plt.cm.coolwarm(np.linspace(0, 1, len(entries))) + for (p_bar, fname), color in zip(entries, colors): + d = load(fname) + vols = d["volumes"] + ax.hist(vols, bins=50, density=True, alpha=0.5, color=color, + label=f"P={p_bar:.0f} bar") + ax.axvline(vols.mean(), color=color, ls="--", lw=1, alpha=0.7) + ax.set_title(name, fontsize=10) + ax.set_xlabel("Volume (\u00c5\u00b3)") + ax.legend(fontsize=8) + + axes[0].set_ylabel("Probability Density") + fig.tight_layout() + plt.savefig("plots/npt_pressure_volume_distributions.png", dpi=150, bbox_inches="tight") + plt.show() +else: + print("No pressure-sweep data found. Run: python run_npt.py --mode pressure") + +# %% [markdown] +# ### physical_validation built-in pressure ensemble check +# +# Uses `physical_validation.ensemble.check()` with two simulations at the same temperature +# but different pressures. The library auto-detects this case and performs a 1D volume-based check. + +# %% +for name in integrators_with_pressure: + entries = pressure_data[name] + if len(entries) < 2: + print(f"{name}: need at least 2 pressures, got {len(entries)}") + continue + + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_lo = load(fname_lo) + d_hi = load(fname_hi) + + p_lo_eva3 = float(d_lo["external_pressure"]) + p_hi_eva3 = float(d_hi["external_pressure"]) + temp = float(d_lo["target_temperature"]) + print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " + f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") + + print(f"\n{'='*60}") + print(f" {name} \u2014 Pressure ensemble check") + print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") + print(f"{'='*60}") + + # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) + # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) + + sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) + sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) + + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, #data_is_uncorrelated=True, + screen=True, verbosity=2, + filename=f"plots/{name}_pressure_ensemble_check.png", + ) + print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: + print(f" Error: {e}") + +plt.show() + +# %% +for name in integrators_with_pressure: + entries = pressure_data[name] + if len(entries) < 2: + print(f"{name}: need at least 2 pressures, got {len(entries)}") + continue + + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_lo = load(fname_lo) + d_hi = load(fname_hi) + + p_lo_eva3 = float(d_lo["external_pressure"]) + p_hi_eva3 = float(d_hi["external_pressure"]) + temp = float(d_lo["target_temperature"]) + print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " + f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") + + print(f"\n{'='*60}") + print(f" {name} \u2014 Pressure ensemble check") + print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") + print(f"{'='*60}") + + # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) + # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) + + sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) + sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) + + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, data_is_uncorrelated=True, + screen=True, verbosity=2, + filename=f"plots/{name}_pressure_ensemble_check.png", + ) + print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: + print(f" Error: {e}") + +plt.show() + +# %% [markdown] +# ## 8. NVE Integrator Convergence +# +# RMSD of the conserved quantity (total energy) vs timestep on a log-log scale. +# For velocity Verlet, the RMSD should scale as $\text{dt}^2$ (slope = 2 on log-log). + +# %% +try: + nve = load("nve_convergence") + timesteps_ps = nve["timesteps_ps"] + + dts, rmsds, drifts = [], [], [] + for dt_ps in timesteps_ps: + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + dts.append(dt_ps) + rmsds.append(com.std()) + drifts.append(com[-1] - com[0]) + + dts = np.array(dts) + rmsds = np.array(rmsds) + drifts = np.array(drifts) + + fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 5)) + fig.suptitle("NVE Integrator Convergence (4-atom Ar, T=5K)", fontsize=13) + + # --- Panel 1: RMSD vs dt (log-log) --- + ax1.loglog(dts, rmsds, "ko-", ms=8, lw=2, label="RMSD(E_tot)") + + # dt^2 reference line through middle point + mid = len(dts) // 2 + ref = rmsds[mid] * (dts / dts[mid]) ** 2 + ax1.loglog(dts, ref, "r--", lw=1.5, alpha=0.7, label="$\\propto dt^2$ (reference)") + + # Fit slope + log_dts = np.log(dts) + log_rmsds = np.log(rmsds) + # Only fit points where RMSD is clearly above noise floor + mask = rmsds > 1.5 * rmsds.min() + if mask.sum() >= 2: + slope, intercept = np.polyfit(log_dts[mask], log_rmsds[mask], 1) + ax1.set_title(f"Log-log slope = {slope:.2f} (expected: 2.0)") + else: + ax1.set_title("RMSD vs dt") + + ax1.set_xlabel("Timestep (ps)") + ax1.set_ylabel("RMSD of E_total (eV)") + ax1.legend() + + # --- Panel 2: Energy time series for each dt --- + for i, dt_ps in enumerate(timesteps_ps): + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + steps = np.arange(len(com)) + ax2.plot(steps, com - com[0], alpha=0.7, lw=0.8, label=f"dt={dt_ps:.4f} ps") + + ax2.set_xlabel("Step") + ax2.set_ylabel("E_total - E_total[0] (eV)") + ax2.set_title("Energy drift per timestep") + ax2.legend(fontsize=7, ncol=2) + + # --- Panel 3: RMSD ratio table --- + ax3.axis("off") + headers = ["dt1 (ps)", "dt2 (ps)", "dt1/dt2", "RMSD1/RMSD2", "(dt1/dt2)²"] + rows = [] + for i in range(len(dts) - 1): + dt_ratio = dts[i] / dts[i + 1] + rmsd_ratio = rmsds[i] / rmsds[i + 1] + expected = dt_ratio ** 2 + rows.append([f"{dts[i]:.4f}", f"{dts[i+1]:.4f}", + f"{dt_ratio:.3f}", f"{rmsd_ratio:.3f}", f"{expected:.3f}"]) + + table = ax3.table(cellText=rows, colLabels=headers, loc="center", cellLoc="center") + table.auto_set_font_size(False) + table.set_fontsize(9) + table.scale(1.0, 1.5) + ax3.set_title("Convergence ratios", pad=20) + + fig.tight_layout() + plt.show() + +except FileNotFoundError: + print("No NVE data found. Run: python run_nve.py") + +# %% [markdown] +# ### physical_validation built-in integrator convergence plot +# +# Uses `physical_validation.integrator.convergence(..., screen=True)` which plots RMSD vs dt with the expected dt^2 reference line. + +# %% +try: + nve = load("nve_convergence") + timesteps_ps = nve["timesteps_ps"] + natoms = int(nve["natoms"]) + masses = nve["masses"] + volume = float(nve["volume"]) + + # Pick 3 timesteps in the dt^2 regime (avoid noise floor and nonlinear regime) + # Use [0.007, 0.005, 0.004] which showed good convergence + selected_dts = [0.007, 0.005, 0.004] + + simulations = [] + for dt_ps in selected_dts: + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + dt_internal = float(nve[f"{key}_dt_internal"]) + + system = physical_validation.data.SystemData( + natoms=natoms, nconstraints=0, + ndof_reduction_tra=3, ndof_reduction_rot=0, mass=masses) + ensemble = physical_validation.data.EnsembleData( + ensemble="NVE", natoms=natoms, volume=volume) + obs = physical_validation.data.ObservableData(constant_of_motion=com) + sd = physical_validation.data.SimulationData( + units=make_unit_data(), dt=dt_internal, + system=system, ensemble=ensemble, observables=obs) + simulations.append(sd) + + result = physical_validation.integrator.convergence( + simulations, verbose=True, screen=True, + ) + print(f"\nConvergence deviation: {result:.3f} ({'PASS' if result < 0.5 else 'FAIL'})") + +except FileNotFoundError: + print("No NVE data found. Run: python run_nve.py") + +plt.show() + +# %% [markdown] +# ## 9. Summary Table +# +# Run `physical_validation` checks programmatically on all available data and collect pass/fail results. + +# %% +# Collect results +results = [] +all_integrators = NVT_INTEGRATORS + NPT_INTEGRATORS + +for name in all_integrators: + is_npt = name.startswith("npt") + ensemble = "NPT" if is_npt else "NVT" + + # --- KE Distribution at 100K --- + label_100 = f"{name}_T100K" + ke_status = "no data" + try: + d100 = load(label_100) + sd = build_sim_data(d100, 100.0, ensemble=ensemble) + d_mean, d_width = physical_validation.kinetic_energy.distribution( + sd, strict=False, verbosity=0) + ke_pass = abs(d_mean) < THRESHOLD and abs(d_width) < THRESHOLD + ke_status = f"PASS (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" if ke_pass else f"FAIL (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" + except Exception as e: + ke_status = f"ERROR: {e}" + + # --- Ensemble Check 88K vs 100K (temperature) --- + ens_status = "no data" + try: + d88 = load(f"{name}_T88K") + d100 = load(f"{name}_T100K") + sd_lo = build_sim_data(d88, 88.0, ensemble=ensemble) + sd_hi = build_sim_data(d100, 100.0, ensemble=ensemble) + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, total_energy=True, data_is_uncorrelated=True, verbosity=0) + max_q = max(abs(q) for q in quantiles) + ens_pass = max_q < THRESHOLD + ens_status = f"PASS (max |q|={max_q:.2f}\u03c3)" if ens_pass else f"FAIL (max |q|={max_q:.2f}\u03c3)" + except Exception as e: + ens_status = f"ERROR: {e}" + + # --- Pressure Ensemble Check (NPT only, same T, different P) --- + press_status = "N/A" + if is_npt and name in pressure_data: + entries = pressure_data[name] + if len(entries) >= 2: + try: + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_plo = load(fname_lo) + d_phi = load(fname_hi) + p_lo_eva3 = float(d_plo["external_pressure"]) + p_hi_eva3 = float(d_phi["external_pressure"]) + temp_p = float(d_plo["target_temperature"]) + sd_plo = build_sim_data(d_plo, temp_p, ensemble="NPT", pressure=p_lo_eva3) + sd_phi = build_sim_data(d_phi, temp_p, ensemble="NPT", pressure=p_hi_eva3) + quantiles_p = physical_validation.ensemble.check( + sd_plo, sd_phi, total_energy=False, data_is_uncorrelated=True, verbosity=0) + max_qp = max(abs(q) for q in quantiles_p) + press_pass = max_qp < THRESHOLD + press_status = f"PASS (max |q|={max_qp:.2f}\u03c3)" if press_pass else f"FAIL (max |q|={max_qp:.2f}\u03c3)" + except Exception as e: + press_status = f"ERROR: {e}" + else: + press_status = "need 2 pressures" + + results.append((name, ke_status, ens_status, press_status)) + +# Print table +print(f"{'Integrator':<30} {'KE Distribution':<40} {'Ensemble (T)':<35} {'Ensemble (P)':<35}") +print("-" * 140) +for name, ke, ens, press in results: + print(f"{name:<30} {ke:<40} {ens:<35} {press:<35}") + + diff --git a/fast_integrator_tests/common.py b/fast_integrator_tests/common.py new file mode 100644 index 000000000..69e21fd06 --- /dev/null +++ b/fast_integrator_tests/common.py @@ -0,0 +1,48 @@ +"""Shared constants and helpers for integrator validation scripts.""" + +import numpy as np +import torch +from ase.build import bulk +from pathlib import Path + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.units import MetalUnits + +DEVICE = torch.device("cpu") +# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +DTYPE = torch.float64 +print(f"Using device: {DEVICE}") + +# LJ Argon +SIGMA = 3.405 +EPSILON = 0.0104 +CUTOFF = 2.5 * SIGMA + +DATA_DIR = Path(__file__).parent / "data" +torch.set_num_threads(4) + +def make_lj_model(compute_stress=False): + return LennardJonesModel( + use_neighbor_list=False, + sigma=SIGMA, + epsilon=EPSILON, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=compute_stress, + cutoff=CUTOFF, + ) + + +def make_ar_supercell(repeat=(2, 2, 2)): + atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) + return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) + + +def to_kT(temperature_K): + return temperature_K * float(MetalUnits.temperature) + + +def to_dt(timestep_ps): + return timestep_ps * float(MetalUnits.time) diff --git a/fast_integrator_tests/run_npt.py b/fast_integrator_tests/run_npt.py new file mode 100644 index 000000000..0b18f80ca --- /dev/null +++ b/fast_integrator_tests/run_npt.py @@ -0,0 +1,174 @@ +"""Run NPT simulations for all NPT integrators and save observables. + +Usage: + python run_npt.py # temperature sweep (default) + python run_npt.py --mode pressure # pressure sweep at fixed T + python run_npt.py --mode all # both sweeps + python run_npt.py --integrator npt_langevin +""" + +import argparse +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, DEVICE, DTYPE, make_ar_supercell, make_lj_model, to_dt, to_kT +from torch_sim.units import MetalUnits + +NPT_INTEGRATORS = [ + "npt_langevin", + "npt_nose_hoover", + "npt_isotropic_crescale", + "npt_anisotropic_crescale", +] + +TEMPERATURES = [88.0, 100.0] +TIMESTEP_PS = 0.004 +EXTERNAL_PRESSURE = 0.0 +N_STEPS = 20_000 +N_EQUILIBRATION = 3_000 + +# Pressure sweep: two pressures at a fixed temperature. +# physical_validation compares volume distributions at same T, different P. +PRESSURE_SWEEP_TEMP = 100.0 # K +PRESSURES_EVA3 = [0.0, 0.0001] # eV/ų (0 bar and ~160 bar) + + +def run_npt(integrator_name, sim_state, model, temperature, external_pressure=0.0, + seed=42): + kT = to_kT(temperature) + dt = torch.tensor(to_dt(TIMESTEP_PS), device=DEVICE, dtype=DTYPE) + ext_p = torch.tensor(external_pressure, device=DEVICE, dtype=DTYPE) + natoms = int(sim_state.positions.shape[0]) + + sim_state = sim_state.clone() + sim_state.rng = seed + + # Initialize + if integrator_name == "npt_langevin": + # state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt) + state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt, b_tau = 1 * dt, alpha= 5 * dt) # Better parameters + elif integrator_name == "npt_nose_hoover": + state = ts.npt_nose_hoover_init(sim_state, model, kT=kT, dt=dt, t_tau=10 * dt, b_tau=100 * dt) + elif integrator_name == "npt_isotropic_crescale": + state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt, tau_p = 10 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) + elif integrator_name == "npt_anisotropic_crescale": + state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt/2, tau_p = 100 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) + else: + raise ValueError(f"Unknown: {integrator_name}") + + def step(s): + if integrator_name == "npt_langevin": + return ts.npt_langevin_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) + if integrator_name == "npt_nose_hoover": + return ts.npt_nose_hoover_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) + if integrator_name == "npt_isotropic_crescale": + return ts.npt_crescale_isotropic_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, tau = 5 * dt + ) + return ts.npt_crescale_anisotropic_step( + s, model, dt=dt/2, kT=kT, external_pressure=ext_p, tau = 5 * dt + ) + + # Equilibration + print(f" Equilibrating {N_EQUILIBRATION} steps...") + for _ in range(N_EQUILIBRATION): + state = step(state) + + # Production + print(f" Producing {N_STEPS} steps...") + ke_list, pe_list, total_e_list = [], [], [] + temp_list, volume_list = [], [] + + for i in range(N_STEPS): + state = step(state) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + temp = float( + ts.calc_temperature(masses=state.masses, momenta=state.momenta) + ) + cell = state.cell[0].detach().cpu().numpy() + vol = float(np.abs(np.linalg.det(cell))) + + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + temp_list.append(temp) + volume_list.append(vol) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "temperature": np.array(temp_list), + "volumes": np.array(volume_list), + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": to_dt(TIMESTEP_PS), + "natoms": natoms, + "target_temperature": temperature, + "external_pressure": external_pressure, + "timestep_ps": TIMESTEP_PS, + "integrator": integrator_name, + } + + +def pressure_to_bar(p_eva3): + """Convert eV/ų to bar for display.""" + return p_eva3 / float(MetalUnits.pressure) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--integrator", choices=NPT_INTEGRATORS, default=None) + parser.add_argument( + "--mode", choices=["temperature", "pressure", "all"], default="all", + help="'temperature' = vary T at P=0 (default), " + "'pressure' = vary P at fixed T, 'all' = both", + ) + args = parser.parse_args() + + integrators = [args.integrator] if args.integrator else NPT_INTEGRATORS + DATA_DIR.mkdir(exist_ok=True) + + sim_state = make_ar_supercell(repeat=(3, 3, 3)) + model = make_lj_model(compute_stress=True) + + # --- Temperature sweep (same P=0, vary T) --- + if args.mode in ("temperature", "all"): + print("=== Temperature sweep ===") + for name in integrators: + for temp in TEMPERATURES: + seed = 42 if temp == TEMPERATURES[0] else 123 + label = f"{name}_T{temp:.0f}K" + print(f" Running {label} ...") + t0 = time.time() + data = run_npt(name, sim_state, model, temp, seed=seed) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + # --- Pressure sweep (same T, vary P) --- + if args.mode in ("pressure", "all"): + print(f"\n=== Pressure sweep at T={PRESSURE_SWEEP_TEMP:.0f}K ===") + for name in integrators: + for i, p_eva3 in enumerate(PRESSURES_EVA3): + p_bar = pressure_to_bar(p_eva3) + seed = 42 if i == 0 else 123 + label = f"{name}_T{PRESSURE_SWEEP_TEMP:.0f}K_P{p_bar:.0f}bar" + print(f" Running {label} ...") + t0 = time.time() + data = run_npt( + name, sim_state, model, PRESSURE_SWEEP_TEMP, + external_pressure=p_eva3, seed=seed, + ) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + +if __name__ == "__main__": + main() diff --git a/fast_integrator_tests/run_nve.py b/fast_integrator_tests/run_nve.py new file mode 100644 index 000000000..a0d60ca9b --- /dev/null +++ b/fast_integrator_tests/run_nve.py @@ -0,0 +1,91 @@ +"""Run NVE simulations at multiple timesteps for convergence analysis. + +Usage: + python run_nve.py +""" + +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT + +# Timesteps: sweep a range so the notebook can pick the best subset +TIMESTEPS_PS = [0.010, 0.008, 0.007, 0.006, 0.005, 0.004, 0.003, 0.002] +TEMPERATURE = 5.0 # K - low so integration error dominates +N_STEPS = 10_000 +SEED = 42 + + +def run_nve(sim_state, model, kT_init, timestep_ps): + dt = to_dt(timestep_ps) + + sim_state = sim_state.clone() + sim_state.rng = SEED + state = ts.nve_init(sim_state, model, kT=kT_init) + + ke_list, pe_list, com_list = [], [], [] + + for _ in range(N_STEPS): + state = ts.nve_step(state, model, dt=dt) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + ke_list.append(ke) + pe_list.append(pe) + com_list.append(ke + pe) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "constant_of_motion": np.array(com_list), + "dt_internal": dt, + "timestep_ps": timestep_ps, + } + + +def main(): + DATA_DIR.mkdir(exist_ok=True) + + # sim_state = make_ar_supercell(repeat=(1, 1, 1)) + sim_state = make_ar_supercell(repeat=(6, 6, 6)) + model = make_lj_model() + kT_init = to_kT(TEMPERATURE) + natoms = int(sim_state.positions.shape[0]) + masses = sim_state.masses.detach().cpu().numpy() + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + all_results = {} + for dt_ps in TIMESTEPS_PS: + label = f"nve_dt{dt_ps:.4f}" + print(f" Running {label} ...") + t0 = time.time() + data = run_nve(sim_state, model, kT_init, dt_ps) + elapsed = time.time() - t0 + all_results[f"dt_{dt_ps}"] = data + print(f" std(E_tot) = {data['constant_of_motion'].std():.3e} ({elapsed:.1f}s)") + + # Save everything into one npz + save_dict = { + "timesteps_ps": np.array(TIMESTEPS_PS), + "temperature": TEMPERATURE, + "natoms": natoms, + "masses": masses, + "volume": volume, + "n_steps": N_STEPS, + } + for dt_ps in TIMESTEPS_PS: + key = f"dt_{dt_ps}" + for field in ("constant_of_motion", "kinetic_energy", "potential_energy"): + save_dict[f"{key}_{field}"] = all_results[key][field] + save_dict[f"{key}_dt_internal"] = all_results[key]["dt_internal"] + + outpath = DATA_DIR / "nve_convergence.npz" + np.savez(outpath, **save_dict) + print(f"\n Saved {outpath.name}") + + +if __name__ == "__main__": + main() diff --git a/fast_integrator_tests/run_nvt.py b/fast_integrator_tests/run_nvt.py new file mode 100644 index 000000000..ba8d141cf --- /dev/null +++ b/fast_integrator_tests/run_nvt.py @@ -0,0 +1,116 @@ +"""Run NVT simulations for all NVT integrators and save observables. + +Usage: + python run_nvt.py + python run_nvt.py --integrator nvt_langevin +""" + +import argparse +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT + +NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] + +# Two temperatures for ensemble check +TEMPERATURES = [95.0, 100.0] +TIMESTEP_PS = 0.004 +N_STEPS = 10_000 +N_EQUILIBRATION = 2_000 + + +def run_nvt(integrator_name, sim_state, model, temperature, seed=42): + kT = to_kT(temperature) + dt = to_dt(TIMESTEP_PS) + natoms = int(sim_state.positions.shape[0]) + + sim_state = sim_state.clone() + sim_state.rng = seed + + # Initialize + if integrator_name == "nvt_langevin": + state = ts.nvt_langevin_init(sim_state, model, kT=kT) + elif integrator_name == "nvt_nose_hoover": + state = ts.nvt_nose_hoover_init(sim_state, model, kT=kT, dt=dt) + elif integrator_name == "nvt_vrescale": + state = ts.nvt_vrescale_init(sim_state, model, kT=kT) + else: + raise ValueError(f"Unknown: {integrator_name}") + + def step(s): + if integrator_name == "nvt_langevin": + return ts.nvt_langevin_step(s, model, dt=dt, kT=kT) + if integrator_name == "nvt_nose_hoover": + return ts.nvt_nose_hoover_step(s, model, dt=dt, kT=kT) + return ts.nvt_vrescale_step(model, s, dt=dt, kT=kT) + + # Equilibration + print(f" Equilibrating {N_EQUILIBRATION} steps...") + for _ in range(N_EQUILIBRATION): + state = step(state) + + # Production + print(f" Producing {N_STEPS} steps...") + ke_list, pe_list, total_e_list = [], [], [] + temp_list = [] + + for i in range(N_STEPS): + state = step(state) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + temp = float( + ts.calc_temperature(masses=state.masses, momenta=state.momenta) + ) + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + temp_list.append(temp) + + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "temperature": np.array(temp_list), + "volume": volume, + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": to_dt(TIMESTEP_PS), + "natoms": natoms, + "target_temperature": temperature, + "timestep_ps": TIMESTEP_PS, + "integrator": integrator_name, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--integrator", choices=NVT_INTEGRATORS, default=None) + args = parser.parse_args() + + integrators = [args.integrator] if args.integrator else NVT_INTEGRATORS + DATA_DIR.mkdir(exist_ok=True) + + sim_state = make_ar_supercell(repeat=(6, 6, 6)) + model = make_lj_model() + + for name in integrators: + for temp in TEMPERATURES: + seed = 42 if temp == TEMPERATURES[0] else 123 + label = f"{name}_T{temp:.0f}K" + print(f" Running {label} ...") + t0 = time.time() + data = run_nvt(name, sim_state, model, temp, seed=seed) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + +if __name__ == "__main__": + main() From e914b3f294ac8e8a485acc5ba536c7266794f445 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 13 Feb 2026 16:14:32 +0000 Subject: [PATCH 03/15] vibecoded (need to check) --- pyproject.toml | 6 +- tests/test_physical_validation.py | 345 ++++++++++++++++++++++++++++++ 2 files changed, 350 insertions(+), 1 deletion(-) create mode 100644 tests/test_physical_validation.py diff --git a/pyproject.toml b/pyproject.toml index 90de3c4a6..a56b49bca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ [project.optional-dependencies] test = [ "torch-sim-atomistic[io,symmetry,vesin]", + "physical-validation>=1.0.5", "platformdirs>=4.0.0", "psutil>=7.0.0", "pymatgen>=2025.6.14", @@ -139,8 +140,11 @@ check-filenames = true ignore-words-list = ["convertor"] # codespell:ignore convertor [tool.pytest.ini_options] -addopts = ["-p no:warnings"] +addopts = ["-p no:warnings", "-m not physical_validation"] testpaths = ["tests"] +markers = [ + "physical_validation: long-running physical validation tests (run with: pytest -m physical_validation)", +] [tool.uv] # make these dependencies mutually exclusive since they use incompatible e3nn versions diff --git a/tests/test_physical_validation.py b/tests/test_physical_validation.py new file mode 100644 index 000000000..e57f85c8b --- /dev/null +++ b/tests/test_physical_validation.py @@ -0,0 +1,345 @@ +"""Physical validation tests for torch-sim MD integrators. + +Uses the physical_validation library (https://github.com/shirtsgroup/physical_validation) +to verify that integrators produce physically correct results. These tests are +long-running (~5 min total) and excluded by default. Run with: + + pytest -m physical_validation -v +""" + +import numpy as np +import pytest +import torch +from ase.build import bulk + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.units import MetalUnits + +physical_validation = pytest.importorskip("physical_validation") + +DEVICE = torch.device("cpu") +DTYPE = torch.float64 + +# LJ Argon parameters +SIGMA = 3.405 +EPSILON = 0.0104 +CUTOFF = 2.5 * SIGMA + + +def _make_unit_data(): + """Create UnitData for torch-sim's MetalUnits system.""" + return physical_validation.data.UnitData( + kb=float(MetalUnits.temperature), # k_B in eV/K = 8.617e-5 + energy_str="eV", + energy_conversion=1.0, + length_str="Ang", + length_conversion=1.0, + volume_str="Ang^3", + volume_conversion=1.0, + temperature_str="K", + temperature_conversion=1.0, + pressure_str="eV/Ang^3", + pressure_conversion=1.0, + time_str="internal", + time_conversion=1.0, + ) + + +def _make_lj_model(): + """Create a Lennard-Jones model for Argon.""" + return LennardJonesModel( + use_neighbor_list=False, + sigma=SIGMA, + epsilon=EPSILON, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=False, + cutoff=CUTOFF, + ) + + +def _make_ar_supercell(repeat=(2, 2, 2)): + """Create an FCC Argon supercell SimState.""" + atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) + return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) + + +def _run_nvt_langevin( + sim_state, + model, + temperature, + timestep_ps, + n_steps, + n_equilibration, + seed=42, +): + """Run NVT Langevin simulation and collect per-step observables.""" + kT = temperature * float(MetalUnits.temperature) + dt_internal = timestep_ps * float(MetalUnits.time) + natoms = int(sim_state.positions.shape[0]) + + state = ts.nvt_langevin_init(sim_state, model, kT=kT, seed=seed) + + # Equilibration + for _ in range(n_equilibration): + state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) + + # Production - collect observables + ke_list = [] + pe_list = [] + total_e_list = [] + position_list = [] + velocity_list = [] + + for _ in range(n_steps): + state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) + + ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + pe = state.energy.sum() + ke_list.append(float(ke)) + pe_list.append(float(pe)) + total_e_list.append(float(ke + pe)) + position_list.append(state.positions.detach().cpu().numpy().copy()) + velocity_list.append(state.velocities.detach().cpu().numpy().copy()) + + # Compute volume from cell + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "positions": np.array(position_list), + "velocities": np.array(velocity_list), + "volume": volume, + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": dt_internal, + "natoms": natoms, + } + + +def _run_nve(sim_state, model, kT_init, timestep_ps, n_steps, seed=42): + """Run NVE simulation and collect constant of motion.""" + dt_internal = timestep_ps * float(MetalUnits.time) + + state = ts.nve_init(sim_state, model, kT=kT_init, seed=seed) + + com_list = [] + for _ in range(n_steps): + state = ts.nve_step(state, model, dt=dt_internal) + ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + pe = state.energy.sum() + com_list.append(float(ke + pe)) + + return { + "constant_of_motion": np.array(com_list), + "dt_internal": dt_internal, + } + + +def _build_nvt_simulation_data(run_data, temperature): + """Build a physical_validation SimulationData from NVT run results.""" + units = _make_unit_data() + + system = physical_validation.data.SystemData( + natoms=run_data["natoms"], + nconstraints=0, + ndof_reduction_tra=3, + ndof_reduction_rot=0, + mass=run_data["masses"], + ) + + ensemble_data = physical_validation.data.EnsembleData( + ensemble="NVT", + natoms=run_data["natoms"], + volume=run_data["volume"], + temperature=temperature, + ) + + observables = physical_validation.data.ObservableData( + kinetic_energy=run_data["kinetic_energy"], + potential_energy=run_data["potential_energy"], + total_energy=run_data["total_energy"], + ) + + trajectory = physical_validation.data.TrajectoryData( + position=run_data["positions"], + velocity=run_data["velocities"], + ) + + return physical_validation.data.SimulationData( + units=units, + dt=run_data["dt_internal"], + system=system, + ensemble=ensemble_data, + observables=observables, + trajectory=trajectory, + ) + + +@pytest.mark.physical_validation +def test_ke_distribution(): + """Test that kinetic energy follows the Maxwell-Boltzmann distribution. + + Runs NVT Langevin at 100K on a 2x2x2 Ar supercell (32 atoms) and checks + that the KE distribution matches the analytical Maxwell-Boltzmann prediction. + """ + sim_state = _make_ar_supercell(repeat=(2, 2, 2)) + model = _make_lj_model() + temperature = 100.0 # K + + run_data = _run_nvt_langevin( + sim_state, + model, + temperature=temperature, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=42, + ) + + data = _build_nvt_simulation_data(run_data, temperature) + + result = physical_validation.kinetic_energy.distribution( + data, + strict=False, + verbosity=0, + ) + # strict=False returns (d_mean, d_width) in sigma units + d_mean, d_width = result + + assert abs(d_mean) < 3, ( + f"KE mean deviation {d_mean:.2f} sigma exceeds threshold" + ) + assert abs(d_width) < 3, ( + f"KE width deviation {d_width:.2f} sigma exceeds threshold" + ) + + +@pytest.mark.physical_validation +def test_integrator_convergence(): + """Test that NVE energy error scales as dt^2 (velocity Verlet). + + Runs NVE at 3 different timesteps from identical initial conditions on a + 4-atom Ar unit cell at low temperature (5K). Low temperature minimizes + thermal fluctuations so the integration error dominates the RMSD of the + conserved quantity, allowing the dt^2 convergence to be observed. + """ + sim_state = _make_ar_supercell(repeat=(1, 1, 1)) # 4 atoms + model = _make_lj_model() + temperature = 5.0 # K, low T so integration error dominates + kT_init = temperature * float(MetalUnits.temperature) + + # Timesteps chosen so integration error >> thermal fluctuations at all dt. + # Factor of ~sqrt(2) spacing gives dt^2 ratio of ~2.0 per step. + timesteps = [0.008, 0.00566, 0.004] # ps + n_steps = 5_000 + seed = 42 + + natoms = int(sim_state.positions.shape[0]) + masses = sim_state.masses.detach().cpu().numpy() + volume = float( + np.abs(np.linalg.det(sim_state.cell[0].detach().cpu().numpy())) + ) + units = _make_unit_data() + + simulations = [] + for dt_ps in timesteps: + run_data = _run_nve( + sim_state, + model, + kT_init=kT_init, + timestep_ps=dt_ps, + n_steps=n_steps, + seed=seed, + ) + + system = physical_validation.data.SystemData( + natoms=natoms, + nconstraints=0, + ndof_reduction_tra=3, + ndof_reduction_rot=0, + mass=masses, + ) + + ensemble_data = physical_validation.data.EnsembleData( + ensemble="NVE", + natoms=natoms, + volume=volume, + ) + + observables = physical_validation.data.ObservableData( + constant_of_motion=run_data["constant_of_motion"], + ) + + sim_data = physical_validation.data.SimulationData( + units=units, + dt=run_data["dt_internal"], + system=system, + ensemble=ensemble_data, + observables=observables, + ) + simulations.append(sim_data) + + result = physical_validation.integrator.convergence( + simulations, + verbose=False, + ) + + assert result < 0.5, ( + f"Integrator convergence deviation {result:.3f} exceeds threshold 0.5" + ) + + +@pytest.mark.physical_validation +def test_ensemble_check(): + """Test NVT ensemble validity via Boltzmann weight ratio at two temperatures. + + Runs NVT Langevin at 80K and 100K on a 2x2x2 Ar supercell (32 atoms), + then checks that the total energy distributions satisfy the expected + Boltzmann weight relationship. Uses total_energy=True for the ensemble + check which includes both kinetic and potential energy contributions. + """ + sim_state = _make_ar_supercell(repeat=(2, 2, 2)) + model = _make_lj_model() + + temp_low = 80.0 + temp_high = 100.0 + + run_low = _run_nvt_langevin( + sim_state, + model, + temperature=temp_low, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=42, + ) + run_high = _run_nvt_langevin( + sim_state, + model, + temperature=temp_high, + timestep_ps=0.004, + n_steps=10_000, + n_equilibration=2_000, + seed=123, + ) + + data_low = _build_nvt_simulation_data(run_low, temp_low) + data_high = _build_nvt_simulation_data(run_high, temp_high) + + quantiles = physical_validation.ensemble.check( + data_low, + data_high, + total_energy=True, + data_is_uncorrelated=True, + verbosity=0, + ) + + for i, q in enumerate(quantiles): + assert abs(q) < 3, ( + f"Ensemble quantile {i} = {q:.2f} sigma exceeds threshold" + ) From 5bf9d5958f5e18463d9589649fd3e038510f0066 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Tue, 24 Mar 2026 17:20:28 +0000 Subject: [PATCH 04/15] code for physical_validation --- fast_integrator_tests/README.md | 70 +++ fast_integrator_tests/analyze.py | 871 +++++++++++++++++++++++++++++++ fast_integrator_tests/common.py | 48 ++ fast_integrator_tests/run_npt.py | 174 ++++++ fast_integrator_tests/run_nve.py | 91 ++++ fast_integrator_tests/run_nvt.py | 116 ++++ 6 files changed, 1370 insertions(+) create mode 100644 fast_integrator_tests/README.md create mode 100644 fast_integrator_tests/analyze.py create mode 100644 fast_integrator_tests/common.py create mode 100644 fast_integrator_tests/run_npt.py create mode 100644 fast_integrator_tests/run_nve.py create mode 100644 fast_integrator_tests/run_nvt.py diff --git a/fast_integrator_tests/README.md b/fast_integrator_tests/README.md new file mode 100644 index 000000000..a60ee9446 --- /dev/null +++ b/fast_integrator_tests/README.md @@ -0,0 +1,70 @@ +# Integrator Physical Validation Tests + +Physical validation of all torch-sim MD integrators using the +[physical_validation](https://physical-validation.readthedocs.io/) library +and LJ Argon as a test system. + +## What's here + +| File | Purpose | +|---|---| +| `common.py` | Shared constants, model/structure builders | +| `run_nvt.py` | Run NVT simulations (Langevin, Nose-Hoover, V-Rescale) at 80K and 100K | +| `run_npt.py` | Run NPT simulations (Langevin, Nose-Hoover, iso/aniso C-Rescale) at 80K and 100K | +| `run_nve.py` | Run NVE simulations at 8 timesteps for convergence analysis | +| `analyze.ipynb` | Jupyter notebook with all diagnostic plots and `physical_validation` checks | +| `data/` | Output directory for `.npz` trajectory data (gitignored) | + +## Quick start + +```bash +# 1. Generate trajectory data (from this directory) +python run_nve.py +python run_nvt.py +python run_npt.py + +# Run a single integrator if needed +python run_nvt.py --integrator nvt_langevin +python run_npt.py --integrator npt_nose_hoover + +# NPT runs both a pressure and temperature sweep, so you can specify which to do: +python run_npt.py --mode temperature # Vary T at P=0 +python run_npt.py --mode pressure # Vary P at fixed T + +# 2. Open the notebook +jupyter notebook analyze.ipynb +``` + +## What the notebook shows + +### Custom plots +- **Time series**: Temperature, total energy, volume (NPT) vs step +- **KE distribution**: Observed histogram vs theoretical Gamma(Nf/2, k_BT) distribution +- **Ensemble check**: Overlapping energy distributions at two temperatures with log-ratio inset +- **NVE convergence**: Log-log RMSD vs dt plot, energy drift traces, convergence ratio table + +### physical_validation native plots +The notebook also calls `physical_validation` functions with `screen=True` to produce +their built-in diagnostic figures: +- `kinetic_energy.distribution()` — KS test or mean/width comparison plot +- `ensemble.check()` — Forward/reverse work distributions and linear fit +- `integrator.convergence()` — RMSD vs dt with reference line + +### Summary table +Final cell runs all checks programmatically and prints a PASS/FAIL table for every integrator. + +## Validation checks + +| Check | Integrators | What it tests | +|---|---|---| +| KE distribution | All NVT + NPT | Kinetic energy follows Maxwell-Boltzmann (gamma) distribution | +| Ensemble check | All NVT + NPT | Energy distributions at T=80K and T=100K satisfy Boltzmann weight ratio | +| NVE convergence | NVE (velocity Verlet) | Energy drift RMSD scales as dt^2 | + +## System details + +- **Structure**: FCC Argon (a=5.26 A), 2x2x2 supercell (32 atoms) for NVT/NPT, unit cell (4 atoms) for NVE +- **Model**: Lennard-Jones (sigma=3.405, epsilon=0.0104 eV, cutoff=8.5125 A), no neighbor list +- **Timestep**: 0.004 ps for NVT/NPT; sweep 0.002-0.010 ps for NVE +- **Production**: 10,000 steps after 2,000 (NVT) or 3,000 (NPT) equilibration steps +- **Threshold**: |deviation| < 3 sigma for all statistical checks diff --git a/fast_integrator_tests/analyze.py b/fast_integrator_tests/analyze.py new file mode 100644 index 000000000..444948d65 --- /dev/null +++ b/fast_integrator_tests/analyze.py @@ -0,0 +1,871 @@ +# %% [markdown] +# # Physical Validation of torch-sim Integrators +# +# This notebook analyzes MD trajectory data produced by `run_nvt.py`, `run_npt.py`, and `run_nve.py`. +# +# **Tests performed:** +# 1. **KE Distribution** — Does kinetic energy follow the Maxwell-Boltzmann (gamma) distribution? +# 2. **Ensemble Check** — Do energy distributions at two temperatures satisfy the Boltzmann weight relationship? +# 3. **Pressure Ensemble Check** — Do volume distributions at different pressures satisfy the expected Boltzmann weight relationship? +# 4. **NVE Convergence** — Does the energy drift RMSD scale as dt² (velocity Verlet)? +# 5. **Time Series** — Visual inspection of temperature, energy, and volume stability. + +# %% +# Create a folder to save plot +import os +os.makedirs("plots", exist_ok=True) + +# %% +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +from scipy import stats + +plt.rcParams.update({ + "figure.dpi": 120, + "axes.grid": True, + "grid.alpha": 0.3, + "figure.facecolor": "white", +}) + +DATA_DIR = Path("fast_integrator_tests/data") + +def load(name): + """Load an npz file from the data directory.""" + return dict(np.load(DATA_DIR / f"{name}.npz", allow_pickle=True)) + +# Check what data is available +available = sorted(p.stem for p in DATA_DIR.glob("*.npz")) +print("Available datasets:") +for a in available: + print(f" {a}") + +# %% +import physical_validation +from torch_sim.units import MetalUnits + +k_B_eV = float(MetalUnits.temperature) # 8.617333e-5 eV/K +THRESHOLD = 3.0 + +# def make_unit_data(): +# return physical_validation.data.UnitData( +# kb=k_B_eV, +# energy_str="eV", energy_conversion=1.0, +# length_str="Ang", length_conversion=1.0, +# volume_str="Ang^3", volume_conversion=1.0, +# temperature_str="K", temperature_conversion=1.0, +# pressure_str="eV/Ang^3", pressure_conversion=1.0, +# time_str="internal", time_conversion=1.0, +# ) +def make_unit_data(): + return physical_validation.data.UnitData( + kb=k_B_eV, + energy_str="eV", energy_conversion=96.485, # Convert to kJ/mol + length_str="Ang", length_conversion=1e-1, # Convert to nm + volume_str="Ang^3", volume_conversion=1e-3, # Convert to nm^3 + temperature_str="K", temperature_conversion=1.0, + # pressure_str="eV/Ang^3", pressure_conversion=1.6e6, # Convert to bar + pressure_str="bar", pressure_conversion=1, + time_str="fs", time_conversion=1e-3, # Convert to ps + ) + +def build_sim_data(d, temperature, ensemble="NVT", pressure=None): + """Build physical_validation.SimulationData from a loaded npz dict.""" + units = make_unit_data() + natoms = int(d["natoms"]) + system = physical_validation.data.SystemData( + natoms=natoms, nconstraints=0, + ndof_reduction_tra=3, ndof_reduction_rot=0, + mass=d["masses"], + ) + ens_kw = dict(natoms=natoms, temperature=temperature) + if ensemble == "NVT": + ens_kw["ensemble"] = "NVT" + ens_kw["volume"] = float(d.get("volume", np.mean(d.get("volumes", [0])))) + else: + ens_kw["ensemble"] = "NPT" + ens_kw["pressure"] = pressure if pressure is not None else 0.0 + + obs_kw = dict( + kinetic_energy=d["kinetic_energy"], + potential_energy=d["potential_energy"], + total_energy=d["total_energy"], + ) + if "volumes" in d: + obs_kw["volume"] = d["volumes"] + + return physical_validation.data.SimulationData( + units=units, + dt=float(d["dt_internal"]), + system=system, + ensemble=physical_validation.data.EnsembleData(**ens_kw), + observables=physical_validation.data.ObservableData(**obs_kw), + ) + +print("physical_validation helpers loaded") + +# %% [markdown] +# ## 1. NVT Time Series +# +# Temperature and energy vs step for each NVT integrator at both temperatures. + +# %% +NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] +TEMPS = [88.0, 100.0] + +fig, axes = plt.subplots(len(NVT_INTEGRATORS), 2, figsize=(14, 3.5 * len(NVT_INTEGRATORS)), + squeeze=False, sharex=True) +fig.suptitle("NVT Integrators — Time Series", fontsize=14, y=1.01) + +for row, name in enumerate(NVT_INTEGRATORS): + for temp in TEMPS: + label = f"{name}_T{temp:.0f}K" + try: + d = load(label) + except FileNotFoundError: + continue + + steps = np.arange(len(d["temperature"])) + target_T = float(d["target_temperature"]) + + # Temperature + ax = axes[row, 0] + ax.plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") + ax.axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) + ax.set_ylabel("Temperature (K)") + ax.set_title(name) + ax.legend(fontsize=8) + + # Total energy + ax = axes[row, 1] + ax.plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") + ax.set_ylabel("Total Energy (eV)") + ax.set_title(name) + ax.legend(fontsize=8) + +axes[-1, 0].set_xlabel("Step") +axes[-1, 1].set_xlabel("Step") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ## 2. KE Distribution — NVT Integrators +# +# For each integrator at 100K, compare the observed KE distribution against the theoretical gamma distribution. +# The KE of an ideal gas with $N_f$ degrees of freedom at temperature $T$ follows $\text{Gamma}(N_f/2, k_BT)$. + +# %% +fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NVT, T=100K)", fontsize=13) + +for ax, name in zip(axes, NVT_INTEGRATORS): + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + ke = d["kinetic_energy"] + natoms = int(d["natoms"]) + target_T = float(d["target_temperature"]) + + # Degrees of freedom: 3*N - 3 (COM removed) + if name == "nvt_langevin": + ndof = 3 * natoms + else: + ndof = 3 * natoms - 3 + # Theoretical gamma: shape=ndof/2, scale=k_B*T + shape = ndof / 2 + scale = k_B_eV * target_T + + # Histogram of observed KE + ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") + + # Theoretical curve + x = np.linspace(ke.min(), ke.max(), 300) + pdf = stats.gamma.pdf(x, a=shape, scale=scale) + ax.plot(x, pdf, "r-", lw=2, label="Theory") + + # Stats + d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) + d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) + ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=10) + ax.set_xlabel("Kinetic Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in KE distribution plots (NVT) +# +# Uses `physical_validation.kinetic_energy.distribution(..., screen=True)` which overlays the observed and theoretical distributions with a KS-test or mean/width comparison. + +# %% +for name in NVT_INTEGRATORS: + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — KE distribution (physical_validation)") + print(f"{'='*60}") + sd = build_sim_data(d, 100.0, ensemble="NVT") + result = physical_validation.kinetic_energy.distribution( + sd, strict=False, screen=True, verbosity=2, filename=f"plots/{name}_ke_distribution.png" + ) + print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") + +plt.show() + +# %% [markdown] +# ## 3. Ensemble Check — NVT Integrators +# +# For each integrator, compare total energy distributions at T_low=88K and T_high=100K. +# The log ratio of the energy histograms should be linear with slope $\Delta\beta = 1/k_BT_\text{low} - 1/k_BT_\text{high}$. + +# %% +fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("NVT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) + +temp_low, temp_high = 88.0, 100.0 +delta_beta = 1 / (k_B_eV * temp_low) - 1 / (k_B_eV * temp_high) + +for ax, name in zip(axes, NVT_INTEGRATORS): + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + e_lo = d_lo["total_energy"] + e_hi = d_hi["total_energy"] + + # Overlapping histogram bins + e_min = min(e_lo.min(), e_hi.min()) + e_max = max(e_lo.max(), e_hi.max()) + bins = np.linspace(e_min, e_max, 50) + bin_centers = (bins[:-1] + bins[1:]) / 2 + + h_lo, _ = np.histogram(e_lo, bins=bins, density=True) + h_hi, _ = np.histogram(e_hi, bins=bins, density=True) + + # Plot overlapping distributions + ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") + ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") + + # Inset: log ratio + mask = (h_lo > 0) & (h_hi > 0) + if mask.sum() > 2: + log_ratio = np.log(h_lo[mask] / h_hi[mask]) + bc = bin_centers[mask] + # Linear fit + slope, intercept = np.polyfit(bc, log_ratio, 1) + inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) + inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) + inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, + label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") + inset.set_xlabel("E (eV)", fontsize=7) + inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) + inset.legend(fontsize=6) + inset.tick_params(labelsize=6) + + ax.set_title(name, fontsize=10) + ax.set_xlabel("Total Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in ensemble check plots (NVT) +# +# Uses `physical_validation.ensemble.check(..., screen=True)` which shows the forward/reverse work distributions and the linear fit of the log-likelihood ratio. + +# %% +for name in NVT_INTEGRATORS: + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — Ensemble check (physical_validation)") + print(f"{'='*60}") + sd_lo = build_sim_data(d_lo, temp_low, ensemble="NVT") + sd_hi = build_sim_data(d_hi, temp_high, ensemble="NVT") + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=True, data_is_uncorrelated=True, + screen=True, verbosity=2, filename=f"plots/{name}_ensemble_check.png" + ) + print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") + +plt.show() + +# %% [markdown] +# ## 4. NPT Time Series +# +# Temperature, energy, and volume vs step for each NPT integrator. + +# %% +NPT_INTEGRATORS = ["npt_langevin", "npt_nose_hoover", "npt_isotropic_crescale", "npt_anisotropic_crescale"] + +fig, axes = plt.subplots(len(NPT_INTEGRATORS), 3, figsize=(16, 3.5 * len(NPT_INTEGRATORS)), + squeeze=False, sharex=True) +fig.suptitle("NPT Integrators — Time Series", fontsize=14, y=1.01) + +for row, name in enumerate(NPT_INTEGRATORS): + for temp in TEMPS: + label = f"{name}_T{temp:.0f}K" + try: + d = load(label) + except FileNotFoundError: + continue + + steps = np.arange(len(d["temperature"])) + target_T = float(d["target_temperature"]) + tag = f"T={target_T:.0f}K" + + # Temperature + axes[row, 0].plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=tag) + axes[row, 0].axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) + axes[row, 0].set_ylabel("Temperature (K)") + axes[row, 0].set_title(name) + axes[row, 0].legend(fontsize=7) + + # Total energy + axes[row, 1].plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=tag) + axes[row, 1].set_ylabel("Total Energy (eV)") + axes[row, 1].set_title(name) + axes[row, 1].legend(fontsize=7) + + # Volume + axes[row, 2].plot(steps, d["volumes"], alpha=0.5, lw=0.5, label=tag) + axes[row, 2].set_ylabel("Volume (ų)") + axes[row, 2].set_title(name) + axes[row, 2].legend(fontsize=7) + +for j in range(3): + axes[-1, j].set_xlabel("Step") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ## 5. KE Distribution — NPT Integrators +# +# Same gamma distribution check, but for NPT integrators at 100K. + +# %% +fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NPT, T=100K)", fontsize=13) + +for ax, name in zip(axes, NPT_INTEGRATORS): + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + ke = d["kinetic_energy"] + natoms = int(d["natoms"]) + target_T = float(d["target_temperature"]) + + ndof = 3 * natoms - 3 + shape = ndof / 2 + scale = k_B_eV * target_T + + ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") + + x = np.linspace(ke.min(), ke.max(), 300) + pdf = stats.gamma.pdf(x, a=shape, scale=scale) + ax.plot(x, pdf, "r-", lw=2, label="Theory") + + d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) + d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) + ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=9) + ax.set_xlabel("Kinetic Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in KE distribution plots (NPT) + +# %% +for name in NPT_INTEGRATORS: + label = f"{name}_T100K" + try: + d = load(label) + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — KE distribution (physical_validation)") + print(f"{'='*60}") + sd = build_sim_data(d, 100.0, ensemble="NPT") + result = physical_validation.kinetic_energy.distribution( + sd, strict=False, screen=True, verbosity=2, + ) + print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") + +plt.show() + +# %% [markdown] +# ## 6. Ensemble Check — NPT Integrators +# +# Same Boltzmann weight ratio check at T=88K vs T=100K for NPT integrators. + +# %% +fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), + sharey=True) +fig.suptitle("NPT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) + +for ax, name in zip(axes, NPT_INTEGRATORS): + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + ax.set_title(f"{name}\n(no data)") + continue + + e_lo = d_lo["total_energy"] + e_hi = d_hi["total_energy"] + + e_min = min(e_lo.min(), e_hi.min()) + e_max = max(e_lo.max(), e_hi.max()) + bins = np.linspace(e_min, e_max, 50) + bin_centers = (bins[:-1] + bins[1:]) / 2 + + h_lo, _ = np.histogram(e_lo, bins=bins, density=True) + h_hi, _ = np.histogram(e_hi, bins=bins, density=True) + + ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") + ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") + + mask = (h_lo > 0) & (h_hi > 0) + if mask.sum() > 2: + log_ratio = np.log(h_lo[mask] / h_hi[mask]) + bc = bin_centers[mask] + slope, intercept = np.polyfit(bc, log_ratio, 1) + inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) + inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) + inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, + label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") + inset.set_xlabel("E (eV)", fontsize=7) + inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) + inset.legend(fontsize=6) + inset.tick_params(labelsize=6) + + ax.set_title(name, fontsize=9) + ax.set_xlabel("Total Energy (eV)") + ax.legend(fontsize=8) + +axes[0].set_ylabel("Probability Density") +fig.tight_layout() +plt.show() + +# %% [markdown] +# ### physical_validation built-in ensemble check plots (NPT) + +# %% +for name in NPT_INTEGRATORS: + try: + d_lo = load(f"{name}_T{temp_low:.0f}K") + d_hi = load(f"{name}_T{temp_high:.0f}K") + except FileNotFoundError: + print(f"{name}: no data") + continue + + print(f"\n{'='*60}") + print(f" {name} — Ensemble check (physical_validation)") + print(f"{'='*60}") + sd_lo = build_sim_data(d_lo, temp_low, ensemble="NPT") + sd_hi = build_sim_data(d_hi, temp_high, ensemble="NPT") + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, # data_is_uncorrelated=True, + screen=True, verbosity=2, + ) + print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: #ConvergenceError + print(f" ConvergenceError: {e}") + +plt.show() + +# %% [markdown] +# ## 6b. NPT Pressure Ensemble Check +# +# Compare volume distributions at two pressures (same temperature). +# For correct NPT sampling, `physical_validation` checks that volume distributions +# at different pressures satisfy the expected Boltzmann weight relationship. +# This is a 1D check on volumes only — no energy data needed. + +# %% +from torch_sim.units import MetalUnits + +P_BAR_CONVERSION = float(MetalUnits.pressure) # 1 bar in eV/Ang^3 + +# Discover available pressure-sweep data files +pressure_files = sorted(DATA_DIR.glob("*_P*bar.npz")) +print("Pressure-sweep data files:") +for f in pressure_files: + print(f" {f.name}") + +# Parse integrator -> {pressure_bar: filename} mapping +pressure_data = {} # integrator_name -> [(p_bar, filename), ...] +for f in pressure_files: + stem = f.stem # e.g. npt_langevin_T100K_P0bar + parts = stem.rsplit("_P", 1) + if len(parts) != 2: + continue + prefix = parts[0] # e.g. npt_langevin_T100K + p_bar = float(parts[1].replace("bar", "")) + int_prefix = prefix.rsplit("_T", 1)[0] # e.g. npt_langevin + pressure_data.setdefault(int_prefix, []).append((p_bar, f.stem)) + +for name, entries in pressure_data.items(): + entries.sort(key=lambda x: x[0]) + print(f"\n{name}: {[(p, fn) for p, fn in entries]}") + +# %% +# Volume distributions at different pressures for each NPT integrator +integrators_with_pressure = [n for n in NPT_INTEGRATORS if n in pressure_data] + +if integrators_with_pressure: + fig, axes = plt.subplots(1, len(integrators_with_pressure), + figsize=(5 * len(integrators_with_pressure), 4), squeeze=False) + axes = axes[0] + fig.suptitle("NPT Pressure Check \u2014 Volume Distributions at Same T, Different P", fontsize=13) + + for ax, name in zip(axes, integrators_with_pressure): + entries = pressure_data[name] + colors = plt.cm.coolwarm(np.linspace(0, 1, len(entries))) + for (p_bar, fname), color in zip(entries, colors): + d = load(fname) + vols = d["volumes"] + ax.hist(vols, bins=50, density=True, alpha=0.5, color=color, + label=f"P={p_bar:.0f} bar") + ax.axvline(vols.mean(), color=color, ls="--", lw=1, alpha=0.7) + ax.set_title(name, fontsize=10) + ax.set_xlabel("Volume (\u00c5\u00b3)") + ax.legend(fontsize=8) + + axes[0].set_ylabel("Probability Density") + fig.tight_layout() + plt.savefig("plots/npt_pressure_volume_distributions.png", dpi=150, bbox_inches="tight") + plt.show() +else: + print("No pressure-sweep data found. Run: python run_npt.py --mode pressure") + +# %% [markdown] +# ### physical_validation built-in pressure ensemble check +# +# Uses `physical_validation.ensemble.check()` with two simulations at the same temperature +# but different pressures. The library auto-detects this case and performs a 1D volume-based check. + +# %% +for name in integrators_with_pressure: + entries = pressure_data[name] + if len(entries) < 2: + print(f"{name}: need at least 2 pressures, got {len(entries)}") + continue + + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_lo = load(fname_lo) + d_hi = load(fname_hi) + + p_lo_eva3 = float(d_lo["external_pressure"]) + p_hi_eva3 = float(d_hi["external_pressure"]) + temp = float(d_lo["target_temperature"]) + print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " + f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") + + print(f"\n{'='*60}") + print(f" {name} \u2014 Pressure ensemble check") + print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") + print(f"{'='*60}") + + # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) + # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) + + sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) + sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) + + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, #data_is_uncorrelated=True, + screen=True, verbosity=2, + filename=f"plots/{name}_pressure_ensemble_check.png", + ) + print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: + print(f" Error: {e}") + +plt.show() + +# %% +for name in integrators_with_pressure: + entries = pressure_data[name] + if len(entries) < 2: + print(f"{name}: need at least 2 pressures, got {len(entries)}") + continue + + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_lo = load(fname_lo) + d_hi = load(fname_hi) + + p_lo_eva3 = float(d_lo["external_pressure"]) + p_hi_eva3 = float(d_hi["external_pressure"]) + temp = float(d_lo["target_temperature"]) + print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " + f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") + + print(f"\n{'='*60}") + print(f" {name} \u2014 Pressure ensemble check") + print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") + print(f"{'='*60}") + + # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) + # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) + + sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) + sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) + + try: + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, + total_energy=False, data_is_uncorrelated=True, + screen=True, verbosity=2, + filename=f"plots/{name}_pressure_ensemble_check.png", + ) + print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") + except Exception as e: + print(f" Error: {e}") + +plt.show() + +# %% [markdown] +# ## 8. NVE Integrator Convergence +# +# RMSD of the conserved quantity (total energy) vs timestep on a log-log scale. +# For velocity Verlet, the RMSD should scale as $\text{dt}^2$ (slope = 2 on log-log). + +# %% +try: + nve = load("nve_convergence") + timesteps_ps = nve["timesteps_ps"] + + dts, rmsds, drifts = [], [], [] + for dt_ps in timesteps_ps: + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + dts.append(dt_ps) + rmsds.append(com.std()) + drifts.append(com[-1] - com[0]) + + dts = np.array(dts) + rmsds = np.array(rmsds) + drifts = np.array(drifts) + + fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 5)) + fig.suptitle("NVE Integrator Convergence (4-atom Ar, T=5K)", fontsize=13) + + # --- Panel 1: RMSD vs dt (log-log) --- + ax1.loglog(dts, rmsds, "ko-", ms=8, lw=2, label="RMSD(E_tot)") + + # dt^2 reference line through middle point + mid = len(dts) // 2 + ref = rmsds[mid] * (dts / dts[mid]) ** 2 + ax1.loglog(dts, ref, "r--", lw=1.5, alpha=0.7, label="$\\propto dt^2$ (reference)") + + # Fit slope + log_dts = np.log(dts) + log_rmsds = np.log(rmsds) + # Only fit points where RMSD is clearly above noise floor + mask = rmsds > 1.5 * rmsds.min() + if mask.sum() >= 2: + slope, intercept = np.polyfit(log_dts[mask], log_rmsds[mask], 1) + ax1.set_title(f"Log-log slope = {slope:.2f} (expected: 2.0)") + else: + ax1.set_title("RMSD vs dt") + + ax1.set_xlabel("Timestep (ps)") + ax1.set_ylabel("RMSD of E_total (eV)") + ax1.legend() + + # --- Panel 2: Energy time series for each dt --- + for i, dt_ps in enumerate(timesteps_ps): + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + steps = np.arange(len(com)) + ax2.plot(steps, com - com[0], alpha=0.7, lw=0.8, label=f"dt={dt_ps:.4f} ps") + + ax2.set_xlabel("Step") + ax2.set_ylabel("E_total - E_total[0] (eV)") + ax2.set_title("Energy drift per timestep") + ax2.legend(fontsize=7, ncol=2) + + # --- Panel 3: RMSD ratio table --- + ax3.axis("off") + headers = ["dt1 (ps)", "dt2 (ps)", "dt1/dt2", "RMSD1/RMSD2", "(dt1/dt2)²"] + rows = [] + for i in range(len(dts) - 1): + dt_ratio = dts[i] / dts[i + 1] + rmsd_ratio = rmsds[i] / rmsds[i + 1] + expected = dt_ratio ** 2 + rows.append([f"{dts[i]:.4f}", f"{dts[i+1]:.4f}", + f"{dt_ratio:.3f}", f"{rmsd_ratio:.3f}", f"{expected:.3f}"]) + + table = ax3.table(cellText=rows, colLabels=headers, loc="center", cellLoc="center") + table.auto_set_font_size(False) + table.set_fontsize(9) + table.scale(1.0, 1.5) + ax3.set_title("Convergence ratios", pad=20) + + fig.tight_layout() + plt.show() + +except FileNotFoundError: + print("No NVE data found. Run: python run_nve.py") + +# %% [markdown] +# ### physical_validation built-in integrator convergence plot +# +# Uses `physical_validation.integrator.convergence(..., screen=True)` which plots RMSD vs dt with the expected dt^2 reference line. + +# %% +try: + nve = load("nve_convergence") + timesteps_ps = nve["timesteps_ps"] + natoms = int(nve["natoms"]) + masses = nve["masses"] + volume = float(nve["volume"]) + + # Pick 3 timesteps in the dt^2 regime (avoid noise floor and nonlinear regime) + # Use [0.007, 0.005, 0.004] which showed good convergence + selected_dts = [0.007, 0.005, 0.004] + + simulations = [] + for dt_ps in selected_dts: + key = f"dt_{dt_ps}" + com = nve[f"{key}_constant_of_motion"] + dt_internal = float(nve[f"{key}_dt_internal"]) + + system = physical_validation.data.SystemData( + natoms=natoms, nconstraints=0, + ndof_reduction_tra=3, ndof_reduction_rot=0, mass=masses) + ensemble = physical_validation.data.EnsembleData( + ensemble="NVE", natoms=natoms, volume=volume) + obs = physical_validation.data.ObservableData(constant_of_motion=com) + sd = physical_validation.data.SimulationData( + units=make_unit_data(), dt=dt_internal, + system=system, ensemble=ensemble, observables=obs) + simulations.append(sd) + + result = physical_validation.integrator.convergence( + simulations, verbose=True, screen=True, + ) + print(f"\nConvergence deviation: {result:.3f} ({'PASS' if result < 0.5 else 'FAIL'})") + +except FileNotFoundError: + print("No NVE data found. Run: python run_nve.py") + +plt.show() + +# %% [markdown] +# ## 9. Summary Table +# +# Run `physical_validation` checks programmatically on all available data and collect pass/fail results. + +# %% +# Collect results +results = [] +all_integrators = NVT_INTEGRATORS + NPT_INTEGRATORS + +for name in all_integrators: + is_npt = name.startswith("npt") + ensemble = "NPT" if is_npt else "NVT" + + # --- KE Distribution at 100K --- + label_100 = f"{name}_T100K" + ke_status = "no data" + try: + d100 = load(label_100) + sd = build_sim_data(d100, 100.0, ensemble=ensemble) + d_mean, d_width = physical_validation.kinetic_energy.distribution( + sd, strict=False, verbosity=0) + ke_pass = abs(d_mean) < THRESHOLD and abs(d_width) < THRESHOLD + ke_status = f"PASS (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" if ke_pass else f"FAIL (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" + except Exception as e: + ke_status = f"ERROR: {e}" + + # --- Ensemble Check 88K vs 100K (temperature) --- + ens_status = "no data" + try: + d88 = load(f"{name}_T88K") + d100 = load(f"{name}_T100K") + sd_lo = build_sim_data(d88, 88.0, ensemble=ensemble) + sd_hi = build_sim_data(d100, 100.0, ensemble=ensemble) + quantiles = physical_validation.ensemble.check( + sd_lo, sd_hi, total_energy=True, data_is_uncorrelated=True, verbosity=0) + max_q = max(abs(q) for q in quantiles) + ens_pass = max_q < THRESHOLD + ens_status = f"PASS (max |q|={max_q:.2f}\u03c3)" if ens_pass else f"FAIL (max |q|={max_q:.2f}\u03c3)" + except Exception as e: + ens_status = f"ERROR: {e}" + + # --- Pressure Ensemble Check (NPT only, same T, different P) --- + press_status = "N/A" + if is_npt and name in pressure_data: + entries = pressure_data[name] + if len(entries) >= 2: + try: + p_lo_bar, fname_lo = entries[0] + p_hi_bar, fname_hi = entries[-1] + d_plo = load(fname_lo) + d_phi = load(fname_hi) + p_lo_eva3 = float(d_plo["external_pressure"]) + p_hi_eva3 = float(d_phi["external_pressure"]) + temp_p = float(d_plo["target_temperature"]) + sd_plo = build_sim_data(d_plo, temp_p, ensemble="NPT", pressure=p_lo_eva3) + sd_phi = build_sim_data(d_phi, temp_p, ensemble="NPT", pressure=p_hi_eva3) + quantiles_p = physical_validation.ensemble.check( + sd_plo, sd_phi, total_energy=False, data_is_uncorrelated=True, verbosity=0) + max_qp = max(abs(q) for q in quantiles_p) + press_pass = max_qp < THRESHOLD + press_status = f"PASS (max |q|={max_qp:.2f}\u03c3)" if press_pass else f"FAIL (max |q|={max_qp:.2f}\u03c3)" + except Exception as e: + press_status = f"ERROR: {e}" + else: + press_status = "need 2 pressures" + + results.append((name, ke_status, ens_status, press_status)) + +# Print table +print(f"{'Integrator':<30} {'KE Distribution':<40} {'Ensemble (T)':<35} {'Ensemble (P)':<35}") +print("-" * 140) +for name, ke, ens, press in results: + print(f"{name:<30} {ke:<40} {ens:<35} {press:<35}") + + diff --git a/fast_integrator_tests/common.py b/fast_integrator_tests/common.py new file mode 100644 index 000000000..69e21fd06 --- /dev/null +++ b/fast_integrator_tests/common.py @@ -0,0 +1,48 @@ +"""Shared constants and helpers for integrator validation scripts.""" + +import numpy as np +import torch +from ase.build import bulk +from pathlib import Path + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.units import MetalUnits + +DEVICE = torch.device("cpu") +# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +DTYPE = torch.float64 +print(f"Using device: {DEVICE}") + +# LJ Argon +SIGMA = 3.405 +EPSILON = 0.0104 +CUTOFF = 2.5 * SIGMA + +DATA_DIR = Path(__file__).parent / "data" +torch.set_num_threads(4) + +def make_lj_model(compute_stress=False): + return LennardJonesModel( + use_neighbor_list=False, + sigma=SIGMA, + epsilon=EPSILON, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=compute_stress, + cutoff=CUTOFF, + ) + + +def make_ar_supercell(repeat=(2, 2, 2)): + atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) + return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) + + +def to_kT(temperature_K): + return temperature_K * float(MetalUnits.temperature) + + +def to_dt(timestep_ps): + return timestep_ps * float(MetalUnits.time) diff --git a/fast_integrator_tests/run_npt.py b/fast_integrator_tests/run_npt.py new file mode 100644 index 000000000..0b18f80ca --- /dev/null +++ b/fast_integrator_tests/run_npt.py @@ -0,0 +1,174 @@ +"""Run NPT simulations for all NPT integrators and save observables. + +Usage: + python run_npt.py # temperature sweep (default) + python run_npt.py --mode pressure # pressure sweep at fixed T + python run_npt.py --mode all # both sweeps + python run_npt.py --integrator npt_langevin +""" + +import argparse +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, DEVICE, DTYPE, make_ar_supercell, make_lj_model, to_dt, to_kT +from torch_sim.units import MetalUnits + +NPT_INTEGRATORS = [ + "npt_langevin", + "npt_nose_hoover", + "npt_isotropic_crescale", + "npt_anisotropic_crescale", +] + +TEMPERATURES = [88.0, 100.0] +TIMESTEP_PS = 0.004 +EXTERNAL_PRESSURE = 0.0 +N_STEPS = 20_000 +N_EQUILIBRATION = 3_000 + +# Pressure sweep: two pressures at a fixed temperature. +# physical_validation compares volume distributions at same T, different P. +PRESSURE_SWEEP_TEMP = 100.0 # K +PRESSURES_EVA3 = [0.0, 0.0001] # eV/ų (0 bar and ~160 bar) + + +def run_npt(integrator_name, sim_state, model, temperature, external_pressure=0.0, + seed=42): + kT = to_kT(temperature) + dt = torch.tensor(to_dt(TIMESTEP_PS), device=DEVICE, dtype=DTYPE) + ext_p = torch.tensor(external_pressure, device=DEVICE, dtype=DTYPE) + natoms = int(sim_state.positions.shape[0]) + + sim_state = sim_state.clone() + sim_state.rng = seed + + # Initialize + if integrator_name == "npt_langevin": + # state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt) + state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt, b_tau = 1 * dt, alpha= 5 * dt) # Better parameters + elif integrator_name == "npt_nose_hoover": + state = ts.npt_nose_hoover_init(sim_state, model, kT=kT, dt=dt, t_tau=10 * dt, b_tau=100 * dt) + elif integrator_name == "npt_isotropic_crescale": + state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt, tau_p = 10 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) + elif integrator_name == "npt_anisotropic_crescale": + state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt/2, tau_p = 100 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) + else: + raise ValueError(f"Unknown: {integrator_name}") + + def step(s): + if integrator_name == "npt_langevin": + return ts.npt_langevin_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) + if integrator_name == "npt_nose_hoover": + return ts.npt_nose_hoover_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) + if integrator_name == "npt_isotropic_crescale": + return ts.npt_crescale_isotropic_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, tau = 5 * dt + ) + return ts.npt_crescale_anisotropic_step( + s, model, dt=dt/2, kT=kT, external_pressure=ext_p, tau = 5 * dt + ) + + # Equilibration + print(f" Equilibrating {N_EQUILIBRATION} steps...") + for _ in range(N_EQUILIBRATION): + state = step(state) + + # Production + print(f" Producing {N_STEPS} steps...") + ke_list, pe_list, total_e_list = [], [], [] + temp_list, volume_list = [], [] + + for i in range(N_STEPS): + state = step(state) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + temp = float( + ts.calc_temperature(masses=state.masses, momenta=state.momenta) + ) + cell = state.cell[0].detach().cpu().numpy() + vol = float(np.abs(np.linalg.det(cell))) + + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + temp_list.append(temp) + volume_list.append(vol) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "temperature": np.array(temp_list), + "volumes": np.array(volume_list), + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": to_dt(TIMESTEP_PS), + "natoms": natoms, + "target_temperature": temperature, + "external_pressure": external_pressure, + "timestep_ps": TIMESTEP_PS, + "integrator": integrator_name, + } + + +def pressure_to_bar(p_eva3): + """Convert eV/ų to bar for display.""" + return p_eva3 / float(MetalUnits.pressure) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--integrator", choices=NPT_INTEGRATORS, default=None) + parser.add_argument( + "--mode", choices=["temperature", "pressure", "all"], default="all", + help="'temperature' = vary T at P=0 (default), " + "'pressure' = vary P at fixed T, 'all' = both", + ) + args = parser.parse_args() + + integrators = [args.integrator] if args.integrator else NPT_INTEGRATORS + DATA_DIR.mkdir(exist_ok=True) + + sim_state = make_ar_supercell(repeat=(3, 3, 3)) + model = make_lj_model(compute_stress=True) + + # --- Temperature sweep (same P=0, vary T) --- + if args.mode in ("temperature", "all"): + print("=== Temperature sweep ===") + for name in integrators: + for temp in TEMPERATURES: + seed = 42 if temp == TEMPERATURES[0] else 123 + label = f"{name}_T{temp:.0f}K" + print(f" Running {label} ...") + t0 = time.time() + data = run_npt(name, sim_state, model, temp, seed=seed) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + # --- Pressure sweep (same T, vary P) --- + if args.mode in ("pressure", "all"): + print(f"\n=== Pressure sweep at T={PRESSURE_SWEEP_TEMP:.0f}K ===") + for name in integrators: + for i, p_eva3 in enumerate(PRESSURES_EVA3): + p_bar = pressure_to_bar(p_eva3) + seed = 42 if i == 0 else 123 + label = f"{name}_T{PRESSURE_SWEEP_TEMP:.0f}K_P{p_bar:.0f}bar" + print(f" Running {label} ...") + t0 = time.time() + data = run_npt( + name, sim_state, model, PRESSURE_SWEEP_TEMP, + external_pressure=p_eva3, seed=seed, + ) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + +if __name__ == "__main__": + main() diff --git a/fast_integrator_tests/run_nve.py b/fast_integrator_tests/run_nve.py new file mode 100644 index 000000000..a0d60ca9b --- /dev/null +++ b/fast_integrator_tests/run_nve.py @@ -0,0 +1,91 @@ +"""Run NVE simulations at multiple timesteps for convergence analysis. + +Usage: + python run_nve.py +""" + +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT + +# Timesteps: sweep a range so the notebook can pick the best subset +TIMESTEPS_PS = [0.010, 0.008, 0.007, 0.006, 0.005, 0.004, 0.003, 0.002] +TEMPERATURE = 5.0 # K - low so integration error dominates +N_STEPS = 10_000 +SEED = 42 + + +def run_nve(sim_state, model, kT_init, timestep_ps): + dt = to_dt(timestep_ps) + + sim_state = sim_state.clone() + sim_state.rng = SEED + state = ts.nve_init(sim_state, model, kT=kT_init) + + ke_list, pe_list, com_list = [], [], [] + + for _ in range(N_STEPS): + state = ts.nve_step(state, model, dt=dt) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + ke_list.append(ke) + pe_list.append(pe) + com_list.append(ke + pe) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "constant_of_motion": np.array(com_list), + "dt_internal": dt, + "timestep_ps": timestep_ps, + } + + +def main(): + DATA_DIR.mkdir(exist_ok=True) + + # sim_state = make_ar_supercell(repeat=(1, 1, 1)) + sim_state = make_ar_supercell(repeat=(6, 6, 6)) + model = make_lj_model() + kT_init = to_kT(TEMPERATURE) + natoms = int(sim_state.positions.shape[0]) + masses = sim_state.masses.detach().cpu().numpy() + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + all_results = {} + for dt_ps in TIMESTEPS_PS: + label = f"nve_dt{dt_ps:.4f}" + print(f" Running {label} ...") + t0 = time.time() + data = run_nve(sim_state, model, kT_init, dt_ps) + elapsed = time.time() - t0 + all_results[f"dt_{dt_ps}"] = data + print(f" std(E_tot) = {data['constant_of_motion'].std():.3e} ({elapsed:.1f}s)") + + # Save everything into one npz + save_dict = { + "timesteps_ps": np.array(TIMESTEPS_PS), + "temperature": TEMPERATURE, + "natoms": natoms, + "masses": masses, + "volume": volume, + "n_steps": N_STEPS, + } + for dt_ps in TIMESTEPS_PS: + key = f"dt_{dt_ps}" + for field in ("constant_of_motion", "kinetic_energy", "potential_energy"): + save_dict[f"{key}_{field}"] = all_results[key][field] + save_dict[f"{key}_dt_internal"] = all_results[key]["dt_internal"] + + outpath = DATA_DIR / "nve_convergence.npz" + np.savez(outpath, **save_dict) + print(f"\n Saved {outpath.name}") + + +if __name__ == "__main__": + main() diff --git a/fast_integrator_tests/run_nvt.py b/fast_integrator_tests/run_nvt.py new file mode 100644 index 000000000..ba8d141cf --- /dev/null +++ b/fast_integrator_tests/run_nvt.py @@ -0,0 +1,116 @@ +"""Run NVT simulations for all NVT integrators and save observables. + +Usage: + python run_nvt.py + python run_nvt.py --integrator nvt_langevin +""" + +import argparse +import time + +import numpy as np +import torch + +import torch_sim as ts +from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT + +NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] + +# Two temperatures for ensemble check +TEMPERATURES = [95.0, 100.0] +TIMESTEP_PS = 0.004 +N_STEPS = 10_000 +N_EQUILIBRATION = 2_000 + + +def run_nvt(integrator_name, sim_state, model, temperature, seed=42): + kT = to_kT(temperature) + dt = to_dt(TIMESTEP_PS) + natoms = int(sim_state.positions.shape[0]) + + sim_state = sim_state.clone() + sim_state.rng = seed + + # Initialize + if integrator_name == "nvt_langevin": + state = ts.nvt_langevin_init(sim_state, model, kT=kT) + elif integrator_name == "nvt_nose_hoover": + state = ts.nvt_nose_hoover_init(sim_state, model, kT=kT, dt=dt) + elif integrator_name == "nvt_vrescale": + state = ts.nvt_vrescale_init(sim_state, model, kT=kT) + else: + raise ValueError(f"Unknown: {integrator_name}") + + def step(s): + if integrator_name == "nvt_langevin": + return ts.nvt_langevin_step(s, model, dt=dt, kT=kT) + if integrator_name == "nvt_nose_hoover": + return ts.nvt_nose_hoover_step(s, model, dt=dt, kT=kT) + return ts.nvt_vrescale_step(model, s, dt=dt, kT=kT) + + # Equilibration + print(f" Equilibrating {N_EQUILIBRATION} steps...") + for _ in range(N_EQUILIBRATION): + state = step(state) + + # Production + print(f" Producing {N_STEPS} steps...") + ke_list, pe_list, total_e_list = [], [], [] + temp_list = [] + + for i in range(N_STEPS): + state = step(state) + ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) + pe = float(state.energy.sum()) + temp = float( + ts.calc_temperature(masses=state.masses, momenta=state.momenta) + ) + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + temp_list.append(temp) + + cell = sim_state.cell[0].detach().cpu().numpy() + volume = float(np.abs(np.linalg.det(cell))) + + return { + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "temperature": np.array(temp_list), + "volume": volume, + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": to_dt(TIMESTEP_PS), + "natoms": natoms, + "target_temperature": temperature, + "timestep_ps": TIMESTEP_PS, + "integrator": integrator_name, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--integrator", choices=NVT_INTEGRATORS, default=None) + args = parser.parse_args() + + integrators = [args.integrator] if args.integrator else NVT_INTEGRATORS + DATA_DIR.mkdir(exist_ok=True) + + sim_state = make_ar_supercell(repeat=(6, 6, 6)) + model = make_lj_model() + + for name in integrators: + for temp in TEMPERATURES: + seed = 42 if temp == TEMPERATURES[0] else 123 + label = f"{name}_T{temp:.0f}K" + print(f" Running {label} ...") + t0 = time.time() + data = run_nvt(name, sim_state, model, temp, seed=seed) + elapsed = time.time() - t0 + outpath = DATA_DIR / f"{label}.npz" + np.savez(outpath, **data) + print(f" Saved {outpath.name} ({elapsed:.1f}s)") + + +if __name__ == "__main__": + main() From 275bf529c8516b592dfa989efa860beae0bbbd24 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 2 Apr 2026 15:39:52 +0000 Subject: [PATCH 05/15] fix docs layout --- torch_sim/integrators/__init__.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py index 950f16644..fb18a7960 100644 --- a/torch_sim/integrators/__init__.py +++ b/torch_sim/integrators/__init__.py @@ -17,33 +17,40 @@ - Langevin barostat integrator :func:`npt.npt_langevin_step` [4, 5] - Nosé-Hoover barostat integrator :func:`npt.npt_nose_hoover_step` from [10] - Isotropic C-Rescale barostat integrator :func:`npt.npt_crescale_isotropic_step` - from [6, 8, 9] + from [6, 8, 9] - C-Rescale barostat integrator :func:`npt.npt_crescale_anisotropic_step` - from [7, 8, 9]. Available implementations include isotropic and - anisotropic cell rescaling, allowing to change cell lengths, and potentially angles - as well. + from [7, 8, 9]. Anisotropic NPT allows to change cell lengths as well as angles. References: [1] Bussi G, Donadio D, Parrinello M. "Canonical sampling through velocity rescaling." The Journal of chemical physics, 126(1), 014101 (2007). + [2] Leimkuhler B, Matthews C.2016 Efficient molecular dynamics using geodesic integration and solvent-solute splitting. Proc. R. Soc. A 472: 20160138 + [3] Martyna, G. J., Tuckerman, M. E., Tobias, D. J., & Klein, M. L. (1996). Explicit reversible integrators for extended systems dynamics. Molecular Physics, 87(5), 1117-1157. + [4] Grønbech-Jensen, N., & Farago, O. (2014). Constant pressure and temperature discrete-time Langevin molecular dynamics. The Journal of chemical physics, 141(19). + [5] LAMMPS: https://docs.lammps.org/fix_press_langevin.html + [6] Bernetti, Mattia, and Giovanni Bussi. "Pressure control using stochastic cell rescaling." The Journal of Chemical Physics 153.11 (2020). + [7] Del Tatto, Vittorio, et al. "Molecular dynamics of solids at constant pressure and stress using anisotropic stochastic cell rescaling." Applied Sciences 12.3 (2022): 1139. + [8] Bussi Anisotropic C-Rescale SimpleMD implementation: https://github.com/bussilab/crescale/blob/master/simplemd_anisotropic/simplemd.cpp + [9] Supplementary Information for [6]. + [10]Tuckerman, Mark E., et al. "A Liouville-operator derived measure-preserving integrator for molecular dynamics simulations in the isothermal-isobaric ensemble." Journal of Physics A: Mathematical and General 39.19 (2006): 5629-5651. From 45d595337e24bd7e471c593a0f8cf651520130aa Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 2 Apr 2026 15:41:49 +0000 Subject: [PATCH 06/15] remove args in docstrings that are not anymore in function args --- torch_sim/runners.py | 1 - torch_sim/trajectory.py | 6 ------ 2 files changed, 7 deletions(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 72c25c070..e2028f801 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -751,7 +751,6 @@ def static( Args: system (StateLike): Input system to calculate properties for model (ModelInterface): Neural network model module - unit_system (UnitSystem): Unit system for energy and forces trajectory_reporter (TrajectoryReporter | dict | None): Optional reporter for tracking trajectory. If a dict, will be passed to the TrajectoryReporter constructor and must include at least the "filenames" key. Any prop diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py index 5a73d5802..60c3a2f2e 100644 --- a/torch_sim/trajectory.py +++ b/torch_sim/trajectory.py @@ -283,9 +283,6 @@ def report( model (ModelInterface, optional): Model used for simulation. Defaults to None. Must be provided if any prop_calculators are provided. - write_to_file (bool, optional): Whether to write the state to the trajectory - files. Defaults to True. Should only be set to `False` if the props - are being collected separately. Returns: list[dict[str, torch.Tensor]]: Map of property names to tensors for each @@ -835,9 +832,6 @@ def get_steps( Args: name (str): Name of the array - start (int, optional): Starting frame index. Defaults to None. - stop (int, optional): Ending frame index (exclusive). Defaults to None. - step (int, optional): Step size between frames. Defaults to 1. Returns: np.ndarray: Array of step numbers with shape [n_selected_frames] From 64838da1ed7a716287633e4d0a17ccc95805fe19 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 2 Apr 2026 15:43:25 +0000 Subject: [PATCH 07/15] docstrings remove removed args --- torch_sim/workflows/a2c.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torch_sim/workflows/a2c.py b/torch_sim/workflows/a2c.py index 81d22078c..1a360da14 100644 --- a/torch_sim/workflows/a2c.py +++ b/torch_sim/workflows/a2c.py @@ -241,7 +241,6 @@ def random_packed_structure( position when computing minimum distances. device: PyTorch device for calculations (CPU/GPU). dtype: PyTorch data type for numerical precision. - log: List to store positions at each iteration. Returns: FIREState: The optimized structure state containing positions, forces, From fd3e295ddaf1de79b19afb764a86edd949fb47e9 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 2 Apr 2026 15:44:39 +0000 Subject: [PATCH 08/15] - update docs to include equations - change default values corresponding to args used in tests with LJ argon - correct severe errors in equations (npt_langevin, npt_anisotropic_crescale) --- torch_sim/integrators/npt.py | 465 +++++++++++++++++++++++++---------- torch_sim/integrators/nve.py | 39 ++- torch_sim/integrators/nvt.py | 212 ++++++++++++---- 3 files changed, 528 insertions(+), 188 deletions(-) diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index c236bd452..15c591c9a 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -21,6 +21,7 @@ from torch_sim.integrators.nvt import _vrescale_update from torch_sim.models.interface import ModelInterface from torch_sim.state import SimState +from torch_sim.units import MetalUnits logger = logging.getLogger(__name__) @@ -121,8 +122,6 @@ def _npt_langevin_beta( Args: state (NPTLangevinState): Current NPT state - alpha (torch.Tensor): Friction coefficient, either scalar or - shape [n_systems] kT (torch.Tensor): Temperature in energy units, either scalar or shape [n_systems] dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] @@ -161,13 +160,9 @@ def _npt_langevin_cell_beta( Args: state (NPTLangevinState): Current NPT state - cell_alpha (torch.Tensor): Cell friction coefficient, either scalar or - with shape [n_systems] kT (torch.Tensor): System temperature in energy units, either scalar or with shape [n_systems] dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] - device (torch.device): Device for tensor operations - dtype (torch.dtype): Data type for tensor operations Returns: torch.Tensor: Scaled random noise for cell dynamics with shape @@ -210,8 +205,6 @@ def _npt_langevin_cell_position_step( [n_systems, n_dim, n_dim] kT (torch.Tensor): Target temperature in energy units, either scalar or with shape [n_systems] - cell_alpha (torch.Tensor): Cell friction coefficient, either scalar or - with shape [n_systems] Returns: NPTLangevinState: Updated state with new cell positions @@ -265,8 +258,6 @@ def _npt_langevin_cell_velocity_step( dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] pressure_force (torch.Tensor): Final pressure force shape [n_systems, n_dim, n_dim] - cell_alpha (torch.Tensor): Cell friction coefficient, either scalar or - shape [n_systems] kT (torch.Tensor): Temperature in energy units, either scalar or shape [n_systems] @@ -287,10 +278,10 @@ def _npt_langevin_cell_velocity_step( cell_masses_expanded = state.cell_masses.view(-1, 1, 1) # shape: (n_systems, 1, 1) # These factors come from the Langevin integration scheme - a = (1 - (cell_alpha_expanded * dt_expanded) / cell_masses_expanded) / ( - 1 + (cell_alpha_expanded * dt_expanded) / cell_masses_expanded + a = (1 - (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded)) / ( + 1 + (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded) ) - b = 1 / (1 + (cell_alpha_expanded * dt_expanded) / cell_masses_expanded) + b = 1 / (1 + (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded)) # Calculate the three terms for velocity update # a will broadcast from (n_systems, 1, 1) to (n_systems, 3, 3) @@ -306,7 +297,7 @@ def _npt_langevin_cell_velocity_step( noise_prefactor = torch.sqrt( 2 * cell_alpha_expanded * kT.view(-1, 1, 1) * dt_expanded ) - noise_term = noise_prefactor * cell_noise / torch.sqrt(cell_masses_expanded) + noise_term = noise_prefactor * cell_noise / cell_masses_expanded # Random noise contribution c_3 = b * noise_term @@ -335,8 +326,6 @@ def _npt_langevin_position_step( dt: Integration timestep, either scalar or with shape [n_systems] kT (torch.Tensor): Target temperature in energy units, either scalar or with shape [n_systems] - alpha (torch.Tensor | None): Friction coefficient, either scalar or with - shape [n_systems]. Returns: NPTLangevinState: Updated state with new positions @@ -413,8 +402,6 @@ def _npt_langevin_velocity_step( dt: Integration timestep, either scalar or with shape [n_systems] kT: Target temperature in energy units, either scalar or with shape [n_systems] - alpha (torch.Tensor | None): Friction coefficient, either scalar or with - shape [n_systems]. Returns: NPTLangevinState: Updated state with new velocities @@ -500,8 +487,12 @@ def _compute_cell_force( 3, device=state.device, dtype=state.dtype ) pressure_tensor = pressure_tensor.unsqueeze(0).expand(state.n_systems, -1, -1) + elif external_pressure.ndim == 1: + # Per-system scalar pressures (n_systems,) - create diagonal tensors + eye = torch.eye(3, device=state.device, dtype=state.dtype) + pressure_tensor = external_pressure.view(-1, 1, 1) * eye.unsqueeze(0) else: - # Already a tensor with shape compatible with n_systems + # Already a tensor with shape (n_systems, 3, 3) pressure_tensor = external_pressure # Calculate virials from stress and external pressure @@ -567,11 +558,11 @@ def npt_langevin_init( # Set default values if not provided if alpha is None: - alpha = 1.0 / (100 * dt) # Default friction based on timestep + alpha = 1.0 / (20 * dt) # Default friction based on timestep if cell_alpha is None: - cell_alpha = alpha # Use same friction for cell by default + cell_alpha = 1.0 / (100 * dt) # Default cell friction based on timestep if b_tau is None: - b_tau = 1 / (1000 * dt) # Default barostat time constant + b_tau = 0.1 * dt # Default barostat time constant # Convert all parameters to tensors with correct device and dtype alpha = torch.as_tensor(alpha, device=device, dtype=dtype) @@ -658,32 +649,86 @@ def npt_langevin_step( kT: float | torch.Tensor, external_pressure: float | torch.Tensor, ) -> NPTLangevinState: - """Perform one complete NPT Langevin dynamics integration step. - - This function implements a modified integration scheme for NPT dynamics, - handling both atomic and cell updates with Langevin thermostats to maintain - constant temperature and pressure. The integration scheme couples particle - motion with cell volume fluctuations. + r"""Perform one complete NPT Langevin dynamics integration step. + + Implements constant-pressure Langevin dynamics based on Gronbech-Jensen & + Farago (2014) [4]_ and the LAMMPS ``fix press/langevin`` scheme [5]_. + + **Particle equations** (GJF-style Langevin with volume coupling): + + .. math:: + + \mathbf{r}_i(t{+}\Delta t) &= \frac{L_{n+1}}{L_n}\,\mathbf{r}_i(t) + + \frac{2L_{n+1}}{L_{n+1}{+}L_n}\,b\,\Delta t + \left(\mathbf{v}_i + \frac{\Delta t\,\mathbf{F}_i}{2m_i} + + \frac{\boldsymbol{\beta}_i}{2m_i}\right) \\ + \mathbf{v}_i(t{+}\Delta t) &= a\,\mathbf{v}_i(t) + + \frac{\Delta t}{2m_i}\bigl(a\,\mathbf{F}_i^n + \mathbf{F}_i^{n+1}\bigr) + + \frac{b\,\boldsymbol{\beta}_i}{m_i} + + with damping coefficients + :math:`a = \frac{1 - \alpha\Delta t/(2m)}{1 + \alpha\Delta t/(2m)}`, + :math:`b = \frac{1}{1 + \alpha\Delta t/(2m)}`, and fluctuation-dissipation + noise :math:`\boldsymbol{\beta}_i = \sqrt{2\alpha k_BT\Delta t}\;\mathbf{R}`, + :math:`\mathbf{R}\sim\mathcal{N}(0,I)`. + + **Cell (volume) equations:** + + .. math:: + + V(t{+}\Delta t) &= V(t) + b_c\,\Delta t\,\dot{V} + + b_c\,\frac{\Delta t^2\,F_p}{2Q} + + b_c\,\frac{\Delta t\,\beta_c}{2Q} \\ + \dot{V}(t{+}\Delta t) &= a_c\,\dot{V}(t) + + \frac{\Delta t}{2Q}\bigl(a_c\,F_p^n + F_p^{n+1}\bigr) + + \frac{b_c\,\beta_c}{Q} + + where :math:`Q = (N+1) k_BT \tau_b^2` is the cell mass and + :math:`F_p = -V(\boldsymbol{\sigma} + P_{\text{ext}}\mathbf{I}) + + Nk_BT\mathbf{I}` is the cell force. + :math:`L = V^{1/3}` is the isotropic cell length scale. + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`\mathbf{r}_i` (positions) ``state.positions`` + :math:`\mathbf{v}_i` (velocities) ``state.velocities`` + :math:`\mathbf{p}_i` (momenta) ``state.momenta`` + :math:`m_i` (masses) ``state.masses`` + :math:`\mathbf{F}_i` (forces) ``state.forces`` + :math:`\alpha` (particle friction) ``state.alpha`` + :math:`\alpha_c` (cell friction) ``state.cell_alpha`` + :math:`V` (volume) ``state.cell_positions`` + :math:`\dot{V}` (cell velocity) ``state.cell_velocities`` + :math:`Q` (cell mass) ``state.cell_masses`` + :math:`\tau_b` (barostat time const) ``state.b_tau`` + :math:`L` (cell length scale) ``L_n`` + :math:`F_p` (cell force) ``F_p_n`` + :math:`P_{\text{ext}}` (target pressure) ``external_pressure`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + ============================================ ============================ Args: - model (ModelInterface): Neural network model that computes energies, forces, + model: Neural network model that computes energies, forces, and stress. Must return a dict with 'energy', 'forces', and 'stress' keys. - state (NPTLangevinState): Current NPT state with particle and cell variables - dt (float | torch.Tensor): Integration timestep, either scalar or - shape [n_systems] - kT (float | torch.Tensor): Target temperature in energy units, either scalar or + state: Current NPT state with particle and cell variables + dt: Integration timestep, either scalar or shape [n_systems] + kT: Target temperature in energy units, either scalar or shape [n_systems] - external_pressure (float | torch.Tensor): Target external pressure, + external_pressure: Target external pressure, either scalar or tensor with shape [n_systems, n_dim, n_dim] - alpha (torch.Tensor): Position friction coefficient, either scalar or - shape [n_systems] - cell_alpha (torch.Tensor): Cell friction coefficient, either scalar or - shape [n_systems] - b_tau (torch.Tensor): Barostat time constant, either scalar or shape [n_systems] Returns: - NPTLangevinState: Updated NPT state after one timestep with new positions, - velocities, cell parameters, forces, energy, and stress + NPTLangevinState: Updated NPT state after one timestep + + References: + .. [4] Gronbech-Jensen, N. & Farago, O. "Constant pressure and temperature + discrete-time Langevin molecular dynamics." J. Chem. Phys. 141(19) (2014). + .. [5] LAMMPS fix press/langevin: + https://docs.lammps.org/fix_press_langevin.html """ device, dtype = model.device, model.dtype @@ -704,11 +749,6 @@ def npt_langevin_step( n_atoms_per_system = torch.bincount(state.system_idx) state.cell_masses = (n_atoms_per_system + 1) * batch_kT * torch.square(state.b_tau) - # Compute model output for current state - model_output = model(state) - state.forces = model_output["forces"] - state.stress = model_output["stress"] - # Store initial values for integration forces = state.forces F_p_n = _compute_cell_force( @@ -1354,10 +1394,10 @@ def npt_nose_hoover_init( dt_tensor = torch.as_tensor(dt, device=device, dtype=dtype) kT_tensor = torch.as_tensor(kT, device=device, dtype=dtype) t_tau_tensor = torch.as_tensor( - 100 * dt_tensor if t_tau is None else t_tau, device=device, dtype=dtype + 10 * dt_tensor if t_tau is None else t_tau, device=device, dtype=dtype ) b_tau_tensor = torch.as_tensor( - 1000 * dt_tensor if b_tau is None else b_tau, device=device, dtype=dtype + 100 * dt_tensor if b_tau is None else b_tau, device=device, dtype=dtype ) # Setup thermostats with appropriate timescales @@ -1470,26 +1510,102 @@ def npt_nose_hoover_step( kT: float | torch.Tensor, external_pressure: float | torch.Tensor, ) -> NPTNoseHooverState: - """Perform a complete NPT integration step with Nose-Hoover chain thermostats. + r"""Perform a complete NPT integration step with Nose-Hoover chain thermostats. + + Implements the MTK (Martyna-Tobias-Klein) NPT scheme from Tuckerman et al. + (2006) [10]_ with Nose-Hoover chains from Martyna et al. (1996) [3]_. + + **Equations of motion** (Tuckerman et al. 2006, Eqs. 1-6): + + .. math:: + + \dot{\mathbf{r}}_i &= \frac{\mathbf{p}_i}{m_i} + + \frac{p_\epsilon}{W}\,\mathbf{r}_i \\ + \dot{\mathbf{p}}_i &= \mathbf{F}_i + - \alpha\,\frac{p_\epsilon}{W}\,\mathbf{p}_i \\ + \dot{\epsilon} &= \frac{p_\epsilon}{W} \\ + \dot{p}_\epsilon &= G_\epsilon + = \alpha\,(2K) + \text{Tr}(\boldsymbol{\sigma}_{\text{int}})\,V + - P_{\text{ext}}\,V\,d + + where :math:`\epsilon = (1/d)\ln(V/V_0)` is the logarithmic cell coordinate, + :math:`\alpha = 1 + d/N_f`, :math:`d=3` is spatial dimension, and + :math:`N_f = 3N - 3` the degrees of freedom. + + **Symmetric propagator** (Trotter factorization): + + .. math:: + + e^{i\mathcal{L}\Delta t} = + e^{i\mathcal{L}_{\text{NHC-baro}}\frac{\Delta t}{2}} + \;e^{i\mathcal{L}_{\text{NHC-part}}\frac{\Delta t}{2}} + \;e^{i\mathcal{L}_2\frac{\Delta t}{2}} + \;e^{i\mathcal{L}_1\Delta t} + \;e^{i\mathcal{L}_2\frac{\Delta t}{2}} + \;e^{i\mathcal{L}_{\text{NHC-part}}\frac{\Delta t}{2}} + \;e^{i\mathcal{L}_{\text{NHC-baro}}\frac{\Delta t}{2}} + + **Position update** :math:`e^{i\mathcal{L}_1\Delta t}`: + + .. math:: + + \mathbf{r}_i \leftarrow \mathbf{r}_i + + \bigl(e^{v_\epsilon\Delta t} - 1\bigr)\,\mathbf{r}_i + + \Delta t\,\mathbf{v}_i\,e^{v_\epsilon\Delta t/2} + \,\frac{\sinh(v_\epsilon\Delta t/2)}{v_\epsilon\Delta t/2} + + **Momentum update** :math:`e^{i\mathcal{L}_2\Delta t/2}`: + + .. math:: + + \mathbf{p}_i \leftarrow \mathbf{p}_i\,e^{-\alpha v_\epsilon\Delta t/2} + + \frac{\Delta t}{2}\,\mathbf{F}_i\, + e^{-\alpha v_\epsilon\Delta t/4} + \,\frac{\sinh(\alpha v_\epsilon\Delta t/4)} + {\alpha v_\epsilon\Delta t/4} + + where :math:`v_\epsilon = p_\epsilon / W` is the cell velocity. + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`\mathbf{r}_i` (positions) ``state.positions`` + :math:`\mathbf{p}_i` (momenta) ``state.momenta`` + :math:`m_i` (masses) ``state.masses`` + :math:`\mathbf{F}_i` (forces) ``state.forces`` + :math:`\epsilon` (log-cell coordinate) ``state.cell_position`` + :math:`p_\epsilon` (cell momentum) ``state.cell_momentum`` + :math:`W` (cell mass) ``state.cell_mass`` + :math:`\alpha` (1 + d/Nf) ``alpha`` (local) + :math:`v_\epsilon` (cell velocity) ``cell_velocities`` (local) + :math:`V_0` (reference volume) ``det(state.reference_cell)`` + :math:`G_\epsilon` (cell force) ``cell_force_val`` + :math:`P_{\text{ext}}` (target pressure) ``external_pressure`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + ============================================ ============================ + If the center of mass motion is removed initially, it remains removed throughout the simulation, so the degrees of freedom decreases by 3. - This function performs a full NPT integration step including: - 1. Mass parameter updates for thermostats and cell - 2. Thermostat chain updates (half step) - 3. Inner NPT dynamics step - 4. Energy updates for thermostats - 5. Final thermostat chain updates (half step) - Args: - model (ModelInterface): Model to compute forces and energies - state (NPTNoseHooverState): Current system state - dt (float | torch.Tensor): Integration timestep - kT (float | torch.Tensor): Target temperature - external_pressure (float | torch.Tensor): Target external pressure + model: Model to compute forces and energies + state: Current system state + dt: Integration timestep + kT: Target temperature + external_pressure: Target external pressure Returns: NPTNoseHooverState: Updated state after complete integration step + + References: + .. [10] Tuckerman, M. E., et al. "A Liouville-operator derived + measure-preserving integrator for molecular dynamics simulations in + the isothermal-isobaric ensemble." J. Phys. A 39(19), 5629-5651 (2006). + .. [3] Martyna, G. J., et al. "Explicit reversible integrators for extended + systems dynamics." Mol. Phys. 87(5), 1117-1157 (1996). """ device, dtype = model.device, model.dtype dt_tensor = torch.as_tensor(dt, device=device, dtype=dtype) @@ -1738,7 +1854,7 @@ def _crescale_anisotropic_barostat_step( ## Step 2: compute deformation matrix random_coeff = 2 * state.isothermal_compressibility * kT * dt / (3 * state.tau_p) prefactor_random_matrix = torch.sqrt(random_coeff) / new_sqrt_volume - a_tilde = -(state.isothermal_compressibility / (3 * state.tau_p))[:, None, None] * ( + a_tilde = (state.isothermal_compressibility / (3 * state.tau_p))[:, None, None] * ( P_int - trace_P_int[:, None, None] / 3 @@ -1780,7 +1896,8 @@ def _crescale_anisotropic_barostat_step( (vscaling + rscaling)[state.system_idx], state.momenta ) * dt / (2 * state.masses.unsqueeze(-1)) state.momenta = batch_matrix_vector(vscaling[state.system_idx], state.momenta) - state.cell = rscaling.mT @ state.cell + # Right multiply: cell @ rscaling^T preserves fractional coordinates + state.cell = state.cell @ rscaling.mT return state @@ -1815,7 +1932,7 @@ def _crescale_independent_lengths_barostat_step( ) # Note: it corresponds to using a diagonal isothermal compressibility tensor P_int_diagonal = torch.diagonal(P_int, dim1=-2, dim2=-1) - a_tilde = -(state.isothermal_compressibility / (3 * state.tau_p))[:, None] * ( + a_tilde = (state.isothermal_compressibility / (3 * state.tau_p))[:, None] * ( P_int_diagonal - trace_P_int[:, None] / 3 ) @@ -1886,8 +2003,7 @@ def _crescale_average_anisotropic_barostat_step( ) -> NPTCRescaleState: volume = torch.det(state.cell) # shape: (n_systems,) P_int = compute_average_pressure_tensor( - # Should it be degrees_of_freedom=state.get_number_of_degrees_of_freedom() / 3, - degrees_of_freedom=state.n_atoms_per_system, + degrees_of_freedom=state.get_number_of_degrees_of_freedom() / 3, kT=kT, stress=state.stress, volumes=volume, @@ -1907,7 +2023,7 @@ def _crescale_average_anisotropic_barostat_step( torch.sqrt(2 * state.isothermal_compressibility * kT * dt / (3 * state.tau_p)) / new_sqrt_volume ) - a_tilde = -(state.isothermal_compressibility / (3 * state.tau_p))[:, None, None] * ( + a_tilde = (state.isothermal_compressibility / (3 * state.tau_p))[:, None, None] * ( P_int - trace_P_int[:, None, None] / 3 @@ -1953,7 +2069,8 @@ def _crescale_average_anisotropic_barostat_step( )[state.system_idx], state.momenta, ) * dt / (2 * state.masses.unsqueeze(-1)) - state.cell = rscaling.mT @ state.cell + # Right multiply: cell @ rscaling^T preserves fractional coordinates + state.cell = state.cell @ rscaling.mT return state @@ -1979,7 +2096,9 @@ def _crescale_isotropic_barostat_step( prefactor = state.isothermal_compressibility * sqrt_vol / (2 * state.tau_p) change_sqrt_vol = -prefactor * ( external_pressure - trace_P_int / 3 - kT / (2 * volume) - ) * dt + prefactor_random * _randn_for_state(state, sqrt_vol.shape) + ) * dt + torch.sqrt( + 2 * torch.ones_like(sqrt_vol) + ) * prefactor_random * _randn_for_state(state, sqrt_vol.shape) new_sqrt_volume = sqrt_vol + change_sqrt_vol # Update positions and momenta (barostat + half momentum step) @@ -2010,7 +2129,7 @@ def _coerce_crescale_step_inputs( external_pressure, device=device, dtype=dtype ) tau_tensor = torch.as_tensor( - 100 * dt_tensor if tau is None else tau, device=device, dtype=dtype + 1 * dt_tensor if tau is None else tau, device=device, dtype=dtype ) return dt_tensor, kT_tensor, external_pressure_tensor, tau_tensor @@ -2026,41 +2145,96 @@ def npt_crescale_anisotropic_step( external_pressure: float | torch.Tensor, tau: float | torch.Tensor | None = None, ) -> NPTCRescaleState: - """Perform one NPT integration step with cell rescaling barostat. + r"""Perform one NPT integration step with anisotropic stochastic cell rescaling. - This function performs a single integration step for NPT dynamics using - a cell rescaling barostat. It updates particle positions, momenta, and - the simulation cell based on the target temperature and pressure. + Implements the anisotropic C-Rescale barostat from Del Tatto et al. + (2022) [7]_ extending the isotropic scheme of Bernetti & Bussi (2020) [6]_. + Cell lengths and angles can change independently. Uses instantaneous kinetic + energy. Both positions and momenta are scaled. - Trotter based splitting: - 1. Half Thermostat (velocity scaling) - 2. Half Update momenta with forces - 3. Barostat (cell rescaling) - 4. Update positions (from barostat + half momenta) - 5. Update forces with new positions and cell - 6. Compute forces - 7. Half Update momenta with forces - 8. Half Thermostat (velocity scaling) + **Trotter splitting:** - Only allow isotropic external stress. This method performs anisotropic - cell rescaling. Lengths and angles can change independently. Based on - pressure using kinetic energy. Positions and momenta are scaled when scaling the cell. + V-Rescale(dt/2) -> B(dt/2) -> Barostat(dt) -> Force eval -> B(dt/2) -> + V-Rescale(dt/2) - Inspired from: https://github.com/bussilab/crescale/blob/master/simplemd_anisotropic/simplemd.cpp - - Time reversible integrator - - Instantaneous kinetic energy (not not the average from equipartition) + **Barostat sub-steps** (3-step volume + deformation update): + + Step 1 -- Propagate :math:`\sqrt{V}` for :math:`\Delta t/2` (same SDE as + isotropic, Eq. 7 of [6]_): + + .. math:: + + \Delta\lambda = -\frac{\beta_T\lambda}{2\tau_p} + \left(P_0 - \frac{\text{Tr}(\mathbf{P}_{\text{int}})}{3} + - \frac{k_BT}{2V}\right)\frac{\Delta t}{2} + + \sqrt{\frac{k_BT\beta_T\Delta t}{4\tau_p}}\;R + + Step 2 -- Compute deviatoric deformation matrix: + + .. math:: + + \tilde{\mathbf{A}} &= \frac{\beta_T}{3\tau_p} + \left(\mathbf{P}_{\text{int}} + - \frac{\text{Tr}(\mathbf{P}_{\text{int}})}{3}\,\mathbf{I}\right) \\ + \boldsymbol{\mu}_{\text{dev}} &= \exp\bigl(\tilde{\mathbf{A}}\,\Delta t + + \sigma\,\tilde{\mathbf{R}}\bigr) + + where :math:`\sigma = \sqrt{2\beta_T k_BT\Delta t/(3\tau_p)}\;/\;\sqrt{V'}` + and :math:`\tilde{\mathbf{R}}` is a traceless random matrix. + + Step 3 -- Propagate :math:`\sqrt{V}` for :math:`\Delta t/2` (same as step 1). + + **Total scaling and update:** + + .. math:: + + \boldsymbol{\mu} &= \boldsymbol{\mu}_{\text{dev}} + \cdot (V'/V)^{1/3} \\ + \mathbf{r}_i &\leftarrow \boldsymbol{\mu}\,\mathbf{r}_i + + (\boldsymbol{\mu}^{-T} + \boldsymbol{\mu})\, + \frac{\mathbf{p}_i}{2m_i}\,\Delta t \\ + \mathbf{p}_i &\leftarrow \boldsymbol{\mu}^{-T}\,\mathbf{p}_i \\ + \mathbf{h} &\leftarrow \mathbf{h}\,\boldsymbol{\mu}^T + + **Variable mapping (equation -> code):** + + ============================================ ================================ + Equation symbol Code variable + ============================================ ================================ + :math:`V` (volume) ``volume`` + :math:`\lambda` (:math:`\sqrt{V}`) ``sqrt_vol`` + :math:`\beta_T` (compressibility) ``state.isothermal_compressibility`` + :math:`\tau_p` (barostat relax. time) ``state.tau_p`` + :math:`P_0` (target pressure) ``external_pressure`` + :math:`\mathbf{P}_{\text{int}}` (press. tensor) ``P_int`` + :math:`\tilde{\mathbf{A}}` (deviator drive) ``a_tilde`` + :math:`\boldsymbol{\mu}_{\text{dev}}` ``deformation_matrix`` + :math:`\boldsymbol{\mu}` (total scaling) ``rscaling`` + :math:`\boldsymbol{\mu}^{-T}` (mom. scaling) ``vscaling`` + :math:`\tilde{\mathbf{R}}` (traceless noise) ``random_matrix_tilde`` + :math:`\sigma` (noise prefactor) ``prefactor_random_matrix`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + :math:`\tau` (thermostat relax.) ``tau`` (V-Rescale) + ============================================ ================================ Args: - model (ModelInterface): Model to compute forces and energies - state (NPTCRescaleState): Current system state - dt (torch.Tensor): Integration timestep - kT (torch.Tensor): Target temperature - external_pressure (torch.Tensor): Target external pressure - tau (torch.Tensor | None): V-Rescale thermostat relaxation time. If None, - defaults to 100*dt + model: Model to compute forces and energies + state: Current system state + dt: Integration timestep + kT: Target temperature + external_pressure: Target external pressure + tau: V-Rescale thermostat relaxation time. If None, defaults to 100*dt Returns: NPTCRescaleState: Updated state after one integration step + + References: + .. [7] Del Tatto, V., et al. "Molecular dynamics of solids at constant + pressure and stress using anisotropic stochastic cell rescaling." + Applied Sciences 12(3), 1139 (2022). + .. [6] Bernetti, M. & Bussi, G. "Pressure control using stochastic cell + rescaling." J. Chem. Phys. 153, 114107 (2020). """ dt_tensor, kT_tensor, external_pressure_tensor, tau_tensor = ( _coerce_crescale_step_inputs(state, dt, kT, external_pressure, tau) @@ -2216,7 +2390,7 @@ def npt_crescale_average_anisotropic_step( external_pressure = torch.as_tensor(external_pressure, device=device, dtype=dtype) # Note: would probably be better to have tau in NVTCRescaleState - tau = torch.as_tensor(tau or 100 * dt, device=device, dtype=dtype) + tau = torch.as_tensor(tau or 1 * dt, device=device, dtype=dtype) state = _vrescale_update(state, tau, kT, dt / 2) @@ -2248,44 +2422,74 @@ def npt_crescale_isotropic_step( external_pressure: float | torch.Tensor, tau: float | torch.Tensor | None = None, ) -> NPTCRescaleState: - """Perform one NPT integration step with cell rescaling barostat. + r"""Perform one NPT integration step with isotropic stochastic cell rescaling. - This function performs a single integration step for NPT dynamics using - a cell rescaling barostat. It updates particle positions, momenta, and - the simulation cell based on the target temperature and pressure. + Implements isotropic C-Rescale from Bernetti & Bussi (2020) [6]_. + Cell shape is preserved; cell lengths are scaled equally. - Trotter based splitting: - 1. Half Thermostat (velocity scaling) - 2. Half Update momenta with forces - 3. Barostat (cell rescaling) - 4. Update positions (from barostat + half momenta) - 5. Update forces with new positions and cell - 6. Compute forces - 7. Half Update momenta with forces - 8. Half Thermostat (velocity scaling) + **Trotter splitting:** - Only allow isotropic external stress. This performs isotropic - cell rescaling: cell shape is preserved, cell lengths are scaled equally. - For anisotropic cell rescaling, use npt_crescale_anisotropic_step. + V-Rescale(dt/2) -> B(dt/2) -> Barostat(dt) -> Force eval -> B(dt/2) -> + V-Rescale(dt/2) - References: - - Bernetti, Mattia, and Giovanni Bussi. - "Pressure control using stochastic cell rescaling." - The Journal of Chemical Physics 153.11 (2020). - - And the corresponding Supplementary Information which details - the integration scheme. Notice an error in scaling of positions in SI Eq. S13a. + **Isotropic volume SDE** (Eq. 7 of [6]_, using :math:`\lambda = \sqrt{V}`): + + .. math:: + + d\lambda = -\frac{\beta_T\lambda}{2\tau_p} + \left(P_0 - \frac{\text{Tr}(\mathbf{P}_{\text{int}})}{3} + - \frac{k_BT}{2V}\right) dt + + \sqrt{\frac{k_BT\,\beta_T}{2\tau_p}}\;dW + + where :math:`\beta_T` is the isothermal compressibility and + :math:`\mathbf{P}_{\text{int}}` is the instantaneous pressure tensor + (including the kinetic contribution). + + **Position and momentum scaling** (SI Eqs. S13a-b of [6]_, corrected): + + .. math:: + + \mathbf{r}_i &\leftarrow \mu\,\mathbf{r}_i + + (\mu + \mu^{-1})\,\frac{\mathbf{p}_i}{2m_i}\,\Delta t \\ + \mathbf{p}_i &\leftarrow \mu^{-1}\,\mathbf{p}_i \\ + \mathbf{h} &\leftarrow \mu\,\mathbf{h} + + where :math:`\mu = (V'/V)^{1/3}` is the isotropic scaling factor and + :math:`\mathbf{h}` is the cell matrix. + + **Variable mapping (equation -> code):** + + ============================================ ================================ + Equation symbol Code variable + ============================================ ================================ + :math:`V` (volume) ``volume`` + :math:`\lambda` (:math:`\sqrt{V}`) ``sqrt_vol`` + :math:`\beta_T` (compressibility) ``state.isothermal_compressibility`` + :math:`\tau_p` (barostat relax. time) ``state.tau_p`` + :math:`P_0` (target pressure) ``external_pressure`` + :math:`\mathbf{P}_{\text{int}}` (press. tensor) ``P_int`` + :math:`\text{Tr}(\mathbf{P}_{\text{int}})` ``trace_P_int`` + :math:`\mu` (scaling factor) ``rscaling`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + :math:`\tau` (thermostat relax.) ``tau`` (V-Rescale) + ============================================ ================================ Args: - model (ModelInterface): Model to compute forces and energies - state (NPTCRescaleState): Current system state - dt (torch.Tensor): Integration timestep - kT (torch.Tensor): Target temperature - external_pressure (torch.Tensor): Target external pressure - tau (torch.Tensor | None): V-Rescale thermostat relaxation time. If None, - defaults to 100*dt + model: Model to compute forces and energies + state: Current system state + dt: Integration timestep + kT: Target temperature + external_pressure: Target external pressure + tau: V-Rescale thermostat relaxation time. If None, defaults to 100*dt Returns: NPTCRescaleState: Updated state after one integration step + + References: + .. [6] Bernetti, M. & Bussi, G. "Pressure control using stochastic cell + rescaling." J. Chem. Phys. 153, 114107 (2020). Note: SI Eq. S13a has a + typo (positions must also be scaled by mu). """ device, dtype = model.device, model.dtype dt = torch.as_tensor(dt, device=device, dtype=dtype) @@ -2293,7 +2497,7 @@ def npt_crescale_isotropic_step( external_pressure = torch.as_tensor(external_pressure, device=device, dtype=dtype) # Note: would probably be better to have tau in NVTCRescaleState - tau = torch.as_tensor(tau or 100 * dt, device=device, dtype=dtype) + tau = torch.as_tensor(tau or 1 * dt, device=device, dtype=dtype) state = _vrescale_update(state, tau, kT, dt / 2) @@ -2351,11 +2555,10 @@ def npt_crescale_init( kT = torch.as_tensor(kT, device=device, dtype=dtype) # Set default values if not provided - tau_p = torch.as_tensor( - tau_p or 5000 * dt, device=device, dtype=dtype - ) # 5ps for dt=1fs + tau_p = torch.as_tensor(tau_p or 3 * dt, device=device, dtype=dtype) # 5ps for dt=1fs isothermal_compressibility = torch.as_tensor( - isothermal_compressibility or 1e-1, + isothermal_compressibility + or 1e-6 / MetalUnits.pressure, # 1e-6 bar^-1 for metals device=device, dtype=dtype, # (eV/A^3)^-1 ) diff --git a/torch_sim/integrators/nve.py b/torch_sim/integrators/nve.py index 07f3064bb..d37ff7ed8 100644 --- a/torch_sim/integrators/nve.py +++ b/torch_sim/integrators/nve.py @@ -68,14 +68,35 @@ def nve_init( def nve_step( state: MDState, model: ModelInterface, *, dt: float | torch.Tensor, **_kwargs: Any ) -> MDState: - """Perform one complete NVE (microcanonical) integration step. + r"""Perform one complete NVE (microcanonical) integration step. - This function implements the velocity Verlet algorithm for NVE dynamics, - which provides energy-conserving time evolution. The integration sequence is: - 1. Half momentum update using current forces - 2. Full position update using updated momenta - 3. Force update at new positions - 4. Half momentum update using new forces + Implements the velocity Verlet algorithm for NVE dynamics, which provides + energy-conserving, time-reversible integration of Hamilton's equations of motion. + + **Equations** (standard velocity Verlet): + + .. math:: + + \mathbf{p}_i(t + \Delta t/2) &= \mathbf{p}_i(t) + + \frac{\Delta t}{2}\,\mathbf{F}_i(t) \\ + \mathbf{r}_i(t + \Delta t) &= \mathbf{r}_i(t) + + \Delta t\,\frac{\mathbf{p}_i(t + \Delta t/2)}{m_i} \\ + \mathbf{F}_i(t + \Delta t) &= -\nabla_{\mathbf{r}_i} U\bigl( + \mathbf{r}(t + \Delta t)\bigr) \\ + \mathbf{p}_i(t + \Delta t) &= \mathbf{p}_i(t + \Delta t/2) + + \frac{\Delta t}{2}\,\mathbf{F}_i(t + \Delta t) + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`\mathbf{r}_i` (positions) ``state.positions`` + :math:`\mathbf{p}_i` (momenta) ``state.momenta`` + :math:`m_i` (masses) ``state.masses`` + :math:`\mathbf{F}_i` (forces) ``state.forces`` + :math:`\Delta t` (timestep) ``dt`` + ============================================ ============================ Args: model: Neural network model that computes energies and forces. @@ -88,10 +109,8 @@ def nve_step( momenta, forces, and energy Notes: - - Uses velocity Verlet algorithm for time reversible integration + - Symplectic, time-reversible integrator of second order accuracy O(dt^2) - Conserves energy in the absence of numerical errors - - Handles periodic boundary conditions if enabled in state - - Symplectic integrator preserving phase space volume """ dt = torch.as_tensor(dt, device=state.device, dtype=state.dtype) state = momentum_step(state, dt / 2) diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index 1b8017279..eadfd4ad1 100644 --- a/torch_sim/integrators/nvt.py +++ b/torch_sim/integrators/nvt.py @@ -143,16 +143,54 @@ def nvt_langevin_step( kT: float | torch.Tensor, gamma: float | torch.Tensor | None = None, ) -> MDState: - """Perform one complete Langevin dynamics integration step. - - This function implements the BAOAB splitting scheme for Langevin dynamics, - which provides accurate sampling of the canonical ensemble. The integration - sequence is: - 1. Half momentum update using forces (B step) - 2. Half position update using updated momenta (A step) - 3. Full stochastic update with noise and friction (O step) - 4. Half position update using updated momenta (A step) - 5. Half momentum update using new forces (B step) + r"""Perform one complete Langevin dynamics integration step using the BAOAB scheme. + + Implements the BAOAB splitting of the Langevin equation from + Leimkuhler & Matthews (2013, 2016) [2]_. The Langevin SDE is: + + .. math:: + + d\mathbf{q} &= M^{-1}\mathbf{p}\,dt \\ + d\mathbf{p} &= -\nabla U(\mathbf{q})\,dt + - \gamma\,\mathbf{p}\,dt + + \sigma M^{1/2}\,d\mathbf{W} + + where :math:`\sigma = \sqrt{2\gamma k_BT}` (fluctuation-dissipation relation). + + **BAOAB splitting** (B = kick, A = drift, O = Ornstein-Uhlenbeck): + + .. math:: + + \text{B:}\quad \mathbf{p} &\leftarrow \mathbf{p} + + \tfrac{\Delta t}{2}\,\mathbf{F}(\mathbf{q}) \\ + \text{A:}\quad \mathbf{q} &\leftarrow \mathbf{q} + + \tfrac{\Delta t}{2}\,M^{-1}\mathbf{p} \\ + \text{O:}\quad \mathbf{p} &\leftarrow + c_1\,\mathbf{p} + c_2\,M^{1/2}\mathbf{R}, + \quad \mathbf{R}\sim\mathcal{N}(0,I) \\ + \text{A:}\quad \mathbf{q} &\leftarrow \mathbf{q} + + \tfrac{\Delta t}{2}\,M^{-1}\mathbf{p} \\ + \text{B:}\quad \mathbf{p} &\leftarrow \mathbf{p} + + \tfrac{\Delta t}{2}\,\mathbf{F}(\mathbf{q}) + + with :math:`c_1 = e^{-\gamma\Delta t}` and + :math:`c_2 = \sqrt{k_BT\,(1-c_1^2)}`. + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`\mathbf{q}` (positions) ``state.positions`` + :math:`\mathbf{p}` (momenta) ``state.momenta`` + :math:`M` (mass matrix) ``state.masses`` + :math:`\mathbf{F}` (forces) ``state.forces`` + :math:`\gamma` (friction coefficient) ``gamma`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + :math:`c_1` ``c1`` in ``_ou_step`` + :math:`c_2` ``c2`` in ``_ou_step`` + ============================================ ============================ Args: state: Current system state containing positions, momenta, forces @@ -168,13 +206,12 @@ def nvt_langevin_step( MDState: Updated state after one complete Langevin step with new positions, momenta, forces, and energy - Notes: - - Uses BAOAB splitting scheme for Langevin dynamics - - Preserves detailed balance for correct NVT sampling - - Handles periodic boundary conditions if enabled in state - - Friction coefficient gamma controls the thermostat coupling strength - - Weak coupling (small gamma) preserves dynamics but with slower thermalization - - Strong coupling (large gamma) faster thermalization but may distort dynamics + References: + .. [2] Leimkuhler B, Matthews C. "Efficient molecular dynamics using geodesic + integration and solvent-solute splitting." Proc. R. Soc. A 472: 20160138 + (2016). Original BAOAB analysis in: Leimkuhler B, Matthews C. "Rational + construction of stochastic numerical methods for molecular sampling." + Appl. Math. Res. Express 2013(1), 34-56 (2013). """ device, dtype = model.device, model.dtype @@ -290,7 +327,7 @@ def nvt_nose_hoover_init( dt_tensor = torch.as_tensor(dt, device=state.device, dtype=state.dtype) kT_tensor = torch.as_tensor(kT, device=state.device, dtype=state.dtype) tau_tensor = torch.as_tensor( - 100.0 * dt_tensor if tau is None else tau, device=state.device, dtype=state.dtype + 10.0 * dt_tensor if tau is None else tau, device=state.device, dtype=state.dtype ) # Create thermostat functions @@ -338,13 +375,63 @@ def nvt_nose_hoover_step( dt: float | torch.Tensor, kT: float | torch.Tensor, ) -> NVTNoseHooverState: - """Perform one complete Nose-Hoover chain integration step. - - This function performs one integration step for an NVT system using a Nose-Hoover - chain thermostat. The integration scheme is time-reversible and conserves an - extended energy quantity. If the center of mass motion is removed initially, - it remains removed throughout the simulation, so the degrees of freedom decreases - by 3. + r"""Perform one complete Nose-Hoover chain (NHC) integration step. + + Implements the NHC thermostat from Martyna et al. (1996) [3]_ with + Suzuki-Yoshida integration of the chain variables. + + **Equations of motion** (Martyna et al. 1996, Eqs. 13-18): + + .. math:: + + \dot{\mathbf{r}}_i &= \mathbf{p}_i / m_i \\ + \dot{\mathbf{p}}_i &= \mathbf{F}_i + - \frac{p_{\xi_1}}{Q_1}\,\mathbf{p}_i \\ + \dot{\xi}_j &= p_{\xi_j} / Q_j \\ + \dot{p}_{\xi_1} &= \bigl(2K - N_f k_BT\bigr) + - \frac{p_{\xi_2}}{Q_2}\,p_{\xi_1} \\ + \dot{p}_{\xi_j} &= \left(\frac{p_{\xi_{j-1}}^2}{Q_{j-1}} + - k_BT\right) - \frac{p_{\xi_{j+1}}}{Q_{j+1}}\,p_{\xi_j} + \quad (j = 2,\ldots,M{-}1) \\ + \dot{p}_{\xi_M} &= \frac{p_{\xi_{M-1}}^2}{Q_{M-1}} - k_BT + + where :math:`K = \sum_i p_i^2/(2m_i)` is the kinetic energy, + :math:`N_f = 3N - 3` the degrees of freedom, and :math:`Q_j = k_BT\tau^2` + (with :math:`Q_1 = N_f k_BT\tau^2`) are the chain masses. + + **Symmetric propagator** (Trotter factorization): + + .. math:: + + e^{i\mathcal{L}\Delta t} = e^{i\mathcal{L}_{\text{NHC}}\Delta t/2} + \;e^{i\mathcal{L}_{\text{VV}}\Delta t} + \;e^{i\mathcal{L}_{\text{NHC}}\Delta t/2} + + where :math:`i\mathcal{L}_{\text{VV}}` is the velocity Verlet propagator + and :math:`i\mathcal{L}_{\text{NHC}}` integrates the chain with + :math:`n_c \times n_{\text{sy}}` sub-steps (Suzuki-Yoshida decomposition). + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`\mathbf{r}_i` (positions) ``state.positions`` + :math:`\mathbf{p}_i` (momenta) ``state.momenta`` + :math:`m_i` (masses) ``state.masses`` + :math:`\mathbf{F}_i` (forces) ``state.forces`` + :math:`\xi_j` (chain positions) ``state.chain.positions`` + :math:`p_{\xi_j}` (chain momenta) ``state.chain.momenta`` + :math:`Q_j` (chain masses) ``state.chain.masses`` + :math:`K` (kinetic energy) ``state.chain.kinetic_energy`` + :math:`N_f` (degrees of freedom) ``state.chain.degrees_of_freedom`` + :math:`\tau` (relaxation time) ``state.chain.tau`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + :math:`M` (chain length) ``chain_length`` + :math:`n_c` (chain substeps) ``chain_steps`` + :math:`n_{\text{sy}}` (SY steps) ``sy_steps`` + ============================================ ============================ Args: state: Current system state containing positions, momenta, forces, and chain @@ -355,13 +442,10 @@ def nvt_nose_hoover_step( Returns: Updated state after one complete Nose-Hoover step - Notes: - Integration sequence: - 1. Update chain masses based on target temperature - 2. First half-step of chain evolution - 3. Full velocity Verlet step - 4. Update chain kinetic energy - 5. Second half-step of chain evolution + References: + .. [3] Martyna, G. J., Tuckerman, M. E., Tobias, D. J. & Klein, M. L. + "Explicit reversible integrators for extended systems dynamics." + Molecular Physics 87(5), 1117-1157 (1996). """ # Get chain functions from state chain_fns = state._chain_fns # noqa: SLF001 @@ -410,7 +494,6 @@ def nvt_nose_hoover_invariant( useful for validating the thermostat implementation. Args: - energy_fn: Function that computes system potential energy given positions state: Current state of the system including chain variables kT: Target temperature in energy units @@ -623,12 +706,53 @@ def nvt_vrescale_step( kT: float | torch.Tensor, tau: float | torch.Tensor | None = None, ) -> NVTVRescaleState: - """Perform one complete V-Rescale dynamics integration step. + r"""Perform one complete V-Rescale (CSVR) dynamics integration step. + + Implements canonical sampling through velocity rescaling from + Bussi, Donadio & Parrinello (2007) [1]_. + + **Stochastic differential equation** for kinetic energy (Eq. 7 of [1]_): + + .. math:: - This function implements the canonical sampling through velocity rescaling (V-Rescale) - thermostat combined with velocity Verlet integration. The V-Rescale thermostat samples - the canonical distribution by rescaling velocities with a properly chosen random - factor that ensures correct canonical sampling. + dK = \frac{\bar{K} - K}{\tau}\,dt + + 2\sqrt{\frac{K\bar{K}}{N_f\tau}}\,dW + + where :math:`\bar{K} = N_f k_BT/2` is the target kinetic energy. + + **Discrete rescaling factor** :math:`\alpha^2 = K'/K` (Eq. 22 of [1]_): + + .. math:: + + \alpha^2 = e^{-\Delta t/\tau} + + \frac{\bar{K}}{N_f K}\bigl(1-e^{-\Delta t/\tau}\bigr) + \Bigl(R_1^2 + \sum_{i=2}^{N_f} R_i^2\Bigr) + + 2\,e^{-\Delta t/(2\tau)} + \sqrt{\frac{\bar{K}}{N_f K} + \bigl(1-e^{-\Delta t/\tau}\bigr)}\;R_1 + + where :math:`R_1 \sim \mathcal{N}(0,1)` and + :math:`\sum_{i=2}^{N_f} R_i^2 \sim \text{Gamma}\bigl((N_f-1)/2,\,2\bigr)`. + Momenta are then rescaled as :math:`\mathbf{p} \leftarrow \alpha\,\mathbf{p}`. + + **Variable mapping (equation -> code):** + + ============================================ ============================ + Equation symbol Code variable + ============================================ ============================ + :math:`K` (kinetic energy) ``KE_old`` + :math:`\bar{K}` (target KE) ``KE_new`` + :math:`N_f` (degrees of freedom) ``dof`` + :math:`\tau` (relaxation time) ``tau`` + :math:`k_BT` (thermal energy) ``kT`` + :math:`\Delta t` (timestep) ``dt`` + :math:`e^{-\Delta t/\tau}` ``c1`` + :math:`(1-c_1)\bar{K}/(N_f K)` ``c2`` + :math:`R_1` ``r1`` + :math:`\sum_{i=2}^{N_f} R_i^2` ``r2`` + :math:`\alpha^2` (scale factor) ``scale`` + :math:`\alpha` (velocity rescaling) ``lam`` + ============================================ ============================ Args: model: Neural network model that computes energies and forces. @@ -644,19 +768,13 @@ def nvt_vrescale_step( MDState: Updated state after one complete V-Rescale step with new positions, momenta, forces, and energy - Notes: - - Uses V-Rescale thermostat for proper canonical ensemble sampling - - Unlike Berendsen thermostat, V-Rescale samples the true canonical distribution - - Integration sequence: V-Rescale rescaling + Velocity Verlet step - - The rescaling factor follows the distribution derived in Bussi et al. - References: - Bussi G, Donadio D, Parrinello M. "Canonical sampling through velocity rescaling." - The Journal of chemical physics, 126(1), 014101 (2007). + .. [1] Bussi G, Donadio D, Parrinello M. "Canonical sampling through velocity + rescaling." J. Chem. Phys. 126(1), 014101 (2007). """ device, dtype = model.device, model.dtype - tau = torch.as_tensor(100 * dt if tau is None else tau, device=device, dtype=dtype) + tau = torch.as_tensor(10 * dt if tau is None else tau, device=device, dtype=dtype) dt = torch.as_tensor(dt, device=device, dtype=dtype) kT = torch.as_tensor(kT, device=device, dtype=dtype) From 349765effa3161e75208b1f8b406605f74fa13bb Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 2 Apr 2026 18:31:20 +0000 Subject: [PATCH 09/15] remove testing files --- fast_integrator_tests/README.md | 70 --- fast_integrator_tests/analyze.py | 871 ------------------------------- fast_integrator_tests/common.py | 48 -- fast_integrator_tests/run_npt.py | 174 ------ fast_integrator_tests/run_nve.py | 91 ---- fast_integrator_tests/run_nvt.py | 116 ---- 6 files changed, 1370 deletions(-) delete mode 100644 fast_integrator_tests/README.md delete mode 100644 fast_integrator_tests/analyze.py delete mode 100644 fast_integrator_tests/common.py delete mode 100644 fast_integrator_tests/run_npt.py delete mode 100644 fast_integrator_tests/run_nve.py delete mode 100644 fast_integrator_tests/run_nvt.py diff --git a/fast_integrator_tests/README.md b/fast_integrator_tests/README.md deleted file mode 100644 index a60ee9446..000000000 --- a/fast_integrator_tests/README.md +++ /dev/null @@ -1,70 +0,0 @@ -# Integrator Physical Validation Tests - -Physical validation of all torch-sim MD integrators using the -[physical_validation](https://physical-validation.readthedocs.io/) library -and LJ Argon as a test system. - -## What's here - -| File | Purpose | -|---|---| -| `common.py` | Shared constants, model/structure builders | -| `run_nvt.py` | Run NVT simulations (Langevin, Nose-Hoover, V-Rescale) at 80K and 100K | -| `run_npt.py` | Run NPT simulations (Langevin, Nose-Hoover, iso/aniso C-Rescale) at 80K and 100K | -| `run_nve.py` | Run NVE simulations at 8 timesteps for convergence analysis | -| `analyze.ipynb` | Jupyter notebook with all diagnostic plots and `physical_validation` checks | -| `data/` | Output directory for `.npz` trajectory data (gitignored) | - -## Quick start - -```bash -# 1. Generate trajectory data (from this directory) -python run_nve.py -python run_nvt.py -python run_npt.py - -# Run a single integrator if needed -python run_nvt.py --integrator nvt_langevin -python run_npt.py --integrator npt_nose_hoover - -# NPT runs both a pressure and temperature sweep, so you can specify which to do: -python run_npt.py --mode temperature # Vary T at P=0 -python run_npt.py --mode pressure # Vary P at fixed T - -# 2. Open the notebook -jupyter notebook analyze.ipynb -``` - -## What the notebook shows - -### Custom plots -- **Time series**: Temperature, total energy, volume (NPT) vs step -- **KE distribution**: Observed histogram vs theoretical Gamma(Nf/2, k_BT) distribution -- **Ensemble check**: Overlapping energy distributions at two temperatures with log-ratio inset -- **NVE convergence**: Log-log RMSD vs dt plot, energy drift traces, convergence ratio table - -### physical_validation native plots -The notebook also calls `physical_validation` functions with `screen=True` to produce -their built-in diagnostic figures: -- `kinetic_energy.distribution()` — KS test or mean/width comparison plot -- `ensemble.check()` — Forward/reverse work distributions and linear fit -- `integrator.convergence()` — RMSD vs dt with reference line - -### Summary table -Final cell runs all checks programmatically and prints a PASS/FAIL table for every integrator. - -## Validation checks - -| Check | Integrators | What it tests | -|---|---|---| -| KE distribution | All NVT + NPT | Kinetic energy follows Maxwell-Boltzmann (gamma) distribution | -| Ensemble check | All NVT + NPT | Energy distributions at T=80K and T=100K satisfy Boltzmann weight ratio | -| NVE convergence | NVE (velocity Verlet) | Energy drift RMSD scales as dt^2 | - -## System details - -- **Structure**: FCC Argon (a=5.26 A), 2x2x2 supercell (32 atoms) for NVT/NPT, unit cell (4 atoms) for NVE -- **Model**: Lennard-Jones (sigma=3.405, epsilon=0.0104 eV, cutoff=8.5125 A), no neighbor list -- **Timestep**: 0.004 ps for NVT/NPT; sweep 0.002-0.010 ps for NVE -- **Production**: 10,000 steps after 2,000 (NVT) or 3,000 (NPT) equilibration steps -- **Threshold**: |deviation| < 3 sigma for all statistical checks diff --git a/fast_integrator_tests/analyze.py b/fast_integrator_tests/analyze.py deleted file mode 100644 index 444948d65..000000000 --- a/fast_integrator_tests/analyze.py +++ /dev/null @@ -1,871 +0,0 @@ -# %% [markdown] -# # Physical Validation of torch-sim Integrators -# -# This notebook analyzes MD trajectory data produced by `run_nvt.py`, `run_npt.py`, and `run_nve.py`. -# -# **Tests performed:** -# 1. **KE Distribution** — Does kinetic energy follow the Maxwell-Boltzmann (gamma) distribution? -# 2. **Ensemble Check** — Do energy distributions at two temperatures satisfy the Boltzmann weight relationship? -# 3. **Pressure Ensemble Check** — Do volume distributions at different pressures satisfy the expected Boltzmann weight relationship? -# 4. **NVE Convergence** — Does the energy drift RMSD scale as dt² (velocity Verlet)? -# 5. **Time Series** — Visual inspection of temperature, energy, and volume stability. - -# %% -# Create a folder to save plot -import os -os.makedirs("plots", exist_ok=True) - -# %% -import numpy as np -import matplotlib.pyplot as plt -from pathlib import Path -from scipy import stats - -plt.rcParams.update({ - "figure.dpi": 120, - "axes.grid": True, - "grid.alpha": 0.3, - "figure.facecolor": "white", -}) - -DATA_DIR = Path("fast_integrator_tests/data") - -def load(name): - """Load an npz file from the data directory.""" - return dict(np.load(DATA_DIR / f"{name}.npz", allow_pickle=True)) - -# Check what data is available -available = sorted(p.stem for p in DATA_DIR.glob("*.npz")) -print("Available datasets:") -for a in available: - print(f" {a}") - -# %% -import physical_validation -from torch_sim.units import MetalUnits - -k_B_eV = float(MetalUnits.temperature) # 8.617333e-5 eV/K -THRESHOLD = 3.0 - -# def make_unit_data(): -# return physical_validation.data.UnitData( -# kb=k_B_eV, -# energy_str="eV", energy_conversion=1.0, -# length_str="Ang", length_conversion=1.0, -# volume_str="Ang^3", volume_conversion=1.0, -# temperature_str="K", temperature_conversion=1.0, -# pressure_str="eV/Ang^3", pressure_conversion=1.0, -# time_str="internal", time_conversion=1.0, -# ) -def make_unit_data(): - return physical_validation.data.UnitData( - kb=k_B_eV, - energy_str="eV", energy_conversion=96.485, # Convert to kJ/mol - length_str="Ang", length_conversion=1e-1, # Convert to nm - volume_str="Ang^3", volume_conversion=1e-3, # Convert to nm^3 - temperature_str="K", temperature_conversion=1.0, - # pressure_str="eV/Ang^3", pressure_conversion=1.6e6, # Convert to bar - pressure_str="bar", pressure_conversion=1, - time_str="fs", time_conversion=1e-3, # Convert to ps - ) - -def build_sim_data(d, temperature, ensemble="NVT", pressure=None): - """Build physical_validation.SimulationData from a loaded npz dict.""" - units = make_unit_data() - natoms = int(d["natoms"]) - system = physical_validation.data.SystemData( - natoms=natoms, nconstraints=0, - ndof_reduction_tra=3, ndof_reduction_rot=0, - mass=d["masses"], - ) - ens_kw = dict(natoms=natoms, temperature=temperature) - if ensemble == "NVT": - ens_kw["ensemble"] = "NVT" - ens_kw["volume"] = float(d.get("volume", np.mean(d.get("volumes", [0])))) - else: - ens_kw["ensemble"] = "NPT" - ens_kw["pressure"] = pressure if pressure is not None else 0.0 - - obs_kw = dict( - kinetic_energy=d["kinetic_energy"], - potential_energy=d["potential_energy"], - total_energy=d["total_energy"], - ) - if "volumes" in d: - obs_kw["volume"] = d["volumes"] - - return physical_validation.data.SimulationData( - units=units, - dt=float(d["dt_internal"]), - system=system, - ensemble=physical_validation.data.EnsembleData(**ens_kw), - observables=physical_validation.data.ObservableData(**obs_kw), - ) - -print("physical_validation helpers loaded") - -# %% [markdown] -# ## 1. NVT Time Series -# -# Temperature and energy vs step for each NVT integrator at both temperatures. - -# %% -NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] -TEMPS = [88.0, 100.0] - -fig, axes = plt.subplots(len(NVT_INTEGRATORS), 2, figsize=(14, 3.5 * len(NVT_INTEGRATORS)), - squeeze=False, sharex=True) -fig.suptitle("NVT Integrators — Time Series", fontsize=14, y=1.01) - -for row, name in enumerate(NVT_INTEGRATORS): - for temp in TEMPS: - label = f"{name}_T{temp:.0f}K" - try: - d = load(label) - except FileNotFoundError: - continue - - steps = np.arange(len(d["temperature"])) - target_T = float(d["target_temperature"]) - - # Temperature - ax = axes[row, 0] - ax.plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") - ax.axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) - ax.set_ylabel("Temperature (K)") - ax.set_title(name) - ax.legend(fontsize=8) - - # Total energy - ax = axes[row, 1] - ax.plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=f"T={target_T:.0f}K") - ax.set_ylabel("Total Energy (eV)") - ax.set_title(name) - ax.legend(fontsize=8) - -axes[-1, 0].set_xlabel("Step") -axes[-1, 1].set_xlabel("Step") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ## 2. KE Distribution — NVT Integrators -# -# For each integrator at 100K, compare the observed KE distribution against the theoretical gamma distribution. -# The KE of an ideal gas with $N_f$ degrees of freedom at temperature $T$ follows $\text{Gamma}(N_f/2, k_BT)$. - -# %% -fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), - sharey=True) -fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NVT, T=100K)", fontsize=13) - -for ax, name in zip(axes, NVT_INTEGRATORS): - label = f"{name}_T100K" - try: - d = load(label) - except FileNotFoundError: - ax.set_title(f"{name}\n(no data)") - continue - - ke = d["kinetic_energy"] - natoms = int(d["natoms"]) - target_T = float(d["target_temperature"]) - - # Degrees of freedom: 3*N - 3 (COM removed) - if name == "nvt_langevin": - ndof = 3 * natoms - else: - ndof = 3 * natoms - 3 - # Theoretical gamma: shape=ndof/2, scale=k_B*T - shape = ndof / 2 - scale = k_B_eV * target_T - - # Histogram of observed KE - ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") - - # Theoretical curve - x = np.linspace(ke.min(), ke.max(), 300) - pdf = stats.gamma.pdf(x, a=shape, scale=scale) - ax.plot(x, pdf, "r-", lw=2, label="Theory") - - # Stats - d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) - d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) - ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=10) - ax.set_xlabel("Kinetic Energy (eV)") - ax.legend(fontsize=8) - -axes[0].set_ylabel("Probability Density") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ### physical_validation built-in KE distribution plots (NVT) -# -# Uses `physical_validation.kinetic_energy.distribution(..., screen=True)` which overlays the observed and theoretical distributions with a KS-test or mean/width comparison. - -# %% -for name in NVT_INTEGRATORS: - label = f"{name}_T100K" - try: - d = load(label) - except FileNotFoundError: - print(f"{name}: no data") - continue - - print(f"\n{'='*60}") - print(f" {name} — KE distribution (physical_validation)") - print(f"{'='*60}") - sd = build_sim_data(d, 100.0, ensemble="NVT") - result = physical_validation.kinetic_energy.distribution( - sd, strict=False, screen=True, verbosity=2, filename=f"plots/{name}_ke_distribution.png" - ) - print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") - -plt.show() - -# %% [markdown] -# ## 3. Ensemble Check — NVT Integrators -# -# For each integrator, compare total energy distributions at T_low=88K and T_high=100K. -# The log ratio of the energy histograms should be linear with slope $\Delta\beta = 1/k_BT_\text{low} - 1/k_BT_\text{high}$. - -# %% -fig, axes = plt.subplots(1, len(NVT_INTEGRATORS), figsize=(5 * len(NVT_INTEGRATORS), 4), - sharey=True) -fig.suptitle("NVT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) - -temp_low, temp_high = 88.0, 100.0 -delta_beta = 1 / (k_B_eV * temp_low) - 1 / (k_B_eV * temp_high) - -for ax, name in zip(axes, NVT_INTEGRATORS): - try: - d_lo = load(f"{name}_T{temp_low:.0f}K") - d_hi = load(f"{name}_T{temp_high:.0f}K") - except FileNotFoundError: - ax.set_title(f"{name}\n(no data)") - continue - - e_lo = d_lo["total_energy"] - e_hi = d_hi["total_energy"] - - # Overlapping histogram bins - e_min = min(e_lo.min(), e_hi.min()) - e_max = max(e_lo.max(), e_hi.max()) - bins = np.linspace(e_min, e_max, 50) - bin_centers = (bins[:-1] + bins[1:]) / 2 - - h_lo, _ = np.histogram(e_lo, bins=bins, density=True) - h_hi, _ = np.histogram(e_hi, bins=bins, density=True) - - # Plot overlapping distributions - ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") - ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") - - # Inset: log ratio - mask = (h_lo > 0) & (h_hi > 0) - if mask.sum() > 2: - log_ratio = np.log(h_lo[mask] / h_hi[mask]) - bc = bin_centers[mask] - # Linear fit - slope, intercept = np.polyfit(bc, log_ratio, 1) - inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) - inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) - inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, - label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") - inset.set_xlabel("E (eV)", fontsize=7) - inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) - inset.legend(fontsize=6) - inset.tick_params(labelsize=6) - - ax.set_title(name, fontsize=10) - ax.set_xlabel("Total Energy (eV)") - ax.legend(fontsize=8) - -axes[0].set_ylabel("Probability Density") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ### physical_validation built-in ensemble check plots (NVT) -# -# Uses `physical_validation.ensemble.check(..., screen=True)` which shows the forward/reverse work distributions and the linear fit of the log-likelihood ratio. - -# %% -for name in NVT_INTEGRATORS: - try: - d_lo = load(f"{name}_T{temp_low:.0f}K") - d_hi = load(f"{name}_T{temp_high:.0f}K") - except FileNotFoundError: - print(f"{name}: no data") - continue - - print(f"\n{'='*60}") - print(f" {name} — Ensemble check (physical_validation)") - print(f"{'='*60}") - sd_lo = build_sim_data(d_lo, temp_low, ensemble="NVT") - sd_hi = build_sim_data(d_hi, temp_high, ensemble="NVT") - quantiles = physical_validation.ensemble.check( - sd_lo, sd_hi, - total_energy=True, data_is_uncorrelated=True, - screen=True, verbosity=2, filename=f"plots/{name}_ensemble_check.png" - ) - print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") - -plt.show() - -# %% [markdown] -# ## 4. NPT Time Series -# -# Temperature, energy, and volume vs step for each NPT integrator. - -# %% -NPT_INTEGRATORS = ["npt_langevin", "npt_nose_hoover", "npt_isotropic_crescale", "npt_anisotropic_crescale"] - -fig, axes = plt.subplots(len(NPT_INTEGRATORS), 3, figsize=(16, 3.5 * len(NPT_INTEGRATORS)), - squeeze=False, sharex=True) -fig.suptitle("NPT Integrators — Time Series", fontsize=14, y=1.01) - -for row, name in enumerate(NPT_INTEGRATORS): - for temp in TEMPS: - label = f"{name}_T{temp:.0f}K" - try: - d = load(label) - except FileNotFoundError: - continue - - steps = np.arange(len(d["temperature"])) - target_T = float(d["target_temperature"]) - tag = f"T={target_T:.0f}K" - - # Temperature - axes[row, 0].plot(steps, d["temperature"], alpha=0.5, lw=0.5, label=tag) - axes[row, 0].axhline(target_T, color="k", ls="--", lw=0.8, alpha=0.5) - axes[row, 0].set_ylabel("Temperature (K)") - axes[row, 0].set_title(name) - axes[row, 0].legend(fontsize=7) - - # Total energy - axes[row, 1].plot(steps, d["total_energy"], alpha=0.5, lw=0.5, label=tag) - axes[row, 1].set_ylabel("Total Energy (eV)") - axes[row, 1].set_title(name) - axes[row, 1].legend(fontsize=7) - - # Volume - axes[row, 2].plot(steps, d["volumes"], alpha=0.5, lw=0.5, label=tag) - axes[row, 2].set_ylabel("Volume (ų)") - axes[row, 2].set_title(name) - axes[row, 2].legend(fontsize=7) - -for j in range(3): - axes[-1, j].set_xlabel("Step") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ## 5. KE Distribution — NPT Integrators -# -# Same gamma distribution check, but for NPT integrators at 100K. - -# %% -fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), - sharey=True) -fig.suptitle("KE Distribution vs Maxwell-Boltzmann (NPT, T=100K)", fontsize=13) - -for ax, name in zip(axes, NPT_INTEGRATORS): - label = f"{name}_T100K" - try: - d = load(label) - except FileNotFoundError: - ax.set_title(f"{name}\n(no data)") - continue - - ke = d["kinetic_energy"] - natoms = int(d["natoms"]) - target_T = float(d["target_temperature"]) - - ndof = 3 * natoms - 3 - shape = ndof / 2 - scale = k_B_eV * target_T - - ax.hist(ke, bins=60, density=True, alpha=0.6, color="steelblue", label="Observed") - - x = np.linspace(ke.min(), ke.max(), 300) - pdf = stats.gamma.pdf(x, a=shape, scale=scale) - ax.plot(x, pdf, "r-", lw=2, label="Theory") - - d_mean = (ke.mean() - shape * scale) / (scale * np.sqrt(2 / ndof)) - d_width = (ke.std() - scale * np.sqrt(shape)) / (scale * np.sqrt(0.5)) - ax.set_title(f"{name}\nd_mean={d_mean:.2f}σ d_width={d_width:.2f}σ", fontsize=9) - ax.set_xlabel("Kinetic Energy (eV)") - ax.legend(fontsize=8) - -axes[0].set_ylabel("Probability Density") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ### physical_validation built-in KE distribution plots (NPT) - -# %% -for name in NPT_INTEGRATORS: - label = f"{name}_T100K" - try: - d = load(label) - except FileNotFoundError: - print(f"{name}: no data") - continue - - print(f"\n{'='*60}") - print(f" {name} — KE distribution (physical_validation)") - print(f"{'='*60}") - sd = build_sim_data(d, 100.0, ensemble="NPT") - result = physical_validation.kinetic_energy.distribution( - sd, strict=False, screen=True, verbosity=2, - ) - print(f" Result: d_mean={result[0]:.3f}σ, d_width={result[1]:.3f}σ") - -plt.show() - -# %% [markdown] -# ## 6. Ensemble Check — NPT Integrators -# -# Same Boltzmann weight ratio check at T=88K vs T=100K for NPT integrators. - -# %% -fig, axes = plt.subplots(1, len(NPT_INTEGRATORS), figsize=(5 * len(NPT_INTEGRATORS), 4), - sharey=True) -fig.suptitle("NPT Ensemble Check — Energy Distributions at T=88K vs T=100K", fontsize=13) - -for ax, name in zip(axes, NPT_INTEGRATORS): - try: - d_lo = load(f"{name}_T{temp_low:.0f}K") - d_hi = load(f"{name}_T{temp_high:.0f}K") - except FileNotFoundError: - ax.set_title(f"{name}\n(no data)") - continue - - e_lo = d_lo["total_energy"] - e_hi = d_hi["total_energy"] - - e_min = min(e_lo.min(), e_hi.min()) - e_max = max(e_lo.max(), e_hi.max()) - bins = np.linspace(e_min, e_max, 50) - bin_centers = (bins[:-1] + bins[1:]) / 2 - - h_lo, _ = np.histogram(e_lo, bins=bins, density=True) - h_hi, _ = np.histogram(e_hi, bins=bins, density=True) - - ax.hist(e_lo, bins=bins, density=True, alpha=0.5, color="blue", label=f"T={temp_low:.0f}K") - ax.hist(e_hi, bins=bins, density=True, alpha=0.5, color="red", label=f"T={temp_high:.0f}K") - - mask = (h_lo > 0) & (h_hi > 0) - if mask.sum() > 2: - log_ratio = np.log(h_lo[mask] / h_hi[mask]) - bc = bin_centers[mask] - slope, intercept = np.polyfit(bc, log_ratio, 1) - inset = ax.inset_axes([0.55, 0.55, 0.42, 0.42]) - inset.scatter(bc, log_ratio, s=8, color="k", zorder=3) - inset.plot(bc, slope * bc + intercept, "r-", lw=1.5, - label=f"fit: {slope:.1f}\ntheory: {delta_beta:.1f}") - inset.set_xlabel("E (eV)", fontsize=7) - inset.set_ylabel("ln(P_lo/P_hi)", fontsize=7) - inset.legend(fontsize=6) - inset.tick_params(labelsize=6) - - ax.set_title(name, fontsize=9) - ax.set_xlabel("Total Energy (eV)") - ax.legend(fontsize=8) - -axes[0].set_ylabel("Probability Density") -fig.tight_layout() -plt.show() - -# %% [markdown] -# ### physical_validation built-in ensemble check plots (NPT) - -# %% -for name in NPT_INTEGRATORS: - try: - d_lo = load(f"{name}_T{temp_low:.0f}K") - d_hi = load(f"{name}_T{temp_high:.0f}K") - except FileNotFoundError: - print(f"{name}: no data") - continue - - print(f"\n{'='*60}") - print(f" {name} — Ensemble check (physical_validation)") - print(f"{'='*60}") - sd_lo = build_sim_data(d_lo, temp_low, ensemble="NPT") - sd_hi = build_sim_data(d_hi, temp_high, ensemble="NPT") - try: - quantiles = physical_validation.ensemble.check( - sd_lo, sd_hi, - total_energy=False, # data_is_uncorrelated=True, - screen=True, verbosity=2, - ) - print(f" Quantiles (σ): {[f'{q:.3f}' for q in quantiles]}") - except Exception as e: #ConvergenceError - print(f" ConvergenceError: {e}") - -plt.show() - -# %% [markdown] -# ## 6b. NPT Pressure Ensemble Check -# -# Compare volume distributions at two pressures (same temperature). -# For correct NPT sampling, `physical_validation` checks that volume distributions -# at different pressures satisfy the expected Boltzmann weight relationship. -# This is a 1D check on volumes only — no energy data needed. - -# %% -from torch_sim.units import MetalUnits - -P_BAR_CONVERSION = float(MetalUnits.pressure) # 1 bar in eV/Ang^3 - -# Discover available pressure-sweep data files -pressure_files = sorted(DATA_DIR.glob("*_P*bar.npz")) -print("Pressure-sweep data files:") -for f in pressure_files: - print(f" {f.name}") - -# Parse integrator -> {pressure_bar: filename} mapping -pressure_data = {} # integrator_name -> [(p_bar, filename), ...] -for f in pressure_files: - stem = f.stem # e.g. npt_langevin_T100K_P0bar - parts = stem.rsplit("_P", 1) - if len(parts) != 2: - continue - prefix = parts[0] # e.g. npt_langevin_T100K - p_bar = float(parts[1].replace("bar", "")) - int_prefix = prefix.rsplit("_T", 1)[0] # e.g. npt_langevin - pressure_data.setdefault(int_prefix, []).append((p_bar, f.stem)) - -for name, entries in pressure_data.items(): - entries.sort(key=lambda x: x[0]) - print(f"\n{name}: {[(p, fn) for p, fn in entries]}") - -# %% -# Volume distributions at different pressures for each NPT integrator -integrators_with_pressure = [n for n in NPT_INTEGRATORS if n in pressure_data] - -if integrators_with_pressure: - fig, axes = plt.subplots(1, len(integrators_with_pressure), - figsize=(5 * len(integrators_with_pressure), 4), squeeze=False) - axes = axes[0] - fig.suptitle("NPT Pressure Check \u2014 Volume Distributions at Same T, Different P", fontsize=13) - - for ax, name in zip(axes, integrators_with_pressure): - entries = pressure_data[name] - colors = plt.cm.coolwarm(np.linspace(0, 1, len(entries))) - for (p_bar, fname), color in zip(entries, colors): - d = load(fname) - vols = d["volumes"] - ax.hist(vols, bins=50, density=True, alpha=0.5, color=color, - label=f"P={p_bar:.0f} bar") - ax.axvline(vols.mean(), color=color, ls="--", lw=1, alpha=0.7) - ax.set_title(name, fontsize=10) - ax.set_xlabel("Volume (\u00c5\u00b3)") - ax.legend(fontsize=8) - - axes[0].set_ylabel("Probability Density") - fig.tight_layout() - plt.savefig("plots/npt_pressure_volume_distributions.png", dpi=150, bbox_inches="tight") - plt.show() -else: - print("No pressure-sweep data found. Run: python run_npt.py --mode pressure") - -# %% [markdown] -# ### physical_validation built-in pressure ensemble check -# -# Uses `physical_validation.ensemble.check()` with two simulations at the same temperature -# but different pressures. The library auto-detects this case and performs a 1D volume-based check. - -# %% -for name in integrators_with_pressure: - entries = pressure_data[name] - if len(entries) < 2: - print(f"{name}: need at least 2 pressures, got {len(entries)}") - continue - - p_lo_bar, fname_lo = entries[0] - p_hi_bar, fname_hi = entries[-1] - d_lo = load(fname_lo) - d_hi = load(fname_hi) - - p_lo_eva3 = float(d_lo["external_pressure"]) - p_hi_eva3 = float(d_hi["external_pressure"]) - temp = float(d_lo["target_temperature"]) - print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " - f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") - - print(f"\n{'='*60}") - print(f" {name} \u2014 Pressure ensemble check") - print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") - print(f"{'='*60}") - - # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) - # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) - - sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) - sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) - - try: - quantiles = physical_validation.ensemble.check( - sd_lo, sd_hi, - total_energy=False, #data_is_uncorrelated=True, - screen=True, verbosity=2, - filename=f"plots/{name}_pressure_ensemble_check.png", - ) - print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") - except Exception as e: - print(f" Error: {e}") - -plt.show() - -# %% -for name in integrators_with_pressure: - entries = pressure_data[name] - if len(entries) < 2: - print(f"{name}: need at least 2 pressures, got {len(entries)}") - continue - - p_lo_bar, fname_lo = entries[0] - p_hi_bar, fname_hi = entries[-1] - d_lo = load(fname_lo) - d_hi = load(fname_hi) - - p_lo_eva3 = float(d_lo["external_pressure"]) - p_hi_eva3 = float(d_hi["external_pressure"]) - temp = float(d_lo["target_temperature"]) - print(f"{name}: p_lo={p_lo_bar:.1f} bar (eva3={p_lo_eva3:.3e}), " - f"p_hi={p_hi_bar:.1f} bar (eva3={p_hi_eva3:.3e}), T={temp:.1f}K") - - print(f"\n{'='*60}") - print(f" {name} \u2014 Pressure ensemble check") - print(f" T={temp:.0f}K, P_lo={p_lo_bar:.0f} bar, P_hi={p_hi_bar:.0f} bar") - print(f"{'='*60}") - - # sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_eva3) - # sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_eva3) - - sd_lo = build_sim_data(d_lo, temp, ensemble="NPT", pressure=p_lo_bar) - sd_hi = build_sim_data(d_hi, temp, ensemble="NPT", pressure=p_hi_bar) - - try: - quantiles = physical_validation.ensemble.check( - sd_lo, sd_hi, - total_energy=False, data_is_uncorrelated=True, - screen=True, verbosity=2, - filename=f"plots/{name}_pressure_ensemble_check.png", - ) - print(f" Quantiles (\u03c3): {[f'{q:.3f}' for q in quantiles]}") - except Exception as e: - print(f" Error: {e}") - -plt.show() - -# %% [markdown] -# ## 8. NVE Integrator Convergence -# -# RMSD of the conserved quantity (total energy) vs timestep on a log-log scale. -# For velocity Verlet, the RMSD should scale as $\text{dt}^2$ (slope = 2 on log-log). - -# %% -try: - nve = load("nve_convergence") - timesteps_ps = nve["timesteps_ps"] - - dts, rmsds, drifts = [], [], [] - for dt_ps in timesteps_ps: - key = f"dt_{dt_ps}" - com = nve[f"{key}_constant_of_motion"] - dts.append(dt_ps) - rmsds.append(com.std()) - drifts.append(com[-1] - com[0]) - - dts = np.array(dts) - rmsds = np.array(rmsds) - drifts = np.array(drifts) - - fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 5)) - fig.suptitle("NVE Integrator Convergence (4-atom Ar, T=5K)", fontsize=13) - - # --- Panel 1: RMSD vs dt (log-log) --- - ax1.loglog(dts, rmsds, "ko-", ms=8, lw=2, label="RMSD(E_tot)") - - # dt^2 reference line through middle point - mid = len(dts) // 2 - ref = rmsds[mid] * (dts / dts[mid]) ** 2 - ax1.loglog(dts, ref, "r--", lw=1.5, alpha=0.7, label="$\\propto dt^2$ (reference)") - - # Fit slope - log_dts = np.log(dts) - log_rmsds = np.log(rmsds) - # Only fit points where RMSD is clearly above noise floor - mask = rmsds > 1.5 * rmsds.min() - if mask.sum() >= 2: - slope, intercept = np.polyfit(log_dts[mask], log_rmsds[mask], 1) - ax1.set_title(f"Log-log slope = {slope:.2f} (expected: 2.0)") - else: - ax1.set_title("RMSD vs dt") - - ax1.set_xlabel("Timestep (ps)") - ax1.set_ylabel("RMSD of E_total (eV)") - ax1.legend() - - # --- Panel 2: Energy time series for each dt --- - for i, dt_ps in enumerate(timesteps_ps): - key = f"dt_{dt_ps}" - com = nve[f"{key}_constant_of_motion"] - steps = np.arange(len(com)) - ax2.plot(steps, com - com[0], alpha=0.7, lw=0.8, label=f"dt={dt_ps:.4f} ps") - - ax2.set_xlabel("Step") - ax2.set_ylabel("E_total - E_total[0] (eV)") - ax2.set_title("Energy drift per timestep") - ax2.legend(fontsize=7, ncol=2) - - # --- Panel 3: RMSD ratio table --- - ax3.axis("off") - headers = ["dt1 (ps)", "dt2 (ps)", "dt1/dt2", "RMSD1/RMSD2", "(dt1/dt2)²"] - rows = [] - for i in range(len(dts) - 1): - dt_ratio = dts[i] / dts[i + 1] - rmsd_ratio = rmsds[i] / rmsds[i + 1] - expected = dt_ratio ** 2 - rows.append([f"{dts[i]:.4f}", f"{dts[i+1]:.4f}", - f"{dt_ratio:.3f}", f"{rmsd_ratio:.3f}", f"{expected:.3f}"]) - - table = ax3.table(cellText=rows, colLabels=headers, loc="center", cellLoc="center") - table.auto_set_font_size(False) - table.set_fontsize(9) - table.scale(1.0, 1.5) - ax3.set_title("Convergence ratios", pad=20) - - fig.tight_layout() - plt.show() - -except FileNotFoundError: - print("No NVE data found. Run: python run_nve.py") - -# %% [markdown] -# ### physical_validation built-in integrator convergence plot -# -# Uses `physical_validation.integrator.convergence(..., screen=True)` which plots RMSD vs dt with the expected dt^2 reference line. - -# %% -try: - nve = load("nve_convergence") - timesteps_ps = nve["timesteps_ps"] - natoms = int(nve["natoms"]) - masses = nve["masses"] - volume = float(nve["volume"]) - - # Pick 3 timesteps in the dt^2 regime (avoid noise floor and nonlinear regime) - # Use [0.007, 0.005, 0.004] which showed good convergence - selected_dts = [0.007, 0.005, 0.004] - - simulations = [] - for dt_ps in selected_dts: - key = f"dt_{dt_ps}" - com = nve[f"{key}_constant_of_motion"] - dt_internal = float(nve[f"{key}_dt_internal"]) - - system = physical_validation.data.SystemData( - natoms=natoms, nconstraints=0, - ndof_reduction_tra=3, ndof_reduction_rot=0, mass=masses) - ensemble = physical_validation.data.EnsembleData( - ensemble="NVE", natoms=natoms, volume=volume) - obs = physical_validation.data.ObservableData(constant_of_motion=com) - sd = physical_validation.data.SimulationData( - units=make_unit_data(), dt=dt_internal, - system=system, ensemble=ensemble, observables=obs) - simulations.append(sd) - - result = physical_validation.integrator.convergence( - simulations, verbose=True, screen=True, - ) - print(f"\nConvergence deviation: {result:.3f} ({'PASS' if result < 0.5 else 'FAIL'})") - -except FileNotFoundError: - print("No NVE data found. Run: python run_nve.py") - -plt.show() - -# %% [markdown] -# ## 9. Summary Table -# -# Run `physical_validation` checks programmatically on all available data and collect pass/fail results. - -# %% -# Collect results -results = [] -all_integrators = NVT_INTEGRATORS + NPT_INTEGRATORS - -for name in all_integrators: - is_npt = name.startswith("npt") - ensemble = "NPT" if is_npt else "NVT" - - # --- KE Distribution at 100K --- - label_100 = f"{name}_T100K" - ke_status = "no data" - try: - d100 = load(label_100) - sd = build_sim_data(d100, 100.0, ensemble=ensemble) - d_mean, d_width = physical_validation.kinetic_energy.distribution( - sd, strict=False, verbosity=0) - ke_pass = abs(d_mean) < THRESHOLD and abs(d_width) < THRESHOLD - ke_status = f"PASS (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" if ke_pass else f"FAIL (\u03bc={d_mean:.2f}\u03c3, w={d_width:.2f}\u03c3)" - except Exception as e: - ke_status = f"ERROR: {e}" - - # --- Ensemble Check 88K vs 100K (temperature) --- - ens_status = "no data" - try: - d88 = load(f"{name}_T88K") - d100 = load(f"{name}_T100K") - sd_lo = build_sim_data(d88, 88.0, ensemble=ensemble) - sd_hi = build_sim_data(d100, 100.0, ensemble=ensemble) - quantiles = physical_validation.ensemble.check( - sd_lo, sd_hi, total_energy=True, data_is_uncorrelated=True, verbosity=0) - max_q = max(abs(q) for q in quantiles) - ens_pass = max_q < THRESHOLD - ens_status = f"PASS (max |q|={max_q:.2f}\u03c3)" if ens_pass else f"FAIL (max |q|={max_q:.2f}\u03c3)" - except Exception as e: - ens_status = f"ERROR: {e}" - - # --- Pressure Ensemble Check (NPT only, same T, different P) --- - press_status = "N/A" - if is_npt and name in pressure_data: - entries = pressure_data[name] - if len(entries) >= 2: - try: - p_lo_bar, fname_lo = entries[0] - p_hi_bar, fname_hi = entries[-1] - d_plo = load(fname_lo) - d_phi = load(fname_hi) - p_lo_eva3 = float(d_plo["external_pressure"]) - p_hi_eva3 = float(d_phi["external_pressure"]) - temp_p = float(d_plo["target_temperature"]) - sd_plo = build_sim_data(d_plo, temp_p, ensemble="NPT", pressure=p_lo_eva3) - sd_phi = build_sim_data(d_phi, temp_p, ensemble="NPT", pressure=p_hi_eva3) - quantiles_p = physical_validation.ensemble.check( - sd_plo, sd_phi, total_energy=False, data_is_uncorrelated=True, verbosity=0) - max_qp = max(abs(q) for q in quantiles_p) - press_pass = max_qp < THRESHOLD - press_status = f"PASS (max |q|={max_qp:.2f}\u03c3)" if press_pass else f"FAIL (max |q|={max_qp:.2f}\u03c3)" - except Exception as e: - press_status = f"ERROR: {e}" - else: - press_status = "need 2 pressures" - - results.append((name, ke_status, ens_status, press_status)) - -# Print table -print(f"{'Integrator':<30} {'KE Distribution':<40} {'Ensemble (T)':<35} {'Ensemble (P)':<35}") -print("-" * 140) -for name, ke, ens, press in results: - print(f"{name:<30} {ke:<40} {ens:<35} {press:<35}") - - diff --git a/fast_integrator_tests/common.py b/fast_integrator_tests/common.py deleted file mode 100644 index 69e21fd06..000000000 --- a/fast_integrator_tests/common.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Shared constants and helpers for integrator validation scripts.""" - -import numpy as np -import torch -from ase.build import bulk -from pathlib import Path - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.units import MetalUnits - -DEVICE = torch.device("cpu") -# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -DTYPE = torch.float64 -print(f"Using device: {DEVICE}") - -# LJ Argon -SIGMA = 3.405 -EPSILON = 0.0104 -CUTOFF = 2.5 * SIGMA - -DATA_DIR = Path(__file__).parent / "data" -torch.set_num_threads(4) - -def make_lj_model(compute_stress=False): - return LennardJonesModel( - use_neighbor_list=False, - sigma=SIGMA, - epsilon=EPSILON, - device=DEVICE, - dtype=DTYPE, - compute_forces=True, - compute_stress=compute_stress, - cutoff=CUTOFF, - ) - - -def make_ar_supercell(repeat=(2, 2, 2)): - atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) - return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) - - -def to_kT(temperature_K): - return temperature_K * float(MetalUnits.temperature) - - -def to_dt(timestep_ps): - return timestep_ps * float(MetalUnits.time) diff --git a/fast_integrator_tests/run_npt.py b/fast_integrator_tests/run_npt.py deleted file mode 100644 index 0b18f80ca..000000000 --- a/fast_integrator_tests/run_npt.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Run NPT simulations for all NPT integrators and save observables. - -Usage: - python run_npt.py # temperature sweep (default) - python run_npt.py --mode pressure # pressure sweep at fixed T - python run_npt.py --mode all # both sweeps - python run_npt.py --integrator npt_langevin -""" - -import argparse -import time - -import numpy as np -import torch - -import torch_sim as ts -from common import DATA_DIR, DEVICE, DTYPE, make_ar_supercell, make_lj_model, to_dt, to_kT -from torch_sim.units import MetalUnits - -NPT_INTEGRATORS = [ - "npt_langevin", - "npt_nose_hoover", - "npt_isotropic_crescale", - "npt_anisotropic_crescale", -] - -TEMPERATURES = [88.0, 100.0] -TIMESTEP_PS = 0.004 -EXTERNAL_PRESSURE = 0.0 -N_STEPS = 20_000 -N_EQUILIBRATION = 3_000 - -# Pressure sweep: two pressures at a fixed temperature. -# physical_validation compares volume distributions at same T, different P. -PRESSURE_SWEEP_TEMP = 100.0 # K -PRESSURES_EVA3 = [0.0, 0.0001] # eV/ų (0 bar and ~160 bar) - - -def run_npt(integrator_name, sim_state, model, temperature, external_pressure=0.0, - seed=42): - kT = to_kT(temperature) - dt = torch.tensor(to_dt(TIMESTEP_PS), device=DEVICE, dtype=DTYPE) - ext_p = torch.tensor(external_pressure, device=DEVICE, dtype=DTYPE) - natoms = int(sim_state.positions.shape[0]) - - sim_state = sim_state.clone() - sim_state.rng = seed - - # Initialize - if integrator_name == "npt_langevin": - # state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt) - state = ts.npt_langevin_init(sim_state, model, kT=kT, dt=dt, b_tau = 1 * dt, alpha= 5 * dt) # Better parameters - elif integrator_name == "npt_nose_hoover": - state = ts.npt_nose_hoover_init(sim_state, model, kT=kT, dt=dt, t_tau=10 * dt, b_tau=100 * dt) - elif integrator_name == "npt_isotropic_crescale": - state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt, tau_p = 10 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) - elif integrator_name == "npt_anisotropic_crescale": - state = ts.npt_crescale_init(sim_state, model, kT=kT, dt=dt/2, tau_p = 100 * dt, isothermal_compressibility = 1e-6 / MetalUnits.pressure) - else: - raise ValueError(f"Unknown: {integrator_name}") - - def step(s): - if integrator_name == "npt_langevin": - return ts.npt_langevin_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) - if integrator_name == "npt_nose_hoover": - return ts.npt_nose_hoover_step(s, model, dt=dt, kT=kT, external_pressure=ext_p) - if integrator_name == "npt_isotropic_crescale": - return ts.npt_crescale_isotropic_step( - s, model, dt=dt, kT=kT, external_pressure=ext_p, tau = 5 * dt - ) - return ts.npt_crescale_anisotropic_step( - s, model, dt=dt/2, kT=kT, external_pressure=ext_p, tau = 5 * dt - ) - - # Equilibration - print(f" Equilibrating {N_EQUILIBRATION} steps...") - for _ in range(N_EQUILIBRATION): - state = step(state) - - # Production - print(f" Producing {N_STEPS} steps...") - ke_list, pe_list, total_e_list = [], [], [] - temp_list, volume_list = [], [] - - for i in range(N_STEPS): - state = step(state) - ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) - pe = float(state.energy.sum()) - temp = float( - ts.calc_temperature(masses=state.masses, momenta=state.momenta) - ) - cell = state.cell[0].detach().cpu().numpy() - vol = float(np.abs(np.linalg.det(cell))) - - ke_list.append(ke) - pe_list.append(pe) - total_e_list.append(ke + pe) - temp_list.append(temp) - volume_list.append(vol) - - return { - "kinetic_energy": np.array(ke_list), - "potential_energy": np.array(pe_list), - "total_energy": np.array(total_e_list), - "temperature": np.array(temp_list), - "volumes": np.array(volume_list), - "masses": sim_state.masses.detach().cpu().numpy(), - "dt_internal": to_dt(TIMESTEP_PS), - "natoms": natoms, - "target_temperature": temperature, - "external_pressure": external_pressure, - "timestep_ps": TIMESTEP_PS, - "integrator": integrator_name, - } - - -def pressure_to_bar(p_eva3): - """Convert eV/ų to bar for display.""" - return p_eva3 / float(MetalUnits.pressure) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--integrator", choices=NPT_INTEGRATORS, default=None) - parser.add_argument( - "--mode", choices=["temperature", "pressure", "all"], default="all", - help="'temperature' = vary T at P=0 (default), " - "'pressure' = vary P at fixed T, 'all' = both", - ) - args = parser.parse_args() - - integrators = [args.integrator] if args.integrator else NPT_INTEGRATORS - DATA_DIR.mkdir(exist_ok=True) - - sim_state = make_ar_supercell(repeat=(3, 3, 3)) - model = make_lj_model(compute_stress=True) - - # --- Temperature sweep (same P=0, vary T) --- - if args.mode in ("temperature", "all"): - print("=== Temperature sweep ===") - for name in integrators: - for temp in TEMPERATURES: - seed = 42 if temp == TEMPERATURES[0] else 123 - label = f"{name}_T{temp:.0f}K" - print(f" Running {label} ...") - t0 = time.time() - data = run_npt(name, sim_state, model, temp, seed=seed) - elapsed = time.time() - t0 - outpath = DATA_DIR / f"{label}.npz" - np.savez(outpath, **data) - print(f" Saved {outpath.name} ({elapsed:.1f}s)") - - # --- Pressure sweep (same T, vary P) --- - if args.mode in ("pressure", "all"): - print(f"\n=== Pressure sweep at T={PRESSURE_SWEEP_TEMP:.0f}K ===") - for name in integrators: - for i, p_eva3 in enumerate(PRESSURES_EVA3): - p_bar = pressure_to_bar(p_eva3) - seed = 42 if i == 0 else 123 - label = f"{name}_T{PRESSURE_SWEEP_TEMP:.0f}K_P{p_bar:.0f}bar" - print(f" Running {label} ...") - t0 = time.time() - data = run_npt( - name, sim_state, model, PRESSURE_SWEEP_TEMP, - external_pressure=p_eva3, seed=seed, - ) - elapsed = time.time() - t0 - outpath = DATA_DIR / f"{label}.npz" - np.savez(outpath, **data) - print(f" Saved {outpath.name} ({elapsed:.1f}s)") - - -if __name__ == "__main__": - main() diff --git a/fast_integrator_tests/run_nve.py b/fast_integrator_tests/run_nve.py deleted file mode 100644 index a0d60ca9b..000000000 --- a/fast_integrator_tests/run_nve.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Run NVE simulations at multiple timesteps for convergence analysis. - -Usage: - python run_nve.py -""" - -import time - -import numpy as np -import torch - -import torch_sim as ts -from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT - -# Timesteps: sweep a range so the notebook can pick the best subset -TIMESTEPS_PS = [0.010, 0.008, 0.007, 0.006, 0.005, 0.004, 0.003, 0.002] -TEMPERATURE = 5.0 # K - low so integration error dominates -N_STEPS = 10_000 -SEED = 42 - - -def run_nve(sim_state, model, kT_init, timestep_ps): - dt = to_dt(timestep_ps) - - sim_state = sim_state.clone() - sim_state.rng = SEED - state = ts.nve_init(sim_state, model, kT=kT_init) - - ke_list, pe_list, com_list = [], [], [] - - for _ in range(N_STEPS): - state = ts.nve_step(state, model, dt=dt) - ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) - pe = float(state.energy.sum()) - ke_list.append(ke) - pe_list.append(pe) - com_list.append(ke + pe) - - return { - "kinetic_energy": np.array(ke_list), - "potential_energy": np.array(pe_list), - "constant_of_motion": np.array(com_list), - "dt_internal": dt, - "timestep_ps": timestep_ps, - } - - -def main(): - DATA_DIR.mkdir(exist_ok=True) - - # sim_state = make_ar_supercell(repeat=(1, 1, 1)) - sim_state = make_ar_supercell(repeat=(6, 6, 6)) - model = make_lj_model() - kT_init = to_kT(TEMPERATURE) - natoms = int(sim_state.positions.shape[0]) - masses = sim_state.masses.detach().cpu().numpy() - cell = sim_state.cell[0].detach().cpu().numpy() - volume = float(np.abs(np.linalg.det(cell))) - - all_results = {} - for dt_ps in TIMESTEPS_PS: - label = f"nve_dt{dt_ps:.4f}" - print(f" Running {label} ...") - t0 = time.time() - data = run_nve(sim_state, model, kT_init, dt_ps) - elapsed = time.time() - t0 - all_results[f"dt_{dt_ps}"] = data - print(f" std(E_tot) = {data['constant_of_motion'].std():.3e} ({elapsed:.1f}s)") - - # Save everything into one npz - save_dict = { - "timesteps_ps": np.array(TIMESTEPS_PS), - "temperature": TEMPERATURE, - "natoms": natoms, - "masses": masses, - "volume": volume, - "n_steps": N_STEPS, - } - for dt_ps in TIMESTEPS_PS: - key = f"dt_{dt_ps}" - for field in ("constant_of_motion", "kinetic_energy", "potential_energy"): - save_dict[f"{key}_{field}"] = all_results[key][field] - save_dict[f"{key}_dt_internal"] = all_results[key]["dt_internal"] - - outpath = DATA_DIR / "nve_convergence.npz" - np.savez(outpath, **save_dict) - print(f"\n Saved {outpath.name}") - - -if __name__ == "__main__": - main() diff --git a/fast_integrator_tests/run_nvt.py b/fast_integrator_tests/run_nvt.py deleted file mode 100644 index ba8d141cf..000000000 --- a/fast_integrator_tests/run_nvt.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Run NVT simulations for all NVT integrators and save observables. - -Usage: - python run_nvt.py - python run_nvt.py --integrator nvt_langevin -""" - -import argparse -import time - -import numpy as np -import torch - -import torch_sim as ts -from common import DATA_DIR, make_ar_supercell, make_lj_model, to_dt, to_kT - -NVT_INTEGRATORS = ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"] - -# Two temperatures for ensemble check -TEMPERATURES = [95.0, 100.0] -TIMESTEP_PS = 0.004 -N_STEPS = 10_000 -N_EQUILIBRATION = 2_000 - - -def run_nvt(integrator_name, sim_state, model, temperature, seed=42): - kT = to_kT(temperature) - dt = to_dt(TIMESTEP_PS) - natoms = int(sim_state.positions.shape[0]) - - sim_state = sim_state.clone() - sim_state.rng = seed - - # Initialize - if integrator_name == "nvt_langevin": - state = ts.nvt_langevin_init(sim_state, model, kT=kT) - elif integrator_name == "nvt_nose_hoover": - state = ts.nvt_nose_hoover_init(sim_state, model, kT=kT, dt=dt) - elif integrator_name == "nvt_vrescale": - state = ts.nvt_vrescale_init(sim_state, model, kT=kT) - else: - raise ValueError(f"Unknown: {integrator_name}") - - def step(s): - if integrator_name == "nvt_langevin": - return ts.nvt_langevin_step(s, model, dt=dt, kT=kT) - if integrator_name == "nvt_nose_hoover": - return ts.nvt_nose_hoover_step(s, model, dt=dt, kT=kT) - return ts.nvt_vrescale_step(model, s, dt=dt, kT=kT) - - # Equilibration - print(f" Equilibrating {N_EQUILIBRATION} steps...") - for _ in range(N_EQUILIBRATION): - state = step(state) - - # Production - print(f" Producing {N_STEPS} steps...") - ke_list, pe_list, total_e_list = [], [], [] - temp_list = [] - - for i in range(N_STEPS): - state = step(state) - ke = float(ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta)) - pe = float(state.energy.sum()) - temp = float( - ts.calc_temperature(masses=state.masses, momenta=state.momenta) - ) - ke_list.append(ke) - pe_list.append(pe) - total_e_list.append(ke + pe) - temp_list.append(temp) - - cell = sim_state.cell[0].detach().cpu().numpy() - volume = float(np.abs(np.linalg.det(cell))) - - return { - "kinetic_energy": np.array(ke_list), - "potential_energy": np.array(pe_list), - "total_energy": np.array(total_e_list), - "temperature": np.array(temp_list), - "volume": volume, - "masses": sim_state.masses.detach().cpu().numpy(), - "dt_internal": to_dt(TIMESTEP_PS), - "natoms": natoms, - "target_temperature": temperature, - "timestep_ps": TIMESTEP_PS, - "integrator": integrator_name, - } - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--integrator", choices=NVT_INTEGRATORS, default=None) - args = parser.parse_args() - - integrators = [args.integrator] if args.integrator else NVT_INTEGRATORS - DATA_DIR.mkdir(exist_ok=True) - - sim_state = make_ar_supercell(repeat=(6, 6, 6)) - model = make_lj_model() - - for name in integrators: - for temp in TEMPERATURES: - seed = 42 if temp == TEMPERATURES[0] else 123 - label = f"{name}_T{temp:.0f}K" - print(f" Running {label} ...") - t0 = time.time() - data = run_nvt(name, sim_state, model, temp, seed=seed) - elapsed = time.time() - t0 - outpath = DATA_DIR / f"{label}.npz" - np.savez(outpath, **data) - print(f" Saved {outpath.name} ({elapsed:.1f}s)") - - -if __name__ == "__main__": - main() From aa879250baa204cc50e2f009b7fba2d640d871f8 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:46:38 +0000 Subject: [PATCH 10/15] change to strain volume equation, add a isotropic and a version with 3 strains for x,y,z independant control --- torch_sim/__init__.py | 6 + torch_sim/integrators/__init__.py | 5 + torch_sim/integrators/npt.py | 1039 ++++++++++++++++++----------- 3 files changed, 664 insertions(+), 386 deletions(-) diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py index d9536754f..a9d939507 100644 --- a/torch_sim/__init__.py +++ b/torch_sim/__init__.py @@ -39,12 +39,15 @@ ) from torch_sim.integrators.npt import ( NPTLangevinState, + NPTLangevinStrainState, NPTNoseHooverState, npt_crescale_anisotropic_step, npt_crescale_init, npt_crescale_isotropic_step, npt_langevin_init, npt_langevin_step, + npt_langevin_strain_init, + npt_langevin_strain_step, npt_nose_hoover_init, npt_nose_hoover_invariant, npt_nose_hoover_step, @@ -124,6 +127,7 @@ "LBFGSState", "NPTLangevinState", "NPTLangevinState", + "NPTLangevinStrainState", "NPTNoseHooverState", "NPTNoseHooverState", "NVTNoseHooverState", @@ -171,6 +175,8 @@ "npt_langevin_init", "npt_langevin_step", "npt_langevin_step", + "npt_langevin_strain_init", + "npt_langevin_strain_step", "npt_nose_hoover_init", "npt_nose_hoover_init", "npt_nose_hoover_invariant", diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py index fb18a7960..7ef072456 100644 --- a/torch_sim/integrators/__init__.py +++ b/torch_sim/integrators/__init__.py @@ -81,12 +81,15 @@ from .md import MDState, initialize_momenta, momentum_step, position_step, velocity_verlet from .npt import ( NPTLangevinState, + NPTLangevinStrainState, NPTNoseHooverState, npt_crescale_anisotropic_step, npt_crescale_init, npt_crescale_isotropic_step, npt_langevin_init, npt_langevin_step, + npt_langevin_strain_init, + npt_langevin_strain_step, npt_nose_hoover_init, npt_nose_hoover_invariant, npt_nose_hoover_step, @@ -136,6 +139,7 @@ class Integrator(StrEnum): nvt_langevin = "nvt_langevin" nvt_nose_hoover = "nvt_nose_hoover" npt_langevin = "npt_langevin" + npt_langevin_strain = "npt_langevin_strain" npt_nose_hoover = "npt_nose_hoover" npt_isotropic_crescale = "npt_isotropic_crescale" npt_anisotropic_crescale = "npt_anisotropic_crescale" @@ -175,6 +179,7 @@ class Integrator(StrEnum): Integrator.nvt_langevin: (nvt_langevin_init, nvt_langevin_step), Integrator.nvt_nose_hoover: (nvt_nose_hoover_init, nvt_nose_hoover_step), Integrator.npt_langevin: (npt_langevin_init, npt_langevin_step), + Integrator.npt_langevin_strain: (npt_langevin_strain_init, npt_langevin_strain_step), Integrator.npt_nose_hoover: (npt_nose_hoover_init, npt_nose_hoover_step), Integrator.npt_isotropic_crescale: (npt_crescale_init, npt_crescale_isotropic_step), Integrator.npt_anisotropic_crescale: ( diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index 15c591c9a..3fa79d412 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -55,12 +55,15 @@ class and add their own auxiliary variables. @dataclass(kw_only=True) class NPTLangevinState(NPTState): - """State information for an NPT system with Langevin dynamics. + """State for NPT Langevin dynamics with independent per-dimension cell lengths. - This class represents the complete state of a molecular system being integrated - in the NPT (constant particle number, pressure, temperature) ensemble using - Langevin dynamics. In addition to particle positions and momenta, it tracks - cell dimensions and their dynamics for volume fluctuations. + Each spatial dimension has its own logarithmic strain coordinate + εi = ln(Li/Li0), driven by the corresponding diagonal pressure + component P_ii. This is analogous to LAMMPS ``fix press/langevin`` + with ``couple none``. + + With three identical target pressures the sum of forces equals the + isotropic strain force, so the isotropic limit is recovered. Attributes: positions (torch.Tensor): Particle positions [n_particles, n_dim] @@ -73,16 +76,19 @@ class NPTLangevinState(NPTState): system_idx (torch.Tensor): System indices [n_particles] atomic_numbers (torch.Tensor): Atomic numbers [n_particles] stress (torch.Tensor): Stress tensor [n_systems, n_dim, n_dim] - reference_cell (torch.Tensor): Original cell vectors used as reference for - scaling [n_systems, n_dim, n_dim] - cell_positions (torch.Tensor): Cell positions [n_systems, n_dim, n_dim] - cell_velocities (torch.Tensor): Cell velocities [n_systems, n_dim, n_dim] - cell_masses (torch.Tensor): Masses associated with the cell degrees of freedom - shape [n_systems] + reference_cell (torch.Tensor): Original cell [n_systems, d, d] + cell_positions (torch.Tensor): Per-dimension strain εi [n_systems, 3] + cell_velocities (torch.Tensor): dεi/dt [n_systems, 3] + cell_masses (torch.Tensor): Mass for strain DOFs [n_systems] + alpha (torch.Tensor): Particle friction [n_systems] + cell_alpha (torch.Tensor): Cell friction [n_systems] + b_tau (torch.Tensor): Barostat time constant [n_systems] Properties: momenta (torch.Tensor): Particle momenta calculated as velocities*masses with shape [n_particles, n_dimensions] + current_cell (torch.Tensor): Cell reconstructed from strain and reference_cell + volume (torch.Tensor): Current volume from cell determinant n_systems (int): Number of independent systems in the batch device (torch.device): Device on which tensors are stored dtype (torch.dtype): Data type of tensors @@ -94,8 +100,8 @@ class NPTLangevinState(NPTState): # Cell variables reference_cell: torch.Tensor - cell_positions: torch.Tensor - cell_velocities: torch.Tensor + cell_positions: torch.Tensor # (n_systems, 3) per-dimension strain + cell_velocities: torch.Tensor # (n_systems, 3) cell_masses: torch.Tensor _system_attributes = NPTState._system_attributes | { # noqa: SLF001 @@ -108,6 +114,17 @@ class NPTLangevinState(NPTState): "b_tau", } + @property + def current_cell(self) -> torch.Tensor: + """Compute cell from per-dimension strain: cell[i,:] = exp(εi) · ref[i,:].""" + scale = torch.exp(self.cell_positions) # (n_systems, 3) + return scale.unsqueeze(-1) * self.reference_cell + + @property + def volume(self) -> torch.Tensor: + """Current volume from cell determinant.""" + return torch.linalg.det(self.cell) + def _npt_langevin_beta( state: NPTLangevinState, @@ -139,10 +156,15 @@ def _npt_langevin_beta( # Map system kT to atoms atom_kT = batch_kT[state.system_idx] + atom_alpha = state.alpha[state.system_idx] + + atom_dt = dt + if dt.ndim == 0: + atom_dt = dt.expand(state.n_systems)[state.system_idx] # Calculate the prefactor for each atom # The standard deviation should be sqrt(2*alpha*kB*T*dt) - prefactor = torch.sqrt(2 * state.alpha * atom_kT * dt) + prefactor = torch.sqrt(2 * atom_alpha * atom_kT * atom_dt) return prefactor.unsqueeze(-1) * noise @@ -152,204 +174,115 @@ def _npt_langevin_cell_beta( kT: torch.Tensor, dt: torch.Tensor, ) -> torch.Tensor: - """Generate random noise for cell fluctuations in NPT dynamics. - - This function creates properly scaled random noise for cell dynamics in NPT - simulations, following the fluctuation-dissipation theorem to ensure correct - thermal sampling of cell degrees of freedom. + """Generate per-dimension noise for cell length fluctuations. Args: - state (NPTLangevinState): Current NPT state - kT (torch.Tensor): System temperature in energy units, either scalar or - with shape [n_systems] - dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] + state: Current NPT state + kT: Temperature in energy units (scalar or [n_systems]) + dt: Timestep (scalar or [n_systems]) Returns: - torch.Tensor: Scaled random noise for cell dynamics with shape - [n_systems, n_dimensions, n_dimensions] + torch.Tensor: Noise [n_systems, 3] """ - # Generate standard normal distribution (zero mean, unit variance) - noise = _randn_for_state(state, state.cell_positions.shape) - - if kT.ndim == 0: - kT = kT.expand(state.n_systems) - - # Reshape for broadcasting - cell_alpha_expanded = state.cell_alpha.view(-1, 1, 1) # shape: (n_systems, 1, 1) - kT = kT.view(-1, 1, 1) # shape: (n_systems, 1, 1) - dt = dt.expand(state.n_systems).view(-1, 1, 1) if dt.ndim == 0 else dt.view(-1, 1, 1) - - # Scale to satisfy the fluctuation-dissipation theorem - # The standard deviation should be sqrt(2*alpha*kB*T*dt) - scaling_factor = torch.sqrt(2.0 * cell_alpha_expanded * kT * dt) - - return scaling_factor * noise + noise = _randn_for_state(state, (state.n_systems, 3)) + batch_kT = kT if kT.ndim > 0 else kT.expand(state.n_systems) + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + scaling = torch.sqrt(2.0 * state.cell_alpha * batch_kT * dt_expanded) + return scaling.unsqueeze(-1) * noise def _npt_langevin_cell_position_step( state: NPTLangevinState, dt: torch.Tensor, - pressure_force: torch.Tensor, - kT: torch.Tensor, + strain_force: torch.Tensor, + cell_beta: torch.Tensor, ) -> NPTLangevinState: - """Update the cell position in NPT dynamics. - - This function updates the cell position (effectively the volume) in NPT dynamics - using the current cell velocities, pressure forces, and thermal noise. It - implements the position update part of the Langevin barostat algorithm. + """GJF position step for per-dimension strain εi. Args: - state (NPTLangevinState): Current NPT state - dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] - pressure_force (torch.Tensor): Pressure force for barostat - [n_systems, n_dim, n_dim] - kT (torch.Tensor): Target temperature in energy units, either scalar or - with shape [n_systems] + state: Current NPT state + dt: Timestep + strain_force: F_εi [n_systems, 3] + cell_beta: Noise [n_systems, 3] Returns: - NPTLangevinState: Updated state with new cell positions + Updated state with new cell_positions (strain) """ - # Calculate effective mass term - Q_2 = 2 * state.cell_masses.view(-1, 1, 1) # shape: (n_systems, 1, 1) + Q_2 = (2 * state.cell_masses).unsqueeze(-1) # (n_systems, 1) + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + dt_3 = dt_expanded.unsqueeze(-1) if dt_expanded.ndim > 0 else dt_expanded - # Ensure parameters have batch dimension - if dt.ndim == 0: - dt = dt.expand(state.n_systems) - - # Reshape for broadcasting - dt_expanded = dt.view(-1, 1, 1) - cell_alpha_expanded = state.cell_alpha.view(-1, 1, 1) - - # Calculate damping factor for cell position update - cell_b = 1 / (1 + ((cell_alpha_expanded * dt_expanded) / Q_2)) - - # Deterministic velocity contribution - c_1 = cell_b * dt_expanded * state.cell_velocities - - # Force contribution - c_2 = cell_b * dt_expanded * dt_expanded * pressure_force / Q_2 + cell_b = 1 / (1 + (state.cell_alpha.unsqueeze(-1) * dt_3) / Q_2) - # Random noise contribution (thermal fluctuations) - c_3 = cell_b * dt_expanded * _npt_langevin_cell_beta(state, kT, dt) / Q_2 + c_1 = cell_b * dt_3 * state.cell_velocities + c_2 = cell_b * dt_3 * dt_3 * strain_force / Q_2 + c_3 = cell_b * dt_3 * cell_beta / Q_2 - # Update cell positions with all contributions state.cell_positions = state.cell_positions + c_1 + c_2 + c_3 return state def _npt_langevin_cell_velocity_step( state: NPTLangevinState, - F_p_n: torch.Tensor, + F_eps_n: torch.Tensor, dt: torch.Tensor, - pressure_force: torch.Tensor, - kT: torch.Tensor, + strain_force: torch.Tensor, + cell_beta: torch.Tensor, ) -> NPTLangevinState: - """Update the cell velocities in NPT dynamics. - - This function updates the cell velocities using a Langevin-type integrator, - accounting for both deterministic forces from pressure differences and - stochastic thermal noise. It implements the velocity update part of the - Langevin barostat algorithm. + """GJF velocity step for per-dimension strain εi. Args: - state (NPTLangevinState): Current NPT state - F_p_n (torch.Tensor): Initial pressure force with shape - [n_systems, n_dimensions, n_dimensions] - dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] - pressure_force (torch.Tensor): Final pressure force - shape [n_systems, n_dim, n_dim] - kT (torch.Tensor): Temperature in energy units, either scalar or - shape [n_systems] + state: Current NPT state + F_eps_n: Initial strain force [n_systems, 3] + dt: Timestep + strain_force: Final strain force [n_systems, 3] + cell_beta: Noise (SAME as in position step) [n_systems, 3] Returns: - NPTLangevinState: Updated state with new cell velocities + Updated state with new cell_velocities """ - # Ensure parameters have batch dimension - if dt.ndim == 0: - dt = dt.expand(state.n_systems) - if kT.ndim == 0: - kT = kT.expand(state.n_systems) - - # Reshape for broadcasting - need to maintain 3x3 dimensions - dt_expanded = dt.view(-1, 1, 1) # shape: (n_systems, 1, 1) - cell_alpha_expanded = state.cell_alpha.view(-1, 1, 1) # shape: (n_systems, 1, 1) - - # Calculate cell masses per system - reshape to match 3x3 cell matrices - cell_masses_expanded = state.cell_masses.view(-1, 1, 1) # shape: (n_systems, 1, 1) - - # These factors come from the Langevin integration scheme - a = (1 - (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded)) / ( - 1 + (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded) - ) - b = 1 / (1 + (cell_alpha_expanded * dt_expanded) / (2 * cell_masses_expanded)) + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + dt_3 = dt_expanded.unsqueeze(-1) if dt_expanded.ndim > 0 else dt_expanded - # Calculate the three terms for velocity update - # a will broadcast from (n_systems, 1, 1) to (n_systems, 3, 3) - c_1 = a * state.cell_velocities # Damped old velocity + Q = state.cell_masses.unsqueeze(-1) # (n_systems, 1) + alpha_c = state.cell_alpha.unsqueeze(-1) # (n_systems, 1) + a = (1 - (alpha_c * dt_3) / (2 * Q)) / (1 + (alpha_c * dt_3) / (2 * Q)) + b = 1 / (1 + (alpha_c * dt_3) / (2 * Q)) - # Force contribution (average of initial and final forces) - c_2 = dt_expanded * ((a * F_p_n) + pressure_force) / (2 * cell_masses_expanded) - - # Generate system-specific cell noise with correct shape (n_systems, 3, 3) - cell_noise = _randn_for_state(state, state.cell_velocities.shape) - - # Calculate thermal noise amplitude - noise_prefactor = torch.sqrt( - 2 * cell_alpha_expanded * kT.view(-1, 1, 1) * dt_expanded - ) - noise_term = noise_prefactor * cell_noise / cell_masses_expanded - - # Random noise contribution - c_3 = b * noise_term + c_1 = a * state.cell_velocities + c_2 = dt_3 * ((a * F_eps_n) + strain_force) / (2 * Q) + c_3 = b * cell_beta / Q - # Update velocities with all contributions state.cell_velocities = c_1 + c_2 + c_3 return state def _npt_langevin_position_step( state: NPTLangevinState, - L_n: torch.Tensor, # This should be shape (n_systems,) + eps_old: torch.Tensor, dt: torch.Tensor, - kT: torch.Tensor, + particle_beta: torch.Tensor, ) -> NPTLangevinState: - """Update the particle positions in NPT dynamics. + """Update particle positions with per-dimension strain scaling. - This function updates particle positions accounting for both the changing - cell dimensions and the particle velocities/forces. It handles the scaling - of positions due to volume changes as well as the normal position updates - from velocities. + Each component of position is scaled by exp(εi_new - εi_old). Args: - state (NPTLangevinState): Current NPT state - L_n (torch.Tensor): Previous cell length scale with shape [n_systems] - dt: Integration timestep, either scalar or with shape [n_systems] - kT (torch.Tensor): Target temperature in energy units, either scalar or - with shape [n_systems] + state: Current state (cell_positions already updated) + eps_old: Previous strain [n_systems, 3] + dt: Timestep + particle_beta: Noise [n_particles, n_dim] Returns: - NPTLangevinState: Updated state with new positions + Updated state with new positions """ - # Calculate effective mass term by system - # Map masses to have batch dimension - M_2 = 2 * state.masses.unsqueeze(-1) # shape: (n_atoms, 1) - - # Calculate new cell length scale (cube root of volume for isotropic scaling) - L_n_new = torch.pow( - state.cell_positions.reshape(state.n_systems, -1)[:, 0], 1 / 3 - ) # shape: (n_systems,) - - # Map system-specific L_n and L_n_new to atom-level using system indices - # Make sure L_n is the right shape (n_systems,) before indexing - if L_n.ndim != 1 or L_n.shape[0] != state.n_systems: - # If L_n has wrong shape, calculate it again to ensure correct shape - L_n = torch.pow(state.cell_positions.reshape(state.n_systems, -1)[:, 0], 1 / 3) - - # Map system-specific values to atoms using system indices - L_n_atoms = L_n[state.system_idx] # shape: (n_atoms,) - L_n_new_atoms = L_n_new[state.system_idx] # shape: (n_atoms,) - - # Calculate damping factor + M_2 = 2 * state.masses.unsqueeze(-1) # (n_atoms, 1) + + # Per-dimension scale factor + scale = torch.exp(state.cell_positions - eps_old) # (n_systems, 3) + scale_atoms = scale[state.system_idx] # (n_atoms, 3) + + # Damping factor alpha_atoms = state.alpha[state.system_idx] dt_atoms = dt if dt.ndim > 0: @@ -357,30 +290,19 @@ def _npt_langevin_position_step( b = 1 / (1 + ((alpha_atoms * dt_atoms) / (2 * state.masses))) - # Scale positions due to cell volume change - c_1 = (L_n_new_atoms / L_n_atoms).unsqueeze(-1) * state.positions + # Scale each position component independently + c_1 = scale_atoms * state.positions # (n_atoms, 3) - # Time step factor with average length scale - c_2 = (2 * L_n_new_atoms / (L_n_new_atoms + L_n_atoms)) * b * dt_atoms + # Time step factor: 2·s/(s+1) per dimension + c_2 = (2 * scale_atoms / (scale_atoms + 1)) * b.unsqueeze(-1) * dt_atoms.unsqueeze(-1) - # Generate atom-specific noise - noise = _randn_for_state(state, state.momenta.shape) - batch_kT = kT - if kT.ndim == 0: - batch_kT = kT.expand(state.n_systems) - atom_kT = batch_kT[state.system_idx] - - # Calculate noise prefactor according to fluctuation-dissipation theorem - noise_prefactor = torch.sqrt(2 * alpha_atoms * atom_kT * dt_atoms) - noise_term = noise_prefactor.unsqueeze(-1) * noise - - # Velocity and force contributions with random noise c_3 = ( - state.velocities + dt_atoms.unsqueeze(-1) * state.forces / M_2 + noise_term / M_2 + state.velocities + + dt_atoms.unsqueeze(-1) * state.forces / M_2 + + particle_beta / M_2 ) - # Update positions with all contributions - state.set_constrained_positions(c_1 + c_2.unsqueeze(-1) * c_3) + state.set_constrained_positions(c_1 + c_2 * c_3) return state @@ -388,20 +310,20 @@ def _npt_langevin_velocity_step( state: NPTLangevinState, forces: torch.Tensor, dt: torch.Tensor, - kT: torch.Tensor, + particle_beta: torch.Tensor, ) -> NPTLangevinState: """Update the particle velocities in NPT dynamics. This function updates particle velocities using a Langevin-type integrator, - accounting for both deterministic forces and stochastic thermal noise. - It implements the velocity update part of the Langevin thermostat algorithm. + accounting for both deterministic forces and pre-generated thermal noise. Args: state (NPTLangevinState): Current NPT state - forces: Forces on particles + forces: Forces on particles (from before position update) dt: Integration timestep, either scalar or with shape [n_systems] - kT: Target temperature in energy units, either scalar or - with shape [n_systems] + particle_beta (torch.Tensor): Pre-generated GJF noise term β for particle + dynamics. Must be the SAME realization used in the position step. + Shape [n_particles, n_dim] Returns: NPTLangevinState: Updated state with new velocities @@ -426,19 +348,8 @@ def _npt_langevin_velocity_step( # Force contribution (average of initial and final forces) c_2 = dt_atoms.unsqueeze(-1) * ((a * forces) + state.forces) / M_2.unsqueeze(-1) - # Generate atom-specific noise - noise = _randn_for_state(state, state.momenta.shape) - batch_kT = kT - if kT.ndim == 0: - batch_kT = kT.expand(state.n_systems) - atom_kT = batch_kT[state.system_idx] - - # Calculate noise prefactor according to fluctuation-dissipation theorem - noise_prefactor = torch.sqrt(2 * alpha_atoms * atom_kT * dt_atoms) - noise_term = noise_prefactor.unsqueeze(-1) * noise - - # Random noise contribution - c_3 = b * noise_term / state.masses.unsqueeze(-1) + # GJF noise term: b * β / m + c_3 = b * particle_beta / state.masses.unsqueeze(-1) # Update momenta (velocities * masses) with all contributions new_velocities = c_1 + c_2 + c_3 @@ -449,67 +360,41 @@ def _npt_langevin_velocity_step( def _compute_cell_force( state: NPTLangevinState, - external_pressure: float | torch.Tensor, - kT: float | torch.Tensor, + external_pressure: torch.Tensor, + kT: torch.Tensor, ) -> torch.Tensor: - """Compute forces on the cell for NPT dynamics. + """Compute per-dimension force on the strain coordinates. + + F_εi = V · (P_ii - P_ext_i) - This function calculates the forces acting on the simulation cell - based on the difference between internal stress and external pressure, - plus a kinetic contribution. These forces drive the volume changes - needed to maintain constant pressure. + where P_ii = -σ_ii + N·kT/V is the ii diagonal pressure component. + The force is in energy units (eV). Args: - state (NPTLangevinState): Current NPT state - external_pressure (torch.Tensor): Target external pressure, either scalar or - tensor with shape [n_systems, n_dimensions, n_dimensions] - kT (torch.Tensor): Temperature in energy units, either scalar or - shape [n_systems] + state: Current NPT state + external_pressure: Target pressure per dimension [3] or [n_systems, 3] + kT: Temperature in energy units (scalar or [n_systems]) Returns: - torch.Tensor: Force acting on the cell [n_systems, n_dim, n_dim] + torch.Tensor: Force per dimension [n_systems, 3] """ - external_pressure = torch.as_tensor( - external_pressure, device=state.device, dtype=state.dtype - ) - kT = torch.as_tensor(kT, device=state.device, dtype=state.dtype) + volumes = state.volume # (n_systems,) - # Get current volumes for each batch - volumes = torch.linalg.det(state.cell) # shape: (n_systems,) - - # Reshape for broadcasting - volumes = volumes.view(-1, 1, 1) # shape: (n_systems, 1, 1) - - # Create pressure tensor (diagonal with external pressure) - if external_pressure.ndim == 0: - # Scalar pressure - create diagonal pressure tensors for each batch - pressure_tensor = external_pressure * torch.eye( - 3, device=state.device, dtype=state.dtype - ) - pressure_tensor = pressure_tensor.unsqueeze(0).expand(state.n_systems, -1, -1) - elif external_pressure.ndim == 1: - # Per-system scalar pressures (n_systems,) - create diagonal tensors - eye = torch.eye(3, device=state.device, dtype=state.dtype) - pressure_tensor = external_pressure.view(-1, 1, 1) * eye.unsqueeze(0) - else: - # Already a tensor with shape (n_systems, 3, 3) - pressure_tensor = external_pressure + # Diagonal stress components \sigma_ii + stress_diag = torch.diagonal(state.stress, dim1=-2, dim2=-1) # (n_systems, 3) - # Calculate virials from stress and external pressure - # Internal stress is negative of virial tensor divided by volume - virial = -volumes * (state.stress + pressure_tensor) + # P_ii = -\sigma_ii (virial part) + P_virial_diag = -stress_diag # (n_systems, 3) - # Add kinetic contribution (kT * Identity) - batch_kT = kT - if kT.ndim == 0: - batch_kT = kT.expand(state.n_systems) + # Kinetic contribution per dimension: N·kT/V (target temperature) + batch_kT = kT if kT.ndim > 0 else kT.expand(state.n_systems) + n_atoms = state.n_atoms_per_system.to(dtype=state.dtype) + kinetic_pressure = (n_atoms * batch_kT / volumes).unsqueeze(-1) # (n_systems, 1) - e_kin_per_atom = batch_kT.view(-1, 1, 1) * torch.eye( - 3, device=state.device, dtype=state.dtype - ).unsqueeze(0) + P_diag = P_virial_diag + kinetic_pressure # (n_systems, 3) - # Correct implementation with scaling by n_atoms_per_system - return virial + e_kin_per_atom * state.n_atoms_per_system.view(-1, 1, 1) + # F_εi = V · (P_ii - P_ext_i) + return volumes.unsqueeze(-1) * (P_diag - external_pressure) def npt_langevin_init( @@ -523,48 +408,34 @@ def npt_langevin_init( b_tau: float | torch.Tensor | None = None, **_kwargs: Any, ) -> NPTLangevinState: - """Initialize an NPT Langevin state from input data. + """Initialize NPT Langevin state with independent per-dimension cell lengths. - This function creates the initial state for NPT Langevin dynamics, - setting up all necessary variables including particle velocities, - cell parameters, and barostat variables. It computes initial forces - and stress using the provided model. + Each spatial dimension gets its own strain DOF εi = ln(Li/Li0), + driven by the corresponding diagonal pressure component. To seed the RNG set ``state.rng = seed`` before calling. Args: - model (ModelInterface): Neural network model that computes energies, forces, - and stress. Must return a dict with 'energy', 'forces', and 'stress' keys. - state (SimState): SimState containing positions, masses, cell, pbc - kT (torch.Tensor): Target temperature in energy units, either scalar or - with shape [n_systems] - dt (torch.Tensor): Integration timestep, either scalar or shape [n_systems] - alpha (torch.Tensor, optional): Friction coefficient for particle Langevin - thermostat, either scalar or shape [n_systems]. Defaults to 1/(100*dt). - cell_alpha (torch.Tensor, optional): Friction coefficient for cell Langevin - thermostat, either scalar or shape [n_systems]. Defaults to same as alpha. - b_tau (torch.Tensor, optional): Barostat time constant controlling how quickly - the system responds to pressure differences, either scalar or shape - [n_systems]. Defaults to 1/(1000*dt). + state: SimState containing positions, masses, cell, pbc + model: Model computing energy, forces, stress + kT: Target temperature in energy units + dt: Integration timestep + alpha: Particle friction. Defaults to 1/(5·dt). + cell_alpha: Cell friction. Defaults to 1/(30·dt). + b_tau: Barostat time constant. Defaults to 300·dt. Returns: - NPTLangevinState: Initialized state for NPT Langevin integration containing - all required attributes for particle and cell dynamics - - Notes: - - The model must provide stress tensor calculations for proper pressure coupling + NPTLangevinState with εi = 0 for all dimensions """ device, dtype = model.device, model.dtype - # Set default values if not provided if alpha is None: - alpha = 1.0 / (20 * dt) # Default friction based on timestep + alpha = 1.0 / (5 * dt) if cell_alpha is None: - cell_alpha = 1.0 / (100 * dt) # Default cell friction based on timestep + cell_alpha = 1.0 / (30 * dt) if b_tau is None: - b_tau = 0.1 * dt # Default barostat time constant + b_tau = 300 * dt - # Convert all parameters to tensors with correct device and dtype alpha = torch.as_tensor(alpha, device=device, dtype=dtype) cell_alpha = torch.as_tensor(cell_alpha, device=device, dtype=dtype) b_tau = torch.as_tensor(b_tau, device=device, dtype=dtype) @@ -578,10 +449,8 @@ def npt_langevin_init( if b_tau.ndim == 0: b_tau = b_tau.expand(state.n_systems) - # Get model output to initialize forces and stress model_output = model(state) - # Initialize momenta if not provided momenta = getattr(state, "momenta", None) if momenta is None: momenta = initialize_momenta( @@ -592,38 +461,26 @@ def npt_langevin_init( state.rng, ) - # Initialize cell parameters reference_cell = state.cell.clone() + dim = state.positions.shape[1] - # Calculate initial cell_positions (volume) - cell_positions = ( - torch.linalg.det(state.cell).unsqueeze(-1).unsqueeze(-1) - ) # shape: (n_systems, 1, 1) - - # Initialize cell velocities to zero - cell_velocities = torch.zeros((state.n_systems, 3, 3), device=device, dtype=dtype) + # εi = 0 at initialization (V = V₀) + cell_positions = torch.zeros(state.n_systems, dim, device=device, dtype=dtype) + cell_velocities = torch.zeros(state.n_systems, dim, device=device, dtype=dtype) - # Calculate cell masses based on system size and temperature - # This follows standard NPT barostat mass scaling + batch_kT = kT.expand(state.n_systems) if kT.ndim == 0 else kT n_atoms_per_system = torch.bincount(state.system_idx) - batch_kT = ( - kT.expand(state.n_systems) - if isinstance(kT, torch.Tensor) and kT.ndim == 0 - else kT - ) cell_masses = (n_atoms_per_system + 1) * batch_kT * b_tau * b_tau if state.constraints: - # warn if constraints are present msg = ( "Constraints are present in the system. " - "Make sure they are compatible with NPT Langevin dynamics." + "Make sure they are compatible with NPT Langevin dynamics. " "We recommend not using constraints with NPT dynamics for now." ) warnings.warn(msg, UserWarning, stacklevel=3) logger.warning(msg) - # Create the initial state return NPTLangevinState.from_state( state, momenta=momenta, @@ -649,80 +506,48 @@ def npt_langevin_step( kT: float | torch.Tensor, external_pressure: float | torch.Tensor, ) -> NPTLangevinState: - r"""Perform one complete NPT Langevin dynamics integration step. + r"""Perform one NPT Langevin step with independent per-dimension cell lengths. Implements constant-pressure Langevin dynamics based on Gronbech-Jensen & Farago (2014) [4]_ and the LAMMPS ``fix press/langevin`` scheme [5]_. - **Particle equations** (GJF-style Langevin with volume coupling): + Each spatial dimension *i* has its own logarithmic strain + :math:`\varepsilon_i = \ln(L_i/L_{i,0})` driven by the diagonal + pressure component :math:`P_{ii}`. + + **Per-dimension strain force:** .. math:: - \mathbf{r}_i(t{+}\Delta t) &= \frac{L_{n+1}}{L_n}\,\mathbf{r}_i(t) - + \frac{2L_{n+1}}{L_{n+1}{+}L_n}\,b\,\Delta t - \left(\mathbf{v}_i + \frac{\Delta t\,\mathbf{F}_i}{2m_i} - + \frac{\boldsymbol{\beta}_i}{2m_i}\right) \\ - \mathbf{v}_i(t{+}\Delta t) &= a\,\mathbf{v}_i(t) - + \frac{\Delta t}{2m_i}\bigl(a\,\mathbf{F}_i^n + \mathbf{F}_i^{n+1}\bigr) - + \frac{b\,\boldsymbol{\beta}_i}{m_i} + F_{\varepsilon_i} = V \cdot (P_{ii} - P_{\text{ext},i}) + + where :math:`P_{ii} = -\sigma_{ii} + N k_B T / V`. - with damping coefficients - :math:`a = \frac{1 - \alpha\Delta t/(2m)}{1 + \alpha\Delta t/(2m)}`, - :math:`b = \frac{1}{1 + \alpha\Delta t/(2m)}`, and fluctuation-dissipation - noise :math:`\boldsymbol{\beta}_i = \sqrt{2\alpha k_BT\Delta t}\;\mathbf{R}`, - :math:`\mathbf{R}\sim\mathcal{N}(0,I)`. + With three identical target pressures the sum + :math:`\sum_i F_{\varepsilon_i}` equals the isotropic strain force. - **Cell (volume) equations:** + **Cell reconstruction:** .. math:: - V(t{+}\Delta t) &= V(t) + b_c\,\Delta t\,\dot{V} - + b_c\,\frac{\Delta t^2\,F_p}{2Q} - + b_c\,\frac{\Delta t\,\beta_c}{2Q} \\ - \dot{V}(t{+}\Delta t) &= a_c\,\dot{V}(t) - + \frac{\Delta t}{2Q}\bigl(a_c\,F_p^n + F_p^{n+1}\bigr) - + \frac{b_c\,\beta_c}{Q} + \mathbf{h}_i = e^{\varepsilon_i}\,\mathbf{h}_{i,0} - where :math:`Q = (N+1) k_BT \tau_b^2` is the cell mass and - :math:`F_p = -V(\boldsymbol{\sigma} + P_{\text{ext}}\mathbf{I}) - + Nk_BT\mathbf{I}` is the cell force. - :math:`L = V^{1/3}` is the isotropic cell length scale. + **Particle scaling (per component):** - **Variable mapping (equation -> code):** + .. math:: - ============================================ ============================ - Equation symbol Code variable - ============================================ ============================ - :math:`\mathbf{r}_i` (positions) ``state.positions`` - :math:`\mathbf{v}_i` (velocities) ``state.velocities`` - :math:`\mathbf{p}_i` (momenta) ``state.momenta`` - :math:`m_i` (masses) ``state.masses`` - :math:`\mathbf{F}_i` (forces) ``state.forces`` - :math:`\alpha` (particle friction) ``state.alpha`` - :math:`\alpha_c` (cell friction) ``state.cell_alpha`` - :math:`V` (volume) ``state.cell_positions`` - :math:`\dot{V}` (cell velocity) ``state.cell_velocities`` - :math:`Q` (cell mass) ``state.cell_masses`` - :math:`\tau_b` (barostat time const) ``state.b_tau`` - :math:`L` (cell length scale) ``L_n`` - :math:`F_p` (cell force) ``F_p_n`` - :math:`P_{\text{ext}}` (target pressure) ``external_pressure`` - :math:`k_BT` (thermal energy) ``kT`` - :math:`\Delta t` (timestep) ``dt`` - ============================================ ============================ + r_{k,i} \to e^{\varepsilon_i^{n+1} - \varepsilon_i^n}\, r_{k,i} Args: - model: Neural network model that computes energies, forces, - and stress. Must return a dict with 'energy', 'forces', and 'stress' keys. - state: Current NPT state with particle and cell variables - dt: Integration timestep, either scalar or shape [n_systems] - kT: Target temperature in energy units, either scalar or - shape [n_systems] - external_pressure: Target external pressure, - either scalar or tensor with shape [n_systems, n_dim, n_dim] + state: Current NPT state + model: Model computing energy, forces, stress + dt: Integration timestep + kT: Target temperature in energy units + external_pressure: Target pressure — scalar (same for all dims), + shape [3] (per-dimension), or [n_systems, 3] Returns: - NPTLangevinState: Updated NPT state after one timestep + NPTLangevinState: Updated state References: .. [4] Gronbech-Jensen, N. & Farago, O. "Constant pressure and temperature @@ -732,7 +557,6 @@ def npt_langevin_step( """ device, dtype = model.device, model.dtype - # Convert any scalar parameters to tensors with batch dimension if needed state.alpha = torch.as_tensor(state.alpha, device=device, dtype=dtype) kT_tensor = torch.as_tensor(kT, device=device, dtype=dtype) state.cell_alpha = torch.as_tensor(state.cell_alpha, device=device, dtype=dtype) @@ -741,67 +565,510 @@ def npt_langevin_step( external_pressure, device=device, dtype=dtype ) - # Make sure parameters have batch dimension if they're scalars + # Broadcast external_pressure to (n_systems, 3) + if external_pressure_tensor.ndim == 0: + external_pressure_tensor = external_pressure_tensor.expand(state.n_systems, 3) + elif external_pressure_tensor.ndim == 1 and external_pressure_tensor.shape[0] == 3: + external_pressure_tensor = external_pressure_tensor.unsqueeze(0).expand( + state.n_systems, 3 + ) + batch_kT = kT_tensor.expand(state.n_systems) if kT_tensor.ndim == 0 else kT_tensor - # Update barostat mass based on current temperature - # This ensures proper coupling between system and barostat + # Update barostat mass n_atoms_per_system = torch.bincount(state.system_idx) state.cell_masses = (n_atoms_per_system + 1) * batch_kT * torch.square(state.b_tau) - # Store initial values for integration + # Store initial values forces = state.forces - F_p_n = _compute_cell_force( + eps_old = state.cell_positions.clone() + + F_eps_n = _compute_cell_force( state=state, external_pressure=external_pressure_tensor, kT=kT_tensor, ) - L_n = torch.pow( - state.cell_positions.reshape(state.n_systems, -1)[:, 0], 1 / 3 - ) # shape: (n_systems,) - # Step 1: Update cell position - state = _npt_langevin_cell_position_step(state, dt_tensor, F_p_n, kT_tensor) + # Generate GJF noise ONCE + cell_beta = _npt_langevin_cell_beta(state, kT_tensor, dt_tensor) + particle_beta = _npt_langevin_beta(state, kT_tensor, dt_tensor) - # Update cell (currently only isotropic fluctuations) - dim = state.positions.shape[1] # Usually 3 for 3D - # V_0 and V are shape: (n_systems,) - V_0 = torch.linalg.det(state.reference_cell) - V = state.cell_positions.reshape(state.n_systems, -1)[:, 0] + # Step 1: Update per-dimension strain + state = _npt_langevin_cell_position_step(state, dt_tensor, F_eps_n, cell_beta) - # Scale cell uniformly in all dimensions - scaling = (V / V_0) ** (1.0 / dim) # shape: (n_systems,) - - # Apply scaling to reference cell to get new cell - new_cell = torch.zeros_like(state.cell) - for sys_idx in range(state.n_systems): - new_cell[sys_idx] = scaling[sys_idx] * state.reference_cell[sys_idx] - - state.cell = new_cell + # Reconstruct cell from updated strain + state.cell = state.current_cell # Step 2: Update particle positions - state = _npt_langevin_position_step(state, L_n, dt_tensor, kT_tensor) + state = _npt_langevin_position_step(state, eps_old, dt_tensor, particle_beta) - # Recompute model output after position updates + # Recompute model output model_output = model(state) state.energy = model_output["energy"] state.forces = model_output["forces"] state.stress = model_output["stress"] - # Compute updated pressure force - F_p_n_new = _compute_cell_force( + # Updated strain force + F_eps_new = _compute_cell_force( state=state, external_pressure=external_pressure_tensor, kT=kT_tensor, ) - # Step 3: Update cell velocities + # Step 3: Update strain velocities (uses SAME cell_beta) state = _npt_langevin_cell_velocity_step( - state, F_p_n, dt_tensor, F_p_n_new, kT_tensor + state, F_eps_n, dt_tensor, F_eps_new, cell_beta + ) + + # Step 4: Update particle velocities (uses SAME particle_beta) + return _npt_langevin_velocity_step(state, forces, dt_tensor, particle_beta) + + +# ============================================================================= +# NPT Langevin Strain integrator — isotropic logarithmic strain coordinate +# ============================================================================= + + +@dataclass(kw_only=True) +class NPTLangevinStrainState(NPTState): + """State for NPT Langevin dynamics using logarithmic strain coordinate. + + The cell degree of freedom is the isotropic logarithmic strain + ε = (1/d)·ln(V/V₀), which is dimensionless. This guarantees V > 0 + and gives the conjugate force F_ε = d·V·(P_avg - P_ext) in energy units, + providing numerically well-scaled dynamics. + + Attributes: + reference_cell (torch.Tensor): Original cell [n_systems, d, d] + cell_positions (torch.Tensor): Strain ε = (1/d)·ln(V/V₀) [n_systems] + cell_velocities (torch.Tensor): dε/dt [n_systems] + cell_masses (torch.Tensor): Mass for strain DOF [n_systems] + alpha (torch.Tensor): Particle friction [n_systems] + cell_alpha (torch.Tensor): Cell friction [n_systems] + b_tau (torch.Tensor): Barostat time constant [n_systems] + """ + + alpha: torch.Tensor + cell_alpha: torch.Tensor + b_tau: torch.Tensor + + reference_cell: torch.Tensor + cell_positions: torch.Tensor # strain ε (dimensionless) + cell_velocities: torch.Tensor # dε/dt + cell_masses: torch.Tensor + + _system_attributes = NPTState._system_attributes | { # noqa: SLF001 + "cell_positions", + "cell_velocities", + "cell_masses", + "reference_cell", + "alpha", + "cell_alpha", + "b_tau", + } + + @property + def current_cell(self) -> torch.Tensor: + """Compute cell from strain: cell = exp(ε) · reference_cell.""" + scale = torch.exp(self.cell_positions) # exp(ε), shape (n_systems,) + return scale.unsqueeze(-1).unsqueeze(-1) * self.reference_cell + + @property + def volume(self) -> torch.Tensor: + """Current volume V = V₀ · exp(d·ε).""" + dim = self.positions.shape[1] + V_0 = torch.linalg.det(self.reference_cell) + return V_0 * torch.exp(dim * self.cell_positions) + + +def _compute_strain_cell_force( + state: NPTLangevinStrainState, + external_pressure: float | torch.Tensor, + kT: float | torch.Tensor, +) -> torch.Tensor: + """Compute force on the strain coordinate ε. + + F_ε = d · V · (P_avg - P_ext) + + where P_avg = -(1/3)Tr(σ) + NkT/V and d·V is the Jacobian dV/dε. + This force is in energy units (eV), making it numerically well-scaled. + + Args: + state: Current strain-based NPT state + external_pressure: Target pressure (scalar or [n_systems]) + kT: Temperature in energy units (scalar or [n_systems]) + + Returns: + torch.Tensor: Force on strain per system [n_systems] + """ + external_pressure = torch.as_tensor( + external_pressure, device=state.device, dtype=state.dtype + ) + kT = torch.as_tensor(kT, device=state.device, dtype=state.dtype) + + dim = state.positions.shape[1] + volumes = state.volume # (n_systems,) + + # Isotropic virial pressure: P_virial = -(1/3)Tr(stress) + stress_trace = torch.einsum("nii->n", state.stress) + avg_virial_pressure = -stress_trace / 3 # (n_systems,) + + # Kinetic contribution: NkT/V + batch_kT = kT if kT.ndim > 0 else kT.expand(state.n_systems) + n_atoms = state.n_atoms_per_system.to(dtype=state.dtype) + kinetic_pressure = n_atoms * batch_kT / volumes # (n_systems,) + + if external_pressure.ndim >= 2: + raise ValueError( + f"External pressure tensor provided with shape {external_pressure.shape}. " + "Only scalar or per-system external pressure is supported." + ) + + P_avg = avg_virial_pressure + kinetic_pressure + # F_ε = d · V · (P_avg - P_ext) + return dim * volumes * (P_avg - external_pressure) + + +def _npt_langevin_strain_cell_beta( + state: NPTLangevinStrainState, + kT: torch.Tensor, + dt: torch.Tensor, +) -> torch.Tensor: + """Generate scalar random noise for isotropic strain fluctuations. + + Returns: + torch.Tensor: Noise [n_systems] + """ + noise = _randn_for_state(state, (state.n_systems,)) + batch_kT = kT if kT.ndim > 0 else kT.expand(state.n_systems) + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + scaling = torch.sqrt(2.0 * state.cell_alpha * batch_kT * dt_expanded) + return scaling * noise + + +def _npt_langevin_strain_cell_position_step( + state: NPTLangevinStrainState, + dt: torch.Tensor, + strain_force: torch.Tensor, + cell_beta: torch.Tensor, +) -> NPTLangevinStrainState: + """GJF position step for the strain coordinate ε. + + ε_{n+1} = ε_n + b·dt·dε/dt + b·dt²·F_ε/(2Q) + b·dt·β/(2Q) + + Args: + state: Current state + dt: Timestep + strain_force: F_ε [n_systems] + cell_beta: Noise term β_c [n_systems] + + Returns: + Updated state with new cell_positions (strain) + """ + Q_2 = 2 * state.cell_masses + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + + cell_b = 1 / (1 + (state.cell_alpha * dt_expanded) / Q_2) + + c_1 = cell_b * dt_expanded * state.cell_velocities + c_2 = cell_b * dt_expanded * dt_expanded * strain_force / Q_2 + c_3 = cell_b * dt_expanded * cell_beta / Q_2 + + state.cell_positions = state.cell_positions + c_1 + c_2 + c_3 + return state + + +def _npt_langevin_strain_cell_velocity_step( + state: NPTLangevinStrainState, + F_eps_n: torch.Tensor, + dt: torch.Tensor, + strain_force: torch.Tensor, + cell_beta: torch.Tensor, +) -> NPTLangevinStrainState: + """GJF velocity step for the strain coordinate ε. + + dε/dt_{n+1} = a·dε/dt_n + dt/(2Q)·(a·F_ε^n + F_ε^{n+1}) + b·β/Q + + Args: + state: Current state + F_eps_n: Initial strain force [n_systems] + dt: Timestep + strain_force: Final strain force [n_systems] + cell_beta: Noise term β_c (SAME as in position step) [n_systems] + + Returns: + Updated state with new cell_velocities (dε/dt) + """ + dt_expanded = dt if dt.ndim > 0 else dt.expand(state.n_systems) + + Q = state.cell_masses + a = (1 - (state.cell_alpha * dt_expanded) / (2 * Q)) / ( + 1 + (state.cell_alpha * dt_expanded) / (2 * Q) + ) + b = 1 / (1 + (state.cell_alpha * dt_expanded) / (2 * Q)) + + c_1 = a * state.cell_velocities + c_2 = dt_expanded * ((a * F_eps_n) + strain_force) / (2 * Q) + c_3 = b * cell_beta / Q + + state.cell_velocities = c_1 + c_2 + c_3 + return state + + +def _npt_langevin_strain_position_step( + state: NPTLangevinStrainState, + eps_old: torch.Tensor, + dt: torch.Tensor, + particle_beta: torch.Tensor, +) -> NPTLangevinStrainState: + """Update particle positions accounting for strain change. + + Positions are scaled by exp(ε_new - ε_old) for the volume change, + then the standard GJF position update is applied. + + Args: + state: Current state (cell_positions already updated to ε_new) + eps_old: Strain before the cell position step [n_systems] + dt: Timestep + particle_beta: Noise [n_particles, n_dim] + + Returns: + Updated state with new positions + """ + M_2 = 2 * state.masses.unsqueeze(-1) # (n_atoms, 1) + + # Scale factor from strain change: L_new/L_old = exp(ε_new - ε_old) + scale = torch.exp(state.cell_positions - eps_old) # (n_systems,) + scale_atoms = scale[state.system_idx] # (n_atoms,) + + # Damping factor + alpha_atoms = state.alpha[state.system_idx] + dt_atoms = dt + if dt.ndim > 0: + dt_atoms = dt[state.system_idx] + + b = 1 / (1 + ((alpha_atoms * dt_atoms) / (2 * state.masses))) + + # Scale positions due to volume change + c_1 = scale_atoms.unsqueeze(-1) * state.positions + + # Time step factor: 2·s/(s+1) where s = scale + c_2 = (2 * scale_atoms / (scale_atoms + 1)) * b * dt_atoms + + c_3 = ( + state.velocities + + dt_atoms.unsqueeze(-1) * state.forces / M_2 + + particle_beta / M_2 + ) + + state.set_constrained_positions(c_1 + c_2.unsqueeze(-1) * c_3) + return state + + +def npt_langevin_strain_init( + state: SimState, + model: ModelInterface, + *, + kT: float | torch.Tensor, + dt: float | torch.Tensor, + alpha: float | torch.Tensor | None = None, + cell_alpha: float | torch.Tensor | None = None, + b_tau: float | torch.Tensor | None = None, + **_kwargs: Any, +) -> NPTLangevinStrainState: + """Initialize an NPT Langevin state using logarithmic strain coordinate. + + The strain coordinate ε = (1/d)·ln(V/V₀) provides well-scaled dynamics + where the conjugate force F_ε = d·V·(P_avg - P_ext) is in energy units. + + Args: + state: Initial SimState + model: Model that computes energy, forces, stress + kT: Target temperature in energy units + dt: Integration timestep + alpha: Particle friction coefficient. Defaults to 1/(5·dt). + cell_alpha: Cell friction coefficient. Defaults to 1/(30·dt). + b_tau: Barostat time constant. Defaults to 300·dt. + + Returns: + NPTLangevinStrainState: Initialized state with ε = 0 + """ + device, dtype = model.device, model.dtype + + if alpha is None: + alpha = 1.0 / (5 * dt) + if cell_alpha is None: + cell_alpha = 1.0 / (30 * dt) + if b_tau is None: + b_tau = 300 * dt + + alpha = torch.as_tensor(alpha, device=device, dtype=dtype) + cell_alpha = torch.as_tensor(cell_alpha, device=device, dtype=dtype) + b_tau = torch.as_tensor(b_tau, device=device, dtype=dtype) + kT = torch.as_tensor(kT, device=device, dtype=dtype) + dt = torch.as_tensor(dt, device=device, dtype=dtype) + + if alpha.ndim == 0: + alpha = alpha.expand(state.n_systems) + if cell_alpha.ndim == 0: + cell_alpha = cell_alpha.expand(state.n_systems) + if b_tau.ndim == 0: + b_tau = b_tau.expand(state.n_systems) + + model_output = model(state) + + momenta = getattr(state, "momenta", None) + if momenta is None: + momenta = initialize_momenta( + state.positions, + state.masses, + state.system_idx, + kT, + state.rng, + ) + + reference_cell = state.cell.clone() + + # ε = 0 at initialization (V = V₀) + cell_positions = torch.zeros(state.n_systems, device=device, dtype=dtype) + cell_velocities = torch.zeros(state.n_systems, device=device, dtype=dtype) + + batch_kT = kT.expand(state.n_systems) if kT.ndim == 0 else kT + n_atoms_per_system = torch.bincount(state.system_idx) + cell_masses = (n_atoms_per_system + 1) * batch_kT * b_tau * b_tau + + if state.constraints: + msg = ( + "Constraints are present in the system. " + "Make sure they are compatible with NPT Langevin dynamics. " + "We recommend not using constraints with NPT dynamics for now." + ) + warnings.warn(msg, UserWarning, stacklevel=3) + logger.warning(msg) + + return NPTLangevinStrainState.from_state( + state, + momenta=momenta, + energy=model_output["energy"], + forces=model_output["forces"], + stress=model_output["stress"], + alpha=alpha, + b_tau=b_tau, + reference_cell=reference_cell, + cell_positions=cell_positions, + cell_velocities=cell_velocities, + cell_masses=cell_masses, + cell_alpha=cell_alpha, + ) + + +@dcite("10.1063/1.4901303") +def npt_langevin_strain_step( + state: NPTLangevinStrainState, + model: ModelInterface, + *, + dt: float | torch.Tensor, + kT: float | torch.Tensor, + external_pressure: float | torch.Tensor, +) -> NPTLangevinStrainState: + r"""Perform one NPT Langevin step using logarithmic strain coordinate. + + Uses the same GJF integrator as :func:`npt_langevin_step` but with the + cell degree of freedom being the isotropic logarithmic strain + :math:`\varepsilon = \frac{1}{d}\ln(V/V_0)` instead of the raw volume. + + **Strain force:** + + .. math:: + + F_\varepsilon = d \cdot V \cdot (P_{\text{avg}} - P_{\text{ext}}) + + where the Jacobian :math:`dV/d\varepsilon = d \cdot V` naturally provides + a volume factor that makes :math:`F_\varepsilon` an energy (eV), giving + numerically well-scaled dynamics. + + **Cell reconstruction:** + + .. math:: + + V = V_0 \exp(d\,\varepsilon), \quad + \mathbf{h} = e^\varepsilon \, \mathbf{h}_0 + + **Particle scaling:** + + .. math:: + + \mathbf{r}_i \to e^{\varepsilon_{n+1} - \varepsilon_n} \, \mathbf{r}_i + + Args: + state: Current strain-based NPT state + model: Model computing energy, forces, stress + dt: Integration timestep + kT: Target temperature in energy units + external_pressure: Target pressure + + Returns: + NPTLangevinStrainState: Updated state + """ + device, dtype = model.device, model.dtype + + state.alpha = torch.as_tensor(state.alpha, device=device, dtype=dtype) + kT_tensor = torch.as_tensor(kT, device=device, dtype=dtype) + state.cell_alpha = torch.as_tensor(state.cell_alpha, device=device, dtype=dtype) + dt_tensor = torch.as_tensor(dt, device=device, dtype=dtype) + external_pressure_tensor = torch.as_tensor( + external_pressure, device=device, dtype=dtype + ) + + batch_kT = kT_tensor.expand(state.n_systems) if kT_tensor.ndim == 0 else kT_tensor + + # Update barostat mass + n_atoms_per_system = torch.bincount(state.system_idx) + state.cell_masses = (n_atoms_per_system + 1) * batch_kT * torch.square(state.b_tau) + + # Store initial values + forces = state.forces + eps_old = state.cell_positions.clone() + + F_eps_n = _compute_strain_cell_force( + state=state, + external_pressure=external_pressure_tensor, + kT=kT_tensor, + ) + + # Generate GJF noise ONCE + cell_beta = _npt_langevin_strain_cell_beta(state, kT_tensor, dt_tensor) + particle_beta = _npt_langevin_beta(state, kT_tensor, dt_tensor) + + # Step 1: Update strain (cell position step) + state = _npt_langevin_strain_cell_position_step(state, dt_tensor, F_eps_n, cell_beta) + + # Reconstruct cell from updated strain + state.cell = state.current_cell + + # Step 2: Update particle positions (with strain-based scaling) + state = _npt_langevin_strain_position_step(state, eps_old, dt_tensor, particle_beta) + + # Recompute model output + model_output = model(state) + state.energy = model_output["energy"] + state.forces = model_output["forces"] + state.stress = model_output["stress"] + + # Compute updated strain force + F_eps_new = _compute_strain_cell_force( + state=state, + external_pressure=external_pressure_tensor, + kT=kT_tensor, + ) + + # Step 3: Update strain velocity (uses SAME cell_beta) + state = _npt_langevin_strain_cell_velocity_step( + state, F_eps_n, dt_tensor, F_eps_new, cell_beta ) - # Step 4: Update particle velocities - return _npt_langevin_velocity_step(state, forces, dt_tensor, kT_tensor) + # Step 4: Update particle velocities (uses SAME particle_beta) + return _npt_langevin_velocity_step(state, forces, dt_tensor, particle_beta) @dataclass(kw_only=True) From 9e2588442e39d4fc8d8ddcc99f4f5f2665e83a0a Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:47:41 +0000 Subject: [PATCH 11/15] add test langevin_strain and update values for nose_hoover (I changed the default parameters) --- tests/test_integrators.py | 78 ++++++++++++++++++++++++++++++++++----- 1 file changed, 69 insertions(+), 9 deletions(-) diff --git a/tests/test_integrators.py b/tests/test_integrators.py index 27ea5dbae..0ada0f06a 100644 --- a/tests/test_integrators.py +++ b/tests/test_integrators.py @@ -155,6 +155,66 @@ def test_npt_langevin( assert pos_diff > 0.0001 # Systems should remain separated +def test_npt_langevin_strain( + ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel +) -> None: + n_steps = 200 + dt = torch.tensor(0.001, dtype=DTYPE) * MetalUnits.time + kT = torch.tensor(300.0, dtype=DTYPE) * MetalUnits.temperature + external_pressure = torch.tensor(10.0, dtype=DTYPE) * MetalUnits.pressure + alpha = 1 * dt + cell_alpha = 10 * dt + b_tau = 30 * dt + + ar_double_sim_state.rng = 42 + state = ts.npt_langevin_strain_init( + state=ar_double_sim_state, + model=lj_model, + dt=dt, + kT=kT, + alpha=alpha, + cell_alpha=cell_alpha, + b_tau=b_tau, + ) + + # Check strain state shape + assert state.cell_positions.shape == (2,) # scalar strain per system + + energies = [] + temperatures = [] + for _step in range(n_steps): + state = ts.npt_langevin_strain_step( + state=state, + model=lj_model, + dt=dt, + kT=kT, + external_pressure=external_pressure, + ) + + temp = ts.calc_kT( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + energies.append(state.energy) + temperatures.append(temp / MetalUnits.temperature) + + temperatures_tensor = torch.stack(temperatures) + energies_tensor = torch.stack(energies) + energies_list = [t.tolist() for t in energies_tensor.T] + + assert len(energies_list[0]) == n_steps + + mean_temps = torch.mean(temperatures_tensor, dim=0) + for mean_temp in mean_temps: + assert abs(mean_temp - kT.item() / MetalUnits.temperature) < 150.0 + + for traj in energies_list: + energy_std = torch.tensor(traj).std() + assert energy_std < 1.0 + + # Cell reconstruction is consistent + assert torch.allclose(state.cell, state.current_cell) + + def test_npt_langevin_multi_kt( ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel ): @@ -339,7 +399,7 @@ def test_nvt_nose_hoover(ar_double_sim_state: ts.SimState, lj_model: LennardJone temperatures_list = [t.tolist() for t in temperatures_tensor.T] assert torch.allclose( temperatures_tensor[-1], - torch.tensor([290.3553, 289.9699], dtype=dtype), + torch.tensor([305.6400, 305.4556], dtype=dtype), ) energies_tensor = torch.stack(energies) @@ -728,7 +788,7 @@ def test_npt_nose_hoover(ar_double_sim_state: ts.SimState, lj_model: LennardJone temperatures_list = [t.tolist() for t in temperatures_tensor.T] assert torch.allclose( temperatures_tensor[-1], - torch.tensor([287.5729, 287.1330], dtype=dtype), + torch.tensor([283.1162, 313.1624], dtype=dtype), ) energies_tensor = torch.stack(energies) @@ -1023,19 +1083,19 @@ def test_compute_cell_force_atoms_per_system(): atomic_numbers=torch.ones(72, dtype=torch.long), stress=torch.zeros((2, 3, 3)), reference_cell=torch.eye(3).repeat(2, 1, 1), - cell_positions=torch.ones((2, 3, 3)), - cell_velocities=torch.zeros((2, 3, 3)), + cell_positions=torch.zeros(2, 3), + cell_velocities=torch.zeros(2, 3), cell_masses=torch.ones(2), alpha=torch.ones(2), cell_alpha=torch.ones(2), b_tau=torch.ones(2), ) - # Get forces and compare ratio - cell_force = _compute_cell_force(state, torch.tensor(0.0), torch.tensor([1.0, 1.0])) - force_ratio = ( - torch.diagonal(cell_force[1]).mean() / torch.diagonal(cell_force[0]).mean() - ) + # Get forces and compare ratio (per-dimension force) + P_ext = torch.zeros(2, 3) + cell_force = _compute_cell_force(state, P_ext, torch.tensor([1.0, 1.0])) + # Check the first dimension's force ratio + force_ratio = cell_force[1, 0] / cell_force[0, 0] # Force ratio should match atom ratio (8:1) with the fix assert abs(force_ratio - 8.0) / 8.0 < 0.1 From 44161226c30cac183896db38a0e96aa55230909e Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:48:16 +0000 Subject: [PATCH 12/15] update this test that fails on gpu because device(cuda:0) is different than device(cuda) (which I believe is not a problem) --- tests/test_nbody.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_nbody.py b/tests/test_nbody.py index e235cd629..5da5a45c5 100644 --- a/tests/test_nbody.py +++ b/tests/test_nbody.py @@ -481,9 +481,9 @@ def test_build_triplets_device(device: str) -> None: result = build_triplets(edge_index, n_atoms) - assert result["trip_in"].device == dev - assert result["trip_out"].device == dev - assert result["center_atom"].device == dev + assert result["trip_in"].device.type == dev.type + assert result["trip_out"].device.type == dev.type + assert result["center_atom"].device.type == dev.type @pytest.mark.parametrize( @@ -507,10 +507,10 @@ def test_build_quadruplets_device(device: str) -> None: internal_cell_offsets, ) - assert result["quad_c_to_a_edge"].device == dev - assert result["quad_d_to_b_trip_idx"].device == dev - assert result["d_to_b_edge"].device == dev - assert result["c_to_a_edge"].device == dev + assert result["quad_c_to_a_edge"].device.type == dev.type + assert result["quad_d_to_b_trip_idx"].device.type == dev.type + assert result["d_to_b_edge"].device.type == dev.type + assert result["c_to_a_edge"].device.type == dev.type def test_build_triplets_jit_script() -> None: From 106e853f29a2f95e1c8976402ea15589e29a9c25 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:48:43 +0000 Subject: [PATCH 13/15] correct unwrap_positions that was wrong (because tests was wrong), error on unwrapping npt simulation --- tests/test_transforms.py | 6 +++--- torch_sim/transforms.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index bac5ec985..d43ffd93e 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1396,13 +1396,13 @@ def test_build_linked_cell_neighborhood_basic() -> None: def test_unwrap_positions(ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel): n_steps = 50 - dt = torch.tensor(0.001, dtype=DTYPE) + dt = torch.tensor(0.001, dtype=DTYPE) * MetalUnits.time kT = torch.tensor(300, dtype=DTYPE) * MetalUnits.temperature # Same cell state = ts.nvt_langevin_init(state=ar_double_sim_state, model=lj_model, kT=kT) state.positions = tst.pbc_wrap_batched(state.positions, state.cell, state.system_idx) - positions = [state.positions.detach().clone()] + positions = [state.positions.detach().clone()] # for _step in range(n_steps): state = ts.nvt_langevin_step(model=lj_model, state=state, dt=dt, kT=kT) positions.append(state.positions.detach().clone()) @@ -1446,7 +1446,7 @@ def test_unwrap_positions(ar_double_sim_state: ts.SimState, lj_model: LennardJon ) unwrapped_positions = tst.unwrap_positions( wrapped_positions, - state.cell, + torch.stack(cells), state.system_idx, ) assert torch.allclose(unwrapped_positions, positions, atol=1e-4) diff --git a/torch_sim/transforms.py b/torch_sim/transforms.py index bdc41f1e6..35fbaa458 100644 --- a/torch_sim/transforms.py +++ b/torch_sim/transforms.py @@ -1260,12 +1260,18 @@ def unwrap_positions( dfrac = frac[1:] - frac[:-1] dfrac -= torch.round(dfrac) - dcart = torch.einsum("tni,tnij->tnj", dfrac, box_atoms[:-1]) + # Reconstruct unwrapped fractional trajectory + unwrapped_frac = torch.empty_like(frac) + unwrapped_frac[0] = frac[0] + unwrapped_frac[1:] = torch.cumsum(dfrac, dim=0) + frac[0] + + # Convert back to Cartesian using each frame's cell + return torch.einsum("tni,tnij->tnj", unwrapped_frac, box_atoms) else: raise ValueError("box must have shape (n_systems,3,3) or (T,n_systems,3,3)") - # Cumulative reconstruction + # Cumulative reconstruction (constant cell path) unwrapped = torch.empty_like(positions) unwrapped[0] = positions[0] unwrapped[1:] = torch.cumsum(dcart, dim=0) + unwrapped[0] From ecabf25fd78148825d6f3e5c4999d2df29e5752d Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:49:08 +0000 Subject: [PATCH 14/15] update shape of test for npt_langevin variables that got changed --- tests/test_state.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_state.py b/tests/test_state.py index ed1b3c29d..9d31c492f 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1096,8 +1096,8 @@ def test_nptlangevinstate_instantiation() -> None: cell_alpha=torch.ones(1), b_tau=torch.ones(1), reference_cell=torch.eye(3).unsqueeze(0), - cell_positions=torch.zeros(1, 3, 3), - cell_velocities=torch.zeros(1, 3, 3), + cell_positions=torch.zeros(1, 3), + cell_velocities=torch.zeros(1, 3), cell_masses=torch.ones(1), ) _check_coercion(state) From b44a445e3f3d1d41ac25fbe125876d95c5340f68 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 3 Apr 2026 14:52:50 +0000 Subject: [PATCH 15/15] updated test_physical_validation --- tests/conftest.py | 15 + tests/test_physical_validation.py | 777 ++++++++++++++++++++++-------- 2 files changed, 604 insertions(+), 188 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 79b39b48a..750edea9c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,6 +15,21 @@ torch.set_num_threads(4) + +def pytest_addoption(parser): + parser.addoption( + "--validation-plots", + action="store_true", + default=False, + help="Save physical validation plots to tests/physical_validation_data/plots/", + ) + parser.addoption( + "--clean-validation-data", + action="store_true", + default=False, + help="Delete saved physical validation data before running tests", + ) + DEVICE = torch.device("cpu") DTYPE = torch.float64 diff --git a/tests/test_physical_validation.py b/tests/test_physical_validation.py index e57f85c8b..a25504c15 100644 --- a/tests/test_physical_validation.py +++ b/tests/test_physical_validation.py @@ -1,110 +1,246 @@ """Physical validation tests for torch-sim MD integrators. Uses the physical_validation library (https://github.com/shirtsgroup/physical_validation) -to verify that integrators produce physically correct results. These tests are -long-running (~5 min total) and excluded by default. Run with: +to verify that integrators produce physically correct results. These tests require CUDA +and are long-running. Run with: pytest -m physical_validation -v + +Options: + --validation-plots Save plots to tests/physical_validation_data/plots/ + --clean-validation-data Delete saved validation data before running + +Run a specific integrator: + + pytest -m physical_validation -v -k "nvt_langevin" + +Tested integrators: + + NVT: + - nvt_langevin + - nvt_nose_hoover + - nvt_vrescale + + NPT: + - npt_langevin (independent per-dimension strain, like LAMMPS couple=none) + - npt_langevin_strain (isotropic logarithmic strain) + - npt_nose_hoover + - npt_isotropic_crescale + - npt_anisotropic_crescale + +Clean up saved data programmatically: + + from tests.test_physical_validation import clean_validation_data + clean_validation_data() """ +import shutil +import warnings +from pathlib import Path + import numpy as np import pytest import torch from ase.build import bulk +from numpy.typing import NDArray import torch_sim as ts +from torch_sim.integrators.npt import npt_crescale_average_anisotropic_step from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.units import MetalUnits + physical_validation = pytest.importorskip("physical_validation") -DEVICE = torch.device("cpu") +if not torch.cuda.is_available(): + pytest.skip("CUDA not available", allow_module_level=True) + +# --------------------------------------------------------------------------- +# Device & dtype — CUDA required +# --------------------------------------------------------------------------- +DEVICE = torch.device("cuda") DTYPE = torch.float64 +# --------------------------------------------------------------------------- # LJ Argon parameters +# --------------------------------------------------------------------------- SIGMA = 3.405 EPSILON = 0.0104 CUTOFF = 2.5 * SIGMA +# --------------------------------------------------------------------------- +# Simulation parameters (matched to fast_integrator_tests_batch) +# --------------------------------------------------------------------------- +TIMESTEP_PS = 0.005 +N_STEPS_NVT = 10_000 +N_STEPS_NPT = 10_000 +N_EQUILIBRATION_NVT = 4_000 +N_EQUILIBRATION_NPT = 5_000 +LOG_EVERY = 5 + +# Ensemble check temperatures and pressures (matched to fast_integrator_tests_batch) +TEMPERATURES = [58.3, 60.0] +EXTERNAL_PRESSURE = 0.0 +PRESSURE_SWEEP_TEMP = 60.0 +PRESSURE_SWEEP_BAR = 90.0 +PRESSURE_SWEEP_EVA3 = PRESSURE_SWEEP_BAR * float(MetalUnits.pressure) + +# Physical validation thresholds (in sigma units) +KE_SIGMA_WARNING = 2.0 +KE_SIGMA_THRESHOLD = 3.0 +ENSEMBLE_SIGMA_WARNING = 2.0 +ENSEMBLE_SIGMA_THRESHOLD = 3.0 + +# Data & plot directories +DATA_DIR = Path(__file__).parent / "physical_validation_data" +PLOTS_DIR = DATA_DIR / "plots" + +RunData = dict[str, NDArray[np.floating] | float | int] + +torch.set_num_threads(4) -def _make_unit_data(): + +# --------------------------------------------------------------------------- +# Cleanup utility +# --------------------------------------------------------------------------- +def clean_validation_data() -> None: + """Delete all saved physical validation data and plots.""" + if DATA_DIR.exists(): + shutil.rmtree(DATA_DIR) + print(f"Removed {DATA_DIR}") + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- +def _to_kT(temperature_K: float) -> float: + return temperature_K * float(MetalUnits.temperature) + + +def _to_dt(timestep_ps: float) -> float: + return timestep_ps * float(MetalUnits.time) + + +def _save_run_data(data: RunData, label: str) -> Path: + """Save run data to a .npz file and return the path.""" + DATA_DIR.mkdir(parents=True, exist_ok=True) + path = DATA_DIR / f"{label}.npz" + np.savez(path, **data) + return path + + +def _get_plot_path(request: pytest.FixtureRequest, name: str) -> str | None: + """Return plot file path if --validation-plots is enabled, else None.""" + if not request.config.getoption("--validation-plots", default=False): + return None + PLOTS_DIR.mkdir(parents=True, exist_ok=True) + return str(PLOTS_DIR / f"{name}.png") + + +def _pressure_to_bar(p_eva3: float) -> float: + """Convert eV/Ang^3 to bar.""" + return p_eva3 / float(MetalUnits.pressure) + + +# --------------------------------------------------------------------------- +# Helpers: unit data, model, structure +# --------------------------------------------------------------------------- +def _make_unit_data() -> physical_validation.data.UnitData: """Create UnitData for torch-sim's MetalUnits system.""" return physical_validation.data.UnitData( kb=float(MetalUnits.temperature), # k_B in eV/K = 8.617e-5 energy_str="eV", - energy_conversion=1.0, + energy_conversion=96.485, # Convert to kJ/mol length_str="Ang", - length_conversion=1.0, + length_conversion=1e-1, # Convert to nm volume_str="Ang^3", - volume_conversion=1.0, + volume_conversion=1e-3, # Convert to nm^3 temperature_str="K", temperature_conversion=1.0, - pressure_str="eV/Ang^3", + pressure_str="bar", pressure_conversion=1.0, - time_str="internal", - time_conversion=1.0, ) -def _make_lj_model(): +def _make_lj_model(*, compute_stress: bool = False) -> LennardJonesModel: """Create a Lennard-Jones model for Argon.""" return LennardJonesModel( - use_neighbor_list=False, sigma=SIGMA, epsilon=EPSILON, device=DEVICE, dtype=DTYPE, compute_forces=True, - compute_stress=False, + compute_stress=compute_stress, cutoff=CUTOFF, ) -def _make_ar_supercell(repeat=(2, 2, 2)): +def _make_ar_supercell( + repeat: tuple[int, int, int] = (8, 8, 8), +) -> ts.SimState: """Create an FCC Argon supercell SimState.""" atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat(repeat) return ts.io.atoms_to_state(atoms, DEVICE, DTYPE) -def _run_nvt_langevin( - sim_state, - model, - temperature, - timestep_ps, - n_steps, - n_equilibration, - seed=42, -): - """Run NVT Langevin simulation and collect per-step observables.""" - kT = temperature * float(MetalUnits.temperature) - dt_internal = timestep_ps * float(MetalUnits.time) +# --------------------------------------------------------------------------- +# Generic NVT runner +# --------------------------------------------------------------------------- +def _run_nvt( + integrator_name: str, + sim_state: ts.SimState, + model: LennardJonesModel, + temperature: float, + timestep_ps: float = TIMESTEP_PS, + n_steps: int = N_STEPS_NVT, + n_equilibration: int = N_EQUILIBRATION_NVT, + log_every: int = LOG_EVERY, + seed: int = 42, +) -> RunData: + """Run an NVT simulation with the specified integrator.""" + kT = _to_kT(temperature) + dt = _to_dt(timestep_ps) natoms = int(sim_state.positions.shape[0]) - state = ts.nvt_langevin_init(sim_state, model, kT=kT, seed=seed) + sim_state = sim_state.clone() + sim_state.rng = seed + + # Initialize (params matched to fast_integrator_tests_batch) + if integrator_name == "nvt_langevin": + state = ts.nvt_langevin_init(sim_state, model, kT=kT) + elif integrator_name == "nvt_nose_hoover": + state = ts.nvt_nose_hoover_init(sim_state, model, kT=kT, dt=dt, tau=10 * dt) + elif integrator_name == "nvt_vrescale": + state = ts.nvt_vrescale_init(sim_state, model, kT=kT) + else: + msg = f"Unknown NVT integrator: {integrator_name}" + raise ValueError(msg) + + def _step(s): + if integrator_name == "nvt_langevin": + return ts.nvt_langevin_step(s, model, dt=dt, kT=kT, gamma=1 / (50 * dt)) + if integrator_name == "nvt_nose_hoover": + return ts.nvt_nose_hoover_step(s, model, dt=dt, kT=kT) + return ts.nvt_vrescale_step(model, s, dt=dt, kT=kT, tau=10 * dt) # Equilibration for _ in range(n_equilibration): - state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) - - # Production - collect observables - ke_list = [] - pe_list = [] - total_e_list = [] - position_list = [] - velocity_list = [] - - for _ in range(n_steps): - state = ts.nvt_langevin_step(state, model, dt=dt_internal, kT=kT) - - ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) - pe = state.energy.sum() - ke_list.append(float(ke)) - pe_list.append(float(pe)) - total_e_list.append(float(ke + pe)) - position_list.append(state.positions.detach().cpu().numpy().copy()) - velocity_list.append(state.velocities.detach().cpu().numpy().copy()) - - # Compute volume from cell + state = _step(state) + + # Production (subsampled every log_every steps) + ke_list, pe_list, total_e_list = [], [], [] + + for i in range(n_steps): + state = _step(state) + if (i + 1) % log_every == 0: + ke = float( + ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + ) + pe = float(state.energy.sum()) + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + cell = sim_state.cell[0].detach().cpu().numpy() volume = float(np.abs(np.linalg.det(cell))) @@ -112,35 +248,137 @@ def _run_nvt_langevin( "kinetic_energy": np.array(ke_list), "potential_energy": np.array(pe_list), "total_energy": np.array(total_e_list), - "positions": np.array(position_list), - "velocities": np.array(velocity_list), "volume": volume, "masses": sim_state.masses.detach().cpu().numpy(), - "dt_internal": dt_internal, + "dt_internal": dt, "natoms": natoms, + "target_temperature": temperature, + "timestep_ps": timestep_ps, + "integrator": integrator_name, } -def _run_nve(sim_state, model, kT_init, timestep_ps, n_steps, seed=42): - """Run NVE simulation and collect constant of motion.""" - dt_internal = timestep_ps * float(MetalUnits.time) +# --------------------------------------------------------------------------- +# Generic NPT runner +# --------------------------------------------------------------------------- +def _run_npt( + integrator_name: str, + sim_state: ts.SimState, + model: LennardJonesModel, + temperature: float, + external_pressure: float = 0.0, + timestep_ps: float = TIMESTEP_PS, + n_steps: int = N_STEPS_NPT, + n_equilibration: int = N_EQUILIBRATION_NPT, + log_every: int = LOG_EVERY, + seed: int = 42, +) -> RunData: + """Run an NPT simulation with the specified integrator.""" + kT = _to_kT(temperature) + dt = torch.tensor(_to_dt(timestep_ps), device=DEVICE, dtype=DTYPE) + ext_p = torch.tensor(external_pressure, device=DEVICE, dtype=DTYPE) + natoms = int(sim_state.positions.shape[0]) + + sim_state = sim_state.clone() + sim_state.rng = seed - state = ts.nve_init(sim_state, model, kT=kT_init, seed=seed) + # Initialize (params matched to fast_integrator_tests_batch) + if integrator_name == "npt_langevin": + state = ts.npt_langevin_init( + sim_state, model, kT=kT, dt=dt, + alpha=1 / (5 * dt), cell_alpha=1 / (30 * dt), b_tau=300 * dt, + ) + elif integrator_name == "npt_langevin_strain": + state = ts.npt_langevin_strain_init( + sim_state, model, kT=kT, dt=dt, + alpha=1 / (5 * dt), cell_alpha=1 / (30 * dt), b_tau=300 * dt, + ) + elif integrator_name == "npt_nose_hoover": + state = ts.npt_nose_hoover_init( + sim_state, model, kT=kT, dt=dt, + t_tau=10 * dt, b_tau=100 * dt, + ) + elif integrator_name == "npt_isotropic_crescale": + state = ts.npt_crescale_init( + sim_state, model, kT=kT, dt=dt, + tau_p=3 * dt, + isothermal_compressibility=1e-6 / MetalUnits.pressure, + ) + elif integrator_name == "npt_anisotropic_crescale": + state = ts.npt_crescale_init( + sim_state, model, kT=kT, dt=dt, + tau_p=3 * dt, + isothermal_compressibility=1e-6 / MetalUnits.pressure, + ) + else: + msg = f"Unknown NPT integrator: {integrator_name}" + raise ValueError(msg) + + def _step(s): + if integrator_name == "npt_langevin": + return ts.npt_langevin_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, + ) + if integrator_name == "npt_langevin_strain": + return ts.npt_langevin_strain_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, + ) + if integrator_name == "npt_nose_hoover": + return ts.npt_nose_hoover_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, + ) + if integrator_name == "npt_anisotropic_crescale": + return npt_crescale_average_anisotropic_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, tau=1 * dt, + ) + return ts.npt_crescale_isotropic_step( + s, model, dt=dt, kT=kT, external_pressure=ext_p, tau=1 * dt, + ) - com_list = [] - for _ in range(n_steps): - state = ts.nve_step(state, model, dt=dt_internal) - ke = ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) - pe = state.energy.sum() - com_list.append(float(ke + pe)) + # Equilibration + for _ in range(n_equilibration): + state = _step(state) + + # Production (subsampled every log_every steps) + ke_list, pe_list, total_e_list = [], [], [] + volume_list = [] + + for i in range(n_steps): + state = _step(state) + if (i + 1) % log_every == 0: + ke = float( + ts.calc_kinetic_energy(masses=state.masses, momenta=state.momenta) + ) + pe = float(state.energy.sum()) + cell = state.cell[0].detach().cpu().numpy() + vol = float(np.abs(np.linalg.det(cell))) + ke_list.append(ke) + pe_list.append(pe) + total_e_list.append(ke + pe) + volume_list.append(vol) return { - "constant_of_motion": np.array(com_list), - "dt_internal": dt_internal, + "kinetic_energy": np.array(ke_list), + "potential_energy": np.array(pe_list), + "total_energy": np.array(total_e_list), + "volumes": np.array(volume_list), + "masses": sim_state.masses.detach().cpu().numpy(), + "dt_internal": float(dt), + "natoms": natoms, + "target_temperature": temperature, + "external_pressure": external_pressure, + "timestep_ps": timestep_ps, + "integrator": integrator_name, } -def _build_nvt_simulation_data(run_data, temperature): +# --------------------------------------------------------------------------- +# SimulationData builders +# --------------------------------------------------------------------------- +def _build_nvt_simulation_data( + run_data: RunData, + temperature: float, +) -> physical_validation.data.SimulationData: """Build a physical_validation SimulationData from NVT run results.""" units = _make_unit_data() @@ -165,181 +403,344 @@ def _build_nvt_simulation_data(run_data, temperature): total_energy=run_data["total_energy"], ) - trajectory = physical_validation.data.TrajectoryData( - position=run_data["positions"], - velocity=run_data["velocities"], + return physical_validation.data.SimulationData( + units=units, + dt=run_data["timestep_ps"], + system=system, + ensemble=ensemble_data, + observables=observables, + ) + + +def _build_npt_simulation_data( + run_data: RunData, + temperature: float, + pressure: float, +) -> physical_validation.data.SimulationData: + """Build a physical_validation SimulationData from NPT run results.""" + units = _make_unit_data() + + system = physical_validation.data.SystemData( + natoms=run_data["natoms"], + nconstraints=0, + ndof_reduction_tra=3, + ndof_reduction_rot=0, + mass=run_data["masses"], + ) + + ensemble_data = physical_validation.data.EnsembleData( + ensemble="NPT", + natoms=run_data["natoms"], + pressure=pressure, + temperature=temperature, + ) + + observables = physical_validation.data.ObservableData( + kinetic_energy=run_data["kinetic_energy"], + potential_energy=run_data["potential_energy"], + total_energy=run_data["total_energy"], + volume=run_data["volumes"], ) return physical_validation.data.SimulationData( units=units, - dt=run_data["dt_internal"], + dt=run_data["timestep_ps"], system=system, ensemble=ensemble_data, observables=observables, - trajectory=trajectory, ) -@pytest.mark.physical_validation -def test_ke_distribution(): - """Test that kinetic energy follows the Maxwell-Boltzmann distribution. +# =========================================================================== +# Session fixture: cleanup saved data +# =========================================================================== +@pytest.fixture(autouse=True, scope="session") +def _manage_validation_data(request): + """Clean data directory if --clean-validation-data is set.""" + if request.config.getoption("--clean-validation-data", default=False): + clean_validation_data() + yield - Runs NVT Langevin at 100K on a 2x2x2 Ar supercell (32 atoms) and checks - that the KE distribution matches the analytical Maxwell-Boltzmann prediction. - """ - sim_state = _make_ar_supercell(repeat=(2, 2, 2)) + +# =========================================================================== +# Tests: KE distribution (Maxwell-Boltzmann) +# =========================================================================== +@pytest.mark.physical_validation +@pytest.mark.parametrize( + "integrator_name", + ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"], +) +def test_nvt_ke_distribution(integrator_name: str, request) -> None: + """Test that KE follows the Maxwell-Boltzmann distribution for NVT.""" + sim_state = _make_ar_supercell(repeat=(8, 8, 8)) model = _make_lj_model() - temperature = 100.0 # K + temperature = TEMPERATURES[1] - run_data = _run_nvt_langevin( - sim_state, - model, - temperature=temperature, - timestep_ps=0.004, - n_steps=10_000, - n_equilibration=2_000, - seed=42, + run_data = _run_nvt( + integrator_name, sim_state, model, temperature=temperature, seed=42, ) + _save_run_data(run_data, f"{integrator_name}_T{temperature:.1f}K_ke") data = _build_nvt_simulation_data(run_data, temperature) + plot_path = _get_plot_path(request, f"{integrator_name}_nvt_ke") - result = physical_validation.kinetic_energy.distribution( - data, - strict=False, - verbosity=0, + kwargs = {} + if plot_path: + kwargs["filename"] = plot_path + d_mean, d_width = physical_validation.kinetic_energy.distribution( + data, strict=False, verbosity=0, **kwargs, ) - # strict=False returns (d_mean, d_width) in sigma units - d_mean, d_width = result - assert abs(d_mean) < 3, ( - f"KE mean deviation {d_mean:.2f} sigma exceeds threshold" + if abs(d_mean) > KE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] KE mean deviation {d_mean:.2f} sigma exceeds " + f"{KE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + if abs(d_width) > KE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] KE width deviation {d_width:.2f} sigma exceeds " + f"{KE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + assert abs(d_mean) < KE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] KE mean deviation {d_mean:.2f} sigma" ) - assert abs(d_width) < 3, ( - f"KE width deviation {d_width:.2f} sigma exceeds threshold" + assert abs(d_width) < KE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] KE width deviation {d_width:.2f} sigma" ) @pytest.mark.physical_validation -def test_integrator_convergence(): - """Test that NVE energy error scales as dt^2 (velocity Verlet). +@pytest.mark.parametrize( + "integrator_name", + [ + "npt_langevin", + "npt_langevin_strain", + "npt_nose_hoover", + "npt_isotropic_crescale", + "npt_anisotropic_crescale", + ], +) +def test_npt_ke_distribution(integrator_name: str, request) -> None: + """Test that KE follows the Maxwell-Boltzmann distribution for NPT.""" + sim_state = _make_ar_supercell(repeat=(8, 8, 8)) + model = _make_lj_model(compute_stress=True) + temperature = TEMPERATURES[1] + + run_data = _run_npt( + integrator_name, sim_state, model, + temperature=temperature, external_pressure=EXTERNAL_PRESSURE, seed=42, + ) + _save_run_data(run_data, f"{integrator_name}_T{temperature:.1f}K_ke") + + # Use NVT builder with mean volume for KE distribution check + run_data_nvt = {**run_data, "volume": float(np.mean(run_data["volumes"]))} + data = _build_nvt_simulation_data(run_data_nvt, temperature) + plot_path = _get_plot_path(request, f"{integrator_name}_npt_ke") + + kwargs = {} + if plot_path: + kwargs["filename"] = plot_path + d_mean, d_width = physical_validation.kinetic_energy.distribution( + data, strict=False, verbosity=0, **kwargs, + ) - Runs NVE at 3 different timesteps from identical initial conditions on a - 4-atom Ar unit cell at low temperature (5K). Low temperature minimizes - thermal fluctuations so the integration error dominates the RMSD of the - conserved quantity, allowing the dt^2 convergence to be observed. - """ - sim_state = _make_ar_supercell(repeat=(1, 1, 1)) # 4 atoms + if abs(d_mean) > KE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] KE mean deviation {d_mean:.2f} sigma exceeds " + f"{KE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + if abs(d_width) > KE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] KE width deviation {d_width:.2f} sigma exceeds " + f"{KE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + assert abs(d_mean) < KE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] KE mean deviation {d_mean:.2f} sigma" + ) + assert abs(d_width) < KE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] KE width deviation {d_width:.2f} sigma" + ) + + +# =========================================================================== +# Tests: ensemble validity (Boltzmann weight ratio at two temperatures) +# =========================================================================== +@pytest.mark.physical_validation +@pytest.mark.parametrize( + "integrator_name", + ["nvt_langevin", "nvt_nose_hoover", "nvt_vrescale"], +) +def test_nvt_ensemble_check(integrator_name: str, request) -> None: + """Test NVT ensemble validity at two temperatures.""" + sim_state = _make_ar_supercell(repeat=(8, 8, 8)) model = _make_lj_model() - temperature = 5.0 # K, low T so integration error dominates - kT_init = temperature * float(MetalUnits.temperature) - # Timesteps chosen so integration error >> thermal fluctuations at all dt. - # Factor of ~sqrt(2) spacing gives dt^2 ratio of ~2.0 per step. - timesteps = [0.008, 0.00566, 0.004] # ps - n_steps = 5_000 - seed = 42 + temp_low, temp_high = TEMPERATURES - natoms = int(sim_state.positions.shape[0]) - masses = sim_state.masses.detach().cpu().numpy() - volume = float( - np.abs(np.linalg.det(sim_state.cell[0].detach().cpu().numpy())) + run_low = _run_nvt( + integrator_name, sim_state, model, temperature=temp_low, seed=42, ) - units = _make_unit_data() + run_high = _run_nvt( + integrator_name, sim_state, model, temperature=temp_high, seed=123, + ) + _save_run_data(run_low, f"{integrator_name}_T{temp_low:.1f}K_ens") + _save_run_data(run_high, f"{integrator_name}_T{temp_high:.1f}K_ens") - simulations = [] - for dt_ps in timesteps: - run_data = _run_nve( - sim_state, - model, - kT_init=kT_init, - timestep_ps=dt_ps, - n_steps=n_steps, - seed=seed, - ) + data_low = _build_nvt_simulation_data(run_low, temp_low) + data_high = _build_nvt_simulation_data(run_high, temp_high) + plot_path = _get_plot_path(request, f"{integrator_name}_nvt_ens") - system = physical_validation.data.SystemData( - natoms=natoms, - nconstraints=0, - ndof_reduction_tra=3, - ndof_reduction_rot=0, - mass=masses, - ) + kwargs = {} + if plot_path: + kwargs["filename"] = plot_path + quantiles = physical_validation.ensemble.check( + data_low, data_high, + total_energy=True, + data_is_uncorrelated=True, + verbosity=0, + **kwargs, + ) - ensemble_data = physical_validation.data.EnsembleData( - ensemble="NVE", - natoms=natoms, - volume=volume, + for i, q in enumerate(quantiles): + if abs(q) > ENSEMBLE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] Ensemble quantile {i} = {q:.2f} sigma exceeds " + f"{ENSEMBLE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + assert abs(q) < ENSEMBLE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] Ensemble quantile {i} = {q:.2f} sigma" ) - observables = physical_validation.data.ObservableData( - constant_of_motion=run_data["constant_of_motion"], - ) - sim_data = physical_validation.data.SimulationData( - units=units, - dt=run_data["dt_internal"], - system=system, - ensemble=ensemble_data, - observables=observables, - ) - simulations.append(sim_data) +@pytest.mark.physical_validation +@pytest.mark.parametrize( + "integrator_name", + [ + "npt_langevin", + "npt_langevin_strain", + "npt_nose_hoover", + "npt_isotropic_crescale", + "npt_anisotropic_crescale", + ], +) +def test_npt_ensemble_check(integrator_name: str, request) -> None: + """Test NPT ensemble validity at two temperatures. + + Uses temperatures both in the solid phase (below LJ Ar melting point ~84K) + to avoid the solid-liquid phase transition which causes non-overlapping + energy distributions in the NPT ensemble. + """ + sim_state = _make_ar_supercell(repeat=(8, 8, 8)) + model = _make_lj_model(compute_stress=True) - result = physical_validation.integrator.convergence( - simulations, - verbose=False, - ) + temp_low, temp_high = TEMPERATURES - assert result < 0.5, ( - f"Integrator convergence deviation {result:.3f} exceeds threshold 0.5" + run_low = _run_npt( + integrator_name, sim_state, model, + temperature=temp_low, external_pressure=EXTERNAL_PRESSURE, seed=42, ) + run_high = _run_npt( + integrator_name, sim_state, model, + temperature=temp_high, external_pressure=EXTERNAL_PRESSURE, seed=123, + ) + _save_run_data(run_low, f"{integrator_name}_T{temp_low:.1f}K_ens") + _save_run_data(run_high, f"{integrator_name}_T{temp_high:.1f}K_ens") + data_low = _build_npt_simulation_data(run_low, temp_low, EXTERNAL_PRESSURE) + data_high = _build_npt_simulation_data(run_high, temp_high, EXTERNAL_PRESSURE) + plot_path = _get_plot_path(request, f"{integrator_name}_npt_ens_temp") -@pytest.mark.physical_validation -def test_ensemble_check(): - """Test NVT ensemble validity via Boltzmann weight ratio at two temperatures. + kwargs = {} + if plot_path: + kwargs["filename"] = plot_path + quantiles = physical_validation.ensemble.check( + data_low, data_high, + total_energy=True, + data_is_uncorrelated=True, + verbosity=0, + **kwargs, + ) + + for i, q in enumerate(quantiles): + if abs(q) > ENSEMBLE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] Ensemble quantile {i} = {q:.2f} sigma exceeds " + f"{ENSEMBLE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + assert abs(q) < ENSEMBLE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] Ensemble quantile {i} = {q:.2f} sigma" + ) - Runs NVT Langevin at 80K and 100K on a 2x2x2 Ar supercell (32 atoms), - then checks that the total energy distributions satisfy the expected - Boltzmann weight relationship. Uses total_energy=True for the ensemble - check which includes both kinetic and potential energy contributions. - """ - sim_state = _make_ar_supercell(repeat=(2, 2, 2)) - model = _make_lj_model() - temp_low = 80.0 - temp_high = 100.0 - - run_low = _run_nvt_langevin( - sim_state, - model, - temperature=temp_low, - timestep_ps=0.004, - n_steps=10_000, - n_equilibration=2_000, - seed=42, +# =========================================================================== +# Tests: ensemble validity (Boltzmann weight ratio at two pressures) +# =========================================================================== +@pytest.mark.physical_validation +@pytest.mark.parametrize( + "integrator_name", + [ + "npt_langevin", + "npt_langevin_strain", + "npt_nose_hoover", + "npt_isotropic_crescale", + "npt_anisotropic_crescale", + ], +) +def test_npt_pressure_ensemble_check(integrator_name: str, request) -> None: + """Test NPT ensemble validity at two pressures (fixed temperature).""" + sim_state = _make_ar_supercell(repeat=(8, 8, 8)) + model = _make_lj_model(compute_stress=True) + + p_low = EXTERNAL_PRESSURE + p_high = PRESSURE_SWEEP_EVA3 + + run_low = _run_npt( + integrator_name, sim_state, model, + temperature=PRESSURE_SWEEP_TEMP, external_pressure=p_low, seed=42, + ) + run_high = _run_npt( + integrator_name, sim_state, model, + temperature=PRESSURE_SWEEP_TEMP, external_pressure=p_high, seed=123, ) - run_high = _run_nvt_langevin( - sim_state, - model, - temperature=temp_high, - timestep_ps=0.004, - n_steps=10_000, - n_equilibration=2_000, - seed=123, + _save_run_data( + run_low, + f"{integrator_name}_T{PRESSURE_SWEEP_TEMP:.1f}K_P0bar", + ) + _save_run_data( + run_high, + f"{integrator_name}_T{PRESSURE_SWEEP_TEMP:.1f}K_P{PRESSURE_SWEEP_BAR:.0f}bar", ) - data_low = _build_nvt_simulation_data(run_low, temp_low) - data_high = _build_nvt_simulation_data(run_high, temp_high) + data_low = _build_npt_simulation_data(run_low, PRESSURE_SWEEP_TEMP, p_low) + data_high = _build_npt_simulation_data(run_high, PRESSURE_SWEEP_TEMP, p_high) + plot_path = _get_plot_path(request, f"{integrator_name}_npt_ens_press") + kwargs = {} + if plot_path: + kwargs["filename"] = plot_path quantiles = physical_validation.ensemble.check( - data_low, - data_high, + data_low, data_high, total_energy=True, data_is_uncorrelated=True, verbosity=0, + **kwargs, ) for i, q in enumerate(quantiles): - assert abs(q) < 3, ( - f"Ensemble quantile {i} = {q:.2f} sigma exceeds threshold" + if abs(q) > ENSEMBLE_SIGMA_WARNING: + warnings.warn( + f"[{integrator_name}] Pressure ensemble quantile {i} = {q:.2f} sigma " + f"exceeds {ENSEMBLE_SIGMA_WARNING} sigma warning threshold", + stacklevel=1, + ) + assert abs(q) < ENSEMBLE_SIGMA_THRESHOLD, ( + f"[{integrator_name}] Pressure ensemble quantile {i} = {q:.2f} sigma" ) + +