NPH Dynamics (MTK) with Lennard-Jones Potential

This example demonstrates isothermal-free (no thermostat) NPH dynamics using the Martyna-Tobias-Klein (MTK) barostat and the Lennard-Jones (LJ) potential.

NPH maintains constant enthalpy and pressure (in the extended-system sense), and is useful for testing barostat mechanics and pressure coupling.

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
import warp as wp
from _dynamics_utils import (
    DEFAULT_CUTOFF,
    DEFAULT_SKIN,
    EPSILON_AR,
    SIGMA_AR,
    MDSystem,
    create_fcc_argon,
    pressure_atm_to_ev_per_a3,
    run_nph_mtk,
)

print("=" * 95)
print("NPH (MTK) DYNAMICS WITH LENNARD-JONES POTENTIAL")
print("=" * 95)
print()

device = "cuda:0" if wp.is_cuda_available() else "cpu"
print(f"Using device: {device}")

print("\n--- Creating FCC Argon System ---")
positions, cell = create_fcc_argon(num_unit_cells=4, a=5.26)  # 256 atoms
print(f"Created {len(positions)} atoms in {cell[0, 0]:.2f} ų box")

print("\n--- Initializing MD System ---")
system = MDSystem(
    positions=positions,
    cell=cell,
    epsilon=EPSILON_AR,
    sigma=SIGMA_AR,
    cutoff=DEFAULT_CUTOFF,
    skin=DEFAULT_SKIN,
    switch_width=0.0,
    device=device,
    dtype=np.float64,
)

print("\n--- Setting Initial Temperature ---")
system.initialize_temperature(temperature=94.4, seed=42)

===============================================================================================
NPH (MTK) DYNAMICS WITH LENNARD-JONES POTENTIAL
===============================================================================================

Using device: cuda:0

--- Creating FCC Argon System ---
Created 256 atoms in 21.04 ų box

--- Initializing MD System ---
Initialized MD system with 256 atoms
  Cell: 21.04 x 21.04 x 21.04 Å
  Cutoff: 8.50 Å (+ 0.50 Å skin)
  LJ: ε = 0.0104 eV, σ = 3.40 Å
  Device: cuda:0, dtype: <class 'numpy.float64'>
  Units: x [Å], t [fs], E [eV], m [eV·fs²/Ų] (from amu), v [Å/fs]

--- Setting Initial Temperature ---
Initialized velocities: target=94.4 K, actual=90.4 K

print("\n--- NPH Run (3000 steps) ---")
print(f"Pressure units sanity: 1 atm = {pressure_atm_to_ev_per_a3(1.0):.6e} eV/ų")
stats = run_nph_mtk(
    system=system,
    num_steps=3000,
    dt_fs=1.0,
    target_pressure_atm=1.0,
    pdamp_fs=5000.0,
    reference_temperature_K=94.4,
    log_interval=200,
)

--- NPH Run (3000 steps) ---
Pressure units sanity: 1 atm = 6.324209e-07 eV/ų

Running 3000 NPH (MTK) steps at P=1.000 atm
  dt = 1.000 fs, pdamp = 5000.0 fs (barostat timescale)
========================================================================================================================
    Step      KE (eV)      PE (eV)   Total (eV)      T (K)    P (atm)    V (Å^3)  Neighbors  min r (Å)     max|F|
========================================================================================================================
       0       2.9787     -21.5621     -18.5833      90.37    483.899    9314.02       9984      3.713  2.313e-03
     200       1.2314     -19.7368     -18.5055      37.36   1776.854    9340.55      10284      3.326  1.403e-01
     400       1.4830     -20.0511     -18.5682      44.99   1205.924    9439.85      10192      3.322  1.694e-01
     600       1.2378     -19.8628     -18.6249      37.55    870.716    9592.65      10132      3.358  1.480e-01
     800       1.3245     -19.9646     -18.6402      40.18    225.838    9783.32      10059      3.253  2.009e-01
    1000       1.2877     -19.8669     -18.5792      39.07   -245.812    9989.04      10000      3.408  1.109e-01
    1200       1.1965     -19.6774     -18.4809      36.30   -581.964   10187.51       9937      3.403  1.097e-01
    1400       1.0739     -19.4352     -18.3613      32.58   -810.741   10360.69       9825      3.378  1.070e-01
    1600       1.0533     -19.3197     -18.2663      31.96  -1008.937   10493.73       9749      3.348  1.220e-01
    1800       1.0086     -19.2108     -18.2022      30.60  -1087.767   10574.99       9680      3.379  1.175e-01
    2000       1.0432     -19.2176     -18.1744      31.65  -1147.906   10597.64       9677      3.382  1.195e-01
    2200       1.0819     -19.2963     -18.2144      32.82  -1125.141   10557.74       9693      3.362  1.151e-01
    2400       1.0197     -19.3096     -18.2898      30.94   -913.656   10458.96       9794      3.383  1.163e-01
    2600       1.1845     -19.5737     -18.3892      35.94   -804.347   10312.34       9913      3.411  1.100e-01
    2800       1.1866     -19.7001     -18.5135      36.00   -470.461   10128.22       9955      3.393  9.997e-02
    2999       1.2248     -19.8286     -18.6038      37.16    -52.197    9926.14      10012      3.402  1.492e-01

print("\n--- Analysis ---")
temps = np.array([s.temperature for s in stats])
pressures_atm = np.array([s.pressure for s in stats])
volumes = np.array([s.volume for s in stats])
steps = np.array([s.step for s in stats])

print(f"  Mean Temperature: {temps.mean():.2f} ± {temps.std():.2f} K")
print(f"  Mean Pressure:    {pressures_atm.mean():.3f} ± {pressures_atm.std():.3f} atm")
print(f"  Mean Volume:      {volumes.mean():.2f} ± {volumes.std():.2f} ų")

fig, ax = plt.subplots(3, 1, figsize=(7.0, 6.5), sharex=True, constrained_layout=True)
ax[0].plot(steps, temps, lw=1.5)
ax[0].set_ylabel("Temperature (K)")
ax[1].plot(steps, pressures_atm, lw=1.5)
ax[1].axhline(1.0, color="k", ls="--", lw=1.0, label="target")
ax[1].set_ylabel("Pressure (atm)")
ax[1].legend(frameon=False, loc="best")
ax[2].plot(steps, volumes, lw=1.5)
ax[2].set_xlabel("Step")
ax[2].set_ylabel(r"Volume ($\AA^3$)")
fig.suptitle("NPH (MTK): Temperature, Pressure, and Volume")

print("\n" + "=" * 95)
print("SIMULATION COMPLETE")
print("=" * 95)

--- Analysis ---
  Mean Temperature: 39.09 ± 13.76 K
  Mean Pressure:    -230.356 ± 878.663 atm
  Mean Volume:      10066.09 ± 440.19 ų

===============================================================================================
SIMULATION COMPLETE
===============================================================================================