🐛[BUG]: NVE energy drift with NeighborListHook(skin>0) + WrapPeriodicHook: position folds bypass the Verlet-skin rebuild check

[FabioFalcioni]

Version

0.2.0-rc (commit 3b6e3f7); the affected code is identical on main (bdc7fb6).

On which installation method(s) does this occur?

Pip, Source

Describe the issue

NVE molecular dynamics does not conserve energy when NeighborListHook(skin > 0) is combined with per-step position wrapping (WrapPeriodicHook), the pairing shown in examples/basic/04_nve_energy_conservation.py. Expected: with a smooth-cutoff potential, a correct Verlet-skin implementation conserves total energy to integrator noise. Observed: total energy rises steadily; in a 128-atom liquid MD run with a MACE-MP potential the drift was ~1.5 meV/atom/step (+26 eV over 20 ps), heating a 600 K replica to 2713 K, and the minimal LJ example below reproduces the same behavior on CPU.

Root cause. The Verlet-skin rebuild check in nvalchemi/hooks/neighbor_list.py (_rebuild, the _batch_nl_rebuild_inplace(...) call) passes cell / cell_inv / pbc, so displacement since the last build is measured with the minimum-image convention. A position fold by WrapPeriodicHook moves an atom by exactly one lattice vector, so its MIC displacement is ~0 and no rebuild fires. But the neighbor_matrix_shifts stored at the last build are now stale by one lattice vector for every pair involving the folded atom, so the model computes wrong edge vectors (r_ij = pos_j - pos_i + S @ cell) until the next rebuild. Each periodic boundary crossing injects a discontinuous quantity of energy; in a diffusive system that is a steady heating rate. The nvalchemiops kernels behave exactly as documented (the build searches to cutoff + skin; the detection kernel computes MIC displacement correctly) - the defect is that MIC is the wrong metric at this call site, because list validity depends on the coordinate representation the shifts were built from, not on physical displacement modulo the lattice.

Isolation matrix (64-atom periodic LJ argon fluid, C2-switched cutoff, NVE, dt = 2 fs, 5 ps; script below):

skin (Å)	wrap	max drift (eV)	eV/atom/step
0.0	on	0.000165	1.0e-09
0.0	off	0.000084	5.2e-10
0.5	on	1.191308	7.5e-06
0.5	off	0.000084	5.2e-10

An instrumented run confirms drift is proportional to the number of fold events (~0.012-0.015 eV per boundary crossing here); a stable fcc solid with zero crossings conserves exactly. Note that lattices built from the origin place many atoms exactly on the cell boundary, so even cold solids drift when constructed that way - no diffusion required.

Aggravating factor: the safety net is blind to this failure mode. The 128-atom production run (not the below MRE but my initial tests) had EnergyDriftMonitorHook(threshold=1e-3, metric="per_atom_per_step") registered and it never fired: the leak (~1e-5 eV/atom/step) sits far below any sensible per-step threshold while being catastrophic cumulatively (+26 eV, 600 K to 2713 K). The monitor currently has no criterion that accumulates drift over the run, so this class of slow leak passes every check the toolkit offers. The drift is also invisible at short test scale (a 1,500-step smoke run barely warms the cell) and dominant at production scale (20,000 steps).

Recommended course of action (in preference order; happy to open a PR for either):

1. Call-site fix in nvalchemi-toolkit (validated): measure raw Cartesian displacement in the skin check (drop cell/cell_inv/pbc from the detection call). A fold then trivially exceeds skin / 2 and forces a rebuild with fresh shifts. Validated locally: drift falls from 7.5e-06 to 1.1e-09 eV/atom/step (the skin = 0 control level); fold-free trajectories show byte-identical drift and identical rebuild counts (raw and MIC displacement agree when no fold occurred); cost was +62 rebuild steps out of 2500 in a hot liquid. Also removes a per-step 3x3 cell inverse. Two regression tests are ready (a fold must set the rebuild flag; stored shifts must still give in-cutoff edge vectors after a fold) - both fail on current main and pass with the fix.
2. Fold-aware detection in nvalchemi-toolkit-ops (optimization, later): a mode that recognizes a pure fold and updates the stored reference positions and shift vectors in place instead of rebuilding would legitimately restore the rebuild-avoidance optimization that MIC presumably intended. Only sound if the shifts are actually updated, which nothing does today.
3. Defense in depth in nvalchemi-toolkit: give EnergyDriftMonitorHook a cumulative-drift criterion (e.g. total |E - E0| per atom since the start of the run, or a windowed slope) alongside the per-step metric, so slow leaks of any origin are caught in NVE runs regardless of cause.

Workarounds for users today: leave skin at its default 0.0 (rebuild every step), or do not wrap positions every step while using skin > 0.

Related minor doc issue: the comment in examples/basic/04_nve_energy_conservation.py states the Verlet skin is "0.5 Å by default"; the actual NeighborConfig.skin default is 0.0 (which is why the example itself does not drift).

Minimum reproducible example

