# 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.
"""JAX Particle Mesh Ewald (PME) implementation.
This module provides JAX bindings for PME long-range electrostatics calculations.
PME achieves :math:`O(N \\log N)` scaling through FFT-based reciprocal space
computation combined with real-space Ewald summation.
The implementation uses:
- JAX FFT operations (jnp.fft.rfftn/irfftn)
- B-spline interpolation from nvalchemiops.jax.spline
- Ewald real-space from nvalchemiops.jax.interactions.electrostatics.ewald
- Warp kernels for Green's function and energy corrections
Key Functions
-------------
particle_mesh_ewald : Complete PME calculation (real + reciprocal space)
pme_reciprocal_space : Reciprocal-space component only
pme_green_structure_factor : Green's function and structure factor
pme_energy_corrections : Self-energy and background corrections
See Also
--------
nvalchemiops.jax.interactions.electrostatics.ewald : Ewald real-space
nvalchemiops.jax.spline : B-spline interpolation
"""
from __future__ import annotations
import functools
import math
import warnings
from typing import Literal
import jax
import jax.numpy as jnp
import warp as wp
from jax.custom_derivatives import SymbolicZero
from warp import JaxCallableGraphMode, jax_callable
from nvalchemiops.interactions.electrostatics._pme_mesh import (
_DEFAULT_PME_SPLINE_ORDER,
)
from nvalchemiops.interactions.electrostatics.pme_factory import get_pme_kernel
from nvalchemiops.interactions.electrostatics.pme_kernels import (
_batch_pme_green_structure_factor_kernel_overload,
_pme_green_structure_factor_kernel_overload,
_pme_virial_bg_apply_kernel_overload,
_pme_virial_bg_reduce_kernel_overload,
)
from nvalchemiops.jax._lazy_jax_kernels import (
make_jax_kernel_factory,
make_jax_kernels,
)
from nvalchemiops.jax.interactions.electrostatics._autograd import (
_inject_charge_grad,
)
from nvalchemiops.jax.interactions.electrostatics._utils import (
_apply_energy_reduction,
_build_electrostatic_result,
_component_direct_output_deprecation_msg,
_direct_output_deprecation_msg,
_prepare_cell,
_system_sum_from_atoms,
_validate_energy_reduction,
)
from nvalchemiops.jax.interactions.electrostatics.ewald import (
ewald_real_space,
)
from nvalchemiops.jax.interactions.electrostatics.k_vectors import (
generate_k_squared_pme,
generate_k_vectors_pme,
)
from nvalchemiops.jax.interactions.electrostatics.parameters import (
estimate_pme_mesh_dimensions,
estimate_pme_parameters,
mesh_spacing_to_dimensions,
)
from nvalchemiops.jax.interactions.electrostatics.slab import (
_prepare_pbc_for_slab,
_slab_correction_energy_autodiff,
)
from nvalchemiops.jax.interactions.electrostatics.slab import (
compute_slab_correction as _compute_slab_correction,
)
from nvalchemiops.jax.spline import (
_spline_gather_gradient_position_hessian,
_spline_gather_with_force,
_spline_spread_gradient_weights,
spline_gather,
spline_gather_gradient,
spline_spread,
)
from nvalchemiops.jax.types import normalize_float_dtype
__all__ = [
"particle_mesh_ewald",
"pme_reciprocal_space",
"pme_green_structure_factor",
"pme_energy_corrections",
"pme_energy_corrections_with_charge_grad",
"compute_bspline_moduli_1d",
]
# ==============================================================================
# Helper Function for JAX Kernel Creation
# ==============================================================================
# ``make_jax_kernels`` returns a lazy dict (see nvalchemiops.jax._lazy_jax_kernels) that
# materializes its ``jax_kernel`` entries on first __getitem__. Prefer
# ``jax_kernel`` for single-launch ops; use ``jax_callable`` only when
# fusing multiple wp.launch calls into one FFI thunk
# (see :func:`_make_jax_pme_virial_bg_fused`).
def _jax_pme_factory_component(
component: str,
output_names: list[str],
*,
batched: bool = False,
charge_grad: bool = False,
):
"""Return a lazy JAX wrapper for a factory-backed PME component."""
return make_jax_kernel_factory(
lambda wp_dtype: get_pme_kernel(
wp_dtype,
component=component,
batched=batched,
charge_grad=charge_grad,
),
len(output_names),
output_names,
)
# ==============================================================================
# JAX Kernel Wrappers
# ==============================================================================
# Single-system kernels
_jax_pme_green_sf = make_jax_kernels(
_pme_green_structure_factor_kernel_overload,
2,
["green_function", "structure_factor_sq"],
)
_jax_pme_energy_corrections = _jax_pme_factory_component(
"pme_corrections",
["corrected_energies"],
)
_jax_pme_energy_corrections_charge_grad = _jax_pme_factory_component(
"pme_corrections",
["corrected_energies", "charge_gradients"],
charge_grad=True,
)
# Batch kernels
_jax_batch_pme_green_sf = make_jax_kernels(
_batch_pme_green_structure_factor_kernel_overload,
2,
["green_function", "structure_factor_sq"],
)
_jax_batch_pme_energy_corrections = _jax_pme_factory_component(
"pme_corrections",
["corrected_energies"],
batched=True,
)
_jax_batch_pme_energy_corrections_charge_grad = _jax_pme_factory_component(
"pme_corrections",
["corrected_energies", "charge_gradients"],
batched=True,
charge_grad=True,
)
# Fused convolve — replaces the older Green's-function + multiply path with
# a single warp kernel that computes G(k), the B-spline structure factor
# correction C^2(k), and multiplies mesh_fft → convolved_mesh in one launch.
# (mirrors the fused-convolve path in the torch bindings.)
_jax_pme_convolve = _jax_pme_factory_component(
"pme_convolve",
["convolved_mesh"],
)
_jax_batch_pme_convolve = _jax_pme_factory_component(
"pme_convolve",
["convolved_mesh"],
batched=True,
)
# Two-pass virial background correction. Pass 1 reduces per-atom
# charges into per-system total charges (atomic_add). Pass 2 computes
# E_bg = π Q² / (2 α² V) per system and subtracts it from the three diagonal
# entries of ``virial_in``. Mirrors the torch path's ``pme_virial_bg_correction``.
_jax_pme_virial_bg_reduce = make_jax_kernels(
_pme_virial_bg_reduce_kernel_overload,
1,
["total_charges"],
)
_jax_pme_virial_bg_apply = make_jax_kernels(
_pme_virial_bg_apply_kernel_overload,
1,
["virial_out"],
)
# Fuse the two-pass virial bg correction (reduce + apply) into one XLA FFI
# call via jax_callable so JAX can CUDA-graph it with the surrounding warp
# ops + FFTs. Per-pass jax_kernel thunks above remain for direct test use.
def _make_jax_pme_virial_bg_fused(wp_dtype):
reduce_overload = _pme_virial_bg_reduce_kernel_overload[wp_dtype]
apply_overload = _pme_virial_bg_apply_kernel_overload[wp_dtype]
def _fn(
# inputs
charges: wp.array(dtype=wp_dtype),
batch_idx: wp.array(dtype=wp.int32),
cell: wp.array3d(dtype=wp_dtype),
alpha: wp.array(dtype=wp_dtype),
virial_in: wp.array3d(dtype=wp_dtype),
# in-out: zero-initialized by caller, scatter-add target for pass 1
total_charges: wp.array(dtype=wp_dtype),
# outputs
virial_out: wp.array3d(dtype=wp_dtype),
):
# Reference closure-captured names so they survive as
# ``__closure__`` cells. ``from __future__ import annotations``
# stringifies the annotations, so they don't on their own pull
# ``wp_dtype`` into the closure -- and warp's annotation eval
# then can't resolve it. Touching the names in the body fixes that.
_ = wp_dtype
wp.launch(
reduce_overload,
dim=charges.shape,
inputs=[charges, batch_idx],
outputs=[total_charges],
)
wp.launch(
apply_overload,
dim=total_charges.shape,
inputs=[total_charges, cell, alpha, virial_in],
outputs=[virial_out],
)
return jax_callable(
_fn,
num_outputs=2,
in_out_argnames=["total_charges"],
graph_mode=JaxCallableGraphMode.JAX,
)
_jax_pme_virial_bg_fused = {
jnp.float32: _make_jax_pme_virial_bg_fused(wp.float32),
jnp.float64: _make_jax_pme_virial_bg_fused(wp.float64),
}
# ==============================================================================
# Public API Functions
# ==============================================================================
[docs]
def compute_bspline_moduli_1d(
miller_indices: jax.Array,
mesh_N: int,
spline_order: int,
) -> jax.Array:
"""Precompute a 1D B-spline modulus LUT for one PME mesh axis.
Returns ``b[i] = sinc(m_i / N)^spline_order`` for each Miller index
``m_i`` (with ``sinc(x) = sin(pi*x)/(pi*x)``, ``sinc(0) = 1``). The
three-axis product ``b_x[i] * b_y[j] * b_z[k]`` is the B-spline
structure factor consumed by :func:`pme_fused_convolve`. Precomputing the
LUT lets the convolve kernel replace three sinc transcendentals + an
order-dependent power loop per (i, j, k) thread with three reads + two
multiplies.
Parameters
----------
miller_indices : jax.Array, shape (N,)
Integer Miller indices for one mesh axis, e.g. from
``jnp.fft.fftfreq(N, d=1.0/N)`` or ``jnp.fft.rfftfreq(N, d=1.0/N)``.
mesh_N : int
Number of mesh points along this axis.
spline_order : int
B-spline interpolation order (e.g. 4 for cubic B-splines).
Returns
-------
jax.Array, shape (N,)
Per-Miller-index B-spline modulus ``sinc(m/N)^spline_order``.
"""
# sinc(x) for x in [-0.5, 0.5] is bounded in [2/pi, 1], so s^spline_order
# (for orders 2-6) stays well within fp32 range. Stay in the input dtype
# to avoid an fp32 -> fp64 -> fp32 round-trip per call. jax.numpy.sinc uses
# the normalized convention sinc(pi*x)/(pi*x); matches torch.
arg = miller_indices / float(mesh_N)
s = jnp.sinc(arg)
return s**spline_order
def pme_fused_convolve(
mesh_fft: jax.Array,
k_squared: jax.Array,
moduli_x: jax.Array,
moduli_y: jax.Array,
moduli_z: jax.Array,
alpha: jax.Array,
volume: jax.Array,
is_batch: bool,
) -> jax.Array:
"""Fused Green's function + structure-factor multiply, single launch.
Replaces (compute G(k), compute C^2(k), divide mesh_fft by C^2, multiply
by G(k)) with a single warp kernel. ``moduli_x/y/z`` are precomputed
1D B-spline modulus LUTs (``sinc(m/N)^spline_order`` per axis); see
``compute_bspline_moduli_1d``.
Parameters
----------
mesh_fft : complex64 or complex128
FFT of the charge mesh. Shape (Nx, Ny, Nz_rfft) for single system
or (B, Nx, Ny, Nz_rfft) for batch.
k_squared : float32 or float64
|k|^2 at each grid point. Same leading shape as mesh_fft, or a
single shared grid when every system uses the same reciprocal cell.
moduli_x, moduli_y, moduli_z : float32 or float64
Per-axis B-spline modulus LUTs.
alpha : float32 or float64
Ewald splitting parameter. Shape (1,) or (B,).
volume : float32 or float64
Cell volume. Shape (1,) or (B,).
is_batch : bool
Whether this is a batched call.
Returns
-------
convolved_mesh : complex64 or complex128, same shape as mesh_fft.
"""
real_dtype = jnp.float32 if mesh_fft.dtype == jnp.complex64 else jnp.float64
complex_dtype = mesh_fft.dtype
input_dtype = normalize_float_dtype(real_dtype)
# generate_k_vectors_pme squeezes the batch dim when B=1 — restore it
# for the batch kernel which expects (B, nx, ny, nz_r).
squeeze_output = False
if is_batch and mesh_fft.ndim == 3:
mesh_fft = mesh_fft[jnp.newaxis, ...]
squeeze_output = True
if is_batch:
k_squared = jnp.broadcast_to(k_squared, mesh_fft.shape)
# Reinterpret complex (N, ..., M) -> real (N, ..., M, 2) for the
# vec2-typed warp kernel. jax.lax.bitcast_convert_type doesn't accept
# complex→float, so use .view() (doubles the trailing dim) and then
# reshape to add the explicit vec2 axis.
mesh_fft_real = mesh_fft.view(real_dtype).reshape(*mesh_fft.shape, 2)
# Ensure alpha / volume are 1-D arrays of the right dtype.
alpha = alpha.astype(real_dtype)
volume = volume.astype(real_dtype)
if alpha.ndim == 0:
alpha = alpha.reshape(1)
if volume.ndim == 0:
volume = volume.reshape(1)
moduli_x = moduli_x.astype(real_dtype)
moduli_y = moduli_y.astype(real_dtype)
moduli_z = moduli_z.astype(real_dtype)
k_squared = k_squared.astype(real_dtype)
# Pre-allocate output (same shape as mesh_fft_real, treated as in-out
# by the warp kernel).
convolved_real = jnp.zeros_like(mesh_fft_real)
if is_batch:
kernel = _jax_batch_pme_convolve[input_dtype]
else:
kernel = _jax_pme_convolve[input_dtype]
# Launch dims match the spectrum shape (drop the trailing length-2 vec2 dim).
launch_dims = mesh_fft.shape
(convolved_real,) = kernel(
mesh_fft_real,
k_squared,
moduli_x,
moduli_y,
moduli_z,
alpha,
volume,
convolved_real,
launch_dims=launch_dims,
)
# Reverse the reshape+view: collapse the trailing vec2 dim, then view
# as complex (which halves the trailing dim back to the original).
convolved_flat = convolved_real.reshape(*mesh_fft.shape[:-1], -1)
convolved = convolved_flat.view(complex_dtype)
if squeeze_output:
convolved = convolved.squeeze(0)
return convolved
def pme_green_structure_factor(
k_squared: jax.Array,
mesh_dimensions: tuple[int, int, int],
alpha: jax.Array,
cell: jax.Array,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
batch_idx: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]:
r"""Compute Green's function and B-spline structure factor correction.
Computes the Coulomb Green's function with volume normalization and the
B-spline aliasing correction factor for PME.
Green's function (volume-normalized):
.. math::
G(\mathbf{k}) = \frac{2\pi}{V} \frac{\exp(-k^2 / (4\alpha^2))}{k^2}
Structure factor correction (for B-spline deconvolution):
.. math::
C^2(\mathbf{k}) = \left[\operatorname{sinc}(m_x/N_x) \cdot
\operatorname{sinc}(m_y/N_y) \cdot \operatorname{sinc}(m_z/N_z)\right]^{2p}
where :math:`p` is the spline order.
Parameters
----------
k_squared : jax.Array
:math:`|k|^2` values at each FFT grid point.
- Single-system: shape (Nx, Ny, Nz_rfft)
- Batch: shape (B, Nx, Ny, Nz_rfft)
mesh_dimensions : tuple[int, int, int]
Full mesh dimensions (Nx, Ny, Nz) before rfft.
alpha : jax.Array
Ewald splitting parameter.
- Single-system: shape (1,) or scalar
- Batch: shape (B,)
cell : jax.Array
Unit cell matrices.
- Single-system: shape (3, 3) or (1, 3, 3)
- Batch: shape (B, 3, 3)
spline_order : int, default=5
B-spline interpolation order (5 for quartic B-splines by default).
batch_idx : jax.Array | None, default=None
If provided, dispatches to batch kernels.
Returns
-------
green_function : jax.Array
Volume-normalized Green's function G(k).
- Single-system: shape (Nx, Ny, Nz_rfft)
- Batch: shape (B, Nx, Ny, Nz_rfft)
structure_factor_sq : jax.Array
Squared structure factor :math:`C^2(k)` for B-spline deconvolution.
Shape (Nx, Ny, Nz_rfft), shared across batch.
Notes
-----
- G(k=0) is set to zero to avoid singularity
- The volume normalization in G(k) eliminates later divisions
- Structure factor is mesh-dependent only, so shared across batch
- This compatibility helper returns raw ``G(k)``. The fused convolve path
folds deconvolution internally as ``G(k) / C^2(k)``.
"""
mesh_nx, mesh_ny, mesh_nz = mesh_dimensions
input_dtype = normalize_float_dtype(k_squared.dtype)
# Ensure cell is correct shape
if cell.ndim == 2:
cell = cell[jnp.newaxis, :, :]
volume = jnp.abs(jnp.linalg.det(cell)).astype(input_dtype)
# Generate Miller indices using JAX FFT frequency functions
# Use d=1.0/n to get integer Miller indices
miller_x = jnp.fft.fftfreq(mesh_nx, d=1.0 / mesh_nx).astype(input_dtype)
miller_y = jnp.fft.fftfreq(mesh_ny, d=1.0 / mesh_ny).astype(input_dtype)
miller_z = jnp.fft.rfftfreq(mesh_nz, d=1.0 / mesh_nz).astype(input_dtype)
# Ensure alpha is 1D array
if alpha.ndim == 0:
alpha = alpha.reshape(1)
alpha = alpha.astype(input_dtype)
# Get kernel for input dtype
if batch_idx is None:
# Single system
kernel = _jax_pme_green_sf[input_dtype]
# Allocate outputs
green_function = jnp.zeros(
(mesh_nx, mesh_ny, mesh_nz // 2 + 1), dtype=input_dtype
)
structure_factor_sq = jnp.zeros(
(mesh_nx, mesh_ny, mesh_nz // 2 + 1), dtype=input_dtype
)
# Launch kernel
green_out, sf_out = kernel(
k_squared.astype(input_dtype),
miller_x,
miller_y,
miller_z,
alpha,
volume,
int(mesh_nx),
int(mesh_ny),
int(mesh_nz),
int(spline_order),
green_function,
structure_factor_sq,
launch_dims=(mesh_nx, mesh_ny, mesh_nz // 2 + 1),
)
return green_out, sf_out
else:
# Batch
num_systems = cell.shape[0]
kernel = _jax_batch_pme_green_sf[input_dtype]
# Ensure k_squared has batch dimension for batch kernels
k_sq = k_squared.astype(input_dtype)
if k_sq.ndim == 3:
k_sq = jnp.broadcast_to(
k_sq[jnp.newaxis], (num_systems, mesh_nx, mesh_ny, mesh_nz // 2 + 1)
)
# Allocate outputs
green_function = jnp.zeros(
(num_systems, mesh_nx, mesh_ny, mesh_nz // 2 + 1), dtype=input_dtype
)
structure_factor_sq = jnp.zeros(
(mesh_nx, mesh_ny, mesh_nz // 2 + 1), dtype=input_dtype
)
# Launch kernel
green_out, sf_out = kernel(
k_sq,
miller_x,
miller_y,
miller_z,
alpha,
volume,
int(mesh_nx),
int(mesh_ny),
int(mesh_nz),
int(spline_order),
green_function,
structure_factor_sq,
launch_dims=(num_systems, mesh_nx, mesh_ny, mesh_nz // 2 + 1),
)
return green_out, sf_out
def pme_virial_bg_correction(
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
virial: jax.Array,
batch_idx: jax.Array | None = None,
volume: jax.Array | None = None,
) -> jax.Array:
r"""Apply non-neutral background virial correction in a single Warp launch.
Two-pass fused kernel:
1. Reduce per-atom ``charges`` into per-system totals (atomic_add).
2. Compute :math:`E_\text{bg} = \pi Q^2 / (2 \alpha^2 V)` per system and subtract it from
the three diagonal entries of ``virial`` (off-diagonal unchanged).
Single-system inputs are fanned out via ``batch_idx`` filled with zeros.
Parameters
----------
charges : jax.Array, shape (N,)
Per-atom charges.
cell : jax.Array, shape (3, 3) or (B, 3, 3)
Unit cell. Single-system 2D is promoted to (1, 3, 3).
alpha : jax.Array, shape () / (1,) / (B,)
Per-system Ewald splitting parameter.
virial : jax.Array, shape (1, 3, 3) or (B, 3, 3)
Virial tensor to correct in place (functional return).
batch_idx : jax.Array | None, shape (N,), dtype=int32, optional
System index per atom. If None, every atom maps to system 0.
volume : jax.Array | None, optional
Precomputed per-system cell volume. Treated as static metadata when
supplied.
Returns
-------
virial_out : jax.Array, same shape as ``virial``
Background-corrected virial.
"""
input_dtype = normalize_float_dtype(charges.dtype)
cell_w = cell.astype(input_dtype)
if cell_w.ndim == 2:
cell_w = cell_w[jnp.newaxis, :, :]
num_systems = cell_w.shape[0]
num_atoms = charges.shape[0]
alpha_w = alpha.astype(input_dtype)
if alpha_w.ndim == 0:
alpha_w = alpha_w.reshape(1)
if alpha_w.shape[0] == 1 and num_systems > 1:
alpha_w = jnp.broadcast_to(alpha_w, (num_systems,))
if batch_idx is None:
bidx = jnp.zeros(num_atoms, dtype=jnp.int32)
else:
bidx = batch_idx.astype(jnp.int32)
virial_w = virial.astype(input_dtype)
if virial_w.ndim == 2:
virial_w = virial_w[jnp.newaxis, :, :]
total_charges = jnp.zeros(num_systems, dtype=input_dtype)
if volume is not None:
total_charges = total_charges.at[bidx].add(charges.astype(input_dtype))
volume_w = volume.astype(input_dtype)
if volume_w.ndim == 0:
volume_w = volume_w.reshape(1)
if volume_w.shape[0] == 1 and num_systems > 1:
volume_w = jnp.broadcast_to(volume_w, (num_systems,))
bg_energy = (
jnp.pi
* total_charges
* total_charges
/ (2.0 * alpha_w * alpha_w * volume_w)
)
eye = jnp.eye(3, dtype=input_dtype)
diag_delta = bg_energy[:, None, None] * eye[jnp.newaxis, :, :]
return virial_w - diag_delta
# Fused two-pass via jax_callable: a single XLA FFI thunk that runs both
# the scatter-add reduce (pass 1) and the per-system E_bg apply (pass 2).
# This collapses what used to be 2 jax_kernel thunks into 1 and lets JAX
# CUDA-graph-capture both passes together.
fused = _jax_pme_virial_bg_fused[input_dtype]
_total_charges, virial_out = fused(
charges.astype(input_dtype),
bidx,
cell_w,
alpha_w,
virial_w,
total_charges, # in/out — zero-initialized scatter-add target
output_dims={"virial_out": virial_w.shape},
)
return virial_out
def pme_energy_corrections(
raw_energies: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
batch_idx: jax.Array | None = None,
volume: jax.Array | None = None,
) -> jax.Array:
r"""Apply self-energy and background corrections to PME energies.
Converts raw interpolated potential to energy and subtracts corrections:
.. math::
E_i = q_i \varphi_i - E_{\text{self},i} - E_{\text{bg},i}
Self-energy correction (removes Gaussian self-interaction):
.. math::
E_{\text{self},i} = \frac{\alpha}{\sqrt{\pi}} q_i^2
Background correction (for non-neutral systems):
.. math::
E_{\text{bg},i} = \frac{\pi}{2\alpha^2 V} q_i Q_\text{total}
Parameters
----------
raw_energies : jax.Array, shape (N,) or (N_total,)
Raw potential values :math:`\varphi_i` from mesh interpolation.
charges : jax.Array, shape (N,) or (N_total,)
Atomic charges.
cell : jax.Array
Unit cell matrices.
- Single-system: shape (3, 3) or (1, 3, 3)
- Batch: shape (B, 3, 3)
alpha : jax.Array
Ewald splitting parameter.
- Single-system: shape (1,) or scalar
- Batch: shape (B,)
batch_idx : jax.Array | None, default=None
System index for each atom. If provided, uses batch kernels. Atoms must
be grouped by system: ``batch_idx`` must be contiguous, nondecreasing,
and use system IDs ``0..B-1``.
volume : jax.Array | None, optional
Precomputed per-system cell volume. Treated as static setup metadata
when supplied.
Returns
-------
corrected_energies : jax.Array, shape (N,) or (N_total,)
Final per-atom reciprocal-space energy with corrections applied.
Notes
-----
- For neutral systems, background correction is zero
- Supports both float32 and float64 dtypes
"""
input_dtype = normalize_float_dtype(raw_energies.dtype)
num_atoms = raw_energies.shape[0]
# Ensure alpha is 1D array
if alpha.ndim == 0:
alpha = alpha.reshape(1)
alpha = alpha.astype(input_dtype)
if batch_idx is None:
# Single system
kernel = _jax_pme_energy_corrections[input_dtype]
# Ensure cell is correct shape
if cell.ndim == 2:
cell = cell[jnp.newaxis, :, :]
if volume is None:
volume = jnp.abs(jnp.linalg.det(cell)).astype(input_dtype)
else:
volume = volume.astype(input_dtype)
if volume.ndim == 0:
volume = volume.reshape(1)
total_charge = charges.sum().reshape(1).astype(input_dtype)
# Allocate output
corrected_energies = jnp.zeros(num_atoms, dtype=input_dtype)
batch_idx_dummy = jnp.zeros((num_atoms,), dtype=jnp.int32)
charge_gradients_dummy = jnp.zeros(num_atoms, dtype=input_dtype)
# Launch kernel
(corrected_out,) = kernel(
raw_energies.astype(input_dtype),
charges.astype(input_dtype),
batch_idx_dummy,
volume,
alpha,
total_charge,
corrected_energies,
charge_gradients_dummy,
launch_dims=(num_atoms,),
)
return corrected_out
else:
# Batch
kernel = _jax_batch_pme_energy_corrections[input_dtype]
num_systems = cell.shape[0] if cell.ndim == 3 else 1
if cell.ndim == 2:
cell = cell[jnp.newaxis, :, :]
if volume is None:
volumes = jnp.abs(jnp.linalg.det(cell)).astype(input_dtype)
else:
volumes = volume.astype(input_dtype)
if volumes.ndim == 0:
volumes = volumes.reshape(1)
if volumes.shape[0] == 1 and num_systems > 1:
volumes = jnp.broadcast_to(volumes, (num_systems,))
# Compute total charge per system
total_charges = jnp.zeros(num_systems, dtype=input_dtype)
total_charges = total_charges.at[batch_idx].add(charges.astype(input_dtype))
# Allocate output
corrected_energies = jnp.zeros(num_atoms, dtype=input_dtype)
charge_gradients_dummy = jnp.zeros(num_atoms, dtype=input_dtype)
# Launch kernel
(corrected_out,) = kernel(
raw_energies.astype(input_dtype),
charges.astype(input_dtype),
batch_idx.astype(jnp.int32),
volumes,
alpha,
total_charges,
corrected_energies,
charge_gradients_dummy,
launch_dims=(num_atoms,),
)
return corrected_out
def pme_energy_corrections_with_charge_grad(
raw_energies: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
batch_idx: jax.Array | None = None,
volume: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]:
r"""Apply energy corrections and compute charge gradients.
Same as pme_energy_corrections but also returns :math:`\partial E/\partial q`
for each atom.
Parameters
----------
raw_energies : jax.Array, shape (N,) or (N_total,)
Raw potential values :math:`\varphi_i` from mesh interpolation.
charges : jax.Array, shape (N,) or (N_total,)
Atomic charges.
cell : jax.Array
Unit cell matrices.
- Single-system: shape (3, 3) or (1, 3, 3)
- Batch: shape (B, 3, 3)
alpha : jax.Array
Ewald splitting parameter.
- Single-system: shape (1,) or scalar
- Batch: shape (B,)
batch_idx : jax.Array | None, default=None
System index for each atom. If provided, uses batch kernels. Atoms must
be grouped by system: ``batch_idx`` must be contiguous, nondecreasing,
and use system IDs ``0..B-1``.
volume : jax.Array | None, optional
Precomputed per-system cell volume. Treated as static setup metadata
when supplied.
Returns
-------
corrected_energies : jax.Array, shape (N,) or (N_total,)
Final per-atom reciprocal-space energy with corrections applied.
charge_gradients : jax.Array, shape (N,) or (N_total,)
Per-atom charge gradients :math:`\partial E/\partial q`.
Notes
-----
- Useful for training models that predict partial charges
- Supports both float32 and float64 dtypes
"""
input_dtype = normalize_float_dtype(raw_energies.dtype)
num_atoms = raw_energies.shape[0]
# Ensure alpha is 1D array
if alpha.ndim == 0:
alpha = alpha.reshape(1)
alpha = alpha.astype(input_dtype)
if batch_idx is None:
# Single system
kernel = _jax_pme_energy_corrections_charge_grad[input_dtype]
# Ensure cell is correct shape
if cell.ndim == 2:
cell = cell[jnp.newaxis, :, :]
if volume is None:
volume = jnp.abs(jnp.linalg.det(cell)).astype(input_dtype)
else:
volume = volume.astype(input_dtype)
if volume.ndim == 0:
volume = volume.reshape(1)
total_charge = charges.sum().reshape(1).astype(input_dtype)
# Allocate outputs
corrected_energies = jnp.zeros(num_atoms, dtype=input_dtype)
charge_gradients = jnp.zeros(num_atoms, dtype=input_dtype)
batch_idx_dummy = jnp.zeros((num_atoms,), dtype=jnp.int32)
# Launch kernel
corrected_out, charge_grad_out = kernel(
raw_energies.astype(input_dtype),
charges.astype(input_dtype),
batch_idx_dummy,
volume,
alpha,
total_charge,
corrected_energies,
charge_gradients,
launch_dims=(num_atoms,),
)
return corrected_out, charge_grad_out
else:
# Batch
kernel = _jax_batch_pme_energy_corrections_charge_grad[input_dtype]
num_systems = cell.shape[0] if cell.ndim == 3 else 1
if cell.ndim == 2:
cell = cell[jnp.newaxis, :, :]
if volume is None:
volumes = jnp.abs(jnp.linalg.det(cell)).astype(input_dtype)
else:
volumes = volume.astype(input_dtype)
if volumes.ndim == 0:
volumes = volumes.reshape(1)
if volumes.shape[0] == 1 and num_systems > 1:
volumes = jnp.broadcast_to(volumes, (num_systems,))
# Compute total charge per system
total_charges = jnp.zeros(num_systems, dtype=input_dtype)
total_charges = total_charges.at[batch_idx].add(charges.astype(input_dtype))
# Allocate outputs
corrected_energies = jnp.zeros(num_atoms, dtype=input_dtype)
charge_gradients = jnp.zeros(num_atoms, dtype=input_dtype)
# Launch kernel
corrected_out, charge_grad_out = kernel(
raw_energies.astype(input_dtype),
charges.astype(input_dtype),
batch_idx.astype(jnp.int32),
volumes,
alpha,
total_charges,
corrected_energies,
charge_gradients,
launch_dims=(num_atoms,),
)
return corrected_out, charge_grad_out
def _compute_pme_reciprocal_virial(
mesh_fft_raw: jax.Array,
convolved_mesh: jax.Array,
k_vectors: jax.Array,
k_squared: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int],
is_batch: bool,
) -> jax.Array:
r"""Compute PME reciprocal-space virial tensor in k-space.
Uses the exact spectral pair from the pipeline (mesh_fft_raw before
deconvolution, and convolved_mesh after Green's function multiplication)
to compute the per-k energy density directly via Parseval's theorem.
The virial per k-point is W_ab(k) = E_k * sigma_ab(k) where:
- E_k = prefactor * weight(k) * Re(mesh_fft_raw(k) * convolved_mesh(k)*)
- sigma_ab(k) = delta_ab - 2*k_a*k_b/k^2 * (1 + k^2/(4*alpha^2))
(sign reflects :math:`W = -dE/d\varepsilon` convention)
Parameters
----------
mesh_fft_raw : jax.Array
Raw rfftn output before B-spline deconvolution.
Shape (nx, ny, nz//2+1) or (B, nx, ny, nz//2+1), complex.
convolved_mesh : jax.Array
Deconvolved mesh FFT multiplied by Green's function: (mesh_fft/B^2)*G.
Shape matching mesh_fft_raw.
k_vectors : jax.Array
k-vectors on the mesh. Shape (..., nx, ny, nz//2+1, 3).
k_squared : jax.Array
|k|^2. Shape (..., nx, ny, nz//2+1).
alpha : jax.Array
Ewald splitting parameter.
mesh_dimensions : tuple
(nx, ny, nz).
is_batch : bool
Whether this is a batched calculation.
Returns
-------
virial : jax.Array, shape (B, 3, 3) or (1, 3, 3)
Per-system virial tensor.
"""
mesh_nx, mesh_ny, mesh_nz = mesh_dimensions
# Determine accumulation dtype from k_squared (float32 or float64)
acc_dtype = normalize_float_dtype(k_squared.dtype)
complex_dtype = jnp.complex64 if acc_dtype == jnp.float32 else jnp.complex128
# Per-k energy density from exact pipeline spectral pair.
# Re(mesh_fft_raw * convolved_mesh*) = |mesh_fft_raw|^2 * G / B^2
fft_raw_cast = mesh_fft_raw.astype(complex_dtype)
conv_cast = convolved_mesh.astype(complex_dtype)
energy_density = (fft_raw_cast * jnp.conj(conv_cast)).real
# Weight for rfft symmetry: 2 for interior k_z, 1 for boundary
weight = jnp.full_like(energy_density, 2.0)
weight = weight.at[..., 0].set(1.0) # k_z = 0
if mesh_nz % 2 == 0:
weight = weight.at[..., -1].set(1.0) # k_z = nz//2 (Nyquist)
# Weighted energy density
weighted_energy = weight * energy_density
# Virial W = -dE/dε, so sigma_ab = delta_ab - 2*k_a*k_b/k^2 * (1 + k^2/(4*alpha^2))
k_sq_acc = k_squared.astype(acc_dtype)
alpha_acc = alpha.astype(acc_dtype)
# generate_k_vectors_pme squeezes the batch dim when B=1; restore it so
# the batched einsum and sum_dims=(1,2,3) operate on the correct axes.
if is_batch and k_sq_acc.ndim == 3:
k_sq_acc = jnp.expand_dims(k_sq_acc, axis=0)
# Handle alpha broadcasting: alpha may be (B,) for batch
if is_batch and alpha_acc.ndim == 1:
alpha_view = alpha_acc.reshape(-1, 1, 1, 1)
else:
alpha_view = alpha_acc.reshape(-1) if alpha_acc.ndim == 0 else alpha_acc
exp_factor = 0.25 / (alpha_view**2)
# Avoid division by zero at k=0
safe_k_sq = jnp.maximum(k_sq_acc, 1e-30)
k_factor = 2.0 * (1.0 + k_sq_acc * exp_factor) / safe_k_sq
# Zero out k=0 contribution (no virial from k=0)
k_mask = k_sq_acc > 1e-10
# Six per-component weighted reductions instead of einsum: XLA lowers
# einsum(k,k,m) to a slow small-MN / large-K cuBLAS GEMM.
k_vecs_acc = k_vectors.astype(acc_dtype) # (..., nx, ny, nz//2+1, 3)
if is_batch and k_vecs_acc.ndim == 4:
k_vecs_acc = jnp.expand_dims(k_vecs_acc, axis=0)
masked_energy = weighted_energy * k_mask # (..., nx, ny, nz//2+1)
masked_energy_kf = masked_energy * k_factor # (..., nx, ny, nz//2+1)
if is_batch:
sum_dims = (1, 2, 3)
else:
sum_dims = (0, 1, 2)
trace_term = masked_energy.sum(axis=sum_dims) # scalar or (B,)
kx = k_vecs_acc[..., 0]
ky = k_vecs_acc[..., 1]
kz = k_vecs_acc[..., 2]
xx = (kx * kx * masked_energy_kf).sum(axis=sum_dims)
yy = (ky * ky * masked_energy_kf).sum(axis=sum_dims)
zz = (kz * kz * masked_energy_kf).sum(axis=sum_dims)
xy = (kx * ky * masked_energy_kf).sum(axis=sum_dims)
xz = (kx * kz * masked_energy_kf).sum(axis=sum_dims)
yz = (ky * kz * masked_energy_kf).sum(axis=sum_dims)
eye = jnp.eye(3, dtype=acc_dtype)
if is_batch:
kk_term = jnp.stack(
[
jnp.stack([xx, xy, xz], axis=-1),
jnp.stack([xy, yy, yz], axis=-1),
jnp.stack([xz, yz, zz], axis=-1),
],
axis=-2,
)
virial = eye * trace_term[:, jnp.newaxis, jnp.newaxis] - kk_term # (B, 3, 3)
else:
kk_term = jnp.stack(
[jnp.stack([xx, xy, xz]), jnp.stack([xy, yy, yz]), jnp.stack([xz, yz, zz])],
) # (3, 3)
virial = (eye * trace_term - kk_term)[jnp.newaxis, :, :] # (1, 3, 3)
return virial.astype(acc_dtype)
def _pme_reciprocal_space_impl(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None = None,
mesh_spacing: float | None = None,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
batch_idx: jax.Array | None = None,
k_vectors: jax.Array | None = None,
k_squared: jax.Array | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
hybrid_forces: bool = False,
volume: jax.Array | None = None,
cell_inv_t: jax.Array | None = None,
moduli_x: jax.Array | None = None,
moduli_y: jax.Array | None = None,
moduli_z: jax.Array | None = None,
) -> (
jax.Array
| tuple[jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array, jax.Array]
):
r"""Compute PME reciprocal-space contribution implementation.
Implements the FFT-based long-range component of PME using B-spline
interpolation and convolution with the Green's function.
Pipeline:
1. Spread charges to mesh (spline_spread)
2. FFT -> frequency space
3. Compute Green's function and structure factor
4. Convolve: mesh_fft * G(k) / C^2(k)
5. IFFT -> potential mesh
6. Gather potential at atoms (spline_gather)
7. Apply self-energy and background corrections
8. (Optional) Compute forces via Fourier gradient
Parameters
----------
positions : jax.Array, shape (N, 3)
Atomic coordinates.
charges : jax.Array, shape (N,)
Atomic partial charges.
cell : jax.Array, shape (3, 3) or (B, 3, 3)
Unit cell matrices with lattice vectors as rows.
alpha : jax.Array
Ewald splitting parameter.
- Single-system: shape (1,) or scalar
- Batch: shape (B,)
mesh_dimensions : tuple[int, int, int], optional
FFT mesh dimensions (nx, ny, nz).
mesh_spacing : float, optional
Target mesh spacing. Used to compute mesh_dimensions if not provided.
spline_order : int, default=5
B-spline interpolation order (4 = cubic).
batch_idx : jax.Array | None, default=None
System index for each atom. When provided, atoms must be grouped by
system: ``batch_idx`` must be contiguous, nondecreasing, and use system
IDs ``0..B-1``.
k_vectors : jax.Array, optional
Precomputed k-vectors from generate_k_vectors_pme.
k_squared : jax.Array, optional
Precomputed :math:`k^2` values from generate_k_vectors_pme.
compute_forces : bool, default=False
If True, compute forces via Fourier gradient.
compute_charge_gradients : bool, default=False
If True, compute charge gradients :math:`\partial E/\partial q`.
compute_virial : bool, default=False
If True, compute the virial tensor ``W = -dE/d(displacement)`` for the
row-vector displacement recipe.
Stress = -virial / volume.
hybrid_forces : bool, default=False
If True, detach ``positions``/``charges``/``cell`` from the autograd
graph through the spline/FFT chain (forces and virial become
forward-only), and inject analytical :math:`\partial E/\partial q` via a custom-VJP straight-
through trick so ``jax.grad`` w.r.t. charges propagates correctly.
Forces ``compute_charge_gradients=True``.
Returns
-------
energies : jax.Array, shape (N,)
Per-atom reciprocal-space energies.
forces : jax.Array, shape (N, 3), optional
Per-atom forces (only if compute_forces=True).
charge_gradients : jax.Array, shape (N,), optional
Per-atom charge gradients (only if compute_charge_gradients=True).
virial : jax.Array, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor (only if compute_virial=True). Always last in the return tuple.
Notes
-----
- Output dtype for energy/forces matches the input positions dtype
- FFT/convolution and spline operations all respect the input dtype
- Automatically determines mesh_dimensions if not provided
- Virial is computed in k-space and uses the same dtype as k_squared
- Energy-derived gradients are supported for positions, charges, and
strain-first virials. Reverse-mode higher-order reciprocal position and
charge losses use the private PME mesh HVP path.
"""
num_atoms = positions.shape[0]
input_dtype = normalize_float_dtype(positions.dtype)
is_batch = batch_idx is not None
fft_dims = (1, 2, 3) if is_batch else (0, 1, 2)
k_vectors_are_supplied = k_vectors is not None
k_squared_is_supplied = k_squared is not None
volume_is_supplied = volume is not None
# hybrid_forces: forward-only spline/FFT chain. We sever ∂/∂{positions,
# cell} (and the spline/FFT path through charges) via lax.stop_gradient,
# then re-attach the analytical ∂E/∂q at the end. Charges still need to
# be a tracer for the custom-VJP injector to see them, so save the
# original handle.
charges_orig = charges
need_charge_gradients = compute_charge_gradients or hybrid_forces
if hybrid_forces:
positions = jax.lax.stop_gradient(positions)
charges = jax.lax.stop_gradient(charges)
cell = jax.lax.stop_gradient(cell)
alpha = jax.lax.stop_gradient(alpha)
if cell_inv_t is not None:
cell_inv_t = jax.lax.stop_gradient(cell_inv_t)
# Ensure cell is correct shape for num_systems calculation
if cell.ndim == 2:
num_systems = 1
else:
num_systems = cell.shape[0]
# Handle empty systems
if num_atoms == 0:
energies = jnp.zeros(num_atoms, dtype=input_dtype)
forces = (
jnp.zeros((num_atoms, 3), dtype=input_dtype) if compute_forces else None
)
charge_grads = (
jnp.zeros(num_atoms, dtype=input_dtype) if need_charge_gradients else None
)
virial = (
jnp.zeros((num_systems, 3, 3), dtype=input_dtype)
if compute_virial
else None
)
return _build_electrostatic_result(
energies,
forces,
charge_grads,
virial,
compute_forces,
compute_charge_gradients,
compute_virial,
)
_require_explicit_mesh_dimensions_in_tracing(
mesh_dimensions=mesh_dimensions,
cell=cell,
alpha=alpha,
batch_idx=batch_idx,
)
# Determine mesh dimensions
if mesh_dimensions is None:
if mesh_spacing is not None:
mesh_dimensions = mesh_spacing_to_dimensions(
cell, mesh_spacing, spline_order=spline_order
)
else:
# Default estimation
mesh_dimensions = estimate_pme_mesh_dimensions(
cell, alpha, accuracy=1e-6, spline_order=spline_order
)
mesh_nx, mesh_ny, mesh_nz = mesh_dimensions
# cell_inv_t cache: MD callers can pass cell_inv_t in (NVT case)
# to skip the per-step linalg.inv + transpose. We compute it once here
# and forward to spline_spread / spline_gather(_with_force), so any
# work upstream of the kernel doesn't get duplicated.
cell_3d = cell if cell.ndim == 3 else cell[jnp.newaxis, :, :]
if cell_inv_t is None:
cell_inv = jnp.linalg.inv(cell_3d)
cell_inv_t = jnp.transpose(cell_inv, (0, 2, 1)).astype(input_dtype)
else:
cell_inv_t = cell_inv_t.astype(input_dtype)
if cell_inv_t.ndim == 2:
cell_inv_t = cell_inv_t[jnp.newaxis, :, :]
cell_inv_t = jax.lax.stop_gradient(cell_inv_t)
cell_inv = jnp.transpose(cell_inv_t, (0, 2, 1))
# Step 1: Spread charges to mesh
mesh_grid = spline_spread(
positions,
charges,
cell,
mesh_dims=mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t,
)
# Step 2: FFT of charge mesh
mesh_fft = jnp.fft.rfftn(mesh_grid, axes=fft_dims, norm="backward")
# Step 3: Generate k-space grid and compute Green's function + structure factor.
# When cell_inv_t is supplied, derive reciprocal_cell = 2π · cell_inv from
# the cached transpose so generate_k_vectors_pme skips its own inv.
if k_squared is None:
reciprocal_cell = (2.0 * jnp.pi) * cell_inv
if compute_virial:
k_vectors, k_squared = generate_k_vectors_pme(
cell,
mesh_dimensions,
reciprocal_cell=reciprocal_cell,
)
else:
k_squared = generate_k_squared_pme(
cell,
mesh_dimensions,
reciprocal_cell=reciprocal_cell,
)
elif compute_virial and k_vectors is None:
reciprocal_cell = (2.0 * jnp.pi) * cell_inv
k_vectors, _generated_k_squared = generate_k_vectors_pme(
cell,
mesh_dimensions,
reciprocal_cell=reciprocal_cell,
)
if k_vectors is not None and (hybrid_forces or k_vectors_are_supplied):
k_vectors = jax.lax.stop_gradient(k_vectors)
if hybrid_forces or k_squared_is_supplied:
k_squared = jax.lax.stop_gradient(k_squared)
# Step 4: Fused Green's function + B-spline deconvolution + multiply in a
# single warp kernel. Replaces the prior 2-pass path that
# called pme_green_structure_factor then divided/multiplied in JAX.
# Caller can supply moduli_x/y/z to skip the per-call fftfreq + sinc^p
# rebuild (they only depend on mesh + spline_order).
if moduli_x is None or moduli_y is None or moduli_z is None:
miller_x = jnp.fft.fftfreq(mesh_nx, d=1.0 / mesh_nx).astype(input_dtype)
miller_y = jnp.fft.fftfreq(mesh_ny, d=1.0 / mesh_ny).astype(input_dtype)
miller_z = jnp.fft.rfftfreq(mesh_nz, d=1.0 / mesh_nz).astype(input_dtype)
moduli_x = compute_bspline_moduli_1d(miller_x, mesh_nx, spline_order)
moduli_y = compute_bspline_moduli_1d(miller_y, mesh_ny, spline_order)
moduli_z = compute_bspline_moduli_1d(miller_z, mesh_nz, spline_order)
else:
moduli_x = jax.lax.stop_gradient(moduli_x)
moduli_y = jax.lax.stop_gradient(moduli_y)
moduli_z = jax.lax.stop_gradient(moduli_z)
# Caller-supplied volume= short-circuits the linalg.det.
if volume is None:
volume = jnp.abs(jnp.linalg.det(cell_3d)).astype(input_dtype)
else:
volume = volume.astype(input_dtype)
if volume.ndim == 0:
volume = volume.reshape(1)
if hybrid_forces or volume_is_supplied:
volume = jax.lax.stop_gradient(volume)
complex_dtype = jnp.complex64 if input_dtype == jnp.float32 else jnp.complex128
mesh_fft = mesh_fft.astype(complex_dtype)
# Raw FFT (before convolve) is what the virial path needs.
mesh_fft_raw = mesh_fft if compute_virial else None
convolved_mesh = pme_fused_convolve(
mesh_fft,
k_squared.astype(input_dtype),
moduli_x,
moduli_y,
moduli_z,
alpha,
volume,
is_batch,
)
# Step 5: Compute virial before forces to allow early release of mesh_fft_raw
# (virial needs mesh_fft_raw; forces only need convolved_mesh)
virial = None
if compute_virial:
virial = _compute_pme_reciprocal_virial(
mesh_fft_raw=mesh_fft_raw,
convolved_mesh=convolved_mesh,
k_vectors=k_vectors,
k_squared=k_squared,
alpha=alpha,
mesh_dimensions=mesh_dimensions,
is_batch=is_batch,
)
del mesh_fft_raw # Free before force field meshes are allocated
# Fused 2-pass warp kernel: reduce per-atom charges → per-system Q,
# then subtract E_bg = π Q² / (2 α² V) from the virial diagonal.
# Matches the torch ``pme_virial_bg_correction`` path. Single-system
# is fanned out internally via batch_idx=zeros.
virial = pme_virial_bg_correction(
charges=charges,
cell=cell,
alpha=alpha,
virial=virial,
batch_idx=batch_idx,
volume=volume,
)
# Step 6: Inverse FFT to get potential mesh
potential_mesh = jnp.fft.irfftn(
convolved_mesh, s=mesh_dimensions, axes=fft_dims, norm="forward"
)
# Step 6: Interpolate potential to atomic positions. With forces requested,
# use the fused gather kernel that walks the spline stencil
# ONCE per atom and emits both potential AND spline-derivative force,
# avoiding the 3 extra IFFTs + spline_gather_vec3 of the Fourier-gradient
# path. Matches the torch reciprocal-space path.
if compute_forces:
raw_energies, gathered_force = _spline_gather_with_force(
positions,
charges,
potential_mesh,
cell,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t,
)
else:
raw_energies = spline_gather(
positions,
potential_mesh,
cell,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t,
)
gathered_force = None
# Step 7: Apply corrections
if need_charge_gradients:
energies, charge_grads = pme_energy_corrections_with_charge_grad(
raw_energies, charges, cell, alpha, batch_idx, volume=volume
)
else:
energies = pme_energy_corrections(
raw_energies, charges, cell, alpha, batch_idx, volume=volume
)
charge_grads = None
# Step 8: Forces from the fused gather above. The 2× scaling absorbs the
# 1/2 pair-counting factor baked into the Green's function
# (G = 2π/(V k²) instead of 4π/(V k²)).
forces = 2.0 * gathered_force if compute_forces else None
# Hybrid-forces: route ∂E/∂q through the analytical kernel-computed
# ``charge_grads`` via a custom-VJP straight-through, using the original
# (non-detached) ``charges`` so jax.grad reaches them.
if hybrid_forces:
bidx_for_inject = (
batch_idx
if batch_idx is not None
else jnp.zeros(num_atoms, dtype=jnp.int32)
)
energies = _inject_charge_grad(
energies,
charges_orig,
charge_grads,
batch_idx is not None,
bidx_for_inject,
num_systems,
)
return _build_electrostatic_result(
energies,
forces,
charge_grads,
virial,
compute_forces,
compute_charge_gradients,
compute_virial,
)
def _stop_optional(value: jax.Array | None) -> jax.Array | None:
"""Stop gradients through an optional residual."""
if value is None:
return None
return jax.lax.stop_gradient(value)
def _is_traced_array(value) -> bool:
"""Return whether ``value`` is a JAX tracer inside transformations."""
return isinstance(value, jax.core.Tracer)
def _require_explicit_mesh_dimensions_in_tracing(
*,
mesh_dimensions: tuple[int, int, int] | None,
cell: jax.Array,
alpha: jax.Array | float | None,
batch_idx: jax.Array | None = None,
) -> None:
"""Reject auto PME mesh sizing under JAX tracing with a clear message."""
is_traced = (
_is_traced_array(cell) or _is_traced_array(alpha) or _is_traced_array(batch_idx)
)
if is_traced and alpha is None and mesh_dimensions is None:
raise ValueError(
"JAX PME requires explicit alpha and explicit mesh_dimensions inside "
"jax.jit or other JAX transformations. Compute PME parameters outside "
"the transformed function and pass alpha and mesh_dimensions=(nx, ny, nz) "
"explicitly."
)
if is_traced and alpha is None:
raise ValueError(
"JAX PME requires explicit alpha inside jax.jit or other JAX "
"transformations. Compute PME parameters outside the transformed "
"function and pass alpha explicitly."
)
if mesh_dimensions is not None:
return
if is_traced:
raise ValueError(
"JAX PME requires explicit mesh_dimensions inside jax.jit or other "
"JAX transformations. Compute mesh_spacing/accuracy-based mesh sizing "
"outside the transformed function and pass mesh_dimensions=(nx, ny, nz)."
)
def _tangent_or_zeros(tangent, primal: jax.Array, dtype=None) -> jax.Array:
"""Materialize a custom-JVP tangent, replacing symbolic zeros."""
out_dtype = primal.dtype if dtype is None else dtype
if _is_symbolic_zero(tangent):
return jnp.zeros(primal.shape, dtype=out_dtype)
return tangent.astype(out_dtype)
def _is_symbolic_zero(tangent) -> bool:
"""Return whether a custom-JVP tangent is JAX's symbolic zero sentinel."""
return tangent is None or isinstance(tangent, SymbolicZero)
def _bspline_weight_reference(u: jax.Array, order: int) -> jax.Array:
"""Pure-JAX cardinal B-spline basis for custom-JVP reference tangents."""
dtype = u.dtype
if order == 1:
return jnp.where((u >= 0.0) & (u < 1.0), 1.0, 0.0).astype(dtype)
result = jnp.zeros_like(u)
for j in range(order + 1):
coeff = (-1.0 if j % 2 else 1.0) * float(math.comb(order, j))
result = result + jnp.asarray(coeff, dtype=dtype) * jnp.maximum(
u - jnp.asarray(float(j), dtype=dtype), 0.0
) ** (order - 1)
return result / jnp.asarray(float(math.factorial(order - 1)), dtype=dtype)
def _reference_cell_inv_t(
cell: jax.Array,
cell_inv_t: jax.Array | None,
dtype,
) -> jax.Array:
"""Return cell inverse-transpose using supplied static metadata if present."""
if cell_inv_t is None:
cell_3d = cell if cell.ndim == 3 else cell[jnp.newaxis, :, :]
return jnp.transpose(jnp.linalg.inv(cell_3d.astype(dtype)), (0, 2, 1))
cell_inv_t = jax.lax.stop_gradient(cell_inv_t).astype(dtype)
if cell_inv_t.ndim == 2:
cell_inv_t = cell_inv_t[jnp.newaxis, :, :]
return cell_inv_t
def _reference_atom_systems(
positions: jax.Array,
batch_idx: jax.Array | None,
) -> jax.Array:
"""Return an int32 system id per atom."""
if batch_idx is None:
return jnp.zeros((positions.shape[0],), dtype=jnp.int32)
return batch_idx.astype(jnp.int32)
def _spline_spread_reference(
positions: jax.Array,
values: jax.Array,
cell_inv_t: jax.Array,
mesh_dimensions: tuple[int, int, int],
spline_order: int,
batch_idx: jax.Array | None,
num_systems: int,
) -> jax.Array:
"""Pure-JAX charge spread with the production B-spline stencil."""
dtype = normalize_float_dtype(positions.dtype)
nx, ny, nz = mesh_dimensions
atom_system = _reference_atom_systems(positions, batch_idx)
dims = jnp.asarray(mesh_dimensions, dtype=dtype)
frac = jnp.einsum(
"nij,nj->ni",
cell_inv_t[atom_system].astype(dtype),
positions.astype(dtype),
)
mesh_coords = frac * dims
base = jnp.floor(mesh_coords).astype(jnp.int32)
theta = mesh_coords - base.astype(dtype)
mesh_shape = (nx, ny, nz) if batch_idx is None else (num_systems, nx, ny, nz)
mesh = jnp.zeros(mesh_shape, dtype=dtype)
half_n_minus_1 = jnp.asarray(0.5 * float(spline_order - 2), dtype=dtype)
half_order = jnp.asarray(0.5 * float(spline_order), dtype=dtype)
starts = jnp.floor(theta - half_n_minus_1).astype(jnp.int32)
for ox in range(spline_order):
off_x = starts[:, 0] + ox
gx = jnp.mod(base[:, 0] + off_x, nx)
wx = _bspline_weight_reference(
half_order + theta[:, 0] - off_x.astype(dtype), spline_order
)
for oy in range(spline_order):
off_y = starts[:, 1] + oy
gy = jnp.mod(base[:, 1] + off_y, ny)
wy = _bspline_weight_reference(
half_order + theta[:, 1] - off_y.astype(dtype), spline_order
)
for oz in range(spline_order):
off_z = starts[:, 2] + oz
gz = jnp.mod(base[:, 2] + off_z, nz)
wz = _bspline_weight_reference(
half_order + theta[:, 2] - off_z.astype(dtype), spline_order
)
contrib = values.astype(dtype) * wx * wy * wz
if batch_idx is None:
mesh = mesh.at[gx, gy, gz].add(contrib)
else:
mesh = mesh.at[atom_system, gx, gy, gz].add(contrib)
return mesh
def _spline_gather_reference(
positions: jax.Array,
mesh: jax.Array,
cell_inv_t: jax.Array,
spline_order: int,
batch_idx: jax.Array | None,
) -> jax.Array:
"""Pure-JAX mesh gather with the production B-spline stencil."""
dtype = normalize_float_dtype(positions.dtype)
nx, ny, nz = mesh.shape[-3:]
atom_system = _reference_atom_systems(positions, batch_idx)
dims = jnp.asarray((nx, ny, nz), dtype=dtype)
frac = jnp.einsum(
"nij,nj->ni",
cell_inv_t[atom_system].astype(dtype),
positions.astype(dtype),
)
mesh_coords = frac * dims
base = jnp.floor(mesh_coords).astype(jnp.int32)
theta = mesh_coords - base.astype(dtype)
output = jnp.zeros((positions.shape[0],), dtype=dtype)
half_n_minus_1 = jnp.asarray(0.5 * float(spline_order - 2), dtype=dtype)
half_order = jnp.asarray(0.5 * float(spline_order), dtype=dtype)
starts = jnp.floor(theta - half_n_minus_1).astype(jnp.int32)
for ox in range(spline_order):
off_x = starts[:, 0] + ox
gx = jnp.mod(base[:, 0] + off_x, nx)
wx = _bspline_weight_reference(
half_order + theta[:, 0] - off_x.astype(dtype), spline_order
)
for oy in range(spline_order):
off_y = starts[:, 1] + oy
gy = jnp.mod(base[:, 1] + off_y, ny)
wy = _bspline_weight_reference(
half_order + theta[:, 1] - off_y.astype(dtype), spline_order
)
for oz in range(spline_order):
off_z = starts[:, 2] + oz
gz = jnp.mod(base[:, 2] + off_z, nz)
wz = _bspline_weight_reference(
half_order + theta[:, 2] - off_z.astype(dtype), spline_order
)
if batch_idx is None:
mesh_values = mesh[gx, gy, gz]
else:
mesh_values = mesh[atom_system, gx, gy, gz]
output = output + mesh_values.astype(dtype) * wx * wy * wz
return output
def _pme_energy_corrections_reference(
raw_energies: jax.Array,
charges: jax.Array,
alpha: jax.Array,
volume: jax.Array,
batch_idx: jax.Array | None,
num_systems: int,
) -> jax.Array:
"""Pure-JAX PME self/background correction for reference tangents."""
dtype = normalize_float_dtype(raw_energies.dtype)
charges = charges.astype(jnp.float64)
raw = raw_energies.astype(jnp.float64)
alpha_arr = _pme_alpha_array(alpha, jnp.float64, num_systems)
volume = volume.astype(jnp.float64)
if volume.ndim == 0:
volume = volume.reshape(1)
if volume.shape[0] == 1 and num_systems > 1:
volume = jnp.broadcast_to(volume, (num_systems,))
atom_system = _reference_atom_systems(raw_energies, batch_idx)
total_charges = _system_sum_from_atoms(charges, batch_idx, num_systems)
alpha_atom = alpha_arr[atom_system]
volume_atom = volume[atom_system]
total_atom = total_charges[atom_system]
self_energy = alpha_atom * charges * charges / jnp.sqrt(jnp.pi)
background = (
jnp.pi * charges * total_atom / (2.0 * alpha_atom * alpha_atom * volume_atom)
)
return (raw * charges - self_energy - background).astype(dtype)
def _pme_convolve_reference(
mesh_fft: jax.Array,
k_squared: jax.Array,
moduli_x: jax.Array,
moduli_y: jax.Array,
moduli_z: jax.Array,
alpha: jax.Array,
volume: jax.Array,
is_batch: bool,
) -> jax.Array:
"""Pure-JAX equivalent of the fused PME convolve multiplier."""
real_dtype = jnp.float32 if mesh_fft.dtype == jnp.complex64 else jnp.float64
alpha = _pme_alpha_array(alpha, real_dtype, volume.shape[0])
volume = volume.astype(real_dtype)
if volume.ndim == 0:
volume = volume.reshape(1)
k_sq = k_squared.astype(real_dtype)
if is_batch and k_sq.ndim == 3:
k_sq = k_sq[jnp.newaxis, ...]
safe_k_sq = jnp.where(k_sq > 1e-10, k_sq, 1.0)
if is_batch:
alpha_view = alpha.reshape(-1, 1, 1, 1)
volume_view = volume.reshape(-1, 1, 1, 1)
else:
alpha_view = alpha.reshape(-1)[0]
volume_view = volume.reshape(-1)[0]
green = (
2.0
* jnp.pi
* jnp.exp(-safe_k_sq / (4.0 * alpha_view * alpha_view))
/ (volume_view * safe_k_sq)
)
green = jnp.where(k_sq > 1e-10, green, 0.0)
sf = (
moduli_x.astype(real_dtype)[:, None, None]
* moduli_y.astype(real_dtype)[None, :, None]
* moduli_z.astype(real_dtype)[None, None, :]
)
sf_sq = jnp.maximum(sf * sf, jnp.asarray(1e-10, dtype=real_dtype))
return mesh_fft * (green / sf_sq).astype(mesh_fft.dtype)
def _pme_reciprocal_energy_reference(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> jax.Array:
"""Pure-JAX PME reciprocal per-atom energy for weighted-loss tangents."""
dtype = normalize_float_dtype(positions.dtype)
cell_3d = cell.astype(dtype)
if cell_3d.ndim == 2:
cell_3d = cell_3d[jnp.newaxis, :, :]
num_systems = cell_3d.shape[0]
is_batch = batch_idx is not None
if mesh_dimensions is None:
if mesh_spacing is None:
mesh_dimensions = estimate_pme_mesh_dimensions(
cell, alpha, accuracy=1e-6, spline_order=spline_order
)
else:
mesh_dimensions = mesh_spacing_to_dimensions(
cell, mesh_spacing, spline_order=spline_order
)
cell_inv_t_ref = _reference_cell_inv_t(cell, cell_inv_t, dtype)
cell_inv_ref = jnp.transpose(cell_inv_t_ref, (0, 2, 1))
mesh_grid = _spline_spread_reference(
positions,
charges,
cell_inv_t_ref,
mesh_dimensions,
spline_order,
batch_idx,
num_systems,
)
fft_dims = (1, 2, 3) if is_batch else (0, 1, 2)
complex_dtype = jnp.complex64 if dtype == jnp.float32 else jnp.complex128
mesh_fft = jnp.fft.rfftn(mesh_grid, axes=fft_dims, norm="backward").astype(
complex_dtype
)
if k_squared is None:
reciprocal_cell = (2.0 * jnp.pi) * cell_inv_ref
k_squared = generate_k_squared_pme(
cell,
mesh_dimensions,
reciprocal_cell=reciprocal_cell,
)
else:
k_squared = jax.lax.stop_gradient(k_squared)
nx, ny, nz = mesh_dimensions
if moduli_x is None or moduli_y is None or moduli_z is None:
miller_x = jnp.fft.fftfreq(nx, d=1.0 / nx).astype(dtype)
miller_y = jnp.fft.fftfreq(ny, d=1.0 / ny).astype(dtype)
miller_z = jnp.fft.rfftfreq(nz, d=1.0 / nz).astype(dtype)
moduli_x = compute_bspline_moduli_1d(miller_x, nx, spline_order)
moduli_y = compute_bspline_moduli_1d(miller_y, ny, spline_order)
moduli_z = compute_bspline_moduli_1d(miller_z, nz, spline_order)
else:
moduli_x = jax.lax.stop_gradient(moduli_x)
moduli_y = jax.lax.stop_gradient(moduli_y)
moduli_z = jax.lax.stop_gradient(moduli_z)
if volume is None:
volume_ref = jnp.abs(jnp.linalg.det(cell_3d)).astype(dtype)
else:
volume_ref = jax.lax.stop_gradient(volume).astype(dtype)
if volume_ref.ndim == 0:
volume_ref = volume_ref.reshape(1)
if volume_ref.shape[0] == 1 and num_systems > 1:
volume_ref = jnp.broadcast_to(volume_ref, (num_systems,))
convolved = _pme_convolve_reference(
mesh_fft,
k_squared,
moduli_x,
moduli_y,
moduli_z,
alpha,
volume_ref,
is_batch,
)
potential_mesh = jnp.fft.irfftn(
convolved, s=mesh_dimensions, axes=fft_dims, norm="forward"
).astype(dtype)
raw_energies = _spline_gather_reference(
positions,
potential_mesh,
cell_inv_t_ref,
spline_order,
batch_idx,
)
return _pme_energy_corrections_reference(
raw_energies,
charges,
alpha,
volume_ref,
batch_idx,
num_systems,
)
@jax.custom_jvp
def _unsupported_pme_cell_hvp_gradient(grad_cell: jax.Array) -> jax.Array:
"""Identity wrapper that rejects cell/stress higher-order PME derivatives."""
return grad_cell
@_unsupported_pme_cell_hvp_gradient.defjvp
def _unsupported_pme_cell_hvp_gradient_jvp(
primals: tuple[jax.Array],
tangents: tuple[jax.Array],
) -> tuple[jax.Array, jax.Array]:
"""Reject JVPs through the first-order PME cell-gradient adapter."""
del tangents
(grad_cell,) = primals
raise NotImplementedError(
"JAX PME stress/cell/strain HVPs are unsupported. Differentiate "
"position or charge losses, or use first-order cell/strain gradients only."
)
@jax.custom_jvp
def _unsupported_pme_cell_hvp_primal(cell: jax.Array) -> jax.Array:
"""Identity wrapper that marks PME cell JVP state as first-order only."""
return cell
@_unsupported_pme_cell_hvp_primal.defjvp
def _unsupported_pme_cell_hvp_primal_jvp(
primals: tuple[jax.Array],
tangents: tuple[jax.Array],
) -> tuple[jax.Array, jax.Array]:
"""Reject JVPs through a PME cell-gradient computation."""
del tangents
(cell,) = primals
raise NotImplementedError(
"JAX PME stress/cell/strain HVPs are unsupported. Differentiate "
"position or charge losses, or use first-order cell/strain gradients only."
)
def _cell_tangent_system_values(
grad_cell: jax.Array,
tangent_cell,
) -> jax.Array:
"""Contract a cell cotangent with a cell tangent per system."""
tcell = _tangent_or_zeros(tangent_cell, grad_cell, dtype=jnp.float64)
values = grad_cell.astype(jnp.float64) * tcell.astype(jnp.float64)
if values.ndim == 2:
return jnp.array([values.sum()], dtype=jnp.float64)
return values.sum(axis=(1, 2))
def _per_atom_cell_inv_t_matvec(
cell_inv_t: jax.Array,
vectors: jax.Array,
batch_idx: jax.Array | None,
) -> jax.Array:
"""Apply the per-system ``cell_inv_t`` matrix to per-atom vectors."""
if batch_idx is None:
return jnp.einsum("ij,nj->ni", cell_inv_t[0], vectors)
return jnp.einsum(
"nij,nj->ni",
cell_inv_t[batch_idx.astype(jnp.int32)],
vectors,
)
def _pme_alpha_array(
alpha: jax.Array,
dtype,
num_systems: int,
) -> jax.Array:
"""Return PME alpha as a length-``num_systems`` array."""
alpha_arr = alpha.astype(dtype)
if alpha_arr.ndim == 0:
alpha_arr = alpha_arr.reshape(1)
if alpha_arr.shape[0] == 1 and num_systems > 1:
alpha_arr = jnp.broadcast_to(alpha_arr, (num_systems,))
return alpha_arr
def _pme_charge_background_hvp(
h_charges: jax.Array,
alpha: jax.Array,
volume: jax.Array,
batch_idx: jax.Array | None,
num_systems: int,
) -> tuple[jax.Array, jax.Array]:
"""Return per-atom ``alpha`` and background charge-Hessian coefficient."""
d_qtotal = _system_sum_from_atoms(h_charges, batch_idx, num_systems)
if batch_idx is None:
alpha_atom = alpha[0]
bg_coeff = jnp.full_like(
h_charges, jnp.pi / (alpha_atom * alpha_atom * volume[0])
)
dqtotal_atom = jnp.full_like(h_charges, d_qtotal[0])
return alpha_atom, bg_coeff * dqtotal_atom
batch_i32 = batch_idx.astype(jnp.int32)
bg_coeff = jnp.pi / (alpha * alpha * volume)
return alpha[batch_i32], bg_coeff[batch_i32] * d_qtotal[batch_i32]
def _pme_reciprocal_hvp_state(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> tuple[jax.Array, ...]:
"""Build reusable PME reciprocal state for fixed-cell HVP evaluation."""
if mesh_dimensions is None:
if mesh_spacing is None:
raise ValueError("mesh_dimensions must be resolved before PME HVP")
mesh_dimensions = mesh_spacing_to_dimensions(
cell, mesh_spacing, spline_order=spline_order
)
dtype = normalize_float_dtype(positions.dtype)
positions_cast = positions.astype(dtype)
charges_cast = charges.astype(dtype)
cell_cast = cell.astype(dtype)
cell_3d = cell_cast if cell_cast.ndim == 3 else cell_cast[jnp.newaxis, :, :]
num_systems = cell_3d.shape[0]
is_batch = batch_idx is not None
fft_dims = (1, 2, 3) if is_batch else (0, 1, 2)
if cell_inv_t is None:
cell_inv = jnp.linalg.inv(cell_3d)
cell_inv_t_work = jnp.transpose(cell_inv, (0, 2, 1)).astype(dtype)
else:
cell_inv_t_work = cell_inv_t.astype(dtype)
if cell_inv_t_work.ndim == 2:
cell_inv_t_work = cell_inv_t_work[jnp.newaxis, :, :]
cell_inv = jnp.transpose(cell_inv_t_work, (0, 2, 1))
if volume is None:
volume_work = jnp.abs(jnp.linalg.det(cell_3d)).astype(dtype)
else:
volume_work = volume.astype(dtype)
if volume_work.ndim == 0:
volume_work = volume_work.reshape(1)
alpha_work = _pme_alpha_array(alpha, dtype, num_systems)
mesh_nx, mesh_ny, mesh_nz = mesh_dimensions
if k_squared is None:
if k_vectors is not None:
k_vectors_work = k_vectors.astype(dtype)
k_squared_work = jnp.sum(k_vectors_work * k_vectors_work, axis=-1)
else:
reciprocal_cell = (2.0 * jnp.pi) * cell_inv
_k_vectors, k_squared_work = generate_k_vectors_pme(
cell_cast,
mesh_dimensions,
reciprocal_cell=reciprocal_cell,
)
else:
k_squared_work = k_squared.astype(dtype)
if moduli_x is None or moduli_y is None or moduli_z is None:
miller_x = jnp.fft.fftfreq(mesh_nx, d=1.0 / mesh_nx).astype(dtype)
miller_y = jnp.fft.fftfreq(mesh_ny, d=1.0 / mesh_ny).astype(dtype)
miller_z = jnp.fft.rfftfreq(mesh_nz, d=1.0 / mesh_nz).astype(dtype)
moduli_x_work = compute_bspline_moduli_1d(miller_x, mesh_nx, spline_order)
moduli_y_work = compute_bspline_moduli_1d(miller_y, mesh_ny, spline_order)
moduli_z_work = compute_bspline_moduli_1d(miller_z, mesh_nz, spline_order)
else:
moduli_x_work = moduli_x.astype(dtype)
moduli_y_work = moduli_y.astype(dtype)
moduli_z_work = moduli_z.astype(dtype)
mesh_grid = spline_spread(
positions_cast,
charges_cast,
cell_cast,
mesh_dims=mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
complex_dtype = jnp.complex64 if dtype == jnp.float32 else jnp.complex128
mesh_fft = jnp.fft.rfftn(mesh_grid, axes=fft_dims, norm="backward").astype(
complex_dtype
)
convolved_mesh = pme_fused_convolve(
mesh_fft,
k_squared_work,
moduli_x_work,
moduli_y_work,
moduli_z_work,
alpha_work,
volume_work,
is_batch,
)
potential_mesh = jnp.fft.irfftn(
convolved_mesh,
s=mesh_dimensions,
axes=fft_dims,
norm="forward",
)
return (
positions_cast,
charges_cast,
cell_cast,
alpha_work,
batch_idx,
cell_inv_t_work,
volume_work,
k_squared_work,
moduli_x_work,
moduli_y_work,
moduli_z_work,
potential_mesh,
)
def _pme_reciprocal_energy_hvp_from_state(
v_positions: jax.Array,
v_charges: jax.Array,
positions_cast: jax.Array,
charges_cast: jax.Array,
cell_cast: jax.Array,
alpha_work: jax.Array,
batch_idx: jax.Array | None,
cell_inv_t_work: jax.Array,
volume_work: jax.Array,
k_squared_work: jax.Array,
moduli_x_work: jax.Array,
moduli_y_work: jax.Array,
moduli_z_work: jax.Array,
potential_mesh: jax.Array,
spline_order: int,
) -> tuple[jax.Array, jax.Array]:
"""Evaluate the linear PME reciprocal HVP from saved mesh state."""
dtype = normalize_float_dtype(positions_cast.dtype)
is_batch = batch_idx is not None
mesh_dimensions = (
tuple(potential_mesh.shape[1:]) if is_batch else tuple(potential_mesh.shape)
)
fft_dims = (1, 2, 3) if is_batch else (0, 1, 2)
v_positions_cast = v_positions.astype(dtype)
v_charges_cast = v_charges.astype(dtype)
complex_dtype = jnp.complex64 if dtype == jnp.float32 else jnp.complex128
cell_3d = cell_cast if cell_cast.ndim == 3 else cell_cast[jnp.newaxis, :, :]
num_systems = cell_3d.shape[0]
v_frac = _per_atom_cell_inv_t_matvec(
cell_inv_t_work,
v_positions_cast,
batch_idx,
)
dmesh_charge = spline_spread(
positions_cast,
v_charges_cast,
cell_cast,
mesh_dims=mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
dmesh_position = _spline_spread_gradient_weights(
positions_cast,
charges_cast[:, jnp.newaxis] * v_frac,
cell_cast,
mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
dmesh = dmesh_charge + dmesh_position
dmesh_fft = jnp.fft.rfftn(dmesh, axes=fft_dims, norm="backward").astype(
complex_dtype
)
dconvolved_mesh = pme_fused_convolve(
dmesh_fft,
k_squared_work,
moduli_x_work,
moduli_y_work,
moduli_z_work,
alpha_work,
volume_work,
is_batch,
)
dpotential_mesh = jnp.fft.irfftn(
dconvolved_mesh,
s=mesh_dimensions,
axes=fft_dims,
norm="forward",
)
grad_raw_cart = spline_gather_gradient(
positions_cast,
-jnp.ones_like(charges_cast),
potential_mesh,
cell_cast,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
dforce_mesh = spline_gather_gradient(
positions_cast,
charges_cast,
dpotential_mesh,
cell_cast,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
dforce_position = _spline_gather_gradient_position_hessian(
positions_cast,
charges_cast,
v_frac,
cell_cast,
potential_mesh,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
hvp_positions = (
2.0 * v_charges_cast[:, jnp.newaxis] * grad_raw_cart
- 2.0 * dforce_mesh
- 2.0 * dforce_position
)
draw_mesh = spline_gather(
positions_cast,
dpotential_mesh,
cell_cast,
spline_order=spline_order,
batch_idx=batch_idx,
cell_inv_t=cell_inv_t_work,
)
draw_position = (grad_raw_cart * v_positions_cast).sum(axis=1)
alpha_atom, background_hvp = _pme_charge_background_hvp(
v_charges_cast,
alpha_work,
volume_work,
batch_idx,
num_systems,
)
hvp_charges = (
2.0 * (draw_mesh + draw_position)
- 2.0 * alpha_atom / jnp.sqrt(jnp.pi) * v_charges_cast
- background_hvp
)
return (
hvp_positions.astype(positions_cast.dtype),
hvp_charges.astype(charges_cast.dtype),
)
def _pme_reciprocal_energy_hvp_raw(
v_positions: jax.Array,
v_charges: jax.Array,
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> tuple[jax.Array, jax.Array]:
"""Evaluate fixed-cell PME reciprocal HVPs for positions and charges."""
if positions.shape[0] == 0:
dtype = normalize_float_dtype(positions.dtype)
return (
jnp.zeros_like(positions, dtype=dtype),
jnp.zeros_like(charges, dtype=dtype),
)
state = _pme_reciprocal_hvp_state(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
return _pme_reciprocal_energy_hvp_from_state(
v_positions,
v_charges,
*state,
spline_order,
)
def _pme_reciprocal_energy_derivative_values(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> tuple[jax.Array, jax.Array]:
"""Return raw PME reciprocal ``dE/dR`` and ``dE/dq`` direct outputs."""
_energy, forces, charge_grads = _pme_reciprocal_space_impl(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
mesh_dimensions=mesh_dimensions,
mesh_spacing=mesh_spacing,
spline_order=spline_order,
batch_idx=batch_idx,
k_vectors=k_vectors,
k_squared=k_squared,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
compute_forces=True,
compute_charge_gradients=True,
compute_virial=False,
hybrid_forces=False,
)
return -forces, charge_grads
@functools.partial(jax.custom_vjp, nondiff_argnums=(4, 5, 6))
def _pme_reciprocal_energy_derivatives(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> tuple[jax.Array, jax.Array]:
"""PME reciprocal ``(dE/dR, dE/dq)`` first-derivative values.
Primal returns the PME *mesh* first derivatives (so forces/charge gradients are
bit-identical to the direct-output path). The custom VJP below supplies the
private PME-native HVP needed for reverse-mode higher-order
position/charge losses.
"""
dpos, charge_grads = _pme_reciprocal_energy_derivative_values(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
return jax.lax.stop_gradient(dpos), jax.lax.stop_gradient(charge_grads)
def _pme_reciprocal_energy_derivatives_fwd(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> tuple[tuple[jax.Array, jax.Array], tuple[jax.Array, ...]]:
"""Forward rule for PME reciprocal first derivatives."""
primal_out = _pme_reciprocal_energy_derivatives(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
state = _pme_reciprocal_hvp_state(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
return primal_out, state
def _pme_reciprocal_energy_derivatives_bwd(
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
residuals: tuple[jax.Array, ...],
ct_out: tuple[jax.Array, jax.Array],
) -> tuple[jax.Array | None, ...]:
"""Backward rule for PME reciprocal first derivatives."""
del mesh_dimensions, mesh_spacing
positions_cast, charges_cast, *_rest = residuals
ct_positions, ct_charges = ct_out
grad_positions, grad_charges = _pme_reciprocal_energy_hvp_from_state(
_tangent_or_zeros(
ct_positions,
positions_cast,
dtype=positions_cast.dtype,
),
_tangent_or_zeros(
ct_charges,
charges_cast,
dtype=charges_cast.dtype,
),
*residuals,
spline_order,
)
return (
grad_positions,
grad_charges,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
)
_pme_reciprocal_energy_derivatives.defvjp(
_pme_reciprocal_energy_derivatives_fwd,
_pme_reciprocal_energy_derivatives_bwd,
)
@functools.partial(jax.custom_jvp, nondiff_argnums=(4, 5, 6))
def _pme_reciprocal_energy_jvp(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
batch_idx: jax.Array | None,
k_vectors: jax.Array | None,
k_squared: jax.Array | None,
volume: jax.Array | None,
cell_inv_t: jax.Array | None,
moduli_x: jax.Array | None,
moduli_y: jax.Array | None,
moduli_z: jax.Array | None,
) -> jax.Array:
"""Energy-only PME reciprocal wrapper with a custom JVP."""
energy = _pme_reciprocal_space_impl(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
mesh_dimensions=mesh_dimensions,
mesh_spacing=mesh_spacing,
spline_order=spline_order,
batch_idx=batch_idx,
k_vectors=k_vectors,
k_squared=k_squared,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
compute_forces=False,
compute_charge_gradients=False,
compute_virial=False,
hybrid_forces=False,
)
return jax.lax.stop_gradient(energy)
def _pme_reciprocal_energy_jvp_rule(
mesh_dimensions: tuple[int, int, int] | None,
mesh_spacing: float | None,
spline_order: int,
primals: tuple[jax.Array | None, ...],
tangents: tuple[jax.Array | None, ...],
) -> tuple[jax.Array, jax.Array]:
"""JVP rule for the reciprocal PME per-atom energy vector."""
(
positions,
charges,
cell,
alpha,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
) = primals
(
t_positions,
t_charges,
t_cell,
_t_alpha,
_t_batch_idx,
_t_k_vectors,
_t_k_squared,
_t_volume,
_t_cell_inv_t,
_t_moduli_x,
_t_moduli_y,
_t_moduli_z,
) = tangents
del (
_t_alpha,
_t_batch_idx,
_t_k_vectors,
_t_k_squared,
_t_volume,
_t_cell_inv_t,
_t_moduli_x,
_t_moduli_y,
_t_moduli_z,
)
primal_out = _pme_reciprocal_energy_jvp(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
tpos = _tangent_or_zeros(t_positions, positions, dtype=positions.dtype)
tq = _tangent_or_zeros(t_charges, charges, dtype=charges.dtype)
tcell = _tangent_or_zeros(t_cell, cell, dtype=cell.dtype)
charges_ref = charges.astype(jnp.float64)
tq_ref = tq.astype(jnp.float64)
# Keep the energy tangent in pure JAX. The primal still uses the Warp PME
# mesh path above, but JAX cannot safely transpose the current Warp spline
# FFI boundary for PME cell/strain HVPs. Position and charge tangents are
# therefore evaluated through a pure-JAX PME reference, while nonzero cell
# tangents are tagged below so higher-order cell/strain requests reject
# explicitly instead of silently dropping terms or falling back to Ewald.
reference_cell = (
_unsupported_pme_cell_hvp_primal(cell)
if not _is_symbolic_zero(t_cell)
else cell
)
if not _is_symbolic_zero(t_cell):
# Cell tangents use the current reciprocal metric and assignment
# geometry. Fixed-cell position and charge tangents keep setup caches.
k_vectors = None
k_squared = None
volume = None
cell_inv_t = None
_reference_out, tangent_out = jax.jvp(
lambda p, q, c: _pme_reciprocal_energy_reference(
p,
q,
c,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
),
(positions, charges_ref, reference_cell),
(tpos, tq_ref, tcell),
)
if not _is_symbolic_zero(t_cell):
# First-order cell gradients are allowed, but differentiating those
# gradients again requires a native transposable PME cell HVP.
tangent_out = _unsupported_pme_cell_hvp_gradient(tangent_out)
return primal_out, tangent_out.astype(primal_out.dtype)
_pme_reciprocal_energy_jvp.defjvp(
_pme_reciprocal_energy_jvp_rule,
symbolic_zeros=True,
)
[docs]
def pme_reciprocal_space(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: jax.Array,
mesh_dimensions: tuple[int, int, int] | None = None,
mesh_spacing: float | None = None,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
batch_idx: jax.Array | None = None,
k_vectors: jax.Array | None = None,
k_squared: jax.Array | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
hybrid_forces: bool = False,
*,
energy_reduction: Literal["atom", "system"] = "atom",
cell_inv_t: jax.Array | None = None,
volume: jax.Array | None = None,
moduli_x: jax.Array | None = None,
moduli_y: jax.Array | None = None,
moduli_z: jax.Array | None = None,
) -> (
jax.Array
| tuple[jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array, jax.Array]
):
"""Compute PME reciprocal-space contribution.
Energy-only calls use a custom JVP so JAX does not attempt to
differentiate the Warp spline/FFT FFI path. ``compute_forces=True`` remains
a forward/direct escape hatch for no-autograd MD/inference loops; charge
gradients, virial, and hybrid direct outputs are deprecated training-style
outputs and warn.
Parameters
----------
positions : jax.Array, shape (N, 3)
Atomic coordinates.
charges : jax.Array, shape (N,)
Atomic partial charges.
cell : jax.Array, shape (3, 3) or (B, 3, 3)
Unit cell matrices with lattice vectors as rows.
alpha : jax.Array
Ewald splitting parameter.
mesh_dimensions : tuple[int, int, int] or None, default=None
Explicit FFT mesh dimensions. Required when ``cell``, ``alpha``, or
batch metadata are traced by ``jax.jit`` or other JAX transformations.
mesh_spacing : float or None, default=None
Target mesh spacing for eager-only mesh-size inference.
spline_order : int, default=5
B-spline interpolation order.
batch_idx : jax.Array or None, default=None
System index for each atom. When provided, atoms must be grouped by
system: ``batch_idx`` must be contiguous, nondecreasing, and use system
IDs ``0..B-1``.
k_vectors, k_squared : jax.Array or None
Optional precomputed reciprocal grid values. ``k_squared`` may be
supplied without ``k_vectors``. For batches with identical cells,
``k_squared`` may contain one shared three-dimensional grid.
If ``compute_virial=True`` and ``k_vectors`` is omitted, Cartesian
vectors are regenerated. These are
setup constants for the JAX custom-JVP path; tangents through them
are ignored. When supplied while differentiating with respect to
``cell``, they are assumed to correspond to the current ``cell``.
compute_forces, compute_charge_gradients, compute_virial : bool
Direct-output flags. ``compute_forces=True`` remains supported for
no-autograd MD/inference use; charge-gradient and virial direct outputs
are deprecated for differentiable training.
hybrid_forces : bool, default=False
Deprecated charge-gradient injection mode for compatibility.
energy_reduction : {"atom", "system"}, keyword-only, default="atom"
Public energy layout. ``"atom"`` returns per-atom energies of shape
(N,) (default, unchanged behavior). ``"system"`` returns per-system
energies of shape (B,) (or (1,) for a single system), computed as a
differentiable, uniform scatter-sum of the per-atom energies; it does
not support nonuniform per-atom weighting. Only the energy field of
the return value changes; forces, charge gradients, and the virial
keep their existing atom/system layouts.
cell_inv_t, volume, moduli_x, moduli_y, moduli_z : jax.Array or None
Optional precomputed PME intermediates. These are setup constants for
JAX and are not differentiable inputs. Cell-derived metadata such as
``cell_inv_t`` and ``volume`` is accepted while differentiating with
respect to ``cell`` and is assumed to correspond to the current
``cell``.
Returns
-------
energies : jax.Array, shape (N,) or (B,)
Reciprocal-space energies: per-atom when ``energy_reduction="atom"``,
per-system when ``energy_reduction="system"``.
forces : jax.Array, shape (N, 3), optional
Per-atom forces. Only present when ``compute_forces=True``.
charge_gradients : jax.Array, shape (N,), optional
Per-atom charge gradients :math:`\\partial E/\\partial q`. Only present when
``compute_charge_gradients=True`` (deprecated direct-output flag).
virial : jax.Array, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor. Only present when ``compute_virial=True`` (deprecated
direct-output flag). Always last in the return tuple.
Notes
-----
When ``cell`` or batch metadata are traced by ``jax.jit`` or other JAX
transformations, pass explicit ``mesh_dimensions``. ``mesh_spacing`` and
component-only accuracy-based mesh sizing depend on concrete setup values.
JAX PME higher-order support is limited to tested position and charge
losses. Stress/cell/strain HVPs, alpha HVPs, and precomputed-metadata HVPs
are unsupported until explicitly implemented and tested.
"""
_validate_energy_reduction(energy_reduction)
_require_explicit_mesh_dimensions_in_tracing(
mesh_dimensions=mesh_dimensions,
cell=cell,
alpha=alpha,
batch_idx=batch_idx,
)
component_deprecated_flags = tuple(
name
for name, enabled in (
("compute_charge_gradients", compute_charge_gradients),
("compute_virial", compute_virial),
("hybrid_forces", hybrid_forces),
)
if enabled
)
if component_deprecated_flags:
warnings.warn(
_component_direct_output_deprecation_msg(
"pme_reciprocal_space", component_deprecated_flags
),
DeprecationWarning,
stacklevel=2,
)
if compute_forces or compute_charge_gradients or compute_virial or hybrid_forces:
result = _pme_reciprocal_space_impl(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
mesh_dimensions=mesh_dimensions,
mesh_spacing=mesh_spacing,
spline_order=spline_order,
batch_idx=batch_idx,
k_vectors=k_vectors,
k_squared=k_squared,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
compute_forces=compute_forces,
compute_charge_gradients=compute_charge_gradients,
compute_virial=compute_virial,
hybrid_forces=hybrid_forces,
)
return _apply_energy_reduction(result, energy_reduction, batch_idx, cell)
if mesh_dimensions is None and mesh_spacing is not None:
mesh_dimensions = mesh_spacing_to_dimensions(
cell, mesh_spacing, spline_order=spline_order
)
mesh_spacing = None
result = _pme_reciprocal_energy_jvp(
positions,
charges,
cell,
alpha,
mesh_dimensions,
mesh_spacing,
spline_order,
batch_idx,
k_vectors,
k_squared,
volume,
cell_inv_t,
moduli_x,
moduli_y,
moduli_z,
)
return _apply_energy_reduction(result, energy_reduction, batch_idx, cell)
def _particle_mesh_ewald_impl(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: float | jax.Array | None = None,
mesh_spacing: float | None = None,
mesh_dimensions: tuple[int, int, int] | None = None,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
batch_idx: jax.Array | None = None,
k_vectors: jax.Array | None = None,
k_squared: jax.Array | None = None,
neighbor_list: jax.Array | None = None,
neighbor_ptr: jax.Array | None = None,
neighbor_shifts: jax.Array | None = None,
neighbor_matrix: jax.Array | None = None,
neighbor_matrix_shifts: jax.Array | None = None,
mask_value: int | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
accuracy: float = 1e-6,
pbc: jax.Array | None = None,
slab_correction: bool = False,
hybrid_forces: bool = False,
volume: jax.Array | None = None,
cell_inv_t: jax.Array | None = None,
moduli_x: jax.Array | None = None,
moduli_y: jax.Array | None = None,
moduli_z: jax.Array | None = None,
) -> (
jax.Array
| tuple[jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array, jax.Array]
):
r"""Complete Particle Mesh Ewald (PME) calculation for long-range electrostatics.
Computes total Coulomb energy using the PME method, which achieves
:math:`O(N \log N)` scaling through FFT-based reciprocal space calculations.
Combines:
1. Real-space contribution (short-range, erfc-damped)
2. Reciprocal-space contribution (long-range, FFT + B-spline interpolation)
3. Self-energy and background corrections
Total Energy Formula:
.. math::
E_{\text{total}} = E_{\text{real}} + E_{\text{reciprocal}}
- E_{\text{self}} - E_{\text{background}}
where:
.. math::
\begin{aligned}
E_{\text{real}} &= \frac{1}{2} \sum_{i \neq j} q_i q_j
\frac{\operatorname{erfc}(\alpha r_{ij})}{r_{ij}} \\
E_{\text{reciprocal}} &= \text{FFT-based smooth long-range contribution} \\
E_{\text{self}} &= \sum_i \frac{\alpha}{\sqrt{\pi}} q_i^2 \\
E_{\text{background}} &= \frac{\pi}{2\alpha^2 V} Q_{\text{total}}^2
\end{aligned}
Parameters
----------
positions : jax.Array, shape (N, 3)
Atomic coordinates.
charges : jax.Array, shape (N,)
Atomic partial charges in elementary charge units.
cell : jax.Array, shape (3, 3) or (B, 3, 3)
Unit cell matrices with lattice vectors as rows. Shape (3, 3) is
automatically promoted to (1, 3, 3) for single-system mode.
alpha : float, jax.Array, or None, default=None
Ewald splitting parameter controlling real/reciprocal space balance.
- float: Same :math:`\alpha` for all systems
- Array shape (B,): Per-system :math:`\alpha` values
The public entry point resolves this value before calling the
implementation.
mesh_spacing : float, optional
Target mesh spacing. Mesh dimensions are computed as
``ceil(cell_length / mesh_spacing)`` when explicit dimensions are absent.
mesh_dimensions : tuple[int, int, int], optional
Explicit FFT mesh dimensions (nx, ny, nz). This or ``mesh_spacing`` is
required by the full-PME implementation.
spline_order : int, default=5
B-spline interpolation order (5 = quartic B-splines, the PME default).
batch_idx : jax.Array, shape (N,), dtype=int32, optional
System index for each atom (0 to B-1). Determines execution mode:
- None: Single-system optimized kernels
- Provided: Batched kernels for multiple independent systems
When provided, atoms must be grouped by system: ``batch_idx`` must be
contiguous, nondecreasing, and use system IDs ``0..B-1``.
k_vectors : jax.Array, optional
Precomputed k-vectors from generate_k_vectors_pme. Providing this
along with k_squared skips k-vector generation.
k_squared : jax.Array, optional
Precomputed :math:`k^2` values from generate_k_vectors_pme.
neighbor_list : jax.Array, optional
CSR-format neighbor list indices. See ewald_real_space.
neighbor_ptr : jax.Array, optional
CSR-format neighbor list pointers. See ewald_real_space.
neighbor_shifts : jax.Array, optional
Periodic image shifts for neighbor list. See ewald_real_space.
neighbor_matrix : jax.Array, optional
Dense neighbor matrix. Alternative to CSR format.
neighbor_matrix_shifts : jax.Array, optional
Shifts for dense neighbor matrix.
mask_value : int, optional
Mask value for invalid neighbors in dense format.
compute_forces : bool, default=False
If True, compute per-atom forces.
compute_charge_gradients : bool, default=False
If True, compute per-atom charge gradients :math:`\partial E/\partial q`.
compute_virial : bool, default=False
If True, compute the virial tensor ``W = -dE/d(displacement)`` for the
row-vector displacement recipe.
Stress = -virial / volume.
accuracy : float, default=1e-6
Retained for public API compatibility. Full-PME parameter selection is
performed by the public entry point before this implementation is called.
pbc : jax.Array, shape (3,) or (B, 3), dtype=bool, optional
Per-system periodic boundary conditions. Required when
``slab_correction=True``. True marks periodic directions and False
marks the non-periodic slab direction.
slab_correction : bool, default=False
If True, add the Yeh-Berkowitz/Ballenegger slab correction to the
3D-periodic PME outputs.
Returns
-------
energies : jax.Array, shape (N,)
Per-atom total electrostatic energies.
forces : jax.Array, shape (N, 3), optional
Per-atom forces (only if compute_forces=True).
charge_gradients : jax.Array, shape (N,), optional
Per-atom charge gradients (only if compute_charge_gradients=True).
virial : jax.Array, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor (only if compute_virial=True). Always last in the return tuple.
Notes
-----
Energy-derived first-order gradients are supported. Higher-order PME
reverse-mode higher-order position and charge losses use the private PME
mesh HVP path.
The public entry point resolves alpha, mesh dimensions, and spline order
before calling this implementation. ``mesh_spacing`` conversion requires
concrete setup values outside JAX transformations.
Examples
--------
Basic usage:
>>> energies = particle_mesh_ewald(
... positions, charges, cell, alpha=0.3,
... mesh_dimensions=(32, 32, 32),
... neighbor_list=nl, neighbor_ptr=ptr, neighbor_shifts=shifts,
... )
With forces and explicit parameters:
>>> energies, forces = particle_mesh_ewald(
... positions, charges, cell, alpha=0.3,
... mesh_spacing=1.0,
... neighbor_list=nl, neighbor_ptr=ptr, neighbor_shifts=shifts,
... compute_forces=True,
... )
Batched systems:
>>> energies = particle_mesh_ewald(
... positions, charges, cell, alpha=0.3,
... mesh_dimensions=(32, 32, 32),
... batch_idx=batch_idx,
... neighbor_list=nl, neighbor_ptr=ptr, neighbor_shifts=shifts,
... )
See Also
--------
pme_reciprocal_space : Reciprocal-space component only
ewald_real_space : Real-space component
estimate_pme_parameters : PME parameter setup
"""
num_atoms = positions.shape[0]
# Prepare cell and slab pbc
cell, num_systems = _prepare_cell(cell)
if batch_idx is not None:
batch_idx = batch_idx.astype(jnp.int32)
if slab_correction:
pbc = _prepare_pbc_for_slab(pbc, num_systems)
_require_explicit_mesh_dimensions_in_tracing(
mesh_dimensions=mesh_dimensions,
cell=cell,
alpha=alpha,
batch_idx=batch_idx,
)
# Prepare alpha
if isinstance(alpha, (int, float)):
alpha = jnp.array([alpha] * num_systems, dtype=positions.dtype)
elif alpha.ndim == 0:
alpha = alpha.reshape(1)
if mask_value is None:
mask_value = num_atoms
# Determine mesh dimensions
if mesh_dimensions is None:
if mesh_spacing is not None:
mesh_dimensions = mesh_spacing_to_dimensions(
cell, mesh_spacing, spline_order=spline_order
)
else:
raise ValueError(
"full particle_mesh_ewald with explicit alpha requires "
"mesh_dimensions or mesh_spacing; use estimate_pme_parameters before "
"building the neighbor list for automatic parameter selection"
)
charges_orig = charges
need_charge_gradients = compute_charge_gradients or hybrid_forces
if hybrid_forces:
positions = jax.lax.stop_gradient(positions)
charges = jax.lax.stop_gradient(charges)
cell = jax.lax.stop_gradient(cell)
alpha = jax.lax.stop_gradient(alpha)
if k_vectors is not None:
k_vectors = jax.lax.stop_gradient(k_vectors)
if k_squared is not None:
k_squared = jax.lax.stop_gradient(k_squared)
if volume is not None:
volume = jax.lax.stop_gradient(volume)
if cell_inv_t is not None:
cell_inv_t = jax.lax.stop_gradient(cell_inv_t)
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message=r"The component direct-output flag\(s\).*",
category=DeprecationWarning,
)
# Compute real-space contribution
rs = ewald_real_space(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
neighbor_list=neighbor_list,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=need_charge_gradients,
compute_virial=compute_virial,
)
# Compute reciprocal-space contribution
rec = pme_reciprocal_space(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
mesh_dimensions=mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=need_charge_gradients,
compute_virial=compute_virial,
k_vectors=k_vectors,
k_squared=k_squared,
hybrid_forces=False,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
)
slab = None
if slab_correction:
if compute_forces or need_charge_gradients or compute_virial:
slab = _compute_slab_correction(
positions,
charges,
cell,
pbc,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=need_charge_gradients,
compute_virial=compute_virial,
)
else:
slab = _slab_correction_energy_autodiff(
positions,
charges,
cell,
pbc,
batch_idx=batch_idx,
)
component_tuples = [
rs if isinstance(rs, tuple) else (rs,),
rec if isinstance(rec, tuple) else (rec,),
]
if slab is not None:
component_tuples.append(slab if isinstance(slab, tuple) else (slab,))
def _sum_component(tuple_index: int) -> jax.Array:
total = component_tuples[0][tuple_index]
for component in component_tuples[1:]:
total = total + component[tuple_index]
return total
tuple_index = 0
total_energies = _sum_component(tuple_index)
tuple_index += 1
total_charge_grads = None
results: tuple[jax.Array, ...] = (total_energies,)
if compute_forces:
total_forces = _sum_component(tuple_index)
results += (total_forces,)
tuple_index += 1
if need_charge_gradients:
total_charge_grads = _sum_component(tuple_index)
tuple_index += 1
if compute_charge_gradients:
results += (total_charge_grads,)
if compute_virial:
total_virial = _sum_component(tuple_index)
results += (total_virial,)
if hybrid_forces and total_charge_grads is not None:
bidx_for_inject = (
batch_idx
if batch_idx is not None
else jnp.zeros(num_atoms, dtype=jnp.int32)
)
total_energies = _inject_charge_grad(
total_energies,
charges_orig,
total_charge_grads,
batch_idx is not None,
bidx_for_inject,
num_systems,
)
results = (total_energies, *results[1:])
return results[0] if len(results) == 1 else results
def _resolve_particle_mesh_ewald_parameters(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: float | jax.Array | None,
mesh_spacing: float | None,
mesh_dimensions: tuple[int, int, int] | None,
spline_order: int | None,
batch_idx: jax.Array | None,
accuracy: float,
) -> tuple[jax.Array, jax.Array, tuple[int, int, int], int]:
"""Resolve one full-PME execution configuration before choosing an output path."""
cell_3d, num_systems = _prepare_cell(cell)
resolved_spline_order = (
_DEFAULT_PME_SPLINE_ORDER if spline_order is None else spline_order
)
if alpha is None:
# The heuristic reads the cell and the per-system atom counts. Traced
# positions or charges are fine here, so jax.grad over positions with
# accuracy-based parameters keeps working.
traced_inputs = (cell, batch_idx)
if any(_is_traced_array(value) for value in traced_inputs if value is not None):
raise ValueError(
"JAX PME requires explicit alpha and explicit mesh_dimensions "
"inside jax.jit, jax.grad, or other JAX transformations. Run "
"estimate_pme_parameters outside the transformation and before building "
"the neighbor list, then pass its alpha and mesh dimensions explicitly."
)
params = estimate_pme_parameters(
positions,
cell_3d,
batch_idx,
accuracy,
spline_order=resolved_spline_order,
)
alpha = params.alpha
if mesh_dimensions is None and mesh_spacing is None:
mesh_dimensions = tuple(int(value) for value in params.mesh_dimensions)
_require_explicit_mesh_dimensions_in_tracing(
mesh_dimensions=mesh_dimensions,
cell=cell_3d,
alpha=alpha,
batch_idx=batch_idx,
)
if isinstance(alpha, (int, float)):
alpha = jnp.array([alpha] * num_systems, dtype=positions.dtype)
elif alpha.ndim == 0:
alpha = alpha.reshape(1)
if mesh_dimensions is None:
if mesh_spacing is not None:
mesh_dimensions = mesh_spacing_to_dimensions(
cell_3d, mesh_spacing, spline_order=resolved_spline_order
)
else:
mesh_dimensions = estimate_pme_mesh_dimensions(
cell_3d,
alpha,
accuracy,
spline_order=resolved_spline_order,
)
return cell_3d, alpha, mesh_dimensions, resolved_spline_order
[docs]
def particle_mesh_ewald(
positions: jax.Array,
charges: jax.Array,
cell: jax.Array,
alpha: float | jax.Array | None = None,
mesh_spacing: float | None = None,
mesh_dimensions: tuple[int, int, int] | None = None,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
batch_idx: jax.Array | None = None,
k_vectors: jax.Array | None = None,
k_squared: jax.Array | None = None,
neighbor_list: jax.Array | None = None,
neighbor_ptr: jax.Array | None = None,
neighbor_shifts: jax.Array | None = None,
neighbor_matrix: jax.Array | None = None,
neighbor_matrix_shifts: jax.Array | None = None,
mask_value: int | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
accuracy: float = 1e-6,
hybrid_forces: bool = False,
pbc: jax.Array | None = None,
slab_correction: bool = False,
*,
energy_reduction: Literal["atom", "system"] = "atom",
cell_inv_t: jax.Array | None = None,
volume: jax.Array | None = None,
moduli_x: jax.Array | None = None,
moduli_y: jax.Array | None = None,
moduli_z: jax.Array | None = None,
) -> (
jax.Array
| tuple[jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array]
| tuple[jax.Array, jax.Array, jax.Array, jax.Array]
):
"""Complete Particle Mesh Ewald calculation for long-range electrostatics.
Computes the total Coulomb energy via the PME method, which achieves
:math:`O(N \\log N)` scaling through FFT-based reciprocal-space calculations:
.. math::
E_{\\text{total}} = E_{\\text{real}} + E_{\\text{reciprocal}}
- E_{\\text{self}} - E_{\\text{background}}
Parameters
----------
positions : jax.Array, shape (N, 3)
Atomic coordinates.
charges : jax.Array, shape (N,)
Atomic partial charges.
cell : jax.Array, shape (3, 3) or (B, 3, 3)
Unit cell matrices with lattice vectors as rows.
alpha : float, jax.Array, or None, default=None
Ewald splitting parameter. If ``None``, alpha and the mesh dimensions
are estimated in the call from ``accuracy`` outside JAX
transformations. Pass ``params.alpha`` from ``estimate_pme_parameters``
to use prepared parameters.
mesh_spacing : float or None, default=None
Target mesh spacing used when ``mesh_dimensions`` is omitted.
mesh_dimensions : tuple[int, int, int] or None, default=None
Explicit FFT mesh dimensions. When both mesh controls are omitted, the
mesh is estimated in the call from ``accuracy`` outside JAX
transformations. Pass ``params.mesh_dimensions`` to use prepared parameters.
spline_order : int, default=5
B-spline interpolation order. Implicit meshes provide assignment-order support.
batch_idx : jax.Array or None, default=None
System index for each atom. When provided, atoms must be grouped by
system: ``batch_idx`` must be contiguous, nondecreasing, and use system
IDs ``0..B-1``.
k_vectors, k_squared : jax.Array or None
Precomputed PME reciprocal grid values. ``k_squared`` may be supplied
without ``k_vectors``. If ``compute_virial=True`` and ``k_vectors`` is
omitted, Cartesian vectors are regenerated.
neighbor_list, neighbor_ptr, neighbor_shifts : jax.Array or None
CSR neighbor-list inputs for the real-space component.
neighbor_matrix, neighbor_matrix_shifts : jax.Array or None
Dense neighbor-matrix inputs for the real-space component.
mask_value : int or None, default=None
Sentinel value for invalid neighbor-matrix entries.
compute_forces, compute_charge_gradients, compute_virial : bool
Deprecated direct-output flags. Compute energy and use JAX autodiff for
differentiable forces, charge gradients, and strain virials.
accuracy : float, default=1e-6
Target used by the in-call parameter estimators when ``alpha`` or the
mesh controls are omitted.
hybrid_forces : bool, default=False
Deprecated Torch-compatibility escape hatch for charge-gradient routing.
pbc : jax.Array, optional
Per-system periodic boundary conditions for slab correction.
slab_correction : bool, default=False
If True, add the Yeh-Berkowitz/Ballenegger slab correction.
energy_reduction : {"atom", "system"}, keyword-only, default="atom"
Public energy layout. ``"atom"`` returns per-atom energies of shape
(N,) (default, unchanged behavior). ``"system"`` returns per-system
energies of shape (B,) (or (1,) for a single system), computed as a
differentiable, uniform scatter-sum of the per-atom energies; it does
not support nonuniform per-atom weighting. Only the energy field of
the return value changes; forces, charge gradients, and the virial
keep their existing atom/system layouts.
volume, cell_inv_t, moduli_x, moduli_y, moduli_z : jax.Array or None
Optional precomputed PME intermediates. Cell-derived values supplied
while differentiating with respect to ``cell`` are treated as static
metadata that corresponds to the current ``cell``.
Returns
-------
energies : jax.Array, shape (N,) or (B,)
Total electrostatic energies (real + reciprocal + slab): per-atom
when ``energy_reduction="atom"``, per-system when
``energy_reduction="system"``.
forces : jax.Array, shape (N, 3), optional
Per-atom forces. Only present when ``compute_forces=True`` (deprecated).
charge_gradients : jax.Array, shape (N,), optional
Per-atom charge gradients :math:`\\partial E/\\partial q`. Only present when
``compute_charge_gradients=True`` (deprecated).
virial : jax.Array, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor. Only present when ``compute_virial=True`` (deprecated).
Always last in the return tuple.
Notes
-----
Call ``estimate_pme_parameters`` before building the neighbor list, build
that list with its ``real_space_cutoff``, then pass its ``alpha`` and
``mesh_dimensions`` here. Use the same ``spline_order`` during setup and
execution. Reuse these values while the cell and atom counts remain fixed,
inside or outside ``jax.jit``.
When cell, alpha, or batch metadata are traced by a JAX transformation,
pass explicit mesh dimensions. When cell or batch metadata are traced and
alpha would be estimated, precompute alpha as well. Position and charge
derivatives can use implicit setup with concrete captured cell and batch
values.
"""
_validate_energy_reduction(energy_reduction)
if compute_forces or compute_virial or compute_charge_gradients or hybrid_forces:
warnings.warn(
_direct_output_deprecation_msg("particle_mesh_ewald"),
DeprecationWarning,
stacklevel=2,
)
cell_3d, alpha_arr, mesh_dims, resolved_spline_order = (
_resolve_particle_mesh_ewald_parameters(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha,
mesh_spacing=mesh_spacing,
mesh_dimensions=mesh_dimensions,
spline_order=spline_order,
batch_idx=batch_idx,
accuracy=accuracy,
)
)
if mask_value is None:
mask_value = positions.shape[0]
if compute_forces or compute_charge_gradients or compute_virial or hybrid_forces:
result = _particle_mesh_ewald_impl(
positions=positions,
charges=charges,
cell=cell_3d,
alpha=alpha_arr,
mesh_dimensions=mesh_dims,
spline_order=resolved_spline_order,
batch_idx=batch_idx,
k_vectors=k_vectors,
k_squared=k_squared,
neighbor_list=neighbor_list,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
compute_forces=compute_forces,
compute_charge_gradients=compute_charge_gradients,
compute_virial=compute_virial,
accuracy=accuracy,
hybrid_forces=hybrid_forces,
pbc=pbc,
slab_correction=slab_correction,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
)
return _apply_energy_reduction(result, energy_reduction, batch_idx, cell_3d)
# Energy-only path: call the impl directly so the full energy is the sum of
# real-space and reciprocal terms. Component custom derivative rules provide
# energy gradients and reverse-mode position/charge higher-order losses.
result = _particle_mesh_ewald_impl(
positions=positions,
charges=charges,
cell=cell_3d,
alpha=alpha_arr,
mesh_dimensions=mesh_dims,
spline_order=resolved_spline_order,
batch_idx=batch_idx,
k_vectors=k_vectors,
k_squared=k_squared,
neighbor_list=neighbor_list,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
compute_forces=False,
compute_charge_gradients=False,
compute_virial=False,
accuracy=accuracy,
hybrid_forces=False,
pbc=pbc,
slab_correction=slab_correction,
volume=volume,
cell_inv_t=cell_inv_t,
moduli_x=moduli_x,
moduli_y=moduli_y,
moduli_z=moduli_z,
)
return _apply_energy_reduction(result, energy_reduction, batch_idx, cell_3d)
