NVIDIA ALCHEMI Toolkit-Ops

https://nvidia.github.io/nvalchemi-toolkit-ops/
Admin
GPU-accelerated computational kernels for molecular dynamics simulations and atomic-scale modeling,...
Tokens:1,286,505
Snippets:6,223
Trust Score:9.7
Update:5 months ago
Show doc for...
Tokens:
Raw
### Example Output

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/neighbors/01_simple_neighbor_list.rst.txt

This is the output of the example script. It confirms successful completion.

```none
Example completed successfully!
```

--------------------------------

### Install PyTorch Backend Extras

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Install ALCHEMI Toolkit-Ops with PyTorch bindings. Verify the installation afterwards.

```bash
$ pip install 'nvalchemi-toolkit-ops[torch]'
```

```python
from nvalchemiops.torch import neighbors; print('PyTorch bindings available')
```

--------------------------------

### Install JAX Backend Extras

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Install ALCHEMI Toolkit-Ops with JAX bindings, including CUDA 12 support. Verify the installation afterwards.

```bash
$ pip install 'nvalchemi-toolkit-ops[jax]'
```

```python
from nvalchemiops.jax import neighbors; print('JAX bindings available')
```

--------------------------------

### Full Installation (x86, without cloning)

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Installs nvalchemi-toolkit-ops, Warp, PyTorch, and JAX for x86 architecture without cloning the repository.

```bash
$ uv venv --seed --python 3.12
$ uv pip install nvalchemi-toolkit-ops \
    https://github.com/NVIDIA/warp/releases/download/v1.12.1/warp_lang-1.12.1+cu13-py3-none-manylinux_2_34_x86_64.whl \
    torch==2.11.0 \
    'jax[cuda13]'
```

--------------------------------

### Install ALCHEMI Toolkit-Ops from PyPI

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Use this command for the most straightforward installation of ALCHEMI Toolkit-Ops.

```bash
$ pip install nvalchemi-toolkit-ops
```

--------------------------------

### Run Batched FIRE Optimization Example

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/dynamics/index.rst.txt

Execute the batched FIRE optimization example script to demonstrate optimizing multiple systems simultaneously.

```bash
python examples/dynamics/08_fire_batched.py
```

--------------------------------

### Run Langevin Integration Example

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/dynamics/index.rst.txt

Execute the Langevin integration example script to simulate molecular dynamics with a Lennard-Jones potential.

```bash
python examples/dynamics/01_langevin_integration.py
```

--------------------------------

### Install ALCHEMI Toolkit-Ops from GitHub Source

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Install directly from the source repository to obtain nightly builds.

```bash
$ pip install git+https://www.github.com/NVIDIA/nvalchemi-toolkit-ops.git
```

--------------------------------

### Install Stable ALCHEMI Toolkit-Ops using uv

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Recommended for production use-cases. This command sets up a virtual environment and installs the stable version.

```bash
$ uv venv --seed --python 3.12
$ uv pip install nvalchemi-toolkit-ops
```

--------------------------------

### Install Nightly ALCHEMI Toolkit-Ops using uv without cloning

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Installs nightly versions directly from GitHub without cloning. Not recommended for production.

```bash
$ uv venv --seed --python 3.12
$ uv pip install git+https://www.github.com/NVIDIA/nvalchemi-toolkit-ops.git
```

--------------------------------

### Install warp-lang

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/dynamics/index.html

Install the warp-lang package using pip. This is a prerequisite for running the provided examples.

```bash
pip install warp-lang

```

--------------------------------

### System Setup and Initialization

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/dynamics/02_velocity_verlet_integration.rst.txt

Sets up the simulation device, creates an FCC lattice of argon atoms, and initializes the MDSystem with LJ parameters and temperature.

```Python
print("=" * 95)
print("VELOCITY VERLET DYNAMICS WITH LENNARD-JONES POTENTIAL (NVE)")
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,
    # NVE can benefit from smooth cutoffs; keep hard cutoff by default for now.
    switch_width=0.0,
    device=device,
    dtype=np.float64,
)

print("\n--- Setting Initial Temperature ---")
system.initialize_temperature(temperature=94.4, seed=42)
```

