Fix 579 2

[curtischong]

Summary

• Feature 1
• Fix 1

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

fixed version of TorchSim#586

MD completed
  systems: 16
  precondition wall: 79.293 s
  md wall: 386.878 s
  system-steps/s: 41.357
  peak CUDA memory: 25.090 GB
  bad indices: []
  bad reasons: []
  final temperature K: [315.43124 307.20044 312.35632 292.31323 306.58423 276.70398 300.48325
 306.9538  316.2425  282.4878  304.26224 306.15756 254.9931  301.22366
 269.04944 309.7748 ]
  final energy eV: [-559.7459  -558.1463  -557.8347  -557.28534 -557.3563  -560.15857
 -558.66504 -560.3809  -559.75604 -557.13654 -557.65424 -558.678
 -559.63214 -557.4609  -558.0716  -556.77484]
  final min distance A: [2.1968294902802037, 2.220445678908415, 2.0329850879817677, 2.195226021556654, 1.981977865951136, 2.193242034845892, 2.2171705179385155, 2.1947683840575833, 2.173143383432847, 2.1819499283504364, 2.1929907916179383, 2.2170552769413563, 2.231164529097381, 2.2007128872956376, 2.2006608344266665, 2.2029369811994153]

"""Minimal reproduction of torch-sim issue #579 (Nose-Hoover `tau` units).

Run:  uv run python repro_579_min.py

The reporter ran batched NVT MD (16 amorphous NaTaCl6 systems, SevenNet-OMNI,
NVT Nose-Hoover, 2 fs, tdamp 200 fs) and got NaN energy/temperature in *some*
batch elements -- with D3 on AND off, and with MatterSim too. ASE stayed finite
under the same physical settings, so it is a torch-sim integrator bug, not a
potential bug.

ROOT CAUSE
    ``ts.integrate`` converts ``timestep`` to internal units but used to forward
    the Nose-Hoover relaxation time ``tau`` unconverted. The reporter passed
    ``tau = 200 fs = 0.2 ps`` and ``timestep = 2 fs = 0.002 ps`` (same
    convention). In metal units time is scaled ~98x, so ``dt`` became ~0.196
    while ``tau`` stayed 0.2 -- i.e. ``tau ~ dt``. A Nose-Hoover chain with
    relaxation time ~= one step has a near-zero thermostat mass and is
    *marginally* unstable: depending on the initial velocity draw, some batch
    systems tip into NaN -- exactly the "some batch elements end with NaN"
    symptom. The fix converts ``tau``/``b_tau`` with the same factor as
    ``timestep`` (torch_sim/runners.py).

Two demos, each comparing BUGGY tau (pre-fix, tau~dt) vs FIXED tau (converted):
  1. Lennard-Jones, single system -- fast, no download, fully deterministic.
  2. SevenNet (7net-0) on the reporter's real 16 structures, batched -- faithful;
     downloads a ~100 MB checkpoint and reproduces the subset-NaN symptom.
"""

from __future__ import annotations

import numpy as np
import torch
from ase import Atoms

import torch_sim as ts
from torch_sim.integrators import nvt_nose_hoover_init, nvt_nose_hoover_step
from torch_sim.models.lennard_jones import LennardJonesModel
from torch_sim.units import MetalUnits

DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float32
DT_PS = 0.002  # 2 fs
TAU_PS = 0.2  # 200 fs
TIME = float(MetalUnits.time)  # ps -> internal (~98.23)


def reporter_structures() -> list[Atoms]:
    """16 packmol-like NaTaCl6 structures (128 atoms, 2.9 g/cm^3), generated in memory.

    Reproduces the reporter's batch: dense NaTaCl6 at box 15.62 A (= 2.9 g/cm^3),
    one structure per seed so divergence hits a marginal subset, not all 16.
    """
    return [lj_single_structure(seed=11001 + i) for i in range(16)]


def lj_single_structure(seed: int = 11001) -> Atoms:
    """Author-like NaTaCl6 (128 atoms, box 15.62, min spacing 2.26 A, no overlap)."""
    box, rng = 15.62, np.random.default_rng(seed)
    numbers = [11] * 16 + [73] * 16 + [17] * 96
    pos: list[np.ndarray] = []
    while len(pos) < len(numbers):
        cand = rng.uniform(0.0, box, size=3)
        if pos:
            d = np.array(pos) - cand
            d -= box * np.round(d / box)
            if np.linalg.norm(d, axis=1).min() < 2.26:
                continue
        pos.append(cand)
    return Atoms(numbers=numbers, positions=np.array(pos), cell=[box] * 3, pbc=True)


def bad_systems(state: ts.state.SimState) -> list[int]:
    return (~torch.isfinite(state.energy)).nonzero().flatten().tolist()


def run_integrate(
    systems: list[Atoms], model, *, tau_ps: float, steps: int
) -> list[int]:
    final = ts.integrate(
        system=systems,
        model=model,
        integrator=ts.Integrator.nvt_nose_hoover,
        n_steps=steps,
        temperature=300.0,
        timestep=DT_PS,
        init_kwargs={"tau": tau_ps, "chain_length": 3, "chain_steps": 1, "sy_steps": 3},
        pbar=False,
    )
    return bad_systems(final)


def lj_demo() -> None:
    print("[1] Lennard-Jones, single system (deterministic, no download)")
    atoms = lj_single_structure()
    # sigma > 2.26 A spacing -> every ion in the repulsive wall (strong forces),
    # a stand-in for the dense ionic melt; divergence is distributed, not a contact.
    model = LennardJonesModel(sigma=3.0, epsilon=0.3, cutoff=7.5, device=DEV, dtype=DTYPE)
    dt = torch.tensor(DT_PS * TIME)
    kT = torch.tensor(300.0 * MetalUnits.temperature)
    # buggy: tau forwarded UNCONVERTED -> 0.2 internal ~= dt
    torch.manual_seed(0)
    s = nvt_nose_hoover_init(
        ts.io.atoms_to_state([atoms], device=DEV, dtype=DTYPE),
        model,
        kT=kT,
        dt=dt,
        tau=TAU_PS,
    )
    for _ in range(200):
        s = nvt_nose_hoover_step(s, model, dt=dt, kT=kT)
    buggy = bad_systems(s)
    # fixed: ts.integrate converts tau like timestep
    fixed = run_integrate([atoms], model, tau_ps=TAU_PS, steps=200)
    print(f"    buggy (tau~dt)   diverged systems = {buggy}  (expect [0])")
    print(f"    fixed (tau*time) diverged systems = {fixed}  (expect [])\n")


def sevennet_demo() -> None:
    print("[2] SevenNet 7net-0, reporter's 16 real structures, batched")
    from torch_sim.models.sevennet import SevenNetModel

    systems = reporter_structures()
    model = SevenNetModel(model="7net-0", device=DEV, dtype=DTYPE)
    # buggy: reconstruct pre-fix internal tau (0.2) by undoing integrate's conversion
    buggy = run_integrate(systems, model, tau_ps=TAU_PS / TIME, steps=200)
    fixed = run_integrate(systems, model, tau_ps=TAU_PS, steps=200)
    print(f"    buggy (tau~dt)   diverged systems = {buggy}  (a marginal subset NaNs)")
    print(f"    fixed (tau*time) diverged systems = {fixed}  (expect [])\n")


def main() -> None:
    print(f"device={DEV}  timestep=2 fs  tau=200 fs  tau/dt(buggy)~1  steps=200\n")
    lj_demo()
    sevennet_demo()


if __name__ == "__main__":
    main()

[curtischong]

closing since we merged this into main TorchSim#587