metal_kernels

Fused Metal kernels for recurring molecular force paths.

Collapses the per-step pairwise LJ force op-chain (gather -> minimum image -> r^2 -> LJ scalar -> scatter-add) into a single mx.fast.metal_kernel dispatch. Diagnostic kernels write per-pair energy without contention, while force-only kernels omit those outputs and reductions entirely.

The simple kernel covers scalar reduced-unit LJ. The parameterized kernel covers per-atom Lorentz-Berthelot parameters, topology scales, shifts, and smooth switching for the production biomolecular path. A separate force-only kernel combines standard bond, angle, periodic-torsion, and improper interactions into one output. Unsupported cases fall back transparently.

Because tests/conftest.py forces MLX_ATOMISTIC_DEVICE=cpu, the kernel is built lazily on first use (not at import) so importing this module never triggers a Metal device load.

import mlx_atomistic.metal_kernels

Functions

aligned_topology_lj_scales

def aligned_topology_lj_scales(pairs: mx.array, excluded_pairs: mx.array, one_four_pairs: mx.array, *, one_four_scale: float) -> mx.array

Build pair-aligned topology LJ scales on Metal.

Parameters

Name	Type	Default	Description
pairs	mx.array		Candidate atom pairs with shape (n_pairs, 2).
excluded_pairs	mx.array		Sorted excluded pairs with shape (n_excluded, 2).
one_four_pairs	mx.array		Sorted non-excluded 1-4 pairs with shape (n_one_four, 2).
one_four_scale	float		LJ scale assigned to 1-4 pairs.

Returns

• mx.array — Float32 scales with shape (n_pairs,). Excluded pairs are zero,
• mx.array — ordinary pairs are one, and 1-4 pairs use one_four_scale.

Raises

• ValueError — If any pair array has the wrong shape or the scale is invalid.

fused_lj_forces

def fused_lj_forces(positions: mx.array, pairs: mx.array, box_lengths: mx.array, *, epsilon: float, sigma: float, cutoff: float, shift: bool) -> tuple[mx.array, mx.array]

Fused LJ energy + forces via a single Metal kernel (orthorhombic, scalar LJ).

Mirrors LennardJonesPotential._pair_energy_forces semantics: a half neighbor list pairs of shape (M, 2), an r^2 cutoff mask, and an optional energy shift at the cutoff. box_lengths are the orthorhombic edge lengths (mx.diag(cell.matrix)). Returns (energy_scalar, forces) with forces shape (N, 3).

Parameters

Name	Type	Default	Description
positions	mx.array
pairs	mx.array
box_lengths	mx.array
epsilon	float
sigma	float
cutoff	float
shift	bool

Returns

• tuple[mx.array, mx.array]

fused_parameterized_lj_forces

def fused_parameterized_lj_forces(positions: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, scales: mx.array, *, cutoff: float, shift: bool, switch_distance: float | None) -> tuple[mx.array, mx.array]

Evaluate parameterized LJ energy and forces with one Metal dispatch.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs	mx.array		Half-neighbor pairs with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
scales	mx.array		Either one shared scale or one scale per pair.
cutoff	float		Finite LJ cutoff.
shift	bool		Whether to subtract each pair’s cutoff energy.
switch_distance	float | None		Optional start of the smooth potential switch.

Returns

• tuple[mx.array, mx.array] — Scalar LJ energy and an (n_atoms, 3) force array.

Raises

• ValueError — If the cutoff or scale count is invalid.

fused_parameterized_pme_direct_components

def fused_parameterized_pme_direct_components(positions: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, charges: mx.array, lj_scales: mx.array, *, cutoff: float, shift: bool, switch_distance: float | None, coulomb_constant: float, alpha: float) -> tuple[mx.array, mx.array, mx.array, mx.array]

Evaluate separated LJ/PME direct energies and combined forces in one dispatch.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs	mx.array		Shared half-neighbor candidates with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
charges	mx.array		Per-atom partial charges.
lj_scales	mx.array		One aligned LJ scale per candidate; zero excludes LJ only.
cutoff	float		Shared finite LJ and PME real-space cutoff.
shift	bool		Whether to subtract each LJ pair’s cutoff energy.
switch_distance	float | None		Optional start of the smooth LJ potential switch.
coulomb_constant	float		Coulomb prefactor in the configured units.
alpha	float		Ewald splitting parameter.

Returns