--------------------------------

### Install Warp

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/dynamics/index.rst.txt

Install the warp-lang package using pip. This is a prerequisite for running the dynamics examples.

```bash
pip install warp-lang
```

--------------------------------

### System Setup and Initialization

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/dynamics/01_langevin_integration.html

Sets up the simulation environment, creates an FCC lattice of argon atoms, initializes the MD system with Lennard-Jones parameters, and sets the initial temperature.

```python
print("=" * 95)
print("LANGEVIN DYNAMICS WITH LENNARD-JONES POTENTIAL")
print("=" * 95)
print()

device = "cuda:0" if wp.is_cuda_available() else "cpu"
print(f"Using device: {device}")

# Create FCC argon system (4x4x4 unit cells = 256 atoms)
print("\n--- Creating FCC Argon System ---")
positions, cell = create_fcc_argon(num_unit_cells=4, a=5.26)
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,
    device=device,
    dtype=np.float64,
)

print("\n--- Setting Initial Temperature ---")
system.initialize_temperature(temperature=94.4, seed=42)

print("\n--- Initial Energy Calculation ---")
wp_energies = system.compute_forces()
pe = float(wp_energies.numpy().sum())
ke_arr = wp.zeros(1, dtype=wp.float64, device=device)
compute_kinetic_energy(
    velocities=system.wp_velocities,
    masses=system.wp_masses,
    kinetic_energy=ke_arr,
    device=device,
)
ke = float(ke_arr.numpy()[0])
print(f"  Kinetic Energy:   {ke:>12.4f} eV")
print(f"  Potential Energy: {pe:>12.4f} eV")
print(f"  Total Energy:     {ke + pe:>12.4f} eV")
print(f"  Neighbors:        {system.neighbor_manager.total_neighbors}())

```

--------------------------------

### System Setup and Initialization

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/79b90d37d46cffabc1dcc7d04a2871e6/01_langevin_integration.ipynb

Sets up the molecular dynamics system, including creating an FCC argon lattice, initializing the system with LJ parameters, and setting the initial temperature. It also performs an initial energy calculation.

```python
print("=" * 95)
print("LANGEVIN DYNAMICS WITH LENNARD-JONES POTENTIAL")
print("=" * 95)
print()

device = "cuda:0" if wp.is_cuda_available() else "cpu"
print(f"Using device: {device}")

# Create FCC argon system (4x4x4 unit cells = 256 atoms)
print("\n--- Creating FCC Argon System ---")
positions, cell = create_fcc_argon(num_unit_cells=4, a=5.26)
print(f"Created {len(positions)} atoms in {cell[0, 0]:.2f} R box")

print("\n--- Initializing MD System ---")
system = MDSystem(
    positions=positions,
    cell=cell,
    epsilon=EPSILON_AR,
    sigma=SIGMA_AR,
    cutoff=DEFAULT_CUTOFF,
    skin=DEFAULT_SKIN,
    device=device,
    dtype=np.float64,
)

print("\n--- Setting Initial Temperature ---")
system.initialize_temperature(temperature=94.4, seed=42)

print("\n--- Initial Energy Calculation ---")
wp_energies = system.compute_forces()
pe = float(wp_energies.numpy().sum())
ke_arr = wp.zeros(1, dtype=wp.float64, device=device)
compute_kinetic_energy(
    velocities=system.wp_velocities,
    masses=system.wp_masses,
    kinetic_energy=ke_arr,
    device=device,
)
ke = float(ke_arr.numpy()[0])
print(f"  Kinetic Energy:   {ke:>12.4f} eV")
print(f"  Potential Energy: {pe:>12.4f} eV")
print(f"  Total Energy:     {ke + pe:>12.4f} eV")
print(f"  Neighbors:        {system.neighbor_manager.total_neighbors()}")
```

