md

Molecular dynamics primitives.

import mlx_atomistic.md

Classes

`VelocityVerlet`

class VelocityVerlet
    def __init__(dt: float)

Velocity Verlet integrator.

Parameters

Name	Type	Default	Description
`dt`	`float`

Methods

`step`

def step(positions: mx.array, velocities: mx.array, masses: mx.array, potential: LennardJonesPotential, *, cell: Cell | None = None, forces: mx.array | None = None, pairs: object | None = None) -> StepState

Advance one MD step.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Current coordinates, shape `(n_particles, 3)`.
`velocities`	`mx.array`		Current velocities, shape `(n_particles, 3)`.
`masses`	`mx.array`		Per-particle masses, shape `(n_particles,)`.
`potential`	`LennardJonesPotential`		Force model providing `energy_forces`.
`cell`	`Cell \| None`	`None`	Optional periodic cell; positions are wrapped into it when given. Defaults to `None`.
`forces`	`mx.array \| None`	`None`	Optional forces at the current positions to skip a recompute; `None` evaluates them. Defaults to `None`.
`pairs`	`object \| None`	`None`	Optional neighbor/pair structure passed to the potential. Defaults to `None`.

Returns

StepState — The StepState after one Velocity Verlet step (new positions, velocities, forces, and energies).

Functions

`analytic_configurational_virial_tensor`

def analytic_configurational_virial_tensor(positions: mx.array, forces: mx.array, force_terms: tuple[ForceTerm, ...], *, cell: Cell | None, pairs: object | None, masses: mx.array | None = None, molecule_ids: object | None = None) -> mx.array

Return the production analytic diagonal configurational virial.

Parameters

Name	Type	Default
`positions`	`mx.array`
`forces`	`mx.array`
`force_terms`	`tuple[ForceTerm, ...]`
`cell`	`Cell \| None`
`pairs`	`object \| None`
`masses`	`mx.array \| None`	`None`
`molecule_ids`	`object \| None`	`None`

Returns

mx.array

`configurational_virial_tensor`

def configurational_virial_tensor(positions: mx.array, forces: mx.array, force_terms: tuple[ForceTerm, ...], *, cell: Cell | None, pairs: object | None, virtual_sites: VirtualSiteManager | None = None, strain_epsilon: float = 0.001, masses: mx.array | None = None, molecule_ids: object | None = None, virial_mode: str = VIRIAL_SUPPORT_ANALYTIC) -> mx.array

Return an analytic production virial or the named validation oracle.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`forces`	`mx.array`		Forces used for the non-periodic fallback.
`force_terms`	`tuple[ForceTerm, ...]`		Force terms contributing to the configurational virial.
`cell`	`Cell \| None`		Periodic cell; `None` uses the non-periodic force virial.
`pairs`	`object \| None`		Optional neighbor/pair structure forwarded to analytic terms.
`virtual_sites`	`VirtualSiteManager \| None`	`None`	Optional virtual-site manager applied before energy evaluation. Defaults to `None`.
`strain_epsilon`	`float`	`0.001`	Half-width used only by the finite-difference oracle.
`masses`	`mx.array \| None`	`None`	Optional particle masses for molecular center construction.
`molecule_ids`	`object \| None`	`None`	Optional contiguous per-particle molecule identifiers; absent identifiers treat every particle as its own molecule.
`virial_mode`	`str`	`VIRIAL_SUPPORT_ANALYTIC`	`analytic` for production or `finite_difference_oracle` for validation.

Returns

mx.array — The diagonal (3, 3) configurational virial tensor.

Raises

ValueError — If the requested support level or cell is invalid.

`finite_difference_configurational_virial_oracle`

def finite_difference_configurational_virial_oracle(positions: mx.array, forces: mx.array, force_terms: tuple[ForceTerm, ...], *, cell: Cell | None, pairs: object | None, virtual_sites: VirtualSiteManager | None = None, strain_epsilon: float = 0.001, masses: mx.array | None = None, molecule_ids: object | None = None) -> mx.array

