Ctrl+K

nvalchemiops.torch: Dynamics Optimizers#

The dynamics module provides PyTorch bindings for GPU-accelerated geometry optimization algorithms.

Tip

For the underlying framework-agnostic Warp kernels and full MD integrators, see nvalchemiops.dynamics: Molecular Dynamics.

PyTorch Adapters for nvalchemiops#

Thin wrappers that convert PyTorch tensors to Warp arrays, call the Warp-first core API, and manage scratch buffer allocation via PyTorch’s CUDA caching allocator.

Submodules#

fire2

PyTorch adapter for the FIRE2 optimizer (coordinate-only and variable-cell).

segment_ops

Differentiable PyTorch wrappers for segmented reductions and broadcasts with explicit first- and second-order backward support.

FIRE2 Optimizer#

PyTorch adapter for the FIRE2 (Fast Inertial Relaxation Engine v2) geometry optimizer. These functions accept PyTorch tensors, allocate scratch buffers via PyTorch’s CUDA caching allocator, and call the pure-Warp FIRE2 kernels.

Coordinate-Only Optimization#

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

FIRE2 coordinate-only optimization step.

Converts PyTorch tensors to Warp arrays (zero-copy) and delegates to the pure-Warp fire2_step().

The deferred half-step is folded into the reduction, so the per-system power is measured on the post-kick velocity, and the mixing fuses that half-step with the FIRE2 rule:

\[P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\qquad \mathbf{v} = (1-\alpha)\,\mathbf{v} + \left[(1-\alpha)\,\Delta t + \alpha\sqrt{\tfrac{\mathbf{v}\cdot\mathbf{v}} {\mathbf{F}\cdot\mathbf{F}}}\right]\mathbf{F}\]

The displacement is \(\Delta\mathbf{r} = \Delta t\,\mathbf{v}\) downhill and \(-\tfrac{1}{2}\Delta t\,\mathbf{v}\) uphill (\(P \le 0\), velocity then zeroed), capped by \(\min(1, \text{maxstep}/\lVert\Delta\mathbf{r}\rVert)\) per system with \(\Delta t\) scaled by the same factor.

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

Parameters:

  • positions (Tensor, shape (N, 3), dtype float32/float64) – Atomic positions.

  • velocities (Tensor, shape (N, 3), dtype float32/float64) – Atomic velocities.

  • forces (Tensor, shape (N, 3), dtype float32/float64) – Forces on atoms (read-only).

  • batch_idx (Tensor, shape (N,), dtype int32) – Sorted system index per atom. Must be non-decreasing; segmented reductions rely on contiguous atom ranges.

  • alpha (Tensor, shape (M,), dtype float32/float64) – FIRE2 mixing parameter (one per system).

  • dt (Tensor, shape (M,), dtype float32/float64) – Per-system timestep.

  • nsteps_inc (Tensor, shape (M,), dtype int32) – Consecutive positive-power step counter.

  • vf (Tensor, shape (M,), optional) –

    Scratch buffers for per-system reductions. Allocated and zeroed if None; zeroed in-place if provided. Pre-allocate and pass them in tight loops to avoid repeated allocation:

    M = alpha.shape[0]
vf = torch.empty(M, dtype=positions.dtype,
                 device=positions.device)
v_sumsq = torch.empty_like(vf)
f_sumsq = torch.empty_like(vf)
max_norm = torch.empty_like(vf)

  • v_sumsq (Tensor, shape (M,), optional) –

    Scratch buffers for per-system reductions. Allocated and zeroed if None; zeroed in-place if provided. Pre-allocate and pass them in tight loops to avoid repeated allocation:

    M = alpha.shape[0]
vf = torch.empty(M, dtype=positions.dtype,
                 device=positions.device)
v_sumsq = torch.empty_like(vf)
f_sumsq = torch.empty_like(vf)
max_norm = torch.empty_like(vf)

  • f_sumsq (Tensor, shape (M,), optional) –

    Scratch buffers for per-system reductions. Allocated and zeroed if None; zeroed in-place if provided. Pre-allocate and pass them in tight loops to avoid repeated allocation:

    M = alpha.shape[0]
vf = torch.empty(M, dtype=positions.dtype,
                 device=positions.device)
