Maintains the input chain ids in RFdiffusion output

The output was previously renumbering all of the chains, making
comparisons to the input structures and handling of multi-chain inputs
challenging.  This commit maintains the input chain ids in the
output.

Author: BYachnin

rfdiffusion/contigs.py

@@ -80,6 +80,7 @@ def __init__(
                 self.inpaint,
                 self.inpaint_hal,
                 self.inpaint_rf,
+                self.sampled_mask_length_bound,
             ) = self.expand_sampled_mask()
             self.ref = self.inpaint + self.receptor
             self.hal = self.inpaint_hal + self.receptor_hal
@@ -241,6 +242,8 @@ def expand_sampled_mask(self):
         inpaint_chain_idx = -1
         receptor_chain_break = []
         inpaint_chain_break = []
+        _receptor_mask_length_bound = []
+        _inpaint_mask_length_bound = []
         for con in self.sampled_mask:
             if (
                 all([i[0].isalpha() for i in con.split("/")[:-1]])
@@ -286,6 +289,7 @@ def expand_sampled_mask(self):
                         receptor_chain_break.append(
                             (receptor_idx - 1, 200)
                         )  # 200 aa chain break
+                _receptor_mask_length_bound.append(len(receptor))
             else:
                 inpaint_chain_idx += 1
                 for subcon in con.split("/"):
@@ -320,6 +324,7 @@ def expand_sampled_mask(self):
                         )
                         inpaint_idx += int(subcon.split("-")[0])
                 inpaint_chain_break.append((inpaint_idx - 1, 200))
+                _inpaint_mask_length_bound.append(len(inpaint))
 
         if self.topo is True or inpaint_hal == []:
             receptor_hal = [(i[0], i[1]) for i in receptor_hal]
@@ -335,14 +340,21 @@ def expand_sampled_mask(self):
             inpaint_rf[ch_break[0] :] += ch_break[1]
         for ch_break in receptor_chain_break[:-1]:
             receptor_rf[ch_break[0] :] += ch_break[1]
-
+        sampled_mask_length_bound = []
+        sampled_mask_length_bound.extend(_inpaint_mask_length_bound)
+        if _inpaint_mask_length_bound:
+            inpaint_last_bound = _inpaint_mask_length_bound[-1]
+        else:
+            inpaint_last_bound = 0
+        sampled_mask_length_bound.extend(map(lambda x: x + inpaint_last_bound, _receptor_mask_length_bound))
         return (
             receptor,
             receptor_hal,
             receptor_rf.tolist(),
             inpaint,
             inpaint_hal,
             inpaint_rf.tolist(),
+            sampled_mask_length_bound
         )
 
     def get_inpaint_seq_str(self, inpaint_s, ss=False):

rfdiffusion/inference/model_runners.py

@@ -309,8 +309,16 @@ def sample_init(self, return_forward_trajectory=False):
         contig_map=self.contig_map
 
         self.diffusion_mask = self.mask_str
-        self.chain_idx=['A' if i < self.binderlen else 'B' for i in range(L_mapped)]
-        
+        length_bound = self.contig_map.sampled_mask_length_bound.copy()
+
+        first_res = 0
+        self.chain_idx = []
+        for last_res in length_bound:
+            chain_ids = {contig_ref[0] for contig_ref in self.contig_map.ref[first_res: last_res]} - {"_"}
+            assert len(chain_ids) == 1, f"Error: Multiple chain IDs in chain: {chain_ids}"
+            self.chain_idx += [list(chain_ids)[0]] * (last_res - first_res)
+            first_res = last_res
+
         ####################################
         ### Generate initial coordinates ###
         ####################################