The tmol How To guide notebook

tmol/__init__.py

@@ -39,6 +39,7 @@ def include_paths():
     default_canonical_ordering,
     default_packed_block_types,
     pose_stack_from_canonical_form,
+    pose_stack_from_biotite,
     canonical_form_from_openfold,
     canonical_ordering_for_openfold,
     packed_block_types_for_openfold,
@@ -65,9 +66,11 @@ def include_paths():
     run_kin_min,
     run_min,
 )
+
 from tmol.pose import (
     PackedBlockTypes,
     PoseStack,
+    PoseStackBuilder,
     ConstraintSet,
     get_named_torsions,
     get_torsion_names,
@@ -79,9 +82,29 @@ def include_paths():
 )
 from tmol.score.constraint import (
     ConstraintEnergyTerm,
+    constrain_all_ca,
     create_mainchain_coordinate_constraints,
 )
-from tmol.relax import fast_relax
+from tmol.pack import (
+    pack_rotamers,
+    PackerPalette,
+    PackerTask,
+)
+from tmol.pack.rotamer import (
+    FixedAAChiSampler,
+    IncludeCurrentSampler,
+    OptHSampler,
+)
+from tmol.pack.rotamer.dunbrack import (
+    create_dunbrack_sampler_from_database,
+    DunbrackChiSampler,
+)
+
+from tmol.relax import (
+    fast_relax,
+    kin_fast_relax,
+    cartesian_fast_relax,
+)
 
 try:
     __version__ = version("tmol")
@@ -100,6 +123,7 @@ def include_paths():
     "PackedBlockTypes",
     "ParameterDatabase",
     "PoseStack",
+    "PoseStackBuilder",
     "ScoreFunction",
     "ScoreType",
     "atom_records_from_pose_stack",
@@ -110,6 +134,7 @@ def include_paths():
     "canonical_form_from_rosettafold2",
     "canonical_ordering_for_openfold",
     "canonical_ordering_for_rosettafold2",
+    "constrain_all_ca",
     "create_mainchain_coordinate_constraints",
     "default_canonical_ordering",
     "default_packed_block_types",
@@ -121,6 +146,7 @@ def include_paths():
     "one2three",
     "packed_block_types_for_openfold",
     "packed_block_types_for_rosettafold2",
+    "pose_stack_from_biotite",
     "pose_stack_from_canonical_form",
     "pose_stack_from_openfold",
     "pose_stack_from_pdb",
@@ -135,4 +161,14 @@ def include_paths():
     "three2one",
     "view",
     "write_pose_stack_pdb",
+    "pack_rotamers",
+    "PackerPalette",
+    "PackerTask",
+    "FixedAAChiSampler",
+    "IncludeCurrentSampler",
+    "OptHSampler",
+    "create_dunbrack_sampler_from_database",
+    "DunbrackChiSampler",
+    "kin_fast_relax",
+    "cartesian_fast_relax",
 ]

tmol/pose/_pose_stack.py

@@ -1,6 +1,6 @@
 import attr
 import torch
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING, Union
 
 from tmol.types import Tensor
 from tmol.chemical import RefinedResidueType
@@ -288,6 +288,12 @@ def split(self, index) -> "PoseStack":
             ),
         )
 
+    def to(self, dtype=Union[torch.float32, torch.float64]) -> "PoseStack":
+        """Create a new PoseStack with the dtype of the coords tensor changed to the requested dtype."""
+        if self.coords.dtype == dtype:
+            return self
+        return attr.evolve(self, coords=self.coords.to(dtype=dtype))
+
     def expand_coords(self):
         """Load the coordinates into a 4D tensor:
         n_poses x max_n_blocks x max_n_atoms_per_block x 3

tmol/relax/__init__.py

@@ -1,14 +1,19 @@
 from ._fast_relax import (  # noqa: F401
     DEFAULT_RELAX_SCHEDULE,
-    _default_cart_min_fn,
+    default_cart_min_fn,
+    default_kin_min_fn,
     accept_best,
     fast_relax,
+    kin_fast_relax,
+    cartesian_fast_relax,
     relax_pack_min_step,
 )
 
 __all__ = [
     "DEFAULT_RELAX_SCHEDULE",
     "accept_best",
     "fast_relax",
+    "kin_fast_relax",
+    "cartesian_fast_relax",
     "relax_pack_min_step",
 ]

tmol/relax/_fast_relax.py

@@ -120,7 +120,7 @@ def constraint_fraction(step_index):
     return normalized
 
 
-def _default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose):
+def default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose):
     """Default minimization function: kinematic (torsion-space) LBFGS."""
     return run_kin_min(
         pose_stack,
@@ -132,7 +132,7 @@ def _default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose):
     )
 
 
-def _default_cart_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose):
+def default_cart_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose):
     """Default Cartesian minimization function for use as fast_relax min_fn.
 
     Extracts ``coord_mask`` from ``move_map`` if it is a