v_sumsq = torch.empty_like(vf)
f_sumsq = torch.empty_like(vf)
max_norm = torch.empty_like(vf)

  • max_norm (Tensor, shape (M,), optional) –

    Scratch buffers for per-system reductions. Allocated and zeroed if None; zeroed in-place if provided. Pre-allocate and pass them in tight loops to avoid repeated allocation:

    M = alpha.shape[0]
vf = torch.empty(M, dtype=positions.dtype,
                 device=positions.device)
v_sumsq = torch.empty_like(vf)
f_sumsq = torch.empty_like(vf)
max_norm = torch.empty_like(vf)

  • delaystep (int) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • dtgrow (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • dtshrink (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • alphashrink (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • alpha0 (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • tmax (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • tmin (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • maxstep (float) – FIRE2 hyperparameters. See fire2_step() for defaults and descriptions.

  • compute_reductions (bool, default True) – If True, recompute vf/v_sumsq/f_sumsq internally. If False, use the caller-supplied values in those buffers instead of recomputing them (they must be provided and are not zeroed); the maxstep clamp still uses this call’s internally-computed max_norm.

Return type:

None

Notes

Default hyperparameters (from the underlying Warp implementation): delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, maxstep=0.1. See nvalchemiops.dynamics.optimizers.fire2.fire2_step() for full descriptions.

For variable-cell optimization (coordinates + cell DOFs), use fire2_step_coord_cell() instead.

Examples

Minimal single-step call:

>>> fire2_step_coord(
...     positions, velocities, forces,
...     batch_idx, alpha, dt, nsteps_inc,
... )

Tight optimization loop with pre-allocated scratch buffers:

>>> M = alpha.shape[0]
>>> vf = torch.empty(M, dtype=positions.dtype, device=positions.device)
>>> v_sumsq = torch.empty_like(vf)
>>> f_sumsq = torch.empty_like(vf)
>>> max_norm = torch.empty_like(vf)
>>> for step in range(num_steps):
...     fire2_step_coord(
...         positions, velocities, forces,
...         batch_idx, alpha, dt, nsteps_inc,
...         vf=vf, v_sumsq=v_sumsq,
...         f_sumsq=f_sumsq, max_norm=max_norm,
...     )

Variable-Cell Optimization#

For optimizing both atomic coordinates and simulation cell parameters simultaneously.

nvalchemiops.torch.fire2.fire2_step_coord_cell(positions, velocities, forces, cell, cell_velocities, cell_force, batch_idx, alpha, dt, nsteps_inc, *, atom_ptr=None, ext_atom_ptr=None, ext_positions=None, ext_velocities=None, ext_forces=None, ext_batch_idx=None, vf=None, v_sumsq=None, f_sumsq=None, max_norm=None, delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, maxstep=0.1, cell_force_scale=1.0, compute_reductions=True)[source]#

FIRE2 variable-cell optimization step.

Performs a FIRE2 step on both atomic coordinates and cell degrees of freedom. Internally packs atomic + cell velocity/force DOFs into an interleaved layout (each system’s atoms followed by its 2 cell vec3s), runs the FIRE2 reduction + mixing phase on those generalized DOFs, unpacks the mixed velocities, and then applies the physically coupled atomic/cell update directly on the caller’s coordinate and cell tensors.

FIRE2 runs on the generalized DOF vectors \(\mathbf{v}\), \(\mathbf{F}\) that stack the atomic velocities/forces and the (scaled) cell velocity/force rows. The power and mixing are taken over those generalized DOFs:

\[P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\qquad \mathbf{v} = (1-\alpha)\,\mathbf{v} + \left[(1-\alpha)\,\Delta t + \alpha\sqrt{\tfrac{\mathbf{v}\cdot\mathbf{v}} {\mathbf{F}\cdot\mathbf{F}}}\right]\mathbf{F}\]

where \(\mathbf{v}\cdot\mathbf{v}\) and \(\mathbf{F}\cdot\mathbf{F}\) sum over each system’s atomic and cell DOFs. The mixed generalized velocity is then applied as a coupled atomic/cell position update: cell DOFs deform the cell while atoms follow the affine cell remap plus their own coordinate motion. The resulting Cartesian atomic displacement is capped by \(\min(1, \text{maxstep}/\lVert\Delta\mathbf{r}\rVert)\) per system (uphill systems, \(P \le 0\), take the \(-\tfrac{1}{2}\Delta t\) correction and zero their velocities), with \(\Delta t\) scaled by the same factor.

The cell must be pre-aligned to upper-triangular form via nvalchemiops.dynamics.utils.cell_filter.align_cell() before the first call.

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

Parameters:

  • positions (Tensor, shape (N, 3), dtype float32/float64) – Atomic positions.

  • velocities (Tensor, shape (N, 3), dtype float32/float64) – Atomic velocities.

  • forces (Tensor, shape (N, 3), dtype float32/float64) – Forces on atoms (read-only).

  • cell (Tensor, shape (M, 3, 3), dtype float32/float64) – Cell matrices (upper-triangular from align_cell()).

  • cell_velocities (Tensor, shape (M, 3, 3), dtype float32/float64) – Cell velocity matrices.

  • cell_force (Tensor, shape (M, 3, 3), dtype float32/float64) – Raw cell force matrices from stress_to_cell_force() (read-only). These are divided by atoms_per_system * cell_force_scale before FIRE2 velocity mixing.

  • batch_idx (Tensor, shape (N,), dtype int32) – Sorted system index per atom.

  • alpha (Tensor, shape (M,), dtype float32/float64) – FIRE2 mixing parameter.

  • dt (Tensor, shape (M,), dtype float32/float64) – Per-system timestep.

  • nsteps_inc (Tensor, shape (M,), dtype int32) – Consecutive positive-power counter.

  • atom_ptr (Tensor, shape (M+1,), dtype int32, optional) – CSR-style atom pointers derived from batch_idx. If None, computed internally each call via batch_idx_to_atom_ptr(). Pre-compute once and pass in tight loops to avoid repeated allocation. See Notes for how to compute.

  • ext_atom_ptr (Tensor, shape (M+1,), dtype int32, optional) – Extended atom pointers (accounts for 2 cell DOFs per system). If None, computed from atom_ptr each call via extend_atom_ptr(). See Notes for how to compute.

  • ext_positions (Tensor, shape (N+2M, 3), optional) – Deprecated and ignored. ext_positions is accepted only for backward compatibility. The coupled variable-cell FIRE2 path no longer updates packed positions directly.

  • ext_velocities (Tensor, shape (N+2M, 3), optional) – Pre-allocated extended working arrays for the FIRE2 generalized-DOF reduction and mixing phase. Allocated if None; contents are overwritten each call.

  • ext_forces (Tensor, shape (N+2M, 3), optional) – Pre-allocated extended working arrays for the FIRE2 generalized-DOF reduction and mixing phase. Allocated if None; contents are overwritten each call.

  • ext_batch_idx (Tensor, shape (N+2M,), dtype int32, optional) – Pre-computed extended batch index (sorted, matching interleaved pack layout). If None, computed from ext_atom_ptr each call via atom_ptr_to_batch_idx(). If provided, assumed correct and reused without recomputation. See Notes for how to compute.

  • vf (Tensor, shape (M,), optional) – Scratch buffers for reductions. Allocated and zeroed if None; zeroed in-place if provided. In this coupled cell adapter, max_norm is the final physical Cartesian atomic displacement norm, recomputed after cell motion is coupled back to the atoms.

  • v_sumsq (Tensor, shape (M,), optional) – Scratch buffers for reductions. Allocated and zeroed if None; zeroed in-place if provided. In this coupled cell adapter, max_norm is the final physical Cartesian atomic displacement norm, recomputed after cell motion is coupled back to the atoms.

  • f_sumsq (Tensor, shape (M,), optional) – Scratch buffers for reductions. Allocated and zeroed if None; zeroed in-place if provided. In this coupled cell adapter, max_norm is the final physical Cartesian atomic displacement norm, recomputed after cell motion is coupled back to the atoms.

  • max_norm (Tensor, shape (M,), optional) – Scratch buffers for reductions. Allocated and zeroed if None; zeroed in-place if provided. In this coupled cell adapter, max_norm is the final physical Cartesian atomic displacement norm, recomputed after cell motion is coupled back to the atoms.

  • delaystep (int) – FIRE2 hyperparameters.

  • dtgrow (float) – FIRE2 hyperparameters.

  • dtshrink (float) – FIRE2 hyperparameters.

  • alphashrink (float) – FIRE2 hyperparameters.

  • alpha0 (float) – FIRE2 hyperparameters.

  • tmax (float) – FIRE2 hyperparameters.

  • tmin (float) – FIRE2 hyperparameters.

  • maxstep (float) – FIRE2 hyperparameters.

  • cell_force_scale (float, default=1.0) – Extra positive multiplier for stress-derived cell-force normalization. Cell forces are divided by atoms_per_system * cell_force_scale.

  • compute_reductions (bool, default True) – If True, recompute vf/v_sumsq/f_sumsq internally. If False, use the caller-supplied values instead (they must be provided and are not zeroed). These reductions are over the generalized (atom + cell) DOFs of each system, so a caller assembling them across a partition must include the (replicated) cell contribution exactly once. The maxstep clamp still uses this call’s internally-recomputed max_norm (the physical Cartesian displacement after cell coupling).

Return type:

None

Notes

Default hyperparameters (from the underlying Warp implementation): delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, maxstep=0.1. See nvalchemiops.dynamics.optimizers.fire2.fire2_step() for full descriptions.

The high-level variable-cell adapter normalizes raw stress-derived cell forces by the number of atoms in each system. cell_force_scale is an extra multiplier on top of that per-system atom-count normalization.

Pre-computing static metadata for tight loops

When batch_idx does not change between steps (fixed system sizes), atom_ptr, ext_atom_ptr, and ext_batch_idx are constant and can be pre-computed once to eliminate per-step allocation and kernel launches:

import warp as wp
from nvalchemiops.batch_utils import (
    atom_ptr_to_batch_idx,
    batch_idx_to_atom_ptr,
)
from nvalchemiops.dynamics.utils.cell_filter import extend_atom_ptr

N, M = positions.shape[0], alpha.shape[0]
N_ext = N + 2 * M
device = positions.device

# 1) atom_ptr from batch_idx  (CSR pointers into atom array)
atom_ptr = torch.zeros(M + 1, dtype=torch.int32, device=device)
atom_counts = torch.zeros(M, dtype=torch.int32, device=device)
batch_idx_to_atom_ptr(
    wp.from_torch(batch_idx, dtype=wp.int32),
    wp.from_torch(atom_counts, dtype=wp.int32),
    wp.from_torch(atom_ptr, dtype=wp.int32),
)

# 2) ext_atom_ptr  (CSR pointers into extended array,
#    each system's range grows by 2 for the cell DOFs)
ext_atom_ptr = torch.zeros(M + 1, dtype=torch.int32, device=device)
extend_atom_ptr(
    wp.from_torch(atom_ptr, dtype=wp.int32),
    wp.from_torch(ext_atom_ptr, dtype=wp.int32),
)

# 3) ext_batch_idx  (sorted system index for extended array)
ext_batch_idx = torch.empty(N_ext, dtype=torch.int32, device=device)
atom_ptr_to_batch_idx(
    wp.from_torch(ext_atom_ptr, dtype=wp.int32),
    wp.from_torch(ext_batch_idx, dtype=wp.int32),
)

Then pass all three on every step:

fire2_step_coord_cell(
    ...,
    atom_ptr=atom_ptr,
    ext_atom_ptr=ext_atom_ptr,
    ext_batch_idx=ext_batch_idx,
)

Extended array layout (interleaved)

The mixing phase places each system’s cell DOFs immediately after its atoms:

[sys0_atom0, ..., sys0_atomK, sys0_cell_row0, sys0_cell_row1,
 sys1_atom0, ..., sys1_atomJ, sys1_cell_row0, sys1_cell_row1, ...]

This ensures that ext_batch_idx is sorted (all DOFs for system 0 precede all DOFs for system 1, etc.), which is required by fire2_step’s segmented reductions.

Examples

Minimal single-step call (all buffers allocated internally):

>>> fire2_step_coord_cell(
...     positions, velocities, forces,
...     cell, cell_velocities, cell_force,
...     batch_idx, alpha, dt, nsteps_inc,
... )

Tight optimization loop with pre-allocated buffers:

>>> # Pre-compute static metadata once
>>> atom_ptr = ...   # see Notes
>>> ext_atom_ptr = ...
>>> ext_batch_idx = ...
>>> N_ext = positions.shape[0] + 2 * alpha.shape[0]
>>> ext_vel = torch.empty(N_ext, 3, dtype=positions.dtype,
...                       device=positions.device)
>>> ext_forces = torch.empty_like(ext_vel)
>>> M = alpha.shape[0]
>>> vf = torch.empty(M, dtype=positions.dtype, device=positions.device)
>>> v_sumsq = torch.empty_like(vf)
>>> f_sumsq = torch.empty_like(vf)
>>> max_norm = torch.empty_like(vf)
>>> for step in range(num_steps):
...     fire2_step_coord_cell(
...         positions, velocities, forces,
...         cell, cell_velocities, cell_force,
...         batch_idx, alpha, dt, nsteps_inc,
...         atom_ptr=atom_ptr,
...         ext_atom_ptr=ext_atom_ptr,
...         ext_velocities=ext_vel,
...         ext_forces=ext_forces,
...         ext_batch_idx=ext_batch_idx,
...         vf=vf, v_sumsq=v_sumsq,
...         f_sumsq=f_sumsq, max_norm=max_norm,
...     )

Extended Array Interface#

For advanced use cases where you manage packed extended arrays directly.

nvalchemiops.torch.fire2.fire2_step_extended(ext_positions, ext_velocities, ext_forces, ext_batch_idx, alpha, dt, nsteps_inc, *, vf=None, v_sumsq=None, f_sumsq=None, max_norm=None, delaystep=60, dtgrow=1.05, dtshrink=0.75, alphashrink=0.985, alpha0=0.09, tmax=0.08, tmin=0.005, maxstep=0.1)[source]#

Run FIRE2 directly on pre-packed extended arrays (no pack/unpack).

This is a lower-level API for callers that maintain persistent extended arrays (positions + cell DOFs interleaved). The caller is responsible for packing data into the extended layout before the first call and unpacking results after the last call (or as needed).

Runs the standard FIRE2 scheme directly on the packed generalized DOFs \(\mathbf{v}\), \(\mathbf{F}\):

\[P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\qquad \mathbf{v} = (1-\alpha)\,\mathbf{v} + \left[(1-\alpha)\,\Delta t + \alpha\sqrt{\tfrac{\mathbf{v}\cdot\mathbf{v}} {\mathbf{F}\cdot\mathbf{F}}}\right]\mathbf{F}\]

with displacement \(\Delta\mathbf{r} = \Delta t\,\mathbf{v}\) downhill and \(-\tfrac{1}{2}\Delta t\,\mathbf{v}\) uphill (\(P \le 0\), velocity then zeroed), capped by \(\min(1, \text{maxstep}/\lVert\Delta\mathbf{r}\rVert)\) per system with \(\Delta t\) scaled by the same factor.

This eliminates the per-step pack/unpack overhead that fire2_step_coord_cell incurs.

Unlike fire2_step_coord_cell(), this function treats the packed DOFs exactly as provided. It does not add the affine atomic remap implied by variable-cell motion, so it should be used only when the caller explicitly wants generic packed-DOF FIRE2 behavior.

Parameters:

  • ext_positions (torch.Tensor, shape (N_ext, 3)) – Extended position array (atoms + cell DOFs interleaved).

  • ext_velocities (torch.Tensor, shape (N_ext, 3)) – Extended velocity array.

  • ext_forces (torch.Tensor, shape (N_ext, 3)) – Extended force array.

  • ext_batch_idx (torch.Tensor, shape (N_ext,), dtype=int32) – System index for each element in the extended arrays.

  • alpha (torch.Tensor, shape (M,)) – FIRE2 mixing parameter per system.

  • dt (torch.Tensor, shape (M,)) – Timestep per system.

  • nsteps_inc (torch.Tensor, shape (M,), dtype=int32) – Consecutive positive-power step counter per system.

  • vf (torch.Tensor or None) – Per-system scratch buffers, shape (M,). Allocated internally if None.

  • v_sumsq (torch.Tensor or None) – Per-system scratch buffers, shape (M,). Allocated internally if None.

  • f_sumsq (torch.Tensor or None) – Per-system scratch buffers, shape (M,). Allocated internally if None.

  • max_norm (torch.Tensor or None) – Per-system scratch buffers, shape (M,). Allocated internally if None.

  • delaystep (int) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • dtgrow (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • dtshrink (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • alphashrink (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • alpha0 (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • tmax (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • tmin (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

  • maxstep (float) – FIRE2 hyperparameters. See fire2_step_coord_cell for details.

Return type:

None

Notes

Modifies ext_positions, ext_velocities, alpha, dt, and nsteps_inc in-place.

On this page
Show Source