nvalchemi.models.pme.PMEModelWrapper

class nvalchemi.models.pme.PMEModelWrapper(cutoff, mesh_spacing=1.0, mesh_dimensions=None, spline_order=4, alpha=None, accuracy=1e-6, coulomb_constant=14.3996, hybrid_forces=False, slab_correction=False, rtol=1e-5, atol=None)

Particle Mesh Ewald electrostatics potential as a model wrapper.

Computes long-range Coulomb interactions via the PME method using B-spline charge interpolation and FFT-based reciprocal space evaluation, achieving \(O(N \log N)\) computational scaling.

Parameters:

• cutoff (float) – Real-space interaction cutoff in Å.

• mesh_spacing (float, optional) – Target mesh spacing in Å used to determine mesh dimensions when mesh_dimensions is not provided. Defaults to 1.0 Å.

• mesh_dimensions (tuple[int, int, int] or None, optional) – Explicit mesh dimensions (nx, ny, nz). When set, overrides automatic estimation from mesh_spacing. Defaults to None (auto-estimated).

• spline_order (int, optional) – B-spline interpolation order. Higher values give greater accuracy at increased cost. Defaults to 4.

• alpha (float or None, optional) – Ewald splitting parameter (inverse Å). None causes automatic estimation via the Kolafa-Perram formula each time the cell changes. Defaults to None.

• accuracy (float, optional) – Target accuracy for automatic parameter estimation. Defaults to 1e-6.

• coulomb_constant (float, optional) – Coulomb prefactor \(k_e\) in \(\mathrm{eV}\cdot\mathrm{\AA}/e^2\). Defaults to 14.3996 (standard value for Å/e/eV unit system).

• slab_correction (bool, optional) – Whether to enable the two-dimensional slab correction. Defaults to False. When enabled, the input batch must provide data.pbc as a boolean tensor with shape (B, 3). Rows with exactly one False entry mark slab systems, for example [True, True, False] for a non-periodic z axis. Fully periodic rows are no-ops, so mixed slab and three-dimensional periodic batches are supported.

• rtol (float, optional) – Relative tolerance for cell change detection. See cell_cache_needs_update().

• atol (float or None, optional) – Absolute tolerance for cell change detection. See cell_cache_needs_update().

• hybrid_forces (bool, optional) – When True (default), direct kernel forces (dE/dR|_q) are used and forces is kept in autograd_outputs only to add the charge chain-rule term; when False, forces come entirely from autograd.

model_config

Mutable configuration controlling which outputs are computed. model_config.autograd_outputs includes "forces" so the pipeline accumulates direct kernel forces with charge-path autograd forces in hybrid mode. Include "stress" in model_config.active_outputs to enable virial computation for NPT/NPH simulations. When charges.requires_grad=True, energy.backward() propagates through the injected \(dE/dq\) pathway while the wrapper returns detached direct kernel forces and detached virial/stress.

Type:

ModelConfig

adapt_input(data, **kwargs)

Collect the kernel inputs from data without enabling gradients.

Gathers the required batch attributes, batch indexing tensors, the PBC cell, and optional neighbor shifts into a plain dict. Gradients are not enabled here: forces and stress are produced analytically by the kernel (or, in the charge-dependent pipeline, via autograd on the energy).

Parameters:

• data (Batch) – Batch with positions, charges, cell, neighbor_matrix, and num_neighbors.

• **kwargs – Unused; accepted for interface compatibility.

Returns:

Kernel inputs including positions [N, 3], charges [N], cell [B, 3, 3], batch_idx [N], ptr, num_graphs, fill_value, the neighbor matrix, and neighbor_matrix_shifts (None when non-periodic).

Return type:

dict[str, Any]

Raises:

• TypeError – If data is an AtomicData rather than a Batch.

• KeyError – If a required input key is missing from data.

• ValueError – If data has no cell (PME requires PBC).

adapt_output(model_output, data)

Select the active outputs into the standard output mapping.