--------------------------------

### Initialize Dipole System and Print Details

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/0fd4b7235600e9365bf599505d24d7b7/01_bspline_visualization.ipynb

Initializes a dipole system using the `create_dipole_system` function and prints the positions of the positive and negative charges. This is a setup step for subsequent visualization examples.

```python
positions_dip, charges_dip, cell_dip = create_dipole_system()
mesh_size = 32

print("Dipole system:")
print(f"  Positive charge at: {positions_dip[0].numpy()}")
print(f"  Negative charge at: {positions_dip[1].numpy()}")
```

--------------------------------

### Velocity Verlet (NVE) Simulation Loop

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/components/dynamics.html

Example of a Velocity Verlet integration loop for NVE simulations. Requires setup of positions, velocities, forces, masses, and time step. Assumes a user-defined `compute_forces` function.

```python
import warp as wp
from nvalchemiops.dynamics.integrators import (
    velocity_verlet_position_update,
    velocity_verlet_velocity_finalize
)

# Setup
positions = wp.array(pos_np, dtype=wp.vec3d, device="cuda:0")
velocities = wp.array(vel_np, dtype=wp.vec3d, device="cuda:0")
forces = wp.array(force_np, dtype=wp.vec3d, device="cuda:0")
masses = wp.array(mass_np, dtype=wp.float64, device="cuda:0")
dt = wp.array([0.001], dtype=wp.float64, device="cuda:0")

# MD loop
for step in range(num_steps):
    # Step 1: Update positions and half-step velocities
    velocity_verlet_position_update(positions, velocities, forces, masses, dt)

    # Step 2: Recalculate forces at new positions
    forces = compute_forces(positions)  # User-defined

    # Step 3: Finalize velocity update
    velocity_verlet_velocity_finalize(velocities, forces, masses, dt)

```

--------------------------------

### Benchmark Configuration File Example

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/benchmarks/dynamics.md.txt

Example configuration file for customizing dynamics benchmarks. Edit this file to adjust system sizes, batch sizes, integrators, and optimization parameters.

```yaml
# MD single-system
md_single:
  enabled: true
  system_sizes: [256, 512, 1024, 2048, 4096]
  integrators:
    velocity_verlet:
      steps: 10000
      dt: 0.001  # fs
      warmup_steps: 100
    langevin:
      steps: 10000
      dt: 0.001
      temperature: 300.0  # K
      friction: 0.01  # 1/fs

# MD batched
md_batch:
  enabled: true
  system_sizes: [256, 512, 1024]
  batch_sizes: [1, 2, 4, 8, 16, 32]
  integrators:
    velocity_verlet:
      steps: 10000
      dt: 0.001
      warmup_steps: 100

# Optimization single-system
opt_single:
  enabled: true
  system_sizes: [256, 512, 1024, 2048]
  optimizers:
    fire:
      max_steps: 1000
      force_tolerance: 0.01  # eV/Å

# Optimization batched
opt_batch:
  enabled: true
  system_sizes: [256, 512]
  batch_sizes: [1, 2, 4, 8, 16]
  optimizers:
    fire:
      max_steps: 1000
      force_tolerance: 0.01

# Potential parameters
potential:
  epsilon: 0.0104  # eV
  sigma: 3.40  # Å
  cutoff: 8.5  # Å
  skin: 1.0  # Å
  neighbor_rebuild_interval: 10

```

--------------------------------

### Langevin (NVT) Simulation Loop

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/components/dynamics.html

Example of a Langevin BAOAB integration loop for NVT simulations. Requires setup of positions, velocities, forces, masses, time step, temperature, and friction coefficient. Assumes a user-defined `compute_forces` function.

