md

Molecular dynamics primitives.

import mlx_atomistic.md

Classes

`ForceTerm`

class ForceTerm(Protocol)

Protocol for composable force terms.

Methods

`energy_forces`

def energy_forces(positions: mx.array, cell: Cell | None = None, pairs: object | None = None) -> tuple[mx.array, mx.array]

Return potential energy and forces.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`cell`	`Cell \| None`	`None`	Optional periodic cell for minimum-image distances. Defaults to `None`.
`pairs`	`object \| None`	`None`	Optional precomputed neighbor/pair structure. Defaults to `None`.

Returns

tuple[mx.array, mx.array] — An (energy, forces) tuple: scalar potential energy and forces of shape (n_particles, 3).

`LangevinThermostat`

class LangevinThermostat
    def __init__(temperature: float = 1.0, friction: float = 1.0, seed: int | None = None, rng_step_offset: int | None = None)

Langevin thermostat parameters in reduced units.

Parameters

Name	Type	Default
`temperature`	`float`	`1.0`
`friction`	`float`	`1.0`
`seed`	`int \| None`	`None`
`rng_step_offset`	`int \| None`	`None`

`LennardJonesPotential`

class LennardJonesPotential
    def __init__(epsilon: float = 1.0, sigma: float = 1.0, cutoff: float | None = 2.5, shift: bool = True, topology: Topology | None = None, one_four_scale: float = 1.0, backend: NonbondedBackend = 'auto', tile_size: int = 512, memory_budget_bytes: int | None = DEFAULT_DENSE_MEMORY_BUDGET_BYTES, name: str = 'lj', supports_virial: bool = True, use_fused_kernel: bool = False)

Naive all-pairs Lennard-Jones potential in reduced units.

Parameters

Name	Type	Default
`epsilon`	`float`	`1.0`
`sigma`	`float`	`1.0`
`cutoff`	`float \| None`	`2.5`
`shift`	`bool`	`True`
`topology`	`Topology \| None`	`None`
`one_four_scale`	`float`	`1.0`
`backend`	`NonbondedBackend`	`'auto'`
`tile_size`	`int`	`512`
`memory_budget_bytes`	`int \| None`	`DEFAULT_DENSE_MEMORY_BUDGET_BYTES`
`name`	`str`	`'lj'`
`supports_virial`	`bool`	`True`
`use_fused_kernel`	`bool`	`False`

Methods

`energy_forces`

def energy_forces(positions: mx.array, cell: Cell | None = None, pairs: object | None = None) -> tuple[mx.array, mx.array]

Return potential energy and forces for positions with shape (n_particles, 3).

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`cell`	`Cell \| None`	`None`	Optional periodic cell for minimum-image distances. Defaults to `None`.
`pairs`	`object \| None`	`None`	Optional neighbor list, neighbor blocks, or dense pair array; the nonbonded backend is chosen automatically. Defaults to `None`.

Returns

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

Raises

ValueError — If positions is not (n_particles, 3) or a lazy topology is used without a runtime pair provider.

`MonteCarloBarostat`

class MonteCarloBarostat
    def __init__(pressure: float = 1.0, temperature: float = 1.0, interval: int = 25, max_log_volume_scale: float = 0.02, seed: int | None = 11, mode: str = 'isotropic', axes: tuple[bool, bool, bool] = (True, True, True), membrane_plane: str = 'xy', normal_axis: str = 'z')

Monte Carlo barostat parameters for isotropic, anisotropic, and membrane NPT.

Parameters

Name	Type	Default
`pressure`	`float`	`1.0`
`temperature`	`float`	`1.0`
`interval`	`int`	`25`
`max_log_volume_scale`	`float`	`0.02`
`seed`	`int \| None`	`11`
`mode`	`str`	`'isotropic'`
`axes`	`tuple[bool, bool, bool]`	`(True, True, True)`
`membrane_plane`	`str`	`'xy'`
`normal_axis`	`str`	`'z'`

`NPTResult`

class NPTResult
    def __init__(production: NVTResult, final_state: SimulationState, final_cell: Cell, cell_lengths: mx.array, cell_matrix: mx.array, volume: mx.array, target_pressure: float, barostat_attempts: int, barostat_accepted: int, barostat_metadata: dict[str, Any] = dict())

NPT production result with delegated NVT trajectory fields.

Parameters

Name	Type	Default
`production`	`NVTResult`
`final_state`	`SimulationState`
`final_cell`	`Cell`
`cell_lengths`	`mx.array`
`cell_matrix`	`mx.array`
`volume`	`mx.array`
`target_pressure`	`float`
`barostat_attempts`	`int`
`barostat_accepted`	`int`
`barostat_metadata`	`dict[str, Any]`	`dict()`