@@ -159,7 +159,7 @@ def fast_relax(  # noqa: C901
     sfxn: ScoreFunction,
     packer_pallete: PackerPalette,
     move_map: Union[MoveMap, CartesianMoveMap],
-    fold_forest: FoldForest,
+    fold_forest: Optional[FoldForest],
     *,
     task_operations=None,
     num_repeats=2,
@@ -176,7 +176,9 @@ def fast_relax(  # noqa: C901
     followed by an accept-to-best check.
 
     Args:
-        pose_stack: The input poses to relax.
+        pose_stack: The input poses to relax. Relax will use the precision of
+            the input coords tensor during minimization, but will only use
+            torch.float32 precision for packing.
         sfxn: Score function used for packing and minimization. If you wish
             to use constraints during relax, then the weight on the "constraint"
             score type must already have a non-zero value.
@@ -251,7 +253,7 @@ def my_min(ps, sfxn, *, fold_forest, move_map, **kw):
         The relaxed PoseStack (best-scoring across all repeats).
     """
     if min_fn is None:
-        min_fn = _default_cart_min_fn
+        min_fn = default_cart_min_fn
     if schedule is None:
         schedule = DEFAULT_RELAX_SCHEDULE
 
@@ -352,6 +354,11 @@ def relax_pack_min_step(
     min_fn,
     verbose,
 ):
+    """Perform a single pack-min step of the FastRelax protocol.
+
+    Convert the PoseStack to float32 for packing, then restore
+    it to the input dtype afterwards."""
+    input_pose_dtype = pose_stack.coords.dtype
 
     if verbose and torch.cuda.is_available():
         torch.cuda.synchronize()
@@ -369,7 +376,12 @@ def relax_pack_min_step(
     if verbose and torch.cuda.is_available():
         torch.cuda.synchronize()
     end_time1 = time.perf_counter()
+
+    # convert pose_stack to float32 for packing, and restore it
+    # to the input dtype afterwards
+    pose_stack = pose_stack.to(torch.float32)
     packed_pose_stack = pack_rotamers(pose_stack, sfxn, task, verbose)
+    packed_pose_stack = packed_pose_stack.to(dtype=input_pose_dtype)
 
     sfxn.set_weight(ScoreType.fa_ljrep, fa_rep_min_weight)
     if verbose:
@@ -442,3 +454,68 @@ def select_better(tensor_name):
         return new_best_pose_stack, new_best_pose_score
     else:  # no change
         return best_pose_stack, best_pose_score
+
+
+def kin_fast_relax(
+    pose_stack: PoseStack,
+    sfxn: ScoreFunction,
+    packer_pallete: PackerPalette,
+    move_map: MoveMap,
+    fold_forest: FoldForest,
+    *,
+    task_operations=None,
+    num_repeats=2,
+    ramp_constraints: Optional[bool] = None,  # default True
+    schedule=None,
+    min_fn=default_kin_min_fn,
+    verbose: bool = False,
+):
+    """Run the FastRelax protocol using kinematic (torsion-space) minimization.
+
+    See documentation for fast_relax.
+    """
+    return fast_relax(
+        pose_stack,
+        sfxn,
+        packer_pallete,
+        move_map,
+        fold_forest,
+        task_operations=task_operations,
+        num_repeats=num_repeats,
+        ramp_constraints=ramp_constraints,
+        schedule=schedule,
+        min_fn=min_fn or default_kin_min_fn,
+        verbose=verbose,
+    )
+
+
+def cartesian_fast_relax(
+    pose_stack: PoseStack,
+    sfxn: ScoreFunction,
+    packer_pallete: PackerPalette,
+    move_map: CartesianMoveMap,
+    *,
+    task_operations=None,
+    num_repeats=2,
+    ramp_constraints: Optional[bool] = None,  # default True
+    schedule=None,
+    min_fn=default_cart_min_fn,
+    verbose: bool = False,
+):
+    """Run the FastRelax protocol using Cartesian (coordinate-space) minimization.
+
+    See documentation for fast_relax.
+    """
+    return fast_relax(
+        pose_stack,
+        sfxn,
+        packer_pallete,
+        move_map,
+        None,
+        task_operations=task_operations,
+        num_repeats=num_repeats,
+        ramp_constraints=ramp_constraints,
+        schedule=schedule,
+        min_fn=min_fn or default_kin_min_fn,
+        verbose=verbose,
+    )