```python
import warp as wp
from nvalchemiops.dynamics.integrators import (
    langevin_baoab_half_step,
    langevin_baoab_finalize
)

# Setup (NVT parameters)
temperature = wp.array([1.0], dtype=wp.float64, device="cuda:0")  # kT in energy units
friction = wp.array([1.0], dtype=wp.float64, device="cuda:0")  # friction coefficient

# MD loop
for step in range(num_steps):
    # Step 1: BAOAB half-step (B-A-O-A)
    langevin_baoab_half_step(
        positions, velocities, forces, masses, dt,
        temperature, friction, random_seed=step
    )

    # Step 2: Recalculate forces
    forces = compute_forces(positions)

    # Step 3: Final B step
    langevin_baoab_finalize(velocities, forces, masses, dt)

```

--------------------------------

### Setup JAX System Parameters

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/90b798e785032850e81e9861b6a309cb/05_jax_neighbor_list.ipynb

Configures system parameters and creates random atomic positions and periodic cell using JAX for demonstration.

```python
print("=" * 70)
print("JAX NEIGHBOR LIST EXAMPLE")
print("=" * 70)

# System parameters
num_atoms = 200
box_size = 15.0
cutoff = 5.0

# Create random atomic positions using JAX random
key = jax.random.PRNGKey(42)
positions = jax.random.uniform(key, (num_atoms, 3), dtype=jnp.float32) * box_size

# Create a cubic periodic cell: (1, 3, 3) shape
cell = jnp.eye(3, dtype=jnp.float32)[None, ...] * box_size

# Enable periodic boundary conditions in all directions: (1, 3) shape
pbc = jnp.array([[True, True, True]])

print("\nSystem configuration:")
print(f"  Number of atoms: {num_atoms}")
print(f"  Box size: {box_size} ")
print(f"  Cutoff distance: {cutoff} ")
print(f"  Positions shape: {positions.shape}")
print(f"  Cell shape: {cell.shape}")
print(f"  PBC shape: {pbc.shape}")
```

--------------------------------

### Synchronize dependencies with uv and all extras (Nightly, with cloning)

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/about/install.html

Install project dependencies with all available extras using uv.

```bash
$ uv sync --all-extras

```

--------------------------------

### Example Completion Output

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/neighbors/04_neighbors_list_torch_compile_performance.html

The output indicating that the example has completed successfully.

```text
Example completed successfully!

```

--------------------------------

### Install Nightly ALCHEMI Toolkit-Ops using uv with cloning

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Recommended for local development and testing. Clones the repository and syncs dependencies.

```bash
$ git clone git@github.com/NVIDIA/nvalchemi-toolkit-ops.git
$ cd nvalchemi-toolkit-ops
$ uv sync
# include torch backend
$ uv sync --extra torch
# include jax backend
$ uv sync --extra jax
# include both backends
$ uv sync --all-extras
```

--------------------------------

### Configure Device and Initialize D3 Module

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/dispersion/01_dftd3_molecule.html

Sets up the computation device (GPU if available) and initializes the DFTD3 module with functional-specific parameters for PBE0. This includes scaling factors and damping parameters.

```python
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float32
print(f"Using device: {device}")

# Initialize with PBE0 parameters
d3_module = DFTD3(
    s6=1.0,  # C6 term coefficient
    s8=1.2177,  # C8 term coefficient (PBE0)
    a1=0.4145,  # BJ damping parameter (PBE0)
    a2=4.8593,  # BJ damping radius (PBE0)
    units="conventional",  # Coordinates in Angstrom, energy in eV
)

```

```text
Using device: cuda

```

--------------------------------

### Install ALCHEMI Toolkit-Ops

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/index.html

Install the ALCHEMI Toolkit-Ops package using pip.

```bash
$ pip install nvalchemi-toolkit-ops

```

--------------------------------

### Example Completion Message

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/2475ba2afbb663b76f52e9ac75f4d5d6/04_neighbors_list_torch_compile_performance.ipynb

Prints a success message indicating that the example has completed.

```python
print("\nExample completed successfully!")
```

--------------------------------

### Verify ALCHEMI Toolkit-Ops Installation

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/index.html

Verify the installation by checking the ALCHEMI Toolkit-Ops version.

