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.900    9314.02       9984      3.713  2.313e-03
     200       1.2437     -19.7223     -18.4786      37.73   1874.371    9315.29      10304      3.324  1.421e-01
     400       1.5568     -20.0382     -18.4814      47.23   1616.679    9320.21      10251      3.313  1.767e-01
     600       1.3286     -19.8128     -18.4842      40.31   1759.341    9328.40      10272      3.345  1.635e-01
     800       1.5427     -20.0477     -18.5050      46.80   1535.262    9339.90      10238      3.238  1.957e-01
    1000       1.5328     -20.0559     -18.5231      46.50   1490.899    9354.50      10246      3.319  1.423e-01
    1200       1.4868     -20.0141     -18.5273      45.11   1456.033    9372.10      10230      3.324  1.814e-01
    1400       1.4602     -19.9886     -18.5284      44.30   1399.313    9392.63      10220      3.254  1.923e-01
    1600       1.5267     -20.0742     -18.5475      46.32   1264.287    9416.05      10197      3.296  1.873e-01
    1800       1.5781     -20.1337     -18.5556      47.88   1126.031    9442.13      10184      3.334  1.587e-01
    2000       1.4324     -20.0079     -18.5755      43.46   1150.127    9470.71      10200      3.331  1.761e-01
    2200       1.4594     -20.0437     -18.5843      44.28   1015.622    9501.66      10168      3.305  1.580e-01
    2400       1.3964     -19.9745     -18.5781      42.36    957.259    9534.74      10164      3.332  2.141e-01
    2600       1.4027     -19.9867     -18.5840      42.56    839.253    9569.86      10137      3.321  1.651e-01
    2800       1.4522     -20.0477     -18.5954      44.06    682.781    9606.82      10105      3.376  1.269e-01
    2999       1.4276     -20.0343     -18.6067      43.31    589.422    9645.21      10106      3.323  1.520e-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: 47.04 ± 11.48 K
  Mean Pressure:    1202.536 ± 405.565 atm
  Mean Volume:      9432.77 ± 106.85 ų

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