Ctrl+K

nvalchemiops.interactions.electrostatics: Electrostatic Interactions (Warp)#

Electrostatics Interactions Module#

This module provides GPU-accelerated implementations of various methods for computing long-range electrostatic interactions in molecular simulations.

Architecture#

This module provides framework-agnostic Warp kernel launchers. For PyTorch bindings, see nvalchemiops.torch.interactions.electrostatics.

Available Methods#

  1. Coulomb (coulomb) - Direct Coulomb energy and forces - Damped (erfc) Coulomb for Ewald/PME real-space contribution - Warp launchers: coulomb_energy(), coulomb_energy_forces(), etc. - PyTorch API: nvalchemiops.torch.interactions.electrostatics.coulomb

  2. Ewald Summation (ewald) - Classical method splitting interactions into real-space and reciprocal-space - \(O(N^2)\) scaling for explicit k-vectors, good for small systems - Full autograd support

  3. Particle Mesh Ewald (PME) (pme) - FFT-based method for \(O(N \log N)\) scaling - Uses B-spline interpolation for charge assignment - Full autograd support

  4. Damped Shifted Force (DSF) (dsf) - Pairwise \(O(N)\) electrostatic summation - Both potential and forces smoothly vanish at cutoff - Supports geometry-dependent charges (MLIP) - Warp launchers: dsf_csr(), dsf_matrix() - PyTorch API: nvalchemiops.torch.interactions.electrostatics.dsf

Core Warp Kernels#

This module provides the framework-agnostic NVIDIA Warp kernels for electrostatic interactions. These functions operate directly on warp.array objects and can be used to build custom integrators or bindings for other frameworks.

Note

For PyTorch-compatible functions that accept torch.Tensor inputs and support autograd, please see nvalchemiops.torch.interactions.electrostatics: Electrostatics.

Coulomb Kernels#

Direct pairwise Coulomb interactions.

nvalchemiops.interactions.electrostatics.coulomb_energy(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, cutoff, alpha, energies, device=None)[source]#

Launch Coulomb energy kernel using CSR neighbor list format.

This is a framework-agnostic launcher that accepts warp arrays directly. Framework-specific wrappers (PyTorch, JAX) handle tensor-to-warp conversion.

Parameters:

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

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

  • cell (wp.array, shape (1,), dtype=wp.mat33d) – Unit cell matrix.

  • idx_j (wp.array, shape (num_pairs,), dtype=wp.int32) – Destination atom indices in CSR format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (num_pairs,), dtype=wp.vec3i) – Integer unit cell shifts.

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies array in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.coulomb_energy_forces(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, cutoff, alpha, energies, forces, device=None)[source]#

Launch Coulomb energy and forces kernel using CSR neighbor list format.

