torchref.refinement.optimizers.langevin_sa module

class torchref.refinement.optimizers.langevin_sa.LangevinSA(params, dt=0.01, friction=10.0, T_initial=2500.0, T_final=0.01, total_steps=1000, cooling_schedule='exponential', adaptive_masses=True, mass_beta=0.999, mass_eps=1e-08, gradient_clip=None, max_step_size=0.1)

Bases: Optimizer

BAOAB Langevin dynamics integrator with simulated annealing.

Implements the BAOAB splitting scheme (Leimkuhler & Matthews, 2013) for gradient-guided exploration with thermodynamically correct noise. One gradient evaluation per step via staggered B steps.

Adaptive masses from EMA of squared gradients provide automatic scale invariance across all parameter types (xyz, B-factors, occupancies, torsions, etc.).

Call calibrate() before the main loop to probe parameter stiffness and warm up the adaptive masses without moving the structure.

Args:

params: Iterable of parameters or param groups. dt: Integration timestep. friction: Friction coefficient gamma. Controls thermalization speed. T_initial: Starting temperature. T_final: Final temperature. total_steps: Total number of annealing steps. cooling_schedule: ‘exponential’ or ‘linear’. adaptive_masses: Use EMA of grad² as per-element masses. mass_beta: EMA decay for adaptive masses. mass_eps: Floor for adaptive masses (numerical stability). gradient_clip: Optional max gradient norm (per-parameter). max_step_size: Maximum displacement per element per full step.

Velocities are clamped so |v * dt| <= max_step_size.

__init__(params, dt=0.01, friction=10.0, T_initial=2500.0, T_final=0.01, total_steps=1000, cooling_schedule='exponential', adaptive_masses=True, mass_beta=0.999, mass_eps=1e-08, gradient_clip=None, max_step_size=0.1)

property temperature

Current temperature from the annealing schedule.

property current_step

property total_steps

property kinetic_energy

Sum of 0.5 * m * v^2 over all parameters (diagnostic).

calibrate(closure, n_steps=10)

Probe parameter stiffness over n_steps, then rollback.

Runs small random perturbations to collect gradient statistics, sets the adaptive masses from the observed grad², then restores all parameters to their original values and initialises velocities from Maxwell-Boltzmann with correctly scaled masses.

Args:

closure: Same closure as for step() — must zero_grad,

compute loss, call backward, and return loss.

n_steps: Number of probing steps.

step(closure)

Perform one BAOAB Langevin dynamics step.

Tracks the best-loss configuration and rolls back to it when the loss exceeds loss_rollback_factor times the best loss seen so far. This prevents the dynamics from permanently damaging the structure while still allowing uphill exploration.

Args:

closure: A callable that re-evaluates the model and returns the

loss. The closure must call loss.backward() before returning.

Returns:

The loss value from the closure evaluation.