torchref.refinement.optimizers.internal_coord_sa module

Internal Coordinate Simulated Annealing optimizer.

Implements batch Metropolis SA for internal coordinates with: - Auto-calibrated temperature from energy landscape sampling - Single scale parameter controlling perturbation magnitude - Best structure tracking throughout optimization - Torsion wrapping to [-pi, pi]

Usage:

model.use_internal_coordinates(n_aa_per_segment=3) ic = model.xyz

# Torsion-only refinement (recommended for conformational exploration): state, best_loss = optimize_internal_coord_sa(

state=state, params=[ic.torsions], scale=1.0, …

)

# All internal coordinates: state, best_loss = optimize_internal_coord_sa(

state=state, params=[ic.torsions, ic.bond_lengths, ic.angles,

ic.segment_positions, ic.segment_orientations],

scale=1.0, …

)

torchref.refinement.optimizers.internal_coord_sa.optimize_internal_coord_sa(state, params, scale=1.0, T_initial=None, T_final=0.01, n_steps=5000, cooling_schedule='exponential', verbose=0, callback=None)[source]

Universal Simulated Annealing for internal coordinates.

Perturbs the provided parameter tensors using batch Metropolis: all parameters are perturbed together, single accept/reject decision.

Parameters:

  • state (LossState) – Configured loss state with targets.

  • params (List[torch.Tensor]) –

    Parameter tensors to optimize. For internal coordinates: - Torsion-only: [ic.torsions] - All coords: [ic.torsions, ic.bond_lengths, ic.angles,

    ic.segment_positions, ic.segment_orientations]

  • scale (float) – Master scale for perturbations. Default 1.0. - Angle-like params: scale * 0.1 rad (~6 degrees at scale=1.0) - Position-like params: scale * 0.01 A

  • T_initial (float, optional) – Initial temperature. If None, auto-calibrated from energy landscape.

  • T_final (float) – Final temperature. Default 0.01.

  • n_steps (int) – Number of SA steps. Default 5000.

  • cooling_schedule (str) – ‘exponential’ or ‘linear’. Default ‘exponential’.

  • verbose (int) – Verbosity level.

  • callback (callable, optional) – Called after each step: callback(step, T, loss, best_loss).

Returns:

  • state (LossState) – Updated state.

  • best_loss (float) – Best (lowest) loss found during optimization.

Return type:

Tuple[LossState, float]

torchref.refinement.optimizers.internal_coord_sa.optimize_gradient_sa(loss_fn, params, scale=1.0, T_initial=None, T_final=0.01, n_steps=5000, cooling_schedule='exponential', per_parameter=True, verbose=0, callback=None)[source]

Gradient-based Simulated Annealing.

Uses gradient information for Metropolis acceptance criterion: - ΔE ≈ gradient · perturbation (first-order Taylor approximation) - Accept if g·δ < 0 (moved downhill) - Accept with probability exp(-g·δ / T) if g·δ > 0

Advantages over loss-based SA: - One gradient computation per step (vs one loss eval per perturbation) - Per-parameter acceptance allows partial moves - Gradient guides acceptance decisions

Parameters:

  • loss_fn (callable) – Loss function returning a differentiable scalar tensor.

  • params (List[torch.Tensor]) – Parameter tensors to optimize (must have requires_grad=True).

  • scale (float) – Perturbation scale. Default 1.0. - Angle-like: scale * 0.1 rad - Position-like: scale * 0.01 A

  • T_initial (float, optional) – Initial temperature. Auto-calibrated if None.

  • T_final (float) – Final temperature. Default 0.01.

  • n_steps (int) – Number of SA steps. Default 5000.

  • cooling_schedule (str) – ‘exponential’ or ‘linear’. Default ‘exponential’.

  • per_parameter (bool) – If True, accept/reject each parameter independently. If False, use batch acceptance (sum of g·δ for all params). Default True.

  • verbose (int) – Verbosity level.

  • callback (callable, optional) – Called after each step: callback(step, T, loss, accept_rate).

Returns:

  • best_loss (float) – Best loss found during optimization.

  • final_loss (float) – Loss at the end of optimization.

Return type:

Tuple[float, float]

torchref.refinement.optimizers.internal_coord_sa.refine_sa_lbfgs(model, loss_fn, n_sa_steps=5000, sa_scale=1.0, sa_T_initial=None, sa_T_final=0.01, sa_params=None, n_lbfgs_cycles=20, lbfgs_lr=1.0, lbfgs_max_iter=20, verbose=1, sa_callback=None, lbfgs_callback=None)[source]

Combined SA + LBFGS pipeline for internal coordinate refinement.

Pipeline: 1. SA phase: Global exploration using internal coordinates (torsions by default) 2. LBFGS phase: Local optimization to refine best structure

Parameters:

  • model (ModelFT) – Model with internal coordinates enabled (model.use_internal_coordinates()).

  • loss_fn (callable) – Loss function that returns a scalar tensor: loss = loss_fn()

  • n_sa_steps (int) – Number of SA steps. Default 5000.

  • sa_scale (float) – SA perturbation scale. Default 1.0.

  • sa_T_initial (float, optional) – Initial SA temperature. Auto-calibrated if None.

  • sa_T_final (float) – Final SA temperature. Default 0.01.

  • sa_params (List[str], optional) – Which internal coord params to use for SA. Options: ‘torsions’, ‘bond_lengths’, ‘angles’, ‘segment_positions’, ‘segment_orientations’ Default: [‘torsions’] (torsion-only exploration)

  • n_lbfgs_cycles (int) – Number of LBFGS optimization cycles. Default 20.

  • lbfgs_lr (float) – LBFGS learning rate. Default 1.0.

  • lbfgs_max_iter (int) – Max iterations per LBFGS cycle. Default 20.

  • verbose (int) – Verbosity level. Default 1.

  • sa_callback (callable, optional) – Callback for SA phase: callback(step, T, loss, best_loss)

  • lbfgs_callback (callable, optional) – Callback for LBFGS phase: callback(cycle, loss)

Returns:

‘initial_loss’: Loss before any optimization ‘sa_loss’: Loss after SA phase ‘final_loss’: Loss after LBFGS phase ‘sa_history’: SA history if callback provided ‘lbfgs_history’: LBFGS history if callback provided

Return type:

dict with keys