Implement NPT Langevin

examples/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

@@ -0,0 +1,154 @@
+"""Lennard-Jones simulation in NPT ensemble using Nose-Hoover chain."""
+
+import os
+
+import torch
+
+from torch_sim.quantities import kinetic_energy, temperature
+from torch_sim.state import BaseState
+from torch_sim.unbatched.models.lennard_jones import UnbatchedLennardJonesModel
+from torch_sim.unbatched.unbatched_integrators import npt_langevin
+from torch_sim.units import MetalUnits as Units
+from torch_sim.units import UnitConversion
+
+
+# Set up the device and data type
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float32
+
+# Set random seed and deterministic behavior for reproducibility
+torch.manual_seed(42)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(42)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set up the random number generator
+generator = torch.Generator(device=device)
+generator.manual_seed(42)  # For reproducibility
+
+# Number of steps to run
+N_steps = 100 if os.getenv("CI") else 10_000
+
+# Create face-centered cubic (FCC) Argon
+# 5.26 Å is a typical lattice constant for Ar
+a_len = 5.26  # Lattice constant
+PERIODIC = True  # Flag to use periodic boundary conditions
+
+# Generate base FCC unit cell positions (scaled by lattice constant)
+base_positions = torch.tensor(
+    [
+        [0.0, 0.0, 0.0],  # Corner
+        [0.0, 0.5, 0.5],  # Face centers
+        [0.5, 0.0, 0.5],
+        [0.5, 0.5, 0.0],
+    ],
+    device=device,
+    dtype=dtype,
+)
+
+# Create 4x4x4 supercell of FCC Argon manually
+positions = []
+for i in range(4):
+    for j in range(4):
+        for k in range(4):
+            for base_pos in base_positions:
+                # Add unit cell position + offset for supercell
+                pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype)
+                positions.append(pos)
+
+# Stack the positions into a tensor
+positions = torch.stack(positions)
+
+# Scale by lattice constant
+positions = positions * a_len
+
+# Create the cell tensor
+cell = torch.tensor(
+    [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype
+)
+
+# Create the atomic numbers tensor (Argon = 18)
+atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int)
+# Create the masses tensor (Argon = 39.948 amu)
+masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype)
+
+# Initialize the Lennard-Jones model
+# Parameters:
+#  - sigma: distance at which potential is zero (3.405 Å for Ar)
+#  - epsilon: depth of potential well (0.0104 eV for Ar)
+#  - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma)
+model = UnbatchedLennardJonesModel(
+    use_neighbor_list=False,
+    sigma=3.405,
+    epsilon=0.0104,
+    cutoff=2.5 * 3.405,
+    device=device,
+    dtype=dtype,
+    compute_force=True,
+    compute_stress=True,
+)
+state = BaseState(
+    positions=positions,
+    masses=masses,
+    cell=cell,
+    pbc=PERIODIC,
+    atomic_numbers=atomic_numbers,
+)
+# Run initial simulation and get results
+results = model(state)
+
+dt = 0.001 * Units.time  # Time step (1 ps)
+kT = 200 * Units.temperature  # Temperature (200 K)
+target_pressure = 10000 * Units.pressure  # Target pressure (10 kbar)
+
+npt_init, npt_update = npt_langevin(
+    model=model,
+    dt=dt,
+    kT=kT,
+    external_pressure=target_pressure,
+)
+
+state = npt_init(state=state, seed=1)
+
+
+def get_pressure(
+    stress: torch.Tensor, kinetic_energy: torch.Tensor, volume: torch.Tensor, dim: int = 3
+) -> torch.Tensor:
+    """Compute the pressure from the stress tensor.
+
+    The stress tensor is defined as 1/volume * dU/de_ij
+    So the pressure is -1/volume * trace(dU/de_ij)
+    """
+    return 1 / (dim) * ((2 * kinetic_energy / volume) - torch.trace(stress))
+
+
+# Run the simulation
+for step in range(N_steps):
+    if step % 50 == 0:
+        temp = temperature(masses=state.masses, momenta=state.momenta) / Units.temperature
+        pressure = get_pressure(
+            model(state)["stress"],
+            kinetic_energy(masses=state.masses, momenta=state.momenta),
+            torch.linalg.det(state.cell),
+        )
+        pressure = pressure.item() / Units.pressure
+        xx, yy, zz = state.cell[0, 0], state.cell[1, 1], state.cell[2, 2]
+        print(
+            f"{step=}: Temperature: {temp:.4f}, "
+            f"{pressure=:.4f}, "
+            f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}"
+        )
+    state = npt_update(state, kT=kT, external_pressure=target_pressure)
+
+temp = temperature(masses=state.masses, momenta=state.momenta) / Units.temperature
+print(f"Final temperature: {temp:.4f}")
+
+
+stress = model(state)["stress"]
+kinetic_energy = kinetic_energy(masses=state.masses, momenta=state.momenta)
+volume = torch.linalg.det(state.cell)
+pressure = get_pressure(stress, kinetic_energy, volume)
+pressure = pressure.item() / Units.pressure
+print(f"Final {pressure=:.4f}")
+print(stress * UnitConversion.eV_per_Ang3_to_GPa)

examples/3_Dynamics/3.12_MACE_NPT_Langevin.py

