`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 r(t+dt) and velocities to half-step v(t+dt/2). After calling this function, recalculate forces at the new positions, then call velocity_verlet_velocity_finalize().

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: v(t+dt) = v(t+dt/2) + 0.5*a(t+dt)*dt

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 r(t+dt). 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 r(t+dt) and half-step velocities v(t+dt/2) to output arrays. Input arrays are NOT modified.

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 v(t+dt) to output array. Input arrays are NOT modified.

Parameters:

velocities (wp.array(dtype=wp.vec3f or wp.vec3d)) – Half-step velocities v(t+dt/2). Shape (N,).
forces_new (wp.array(dtype=wp.vec3f or wp.vec3d)) – Forces at new positions r(t+dt). 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 v(t+dt).

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), A (position), O (thermostat), A (position).

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 (kT). Shape (1,) for single, (B,) for batched.
friction (wp.array(dtype=wp.float32 or wp.float64)) – Friction coefficient. 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-step update using forces calculated at the new positions.

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.

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 (kT). Shape (1,) for single, (B,) for batched.
friction (wp.array) – Friction coefficient. 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.

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

v += 0.5 * (F/m) * dt

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

r += v * dt

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)[source]#

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

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 (kT). 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) – Scratch array for 2*KE computation. Zeroed internally before each use. 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.

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

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

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)[source]#

Propagate Nosé-Hoover chain and scale velocities (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,).
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 (kT). 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) – Scratch array for 2*KE computation. Zeroed internally before each use. 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.

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 Nosé-Hoover chain:: Q_0 = ndof * kT * tau^2 Q_k = kT * tau^2 for k > 0

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 (kT) per system. Shape (1,) for single system, (num_systems,) for batched.
tau (wp.array) – Time constant 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_ext = KE_particles + PE + KE_chain + PE_chain
where:: KE_chain = sum_k 0.5 * Q_k * eta_dot_k^2 PE_chain = ndof * kT * eta_0 + kT * sum_{k>0} eta_k

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 (kT). 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).

Applies v_i *= scale_factor to all velocities.

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_target / T_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).

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)[source]#

Compute full pressure tensor.

P = (kinetic + virial) / 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. Shape (B,).
cells (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell matrices. Shape (B,).
kinetic_tensors (wp.array(dtype=scalar, ndim=2)) – Scratch array for kinetic tensor accumulation. Shape (B, 9). Zeroed internally before each use.
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.

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_scalar = (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 thermostat half-step for NPT.

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, eta_dots, dt, device=None)[source]#

Perform barostat half-step for NPT ensemble (Martyna-Tobias-Klein equations).

This function updates the cell velocity matrix ḣ based on the pressure difference between the instantaneous and target pressures, coupled with the Nosé-Hoover thermostat chain. 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 of motion:

Isotropic mode:

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

where P is the scalar trace of the pressure tensor (P = Tr(P)/3).

Anisotropic mode:

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

for i ∈ {x, y, z}, with off-diagonal elements remaining zero.

Triclinic mode:

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

for all i, j ∈ {x, y, z}.

The instantaneous pressure includes a kinetic correction:

\[P_{ij}^{inst} = P_{ij}^{virial} + \frac{2 KE}{N_f V} \delta_{ij}\]

param cell_velocities:

Cell velocity matrices ḣ. Shape (B,) where B is batch size. MODIFIED in-place.

type cell_velocities:

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

param pressure_tensors:

Current pressure tensors from virial. Shape (B,). Components ordered as [xx, xy, xz, yx, yy, yz, zx, zy, zz].

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 (scalar P). 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,). See compute_barostat_mass().

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 eta_dots:

Thermostat chain velocities η̇. Shape (B, chain_length). Only eta_dots[:, 0] (first thermostat) couples to barostat.

type eta_dots:

wp.array2d(dtype=scalar)

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. Default: inferred from cell_velocities.

type device:

str, optional

Examples

Isotropic pressure control (most common):

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

Anisotropic (orthorhombic) pressure control:

>>> # Different pressures for x, y, z axes
>>> target_P = wp.array([[1.0, 2.0, 1.5]], dtype=wp.vec3f, device="cuda:0")
>>> npt_barostat_half_step(
...     cell_velocities, pressure_tensors, target_P,
...     volumes, cell_masses, kinetic_energy,
...     num_atoms_per_system, eta_dots, dt=0.001
... )

Full triclinic stress control:

>>> # Full 3x3 stress tensor (9 components)
>>> target_stress = wp.array(
...     [[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]],
...     dtype=vec9f, device="cuda:0"
... )
>>> npt_barostat_half_step(
...     cell_velocities, pressure_tensors, target_stress,
...     volumes, cell_masses, kinetic_energy,
...     num_atoms_per_system, eta_dots, dt=0.001
... )

See also

nph_barostat_half_step: Barostat without thermostat coupling.
npt_velocity_half_step: Velocity update with barostat/thermostat coupling.
compute_barostat_mass: Compute barostat mass from relaxation time.

References

[MTK1994]

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

[SSM2004]

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

Parameters:

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

Return type:

None

nvalchemiops.dynamics.integrators.npt.npt_velocity_half_step(velocities, masses, forces, cell_velocities, volumes, eta_dots, num_atoms, dt, 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 (Nosé-Hoover chain)

Mathematical Formulation#

The velocity equation of motion in NPT is:

\[\dot{v}_i = \frac{F_i}{m_i} - \left(\gamma \cdot \dot{\varepsilon} + \dot{\eta}_1\right) v_i\]

where:

F_i / m_i is the acceleration from forces
γ = 1 + 1/N_f is the coupling factor (N_f = 3N degrees of freedom)
ε̇ is the strain rate from cell velocity
η̇₁ is the first thermostat chain velocity

Isotropic mode (default):: Uses scalar strain rate ε̇ = Tr(ḣ)/V
Anisotropic mode:: Uses diagonal strain rates ε̇_ii = ḣ_ii/V for direction-dependent drag
Triclinic mode:: Uses full strain rate tensor ε̇ = ḣ @ h⁻¹ for full coupling. Requires cells_inv parameter.

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:

Cell velocity matrices ḣ. Shape (B,).

type cell_velocities:

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

param volumes:

Cell volumes. Shape (B,).

type volumes:

wp.array(dtype=scalar)

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:

Inverse cell matrices h⁻¹. Shape (B,). Required for triclinic mode.

type cells_inv:

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

param mode:

Pressure control mode. One of:

"isotropic" (default): Uniform scalar strain rate coupling
"anisotropic": Diagonal strain rate coupling (orthorhombic)
"triclinic": Full tensor strain rate coupling (requires cells_inv)

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, volumes,
...     eta_dots, num_atoms=100, dt=0.001
... )

Triclinic cell:

>>> npt_velocity_half_step(
...     velocities, masses, forces, cell_velocities, volumes,
...     eta_dots, num_atoms=100, dt=0.001,
...     cells_inv=cells_inv, mode="triclinic"
... )

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, batch_idx=None, device=None)[source]#

Update positions for NPT integration.

Parameters:

positions (wp.array) – Particle positions. MODIFIED in-place.
velocities (wp.array) – Particle velocities.
cells (wp.array) – Cell matrices.
cell_velocities (wp.array) – Cell velocity matrices.
dt (wp.array(dtype=scalar)) – Full time step per system. Shape (B,).
cells_inv (wp.array) – Pre-computed cell inverses. Caller must pre-compute via compute_cell_inverse.
batch_idx (wp.array, optional) – System index for each atom.
device (str, optional) – Warp device.

Return type:

None

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

Update cell matrices: h_new = h + dt * ḣ.

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, eta_dots, num_atoms, dt, velocities_out, 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).

Non-mutating version of npt_velocity_half_step() that returns a new array instead of modifying in-place.

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) – System state arrays.
eta_dots (wp.array) – System state arrays.
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) – Inverse cell matrices. Required for triclinic mode.
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, batch_idx=None, device=None, _skip_validation=False)[source]#

Update positions for NPT integration (non-mutating).

Parameters:

cells_inv (wp.array) – Pre-computed cell inverses. Caller must pre-compute via compute_cell_inverse.
positions (array)
velocities (array)
cells (array)
cell_velocities (array)
dt (array)
positions_out (array)
batch_idx (array)
device (str)
_skip_validation (bool)

Returns:

Updated positions.

Return type:

wp.array

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

Update cell matrices (non-mutating).

Returns:

Updated cell matrices.

Return type:

wp.array

Parameters:

cells (array)
cell_velocities (array)
dt (array)
cells_out (array)
device (str)
_skip_validation (bool)

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, kinetic_tensors, num_atoms_per_system, 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

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(dtype=mat33)) – Scratch array for cell inverses. Shape (B,).
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 ḣ 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_{ext})\]

Anisotropic mode:

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

Triclinic mode:

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

param cell_velocities:

Cell velocity matrices ḣ. 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, num_atoms, dt, 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#

The NPH velocity equation of motion:

\[\dot{v}_i = \frac{F_i}{m_i} - \gamma \cdot \dot{\varepsilon} \cdot v_i\]

where γ = 1 + 1/N_f is the coupling factor.

Isotropic mode: Uses scalar strain rate ε̇ = Tr(ḣ)/V

Anisotropic mode: Uses diagonal strain rates ε̇_ii = ḣ_ii/V

Triclinic mode: Uses full strain rate tensor ε̇ = ḣ @ h⁻¹

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:

Cell velocity matrices.

type cell_velocities:

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

param volumes:

Cell volumes.

type volumes:

wp.array(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.

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:

Inverse cell matrices. Required for triclinic mode.

type cells_inv:

wp.array, optional

param mode:

Pressure control mode:

"isotropic": Uniform scalar strain rate coupling
"anisotropic": Diagonal strain rate coupling (orthorhombic)
"triclinic": Full tensor coupling (requires cells_inv)

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, batch_idx=None, device=None)[source]#

Update positions for NPH integration.

Uses same kernel as NPT since position update is identical.

Parameters:

positions (array)
velocities (array)
cells (array)
cell_velocities (array)
dt (array)
cells_inv (array)
batch_idx (array)
device (str)

Return type:

None

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

Update cell matrices for NPH.

Uses same kernel as NPT since cell update is identical.

Parameters:

cells (array)
cell_velocities (array)
dt (array)
device (str)

Return type:

None

nvalchemiops.dynamics.integrators.npt.nph_velocity_half_step_out(velocities, masses, forces, cell_velocities, volumes, num_atoms, dt, velocities_out, 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).

Non-mutating version of nph_velocity_half_step().

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) – System state arrays.
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) – Inverse cell matrices. Required for triclinic mode.
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, batch_idx=None, device=None)[source]#

Update positions for NPH integration (non-mutating).

Returns:

Updated positions.

Return type:

wp.array

Parameters:

positions (array)
velocities (array)
cells (array)
cell_velocities (array)
dt (array)
positions_out (array)
cells_inv (array)
batch_idx (array)
device (str)

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, kinetic_tensors, num_atoms_per_system, 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(dtype=mat33)) – Scratch array for cell inverses. Shape (B,).
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 to give a characteristic pressure relaxation time τ_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) - τ_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 τ_p accordingly for real units.
For isotropic barostat, a single mass controls all cell dimensions.
For anisotropic/triclinic barostat, 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, device=None)[source]#

Compute kinetic energy of cell degrees of freedom.

KE_cell = 0.5 * W * ||ḣ||²_F

Parameters:

cell_velocities (wp.array(dtype=wp.mat33f or wp.mat33d)) – Cell velocity matrices. Shape (B,).
cell_masses (wp.array) – Barostat masses. Shape (B,).
kinetic_energy (wp.array(dtype=scalar)) – Output cell kinetic energy. Shape (B,).
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_ext * V.

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)[source]#

Unified FIRE optimization step with MD integration.

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

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)[source]#

FIRE parameter update and velocity mixing WITHOUT MD integration.

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

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)[source]#

Complete FIRE2 optimization step.

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

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

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)

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 ε is defined by: cell = (I + ε) @ cell_ref So: ε = cell @ 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(γ), c1] (first 3 cell parameters) - Entry N+1: [b*sin(γ), 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 * σ * (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#

Mathematical Formulation#

Optimizers#

FIRE#

FIRE2#

Thermostat Utilities#

Cell Utilities#

Cell Filter Utilities (Variable-Cell Optimization)#

Constraint Utilities (SHAKE/RATTLE)#

`nvalchemiops.dynamics`: Molecular Dynamics#