Ctrl+K

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

Warp implementation of the divmod utility.

Parameters:

Type Conversion Utilities#

PyTorch dtype and device conversion utilities for Warp interop.

nvalchemiops.torch.types.get_wp_dtype(dtype)[source]#

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)[source]#

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)[source]#

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)[source]

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)[source]

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)[source]

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.

On this page
Show Source