Return the validation-only molecular cell-strain virial oracle.

Parameters

Name	Type	Default
`positions`	`mx.array`
`forces`	`mx.array`
`force_terms`	`tuple[ForceTerm, ...]`
`cell`	`Cell \| None`
`pairs`	`object \| None`
`virtual_sites`	`VirtualSiteManager \| None`	`None`
`strain_epsilon`	`float`	`0.001`
`masses`	`mx.array \| None`	`None`
`molecule_ids`	`object \| None`	`None`

Returns

mx.array

`missing_analytic_virial_support`

def missing_analytic_virial_support(force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]) -> tuple[str, ...]

Return exact force-term names lacking production analytic virial support.

Parameters

Name	Type	Default	Description
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...]`

Returns

tuple[str, ...]

`missing_virial_support`

def missing_virial_support(force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]) -> tuple[str, ...]

Return exact force-term names without a supported virial diagnostics path.

Parameters

Name	Type	Default	Description
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...]`		A single force term or a list/tuple of force terms to inspect.

Returns

tuple[str, ...] — The names of the terms lacking virial support, in input order (empty if all are supported).

`pressure_tensor`

def pressure_tensor(positions: mx.array, velocities: mx.array, masses: mx.array, forces: mx.array, force_terms: tuple[ForceTerm, ...], *, cell: Cell | None, pairs: object | None, kinetic_energy_scale: float = 1.0, virtual_sites: VirtualSiteManager | None = None, molecule_ids: object | None = None, virial_mode: str = VIRIAL_SUPPORT_ANALYTIC) -> tuple[mx.array, mx.array, mx.array]

Return virial tensor, pressure tensor, and scalar pressure diagnostics.

The pressure tensor uses the reduced-unit convention P = (kinetic tensor + configurational virial) / V. Periodic virials are diagonal-only orthorhombic cell-strain diagnostics; non-periodic runs report finite zero pressure diagnostics because no volume is defined.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`velocities`	`mx.array`		Per-particle velocities, shape `(n_particles, 3)`.
`masses`	`mx.array`		Per-particle masses, shape `(n_particles,)`.
`forces`	`mx.array`		Forces on each particle, shape `(n_particles, 3)`.
`force_terms`	`tuple[ForceTerm, ...]`		Force terms supplying the configurational virial; each must support virial diagnostics.
`cell`	`Cell \| None`		Periodic cell; `None` reports zero pressure (no volume defined).
`pairs`	`object \| None`		Optional neighbor/pair structure forwarded to the force terms.
`kinetic_energy_scale`	`float`	`1.0`	Energy-unit factor for the kinetic tensor. Defaults to `1.0`.
`virtual_sites`	`VirtualSiteManager \| None`	`None`	Optional virtual-site manager applied before energy evaluation. Defaults to `None`.
`molecule_ids`	`object \| None`	`None`	Optional exact molecule membership for molecular configurational and kinetic pressure.
`virial_mode`	`str`	`VIRIAL_SUPPORT_ANALYTIC`	`analytic` for production or `finite_difference_oracle` for validation.

Returns

tuple[mx.array, mx.array, mx.array] — A (virial, pressure, scalar) tuple: the (3, 3) configurational virial, the (3, 3) pressure tensor (kinetic + virial) / V, and the scalar pressure tr(P) / 3.

Raises

ValueError — If a periodic cell has non-positive volume.

`simulate`

def simulate(positions, velocities, *, masses = None, cell: Cell | None = None, potential: LennardJonesPotential | None = None, pairs: object | None = None, dt: float = 0.005, steps: int = 100) -> SimulationResult

Run a short NVE MD simulation in reduced units.

Parameters

