replica_exchange

Replica exchange molecular dynamics driver.

import mlx_atomistic.replica_exchange

Classes

ReplicaExchangeResult

class ReplicaExchangeResult
    def __init__(replica_states: tuple[SimulationState, ...], temperatures: tuple[float, ...], lambdas: tuple[float | None, ...], swap_attempts: tuple[SwapAttempt, ...], accepted_swaps: int, energy_history: mx.array, state_index_history: mx.array)

Final states and diagnostics from a replica exchange run.

Parameters

Name	Type	Default	Description
replica_states	tuple[SimulationState, ...]
temperatures	tuple[float, ...]
lambdas	tuple[float | None, ...]
swap_attempts	tuple[SwapAttempt, ...]
accepted_swaps	int
energy_history	mx.array
state_index_history	mx.array

SwapAttempt

class SwapAttempt
    def __init__(step: int, left: int, right: int, probability: float, accepted: bool, log_acceptance: float, random_value: float)

One attempted exchange between adjacent thermodynamic states.

Parameters

Name	Type	Default	Description
step	int
left	int
right	int
probability	float
accepted	bool
log_acceptance	float
random_value	float

Functions

metropolis_acceptance_probability

def metropolis_acceptance_probability(energy_left_at_left: float, energy_right_at_right: float, energy_left_at_right: float, energy_right_at_left: float, beta_left: float, beta_right: float) -> tuple[float, float]

Return generalized replica-exchange Metropolis probability and log value.

Parameters

Name	Type	Default	Description
energy_left_at_left	float
energy_right_at_right	float
energy_left_at_right	float
energy_right_at_left	float
beta_left	float
beta_right	float

Returns

• tuple[float, float]

replica_exchange_initial_state

def replica_exchange_initial_state(positions, velocities, masses) -> SimulationState

Build an initial replica state with zero placeholder forces.

Parameters

Name	Type	Default	Description
positions
velocities
masses

Returns

• SimulationState

replica_exchange_metadata

def replica_exchange_metadata(metadata: dict[str, Any]) -> dict[str, Any]

Return normalized replica-exchange artifact metadata.

Parameters

Name	Type	Default	Description
metadata	dict[str, Any]

Returns

• dict[str, Any]

simulate_replica_exchange

def simulate_replica_exchange(initial_states: list[SimulationState] | tuple[SimulationState, ...], force_terms: ForceTerm | list[ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]] | tuple[ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...], ...], *, temperatures: list[float] | tuple[float, ...], lambdas: list[float | None] | tuple[float | None, ...] | None = None, cell: Cell | None = None, config: SimulationConfig | None = None, neighbor_manager: NeighborListManager | None = None, constraints: DistanceConstraints | None = None, reporters: RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None = None, swap_interval: int = 10, thermostat_friction: float = 1.0, seed: int | None = None) -> ReplicaExchangeResult

Run adjacent-pair temperature or Hamiltonian replica exchange.

force_terms may be one shared force-term set or one explicit nested force-term set per replica, e.g. [(term_a,), (term_b,)]. When lambdas is provided with a shared lambda-scalable nonbonded term, per-replica soft-core nonbonded terms are created for Hamiltonian exchange.

Parameters

Name	Type	Default	Description
initial_states	list[SimulationState] | tuple[SimulationState, ...]
force_terms	ForceTerm | list[ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...]] | tuple[ForceTerm | list[ForceTerm] | tuple[ForceTerm, ...], ...]
temperatures	list[float] | tuple[float, ...]
lambdas	list[float | None] | tuple[float | None, ...] | None	None
cell	Cell | None	None
config	SimulationConfig | None	None
neighbor_manager	NeighborListManager | None	None
constraints	DistanceConstraints | None	None
reporters	RuntimeReporter | list[RuntimeReporter] | tuple[RuntimeReporter, ...] | None	None
swap_interval	int	10
thermostat_friction	float	1.0
seed	int | None	None

Returns

• ReplicaExchangeResult

temperature_exchange_probability

def temperature_exchange_probability(energy_left: float, energy_right: float, beta_left: float, beta_right: float) -> tuple[float, float]

Return the temperature-exchange Metropolis probability.

Parameters

Name	Type	Default	Description
energy_left	float
energy_right	float
beta_left	float
beta_right	float

Returns

• tuple[float, float]