# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""
FIRE2 Optimizer Kernels
=======================
GPU-accelerated Warp kernels for the FIRE2 (Fast Inertial Relaxation
Engine v2) geometry optimizer.
This module provides three highly-fused kernels that implement a complete
FIRE2 step in only **3 kernel launches**, minimizing Python-side and
launch overhead.
The public :func:`fire2_update` helper exposes only the reduction, adaptive
parameter update, and velocity-mixing phase for callers that need a custom
final apply step, such as coupled variable-cell optimization.
FIRE2 ALGORITHM (Guenole et al., 2020)
=======================================
Given positions :math:`\mathbf{r}`, velocities :math:`\mathbf{v}`, and forces
:math:`\mathbf{F}`:
1. Half-step velocity update:
:math:`\mathbf{v} \leftarrow \mathbf{v} + \mathbf{F} \Delta t`
2. Compute power: :math:`P = \sum \mathbf{v} \cdot \mathbf{F}` per system
3. Adaptive parameter update:
- If :math:`P > 0`: increment counter, optionally grow :math:`\Delta t`,
shrink :math:`\alpha`
- If :math:`P \leq 0`: reset counter, shrink :math:`\Delta t`, reset
:math:`\alpha`
4. Velocity mixing:
:math:`\mathbf{v} = (1 - \alpha) \mathbf{v}
+ \alpha \sqrt{(\mathbf{v} \cdot \mathbf{v}) / (\mathbf{F} \cdot \mathbf{F})}\, \mathbf{F}`
5. Compute step: :math:`\Delta \mathbf{r} = \mathbf{v} \Delta t`
6. Uphill correction:
if :math:`P \leq 0`:
:math:`\Delta \mathbf{r} = -\frac{1}{2} \Delta t\, \mathbf{v}_\text{mixed}`;
:math:`\mathbf{v} = 0`
7. Step clamping + position update + coupled :math:`\Delta t` scaling
KERNEL STRUCTURE
================
Kernel 1 (_fire2_reduce_only):
Runs-based triple inner-product reduction (vf, v.v, f.f) with
deferred half-step computed in registers only (no velocity write).
Kernel 2 (_fire2_fused_mix_maxnorm):
Fuses per-system parameter update, deferred half-step, velocity
mixing, and runs-based max-norm reduction into a single launch.
Each thread redundantly computes the parameter update for its
segment from shared read-only inputs, avoiding inter-thread
synchronization.
Kernel 3 (_fire2_clamp_apply_recompute):
Recomputes step from mixed velocities, applies step clamping,
position update, deferred velocity zeroing for uphill systems,
and coupled dt scaling.
REFERENCES
==========
- Guenole et al. (2020). Comp. Mat. Sci. 175, 109584 (FIRE2)
- Bitzek et al. (2006). Phys. Rev. Lett. 97, 170201 (FIRE)
"""
from __future__ import annotations
import os
from typing import Any
import warp as wp
from nvalchemiops.dynamics.utils.kernel_functions import compute_vf_vv_ff
from nvalchemiops.segment_ops import compute_ept
__all__ = ["fire2_step", "fire2_update"]
# =============================================================================
# Kernel 1: Triple inner-product reduction (deferred half-step)
# =============================================================================
# Tile block size for cooperative reductions
TILE_DIM = int(os.getenv("NVALCHEMIOPS_DYNAMICS_TILE_DIM", 256))
@wp.kernel(enable_backward=False)
def _fire2_reduce_only_kernel(
velocities: wp.array(dtype=Any),
forces: wp.array(dtype=Any),
dt: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
vf: wp.array(dtype=Any),
v_sumsq: wp.array(dtype=Any),
f_sumsq: wp.array(dtype=Any),
N: wp.int32,
elems_per_thread: wp.int32,
):
r"""Triple inner-product reduction with deferred velocity half-step.
Computes three inner products per segment without modifying velocities:
- ``vf[s] = sum(dot(v_upd[i], f[i]) for i where batch_idx[i] == s)``
- ``v_sumsq[s] = sum(dot(v_upd[i], v_upd[i]) for i where batch_idx[i] == s)``
- ``f_sumsq[s] = sum(dot(f[i], f[i]) for i where batch_idx[i] == s)``
where ``v_upd[i] = velocities[i] + forces[i] * dt[batch_idx[i]]``.
The half-step velocity update is computed in registers only and NOT written
back to the velocities array. This deferred write is performed by the
subsequent fused mixing kernel, which algebraically combines the half-step
with the velocity mixing operation.
Launch Grid
-----------
dim = ceil(N / elems_per_thread)
Parameters
----------
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Atomic velocities, read-only (not modified by this kernel).
forces : wp.array, shape (N,), dtype vec3f/vec3d
Forces on atoms.
dt : wp.array, shape (M,), dtype float32/float64
Per-system timestep (scalar dtype matching vector precision).
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per atom in [0, M).
vf : wp.array, shape (M,), dtype float32/float64
OUTPUT: :math:`v_\text{upd} \cdot f` per segment. Zeroed internally before each use.
v_sumsq : wp.array, shape (M,), dtype float32/float64
OUTPUT: :math:`v_\text{upd} \cdot v_\text{upd}` per segment. Zeroed internally before each use.
f_sumsq : wp.array, shape (M,), dtype float32/float64
OUTPUT: :math:`f \cdot f` per segment. Zeroed internally before each use.
N : int32
Total number of atoms.
elems_per_thread : int32
Elements processed per thread (auto-tuned based on array size and SM count).
Notes
-----
- batch_idx must be sorted in non-decreasing order for correctness
- Uses run-length encoding to minimize atomic operations
- Part of the FIRE2 3-kernel optimization strategy
- The deferred half-step approach avoids an intermediate velocity write
"""
t = wp.tid()
start = t * elems_per_thread
if start >= N:
return
end = wp.min(start + elems_per_thread, N)
# First element -- compute v_upd in register, do NOT write back
s_cur = batch_idx[start]
v_upd = velocities[start] + forces[start] * dt[s_cur]
acc_vf, acc_vv, acc_ff = compute_vf_vv_ff(v_upd, forces[start])
for i in range(start + 1, end):
s = batch_idx[i]
v_upd = velocities[i] + forces[i] * dt[s]
val_vf, val_vv, val_ff = compute_vf_vv_ff(v_upd, forces[i])
if s == s_cur:
acc_vf = acc_vf + val_vf
acc_vv = acc_vv + val_vv
acc_ff = acc_ff + val_ff
else:
wp.atomic_add(vf, s_cur, acc_vf)
wp.atomic_add(v_sumsq, s_cur, acc_vv)
wp.atomic_add(f_sumsq, s_cur, acc_ff)
s_cur = s
acc_vf = val_vf
acc_vv = val_vv
acc_ff = val_ff
wp.atomic_add(vf, s_cur, acc_vf)
wp.atomic_add(v_sumsq, s_cur, acc_vv)
wp.atomic_add(f_sumsq, s_cur, acc_ff)
@wp.kernel(enable_backward=False)
def _fire2_reduce_only_tiled_kernel(
velocities: wp.array(dtype=Any),
forces: wp.array(dtype=Any),
dt: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
vf: wp.array(dtype=Any),
v_sumsq: wp.array(dtype=Any),
f_sumsq: wp.array(dtype=Any),
):
"""Triple inner-product reduction with tile reductions (per-atom processing).
Computes three inner products per system using block-level tile reductions:
- vf[s] = sum(dot(v_upd[i], f[i]) for i where batch_idx[i] == s)
- v_sumsq[s] = sum(dot(v_upd[i], v_upd[i]) for i where batch_idx[i] == s)
- f_sumsq[s] = sum(dot(f[i], f[i]) for i where batch_idx[i] == s)
where v_upd[i] = velocities[i] + forces[i] * dt[batch_idx[i]].
Launch Grid: dim = [N_atoms], block_dim = TILE_DIM
Notes
-----
- Simpler per-atom processing (no RLE complexity)
- Uses wp.tile() and wp.tile_sum() for cooperative reduction
- Reduces atomics from N to N/TILE_DIM per system
"""
atom_idx = wp.tid()
system_id = batch_idx[atom_idx]
# Compute deferred half-step in register only
v_upd = velocities[atom_idx] + forces[atom_idx] * dt[system_id]
# Compute local contributions
local_vf, local_vv, local_ff = compute_vf_vv_ff(v_upd, forces[atom_idx])
# Convert to tiles for block-level reduction
t_vf = wp.tile(local_vf)
t_vv = wp.tile(local_vv)
t_ff = wp.tile(local_ff)
# Cooperative sum within block
s_vf = wp.tile_sum(t_vf)
s_vv = wp.tile_sum(t_vv)
s_ff = wp.tile_sum(t_ff)
# Extract scalar values from tile sums
sum_vf = s_vf[0]
sum_vv = s_vv[0]
sum_ff = s_ff[0]
# Only first thread in block writes (3 atomics per block)
if atom_idx % TILE_DIM == 0:
wp.atomic_add(vf, system_id, sum_vf)
wp.atomic_add(v_sumsq, system_id, sum_vv)
wp.atomic_add(f_sumsq, system_id, sum_ff)
# =============================================================================
# Kernel 2: Fused param update + deferred halfstep + mix + max-norm
# =============================================================================
@wp.kernel(enable_backward=False)
def _fire2_fused_mix_maxnorm_kernel(
velocities: wp.array(dtype=Any),
forces: wp.array(dtype=Any),
dt: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
vf: wp.array(dtype=Any),
v_sumsq: wp.array(dtype=Any),
f_sumsq: wp.array(dtype=Any),
alpha: wp.array(dtype=Any),
nsteps_inc: wp.array(dtype=wp.int32),
max_norm: wp.array(dtype=Any),
compute_max_norm: wp.bool,
N: wp.int32,
elems_per_thread: wp.int32,
delaystep: wp.int32,
dtgrow: Any,
dtshrink: Any,
alphashrink: Any,
alpha0: Any,
tmax: Any,
tmin: Any,
):
r"""Fused adaptive parameter update, deferred half-step, velocity mixing, and max-norm reduction.
This kernel performs four operations in a single launch:
1. **Adaptive parameter update** (per-segment, redundantly computed):
- If ``vf[s] > 0`` (downhill): increment counter, optionally grow dt, shrink alpha
- If ``vf[s] <= 0`` (uphill): reset counter, shrink dt, reset alpha
2. **Deferred half-step + velocity mixing** (algebraically combined):
``v[i] = mix_a * v[i] + (mix_a * dt_old + mix_b) * f[i]``
where ``mix_a = 1 - alpha``, ``mix_b = alpha * sqrt(v.v / f.f)``
3. **State updates** (first atom per segment writes):
Updates ``alpha[s]``, ``dt[s]``, and ``nsteps_inc[s]``
4. **Max-norm reduction** (run-length encoded):
Computes ``max_norm[s] = max(||step[i]|| for i where batch_idx[i] == s)``
where step depends on uphill/downhill status
Each thread redundantly computes the per-system parameter update from shared
read-only inputs (vf, v_sumsq, f_sumsq), avoiding inter-thread synchronization.
Launch Grid
-----------
dim = ceil(N / elems_per_thread)
Parameters
----------
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Atomic velocities, modified in-place. Must hold pre-halfstep values
(kernel 1 did not modify them).
forces : wp.array, shape (N,), dtype vec3f/vec3d
Forces on atoms (read-only).
dt : wp.array, shape (M,), dtype float32/float64
Per-system timestep. Modified in-place by first atom per segment.
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per atom in [0, M).
vf : wp.array, shape (M,), dtype float32/float64
:math:`v \cdot f` inner product per segment from kernel 1 (read-only).
v_sumsq : wp.array, shape (M,), dtype float32/float64
:math:`v \cdot v` inner product per segment from kernel 1 (read-only).
f_sumsq : wp.array, shape (M,), dtype float32/float64
:math:`f \cdot f` inner product per segment from kernel 1 (read-only).
alpha : wp.array, shape (M,), dtype float32/float64
FIRE2 mixing parameter. Modified in-place by first atom per segment.
nsteps_inc : wp.array, shape (M,), dtype int32
Consecutive positive-power step counter. Modified by first atom per segment.
max_norm : wp.array, shape (M,), dtype float32/float64
OUTPUT: Maximum step norm per segment. Zeroed internally before each use.
compute_max_norm : bool
If true, compute and write ``max_norm``. If false, skip the max-norm
reduction while still applying the FIRE2 velocity and state updates.
N : int32
Total number of atoms.
elems_per_thread : int32
Elements processed per thread (auto-tuned based on array size).
delaystep : int32
Minimum consecutive positive steps before dt growth.
dtgrow : float32/float64
Timestep growth factor (typically 1.05).
dtshrink : float32/float64
Timestep shrink factor (typically 0.75).
alphashrink : float32/float64
Alpha decay factor (typically 0.985).
alpha0 : float32/float64
Alpha reset value (typically 0.09).
tmax : float32/float64
Maximum allowed timestep.
tmin : float32/float64
Minimum allowed timestep.
Notes
-----
- batch_idx must be sorted in non-decreasing order
- Only the first atom in each segment writes to alpha, dt, nsteps_inc
- Parameter updates are computed redundantly by each thread to avoid synchronization
- The algebraic combination of half-step and mixing eliminates one velocity write
- For uphill systems (vf <= 0), step norm uses factor -0.5 for the correction
"""
t = wp.tid()
start = t * elems_per_thread
if start >= N:
return
end = wp.min(start + elems_per_thread, N)
s_cur = batch_idx[start]
# --- Redundant param-update computation for the first segment ---
_vf = vf[s_cur]
_vv = v_sumsq[s_cur]
_ff = f_sumsq[s_cur]
_a = alpha[s_cur]
_dt = dt[s_cur]
dt_old = _dt # pre-update dt for the deferred half-step
zero = type(_dt)(0.0)
one = type(_dt)(1.0)
w_inc = _vf > zero
if w_inc:
_nsi = nsteps_inc[s_cur] + 1
if _nsi > delaystep:
_dt = wp.min(dtgrow * _dt, tmax)
_a = alphashrink * _a
else:
_nsi = 0
_a = alpha0
_dt = wp.max(dtshrink * _dt, tmin)
# First atom per segment writes updated params
if start == 0 or batch_idx[start - 1] != s_cur:
alpha[s_cur] = _a
dt[s_cur] = _dt
nsteps_inc[s_cur] = _nsi
if _ff > zero:
ratio = wp.sqrt(_vv / _ff)
else:
ratio = zero
mix_a = one - _a
mix_b = _a * ratio
w_dec = not w_inc
# --- Process first element: deferred halfstep + mix (algebraic combo) ---
f_coeff = mix_a * dt_old + mix_b
velocities[start] = mix_a * velocities[start] + f_coeff * forces[start]
max_val = zero
if compute_max_norm:
if w_dec:
max_val = wp.length(type(_dt)(-0.5) * _dt * velocities[start])
else:
max_val = wp.length(_dt * velocities[start])
for i in range(start + 1, end):
s = batch_idx[i]
if s != s_cur:
# Flush max_norm for previous segment
if compute_max_norm:
wp.atomic_max(max_norm, s_cur, max_val)
s_cur = s
# --- Redundant param-update computation for new segment ---
_vf = vf[s]
_vv = v_sumsq[s]
_ff = f_sumsq[s]
_a = alpha[s]
_dt = dt[s]
dt_old = _dt
w_inc = _vf > zero
if w_inc:
_nsi = nsteps_inc[s] + 1
if _nsi > delaystep:
_dt = wp.min(dtgrow * _dt, tmax)
_a = alphashrink * _a
else:
_nsi = 0
_a = alpha0
_dt = wp.max(dtshrink * _dt, tmin)
if batch_idx[i - 1] != s:
alpha[s] = _a
dt[s] = _dt
nsteps_inc[s] = _nsi
if _ff > zero:
ratio = wp.sqrt(_vv / _ff)
else:
ratio = zero
mix_a = one - _a
mix_b = _a * ratio
w_dec = not w_inc
f_coeff = mix_a * dt_old + mix_b
max_val = type(_dt)(0.0)
# Deferred halfstep + mix (algebraic combo)
velocities[i] = mix_a * velocities[i] + f_coeff * forces[i]
if compute_max_norm:
if w_dec:
norm = wp.length(type(_dt)(-0.5) * _dt * velocities[i])
else:
norm = wp.length(_dt * velocities[i])
max_val = wp.max(max_val, norm)
if compute_max_norm:
wp.atomic_max(max_norm, s_cur, max_val)
# =============================================================================
# Kernel 3: Step recompute + clamping + position update + velocity zeroing
# =============================================================================
@wp.kernel(enable_backward=False)
def _fire2_clamp_apply_recompute_kernel(
positions: wp.array(dtype=Any),
velocities: wp.array(dtype=Any),
dt: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
max_norm: wp.array(dtype=Any),
vf: wp.array(dtype=Any),
maxstep: Any,
):
r"""Step recomputation, clamping, position update, velocity zeroing, and coupled dt scaling.
This kernel performs the final operations of the FIRE2 step:
1. **Step recomputation** from mixed velocities (avoids storing step buffer)
2. **Uphill correction**: For ``vf[s] <= 0``, applies ``step = -0.5 * dt * v``
3. **Step clamping**: Scales step by ``min(1.0, maxstep / max_norm[s])``
4. **Position update**: ``positions[i] += clamped_step``
5. **Velocity zeroing**: Sets ``velocities[i] = 0`` for uphill systems
6. **Coupled dt scaling**: Scales ``dt[s]`` by the same clamping factor
The algorithm:
```
local_dt = dt[s] (snapshot before any thread modifies it)
inv = min(1.0, maxstep / max_norm[s])
if vf[s] <= 0: # uphill
step = -0.5 * local_dt * v[i]
v[i] = 0
else: # downhill
step = local_dt * v[i]
positions[i] += step * inv
if first_atom_in_segment:
dt[s] = local_dt * inv
```
Launch Grid
-----------
dim = N (total atoms)
Parameters
----------
positions : wp.array, shape (N,), dtype vec3f/vec3d
Atomic positions, modified in-place.
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Atomic velocities, modified in-place (zeroed for uphill systems).
dt : wp.array, shape (M,), dtype float32/float64
Per-system timestep, modified in-place by first atom per segment
(clamped proportionally to step scaling).
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per atom in [0, M).
max_norm : wp.array, shape (M,), dtype float32/float64
Maximum step norm per segment from kernel 2.
vf : wp.array, shape (M,), dtype float32/float64
:math:`v \cdot f` inner product per segment from kernel 1. System is uphill if vf[s] <= 0.
maxstep : float32/float64
Maximum allowed step size (FIRE2 hyperparameter).
Notes
-----
- Each thread reads dt[s] before any thread writes to avoid race conditions
- Only the first atom in each segment (batch_idx[i-1] != batch_idx[i]) writes dt[s]
- Velocity zeroing for uphill systems is deferred to this kernel for efficiency
- Coupled dt scaling ensures consistency between step size and timestep
- The -0.5 factor for uphill correction is part of the FIRE2 algorithm
"""
tid = wp.tid()
s = batch_idx[tid]
# Snapshot dt before any thread writes to it (race-condition guard)
local_dt = dt[s]
mn = max_norm[s]
inv = wp.min(type(mn)(1.0), maxstep / mn)
if vf[s] <= type(mn)(0.0):
local_step = type(mn)(-0.5) * local_dt * velocities[tid]
velocities[tid] = type(velocities[tid])()
else:
local_step = local_dt * velocities[tid]
positions[tid] = positions[tid] + local_step * inv
# Only first atom in segment updates dt (idx is sorted)
if tid == 0 or batch_idx[tid - 1] != s:
dt[s] = local_dt * inv
# =============================================================================
# Overloads
# =============================================================================
_T = [wp.float32, wp.float64]
_V = [wp.vec3f, wp.vec3d]
_fire2_reduce_only_overloads = {}
_fire2_reduce_only_tiled_overloads = {}
_fire2_fused_mix_maxnorm_overloads = {}
_fire2_clamp_apply_recompute_overloads = {}
for _t, _v in zip(_T, _V):
_fire2_reduce_only_overloads[_v] = wp.overload(
_fire2_reduce_only_kernel,
[
wp.array(dtype=_v), # velocities
wp.array(dtype=_v), # forces
wp.array(dtype=_t), # dt
wp.array(dtype=wp.int32), # batch_idx
wp.array(dtype=_t), # vf
wp.array(dtype=_t), # v_sumsq
wp.array(dtype=_t), # f_sumsq
wp.int32, # N
wp.int32, # elems_per_thread
],
)
_fire2_reduce_only_tiled_overloads[_v] = wp.overload(
_fire2_reduce_only_tiled_kernel,
[
wp.array(dtype=_v), # velocities
wp.array(dtype=_v), # forces
wp.array(dtype=_t), # dt
wp.array(dtype=wp.int32), # batch_idx
wp.array(dtype=_t), # vf
wp.array(dtype=_t), # v_sumsq
wp.array(dtype=_t), # f_sumsq
],
)
_fire2_fused_mix_maxnorm_overloads[_v] = wp.overload(
_fire2_fused_mix_maxnorm_kernel,
[
wp.array(dtype=_v), # velocities
wp.array(dtype=_v), # forces
wp.array(dtype=_t), # dt
wp.array(dtype=wp.int32), # batch_idx
wp.array(dtype=_t), # vf
wp.array(dtype=_t), # v_sumsq
wp.array(dtype=_t), # f_sumsq
wp.array(dtype=_t), # alpha
wp.array(dtype=wp.int32), # nsteps_inc
wp.array(dtype=_t), # max_norm
wp.bool, # compute_max_norm
wp.int32, # N
wp.int32, # elems_per_thread
wp.int32, # delaystep
_t, # dtgrow
_t, # dtshrink
_t, # alphashrink
_t, # alpha0
_t, # tmax
_t, # tmin
],
)
_fire2_clamp_apply_recompute_overloads[_v] = wp.overload(
_fire2_clamp_apply_recompute_kernel,
[
wp.array(dtype=_v), # positions
wp.array(dtype=_v), # velocities
wp.array(dtype=_t), # dt
wp.array(dtype=wp.int32), # batch_idx
wp.array(dtype=_t), # max_norm
wp.array(dtype=_t), # vf (v.f inner product)
_t, # maxstep
],
)
# =============================================================================
# Public API
# =============================================================================
[docs]
def fire2_step(
# Per-atom arrays (N,)
positions: wp.array,
velocities: wp.array,
forces: wp.array,
batch_idx: wp.array,
# Per-system state (M,)
alpha: wp.array,
dt: wp.array,
nsteps_inc: wp.array,
# Scratch buffers (M,)
vf: wp.array,
v_sumsq: wp.array,
f_sumsq: wp.array,
max_norm: wp.array,
# Hyperparameters (Python scalars)
delaystep: int = 60,
dtgrow: float = 1.05,
dtshrink: float = 0.75,
alphashrink: float = 0.985,
alpha0: float = 0.09,
tmax: float = 0.08,
tmin: float = 0.005,
maxstep: float = 0.1,
compute_reductions: bool = True,
) -> None:
r"""Complete FIRE2 optimization step.
Modifies *positions*, *velocities*, *alpha*, *dt*, and *nsteps_inc* in-place.
Runs the full FIRE2 scheme of Guenole et al. (2020): the leapfrog half-step
is deferred into the reduction, so the per-system power is measured on the
post-kick velocity :math:`\mathbf{v} + \Delta t\,\mathbf{F}`:
.. math::
P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\quad
vv = \sum_i \lVert \mathbf{v}_i + \Delta t\,\mathbf{F}_i \rVert^2,\quad
ff = \sum_i \mathbf{F}_i \cdot \mathbf{F}_i
The half-step and the FIRE2 mixing are then fused (no intermediate velocity
write):
.. math::
\mathbf{v} = (1-\alpha)\,\mathbf{v}
+ \big[(1-\alpha)\,\Delta t + \alpha\sqrt{vv/ff}\big]\,\mathbf{F}
Displacement is :math:`\Delta\mathbf{r} = \Delta t\,\mathbf{v}` when
:math:`P > 0`; for uphill systems (:math:`P \le 0`) the FIRE2 correction
:math:`\Delta\mathbf{r} = -\tfrac{1}{2}\Delta t\,\mathbf{v}` is applied and
the velocity is then zeroed. The step is capped by
:math:`\min(1, \text{maxstep}/\lVert\text{step}\rVert)` per system, and :math:`\Delta t`
is scaled by the same factor. This differs from :func:`~nvalchemiops.dynamics.optimizers.fire.fire_step`
by the deferred half-step, the modified mixing coefficient, and the
half-back uphill correction (FIRE simply zeroes velocity uphill).
Parameters
----------
positions : wp.array, shape (N,), dtype vec3f/vec3d
Atomic positions.
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Atomic velocities.
forces : wp.array, shape (N,), dtype vec3f/vec3d
Forces on atoms (read-only).
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per atom. Required for non-empty inputs.
alpha : wp.array, shape (M,), dtype float*
FIRE2 mixing parameter.
dt : wp.array, shape (M,), dtype float*
Per-system timestep.
nsteps_inc : wp.array, shape (M,), dtype int32
Consecutive positive-power step counter.
vf, v_sumsq, f_sumsq, 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, dtshrink : float
Timestep growth/shrink factors.
alphashrink : float
Alpha decay factor.
alpha0 : float
Alpha reset value.
tmax, tmin : float
Timestep bounds.
maxstep : float
Maximum step magnitude per system.
compute_reductions : bool, default True
If True, recompute ``vf``/``v_sumsq``/``f_sumsq`` internally. If False,
use the caller-supplied values for the mixing and parameter update. Note
that the final ``maxstep`` clamp still uses the ``max_norm`` produced by
this call's mixing kernel; for a clamp reduced across a caller-side
partition, drive the mixing via ``fire2_update`` and apply the clamp
after reducing ``max_norm``.
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)
"""
N = positions.shape[0]
if velocities.shape[0] != N:
raise ValueError(
f"velocities length {velocities.shape[0]} != positions length {N}"
)
if forces.shape[0] != N:
raise ValueError(f"forces length {forces.shape[0]} != positions length {N}")
if batch_idx is not None and batch_idx.shape[0] != N:
raise ValueError(
f"batch_idx length {batch_idx.shape[0]} != positions length {N}"
)
if N == 0:
vf.zero_()
v_sumsq.zero_()
f_sumsq.zero_()
max_norm.zero_()
return
if batch_idx is None:
raise ValueError("batch_idx is required for fire2_step")
vec_dtype = positions.dtype
device = positions.device
fire2_update(
velocities=velocities,
forces=forces,
batch_idx=batch_idx,
alpha=alpha,
dt=dt,
nsteps_inc=nsteps_inc,
vf=vf,
v_sumsq=v_sumsq,
f_sumsq=f_sumsq,
max_norm=max_norm,
delaystep=delaystep,
dtgrow=dtgrow,
dtshrink=dtshrink,
alphashrink=alphashrink,
alpha0=alpha0,
tmax=tmax,
tmin=tmin,
compute_reductions=compute_reductions,
)
# Kernel 3: recompute step + clamp + position update + velocity zeroing
wp.launch(
_fire2_clamp_apply_recompute_overloads[vec_dtype],
dim=N,
inputs=[
positions,
velocities,
dt,
batch_idx,
max_norm,
vf, # vf holds v.f; uphill if <= 0
maxstep,
],
device=device,
)
def fire2_reduce(
velocities: wp.array,
forces: wp.array,
dt: wp.array,
batch_idx: wp.array,
vf: wp.array,
v_sumsq: wp.array,
f_sumsq: wp.array,
) -> None:
r"""Fill the FIRE2 per-system reductions over the half-stepped velocities.
Computes, per system :math:`s`:
.. math::
vf[s] &= \sum_i (\mathbf{v}_i + \mathbf{F}_i \Delta t[s]) \cdot \mathbf{F}_i \\
v_{sumsq}[s] &= \sum_i (\mathbf{v}_i + \mathbf{F}_i \Delta t[s]) \cdot (\mathbf{v}_i + \mathbf{F}_i \Delta t[s]) \\
f_{sumsq}[s] &= \sum_i \mathbf{F}_i \cdot \mathbf{F}_i
— the same reduction ``fire2_update`` performs internally (the
:math:`\mathbf{v} + \mathbf{F} \Delta t` half-step is applied in registers,
velocities are not modified). Exposed so a caller can build the reductions
over a subset of the degrees of freedom and combine them.
Parameters
----------
velocities, forces : wp.array, shape (N,), dtype vec3f/vec3d
Velocities and forces (read-only).
dt : wp.array, shape (M,), dtype float*
Per-system timestep used for the half-step.
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per element.
vf, v_sumsq, f_sumsq : wp.array, shape (M,), dtype float*
Outputs. Zeroed internally before accumulation.
"""
vf.zero_()
v_sumsq.zero_()
f_sumsq.zero_()
N = velocities.shape[0]
if N == 0:
return
if batch_idx is None:
raise ValueError("batch_idx is required for fire2_reduce")
vec_dtype = velocities.dtype
device = velocities.device
sm = max(device.sm_count, 1)
ept = compute_ept(N, sm, True)
dim = (N + ept - 1) // ept
wp.launch(
_fire2_reduce_only_overloads[vec_dtype],
dim=dim,
inputs=[velocities, forces, dt, batch_idx, vf, v_sumsq, f_sumsq, N, ept],
device=device,
)
def fire2_apply_step(
positions: wp.array,
velocities: wp.array,
dt: wp.array,
batch_idx: wp.array,
max_norm: wp.array,
vf: wp.array,
maxstep: float = 0.1,
) -> None:
r"""Apply the final FIRE2 clamp + position update from mixed velocities.
This is the third phase of :func:`fire2_step`, exposed on its own so a
caller can interpose between the velocity mix and the displacement clamp.
For each atom it recomputes the raw step from the (already mixed) velocity
— :math:`\mathbf{v}\,\Delta t` downhill, :math:`-\tfrac{1}{2}\Delta t\,\mathbf{v}`
uphill — clamps it by
:math:`\min(1, \text{maxstep} / \text{max\_norm}[s])`, applies
.. math::
\mathbf{r} \leftarrow \mathbf{r}
+ \min\!\left(1, \tfrac{\text{maxstep}}{\text{max\_norm}[s]}\right)
\Delta\mathbf{r},
scales ``dt`` by the same factor, and zeroes velocities for uphill systems
(:math:`vf[s] \leq 0`).
Pairing with :func:`fire2_update` reproduces :func:`fire2_step` exactly::
fire2_update(velocities, forces, batch_idx, alpha, dt, nsteps_inc,
vf, v_sumsq, f_sumsq, max_norm, ...) # mix, fill max_norm
fire2_apply_step(positions, velocities, dt, batch_idx, max_norm, vf,
maxstep=maxstep)
Splitting the two lets a caller post-process ``max_norm`` (e.g. combine the
per-partition maxima with a MAX reduction) before the clamp so every
partition scales the step identically.
Parameters
----------
positions : wp.array, shape (N,), dtype vec3f/vec3d
Atomic positions, modified in-place.
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Mixed velocities (from ``fire2_update``), modified in-place (zeroed for
uphill systems).
dt : wp.array, shape (M,), dtype float*
Per-system timestep, scaled in-place by the clamp factor.
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index per atom.
max_norm : wp.array, shape (M,), dtype float*
Maximum step norm per system used for the clamp. Supply the
post-processed value here to override the one ``fire2_update`` produced.
vf : wp.array, shape (M,), dtype float*
Per-system :math:`\mathbf{v} \cdot \mathbf{F}`; a system is uphill when
:math:`vf[s] \leq 0`.
maxstep : float
Maximum allowed step size.
"""
N = positions.shape[0]
if N == 0:
return
if batch_idx is None:
raise ValueError("batch_idx is required for fire2_apply_step")
vec_dtype = positions.dtype
device = positions.device
wp.launch(
_fire2_clamp_apply_recompute_overloads[vec_dtype],
dim=N,
inputs=[positions, velocities, dt, batch_idx, max_norm, vf, maxstep],
device=device,
)
[docs]
def fire2_update(
velocities: wp.array,
forces: wp.array,
batch_idx: wp.array,
alpha: wp.array,
dt: wp.array,
nsteps_inc: wp.array,
vf: wp.array,
v_sumsq: wp.array,
f_sumsq: wp.array,
max_norm: wp.array,
*,
delaystep: int = 60,
dtgrow: float = 1.05,
dtshrink: float = 0.75,
alphashrink: float = 0.985,
alpha0: float = 0.09,
tmax: float = 0.08,
tmin: float = 0.005,
compute_max_norm: bool = True,
compute_reductions: bool = True,
) -> None:
r"""Run FIRE2 reduction, parameter update, and velocity mixing only.
This low-level helper updates ``velocities``, ``alpha``, ``dt``, and
``nsteps_inc`` in-place, and refreshes ``vf``, ``v_sumsq``, ``f_sumsq``,
and, by default, ``max_norm``. It deliberately does not apply positions,
clamp the final displacement, or zero velocities on uphill systems.
The FIRE2 power is measured on the deferred half-step velocity, and mixing
fuses that half-step with the FIRE2 rule:
.. math::
P = \sum_i (\mathbf{v}_i + \Delta t\,\mathbf{F}_i) \cdot \mathbf{F}_i,\qquad
\mathbf{v} = (1-\alpha)\,\mathbf{v}
+ \big[(1-\alpha)\,\Delta t + \alpha\sqrt{vv/ff}\big]\,\mathbf{F}
where :math:`vv = \sum_i \lVert \mathbf{v}_i + \Delta t\,\mathbf{F}_i \rVert^2`
and :math:`ff = \sum_i \mathbf{F}_i \cdot \mathbf{F}_i`. When
``compute_max_norm`` is set, ``max_norm`` records
:math:`\max_i \lVert \Delta t\,\mathbf{v}_i \rVert` downhill and
:math:`\max_i \lVert -\tfrac{1}{2}\Delta t\,\mathbf{v}_i \rVert` uphill for a
later clamp.
Parameters
----------
velocities : wp.array, shape (N,), dtype vec3f/vec3d
Generalized velocities, modified in-place with FIRE2 velocity mixing.
forces : wp.array, shape (N,), dtype vec3f/vec3d
Generalized forces. Read-only.
batch_idx : wp.array, shape (N,), dtype int32
Sorted system index for each generalized degree of freedom. Required
for non-empty inputs.
alpha : wp.array, shape (M,), dtype float32/float64
FIRE2 mixing parameter, modified in-place.
dt : wp.array, shape (M,), dtype float32/float64
Per-system timestep, modified in-place by FIRE2 growth/shrink rules.
nsteps_inc : wp.array, shape (M,), dtype int32
Consecutive positive-power step counter, modified in-place.
vf : wp.array, shape (M,), dtype float32/float64
Scratch buffer for per-system :math:`\mathbf{v} \cdot \mathbf{F}` power. Zeroed internally.
v_sumsq : wp.array, shape (M,), dtype float32/float64
Scratch buffer for per-system :math:`\mathbf{v} \cdot \mathbf{v}` after the half-step. Zeroed internally.
f_sumsq : wp.array, shape (M,), dtype float32/float64
Scratch buffer for per-system :math:`\mathbf{F} \cdot \mathbf{F}`. Zeroed internally.
max_norm : wp.array, shape (M,), dtype float32/float64
Scratch buffer for the maximum raw final-step norm after mixing, using
:math:`\Delta t\, \mathbf{v}` downhill and :math:`-\frac{1}{2} \Delta t\, \mathbf{v}` uphill. Zeroed
internally when ``compute_max_norm`` is true. Custom final apply phases
may pass ``compute_max_norm=False`` and recompute this buffer for their
own displacement definition.
delaystep : int
Minimum positive-power steps before timestep growth.
dtgrow, dtshrink : float
Timestep growth/shrink factors.
alphashrink : float
Alpha decay factor after enough positive-power steps.
alpha0 : float
Alpha reset value for uphill systems.
tmax, tmin : float
Timestep bounds.
compute_max_norm : bool, default True
Whether to compute the extended-DOF maximum raw final-step norm into
``max_norm``. Set to false when a custom final apply phase will
overwrite ``max_norm`` with a different physical displacement norm.
compute_reductions : bool, default True
If True, recompute ``vf``/``v_sumsq``/``f_sumsq`` internally from the
post-half-step velocities and forces. If False, use the caller-supplied
values already in those buffers for the mixing and parameter update
(they are not zeroed). ``max_norm`` is produced by the mixing kernel and
is unaffected by this flag; a caller that needs the displacement clamp
reduced across a caller-side partition should leave the final clamp to
itself (see ``compute_max_norm`` and the Notes).
Notes
-----
Callers that use this helper directly must finish the FIRE2 step
themselves. For systems with :math:`vf \leq 0`, the final apply phase must use
the uphill correction :math:`\Delta \mathbf{r} = -\frac{1}{2} \Delta t\, \mathbf{v}`
and zero the corresponding velocities after applying the step. Downhill
systems use :math:`\Delta \mathbf{r} = \Delta t\, \mathbf{v}`. Any final
``maxstep`` clamp must be applied by
the caller and should scale ``dt`` consistently with the accepted step.
"""
N = velocities.shape[0]
if forces.shape[0] != N:
raise ValueError(f"forces length {forces.shape[0]} != velocities length {N}")
if batch_idx is not None and batch_idx.shape[0] != N:
raise ValueError(
f"batch_idx length {batch_idx.shape[0]} != velocities length {N}"
)
M = alpha.shape[0]
if dt.shape[0] != M:
raise ValueError(f"dt length {dt.shape[0]} != alpha length {M}")
if nsteps_inc.shape[0] != M:
raise ValueError(f"nsteps_inc length {nsteps_inc.shape[0]} != alpha length {M}")
vec_dtype = velocities.dtype
device = velocities.device
if compute_reductions:
vf.zero_()
v_sumsq.zero_()
f_sumsq.zero_()
if compute_max_norm or N == 0:
max_norm.zero_()
if N == 0:
return
if batch_idx is None:
raise ValueError("batch_idx is required for FIRE2 updates")
sm = max(device.sm_count, 1)
if compute_reductions:
ept1 = compute_ept(N, sm, True)
dim1 = (N + ept1 - 1) // ept1
wp.launch(
_fire2_reduce_only_overloads[vec_dtype],
dim=dim1,
inputs=[velocities, forces, dt, batch_idx, vf, v_sumsq, f_sumsq, N, ept1],
device=device,
)
ept2 = compute_ept(N, sm, True)
dim2 = (N + ept2 - 1) // ept2
wp.launch(
_fire2_fused_mix_maxnorm_overloads[vec_dtype],
dim=dim2,
inputs=[
velocities,
forces,
dt,
batch_idx,
vf,
v_sumsq,
f_sumsq,
alpha,
nsteps_inc,
max_norm,
compute_max_norm,
N,
ept2,
delaystep,
dtgrow,
dtshrink,
alphashrink,
alpha0,
tmax,
tmin,
],
device=device,
)