• mx.array — Combined energy, (n_atoms, 3) forces, LJ energy, and direct-space
• mx.array — Coulomb energy.

Raises

• ValueError — If the cutoff, alpha, or aligned scale count is invalid.

fused_parameterized_pme_direct_components_virial

def fused_parameterized_pme_direct_components_virial(positions: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, charges: mx.array, lj_scales: mx.array, *, cutoff: float, shift: bool, switch_distance: float | None, coulomb_constant: float, alpha: float) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]

Evaluate direct components, forces, and atomic diagonal virial.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs	mx.array		Shared half-neighbor candidates with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
charges	mx.array		Per-atom partial charges.
lj_scales	mx.array		One aligned LJ scale per candidate; zero excludes LJ only.
cutoff	float		Shared finite LJ and PME real-space cutoff.
shift	bool		Whether to subtract each LJ pair’s cutoff energy.
switch_distance	float | None		Optional start of the smooth LJ potential switch.
coulomb_constant	float		Coulomb prefactor in the configured units.
alpha	float		Ewald splitting parameter.

Returns

• mx.array — Combined energy, forces, LJ energy, direct Coulomb energy, and the
• mx.array — three diagonal atomic-virial components.

Raises

• ValueError — If the cutoff, alpha, or aligned scale count is invalid.

fused_parameterized_pme_direct_force_only

def fused_parameterized_pme_direct_force_only(positions: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, charges: mx.array, lj_scales: mx.array, *, cutoff: float, shift: bool, switch_distance: float | None, coulomb_constant: float, alpha: float) -> mx.array

Evaluate combined LJ plus PME direct-space forces without energy outputs.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs	mx.array		Shared half-neighbor candidates with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
charges	mx.array		Per-atom partial charges.
lj_scales	mx.array		One aligned LJ scale per candidate; zero excludes LJ only.
cutoff	float		Shared finite LJ and PME real-space cutoff.
shift	bool		Whether to subtract each LJ pair’s cutoff energy.
switch_distance	float | None		Optional start of the smooth LJ potential switch.
coulomb_constant	float		Coulomb prefactor in the configured units.
alpha	float		Ewald splitting parameter.

Returns

• mx.array — An (n_atoms, 3) force array. The kernel allocates no per-pair
• mx.array — energy outputs.

Raises

• ValueError — If the cutoff, alpha, or aligned scale count is invalid.

fused_parameterized_pme_direct_forces

def fused_parameterized_pme_direct_forces(positions: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, charges: mx.array, lj_scales: mx.array, *, cutoff: float, shift: bool, switch_distance: float | None, coulomb_constant: float, alpha: float) -> tuple[mx.array, mx.array]

Evaluate combined LJ plus PME direct-space energy and forces.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs	mx.array		Shared half-neighbor candidates with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
charges	mx.array		Per-atom partial charges.
lj_scales	mx.array		One aligned LJ scale per candidate; zero excludes LJ only.
cutoff	float		Shared finite LJ and PME real-space cutoff.
shift	bool		Whether to subtract each LJ pair’s cutoff energy.
switch_distance	float | None		Optional start of the smooth LJ potential switch.
coulomb_constant	float		Coulomb prefactor in the configured units.
alpha	float		Ewald splitting parameter.

Returns

• tuple[mx.array, mx.array] — Combined scalar direct-space energy and (n_atoms, 3) forces.

fused_pme_cutoff_correction_virial

def fused_pme_cutoff_correction_virial(positions: mx.array, molecule_centers: mx.array, pairs: mx.array, box_lengths: mx.array, sigma: mx.array, epsilon: mx.array, charges: mx.array, lj_scales: mx.array, *, cutoff: float, strain_epsilon: float, coulomb_constant: float, alpha: float, include_lj: bool, include_coulomb: bool) -> mx.array

Evaluate the finite-strain cutoff correction in one Metal dispatch.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
molecule_centers	mx.array		Per-atom geometric molecule centers, shape (n_atoms, 3).
pairs	mx.array		Compact cutoff-shell pairs with shape (n_pairs, 2).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
sigma	mx.array		Per-atom LJ sigma values.
epsilon	mx.array		Per-atom LJ epsilon values.
charges	mx.array		Per-atom partial charges.
lj_scales	mx.array		One aligned LJ scale per candidate; zero excludes LJ.
cutoff	float		Shared finite LJ and PME real-space cutoff.
strain_epsilon	float		Positive central finite-strain displacement.
coulomb_constant	float		Coulomb prefactor in the configured units.
alpha	float		Ewald splitting parameter.
include_lj	bool		Whether to include the unswitched LJ correction.
include_coulomb	bool		Whether to include the PME real-space boundary term.

