nonbonded

MLX execution helpers for pairwise nonbonded interactions.

import mlx_atomistic.nonbonded

Classes

`EwaldReferenceConfig`

class EwaldReferenceConfig
    def __init__(alpha: float = 0.35, real_cutoff: float | None = None, reciprocal_cutoff: int = 5, charge_tolerance: float = 1e-05)

Controls for the small-system Ewald reference electrostatics backend.

Parameters

Name	Type	Default
`alpha`	`float`	`0.35`
`real_cutoff`	`float \| None`	`None`
`reciprocal_cutoff`	`int`	`5`
`charge_tolerance`	`float`	`1e-05`

`ForceScopeReport`

class ForceScopeReport
    def __init__(scope: ForceEvaluationScope, supported: bool, execution_path: str, backend: str | None = None, electrostatics: str | None = None, production_total_only: bool = False, diagnostic_components: bool = False, direct_space: bool = False, reciprocal_space: bool = False, component_work: bool = False, requires_full_system: bool = False, unsupported_reason: str | None = None)

Metadata for a requested force-evaluation scope.

Parameters

Name	Type	Default
`scope`	`ForceEvaluationScope`
`supported`	`bool`
`execution_path`	`str`
`backend`	`str \| None`	`None`
`electrostatics`	`str \| None`	`None`
`production_total_only`	`bool`	`False`
`diagnostic_components`	`bool`	`False`
`direct_space`	`bool`	`False`
`reciprocal_space`	`bool`	`False`
`component_work`	`bool`	`False`
`requires_full_system`	`bool`	`False`
`unsupported_reason`	`str \| None`	`None`

Methods

`to_dict`

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

Return the force-scope report as a plain JSON-serializable dict.

Returns

dict[str, object] — The report’s fields (scope, support flags, execution path, electrostatics, reason) as a dict.

`FullLoopNonbondedPolicy`

class FullLoopNonbondedPolicy
    def __init__(mode: FullLoopNonbondedMode, use_neighbor_list: bool, potential_backend: NonbondedBackend, selected_backend: str, selected_policy: str, evidence: str, neighbor_backend: str | None = None, blocker: str | None = None)

Evidence-backed nonbonded policy for full MD loops.

Parameters

Name	Type	Default
`mode`	`FullLoopNonbondedMode`
`use_neighbor_list`	`bool`
`potential_backend`	`NonbondedBackend`
`selected_backend`	`str`
`selected_policy`	`str`
`evidence`	`str`
`neighbor_backend`	`str \| None`	`None`
`blocker`	`str \| None`	`None`

`NonbondedExecutionConfig`

class NonbondedExecutionConfig
    def __init__(backend: NonbondedBackend = 'auto', electrostatics: NonbondedElectrostatics = 'cutoff', tile_size: int = 512, memory_budget_bytes: int | None = DEFAULT_DENSE_MEMORY_BUDGET_BYTES)

Execution controls for MLX nonbonded pair evaluation.

Parameters

Name	Type	Default
`backend`	`NonbondedBackend`	`'auto'`
`electrostatics`	`NonbondedElectrostatics`	`'cutoff'`
`tile_size`	`int`	`512`
`memory_budget_bytes`	`int \| None`	`DEFAULT_DENSE_MEMORY_BUDGET_BYTES`

Functions

`choose_full_loop_nonbonded_policy`

def choose_full_loop_nonbonded_policy(*, mode: str, n_atoms: int, cutoff: float | None, cell_provided: bool, dense_threshold: int = DEFAULT_FULL_LOOP_DENSE_THRESHOLD) -> FullLoopNonbondedPolicy

Choose the default nonbonded policy for a complete MD loop.

The default is based on S1/S2 full-loop and staged-neighbor evidence: above the dense threshold use a dynamic neighbor list with compacted pairs (mlx_cell_pairs), which is measured to dominate the padded-block path at every tested scale (see DEFAULT_LARGE_SYSTEM_NEIGHBOR_BACKEND), while preserving explicit dense and fallback modes.

