From 3b0055084cfc0c8a9a220c8b921ff0bd310e5b1e Mon Sep 17 00:00:00 2001
From: Andrew Leaver-Fay <aleaverfay@gmail.com>
Date: Fri, 7 Mar 2025 16:17:59 +0000
Subject: [PATCH] First attempt at resetting R3 count-pair style i to i+2
 connectivity

---
 tmol/io/details/select_from_canonical.py | 41 ++++++++++++++++++++++++
 tmol/io/pose_stack_construction.py       |  4 ++-
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/tmol/io/details/select_from_canonical.py b/tmol/io/details/select_from_canonical.py
index a7cac119e..d19f49ddc 100644
--- a/tmol/io/details/select_from_canonical.py
+++ b/tmol/io/details/select_from_canonical.py
@@ -21,6 +21,7 @@ def assign_block_types(
     res_type_variants: Tensor[torch.int32][:, :],
     found_disulfides64: Tensor[torch.int64][:, 3],
     res_not_connected: Optional[Tensor[torch.bool][:, :, 2]] = None,
+    hack_rosetta3_i_ip2_count_pair: bool = False,
 ) -> Tuple[
     Tensor[torch.int64][:, :],
     Tensor[torch.int64][:, :, :, 2],
@@ -193,6 +194,46 @@ def assign_block_types(
     ibb64 = PoseStackBuilder._calculate_interblock_bondsep_from_connectivity_graph(
         pbt, block_n_conn, pose_n_pconn, pconn_matrix
     )
+    if hack_rosetta3_i_ip2_count_pair:
+        # overwrite the i to i+2 bond separation for the upper/lower pair
+        # to reflect rosetta3's count pair logic that treats the C on i
+        # and the N on i+2 at full strength.
+        # NOTE: this only works for pure-alpha-amino-acid monomers; the logic
+        # assumes that i is connected to i+1 and i+1 is connected to i+2.
+        print("ibb64 before")
+        print(
+            ibb64[
+                nz_res_is_poly_and_conn_to_next_pose_ind[:-1][:5],
+                nz_res_is_poly_and_conn_to_next_res_ind[:-1][:5],
+                nz_res_is_poly_and_conn_to_prev_res_ind[1:][:5],
+                connected_up_conn_inds[:-1][:5],
+                connected_down_conn_inds[1:][:5],
+            ]
+        )
+        ibb64[
+            nz_res_is_poly_and_conn_to_next_pose_ind[:-1],
+            nz_res_is_poly_and_conn_to_next_res_ind[:-1],
+            nz_res_is_poly_and_conn_to_prev_res_ind[1:],
+            connected_up_conn_inds[:-1],
+            connected_down_conn_inds[1:],
+        ] = MAX_SIG_BOND_SEPARATION
+        ibb64[
+            nz_res_is_poly_and_conn_to_next_pose_ind[:-1],
+            nz_res_is_poly_and_conn_to_prev_res_ind[1:],
+            nz_res_is_poly_and_conn_to_next_res_ind[:-1],
+            connected_down_conn_inds[1:],
+            connected_up_conn_inds[:-1],
+        ] = MAX_SIG_BOND_SEPARATION
+        print("ibb64 after")
+        print(
+            ibb64[
+                nz_res_is_poly_and_conn_to_next_pose_ind[:-1][:5],
+                nz_res_is_poly_and_conn_to_next_res_ind[:-1][:5],
+                nz_res_is_poly_and_conn_to_prev_res_ind[1:][:5],
+                connected_up_conn_inds[:-1][:5],
+                connected_down_conn_inds[1:][:5],
+            ]
+        )
 
     return (block_type_ind64, inter_residue_connections64, ibb64)
 
diff --git a/tmol/io/pose_stack_construction.py b/tmol/io/pose_stack_construction.py
index 5f9ed0527..8fbc1b98c 100644
--- a/tmol/io/pose_stack_construction.py
+++ b/tmol/io/pose_stack_construction.py
@@ -20,7 +20,8 @@ def pose_stack_from_canonical_form(
     res_not_connected: Optional[Tensor[torch.bool][:, :, 2]] = None,
     return_chain_ind: bool = False,
     return_atom_mapping: bool = False,
-):
+    hack_rosetta3_i_ip2_count_pair: bool = False,
+) -> PoseStack:
     """Create a PoseStack, resolving which block type is requested by the
     presence and absence of the provided atoms for each residue type.
     There are five required arguments and several optional arguments.
@@ -194,6 +195,7 @@ def pose_stack_from_canonical_form(
         res_type_variants,
         found_disulfides,
         res_not_connected,
+        hack_rosetta3_i_ip2_count_pair,
     )
 
     # 6