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-4).

  • 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-4).

  • 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-4).

  • 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-4).

  • 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.eval_gto_density_pytorch(positions, sigma, L_max=2, device=None)[source]#

Evaluate GTO densities from PyTorch tensors.

Parameters:

  • positions (torch.Tensor) – Input positions [N, 3] as float64.

  • sigma (float) – Gaussian width parameter.

  • L_max (int) – Maximum angular momentum (0, 1, or 2). Default: 2.

  • device (torch.device, optional) – Device for computation.

Returns:

GTO density values [N, num_components].

Return type:

torch.Tensor

nvalchemiops.math.eval_gto_fourier_pytorch(k_vectors, sigma, L_max=2, device=None)[source]#

Evaluate GTO Fourier transforms from PyTorch tensors.

Parameters:

  • k_vectors (torch.Tensor) – Input wave vectors [K, 3] as float64.

  • sigma (float) – Gaussian width parameter.

  • L_max (int) – Maximum angular momentum (0, 1, or 2). Default: 2.

  • device (torch.device, optional) – Device for computation.

Returns:

(real_part, imag_part) each of shape [K, num_components].

Return type:

tuple[torch.Tensor, torch.Tensor]

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-4).

  • 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-4).

  • 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-4).

  • 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-4).

  • 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 dtype for a given torch dtype.

Parameters:

dtype (dtype)

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

Get the warp mat dtype for a given torch dtype.

Parameters:

dtype (dtype)

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

Get the warp vec dtype for a given torch dtype.

Parameters:

dtype (dtype)

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

Get the warp dtype for a given torch dtype.

Parameters:

dtype (dtype)

nvalchemiops.torch.types.get_wp_vec_dtype(dtype)[source]

Get the warp vec dtype for a given torch dtype.

Parameters:

dtype (dtype)

nvalchemiops.torch.types.get_wp_mat_dtype(dtype)[source]

Get the warp mat dtype for a given torch dtype.

Parameters:

dtype (dtype)

On this page
Show Source