`nvalchemiops.dynamics`: Molecular Dynamics#

Dynamics Module#

GPU-accelerated Warp kernels for molecular dynamics integrators and geometry optimization algorithms.

All kernels use direct array inputs (no structs) and provide both: - Mutating (in-place) versions for efficiency - Non-mutating versions for gradient tracking

Submodules#

integrators: MD integrators: velocity Verlet (NVE), Langevin (NVT), Nosé-Hoover (NVT)
optimizers: Geometry optimizers: FIRE, FIRE2
utils: Utility functions: temperature computation, velocity initialization

Example

>>> import warp as wp
>>> from nvalchemiops.dynamics.integrators import velocity_verlet_position_update
>>>
>>> # Mutating (in-place) API
>>> velocity_verlet_position_update(
...     positions, velocities, forces, masses, dt
... )
>>>
>>> # Non-mutating API (for gradient tracking)
>>> from nvalchemiops.dynamics.integrators import velocity_verlet_position_update_out
>>> positions_out = wp.zeros_like(positions)
>>> velocities_out = wp.zeros_like(velocities)
>>> positions_out, velocities_out = velocity_verlet_position_update_out(
...     positions, velocities, forces, masses, dt,
...     positions_out, velocities_out,
... )

Warp-Level Interface#

Tip

This is the low-level Warp interface that operates on warp.array objects. For PyTorch tensor support with the FIRE2 optimizer, see nvalchemiops.torch: Dynamics Optimizers.

High-Level Interface#

This module provides GPU-accelerated integrators and thermostats for molecular dynamics simulations. All functions support automatic differentiation through PyTorch’s autograd system. Most functions offer three execution modes: single system, batched with batch_idx, and batched with atom_ptr. NPT/NPH integrators support batch_idx only – see individual function docs for details.

Tip

Check out the Molecular Dynamics Integrators page for usage examples and a conceptual overview of the available integrators and thermostats.

Integrators#

Velocity Verlet#

Time-reversible, symplectic integrator for NVE (microcanonical) ensemble.

nvalchemiops.dynamics.integrators.velocity_verlet.velocity_verlet_position_update(positions, velocities, forces, masses, dt, batch_idx=None, atom_ptr=None, device=None)[source]#

Perform velocity Verlet position update step (in-place).

Updates positions to \(\mathbf{r}(t + \Delta t)\) and velocities to the half step \(\mathbf{v}(t + \Delta t/2)\). After calling this function, recalculate forces at the new positions, then call velocity_verlet_velocity_finalize().

The kernel evaluates, with \(\mathbf{a}(t) = \mathbf{F}(t)/m\):

\[ \begin{align}\begin{aligned}\mathbf{r}(t + \Delta t) = \mathbf{r}(t) + \mathbf{v}(t)\,\Delta t + \tfrac{1}{2}\mathbf{a}(t)\,\Delta t^2\\\mathbf{v}(t + \tfrac{\Delta t}{2}) = \mathbf{v}(t) + \tfrac{1}{2}\mathbf{a}(t)\,\Delta t\end{aligned}\end{align} \]

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). MODIFIED in-place.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single system, (B,) for batched.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Return type:

None

See also

velocity_verlet_velocity_finalize: Complete the velocity update

nvalchemiops.dynamics.integrators.velocity_verlet.velocity_verlet_velocity_finalize(velocities, forces_new, masses, dt, batch_idx=None, atom_ptr=None, device=None)[source]#

Finalize velocity Verlet velocity update (in-place).

Completes the velocity update using forces evaluated at the new positions, with \(\mathbf{a}(t + \Delta t) = \mathbf{F}_\text{new}/m\):

