pme

Standalone particle-mesh Ewald electrostatics for small periodic fixtures.

import mlx_atomistic.pme

Classes

PMEConfig

class PMEConfig
    def __init__(mesh_shape: tuple[int, int, int] = (32, 32, 32), alpha: float = 0.35, real_cutoff: float | None = None, assignment_order: int = 2, charge_tolerance: float = 1e-05, deconvolve_assignment: bool = True, background_policy: PMEBackgroundPolicy = 'reject_non_neutral')

Controls for the standalone PME mesh backend.

Parameters

Name	Type	Default	Description
mesh_shape	tuple[int, int, int]	(32, 32, 32)
alpha	float	0.35
real_cutoff	float | None	None
assignment_order	int	2
charge_tolerance	float	1e-05
deconvolve_assignment	bool	True
background_policy	PMEBackgroundPolicy	'reject_non_neutral'

PMEDiagnostics

class PMEDiagnostics
    def __init__(mesh_shape: tuple[int, int, int], assignment_order: int, alpha: float, real_cutoff: float, net_charge: float, volume: float, charge_grid_sum: float, reciprocal_modes: int, max_charge_grid_abs: float, background_policy: PMEBackgroundPolicy = 'reject_non_neutral', background_energy: float = 0.0, direct_space_policy: str = 'dense', direct_space_representation: str = 'dense', direct_space_pair_count: int | None = None, direct_space_candidate_count: int | None = None, direct_space_fallback_reason: str | None = None, plan_fingerprint: str | None = None, plan_build_count: int = 0, plan_reuse_count: int = 0, plan_setup_seconds: float = 0.0, plan_estimated_resident_bytes: int = 0, plan_backend: str | None = None, plan_device: str | None = None, plan_dtype: str | None = None)

Diagnostics emitted by one standalone PME evaluation.

Parameters

Name	Type	Default	Description
mesh_shape	tuple[int, int, int]
assignment_order	int
alpha	float
real_cutoff	float
net_charge	float
volume	float
charge_grid_sum	float
reciprocal_modes	int
max_charge_grid_abs	float
background_policy	PMEBackgroundPolicy	'reject_non_neutral'
background_energy	float	0.0
direct_space_policy	str	'dense'
direct_space_representation	str	'dense'
direct_space_pair_count	int | None	None
direct_space_candidate_count	int | None	None
direct_space_fallback_reason	str | None	None
plan_fingerprint	str | None	None
plan_build_count	int	0
plan_reuse_count	int	0
plan_setup_seconds	float	0.0
plan_estimated_resident_bytes	int	0
plan_backend	str | None	None
plan_device	str | None	None
plan_dtype	str | None	None

Methods

to_dict

def to_dict() -> dict[str, object]

Return the diagnostics as a plain JSON-serializable dict.

Returns

• dict[str, object] — The diagnostics fields (mesh shape, alpha, cutoff, net charge, grid sums, direct-space policy, …) as a dict.

PMEDirectSpacePolicyReport

class PMEDirectSpacePolicyReport
    def __init__(policy: str, representation: str, uses_shared_neighbor_policy: bool, supported: bool, real_cutoff: float, minimum_image_safe: bool, pair_count: int | None = None, compact_pair_count: int | None = None, candidate_count: int | None = None, candidate_waste_count: int | None = None, fallback_reason: str | None = None)

Execution policy selected for PME real/direct-space evaluation.

Parameters

Name	Type	Default	Description
policy	str
representation	str
uses_shared_neighbor_policy	bool
supported	bool
real_cutoff	float
minimum_image_safe	bool
pair_count	int | None	None
compact_pair_count	int | None	None
candidate_count	int | None	None
candidate_waste_count	int | None	None
fallback_reason	str | None	None

Methods

to_dict

def to_dict() -> dict[str, object]

Return the policy report as a plain JSON-serializable dict.

Returns

• dict[str, object] — The policy report fields (policy, representation, support, cutoff, pair counts, fallback reason) as a dict.

