Fixed issues with designing in scaffoldguided mode, for example: design_ppi_scaffold.sh. The solutions to issues 272 and 273 did not fully address the issue.

Author: OrangeCatzhang

config/inference/base.yaml

@@ -126,7 +126,7 @@ logging:
   inputs: False
 
 scaffoldguided:
-  scaffoldguided: False
+  scaffoldguided_enable: False
   target_pdb: False
   target_path: null
   scaffold_list: null

examples/design_ppi_scaffolded.sh

@@ -7,4 +7,4 @@
 # We then provide a path to a directory of different scaffolds (we've provided some for you to use, from Cao et al., 2022)
 # We generate 10 designs, and reduce the noise added during inference to 0 (which improves the quality of designs)
 
-../scripts/run_inference.py scaffoldguided.target_path=input_pdbs/insulin_target.pdb inference.output_prefix=example_outputs/design_ppi_scaffolded scaffoldguided.scaffoldguided=True 'ppi.hotspot_res=[A59,A83,A91]' scaffoldguided.target_pdb=True scaffoldguided.target_ss=target_folds/insulin_target_ss.pt scaffoldguided.target_adj=target_folds/insulin_target_adj.pt scaffoldguided.scaffold_dir=./ppi_scaffolds/ inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0
+../scripts/run_inference.py scaffoldguided.target_path=input_pdbs/insulin_target.pdb inference.output_prefix=example_outputs/design_ppi_scaffolded scaffoldguided.scaffoldguided_enable=True 'ppi.hotspot_res=[A59,A83,A91]' scaffoldguided.target_pdb=True scaffoldguided.target_ss=target_folds/insulin_target_ss.pt scaffoldguided.target_adj=target_folds/insulin_target_adj.pt scaffoldguided.scaffold_dir=./ppi_scaffolds/ inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0

rfdiffusion/inference/model_runners.py

@@ -77,14 +77,14 @@ def initialize(self, conf: DictConfig) -> None:
                 if conf.contigmap.provide_seq is not None:
                     # this is only used for partial diffusion
                     assert conf.diffuser.partial_T is not None, "The provide_seq input is specifically for partial diffusion"
-                if conf.scaffoldguided.scaffoldguided:
+                if conf.scaffoldguided.scaffoldguided_enable:
                     self.ckpt_path = f'{model_directory}/InpaintSeq_Fold_ckpt.pt'
                 else:
                     self.ckpt_path = f'{model_directory}/InpaintSeq_ckpt.pt'
-            elif conf.ppi.hotspot_res is not None and conf.scaffoldguided.scaffoldguided is False:
+            elif conf.ppi.hotspot_res is not None and conf.scaffoldguided.scaffoldguided_enable is False:
                 # use complex trained model
                 self.ckpt_path = f'{model_directory}/Complex_base_ckpt.pt'
-            elif conf.scaffoldguided.scaffoldguided is True:
+            elif conf.scaffoldguided.scaffoldguided_enable is True:
                 # use complex and secondary structure-guided model
                 self.ckpt_path = f'{model_directory}/Complex_Fold_base_ckpt.pt'
             else:
@@ -279,7 +279,6 @@ def sample_init(self, return_forward_trajectory=False):
         self.mask_seq = torch.from_numpy(self.contig_map.inpaint_seq)[None,:]
         self.mask_str = torch.from_numpy(self.contig_map.inpaint_str)[None,:]
         self.binderlen =  len(self.contig_map.inpaint)    
-        
         #######################################
         ### Resolve cyclic peptide indicies ###
         #######################################
@@ -301,7 +300,7 @@ def sample_init(self, return_forward_trajectory=False):
                 self.cyclic_reses = is_cyclized
         else:
             self.cyclic_reses = torch.zeros_like(self.mask_str).bool().to(self.device).squeeze()
-
+        
         ####################
         ### Get Hotspots ###
         ####################
@@ -681,7 +680,6 @@ def sample_step(self, *, t, x_t, seq_init, final_step):
         ####################
         ### Forward Pass ###
         ####################
-
         with torch.no_grad():
             msa_prev, pair_prev, px0, state_prev, alpha, logits, plddt = self.model(msa_masked,
                                 msa_full,
@@ -771,7 +769,7 @@ def __init__(self, conf: DictConfig):
         else:
             # initialize BlockAdjacency sampling class
             assert all(x is None for x in (conf.contigmap.inpaint_str_helix, conf.contigmap.inpaint_str_strand, conf.contigmap.inpaint_str_loop)), "can't provide scaffold_dir if you're also specifying per-residue ss"
-            self.blockadjacency = iu.BlockAdjacency(conf.scaffoldguided, conf.inference.num_designs)
+            self.blockadjacency = iu.BlockAdjacency(conf, conf.inference.num_designs)
 
 
         #################################################
@@ -945,6 +943,27 @@ def sample_init(self):
 
 
         xT = torch.clone(fa_stack[-1].squeeze()[:,:14,:])
+
+        ################################
+        ### Add to Cyclic_reses init ###
+        ################################
+
+        if self._conf.inference.cyclic:
+            if self._conf.inference.cyc_chains is None:
+                self.cyclic_reses = ~self.mask_str.to(self.device).squeeze()
+            else:
+                assert isinstance(self._conf.inference.cyc_chains, str), 'cyc_chains arg must be string'
+                cyc_chains = self._conf.inference.cyc_chains
+                cyc_chains = [i.upper() for i in cyc_chains]
+                hal_idx = self.contig_map.hal
+                is_cyclized = torch.zeros_like(self.mask_str).bool().to(self.device).squeeze()
+                for ch in cyc_chains:
+                    ch_mask = torch.tensor([idx[0] == ch for idx in hal_idx]).bool()
+                    is_cyclized[ch_mask] = True
+                self.cyclic_reses = is_cyclized
+        else:
+            self.cyclic_reses = torch.zeros_like(self.mask_str).bool().to(self.device).squeeze()
+
         return xT, seq_T
 
     def _preprocess(self, seq, xyz_t, t):

rfdiffusion/inference/utils.py

@@ -502,7 +502,7 @@ def get_next_pose(
 
 
 def sampler_selector(conf: DictConfig):
-    if conf.scaffoldguided.scaffoldguided:
+    if conf.scaffoldguided.scaffoldguided_enable:
         sampler = model_runners.ScaffoldedSampler(conf)
     else:
         if conf.inference.model_runner == "default":
@@ -1012,4 +1012,4 @@ def ss_from_contig(ss_masks: dict):
     for idx, mask in enumerate([ss_masks['helix'],ss_masks['strand'], ss_masks['loop']]):
         ss[mask,idx] = 1
         ss[mask, 3] = 0 # remove the mask token
-    return ss
+    return ss