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, max_neighbors=256)

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).

• max_neighbors (int, optional) – Maximum neighbors per atom for the dense neighbor matrix. Defaults to 256.

model_config

Mutable configuration controlling which outputs are computed. Set model.model_config.compute_stresses = True to enable virial computation for NPT/NPH simulations.

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.

This method should extract meaningful representations from the model at node (atomic), edge (bond), and/or graph/system (structure) levels. The concrete implementation should check if the model supports computing embeddings, as well as perform validation on kwargs to make sure they are valid for the model.

The method should add graph, node, and/or edge embeddings to the Batch data structure in-place.

Parameters:

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

• kwargs (Any)

Returns:

Standardized AtomicData or Batch data structure mutated in place.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – If the model does not support embeddings computation

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

forward(data, **kwargs)

Run the PME kernel and return a ModelOutputs dict.

Parameters:

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

• kwargs (Any)

Returns:

OrderedDict with keys "energies" (shape [B, 1], eV), "forces" (shape [N, 3], eV/Å), and optionally "stress" (shape [B, 3, 3], eV — the raw virial \(W_{phys}\)).

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

property model_card: ModelCard

Retrieves the model card for the model.

The model card is a Pydantic model that contains information about the model’s capabilities and requirements.

output_data()

Return the set of keys that the model produces.

Return type:

set[str]