Properties

cell_history mx.array — Return the sampled cell-matrix history for pressure-coupled runs.

`NVEResult`

class NVEResult
    def __init__(sampled_positions: mx.array, sampled_velocities: mx.array, sampled_steps: mx.array, sampled_time: mx.array, diagnostic_steps: mx.array, diagnostic_time: mx.array, potential_energy: mx.array, kinetic_energy: mx.array, total_energy: mx.array, potential_energy_by_term: dict[str, mx.array], temperature: mx.array, virial_tensor: mx.array, pressure_tensor: mx.array, pressure: mx.array, pair_count: mx.array, rebuild_count: mx.array, constraint_max_error: mx.array, final_state: SimulationState, nonbonded_report: dict[str, int | float | str | None] = dict(), runtime_sync_report: dict[str, int | float] = dict())

Sparse trajectory and per-step diagnostics from an NVE simulation.

Parameters

Name	Type	Default
`sampled_positions`	`mx.array`
`sampled_velocities`	`mx.array`
`sampled_steps`	`mx.array`
`sampled_time`	`mx.array`
`diagnostic_steps`	`mx.array`
`diagnostic_time`	`mx.array`
`potential_energy`	`mx.array`
`kinetic_energy`	`mx.array`
`total_energy`	`mx.array`
`potential_energy_by_term`	`dict[str, mx.array]`
`temperature`	`mx.array`
`virial_tensor`	`mx.array`
`pressure_tensor`	`mx.array`
`pressure`	`mx.array`
`pair_count`	`mx.array`
`rebuild_count`	`mx.array`
`constraint_max_error`	`mx.array`
`final_state`	`SimulationState`
`nonbonded_report`	`dict[str, int \| float \| str \| None]`	`dict()`
`runtime_sync_report`	`dict[str, int \| float]`	`dict()`

Properties

energy_drift mx.array — Total energy minus the initial total energy for each diagnostic step.
max_energy_drift mx.array — Maximum absolute total-energy drift over the run.
relative_energy_drift mx.array — Energy drift normalized by the absolute initial total energy.

`NVTResult`

class NVTResult
    def __init__(sampled_positions: mx.array, sampled_velocities: mx.array, sampled_steps: mx.array, sampled_time: mx.array, diagnostic_steps: mx.array, diagnostic_time: mx.array, potential_energy: mx.array, kinetic_energy: mx.array, total_energy: mx.array, potential_energy_by_term: dict[str, mx.array], temperature: mx.array, virial_tensor: mx.array, pressure_tensor: mx.array, pressure: mx.array, pair_count: mx.array, rebuild_count: mx.array, constraint_max_error: mx.array, final_state: SimulationState, target_temperature: float, thermostat_metadata: dict[str, Any] = dict(), nonbonded_report: dict[str, int | float | str | None] = dict(), runtime_sync_report: dict[str, int | float] = dict())

Sparse trajectory and per-step diagnostics from an NVT simulation.

Parameters

Name	Type	Default
`sampled_positions`	`mx.array`
`sampled_velocities`	`mx.array`
`sampled_steps`	`mx.array`
`sampled_time`	`mx.array`
`diagnostic_steps`	`mx.array`
`diagnostic_time`	`mx.array`
`potential_energy`	`mx.array`
`kinetic_energy`	`mx.array`
`total_energy`	`mx.array`
`potential_energy_by_term`	`dict[str, mx.array]`
`temperature`	`mx.array`
`virial_tensor`	`mx.array`
`pressure_tensor`	`mx.array`
`pressure`	`mx.array`
`pair_count`	`mx.array`
`rebuild_count`	`mx.array`
`constraint_max_error`	`mx.array`
`final_state`	`SimulationState`
`target_temperature`	`float`
`thermostat_metadata`	`dict[str, Any]`	`dict()`
`nonbonded_report`	`dict[str, int \| float \| str \| None]`	`dict()`
`runtime_sync_report`	`dict[str, int \| float]`	`dict()`

Properties

temperature_error mx.array — Instantaneous temperature minus the target thermostat temperature.

`NoseHooverThermostat`

class NoseHooverThermostat
    def __init__(temperature: float = 1.0, relaxation_time: float = 0.1, thermal_mass: float | None = None, chain_position: float = 0.0, chain_velocity: float = 0.0)

Deterministic single-variable Nose-Hoover thermostat parameters.

Parameters

Name	Type	Default
`temperature`	`float`	`1.0`
`relaxation_time`	`float`	`0.1`
`thermal_mass`	`float \| None`	`None`
`chain_position`	`float`	`0.0`
`chain_velocity`	`float`	`0.0`

