PMEModel

class torch_sim.models.electrostatics.PMEModel(cutoff=10.0, *, accuracy=1e-6, mesh_spacing=None, mesh_dimensions=None, spline_order=4, device=None, dtype=torch.float64, compute_forces=True, compute_stress=True, neighbor_list_fn=torchsim_nl)[source]

Bases: ModelInterface

Particle Mesh Ewald electrostatics as a ModelInterface.

Uses the nvalchemiops PME kernel for O(N log N) periodic electrostatics. Returns per-atom energies that are aggregated to per-system. All user-facing quantities are in metal units (Angstrom / eV).

Per-atom partial charges are read from state.partial_charges.

Requires periodic boundary conditions.

Parameters:

  • cutoff (float) – Real-space cutoff in Angstrom.

  • accuracy (float) – Target accuracy for auto-estimated parameters.

  • mesh_spacing (float | None) – Optional mesh spacing (Angstrom) for automatic mesh sizing.

  • mesh_dimensions (tuple[int, int, int] | None) – Explicit FFT mesh dimensions (nx, ny, nz).

  • spline_order (int) – B-spline interpolation order. Defaults to 4.

  • device (device | None) – Compute device. Defaults to CUDA if available, else CPU.

  • dtype (dtype) – Floating-point dtype. Defaults to torch.float64.

  • compute_forces (bool) – Whether to return forces. Defaults to True.

  • compute_stress (bool) – Whether to return stress. Defaults to True.

  • neighbor_list_fn (Callable) – Neighbor-list constructor. Defaults to torchsim_nl.

forward(state, **_kwargs)[source]

Compute PME electrostatic energy, forces, and stress.

Parameters:

  • state (SimState) – Simulation state with partial_charges set as an atom extra (shape [n_atoms]). Returns zeros for non-periodic states.

  • **_kwargs (object) – Unused; accepted for interface compatibility.

Returns:

dict with "energy" [n_systems], "forces" [n_atoms, 3], and (if compute_stress) "stress" [n_systems, 3, 3].

Return type:

dict[str, Tensor]