NPT Dynamics (MTK + Nosé-Hoover Chain) with Lennard-Jones Potential

This example demonstrates isothermal-isobaric (NPT) dynamics using the Martyna-Tobias-Klein (MTK) barostat coupled with a Nosé-Hoover chain (NHC) thermostat, driven by the Lennard-Jones (LJ) potential.

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_npt_mtk,
)

print("=" * 95)
print("NPT (MTK + NHC) 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)

===============================================================================================
NPT (MTK + NHC) 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--- NPT Run (3000 steps) ---")
print(f"Pressure units sanity: 1 atm = {pressure_atm_to_ev_per_a3(1.0):.6e} eV/ų")
stats = run_npt_mtk(
    system=system,
    num_steps=3000,
    dt_fs=1.0,
    target_temperature_K=94.4,
    target_pressure_atm=1.0,
    tdamp_fs=500.0,
    pdamp_fs=5000.0,
    chain_length=3,
    log_interval=200,
)

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

Running 3000 NPT (MTK) steps at T=94.4 K, P=1.000 atm
  dt = 1.000 fs, tdamp = 500.0 fs, pdamp = 5000.0 fs, chain_length = 3
========================================================================================================================
    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.2737     -19.7216     -18.4478      38.64   1794.983    9340.61      10285      3.325  1.415e-01
     400       1.6045     -19.9387     -18.3341      48.68   1313.033    9441.24      10194      3.307  1.811e-01
     600       1.4326     -19.6397     -18.2071      43.46   1064.340    9600.07      10139      3.340  1.629e-01
     800       1.6368     -19.6574     -18.0206      49.66    464.819    9805.32      10062      3.225  2.114e-01
    1000       1.7256     -19.4639     -17.7383      52.35     29.893   10037.40       9983      3.361  1.424e-01
    1200       1.7435     -19.1470     -17.4035      52.90   -268.667   10277.02       9827      3.340  1.292e-01
    1400       1.6746     -18.6999     -17.0253      50.80   -434.878   10509.14       9613      3.324  1.414e-01
    1600       1.7595     -18.3995     -16.6400      53.38   -612.870   10721.88       9391      3.208  2.462e-01
    1800       1.8449     -18.0848     -16.2399      55.97   -677.709   10904.46       9193      3.311  1.612e-01
    2000       1.9823     -17.8171     -15.8348      60.14   -727.791   11050.48       9080      3.274  2.109e-01
    2200       2.0926     -17.5202     -15.4276      63.49   -647.095   11154.74       8932      3.217  2.056e-01
    2400       2.2495     -17.2974     -15.0480      68.25   -545.253   11219.64       8874      3.257  2.229e-01
    2600       2.4450     -17.1019     -14.6569      74.18   -417.020   11251.64       8905      3.243  2.005e-01
    2800       2.7895     -17.0764     -14.2869      84.63   -368.390   11257.39       8792      3.253  1.851e-01
    2999       2.8218     -16.7641     -13.9423      85.61    -87.506   11243.19       8840      3.217  2.354e-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("  Target Temperature: 94.4 K")
print()
print(f"  Mean Pressure:    {pressures_atm.mean():.3f} ± {pressures_atm.std():.3f} atm")
print("  Target Pressure:  1.0 atm")
print()
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].axhline(94.4, color="k", ls="--", lw=1.0, label="target")
ax[0].set_ylabel("Temperature (K)")
ax[0].legend(frameon=False, loc="best")
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("NPT (MTK + NHC): Temperature, Pressure, and Volume")

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

--- Analysis ---
  Mean Temperature: 60.78 ± 15.15 K
  Target Temperature: 94.4 K

  Mean Pressure:    22.737 ± 757.499 atm
  Target Pressure:  1.0 atm

  Mean Volume:      10445.52 ± 730.32 ų

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