# 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.
r"""
Ewald Summation - PyTorch Bindings
==================================
This module provides PyTorch bindings for Ewald summation calculations.
It wraps the framework-agnostic Warp launchers from
``nvalchemiops.interactions.electrostatics.ewald_kernels``.
Public API
----------
- ``ewald_real_space()``: Real-space component of Ewald summation
- ``ewald_reciprocal_space()``: Reciprocal-space component
- ``ewald_summation()``: Complete Ewald summation (real + reciprocal)
Mathematical Formulation
------------------------
The Ewald method splits long-range Coulomb interactions into components:
.. math::
E_{\text{total}} = E_{\text{real}} + E_{\text{reciprocal}} - E_{\text{self}} - E_{\text{background}}
All functions support:
- Both neighbor list (CSR) and neighbor matrix formats
- Batched calculations
- Energy autograd for differentiable training
The full ``ewald_summation`` API treats energy autograd as the differentiable
contract; its direct-output flags warn and are deprecated. The component APIs
(``ewald_real_space`` and ``ewald_reciprocal_space``) intentionally retain direct
forces as no-autograd MD/inference escape hatches. Component charge-gradient,
virial, and hybrid direct outputs are deprecated training-style outputs and warn.
Examples
--------
>>> # Complete Ewald summation
>>> energies = ewald_summation(
... positions, charges, cell,
... neighbor_list=nl, neighbor_ptr=neighbor_ptr, neighbor_shifts=shifts,
... accuracy=1e-6,
... )
>>> forces = -torch.autograd.grad(energies.sum(), positions, create_graph=True)[0]
>>> # Fixed-cell loop: precompute the actual reciprocal vectors once
>>> k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0)
>>> energies = ewald_summation(
... positions, charges, cell,
... alpha=alpha, k_vectors=k_vectors,
... neighbor_list=nl, neighbor_ptr=neighbor_ptr, neighbor_shifts=shifts,
... )
>>> # Changing-cell loop: precompute conservative Miller half-bounds once
>>> energies = ewald_summation(
... positions, charges, cell_t,
... alpha=alpha, k_cutoff=8.0, miller_bounds=(16, 16, 16),
... neighbor_list=nl, neighbor_ptr=neighbor_ptr, neighbor_shifts=shifts,
... )
>>> # Separate real and reciprocal components with direct no-autograd forces
>>> e_real, f_real = ewald_real_space(
... positions, charges, cell, alpha,
... neighbor_list=nl, neighbor_ptr=neighbor_ptr, neighbor_shifts=shifts,
... compute_forces=True,
... )
>>> e_recip, f_recip = ewald_reciprocal_space(
... positions, charges, cell, k_vectors, alpha,
... compute_forces=True,
... )
"""
from __future__ import annotations
import warnings
from typing import Literal
import torch
from nvalchemiops.torch.interactions.electrostatics._ewald_corrections_chain import (
ewald_energy_corrections,
ewald_energy_corrections_batch,
)
from nvalchemiops.torch.interactions.electrostatics._ewald_direct import (
reciprocal_space_direct,
)
from nvalchemiops.torch.interactions.electrostatics._ewald_real_chain import (
real_space_cell_connect,
)
from nvalchemiops.torch.interactions.electrostatics._ewald_recip_chain import (
_recip_ksum_energy_torch,
)
from nvalchemiops.torch.interactions.electrostatics._registration import (
ensure_electrostatics_ops_registered,
)
from nvalchemiops.torch.interactions.electrostatics._util import (
_build_electrostatic_result,
_combine_electrostatic_outputs,
_compiled_direct_output_deprecation_signal,
_component_direct_output_deprecation_msg,
_detach_setup_tensor,
_direct_output_deprecation_msg,
_InjectCachedEvalGrad,
_InjectCachedEvalGradWithFallback,
_InjectChargeGrad,
_reduce_atom_energy,
_unpack_electrostatic_outputs,
_validate_energy_reduction,
)
from nvalchemiops.torch.interactions.electrostatics.k_vectors import (
generate_k_vectors_ewald_summation,
)
from nvalchemiops.torch.interactions.electrostatics.parameters import (
estimate_ewald_parameters,
)
from nvalchemiops.torch.interactions.electrostatics.slab import (
compute_slab_correction as _compute_slab_correction,
)
__all__ = [
"ewald_real_space",
"ewald_reciprocal_space",
"ewald_summation",
]
###########################################################################################
########################### Helper Functions ##############################################
###########################################################################################
def _prepare_alpha(
alpha: float | torch.Tensor,
num_systems: int,
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
"""Convert alpha to a per-system tensor.
Parameters
----------
alpha : float or torch.Tensor
Ewald splitting parameter. Can be:
- A scalar float (broadcast to all systems)
- A 0-d tensor (broadcast to all systems)
- A 1-d tensor of shape (num_systems,) for per-system values
num_systems : int
Number of systems in the batch.
dtype : torch.dtype
Target dtype for the output tensor.
device : torch.device
Target device for the output tensor.
Returns
-------
torch.Tensor, shape (num_systems,)
Per-system alpha values.
"""
if isinstance(alpha, (int, float)):
return torch.full((num_systems,), float(alpha), dtype=dtype, device=device)
elif isinstance(alpha, torch.Tensor):
if alpha.dim() == 0:
return alpha.expand(num_systems).to(dtype=dtype, device=device)
elif alpha.shape[0] != num_systems:
raise ValueError(
f"alpha has {alpha.shape[0]} values but there are {num_systems} systems"
)
return alpha.to(dtype=dtype, device=device)
else:
raise TypeError(f"alpha must be float or torch.Tensor, got {type(alpha)}")
def _prepare_cell(cell: torch.Tensor) -> tuple[torch.Tensor, int]:
"""Ensure cell is 3D (B, 3, 3) and return number of systems.
Parameters
----------
cell : torch.Tensor
Unit cell matrix. Shape (3, 3) for single system or (B, 3, 3) for batch.
Returns
-------
cell : torch.Tensor, shape (B, 3, 3)
Cell with batch dimension.
num_systems : int
Number of systems (B).
"""
if cell.dim() == 2:
cell = cell.unsqueeze(0)
return cell, cell.shape[0]
###########################################################################################
########################### Internal Energy Assembly (explicit chains) ####################
###########################################################################################
#
# The differentiable public energy is assembled from the explicit factory chains
# registered in ``_ewald_real_chain`` / ``_ewald_recip_chain`` (forward energy ->
# backward -> double_backward), plus Torch-native pieces that the
# kernels do not own:
#
# * real-space ``cell`` gradient -- :func:`real_space_cell_connect` (the literal
# ``dE/dcell`` via the differentiable periodic shift ``unit_shifts @ cell``);
# * reciprocal self-energy + background corrections (closed-form in charges /
# alpha / volume) and the ``k_vectors(cell)`` / ``volume(cell)`` maps that carry
# the reciprocal ``cell`` gradient through Torch.
#
# When nothing requires grad, the forward-only ``_DerivState.E`` kernel runs with no
# derivative state (inference performance preserved). The deprecated direct flags
# (``compute_forces`` / ``compute_charge_gradients`` / ``compute_virial`` /
# ``hybrid_forces``) are served by :mod:`_ewald_direct` (tape-free forward kernels).
# Output dtype convention (unchanged):
# - Energies: always float64 for numerical-stability accumulation.
# - Forces / virial / charge gradients: accumulated in float64, returned in the
# input precision (float32 or float64).
def _atom_ranges(
batch_idx: torch.Tensor, num_systems: int
) -> tuple[torch.Tensor, torch.Tensor]:
"""Per-system [start, end) atom index ranges from a sorted ``batch_idx``.
The batched factory kernels assume atoms are grouped contiguously by system
(the existing batched-Ewald contract); ``atom_start``/``atom_end`` are int32.
"""
device = batch_idx.device
if num_systems == 1:
starts = torch.zeros((1,), dtype=torch.int32, device=device)
ends = torch.full(
(1,),
batch_idx.shape[0],
dtype=torch.int32,
device=device,
)
return starts, ends
counts = torch.zeros(num_systems, dtype=torch.long, device=device)
counts = counts.index_add(
0,
batch_idx,
torch.ones(batch_idx.shape[0], dtype=counts.dtype, device=device),
)
ends = torch.cumsum(counts, dim=0)
starts = ends - counts
return starts.to(torch.int32), ends.to(torch.int32)
def _attach_virial_charge_grad(
virial_value: torch.Tensor,
charges: torch.Tensor,
energy_fn,
positions: torch.Tensor,
cell: torch.Tensor,
batch_idx: torch.Tensor | None,
k_vectors_2d: torch.Tensor | None = None,
) -> torch.Tensor:
"""Give the direct (kernel) ``virial`` a charge gradient via strain autograd.
The deprecated ``compute_virial`` output remains differentiable w.r.t.
``charges``. The direct factory kernel output is forward-only, so this
re-attaches the charge gradient with a straight-through:
the value stays the kernel ``virial_value`` while the gradient comes from the
row-vector displacement virial ``W = -dE/dstrain`` of the autograd-connected
energy, recomputed with ``positions`` / ``cell`` detached so only the
``charges`` pathway is live (forces / cell gradients of the direct virial
stay forward-only, as before).
``k_vectors_2d``: when given (reciprocal path), the k-vectors are deformed with
the strain as ``k_s = k @ inv(deform).T`` (the reciprocal lattice transforms
contravariantly with ``cell_s = cell @ deform``) and the 4-argument
``energy_fn(p, q, c, k)`` is called, matching the kernel virial's k-vector
strain response.
"""
if not charges.requires_grad:
return virial_value
num_systems = cell.shape[0]
pos_d = positions.detach()
cell_d = cell.detach()
eye = torch.eye(3, device=positions.device, dtype=positions.dtype).unsqueeze(0)
strain = torch.zeros(
num_systems,
3,
3,
device=positions.device,
dtype=positions.dtype,
requires_grad=True,
)
deform = eye + strain # (S, 3, 3)
atom_sys = (
torch.zeros(positions.shape[0], dtype=torch.int32, device=positions.device)
if batch_idx is None
else batch_idx
)
pos_s = torch.einsum("ni,nij->nj", pos_d, deform[atom_sys])
cell_s = torch.einsum("bij,bjk->bik", cell_d, deform)
if k_vectors_2d is None:
energy = energy_fn(pos_s, charges, cell_s).sum()
else:
k_s = torch.matmul(
k_vectors_2d.detach(), torch.linalg.inv(deform).transpose(1, 2)
)
energy = energy_fn(pos_s, charges, cell_s, k_s).sum()
(dE_dstrain,) = torch.autograd.grad(energy, strain, create_graph=True)
w_torch = (-dE_dstrain).to(virial_value.dtype)
# Straight-through: value from the kernel, charge gradient from ``w_torch``.
return virial_value.detach() + (w_torch - w_torch.detach())
def _real_space_energy_outputs(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
alpha: torch.Tensor,
*,
batch_idx: torch.Tensor | None,
idx_j: torch.Tensor | None,
neighbor_ptr: torch.Tensor | None,
neighbor_shifts: torch.Tensor | None,
neighbor_matrix: torch.Tensor | None,
neighbor_matrix_shifts: torch.Tensor | None,
mask_value: int,
want_forces: bool = False,
want_charge_grad: bool = False,
want_virial: bool = False,
energy_layout: Literal["atom", "system"] = "atom",
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
"""Per-atom real-space Ewald energy plus optional deprecated direct outputs.
Differentiable in ``positions`` / ``charges`` through the explicit chain and in ``cell``
through :func:`real_space_cell_connect` (literal ``dE/dcell``). When no input
requires grad, the plain forward-only chain op runs (no derivative state).
Deprecated direct outputs reuse the same forward launch and remain forward-only.
"""
num_atoms = positions.shape[0]
device = positions.device
use_matrix = neighbor_matrix is not None
if num_atoms == 0:
energy = torch.zeros(
num_atoms if energy_layout == "atom" else cell.shape[0],
device=device,
dtype=torch.float64,
)
forces = (
torch.zeros(num_atoms, 3, device=device, dtype=positions.dtype)
if want_forces
else None
)
charge_grads = (
torch.zeros(num_atoms, device=device, dtype=torch.float64)
if want_charge_grad
else None
)
virial = (
torch.zeros(cell.shape[0], 3, 3, device=device, dtype=positions.dtype)
if want_virial
else None
)
return energy, forces, charge_grads, virial
if use_matrix:
if neighbor_matrix_shifts is None:
raise ValueError(
"neighbor_matrix_shifts is required when using neighbor_matrix format"
)
idx_j_t = torch.zeros(0, dtype=torch.int32, device=device)
neighbor_ptr_t = torch.zeros(0, dtype=torch.int32, device=device)
neighbor_shifts_t = torch.zeros(0, 3, dtype=torch.int32, device=device)
neighbor_matrix_t = neighbor_matrix.to(torch.int32)
neighbor_matrix_shifts_t = neighbor_matrix_shifts.to(torch.int32)
else:
if idx_j is None:
raise ValueError("neighbor_ptr is required when using neighbor_list format")
if neighbor_shifts is None:
raise ValueError(
"neighbor_shifts is required when using neighbor_list format"
)
idx_j_t = idx_j.to(torch.int32)
neighbor_ptr_t = neighbor_ptr.to(torch.int32)
neighbor_shifts_t = neighbor_shifts.to(torch.int32)
neighbor_matrix_t = torch.zeros(num_atoms, 0, dtype=torch.int32, device=device)
neighbor_matrix_shifts_t = torch.zeros(
num_atoms, 0, 3, dtype=torch.int32, device=device
)
# Forward-precompute gating: pick the fused forward specialization by the
# requires-grad set so energy + the dE/dR / dE/dq caches come from ONE launch and
# the first backward is a cheap scale. The cell first-order grad is owned by the
# Torch ``_RealCellGrad`` connector (literal dE/dcell), so the chain caches only
# the position / charge first-order state.
need_pos = bool(positions.requires_grad or want_forces)
need_charge = bool(charges.requires_grad or want_charge_grad)
need_cell = bool(cell.requires_grad)
ensure_electrostatics_ops_registered()
if batch_idx is None:
real_op = (
torch.ops.nvalchemiops.ewald_real_energy_single
if energy_layout == "atom"
else torch.ops.nvalchemiops.ewald_real_system_energy_single
)
energy, dEdR, dEdq, dedcell, direct_virial = real_op(
positions,
charges,
cell,
alpha,
idx_j_t,
neighbor_ptr_t,
neighbor_shifts_t,
neighbor_matrix_t,
neighbor_matrix_shifts_t,
int(mask_value),
use_matrix,
need_pos,
need_charge,
need_cell,
want_virial,
)
else:
real_op = (
torch.ops.nvalchemiops.ewald_real_energy_batch
if energy_layout == "atom"
else torch.ops.nvalchemiops.ewald_real_system_energy_batch
)
energy, dEdR, dEdq, dedcell, direct_virial = real_op(
positions,
charges,
cell,
alpha,
batch_idx.to(torch.int32),
idx_j_t,
neighbor_ptr_t,
neighbor_shifts_t,
neighbor_matrix_t,
neighbor_matrix_shifts_t,
int(mask_value),
use_matrix,
need_pos,
need_charge,
need_cell,
want_virial,
)
forces = -dEdR.to(positions.dtype) if want_forces else None
charge_grads = dEdq if want_charge_grad else None
virial = direct_virial if want_virial else None
# Connect ``cell`` (literal dE/dcell) when it carries grad. The connector is a
# value-zero straight-through term, so the energy value is unchanged. For the
# matrix layout ``dedcell`` is the forward-fused per-atom cache (first-order
# backward is a pure scatter, no kernel, no edge list): the (potentially huge)
# edge-list build is DEFERRED -- the raw neighbor matrix is threaded through and
# flattened to edges only inside the rare double-backward branch. For CSR the
# cheap edges are built here (the edge-kernel path needs them) and the matrix is
# passed empty.
if need_cell:
empty_edges = torch.zeros(0, dtype=torch.long, device=device)
if use_matrix:
edge_i = edge_j = empty_edges
unit_shifts = torch.zeros(0, 3, dtype=torch.int32, device=device)
nm_for_grad = neighbor_matrix_t
nms_for_grad = neighbor_matrix_shifts_t
else:
# CSR -> edges: edge_i is the row owner per edge.
edge_i = torch.repeat_interleave(
torch.arange(num_atoms, device=device),
neighbor_ptr_t.to(torch.long).diff(),
)
edge_j = idx_j_t.to(torch.long)
unit_shifts = neighbor_shifts_t
# Empty matrix -> the connector backward uses the CSR edges directly.
nm_for_grad = torch.zeros(num_atoms, 0, dtype=torch.int32, device=device)
nms_for_grad = torch.zeros(
num_atoms, 0, 3, dtype=torch.int32, device=device
)
energy = real_space_cell_connect(
energy,
positions,
charges,
cell,
alpha,
edge_i,
edge_j,
unit_shifts,
batch_idx,
dedcell.detach(),
nm_for_grad,
nms_for_grad,
int(mask_value),
energy_layout,
)
return energy, forces, charge_grads, virial
def _real_space_energy(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
alpha: torch.Tensor,
*,
batch_idx: torch.Tensor | None,
idx_j: torch.Tensor | None,
neighbor_ptr: torch.Tensor | None,
neighbor_shifts: torch.Tensor | None,
neighbor_matrix: torch.Tensor | None,
neighbor_matrix_shifts: torch.Tensor | None,
mask_value: int,
) -> torch.Tensor:
"""Per-atom real-space Ewald energy, connected to autograd via the explicit chain."""
energy, _, _, _ = _real_space_energy_outputs(
positions,
charges,
cell,
alpha,
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
)
return energy
def _apply_reciprocal_corrections(
e_ksum: torch.Tensor,
charges: torch.Tensor,
volume: torch.Tensor,
alpha: torch.Tensor,
batch_idx: torch.Tensor | None,
) -> torch.Tensor:
"""Apply Ewald reciprocal self-energy + background corrections.
``volume`` and ``total_charge`` carry gradients back to ``cell`` and
``charges``; ``alpha`` is setup-only and is detached from public autograd.
"""
alpha = alpha.detach()
if batch_idx is None:
total_charge = charges.sum().reshape(1)
return ewald_energy_corrections(e_ksum, charges, volume, alpha, total_charge)
if volume.shape[0] == 1:
total_charges = charges.sum().reshape(1)
else:
total_charges = torch.zeros(
volume.shape[0],
dtype=charges.dtype,
device=charges.device,
)
total_charges = total_charges.index_add(0, batch_idx.to(torch.long), charges)
return ewald_energy_corrections_batch(
e_ksum,
charges,
batch_idx.to(torch.int32),
volume,
alpha,
total_charges,
)
def _reciprocal_system_energy_torch(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
k_vectors_2d: torch.Tensor,
alpha: torch.Tensor,
batch_idx: torch.Tensor | None,
num_systems: int,
) -> torch.Tensor:
"""Return pure-Torch reciprocal system energies for terminal facades."""
cell_3d = cell if cell.dim() == 3 else cell.unsqueeze(0)
volume = torch.abs(torch.linalg.det(cell_3d)).to(torch.float64)
ksum = _recip_ksum_energy_torch(
positions,
charges,
k_vectors_2d,
volume,
alpha,
batch_idx,
num_systems,
)
energy = _apply_reciprocal_corrections(
ksum,
charges,
volume,
alpha,
batch_idx,
)
return _reduce_atom_energy(energy, batch_idx, num_systems)
def _reciprocal_space_energy(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
k_vectors: torch.Tensor,
alpha: torch.Tensor,
*,
batch_idx: torch.Tensor | None,
max_atoms_per_system: int | None = None,
energy_reduction: Literal["atom", "system"] = "atom",
preserve_k_vector_grad: bool = False,
) -> torch.Tensor:
"""Per-atom reciprocal-space Ewald energy, connected to autograd via the chain.
Energy = k-sum (explicit chain, differentiable in positions, charges, and
reciprocal vectors; cell-dependent paths compose through reciprocal-vector
and volume graphs) minus the Torch-native self + background corrections.
"""
num_atoms = positions.shape[0]
device = positions.device
if num_atoms == 0:
return torch.zeros(num_atoms, device=device, dtype=torch.float64)
# Forward-precompute gating: the fused recip forward emits energy + the
# detached dE/dR / dE/dq caches in one atom-major pass when positions / charges
# require grad, so the first backward scales them instead of re-running the
# atom-major ``compute`` kernel (the k/V cell grads stay on the cheap k-major
# recompute). ``cell`` flows through ``k_vectors(cell)`` / ``volume`` as before.
need_pos = bool(positions.requires_grad)
need_charge = bool(charges.requires_grad)
# ``need_cell`` gates the k-major grad_kvectors / grad_volume backward owned by
# the recip chain. The public component preserves a supplied k-vector graph when
# both cell and k_vectors require grad; otherwise vectors are fixed Cartesian
# metadata for cell derivatives. Leaf vectors without a cell edge still receive
# dE/dk; physical strain requires vectors generated from the differentiable cell.
need_cell = bool(cell.requires_grad)
ensure_electrostatics_ops_registered()
if batch_idx is None:
num_systems = 1
num_k = k_vectors.shape[-2]
k_vectors_2d = k_vectors.reshape(1, num_k, 3)
volume = torch.abs(
torch.det(cell.reshape(3, 3) if cell.dim() == 2 else cell[0])
).reshape(1)
e_ksum, dEdR, dEdq, _cellgrad_cache = (
torch.ops.nvalchemiops.ewald_recip_energy_single(
positions,
charges,
cell,
k_vectors_2d,
volume.to(torch.float64),
alpha,
need_pos,
need_charge,
need_cell,
)
)
else:
num_systems = cell.shape[0]
num_k = k_vectors.shape[-2]
k_vectors_2d = (
k_vectors
if k_vectors.dim() == 3
else k_vectors.reshape(1, num_k, 3).expand(num_systems, num_k, 3)
)
volume = torch.abs(torch.linalg.det(cell)).to(torch.float64)
atom_start, atom_end = _atom_ranges(batch_idx, num_systems)
max_atoms_bound = 0
if max_atoms_per_system is not None:
max_atoms_bound = int(max_atoms_per_system)
if positions.shape[0] == 0:
if max_atoms_bound < 0:
raise ValueError("max_atoms_per_system must be non-negative")
elif max_atoms_bound <= 0:
raise ValueError(
"max_atoms_per_system must be positive for non-empty batches"
)
elif max_atoms_bound > positions.shape[0]:
raise ValueError(
"max_atoms_per_system cannot exceed the total number of atoms"
)
e_ksum, dEdR, dEdq, _cellgrad_cache = (
torch.ops.nvalchemiops.ewald_recip_energy_batch(
positions,
charges,
cell,
k_vectors_2d,
volume,
alpha,
batch_idx.to(torch.int32),
atom_start,
atom_end,
need_pos,
need_charge,
need_cell,
max_atoms_bound,
)
)
correction = _apply_reciprocal_corrections(
torch.zeros_like(e_ksum), charges, volume, alpha, batch_idx
)
def _system_fallback(
p,
q,
c,
fallback_batch_idx,
fallback_k_vectors,
fallback_alpha,
):
return _reciprocal_system_energy_torch(
p,
q,
c,
fallback_k_vectors,
fallback_alpha,
fallback_batch_idx,
num_systems,
)
if energy_reduction == "system" and cell.requires_grad:
if preserve_k_vector_grad:
return _reduce_atom_energy(e_ksum + correction, batch_idx, num_systems)
return _InjectCachedEvalGradWithFallback.apply(
(e_ksum + correction).detach(),
positions,
charges,
cell,
None,
None,
None,
batch_idx,
_system_fallback,
"system",
num_systems,
True,
True,
k_vectors_2d,
alpha,
)
if (
energy_reduction == "system"
and not cell.requires_grad
and (positions.requires_grad or charges.requires_grad)
):
e_ksum = _InjectCachedEvalGrad.apply(
e_ksum,
positions,
charges,
cell,
dEdR.detach() if positions.requires_grad else None,
dEdq.detach() if charges.requires_grad else None,
None,
batch_idx,
"system",
num_systems,
)
return e_ksum + _reduce_atom_energy(correction, batch_idx, num_systems)
energies = e_ksum + correction
if energy_reduction == "system":
return _reduce_atom_energy(energies, batch_idx, num_systems)
return energies
###########################################################################################
########################### Public Wrapper APIs ###########################################
###########################################################################################
[docs]
def ewald_real_space(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
alpha: torch.Tensor,
neighbor_list: torch.Tensor | None = None,
neighbor_ptr: torch.Tensor | None = None,
neighbor_shifts: torch.Tensor | None = None,
neighbor_matrix: torch.Tensor | None = None,
neighbor_matrix_shifts: torch.Tensor | None = None,
mask_value: int | None = None,
batch_idx: torch.Tensor | 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",
) -> torch.Tensor | tuple[torch.Tensor, ...]:
r"""Compute real-space Ewald energy and optionally forces, charge gradients, and virial.
Computes the damped Coulomb interactions for atom pairs within the real-space
cutoff. The complementary error function (erfc) damping ensures rapid
convergence in real space.
Parameters
----------
positions : torch.Tensor, shape (N, 3)
Atomic coordinates.
charges : torch.Tensor, shape (N,)
Atomic partial charges.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrices.
alpha : torch.Tensor, shape (1,) or (B,)
Ewald splitting parameter(s).
neighbor_list : torch.Tensor, shape (2, M), optional
Neighbor list in COO format.
neighbor_ptr : torch.Tensor, shape (N+1,), optional
CSR row pointers for neighbor list.
neighbor_shifts : torch.Tensor, shape (M, 3), optional
Periodic image shifts for neighbor list.
neighbor_matrix : torch.Tensor, shape (N, max_neighbors), optional
Dense neighbor matrix format.
neighbor_matrix_shifts : torch.Tensor, shape (N, max_neighbors, 3), optional
Periodic image shifts for neighbor_matrix.
mask_value : int, optional
Value indicating invalid entries in neighbor_matrix. Defaults to N.
batch_idx : torch.Tensor, shape (N,), optional
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``.
compute_forces : bool, default=False
Whether to compute explicit component forces. This direct output is kept
for no-autograd MD/inference use; use energy autograd for differentiable
training.
compute_charge_gradients : bool, default=False
Whether to compute explicit component charge gradients. This direct
output follows the same no-autograd contract as ``compute_forces``.
compute_virial : bool, default=False
Whether to compute the component virial tensor
:math:`W = -\partial E / \partial \varepsilon`.
Stress = -virial / volume.
hybrid_forces : bool, default=False
Enables the legacy direct-output path. When ``charges.requires_grad``,
ordinary first-order losses whose cotangent is uniform within each
system use detached positions/cell and cached charge gradients through
a straight-through connector. Non-uniform per-atom losses and
``create_graph=True`` rebuild the eager energy graph with geometry and
charge-chain derivatives. Fixed-charge hybrid calls remain forward-only.
Forces and virial are forward-only. Do not add direct analytical forces
to a fallback-derived full geometry gradient.
energy_reduction : {"atom", "system"}, default="atom"
Return per-atom energies ``(N,)`` or summed per-system energies ``(B,)``.
Returns
-------
energies : torch.Tensor, shape (N,) or (B,)
Real-space Ewald energy: per-atom when ``energy_reduction="atom"``,
per-system when ``energy_reduction="system"``.
forces : torch.Tensor, shape (N, 3), optional
Direct component forces (if compute_forces=True).
charge_gradients : torch.Tensor, shape (N,), optional
Direct component charge gradients (if compute_charge_gradients=True).
virial : torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor (if compute_virial=True). Always last in the tuple.
Note
----
Energies are always float64 for numerical stability during accumulation.
Forces, virial, and charge gradients match the input dtype (float32 or float64).
When ``charges`` is a non-leaf tensor that may depend on ``positions``
(:math:`q = q(R)`), ordinary first-order losses may use cached partial
derivatives and let PyTorch apply
:math:`\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R` once. Weighted
losses and higher-order derivatives recompute safe partials or connected
gradients as needed to avoid double-counting that chain term (issue #115).
Hybrid direct-output mode uses the same cached fallback connector so
weighted :math:`q = q(R)` losses can recover a valid energy gradient when
the forward energy was detached.
"""
_validate_energy_reduction(energy_reduction)
num_systems = int(cell.shape[0]) if cell.dim() == 3 else 1
def _select_energy(energy):
if energy_reduction == "system":
return _reduce_atom_energy(energy, batch_idx, num_systems)
return energy
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 and not torch.compiler.is_compiling():
warnings.warn(
_component_direct_output_deprecation_msg(
"ewald_real_space", component_deprecated_flags
),
DeprecationWarning,
stacklevel=2,
)
if mask_value is None:
mask_value = positions.shape[0]
# The factory kernels index ``alpha[isys]``; accept a 0-d scalar alpha.
if alpha.dim() == 0:
alpha = alpha.reshape(1)
alpha = _detach_setup_tensor(alpha)
if neighbor_list is None and neighbor_matrix is None:
raise ValueError("Either neighbor_list or neighbor_matrix must be provided")
if neighbor_list is not None and neighbor_ptr is None:
raise ValueError("neighbor_ptr is required when using neighbor_list format")
idx_j = neighbor_list[1] if neighbor_list is not None else None
# Helper to build the return tuple from raw outputs using match dispatch.
def _build_result(energies, forces=None, charge_grads=None, virial=None):
match (
compute_forces and forces is not None,
compute_charge_gradients and charge_grads is not None,
compute_virial and virial is not None,
):
case (True, True, True):
return energies, forces, charge_grads, virial
case (True, True, False):
return energies, forces, charge_grads
case (True, False, True):
return energies, forces, virial
case (True, False, False):
return energies, forces
case (False, True, True):
return energies, charge_grads, virial
case (False, True, False):
return energies, charge_grads
case (False, False, True):
return energies, virial
case _:
return energies
want_direct = compute_forces or compute_charge_gradients or compute_virial
if energy_reduction == "system" and not (hybrid_forces and charges.requires_grad):
system_positions = positions.detach() if hybrid_forces else positions
system_cell = cell.detach() if hybrid_forces else cell
energies, forces, charge_grads, virial = _real_space_energy_outputs(
system_positions,
charges,
system_cell,
alpha,
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
want_forces=compute_forces,
want_charge_grad=compute_charge_gradients or charges.requires_grad,
want_virial=compute_virial,
energy_layout="system",
)
if compute_virial and charges.requires_grad and not hybrid_forces:
def _rs_energy_fn(p, q, c):
return _real_space_energy(
p,
q,
c,
alpha,
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
)
virial = _attach_virial_charge_grad(
virial, charges, _rs_energy_fn, positions, cell, batch_idx
)
charge_grads_out = (
charge_grads.to(positions.dtype) if charge_grads is not None else None
)
return _build_result(energies, forces, charge_grads_out, virial)
if hybrid_forces:
# Positions/cell detached from the graph; charge gradients attached via the
# lazy cached-eval trick. Forces/virial forward-only.
energies, forces, charge_grads, virial = _real_space_energy_outputs(
positions.detach(),
charges.detach(),
cell.detach(),
alpha.detach(),
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
want_forces=compute_forces,
want_charge_grad=True,
want_virial=compute_virial,
)
if charges.requires_grad:
def _fallback(
p,
q,
c,
fallback_batch_idx,
fallback_alpha,
fallback_idx_j,
fallback_neighbor_ptr,
fallback_neighbor_shifts,
fallback_neighbor_matrix,
fallback_neighbor_matrix_shifts,
):
return _real_space_energy(
p,
q,
c,
fallback_alpha,
batch_idx=fallback_batch_idx,
idx_j=fallback_idx_j,
neighbor_ptr=fallback_neighbor_ptr,
neighbor_shifts=fallback_neighbor_shifts,
neighbor_matrix=fallback_neighbor_matrix,
neighbor_matrix_shifts=fallback_neighbor_matrix_shifts,
mask_value=mask_value,
)
energies = _InjectCachedEvalGradWithFallback.apply(
energies.detach(),
positions,
charges,
cell,
None,
charge_grads.detach(),
None,
batch_idx,
_fallback,
energy_reduction,
num_systems,
False,
False,
alpha,
idx_j,
neighbor_ptr,
neighbor_shifts,
neighbor_matrix,
neighbor_matrix_shifts,
)
else:
energies = _select_energy(energies)
return _build_result(energies, forces, charge_grads.to(positions.dtype), virial)
if (
not want_direct
and not cell.requires_grad
and not alpha.requires_grad
and (positions.requires_grad or charges.requires_grad)
):
# Ordinary scalar first-derivative evaluations can use detached direct
# caches. Weighted losses and create_graph=True rebuild the true energy
# graph lazily in the custom backward below.
energies, forces, charge_grads, _virial = _real_space_energy_outputs(
positions.detach(),
charges.detach(),
cell.detach(),
alpha.detach(),
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
want_forces=positions.requires_grad,
want_charge_grad=charges.requires_grad,
want_virial=False,
)
def _fallback(
p,
q,
c,
fallback_batch_idx,
fallback_alpha,
fallback_idx_j,
fallback_neighbor_ptr,
fallback_neighbor_shifts,
fallback_neighbor_matrix,
fallback_neighbor_matrix_shifts,
):
return _real_space_energy(
p,
q,
c,
fallback_alpha,
batch_idx=fallback_batch_idx,
idx_j=fallback_idx_j,
neighbor_ptr=fallback_neighbor_ptr,
neighbor_shifts=fallback_neighbor_shifts,
neighbor_matrix=fallback_neighbor_matrix,
neighbor_matrix_shifts=fallback_neighbor_matrix_shifts,
mask_value=mask_value,
)
energies = _InjectCachedEvalGradWithFallback.apply(
energies.detach(),
positions,
charges,
cell,
-forces.detach() if positions.requires_grad else None,
charge_grads.detach() if charges.requires_grad else None,
None,
batch_idx,
_fallback,
energy_reduction,
num_systems,
False,
False,
alpha,
idx_j,
neighbor_ptr,
neighbor_shifts,
neighbor_matrix,
neighbor_matrix_shifts,
)
return energies
# Differentiable energy plus optional deprecated direct outputs from one chain
# forward launch. Autograd still propagates only from the energy output.
energies, forces, charge_grads, virial = _real_space_energy_outputs(
positions,
charges,
cell,
alpha,
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
want_forces=compute_forces,
want_charge_grad=compute_charge_gradients or charges.requires_grad,
want_virial=compute_virial,
)
if not want_direct:
return energies
# The deprecated direct virial is differentiable w.r.t. charges; re-attach that
# gradient (value stays the kernel output) via the strain virial of the
# autograd-connected real-space energy.
if compute_virial and charges.requires_grad:
def _rs_energy_fn(p, q, c):
return _real_space_energy(
p,
q,
c,
alpha,
batch_idx=batch_idx,
idx_j=idx_j,
neighbor_ptr=neighbor_ptr,
neighbor_shifts=neighbor_shifts,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
mask_value=mask_value,
)
virial = _attach_virial_charge_grad(
virial, charges, _rs_energy_fn, positions, cell, batch_idx
)
charge_grads_out = (
charge_grads.to(positions.dtype) if charge_grads is not None else None
)
if (
not cell.requires_grad
and (positions.requires_grad or charges.requires_grad)
and (forces is not None or charge_grads is not None)
and (not positions.requires_grad or forces is not None)
and (not charges.requires_grad or charge_grads is not None)
):
energies = _InjectCachedEvalGrad.apply(
energies,
positions,
charges,
cell,
-forces.detach()
if positions.requires_grad and forces is not None
else None,
charge_grads.detach()
if charges.requires_grad and charge_grads is not None
else None,
None,
batch_idx,
energy_reduction,
num_systems,
)
else:
energies = _select_energy(energies)
return _build_result(energies, forces, charge_grads_out, virial)
[docs]
def ewald_reciprocal_space(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
k_vectors: torch.Tensor,
alpha: torch.Tensor,
batch_idx: torch.Tensor | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
hybrid_forces: bool = False,
*,
max_atoms_per_system: int | None = None,
energy_reduction: Literal["atom", "system"] = "atom",
) -> torch.Tensor | tuple[torch.Tensor, ...]:
r"""Compute reciprocal-space Ewald energy and optionally forces, charge gradients, virial.
Computes the smooth long-range electrostatic contribution using structure
factors in reciprocal space.
Parameters
----------
positions : torch.Tensor, shape (N, 3)
Atomic coordinates.
charges : torch.Tensor, shape (N,)
Atomic partial charges.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrices.
k_vectors : torch.Tensor
Reciprocal lattice vectors. Shape (K, 3) for single system, (B, K, 3) for batch.
When both ``cell`` and ``k_vectors`` require gradients, the supplied
k-vector graph is preserved. Physical strain derivatives require vectors
generated from the same differentiable cell. Non-differentiable vectors,
and grad-bearing leaf vectors without a cell edge, are fixed Cartesian
metadata for cell derivatives.
alpha : torch.Tensor, shape (1,) or (B,)
Ewald splitting parameter(s).
batch_idx : torch.Tensor, shape (N,), optional
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``.
compute_forces : bool, default=False
Whether to compute explicit component forces. This direct output is kept
for no-autograd MD/inference use; use energy autograd for differentiable
training.
compute_charge_gradients : bool, default=False
Whether to compute explicit component charge gradients. This direct
output follows the same no-autograd contract as ``compute_forces``.
compute_virial : bool, default=False
Whether to compute the component virial tensor
:math:`W = -\partial E / \partial \varepsilon`.
Stress = -virial / volume.
hybrid_forces : bool, default=False
Enables the legacy direct-output path. With ``charges.requires_grad``,
uniform first-order cotangents use cached charge gradients; non-uniform
per-atom losses and ``create_graph=True`` rebuild the eager energy graph
with geometry and charge-chain derivatives. Fixed-charge hybrid calls
remain forward-only. See :func:`ewald_real_space` for the complete
contract.
max_atoms_per_system : int, optional, keyword-only
Maximum number of atoms in any single system when ``batch_idx`` is
provided. Passing this host-known upper bound avoids CUDA host
synchronization from launch-size inference in the reciprocal kernel.
Overestimates are safe but may launch extra blocks. When omitted, the
bound is inferred from ``atom_start`` / ``atom_end`` and may
synchronize on CUDA.
energy_reduction : {"atom", "system"}, default="atom"
Return per-atom energies ``(N,)`` or summed per-system energies ``(B,)``.
Returns
-------
energies : torch.Tensor, shape (N,) or (B,)
Reciprocal-space Ewald energy: per-atom when
``energy_reduction="atom"``, per-system when
``energy_reduction="system"``.
forces : torch.Tensor, shape (N, 3), optional
Direct component forces (if compute_forces=True).
charge_gradients : torch.Tensor, shape (N,), optional
Direct component charge gradients (if compute_charge_gradients=True).
virial : torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor (if compute_virial=True). Always last in the tuple.
Note
----
Energies are always float64 for numerical stability during accumulation.
Forces, virial, and charge gradients match the input dtype (float32 or float64).
For eager execution, a differentiable ``cell`` paired with fixed Cartesian
``k_vectors`` emits a warning because its cell derivative is not the physical
Ewald strain virial. This advisory warning is suppressed under
``torch.compile``. Generate vectors from the differentiable cell with fixed
Miller bounds for physical strain derivatives.
When ``charges`` is a non-leaf tensor that may depend on ``positions``
(:math:`q = q(R)`), ordinary first-order losses may use cached partial
derivatives and let PyTorch apply
:math:`\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R` once. Weighted
losses and higher-order derivatives recompute safe partials or connected
gradients as needed to avoid double-counting that chain term (issue #115).
"""
_validate_energy_reduction(energy_reduction)
allow_cell_grad_with_k_vectors = bool(
cell.requires_grad and k_vectors.requires_grad
)
k_vectors_fixed_for_cell = bool(not k_vectors.requires_grad or k_vectors.is_leaf)
if (
torch.is_grad_enabled()
and cell.requires_grad
and k_vectors_fixed_for_cell
and not hybrid_forces
and not torch.compiler.is_compiling()
):
warnings.warn(
"ewald_reciprocal_space received k_vectors that are fixed Cartesian "
"metadata for cell derivatives. If differentiated with respect to "
"cell or strain, this call does not produce the physical Ewald strain "
"virial. Generate k_vectors from the differentiable cell with fixed "
"miller_bounds, or use ewald_summation with k_cutoff.",
UserWarning,
stacklevel=2,
)
return _ewald_reciprocal_space(
positions=positions,
charges=charges,
cell=cell,
k_vectors=k_vectors,
alpha=alpha,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=compute_charge_gradients,
compute_virial=compute_virial,
hybrid_forces=hybrid_forces,
allow_cell_grad_with_k_vectors=allow_cell_grad_with_k_vectors,
preserve_k_vector_grad=allow_cell_grad_with_k_vectors,
max_atoms_per_system=max_atoms_per_system,
energy_reduction=energy_reduction,
)
def _ewald_reciprocal_space(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
k_vectors: torch.Tensor,
alpha: torch.Tensor,
batch_idx: torch.Tensor | None = None,
compute_forces: bool = False,
compute_charge_gradients: bool = False,
compute_virial: bool = False,
hybrid_forces: bool = False,
allow_cell_grad_with_k_vectors: bool = False,
preserve_k_vector_grad: bool = False,
max_atoms_per_system: int | None = None,
energy_reduction: Literal["atom", "system"] = "atom",
) -> torch.Tensor | tuple[torch.Tensor, ...]:
"""Private reciprocal-space implementation with an internal cell-gradient path."""
_validate_energy_reduction(energy_reduction)
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 and not torch.compiler.is_compiling():
warnings.warn(
_component_direct_output_deprecation_msg(
"ewald_reciprocal_space", component_deprecated_flags
),
DeprecationWarning,
stacklevel=2,
)
is_batch = batch_idx is not None
# The factory kernels index ``alpha[isys]``; accept a 0-d scalar alpha.
if alpha.dim() == 0:
alpha = alpha.reshape(1)
alpha = _detach_setup_tensor(alpha)
# Normalize k-vectors to a (S, K, 3) tensor for the factory kernels.
if is_batch:
num_systems = cell.shape[0] if cell.dim() == 3 else 1
if k_vectors.dim() == 2:
k_vectors_2d = k_vectors.unsqueeze(0).expand(num_systems, *k_vectors.shape)
else:
k_vectors_2d = k_vectors
else:
if k_vectors.dim() == 3:
k_vectors_2d = k_vectors[:1]
else:
k_vectors_2d = k_vectors.unsqueeze(0)
num_systems = int(k_vectors_2d.shape[0])
def _select_energy(energy):
if energy_reduction == "system":
return _reduce_atom_energy(energy, batch_idx, num_systems)
return energy
if not allow_cell_grad_with_k_vectors:
k_vectors_2d = k_vectors_2d.detach()
atom_start = atom_end = None
if is_batch:
atom_start, atom_end = _atom_ranges(batch_idx, k_vectors_2d.shape[0])
# Helper to build the return tuple from raw outputs using match dispatch.
def _build_result(energies, forces=None, charge_grads=None, virial=None):
match (
compute_forces and forces is not None,
compute_charge_gradients and charge_grads is not None,
compute_virial and virial is not None,
):
case (True, True, True):
return energies, forces, charge_grads, virial
case (True, True, False):
return energies, forces, charge_grads
case (True, False, True):
return energies, forces, virial
case (True, False, False):
return energies, forces
case (False, True, True):
return energies, charge_grads, virial
case (False, True, False):
return energies, charge_grads
case (False, False, True):
return energies, virial
case _:
return energies
want_direct = compute_forces or compute_charge_gradients or compute_virial
# No atoms have no reciprocal contribution. Empty k-vector sets still need
# self/background corrections below.
num_atoms = positions.shape[0]
if num_atoms == 0:
zeros_e = torch.zeros(num_atoms, device=positions.device, dtype=torch.float64)
zeros_f = torch.zeros(
num_atoms, 3, device=positions.device, dtype=positions.dtype
)
zeros_cg = torch.zeros(
num_atoms, device=positions.device, dtype=positions.dtype
)
zeros_v = torch.zeros(
num_systems, 3, 3, device=positions.device, dtype=positions.dtype
)
return _build_result(_select_energy(zeros_e), zeros_f, zeros_cg, zeros_v)
if hybrid_forces:
k_vectors_hybrid = k_vectors_2d.detach()
e_ksum, forces, charge_grads, virial = reciprocal_space_direct(
positions.detach(),
charges.detach(),
cell.detach(),
k_vectors_hybrid,
alpha.detach(),
batch_idx=batch_idx,
atom_start=atom_start,
atom_end=atom_end,
want_charge_grad=True,
want_virial=compute_virial,
)
volume = torch.abs(torch.linalg.det(cell.detach().to(torch.float64))).reshape(
k_vectors_hybrid.shape[0]
)
energies = _apply_reciprocal_corrections(
e_ksum,
charges.detach(),
volume,
alpha,
batch_idx,
)
if charges.requires_grad:
def _fallback(
p,
q,
c,
fallback_batch_idx,
fallback_k_vectors,
fallback_alpha,
):
return _reciprocal_space_energy(
p,
q,
c,
fallback_k_vectors,
fallback_alpha,
batch_idx=fallback_batch_idx,
max_atoms_per_system=max_atoms_per_system,
)
energies = _InjectCachedEvalGradWithFallback.apply(
energies,
positions,
charges,
cell,
None,
charge_grads.detach(),
None,
batch_idx,
_fallback,
energy_reduction,
num_systems,
False,
False,
k_vectors_hybrid,
alpha,
)
else:
energies = _select_energy(energies)
return _build_result(energies, forces, charge_grads.to(positions.dtype), virial)
differentiable_inputs = (
positions.requires_grad or charges.requires_grad or cell.requires_grad
)
if want_direct and not differentiable_inputs:
e_ksum, forces, charge_grads, virial = reciprocal_space_direct(
positions,
charges,
cell,
k_vectors_2d,
alpha,
batch_idx=batch_idx,
atom_start=atom_start,
atom_end=atom_end,
want_charge_grad=compute_charge_gradients,
want_virial=compute_virial,
)
volume = torch.abs(torch.linalg.det(cell.to(torch.float64))).reshape(
k_vectors_2d.shape[0]
)
energies = _apply_reciprocal_corrections(
e_ksum,
charges,
volume,
alpha,
batch_idx,
)
charge_grads_out = (
charge_grads.to(positions.dtype) if charge_grads is not None else None
)
return _build_result(_select_energy(energies), forces, charge_grads_out, virial)
energies = _reciprocal_space_energy(
positions,
charges,
cell,
k_vectors_2d,
alpha,
batch_idx=batch_idx,
max_atoms_per_system=max_atoms_per_system,
energy_reduction=energy_reduction if not want_direct else "atom",
preserve_k_vector_grad=preserve_k_vector_grad,
)
if not want_direct:
return energies
_e_ksum, forces, charge_grads, virial = reciprocal_space_direct(
positions,
charges,
cell,
k_vectors_2d,
alpha,
batch_idx=batch_idx,
atom_start=atom_start,
atom_end=atom_end,
want_charge_grad=compute_charge_gradients or charges.requires_grad,
want_virial=compute_virial,
)
# Re-attach the direct virial's charge gradient (value stays the kernel
# output) via the strain virial of the reciprocal energy. The strain deforms
# the k-vectors too (``_attach_virial_charge_grad`` with ``k_vectors_2d``), so
# the charge gradient matches the kernel virial's ``k_factor`` term.
if compute_virial and charges.requires_grad:
def _rec_energy_fn(p, q, c, k):
return _reciprocal_space_energy(
p,
q,
c,
k,
alpha,
batch_idx=batch_idx,
max_atoms_per_system=max_atoms_per_system,
)
virial = _attach_virial_charge_grad(
virial,
charges,
_rec_energy_fn,
positions,
cell,
batch_idx,
k_vectors_2d=k_vectors_2d,
)
if energy_reduction == "system" and cell.requires_grad:
if preserve_k_vector_grad:
energies = _select_energy(energies)
else:
def _system_fallback(
p,
q,
c,
fallback_batch_idx,
fallback_k_vectors,
fallback_alpha,
):
return _reciprocal_system_energy_torch(
p,
q,
c,
fallback_k_vectors,
fallback_alpha,
fallback_batch_idx,
num_systems,
)
energies = _InjectCachedEvalGradWithFallback.apply(
energies.detach(),
positions,
charges,
cell,
None,
None,
None,
batch_idx,
_system_fallback,
"system",
num_systems,
True,
True,
k_vectors_2d,
alpha,
)
elif (
not cell.requires_grad
and (positions.requires_grad or charges.requires_grad)
and (forces is not None or charge_grads is not None)
):
energies = _InjectCachedEvalGrad.apply(
energies,
positions,
charges,
cell,
-forces.detach()
if positions.requires_grad and forces is not None
else None,
charge_grads.detach()
if charges.requires_grad and charge_grads is not None
else None,
None,
batch_idx,
energy_reduction,
num_systems,
)
else:
energies = _select_energy(energies)
charge_grads_out = (
charge_grads.to(positions.dtype) if charge_grads is not None else None
)
return _build_result(energies, forces, charge_grads_out, virial)
def ewald_summation(
positions: torch.Tensor,
charges: torch.Tensor,
cell: torch.Tensor,
alpha: float | torch.Tensor | None = None,
k_vectors: torch.Tensor | None = None,
k_cutoff: float | None = None,
batch_idx: torch.Tensor | None = None,
neighbor_list: torch.Tensor | None = None,
neighbor_ptr: torch.Tensor | None = None,
neighbor_shifts: torch.Tensor | None = None,
neighbor_matrix: torch.Tensor | None = None,
neighbor_matrix_shifts: torch.Tensor | 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: torch.Tensor | None = None,
slab_correction: bool = False,
*,
miller_bounds: tuple[int, int, int] | torch.Tensor | None = None,
max_atoms_per_system: int | None = None,
energy_reduction: Literal["atom", "system"] = "atom",
) -> tuple[torch.Tensor, ...] | torch.Tensor:
r"""Complete Ewald summation for long-range electrostatics.
Computes total Coulomb energy by combining real-space and reciprocal-space
contributions with self-energy and background corrections.
Parameters
----------
positions : torch.Tensor, shape (N, 3)
Atomic coordinates.
charges : torch.Tensor, shape (N,)
Atomic partial charges.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrices.
alpha : float, torch.Tensor, or None, default=None
Ewald splitting parameter. Auto-estimated if None.
k_vectors : torch.Tensor, optional
Pre-computed reciprocal lattice vectors for fixed-cell reuse.
Caller-supplied vectors are static metadata assumed to correspond to the
current ``cell``; cache-generation derivatives are not recovered. For
physical cell/strain derivatives, omit ``k_vectors`` so vectors are
regenerated from the differentiable ``cell``.
k_cutoff : float, optional
K-space cutoff for generating k_vectors.
miller_bounds : tuple[int, int, int] or torch.Tensor, optional, keyword-only
Precomputed Miller-index half-bounds used when ``k_vectors`` is not
supplied. Passing Python integer bounds avoids deriving range sizes from
device tensors inside regenerated-k-vector loops.
max_atoms_per_system : int, optional, keyword-only
Maximum number of atoms in any single system when ``batch_idx`` is
provided. See :func:`ewald_reciprocal_space` for the sync-free launch
contract.
batch_idx : torch.Tensor, shape (N,), optional
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``.
neighbor_list : torch.Tensor, shape (2, M), optional
Neighbor pairs in COO format.
neighbor_ptr : torch.Tensor, shape (N+1,), optional
CSR row pointers.
neighbor_shifts : torch.Tensor, shape (M, 3), optional
Periodic image shifts for neighbor list.
neighbor_matrix : torch.Tensor, shape (N, max_neighbors), optional
Dense neighbor matrix.
neighbor_matrix_shifts : torch.Tensor, shape (N, max_neighbors, 3), optional
Periodic image shifts for neighbor_matrix.
mask_value : int, optional
Value indicating invalid entries. Defaults to N.
compute_forces : bool, default=False
.. deprecated:: 0.4.0
Deprecated direct-output flag. Compute energy and use
``torch.autograd.grad`` for differentiable forces.
compute_charge_gradients : bool, default=False
.. deprecated:: 0.4.0
Deprecated direct-output flag. Compute energy and use
``torch.autograd.grad`` for :math:`\partial E / \partial q_i`.
compute_virial : bool, default=False
.. deprecated:: 0.4.0
Deprecated direct-output flag for the virial tensor
:math:`W = -\partial E / \partial \varepsilon`.
Stress = -virial / volume.
accuracy : float, default=1e-6
Target accuracy for parameter estimation.
hybrid_forces : bool, default=False
Enables the legacy direct-output path. With ``charges.requires_grad``,
uniform first-order cotangents use cached charge gradients; non-uniform
per-atom losses and ``create_graph=True`` rebuild the eager energy graph
with geometry and charge-chain derivatives. Fixed-charge hybrid calls
remain forward-only. See :func:`ewald_real_space` for the complete
contract.
pbc : torch.Tensor, shape (3,) or (B, 3), dtype=bool, optional
Per-system periodic boundary conditions. Required when
``slab_correction=True``. Each row has True for periodic directions
and False for the non-periodic (slab) direction. A (3,) tensor is
accepted only for single-system calls; batched calls require explicit
(B, 3) per-system pbc. This argument controls the slab correction
geometry; real-space periodic images are determined by the neighbor
list supplied to the Ewald real-space term.
slab_correction : bool, default=False
When True, apply the Yeh-Berkowitz slab correction (with the
Ballenegger et al. 2009 Eq. 29 non-neutral extension) to the total
energy and to forces/charge_grads/virial when those are requested.
Orthorhombic and triclinic slab cells are supported.
energy_reduction : {"atom", "system"}, default="atom"
Return per-atom energies ``(N,)`` or summed per-system energies ``(B,)``.
Returns
-------
energies : torch.Tensor, shape (N,) or (B,)
Total Ewald energy: per-atom when ``energy_reduction="atom"``,
per-system when ``energy_reduction="system"``.
forces : torch.Tensor, shape (N, 3), optional
.. deprecated:: 0.4.0
Deprecated direct forces (if compute_forces=True).
charge_gradients : torch.Tensor, shape (N,), optional
.. deprecated:: 0.4.0
Deprecated direct charge gradients (if compute_charge_gradients=True).
virial : torch.Tensor, shape (1, 3, 3) or (B, 3, 3), optional
Virial tensor (if compute_virial=True). Always last in the tuple.
Note
----
Energies are accumulated in float64 for numerical stability. Deprecated
direct forces, charge gradients, and virials match the input dtype where the
underlying component path returns typed outputs.
When ``charges`` is a non-leaf tensor that may depend on ``positions``
(:math:`q = q(R)`), ordinary first-order losses may use cached partial
derivatives and let PyTorch apply
:math:`\partial E/\partial q \cdot \mathrm{d}q/\mathrm{d}R` once. Weighted
losses and higher-order derivatives recompute safe partials or connected
gradients as needed to avoid double-counting that chain term (issue #115).
Enabled output flags are appended in order: energies, [forces],
[charge_gradients], [virial]. A single output is returned unwrapped;
multiple outputs are returned as a tuple.
Examples
--------
Automatic parameter estimation (recommended for most cases)::
>>> energies = ewald_summation(
... positions, charges, cell,
... neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... accuracy=1e-6,
... )
>>> total_energy = energies.sum()
Explicit parameters with forces::
>>> energies, forces = ewald_summation(
... positions, charges, cell,
... alpha=0.3, k_cutoff=8.0,
... neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... compute_forces=True,
... )
Slab correction for two-dimensional periodic systems::
>>> pbc_slab = torch.tensor([[True, True, False]], device=positions.device)
>>> energies, forces = ewald_summation(
... positions, charges, cell,
... alpha=0.3, k_cutoff=8.0,
... neighbor_list=nl, neighbor_ptr=nptr, neighbor_shifts=shifts,
... pbc=pbc_slab, slab_correction=True,
... compute_forces=True,
... )
"""
_validate_energy_reduction(energy_reduction)
if compute_forces or compute_virial or compute_charge_gradients or hybrid_forces:
if torch.compiler.is_compiling():
_compiled_direct_output_deprecation_signal("ewald_summation")
else:
warnings.warn(
_direct_output_deprecation_msg("ewald_summation"),
DeprecationWarning,
stacklevel=2,
)
device = positions.device
dtype = positions.dtype
num_atoms = positions.shape[0]
cell, num_systems = _prepare_cell(cell)
if alpha is None or (k_cutoff is None and k_vectors is None):
params = estimate_ewald_parameters(positions, cell, batch_idx, accuracy)
if alpha is None:
alpha = params.alpha
if k_cutoff is None:
k_cutoff = params.reciprocal_space_cutoff
alpha_tensor = _detach_setup_tensor(
_prepare_alpha(alpha, num_systems, dtype, device)
)
generated_k_vectors = k_vectors is None
if k_vectors is None:
k_vectors = generate_k_vectors_ewald_summation(
cell, k_cutoff, miller_bounds=miller_bounds
)
if mask_value is None:
mask_value = num_atoms
def _compute_components():
# Compute real-space
real = ewald_real_space(
positions=positions,
charges=charges,
cell=cell,
alpha=alpha_tensor,
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=compute_charge_gradients,
compute_virial=compute_virial,
hybrid_forces=hybrid_forces,
energy_reduction=energy_reduction,
)
# Compute reciprocal-space
reciprocal = _ewald_reciprocal_space(
positions=positions,
charges=charges,
cell=cell,
k_vectors=k_vectors,
alpha=alpha_tensor,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=compute_charge_gradients,
compute_virial=compute_virial,
hybrid_forces=hybrid_forces,
allow_cell_grad_with_k_vectors=generated_k_vectors,
max_atoms_per_system=max_atoms_per_system,
energy_reduction=energy_reduction,
)
return real, reciprocal
if torch.compiler.is_compiling():
rs, rec = _compute_components()
else:
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message=r"The component direct-output flag\(s\).*",
category=DeprecationWarning,
)
rs, rec = _compute_components()
# Optional slab correction: returns a same-shape tuple as rs/rec,
# so it composes uniformly with the named-field combination below.
slab_result: torch.Tensor | tuple[torch.Tensor, ...] | None = None
if slab_correction:
if hybrid_forces:
slab_out = _compute_slab_correction(
positions.detach(),
charges.detach(),
cell.detach(),
pbc,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=True,
compute_virial=compute_virial,
energy_reduction=energy_reduction,
)
slab_energies, slab_forces, slab_charge_grads, slab_virial = (
_unpack_electrostatic_outputs(
slab_out,
compute_forces,
compute_charge_gradients=True,
compute_virial=compute_virial,
)
)
if charges.requires_grad:
slab_energies = _compute_slab_correction(
positions,
charges,
cell,
pbc,
batch_idx=batch_idx,
compute_forces=False,
compute_charge_gradients=False,
compute_virial=False,
energy_reduction=energy_reduction,
)
slab_energies = _InjectChargeGrad.apply(
slab_energies,
charges,
slab_charge_grads,
batch_idx,
energy_reduction,
num_systems,
energy_reduction == "system",
)
slab_result = _build_electrostatic_result(
slab_energies,
slab_forces,
slab_charge_grads,
slab_virial,
compute_forces,
compute_charge_gradients,
compute_virial,
)
else:
slab_result = _compute_slab_correction(
positions,
charges,
cell,
pbc,
batch_idx=batch_idx,
compute_forces=compute_forces,
compute_charge_gradients=compute_charge_gradients,
compute_virial=compute_virial,
energy_reduction=energy_reduction,
)
return _combine_electrostatic_outputs(
rs,
rec,
slab_result,
compute_forces,
compute_charge_gradients,
compute_virial,
)