```bash
$ python -c "import nvalchemiops; print(nvalchemiops.__version__)"

```

--------------------------------

### Full Installation (Arm, with cloning)

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Installs nvalchemi-toolkit-ops, Warp, PyTorch, and JAX for Arm architecture after cloning the repository. Includes extra index URL for PyTorch and uninstalls CUDA 12 JAX plugins.

```bash
$ git clone git@github.com:NVIDIA/nvalchemi-toolkit-ops.git
$ cd nvalchemi-toolkit-ops
$ uv sync --group dev
# Replace the default CUDA 12 wheels with CUDA 13 builds
$ uv pip install \
    https://github.com/NVIDIA/warp/releases/download/v1.12.1/warp_lang-1.12.1+cu13-py3-none-manylinux_2_34_aarch64.whl \
    torch==2.11.0+cu130 \
    'jax[cuda13]' \
    --force-reinstall \
    --extra-index-url https://download.pytorch.org/whl/cu130
# Remove the CUDA 12 JAX plugins to avoid a plugin conflict
$ uv pip uninstall jax-cuda12-pjrt jax-cuda12-plugin
```

--------------------------------

### Verify ALCHEMI Toolkit-Ops Installation

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/index.md.txt

Verify the installation by importing the package and printing its version.

```bash
$ python -c "import nvalchemiops; print(nvalchemiops.__version__)"
```

--------------------------------

### Batch Rebuild Detection Setup

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/8f452359dde252256b4f0a0ed9aeba7e/03_rebuild_neighborlist_detection.ipynb

Sets up a batch of systems with varying atom counts and creates initial batch neighbor lists. This demonstrates the initial state before testing batch rebuild detection.

```python
print("\n" + "=" * 70)
print("BATCH REBUILD DETECTION")
print("=" * 70)

# Set up a batch of systems with different atom counts
batch_sizes = [32, 48, 40]
batch_size = sum(batch_sizes)
num_systems_batch = len(batch_sizes)
batch_box_size = 5.0
batch_cutoff = 1.5
batch_skin = 0.4
batch_total_cutoff = batch_cutoff + batch_skin

# Create per-atom batch index and batch pointer
batch_idx = torch.repeat_interleave(
    torch.arange(num_systems_batch, dtype=torch.int32, device=device),
    torch.tensor(batch_sizes, dtype=torch.int32, device=device),
)
ptr_vals = [0] + [sum(batch_sizes[: i + 1]) for i in range(num_systems_batch)]
batch_ptr = torch.tensor(ptr_vals, dtype=torch.int32, device=device)

# Per-system cells and PBCs
batch_cell = (torch.eye(3, dtype=dtype, device=device) * batch_box_size).unsqueeze(0)
batch_cell = batch_cell.expand(num_systems_batch, -1, -1).contiguous()
batch_pbc = torch.zeros(num_systems_batch, 3, dtype=torch.bool, device=device)

# Random initial positions for each system
torch.manual_seed(1234)
batch_positions = torch.rand(batch_size, 3, dtype=dtype, device=device) * batch_box_size

print(f"\nBatch of {num_systems_batch} systems, {batch_sizes} atoms each")
print(f"  Cutoff: {batch_cutoff}, skin: {batch_skin}")

# Build initial batch neighbor lists
batch_max_neighbors = estimate_max_neighbors(batch_total_cutoff)
batch_nm, batch_nn = batch_naive_neighbor_list(
    positions=batch_positions,
    cutoff=batch_total_cutoff,
    batch_idx=batch_idx,
    batch_ptr=batch_ptr,
    max_neighbors=batch_max_neighbors,
)
print(f"\nInitial batch neighbor list built (max_neighbors={batch_max_neighbors})")
for s in range(num_systems_batch):
    sys_mask = batch_idx == s
    avg_nn = batch_nn[sys_mask].float().mean().item()
    print(f"  System {s}: avg {avg_nn:.1f} neighbors")
```

--------------------------------

