Optimization

Cartesian and kinematic minimization entry points.

Minimizers and scoring-function optimization adapters.

class tmol.optimization.CartesianSfxnNetwork(score_function: ScoreFunction, pose_stack: PoseStack, coord_mask=None, cuda_graph: bool | str = False)

Bases: Module

Differentiable score network over selected Cartesian coordinates.

forward() → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class tmol.optimization.KinForestSfxnNetwork(score_function: ScoreFunction, pose_stack: PoseStack, kin_module: PoseStackKinematicsModule, dof_mask=None, kin_dtype=torch.float32)

Bases: Module

Differentiable score network over selected kinematic degrees of freedom.

forward() → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

tmol.optimization.build_kinforest_network(pose_stack: PoseStack, sfxn: ScoreFunction, ff: FoldForest, mm: MoveMap, verbose: bool = False, kin_dtype: torch.dtype = torch.float32) → KinForestSfxnNetwork

Build a differentiable kinematic scoring network for a pose stack.

Parameters:

• pose_stack – Structures and topology to score.

• sfxn – Score function to render against pose_stack.

• ff – Fold forest defining the kinematic tree.

• mm – Internal-coordinate degrees of freedom allowed to move.

• verbose – Print synchronized setup timings.

• kin_dtype – Floating-point dtype for kinematic degrees of freedom.

Returns:

A network mapping movable kinematic degrees of freedom to pose energies.

tmol.optimization.run_cart_min(pose_stack: ~tmol.pose._pose_stack.PoseStack, sfxn: ~tmol.score._score_function.ScoreFunction, coord_mask: ~tmol.types._torch.Tensor[slice(None, None, None), slice(None, None, None)] | None = None, optimizer_cls: type[~torch.optim.optimizer.Optimizer] = <class 'tmol.optimization._lbfgs_armijo.LBFGS_Armijo'>, optimizer_kwargs: dict[str, object] | None = None, verbose: bool = False, cuda_graph: bool = False) → PoseStack

Run minimization on a PoseStack in Cartesian coordinate space.

Builds a CartesianSfxnNetwork and delegates to run_min().

Parameters:

• pose_stack – Structures and coordinates to minimize.

• sfxn – Score function defining the objective.

• coord_mask – Movable atoms shaped [pose, atom]; None moves all.

• optimizer_cls – Closure-based PyTorch optimizer class.

• optimizer_kwargs – Optional optimizer constructor arguments.

• verbose – Print synchronized setup and minimization timings.

• cuda_graph – Capture the fixed-shape CUDA scoring forward/backward path.

Returns:

A new pose stack containing the minimized coordinates.

tmol.optimization.run_kin_min(pose_stack: ~tmol.pose._pose_stack.PoseStack, sfxn: ~tmol.score._score_function.ScoreFunction, ff: ~tmol.kinematics._fold_forest.FoldForest, mm: ~tmol.kinematics._move_map.MoveMap, optimizer_cls: type[~torch.optim.optimizer.Optimizer] = <class 'tmol.optimization._lbfgs_armijo.LBFGS_Armijo'>, optimizer_kwargs: dict[str, object] | None = None, verbose: bool = False, kin_dtype: ~torch.dtype = torch.float32) → PoseStack

Run minimization on a PoseStack in internal DOF space.

Builds a KinForestSfxnNetwork and delegates to run_min().

Parameters:

• pose_stack – Structures and coordinates to minimize.

• sfxn – Score function defining the objective.

• ff – Fold forest defining the kinematic tree.

• mm – Internal-coordinate degrees of freedom allowed to move.

• optimizer_cls – Closure-based PyTorch optimizer class.

• optimizer_kwargs – Optional optimizer constructor arguments.

• verbose – Print synchronized setup and minimization timings.

• kin_dtype – Floating-point dtype for kinematic degrees of freedom.

Returns:

A new pose stack containing the minimized coordinates.

tmol.optimization.run_min(sfxn_module: CartesianSfxnNetwork | KinForestSfxnNetwork, optimizer_cls: type[torch.optim.Optimizer] = <class 'tmol.optimization._lbfgs_armijo.LBFGS_Armijo'>, optimizer_kwargs: dict[str, object] | None = None, verbose: bool = False, per_pose: bool = True) → PoseStack

Run minimization on any sfxn module (Cartesian or KinForest).

The sfxn_module must be a torch.nn.Module whose forward() returns per-pose energies and which provides a pose_stack_from_dofs() method to extract the optimized PoseStack.

Parameters:

• sfxn_module – A CartesianSfxnNetwork, KinForestSfxnNetwork, or any nn.Module with a compatible interface.

• optimizer_cls – A torch.optim.Optimizer class. Must support a closure-based step() call (e.g. LBFGS_Armijo, torch LBFGS).

• optimizer_kwargs – Dict of keyword arguments passed to the optimizer constructor.

• verbose – Print timing information.

• per_pose – Give each pose its own inverse-Hessian estimate and convergence test, so that minimizing a stack matches minimizing its poses one at a time. Ignored by optimizers that do not support it.

Returns:

A new PoseStack with optimized coordinates.