PMEExecutionPlan

class PMEExecutionPlan
    def __init__(cell: Cell, *, config: PMEConfig | None = None, coulomb_constant: float = 1.0, dtype: object = mx.float32, backend: str = PME_EXECUTION_BACKEND)

Reusable cell-invariant reciprocal state for fixed-cell PME evaluation.

Parameters

Name	Type	Default	Description
cell	Cell		Periodic orthorhombic cell whose reciprocal state is planned.
config	PMEConfig | None	None	PME parameters; None uses PMEConfig defaults.
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system.
dtype	object	mx.float32	Reciprocal-array dtype. The current backend supports float32.
backend	str	PME_EXECUTION_BACKEND	PME execution backend identifier. The current backend is PME_EXECUTION_BACKEND.

Properties

• diagnostics dict[str, object] — Return JSON-serializable setup and reuse diagnostics for this plan.

Methods

build

def build(cell: Cell, *, config: PMEConfig | None = None, coulomb_constant: float = 1.0, dtype: object = mx.float32, backend: str = PME_EXECUTION_BACKEND) -> PMEExecutionPlan

Build a reusable PME execution plan.

Parameters

Name	Type	Default	Description
cell	Cell		Periodic orthorhombic cell whose reciprocal state is planned.
config	PMEConfig | None	None	PME parameters; None uses PMEConfig defaults.
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system.
dtype	object	mx.float32	Reciprocal-array dtype. The current backend supports float32.
backend	str	PME_EXECUTION_BACKEND	PME execution backend identifier.

Returns

• PMEExecutionPlan — A materialized reusable execution plan.

rebuild

def rebuild(*, cell: Cell | None = None, config: PMEConfig | None = None, coulomb_constant: float | None = None, dtype: object | None = None, backend: str | None = None) -> PMEExecutionPlan

Build a distinct plan after an explicit defining-input change.

Parameters

Name	Type	Default	Description
cell	Cell | None	None	Replacement cell; None retains the current cell.
config	PMEConfig | None	None	Replacement configuration; None retains the current config.
coulomb_constant	float | None	None	Replacement Coulomb prefactor; None retains the current value.
dtype	object | None	None	Replacement reciprocal dtype; None retains the current dtype.
backend	str | None	None	Replacement execution backend; None retains the current backend.

Returns

• PMEExecutionPlan — A new materialized execution plan with counters independent of this plan.

to_dict

def to_dict() -> dict[str, object]

Return the plan diagnostics as a plain JSON-serializable dict.

Returns

• dict[str, object] — Setup, fingerprint, resident-size, backend, and reuse metadata.

validate

def validate(cell: Cell, *, config: PMEConfig | None = None, coulomb_constant: float | None = None, dtype: object = mx.float32, backend: str = PME_EXECUTION_BACKEND) -> None

Validate that requested PME inputs exactly match this plan.

Parameters

Name	Type	Default	Description
cell	Cell		Cell requested by the evaluation.
config	PMEConfig | None	None	Requested PME configuration; None uses this plan’s config.
coulomb_constant	float | None	None	Requested Coulomb prefactor; None uses this plan’s value.
dtype	object	mx.float32	Requested reciprocal-array dtype.
backend	str	PME_EXECUTION_BACKEND	Requested PME execution backend.

Returns

• None

Raises

• PMEPlanMismatchError — If any plan-defining input differs.
• ValueError — If an input is invalid independently of the stored plan.

PMEPlanMismatchError

class PMEPlanMismatchError(ValueError)
    def __init__(mismatches: tuple[str, ...], *, expected_fingerprint: str, actual_fingerprint: str)

Raised when a PME execution plan does not match an evaluation request.

Parameters

Name	Type	Default	Description
mismatches	tuple[str, ...]
expected_fingerprint	str
actual_fingerprint	str

Functions

assign_charges_bspline

def assign_charges_bspline(positions: mx.array, charges: mx.array, cell: Cell, mesh_shape: tuple[int, int, int], *, assignment_order: int = 2) -> mx.array