### Full Installation (Arm, without cloning)

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/userguide/about/install.md.txt

Installs nvalchemi-toolkit-ops, Warp, PyTorch, and JAX for Arm architecture without cloning the repository, including the necessary extra index URL for PyTorch.

```bash
$ uv venv --seed --python 3.12
$ uv pip install nvalchemi-toolkit-ops \
    https://github.com/NVIDIA/warp/releases/download/v1.12.1/warp_lang-1.12.1+cu13-py3-none-manylinux_2_34_aarch64.whl \
    torch==2.11.0+cu130 \
    'jax[cuda13]' \
    --extra-index-url https://download.pytorch.org/whl/cu130
```

--------------------------------

### PME Example Completion Output

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/electrostatics/03_pme_example.html

This is the output message confirming the completion of the PME example.

```text
PME example complete!

```

--------------------------------

### Set up and initialize the molecular dynamics system

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_sources/examples/dynamics/01_langevin_integration.rst.txt

Configures the simulation device, creates an FCC lattice of argon atoms, initializes the MD system with Lennard-Jones parameters, and sets the initial temperature.

```Python
print("=" * 95)
print("LANGEVIN DYNAMICS WITH LENNARD-JONES POTENTIAL")
print("=" * 95)
print()

device = "cuda:0" if wp.is_cuda_available() else "cpu"
print(f"Using device: {device}")

# Create FCC argon system (4x4x4 unit cells = 256 atoms)
print("\n--- Creating FCC Argon System ---")
positions, cell = create_fcc_argon(num_unit_cells=4, a=5.26)
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,
    device=device,
    dtype=np.float64,
)

print("\n--- Setting Initial Temperature ---")
system.initialize_temperature(temperature=94.4, seed=42)

print("\n--- Initial Energy Calculation ---")
wp_energies = system.compute_forces()
pe = float(wp_energies.numpy().sum())
ke_arr = wp.zeros(1, dtype=wp.float64, device=device)
compute_kinetic_energy(
    velocities=system.wp_velocities,
    masses=system.wp_masses,
    kinetic_energy=ke_arr,
    device=device,
)
ke = float(ke_arr.numpy()[0])
print(f"  Kinetic Energy:   {ke:>12.4f} eV")
print(f"  Potential Energy: {pe:>12.4f} eV")
print(f"  Total Energy:     {ke + pe:>12.4f} eV")
print(f"  Neighbors:        {system.neighbor_manager.total_neighbors()}")
```

--------------------------------

### JAX PME Example Summary Output

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/electrostatics/04_jax_pme_example.html

This snippet shows the summary output of the JAX PME example, highlighting key takeaways and successful completion. It is useful for understanding the main features demonstrated in the example.

```python
print("\n" + "=" * 70)
print("SUMMARY")
print("=" * 70)
print("\nKey takeaways:")
print("  - Use estimate_pme_parameters() for automatic parameter selection")
print("  - Both COO and dense neighbor formats produce identical results")
print("  - Real and reciprocal components can be computed separately")
print("  - Charge gradients are available for ML potential training")
print("  - Use jax.jit to fuse neighbor list + PME into one compiled function")
print("\nJAX PME example completed successfully!")

```

```text
======================================================================
SUMMARY
======================================================================

Key takeaways:
  - Use estimate_pme_parameters() for automatic parameter selection
  - Both COO and dense neighbor formats produce identical results
  - Real and reciprocal components can be computed separately
  - Charge gradients are available for ML potential training
  - Use jax.jit to fuse neighbor list + PME into one compiled function

JAX PME example completed successfully!

```

--------------------------------

### Compute Cell Atom Start Indices

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_modules/nvalchemiops/neighbors/batch_cell_list.html

Performs an exclusive scan on atom counts per cell to determine the starting index for each cell's atom list. This converts counts into starting positions.

```python
wp.utils.array_scan(atoms_per_cell_count, cell_atom_start_indices, inclusive=False)
```

--------------------------------