Parameters:

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

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

  • cell (wp.array, shape (1,), dtype=wp.mat33d) – Unit cell matrix.

  • idx_j (wp.array, shape (num_pairs,), dtype=wp.int32) – Destination atom indices in CSR format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (num_pairs,), dtype=wp.vec3i) – Integer unit cell shifts.

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • forces (wp.array, shape (N,), dtype=wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies and forces arrays in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.coulomb_energy_matrix(positions, charges, cell, neighbor_matrix, neighbor_matrix_shifts, cutoff, alpha, fill_value, energies, device=None)[source]#

Launch Coulomb energy kernel using neighbor matrix format.

Parameters:

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

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

  • cell (wp.array, shape (1,), dtype=wp.mat33d) – Unit cell matrix.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • neighbor_matrix_shifts (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Integer unit cell shifts.

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • fill_value (int) – Value indicating padding in neighbor_matrix.

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies array in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.coulomb_energy_forces_matrix(positions, charges, cell, neighbor_matrix, neighbor_matrix_shifts, cutoff, alpha, fill_value, energies, forces, device=None)[source]#

Launch Coulomb energy and forces kernel using neighbor matrix format.

Parameters:

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

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

  • cell (wp.array, shape (1,), dtype=wp.mat33d) – Unit cell matrix.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • neighbor_matrix_shifts (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Integer unit cell shifts.

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • fill_value (int) – Value indicating padding in neighbor_matrix.

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • forces (wp.array, shape (N,), dtype=wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies and forces arrays in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_coulomb_energy(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, batch_idx, cutoff, alpha, energies, device=None)[source]#

Launch batched Coulomb energy kernel using CSR neighbor list format.

Parameters:

  • positions (wp.array, shape (N,), dtype=wp.vec3d) – Atomic positions (all systems concatenated).

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

  • cell (wp.array, shape (B,), dtype=wp.mat33d) – Unit cell matrices for each system.

  • idx_j (wp.array, shape (num_pairs,), dtype=wp.int32) – Destination atom indices in CSR format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (num_pairs,), dtype=wp.vec3i) – Integer unit cell shifts.

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

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies array in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_coulomb_energy_forces(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, batch_idx, cutoff, alpha, energies, forces, device=None)[source]#

Launch batched Coulomb energy and forces kernel using CSR neighbor list format.

Parameters:

  • positions (wp.array, shape (N,), dtype=wp.vec3d) – Atomic positions (all systems concatenated).

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

  • cell (wp.array, shape (B,), dtype=wp.mat33d) – Unit cell matrices for each system.

  • idx_j (wp.array, shape (num_pairs,), dtype=wp.int32) – Destination atom indices in CSR format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (num_pairs,), dtype=wp.vec3i) – Integer unit cell shifts.

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

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • forces (wp.array, shape (N,), dtype=wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies and forces arrays in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_coulomb_energy_matrix(positions, charges, cell, neighbor_matrix, neighbor_matrix_shifts, batch_idx, cutoff, alpha, fill_value, energies, device=None)[source]#

Launch batched Coulomb energy kernel using neighbor matrix format.

Parameters:

  • positions (wp.array, shape (N,), dtype=wp.vec3d) – Atomic positions (all systems concatenated).

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

  • cell (wp.array, shape (B,), dtype=wp.mat33d) – Unit cell matrices for each system.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • neighbor_matrix_shifts (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Integer unit cell shifts.

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

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • fill_value (int) – Value indicating padding in neighbor_matrix.

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies array in-place.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_coulomb_energy_forces_matrix(positions, charges, cell, neighbor_matrix, neighbor_matrix_shifts, batch_idx, cutoff, alpha, fill_value, energies, forces, device=None)[source]#

Launch batched Coulomb energy and forces kernel using neighbor matrix format.

Parameters:

  • positions (wp.array, shape (N,), dtype=wp.vec3d) – Atomic positions (all systems concatenated).

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

  • cell (wp.array, shape (B,), dtype=wp.mat33d) – Unit cell matrices for each system.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • neighbor_matrix_shifts (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Integer unit cell shifts.

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

  • cutoff (float) – Cutoff distance.

  • alpha (float) – Ewald splitting parameter (0.0 for undamped).

  • fill_value (int) – Value indicating padding in neighbor_matrix.

  • energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

  • forces (wp.array, shape (N,), dtype=wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated and zeroed.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Results are written to energies and forces arrays in-place.

Return type:

None

DSF Kernels#

Damped Shifted Force (DSF) pairwise electrostatic summation.

nvalchemiops.interactions.electrostatics.dsf_csr(positions, charges, idx_j, neighbor_ptr, cutoff, alpha, energy, forces, virial, charge_grad, cell=None, unit_shifts=None, device=None, batch_idx=None, compute_forces=True, compute_virial=False, compute_charge_grad=False, wp_scalar_type=None)[source]#

Launch DSF calculation using CSR neighbor list format.

Handles both periodic and non-periodic systems via optional cell and unit_shifts parameters. Dispatches to single-system or batched kernel based on batch_idx.

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.

  • idx_j (wp.array, shape (M,), dtype=wp.int32) – Destination atom indices in CSR format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • cutoff (float) – Cutoff radius.

  • alpha (float) – Damping parameter (0.0 for shifted-force bare Coulomb).

  • energy (wp.array, shape (num_systems,), dtype=wp.float64) – OUTPUT: Per-system energies. Must be pre-allocated.

  • forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated.

  • virial (wp.array, shape (num_systems,), dtype=wp.mat33f or wp.mat33d) – OUTPUT: Per-system virial tensor. Must be pre-allocated.

  • charge_grad (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Per-atom charge gradients dE/dq_i.

  • cell (wp.array or None) – Unit cell matrices for PBC. If None, non-periodic.

  • unit_shifts (wp.array or None) – Integer unit cell shifts for PBC. Required when cell is not None.

  • device (str, optional) – Warp device. If None, inferred from positions.

  • batch_idx (wp.array or None) – System index for each atom. Must be sorted. If None, single system.

  • compute_forces (bool, default True) – Whether to compute forces.

  • compute_virial (bool, default False) – Whether to compute virial tensor.

  • compute_charge_grad (bool, default False) – Whether to compute charge gradients dE/dq_i.

  • wp_scalar_type (type, optional) – Warp scalar type. If None, inferred from positions.dtype.

Return type:

None

nvalchemiops.interactions.electrostatics.dsf_matrix(positions, charges, neighbor_matrix, cutoff, alpha, fill_value, energy, forces, virial, charge_grad, cell=None, neighbor_matrix_shifts=None, device=None, batch_idx=None, compute_forces=True, compute_virial=False, compute_charge_grad=False, wp_scalar_type=None)[source]#

Launch DSF calculation using neighbor matrix format.

Handles both periodic and non-periodic systems via optional cell and neighbor_matrix_shifts parameters. Dispatches to single-system or batched kernel based on batch_idx.

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.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices. Padding entries have values >= fill_value.

  • cutoff (float) – Cutoff radius.

  • alpha (float) – Damping parameter.

  • fill_value (int) – Value indicating padding in neighbor_matrix.

  • energy (wp.array, shape (num_systems,), dtype=wp.float64) – OUTPUT: Per-system energies. Must be pre-allocated.

  • forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces. Must be pre-allocated.

  • virial (wp.array, shape (num_systems,), dtype=wp.mat33f or wp.mat33d) – OUTPUT: Per-system virial tensor. Must be pre-allocated.

  • charge_grad (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Per-atom charge gradients dE/dq_i.

  • cell (wp.array or None) – Unit cell matrices for PBC. If None, non-periodic.

  • neighbor_matrix_shifts (wp.array2d or None) – Integer unit cell shifts for PBC. Required when cell is not None.

  • device (str, optional) – Warp device. If None, inferred from positions.

  • batch_idx (wp.array or None) – System index for each atom. Must be sorted. If None, single system.

  • compute_forces (bool, default True) – Whether to compute forces.

  • compute_virial (bool, default False) – Whether to compute virial tensor.

  • compute_charge_grad (bool, default False) – Whether to compute charge gradients dE/dq_i.

  • wp_scalar_type (type, optional) – Warp scalar type. If None, inferred from positions.dtype.

Return type:

None

Ewald Kernels#

Ewald summation kernels for real-space and reciprocal-space.

Real-Space:

nvalchemiops.interactions.electrostatics.ewald_real_space_energy(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, alpha, pair_energies, wp_dtype, device=None)[source]#

Launch Ewald real-space energy kernel using CSR neighbor list format.

This is a framework-agnostic launcher that accepts warp arrays directly.

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 (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • idx_j (wp.array, shape (M,), dtype=wp.int32) – Target atom indices (CSR data).

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (M,), dtype=wp.vec3i) – Periodic image shifts.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • pair_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies. Must be pre-allocated and zeroed.

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

  • device (str, optional) – Warp device. If None, inferred from positions.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_real_space_energy_forces(positions, charges, cell, idx_j, neighbor_ptr, unit_shifts, alpha, pair_energies, atomic_forces, virial, wp_dtype, device=None, compute_virial=False)[source]#

Launch Ewald real-space energy and forces kernel using CSR neighbor list.

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 (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • idx_j (wp.array, shape (M,), dtype=wp.int32) – Target atom indices (CSR data).

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (M,), dtype=wp.vec3i) – Periodic image shifts.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • pair_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

  • virial (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – OUTPUT: Virial tensor. Only written when compute_virial=True. Must be pre-allocated and zeroed by caller.

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

  • device (str, optional) – Warp device.

  • compute_virial (bool, optional) – Whether to compute the virial tensor. Default False.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_real_space_energy_matrix(positions, charges, cell, neighbor_matrix, unit_shifts_matrix, mask_value, alpha, pair_energies, wp_dtype, device=None)[source]#

Launch Ewald real-space energy kernel using neighbor matrix format.

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 (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • unit_shifts_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Periodic image shifts.

  • mask_value (int) – Value indicating invalid/padded entries.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • pair_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies.

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

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_real_space_energy_forces_matrix(positions, charges, cell, neighbor_matrix, unit_shifts_matrix, mask_value, alpha, pair_energies, atomic_forces, virial, wp_dtype, device=None, compute_virial=False)[source]#

Launch Ewald real-space energy and forces kernel using neighbor matrix.

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 (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • neighbor_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • unit_shifts_matrix (wp.array2d, shape (N, max_neighbors), dtype=wp.vec3i) – Periodic image shifts.

  • mask_value (int) – Value indicating invalid/padded entries.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • pair_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

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

  • device (str, optional) – Warp device.

  • compute_virial (bool, optional) – Whether to compute the virial tensor. Default False.

  • virial (wp.array, optional) – OUTPUT: Virial tensor. Must be pre-allocated by caller.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_real_space_energy(positions, charges, cell, batch_id, idx_j, neighbor_ptr, unit_shifts, alpha, pair_energies, wp_dtype, device=None)[source]#

Launch batched Ewald real-space energy kernel using CSR neighbor list.

Parameters:

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

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

  • cell (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrices for each system.

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

  • idx_j (wp.array, shape (M,), dtype=wp.int32) – Target atom indices (CSR data).

  • neighbor_ptr (wp.array, shape (N_total+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (M,), dtype=wp.vec3i) – Periodic image shifts.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • pair_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

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

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_real_space_energy_forces(positions, charges, cell, batch_id, idx_j, neighbor_ptr, unit_shifts, alpha, pair_energies, atomic_forces, virial, wp_dtype, device=None, compute_virial=False)[source]#

Launch batched Ewald real-space energy and forces kernel (CSR).

Parameters:

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

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

  • cell (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrices.

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

  • idx_j (wp.array, shape (M,), dtype=wp.int32) – Target atom indices.

  • neighbor_ptr (wp.array, shape (N_total+1,), dtype=wp.int32) – CSR row pointers.

  • unit_shifts (wp.array, shape (M,), dtype=wp.vec3i) – Periodic image shifts.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • pair_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

  • compute_virial (bool, optional) – Whether to compute the virial tensor. Default False.

  • virial (wp.array, optional) – OUTPUT: Virial tensor, shape (B,). If None, a dummy array is created.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_real_space_energy_matrix(positions, charges, cell, batch_id, neighbor_matrix, unit_shifts_matrix, mask_value, alpha, pair_energies, wp_dtype, device=None)[source]#

Launch batched Ewald real-space energy kernel using neighbor matrix.

Parameters:

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

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

  • cell (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrices.

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

  • neighbor_matrix (wp.array2d, shape (N_total, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • unit_shifts_matrix (wp.array2d, shape (N_total, max_neighbors), dtype=wp.vec3i) – Periodic image shifts.

  • mask_value (int) – Value indicating invalid entries.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • pair_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_real_space_energy_forces_matrix(positions, charges, cell, batch_id, neighbor_matrix, unit_shifts_matrix, mask_value, alpha, pair_energies, atomic_forces, virial, wp_dtype, device=None, compute_virial=False)[source]#

Launch batched Ewald real-space energy and forces kernel (matrix).

Parameters:

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

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

  • cell (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrices.

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

  • neighbor_matrix (wp.array2d, shape (N_total, max_neighbors), dtype=wp.int32) – Neighbor indices.

  • unit_shifts_matrix (wp.array2d, shape (N_total, max_neighbors), dtype=wp.vec3i) – Periodic image shifts.

  • mask_value (int) – Value indicating invalid entries.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • pair_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

  • compute_virial (bool, optional) – Whether to compute the virial tensor. Default False.

  • virial (wp.array, optional) – OUTPUT: Virial tensor, shape (B,). If None, a dummy array is created.

Return type:

None

Reciprocal-Space:

nvalchemiops.interactions.electrostatics.ewald_reciprocal_space_fill_structure_factors(positions, charges, k_vectors, cell, alpha, total_charge, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, wp_dtype, device=None)[source]#

Launch kernel to compute structure factors for reciprocal-space Ewald.

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.

  • k_vectors (wp.array, shape (K,), dtype=wp.vec3f or wp.vec3d) – Half-space reciprocal lattice vectors.

  • cell (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • total_charge (wp.array, shape (1,), dtype=wp.float64) – OUTPUT: Q_total/V for background correction.

  • cos_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – OUTPUT: cos(k.r) for each (k, atom) pair.

  • sin_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – OUTPUT: sin(k.r) for each (k, atom) pair.

  • real_structure_factors (wp.array, shape (K,), dtype=wp.float64) – OUTPUT: Real part of weighted structure factors.

  • imag_structure_factors (wp.array, shape (K,), dtype=wp.float64) – OUTPUT: Imaginary part of weighted structure factors.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_reciprocal_space_compute_energy(charges, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, reciprocal_energies, wp_dtype, device=None)[source]#

Launch kernel to compute per-atom reciprocal-space energies.

Parameters:

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

  • cos_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – cos(k.r) from structure factor computation.

  • sin_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – sin(k.r) from structure factor computation.

  • real_structure_factors (wp.array, shape (K,), dtype=wp.float64) – Real structure factors.

  • imag_structure_factors (wp.array, shape (K,), dtype=wp.float64) – Imaginary structure factors.

  • reciprocal_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_subtract_self_energy(charges, alpha, total_charge, energy_in, energy_out, wp_dtype, device=None)[source]#

Launch kernel to apply self-energy and background corrections.

Parameters:

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

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • total_charge (wp.array, shape (1,), dtype=wp.float64) – Q_total/V from structure factor computation.

  • energy_in (wp.array, shape (N,), dtype=wp.float64) – Raw reciprocal-space energies.

  • energy_out (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Corrected energies.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.ewald_reciprocal_space_energy_forces(charges, k_vectors, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, reciprocal_energies, atomic_forces, wp_dtype, device=None)[source]#

Launch kernel to compute reciprocal-space energies and forces.

Parameters:

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

  • k_vectors (wp.array, shape (K,), dtype=wp.vec3f or wp.vec3d) – Reciprocal lattice vectors.

  • cos_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – cos(k.r) from structure factor computation.

  • sin_k_dot_r (wp.array2d, shape (K, N), dtype=wp.float64) – sin(k.r) from structure factor computation.

  • real_structure_factors (wp.array, shape (K,), dtype=wp.float64) – Real structure factors.

  • imag_structure_factors (wp.array, shape (K,), dtype=wp.float64) – Imaginary structure factors.

  • reciprocal_energies (wp.array, shape (N,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_reciprocal_space_fill_structure_factors(positions, charges, k_vectors, cell, alpha, atom_start, atom_end, total_charges, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, num_k, num_systems, max_blocks_per_system, wp_dtype, device=None)[source]#

Launch batched kernel to compute structure factors for reciprocal-space Ewald.

Parameters:

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

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

  • k_vectors (wp.array2d, shape (B, K), dtype=wp.vec3f or wp.vec3d) – Per-system reciprocal lattice vectors.

  • cell (wp.array, shape (B,), dtype=wp.mat33f or wp.mat33d) – Per-system unit cell matrices.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • atom_start (wp.array, shape (B,), dtype=wp.int32) – First atom index for each system.

  • atom_end (wp.array, shape (B,), dtype=wp.int32) – Last atom index (exclusive) for each system.

  • total_charges (wp.array, shape (B,), dtype=wp.float64) – OUTPUT: Per-system Q_total/V.

  • cos_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – OUTPUT: cos(k.r) for each (k, atom) pair.

  • sin_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – OUTPUT: sin(k.r) for each (k, atom) pair.

  • real_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – OUTPUT: Per-system real structure factors.

  • imag_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – OUTPUT: Per-system imaginary structure factors.

  • num_k (int) – Number of k-vectors per system.

  • num_systems (int) – Number of systems in the batch.

  • max_blocks_per_system (int) – Maximum atom blocks per system.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_reciprocal_space_compute_energy(charges, batch_id, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, reciprocal_energies, wp_dtype, device=None)[source]#

Launch batched kernel to compute per-atom reciprocal-space energies.

Parameters:

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

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

  • cos_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – cos(k.r) from structure factor computation.

  • sin_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – sin(k.r) from structure factor computation.

  • real_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – Per-system real structure factors.

  • imag_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – Per-system imaginary structure factors.

  • reciprocal_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_subtract_self_energy(charges, batch_idx, alpha, total_charges, energy_in, energy_out, wp_dtype, device=None)[source]#

Launch batched kernel to apply self-energy and background corrections.

Parameters:

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

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • total_charges (wp.array, shape (B,), dtype=wp.float64) – Per-system Q_total/V.

  • energy_in (wp.array, shape (N_total,), dtype=wp.float64) – Raw reciprocal-space energies.

  • energy_out (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Corrected energies.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

nvalchemiops.interactions.electrostatics.batch_ewald_reciprocal_space_energy_forces(charges, batch_id, k_vectors, cos_k_dot_r, sin_k_dot_r, real_structure_factors, imag_structure_factors, reciprocal_energies, atomic_forces, wp_dtype, device=None)[source]#

Launch batched kernel to compute reciprocal-space energies and forces.

Parameters:

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

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

  • k_vectors (wp.array2d, shape (B, K), dtype=wp.vec3f or wp.vec3d) – Per-system reciprocal lattice vectors.

  • cos_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – cos(k.r) from structure factor computation.

  • sin_k_dot_r (wp.array2d, shape (K, N_total), dtype=wp.float64) – sin(k.r) from structure factor computation.

  • real_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – Per-system real structure factors.

  • imag_structure_factors (wp.array2d, shape (B, K), dtype=wp.float64) – Per-system imaginary structure factors.

  • reciprocal_energies (wp.array, shape (N_total,), dtype=wp.float64) – OUTPUT: Per-atom energies.

  • atomic_forces (wp.array, shape (N_total,), dtype=wp.vec3f or wp.vec3d) – OUTPUT: Per-atom forces.

  • wp_dtype (type) – Warp scalar type.

  • device (str, optional) – Warp device.

Return type:

None

PME Kernels#

Particle Mesh Ewald kernels. Note that FFT operations are typically offloaded to the host framework.

Warning

Keep in mind that currently, convolution required by the PME algorithm needs an FFT interface, and as of the current release, warp does not have an exact analogous method to the API used in PyTorch (e.g. ffttreq). For this reason, the warp kernel set cannot be run completely end-to-end in the same way as the other kernels can be. See the PyTorch bindings instead.

nvalchemiops.interactions.electrostatics.pme_green_structure_factor(k_squared, miller_x, miller_y, miller_z, alpha, volume, mesh_nx, mesh_ny, mesh_nz, spline_order, green_function, structure_factor_sq, wp_dtype, device=None)[source]#

Compute PME Green’s function and B-spline structure factor correction.

Framework-agnostic launcher for single-system Green’s function computation.

Note: FFT Operations Offloaded to Framework#

This kernel computes the Green’s function multipliers for PME. The complete PME reciprocal-space workflow requires FFT operations that are not available in Warp and must be performed by the calling framework. The typical workflow is:

  1. Spread charges to mesh: spline_spread()

  2. Forward FFT: framework.fft.rfftn(mesh) <– Framework-specific

  3. Compute Green’s function: pme_green_structure_factor()

  4. Convolution: mesh_fft * green_function / structure_factor_sq

  5. Inverse FFT: framework.fft.irfftn(…) <– Framework-specific

  6. Gather potential: spline_gather()

  7. Apply corrections: pme_energy_corrections()

param k_squared:

Squared magnitude of k-vectors at each grid point.

type k_squared:

wp.array, shape (Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64

param miller_x:

Miller indices in x direction (from fftfreq).

type miller_x:

wp.array, shape (Nx,), dtype=wp.float32 or wp.float64

param miller_y:

Miller indices in y direction (from fftfreq).

type miller_y:

wp.array, shape (Ny,), dtype=wp.float32 or wp.float64

param miller_z:

Miller indices in z direction (from rfftfreq).

type miller_z:

wp.array, shape (Nz_rfft,), dtype=wp.float32 or wp.float64

param alpha:

Ewald splitting parameter.

type alpha:

wp.array, shape (1,), dtype=wp.float32 or wp.float64

param volume:

Unit cell volume.

type volume:

wp.array, shape (1,), dtype=wp.float32 or wp.float64

param mesh_nx:

Full mesh dimensions (Nz is the full size, not rfft size).

type mesh_nx:

int

param mesh_ny:

Full mesh dimensions (Nz is the full size, not rfft size).

type mesh_ny:

int

param mesh_nz:

Full mesh dimensions (Nz is the full size, not rfft size).

type mesh_nz:

int

param spline_order:

B-spline order (1-4). Order 4 (cubic) recommended.

type spline_order:

int

param green_function:

OUTPUT: Green’s function G(k) at each grid point.

type green_function:

wp.array, shape (Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64

param structure_factor_sq:

OUTPUT: \(|B(k)|^2\) structure factor squared at each grid point.

type structure_factor_sq:

wp.array, shape (Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64

param wp_dtype:

Warp scalar dtype (wp.float32 or wp.float64).

type wp_dtype:

type

param device:

Warp device string. If None, inferred from arrays.

type device:

str | None

See also

nvalchemiops.torch.interactions.electrostatics.pme

Complete PyTorch implementation

Parameters:

  • k_squared (array)

  • miller_x (array)

  • miller_y (array)

  • miller_z (array)

  • alpha (array)

  • volume (array)

  • mesh_nx (int)

  • mesh_ny (int)

  • mesh_nz (int)

  • spline_order (int)

  • green_function (array)

  • structure_factor_sq (array)

  • wp_dtype (type)

  • device (str | None)

Return type:

None

nvalchemiops.interactions.electrostatics.batch_pme_green_structure_factor(k_squared, miller_x, miller_y, miller_z, alpha, volumes, mesh_nx, mesh_ny, mesh_nz, spline_order, green_function, structure_factor_sq, wp_dtype, device=None)[source]#

Compute PME Green’s function and B-spline structure factor for batched systems.

Framework-agnostic launcher for batched Green’s function computation. Each system can have different alpha and volume values, but shares the same mesh dimensions.

Parameters:

  • k_squared (wp.array, shape (B, Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64) – Per-system squared magnitude of k-vectors at each grid point.

  • miller_x (wp.array, shape (Nx,), dtype=wp.float32 or wp.float64) – Miller indices in x direction (shared across systems).

  • miller_y (wp.array, shape (Ny,), dtype=wp.float32 or wp.float64) – Miller indices in y direction (shared across systems).

  • miller_z (wp.array, shape (Nz_rfft,), dtype=wp.float32 or wp.float64) – Miller indices in z direction (shared across systems).

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • volumes (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system unit cell volume.

  • mesh_nx (int) – Full mesh dimensions (Nz is the full size, not rfft size).

  • mesh_ny (int) – Full mesh dimensions (Nz is the full size, not rfft size).

  • mesh_nz (int) – Full mesh dimensions (Nz is the full size, not rfft size).

  • spline_order (int) – B-spline order (1-4). Order 4 (cubic) recommended.

  • green_function (wp.array, shape (B, Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64) – OUTPUT: Per-system Green’s function G_s(k) at each grid point.

  • structure_factor_sq (wp.array, shape (Nx, Ny, Nz_rfft), dtype=wp.float32 or wp.float64) – OUTPUT: \(|B(k)|^2\) structure factor squared (computed only at batch_idx=0).

  • 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

See also

nvalchemiops.torch.interactions.electrostatics.pme

Complete PyTorch implementation

nvalchemiops.interactions.electrostatics.pme_energy_corrections(raw_energies, charges, volume, alpha, total_charge, corrected_energies, wp_dtype, device=None)[source]#

Apply self-energy and background corrections to PME energies.

Framework-agnostic launcher for single-system energy corrections.

Converts raw potential values (φ_i) to corrected per-atom energies by: 1. Multiplying potential by charge: E_pot = q_i * φ_i 2. Subtracting self-energy: E_self = (α/√π) * q_i² 3. Subtracting background: E_bg = (π/(2α²V)) * q_i * Q_total

Final: E_i = q_i * φ_i - (α/√π) * q_i² - (π/(2α²V)) * q_i * Q_total

Parameters:

  • raw_energies (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – Raw potential values φ_i from mesh interpolation.

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

  • volume (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Unit cell volume.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • total_charge (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Sum of all charges (Q_total = ∑_i q_i).

  • corrected_energies (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Corrected per-atom energies.

  • 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.interactions.electrostatics.pme_energy_corrections_with_charge_grad(raw_energies, charges, volume, alpha, total_charge, corrected_energies, charge_gradients, wp_dtype, device=None)[source]#

Apply corrections and compute charge gradients for PME energies.

Framework-agnostic launcher for single-system energy corrections with analytical charge gradient computation.

Computes both corrected energies and analytical charge gradients: - Energy: E_i = q_i * φ_i - (α/√π) * q_i² - (π/(2α²V)) * q_i * Q_total - Charge gradient: ∂E_total/∂q_i = 2*φ_i - 2*(α/√π)*q_i - (π/(α²V))*Q_total

Parameters:

  • raw_energies (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – Raw potential values φ_i from mesh interpolation.

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

  • volume (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Unit cell volume.

  • alpha (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Ewald splitting parameter.

  • total_charge (wp.array, shape (1,), dtype=wp.float32 or wp.float64) – Sum of all charges (Q_total = ∑_i q_i).

  • corrected_energies (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Corrected per-atom energies.

  • charge_gradients (wp.array, shape (N,), dtype=wp.float32 or wp.float64) – OUTPUT: Analytical charge gradients ∂E_total/∂q_i.

  • 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.interactions.electrostatics.batch_pme_energy_corrections(raw_energies, charges, batch_idx, volumes, alpha, total_charges, corrected_energies, wp_dtype, device=None)[source]#

Apply self-energy and background corrections for batched PME.

Framework-agnostic launcher for batched energy corrections. Each atom looks up its system’s parameters via batch_idx.

Parameters:

  • raw_energies (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – Raw potential values φ_i from mesh interpolation.

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

  • batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom (0 to B-1).

  • volumes (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system unit cell volume.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • total_charges (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system sum of charges (Q_s = ∑_{i∈s} q_i).

  • corrected_energies (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – OUTPUT: Corrected per-atom energies.

  • 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.interactions.electrostatics.batch_pme_energy_corrections_with_charge_grad(raw_energies, charges, batch_idx, volumes, alpha, total_charges, corrected_energies, charge_gradients, wp_dtype, device=None)[source]#

Apply corrections and compute charge gradients for batched PME.

Framework-agnostic launcher for batched energy corrections with analytical charge gradient computation.

Parameters:

  • raw_energies (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – Raw potential values φ_i from mesh interpolation.

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

  • batch_idx (wp.array, shape (N_total,), dtype=wp.int32) – System index for each atom (0 to B-1).

  • volumes (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system unit cell volume.

  • alpha (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system Ewald splitting parameter.

  • total_charges (wp.array, shape (B,), dtype=wp.float32 or wp.float64) – Per-system sum of charges (Q_s = ∑_{i∈s} q_i).

  • corrected_energies (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – OUTPUT: Corrected per-atom energies.

  • charge_gradients (wp.array, shape (N_total,), dtype=wp.float32 or wp.float64) – OUTPUT: Analytical charge gradients ∂E_total/∂q_i.

  • 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

On this page
Show Source