Assign charges to a periodic mesh with cardinal B-spline weights.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom charges, shape (n_atoms,).
cell	Cell		Periodic orthorhombic cell.
mesh_shape	tuple[int, int, int]		Reciprocal-space mesh dimensions (nx, ny, nz).
assignment_order	int	2	Cardinal B-spline order (2 = cloud-in-cell). Defaults to 2.

Returns

• mx.array — The charge density on the mesh, shape mesh_shape.

Raises

• ValueError — If the cell is non-orthorhombic or the shapes are invalid.

assign_charges_cic

def assign_charges_cic(positions: mx.array, charges: mx.array, cell: Cell, mesh_shape: tuple[int, int, int]) -> mx.array

Assign charges to a periodic mesh with cloud-in-cell weights.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom charges, shape (n_atoms,).
cell	Cell		Periodic orthorhombic cell.
mesh_shape	tuple[int, int, int]		Reciprocal-space mesh dimensions (nx, ny, nz).

Returns

• mx.array — The charge density on the mesh, shape mesh_shape.

Raises

• ValueError — If the cell is non-orthorhombic or the shapes are invalid.

normalize_pme_background_policy

def normalize_pme_background_policy(value: object) -> PMEBackgroundPolicy

Return a supported PME background-charge policy.

Parameters

Name	Type	Default	Description
value	object		Policy name to normalize.

Returns

• PMEBackgroundPolicy — A supported normalized policy string.

Raises

• ValueError — If value is not a supported policy.

pme_coulomb_direct_space_energy_forces

def pme_coulomb_direct_space_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, pairs: mx.array | NeighborBlocks | None = None, plan: PMEExecutionPlan | None = None) -> tuple[mx.array, mx.array]

Evaluate only the PME real/direct-space Coulomb energy and forces.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,); charged systems require the explicit uniform-plasma background policy.
cell	Cell		Periodic orthorhombic Cell (triclinic is unsupported).
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system. Defaults to 1.0.
config	PMEConfig | None	None	PME parameters; None uses defaults. Defaults to None.
pairs	mx.array | NeighborBlocks | None	None	Optional real-space pairs or neighbor blocks. Defaults to None.
plan	PMEExecutionPlan | None	None	Optional reusable fixed-cell execution plan. A one-shot plan is built when omitted. Defaults to None.

Returns

• tuple[mx.array, mx.array] — An (energy, forces) tuple: scalar real-space energy and (n_atoms, 3) forces.

Raises

• ValueError — If shapes/cell or the configured charge policy are invalid.

pme_coulomb_energy_forces

def pme_coulomb_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, direct_space_pairs: mx.array | NeighborBlocks | None = None, plan: PMEExecutionPlan | None = None) -> tuple[mx.array, mx.array, dict[str, mx.array | PMEDiagnostics]]

Evaluate neutral orthorhombic Coulomb energy and forces with PME.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,); charged systems require the explicit uniform-plasma background policy.
cell	Cell		Periodic orthorhombic Cell (triclinic is unsupported).
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system. Defaults to 1.0.
config	PMEConfig | None	None	PME parameters; None uses defaults. Defaults to None.
direct_space_pairs	mx.array | NeighborBlocks | None	None	Optional precomputed real-space pairs or neighbor blocks. Defaults to None.
plan	PMEExecutionPlan | None	None	Optional reusable fixed-cell execution plan. A one-shot plan is built when omitted. Defaults to None.

Returns

• tuple[mx.array, mx.array, dict[str, mx.array | PMEDiagnostics]] — An (energy, forces, components) tuple: scalar total Coulomb energy, (n_atoms, 3) forces, and a components dict including a PMEDiagnostics entry.

Raises

• ValueError — If shapes are wrong, the cell is non-orthorhombic, or charge is incompatible with the configured background policy.

pme_coulomb_reciprocal_space_energy_forces

