Optimization

Cartesian and kinematic minimization entry points.

Minimizers and scoring-function optimization adapters.

class tmol.optimization.CartesianMinimizer(cuda_graph: bool = False)

Bases: object

Reusable Cartesian minimizer for pose stacks with stable topology.

Rendering a score function builds topology-dependent tensors and, when requested, captures a CUDA graph. Those are setup costs rather than part of the numerical minimization. This class retains that rendered network across calls and resets only its coordinates when the score function, topology, shape, and coordinate mask remain compatible.

Create one instance per reusable workflow. A topology or score-function change is detected automatically and safely rebuilds the network.

Parameters:

cuda_graph – Capture and replay the fixed-shape CUDA forward/backward score path. This has a one-time setup and memory cost and requires CUDA, so it is most useful when the instance will be reused.

release_retained_state() → None

Release topology-specific scoring and optimizer state.

The returned pose from __call__() owns its coordinates, so callers may drop the rendered scorer, CUDA graph pools, gradients, and L-BFGS history before entering another memory-intensive phase.

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: PoseStack, sfxn: ScoreFunction, coord_mask: Tensor[slice(None, None, None), slice(None, None, None)] | None=None, optimizer_cls: type[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: PoseStack, sfxn: ScoreFunction, ff: FoldForest, mm: MoveMap, optimizer_cls: type[Optimizer] = <class 'tmol.optimization._lbfgs_armijo.LBFGS_Armijo'>, optimizer_kwargs: dict[str, object] | None=None, verbose: bool = False, kin_dtype: 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.