### Create initial system configuration

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/_downloads/2475ba2afbb663b76f52e9ac75f4d5d6/04_neighbors_list_torch_compile_performance.ipynb

Generates the initial configuration of atoms in the simulation box using create_bulk_structure.

```python
structure = create_bulk_structure(
    num_atoms=num_atoms,
    box_size=box_size,
    device=device,
    dtype=dtype,
)

print(f"\nInitial structure created with {structure.n_atoms} atoms.")
```

--------------------------------

### Install ALCHEMI Toolkit-Ops with JAX Backend

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/index.html

Install the ALCHEMI Toolkit-Ops package with JAX support.

```bash
$ pip install 'nvalchemi-toolkit-ops[jax]'

```

--------------------------------

### Setup and Parameter Loading with JAX

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/examples/dispersion/03_jax_dftd3_molecule.html

Imports necessary JAX and NVALCHEMI-OPS modules, loads DFT-D3 parameters, and converts them to JAX arrays. Includes error handling for missing JAX/Warp backends and downloads parameters if not cached.

```python
from __future__ import annotations

import os
import sys
from pathlib import Path

# Prevent JAX from pre-allocating 75% of GPU memory at init, which fails
# when PyTorch has already reserved memory in the same process (e.g. during
# sphinx-gallery builds that run all examples sequentially).
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")

try:
    import jax
    import jax.numpy as jnp
except ImportError:
    print(
        "This example requires JAX. Install with: pip install 'nvalchemi-toolkit-ops[jax]'"
    )
    sys.exit(0)

import numpy as np
import torch

try:
    from nvalchemiops.jax.interactions.dispersion import D3Parameters, dftd3
    from nvalchemiops.jax.neighbors import neighbor_list
    from nvalchemiops.jax.neighbors.naive import naive_neighbor_list
except Exception as exc:
    print(
        f"JAX/Warp backend unavailable ({exc}). This example requires a CUDA-backed runtime."
    )
    sys.exit(0)

# Unit conversion constants (CODATA 2022)
BOHR_TO_ANGSTROM = 0.529177210544
HARTREE_TO_EV = 27.211386245981
ANGSTROM_TO_BOHR = 1.0 / BOHR_TO_ANGSTROM

# Check for cached parameters, download if needed
param_file = (
    Path(os.path.expanduser("~")) / ".cache" / "nvalchemiops" / "dftd3_parameters.pt"
)
if not param_file.exists():
    print("Downloading DFT-D3 parameters...")
    try:
        _import_dir = str(Path(__file__).parent)
    except NameError:
        _import_dir = str(Path.cwd())
    sys.path.insert(0, _import_dir)
    from utils import extract_dftd3_parameters, save_dftd3_parameters

    params_torch = extract_dftd3_parameters()
    save_dftd3_parameters(params_torch)
else:
    params_torch = torch.load(param_file, weights_only=True)
    print("Loaded cached DFT-D3 parameters")

# Convert PyTorch tensors to JAX arrays
d3_params = D3Parameters(
    rcov=jnp.array(params_torch["rcov"].numpy(), dtype=jnp.float32),
    r4r2=jnp.array(params_torch["r4r2"].numpy(), dtype=jnp.float32),
    c6ab=jnp.array(params_torch["c6ab"].numpy(), dtype=jnp.float32),
    cn_ref=jnp.array(params_torch["cn_ref"].numpy(), dtype=jnp.float32),
)

print(f"Loaded D3 parameters for elements 1-{d3_params.max_z}")

```

```text
Loaded cached DFT-D3 parameters
Loaded D3 parameters for elements 1-94

```

--------------------------------

### Install ALCHEMI Toolkit-Ops with PyTorch Backend

Source: https://nvidia.github.io/nvalchemi-toolkit-ops/userguide/index.html

Install the ALCHEMI Toolkit-Ops package with PyTorch support.

```bash
$ pip install 'nvalchemi-toolkit-ops[torch]'

```
Tokens:
Raw