def pme_coulomb_reciprocal_space_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, include_self_correction: bool = True, plan: PMEExecutionPlan | None = None) -> tuple[mx.array, mx.array]

Evaluate PME mesh/reciprocal-space Coulomb forces.

By default the scalar self correction is included so direct-space plus reciprocal-space equals the total PME Coulomb energy.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,); charged systems require the explicit uniform-plasma background policy.
cell	Cell		Periodic orthorhombic Cell (triclinic is unsupported).
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system. Defaults to 1.0.
config	PMEConfig | None	None	PME parameters; None uses defaults. Defaults to None.
include_self_correction	bool	True	Whether to add the scalar self-energy correction so direct + reciprocal equals the total. Defaults to True.
plan	PMEExecutionPlan | None	None	Optional reusable fixed-cell execution plan. A one-shot plan is built when omitted. Defaults to None.

Returns

• tuple[mx.array, mx.array] — An (energy, forces) tuple: scalar reciprocal-space energy and (n_atoms, 3) forces.

Raises

• ValueError — If shapes/cell or the configured charge policy are invalid.

pme_coulomb_strain_components

def pme_coulomb_strain_components(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, direct_space_pairs: mx.array | NeighborBlocks | None = None) -> dict[str, mx.array]

Return cell-differentiable PME components for analytic virials.

This path rebuilds the cell-dependent reciprocal arrays from MLX values so automatic differentiation sees the complete diagonal cell derivative. It is intentionally separate from the reusable fixed-cell execution plan.

Parameters

Name	Type	Default	Description
positions	mx.array		Molecularly strained coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,).
cell	Cell		Strained orthorhombic periodic cell.
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system.
config	PMEConfig | None	None	PME parameters with an explicit real-space cutoff.
direct_space_pairs	mx.array | NeighborBlocks | None	None	Compact pairs or neighbor blocks held fixed during the infinitesimal strain derivative.

Returns

• dict[str, mx.array] — Differentiable real, reciprocal, self, and background energies.

Raises

• ValueError — If the cell, shapes, cutoff, charge policy, or direct-space pair representation is outside the analytic-virial envelope.

pme_coulomb_strain_energy

def pme_coulomb_strain_energy(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, direct_space_pairs: mx.array | NeighborBlocks | None = None) -> mx.array

Return the complete cell-differentiable PME energy for analytic virials.

Parameters

Name	Type	Default	Description
positions	mx.array		Molecularly strained coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,).
cell	Cell		Strained orthorhombic periodic cell.
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system.
config	PMEConfig | None	None	PME parameters with an explicit real-space cutoff.
direct_space_pairs	mx.array | NeighborBlocks | None	None	Compact pairs or neighbor blocks held fixed during the infinitesimal strain derivative.

Returns

• mx.array — The sum of real, reciprocal, self, and background PME energies.

pme_coulomb_total_energy_forces

def pme_coulomb_total_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: PMEConfig | None = None, direct_space_pairs: mx.array | NeighborBlocks | None = None, plan: PMEExecutionPlan | None = None) -> tuple[mx.array, mx.array]

Evaluate neutral orthorhombic Coulomb total energy and forces with PME.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
charges	mx.array		Per-atom partial charges, shape (n_atoms,); charged systems require the explicit uniform-plasma background policy.
cell	Cell		Periodic orthorhombic Cell (triclinic is unsupported).
coulomb_constant	float	1.0	Coulomb prefactor in the configured unit system. Defaults to 1.0.
config	PMEConfig | None	None	PME parameters; None uses defaults. Defaults to None.
direct_space_pairs	mx.array | NeighborBlocks | None	None	Optional precomputed real-space pairs or neighbor blocks. Defaults to None.
plan	PMEExecutionPlan | None	None	Optional reusable fixed-cell execution plan. A one-shot plan is built when omitted. Defaults to None.

Returns

• tuple[mx.array, mx.array] — An (energy, forces) tuple: scalar total Coulomb energy and (n_atoms, 3) forces.

Raises