Name	Type	Default	Description
`positions`			Initial coordinates, shape `(n_particles, 3)`.
`velocities`			Initial velocities, shape `(n_particles, 3)`.
`masses`		`None`	Per-particle masses, shape `(n_particles,)`; `None` uses unit masses. Defaults to `None`.
`cell`	`Cell \| None`	`None`	Optional periodic cell for minimum-image distances and wrapping. Defaults to `None`.
`potential`	`LennardJonesPotential \| None`	`None`	Force model to integrate; `None` uses a default `LennardJonesPotential`. Defaults to `None`.
`pairs`	`object \| None`	`None`	Optional precomputed neighbor/pair structure passed to the potential. Defaults to `None`.
`dt`	`float`	`0.005`	Integration time step. Defaults to `0.005`.
`steps`	`int`	`100`	Number of Velocity Verlet steps. Defaults to `100`.

Returns

SimulationResult — A SimulationResult with stacked per-frame positions, velocities, and energy/temperature series (steps + 1 frames).

`simulate_npt`

def simulate_npt(positions, velocities, *, masses = None, cell: Cell | None = None, force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...] | None = None, neighbor_manager: NeighborListManager | None = None, config: SimulationConfig | None = None, thermostat: Thermostat | None = None, barostat: MonteCarloBarostat | None = None, barostat_state: dict[str, Any] | None = None, constraints: DistanceConstraints | None = None, molecule_ids: object | None = None, reporters: RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None = None) -> NPTResult

Run molecular Monte Carlo pressure coupling at exact in-loop intervals.

Parameters

Name	Type	Default	Description
`positions`			Initial coordinates, shape `(n_particles, 3)`.
`velocities`			Initial velocities, shape `(n_particles, 3)`.
`masses`		`None`	Per-particle masses, shape `(n_particles,)`; `None` uses unit masses. Defaults to `None`.
`cell`	`Cell \| None`	`None`	Periodic cell (required for NPT). Defaults to `None`.
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...] \| None`	`None`	One or more force terms; `None` uses a default `LennardJonesPotential`. Defaults to `None`.
`neighbor_manager`	`NeighborListManager \| None`	`None`	Optional neighbor-list manager. Defaults to `None`.
`config`	`SimulationConfig \| None`	`None`	Run configuration; `None` uses defaults. Defaults to `None`.
`thermostat`	`Thermostat \| None`	`None`	Thermostat for the NVT stage; `None` uses a default `LangevinThermostat`. Defaults to `None`.
`barostat`	`MonteCarloBarostat \| None`	`None`	Monte Carlo barostat; `None` uses one matched to the thermostat temperature. Defaults to `None`.
`barostat_state`	`dict[str, Any] \| None`	`None`	Optional serialized persistent barostat state from a prior committed NPT boundary. Defaults to `None`.
`constraints`	`DistanceConstraints \| None`	`None`	Optional distance constraints applied each step. Defaults to `None`.
`molecule_ids`	`object \| None`	`None`	Optional contiguous per-particle molecule identifiers. When omitted, each particle is treated as a separate molecule.
`reporters`	`RuntimeReporter \| list[RuntimeReporter] \| tuple[RuntimeReporter, ...] \| None`	`None`	Optional runtime reporter(s). Defaults to `None`.

Returns

NPTResult — An NPTResult with the integrated trajectory, sampled cell history, and persistent barostat counters.

Raises

ValueError — If cell is None (NPT requires a periodic cell).

`simulate_nve`

def simulate_nve(positions, velocities, *, masses = None, cell: Cell | None = None, force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...] | None = None, neighbor_manager: NeighborListManager | None = None, config: SimulationConfig | None = None, constraints: DistanceConstraints | None = None, reporters: RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None = None) -> NVEResult

Run NVE molecular dynamics with sparse trajectory and configurable diagnostics.

sample_interval controls trajectory storage. diagnostic_interval controls energy, temperature, pair-count, and constraint diagnostics.

Parameters

Name	Type	Default	Description
`positions`			Initial coordinates, shape `(n_particles, 3)`.
`velocities`			Initial velocities, shape `(n_particles, 3)`.
`masses`		`None`	Per-particle masses, shape `(n_particles,)`; `None` uses unit masses. Defaults to `None`.
`cell`	`Cell \| None`	`None`	Optional periodic cell. Defaults to `None`.
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...] \| None`	`None`	One or more force terms; `None` uses a default `LennardJonesPotential`. Defaults to `None`.
`neighbor_manager`	`NeighborListManager \| None`	`None`	Optional neighbor-list manager for compact nonbonded backends. Defaults to `None`.
`config`	`SimulationConfig \| None`	`None`	Run configuration (step count, sampling/diagnostic intervals, virtual sites); `None` uses defaults. Defaults to `None`.
`constraints`	`DistanceConstraints \| None`	`None`	Optional distance constraints applied each step. Defaults to `None`.
`reporters`	`RuntimeReporter \| list[RuntimeReporter] \| tuple[RuntimeReporter, ...] \| None`	`None`	Optional runtime reporter(s) invoked on diagnostic events. Defaults to `None`.

