# 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"""
Nosé-Hoover Chain (NHC) Thermostat for NVT Ensemble.
This module implements the Nosé-Hoover chain thermostat following the
Martyna-Tobias-Klein (MTK) equations of motion with time-reversible
integration based on Tuckerman et al.
References
----------
- Martyna, Tobias, Klein. J Chem Phys, 101, 4177 (1994)
- Tuckerman et al. J Phys A: Math Gen, 39, 5629 (2006)
The Nosé-Hoover chain equations of motion:
.. math::
\dot{r}_i = v_i
\dot{v}_i = F_i/m_i - \dot{\eta}_1 v_i
\dot{\eta}_1 = (2 \cdot \text{KE} - N_\text{dof} k T) / Q_1
\dot{\eta}_k = (Q_{k-1} \dot{\eta}_{k-1}^2 - kT) / Q_k \quad \text{for } k > 1
Where:
- :math:`\eta` : thermostat chain positions (unitless)
- :math:`\dot{\eta}` : thermostat chain velocities (1/time)
- :math:`Q` : thermostat chain masses (energy*time^2)
- :math:`N_\text{dof}` : degrees of freedom (typically 3N - 3)
- :math:`kT` : target temperature in energy units (k_B = 1)
BATCH MODE
==========
All functions in this module support three execution modes:
**Single System Mode**::
# Simple position and velocity updates
nhc_velocity_half_step(velocities, forces, masses, dt)
nhc_position_update(positions, velocities, dt)
**Batch Mode with batch_idx** (atomic operations)::
batch_idx = wp.array([0]*N0 + [1]*N1 + [2]*N2, dtype=wp.int32, device="cuda:0")
dt = wp.array([dt0, dt1, dt2], dtype=wp.float64, device="cuda:0")
nhc_velocity_half_step(velocities, forces, masses, dt, batch_idx=batch_idx)
nhc_position_update(positions, velocities, dt, batch_idx=batch_idx)
**Batch Mode with atom_ptr** (sequential per-system)::
atom_ptr = wp.array([0, N0, N0+N1, N0+N1+N2], dtype=wp.int32, device="cuda:0")
nhc_velocity_half_step(velocities, forces, masses, dt, atom_ptr=atom_ptr)
nhc_position_update(positions, velocities, dt, atom_ptr=atom_ptr)
"""
from __future__ import annotations
import os
from typing import Any
import warp as wp
from nvalchemiops.dynamics.utils.launch_helpers import dispatch_family
from nvalchemiops.dynamics.utils.shared_kernels import (
position_update_families,
velocity_kick_families,
)
from nvalchemiops.warp_dispatch import validate_out_array
__all__ = [
# Mutating (in-place) APIs
"nhc_thermostat_chain_update",
"nhc_velocity_half_step",
"nhc_position_update",
"nhc_compute_chain_energy",
# Non-mutating (output) APIs
"nhc_thermostat_chain_update_out",
"nhc_velocity_half_step_out",
"nhc_position_update_out",
# Utility functions
"nhc_compute_masses",
]
# ==============================================================================
# Constants
# ==============================================================================
# Maximum supported chain length (typically 3-5 in practice)
MAX_CHAIN_LENGTH = 8
# Yoshida-Suzuki Integration Weights
# These weights provide 4th-order accurate, time-reversible integration
# for the thermostat chain propagation.
# 3-step Yoshida-Suzuki weights
_YS3_W0 = 1.0 / (2.0 - 2.0 ** (1.0 / 3.0))
_YS3_W1 = 1.0 - 2.0 * _YS3_W0
YOSHIDA_SUZUKI_3 = [_YS3_W0, _YS3_W1, _YS3_W0]
# 5-step Yoshida-Suzuki weights (higher accuracy)
_YS5_W0 = 1.0 / (4.0 - 4.0 ** (1.0 / 3.0))
_YS5_W1 = _YS5_W0
_YS5_W2 = 1.0 - 4.0 * _YS5_W0
YOSHIDA_SUZUKI_5 = [_YS5_W0, _YS5_W1, _YS5_W2, _YS5_W1, _YS5_W0]
# ==============================================================================
# Diagnostic Kernels
# ==============================================================================
# Tile block size for cooperative reductions
TILE_DIM = int(os.getenv("NVALCHEMIOPS_DYNAMICS_TILE_DIM", 256))
@wp.kernel
def _compute_2ke_kernel(
velocities: wp.array(dtype=Any),
masses: wp.array(dtype=Any),
ke2: wp.array(dtype=Any),
):
"""Compute 2*KE = sum(m * v^2) for thermostat forcing.
Launch Grid
-----------
dim = [num_atoms]
"""
atom_idx = wp.tid()
v = velocities[atom_idx]
m = masses[atom_idx]
v_sq = wp.dot(v, v)
wp.atomic_add(ke2, 0, type(ke2[0])(m * v_sq))
@wp.kernel
def _compute_2ke_tiled_kernel(
velocities: wp.array(dtype=Any),
masses: wp.array(dtype=Any),
ke2: wp.array(dtype=Any),
):
"""Compute 2*KE with tile reductions (single system).
Launch Grid: dim = [num_atoms], block_dim = TILE_DIM
"""
atom_idx = wp.tid()
v = velocities[atom_idx]
m = masses[atom_idx]
v_sq = wp.dot(v, v)
local_2ke = type(ke2[0])(m * v_sq)
# Convert to tile for block-level reduction
t = wp.tile(local_2ke)
# Cooperative sum within block
s = wp.tile_sum(t)
# Extract scalar from tile sum
sum_2ke = s[0]
# Only first thread in block writes
if atom_idx % TILE_DIM == 0:
wp.atomic_add(ke2, 0, sum_2ke)
@wp.kernel(enable_backward=False)
def _batch_compute_2ke_kernel(
velocities: wp.array(dtype=Any),
masses: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
ke2: wp.array(dtype=Any),
):
"""Compute 2*KE per system for batched simulations.
Launch Grid
-----------
dim = [num_atoms_total]
"""
atom_idx = wp.tid()
system_id = batch_idx[atom_idx]
v = velocities[atom_idx]
m = masses[atom_idx]
v_sq = wp.dot(v, v)
wp.atomic_add(ke2, system_id, type(ke2[system_id])(m * v_sq))
@wp.kernel
def _nhc_compute_masses_kernel(
ndof: wp.array(dtype=wp.int32),
target_temp: wp.array(dtype=Any),
tau: wp.array(dtype=Any),
masses: wp.array(dtype=Any),
):
"""Compute Nosé-Hoover chain masses (single system).
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. Shape (1,).
target_temp : wp.array(dtype=Any)
Target temperature (kT). Shape (1,).
tau : wp.array(dtype=Any)
Time constant. Shape (1,).
masses : wp.array(dtype=Any)
Chain masses output. Shape (chain_length,).
Launch Grid
-----------
dim = [chain_length]
"""
k = wp.tid()
tau_val = tau[0]
tau_sq = tau_val * tau_val
if k == 0:
masses[k] = type(tau_val)(ndof[0]) * target_temp[0] * tau_sq
else:
masses[k] = target_temp[0] * tau_sq
@wp.kernel
def _batch_nhc_compute_masses_kernel(
ndof: wp.array(dtype=wp.int32),
target_temp: wp.array(dtype=Any),
tau: wp.array(dtype=Any),
masses: wp.array2d(dtype=Any),
):
"""Compute Nosé-Hoover chain masses for batched simulations.
Parameters
----------
ndof : wp.array(dtype=wp.int32)
Number of degrees of freedom per system. Shape (num_systems,).
target_temp : wp.array(dtype=Any)
Target temperature (kT) per system. Shape (num_systems,).
tau : wp.array(dtype=Any)
Time constant per system. Shape (num_systems,).
masses : wp.array2d(dtype=Any)
Chain masses output. Shape (num_systems, chain_length).
Launch Grid
-----------
dim = [num_systems, chain_length]
"""
sys_id, k = wp.tid()
tau_val = tau[sys_id]
tau_sq = tau_val * tau_val
if k == 0:
masses[sys_id, k] = type(tau_val)(ndof[sys_id]) * target_temp[sys_id] * tau_sq
else:
masses[sys_id, k] = target_temp[sys_id] * tau_sq
# ==============================================================================
# Chain Propagation Kernels (Pure Warp Implementation)
# ==============================================================================
@wp.kernel
def _nhc_chain_propagate_kernel(
eta: wp.array2d(dtype=Any),
eta_dot: wp.array2d(dtype=Any),
eta_mass: wp.array2d(dtype=Any),
ke2: wp.array(dtype=Any),
target_temp: wp.array(dtype=Any),
ndof: wp.array(dtype=Any),
dt_chain: wp.array(dtype=Any),
chain_length: int,
vel_scale: wp.array(dtype=Any),
):
r"""Propagate Nosé-Hoover chain for one Yoshida-Suzuki sub-step.
This kernel implements the time-reversible Martyna-Tobias-Klein (MTK)
integration scheme for Nosé-Hoover chains.
Algorithm (for each system):
1. Half-step position update: :math:`\eta_k \mathrel{+}= 0.5 \cdot dt \cdot \dot{\eta}_k`
2. Backward sweep: Update :math:`\dot{\eta}` from chain end to start with friction
3. Compute velocity scale factor: :math:`\exp(-0.5 \cdot dt \cdot \dot{\eta}_0)`
4. Forward sweep: Update :math:`\dot{\eta}` from start to chain end with new forces
5. Half-step position update: :math:`\eta_k \mathrel{+}= 0.5 \cdot dt \cdot \dot{\eta}_k`
Launch Grid
-----------
dim = [num_systems]
Parameters
----------
eta : wp.array2d(dtype=Any)
Chain positions, shape (num_systems, chain_length). MODIFIED in-place.
eta_dot : wp.array2d(dtype=Any)
Chain velocities, shape (num_systems, chain_length). MODIFIED in-place.
eta_mass : wp.array2d(dtype=Any)
Chain masses, shape (num_systems, chain_length).
ke2 : wp.array(dtype=Any)
2*KE for each system, shape (num_systems,). MODIFIED to reflect scaled KE.
target_temp : wp.array(dtype=wp.float64)
Target temperature (kT), shape (num_systems,).
ndof : wp.array(dtype=wp.float64)
Degrees of freedom, shape (num_systems,).
dt_chain : wp.array(dtype=wp.float64)
Time step for this sub-step (weight * dt), shape (num_systems,).
chain_length : int
Number of thermostats in the chain.
vel_scale : wp.array(dtype=wp.float64)
Output velocity scale factors, shape (num_systems,). MODIFIED.
"""
sys_id = wp.tid()
kT = target_temp[sys_id]
ndof_sys = ndof[sys_id]
dt = dt_chain[sys_id]
half_dt = type(dt)(0.5) * dt
quarter_dt = type(dt)(0.25) * dt
eighth_dt = type(dt)(0.125) * dt
ke2_sys = ke2[sys_id]
# Local copies for chain state (we'll write back at the end)
# Using fixed-size local arrays for the chain
eta_local = wp.vector(dtype=eta.dtype, length=MAX_CHAIN_LENGTH)
eta_dot_local = wp.vector(dtype=eta_dot.dtype, length=MAX_CHAIN_LENGTH)
eta_mass_local = wp.vector(dtype=eta_mass.dtype, length=MAX_CHAIN_LENGTH)
# Load chain state
for k in range(chain_length):
eta_local[k] = eta[sys_id, k]
eta_dot_local[k] = eta_dot[sys_id, k]
eta_mass_local[k] = eta_mass[sys_id, k]
# ========== Step 1: Half-step position update ==========
for k in range(chain_length):
eta_local[k] = eta_local[k] + half_dt * eta_dot_local[k]
# ========== Step 2: Backward sweep (chain_length-1 down to 0) ==========
# Update last thermostat (no friction from above)
if chain_length > 1:
G_last = (
eta_mass_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
- kT
) / eta_mass_local[chain_length - 1]
eta_dot_local[chain_length - 1] = (
eta_dot_local[chain_length - 1] + quarter_dt * G_last
)
# Update intermediate thermostats (chain_length-2 down to 1)
for k in range(chain_length - 2, 0, -1):
G_k = (
eta_mass_local[k - 1] * eta_dot_local[k - 1] * eta_dot_local[k - 1] - kT
) / eta_mass_local[k]
# Apply friction from k+1
scale = wp.exp(-eighth_dt * eta_dot_local[k + 1])
eta_dot_local[k] = eta_dot_local[k] * scale
eta_dot_local[k] = eta_dot_local[k] + quarter_dt * G_k
eta_dot_local[k] = eta_dot_local[k] * scale
# Update first thermostat (couples to particle KE)
G_0 = (ke2_sys - ndof_sys * kT) / eta_mass_local[0]
if chain_length > 1:
scale = wp.exp(-eighth_dt * eta_dot_local[1])
eta_dot_local[0] = eta_dot_local[0] * scale
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0
eta_dot_local[0] = eta_dot_local[0] * scale
else:
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0
# ========== Step 3: Compute velocity scale factor ==========
vel_scale_factor = wp.exp(-half_dt * eta_dot_local[0])
vel_scale[sys_id] = vel_scale_factor
# Update ke2 for the forward sweep
ke2_sys = ke2_sys * vel_scale_factor * vel_scale_factor
ke2[sys_id] = ke2_sys
# ========== Step 4: Forward sweep (0 to chain_length-1) ==========
# Update first thermostat with new force
G_0_new = (ke2_sys - ndof_sys * kT) / eta_mass_local[0]
if chain_length > 1:
scale = wp.exp(-eighth_dt * eta_dot_local[1])
eta_dot_local[0] = eta_dot_local[0] * scale
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0_new
eta_dot_local[0] = eta_dot_local[0] * scale
else:
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0_new
# Update intermediate thermostats (1 to chain_length-2)
for k in range(1, chain_length - 1):
G_k = (
eta_mass_local[k - 1] * eta_dot_local[k - 1] * eta_dot_local[k - 1] - kT
) / eta_mass_local[k]
scale = wp.exp(-eighth_dt * eta_dot_local[k + 1])
eta_dot_local[k] = eta_dot_local[k] * scale
eta_dot_local[k] = eta_dot_local[k] + quarter_dt * G_k
eta_dot_local[k] = eta_dot_local[k] * scale
# Update last thermostat
if chain_length > 1:
G_last = (
eta_mass_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
- kT
) / eta_mass_local[chain_length - 1]
eta_dot_local[chain_length - 1] = (
eta_dot_local[chain_length - 1] + quarter_dt * G_last
)
# ========== Step 5: Second half-step position update ==========
for k in range(chain_length):
eta_local[k] = eta_local[k] + half_dt * eta_dot_local[k]
# ========== Write back chain state ==========
for k in range(chain_length):
eta[sys_id, k] = eta_local[k]
eta_dot[sys_id, k] = eta_dot_local[k]
@wp.kernel
def _nhc_chain_propagate_single_kernel(
eta: wp.array(dtype=Any),
eta_dot: wp.array(dtype=Any),
eta_mass: wp.array(dtype=Any),
ke2: wp.array(dtype=Any),
target_temp: wp.array(dtype=Any),
ndof: wp.array(dtype=Any),
dt_chain: wp.array(dtype=Any),
chain_length: int,
vel_scale: wp.array(dtype=Any),
):
"""Propagate Nosé-Hoover chain for single system (non-batched).
Same algorithm as batched version but for 1D arrays.
Parameters
----------
eta : wp.array(dtype=Any)
Chain positions, shape (chain_length,). MODIFIED in-place.
eta_dot : wp.array(dtype=Any)
Chain velocities, shape (chain_length,). MODIFIED in-place.
eta_mass : wp.array(dtype=Any)
Chain masses, shape (chain_length,).
ke2 : wp.array(dtype=Any)
2*KE for the system, shape (1,). MODIFIED to reflect scaled KE.
target_temp : wp.array(dtype=Any)
Target temperature (kT), shape (1,).
ndof : wp.array(dtype=Any)
Degrees of freedom, shape (1,).
dt_chain : wp.array(dtype=Any)
Time step for this sub-step (weight * dt), shape (1,).
chain_length : int
Number of thermostats in the chain.
vel_scale : wp.array(dtype=Any)
Output velocity scale factors, shape (1,). MODIFIED.
Launch Grid
-----------
dim = [1]
"""
kT = target_temp[0]
ndof_sys = ndof[0]
dt = dt_chain[0]
half_dt = type(dt)(0.5) * dt
quarter_dt = type(dt)(0.25) * dt
eighth_dt = type(dt)(0.125) * dt
ke2_sys = ke2[0]
# Local copies for chain state
eta_local = wp.vector(dtype=eta.dtype, length=MAX_CHAIN_LENGTH)
eta_dot_local = wp.vector(dtype=eta_dot.dtype, length=MAX_CHAIN_LENGTH)
eta_mass_local = wp.vector(dtype=eta_mass.dtype, length=MAX_CHAIN_LENGTH)
# Load chain state
for k in range(chain_length):
eta_local[k] = eta[k]
eta_dot_local[k] = eta_dot[k]
eta_mass_local[k] = eta_mass[k]
# ========== Step 1: Half-step position update ==========
for k in range(chain_length):
eta_local[k] = eta_local[k] + half_dt * eta_dot_local[k]
# ========== Step 2: Backward sweep ==========
if chain_length > 1:
G_last = (
eta_mass_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
- kT
) / eta_mass_local[chain_length - 1]
eta_dot_local[chain_length - 1] = (
eta_dot_local[chain_length - 1] + quarter_dt * G_last
)
for k in range(chain_length - 2, 0, -1):
G_k = (
eta_mass_local[k - 1] * eta_dot_local[k - 1] * eta_dot_local[k - 1] - kT
) / eta_mass_local[k]
scale = wp.exp(-eighth_dt * eta_dot_local[k + 1])
eta_dot_local[k] = eta_dot_local[k] * scale
eta_dot_local[k] = eta_dot_local[k] + quarter_dt * G_k
eta_dot_local[k] = eta_dot_local[k] * scale
G_0 = (ke2_sys - ndof_sys * kT) / eta_mass_local[0]
if chain_length > 1:
scale = wp.exp(-eighth_dt * eta_dot_local[1])
eta_dot_local[0] = eta_dot_local[0] * scale
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0
eta_dot_local[0] = eta_dot_local[0] * scale
else:
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0
# ========== Step 3: Compute velocity scale factor ==========
vel_scale_factor = wp.exp(-half_dt * eta_dot_local[0])
vel_scale[0] = vel_scale_factor
ke2_sys = ke2_sys * vel_scale_factor * vel_scale_factor
ke2[0] = ke2_sys
# ========== Step 4: Forward sweep ==========
G_0_new = (ke2_sys - ndof_sys * kT) / eta_mass_local[0]
if chain_length > 1:
scale = wp.exp(-eighth_dt * eta_dot_local[1])
eta_dot_local[0] = eta_dot_local[0] * scale
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0_new
eta_dot_local[0] = eta_dot_local[0] * scale
else:
eta_dot_local[0] = eta_dot_local[0] + quarter_dt * G_0_new
for k in range(1, chain_length - 1):
G_k = (
eta_mass_local[k - 1] * eta_dot_local[k - 1] * eta_dot_local[k - 1] - kT
) / eta_mass_local[k]
scale = wp.exp(-eighth_dt * eta_dot_local[k + 1])
eta_dot_local[k] = eta_dot_local[k] * scale
eta_dot_local[k] = eta_dot_local[k] + quarter_dt * G_k
eta_dot_local[k] = eta_dot_local[k] * scale
if chain_length > 1:
G_last = (
eta_mass_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
* eta_dot_local[chain_length - 2]
- kT
) / eta_mass_local[chain_length - 1]
eta_dot_local[chain_length - 1] = (
eta_dot_local[chain_length - 1] + quarter_dt * G_last
)
# ========== Step 5: Second half-step position update ==========
for k in range(chain_length):
eta_local[k] = eta_local[k] + half_dt * eta_dot_local[k]
for k in range(chain_length):
eta[k] = eta_local[k]
eta_dot[k] = eta_dot_local[k]
# ==============================================================================
# Chain Energy Kernels
# ==============================================================================
@wp.kernel
def _nhc_compute_chain_energy_kernel(
eta: wp.array(dtype=Any),
eta_dot: wp.array(dtype=Any),
eta_mass: wp.array(dtype=Any),
target_temp: wp.array(dtype=Any),
ndof: wp.array(dtype=Any),
chain_length: int,
ke_chain: wp.array(dtype=Any),
pe_chain: wp.array(dtype=Any),
):
r"""Compute NHC kinetic and potential energy for single system.
.. math::
\text{KE}_\text{chain} = \sum_k \tfrac{1}{2} Q_k \dot{\eta}_k^2
\text{PE}_\text{chain} = N_\text{dof} \cdot kT \cdot \eta_0 + kT \sum_{k>0} \eta_k
Parameters
----------
eta : wp.array(dtype=Any)
Chain positions, shape (chain_length,).
eta_dot : wp.array(dtype=Any)
Chain velocities, shape (chain_length,).
eta_mass : wp.array(dtype=Any)
Chain masses, shape (chain_length,).
target_temp : wp.array(dtype=Any)
Target temperature (kT), shape (1,).
ndof : wp.array(dtype=Any)
Degrees of freedom, shape (1,).
chain_length : int
Number of thermostats in the chain.
ke_chain : wp.array(dtype=Any)
Kinetic energy of the chain, shape (1,).
pe_chain : wp.array(dtype=Any)
Potential energy of the chain, shape (1,).
Launch Grid
-----------
dim = [1]
"""
kT = target_temp[0]
ndof_sys = ndof[0]
ke = type(eta[0])(0.0)
pe = type(eta[0])(0.0)
for k in range(chain_length):
ke = ke + type(eta[0])(0.5) * eta_mass[k] * eta_dot[k] * eta_dot[k]
if k == 0:
pe = pe + type(eta[0])(ndof_sys) * type(eta[0])(kT) * eta[k]
else:
pe = pe + type(eta[0])(kT) * eta[k]
ke_chain[0] = type(eta[0])(ke)
pe_chain[0] = type(eta[0])(pe)
@wp.kernel
def _batch_nhc_compute_chain_energy_kernel(
eta: wp.array2d(dtype=Any),
eta_dot: wp.array2d(dtype=Any),
eta_mass: wp.array2d(dtype=Any),
target_temp: wp.array(dtype=Any),
ndof: wp.array(dtype=Any),
chain_length: int,
ke_chain: wp.array(dtype=Any),
pe_chain: wp.array(dtype=Any),
):
"""Compute NHC kinetic and potential energy for batched systems.
Parameters
----------
eta : wp.array2d(dtype=Any)
Chain positions, shape (num_systems, chain_length).
eta_dot : wp.array2d(dtype=Any)
Chain velocities, shape (num_systems, chain_length).
eta_mass : wp.array2d(dtype=Any)
Chain masses, shape (num_systems, chain_length).
target_temp : wp.array(dtype=Any)
Target temperature (kT), shape (num_systems,).
ndof : wp.array(dtype=Any)
Degrees of freedom, shape (num_systems,).
chain_length : int
Number of thermostats in the chain.
ke_chain : wp.array(dtype=Any)
Kinetic energy of the chain, shape (num_systems,).
pe_chain : wp.array(dtype=Any)
Potential energy of the chain, shape (num_systems,).
Launch Grid
-----------
dim = [num_systems]
"""
sys_id = wp.tid()
# Use eta[sys_id, 0] to get a scalar for type inference (not eta[0] which returns a row)
kT = type(eta[sys_id, 0])(target_temp[sys_id])
ndof_sys = type(eta[sys_id, 0])(ndof[sys_id])
ke = type(kT)(0.0)
pe = type(kT)(0.0)
for k in range(chain_length):
ke = (
ke
+ type(kT)(0.5)
* eta_mass[sys_id, k]
* eta_dot[sys_id, k]
* eta_dot[sys_id, k]
)
if k == 0:
pe = pe + ndof_sys * kT * eta[sys_id, k]
else:
pe = pe + kT * eta[sys_id, k]
ke_chain[sys_id] = ke
pe_chain[sys_id] = pe
# ==============================================================================
# Velocity Scaling Kernels
# ==============================================================================
@wp.kernel
def _scale_velocities_kernel(
velocities: wp.array(dtype=Any),
scale_factor: wp.array(dtype=Any),
):
"""Scale all velocities by a single factor.
Launch Grid
-----------
dim = [num_atoms]
"""
atom_idx = wp.tid()
v = velocities[atom_idx]
s = type(v[0])(scale_factor[0])
new_vx = type(v[0])(v[0] * s)
new_vy = type(v[1])(v[1] * s)
new_vz = type(v[2])(v[2] * s)
velocities[atom_idx] = type(v)(new_vx, new_vy, new_vz)
@wp.kernel
def _batch_scale_velocities_kernel(
velocities: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
scale_factors: wp.array(dtype=Any),
):
"""Scale velocities with per-system factors.
Launch Grid
-----------
dim = [num_atoms_total]
"""
atom_idx = wp.tid()
system_id = batch_idx[atom_idx]
v = velocities[atom_idx]
s = type(v[0])(scale_factors[system_id])
new_vx = type(v[0])(v[0] * s)
new_vy = type(v[1])(v[1] * s)
new_vz = type(v[2])(v[2] * s)
velocities[atom_idx] = type(v)(new_vx, new_vy, new_vz)
@wp.kernel
def _scale_velocities_out_kernel(
velocities: wp.array(dtype=Any),
scale_factor: wp.array(dtype=Any),
velocities_out: wp.array(dtype=Any),
):
"""Scale velocities to output array.
Launch Grid
-----------
dim = [num_atoms]
"""
atom_idx = wp.tid()
v = velocities[atom_idx]
s = type(v[0])(scale_factor[0])
new_vx = type(v[0])(v[0] * s)
new_vy = type(v[1])(v[1] * s)
new_vz = type(v[2])(v[2] * s)
velocities_out[atom_idx] = type(v)(new_vx, new_vy, new_vz)
@wp.kernel
def _batch_scale_velocities_out_kernel(
velocities: wp.array(dtype=Any),
batch_idx: wp.array(dtype=wp.int32),
scale_factors: wp.array(dtype=Any),
velocities_out: wp.array(dtype=Any),
):
"""Scale velocities with per-system factors to output array.
Launch Grid
-----------
dim = [num_atoms_total]
"""
atom_idx = wp.tid()
system_id = batch_idx[atom_idx]
v = velocities[atom_idx]
s = type(v[0])(scale_factors[system_id])
new_vx = type(v[0])(v[0] * s)
new_vy = type(v[1])(v[1] * s)
new_vz = type(v[2])(v[2] * s)
velocities_out[atom_idx] = type(v)(new_vx, new_vy, new_vz)
# ==============================================================================
# Multiply Scale Factors Kernel
# ==============================================================================
@wp.kernel
def _multiply_scale_factors_kernel(
total_scale: wp.array(dtype=Any),
step_scale: wp.array(dtype=Any),
):
"""Multiply total scale factors by step scale factors.
Launch Grid
-----------
dim = [num_systems]
"""
sys_id = wp.tid()
total_scale[sys_id] = total_scale[sys_id] * step_scale[sys_id]
# Overloads for _multiply_scale_factors_kernel
_multiply_scale_factors_kernel_overload = {}
for t in [wp.float32, wp.float64]:
_multiply_scale_factors_kernel_overload[t] = wp.overload(
_multiply_scale_factors_kernel,
[wp.array(dtype=t), wp.array(dtype=t)],
)
# ==============================================================================
# Kernel Overloads for Explicit Typing
# ==============================================================================
# These overloads provide explicit type annotations for each kernel, avoiding
# Warp's type inference issues that can occur with dtype=Any parameters.
_T = [wp.float32, wp.float64] # Scalar types
_V = [wp.vec3f, wp.vec3d] # Vector types
# Diagnostic kernel overloads
_compute_2ke_kernel_overload = {}
_compute_2ke_tiled_kernel_overload = {}
_batch_compute_2ke_kernel_overload = {}
# Velocity scaling kernel overloads
_scale_velocities_kernel_overload = {}
_batch_scale_velocities_kernel_overload = {}
_scale_velocities_out_kernel_overload = {}
_batch_scale_velocities_out_kernel_overload = {}
# Compute masses kernel overloads (keyed by scalar type)
_nhc_compute_masses_kernel_overload = {}
_batch_nhc_compute_masses_kernel_overload = {}
for t in _T:
_nhc_compute_masses_kernel_overload[t] = wp.overload(
_nhc_compute_masses_kernel,
[
wp.array(dtype=wp.int32),
wp.array(dtype=t),
wp.array(dtype=t),
wp.array(dtype=t),
],
)
_batch_nhc_compute_masses_kernel_overload[t] = wp.overload(
_batch_nhc_compute_masses_kernel,
[
wp.array(dtype=wp.int32),
wp.array(dtype=t),
wp.array(dtype=t),
wp.array2d(dtype=t),
],
)
# Create overloads for all combinations of velocity type (v) and output type (t_out)
# The ke2 output needs to match the chain state dtype, not the velocity dtype
for v in _V:
# Determine the scalar type for masses from the vector type
t_mass = wp.float32 if v == wp.vec3f else wp.float64
for t_out in _T:
# Key by (velocity_type, output_type)
_compute_2ke_kernel_overload[(v, t_out)] = wp.overload(
_compute_2ke_kernel,
[wp.array(dtype=v), wp.array(dtype=t_mass), wp.array(dtype=t_out)],
)
_compute_2ke_tiled_kernel_overload[(v, t_out)] = wp.overload(
_compute_2ke_tiled_kernel,
[wp.array(dtype=v), wp.array(dtype=t_mass), wp.array(dtype=t_out)],
)
_batch_compute_2ke_kernel_overload[(v, t_out)] = wp.overload(
_batch_compute_2ke_kernel,
[
wp.array(dtype=v),
wp.array(dtype=t_mass),
wp.array(dtype=wp.int32),
wp.array(dtype=t_out),
],
)
# Create velocity scaling kernel overloads for all combinations of velocity type and scale factor dtype
for v in _V:
for t_scale in _T:
# Key by (velocity_type, scale_factor_dtype) to support mixed dtypes
_scale_velocities_kernel_overload[(v, t_scale)] = wp.overload(
_scale_velocities_kernel,
[wp.array(dtype=v), wp.array(dtype=t_scale)],
)
_batch_scale_velocities_kernel_overload[(v, t_scale)] = wp.overload(
_batch_scale_velocities_kernel,
[wp.array(dtype=v), wp.array(dtype=wp.int32), wp.array(dtype=t_scale)],
)
_scale_velocities_out_kernel_overload[(v, t_scale)] = wp.overload(
_scale_velocities_out_kernel,
[wp.array(dtype=v), wp.array(dtype=t_scale), wp.array(dtype=v)],
)
_batch_scale_velocities_out_kernel_overload[(v, t_scale)] = wp.overload(
_batch_scale_velocities_out_kernel,
[
wp.array(dtype=v),
wp.array(dtype=wp.int32),
wp.array(dtype=t_scale),
wp.array(dtype=v),
],
)
# NHC chain propagation kernels - keyed by scalar type
_nhc_chain_propagate_single_kernel_overload = {}
for t in _T:
_nhc_chain_propagate_single_kernel_overload[t] = wp.overload(
_nhc_chain_propagate_single_kernel,
[
wp.array(dtype=t), # eta
wp.array(dtype=t), # eta_dot
wp.array(dtype=t), # eta_mass
wp.array(dtype=t), # ke2
wp.array(dtype=t), # target_temp
wp.array(dtype=t), # ndof
wp.array(dtype=t), # dt_chain
wp.int32, # chain_length
wp.array(dtype=t), # vel_scale
],
)
# NHC chain energy kernels - keyed by scalar type
_nhc_compute_chain_energy_kernel_overload = {}
_batch_nhc_compute_chain_energy_kernel_overload = {}
for t in _T:
_nhc_compute_chain_energy_kernel_overload[t] = wp.overload(
_nhc_compute_chain_energy_kernel,
[
wp.array(dtype=t), # eta
wp.array(dtype=t), # eta_dot
wp.array(dtype=t), # eta_mass
wp.array(dtype=t), # target_temp
wp.array(dtype=t), # ndof
wp.int32, # chain_length
wp.array(dtype=t), # ke_chain
wp.array(dtype=t), # pe_chain
],
)
_batch_nhc_compute_chain_energy_kernel_overload[t] = wp.overload(
_batch_nhc_compute_chain_energy_kernel,
[
wp.array2d(dtype=t), # eta
wp.array2d(dtype=t), # eta_dot
wp.array2d(dtype=t), # eta_mass
wp.array(dtype=t), # target_temp
wp.array(dtype=t), # ndof
wp.int32, # chain_length
wp.array(dtype=t), # ke_chain
wp.array(dtype=t), # pe_chain
],
)
# ==============================================================================
# Functional Interfaces
# ==============================================================================
[docs]
def nhc_compute_masses(
ndof: wp.array,
target_temp: wp.array,
tau: wp.array,
chain_length: int,
masses: wp.array,
num_systems: int = 1,
device: str = None,
dtype=wp.float64,
) -> wp.array:
r"""Compute Nosé-Hoover chain masses using GPU kernel.
Computes :math:`Q_k` values for the Nosé-Hoover chain:
.. math::
Q_0 &= N_\text{dof} \cdot k_B T \cdot \tau^2 \\
Q_k &= k_B T \cdot \tau^2 \quad \text{for } k > 0
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 :math:`k_B T` per system. Shape (1,) for single system,
(num_systems,) for batched.
tau : wp.array
Time constant :math:`\tau` per system. Shape (1,) for single system,
(num_systems,) for batched.
chain_length : int
Number of thermostats in the chain.
masses : wp.array
Chain masses output. Caller must pre-allocate.
Shape (chain_length,) for single system,
(num_systems, chain_length) for batched.
num_systems : int, optional
Number of systems for batched mode. Default: 1.
device : str, optional
Warp device. If None, inferred from masses.
dtype : dtype, optional
Data type for the masses. Default: wp.float64.
Returns
-------
wp.array
Chain masses. Shape (chain_length,) for single system,
(num_systems, chain_length) for batched.
"""
if device is None:
device = masses.device
is_batched = masses.ndim == 2
# Select overload based on dtype
scalar_type = dtype
if is_batched:
wp.launch(
_batch_nhc_compute_masses_kernel_overload[scalar_type],
dim=(num_systems, chain_length),
inputs=[ndof, target_temp, tau, masses],
device=device,
)
else:
wp.launch(
_nhc_compute_masses_kernel_overload[scalar_type],
dim=chain_length,
inputs=[ndof, target_temp, tau, masses],
device=device,
)
return masses
[docs]
def nhc_compute_2ke(
velocities: wp.array,
masses: wp.array,
ke2: wp.array,
batch_idx: wp.array = None,
device: str = None,
) -> wp.array:
r"""
Fill ``ke2[s]`` with :math:`\sum_{i \in s} m_i \lVert \mathbf{v}_i \rVert^2`
(i.e. :math:`2\,\mathrm{KE}` per system) — the same reduction that
``nhc_thermostat_chain_update`` performs internally, exposed standalone.
Lets a caller compute :math:`2\,\mathrm{KE}` separately, optionally
post-process it, and feed it back via
``nhc_thermostat_chain_update(..., compute_ke=False)``. Using this kernel for
the reduction gives machine-precision parity with the value the chain update
would otherwise compute itself.
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,).
ke2 : wp.array
Output :math:`2\,\mathrm{KE}`. Shape (1,) for single system,
(num_systems,) for batched (its length sets the number of systems).
Zeroed internally before accumulation. MODIFIED in-place.
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. Required for batched mode.
device : str, optional
Warp device. If None, inferred from velocities.
Returns
-------
wp.array
The ``ke2`` array, filled with :math:`2\,\mathrm{KE}` per system.
"""
if device is None:
device = velocities.device
num_atoms = velocities.shape[0]
is_batched = batch_idx is not None
vec_dtype = velocities.dtype
out_dtype = ke2.dtype
ke2.zero_()
if is_batched:
wp.launch(
_batch_compute_2ke_kernel_overload[(vec_dtype, out_dtype)],
dim=num_atoms,
inputs=[velocities, masses, batch_idx, ke2],
device=device,
)
else:
wp.launch(
_compute_2ke_kernel_overload[(vec_dtype, out_dtype)],
dim=num_atoms,
inputs=[velocities, masses, ke2],
device=device,
)
return ke2
[docs]
def nhc_thermostat_chain_update(
velocities: wp.array,
masses: wp.array,
eta: wp.array,
eta_dot: wp.array,
eta_mass: wp.array,
target_temp: wp.array,
dt: wp.array,
ndof: wp.array,
ke2: wp.array,
total_scale: wp.array,
step_scale: wp.array,
dt_chain: wp.array,
nloops: int = 1,
batch_idx: wp.array = None,
num_systems: int = 1,
device: str = None,
compute_ke: bool = True,
) -> None:
r"""
Propagate Nosé-Hoover chain and scale velocities (in-place).
Propagates the chain over the Yoshida-Suzuki sub-steps, accumulating the
per-step velocity scale factor :math:`\exp(-\tfrac{1}{2}\,\Delta t_c\,\dot{\eta}_0)`
into ``total_scale``, then rescales every particle velocity by it:
.. math::
\mathbf{v}_i \leftarrow s\,\mathbf{v}_i,
\qquad
s = \prod_j \exp\!\left(-\tfrac{1}{2}\,\Delta t_{c,j}\,\dot{\eta}_0\right)
Uses Yoshida-Suzuki factorization for time-reversible integration.
All computations are performed on GPU using Warp kernels.
Parameters
----------
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,). MODIFIED in-place.
masses : wp.array(dtype=wp.float32 or wp.float64)
Atomic masses. Shape (N,).
eta : wp.array(dtype=wp.float64)
Thermostat chain positions.
Non-batched: Shape (chain_length,).
Batched: Shape (num_systems, chain_length) as wp.array2d.
MODIFIED in-place.
eta_dot : wp.array(dtype=wp.float64)
Thermostat chain velocities. Same shape as eta. MODIFIED in-place.
eta_mass : wp.array(dtype=wp.float64)
Thermostat chain masses. Same shape as eta.
target_temp : wp.array(dtype=wp.float64)
Target temperature :math:`k_B T`. Shape (1,) or (num_systems,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (num_systems,).
ndof : wp.array(dtype=wp.float64)
Degrees of freedom. Shape (1,) or (num_systems,).
ke2 : wp.array
:math:`2\,\mathrm{KE} = \sum_i m_i \lVert \mathbf{v}_i \rVert^2` per
system. When compute_ke is True (default) this
is a scratch array zeroed and filled internally from velocities/masses
before use. When compute_ke is False it is an INPUT supplying a
caller-precomputed :math:`2\,\mathrm{KE}`; it is not zeroed. Either way the chain forward
sweep rescales it in place.
Shape (1,) for single system, (num_systems,) for batched.
total_scale : wp.array
Scratch array for accumulated velocity scale factor.
Must be initialized to ones by caller (wp.ones).
Shape (1,) for single system, (num_systems,) for batched.
step_scale : wp.array
Scratch array for per-step velocity scale factor.
Shape (1,) for single system, (num_systems,) for batched.
dt_chain : wp.array
Scratch array for weighted time steps.
Shape (1,) for single system, (num_systems,) for batched.
nloops : int, optional
Number of Yoshida-Suzuki integration sub-steps. Default: 1.
Use nloops=3 or 5 for higher accuracy.
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. Required for batched mode.
num_systems : int, optional
Number of systems for batched mode. Default: 1.
device : str, optional
Warp device. If None, inferred from velocities.
compute_ke : bool, optional
If True (default), recompute :math:`2\,\mathrm{KE}` from velocities/masses
internally (byte-identical to legacy behavior). If False, use the
caller-supplied value already in ``ke2`` instead of recomputing it. See
``nhc_compute_2ke`` for the matching reduction helper.
"""
if device is None:
device = velocities.device
num_atoms = velocities.shape[0]
is_batched = batch_idx is not None
# Get Yoshida-Suzuki weights
if nloops == 1:
weights = [1.0]
elif nloops == 3:
weights = YOSHIDA_SUZUKI_3
elif nloops == 5:
weights = YOSHIDA_SUZUKI_5
else:
# Simple equal weights for other values
weights = [1.0 / nloops] * nloops
# Determine chain length
if is_batched:
chain_length = eta.shape[1]
else:
chain_length = eta.shape[0]
if chain_length > MAX_CHAIN_LENGTH:
raise ValueError(
f"Chain length {chain_length} exceeds maximum {MAX_CHAIN_LENGTH}"
)
# Compute 2*KE - ke2 is zeroed internally before each use UNLESS the caller
# supplied it. When compute_ke is False, ke2 arrives pre-filled and is NOT
# zeroed — the caller owns it.
vec_dtype = velocities.dtype
chain_dtype = eta.dtype
n_scale = num_systems if is_batched else 1
if compute_ke:
ke2.zero_()
if is_batched:
wp.launch(
_batch_compute_2ke_kernel_overload[(vec_dtype, chain_dtype)],
dim=num_atoms,
inputs=[velocities, masses, batch_idx, ke2],
device=device,
)
else:
wp.launch(
_compute_2ke_kernel_overload[(vec_dtype, chain_dtype)],
dim=num_atoms,
inputs=[velocities, masses, ke2],
device=device,
)
# Run Yoshida-Suzuki sub-steps
for w in weights:
# Compute weighted time step: dt_chain = w * dt
if is_batched:
# For batched case, we need to scale each system's dt by the weight
_compute_weighted_dt(dt, dt_chain, w, num_systems, device)
else:
_compute_weighted_dt(dt, dt_chain, w, 1, device)
# Propagate chain
if is_batched:
wp.launch(
_nhc_chain_propagate_kernel,
dim=num_systems,
inputs=[
eta,
eta_dot,
eta_mass,
ke2,
target_temp,
ndof,
dt_chain,
chain_length,
step_scale,
],
device=device,
)
else:
wp.launch(
_nhc_chain_propagate_single_kernel_overload[chain_dtype],
dim=1,
inputs=[
eta,
eta_dot,
eta_mass,
ke2,
target_temp,
ndof,
dt_chain,
chain_length,
step_scale,
],
device=device,
)
# Accumulate total scale factor
wp.launch(
_multiply_scale_factors_kernel_overload[chain_dtype],
dim=n_scale,
inputs=[total_scale, step_scale],
device=device,
)
# Scale velocities
if is_batched:
wp.launch(
_batch_scale_velocities_kernel_overload[(vec_dtype, chain_dtype)],
dim=num_atoms,
inputs=[velocities, batch_idx, total_scale],
device=device,
)
else:
wp.launch(
_scale_velocities_kernel_overload[(vec_dtype, chain_dtype)],
dim=num_atoms,
inputs=[velocities, total_scale],
device=device,
)
@wp.kernel
def _compute_weighted_dt_kernel(
dt: wp.array(dtype=Any),
dt_chain: wp.array(dtype=Any),
weight: Any,
):
"""Compute weighted time step: dt_chain = weight * dt.
Launch Grid
-----------
dim = [num_systems]
"""
sys_id = wp.tid()
dt_chain[sys_id] = type(dt_chain[sys_id])(dt[sys_id]) * type(dt_chain[sys_id])(
weight
)
# Overloads for _compute_weighted_dt_kernel - support all combinations of dt and dt_chain dtypes
_compute_weighted_dt_kernel_overload = {}
for t_in in _T:
for t_out in _T:
_compute_weighted_dt_kernel_overload[(t_in, t_out)] = wp.overload(
_compute_weighted_dt_kernel,
[
wp.array(dtype=t_in),
wp.array(dtype=t_out),
t_out,
], # weight uses output type
)
def _compute_weighted_dt(
dt: wp.array,
dt_chain: wp.array,
weight: float,
num_systems: int,
device: str,
):
"""Helper to compute weighted dt using appropriate kernel overload."""
dt_dtype = dt.dtype
dt_chain_dtype = dt_chain.dtype
weight_typed = dt_chain_dtype(weight)
wp.launch(
_compute_weighted_dt_kernel_overload[(dt_dtype, dt_chain_dtype)],
dim=num_systems,
inputs=[dt, dt_chain, weight_typed],
device=device,
)
def nhc_thermostat_chain_update_out(
velocities: wp.array,
masses: wp.array,
eta: wp.array,
eta_dot: wp.array,
eta_mass: wp.array,
target_temp: wp.array,
dt: wp.array,
ndof: wp.array,
ke2: wp.array,
total_scale: wp.array,
step_scale: wp.array,
dt_chain: wp.array,
velocities_out: wp.array,
eta_out: wp.array,
eta_dot_out: wp.array,
nloops: int = 1,
batch_idx: wp.array = None,
num_systems: int = 1,
device: str = None,
compute_ke: bool = True,
) -> tuple[wp.array, wp.array, wp.array]:
r"""
Propagate Nosé-Hoover chain and scale velocities (non-mutating).
Out-of-place counterpart of :func:`nhc_thermostat_chain_update`: copies
``velocities``/``eta``/``eta_dot`` into the ``*_out`` buffers and applies the
same chain propagation and velocity rescaling
:math:`\mathbf{v}_i \leftarrow s\,\mathbf{v}_i` to the copies.
Parameters
----------
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,).
masses : wp.array(dtype=wp.float32 or wp.float64)
Atomic masses. Shape (N,).
eta : wp.array(dtype=wp.float64)
Thermostat chain positions. Shape (M,) or (B, M).
eta_dot : wp.array(dtype=wp.float64)
Thermostat chain velocities. Shape (M,) or (B, M).
eta_mass : wp.array(dtype=wp.float64)
Thermostat chain masses. Shape (M,) or (B, M).
target_temp : wp.array(dtype=wp.float64)
Target temperature :math:`k_B T`. Shape (1,) or (B,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (B,).
ndof : wp.array(dtype=wp.float64)
Degrees of freedom. Shape (1,) or (B,).
ke2 : wp.array
:math:`2\,\mathrm{KE} = \sum_i m_i \lVert \mathbf{v}_i \rVert^2` per
system. Zeroed and filled internally when
compute_ke is True (default); when compute_ke is False it is an INPUT
supplying a caller-precomputed :math:`2\,\mathrm{KE}` that is not zeroed. See
``nhc_thermostat_chain_update``.
Shape (1,) for single system, (num_systems,) for batched.
total_scale : wp.array
Scratch array for accumulated velocity scale factor.
Must be initialized to ones by caller (wp.ones).
Shape (1,) for single system, (num_systems,) for batched.
step_scale : wp.array
Scratch array for per-step velocity scale factor.
Shape (1,) for single system, (num_systems,) for batched.
dt_chain : wp.array
Scratch array for weighted time steps.
Shape (1,) for single system, (num_systems,) for batched.
velocities_out : wp.array
Output velocities. Must be pre-allocated with same shape/dtype/device as velocities.
eta_out : wp.array
Output eta. Must be pre-allocated with same shape/dtype/device as eta.
eta_dot_out : wp.array
Output eta_dot. Must be pre-allocated with same shape/dtype/device as eta_dot.
nloops : int, optional
Number of Yoshida-Suzuki integration sub-steps. Default: 1.
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. Required for batched mode.
num_systems : int, optional
Number of systems for batched mode. Default: 1.
device : str, optional
Warp device. If None, inferred from velocities.
compute_ke : bool, optional
If True (default), recompute :math:`2\,\mathrm{KE}` internally. If False, use the
caller-supplied value in ``ke2``. See
``nhc_thermostat_chain_update`` for details.
Returns
-------
tuple[wp.array, wp.array, wp.array]
(velocities_out, eta_out, eta_dot_out)
"""
if device is None:
device = velocities.device
validate_out_array(velocities_out, velocities, "velocities_out")
validate_out_array(eta_out, eta, "eta_out")
validate_out_array(eta_dot_out, eta_dot, "eta_dot_out")
# Copy inputs to outputs
wp.copy(velocities_out, velocities)
wp.copy(eta_out, eta)
wp.copy(eta_dot_out, eta_dot)
# Run in-place update on copies
nhc_thermostat_chain_update(
velocities_out,
masses,
eta_out,
eta_dot_out,
eta_mass,
target_temp,
dt,
ndof,
ke2,
total_scale,
step_scale,
dt_chain,
nloops=nloops,
batch_idx=batch_idx,
num_systems=num_systems,
device=device,
compute_ke=compute_ke,
)
return velocities_out, eta_out, eta_dot_out
[docs]
def nhc_velocity_half_step(
velocities: wp.array,
forces: wp.array,
masses: wp.array,
dt: wp.array,
batch_idx: wp.array = None,
atom_ptr: wp.array = None,
device: str = None,
) -> None:
r"""
Half-step velocity update (in-place).
.. math::
\mathbf{v}_i \leftarrow \mathbf{v}_i + \tfrac{1}{2}\,\frac{\mathbf{F}_i}{m_i}\,\Delta t
Parameters
----------
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,). MODIFIED in-place.
forces : wp.array(dtype=wp.vec3f or wp.vec3d)
Forces on atoms. Shape (N,).
masses : wp.array(dtype=wp.float32 or wp.float64)
Atomic masses. Shape (N,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (B,).
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. For batched mode (atomic operations).
atom_ptr : wp.array(dtype=wp.int32), optional
CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device : str, optional
Warp device. If None, inferred from velocities.
"""
dispatch_family(
velocity_kick_families,
velocities,
batch_idx=batch_idx,
atom_ptr=atom_ptr,
device=device,
inputs_single=[velocities, forces, masses, dt, velocities],
inputs_batch=[velocities, forces, masses, batch_idx, dt, velocities],
inputs_ptr=[velocities, forces, masses, atom_ptr, dt, velocities],
)
def nhc_velocity_half_step_out(
velocities: wp.array,
forces: wp.array,
masses: wp.array,
dt: wp.array,
velocities_out: wp.array,
batch_idx: wp.array = None,
atom_ptr: wp.array = None,
device: str = None,
) -> wp.array:
r"""
Half-step velocity update (non-mutating).
Writes :math:`\mathbf{v}_i + \tfrac{1}{2}\,\frac{\mathbf{F}_i}{m_i}\,\Delta t`
into ``velocities_out``.
Parameters
----------
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,).
forces : wp.array(dtype=wp.vec3f or wp.vec3d)
Forces on atoms. Shape (N,).
masses : wp.array(dtype=wp.float32 or wp.float64)
Atomic masses. Shape (N,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (B,).
velocities_out : wp.array
Output velocities. Must be pre-allocated with same shape/dtype/device as velocities.
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. For batched mode (atomic operations).
atom_ptr : wp.array(dtype=wp.int32), optional
CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device : str, optional
Warp device. If None, inferred from velocities.
Returns
-------
wp.array
Updated velocities.
"""
validate_out_array(velocities_out, velocities, "velocities_out")
dispatch_family(
velocity_kick_families,
velocities,
batch_idx=batch_idx,
atom_ptr=atom_ptr,
device=device,
inputs_single=[velocities, forces, masses, dt, velocities_out],
inputs_batch=[velocities, forces, masses, batch_idx, dt, velocities_out],
inputs_ptr=[velocities, forces, masses, atom_ptr, dt, velocities_out],
)
return velocities_out
[docs]
def nhc_position_update(
positions: wp.array,
velocities: wp.array,
dt: wp.array,
batch_idx: wp.array = None,
atom_ptr: wp.array = None,
device: str = None,
) -> None:
r"""
Full-step position update (in-place).
.. math::
\mathbf{r}_i \leftarrow \mathbf{r}_i + \mathbf{v}_i\,\Delta t
Parameters
----------
positions : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic positions. Shape (N,). MODIFIED in-place.
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (B,).
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. For batched mode (atomic operations).
atom_ptr : wp.array(dtype=wp.int32), optional
CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device : str, optional
Warp device. If None, inferred from positions.
"""
dispatch_family(
position_update_families,
positions,
batch_idx=batch_idx,
atom_ptr=atom_ptr,
device=device,
inputs_single=[positions, velocities, dt, positions],
inputs_batch=[positions, velocities, batch_idx, dt, positions],
inputs_ptr=[positions, velocities, atom_ptr, dt, positions],
)
def nhc_position_update_out(
positions: wp.array,
velocities: wp.array,
dt: wp.array,
positions_out: wp.array,
batch_idx: wp.array = None,
atom_ptr: wp.array = None,
device: str = None,
) -> wp.array:
r"""
Full-step position update (non-mutating).
Writes :math:`\mathbf{r}_i + \mathbf{v}_i\,\Delta t` into ``positions_out``.
Parameters
----------
positions : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic positions. Shape (N,).
velocities : wp.array(dtype=wp.vec3f or wp.vec3d)
Atomic velocities. Shape (N,).
dt : wp.array(dtype=wp.float32 or wp.float64)
Time step. Shape (1,) or (B,).
positions_out : wp.array
Output positions. Must be pre-allocated with same shape/dtype/device as positions.
batch_idx : wp.array(dtype=wp.int32), optional
System index for each atom. For batched mode (atomic operations).
atom_ptr : wp.array(dtype=wp.int32), optional
CSR-style pointers. Shape (num_systems + 1,). For batched mode (sequential per-system).
device : str, optional
Warp device. If None, inferred from positions.
Returns
-------
wp.array
Updated positions.
"""
validate_out_array(positions_out, positions, "positions_out")
dispatch_family(
position_update_families,
positions,
batch_idx=batch_idx,
atom_ptr=atom_ptr,
device=device,
inputs_single=[positions, velocities, dt, positions_out],
inputs_batch=[positions, velocities, batch_idx, dt, positions_out],
inputs_ptr=[positions, velocities, atom_ptr, dt, positions_out],
)
return positions_out
[docs]
def nhc_compute_chain_energy(
eta: wp.array,
eta_dot: wp.array,
eta_mass: wp.array,
target_temp: wp.array,
ndof: wp.array,
ke_chain: wp.array,
pe_chain: wp.array,
batch_idx: wp.array = None,
num_systems: int = 1,
device: str = None,
) -> tuple[wp.array, wp.array]:
r"""
Compute Nosé-Hoover chain kinetic and potential energy.
For conservation checks, the extended system Hamiltonian is:
.. math::
H_\text{ext} = \text{KE}_\text{particles} + \text{PE} + \text{KE}_\text{chain} + \text{PE}_\text{chain}
where:
.. math::
\text{KE}_\text{chain} = \sum_k \tfrac{1}{2} Q_k \dot{\eta}_k^2
\text{PE}_\text{chain} = N_\text{dof} \cdot kT \cdot \eta_0 + kT \sum_{k>0} \eta_k
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 :math:`k_B T`. Shape (1,) or (B,).
ndof : wp.array(dtype=wp.float64)
Degrees of freedom. Shape (1,) or (B,).
ke_chain : wp.array
Output kinetic energy of the chain.
Shape (1,) for single system, (num_systems,) for batched.
pe_chain : wp.array
Output potential energy of the chain.
Shape (1,) for single system, (num_systems,) for batched.
batch_idx : wp.array(dtype=wp.int32), optional
Not used directly, but included for API consistency.
num_systems : int, optional
Number of systems for batched mode. Default: 1.
device : str, optional
Warp device. If None, inferred from eta.
Returns
-------
tuple[wp.array, wp.array]
(ke_chain, pe_chain) each with shape (1,) or (B,).
"""
if device is None:
device = eta.device
is_batched = num_systems > 1 or (batch_idx is not None)
# Determine chain length
if is_batched:
chain_length = eta.shape[1]
else:
chain_length = eta.shape[0]
chain_dtype = eta.dtype
if is_batched:
wp.launch(
_batch_nhc_compute_chain_energy_kernel_overload[chain_dtype],
dim=num_systems,
inputs=[
eta,
eta_dot,
eta_mass,
target_temp,
ndof,
chain_length,
ke_chain,
pe_chain,
],
device=device,
)
else:
wp.launch(
_nhc_compute_chain_energy_kernel_overload[chain_dtype],
dim=1,
inputs=[
eta,
eta_dot,
eta_mass,
target_temp,
ndof,
chain_length,
ke_chain,
pe_chain,
],
device=device,
)
return ke_chain, pe_chain
