# 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.
"""
Parameter Estimation for Ewald and PME Methods (PyTorch)
========================================================
This module provides functions to automatically estimate optimal parameters
for Ewald summation and Particle Mesh Ewald (PME) calculations using PyTorch.
"""
import math
from dataclasses import dataclass
import torch
from nvalchemiops.interactions.electrostatics._pme_mesh import (
_DEFAULT_PME_FFT_PADDING_FRACTION,
_DEFAULT_PME_SPLINE_ORDER,
_round_pme_mesh_dimensions,
)
[docs]
@dataclass
class EwaldParameters:
"""Container for Ewald summation parameters.
All values are tensors of shape (B,), for
single system calculations, the shape is (1,).
Attributes
----------
alpha : torch.Tensor, shape (B,)
Ewald splitting parameter (inverse length units).
real_space_cutoff : torch.Tensor, shape (B,)
Real-space cutoff distance.
reciprocal_space_cutoff : torch.Tensor, shape (B,)
Reciprocal-space cutoff (:math:`|k|` in inverse length units).
"""
alpha: torch.Tensor
real_space_cutoff: torch.Tensor
reciprocal_space_cutoff: torch.Tensor
[docs]
@dataclass
class PMEParameters:
"""Container for PME parameters.
Attributes
----------
alpha : torch.Tensor, shape (B,)
Ewald splitting parameter.
mesh_dimensions : tuple[int, int, int], shape (3,)
Mesh dimensions (nx, ny, nz).
mesh_spacing : torch.Tensor, shape (B, 3)
Actual mesh spacing in each direction.
real_space_cutoff : torch.Tensor, shape (B,)
Real-space cutoff distance.
"""
alpha: torch.Tensor
mesh_dimensions: tuple[int, int, int]
mesh_spacing: torch.Tensor
real_space_cutoff: torch.Tensor
def _count_atoms_per_system(
positions: torch.Tensor, num_systems: int, batch_idx: torch.Tensor | None = None
) -> torch.Tensor:
"""Count number of atoms per system."""
if batch_idx is None:
return torch.tensor(
[positions.shape[0]], dtype=torch.int32, device=positions.device
)
counts = torch.zeros(num_systems, dtype=torch.int32, device=batch_idx.device)
ones = torch.ones_like(batch_idx)
return counts.scatter_add_(0, batch_idx, ones)
[docs]
def estimate_ewald_parameters(
positions: torch.Tensor,
cell: torch.Tensor,
batch_idx: torch.Tensor | None = None,
accuracy: float = 1e-6,
) -> EwaldParameters:
"""Estimate optimal Ewald summation parameters for a given accuracy.
Uses the Kolafa-Perram formula to balance real-space and reciprocal-space
contributions for optimal efficiency at the target accuracy.
Parameters
----------
positions : torch.Tensor, shape (N, 3)
Atomic coordinates.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrix.
batch_idx : torch.Tensor, shape (N,), dtype=int32, optional
System index for each atom. If None, single-system mode.
accuracy : float, default=1e-6
Target accuracy (relative error tolerance).
Returns
-------
EwaldParameters
Dataclass containing alpha, real_space_cutoff, reciprocal_space_cutoff
as ``torch.Tensor`` objects.
"""
if cell.ndim == 2:
cell = cell.unsqueeze(0)
num_systems = cell.shape[0]
# Compute volume per system: (B,)
volume = torch.abs(torch.linalg.det(cell)).squeeze(-1)
# Get number of atoms per system: (B,)
num_atoms = _count_atoms_per_system(positions, num_systems, batch_idx).to(
positions.dtype
)
# Intermediate parameter eta: (B,)
eta = (volume**2 / num_atoms) ** (1.0 / 6.0) / math.sqrt(2.0 * math.pi)
# Error factor from log(accuracy)
error_factor = math.sqrt(-2.0 * math.log(accuracy))
# Real-space cutoff: (B,)
real_space_cutoff = error_factor * eta
# Reciprocal-space cutoff: (B,)
reciprocal_space_cutoff = error_factor / eta
# Splitting parameter alpha: (B,)
alpha = 1.0 / (math.sqrt(2.0) * eta)
return EwaldParameters(
alpha=alpha,
real_space_cutoff=real_space_cutoff,
reciprocal_space_cutoff=reciprocal_space_cutoff,
)
[docs]
def estimate_pme_mesh_dimensions(
cell: torch.Tensor,
alpha: torch.Tensor,
accuracy: float = 1e-6,
mesh_safety_factor: float = 1.0,
*,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
fft_padding_fraction: float = _DEFAULT_PME_FFT_PADDING_FRACTION,
) -> tuple[int, int, int]:
r"""Return accuracy-based FFT mesh dimensions with spline support.
The mesh size along each axis is chosen as
.. math::
K_i = \left\lceil \text{mesh\_safety\_factor} \cdot \frac{2\alpha L_i}{3\varepsilon^{1/5}} \right\rceil
and rounded upward to a mesh with prime factors 2, 3, 5, and 7. For the
default order five, snap to powers of two first, then smooth multiples
of four, within the configured budget of extra total mesh points.
Other spline orders retain upstream power-of-two accuracy sizing. Each
dimension covers the continuous target and assignment support.
``accuracy`` controls the historical fifth-root sizing rule. The estimate
retains its relative sizing convention; independent force checks determine
the measured error for a particular charge distribution.
Parameters
----------
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrix.
alpha : torch.Tensor, shape (B,)
Ewald splitting parameter.
accuracy : float, default=1e-6
Dimensionless input to the legacy mesh-size heuristic.
mesh_safety_factor : float, default=1.0
Caller-requested multiplier on the continuous size estimate. The
default leaves the estimate unchanged.
spline_order : int, default=5
B-spline assignment order. The mesh covers its assignment support.
fft_padding_fraction : float, default=0.25
Maximum fraction of extra total mesh points for snapping to powers
of two, then smooth multiples of four. Zero selects the smallest
2/3/5/7-smooth mesh. Explicit other orders keep dyadic accuracy sizing.
Returns
-------
tuple[int, int, int]
Maximum mesh dimensions (nx, ny, nz) across all systems in batch.
Notes
-----
Measured example on an NVIDIA H100 NVL with Torch: for 131,072 NH3 atoms,
snapping from 486 cubed to 512 cubed increased mesh points by 16.92% and
reduced the energy, forces, and charge-gradient call from 13.96 ms to
11.17 ms. Both runs used order five, float64, a 9 Angstrom cutoff, and
accuracy 1e-6. Times are medians of five groups of ten calls after three
warmups; parameter setup and neighbor construction were excluded.
"""
if cell.ndim == 2:
cell = cell.unsqueeze(0)
cell_lengths = torch.norm(cell, dim=2) # (B, 3)
# Preserve the public estimator's existing relative sizing convention.
accuracy_factor = 3.0 * (accuracy**0.2)
n = (
mesh_safety_factor * 2.0 * alpha[:, None] * cell_lengths / accuracy_factor
) # (B, 3)
max_n = torch.max(n, dim=0).values # (3,)
return _round_pme_mesh_dimensions(
max_n.detach().cpu().numpy(),
spline_order=spline_order,
power_of_two=(spline_order != _DEFAULT_PME_SPLINE_ORDER),
fft_padding_fraction=fft_padding_fraction,
)
[docs]
def estimate_pme_parameters(
positions: torch.Tensor,
cell: torch.Tensor,
batch_idx: torch.Tensor | None = None,
accuracy: float = 1e-6,
real_space_cutoff: float | None = None,
mesh_safety_factor: float = 1.0,
*,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
fft_padding_fraction: float = _DEFAULT_PME_FFT_PADDING_FRACTION,
) -> PMEParameters:
r"""Estimate PME splitting parameters and FFT-friendly mesh dimensions.
Uses the closed-form Essmann/Kolafa-Perram derivation: a single
length scale :math:`\eta = (V^2 / N)^{1/6} / \sqrt{2\pi}` determines both ``rc``
and ``alpha``. Callers who want to pin a specific cutoff (e.g. tied to
neighbor-list update frequency in MD) should pass ``real_space_cutoff``.
Parameters
----------
positions : torch.Tensor, shape (N, 3)
Atomic coordinates.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrix.
batch_idx : torch.Tensor, shape (N,), dtype=int32, optional
System index for each atom.
accuracy : float, default=1e-6
Target accuracy.
real_space_cutoff : float, optional
Caller-supplied cutoff. When given, ``alpha`` is derived from it via
:math:`\alpha = \sqrt{-\log\varepsilon} / r_c`; otherwise ``rc`` and ``alpha`` come from ``eta``.
mesh_safety_factor : float, default=1.0
Caller-requested multiplier on the continuous mesh-size estimate.
spline_order : int, default=5
B-spline assignment order used for mesh sizing.
fft_padding_fraction : float, default=0.25
Maximum fraction of extra total mesh points for snapping to powers
of two, then smooth multiples of four. Zero selects the smallest
2/3/5/7-smooth mesh. Explicit other orders keep dyadic accuracy sizing.
Returns
-------
PMEParameters
Dataclass containing alpha, mesh dimensions, spacing, and cutoffs.
Notes
-----
``accuracy`` retains the public estimator's relative sizing convention.
Independent force checks determine the measured error for a charge
distribution.
"""
if cell.ndim == 2:
cell = cell.unsqueeze(0)
num_systems = cell.shape[0]
volume = torch.abs(torch.linalg.det(cell))
num_atoms = _count_atoms_per_system(positions, num_systems, batch_idx).to(
positions.dtype
)
cell_lengths = torch.norm(cell, dim=2) # (B, 3)
# For batched inputs we share a single rc / α across the batch (one
# neighbor cutoff for all systems). Use median-system properties.
if real_space_cutoff is None:
if num_systems == 1:
n_repr = float(num_atoms[0].item())
v_repr = float(volume[0].item())
else:
n_repr = float(num_atoms.median().item())
v_repr = float(volume.median().item())
eta = (v_repr**2 / n_repr) ** (1.0 / 6.0) / math.sqrt(2.0 * math.pi)
rc_value = math.sqrt(-2.0 * math.log(accuracy)) * eta
alpha_value = 1.0 / (math.sqrt(2.0) * eta)
else:
rc_value = float(real_space_cutoff)
alpha_value = math.sqrt(-math.log(accuracy)) / rc_value
alpha = torch.full(
(num_systems,),
alpha_value,
dtype=positions.dtype,
device=positions.device,
)
rc_tensor = torch.full(
(num_systems,),
rc_value,
dtype=positions.dtype,
device=positions.device,
)
mesh_dims = estimate_pme_mesh_dimensions(
cell,
alpha,
accuracy,
mesh_safety_factor=mesh_safety_factor,
spline_order=spline_order,
fft_padding_fraction=fft_padding_fraction,
)
mesh_dims_tensor = torch.tensor(
mesh_dims, dtype=cell_lengths.dtype, device=cell_lengths.device
)
mesh_spacing = cell_lengths / mesh_dims_tensor # (B, 3)
return PMEParameters(
alpha=alpha,
mesh_dimensions=mesh_dims,
mesh_spacing=mesh_spacing,
real_space_cutoff=rc_tensor,
)
[docs]
@dataclass
class MultipoleEwaldParameters:
"""Container for GTO-Ewald multipole parameters.
Like :class:`EwaldParameters` but with the GTO basis width ``sigma``
propagated through. The Kolafa-Perram balance for the multipole case
has the same ``rcut`` / ``kcut`` formulas as the monopole case, but
``alpha`` differs because the effective Ewald split width is
``sigma_c = sqrt(sigma**2 + 1/(4 alpha**2))`` rather than ``1/(alpha
sqrt(2))``.
Attributes
----------
alpha : torch.Tensor, shape (B,)
Ewald splitting parameter (inverse length units).
sigma : torch.Tensor, shape (B,)
GTO basis width (passed through; physics).
real_space_cutoff : torch.Tensor, shape (B,)
Real-space cutoff distance.
reciprocal_space_cutoff : torch.Tensor, shape (B,)
Reciprocal-space cutoff (``|k|`` in inverse length units).
"""
alpha: torch.Tensor
sigma: torch.Tensor
real_space_cutoff: torch.Tensor
reciprocal_space_cutoff: torch.Tensor
[docs]
@dataclass
class MultipolePMEParameters:
"""Container for GTO-Ewald multipole PME parameters.
Attributes
----------
alpha : torch.Tensor, shape (B,)
Ewald splitting parameter.
sigma : torch.Tensor, shape (B,)
GTO basis width (passed through; physics).
mesh_dimensions : tuple[int, int, int]
Mesh dimensions ``(nx, ny, nz)`` (max across batch).
mesh_spacing : torch.Tensor, shape (B, 3)
Actual mesh spacing per direction.
real_space_cutoff : torch.Tensor, shape (B,)
Real-space cutoff distance.
"""
alpha: torch.Tensor
sigma: torch.Tensor
mesh_dimensions: tuple[int, int, int]
mesh_spacing: torch.Tensor
real_space_cutoff: torch.Tensor
def _kp_eta(volume: torch.Tensor, num_atoms: torch.Tensor) -> torch.Tensor:
r"""Kolafa-Perram cost-balance length scale :math:`\eta = (V^2/N)^{1/6} / \sqrt{2\pi}`.
Common to both the monopole and multipole estimators — captures the
geometric balance between real-space-pair count and reciprocal-space
k-vector count at fixed accuracy.
"""
return (volume**2 / num_atoms) ** (1.0 / 6.0) / math.sqrt(2.0 * math.pi)
def _prepare_sigma(
sigma: float | torch.Tensor,
num_systems: int,
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
"""Broadcast a scalar or per-system ``sigma`` into shape ``(B,)``."""
if isinstance(sigma, torch.Tensor):
sig = sigma.to(dtype=dtype, device=device)
if sig.ndim == 0:
sig = sig.expand(num_systems).clone()
elif sig.shape != (num_systems,):
raise ValueError(
f"sigma must be scalar or shape ({num_systems},); "
f"got {tuple(sig.shape)}"
)
return sig
return torch.full((num_systems,), float(sigma), dtype=dtype, device=device)
[docs]
def estimate_multipole_ewald_parameters(
positions: torch.Tensor,
cell: torch.Tensor,
sigma: float | torch.Tensor,
batch_idx: torch.Tensor | None = None,
accuracy: float = 1e-6,
cost_ratio: float = 1.0,
) -> MultipoleEwaldParameters:
"""Estimate GTO-Ewald multipole parameters at a given target accuracy.
Mirrors :func:`estimate_ewald_parameters` semantics ("relative
energy-error" accuracy via the Kolafa-Perram envelope), adjusted for
the multipole case where the effective Ewald-split width is
``sigma_c = sqrt(sigma**2 + 1/(4 alpha**2))`` rather than
``1/(alpha sqrt(2))``.
Derivation
----------
Both the real-space tail (``erfc(r/(2 sigma_c))``) and the
reciprocal-space envelope (``exp(-k**2 sigma_c**2)``) decay with the
same effective width ``sigma_c * sqrt(2)``. Substituting that for the
monopole's ``eta = 1/(alpha * sqrt(2))`` in Kolafa-Perram gives
``rcut = error_factor * eta``, ``kcut = error_factor / eta``, and
``alpha = 1 / (sqrt(2) * sqrt(eta**2 - 2 * sigma**2))``. The
``sigma -> 0`` limit recovers the monopole formula.
Cost-ratio correction
---------------------
The textbook Kolafa-Perram balance assumes the per-real-space-pair
cost equals the per-k-vector cost. On real hardware (and especially
for the lmax=1 multipole tile kernels) those costs differ — measured
``C_r / C_k`` is in the 20-40x range for our cluster-pair tile
kernels at fp64. The cost-balanced optimum scales as
``eta_eff = eta_KP / cost_ratio**(1/6)``: a 30x cost ratio shrinks
rcut by ~1.76x (and grows kcut by the same factor), which can cut
the real-space pair count by ~5x at the same target accuracy.
The math: with cost ratio ``R = C_r / C_k``, the cost-balanced
formula becomes ``eta_eff = (V^2 / (N * R))^(1/6) / sqrt(2 pi)``.
Setting ``R = 1`` (default) reproduces the canonical KP estimator.
Parameters
----------
positions : torch.Tensor, shape (N, 3) or (N_total, 3)
Atomic coordinates.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrix (matches the multipole-Ewald convention).
sigma : float or torch.Tensor
GTO basis width — same value used by the multipole-Ewald kernel.
Scalar or shape ``(B,)``.
batch_idx : torch.Tensor, shape (N_total,), int32, optional
System index per atom. ``None`` selects single-system mode.
accuracy : float, default 1e-6
Target relative accuracy (matches monopole convention).
cost_ratio : float, default 1.0
Empirical ``C_r / C_k`` ratio — per-pair real-space cost divided
by per-k-vector reciprocal cost on the target hardware. ``1.0``
(default) reproduces canonical Kolafa-Perram. Higher values
shift the optimum toward smaller rcut (fewer pairs) + larger
kcut (more k-vectors), which wins when the per-pair cluster-pair
tile kernel dominates. For the lmax=1 multipole kernels on a
GB10-class GPU, ``cost_ratio`` ~ 30 is a reasonable starting
point — measure on your own hardware via the per-pair / per-k
timing probe (see ``docs/learnings/`` if archived). Setting
below 1 is allowed but rarely useful (the formula is
symmetric).
Returns
-------
MultipoleEwaldParameters
``(alpha, sigma, real_space_cutoff, reciprocal_space_cutoff)``.
Raises
------
ValueError
If any system has ``eta_eff <= sigma * sqrt(2)`` — meaning the
cost-balanced Ewald split is degenerate at this size + sigma
combination. Note: the validity threshold scales as
``cost_ratio**(-1/6)`` — large ``cost_ratio`` makes the split
more likely to be invalid for small/dense systems.
"""
if cost_ratio <= 0.0:
raise ValueError(f"cost_ratio must be positive, got {cost_ratio}")
if cell.ndim == 2:
cell = cell.unsqueeze(0)
num_systems = cell.shape[0]
dtype = positions.dtype
device = positions.device
volume = torch.abs(torch.linalg.det(cell)).squeeze(-1)
num_atoms = _count_atoms_per_system(positions, num_systems, batch_idx).to(dtype)
eta_kp = _kp_eta(volume, num_atoms)
eta = eta_kp / (cost_ratio ** (1.0 / 6.0))
error_factor = math.sqrt(-2.0 * math.log(accuracy))
real_space_cutoff = error_factor * eta
reciprocal_space_cutoff = error_factor / eta
sigma_t = _prepare_sigma(sigma, num_systems, dtype, device)
# alpha = 1 / (sqrt(2) * sqrt(eta**2 - 2 sigma**2)). The validity
# check uses the cost-corrected eta — large cost_ratio shrinks eta,
# which can push borderline systems into the invalid regime.
discriminant = eta * eta - 2.0 * sigma_t * sigma_t
if torch.any(discriminant <= 0.0):
bad = (discriminant <= 0.0).nonzero(as_tuple=False).flatten().tolist()
raise ValueError(
"Multipole Ewald parameter estimation: GTO sigma is too large "
f"relative to the system size for systems {bad} at cost_ratio="
f"{cost_ratio}. The cost-balanced eta={eta.tolist()} satisfies "
f"eta**2 <= 2 sigma**2 (sigma={sigma_t.tolist()}). Either reduce "
"sigma, increase the system, drop cost_ratio toward 1.0, or use "
"direct k-space (multipole_electrostatic_energy)."
)
alpha = 1.0 / (math.sqrt(2.0) * torch.sqrt(discriminant))
return MultipoleEwaldParameters(
alpha=alpha,
sigma=sigma_t,
real_space_cutoff=real_space_cutoff,
reciprocal_space_cutoff=reciprocal_space_cutoff,
)
[docs]
def estimate_multipole_pme_parameters(
positions: torch.Tensor,
cell: torch.Tensor,
sigma: float | torch.Tensor,
batch_idx: torch.Tensor | None = None,
accuracy: float = 1e-6,
cost_ratio: float = 1.0,
*,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
fft_padding_fraction: float = _DEFAULT_PME_FFT_PADDING_FRACTION,
) -> MultipolePMEParameters:
"""Estimate GTO-Ewald multipole PME parameters at a given target accuracy.
Same Kolafa-Perram backbone as
:func:`estimate_multipole_ewald_parameters`. Mesh dimensions follow
the standard B-spline-error formula
``n_per_dim = 2 alpha_eff L / (3 accuracy**0.2)`` with
``alpha_eff = 1 / (sqrt(2) eta)`` — the monopole-equivalent alpha at
the same eta.
The ``cost_ratio`` knob has the same meaning as in
:func:`estimate_multipole_ewald_parameters` (per-pair vs per-k cost
asymmetry), and shifts the rcut/kcut/mesh balance the same way:
``eta_eff = eta_KP / cost_ratio**(1/6)``. A larger ``cost_ratio``
grows the FFT mesh and shrinks the real-space cutoff. Note that
PME's true reciprocal cost is FFT (``M log M``) plus
spread/gather (``N p**3``), which is not the same shape as the
Ewald per-k-vector cost — so the optimal ``cost_ratio`` for PME may
differ from the Ewald optimum even on the same hardware. Default
``1.0`` (canonical KP) is a safe starting point.
Parameters
----------
positions : torch.Tensor, shape (N, 3) or (N_total, 3)
Atomic coordinates.
cell : torch.Tensor, shape (3, 3) or (B, 3, 3)
Unit cell matrix (matches the multipole-PME convention).
sigma : float or torch.Tensor
GTO basis width — same value used by the multipole-PME kernel.
Scalar or shape ``(B,)``.
batch_idx : torch.Tensor, shape (N_total,), int32, optional
System index per atom. ``None`` selects single-system mode.
accuracy : float, default 1e-6
Target relative accuracy (matches monopole convention). Also
drives mesh resolution via the B-spline error envelope.
cost_ratio : float, default 1.0
Empirical ``C_r / C_k`` ratio — per-pair real-space cost divided
by reciprocal-space cost on the target hardware. ``1.0`` (default)
reproduces canonical Kolafa-Perram. Higher values shift the
optimum toward smaller ``real_space_cutoff`` (fewer pairs) and a
finer FFT mesh (larger ``mesh_dimensions``). For PME the true
reciprocal cost is ``O(M log M)`` plus spread/gather, not the
per-k-vector cost of direct Ewald summation, so the optimal
``cost_ratio`` may differ from
:func:`estimate_multipole_ewald_parameters` even on the same
hardware. Setting below 1 is allowed but rarely useful.
spline_order : int, default=5
B-spline assignment order. The mesh covers its assignment support.
fft_padding_fraction : float, default=0.25
Maximum fraction of extra total mesh points for snapping to powers
of two, then smooth multiples of four. Zero selects the smallest
2/3/5/7-smooth mesh. Explicit other orders keep dyadic accuracy sizing.
Returns
-------
MultipolePMEParameters
``(alpha, sigma, mesh_dimensions, mesh_spacing, real_space_cutoff)``.
"""
if cost_ratio <= 0.0:
raise ValueError(f"cost_ratio must be positive, got {cost_ratio}")
if cell.ndim == 2:
cell = cell.unsqueeze(0)
num_systems = cell.shape[0]
dtype = positions.dtype
device = positions.device
volume = torch.abs(torch.linalg.det(cell)).squeeze(-1)
num_atoms = _count_atoms_per_system(positions, num_systems, batch_idx).to(dtype)
eta_kp = _kp_eta(volume, num_atoms)
eta = eta_kp / (cost_ratio ** (1.0 / 6.0))
error_factor = math.sqrt(-2.0 * math.log(accuracy))
real_space_cutoff = error_factor * eta
sigma_t = _prepare_sigma(sigma, num_systems, dtype, device)
discriminant = eta * eta - 2.0 * sigma_t * sigma_t
if torch.any(discriminant <= 0.0):
bad = (discriminant <= 0.0).nonzero(as_tuple=False).flatten().tolist()
raise ValueError(
"Multipole PME parameter estimation: GTO sigma too large for "
f"systems {bad} at cost_ratio={cost_ratio} (eta**2 <= 2 sigma**2). "
"Reduce sigma, drop cost_ratio toward 1.0, or use direct k-space."
)
alpha = 1.0 / (math.sqrt(2.0) * torch.sqrt(discriminant))
# Mesh resolution from the reciprocal-space B-spline error envelope.
# Width is set by eta; alpha_eff = 1/(sqrt(2) eta) (monopole-equivalent).
alpha_eff = torch.full_like(eta, 1.0 / math.sqrt(2.0)) / eta
mesh_dims = estimate_pme_mesh_dimensions(
cell,
alpha_eff,
accuracy,
spline_order=spline_order,
fft_padding_fraction=fft_padding_fraction,
)
cell_lengths = torch.norm(cell, dim=2) # (B, 3)
mesh_dims_tensor = torch.tensor(mesh_dims, dtype=cell_lengths.dtype, device=device)
mesh_spacing = cell_lengths / mesh_dims_tensor
return MultipolePMEParameters(
alpha=alpha,
sigma=sigma_t,
mesh_dimensions=mesh_dims,
mesh_spacing=mesh_spacing,
real_space_cutoff=real_space_cutoff,
)
[docs]
def mesh_spacing_to_dimensions(
cell: torch.Tensor,
mesh_spacing: float | torch.Tensor,
*,
spline_order: int = _DEFAULT_PME_SPLINE_ORDER,
fft_padding_fraction: float = _DEFAULT_PME_FFT_PADDING_FRACTION,
) -> tuple[int, int, int]:
"""Convert mesh spacing to mesh dimensions.
Parameters
----------
cell : torch.Tensor
Unit cell matrix.
mesh_spacing : int, float, or torch.Tensor
Target mesh spacing. Scalars apply to every cell and direction.
Tensor shapes ``(B,)`` and ``(B, 3)`` specify per-system isotropic
and per-system anisotropic spacings, respectively.
spline_order : int, default=5
B-spline assignment order. Each axis covers its assignment support.
fft_padding_fraction : float, default=0.25
Maximum fraction of extra total mesh points for snapping to powers
of two, then smooth multiples of four. Zero selects the smallest
2/3/5/7-smooth mesh. These spacing preferences apply to every spline order.
Returns
-------
tuple[int, int, int]
Maximum dimensions across the batch, rounded upward to 2/3/5/7-smooth
integers that cover the requested spacing and assignment support.
"""
if cell.ndim == 2:
cell = cell.unsqueeze(0)
cell_lengths = torch.norm(cell, dim=2) # (B, 3)
if isinstance(mesh_spacing, (int, float)) or (
isinstance(mesh_spacing, torch.Tensor) and mesh_spacing.ndim == 0
):
mesh_dims = torch.ceil(cell_lengths / mesh_spacing)
elif isinstance(mesh_spacing, torch.Tensor):
if mesh_spacing.ndim == 1:
if mesh_spacing.shape[0] != cell.shape[0]:
raise ValueError(
f"mesh_spacing shape {mesh_spacing.shape} incompatible with "
f"cell batch size {cell.shape[0]}"
)
mesh_dims = torch.ceil(cell_lengths / mesh_spacing[:, None])
else:
if mesh_spacing.shape != cell_lengths.shape:
raise ValueError(
f"mesh_spacing shape {mesh_spacing.shape} incompatible with "
f"cell_lengths shape {cell_lengths.shape}"
)
mesh_dims = torch.ceil(cell_lengths / mesh_spacing)
else:
raise TypeError(
f"mesh_spacing must be int, float, or torch.Tensor, got {type(mesh_spacing)}"
)
max_mesh_dims = torch.max(mesh_dims, dim=0).values
return _round_pme_mesh_dimensions(
max_mesh_dims.detach().cpu().numpy(),
spline_order=spline_order,
fft_padding_fraction=fft_padding_fraction,
)