`ReporterEvent`

class ReporterEvent
    def __init__(ensemble: str, event_type: str, step: int, time: float, state: SimulationState, potential_energy: mx.array | None = None, kinetic_energy: mx.array | None = None, total_energy: mx.array | None = None, temperature: mx.array | None = None, energy_by_term: dict[str, mx.array] = dict(), virial_tensor: mx.array | None = None, pressure_tensor: mx.array | None = None, pressure: mx.array | None = None, pair_count: int | mx.array | None = None, rebuild_count: int | mx.array | None = None, constraint_max_error: mx.array | None = None, thermostat: dict[str, Any] = dict(), barostat: dict[str, Any] = dict())

State exposed to runtime reporter callbacks.

Parameters

Name	Type	Default
`ensemble`	`str`
`event_type`	`str`
`step`	`int`
`time`	`float`
`state`	`SimulationState`
`potential_energy`	`mx.array \| None`	`None`
`kinetic_energy`	`mx.array \| None`	`None`
`total_energy`	`mx.array \| None`	`None`
`temperature`	`mx.array \| None`	`None`
`energy_by_term`	`dict[str, mx.array]`	`dict()`
`virial_tensor`	`mx.array \| None`	`None`
`pressure_tensor`	`mx.array \| None`	`None`
`pressure`	`mx.array \| None`	`None`
`pair_count`	`int \| mx.array \| None`	`None`
`rebuild_count`	`int \| mx.array \| None`	`None`
`constraint_max_error`	`mx.array \| None`	`None`
`thermostat`	`dict[str, Any]`	`dict()`
`barostat`	`dict[str, Any]`	`dict()`

`RuntimeReporter`

class RuntimeReporter(Protocol)

Callable observer for sampled frames and diagnostic state.

`SimulationConfig`

class SimulationConfig
    def __init__(dt: float = 0.005, steps: int = 100, sample_interval: int = 1, kinetic_energy_scale: float = 1.0, force_to_acceleration_scale: float = 1.0, boltzmann_constant: float = 1.0, evaluation_interval: int = 25, diagnostic_interval: int = 1, compile_force_evaluator: bool = False, pressure_diagnostics: bool = True, initial_step: int = 0, initial_time: float = 0.0, virtual_sites: VirtualSiteManager | None = None, block_size: int = 1)

Configuration for molecular dynamics.

Parameters

Name	Type	Default
`dt`	`float`	`0.005`
`steps`	`int`	`100`
`sample_interval`	`int`	`1`
`kinetic_energy_scale`	`float`	`1.0`
`force_to_acceleration_scale`	`float`	`1.0`
`boltzmann_constant`	`float`	`1.0`
`evaluation_interval`	`int`	`25`
`diagnostic_interval`	`int`	`1`
`compile_force_evaluator`	`bool`	`False`
`pressure_diagnostics`	`bool`	`True`
`initial_step`	`int`	`0`
`initial_time`	`float`	`0.0`
`virtual_sites`	`VirtualSiteManager \| None`	`None`
`block_size`	`int`	`1`

`SimulationResult`

class SimulationResult
    def __init__(positions: mx.array, velocities: mx.array, potential_energy: mx.array, kinetic_energy: mx.array, total_energy: mx.array, temperature: mx.array)

Trajectory and diagnostics from an MD run.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`
`velocities`	`mx.array`
`potential_energy`	`mx.array`
`kinetic_energy`	`mx.array`
`total_energy`	`mx.array`
`temperature`	`mx.array`

`SimulationState`

class SimulationState
    def __init__(positions: mx.array, velocities: mx.array, masses: mx.array, forces: mx.array, step: int = 0, time: float = 0.0)

Current NVE simulation state.

Parameters

Name	Type	Default
`positions`	`mx.array`
`velocities`	`mx.array`
`masses`	`mx.array`
`forces`	`mx.array`
`step`	`int`	`0`
`time`	`float`	`0.0`

`StepState`

class StepState
    def __init__(positions: mx.array, velocities: mx.array, forces: mx.array, potential_energy: mx.array, kinetic_energy: mx.array)

Single MD state.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`
`velocities`	`mx.array`
`forces`	`mx.array`
`potential_energy`	`mx.array`
`kinetic_energy`	`mx.array`

Properties

total_energy mx.array — Potential plus kinetic energy.

`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

`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) -> mx.array

Return a configurational virial diagnostic after explicit support validation.