@@ -0,0 +1,144 @@
+"""NPT simulation with MACE and Nose-Hoover thermostat."""
+
+# /// script
+# dependencies = [
+#     "mace-torch>=0.3.11",
+# ]
+# ///
+
+import os
+
+import torch
+from ase.build import bulk
+from mace.calculators.foundations_models import mace_mp
+
+from torch_sim.neighbors import vesin_nl_ts
+from torch_sim.quantities import kinetic_energy, temperature
+from torch_sim.state import BaseState
+from torch_sim.unbatched.models.mace import UnbatchedMaceModel
+from torch_sim.unbatched.unbatched_integrators import (
+    npt_langevin,
+    nvt_nose_hoover,
+    nvt_nose_hoover_invariant,
+)
+from torch_sim.units import MetalUnits as Units
+
+
+# Set device and data type
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float32
+
+# Option 1: Load the raw model from the downloaded model
+mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
+loaded_model = mace_mp(
+    model=mace_checkpoint_url,
+    return_raw_model=True,
+    default_dtype=dtype,
+    device=device,
+)
+
+# Option 2: Load from local file (comment out Option 1 to use this)
+# MODEL_PATH = "../../../checkpoints/MACE/mace-mpa-0-medium.model"
+# loaded_model = torch.load(MODEL_PATH, map_location=device)
+
+PERIODIC = True
+
+# Create diamond cubic Silicon
+si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2))
+
+# Prepare input tensors
+positions = torch.tensor(si_dc.positions, device=device, dtype=dtype)
+cell = torch.tensor(si_dc.cell.array, device=device, dtype=dtype)
+atomic_numbers = torch.tensor(si_dc.get_atomic_numbers(), device=device, dtype=torch.int)
+masses = torch.tensor(si_dc.get_masses(), device=device, dtype=dtype)
+
+# Print shapes for verification
+print(f"Positions: {positions.shape}")
+print(f"Cell: {cell.shape}")
+
+# Initialize the unbatched MACE model
+model = UnbatchedMaceModel(
+    model=loaded_model,
+    device=device,
+    neighbor_list_fn=vesin_nl_ts,
+    periodic=PERIODIC,
+    compute_force=True,
+    compute_stress=True,
+    dtype=dtype,
+    enable_cueq=False,
+)
+state = BaseState(
+    positions=positions,
+    masses=masses,
+    cell=cell,
+    pbc=PERIODIC,
+    atomic_numbers=atomic_numbers,
+)
+# Run initial inference
+results = model(state)
+
+N_steps_nvt = 20 if os.getenv("CI") else 2_000
+N_steps_npt = 20 if os.getenv("CI") else 2_000
+dt = 0.001 * Units.time  # Time step (1 ps)
+kT = 300 * Units.temperature  # Initial temperature (300 K)
+target_pressure = 10000 * Units.pressure  # Target pressure (0 bar)
+
+nvt_init, nvt_update = nvt_nose_hoover(model=model, kT=kT, dt=dt)
+state = nvt_init(state=state, seed=1)
+
+for step in range(N_steps_nvt):
+    if step % 10 == 0:
+        temp = temperature(masses=state.masses, momenta=state.momenta) / Units.temperature
+        invariant = nvt_nose_hoover_invariant(state, kT=kT).item()
+        print(f"{step=}: Temperature: {temp:.4f}: invariant: {invariant:.4f}, ")
+    state = nvt_update(state, kT=kT)
+
+npt_init, npt_update = npt_langevin(
+    model=model, kT=kT, dt=dt, external_pressure=target_pressure
+)
+state = npt_init(state=state, seed=1)
+
+
+def get_pressure(
+    stress: torch.Tensor, kinetic_energy: torch.Tensor, volume: torch.Tensor, dim: int = 3
+) -> torch.Tensor:
+    """Compute the pressure from the stress tensor.
+
+    The stress tensor is defined as 1/volume * dU/de_ij
+    So the pressure is -1/volume * trace(dU/de_ij)
+    """
+    return 1 / dim * ((2 * kinetic_energy / volume) - torch.trace(stress))
+
+
+for step in range(N_steps_npt):
+    if step % 10 == 0:
+        temp = temperature(masses=state.masses, momenta=state.momenta) / Units.temperature
+        stress = model(state)["stress"]
+        volume = torch.det(state.cell)
+        pressure = (
+            get_pressure(
+                stress, kinetic_energy(masses=state.masses, momenta=state.momenta), volume
+            ).item()
+            / Units.pressure
+        )
+        xx, yy, zz = torch.diag(state.cell)
+        print(
+            f"{step=}: Temperature: {temp:.4f}, "
+            f"pressure: {pressure:.4f}, "
+            f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}"
+        )
+    state = npt_update(state, kT=kT, external_pressure=target_pressure)
+
+final_temp = temperature(masses=state.masses, momenta=state.momenta) / Units.temperature
+print(f"Final temperature: {final_temp:.4f} K")
+final_stress = model(state)["stress"]
+final_volume = torch.det(state.cell)
+final_pressure = (
+    get_pressure(
+        final_stress,
+        kinetic_energy(masses=state.masses, momenta=state.momenta),
+        final_volume,
+    ).item()
+    / Units.pressure
+)
+print(f"Final pressure: {final_pressure:.4f} bar")