io

Structure and trajectory I/O helpers.

import mlx_atomistic.io

Classes

RuntimeTraceReporter

class RuntimeTraceReporter
    def __init__(include_samples: bool = True, include_diagnostics: bool = True, events: list[dict[str, Any]] = list())

Collect scalar runtime reporter events for parity and diagnostics traces.

Parameters

Name	Type	Default	Description
include_samples	bool	True
include_diagnostics	bool	True
events	list[dict[str, Any]]	list()

Methods

to_jsonable

def to_jsonable() -> list[dict[str, Any]]

Return the collected events as a list of JSON-serializable dicts.

Returns

• list[dict[str, Any]] — One dict per recorded reporter event, in arrival order.

SimulationCheckpoint

class SimulationCheckpoint
    def __init__(positions: np.ndarray, velocities: np.ndarray, masses: np.ndarray, forces: np.ndarray, step: int, time: float, cell: np.ndarray | None, thermostat: dict[str, Any], neighbor_policy: dict[str, Any], force_terms: tuple[str, ...], diagnostic_cursor: int, metadata: dict[str, Any])

Serializable boundary for resuming a production MD run.

Parameters

Name	Type	Default	Description
positions	np.ndarray
velocities	np.ndarray
masses	np.ndarray
forces	np.ndarray
step	int
time	float
cell	np.ndarray | None
thermostat	dict[str, Any]
neighbor_policy	dict[str, Any]
force_terms	tuple[str, ...]
diagnostic_cursor	int
metadata	dict[str, Any]

Properties

• hmr_state dict[str, Any] — Hydrogen-mass-repartitioning state recovered from the checkpoint metadata.

Methods

state

def state() -> SimulationState

Rebuild the in-memory simulation state from this checkpoint.

Returns

• SimulationState — A SimulationState with MLX-backed arrays.

TrajectoryRecord

class TrajectoryRecord
    def __init__(sampled_positions: np.ndarray, sampled_velocities: np.ndarray, sampled_steps: np.ndarray, sampled_time: np.ndarray, diagnostic_steps: np.ndarray, diagnostic_time: np.ndarray, potential_energy: np.ndarray, kinetic_energy: np.ndarray, total_energy: np.ndarray, potential_energy_by_term: dict[str, np.ndarray], temperature: np.ndarray, pair_count: np.ndarray, rebuild_count: np.ndarray, constraint_max_error: np.ndarray, symbols: tuple[str, ...], cell: np.ndarray | None, metadata: dict[str, Any], virial_tensor: np.ndarray | None = None, pressure_tensor: np.ndarray | None = None, pressure: np.ndarray | None = None)

Native loaded trajectory record.

Parameters

Name	Type	Default	Description
sampled_positions	np.ndarray
sampled_velocities	np.ndarray
sampled_steps	np.ndarray
sampled_time	np.ndarray
diagnostic_steps	np.ndarray
diagnostic_time	np.ndarray
potential_energy	np.ndarray
kinetic_energy	np.ndarray
total_energy	np.ndarray
potential_energy_by_term	dict[str, np.ndarray]
temperature	np.ndarray
pair_count	np.ndarray
rebuild_count	np.ndarray
constraint_max_error	np.ndarray
symbols	tuple[str, ...]
cell	np.ndarray | None
metadata	dict[str, Any]
virial_tensor	np.ndarray | None	None
pressure_tensor	np.ndarray | None	None
pressure	np.ndarray | None	None

Functions

load_npz_trajectory

def load_npz_trajectory(path: str | Path) -> TrajectoryRecord

Load a native trajectory record from a .npz file.

Parameters

Name	Type	Default	Description
path	str | Path		Path to a file written by save_npz_trajectory.

Returns

• TrajectoryRecord — The reconstructed TrajectoryRecord.

load_simulation_checkpoint

def load_simulation_checkpoint(path: str | Path) -> SimulationCheckpoint

Load a restart checkpoint written by save_simulation_checkpoint.

Parameters

Name	Type	Default	Description
path	str | Path		Path to a checkpoint file.

Returns

• SimulationCheckpoint — The reconstructed SimulationCheckpoint.

read_xyz

def read_xyz(path: str | Path) -> tuple[tuple[str, ...], np.ndarray, str]

Read a single-frame XYZ file.

Parameters

Name	Type	Default	Description
path	str | Path		Path to the .xyz file.

Returns

• tuple[tuple[str, ...], np.ndarray, str] — A (symbols, positions, comment) tuple: element symbols, an (n_atoms, 3) float array of coordinates, and the comment line.

Raises

