nvalchemiops.math and nvalchemiops.torch.types: Utilities

Mathematical Utilities

Mathematical Utilities

This module provides low-level mathematical functions implemented in Warp for use in GPU-accelerated scientific computing.

Available Submodules

spherical_harmonics

Real spherical harmonics \(Y_l^m\) for angular momentum \(L \leq 2\). Used in multipole expansions for electrostatics.

gto

Gaussian Type Orbital (GTO) basis functions for multipole charge distributions. Includes real-space densities and Fourier transforms for \(L \leq 2\).

spline

B-spline interpolation kernels for mesh-based calculations (e.g., PME). Provides spread, gather, and gradient operations with framework-agnostic Warp kernels and launcher functions.

nvalchemiops.math.batch_spline_gather(positions, batch_idx, cell_inv_t, order, mesh, output, wp_dtype, device=None)

Gather values from batched mesh to atoms using B-spline interpolation.

Framework-agnostic launcher for batched spline gather.

Parameters:

• positions (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – Atomic positions for all systems.

• batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom.

• cell_inv_t (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Per-system transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (B, nx, ny, nz), dtype=wp.float32 or wp.float64) – Batched mesh to interpolate from.

• output (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – OUTPUT: Interpolated values per atom. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.batch_spline_gather_gradient(positions, charges, batch_idx, cell_inv_t, order, mesh, forces, wp_dtype, device=None)

Compute forces using B-spline gradient interpolation from batched mesh.

Framework-agnostic launcher for batched spline gradient gather.

Parameters:

• positions (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – Atomic positions for all systems.

• charges (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – Atomic charges.

• batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom.

• cell_inv_t (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Per-system transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (B, nx, ny, nz), dtype=wp.float32 or wp.float64) – Batched potential mesh.

• forces (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Forces per atom. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.batch_spline_gather_vec3(positions, charges, batch_idx, cell_inv_t, order, mesh, output, wp_dtype, device=None)

Gather charge-weighted vector values from batched mesh using B-splines.

Framework-agnostic launcher for batched vec3 spline gather.

Parameters:

• positions (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – Atomic positions for all systems.

• charges (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – Atomic charges.

• batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom.

• cell_inv_t (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Per-system transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (B, nx, ny, nz), dtype=wp.vec3f or wp.vec3d) – Batched vector mesh to interpolate from.

• output (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Charge-weighted interpolated vectors. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.batch_spline_spread(positions, values, batch_idx, cell_inv_t, order, mesh, wp_dtype, device=None)

Spread values from atoms to batched mesh using B-spline interpolation.

Framework-agnostic launcher for batched spline spread.

Parameters:

• positions (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – Atomic positions for all systems.

• values (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – Values to spread.

• batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom.

• cell_inv_t (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Per-system transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (B, nx, ny, nz), dtype=wp.float32 or wp.float64) – OUTPUT: Batched mesh to accumulate values. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.spline_gather(positions, cell_inv_t, order, mesh, output, wp_dtype, device=None)

Gather values from mesh to atoms using B-spline interpolation.

Framework-agnostic launcher for single-system spline gather.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic positions.

• cell_inv_t (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (nx, ny, nz), dtype=wp.float32 or wp.float64) – Mesh to interpolate from.

• output (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Interpolated values per atom. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.spline_gather_gradient(positions, charges, cell_inv_t, order, mesh, forces, wp_dtype, device=None)

Compute forces using B-spline gradient interpolation.

Framework-agnostic launcher for single-system spline gradient gather.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic positions.

• charges (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – Atomic charges.

• cell_inv_t (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (nx, ny, nz), dtype=wp.float32 or wp.float64) – Potential mesh.

• forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Forces per atom. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.spline_gather_vec3(positions, charges, cell_inv_t, order, mesh, output, wp_dtype, device=None)

Gather charge-weighted vector values from mesh using B-splines.

Framework-agnostic launcher for single-system vec3 spline gather.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic positions.

• charges (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – Atomic charges.

• cell_inv_t (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (nx, ny, nz), dtype=wp.vec3f or wp.vec3d) – Vector-valued mesh to interpolate from.

• output (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Charge-weighted interpolated vectors. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.spline_spread(positions, values, cell_inv_t, order, mesh, wp_dtype, device=None)

Spread values from atoms to mesh using B-spline interpolation.

Framework-agnostic launcher for single-system spline spread.

Parameters:

• positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic positions.

• values (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – Values to spread (e.g., charges).

• cell_inv_t (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Transpose of inverse cell matrix.

• order (int) – B-spline order (1-6).

• mesh (wp.array, shape (nx, ny, nz), dtype=wp.float32 or wp.float64) – OUTPUT: Mesh to accumulate values. Must be zero-initialized.

• wp_dtype (type) – Warp scalar dtype (wp.float32 or wp.float64).

• device (str | None) – Warp device string. If None, inferred from arrays.

Return type:

None

nvalchemiops.math.wpdivmod(a, b)

Integer quotient and remainder.

Computes quotient = int(a / b) with truncating division toward zero, then adjusts a negative C-style remainder. For the positive-divisor contract used by the neighbor kernels, this matches divmod(), including for negative dividends (for example, wpdivmod(-5, 3) and Python divmod(-5, 3) both return (-2, 1)).

Parameters:

• a (int) – Dividend.

• b (int) – Positive divisor.

Returns:

• quotient (int) – Floor quotient for the positive divisor.

• remainder (int) – Non-negative remainder in [0, |b|) produced by the adjustment above.

nvalchemiops.math.wp_erfc_input_precision(x)

See the wp.float64 overload.

Parameters:

x (float32)

Return type:

float32

nvalchemiops.math.wp_erfc_f32_fast(x)

Complementary error function approximation for float32.

Uses the Abramowitz and Stegun approximation with maximum error ~1.5e-7. erfc(x) = 1 - erf(x) for x >= 0, and erfc(-x) = 2 - erfc(x) for x < 0.

This is a float32-grade approximation: its own error (~1.5e-7) is close to float32 epsilon, so evaluating it in float64 buys no accuracy, only cost. Do not use it for float64 inputs – use native wp.erfc there, or call wp_erfc_input_precision() to dispatch automatically.

Parameters:

x (Any) – Input value

Returns:

erfc(x) approximation

Return type:

Any

nvalchemiops.math.wp_erfc(x)

Deprecated alias for wp_erfc_f32_fast().

Deprecated since version Kept: only for backward compatibility: existing callers that evaluate this float32-grade approximation at float64 get exactly the same (imprecise, for float64) result they always did. New code should call wp_erfc_input_precision(), which dispatches to native wp.erfc for float64 and to wp_erfc_f32_fast() otherwise.

Parameters:

x (Any)

Return type:

Any

Type Conversion Utilities

PyTorch dtype and device conversion utilities for Warp interop.

nvalchemiops.torch.types.get_wp_dtype(dtype)

Get the Warp scalar dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp scalar dtype (wp.float32, wp.float64, or wp.float16).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.

nvalchemiops.torch.types.get_wp_mat_dtype(dtype)

Get the Warp 3x3 matrix dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp mat33 dtype (wp.mat33f, wp.mat33d, or wp.mat33h).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.

nvalchemiops.torch.types.get_wp_vec_dtype(dtype)

Get the Warp 3-vector dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp vec3 dtype (wp.vec3f, wp.vec3d, or wp.vec3h).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.

nvalchemiops.torch.types.get_wp_dtype(dtype)

Get the Warp scalar dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp scalar dtype (wp.float32, wp.float64, or wp.float16).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.

nvalchemiops.torch.types.get_wp_vec_dtype(dtype)

Get the Warp 3-vector dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp vec3 dtype (wp.vec3f, wp.vec3d, or wp.vec3h).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.

nvalchemiops.torch.types.get_wp_mat_dtype(dtype)

Get the Warp 3x3 matrix dtype corresponding to a PyTorch dtype.

Parameters:

dtype (torch.dtype) – PyTorch floating-point dtype. Supported values: torch.float32, torch.float64, torch.float16.

Returns:

Corresponding Warp mat33 dtype (wp.mat33f, wp.mat33d, or wp.mat33h).

Return type:

wp.type

Raises:

ValueError – If dtype is not one of the supported floating-point dtypes.