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)

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

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

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

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.

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.

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) -> 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,); the system must be net-neutral.
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.

Returns

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

Raises

• ValueError — If shapes/cell are invalid or the system is not neutral.

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) -> 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,); the system must be net-neutral.
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.

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 the system is not neutral.

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) -> 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,); the system must be net-neutral.
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.

Returns

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

Raises

• ValueError — If shapes/cell are invalid or the system is not neutral.

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) -> 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,); the system must be net-neutral.
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.

Returns

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

Raises

• ValueError — If shapes/cell are invalid or the system is not neutral.

pme_direct_space_policy_report

def pme_direct_space_policy_report(cell: Cell, *, config: PMEConfig | None = None, pairs: mx.array | NeighborBlocks | 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.

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

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

Returns

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