NVIDIA/nvalchemi-toolkit v0.1.0

v0.1.0

Repository: NVIDIA/nvalchemi-toolkit

Tag: v0.1.0

Published: 2026-04-16T08:42:01Z

Prerelease: no

Release notes: Initial public-beta release of NVIDIA ALCHEMI Toolkit, a GPU-first Python framework for AI-driven atomic simulation workflows.

Core Data Layer

• AtomicData — Pydantic-backed graph representation of atomic systems

(positions, atomic numbers, masses, node/edge properties) with factory constructors from_atoms() (ASE) and from_structure() (pymatgen).

• Batch — GPU-resident graph batch with MultiLevelStorage backend

supporting node-, edge-, and system-level tensors. Lazy batch_idx/batch_ptr, index_select, append, and from_data_list for efficient batching.

• Zarr I/O — AtomicDataZarrWriter and AtomicDataZarrReader with

configurable Zstd compression, chunking, and sharding for high-throughput trajectory storage.

• Dataset & DataLoader — CUDA-stream prefetching, async I/O, and

drop-in DataLoader replacement yielding Batch objects.

Model Wrappers

All wrappers implement BaseModelMixin with a unified ModelConfig for capability declaration and runtime control.

• DemoModelWrapper — Lightweight test/demo model (point-cloud energy +

autograd forces).

• MACEWrapper — MACE equivariant neural network; supports foundation

checkpoints; COO neighbor format; conservative forces via autograd.

• AIMNet2Wrapper — AIMNet2 atom-in-molecule network; energy, forces,

charges, stress; MATRIX neighbor format; NSE auto-detection.

• LennardJonesModelWrapper — Warp-accelerated single-species LJ with

analytical forces and optional virial stress.

• EwaldModelWrapper — Real + reciprocal space Ewald summation for

periodic charged systems; k-vector caching; hybrid analytical forces.

• PMEModelWrapper — Particle Mesh Ewald (FFT-based, O(N log N)) for

large periodic systems.

• DFTD3ModelWrapper — DFT-D3(BJ) dispersion correction with

auto-downloaded reference parameters and cutoff smoothing.

• PipelineModelWrapper — Compose multiple models into groups with

independent derivative strategies (autograd vs. analytical).

Dynamics Engine

• BaseDynamics — Abstract base orchestrating model evaluation, integrator

updates, hook dispatch, convergence detection, and inflight batching.

• 9 hook insertion points per step (DynamicsStage enum): BEFORE_STEP,

BEFORE_PRE_UPDATE, AFTER_PRE_UPDATE, BEFORE_COMPUTE, AFTER_COMPUTE, BEFORE_POST_UPDATE, AFTER_POST_UPDATE, AFTER_STEP, ON_CONVERGE.

• ConvergenceHook — Flexible convergence criteria with from_fmax()

convenience constructor and per-system masking.

Integrators

• NVE — Velocity Verlet; symplectic, time-reversible, energy-conserving.
• NVTLangevin — BAOAB Langevin dynamics with Ornstein-Uhlenbeck

thermostat for canonical sampling.

• NVTNoseHoover — Nosé-Hoover chain thermostat with Yoshida-Suzuki

factorization; deterministic and ergodic.

• NPT — Martyna-Tobias-Klein isothermal-isobaric with dual Nosé-Hoover

chains (particle + cell DOFs).

• NPH — MTK isenthalpic-isobaric without thermostat.

Optimizers

• FIRE — Fast Inertial Relaxation Engine with adaptive timestep.
• FIREVariableCell — FIRE with NPH-like variable-cell propagation.
• FIRE2 — Improved FIRE (Shuang et al. 2020) with better restart

conditions and modified velocity mixing.

• FIRE2VariableCell — FIRE2 with variable-cell structural relaxation.

Built-in Hooks

Dynamics hooks (nvalchemi.dynamics.hooks):

• LoggingHook — Per-graph scalar statistics with thread-pooled I/O and

optional CUDA stream prefetch.

• NaNDetectorHook — Immediate NaN/Inf detection in forces and energy.
• MaxForceClampHook — Clamps force magnitudes to prevent numerical

explosions.

• EnergyDriftMonitorHook — Cumulative energy drift tracking with

configurable thresholds (absolute and per-atom-per-step).

• FreezeAtomsHook — Freezes selected atoms by category during MD.
• SnapshotHook — Periodic full-state snapshots to a DataSink.
• ConvergedSnapshotHook — Snapshot on convergence.
• ProfilerHook — Per-stage wall-clock profiling with NVTX annotations

and CSV output.

• AlignCellHook — Upper-triangular cell alignment for variable-cell

optimization.

General hooks (nvalchemi.hooks):

• NeighborListHook — On-the-fly neighbor list construction/refresh with

Verlet skin buffer; MATRIX and COO formats.

• WrapPeriodicHook — GPU-accelerated PBC wrapping via Warp kernel.
• BiasedPotentialHook — External bias potentials for enhanced sampling

(umbrella sampling, metadynamics, etc.).

Multi-stage Pipelines

• FusedStage (+ operator) — Compose dynamics stages on a single GPU

with shared forward pass and masked updates per sub-stage.

• DistributedPipeline (| operator) — Distribute stages across GPU

ranks with blocking inter-rank communication.

• SizeAwareSampler — Bin-packing inflight batching that respects

max_atoms, max_edges, and max_batch_size constraints.

• Data sinks — HostMemory (CPU), GPUBuffer (device), ZarrData

(persistent disk) for capturing pipeline outputs.

GPU Primitives

All low-level kernels built on `nvalchemi-toolkit-ops` via NVIDIA Warp:

• Velocity Verlet position/velocity updates
• BAOAB Langevin half-steps
• Nosé-Hoover chain integration
• MTK barostat (NPT/NPH) cell and position propagation
• FIRE/FIRE2 coordinate and cell steps
• Kinetic energy and velocity initialization
• Neighbor list rebuild with Verlet skin
• Cell alignment to upper-triangular form

Developer & Agent Experience

• 20 worked examples across four tiers (basic, intermediate, advanced,

distributed) covering data structures, optimization, MD ensembles, Zarr I/O, inflight batching, custom hooks, model composition, Ewald electrostatics, and multi-GPU pipelines.

• 7 Claude Code agent skills (.claude/skills/) for guided workflows:

model wrapping, data structures, data storage, dynamics API, dynamics hooks, dynamics implementation, and engineering scoping.

• OptionalDependency guards for graceful degradation when MACE, AIMNet2,

ASE, or pymatgen are not installed.

Requirements

• Python 3.11–3.13
• PyTorch >= 2.5.1
• nvalchemi-toolkit-ops[torch] >= 0.3.1
• Optional: [mace], [aimnet], [ase], [pymatgen] extras