The tmol How To guide notebook

tmol/__init__.py

@@ -38,6 +38,9 @@ def include_paths():
 from tmol.io.pose_stack_construction import (
     pose_stack_from_canonical_form,
 )
+from tmol.io.pose_stack_from_biotite import (
+    pose_stack_from_biotite,
+)
 from tmol.io.pose_stack_from_openfold import (
     canonical_form_from_openfold,
     canonical_ordering_for_openfold,
@@ -65,6 +68,7 @@ def include_paths():
 )
 from tmol.pose.packed_block_types import PackedBlockTypes
 from tmol.pose.pose_stack import PoseStack
+from tmol.pose.pose_stack_builder import PoseStackBuilder
 from tmol.score import beta2016_score_function
 from tmol.score.score_function import ScoreFunction
 from tmol.score.score_types import ScoreType

tmol/io/pose_stack_from_biotite.py

@@ -17,8 +17,6 @@
 from tmol.pose.pdb_info import DEFAULT_ATOM_B_FACTOR, DEFAULT_ATOM_OCCUPANCY
 from tmol.utility.biotite_util import get_all_residue_positions
 
-from tmol import beta2016_score_function
-
 logger = logging.getLogger(__name__)
 
 
@@ -211,6 +209,7 @@ def pose_stack_from_biotite(
     from tmol.pack.rotamer.dunbrack.dunbrack_chi_sampler import (
         create_dunbrack_sampler_from_database,
     )
+    from tmol import beta2016_score_function
 
     if context is not None:
         if param_db is not None:

tmol/pose/pose_stack.py

@@ -1,6 +1,6 @@
 import attr
 import torch
-from typing import Optional
+from typing import Optional, Union
 
 from tmol.types.torch import Tensor
 from tmol.chemical.restypes import RefinedResidueType
@@ -276,6 +276,12 @@ def split(self, index) -> "PoseStack":
             device=self.device,
         )
 
+    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/fast_relax.py

@@ -166,7 +166,9 @@ def fast_relax(
     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.
@@ -340,6 +342,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()
@@ -357,7 +364,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:

tmol/tests/relax/test_fast_relax.py

@@ -2,7 +2,7 @@
 import numpy
 import pytest
 
-from tmol.relax.fast_relax import _default_cart_min_fn, fast_relax
+from tmol.relax.fast_relax import _default_cart_min_fn, _default_kin_min_fn, fast_relax
 import time
 
 from tmol.pose.pose_stack import PoseStack
@@ -84,6 +84,7 @@ def task_op(task):
         fold_forest,
         task_operations=[task_op],
         num_repeats=1,
+        min_fn=_default_kin_min_fn,
         verbose=verbose,
     )
 
@@ -220,6 +221,7 @@ def task_op(task):
         fold_forest,
         task_operations=[task_op],
         num_repeats=1,
+        min_fn=_default_kin_min_fn,
         verbose=verbose,
     )
 
@@ -293,6 +295,7 @@ def task_op(task):
         fold_forest,
         task_operations=[task_op],
         num_repeats=1,
+        min_fn=_default_kin_min_fn,
         verbose=verbose,
     )
 
@@ -308,3 +311,46 @@ def task_op(task):
     print(
         f"Three differently-shaped PDBs relaxed; Execution time: {elapsed_time:.6f} seconds"
     )
+
+
+def test_fast_relax_with_f64(default_database, ubq_pdb, dun_sampler, torch_device):
+    # if torch_device == torch.device("cpu"):
+    #     return
+
+    pose_stack = pose_stack_from_pdb(
+        ubq_pdb, torch_device, residue_start=0, residue_end=76
+    ).to(torch.float64)
+
+    sfxn = get_relax_sfxn(default_database, torch_device)
+
+    mm = MoveMap.from_pose_stack(pose_stack)
+    mm.move_all_jumps = True
+    mm.move_all_named_torsions = True
+
+    palette = PackerPalette()
+    fold_forest = FoldForest.reasonable_fold_forest(pose_stack)
+
+    def task_op(task):
+        task.restrict_to_repacking()
+
+        fixed_sampler = FixedAAChiSampler()
+        task.add_conformer_sampler(dun_sampler)
+        task.add_conformer_sampler(fixed_sampler)
+        task.add_conformer_sampler(IncludeCurrentSampler())
+
+    # Now let's run fast_relax
+    verbose = True
+    new_pose_stack = fast_relax(
+        pose_stack,
+        sfxn,
+        palette,
+        mm,
+        fold_forest,
+        task_operations=[task_op],
+        num_repeats=1,
+        min_fn=_default_kin_min_fn,
+        verbose=verbose,
+    )
+    assert (
+        new_pose_stack.coords.dtype == torch.float64
+    ), "Output coords dtype should match input coords dtype"