Returns

NVEResult — An NVEResult with the sparse trajectory, diagnostics, and energy-drift metrics.

`simulate_nvt`

def simulate_nvt(positions, velocities, *, masses = None, cell: Cell | None = None, force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...] | None = None, neighbor_manager: NeighborListManager | None = None, config: SimulationConfig | None = None, thermostat: Thermostat | None = None, constraints: DistanceConstraints | None = None, reporters: RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None = None) -> NVTResult

Run NVT molecular dynamics with Langevin BAOAB or Nose-Hoover dynamics.

Parameters

Name	Type	Default	Description
`positions`			Initial coordinates, shape `(n_particles, 3)`.
`velocities`			Initial velocities, shape `(n_particles, 3)`.
`masses`		`None`	Per-particle masses, shape `(n_particles,)`; `None` uses unit masses. Defaults to `None`.
`cell`	`Cell \| None`	`None`	Optional periodic cell. Defaults to `None`.
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...] \| None`	`None`	One or more force terms; `None` uses a default `LennardJonesPotential`. Defaults to `None`.
`neighbor_manager`	`NeighborListManager \| None`	`None`	Optional neighbor-list manager for compact nonbonded backends. Defaults to `None`.
`config`	`SimulationConfig \| None`	`None`	Run configuration; `None` uses defaults. Defaults to `None`.
`thermostat`	`Thermostat \| None`	`None`	Langevin (BAOAB) or Nose-Hoover thermostat; `None` uses a default `LangevinThermostat`. Defaults to `None`.
`constraints`	`DistanceConstraints \| None`	`None`	Optional distance constraints applied each step. Defaults to `None`.
`reporters`	`RuntimeReporter \| list[RuntimeReporter] \| tuple[RuntimeReporter, ...] \| None`	`None`	Optional runtime reporter(s). Defaults to `None`.

Returns

NVTResult — An NVTResult with the trajectory, diagnostics, and temperature-control metrics.

`validate_analytic_virial_support`

def validate_analytic_virial_support(force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]) -> None

Fail closed when production pressure sees an oracle-only force term.

Parameters

Name	Type	Default	Description
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...]`

Returns

None

`validate_virial_support`

def validate_virial_support(force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]) -> None

Fail closed when future pressure-coupled runtimes see unsupported terms.

Parameters

Name	Type	Default	Description
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...]`		A single force term or a list/tuple of force terms to validate.

Returns

None

Raises

ValueError — If any term lacks a supported virial diagnostics path.

`virial_readiness_report`

def virial_readiness_report(force_terms: ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...], *, require_analytic: bool = False) -> ReadinessReport

Return per-term analytic, oracle-only, or unsupported virial readiness.

Parameters

Name	Type	Default	Description
`force_terms`	`ForceTerm \| list[ForceTerm] \| tuple[ForceTerm, ...]`
`require_analytic`	`bool`	`False`

Returns

ReadinessReport

`virial_support_state`

def virial_support_state(term: ForceTerm) -> str

Return one force term’s truthful virial capability.

Parameters

Name	Type	Default	Description
`term`	`ForceTerm`

Returns

str