Periodic orthorhombic cells use diagonal finite differences in cell strain with fractional coordinates held fixed. Off-diagonal strain is not part of this Slice 8 diagnostic convention and is reported as zero.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`forces`	`mx.array`		Forces used for the non-periodic fallback, shape `(n_particles, 3)`.
`force_terms`	`tuple[ForceTerm, ...]`		Force terms whose energy is finite-differenced under cell strain; each must support virial diagnostics.
`cell`	`Cell \| None`		Periodic cell; `None` returns the non-periodic `virial_tensor`.
`pairs`	`object \| None`		Optional neighbor/pair structure forwarded to the force terms; `None` lets each term build its own.
`virtual_sites`	`VirtualSiteManager \| None`	`None`	Optional virtual-site manager applied before energy evaluation. Defaults to `None`.
`strain_epsilon`	`float`	`0.001`	Half-width of the symmetric cell-strain finite difference. Defaults to `1e-3`.

Returns

mx.array — The (3, 3) diagonal virial tensor (off-diagonal entries are zero by this diagnostic’s convention).

Raises

ValueError — If a force term lacks virial support, strain_epsilon is non-positive, or the cell volume is non-positive.

`instantaneous_temperature`

def instantaneous_temperature(velocities: mx.array, masses: mx.array, *, dof: int | None = None, kinetic_energy_scale: float = 1.0, boltzmann_constant: float = 1.0) -> mx.array

Return the instantaneous temperature from the kinetic energy.

Parameters

Name	Type	Default	Description
`velocities`	`mx.array`		Per-particle velocities, shape `(n_particles, 3)`.
`masses`	`mx.array`		Per-particle masses, shape `(n_particles,)`.
`dof`	`int \| None`	`None`	Degrees of freedom in the equipartition denominator; `None` uses `velocities.size` (no constraints removed). Defaults to `None`.
`kinetic_energy_scale`	`float`	`1.0`	Energy-unit factor forwarded to `kinetic_energy`. Defaults to `1.0`.
`boltzmann_constant`	`float`	`1.0`	Boltzmann constant in the configured units. Defaults to `1.0` (reduced units).

Returns

mx.array — Scalar temperature 2·E_kin / (dof · k_B).

`kinetic_energy`

def kinetic_energy(velocities: mx.array, masses: mx.array, *, kinetic_energy_scale: float = 1.0) -> mx.array

Return the total kinetic energy in the configured unit system.

Parameters

Name	Type	Default	Description
`velocities`	`mx.array`		Per-particle velocities, shape `(n_particles, 3)`.
`masses`	`mx.array`		Per-particle masses, shape `(n_particles,)`.
`kinetic_energy_scale`	`float`	`1.0`	Multiplicative factor converting the raw kinetic quantity into the configured energy unit. Defaults to `1.0`.

Returns

mx.array — Scalar kinetic energy ½ · scale · Σ_i m_i |v_i|².

`kinetic_pressure_tensor`

def kinetic_pressure_tensor(velocities: mx.array, masses: mx.array, *, kinetic_energy_scale: float = 1.0) -> mx.array

Return the kinetic (momentum-flux) tensor in the configured units.

Parameters

Name	Type	Default	Description
`velocities`	`mx.array`		Per-particle velocities, shape `(n_particles, 3)`.
`masses`	`mx.array`		Per-particle masses, shape `(n_particles,)`.
`kinetic_energy_scale`	`float`	`1.0`	Multiplicative factor converting the raw kinetic quantity into the configured energy unit. Defaults to `1.0`.

Returns

mx.array — The (3, 3) tensor scale · Σ_i m_i v_i ⊗ v_i.

Raises

ValueError — If velocities is not (n_particles, 3) or masses is not (n_particles,).

`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) -> 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`.

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: LangevinThermostat | None = None, barostat: MonteCarloBarostat | None = None, constraints: DistanceConstraints | None = None, reporters: RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None = None) -> NPTResult

Run NVT dynamics followed by a Monte Carlo pressure-coupling attempt.

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`	`LangevinThermostat \| 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`.
`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

NPTResult — An NPTResult with the NVT trajectory plus cell-volume history from the pressure-coupling attempts.

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: LangevinThermostat | 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`	`LangevinThermostat \| 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_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_tensor`

def virial_tensor(positions: mx.array, forces: mx.array) -> mx.array

Return the non-periodic configurational virial tensor.

Parameters

Name	Type	Default	Description
`positions`	`mx.array`		Particle coordinates, shape `(n_particles, 3)`.
`forces`	`mx.array`		Forces on each particle, shape `(n_particles, 3)`.

Returns

mx.array — The (3, 3) virial tensor positionsᵀ · forces.

Raises

ValueError — If positions and forces do not both have shape (n_particles, 3).