Always forwards energy; adds forces and stress when each is in model_config.active_outputs.

Parameters:

• model_output (dict[str, Any]) – Raw kernel outputs keyed by "energy", "forces", and "stress".

• data (AtomicData | Batch) – The input system the outputs were computed for (unused).

Returns:

OrderedDict with "energy" and any active "forces" / "stress".

Return type:

ModelOutputs

Raises:

RuntimeError – If "stress" is active but absent from model_output.

compute_embeddings(data, **kwargs)

Embeddings are not defined for a PME electrostatics model.

Parameters:

• data (AtomicData | Batch) – The input system (unused).

• **kwargs – Unused; accepted for interface compatibility.

Returns:

Never returned.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – Always; PME produces no learned embeddings.

direct_derivative_keys()

Report which outputs are computed analytically by the kernel.

Returns:

{"forces", "stress"} (intersected with the active outputs) when hybrid_forces=True; an empty set otherwise, in which case forces/stress come from autograd on the energy.

Return type:

set[str]

distributed_setup(ctx)

Enter distributed mode for this wrapper.

Records the distributed context and global atom count, then invalidates the cache so alpha / mesh are re-estimated from the global N rather than a stale per-rank count.

Parameters:

ctx (DistributedContext) – The live distributed context, exposing n_atoms_total and the halo metadata.

Return type:

None

distributed_teardown()

Leave distributed mode and return to single-GPU behaviour.

Clears the distributed context and global atom count and invalidates the cache.

Return type:

None

distribution_spec(strategy=None)

Domain-decomposition spec for the PME wrapper.

Halo-only; the strategy argument is accepted for the framework contract and ignored.

Four ops get owned-slice + all-reduce handlers so the reciprocal-space pathway sees globally-correct quantities: the spline-spread ops all-reduce each rank’s partial charge mesh into a replicated global mesh, and the total-charge ops all-reduce each rank’s partial charge sum into the true global total charge used by the background correction. Every downstream stage (FFT, Green’s function, IFFT, gather, per-atom corrections) then runs identically on every rank, so forward() is distribution-agnostic.

Returns:

Halo-storage spec carrying the spline-spread and total-charge custom_ops, plus output handling: energy and stress per-graph, forces per-node owned-only, atomic_energies per-node.

Return type:

MLIPSpec

Parameters:

strategy (Any)

property embedding_shapes: dict[str, tuple[int, ...]]

Retrieves the expected shapes of the node, edge, and graph embeddings.

export_model(path, as_state_dict=False)

Serialize the model (not supported for the PME wrapper).

Parameters:

• path (Path) – Intended output path (unused).

• as_state_dict (bool, optional) – Whether to save only the state_dict (unused). Defaults to False.

Return type:

None

Raises:

NotImplementedError – Always; the PME wrapper holds no trainable weights to export.

extra_repr()

Format the model config for nn.Module.__repr__.

Parameters:

self (Any)

Return type:

str

forward(data, **kwargs)

Run the PME kernel and return a ModelOutputs dict.

Parameters:

• data (Batch) – Batch containing positions, charges, cell, neighbor_matrix, and num_neighbors (populated by NeighborListHook).

• kwargs (Any)

Returns:

OrderedDict with keys "energy" (shape [B, 1], eV), "forces" (shape [N, 3], eV/Å), and optionally "stress" (shape [B, 3, 3], \(\mathrm{eV}/\mathrm{\AA}^3\) — Cauchy stress -W/V).

Return type:

ModelOutputs

input_data()

List the batch attributes the PME forward reads.

Returns:

{"positions", "charges", "neighbor_matrix", "num_neighbors"}, plus "pbc" when slab_correction=True. Notably excludes atomic_numbers, which PME does not use.

Return type:

set[str]

invalidate_cache()

Force recomputation of PME parameters, k-vectors, and mesh.

Return type:

None

output_data()

List the output keys the forward currently produces.

Returns:

{"energy"} plus "forces" and/or "stress" when each is in model_config.active_outputs.

Return type:

set[str]