• ValueError — If the file has fewer than two lines or the declared atom count does not match the number of coordinate lines.

restart_state_from_trajectory

def restart_state_from_trajectory(record: TrajectoryRecord, masses, force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...], *, cell: Cell | None = None, frame: int = -1) -> SimulationState

Recompute forces for a continuation-ready state from a trajectory frame.

Parameters

Name	Type	Default	Description
record	TrajectoryRecord		The TrajectoryRecord to restart from.
masses			Per-particle masses, shape (n_particles,).
force_terms	ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]		One or more force terms used to recompute forces at the frame.
cell	Cell | None	None	Optional periodic cell. Defaults to None.
frame	int	-1	Frame index into the sampled trajectory. Defaults to -1 (last frame).

Returns

• SimulationState — A SimulationState with the frame’s positions/velocities and freshly recomputed forces.

Raises

• ValueError — If force_terms is empty.

save_npz_trajectory

def save_npz_trajectory(path: str | Path, result, *, symbols: list[str] | tuple[str, ...] | None = None, cell: Cell | None = None, metadata: dict[str, Any] | None = None) -> None

Save a native trajectory record to a compressed .npz file.

Parameters

Name	Type	Default	Description
path	str | Path		Destination .npz path.
result			A finished simulation result (e.g. NVEResult) to serialize.
symbols	list[str] | tuple[str, ...] | None	None	Optional element symbols stored with the trajectory. Defaults to None.
cell	Cell | None	None	Optional periodic cell stored with the frames. Defaults to None.
metadata	dict[str, Any] | None	None	Optional JSON-serializable metadata dict. Defaults to None.

Returns

• None

save_simulation_checkpoint

def save_simulation_checkpoint(path: str | Path, state: SimulationState, *, cell: Cell | None = None, thermostat: dict[str, Any] | None = None, neighbor_policy: dict[str, Any] | None = None, force_terms: tuple[str, ...] | list[str] | None = None, diagnostic_cursor: int | None = None, metadata: dict[str, Any] | None = None, runtime_sync_report: dict[str, int | float] | None = None, runtime_nonbonded_report: dict[str, int | float | str | None] | None = None) -> None

Write a restart checkpoint for a runner-level continuation.

Parameters

Name	Type	Default	Description
path	str | Path		Destination checkpoint path (parent directories are created).
state	SimulationState		The SimulationState to serialize.
cell	Cell | None	None	Optional periodic cell. Defaults to None.
thermostat	dict[str, Any] | None	None	Optional thermostat state (RNG offset or Nose-Hoover state). Defaults to None.
neighbor_policy	dict[str, Any] | None	None	Optional neighbor-list policy dict. Defaults to None.
force_terms	tuple[str, ...] | list[str] | None	None	Optional names of the force terms in effect. Defaults to None.
diagnostic_cursor	int | None	None	Optional index into the diagnostic series for exact resumption. Defaults to None.
metadata	dict[str, Any] | None	None	Optional JSON-serializable metadata dict. Defaults to None.
runtime_sync_report	dict[str, int | float] | None	None	Optional runtime-sync counters to persist. Defaults to None.
runtime_nonbonded_report	dict[str, int | float | str | None] | None	None	Optional nonbonded-runtime report to persist. Defaults to None.

Returns

• None

trajectory_record_from_result

def trajectory_record_from_result(result, *, symbols: list[str] | tuple[str, ...] | None = None, cell: Cell | None = None, metadata: dict[str, Any] | None = None) -> TrajectoryRecord

Create a native loaded-record view from an in-memory simulation result.

Parameters

Name	Type	Default	Description
result			A finished simulation result to convert.
symbols	list[str] | tuple[str, ...] | None	None	Optional element symbols. Defaults to None.
cell	Cell | None	None	Optional periodic cell. Defaults to None.
metadata	dict[str, Any] | None	None	Optional JSON-serializable metadata dict. Defaults to None.

Returns

• TrajectoryRecord — A TrajectoryRecord mirroring what load_npz_trajectory would return for the same run.

write_xyz

def write_xyz(path: str | Path, symbols: list[str] | tuple[str, ...], positions, *, comment: str = '') -> None

Write a single-frame XYZ file.

Parameters

Name	Type	Default	Description
path	str | Path		Destination path for the .xyz file.
symbols	list[str] | tuple[str, ...]		Element symbols, one per atom.
positions			Atomic coordinates, shape (n_atoms, 3).
comment	str	''	Comment written as the file’s second line. Defaults to "".

Returns

• None

Raises

• ValueError — If positions does not have shape (len(symbols), 3).