`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) – Input velocities (not modified when velocities_out differs).
masses (wp.array) – System state arrays.
forces (wp.array) – System state arrays.
cell_velocities (wp.array) – System state arrays.
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)

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)

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}\). 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)

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) – Force computation function.
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) – Input velocities (not modified when velocities_out differs).
masses (wp.array) – System state arrays.
forces (wp.array) – System state arrays.
cell_velocities (wp.array) – System state arrays.
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)

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) – Force computation function.
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)

FIRE2#

Improved FIRE optimizer with adaptive damping and velocity mixing.

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 upper-triangular form and transform positions accordingly.

This is a one-time preprocessing step before variable-cell optimization. The cell is transformed to the standard upper-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 upper-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 * sigma * (H^{-1})^T

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 compression.
cell (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. 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-triangular off-diagonal elements [0,1], [0,2], [1,2] of the cell force. This is essential to prevent the cell from rotating away from the upper-triangular form established by align_cell(). Only set to False if you know what you’re doing.
device (str, optional) – Warp device.

Returns:

Cell force matrices. Shape (B,).

Return type:

wp.array

Notes

The keep_aligned=True behavior zeros out forces on the upper off-diagonal elements of the 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 upper-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#

FIRE2#

Thermostat Utilities#

Cell Utilities#

Cell Filter Utilities (Variable-Cell Optimization)#

Constraint Utilities (SHAKE/RATTLE)#

`nvalchemiops.dynamics`: Molecular Dynamics#