\[\mathbf{v}(t + \Delta t) = \mathbf{v}(t + \tfrac{\Delta t}{2}) + \tfrac{1}{2}\mathbf{a}(t + \Delta t)\,\Delta t\]

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Half-step velocities. Shape (N,). MODIFIED in-place to full-step.
forces_new (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces evaluated at new positions \(\mathbf{r}(t + \Delta t)\). Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single, (B,) for batched.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Return type:

None

See also

velocity_verlet_position_update: First step of velocity Verlet

nvalchemiops.dynamics.integrators.velocity_verlet.velocity_verlet_position_update_out(positions, velocities, forces, masses, dt, positions_out, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Perform velocity Verlet position update step (non-mutating).

Writes new positions \(\mathbf{r}(t + \Delta t)\) and half-step velocities \(\mathbf{v}(t + \Delta t/2)\) to output arrays. Input arrays are NOT modified.

The kernel evaluates, with \(\mathbf{a}(t) = \mathbf{F}(t)/m\):

\[ \begin{align}\begin{aligned}\mathbf{r}(t + \Delta t) = \mathbf{r}(t) + \mathbf{v}(t)\,\Delta t + \tfrac{1}{2}\mathbf{a}(t)\,\Delta t^2\\\mathbf{v}(t + \tfrac{\Delta t}{2}) = \mathbf{v}(t) + \tfrac{1}{2}\mathbf{a}(t)\,\Delta t\end{aligned}\end{align} \]

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions at time t. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities at time t. Shape (N,).
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms at time \(t\). Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single, (B,) for batched.
positions_out (wp.array) – Pre-allocated output array for new positions. Must match positions in shape, dtype, and device.
velocities_out (wp.array) – Pre-allocated output array for half-step velocities. Must match velocities in shape, dtype, and device.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Returns:

(positions_out, velocities_out) - New positions and half-step velocities.

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.integrators.velocity_verlet.velocity_verlet_velocity_finalize_out(velocities, forces_new, masses, dt, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Finalize velocity Verlet velocity update (non-mutating).

Writes full-step velocities \(\mathbf{v}(t + \Delta t)\) to output array. Input arrays are NOT modified.

The kernel evaluates, with \(\mathbf{a}(t + \Delta t) = \mathbf{F}_\text{new}/m\):

\[\mathbf{v}(t + \Delta t) = \mathbf{v}(t + \tfrac{\Delta t}{2}) + \tfrac{1}{2}\mathbf{a}(t + \Delta t)\,\Delta t\]

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Half-step velocities \(\mathbf{v}(t + \Delta t/2)\). Shape (N,).
forces_new (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces at new positions \(\mathbf{r}(t + \Delta t)\). Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single, (B,) for batched.
velocities_out (wp.array) – Pre-allocated output array for full-step velocities. Must match velocities in shape, dtype, and device.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Returns:

Full-step velocities \(\mathbf{v}(t + \Delta t)\).

Return type:

wp.array

Langevin Dynamics#

BAOAB splitting scheme for NVT (canonical) ensemble with stochastic thermostat.

nvalchemiops.dynamics.integrators.langevin.langevin_baoab_half_step(positions, velocities, forces, masses, dt, temperature, friction, random_seed, batch_idx=None, atom_ptr=None, device=None)[source]#

Perform BAOAB Langevin half-step (B-A-O-A sequence) in-place.

This function performs the first four operations of the BAOAB splitting: B (velocity kick), A (drift), O (Ornstein-Uhlenbeck thermostat), A (drift). The kernel evaluates the following updates in order:

\[ \begin{align}\begin{aligned}B:\quad \mathbf{v} \leftarrow \mathbf{v} + \tfrac{\Delta t}{2m}\mathbf{F}\\A:\quad \mathbf{r} \leftarrow \mathbf{r} + \tfrac{\Delta t}{2}\mathbf{v}\\O:\quad \mathbf{v} \leftarrow c_1\,\mathbf{v} + c_2\,\boldsymbol{\xi}\\A:\quad \mathbf{r} \leftarrow \mathbf{r} + \tfrac{\Delta t}{2}\mathbf{v}\end{aligned}\end{align} \]

with damping \(c_1 = e^{-\gamma \Delta t}\), noise amplitude \(c_2 = \sqrt{k_B T\,(1 - c_1^2)/m}\), and \(\boldsymbol{\xi} \sim \mathcal{N}(0, 1)\).

After calling this function, recalculate forces at the new positions, then call langevin_baoab_finalize() to complete the step.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). MODIFIED in-place.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single, (B,) for batched.
temperature (wp.array(dtype=wp.float32 or wp.float64)) – Temperature \(k_B T\). Shape (1,) for single, (B,) for batched.
friction (wp.array(dtype=wp.float32 or wp.float64)) – Friction coefficient \(\gamma\). Shape (1,) for single, (B,) for batched.
random_seed (int) – Random seed for stochastic forces.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Return type:

None

Example

Single system NVT simulation:

import warp as wp
import numpy as np

# Setup
positions = wp.array(np.random.randn(100, 3), dtype=wp.vec3d, device="cuda:0")
velocities = wp.array(np.random.randn(100, 3), dtype=wp.vec3d, device="cuda:0")
forces = wp.array(np.random.randn(100, 3), dtype=wp.vec3d, device="cuda:0")
masses = wp.array(np.ones(100), dtype=wp.float64, device="cuda:0")

dt = wp.array([0.001], dtype=wp.float64, device="cuda:0")
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")

# BAOAB half-step
langevin_baoab_half_step(
    positions, velocities, forces, masses, dt, temperature, friction,
    random_seed=42
)

Complete BAOAB step:

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)

Batched mode:

# With batch_idx (3 systems)
batch_idx = wp.array([0]*30 + [1]*40 + [2]*30, dtype=wp.int32, device="cuda:0")
dt = wp.array([0.001, 0.002, 0.0015], dtype=wp.float64, device="cuda:0")
temperature = wp.array([1.0, 1.5, 1.2], dtype=wp.float64, device="cuda:0")
friction = wp.array([1.0, 1.0, 1.0], dtype=wp.float64, device="cuda:0")

langevin_baoab_half_step(
    positions, velocities, forces, masses, dt, temperature, friction,
    random_seed=42, batch_idx=batch_idx
)

See also

langevin_baoab_finalize: Complete the BAOAB step

nvalchemiops.dynamics.integrators.langevin.langevin_baoab_finalize(velocities, forces_new, masses, dt, batch_idx=None, atom_ptr=None, device=None)[source]#

Finalize BAOAB Langevin step (final B step) in-place.

Completes the BAOAB sequence with the final velocity half-kick using forces \(\mathbf{F}_\text{new}\) calculated at the new positions:

\[B:\quad \mathbf{v} \leftarrow \mathbf{v} + \tfrac{\Delta t}{2m}\mathbf{F}_\text{new}\]

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
forces_new (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces at new positions. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Timestep(s). Shape (1,) for single, (B,) for batched.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Return type:

None

nvalchemiops.dynamics.integrators.langevin.langevin_baoab_half_step_out(positions, velocities, forces, masses, dt, temperature, friction, random_seed, positions_out, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Perform BAOAB Langevin half-step (B-A-O-A sequence) non-mutating.

Writes new positions and velocities to output arrays. Input arrays are NOT modified. The kernel evaluates the following updates in order:

\[ \begin{align}\begin{aligned}B:\quad \mathbf{v} \leftarrow \mathbf{v} + \tfrac{\Delta t}{2m}\mathbf{F}\\A:\quad \mathbf{r} \leftarrow \mathbf{r} + \tfrac{\Delta t}{2}\mathbf{v}\\O:\quad \mathbf{v} \leftarrow c_1\,\mathbf{v} + c_2\,\boldsymbol{\xi}\\A:\quad \mathbf{r} \leftarrow \mathbf{r} + \tfrac{\Delta t}{2}\mathbf{v}\end{aligned}\end{align} \]

with damping \(c_1 = e^{-\gamma \Delta t}\), noise amplitude \(c_2 = \sqrt{k_B T\,(1 - c_1^2)/m}\), and \(\boldsymbol{\xi} \sim \mathcal{N}(0, 1)\).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions at time \(t\). Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities at time \(t\). Shape (N,).
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms at time \(t\). Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array) – Timestep(s). Shape (1,) for single, (B,) for batched.
temperature (wp.array) – Temperature \(k_B T\). Shape (1,) for single, (B,) for batched.
friction (wp.array) – Friction coefficient \(\gamma\). Shape (1,) for single, (B,) for batched.
random_seed (int) – Random seed for stochastic forces.
positions_out (wp.array) – Pre-allocated output array for new positions. Must match positions in shape, dtype, and device.
velocities_out (wp.array) – Pre-allocated output array for new velocities. Must match velocities in shape, dtype, and device.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Returns:

(positions_out, velocities_out) - New positions and velocities.

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.integrators.langevin.langevin_baoab_finalize_out(velocities, forces_new, masses, dt, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Finalize BAOAB Langevin step (final B step) non-mutating.

Writes full-step velocities to output array. Input arrays are NOT modified. The kernel applies the final velocity half-kick using forces \(\mathbf{F}_\text{new}\) calculated at the new positions:

\[B:\quad \mathbf{v} \leftarrow \mathbf{v} + \tfrac{\Delta t}{2m}\mathbf{F}_\text{new}\]

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Velocities after half-step. Shape (N,).
forces_new (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces at new positions. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array) – Timestep(s). Shape (1,) for single, (B,) for batched.
velocities_out (wp.array) – Pre-allocated output array for final velocities. Must match velocities in shape, dtype, and device.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Returns:

Full-step velocities.

Return type:

wp.array

Nosé-Hoover Chain#

Deterministic thermostat for NVT ensemble using Nosé-Hoover chains with Yoshida-Suzuki integration.

nvalchemiops.dynamics.integrators.nose_hoover.nhc_velocity_half_step(velocities, forces, masses, dt, batch_idx=None, atom_ptr=None, device=None)[source]#

Half-step velocity update (in-place).

\[\mathbf{v}_i \leftarrow \mathbf{v}_i + \tfrac{1}{2}\,\frac{\mathbf{F}_i}{m_i}\,\Delta t\]

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (B,).
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Return type:

None

nvalchemiops.dynamics.integrators.nose_hoover.nhc_position_update(positions, velocities, dt, batch_idx=None, atom_ptr=None, device=None)[source]#

Full-step position update (in-place).

\[\mathbf{r}_i \leftarrow \mathbf{r}_i + \mathbf{v}_i\,\Delta t\]

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). MODIFIED in-place.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (B,).
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Return type:

None

nvalchemiops.dynamics.integrators.nose_hoover.nhc_thermostat_chain_update(velocities, masses, eta, eta_dot, eta_mass, target_temp, dt, ndof, ke2, total_scale, step_scale, dt_chain, nloops=1, batch_idx=None, num_systems=1, device=None, compute_ke=True)[source]#

Propagate Nosé-Hoover chain and scale velocities (in-place).

Propagates the chain over the Yoshida-Suzuki sub-steps, accumulating the per-step velocity scale factor \(\exp(-\tfrac{1}{2}\,\Delta t_c\,\dot{\eta}_0)\) into total_scale, then rescales every particle velocity by it:

\[\mathbf{v}_i \leftarrow s\,\mathbf{v}_i, \qquad s = \prod_j \exp\!\left(-\tfrac{1}{2}\,\Delta t_{c,j}\,\dot{\eta}_0\right)\]

Uses Yoshida-Suzuki factorization for time-reversible integration. All computations are performed on GPU using Warp kernels.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
eta (wp.array(dtype=wp.float64)) – Thermostat chain positions. Non-batched: Shape (chain_length,). Batched: Shape (num_systems, chain_length) as wp.array2d. MODIFIED in-place.
eta_dot (wp.array(dtype=wp.float64)) – Thermostat chain velocities. Same shape as eta. MODIFIED in-place.
eta_mass (wp.array(dtype=wp.float64)) – Thermostat chain masses. Same shape as eta.
target_temp (wp.array(dtype=wp.float64)) – Target temperature \(k_B T\). Shape (1,) or (num_systems,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (num_systems,).
ndof (wp.array(dtype=wp.float64)) – Degrees of freedom. Shape (1,) or (num_systems,).
ke2 (wp.array) – \(2\,\mathrm{KE} = \sum_i m_i \lVert \mathbf{v}_i \rVert^2\) per system. When compute_ke is True (default) this is a scratch array zeroed and filled internally from velocities/masses before use. When compute_ke is False it is an INPUT supplying a caller-precomputed \(2\,\mathrm{KE}\); it is not zeroed. Either way the chain forward sweep rescales it in place. Shape (1,) for single system, (num_systems,) for batched.
total_scale (wp.array) – Scratch array for accumulated velocity scale factor. Must be initialized to ones by caller (wp.ones). Shape (1,) for single system, (num_systems,) for batched.
step_scale (wp.array) – Scratch array for per-step velocity scale factor. Shape (1,) for single system, (num_systems,) for batched.
dt_chain (wp.array) – Scratch array for weighted time steps. Shape (1,) for single system, (num_systems,) for batched.
nloops (int, optional) – Number of Yoshida-Suzuki integration sub-steps. Default: 1. Use nloops=3 or 5 for higher accuracy.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched mode.
num_systems (int, optional) – Number of systems for batched mode. Default: 1.
device (str, optional) – Warp device. If None, inferred from velocities.
compute_ke (bool, optional) – If True (default), recompute \(2\,\mathrm{KE}\) from velocities/masses internally (byte-identical to legacy behavior). If False, use the caller-supplied value already in ke2 instead of recomputing it. See nhc_compute_2ke for the matching reduction helper.

Return type:

None

nvalchemiops.dynamics.integrators.nose_hoover.nhc_velocity_half_step_out(velocities, forces, masses, dt, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Half-step velocity update (non-mutating).

Writes \(\mathbf{v}_i + \tfrac{1}{2}\,\frac{\mathbf{F}_i}{m_i}\,\Delta t\) into velocities_out.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on atoms. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (B,).
velocities_out (wp.array) – Output velocities. Must be pre-allocated with same shape/dtype/device as velocities.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Returns:

Updated velocities.

Return type:

wp.array

nvalchemiops.dynamics.integrators.nose_hoover.nhc_position_update_out(positions, velocities, dt, positions_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Full-step position update (non-mutating).

Writes \(\mathbf{r}_i + \mathbf{v}_i\,\Delta t\) into positions_out.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (B,).
positions_out (wp.array) – Output positions. Must be pre-allocated with same shape/dtype/device as positions.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Updated positions.

Return type:

wp.array

nvalchemiops.dynamics.integrators.nose_hoover.nhc_thermostat_chain_update_out(velocities, masses, eta, eta_dot, eta_mass, target_temp, dt, ndof, ke2, total_scale, step_scale, dt_chain, velocities_out, eta_out, eta_dot_out, nloops=1, batch_idx=None, num_systems=1, device=None, compute_ke=True)[source]#

Propagate Nosé-Hoover chain and scale velocities (non-mutating).

Out-of-place counterpart of nhc_thermostat_chain_update(): copies velocities/eta/eta_dot into the *_out buffers and applies the same chain propagation and velocity rescaling \(\mathbf{v}_i \leftarrow s\,\mathbf{v}_i\) to the copies.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
eta (wp.array(dtype=wp.float64)) – Thermostat chain positions. Shape (M,) or (B, M).
eta_dot (wp.array(dtype=wp.float64)) – Thermostat chain velocities. Shape (M,) or (B, M).
eta_mass (wp.array(dtype=wp.float64)) – Thermostat chain masses. Shape (M,) or (B, M).
target_temp (wp.array(dtype=wp.float64)) – Target temperature \(k_B T\). Shape (1,) or (B,).
dt (wp.array(dtype=wp.float32 or wp.float64)) – Time step. Shape (1,) or (B,).
ndof (wp.array(dtype=wp.float64)) – Degrees of freedom. Shape (1,) or (B,).
ke2 (wp.array) – \(2\,\mathrm{KE} = \sum_i m_i \lVert \mathbf{v}_i \rVert^2\) per system. Zeroed and filled internally when compute_ke is True (default); when compute_ke is False it is an INPUT supplying a caller-precomputed \(2\,\mathrm{KE}\) that is not zeroed. See nhc_thermostat_chain_update. Shape (1,) for single system, (num_systems,) for batched.
total_scale (wp.array) – Scratch array for accumulated velocity scale factor. Must be initialized to ones by caller (wp.ones). Shape (1,) for single system, (num_systems,) for batched.
step_scale (wp.array) – Scratch array for per-step velocity scale factor. Shape (1,) for single system, (num_systems,) for batched.
dt_chain (wp.array) – Scratch array for weighted time steps. Shape (1,) for single system, (num_systems,) for batched.
velocities_out (wp.array) – Output velocities. Must be pre-allocated with same shape/dtype/device as velocities.
eta_out (wp.array) – Output eta. Must be pre-allocated with same shape/dtype/device as eta.
eta_dot_out (wp.array) – Output eta_dot. Must be pre-allocated with same shape/dtype/device as eta_dot.
nloops (int, optional) – Number of Yoshida-Suzuki integration sub-steps. Default: 1.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched mode.
num_systems (int, optional) – Number of systems for batched mode. Default: 1.
device (str, optional) – Warp device. If None, inferred from velocities.
compute_ke (bool, optional) – If True (default), recompute \(2\,\mathrm{KE}\) internally. If False, use the caller-supplied value in ke2. See nhc_thermostat_chain_update for details.

Returns:

(velocities_out, eta_out, eta_dot_out)

Return type:

tuple[wp.array, wp.array, wp.array]

nvalchemiops.dynamics.integrators.nose_hoover.nhc_compute_masses(ndof, target_temp, tau, chain_length, masses, num_systems=1, device=None, dtype=wp.float64)[source]#

Compute Nosé-Hoover chain masses using GPU kernel.

Computes \(Q_k\) values for the Nosé-Hoover chain:

\[\begin{split}Q_0 &= N_\text{dof} \cdot k_B T \cdot \tau^2 \\ Q_k &= k_B T \cdot \tau^2 \quad \text{for } k > 0\end{split}\]

Parameters:

ndof (wp.array(dtype=wp.int32)) – Number of degrees of freedom per system. Shape (1,) for single system, (num_systems,) for batched.
target_temp (wp.array) – Target temperature \(k_B T\) per system. Shape (1,) for single system, (num_systems,) for batched.
tau (wp.array) – Time constant \(\tau\) per system. Shape (1,) for single system, (num_systems,) for batched.
chain_length (int) – Number of thermostats in the chain.
masses (wp.array) – Chain masses output. Caller must pre-allocate. Shape (chain_length,) for single system, (num_systems, chain_length) for batched.
num_systems (int, optional) – Number of systems for batched mode. Default: 1.
device (str, optional) – Warp device. If None, inferred from masses.
dtype (dtype, optional) – Data type for the masses. Default: wp.float64.

Returns:

Chain masses. Shape (chain_length,) for single system, (num_systems, chain_length) for batched.

Return type:

wp.array

nvalchemiops.dynamics.integrators.nose_hoover.nhc_compute_chain_energy(eta, eta_dot, eta_mass, target_temp, ndof, ke_chain, pe_chain, batch_idx=None, num_systems=1, device=None)[source]#

Compute Nosé-Hoover chain kinetic and potential energy.

For conservation checks, the extended system Hamiltonian is:

\[H_\text{ext} = \text{KE}_\text{particles} + \text{PE} + \text{KE}_\text{chain} + \text{PE}_\text{chain}\]

where:

\[ \begin{align}\begin{aligned}\text{KE}_\text{chain} = \sum_k \tfrac{1}{2} Q_k \dot{\eta}_k^2\\\text{PE}_\text{chain} = N_\text{dof} \cdot kT \cdot \eta_0 + kT \sum_{k>0} \eta_k\end{aligned}\end{align} \]

Parameters:

eta (wp.array(dtype=wp.float64)) – Thermostat chain positions. Shape (M,) or (B, M).
eta_dot (wp.array(dtype=wp.float64)) – Thermostat chain velocities. Shape (M,) or (B, M).
eta_mass (wp.array(dtype=wp.float64)) – Thermostat chain masses. Shape (M,) or (B, M).
target_temp (wp.array(dtype=wp.float64)) – Target temperature \(k_B T\). Shape (1,) or (B,).
ndof (wp.array(dtype=wp.float64)) – Degrees of freedom. Shape (1,) or (B,).
ke_chain (wp.array) – Output kinetic energy of the chain. Shape (1,) for single system, (num_systems,) for batched.
pe_chain (wp.array) – Output potential energy of the chain. Shape (1,) for single system, (num_systems,) for batched.
batch_idx (wp.array(dtype=wp.int32), optional) – Not used directly, but included for API consistency.
num_systems (int, optional) – Number of systems for batched mode. Default: 1.
device (str, optional) – Warp device. If None, inferred from eta.

Returns:

(ke_chain, pe_chain) each with shape (1,) or (B,).

Return type:

tuple[wp.array, wp.array]

Velocity Rescaling#

Simple velocity rescaling thermostat for quick equilibration (non-canonical).

nvalchemiops.dynamics.integrators.velocity_rescaling.velocity_rescale(velocities, scale_factor, batch_idx=None, atom_ptr=None, device=None)[source]#

Rescale velocities to achieve target temperature (in-place).

\[\mathbf{v}_i \leftarrow s\,\mathbf{v}_i\]

where \(s\) is the scale factor, typically \(s = \sqrt{T_{\text{target}} / T_{\text{current}}}\).

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
scale_factor (wp.array) – Scaling factor(s). Shape (1,) for single system, (B,) for batched. Typically \(\sqrt{T_{\text{target}} / T_{\text{current}}}\).
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device. If None, inferred from velocities.

Return type:

None

Example

>>> # Compute current temperature
>>> ke = wp.empty(1, dtype=wp.float64, device=device)
>>> compute_kinetic_energy(velocities, masses, ke)
>>> T_current = wp.empty(1, dtype=wp.float64, device=device)
>>> compute_temperature(ke, T_current, num_atoms=100)
>>>
>>> # Compute rescaling factor
>>> factor = _compute_rescale_factor(T_current, T_target)
>>> scale = wp.array([factor], dtype=wp.float32, device=device)
>>>
>>> # Apply rescaling
>>> velocity_rescale(velocities, scale)

nvalchemiops.dynamics.integrators.velocity_rescaling.velocity_rescale_out(velocities, scale_factor, velocities_out, batch_idx=None, atom_ptr=None, device=None)[source]#

Rescale velocities to achieve target temperature (non-mutating).

Writes \(s\,\mathbf{v}_i\) into velocities_out, where \(s\) is the scale factor.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
scale_factor (wp.array) – Scaling factor(s). Shape (1,) for single system, (B,) for batched.
velocities_out (wp.array) – Pre-allocated output array. Must match velocities in shape, dtype, and device.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device (str, optional) – Warp device.

Returns:

Rescaled velocities (same object as velocities_out).

Return type:

wp.array

NPT/NPH (Isothermal-Isobaric / Isenthalpic-Isobaric)#

Extended ensemble integrators for constant pressure simulations using Nosé-Hoover chains and Martyna-Tobias-Klein barostat. Supports both isotropic and anisotropic pressure control.

Pressure Calculations

nvalchemiops.dynamics.integrators.npt.compute_pressure_tensor(velocities, masses, virial_tensors, cells, kinetic_tensors, pressure_tensors, volumes, batch_idx=None, device=None, compute_kinetic=True)[source]#

Compute full pressure tensor \(\mathbf{P} = (\mathbf{K} + \mathbf{W}) / V\), with kinetic term \(\mathbf{K}\), virial term \(\mathbf{W}\), and cell volume \(V\).

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle velocities. Shape (N,).
masses (wp.array) – Particle masses. Shape (N,).
virial_tensors (wp.array(dtype=vec9f or vec9d)) – Virial tensor from forces (physics sign). Shape (B,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
kinetic_tensors (wp.array(dtype=scalar, ndim=2)) – Kinetic tensor \(\mathbf{K}_s = \sum_{i \in s} m_i (\mathbf{v}_i \otimes \mathbf{v}_i)\), vec9 layout. Shape (B, 9). When compute_kinetic is True (default) this is a scratch array zeroed and filled internally from velocities/masses. When compute_kinetic is False it is an INPUT supplying a caller-precomputed kinetic tensor; it is not zeroed or recomputed.
pressure_tensors (wp.array(dtype=vec9f or vec9d)) – Output pressure tensor. Shape (B,).
volumes (wp.array(dtype=scalar)) – Pre-computed cell volumes. Shape (B,). Caller must pre-compute via compute_cell_volume.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. If None, assumes single system.
device (str, optional) – Warp device.
compute_kinetic (bool, optional) – If True (default), recompute the kinetic tensor from velocities/masses internally (byte-identical to legacy behavior). If False, use the caller-supplied value already in kinetic_tensors instead of recomputing it. See compute_kinetic_tensor for the matching reduction helper.

Returns:

Pressure tensor. Shape (B,).

Return type:

wp.array(dtype=vec9f or vec9d)

nvalchemiops.dynamics.integrators.npt.compute_scalar_pressure(pressure_tensors, scalar_pressures, device=None)[source]#

Compute scalar pressure from pressure tensor.

\[P = (P_{xx} + P_{yy} + P_{zz}) / 3\]

Parameters:

pressure_tensors (wp.array(dtype=vec9f or vec9d)) – Pressure tensor. Shape (B,).
scalar_pressures (wp.array(dtype=scalar)) – Output scalar pressure. Shape (B,).
device (str, optional) – Warp device.

Returns:

Scalar pressure. Shape (B,).

Return type:

wp.array

NPT Integration (Isothermal-Isobaric)

nvalchemiops.dynamics.integrators.npt.npt_thermostat_half_step(eta, eta_dot, kinetic_energy, target_temperature, thermostat_masses, num_atoms_per_system, chain_length, dt, device=None)[source]#

Perform Nose-Hoover chain thermostat half-step for NPT.

Propagates the thermostat chain with the leading driving force

\[G_1 = \frac{2\,\mathrm{KE} - N_f\,k_B T}{Q_1}, \qquad G_k = \frac{Q_{k-1}\,\dot{\eta}_{k-1}^2 - k_B T}{Q_k}\]

and advances the chain positions by \(\eta_k \mathrel{+}= \tfrac{\Delta t}{2}\,\dot{\eta}_k\), where \(N_f = 3N\), \(Q_k\) are the thermostat_masses and \(\dot{\eta}_k\) the chain velocities.

Parameters:

eta (wp.array2d) – Thermostat positions. Shape (B, chain_length). MODIFIED in-place.
eta_dot (wp.array2d) – Thermostat velocities. Shape (B, chain_length). MODIFIED in-place.
kinetic_energy (wp.array) – Kinetic energy per system. Shape (B,).
target_temperature (wp.array) – Target temperatures. Shape (B,).
thermostat_masses (wp.array2d) – Thermostat masses. Shape (B, chain_length).
num_atoms_per_system (wp.array(dtype=wp.int32)) – Number of atoms per system. Shape (B,).
chain_length (int) – Number of thermostats in chain.
dt (wp.array) – Full time step dt per system. Shape (B,). The half-step and quarter-step factors are applied internally.
device (str, optional) – Warp device.

Return type:

None

nvalchemiops.dynamics.integrators.npt.npt_barostat_half_step(cell_velocities, pressure_tensors, target_pressures, volumes, cell_masses, kinetic_energy, num_atoms_per_system, dt, device=None)[source]#

NPT barostat half-step: \(\dot{\varepsilon} += \frac{\Delta t}{2} \frac{V}{W} (P_{\text{inst}} - P_{\text{ext}})\).

Mode is dispatched by target_pressures dtype (scalar / vec3 / vec9). Barostat-NHC coupling is applied separately by the caller.

Parameters:

cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,). MODIFIED in-place.
pressure_tensors (wp.array(dtype=vec9f or vec9d)) – Current pressure tensors from virial. Shape (B,).
target_pressures (wp.array) – External pressure(s); dtype selects mode (scalar/vec3/vec9).
volumes (wp.array(dtype=scalar)) – Cell volumes V, barostat masses W, system KE. Shape (B,).
cell_masses (wp.array(dtype=scalar)) – Cell volumes V, barostat masses W, system KE. Shape (B,).
kinetic_energy (wp.array(dtype=scalar)) – Cell volumes V, barostat masses W, system KE. Shape (B,).
num_atoms_per_system (wp.array(dtype=wp.int32)) – Atoms per system. Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step; half-step factor applied internally.
device (str, optional) – Warp device. Default: inferred from cell_velocities.

Return type:

None

See also

nph_barostat_half_step: NPH (no thermostat).
npt_velocity_half_step: Particle velocity update with NHC drag.

References

See compute_barostat_mass for the MTK1994 and SSM2004 citations.

nvalchemiops.dynamics.integrators.npt.npt_velocity_half_step(velocities, masses, forces, cell_velocities, volumes=None, eta_dots=None, num_atoms=None, dt=None, batch_idx=None, num_atoms_per_system=None, cells_inv=None, mode='isotropic', device=None)[source]#

Perform half-step velocity update for NPT ensemble.

Updates particle velocities accounting for: 1. Forces from the potential energy surface 2. Coupling to barostat (cell velocity / strain rate) 3. Coupling to thermostat (Nose-Hoover chain)

Mathematical Formulation#

Isotropic mode (default): \(\mathbf{v} \leftarrow \mathbf{v} \exp\!\left(-\Delta t \left((1 + 1/N_{\text{atoms}}) \frac{\operatorname{Tr}(\dot{\varepsilon})}{3} + \dot{\eta}_1\right)\right)\)
Anisotropic / triclinic mode: \(\mathbf{v} \leftarrow \mathbf{v} \exp\!\left(-\Delta t \left(\dot{\varepsilon} + \frac{\operatorname{Tr}(\dot{\varepsilon})}{3 N_{\text{atoms}}} I + \dot{\eta}_1 I\right)\right)\)

where eps_dot = cell_velocities = p_g/W and \(\dot{\eta}_1\) is the first thermostat-chain velocity. This primitive already applies particle thermostat drag through eta_dots[:, 0]; callers that apply particle NHC velocity scaling as a separate Trotter operator should use nph_velocity_half_step() or another no-thermostat velocity primitive instead.

param velocities:

Particle velocities. Shape (N,). MODIFIED in-place.

type velocities:

wp.array(dtype=wp.vec3f or wp.vec3d)

param masses:

Particle masses. Shape (N,).

type masses:

wp.array(dtype=scalar)

param forces:

Forces on particles. Shape (N,).

type forces:

wp.array(dtype=wp.vec3f or wp.vec3d)

param cell_velocities:

Strain-rate matrices \(\dot{\varepsilon} = p_g/W\). Shape (B,).

type cell_velocities:

wp.array(dtype=wp.mat33f or wp.mat33d)

param volumes:

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.

type volumes:

wp.array, optional

param eta_dots:

Thermostat chain velocities. Shape (B, chain_length).

type eta_dots:

wp.array2d(dtype=scalar)

param num_atoms:

Atom count for single-system mode. Shape (1,).

type num_atoms:

wp.array(dtype=wp.int32)

param dt:

Full time step per system. Shape (B,). The half-step factor is applied internally.

type dt:

wp.array(dtype=scalar)

param batch_idx:

System index for each atom. Required for batched simulations.

type batch_idx:

wp.array(dtype=wp.int32), optional

param num_atoms_per_system:

Number of atoms per system. Required for batched simulations.

type num_atoms_per_system:

wp.array(dtype=wp.int32), optional

param cells_inv:

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.

type cells_inv:

wp.array(dtype=wp.mat33f or wp.mat33d), optional

param mode:

Pressure control mode. One of:

"isotropic" (default): scalar coupling on \(\operatorname{Tr}(\dot{\varepsilon})/3\)
"anisotropic": diagonal coupling on \(\dot{\varepsilon}\)
"triclinic": full-tensor coupling on \(\dot{\varepsilon}\)

type mode:

str, optional

param device:

Warp device.

type device:

str, optional

Examples

Single system (isotropic):

>>> npt_velocity_half_step(
...     velocities, masses, forces, cell_velocities,
...     eta_dots=eta_dots, num_atoms=num_atoms, dt=dt,
... )

See also

npt_barostat_half_step: Cell velocity update step.
npt_position_update: Position update step.

Parameters:

velocities (array)
masses (array)
forces (array)
cell_velocities (array)
volumes (array)
eta_dots (array)
num_atoms (array)
dt (array)
batch_idx (array)
num_atoms_per_system (array)
cells_inv (array)
mode (str)
device (str)

Return type:

None

nvalchemiops.dynamics.integrators.npt.npt_position_update(positions, velocities, cells, cell_velocities, dt, cells_inv=None, batch_idx=None, device=None)[source]#

Update particle positions for NPT integration in-place.

Applies \(\mathbf{r} \leftarrow \mathbf{r} + \Delta t\,(\mathbf{v} + \dot{\varepsilon}\,\mathbf{r})\) where cell_velocities is the strain-rate tensor \(\dot{\varepsilon} = p_g / W\).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle positions. Shape (N,). Modified in-place.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle velocities. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched simulations.
device (str, optional) – Warp device. Default: inferred from positions.

Return type:

None

See also

nvalchemiops.dynamics.integrators.npt.npt_position_update_out(): Non-mutating variant.
nvalchemiops.dynamics.integrators.npt.npt_cell_update(): Cell matrix update step.

nvalchemiops.dynamics.integrators.npt.npt_cell_update(cells, cell_velocities, dt, device=None)[source]#

Update cell matrices: \(\mathbf{h}_{\text{new}} = \mathbf{h} + \Delta t\,\dot{\varepsilon}\,\mathbf{h}\).

Parameters:

cells (wp.array) – Cell matrices. MODIFIED in-place.
cell_velocities (wp.array) – Cell velocity matrices.
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
device (str, optional) – Warp device.

Return type:

None

nvalchemiops.dynamics.integrators.npt.npt_velocity_half_step_out(velocities, masses, forces, cell_velocities, volumes=None, eta_dots=None, num_atoms=None, dt=None, velocities_out=None, batch_idx=None, num_atoms_per_system=None, cells_inv=None, mode='isotropic', device=None, _skip_validation=False)[source]#

Perform half-step velocity update for NPT (non-mutating).

Applies the same update as npt_velocity_half_step(),

\[\mathbf{v}(t + \tfrac{\Delta t}{2}) = \mathbf{v} + \frac{\Delta t}{2}\left(\frac{\mathbf{F}}{m} - \mathbf{D}(\dot{\varepsilon}, \dot{\eta}_1)\,\mathbf{v}\right)\]

where \(\mathbf{D}\) is the mode-dependent barostat + thermostat drag (see npt_velocity_half_step()), but writes into velocities_out instead of mutating velocities. Like the in-place variant, this includes particle thermostat drag through eta_dots[:, 0] and should not be combined with separate particle NHC velocity scaling.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Input velocities. Shape (N,). Not modified when velocities_out differs.
masses (wp.array(dtype=scalar)) – Particle masses. Shape (N,).
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on particles. Shape (N,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon} = p_g/W\). Shape (B,).
volumes (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
eta_dots (wp.array) – Thermostat chain velocities.
num_atoms (wp.array(dtype=wp.int32)) – Atom count for single-system mode. Shape (1,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,). The half-step factor is applied internally.
velocities_out (wp.array) – Pre-allocated output array.
batch_idx (wp.array, optional) – System indices for batched simulations.
num_atoms_per_system (wp.array, optional) – Atom counts per system.
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
mode (str, optional) – Pressure control mode: “isotropic”, “anisotropic”, or “triclinic”.
device (str, optional) – Warp device.
_skip_validation (bool, default False) – Internal flag. When True, skip validate_out_array() on velocities_out. Used by the in-place wrapper when input and output are the same array.

Returns:

Updated velocities.

Return type:

wp.array

nvalchemiops.dynamics.integrators.npt.npt_position_update_out(positions, velocities, cells, cell_velocities, dt, positions_out, cells_inv=None, batch_idx=None, device=None, _skip_validation=False)[source]#

Update particle positions for NPT integration, writing results into a pre-allocated output.

Non-mutating variant of npt_position_update(), applying the scaled-coordinate position update

\[\mathbf{r}(t + \Delta t) = \mathbf{r} + \Delta t\,(\mathbf{v} + \dot{\varepsilon}\,\mathbf{r})\]

into positions_out.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Input particle positions. Shape (N,). Not modified when positions_out differs.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle velocities. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
positions_out (wp.array(dtype=wp.vec3f or wp.vec3d)) – Pre-allocated output array for updated positions. Shape (N,).
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched simulations.
device (str, optional) – Warp device. Default: inferred from positions.
_skip_validation (bool, default False) – Internal flag. When True, skip validate_out_array() on positions_out. Used by the in-place wrapper when input and output are the same array.

Returns:

Updated positions array (same object as positions_out). Shape (N,).

Return type:

wp.array(dtype=wp.vec3f or wp.vec3d)

See also

nvalchemiops.dynamics.integrators.npt.npt_position_update(): In-place variant.

nvalchemiops.dynamics.integrators.npt.npt_cell_update_out(cells, cell_velocities, dt, cells_out, device=None, _skip_validation=False)[source]#

Update cell matrices, writing results into a pre-allocated output.

Non-mutating variant of npt_cell_update(), applying the multiplicative cell update

\[\mathbf{h}(t + \Delta t) = \mathbf{h} + \Delta t\,\dot{\varepsilon}\,\mathbf{h}\]

into cells_out.

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Input cell matrices. Shape (B,). Not modified when cells_out differs.
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon} = p_g/W\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
cells_out (wp.array(dtype=wp.mat33f or wp.mat33d)) – Pre-allocated output array for updated cell matrices. Shape (B,).
device (str, optional) – Warp device. Default: inferred from cells.
_skip_validation (bool, default False) – Internal flag. When True, skip validate_out_array() on cells_out. Used by the in-place wrapper when input and output are the same array.

Returns:

Updated cell matrices (same object as cells_out). Shape (B,).

Return type:

wp.array(dtype=wp.mat33f or wp.mat33d)

See also

nvalchemiops.dynamics.integrators.npt.npt_cell_update(): In-place variant.

nvalchemiops.dynamics.integrators.npt.run_npt_step(positions, velocities, forces, masses, cells, cell_velocities, virial_tensors, eta, eta_dot, thermostat_masses, cell_masses, target_temperature, target_pressure, num_atoms, chain_length, dt, pressure_tensors, volumes, kinetic_energy, cells_inv=None, kinetic_tensors=None, num_atoms_per_system=None, compute_forces_fn=None, batch_idx=None, device=None)[source]#

Perform a complete NPT integration step.

Integration order: 1. Thermostat half-step 2. Barostat half-step 3. Velocity half-step 4. Position update 5. Cell update 6. Recompute forces 7. Velocity half-step 8. Barostat half-step 9. Thermostat half-step

Particle thermostat friction is applied inside the two npt_velocity_half_step() calls via eta_dot[:, 0]. The npt_thermostat_half_step() calls update the chain state; callers using a different split with explicit particle velocity scaling should not also use npt_velocity_half_step for those particle half-steps.

Parameters:

positions (wp.array) – Particle state arrays. MODIFIED in-place.
velocities (wp.array) – Particle state arrays. MODIFIED in-place.
forces (wp.array) – Particle state arrays. MODIFIED in-place.
masses (wp.array) – Particle masses.
cells (wp.array) – Cell matrices. MODIFIED in-place.
cell_velocities (wp.array) – Cell velocity matrices. MODIFIED in-place.
virial_tensors (wp.array(dtype=vec9f or vec9d)) – Virial tensor. Updated by compute_forces_fn.
eta (wp.array2d) – Thermostat state. MODIFIED in-place.
eta_dot (wp.array2d) – Thermostat state. MODIFIED in-place.
thermostat_masses (wp.array2d) – Thermostat masses.
cell_masses (wp.array) – Barostat masses.
target_temperature (wp.array) – Target conditions.
target_pressure (wp.array) – Target conditions.
num_atoms (wp.array(dtype=wp.int32)) – Atom count for single-system mode. Shape (1,).
chain_length (int) – Number of thermostats in chain.
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
pressure_tensors (wp.array(dtype=vec9f or vec9d)) – Scratch array for pressure tensor. Shape (B,).
volumes (wp.array(dtype=scalar)) – Scratch array for cell volumes. Shape (B,).
kinetic_energy (wp.array(dtype=scalar)) – Scratch array for kinetic energy. Shape (B,). Zeroed internally before each use.
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
kinetic_tensors (wp.array(dtype=scalar, ndim=2)) – Scratch array for kinetic tensor accumulation. Shape (B, 9). Zeroed internally before each use.
num_atoms_per_system (wp.array(dtype=wp.int32)) – Number of atoms per system. Shape (B,).
compute_forces_fn (callable, optional) –
Callback invoked once after the position and cell update (step 6) to recompute forces at the new geometry. Signature:
```
compute_forces_fn(positions, cells, forces, virial_tensors)
```
positions and cells are the integrator state arrays after the update; forces and virial_tensors are pre-allocated output buffers that must be filled in-place. positions and forces are shape (N,) arrays of wp.vec3f or wp.vec3d; cells is a shape (B,) array of wp.mat33f or wp.mat33d; and virial_tensors is shape (B,) with the matching wp.vec9f/wp.vec9d precision. Any return value is ignored. If None, force recomputation is skipped.
batch_idx (wp.array, optional) – System index for each atom.
device (str, optional) – Warp device.

Return type:

None

NPH Integration (Isenthalpic-Isobaric)

nvalchemiops.dynamics.integrators.npt.nph_barostat_half_step(cell_velocities, pressure_tensors, target_pressures, volumes, cell_masses, kinetic_energy, num_atoms_per_system, dt, device=None)[source]#

Perform barostat half-step for NPH ensemble (no thermostat coupling).

NPH (isenthalpic-isobaric) simulations maintain constant pressure and enthalpy. Unlike NPT, there is no thermostat - the temperature evolves naturally on the constant-enthalpy surface.

This function updates the cell velocity matrix \(\dot{h}\) based on the pressure difference. The pressure control mode is automatically detected from the target_pressures array dtype:

Isotropic (scalar dtype): Uniform scaling in all directions
Anisotropic/Orthorhombic (vec3 dtype): Independent x, y, z control
Triclinic (vec9 dtype): Full stress tensor control

Mathematical Formulation#

The cell velocity follows the MTK equations without thermostat:

Isotropic mode:

\[\ddot{h} = \frac{V}{W}(P - P_{\text{ext}})\]

Anisotropic mode:

\[\ddot{h}_{ii} = \frac{V}{W}(P_{ii} - P_{\text{ext},ii})\]

Triclinic mode:

\[\ddot{h}_{ij} = \frac{V}{W}(P_{ij} - P_{\text{ext},ij})\]

param cell_velocities:

Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,). MODIFIED in-place.

type cell_velocities:

wp.array(dtype=wp.mat33f or wp.mat33d)

param pressure_tensors:

Current pressure tensors from virial. Shape (B,).

type pressure_tensors:

wp.array(dtype=vec9f or vec9d)

param target_pressures:

External/target pressure(s). The dtype determines the mode:

wp.float32 or wp.float64: Isotropic. Shape (B,).
wp.vec3f or wp.vec3d: Anisotropic [Pxx, Pyy, Pzz]. Shape (B,).
vec9f or vec9d: Full stress tensor. Shape (B,).

type target_pressures:

wp.array

param volumes:

Cell volumes V. Shape (B,).

type volumes:

wp.array(dtype=scalar)

param cell_masses:

Barostat masses W. Shape (B,).

type cell_masses:

wp.array(dtype=scalar)

param kinetic_energy:

System kinetic energies. Shape (B,).

type kinetic_energy:

wp.array(dtype=scalar)

param num_atoms_per_system:

Number of atoms per system. Shape (B,).

type num_atoms_per_system:

wp.array(dtype=wp.int32)

param dt:

Full time step per system. Shape (B,). The half-step factor is applied internally.

type dt:

wp.array(dtype=scalar)

param device:

Warp device.

type device:

str, optional

Examples

Isotropic NPH:

>>> target_P = wp.array([1.0], dtype=wp.float32, device="cuda:0")
>>> nph_barostat_half_step(
...     cell_velocities, pressure_tensors, target_P,
...     volumes, cell_masses, kinetic_energy,
...     num_atoms_per_system, dt=0.001
... )

Anisotropic NPH:

>>> target_P = wp.array([[1.0, 2.0, 1.5]], dtype=wp.vec3f, device="cuda:0")
>>> nph_barostat_half_step(
...     cell_velocities, pressure_tensors, target_P, ...
... )

See also

npt_barostat_half_step: Barostat with thermostat coupling.
run_nph_step: Complete NPH integration step.

Parameters:

cell_velocities (array)
pressure_tensors (array)
target_pressures (array)
volumes (array)
cell_masses (array)
kinetic_energy (array)
num_atoms_per_system (array)
dt (array)
device (str)

Return type:

None

nvalchemiops.dynamics.integrators.npt.nph_velocity_half_step(velocities, masses, forces, cell_velocities, volumes=None, num_atoms=None, dt=None, batch_idx=None, num_atoms_per_system=None, cells_inv=None, mode='isotropic', device=None)[source]#

Perform half-step velocity update for NPH ensemble (no thermostat).

Updates particle velocities accounting for: 1. Forces from the potential energy surface 2. Coupling to barostat (cell velocity / strain rate)

Unlike NPT, there is no thermostat coupling - the temperature evolves naturally on the constant-enthalpy surface.

Mathematical Formulation#

Isotropic mode: \(\mathbf{v} \leftarrow \mathbf{v} \exp\!\left(-\Delta t (1 + 1/N_{\text{atoms}}) \frac{\operatorname{Tr}(\dot{\varepsilon})}{3}\right)\)
Anisotropic / triclinic mode: \(\mathbf{v} \leftarrow \mathbf{v} \exp\!\left(-\Delta t \left(\dot{\varepsilon} + \frac{\operatorname{Tr}(\dot{\varepsilon})}{3 N_{\text{atoms}}} I\right)\right)\)

param velocities:

Particle velocities. MODIFIED in-place.

type velocities:

wp.array(dtype=wp.vec3f or wp.vec3d)

param masses:

Particle masses.

type masses:

wp.array(dtype=scalar)

param forces:

Forces on particles.

type forces:

wp.array(dtype=wp.vec3f or wp.vec3d)

param cell_velocities:

Strain-rate matrices \(\dot{\varepsilon} = p_g/W\).

type cell_velocities:

wp.array(dtype=wp.mat33f or wp.mat33d)

param volumes:

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.

type volumes:

wp.array, optional

param num_atoms:

Atom count for single-system mode. Shape (1,).

type num_atoms:

wp.array(dtype=wp.int32)

param dt:

Full time step per system. Shape (B,). The half-step factor is applied internally.

type dt:

wp.array(dtype=scalar)

param batch_idx:

System index for each atom.

type batch_idx:

wp.array, optional

param num_atoms_per_system:

Number of atoms per system.

type num_atoms_per_system:

wp.array, optional

param cells_inv:

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.

type cells_inv:

wp.array(dtype=wp.mat33f or wp.mat33d), optional

param mode:

Pressure control mode:

"isotropic": scalar coupling on \(\operatorname{Tr}(\dot{\varepsilon})/3\)
"anisotropic": diagonal coupling on \(\dot{\varepsilon}\)
"triclinic": full-tensor coupling on \(\dot{\varepsilon}\)

type mode:

str, optional

param device:

Warp device.

type device:

str, optional

See also

nph_barostat_half_step: Cell velocity update.
npt_velocity_half_step: Velocity update with thermostat.

Parameters:

velocities (array)
masses (array)
forces (array)
cell_velocities (array)
volumes (array)
num_atoms (array)
dt (array)
batch_idx (array)
num_atoms_per_system (array)
cells_inv (array)
mode (str)
device (str)

Return type:

None

nvalchemiops.dynamics.integrators.npt.nph_position_update(positions, velocities, cells, cell_velocities, dt, cells_inv=None, batch_idx=None, device=None)[source]#

Update particle positions for NPH integration in-place.

Applies the scaled-coordinate position update

\[\mathbf{r}(t + \Delta t) = \mathbf{r} + \Delta t\,(\mathbf{v} + \dot{\varepsilon}\,\mathbf{r})\]

Delegates to npt_position_update(); the position update kernel is identical for NPT and NPH.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle positions. Shape (N,). Modified in-place.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle velocities. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched simulations.
device (str, optional) – Warp device. Default: inferred from positions.

Return type:

None

See also

nvalchemiops.dynamics.integrators.npt.nph_position_update_out(): Non-mutating variant.
nvalchemiops.dynamics.integrators.npt.npt_position_update(): NPT equivalent.

nvalchemiops.dynamics.integrators.npt.nph_cell_update(cells, cell_velocities, dt, device=None)[source]#

Update cell matrices for NPH integration in-place.

Applies the multiplicative cell update

\[\mathbf{h}(t + \Delta t) = \mathbf{h} + \Delta t\,\dot{\varepsilon}\,\mathbf{h}\]

Delegates to npt_cell_update(); the cell update kernel is identical for NPT and NPH.

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,). Modified in-place.
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
device (str, optional) – Warp device. Default: inferred from cells.

Return type:

None

See also

nvalchemiops.dynamics.integrators.npt.npt_cell_update(): NPT equivalent.

nvalchemiops.dynamics.integrators.npt.nph_velocity_half_step_out(velocities, masses, forces, cell_velocities, volumes=None, num_atoms=None, dt=None, velocities_out=None, batch_idx=None, num_atoms_per_system=None, cells_inv=None, mode='isotropic', device=None, _skip_validation=False)[source]#

Perform half-step velocity update for NPH (non-mutating).

Applies the same update as nph_velocity_half_step(),

\[\mathbf{v}(t + \tfrac{\Delta t}{2}) = \mathbf{v} + \frac{\Delta t}{2}\left(\frac{\mathbf{F}}{m} - \mathbf{D}(\dot{\varepsilon})\,\mathbf{v}\right)\]

where \(\mathbf{D}\) is the mode-dependent barostat drag (no thermostat term), but writes into velocities_out instead of mutating velocities.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Input velocities. Shape (N,). Not modified when velocities_out differs.
masses (wp.array(dtype=scalar)) – Particle masses. Shape (N,).
forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces on particles. Shape (N,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon} = p_g/W\). Shape (B,).
volumes (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
num_atoms (wp.array(dtype=wp.int32)) – Atom count for single-system mode. Shape (1,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,). The half-step factor is applied internally.
velocities_out (wp.array) – Pre-allocated output array.
batch_idx (wp.array, optional) – For batched simulations.
num_atoms_per_system (wp.array, optional) – For batched simulations.
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
mode (str, optional) – Pressure control mode: “isotropic”, “anisotropic”, or “triclinic”.
device (str, optional) – Warp device.
_skip_validation (bool, default False) – Internal flag. When True, skip validate_out_array() on velocities_out. Used by the in-place wrapper when input and output are the same array.

Returns:

Updated velocities.

Return type:

wp.array

nvalchemiops.dynamics.integrators.npt.nph_position_update_out(positions, velocities, cells, cell_velocities, dt, positions_out, cells_inv=None, batch_idx=None, device=None)[source]#

Update particle positions for NPH integration, writing results into a pre-allocated output.

Non-mutating variant of nph_position_update(), applying the scaled-coordinate position update

\[\mathbf{r}(t + \Delta t) = \mathbf{r} + \Delta t\,(\mathbf{v} + \dot{\varepsilon}\,\mathbf{r})\]

Delegates to npt_position_update_out(); the position update kernel is identical for NPT and NPH.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Input particle positions. Shape (N,). Not modified when positions_out differs.
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Particle velocities. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon}\). Shape (B,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
positions_out (wp.array(dtype=wp.vec3f or wp.vec3d)) – Pre-allocated output array for updated positions. Shape (N,).
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Required for batched simulations.
device (str, optional) – Warp device. Default: inferred from positions.

Returns:

Updated positions array (same object as positions_out). Shape (N,).

Return type:

wp.array(dtype=wp.vec3f or wp.vec3d)

See also

nvalchemiops.dynamics.integrators.npt.nph_position_update(): In-place variant.

nvalchemiops.dynamics.integrators.npt.run_nph_step(positions, velocities, forces, masses, cells, cell_velocities, virial_tensors, cell_masses, target_pressure, num_atoms, dt, pressure_tensors, volumes, kinetic_energy, cells_inv=None, kinetic_tensors=None, num_atoms_per_system=None, compute_forces_fn=None, batch_idx=None, device=None)[source]#

Perform a complete NPH integration step.

Integration order (no thermostat): 1. Barostat half-step 2. Velocity half-step 3. Position update 4. Cell update 5. Recompute forces 6. Velocity half-step 7. Barostat half-step

Parameters:

positions (wp.array) – Particle state arrays. MODIFIED in-place.
velocities (wp.array) – Particle state arrays. MODIFIED in-place.
forces (wp.array) – Particle state arrays. MODIFIED in-place.
masses (wp.array) – Particle masses.
cells (wp.array) – Cell matrices. MODIFIED in-place.
cell_velocities (wp.array) – Cell velocity matrices. MODIFIED in-place.
virial_tensors (wp.array(dtype=vec9f or vec9d)) – Virial tensor. Updated by compute_forces_fn.
cell_masses (wp.array) – Barostat masses.
target_pressure (wp.array) – Target/external pressures.
num_atoms (wp.array(dtype=wp.int32)) – Atom count for single-system mode. Shape (1,).
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
pressure_tensors (wp.array(dtype=vec9f or vec9d)) – Scratch array for pressure tensor. Shape (B,).
volumes (wp.array(dtype=scalar)) – Scratch array for cell volumes. Shape (B,).
kinetic_energy (wp.array(dtype=scalar)) – Scratch array for kinetic energy. Shape (B,). Zeroed internally before each use.
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
kinetic_tensors (wp.array(dtype=scalar, ndim=2)) – Scratch array for kinetic tensor accumulation. Shape (B, 9). Zeroed internally before each use.
num_atoms_per_system (wp.array(dtype=wp.int32)) – Number of atoms per system. Shape (B,).
compute_forces_fn (callable, optional) –
Callback invoked once after the position and cell update (step 5) to recompute forces at the new geometry. Signature:
```
compute_forces_fn(positions, cells, forces, virial_tensors)
```
positions and cells are the integrator state arrays after the update; forces and virial_tensors are pre-allocated output buffers that must be filled in-place. positions and forces are shape (N,) arrays of wp.vec3f or wp.vec3d; cells is a shape (B,) array of wp.mat33f or wp.mat33d; and virial_tensors is shape (B,) with the matching wp.vec9f/wp.vec9d precision. Any return value is ignored. If None, force recomputation is skipped.
batch_idx (wp.array, optional) – System index for each atom.
device (str, optional) – Warp device.

Return type:

None

Barostat Utilities

nvalchemiops.dynamics.integrators.npt.compute_barostat_mass(target_temperature, tau_p, num_atoms, masses_out, device=None)[source]#

Compute barostat mass(es) for desired pressure fluctuation timescale.

The barostat mass determines the inertia of cell volume/shape fluctuations in NPT/NPH simulations. It is computed from the Martyna-Tobias-Klein equations [MTK1994] to give a characteristic pressure relaxation time \(\tau_p\):

\[W = (N_f + d) \cdot k_B T \cdot \tau_p^2\]

where: - N_f = 3N is the number of degrees of freedom (N atoms in 3D) - d = 3 is the dimensionality (for 3D simulations) - k_B T is the thermal energy (in reduced units, k_B = 1) - \(\tau_p\) is the pressure relaxation time

Larger W gives slower pressure equilibration (more stable but slower to reach target pressure). Smaller W gives faster equilibration but may cause oscillations.

Parameters:

target_temperature (wp.array) – Per-system target temperatures. Shape (num_systems,). Caller must pre-broadcast if a single value applies to all systems.
tau_p (wp.array) – Per-system pressure relaxation times. Shape (num_systems,). Typical values: 0.5-2.0 ps. Caller must pre-broadcast if a single value applies to all systems.
num_atoms (wp.array(dtype=wp.int32)) – Per-system atom counts. Shape (num_systems,). Caller must pre-broadcast if a single value applies to all systems.
masses_out (wp.array) – Output barostat masses W. Shape (num_systems,).
device (str, optional) – Warp device.

Returns:

Barostat mass(es) W. Shape (num_systems,).

Return type:

wp.array

Examples

Batched systems (using wp.arrays):

>>> temps = wp.array([1.0, 2.0], dtype=wp.float64, device="cuda:0")
>>> tau = wp.array([1.0, 1.0], dtype=wp.float64, device="cuda:0")
>>> n_atoms = wp.array([100, 200], dtype=wp.int32, device="cuda:0")
>>> W = wp.empty(2, dtype=wp.float64, device="cuda:0")
>>> compute_barostat_mass(temps, tau, n_atoms, W)
>>> print(W.numpy())  # [303.0, 1206.0]

Notes

The formula assumes k_B = 1 (reduced units). Scale \(\tau_p\) accordingly for real units.
For isotropic barostat, a single mass controls all cell dimensions.
For anisotropic/triclinic barostat [SSM2004], the same mass is typically used for all cell velocity components.
All input arrays must have the same length (num_systems). The caller is responsible for broadcasting scalar values to arrays before calling.

References

[MTK1994]

Martyna, Tobias, Klein, J. Chem. Phys. 101, 4177 (1994)

[SSM2004]

Shinoda, Shiga, Mikami, Phys. Rev. B 69, 134103 (2004)

nvalchemiops.dynamics.integrators.npt.compute_cell_kinetic_energy(cell_velocities, cell_masses, kinetic_energy, cells_inv=None, volumes=None, device=None)[source]#

Compute kinetic energy of cell degrees of freedom.

\(\mathrm{KE}_{\text{cell}} = \frac{1}{2} W \|\dot{\varepsilon}\|_F^2\)

Parameters:

cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Strain-rate matrices \(\dot{\varepsilon} = p_g/W\). Shape (B,).
cell_masses (wp.array) – Barostat masses. Shape (B,). Note that this is the per-DOF mass \(W = (d + 1) k_B T \tau_p^2\) consumed by the strain-rate update, not the full barostat mass \(d \cdot W\).
kinetic_energy (wp.array(dtype=scalar)) – Output cell kinetic energy. Shape (B,).
cells_inv (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
volumes (wp.array, optional) –

Deprecated since version 0.3.1: Ignored; cell_velocities is the strain rate eps_dot = p_g/W.
device (str, optional) – Warp device.

Returns:

Cell kinetic energy. Shape (B,).

Return type:

wp.array

nvalchemiops.dynamics.integrators.npt.compute_barostat_potential_energy(target_pressures, volumes, potential_energy, device=None)[source]#

Compute barostat potential energy.

\[U = P_{\text{ext}} V\]

where \(P_{\text{ext}}\) is the target pressure and \(V\) the cell volume.

Parameters:

target_pressures (wp.array) – External/target pressures. Shape (B,).
volumes (wp.array) – Cell volumes. Shape (B,).
potential_energy (wp.array) – Output barostat potential energy. Shape (B,).
device (str, optional) – Warp device.

Returns:

Barostat potential energy. Shape (B,).

Return type:

wp.array

Optimizers#

FIRE#

Fast Inertial Relaxation Engine for geometry optimization.

nvalchemiops.dynamics.optimizers.fire.fire_step(positions, velocities, forces, masses, alpha, dt, alpha_start, f_alpha, dt_min, dt_max, maxstep, n_steps_positive, n_min, f_dec, f_inc, uphill_flag, vf=None, vv=None, ff=None, batch_idx=None, atom_ptr=None, energy=None, energy_last=None, positions_last=None, velocities_last=None, compute_reductions=True)[source]#

Unified FIRE optimization step with MD integration.

Per system, reduce the power \(P = \sum_i \mathbf{F}_i \cdot \mathbf{v}_i\) together with \(\sum_i \mathbf{v}_i \cdot \mathbf{v}_i\) and \(\sum_i \mathbf{F}_i \cdot \mathbf{F}_i\), mix each atom’s velocity toward the force direction, take a mass-weighted MD kick, and apply the displacement with a per-step cap:

\[\begin{split}\mathbf{v} &\leftarrow (1-\alpha)\,\mathbf{v} + \alpha \sqrt{\tfrac{\mathbf{v}\cdot\mathbf{v}} {\mathbf{F}\cdot\mathbf{F}}}\,\mathbf{F} \\ \mathbf{v} &\leftarrow \mathbf{v} + \Delta t\,\mathbf{F}/m \\ \Delta\mathbf{r} &= \Delta t\,\mathbf{v},\quad \mathbf{r} \leftarrow \mathbf{r} + \min\!\left(1, \tfrac{\text{maxstep}}{\lVert\Delta\mathbf{r}\rVert}\right) \Delta\mathbf{r}\end{split}\]

When \(P > 0\) the timestep grows (after n_min consecutive positive steps) and \(\alpha\) decays; when \(P \le 0\) (uphill) the velocity is zeroed, \(\Delta t\) shrinks, and \(\alpha\) resets to alpha_start. Mixing uses the fire_velocity_mixing form \(\alpha\sqrt{vv/ff}\,\mathbf{F}\) rather than the textbook \(\alpha\lVert\mathbf{v}\rVert\hat{\mathbf{F}}\); the two coincide only when the pre-mix speed already matches the force-scaled velocity.

This function dispatches to the appropriate kernel based on: - Batching mode: single system, batch_idx, or atom_ptr - Downhill check: enabled if all downhill arrays are provided

Parameters:

positions (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*) – Atomic positions (modified in-place).
velocities (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*) – Atomic velocities (modified in-place).
forces (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*) – Forces on atoms.
masses (wp.array, shape (N,) or (N_total,), dtype=wp.float*) – Per-atom masses.
alpha (wp.array, shape (1,) or (B,), dtype=wp.float*) – FIRE mixing parameter.
dt (wp.array, shape (1,) or (B,), dtype=wp.float*) – FIRE timestep.
alpha_start (wp.array, shape (1,) or (B,), dtype=wp.float*) – Reset value for alpha.
f_alpha (wp.array, shape (1,) or (B,), dtype=wp.float*) – Alpha decay factor.
dt_min (wp.array, shape (1,) or (B,), dtype=wp.float*) – Minimum timestep.
dt_max (wp.array, shape (1,) or (B,), dtype=wp.float*) – Maximum timestep.
maxstep (wp.array, shape (1,) or (B,), dtype=wp.float*) – Maximum displacement per step.
n_steps_positive (wp.array, shape (1,) or (B,), dtype=wp.int32) – Counter for consecutive positive power steps.
n_min (wp.array, shape (1,) or (B,), dtype=wp.int32) – Steps before dt increase / alpha decrease.
f_dec (wp.array, shape (1,) or (B,), dtype=wp.float*) – Timestep decrease factor.
f_inc (wp.array, shape (1,) or (B,), dtype=wp.float*) – Timestep increase factor.
vf (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
vv (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
ff (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
uphill_flag (wp.array, shape (B,), dtype=wp.int32, optional) – Scratch array for uphill detection. Shape (B,) where B = num_systems. Only used when downhill_enabled=True and batch_idx is provided.
batch_idx (wp.array, shape (N_total,), dtype=wp.int32, optional) – System index per atom. If provided, uses batch_idx kernel.
atom_ptr (wp.array, shape (B+1,), dtype=wp.int32, optional) – CSR pointers for atom ranges. If provided, uses ptr kernel.
energy (wp.array, shape (1,) or (B,), dtype=wp.float*, optional) – Current energies (for downhill check).
energy_last (wp.array, shape (1,) or (B,), dtype=wp.float*, optional) – Last accepted energies (for downhill check).
positions_last (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*, optional) – Last accepted positions (for downhill rollback).
velocities_last (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*, optional) – Last accepted velocities (for downhill rollback).
compute_reductions (bool, optional) – If True (default), recompute the per-system reductions vf/vv/ff (\(\sum \mathbf{F} \cdot \mathbf{v}\) / \(\sum \mathbf{v} \cdot \mathbf{v}\) / \(\sum \mathbf{F} \cdot \mathbf{F}\)) internally from velocities/forces. If False, use the caller-supplied values already in vf/vv/ff instead of recomputing them (they are not zeroed); the state roll-back still runs, so the supplied values must reflect the post-roll-back velocities. Only supported for single/batch_idx mode. See fire_compute_vf_vv_ff for the matching reduction helper.

Return type:

None

Examples

Single system (no downhill):

>>> fire_step(positions, velocities, forces, masses,
...           alpha, dt, alpha_start, f_alpha, dt_min, dt_max,
...           maxstep, n_steps_positive, n_min, f_dec, f_inc,
...           vf, vv, ff)

Batched with batch_idx:

>>> fire_step(positions, velocities, forces, masses,
...           alpha, dt, alpha_start, f_alpha, dt_min, dt_max,
...           maxstep, n_steps_positive, n_min, f_dec, f_inc,
...           vf, vv, ff, batch_idx=batch_idx)

Batched with atom_ptr:

>>> fire_step(positions, velocities, forces, masses,
...           alpha, dt, alpha_start, f_alpha, dt_min, dt_max,
...           maxstep, n_steps_positive, n_min, f_dec, f_inc,
...           atom_ptr=atom_ptr)

With downhill check:

>>> fire_step(positions, velocities, forces, masses,
...           alpha, dt, alpha_start, f_alpha, dt_min, dt_max,
...           maxstep, n_steps_positive, n_min, f_dec, f_inc,
...           vf, vv, ff,
...           energy=energy, energy_last=energy_last,
...           positions_last=positions_last, velocities_last=velocities_last)

nvalchemiops.dynamics.optimizers.fire.fire_update(velocities, forces, alpha, dt, alpha_start, f_alpha, dt_min, dt_max, n_steps_positive, n_min, f_dec, f_inc, vf=None, vv=None, ff=None, batch_idx=None, atom_ptr=None, energy=None, energy_last=None, positions=None, positions_last=None, velocities_last=None, compute_reductions=True)[source]#

FIRE parameter update and velocity mixing WITHOUT MD integration.

Reduces the per-system power \(P = \sum_i \mathbf{F}_i \cdot \mathbf{v}_i\) (with \(\sum_i \mathbf{v}_i \cdot \mathbf{v}_i\) and \(\sum_i \mathbf{F}_i \cdot \mathbf{F}_i\)), advances \(\Delta t\), \(\alpha\), and the positive-step counter by the FIRE rules, and mixes each velocity toward the force direction:

\[\mathbf{v} \leftarrow (1-\alpha)\,\mathbf{v} + \alpha \sqrt{\tfrac{\mathbf{v}\cdot\mathbf{v}} {\mathbf{F}\cdot\mathbf{F}}}\,\mathbf{F}\]

Unlike fire_step(), it stops here: no MD kick (\(\mathbf{v} \mathrel{+}= \Delta t\,\mathbf{F}/m\)), no displacement (\(\Delta\mathbf{r} = \Delta t\,\mathbf{v}\)), and no maxstep clamp. Uphill systems (\(P \le 0\)) still have their velocity zeroed and \(\alpha\)/\(\Delta t\) reset.

Use this for variable-cell optimization where you want to: 1. Pack atomic + cell DOFs into extended arrays 2. Apply FIRE velocity mixing to extended velocities 3. Perform your own MD step (e.g., with cell-aware position scaling)

This function dispatches to the appropriate “update params” kernel based on: - Batching mode: single system, batch_idx, or atom_ptr - Downhill check: enabled if all downhill arrays are provided

Parameters:

velocities (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*) – Velocities (modified in-place with FIRE mixing).
forces (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*) – Atomic forces.
alpha (wp.array, shape (1,) or (B,), dtype=wp.float*) – FIRE mixing parameter.
dt (wp.array, shape (1,) or (B,), dtype=wp.float*) – FIRE timestep.
alpha_start (wp.array, shape (1,) or (B,), dtype=wp.float*) – Reset value for alpha.
f_alpha (wp.array, shape (1,) or (B,), dtype=wp.float*) – Alpha decay factor.
dt_min (wp.array, shape (1,) or (B,), dtype=wp.float*) – Minimum timestep.
dt_max (wp.array, shape (1,) or (B,), dtype=wp.float*) – Maximum timestep.
n_steps_positive (wp.array, shape (1,) or (B,), dtype=wp.int32) – Counter for consecutive positive power steps.
n_min (wp.array, shape (1,) or (B,), dtype=wp.int32) – Steps before dt increase / alpha decrease.
f_dec (wp.array, shape (1,) or (B,), dtype=wp.float*) – Timestep decrease factor.
f_inc (wp.array, shape (1,) or (B,), dtype=wp.float*) – Timestep increase factor.
vf (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
vv (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
ff (wp.array, shape (1,) or (B,), dtype=wp.float*) – Accumulators for diagnostics. Zeroed internally before each use. Required for single/batch_idx modes. Ignored for atom_ptr mode.
batch_idx (wp.array, shape (N_total,), dtype=wp.int32, optional) – System index per atom. If provided, uses batch_idx kernel.
atom_ptr (wp.array, shape (B+1,), dtype=wp.int32, optional) – CSR pointers for atom ranges. If provided, uses ptr kernel.
energy (wp.array, shape (1,) or (B,), dtype=wp.float*, optional) – Current energies (for downhill check).
energy_last (wp.array, shape (1,) or (B,), dtype=wp.float*, optional) – Last accepted energies (for downhill check).
positions (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*, optional) – Positions (for downhill rollback). Required if downhill enabled.
positions_last (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*, optional) – Last accepted positions (for downhill rollback).
velocities_last (wp.array, shape (N,) or (N_total,), dtype=wp.vec3*, optional) – Last accepted velocities (for downhill rollback).
compute_reductions (bool, optional) – If True (default), recompute vf/vv/ff internally. If False, use the caller-supplied values already in vf/vv/ff instead of recomputing them (they are not zeroed); any state roll-back still runs. Only supported for single/batch_idx mode. See fire_step and fire_compute_vf_vv_ff.

Return type:

None

Examples

Variable-cell optimization workflow:

>>> # Pack extended arrays (atomic + cell DOFs)
>>> ext_pos = pack_positions_with_cell(positions, cell)
>>> ext_vel = pack_velocities_with_cell(velocities, cell_velocity)
>>> ext_forces = pack_forces_with_cell(forces, cell_force)
>>>
>>> # FIRE velocity mixing only (no position update)
>>> fire_update(ext_vel, ext_forces,
...             alpha, dt, alpha_start, f_alpha, dt_min, dt_max,
...             n_steps_positive, n_min, f_dec, f_inc,
...             vf, vv, ff)
>>>
>>> # Perform your own MD step with cell-aware scaling
>>> ext_vel += dt * ext_forces / ext_masses
>>> ext_pos += dt * ext_vel  # (with maxstep capping)
>>>
>>> # Unpack results
>>> positions, cell = unpack_positions_with_cell(ext_pos, num_atoms)

L-BFGS#

Limited-memory quasi-Newton optimizer with a maxstep trust region. Each lbfgs_step() call consumes exactly one force evaluation: it updates the curvature history, restarts if the direction stops descending, and takes one bounded step. A whole batch relaxes in one stream of kernel launches.

Like FIRE2, it owns no tolerance and has no terminal status – testing convergence and ending the loop are the caller’s. It reads no energy either, so a model whose forces are not the gradient of its energy relaxes just as well.

The public surface is the two states, the preparation helpers, one step per call, and the variable-cell setup below.

The module documentation states the full calling convention, the dtype and shape contract, and the variable-cell contract:

L-BFGS Optimizer Kernels#

GPU-accelerated Warp kernels for batched L-BFGS geometry optimization.

L-BFGS approximates the inverse Hessian from the last history_size position/force differences to pick a direction, then steps along it as far as a maxstep trust region allows. It reaches a given force tolerance in far fewer force evaluations than FIRE – the cost that dominates relaxation with a machine-learned potential.

Calling convention#

You own the loop and the stopping rule. Each lbfgs_step() consumes exactly one force evaluation, and every call is an accepted step forming one curvature pair:

for _ in range(max_steps):
    forces = my_model(positions)
    if np.linalg.norm(forces.numpy(), axis=1).max() < force_tol:
        break
    lbfgs_step(positions, forces, state, batch_idx, maxstep=0.2)

Test before stepping, as above. After a step positions hold a new, unevaluated point and the forces you passed describe the point before it; the optimizer keeps that pair in x_base/force_base. Positions only move forward, so a point you have left is never revisited.

The optimizer owns no tolerance and has no terminal state, mirroring FIRE2. Two things are algorithmic and stay inside: a pair with s . y too small is discarded, and a direction with d0 >= 0 is replaced by steepest descent. Neither stops anything.

Nothing here reads an energy, so there is no line search. maxstep shrinks alpha rather than clamping each displacement, because the stored pair s = alpha * d is defined through that relation. Nothing forces the energy down – it may rise on a step – but the force still converges.

Systems stay in lock step in evaluations while diverging in iterations, so a batch relaxes in one stream of launches with no per-system host control flow. A batch converges when its slowest system does; compact the batch on the host to retire finished systems earlier.

The API is in forces, the algorithm in gradients, with g = -F folded into each kernel: y_history holds force_base - F, and a descent direction satisfies force_base . d > 0.

Relation to FIRE2#

The decomposition mirrors FIRE2’s, so a caller can interpose logic at the same points. The extra phase exists because the step length must be settled before the atoms move.

Phase	FIRE2	L-BFGS
per-system reductions	`fire2_reduce`	`_lbfgs_reduce()`
state update, no motion	`fire2_update`	`_lbfgs_update()`
step-length finalization	–	`_lbfgs_prepare_step()`
move the atoms	`fire2_apply_step`	`_lbfgs_apply_step()`
all of the above	`fire2_step`	`lbfgs_step()`

That is internal structure. The public surface is the state, the preparation helpers and one step per call; the phases are importable by name but are decomposition points, not operations in their own right.

Precision#

Every array follows the coordinate dtype. wp.vec3f gives an fp32 state end to end, wp.vec3d an fp64 one. There is no mixed configuration – the overloads are keyed so a mismatched state is rejected rather than silently half-converted. Widening a reduction on its own buys nothing measurable here, since y = force_base - F is already formed at the coordinate precision.

curvature_eps follows from this. It gates the history on ys > eps * sqrt(ss * yy) – a floor on the cosine between s and y – and ys accumulates at the coordinate precision, so the floor must sit above the level at which that sum is still signal. Defaults: 1e-6 for float32, 1e-10 for float64 (_CURVATURE_EPS). Pass a value to override.

State#

lbfgs_prepare_state() allocates, initializes and validates everything in one call; calling it again is how you reset. Nothing is allocated per step, so the step stays capturable in a CUDA graph.

LBFGSState is a plain dataclass, so you can also build one from buffers you already own and check it with LBFGSState.validate(), which compares shapes, dtypes and device without touching the GPU.

Every array is indexed by its owning entity first: a per-degree-of-freedom buffer leads with num_packed, a per-system one with num_systems, and the history depth is always the trailing axis. That is what lets a caller select a subset of systems, or concatenate two states, with an ordinary gather along dimension zero – the operation batched relaxation drivers perform when they retire converged systems and admit replacements.

Writing num_packed for the degrees of freedom, num_systems for the independent systems and history_size for the stored pairs:

Buffer	Shape	dtype
`x_base`, `force_base`, `direction`	`(num_packed,)`	vec3f/vec3d
`s_history`, `y_history`	`(num_packed, history)`	vec3f/vec3d
`ys`, `yy`, `alpha_hist`, `beta_hist`	`(num_systems, history)`	float32/float64
`ss`, `gg`	`(num_systems,)`	float32/float64
`d0`, `dmax`, `dquad`, `alpha_step`	`(num_systems,)`	float32/float64
`iteration`, `end`	`(num_systems,)`	int32
`n_loop`, `history_count`	`(num_systems,)`	int32

history abbreviates history_size in the table above only.

Building one by hand means matching the initial contents: zero everything, then alpha_step to one and iteration to minus one (the “never evaluated” marker).

Variable-cell relaxation adds an LBFGSCellState from lbfgs_prepare_cell_state(), for the packed path with num_packed = num_atoms + 2 * num_systems:

Buffer	Shape	dtype
`ref_cell`, `ref_cell_inv`, `phi`, `phi_inv`	`(num_systems,)`	mat33f/mat33d
`d_phi`	`(num_systems,)`	mat33f/mat33d
`kappa`	`(num_systems,)`	float32/float64
`cell_dof_a/b`, `cell_force_a/b`	`(num_systems,)`	vec3f/vec3d
`ext_positions`, `ext_forces`	`(num_packed,)`	vec3f/vec3d
`ext_batch_idx`	`(num_packed,)`	int32
`ext_atom_ptr`	`(num_systems + 1,)`	int32

kappa follows the coordinate precision because it scales matrices. ref_cell, ref_cell_inv and kappa are the chart; the rest is scratch.

The variable-cell contract#

The one authoritative statement of these rules; the bindings and the user guide defer here.

1. Align the cell once, as FIRE2 requires.

Prepare geometry with align_cell(), which aligns the cell and rotates positions consistently. The aligned matrix has zeros strictly above the diagonal. Torch and JAX reference-cell setup validates aligned input. Raw Warp setup and subsequent optimization steps assume alignment; the six-component cell updates preserve it.

2. Six components, the same six FIRE2 packs.

(0,0), (1,0), (2,0) | (1,1), (2,1), (2,2) – two vec3 entries per system, in that order, matching pack_positions_with_cell(). The strictly-upper entries are not represented, so the cell cannot drift into a rotation. This is why rule 1 is required rather than advisory: in an unaligned frame those three entries are not the redundant ones.

3. The reference cell is fixed for the whole relaxation.

FIRE2 packs the cell itself; L-BFGS packs the deformation gradient Phi = H H_ref^-1 against a fixed reference, scaled by kappa (ASE’s UnitCellFilter chart). Stored pairs compare cell coordinates across steps, so those coordinates must mean the same thing at every step. Re-referencing invalidates every stored pair: if you rebuild the chart mid-run, reset the state with lbfgs_prepare_state() in the same breath.

4. The packed topology is derived, not carried.

ext_batch_idx and ext_atom_ptr are a function of atom_ptr alone, so lbfgs_prepare_cell_state() builds them from the ordinary atom pointer you already have. They are fields of the state because the kernels index them, not because they are yours to supply – and when batch membership changes they should be rebuilt from the new topology rather than gathered and concatenated along with the rest of the state, since their values are global offsets that a gather would leave stale.

Ragged batches are unaffected by deriving them: atom_ptr already carries each system’s atom count, so nothing assumes an even split. Each system contributes exactly two packed entries whatever its size, so num_packed = num_atoms + 2 * num_systems holds for any split, and LBFGSCellState.validate() checks it.

5. Empty systems are rejected here, and supported on the coordinate path.

kappa scales the cell coordinate against the atomic ones and is divided into the cell force, so a system with no atoms has no scale to give it. lbfgs_prepare_cell_state() rejects the batch, naming the offending systems, in all three layers. The coordinate-only path is a different case and stays supported: with no cell to scale against, zero degrees of freedom is a no-op.

Memory#

With P degrees of freedom, M systems and m stored pairs – letters used only to keep this formula readable:

bytes = (2m + 3) * 3 * sizeof(dof) * P + (4m + 6) * sizeof(dof) * M
        + 4 * 4 * M

At history_size = 6 that is 180 bytes per degree of freedom with float32 coordinates, 360 with float64. The two histories dominate; 3 to 7 is the usual range.

References

Nocedal, J. “Updating Quasi-Newton Matrices with Limited Storage.” Math. Comp. 35 (1980) 773-782.

Liu, D. C. and Nocedal, J. “On the limited memory BFGS method for large scale optimization.” Math. Program. 45 (1989) 503-528.

Nocedal, J. and Wright, S. J. Numerical Optimization, 2nd ed., chapters 3 and 7.

class nvalchemiops.dynamics.optimizers.lbfgs.LBFGSState(x_base, force_base, direction, s_history, y_history, ys, yy, alpha_hist, beta_hist, ss, gg, d0, dmax, dquad, alpha_step, iteration, end, n_loop, history_count)[source]#

The optimizer’s persistent and scratch arrays, one field per array.

Deliberately transparent: every array a step touches is a named field you can read, plot or checkpoint. Build one with lbfgs_prepare_state(), which allocates and initializes it, or construct it directly from arrays you already own – the package never requires you to use its allocation policy.

Field order is the wire ABI. It is what the PyTorch operator’s mutates_args and the JAX callable’s in_out_argnames are generated from, so reordering these fields reorders those.

Parameters:

x_base (Any)
force_base (Any)
direction (Any)
s_history (Any)
y_history (Any)
ys (Any)
yy (Any)
alpha_hist (Any)
beta_hist (Any)
ss (Any)
gg (Any)
d0 (Any)
dmax (Any)
dquad (Any)
alpha_step (Any)
iteration (Any)
end (Any)
n_loop (Any)
history_count (Any)

x_base, force_base

Last accepted point and the FORCES there. The gradient is -force_base; see the module docstring on the sign convention.

Type:: wp.array, shape (num_packed,), dtype vec3f/vec3d

direction#

Search direction; a descent direction satisfies force_base . d > 0.

Type:: wp.array, shape (num_packed,), dtype vec3f/vec3d

s_history, y_history

Ring buffers of position and gradient differences. The degree of freedom leads; the history slot is the trailing axis.

Type:: wp.array, shape (num_packed, history_size), dtype vec3f/vec3d

ys, yy

Type:: wp.array, shape (num_systems, history_size), dtype float32/float64

alpha_hist#

Type:: wp.array, shape (num_systems, history_size), dtype float32/float64

beta_hist#

Per-slot curvature products and two-loop coefficients.

Type:: wp.array, shape (num_systems, history_size), dtype float32/float64

ss, gg

Squared norms of the newest s and of the packed force.

Type:: wp.array, shape (num_systems,), dtype float32/float64

d0, dmax, dquad

Directional derivative and the trust region’s linear and quadratic displacement coefficients.

Type:: wp.array, shape (num_systems,), dtype float32/float64

alpha_step#

Step length, recomputed from the trust region every call.

Type:: wp.array, shape (num_systems,), dtype float32/float64

iteration, end, n_loop, history_count

Per-system control state. There is no status: the optimizer has no terminal state, and deciding when to stop is the caller’s, as it is for FIRE2.

Type:: wp.array, shape (num_systems,), dtype int32

property history_size: int#: Stored curvature pairs.

property num_dofs: int#: Degrees of freedom the optimizer moves.

property num_systems: int#: Independent systems in the batch.

validate()[source]#

Check the fields against each other: shapes, dtypes and device.

Structural only – no device reads – so this is sync-free and safe to call anywhere. Run for you by lbfgs_prepare_state(); call it yourself if you built the state from your own arrays.

Three checks. Only leading dimensions are compared, because the component axis is a framework detail: Warp stores a vec3 array as (num_packed,) while PyTorch and JAX store it as (num_packed, 3) – and likewise s_history is (num_packed, history_size) against (num_packed, history_size, 3). Fields that hold the same kind of thing must share a dtype, which is what catches one array left at the wrong precision. And every field must live on one device.

Sharing a dtype is checked rather than a named one, since the same state is expressed in three frameworks’ type vocabularies. Per-system scalars follow the coordinate precision: float32 for vec3f coordinates and float64 for vec3d coordinates. Every field must live on one device.

Raises:: ValueError – If any field is inconsistent with the rest.
Return type:: None

nvalchemiops.dynamics.optimizers.lbfgs.lbfgs_prepare_state(num_dofs, num_systems, *, dtype=wp.vec3d, history_size=6, device=None)[source]#

Allocate, initialize and validate a complete optimizer state.

Call this once, before the first step. The two fields that do not start at zero are set for you – alpha_step to one and iteration to minus one – which is the part that is easy to get wrong by hand. Calling it again is how you reset.

You are not obliged to use it: LBFGSState is a plain dataclass, so you can build one from arrays you already own and call LBFGSState.validate() yourself.

Parameters:

num_dofs (int) – Degrees of freedom the optimizer moves. On the variable-cell path this is num_atoms + 2 * num_systems, not the atom count.
num_systems (int) – Independent systems in the batch.
dtype (wp.vec3f or wp.vec3d, optional) – Coordinate precision, wp.vec3f or wp.vec3d. Every per-system scalar follows it, so wp.vec3f gives an end-to-end fp32 state and wp.vec3d an end-to-end fp64 one.
history_size (int, optional) – Stored curvature pairs; 3 to 7 is the usual range. Memory is dominated by the two (num_dofs, history_size) history buffers.
device (optional) – Warp device.

Return type:

LBFGSState

nvalchemiops.dynamics.optimizers.lbfgs.lbfgs_step(positions, forces, state, batch_idx, **kwargs)[source]#

Consume one force evaluation and produce the next geometry.

One history update, one algorithmic restart if the direction stops descending, and one maxstep-bounded step. It never stops: testing convergence and ending the loop are yours, as they are for FIRE2.

Parameters:

positions (wp.array, shape (num_dofs,), dtype vec3f/vec3d) – Current geometry and the forces there. Mutated in place.
forces (wp.array, shape (num_dofs,), dtype vec3f/vec3d) – Current geometry and the forces there. Mutated in place.
state (LBFGSState) – From lbfgs_prepare_state(), or built from your own arrays.
batch_idx (wp.array, shape (num_dofs,), dtype int32) – Sorted system index per degree of freedom.
**kwargs – maxstep, curvature_eps and compute_reductions.

Return type:

None

Note

lbfgs_prepare_state() allocates, initializes and validates the whole state in one call; calling it again is how you reset. The arrays remain yours – LBFGSState is a plain dataclass, so you can build one from buffers you already own and check it with LBFGSState.validate().

Variable-cell relaxation maps positions and cell into a single packed coordinate vector, so the two-loop recursion couples them automatically. lbfgs_prepare_cell_state() derives the packed topology (ext_atom_ptr and ext_batch_idx) from the ordinary atom_ptr and the aligned cells; ragged batches work the same way uniform ones do.

class nvalchemiops.dynamics.optimizers.lbfgs.LBFGSCellState(ref_cell, ref_cell_inv, kappa, ext_batch_idx, ext_atom_ptr, phi, phi_inv, d_phi, cell_dof_a, cell_dof_b, cell_force_a, cell_force_b, ext_positions, ext_forces)[source]#

The variable-cell chart and its scratch space, one field per array.

The first five fields are configuration and topology, written once before the first step and read-only thereafter; the rest is scratch the step rewrites every call. Build one with lbfgs_prepare_cell_state() or construct it from your own arrays.

Parameters:

ref_cell (Any)
ref_cell_inv (Any)
kappa (Any)
ext_batch_idx (Any)
ext_atom_ptr (Any)
phi (Any)
phi_inv (Any)
d_phi (Any)
cell_dof_a (Any)
cell_dof_b (Any)
cell_force_a (Any)
cell_force_b (Any)
ext_positions (Any)
ext_forces (Any)

ref_cell, ref_cell_inv

Reference cell H0 and its inverse, from _lbfgs_set_reference_cell().

Type:: wp.array, shape (num_systems,), dtype mat33f/mat33d

kappa#

Cell coordinate scaling, from _lbfgs_cell_kappa(). Matches the coordinate precision, not float64, because it scales matrices.

Type:: wp.array, shape (num_systems,), dtype float32/float64

ext_batch_idx#

Type:: wp.array, shape (num_packed,), dtype int32

ext_atom_ptr#

Packed topology. Build these with the generic batch utilities so ragged batches are expressible.

Type:: wp.array, shape (num_systems + 1,), dtype int32

phi, phi_inv, d_phi

Deformation gradient, its inverse, and its search direction.

Type:: wp.array, shape (num_systems,), dtype mat33f/mat33d

cell_dof_a, cell_dof_b

Type:: wp.array, shape (num_systems,), dtype vec3f/vec3d

cell_force_a, cell_force_b

The cell’s two packed degrees of freedom and their conjugate forces.

Type:: wp.array, shape (num_systems,), dtype vec3f/vec3d

ext_positions, ext_forces

The packed coordinate vector and its forces.

Type:: wp.array, shape (num_packed,), dtype vec3f/vec3d

property num_packed_dofs: int#: Packed degrees of freedom, num_atoms + 2 * num_systems.

property num_systems: int#: Independent systems in the batch.

validate(num_atoms=None)[source]#

Check shapes, dtypes, device and the packed relationship.

Structural only, so no device reads: this confirms that P = num_atoms + 2 * num_systems and that every array is sized for it, but not that ext_atom_ptr holds sensible values – checking that would need a device-to-host sync. As with LBFGSState.validate(), fields of the same kind must share a dtype and the whole state one device. kappa is exempt from the dtype grouping: it scales matrices, so it follows the coordinate precision rather than the float64 of the other scalars.

Parameters:: num_atoms (int, optional) – Atom count. When given, the packed size is checked against it.
Raises:: ValueError – If any field is inconsistent with the rest.
Return type:: None

nvalchemiops.dynamics.optimizers.lbfgs.lbfgs_prepare_cell_state(atom_ptr, cell, *, cell_force_scale=1.0, dtype=wp.vec3d, device=None)[source]#

Build a complete, ready-to-step variable-cell chart from atom topology.

One call. Give it the ordinary atom_ptr you already have and the aligned cells, and it derives the packed topology, captures the reference chart and computes kappa. There is no follow-up call and no state that needs repairing before the first step.

Ragged batches are unaffected by deriving the topology: atom_ptr already carries each system’s atom count, so nothing here assumes an even split. What went away is the caller’s obligation to build ext_atom_ptr and ext_batch_idx by hand.

The cell must already be aligned – see rule 1 of the variable-cell contract. Raw Warp setup assumes this precondition; the Torch and JAX reference-cell setup helpers check it before capturing the reference cell.

Parameters:

atom_ptr (wp.array(dtype=int32), shape (num_systems + 1,)) – CSR-style atom pointer for the batch, the same array the rest of the package takes. Read on the host once, which is what makes the derived topology and the per-system counts available.
cell (wp.array(dtype=mat33), shape (num_systems,)) – Reference lattice per system, lattice vectors in columns, aligned.
cell_force_scale (float, optional) – Multiplier on the atom count in kappa.
dtype (wp.vec3f or wp.vec3d, optional) – Coordinate precision, wp.vec3f or wp.vec3d.
device (optional) – Warp device. Defaults to the device atom_ptr is on.

Returns:

Complete: every field is either filled or zeroed scratch.

Return type:

LBFGSCellState

Raises:

ValueError – If dtype is not a supported vector type, if atom_ptr is not a valid pointer, or if any system has no atoms.

nvalchemiops.dynamics.optimizers.lbfgs.lbfgs_step_coord_cell(positions, forces, cell, stress, state, cell_state, batch_idx, **kwargs)[source]#

Advance one variable-cell step, relaxing coordinates and cell together.

state must be sized for num_atoms + 2 * num_systems degrees of freedom; cell_state carries the chart. Call _lbfgs_set_reference_cell() and _lbfgs_cell_kappa() once before the first step.

Parameters:

positions (wp.array, shape (num_atoms,), dtype vec3f/vec3d) – Cartesian geometry and the forces there, advanced in place. Forces, not gradients.
forces (wp.array, shape (num_atoms,), dtype vec3f/vec3d) – Cartesian geometry and the forces there, advanced in place. Forces, not gradients.
cell (wp.array, shape (num_systems,), dtype mat33f/mat33d) – Cell with lattice vectors as columns, advanced in place, and the Cauchy stress per system, which drives the cell degrees of freedom.
stress (wp.array, shape (num_systems,), dtype mat33f/mat33d) – Cell with lattice vectors as columns, advanced in place, and the Cauchy stress per system, which drives the cell degrees of freedom.
state (LBFGSState) – From lbfgs_prepare_state(), sized for num_atoms + 2 * num_systems degrees of freedom.
cell_state (LBFGSCellState) – From lbfgs_prepare_cell_state(), or built from your own arrays.
batch_idx (wp.array, shape (num_atoms,), dtype int32) – Sorted system index per atom.
**kwargs – maxstep and curvature_eps.

Return type:

None

FIRE2#

Improved FIRE optimizer with adaptive damping and velocity mixing.

FIRE2 Optimizer Kernels#

GPU-accelerated Warp kernels for the FIRE2 (Fast Inertial Relaxation Engine v2) geometry optimizer.

This module provides three highly-fused kernels that implement a complete FIRE2 step in only 3 kernel launches, minimizing Python-side and launch overhead.

The public fire2_update() helper exposes only the reduction, adaptive parameter update, and velocity-mixing phase for callers that need a custom final apply step, such as coupled variable-cell optimization.

FIRE2 ALGORITHM (Guenole et al., 2020)#

Given positions \(\mathbf{r}\), velocities \(\mathbf{v}\), and forces \(\mathbf{F}\):

Half-step velocity update: \(\mathbf{v} \leftarrow \mathbf{v} + \mathbf{F} \Delta t\)
Compute power: \(P = \sum \mathbf{v} \cdot \mathbf{F}\) per system
Adaptive parameter update:
- If \(P > 0\): increment counter, optionally grow \(\Delta t\), shrink \(\alpha\)
- If \(P \leq 0\): reset counter, shrink \(\Delta t\), reset \(\alpha\)
Velocity mixing: \(\mathbf{v} = (1 - \alpha) \mathbf{v} + \alpha \sqrt{(\mathbf{v} \cdot \mathbf{v}) / (\mathbf{F} \cdot \mathbf{F})}\, \mathbf{F}\)
Compute step: \(\Delta \mathbf{r} = \mathbf{v} \Delta t\)
Uphill correction: if \(P \leq 0\): \(\Delta \mathbf{r} = -\frac{1}{2} \Delta t\, \mathbf{v}_\text{mixed}\); \(\mathbf{v} = 0\)
Step clamping + position update + coupled \(\Delta t\) scaling

KERNEL STRUCTURE#

Kernel 1 (_fire2_reduce_only):: Runs-based triple inner-product reduction (vf, v.v, f.f) with deferred half-step computed in registers only (no velocity write).
Kernel 2 (_fire2_fused_mix_maxnorm):: Fuses per-system parameter update, deferred half-step, velocity mixing, and runs-based max-norm reduction into a single launch. Each thread redundantly computes the parameter update for its segment from shared read-only inputs, avoiding inter-thread synchronization.
Kernel 3 (_fire2_clamp_apply_recompute):: Recomputes step from mixed velocities, applies step clamping, position update, deferred velocity zeroing for uphill systems, and coupled dt scaling.

References

Guenole et al. (2020). Comp. Mat. Sci. 175, 109584 (FIRE2)
Bitzek et al. (2006). Phys. Rev. Lett. 97, 170201 (FIRE)

nvalchemiops.dynamics.optimizers.fire2.fire2_step(positions, velocities, forces, batch_idx, alpha, dt, nsteps_inc, vf, v_sumsq, f_sumsq, max_norm, delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, maxstep=0.1, compute_reductions=True)[source]#

Complete FIRE2 optimization step.

Modifies positions, velocities, alpha, dt, and nsteps_inc in-place.

Runs the full FIRE2 scheme of Guenole et al. (2020): the leapfrog half-step is deferred into the reduction, so the per-system power is measured on the post-kick velocity \(\mathbf{v} + \Delta t\,\mathbf{F}\):

\[P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\quad vv = \sum_i \lVert \mathbf{v}_i + \Delta t\,\mathbf{F}_i \rVert^2,\quad ff = \sum_i \mathbf{F}_i \cdot \mathbf{F}_i\]

The half-step and the FIRE2 mixing are then fused (no intermediate velocity write):

\[\mathbf{v} = (1-\alpha)\,\mathbf{v} + \big[(1-\alpha)\,\Delta t + \alpha\sqrt{vv/ff}\big]\,\mathbf{F}\]

Displacement is \(\Delta\mathbf{r} = \Delta t\,\mathbf{v}\) when \(P > 0\); for uphill systems (\(P \le 0\)) the FIRE2 correction \(\Delta\mathbf{r} = -\tfrac{1}{2}\Delta t\,\mathbf{v}\) is applied and the velocity is then zeroed. The step is capped by \(\min(1, \text{maxstep}/\lVert\text{step}\rVert)\) per system, and \(\Delta t\) is scaled by the same factor. This differs from fire_step() by the deferred half-step, the modified mixing coefficient, and the half-back uphill correction (FIRE simply zeroes velocity uphill).

Parameters:

positions (wp.array, shape (N,), dtype vec3f/vec3d) – Atomic positions.
velocities (wp.array, shape (N,), dtype vec3f/vec3d) – Atomic velocities.
forces (wp.array, shape (N,), dtype vec3f/vec3d) – Forces on atoms (read-only).
batch_idx (wp.array, shape (N,), dtype int32) – Sorted system index per atom. Required for non-empty inputs.
alpha (wp.array, shape (M,), dtype float*) – FIRE2 mixing parameter.
dt (wp.array, shape (M,), dtype float*) – Per-system timestep.
nsteps_inc (wp.array, shape (M,), dtype int32) – Consecutive positive-power step counter.
vf (wp.array, shape (M,), dtype float*) – Scratch buffers for reductions. Zeroed internally before each use.
v_sumsq (wp.array, shape (M,), dtype float*) – Scratch buffers for reductions. Zeroed internally before each use.
f_sumsq (wp.array, shape (M,), dtype float*) – Scratch buffers for reductions. Zeroed internally before each use.
max_norm (wp.array, shape (M,), dtype float*) – Scratch buffers for reductions. Zeroed internally before each use.
delaystep (int) – Minimum positive steps before dt growth.
dtgrow (float) – Timestep growth/shrink factors.
dtshrink (float) – Timestep growth/shrink factors.
alphashrink (float) – Alpha decay factor.
alpha0 (float) – Alpha reset value.
tmax (float) – Timestep bounds.
tmin (float) – Timestep bounds.
maxstep (float) – Maximum step magnitude per system.
compute_reductions (bool, default True) – If True, recompute vf/v_sumsq/f_sumsq internally. If False, use the caller-supplied values for the mixing and parameter update. Note that the final maxstep clamp still uses the max_norm produced by this call’s mixing kernel; for a clamp reduced across a caller-side partition, drive the mixing via fire2_update and apply the clamp after reducing max_norm.

Return type:

None

Notes

batch_idx must be sorted; segment reductions assume contiguous atom ranges per system.

Examples

>>> fire2_step(positions, velocities, forces, batch_idx,
...            alpha, dt, nsteps_inc,
...            vf, v_sumsq, f_sumsq, max_norm)

nvalchemiops.dynamics.optimizers.fire2.fire2_update(velocities, forces, batch_idx, alpha, dt, nsteps_inc, vf, v_sumsq, f_sumsq, max_norm, *, delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, compute_max_norm=True, compute_reductions=True)[source]#

Run FIRE2 reduction, parameter update, and velocity mixing only.

This low-level helper updates velocities, alpha, dt, and nsteps_inc in-place, and refreshes vf, v_sumsq, f_sumsq, and, by default, max_norm. It deliberately does not apply positions, clamp the final displacement, or zero velocities on uphill systems.

The FIRE2 power is measured on the deferred half-step velocity, and mixing fuses that half-step with the FIRE2 rule:

\[P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\qquad \mathbf{v} = (1-\alpha)\,\mathbf{v} + \big[(1-\alpha)\,\Delta t + \alpha\sqrt{vv/ff}\big]\,\mathbf{F}\]

where \(vv = \sum_i \lVert \mathbf{v}_i + \Delta t\,\mathbf{F}_i \rVert^2\) and \(ff = \sum_i \mathbf{F}_i \cdot \mathbf{F}_i\). When compute_max_norm is set, max_norm records \(\max_i \lVert \Delta t\,\mathbf{v}_i \rVert\) downhill and \(\max_i \lVert -\tfrac{1}{2}\Delta t\,\mathbf{v}_i \rVert\) uphill for a later clamp.

Parameters:

velocities (wp.array, shape (N,), dtype vec3f/vec3d) – Generalized velocities, modified in-place with FIRE2 velocity mixing.
forces (wp.array, shape (N,), dtype vec3f/vec3d) – Generalized forces. Read-only.
batch_idx (wp.array, shape (N,), dtype int32) – Sorted system index for each generalized degree of freedom. Required for non-empty inputs.
alpha (wp.array, shape (M,), dtype float32/float64) – FIRE2 mixing parameter, modified in-place.
dt (wp.array, shape (M,), dtype float32/float64) – Per-system timestep, modified in-place by FIRE2 growth/shrink rules.
nsteps_inc (wp.array, shape (M,), dtype int32) – Consecutive positive-power step counter, modified in-place.
vf (wp.array, shape (M,), dtype float32/float64) – Scratch buffer for per-system \(\mathbf{v} \cdot \mathbf{F}\) power. Zeroed internally.
v_sumsq (wp.array, shape (M,), dtype float32/float64) – Scratch buffer for per-system \(\mathbf{v} \cdot \mathbf{v}\) after the half-step. Zeroed internally.
f_sumsq (wp.array, shape (M,), dtype float32/float64) – Scratch buffer for per-system \(\mathbf{F} \cdot \mathbf{F}\). Zeroed internally.
max_norm (wp.array, shape (M,), dtype float32/float64) – Scratch buffer for the maximum raw final-step norm after mixing, using \(\Delta t\, \mathbf{v}\) downhill and \(-\frac{1}{2} \Delta t\, \mathbf{v}\) uphill. Zeroed internally when compute_max_norm is true. Custom final apply phases may pass compute_max_norm=False and recompute this buffer for their own displacement definition.
delaystep (int) – Minimum positive-power steps before timestep growth.
dtgrow (float) – Timestep growth/shrink factors.
dtshrink (float) – Timestep growth/shrink factors.
alphashrink (float) – Alpha decay factor after enough positive-power steps.
alpha0 (float) – Alpha reset value for uphill systems.
tmax (float) – Timestep bounds.
tmin (float) – Timestep bounds.
compute_max_norm (bool, default True) – Whether to compute the extended-DOF maximum raw final-step norm into max_norm. Set to false when a custom final apply phase will overwrite max_norm with a different physical displacement norm.
compute_reductions (bool, default True) – If True, recompute vf/v_sumsq/f_sumsq internally from the post-half-step velocities and forces. If False, use the caller-supplied values already in those buffers for the mixing and parameter update (they are not zeroed). max_norm is produced by the mixing kernel and is unaffected by this flag; a caller that needs the displacement clamp reduced across a caller-side partition should leave the final clamp to itself (see compute_max_norm and the Notes).

Return type:

None

Notes

Callers that use this helper directly must finish the FIRE2 step themselves. For systems with \(vf \leq 0\), the final apply phase must use the uphill correction \(\Delta \mathbf{r} = -\frac{1}{2} \Delta t\, \mathbf{v}\) and zero the corresponding velocities after applying the step. Downhill systems use \(\Delta \mathbf{r} = \Delta t\, \mathbf{v}\). Any final maxstep clamp must be applied by the caller and should scale dt consistently with the accepted step.

Thermostat Utilities#

Functions for temperature control and velocity initialization.

nvalchemiops.dynamics.utils.thermostat_utils.compute_kinetic_energy(velocities, masses, kinetic_energy, batch_idx=None, atom_ptr=None, num_systems=1, device=None)[source]#

Compute kinetic energy for single or batched MD systems.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
kinetic_energy (wp.array) – Output array. Same dtype as masses. Shape (1,) for single system, (B,) for batched. Zeroed internally before each use.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
num_systems (int, optional) – Number of systems for batched mode. Default 1.
device (str, optional) – Warp device. If None, inferred from velocities.

Returns:

Kinetic energy (same dtype as masses). Shape (1,) for single, (B,) for batched.

Return type:

wp.array

Example

Single system:

import warp as wp
import numpy as np

velocities = wp.array(np.random.randn(100, 3), dtype=wp.vec3d, device="cuda:0")
masses = wp.array(np.ones(100), dtype=wp.float64, device="cuda:0")

ke = wp.empty(1, dtype=wp.float64, device="cuda:0")
ke = compute_kinetic_energy(velocities, masses, ke)
print(f"Kinetic energy: {ke.numpy()[0]}")

Batched mode with batch_idx:

# 3 systems with different atom counts
batch_idx = wp.array([0]*30 + [1]*40 + [2]*30, dtype=wp.int32, device="cuda:0")
ke = wp.empty(3, dtype=wp.float64, device="cuda:0")
ke = compute_kinetic_energy(
    velocities, masses, ke, batch_idx=batch_idx, num_systems=3
)
# ke.shape = (3,), one KE per system

Batched mode with atom_ptr:

atom_ptr = wp.array([0, 30, 70, 100], dtype=wp.int32, device="cuda:0")
ke = wp.empty(3, dtype=wp.float64, device="cuda:0")
ke = compute_kinetic_energy(
    velocities, masses, ke, atom_ptr=atom_ptr, num_systems=3
)

See also

compute_temperature: Convert kinetic energy to temperature

nvalchemiops.dynamics.utils.thermostat_utils.compute_temperature(kinetic_energy, temperature, num_atoms_per_system)[source]#

Compute instantaneous temperature from kinetic energy.

Temperature is computed as T = 2*KE / (DOF * k_B), where k_B = 1 in natural units (so temperature is in energy units).

Parameters:

kinetic_energy (wp.array) – Pre-computed kinetic energy. Shape (1,) or (B,).
temperature (wp.array) – Output temperature array. Temperature - k_B * T. Shape (1,) or (B,).
num_atoms_per_system (wp.array) – Number of atoms (per system for batched).

Returns:

Temperature in energy units (k_B*T). Shape (1,) or (B,).

Return type:

wp.array

nvalchemiops.dynamics.utils.thermostat_utils.initialize_velocities(velocities, masses, temperature, total_momentum, total_mass, com_velocities, random_seed=42, remove_com=True, batch_idx=None, atom_ptr=None, num_systems=1, device=None)[source]#

Initialize velocities from Maxwell-Boltzmann distribution (in-place).

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
temperature (wp.array(dtype=wp.float32 or wp.float64)) – Target temperature (k_B*T in energy units). Shape (1,) or (B,).
total_momentum (wp.array) – Scratch array for COM removal. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Zeroed internally before each use. Only used when remove_com=True.
total_mass (wp.array) – Scratch array for COM removal. Same scalar dtype as masses. Shape (B,) for batched, (1,) for single. Zeroed internally before each use. Only used when remove_com=True.
com_velocities (wp.array) – Scratch array for COM removal. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Only used when remove_com=True.
random_seed (int, optional) – Random seed for reproducibility. Default: 42.
remove_com (bool, optional) – If True, remove center of mass motion after initialization.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
num_systems (int, optional) – Number of systems. Default 1.
device (str, optional) – Warp device.

Return type:

None

See also

remove_com_motion: Remove center of mass motion
compute_temperature: Compute instantaneous temperature

nvalchemiops.dynamics.utils.thermostat_utils.initialize_velocities_out(masses, temperature, velocities_out, total_momentum, total_mass, com_velocities, random_seed=42, remove_com=True, batch_idx=None, atom_ptr=None, num_systems=1, device=None)[source]#

Initialize velocities from Maxwell-Boltzmann distribution (non-mutating).

Parameters:

masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
temperature (wp.array(dtype=wp.float32 or wp.float64)) – Target temperature (k_B*T in energy units). Shape (1,) or (B,).
velocities_out (wp.array) – Output array for velocities. Shape (N,). Caller must pre-allocate.
total_momentum (wp.array) – Scratch array for COM removal. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Zeroed internally before each use. Only used when remove_com=True.
total_mass (wp.array) – Scratch array for COM removal. Same scalar dtype as masses. Shape (B,) for batched, (1,) for single. Zeroed internally before each use. Only used when remove_com=True.
com_velocities (wp.array) – Scratch array for COM removal. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Only used when remove_com=True.
random_seed (int, optional) – Random seed for reproducibility. Default: 42.
remove_com (bool, optional) – If True, remove center of mass motion after initialization.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
num_systems (int, optional) – Number of systems. Default 1.
device (str, optional) – Warp device.

Returns:

Initialized velocities.

Return type:

wp.array

nvalchemiops.dynamics.utils.thermostat_utils.remove_com_motion(velocities, masses, total_momentum, total_mass, com_velocities, batch_idx=None, atom_ptr=None, num_systems=1, device=None)[source]#

Remove center of mass velocity from the system (in-place).

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,). MODIFIED in-place.
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
total_momentum (wp.array) – Scratch array for momentum accumulation. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Zeroed internally before each use.
total_mass (wp.array) – Scratch array for mass accumulation. Same scalar dtype as masses. Shape (B,) for batched, (1,) for single. Zeroed internally before each use.
com_velocities (wp.array) – Scratch array for COM velocities. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
num_systems (int, optional) – Number of systems. Default 1.
device (str, optional) – Warp device.

Return type:

None

nvalchemiops.dynamics.utils.thermostat_utils.remove_com_motion_out(velocities, masses, total_momentum, total_mass, com_velocities, velocities_out, batch_idx=None, atom_ptr=None, num_systems=1, device=None)[source]#

Remove center of mass velocity from the system (non-mutating).

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,).
masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,).
total_momentum (wp.array) – Scratch array for momentum accumulation. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single. Zeroed internally before each use.
total_mass (wp.array) – Scratch array for mass accumulation. Same scalar dtype as masses. Shape (B,) for batched, (1,) for single. Zeroed internally before each use.
com_velocities (wp.array) – Scratch array for COM velocities. Same vec dtype as velocities. Shape (B,) for batched, (1,) for single.
velocities_out (wp.array) – Output array. Shape (N,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. For batched mode (atomic operations).
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
num_systems (int, optional) – Number of systems. Default 1.
device (str, optional) – Warp device.

Returns:

Velocities with COM motion removed.

Return type:

wp.array

Cell Utilities#

Functions for periodic cell operations and coordinate transformations.

nvalchemiops.dynamics.utils.cell_utils.compute_cell_volume(cells, volumes, device=None)[source]#

Compute cell volume \(V = |\det(cell)|\).

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,) where B is number of systems. Even single systems use shape (1,).
volumes (wp.array) – Output array for volumes. Shape (B,). Caller must pre-allocate.
device (str, optional) – Warp device. If None, inferred from cells.

Returns:

Cell volumes. Shape (B,).

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.compute_cell_inverse(cells, cells_inv, device=None)[source]#

Compute cell inverse matrices for coordinate transformations.

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cells_inv (wp.array) – Output array for inverses. Shape (B,). Caller must pre-allocate.
device (str, optional) – Warp device. If None, inferred from cells.

Returns:

Cell inverse matrices. Shape (B,).

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.compute_strain_tensor(cells, cells_ref_inv, strains, device=None)[source]#

Compute strain tensor from current and reference cells.

The strain tensor \(\varepsilon\) is defined by: \(\text{cell} = (I + \varepsilon) \cdot \text{cell\_ref}\), so \(\varepsilon = \text{cell} \cdot \text{cell\_ref\_inv} - I\).

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Current cell matrices. Shape (B,).
cells_ref_inv (wp.array) – Pre-computed inverse of reference cells. Shape (B,). Caller must pre-compute via compute_cell_inverse.
strains (wp.array) – Output strain tensors. Shape (B,). Caller must pre-allocate.
device (str, optional) – Warp device. If None, inferred from cells.

Returns:

Strain tensors. Shape (B,).

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.apply_strain_to_cell(cells, strains, cells_out, device=None)[source]#

Apply strain tensor to cell: cell_new = (I + strain) @ cell.

Parameters:

cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Current cell matrices. Shape (B,).
strains (wp.array) – Strain tensors to apply. Shape (B,).
cells_out (wp.array) – Output cell matrices. Shape (B,). Caller must pre-allocate.
device (str, optional) – Warp device. If None, inferred from cells.

Returns:

Updated cell matrices. Shape (B,).

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.scale_positions_with_cell(positions, cells_new, cells_old_inv, batch_idx=None, device=None)[source]#

Scale positions when cell changes, maintaining fractional coordinates (in-place).

r_new = cell_new @ cell_old_inv @ r_old

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). MODIFIED in-place.
cells_new (wp.array) – New cell matrices. Shape (B,).
cells_old_inv (wp.array) – Pre-computed inverse of old cell matrices. Shape (B,). Caller must pre-compute via compute_cell_inverse.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device. If None, inferred from positions.

Return type:

None

nvalchemiops.dynamics.utils.cell_utils.scale_positions_with_cell_out(positions, cells_new, cells_old_inv, positions_out, batch_idx=None, device=None)[source]#

Scale positions when cell changes (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,).
cells_new (wp.array) – New cell matrices. Shape (B,).
cells_old_inv (wp.array) – Pre-computed inverse of old cell matrices. Shape (B,). Caller must pre-compute via compute_cell_inverse.
positions_out (wp.array) – Output positions. Shape (N,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device.

Returns:

Scaled positions.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.wrap_positions_to_cell(positions, cells, cells_inv, batch_idx=None, device=None)[source]#

Wrap positions into primary cell [0, 1) in fractional coordinates (in-place).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). MODIFIED in-place.
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cells_inv (wp.array) – Pre-computed inverse of cells. Shape (B,). Caller must pre-compute via compute_cell_inverse.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device.

Return type:

None

nvalchemiops.dynamics.utils.cell_utils.wrap_positions_to_cell_out(positions, cells, cells_inv, positions_out, batch_idx=None, device=None)[source]#

Wrap positions into primary cell (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
cells_inv (wp.array) – Pre-computed inverse of cells. Shape (B,). Caller must pre-compute via compute_cell_inverse.
positions_out (wp.array) – Output positions. Shape (N,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device.

Returns:

Wrapped positions.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.cartesian_to_fractional(positions, cells_inv, fractional, batch_idx=None, device=None)[source]#

Convert Cartesian coordinates to fractional coordinates.

s = cell_inv @ r

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Cartesian positions. Shape (N,).
cells_inv (wp.array) – Pre-computed inverse of cells. Shape (B,). Caller must pre-compute via compute_cell_inverse.
fractional (wp.array) – Output fractional coordinates. Shape (N,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device.

Returns:

Fractional coordinates.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_utils.fractional_to_cartesian(fractional, cells, positions, batch_idx=None, device=None)[source]#

Convert fractional coordinates to Cartesian coordinates.

r = cell @ s

Parameters:

fractional (wp.array(dtype=wp.vec3f or wp.vec3d)) – Fractional coordinates. Shape (N,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
positions (wp.array) – Output Cartesian positions. Shape (N,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device.

Returns:

Cartesian positions.

Return type:

wp.array

Cell Filter Utilities (Variable-Cell Optimization)#

Utilities for packing/unpacking extended arrays for variable-cell optimization.

nvalchemiops.dynamics.utils.cell_filter.align_cell(positions, cell, transform, batch_idx=None, device=None)[source]#

Align cell to lower-triangular form and transform positions accordingly.

This is a one-time preprocessing step before variable-cell optimization. The cell is transformed to the standard lower-triangular form, and positions are rotated to maintain their fractional coordinates.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,). Modified in-place.
cell (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,). Modified in-place.
transform (wp.array(dtype=wp.mat33f or wp.mat33d)) – Scratch array for rotation transform. Shape (B,). Caller must pre-allocate.
batch_idx (wp.array(dtype=wp.int32), optional) – System index for each atom. Shape (N,). If None, assumes single system.
device (str, optional) – Warp device. If None, inferred from positions.

Returns:

(positions, cell) - same arrays, modified in-place for convenience.

Return type:

tuple[wp.array, wp.array]

Example

>>> # Before optimization loop
>>> transform = wp.zeros(1, dtype=wp.mat33d, device=device)
>>> positions, cell = align_cell(positions, cell, transform)

nvalchemiops.dynamics.utils.cell_filter.extend_batch_idx(batch_idx, num_atoms, num_systems, extended_batch_idx, device=None)[source]#

Extend batch_idx to include cell DOFs for variable-cell optimization.

For each system, 2 additional “atoms” (representing the 6 cell DOFs as 2 vec3s) are appended. The extended batch_idx assigns these cell DOFs to their respective systems.

Parameters:

batch_idx (wp.array(dtype=wp.int32)) – Original batch index for atoms. Shape (N,).
num_atoms (int) – Number of atoms (N).
num_systems (int) – Number of systems (B).
extended_batch_idx (wp.array) – Output extended batch index. Shape (N + 2*B,). Caller must pre-allocate.
device (str, optional) – Warp device.

Returns:

Extended batch index. Shape (N + 2*B,).

Return type:

wp.array

Example

>>> # Original: 100 atoms across 2 systems
>>> # Extended: 100 + 4 = 104 "atoms" (2 cell DOFs per system)
>>> ext_batch_idx = wp.zeros(104, dtype=wp.int32, device=device)
>>> extend_batch_idx(batch_idx, num_atoms=100, num_systems=2, extended_batch_idx=ext_batch_idx)

nvalchemiops.dynamics.utils.cell_filter.extend_atom_ptr(atom_ptr, extended_atom_ptr, device=None)[source]#

Extend atom_ptr to include cell DOFs for variable-cell optimization.

Each system gets 2 additional DOFs (representing 6 cell parameters as 2 vec3s), so the CSR pointers are adjusted: extended_atom_ptr[sys] = atom_ptr[sys] + 2*sys.

Parameters:

atom_ptr (wp.array(dtype=wp.int32)) – Original CSR pointers. Shape (B+1,).
extended_atom_ptr (wp.array) – Output extended CSR pointers. Shape (B+1,). Caller must pre-allocate.
device (str, optional) – Warp device.

Returns:

Extended CSR pointers. Shape (B+1,).

Return type:

wp.array

Example

>>> # Original: atom_ptr = [0, 50, 100] (2 systems, 50 atoms each)
>>> # Extended: [0, 52, 104] (50+2=52, 100+4=104)
>>> ext_atom_ptr = wp.zeros(3, dtype=wp.int32, device=device)
>>> extend_atom_ptr(atom_ptr, ext_atom_ptr)

nvalchemiops.dynamics.utils.cell_filter.pack_positions_with_cell(positions, cell, extended, atom_ptr=None, ext_atom_ptr=None, device=None, batch_idx=None)[source]#

Pack atomic positions and cell into extended position array.

Single-system mode (atom_ptr=None):: The extended array has shape (N + 2,) with dtype vec3*, where: - First N entries: atomic positions - Entry N: [a, b*cos(gamma), c1] (first 3 cell parameters) - Entry N+1: [b*sin(gamma), c2, c3] (remaining 3 cell parameters)
Batched mode (atom_ptr provided):: Positions are concatenated across systems, cells have shape (B,). The extended array interleaves each system’s positions with its cell DOFs. Both atom_ptr and ext_atom_ptr must be provided. Optionally pass batch_idx to avoid recomputing it internally.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic positions. Shape (N,) for single system or (total_atoms,) for batched.
cell (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrix (should be lower-triangular from align_cell). Shape (1,) for single system or (B,) for batched.
extended (wp.array) – Output extended array. Caller must pre-allocate. Shape (N+2,) for single, (N+2*B,) for batched.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.
batch_idx (wp.array(dtype=wp.int32), optional) – Sorted system index per atom. Shape (N,). Computed from atom_ptr if not provided.

Returns:

Extended position array.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_filter.pack_velocities_with_cell(velocities, cell_velocity, extended, atom_ptr=None, ext_atom_ptr=None, device=None, batch_idx=None)[source]#

Pack atomic velocities and cell velocity into extended velocity array.

Single-system mode (atom_ptr=None):: Extended array has shape (N + 2,).
Batched mode (atom_ptr provided):: Velocities are concatenated across systems, cell velocities have shape (B,). Both atom_ptr and ext_atom_ptr must be provided. Optionally pass batch_idx to avoid recomputing it internally.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic velocities. Shape (N,) for single system or (total_atoms,) for batched.
cell_velocity (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell velocity matrix. Shape (1,) for single system or (B,) for batched.
extended (wp.array) – Output extended array. Caller must pre-allocate.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.
batch_idx (wp.array(dtype=wp.int32), optional) – Sorted system index per atom. Computed from atom_ptr if not provided.

Returns:

Extended velocity array.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_filter.pack_forces_with_cell(forces, cell_force, extended, atom_ptr=None, ext_atom_ptr=None, device=None, batch_idx=None)[source]#

Pack atomic forces and cell force into extended force array.

Single-system mode (atom_ptr=None):: Extended array has shape (N + 2,).
Batched mode (atom_ptr provided):: Forces are concatenated across systems, cell forces have shape (B,). Both atom_ptr and ext_atom_ptr must be provided. Optionally pass batch_idx to avoid recomputing it internally.

Parameters:

forces (wp.array(dtype=wp.vec3f or wp.vec3d)) – Atomic forces. Shape (N,) for single system or (total_atoms,) for batched.
cell_force (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell force matrix (from stress_to_cell_force). Shape (1,) for single system or (B,) for batched.
extended (wp.array) – Output extended array. Caller must pre-allocate. Shape (N+2,) for single, (N+2*B,) for batched.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.
batch_idx (wp.array(dtype=wp.int32), optional) – Sorted system index per atom. Computed from atom_ptr if not provided.

Returns:

Extended force array.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_filter.pack_masses_with_cell(masses, cell_mass_arr, extended, atom_ptr=None, ext_atom_ptr=None, device=None)[source]#

Pack atomic masses and cell mass into extended mass array.

Single-system mode (atom_ptr=None):: Extended array has shape (N + 2,).
Batched mode (atom_ptr provided):: Masses are concatenated across systems. Cell mass is applied to all systems. Both atom_ptr and ext_atom_ptr must be provided.

Parameters:

masses (wp.array(dtype=wp.float32 or wp.float64)) – Atomic masses. Shape (N,) for single system or (total_atoms,) for batched.
cell_mass_arr (wp.array) – Cell mass as a warp array. Shape (1,) for single system or (B,) for batched. Caller must pre-allocate.
extended (wp.array) – Output extended array. Caller must pre-allocate. Shape (N+2,) for single, (N+2*B,) for batched.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.

Returns:

Extended mass array.

Return type:

wp.array

nvalchemiops.dynamics.utils.cell_filter.unpack_positions_with_cell(extended, positions, cell, num_atoms=None, atom_ptr=None, ext_atom_ptr=None, device=None, batch_idx=None)[source]#

Unpack extended position array to atomic positions and cell.

Single-system mode (atom_ptr=None):: Unpacks extended array of shape (N + 2,) to positions (N,) and cell (1,). Requires num_atoms to be specified.
Batched mode (atom_ptr provided):: Unpacks extended array to concatenated positions (total_atoms,) and cells (B,). Both atom_ptr and ext_atom_ptr must be provided. Optionally pass batch_idx to avoid recomputing it internally.

Parameters:

extended (wp.array(dtype=wp.vec3f or wp.vec3d)) – Extended position array.
positions (wp.array) – Output atomic positions. Caller must pre-allocate.
cell (wp.array) – Output cell matrix. Caller must pre-allocate.
num_atoms (int, optional) – Number of atoms (N). Required for single-system mode.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.
batch_idx (wp.array(dtype=wp.int32), optional) – Sorted system index per atom. Shape (N,). Computed from atom_ptr if not provided.

Returns:

(positions, cell)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.utils.cell_filter.unpack_velocities_with_cell(extended, velocities, cell_velocity, num_atoms=None, atom_ptr=None, ext_atom_ptr=None, device=None, batch_idx=None)[source]#

Unpack extended velocity array to atomic velocities and cell velocity.

Single-system mode (atom_ptr=None):: Unpacks extended array of shape (N + 2,). Requires num_atoms.
Batched mode (atom_ptr provided):: Unpacks to concatenated velocities (total_atoms,) and cell velocities (B,). Both atom_ptr and ext_atom_ptr must be provided. Optionally pass batch_idx to avoid recomputing it internally.

Parameters:

extended (wp.array(dtype=wp.vec3f or wp.vec3d)) – Extended velocity array.
velocities (wp.array) – Output atomic velocities. Caller must pre-allocate.
cell_velocity (wp.array) – Output cell velocity matrix. Caller must pre-allocate.
num_atoms (int, optional) – Number of atoms (N). Required for single-system mode.
atom_ptr (wp.array(dtype=wp.int32), optional) – CSR-style atom pointers. Shape (B+1,). If provided, enables batched mode.
ext_atom_ptr (wp.array(dtype=wp.int32), optional) – Extended atom pointers from extend_atom_ptr(). Shape (B+1,). Required if atom_ptr is provided.
device (str, optional) – Warp device.
batch_idx (wp.array(dtype=wp.int32), optional) – Sorted system index per atom. Computed from atom_ptr if not provided.

Returns:

(velocities, cell_velocity)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.utils.cell_filter.stress_to_cell_force(stress, cell, volume, cell_force, keep_aligned=True, device=None)[source]#

Convert stress tensor to cell force for optimization.

Computes: F_cell = -V * H^{-T} * sigma

This is the “force” on the cell that, when minimized, leads to zero stress (pressure equilibration).

Parameters:

stress (wp.array(dtype=wp.mat33f or wp.mat33d)) – Stress tensor. Shape (B,). Convention: positive values indicate tension; negative values indicate compression.
cell (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices with lattice vectors stored as rows. Shape (B,).
volume (wp.array) – Cell volumes. Shape (B,). Caller must pre-compute via compute_cell_volume.
cell_force (wp.array) – Output cell force matrices. Shape (B,). Caller must pre-allocate.
keep_aligned (bool, default=True) – If True, zero out upper off-diagonal elements [0,1], [0,2], [1,2] of the cell force. These entries are intentionally held fixed because they are outside the lower-triangular six-parameter cell representation.
device (str, optional) – Warp device.

Returns:

Cell force matrices. Shape (B,).

Return type:

wp.array

Notes

The keep_aligned=True behavior intentionally zeros forces on the upper off-diagonal elements of the lower-triangular row-cell matrix:

Cell force structure (keep_aligned=True):
[F00,  0,   0 ]
[F10, F11,  0 ]
[F20, F21, F22]

This prevents the optimizer from introducing rotations that would break the lower-triangular cell representation from align_cell().

Constraint Utilities (SHAKE/RATTLE)#

Holonomic constraint algorithms for bond length constraints.

nvalchemiops.dynamics.utils.constraints.shake_constraints(positions, positions_old, masses, bond_atom_i, bond_atom_j, bond_lengths_sq, max_error, num_iter=10, device=None)[source]#

Apply SHAKE constraints for a fixed number of iterations (in-place).

This function runs a fixed number of SHAKE iterations without convergence checking during the loop. The final error is returned as a wp.array. The caller can check convergence by inspecting the error value.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current positions. Shape (N,). MODIFIED in-place.
positions_old (wp.array(dtype=wp.vec3f or wp.vec3d)) – Positions before unconstrained update. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
bond_lengths_sq (wp.array) – Target bond length squared. Shape (M,).
max_error (wp.array(dtype=wp.float64)) – Scratch array for max error tracking. Shape (1,). Caller must pre-allocate. Zeroed internally between iterations.
num_iter (int, optional) – Number of iterations to run. Default 10. Typical values: 3-20 depending on constraint stiffness.
device (str, optional) – Warp device.

Returns:

Final constraint error \(|r^2_{ij} - d^2_{ij}|\). Shape (1,).

Return type:

wp.array(dtype=wp.float64)

Example

>>> # After unconstrained position update
>>> max_error = wp.empty(1, dtype=wp.float64, device=device)
>>> final_error = shake_constraints(
...     positions, positions_old, masses,
...     bond_i, bond_j, bond_lengths_sq, max_error,
...     num_iter=10
... )

nvalchemiops.dynamics.utils.constraints.shake_iteration(positions, positions_old, masses, bond_atom_i, bond_atom_j, bond_lengths_sq, max_error, device=None)[source]#

Perform single SHAKE iteration (in-place).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current positions. Shape (N,). MODIFIED in-place.
positions_old (wp.array(dtype=wp.vec3f or wp.vec3d)) – Positions before unconstrained update. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
bond_lengths_sq (wp.array) – Target bond length squared. Shape (M,).
max_error (wp.array(dtype=wp.float64)) – Output array to store max error. Shape (1,). Zeroed internally before each use.
device (str, optional) – Warp device.

Returns:

Maximum constraint error \(|r^2_{ij} - d^2_{ij}|\). Shape (1,).

Return type:

wp.array(dtype=wp.float64)

nvalchemiops.dynamics.utils.constraints.shake_constraints_out(positions, positions_old, masses, bond_atom_i, bond_atom_j, bond_lengths_sq, positions_out, max_error, num_iter=10, device=None)[source]#

Apply SHAKE constraints (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current positions. Shape (N,).
positions_old (wp.array(dtype=wp.vec3f or wp.vec3d)) – Positions before unconstrained update. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
bond_lengths_sq (wp.array) – Target bond length squared. Shape (M,).
positions_out (wp.array) – Output positions. Shape (N,). Caller must pre-allocate. Caller must copy input positions into this array before calling (e.g., via wp.copy(positions_out, positions)).
max_error (wp.array(dtype=wp.float64)) – Scratch array for max error tracking. Shape (1,). Caller must pre-allocate. Zeroed internally between iterations.
num_iter (int, optional) – Number of iterations to run. Default 10.
device (str, optional) – Warp device.

Returns:

(positions_out, final_error) final_error is shape (1,)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.utils.constraints.shake_iteration_out(positions, positions_old, masses, bond_atom_i, bond_atom_j, bond_lengths_sq, position_corrections, max_error, device=None)[source]#

Compute SHAKE corrections without applying (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current positions. Shape (N,).
positions_old (wp.array(dtype=wp.vec3f or wp.vec3d)) – Positions before unconstrained update. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
bond_lengths_sq (wp.array) – Target bond length squared. Shape (M,).
position_corrections (wp.array) – Output corrections. Shape (N,). Zeroed internally before each use.
max_error (wp.array(dtype=wp.float64)) – Output max error. Shape (1,). Zeroed internally before each use.
device (str, optional) – Warp device.

Returns:

(position_corrections, max_error) max_error is shape (1,)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.utils.constraints.rattle_constraints(positions, velocities, masses, bond_atom_i, bond_atom_j, max_error, num_iter=10, device=None)[source]#

Apply RATTLE velocity constraints for a fixed number of iterations (in-place).

This function runs a fixed number of RATTLE iterations without convergence checking during the loop. The final error is returned as a wp.array. The caller can check convergence by inspecting the error value.

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current (constrained) positions. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Velocities. Shape (N,). MODIFIED in-place.
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
max_error (wp.array(dtype=wp.float64)) – Scratch array for max error tracking. Shape (1,). Caller must pre-allocate. Zeroed internally between iterations.
num_iter (int, optional) – Number of iterations to run. Default 10.
device (str, optional) – Warp device.

Returns:

Final constraint error \(|v_{ij} \cdot r_{ij}|\). Shape (1,).

Return type:

wp.array(dtype=wp.float64)

nvalchemiops.dynamics.utils.constraints.rattle_iteration(positions, velocities, masses, bond_atom_i, bond_atom_j, max_error, device=None)[source]#

Perform single RATTLE iteration (in-place).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current (constrained) positions. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Velocities. Shape (N,). MODIFIED in-place.
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
max_error (wp.array(dtype=wp.float64)) – Output array to store max error. Shape (1,). Zeroed internally before each use.
device (str, optional) – Warp device.

Returns:

Maximum velocity constraint error \(|v_{ij} \cdot r_{ij}|\). Shape (1,).

Return type:

wp.array(dtype=wp.float64)

nvalchemiops.dynamics.utils.constraints.rattle_constraints_out(positions, velocities, masses, bond_atom_i, bond_atom_j, velocities_out, max_error, num_iter=10, device=None)[source]#

Apply RATTLE velocity constraints (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current (constrained) positions. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Velocities. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
velocities_out (wp.array) – Output velocities. Shape (N,). Caller must pre-allocate. Caller must copy input velocities into this array before calling (e.g., via wp.copy(velocities_out, velocities)).
max_error (wp.array(dtype=wp.float64)) – Scratch array for max error tracking. Shape (1,). Caller must pre-allocate. Zeroed internally between iterations.
num_iter (int, optional) – Number of iterations to run. Default 10.
device (str, optional) – Warp device.

Returns:

(velocities_out, final_error) final_error is shape (1,)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics.utils.constraints.rattle_iteration_out(positions, velocities, masses, bond_atom_i, bond_atom_j, velocity_corrections, max_error, device=None)[source]#

Compute RATTLE corrections without applying (non-mutating).

Parameters:

positions (wp.array(dtype=wp.vec3f or wp.vec3d)) – Current (constrained) positions. Shape (N,).
velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Velocities. Shape (N,).
masses (wp.array) – Atomic masses. Shape (N,).
bond_atom_i (wp.array(dtype=wp.int32)) – First atom index for each bond. Shape (M,).
bond_atom_j (wp.array(dtype=wp.int32)) – Second atom index for each bond. Shape (M,).
velocity_corrections (wp.array) – Output corrections. Shape (N,). Zeroed internally before each use.
max_error (wp.array(dtype=wp.float64)) – Output max error. Shape (1,). Zeroed internally before each use.
device (str, optional) – Warp device.

Returns:

(velocity_corrections, max_error) max_error is shape (1,)

Return type:

tuple[wp.array, wp.array]

nvalchemiops.dynamics: Molecular Dynamics#

Dynamics Module#

Submodules#

Warp-Level Interface#

High-Level Interface#

Integrators#

Velocity Verlet#

Langevin Dynamics#

Nosé-Hoover Chain#

Velocity Rescaling#

NPT/NPH (Isothermal-Isobaric / Isenthalpic-Isobaric)#

Mathematical Formulation#

Mathematical Formulation#

Mathematical Formulation#

Optimizers#

FIRE#

L-BFGS#

L-BFGS Optimizer Kernels#

Calling convention#

Relation to FIRE2#

Precision#

State#

The variable-cell contract#

Memory#

FIRE2#

FIRE2 Optimizer Kernels#

FIRE2 ALGORITHM (Guenole et al., 2020)#

KERNEL STRUCTURE#

Thermostat Utilities#

Cell Utilities#

Cell Filter Utilities (Variable-Cell Optimization)#

Constraint Utilities (SHAKE/RATTLE)#

`nvalchemiops.dynamics`: Molecular Dynamics#