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.7224     -18.4787      37.73   1874.338    9315.29      10304      3.324  1.421e-01
     400       1.5565     -20.0385     -18.4820      47.22   1616.424    9320.21      10251      3.313  1.767e-01
     600       1.3279     -19.8136     -18.4857      40.29   1758.562    9328.40      10272      3.345  1.634e-01
     800       1.5413     -20.0489     -18.5077      46.76   1534.009    9339.90      10238      3.238  1.956e-01
    1000       1.5305     -20.0579     -18.5274      46.43   1488.875    9354.50      10245      3.320  1.421e-01
    1200       1.4836     -20.0173     -18.5336      45.01   1453.012    9372.08      10230      3.324  1.812e-01
    1400       1.4560     -19.9927     -18.5367      44.17   1395.335    9392.60      10221      3.255  1.918e-01
    1600       1.5210     -20.0794     -18.5584      46.14   1259.430    9415.99      10198      3.296  1.870e-01
    1800       1.5707     -20.1399     -18.5692      47.65   1120.123    9442.05      10183      3.334  1.584e-01
    2000       1.4235     -20.0152     -18.5917      43.19   1143.202    9470.57      10199      3.332  1.754e-01
    2200       1.4500     -20.0538     -18.6038      43.99   1006.484    9501.46      10166      3.306  1.569e-01
    2400       1.3852     -19.9863     -18.6011      42.02    946.937    9534.46      10164      3.333  2.134e-01
    2600       1.3896     -19.9992     -18.6097      42.16    828.089    9569.49      10135      3.322  1.641e-01
    2800       1.4376     -20.0617     -18.6240      43.62    670.188    9606.32      10105      3.378  1.259e-01
    2999       1.4092     -20.0481     -18.6389      42.75    577.191    9644.55      10104      3.324  1.515e-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: 46.84 ± 11.53 K
  Mean Pressure:    1197.256 ± 408.620 atm
  Mean Volume:      9432.62 ± 106.66 ų

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