nvalchemi.models.ewald.EwaldModelWrapper

class nvalchemi.models.ewald.EwaldModelWrapper(cutoff, accuracy=1e-6, coulomb_constant=14.3996, hybrid_forces=False, slab_correction=False, rtol=1e-5, atol=None)

Ewald summation electrostatics potential as a model wrapper.

Computes long-range Coulomb interactions via the Ewald method, splitting contributions into real-space (erfc-damped, handled by a neighbor matrix) and reciprocal-space (structure factor summation) components.

Parameters:

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

• accuracy (float, optional) – Target accuracy for automatic Ewald 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.

  Note

  Under domain decomposition the Ewald wrapper runs a split real-space / reciprocal-space path; the slab correction is only wired through the single-call kernel and is therefore inactive on the distributed reciprocal path.

• 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), the Warp kernel computes forces analytically 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. autograd_outputs includes "forces" so the pipeline accumulates direct kernel forces with charge-path autograd forces in hybrid mode. Add "stress" to active_outputs to enable virial computation for NPT/NPH. When charges.requires_grad=True, energy.backward() flows through the injected \(dE/dq\) pathway while forces and virial/stress are returned detached.

Type:

ModelConfig

adapt_input(data, **kwargs)

Build the input dict for the Ewald kernels.

Collects the keys named by input_data(), casts the topology tensors to int32, and attaches cell and optional neighbor_matrix_shifts. Gradients are not enabled here; forces and stress come from the kernel analytically (hybrid) or from energy autograd downstream.

Parameters:

• data (AtomicData | Batch) – The input system; must be a Batch (Ewald has no single-graph path).

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

Returns:

Kernel inputs: positions [N, 3], charges [N], neighbor_matrix, num_neighbors, batch_idx [N], ptr, num_graphs, fill_value, cell [B, 3, 3], and neighbor_matrix_shifts (or None).

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.cell is absent (PBC is required).

adapt_output(model_output, data)

Map raw kernel outputs to nvalchemi standard keys.

Always forwards energy; adds forces and stress when they are in model_config.active_outputs.

Parameters:

• model_output (Any) – The dict produced by forward() (energy [B, 1], and forces / stress when active).

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

Returns:

Ordered dict with energy and, when active, forces [N, 3] and stress [B, 3, 3].

Return type:

ModelOutputs

Raises:

RuntimeError – If stress is active but missing from model_output.

compute_embeddings(data, **kwargs)

Embeddings are not defined for an electrostatics model.

Parameters:

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

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

Returns:

Never returns.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – Always; Ewald has no learned embeddings.

direct_derivative_keys()

Return the outputs the kernel computes analytically.

With hybrid_forces=True the Warp kernel returns analytical forces and stress directly, so the pipeline must not also derive them from energy autograd.

Returns:

{"forces", "stress"} (intersected with the active outputs) when hybrid_forces is set, otherwise an empty set.

Return type:

set[str]

distributed_setup(ctx)

Enter distributed mode: stash the context and global atom count.

Caches the global atom count so _update_cache() estimates Ewald parameters from the global N rather than the per-rank padded count, then invalidates the cache so any params derived from a stale N are rebuilt on the next forward.

Parameters:

ctx (DistributedContext) – The live distributed context for this run.

Return type:

None

distributed_teardown()

Exit distributed mode: clear the context and global atom count.

Return type:

None

distribution_spec(strategy=None)

MLIPSpec for Ewald electrostatics under domain decomposition.

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

The real-space pair kernel runs on halo-padded inputs. Reciprocal-space is declared via the two structure-factor ops (single-system + batched) in custom_ops: each takes owned-slice inputs (every atom contributes once globally) and all-reduces its three partial outputs across ranks before the energy stage consumes them.

Returns:

The Ewald halo spec with structure-factor custom_ops and per-output reduction kinds. forces is owned-only because reciprocal forces come from the all-reduced S(k) and so are exact on every owned row.

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)

Serialization is not supported for the Ewald model.

Parameters:

• path (Path) – Output path (unused).

• as_state_dict (bool, optional) – Unused. Defaults to False.

Returns:

Never returns.

Return type:

None

Raises:

NotImplementedError – Always; Ewald has no checkpointable state.

extra_repr()

Format the model config for nn.Module.__repr__.

Parameters:

self (Any)

Return type:

str

forward(data, **kwargs)

Run the Ewald summation and return a ModelOutputs dict.

Parameters:

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

• **kwargs – Forwarded to adapt_input().

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\) — tensile-positive Cauchy stress -W/V).

Return type:

ModelOutputs

input_data()

Return the batch keys adapt_input() reads.

Overrides the base set to require charges and the matrix-format neighbor data, and to drop atomic_numbers (unused by Ewald). When slab_correction is enabled, pbc is also required.

Returns:

{"positions", "charges", "neighbor_matrix", "num_neighbors"}, plus "pbc" when slab_correction=True.

Return type:

set[str]

invalidate_cache()

Force recomputation of Ewald parameters and k-vectors.

Call this after modifying the unit cell (e.g. an NPT integrator) so the next forward rebuilds alpha and the k-vectors.

Return type:

None

output_data()

Return the output keys this model produces for the active config.

Returns:

{"energy"} plus "forces" and/or "stress" when those are in model_config.active_outputs.

Return type:

set[str]