Returns

• mx.array — A diagonal (3, 3) cutoff-correction virial tensor.

Raises

• ValueError — If shapes or positive scalar parameters are invalid.

fused_sparse_pme_correction_forces

def fused_sparse_pme_correction_forces(positions: mx.array, pairs: mx.array, box_lengths_and_inverses: mx.array, charge_products: mx.array, lj_sigma: mx.array, lj_epsilon: mx.array, *, coulomb_constant: float) -> mx.array

Evaluate sparse PME exclusions, exceptions, and 1-4 force corrections.

Parameters

Name	Type	Default	Description
positions	mx.array
pairs	mx.array
box_lengths_and_inverses	mx.array
charge_products	mx.array
lj_sigma	mx.array
lj_epsilon	mx.array
coulomb_constant	float

Returns

• mx.array

neighbor_pair_cutoff_mask

def neighbor_pair_cutoff_mask(positions: mx.array, pairs_i: mx.array, pairs_j: mx.array, box_lengths: mx.array, *, search_radius: float) -> mx.array

Return the periodic cutoff mask for aligned candidate pairs on Metal.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
pairs_i	mx.array		Left atom indices with shape (n_pairs,).
pairs_j	mx.array		Right atom indices with shape (n_pairs,).
box_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
search_radius	float		Positive neighbor search radius.

Returns

• mx.array — Boolean mask with one entry per candidate pair.

Raises

• ValueError — If an input shape or the search radius is invalid.

neighbor_pair_ordered_scatter

def neighbor_pair_ordered_scatter(pairs_i: mx.array, pairs_j: mx.array, close: mx.array, prefix: mx.array) -> tuple[mx.array, mx.array]

Scatter accepted neighbor candidates by their stable prefix positions.

Parameters

Name	Type	Default	Description
pairs_i	mx.array		Left atom indices with shape (n_pairs,).
pairs_j	mx.array		Right atom indices with shape (n_pairs,).
close	mx.array		Boolean cutoff mask with shape (n_pairs,).
prefix	mx.array		Inclusive integer prefix sum of close.

Returns

• mx.array — Candidate-sized left and right output buffers whose leading accepted
• mx.array — entries preserve the input candidate order.

Raises

• ValueError — If the inputs are not matching one-dimensional arrays.

pme_order5_charge_grid

def pme_order5_charge_grid(positions: mx.array, charges: mx.array, cell_lengths: mx.array, mesh_shape: tuple[int, int, int]) -> mx.array

Spread particle charges onto a PME mesh with one Metal dispatch.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
charges	mx.array		Per-atom charges with shape (n_atoms,).
cell_lengths	mx.array		Orthorhombic cell lengths with shape (3,).
mesh_shape	tuple[int, int, int]		Three positive PME mesh dimensions.

Returns

• mx.array — The float32 order-five B-spline charge grid.

Raises

• ValueError — If input shapes or mesh dimensions are invalid.

pme_order5_energy_forces

def pme_order5_energy_forces(positions: mx.array, charges: mx.array, potential_grid: mx.array, cell_lengths: mx.array) -> tuple[mx.array, mx.array]

Interpolate order-five PME energy and forces with one Metal dispatch.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates with shape (n_atoms, 3).
charges	mx.array		Per-atom charges with shape (n_atoms,).
potential_grid	mx.array		Scalar three-dimensional reciprocal potential mesh.
cell_lengths	mx.array		Orthorhombic cell lengths with shape (3,).

Returns

• tuple[mx.array, mx.array] — Scalar reciprocal energy and (n_atoms, 3) forces.

Raises

• ValueError — If input shapes are invalid.

tile_topology_lj_masks

def tile_topology_lj_masks(atom_blocks: mx.array, tile_blocks: mx.array, member_mask: mx.array, excluded_pairs: mx.array, one_four_pairs: mx.array) -> tuple[mx.array, mx.array]

Build tile-aligned LJ eligibility and 1-4 masks on Metal.

Parameters

Name	Type	Default	Description
atom_blocks	mx.array
tile_blocks	mx.array
member_mask	mx.array
excluded_pairs	mx.array
one_four_pairs	mx.array

Returns

• tuple[mx.array, mx.array]