Relax#

FastRelax alternates packing and minimization under a configurable weight schedule.

Packing and minimization protocols for structural relaxation.

tmol.relax.accept_best(sfxn: ScoreFunction, best_pose_stack: PoseStack, best_pose_score: Tensor, candidate_pose_stack: PoseStack, verbose: bool = False) → tuple[PoseStack, Tensor][source]#

Keep the lower-scoring conformation independently for each pose.

Parameters:

  • sfxn – Score function used for comparison.

  • best_pose_stack – Best poses from previous repeats.

  • best_pose_score – Best scores shaped [n_poses].

  • candidate_pose_stack – Newly minimized poses.

  • verbose – Print accepted scores.

Returns:

Updated best poses and scores shaped [n_poses].

tmol.relax.fast_relax(pose_stack: PoseStack, sfxn: ScoreFunction, packer_pallete: PackerPalette, move_map: MoveMap | CartesianMoveMap, fold_forest: FoldForest, *, task_operations: Sequence[Callable[[PackerTask], None]] | None = None, num_repeats: int = 2, ramp_constraints: bool | None = None, schedule: Sequence[float | int | dict[str, float]] | None = None, min_fn: RelaxMinimizer | None = None, cuda_graph: bool | None = None, release_between_stages: bool = True, verbose: bool = False) → PoseStack[source]#

Relax poses through repeated side-chain packing and minimization.

Each repeat applies the MonomerRelax2019 fa_rep ramp and retains the lowest-scoring conformation independently for every pose in the batch.

Parameters:

  • pose_stack – Input poses to relax.

  • sfxn – Packing and minimization score function. Constraints are active only when its constraint weight is nonzero.

  • packer_pallete – Residue types available to the packer.

  • move_map – Specifies which DOFs are free to move during minimization.

  • fold_forest – Fold forest defining the kinematic connectivity.

  • task_operations – In-place task configuration callbacks. By default, restrict to repacking with the scoring database’s Dunbrack, fixed-AA, current and nucleic-acid rotamers, sampling covalent groups jointly.

  • num_repeats – Number of complete pack-minimize ramps.

  • ramp_constraints – Ramp an active constraint weight to zero. Defaults to true; the input weight is restored after relaxation.

  • schedule – Numeric fa_rep fractions or dictionaries with separate pack, minimize, and optional constraint fractions. Defaults to DEFAULT_RELAX_SCHEDULE.

  • min_fn – Minimizer called with the pose, score function, fold forest, move map, and verbosity. If it provides release_retained_state(), FastRelax calls that hook after extracting each minimized pose so packing does not overlap retained scorer or optimizer storage. Defaults to Cartesian minimization.

  • release_between_stages – Drop the rendered scorer, CUDA graph and L-BFGS history after each stage’s minimization. On by default: a large system is bound by that retained state through the next packing phase, which is why the release exists. Turn it off when the setup cost dominates instead – a small pose on the default four-stage schedule otherwise pays four network constructions, and four graph captures, where one would do.

  • cuda_graph – Capture the default Cartesian minimizer’s repeated CUDA scoring path. By default, enable it automatically for CUDA poses containing DNA or RNA, where replay savings exceed capture setup. Pass True for repeated protein workloads where steady-state throughput matters more than first-call latency, or False to disable it. It cannot be combined with a custom min_fn.

  • verbose – Print timing information for each step.

Returns:

Best-scoring relaxed poses across all repeats.

tmol.relax.relax_pack_min_step(pose_stack: PoseStack, sfxn: ScoreFunction, fold_forest: FoldForest, move_map: MoveMap | CartesianMoveMap, packer_pallete: PackerPalette, fa_rep_pack_weight: float, fa_rep_min_weight: float, cst_weight: float, task_operations: Sequence[Callable[[PackerTask], None]], min_fn: RelaxMinimizer, verbose: bool, release_between_stages: bool = True) → PoseStack[source]#

Execute one weighted packing and minimization stage.

Parameters:

  • pose_stack – Current poses.

  • sfxn – Score function whose repulsive and constraint weights are updated.

  • fold_forest – Connectivity passed to the minimizer.

  • move_map – Movable degrees of freedom passed to the minimizer.

  • packer_pallete – Residue types available during packing.

  • fa_rep_pack_weight – Repulsive weight for packing.

  • fa_rep_min_weight – Repulsive weight for minimization.

  • cst_weight – Constraint weight for both operations.

  • task_operations – In-place task configuration callbacks.

  • min_fn – Minimization callable.

  • verbose – Print synchronized stage timings.

Returns:

The minimized pose stack.