core

Core atomistic data structures.

import mlx_atomistic.core

Classes

Atoms

class Atoms
    def __init__(symbols: tuple[str, ...], positions: mx.array, masses: mx.array, velocities: mx.array | None = None)

A collection of atoms (or particles) with MLX-backed arrays.

Parameters

Name	Type	Default	Description
symbols	tuple[str, ...]
positions	mx.array
masses	mx.array
velocities	mx.array | None	None

Properties

• count int — Number of atoms in the collection.

Methods

from_sequences

def from_sequences(symbols: Sequence[str], positions: Sequence[Sequence[float]], *, masses: Sequence[float] | None = None, velocities: Sequence[Sequence[float]] | None = None) -> Atoms

Build an Atoms collection from plain Python sequences.

Parameters

Name	Type	Default	Description
symbols	Sequence[str]		Chemical symbols, one per atom.
positions	Sequence[Sequence[float]]		(n_atoms, 3) Cartesian coordinates.
masses	Sequence[float] | None	None	Optional per-atom masses; defaults to unit mass for every atom.
velocities	Sequence[Sequence[float]] | None	None	Optional (n_atoms, 3) initial velocities.

Returns

• Atoms — An Atoms instance with MLX-backed arrays.

Cell

class Cell
    def __init__(lengths: Sequence[float] | Sequence[Sequence[float]] | mx.array)

Periodic simulation cell in MD reduced units.

The cell is stored as a row-vector matrix: row i is the i-th cell vector. One-dimensional input preserves the historical orthorhombic API; a full (3, 3) matrix stores an arbitrary triclinic box.

Parameters

Name	Type	Default	Description
lengths	Sequence[float] | Sequence[Sequence[float]] | mx.array		Either three positive box lengths with shape (3,) for an orthorhombic cell, or a (3, 3) row-vector matrix for a triclinic cell.

Properties

• is_orthorhombic bool — Whether the cell is axis-aligned orthorhombic (no off-diagonal terms).
• lengths mx.array — Cell-vector lengths (3,): the Euclidean norm of each row of matrix.
• volume mx.array — Cell volume, computed as the determinant of matrix.

Methods

cartesian_coordinates

def cartesian_coordinates(fractional: mx.array) -> mx.array

Map fractional cell coordinates back to Cartesian coordinates.

Parameters

Name	Type	Default	Description
fractional	mx.array		Fractional row-vector coordinates, shape (..., 3).

Returns

• mx.array — Cartesian coordinates, with the same shape as fractional.

cubic

def cubic(length: float) -> Cell

Create a cubic periodic cell.

Parameters

Name	Type	Default	Description
length	float		Edge length of the cube.

Returns

• Cell — A cubic Cell with all three edges equal to length.

fractional_coordinates

def fractional_coordinates(positions: mx.array) -> mx.array

Map Cartesian coordinates into the fractional cell basis.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian row-vector coordinates, shape (..., 3).

Returns

• mx.array — Fractional coordinates (Cartesian expressed in the cell basis), with
• mx.array — the same shape as positions.

minimum_image

def minimum_image(displacement: mx.array) -> mx.array

Apply the minimum-image convention to displacement vectors.

Returns the shortest periodic image of each displacement, so that pair distances use the nearest copy of each atom across cell boundaries.

Parameters

Name	Type	Default	Description
displacement	mx.array		Cartesian displacement vectors, shape (..., 3).

Returns

• mx.array — Minimum-image displacements, with the same shape as the input.

orthorhombic

def orthorhombic(lengths: Sequence[float]) -> Cell

Create an orthorhombic (axis-aligned) periodic cell.

Parameters

Name	Type	Default	Description
lengths	Sequence[float]		The three box edge lengths (a, b, c).

Returns

• Cell — An axis-aligned Cell.

Raises

• ValueError — If lengths does not contain exactly three values.

triclinic

def triclinic(matrix: Sequence[Sequence[float]]) -> Cell

Create a triclinic periodic cell from row-vector cell vectors.

Parameters

Name	Type	Default	Description
matrix	Sequence[Sequence[float]]		(3, 3) matrix whose rows are the three cell vectors.

Returns

• Cell — A Cell for the given (possibly non-orthogonal) box.

wrap

def wrap(positions: mx.array) -> mx.array

Wrap positions back into the primary periodic cell.

For orthorhombic cells this subtracts an integer number of box lengths directly (x - L*floor(x/L)) instead of round-tripping through fractional coordinates. The round-trip form (x/L - floor(x/L))*L is algebraically identical but, in float32, does not return a position that is exactly x minus an integer multiple of L — it nudges atoms near a cell boundary by up to ~1e-2. Applied every MD step, that spurious displacement does work against the forces and injects energy, breaking energy conservation over long runs (invisible in short-run tests). The direct form keeps the wrap a pure lattice translation, consistent with minimum_image.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian row-vector coordinates, shape (..., 3).

Returns

• mx.array — Positions translated into the primary cell, with the same shape as
• mx.array — positions.

Functions

as_mx_array

def as_mx_array(value: Sequence | np.ndarray | mx.array, *, dtype: mx.Dtype = DEFAULT_DTYPE) -> mx.array

Convert a value to an MLX array in the project’s default dtype.

An existing mx.array is returned unchanged unless its dtype differs. If no Metal device is available the conversion transparently falls back to the CPU device (also forced when the MLX_ATOMISTIC_DEVICE=cpu environment variable is set).

Parameters

Name	Type	Default	Description
value	Sequence | np.ndarray | mx.array		Array-like data to convert (a sequence, NumPy array, or mx.array).
dtype	mx.Dtype	DEFAULT_DTYPE	Target MLX dtype. Defaults to mx.float32.

Returns

• mx.array — The data as an mx.array of dtype.

use_cpu_device

def use_cpu_device() -> None

Route subsequent MLX operations to the CPU device.

Sets the default MLX device and stream to the CPU. Used as a fallback when no Metal GPU is available, e.g. in CI or other headless environments.

Returns

• None