• ValueError — If shapes/cell or the configured charge policy are invalid.

pme_direct_space_policy_report

def pme_direct_space_policy_report(cell: Cell, *, config: PMEConfig | None = None, pairs: mx.array | NeighborBlocks | None = None, plan: PMEExecutionPlan | None = None) -> dict[str, object]

Report whether PME direct-space can use dense, pair, block, or fallback policy.

Parameters

Name	Type	Default	Description
cell	Cell		Periodic orthorhombic cell.
config	PMEConfig | None	None	PME parameters; None uses defaults. Defaults to None.
pairs	mx.array | NeighborBlocks | None	None	Optional candidate real-space pairs or neighbor blocks to evaluate the policy against. Defaults to None.
plan	PMEExecutionPlan | None	None	Optional plan whose configuration and fixed cell must match. The policy query does not advance its reuse counter. Defaults to None.

Returns

• dict[str, object] — A PMEDirectSpacePolicyReport dict (selected policy/representation, support, cutoff, and pair-count provenance).

Raises

• ValueError — If the cell is missing, non-orthorhombic, or has non-positive lengths.

pme_force_scope_report

def pme_force_scope_report(scope: str) -> dict[str, object]

Return PME support metadata for a requested force-evaluation scope.

Parameters

Name	Type	Default	Description
scope	str		Force-evaluation scope to query ("total", "components", "direct_space", "reciprocal_space").

Returns

• dict[str, object] — A ForceScopeReport dict; every PME scope is supported and marked as requiring full-system evaluation.

pme_platform_readiness_report

def pme_platform_readiness_report(*, atom_count: int, charges: object, cell_lengths: object, config: PMEConfig | None, nonbonded_cutoff: float | None, exclusion_count: int, one_four_count: int, explicit_exception_count: int, cell_matrix: object | None = None) -> ReadinessReport

Return PME readiness using the shared platform readiness schema.

Parameters

Name	Type	Default	Description
atom_count	int		Number of atoms in the system.
charges	object		Per-atom partial charges, shape (atom_count,).
cell_lengths	object		Orthorhombic box lengths, shape (3,).
config	PMEConfig | None		PME parameters; None is reported as a blocker.
nonbonded_cutoff	float | None		Real-space nonbonded cutoff, cross-checked against the PME cutoff.
exclusion_count	int		Number of excluded pairs.
one_four_count	int		Number of 1-4 corrected pairs.
explicit_exception_count	int		Number of explicit nonbonded exceptions.
cell_matrix	object | None	None	Optional full cell matrix for the orthorhombic-envelope check.

Returns

• ReadinessReport — A ReadinessReport (name "pme", status, blockers, and the remaining checks as metadata).

pme_readiness_report

def pme_readiness_report(*, atom_count: int, charges: object, cell_lengths: object, config: PMEConfig | None, nonbonded_cutoff: float | None, exclusion_count: int, one_four_count: int, explicit_exception_count: int, cell_matrix: object | None = None) -> dict[str, object]

Return fail-closed PME readiness metadata for production run gates.

Parameters

Name	Type	Default	Description
atom_count	int		Number of atoms in the system.
charges	object		Per-atom partial charges, shape (atom_count,).
cell_lengths	object		Orthorhombic box lengths, shape (3,).
config	PMEConfig | None		PME parameters (mesh, alpha, cutoffs); None is reported as a blocker.
nonbonded_cutoff	float | None		Real-space nonbonded cutoff, cross-checked against the PME cutoff.
exclusion_count	int		Number of excluded pairs.
one_four_count	int		Number of 1-4 corrected pairs.
explicit_exception_count	int		Number of explicit nonbonded exceptions.
cell_matrix	object | None	None	Optional full cell matrix used to assert the orthorhombic execution envelope. None treats valid lengths as orthorhombic.

Returns

• dict[str, object] — A readiness dict with a "status", a "blockers" list, and the individual boolean checks (neutrality, box, mesh, alpha, cutoff, …).