"""NVE energy drift with NeighborListHook(skin > 0) + WrapPeriodicHook.

Run:      python repro_verlet_skin_nve_drift.py
Expected: skin=0.5 drift comparable to the skin=0 control (~1e-9 eV/atom/step).
Actual:   skin=0.5 + wrapping drifts ~7e-6 eV/atom/step (thousands x control).
"""
from __future__ import annotations

import os

os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")  # run eager

import csv
import tempfile
from pathlib import Path

import torch

from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import NVE
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.hooks import LoggingHook
from nvalchemi.hooks import NeighborListHook, WrapPeriodicHook
from nvalchemi.models.lj import LennardJonesModelWrapper

LJ_EPSILON, LJ_SIGMA, LJ_CUTOFF, SWITCH_W = 0.0104, 3.40, 6.0, 1.0
N_SIDE, T_INIT, DT_FS, N_STEPS, SKIN = 4, 300.0, 2.0, 2500, 0.5
KB_EV, M_AR, SEED = 8.617333262e-5, 39.948, 42
_r_min = 2 ** (1 / 6) * LJ_SIGMA
BOX = N_SIDE * _r_min
N_ATOMS = N_SIDE**3
tmpdir = Path(tempfile.mkdtemp(prefix="repro_skin_"))


def make_batch() -> Batch:
    coords = torch.arange(N_SIDE, dtype=torch.float32) * _r_min
    gx, gy, gz = torch.meshgrid(coords, coords, coords, indexing="ij")
    positions = torch.stack([gx.flatten(), gy.flatten(), gz.flatten()], dim=-1)
    torch.manual_seed(SEED)
    v = torch.randn(N_ATOMS, 3) * (KB_EV * T_INIT / M_AR) ** 0.5
    v -= v.mean(dim=0, keepdim=True)
    data = AtomicData(
        positions=positions,
        atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long),
        forces=torch.zeros(N_ATOMS, 3),
        energy=torch.zeros(1, 1),
        cell=torch.eye(3).unsqueeze(0) * BOX,
        pbc=torch.tensor([[True, True, True]]),
    )
    data.add_node_property("velocities", v)
    return Batch.from_data_list([data])


def run_nve(skin: float, wrap: bool) -> float:
    model = LennardJonesModelWrapper(
        epsilon=LJ_EPSILON, sigma=LJ_SIGMA, cutoff=LJ_CUTOFF, switch_width=SWITCH_W
    )
    nve = NVE(model=model, dt=DT_FS, n_steps=N_STEPS)
    nve.register_hook(
        NeighborListHook(
            config=model.model_config.neighbor_config, skin=skin,
            stage=DynamicsStage.BEFORE_COMPUTE,
        ),
        stage=DynamicsStage.BEFORE_COMPUTE,
    )
    if wrap:
        nve.register_hook(WrapPeriodicHook(stage=DynamicsStage.AFTER_POST_UPDATE))
    log_path = tmpdir / f"nve_skin{skin}_wrap{int(wrap)}.csv"
    with LoggingHook(backend="csv", log_path=str(log_path), frequency=10) as log_hook:
        nve.register_hook(log_hook)
        nve.run(make_batch())
    etot = []
    with open(log_path) as f:
        for row in csv.DictReader(f):
            pe = float(row["energy"])
            ke = 1.5 * N_ATOMS * KB_EV * float(row["temperature"])
            etot.append(ke + pe)
    return max(abs(e - etot[0]) for e in etot)


print(f"{'skin (A)':>9}  {'wrap':>5}  {'max|dE| (eV)':>13}  {'eV/atom/step':>13}")
results = {}
for skin in (0.0, SKIN):
    for wrap in (True, False):
        drift = results[(skin, wrap)] = run_nve(skin, wrap)
        print(f"{skin:>9.1f}  {str(wrap):>5}  {drift:>13.6f}  {drift / (N_ATOMS * N_STEPS):>13.2e}")

control, bug = results[(0.0, True)], results[(SKIN, True)]
assert bug <= 5 * max(control, 1e-6), (
    f"NVE energy drift with skin={SKIN} is {bug:.6f} eV vs control "
    f"{control:.6f} eV - Verlet-skin path is not energy-conserving"
)

Relevant log output

 skin (A)   wrap   max|dE| (eV)   eV/atom/step
      0.0   True       0.000165       1.03e-09
      0.0  False       0.000084       5.23e-10
      0.5   True       1.191308       7.45e-06
      0.5  False       0.000084       5.23e-10

AssertionError: NVE energy drift with skin=0.5 is 1.191308 eV vs control
0.000165 eV - Verlet-skin path is not energy-conserving

Environment details

+ Environment location: Bare-metal (macOS, Apple Silicon, CPU-only repro)
+ GPU: none for the MRE; the same drift was first observed on NVIDIA H200 production MD
+ CUDA version: n/a for the CPU MRE (13.0 on the H200 runs where first observed)
+ Python version: 3.11
+ PyTorch version: 2.12.1 (CPU)
+ NVIDIA Warp version: 1.14.0
+ nvalchemi-toolkit-ops version: 0.4.0