Parameters

Name	Type	Default	Description
`mode`	`str`		Requested policy mode (`"auto"`, `"dense"`, `"dynamic-neighbor"`).
`n_atoms`	`int`		Number of atoms in the loop.
`cutoff`	`float \| None`		Nonbonded cutoff; required (finite, positive) for the neighbor-list path.
`cell_provided`	`bool`		Whether a periodic cell is available.
`dense_threshold`	`int`	`DEFAULT_FULL_LOOP_DENSE_THRESHOLD`	Atom count above which `"auto"` switches to a dynamic neighbor list. Defaults to `DEFAULT_FULL_LOOP_DENSE_THRESHOLD`.

Returns

FullLoopNonbondedPolicy — The selected FullLoopNonbondedPolicy (backend, neighbor-list use, and the evidence/blocker provenance).

Raises

ValueError — If counts are negative, or a "dynamic-neighbor" request lacks a periodic cell or a finite positive cutoff.

`choose_nonbonded_backend`

def choose_nonbonded_backend(*, requested: NonbondedBackend, n_atoms: int, pairs_provided: bool, estimated_dense_bytes: int, memory_budget_bytes: int | None) -> NonbondedBackend

Choose the concrete backend for one nonbonded evaluation.

Parameters

Name	Type	Description
`requested`	`NonbondedBackend`	The requested backend; `"auto"` resolves from the inputs below.
`n_atoms`	`int`	Number of atoms in the evaluation.
`pairs_provided`	`bool`	Whether an explicit pair list is supplied.
`estimated_dense_bytes`	`int`	Estimated dense temporary size (see `estimate_dense_nonbonded_bytes`).
`memory_budget_bytes`	`int \| None`	Dense memory budget in bytes; `None` means unbounded.

Returns

NonbondedBackend — The concrete backend to execute.

Raises

MemoryError — If "mlx_dense" is explicitly requested but exceeds the budget.

`dense_combined_energy_forces`

def dense_combined_energy_forces(positions: mx.array, *, sigma: mx.array, epsilon: mx.array, charges: mx.array, coulomb_constant: float, cutoff: float | None, lj_shift: bool, coulomb_shift: bool, cell: Cell | None, topology: Topology | None, lj_one_four_scale: float, coulomb_one_four_scale: float, tile_size: int | None = None) -> tuple[mx.array, mx.array, mx.array]

Evaluate mixed LJ+Coulomb energy components and forces with dense MLX arrays.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Atomic coordinates, shape `(n_atoms, 3)`.
`sigma`	`mx.array`		Per-atom LJ diameters, shape `(n_atoms,)` (Lorentz-Berthelot mixed).
`epsilon`	`mx.array`		Per-atom LJ well depths, shape `(n_atoms,)` (Lorentz-Berthelot mixed).
`charges`	`mx.array`		Per-atom partial charges, shape `(n_atoms,)`.
`coulomb_constant`	`float`		Coulomb prefactor in the configured unit system.
`cutoff`	`float \| None`		Pair cutoff distance; `None` disables the cutoff.
`lj_shift`	`bool`		Whether to energy-shift the LJ term at the cutoff.
`coulomb_shift`	`bool`		Whether to shift the Coulomb term at the cutoff.
`cell`	`Cell \| None`		Optional periodic cell for minimum-image distances.
`topology`	`Topology \| None`		Optional topology supplying exclusions and 1-4 scaling.
`lj_one_four_scale`	`float`		1-4 scale factor for the LJ term.
`coulomb_one_four_scale`	`float`		1-4 scale factor for the Coulomb term.
`tile_size`	`int \| None`	`None`	If set, evaluate in row tiles of this size. Defaults to `None`.

Returns

