For fixed chains, retain residue numbering

For chains that are completely fixed, retain the residue numbering from
the input rather than renumbering.  For chains that are partially or
fully designed by RFdiffusion, it isn't clear to me what the 'correct'
behaviour should be, so these chains will be re-numbered starting at
residue 1.

Author: BYachnin

rfdiffusion/inference/model_runners.py

@@ -14,6 +14,7 @@
 from rfdiffusion import util
 from hydra.core.hydra_config import HydraConfig
 import os
+import string
 
 from rfdiffusion.model_input_logger import pickle_function_call
 import sys
@@ -144,13 +145,14 @@ def initialize(self, conf: DictConfig) -> None:
             self.symmetry = None
 
         self.allatom = ComputeAllAtomCoords().to(self.device)
-        
+
         if self.inf_conf.input_pdb is None:
             # set default pdb
             script_dir=os.path.dirname(os.path.realpath(__file__))
             self.inf_conf.input_pdb=os.path.join(script_dir, '../../examples/input_pdbs/1qys.pdb')
         self.target_feats = iu.process_target(self.inf_conf.input_pdb, parse_hetatom=True, center=False)
         self.chain_idx = None
+        self.idx_pdb = None
 
         ##############################
         ### Handle Partial Noising ###
@@ -313,10 +315,31 @@ def sample_init(self, return_forward_trajectory=False):
 
         first_res = 0
         self.chain_idx = []
+        self.idx_pdb = []
+        all_chains = {contig_ref[0] for contig_ref in self.contig_map.ref}
+        available_chains = sorted(list(set(string.ascii_uppercase) - all_chains))
+        # Iterate over each chain
         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)
+            chain_ids = {contig_ref[0] for contig_ref in self.contig_map.ref[first_res: last_res]}
+            # If we are designing this chain, it will have a '-' in the contig map
+            # Renumber this chain from 1
+            if "_" in chain_ids:
+                self.idx_pdb += [idx + 1 for idx in range(last_res - first_res)]
+                chain_ids = chain_ids - {"_"}
+                # If there are no fixed residues that have a chain id, pick the first available letter
+                if not chain_ids:
+                    chain_id = available_chains[0]
+                    available_chains.remove(chain_id)
+                # Otherwise, use the chain of the fixed (motif) residues
+                else:
+                    assert len(chain_ids) == 1, f"Error: Multiple chain IDs in chain: {chain_ids}"
+                    chain_id = list(chain_ids)[0]
+                self.chain_idx += [chain_id] * (last_res - first_res)
+            # If this is a fixed chain, maintain the chain and residue numbering
+            else:
+                self.idx_pdb += [contig_ref[1] 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
 
         ####################################

scripts/run_inference.py

@@ -141,6 +141,7 @@ def main(conf: HydraConfig) -> None:
             sampler.binderlen,
             chain_idx=sampler.chain_idx,
             bfacts=bfacts,
+            idx_pdb=sampler.idx_pdb
         )
 
         # run metadata