`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, default 60) – Minimum consecutive positive-power steps before timestep growth.
dtgrow (float, default 1.05) – Timestep growth factor applied when nsteps_inc > delaystep.
dtshrink (float, default 0.75) – Timestep shrink factor applied on uphill steps (\(P \le 0\)).
alphashrink (float, default 0.985) – Alpha decay factor applied after enough positive-power steps.
alpha0 (float, default 0.09) – Alpha reset value for uphill systems (\(P \le 0\)).
tmax (float, default 0.08) – Maximum allowed per-system timestep.
tmin (float, default 0.005) – Minimum allowed per-system timestep.
maxstep (float, default 0.1) – Maximum allowed displacement per atom. Steps larger than this are rescaled by maxstep / max_norm[s] per system.
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

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, default 60) – Minimum consecutive positive-power steps before timestep growth.
dtgrow (float, default 1.05) – Timestep growth factor applied when nsteps_inc > delaystep.
dtshrink (float, default 0.75) – Timestep shrink factor applied on uphill steps (\(P \le 0\)).
alphashrink (float, default 0.985) – Alpha decay factor applied after enough positive-power steps.
alpha0 (float, default 0.09) – Alpha reset value for uphill systems (\(P \le 0\)).
tmax (float, default 0.08) – Maximum allowed per-system timestep.
tmin (float, default 0.005) – Minimum allowed per-system timestep.
maxstep (float, default 0.1) – Maximum allowed physical Cartesian displacement per atom after cell coupling. Steps larger than this are rescaled by maxstep / max_norm[s] per system.
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

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, default 60) – Minimum consecutive positive-power steps before timestep growth.
dtgrow (float, default 1.05) – Timestep growth factor applied when nsteps_inc > delaystep.
dtshrink (float, default 0.75) – Timestep shrink factor applied on uphill steps (\(P \le 0\)).
alphashrink (float, default 0.985) – Alpha decay factor applied after enough positive-power steps.
alpha0 (float, default 0.09) – Alpha reset value for uphill systems (\(P \le 0\)).
tmax (float, default 0.08) – Maximum allowed per-system timestep.
tmin (float, default 0.005) – Minimum allowed per-system timestep.
maxstep (float, default 0.1) – Maximum allowed displacement per packed degree of freedom. Steps larger than this are rescaled by maxstep / max_norm[s] per system.

Return type:

None

Notes

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

L-BFGS Optimizer#

Quasi-Newton relaxation with a maxstep trust region. Each step consumes one force evaluation, updates the history, restarts a non-descending direction and takes one bounded step. As with FIRE2, convergence is the caller’s: there is no tolerance and no terminal status. No energy is read either – the step length comes from the trust region.

Note

Call lbfgs_prepare_state() once to allocate, initialize and validate the whole state; calling it again is how you reset. The tensors remain yours – LBFGSState is a plain dataclass, so you can build one from tensors you already own.

nvalchemiops.torch.lbfgs.lbfgs_prepare_state(num_dofs, num_systems, *, dtype=torch.float64, device=None, history_size=6)[source]#

Allocate, initialize and validate a complete optimizer state.

The three fields that do not start at zero – alpha_step, iteration – are set for you. Calling this again is how you reset. LBFGSState is a plain dataclass, so you can equally build one from tensors you already own and call validate().

Parameters:

num_dofs (int) – Degrees of freedom. On the variable-cell path this is num_atoms + 2 * num_systems.
num_systems (int) – Independent systems in the batch.
dtype (torch.dtype, optional) – Coordinate precision. Every per-system scalar follows it, so this selects an end-to-end fp32 or fp64 state.
device (optional) – Torch device.
history_size (int, optional) – Stored curvature pairs.

Return type:

LBFGSState

nvalchemiops.torch.lbfgs.lbfgs_step_coord(positions, forces, state, batch_idx, *, maxstep=0.2, curvature_eps=None, compute_reductions=True)[source]#

Advance one batched L-BFGS step, consuming one force evaluation.

Mutates positions and every field of state in place.

Parameters:

positions (torch.Tensor, shape (num_atoms, 3)) – Current geometry and the forces there. Forces, not gradients.
forces (torch.Tensor, shape (num_atoms, 3)) – Current geometry and the forces there. Forces, not gradients.
state (LBFGSState) – From lbfgs_prepare_state(), or built from your own tensors.
batch_idx (torch.Tensor, shape (num_atoms,), dtype int32) – Sorted system index per atom.
maxstep (float, optional) – Largest distance any atom may move in one step.
curvature_eps (float, optional) – Relative threshold below which a curvature pair is discarded. Defaults to 1e-6 for float32 coordinates and 1e-10 for float64: ys is accumulated at the coordinate precision, so one value cannot serve both.
compute_reductions (bool, optional) – Set False only if you have already filled the reduction fields of state yourself for this geometry.

Raises:

ValueError – If this call’s inputs are incompatible with the prepared state.

Return type:

None

Variable-cell relaxation maps coordinates and cell into one packed coordinate vector, so the two-loop recursion couples them automatically. Build ext_atom_ptr and ext_batch_idx with extend_atom_ptr() and atom_ptr_to_batch_idx(), which handle ragged batches as well as uniform ones.

nvalchemiops.torch.lbfgs.lbfgs_prepare_cell_state(atom_ptr, cell, *, cell_force_scale=1.0, dtype=torch.float64, device=None, alignment_atol=None)[source]#

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

One call: give it the ordinary atom_ptr and aligned cells, and it derives the packed topology, captures the reference chart and computes kappa. Alignment is checked before the reference cell is captured.

Ragged batches are unaffected – atom_ptr already carries each system’s atom count, so nothing here assumes an even split.

Parameters:

atom_ptr (torch.Tensor, shape (num_systems + 1,), int32) – CSR-style atom pointer for the batch.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Reference lattice per system, aligned.
cell_force_scale (float, optional) – Multiplier on the atom count in kappa.
dtype (torch.dtype, optional) – Coordinate precision.
device (optional) – Defaults to atom_ptr.device.
alignment_atol (float, optional) – Absolute tolerance for entries strictly above the cell diagonal. Defaults to 1e-4 for float32 cells and 1e-8 for float64 cells.

Return type:

LBFGSCellState

nvalchemiops.torch: Dynamics Optimizers#

PyTorch Adapters for nvalchemiops#

Submodules#

FIRE2 Optimizer#

Coordinate-Only Optimization#

Variable-Cell Optimization#

Extended Array Interface#

L-BFGS Optimizer#

`nvalchemiops.torch`: Dynamics Optimizers#