tuple[mx.array, mx.array, mx.array] — An (lj_energy, coulomb_energy, forces) tuple: the two scalar energy components and (n_atoms, 3) forces.

`dense_lj_energy_forces`

def dense_lj_energy_forces(positions: mx.array, *, epsilon: float, sigma: float, cutoff: float | None, shift: bool, cell: Cell | None, topology: Topology | None, one_four_scale: float, tile_size: int | None = None) -> tuple[mx.array, mx.array]

Evaluate uniform Lennard-Jones energy and forces with dense MLX arrays.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Atomic coordinates, shape `(n_atoms, 3)`.
`epsilon`	`float`		Uniform LJ well depth ε.
`sigma`	`float`		Uniform LJ diameter σ.
`cutoff`	`float \| None`		Pair cutoff distance; `None` disables the cutoff.
`shift`	`bool`		Whether to energy-shift the potential to zero at the cutoff.
`cell`	`Cell \| None`		Optional periodic cell for minimum-image distances; `None` is non-periodic.
`topology`	`Topology \| None`		Optional topology supplying exclusions and 1-4 scaling; `None` is all-pairs.
`one_four_scale`	`float`		Scale factor applied to 1-4 interactions.
`tile_size`	`int \| None`	`None`	If set, evaluate in row tiles of this size instead of one dense matrix. Defaults to `None`.

Returns

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

`dense_pair_mask_and_scales`

def dense_pair_mask_and_scales(n_atoms: int, *, r2: mx.array, cutoff: float | None, topology: Topology | None, lj_one_four_scale: float, coulomb_one_four_scale: float) -> tuple[mx.array, mx.array, mx.array]

Return dense pair mask and scale matrices.

Parameters

Name	Type	Description
`n_atoms`	`int`	Number of atoms (mask side length).
`r2`	`mx.array`	Pairwise squared distances, shape `(n_atoms, n_atoms)`.
`cutoff`	`float \| None`	Pair cutoff distance; `None` disables the distance mask.
`topology`	`Topology \| None`	Optional topology supplying exclusions and 1-4 scaling.
`lj_one_four_scale`	`float`	1-4 scale factor for the LJ scale matrix.
`coulomb_one_four_scale`	`float`	1-4 scale factor for the Coulomb scale matrix.

Returns

tuple[mx.array, mx.array, mx.array] — A (pair_mask, lj_scales, coulomb_scales) tuple of (n_atoms, n_atoms) arrays: a boolean inclusion mask and the per-pair LJ/Coulomb scales.

`estimate_dense_nonbonded_bytes`

def estimate_dense_nonbonded_bytes(n_atoms: int, *, components: str = 'combined') -> int

Return a conservative estimate for dense MLX nonbonded temporaries.

Parameters

Name	Type	Default	Description
`n_atoms`	`int`		Number of atoms in the dense evaluation.
`components`	`str`	`'combined'`	`"lj"` for Lennard-Jones-only temporaries or `"combined"` for LJ+Coulomb. Defaults to `"combined"`.

Returns

int — The estimated peak temporary size in bytes (grows as n_atoms²).

Raises

ValueError — If n_atoms is negative.

`ewald_reference_coulomb_energy`

def ewald_reference_coulomb_energy(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: EwaldReferenceConfig | None = None) -> tuple[mx.array, dict[str, mx.array]]

Evaluate neutral orthorhombic Ewald Coulomb energy.

This is a correctness reference for small periodic systems. It is not a particle-mesh implementation and is intentionally kept separate from the production nonbonded force path until force validation is complete.

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`	`EwaldReferenceConfig \| None`	`None`	Ewald reference parameters (alpha, cutoffs, charge tolerance); `None` uses defaults. Defaults to `None`.

Returns

tuple[mx.array, dict[str, mx.array]] — An (energy, components) tuple: scalar total Coulomb energy and a dict of coulomb_real/coulomb_reciprocal/coulomb_self parts.

Raises

ValueError — If shapes are wrong, the cell is non-orthorhombic, lengths are non-positive, or the system is not neutral.

`ewald_reference_coulomb_energy_forces`

def ewald_reference_coulomb_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: EwaldReferenceConfig | None = None) -> tuple[mx.array, mx.array, dict[str, mx.array]]

Evaluate neutral orthorhombic Ewald Coulomb energy and analytical 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`	`EwaldReferenceConfig \| None`	`None`	Ewald reference parameters; `None` uses defaults. Defaults to `None`.

