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=True)

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 eV·Å/e². Defaults to 14.3996 (standard value for Å/e/eV unit system).

• hybrid_forces (bool)

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 required inputs from data without enabling gradients.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

dict[str, Any]

adapt_output(model_output, data)

Adapt the model output to the framework output format.

Parameters:

• model_output (Any)

• data (AtomicData | Batch)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

compute_embeddings(data, **kwargs)

Compute embeddings at different levels of a batch of atomic graphs.

Parameters:

• data (AtomicData | Batch) – Input atomic data containing positions, atomic numbers, etc.

• kwargs (Any)

Returns:

Data structure with embeddings attached in-place.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – If the model does not support embeddings computation

direct_derivative_keys()

Analytical force/stress keys when hybrid_forces=True.

Return type:

set[str]

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)

Export model is not implemented for PME models.

Parameters:

• path (Path)

• as_state_dict (bool)

Return type:

None

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], eV/ų — Cauchy stress W/V).

Return type:

ModelOutputs

input_data()

Return required input keys (override to drop atomic_numbers).

Return type:

set[str]

invalidate_cache()

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

Return type:

None

output_data()

Return the set of keys that the model produces.

Return type:

set[str]