From c56612965cfe1e6c63f118dae033977d0973adcc Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 13 Feb 2026 16:14:32 +0000 Subject: [PATCH 01/10] 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 8093d764..09c5e5c0 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 00000000..e57f85c8 --- /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/10] 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 00000000..a60ee944 --- /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 00000000..444948d6 --- /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 00000000..69e21fd0 --- /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 00000000..0b18f80c --- /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 00000000..a0d60ca9 --- /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 00000000..ba8d141c --- /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 cbfbec5554a932520f0a0b61879abd19fab0664d Mon Sep 17 00:00:00 2001 From: "jinmu.you" <98165579+alphalm4@users.noreply.github.com> Date: Tue, 31 Mar 2026 00:12:58 +0900 Subject: [PATCH 03/10] Fix NPT Nose-Hoover things (#520) --- tests/test_integrators.py | 4 ++-- torch_sim/integrators/md.py | 2 +- torch_sim/integrators/npt.py | 23 ++++++++++------------- torch_sim/integrators/nvt.py | 13 +++++-------- 4 files changed, 18 insertions(+), 24 deletions(-) diff --git a/tests/test_integrators.py b/tests/test_integrators.py index dd9b49d0..27ea5dba 100644 --- a/tests/test_integrators.py +++ b/tests/test_integrators.py @@ -339,7 +339,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([300.0096, 299.7024], dtype=dtype), + torch.tensor([290.3553, 289.9699], dtype=dtype), ) energies_tensor = torch.stack(energies) @@ -728,7 +728,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([298.2752, 297.9444], dtype=dtype), + torch.tensor([287.5729, 287.1330], dtype=dtype), ) energies_tensor = torch.stack(energies) diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py index 76ff69c1..4a593c04 100644 --- a/torch_sim/integrators/md.py +++ b/torch_sim/integrators/md.py @@ -409,7 +409,7 @@ def init_fn( Q = ( kT_batched.unsqueeze(-1) - * torch.square(tau_batched).unsqueeze(-1) ** 2 + * torch.square(tau_batched).unsqueeze(-1) * torch.ones((n_systems, chain_length), dtype=dtype, device=device) ) Q[:, 0] *= degrees_of_freedom diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index 4a86084c..c236bd45 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -1179,9 +1179,9 @@ def _npt_nose_hoover_compute_cell_force( internal_pressure = torch.trace(stress).unsqueeze(0).expand(n_systems) # Compute force on cell coordinate per system - # F = alpha * KE - dU/dV - P*V*d + # F = alpha * (2 * KE) - dU/dV - P*V*d return ( - (alpha * KE_per_system) + (alpha * 2 * KE_per_system) - (internal_pressure * volume) - (external_pressure * volume * dim) ) @@ -1226,21 +1226,18 @@ def _npt_nose_hoover_inner_step( volume, volume_to_cell = _npt_nose_hoover_cell_info(state) cell = volume_to_cell(volume) - # Get model output - state.cell = cell - model_output = model(state) - # First half step: Update momenta - n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) - alpha = 1 + 1 / n_atoms_per_system # [n_systems] + # alpha = 1 + dim / degrees_of_freedom (3 * natoms - 3) + alpha = 1 + 3 / state.get_number_of_degrees_of_freedom() # [n_systems] + # Reuse stress from previous step since positions and cell unchanged cell_force_val = _npt_nose_hoover_compute_cell_force( alpha=alpha, volume=volume, positions=positions, momenta=momenta, masses=masses, - stress=model_output["stress"], + stress=state.stress, external_pressure=external_pressure, system_idx=state.system_idx, ) @@ -1406,7 +1403,8 @@ def npt_nose_hoover_init( ) # Compute total DOF for thermostat initialization and a zero KE placeholder - dof_per_system = torch.bincount(state.system_idx, minlength=n_systems) * dim + dof_per_system = state.get_number_of_degrees_of_freedom() - 3 + KE_thermostat = ts.calc_kinetic_energy( masses=state.masses, momenta=momenta, system_idx=state.system_idx ) @@ -1612,13 +1610,12 @@ def npt_nose_hoover_invariant( ) # Calculate degrees of freedom per system - n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) - dof_per_system = n_atoms_per_system * state.positions.shape[-1] # n_atoms * n_dim + dof_per_system = state.get_number_of_degrees_of_freedom() # Initialize total energy with PE + KE e_tot = e_pot + e_kin_per_system - # Add thermostat chain contributions (batched per system, DOF = n_atoms * 3) + # Add thermostat chain contributions (batched per system, DOF = 3 * n_atoms - 3) e_tot += _compute_chain_energy(state.thermostat, kT, e_tot, dof_per_system) # Add barostat chain contributions (batched per system, DOF = 1) diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index 8e74bf85..1b801727 100644 --- a/torch_sim/integrators/nvt.py +++ b/torch_sim/integrators/nvt.py @@ -314,11 +314,9 @@ def nvt_nose_hoover_init( masses=state.masses, momenta=momenta, system_idx=state.system_idx ) - # Calculate degrees of freedom per system - n_atoms_per_system = torch.bincount(state.system_idx) - dof_per_system = ( - n_atoms_per_system * state.positions.shape[-1] - ) # n_atoms * n_dimensions + # Calculate degrees of freedom per system (subtract 3 for COM motion, + # matching LAMMPS compute_temp which uses dof = 3N - 3) + dof_per_system = state.get_number_of_degrees_of_freedom() - 3 # Initialize state return NVTNoseHooverState.from_state( @@ -431,9 +429,8 @@ def nvt_nose_hoover_invariant( masses=state.masses, momenta=state.momenta, system_idx=state.system_idx ) - # Get system degrees of freedom per system - n_atoms_per_system = torch.bincount(state.system_idx) - dof = n_atoms_per_system * state.positions.shape[-1] # n_atoms * n_dimensions + # Get system degrees of freedom per system (3N - 3 for COM correction) + dof = state.get_number_of_degrees_of_freedom() # Start with system energy e_tot = e_pot + e_kin From e914b3f294ac8e8a485acc5ba536c7266794f445 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Fri, 13 Feb 2026 16:14:32 +0000 Subject: [PATCH 04/10] 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 90de3c4a..a56b49bc 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 00000000..e57f85c8 --- /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 05/10] 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 00000000..a60ee944 --- /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 00000000..444948d6 --- /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 00000000..69e21fd0 --- /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 00000000..0b18f80c --- /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 00000000..a0d60ca9 --- /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 00000000..ba8d141c --- /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 06/10] 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 950f1664..fb18a796 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 07/10] 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 72c25c07..e2028f80 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 5a73d580..60c3a2f2 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 08/10] 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 81d22078..1a360da1 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 09/10] - 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 c236bd45..15c591c9 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 07f3064b..d37ff7ed 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 1b801727..eadfd4ad 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 10/10] 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 a60ee944..00000000 --- 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 444948d6..00000000 --- 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 69e21fd0..00000000 --- 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 0b18f80c..00000000 --- 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 a0d60ca9..00000000 --- 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 ba8d141c..00000000 --- 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()