Returns

tuple[mx.array, mx.array, dict[str, mx.array]] — An (energy, forces, components) tuple: scalar energy, (n_atoms, 3) forces, and the real/reciprocal/self component dict.

Raises

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

`ewald_reference_coulomb_total_energy_forces`

def ewald_reference_coulomb_total_energy_forces(positions: mx.array, charges: mx.array, cell: Cell, *, coulomb_constant: float = 1.0, config: EwaldReferenceConfig | None = None) -> tuple[mx.array, mx.array]

Evaluate neutral orthorhombic Ewald Coulomb total 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`	`EwaldReferenceConfig \| None`	`None`	Ewald reference parameters; `None` uses defaults. Defaults to `None`.

Returns

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

Raises

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

`normalize_force_scope`

def normalize_force_scope(scope: str) -> ForceEvaluationScope

Validate and normalize a force-evaluation scope name.

Parameters

Name	Type	Default	Description
`scope`	`str`		Scope name or alias (case/separator-insensitive), e.g. `"total"`, `"components"`, `"direct_space"`, `"reciprocal_space"`.

Returns

ForceEvaluationScope — The canonical force-evaluation scope.

Raises

ValueError — If scope is not a recognized scope.

`normalize_nonbonded_electrostatics`

def normalize_nonbonded_electrostatics(mode: str) -> NonbondedElectrostatics

Normalize an electrostatics mode or known metadata alias.

Parameters

Name	Type	Default	Description
`mode`	`str`		Electrostatics mode or alias (case/separator-insensitive), resolving to `"cutoff"`, `"ewald_reference"`, or `"pme"`.

Returns

NonbondedElectrostatics — The canonical electrostatics mode.

Raises

ValueError — If mode does not resolve to a known electrostatics mode.

`tiled_combined_energy_forces`

def tiled_combined_energy_forces(positions: mx.array, *, sigma: mx.array, epsilon: mx.array, charges: mx.array, coulomb_constant: float, cutoff: float | None, lj_shift: bool, coulomb_shift: bool, cell: Cell | None, topology: Topology | None, lj_one_four_scale: float, coulomb_one_four_scale: float, tile_size: int) -> tuple[mx.array, mx.array, mx.array]

Evaluate mixed LJ+Coulomb energy components and forces by row tiles.

Parameters

Name	Type	Description
`positions`	`mx.array`	Atomic coordinates, shape `(n_atoms, 3)`.
`sigma`	`mx.array`	Per-atom LJ diameters, shape `(n_atoms,)` (Lorentz-Berthelot mixed).
`epsilon`	`mx.array`	Per-atom LJ well depths, shape `(n_atoms,)` (Lorentz-Berthelot mixed).
`charges`	`mx.array`	Per-atom partial charges, shape `(n_atoms,)`.
`coulomb_constant`	`float`	Coulomb prefactor in the configured unit system.
`cutoff`	`float \| None`	Pair cutoff distance; `None` disables the cutoff.
`lj_shift`	`bool`	Whether to energy-shift the LJ term at the cutoff.
`coulomb_shift`	`bool`	Whether to shift the Coulomb term at the cutoff.
`cell`	`Cell \| None`	Optional periodic cell for minimum-image distances.
`topology`	`Topology \| None`	Optional topology supplying exclusions and 1-4 scaling.
`lj_one_four_scale`	`float`	1-4 scale factor for the LJ term.
`coulomb_one_four_scale`	`float`	1-4 scale factor for the Coulomb term.
`tile_size`	`int`	Row-block size for the tiled evaluation.

Returns

tuple[mx.array, mx.array, mx.array] — An (lj_energy, coulomb_energy, forces) tuple: the two scalar energy components and (n_atoms, 3) forces.

`tiled_lj_energy_forces`

def tiled_lj_energy_forces(positions: mx.array, *, epsilon: float, sigma: float, cutoff: float | None, shift: bool, cell: Cell | None, topology: Topology | None, one_four_scale: float, tile_size: int) -> tuple[mx.array, mx.array]

Evaluate uniform Lennard-Jones energy and forces by row tiles.

Parameters

Name	Type	Description
`positions`	`mx.array`	Atomic coordinates, shape `(n_atoms, 3)`.
`epsilon`	`float`	Uniform LJ well depth ε.
`sigma`	`float`	Uniform LJ diameter σ.
`cutoff`	`float \| None`	Pair cutoff distance; `None` disables the cutoff.
`shift`	`bool`	Whether to energy-shift the potential to zero at the cutoff.
`cell`	`Cell \| None`	Optional periodic cell for minimum-image distances; `None` is non-periodic.
`topology`	`Topology \| None`	Optional topology supplying exclusions and 1-4 scaling; `None` is all-pairs.
`one_four_scale`	`float`	Scale factor applied to 1-4 interactions.
`tile_size`	`int`	Row-block size for the tiled evaluation.

Returns

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

`topology_dense_matrices`

def topology_dense_matrices(n_atoms: int, *, topology: Topology | None, lj_one_four_scale: float, coulomb_one_four_scale: float) -> tuple[mx.array | None, mx.array | None, mx.array | None]

Build dense topology masks/scales without creating explicit pair lists.

Parameters

Name	Type	Description
`n_atoms`	`int`	Number of atoms (matrix side length).
`topology`	`Topology \| None`	Topology supplying exclusions and 1-4 pairs; `None` returns all-`None`.
`lj_one_four_scale`	`float`	1-4 scale factor for the LJ scale matrix.
`coulomb_one_four_scale`	`float`	1-4 scale factor for the Coulomb scale matrix.

Returns

tuple[mx.array | None, mx.array | None, mx.array | None] — A (mask, lj_scales, coulomb_scales) tuple of (n_atoms, n_atoms) arrays, or all None when topology is None.

Raises

ValueError — If topology.n_atoms does not equal n_atoms.

`validate_full_loop_nonbonded_mode`

def validate_full_loop_nonbonded_mode(mode: str) -> FullLoopNonbondedMode

Validate a full-loop nonbonded policy mode.

Parameters

Name	Type	Default	Description
`mode`	`str`		Policy mode; one of `"auto"`, `"dense"`, `"dynamic-neighbor"`.

Returns

FullLoopNonbondedMode — The validated mode.

Raises

ValueError — If mode is not a recognized mode.

`validate_nonbonded_backend`

def validate_nonbonded_backend(backend: str) -> NonbondedBackend

Validate and normalize a nonbonded backend name.

Parameters

Name	Type	Default	Description
`backend`	`str`		Backend name; one of `"auto"`, `"mlx_dense"`, `"mlx_tiled"`, `"mlx_pairs"`, `"python_neighbor"`.

Returns

NonbondedBackend — The validated backend name.

Raises

ValueError — If backend is not a recognized backend.

`validate_nonbonded_electrostatics`

def validate_nonbonded_electrostatics(mode: str) -> NonbondedElectrostatics

Validate an electrostatics mode for executable nonbonded evaluation.

Parameters

Name	Type	Default	Description
`mode`	`str`		Electrostatics mode or alias to validate.

Returns

NonbondedElectrostatics — The canonical electrostatics mode.

Raises

ValueError — If mode is not a recognized electrostatics mode.