Added secondary structure specification

Author: Joseph Watson

config/inference/base.yaml

@@ -26,6 +26,9 @@ contigmap:
   contigs: null
   inpaint_seq: null
   inpaint_str: null
+  inpaint_str_helix: null
+  inpaint_str_strand: null
+  inpaint_str_loop: null
   provide_seq: null
   length: null
 

rfdiffusion/contigs.py

@@ -27,6 +27,9 @@ def __init__(
         inpaint_str_tensor=None,
         topo=False,
         provide_seq=None,
+        inpaint_str_strand=None,
+        inpaint_str_helix=None,
+        inpaint_str_loop=None
     ):
         # sanity checks
         if contigs is None and ref_idx is None:
@@ -48,12 +51,15 @@ def __init__(
         self.ref_idx = ref_idx
         self.hal_idx = hal_idx
         self.idx_rf = idx_rf
-        self.inpaint_seq = (
-            "/".join(inpaint_seq).split("/") if inpaint_seq is not None else None
-        )
-        self.inpaint_str = (
-            "/".join(inpaint_str).split("/") if inpaint_str is not None else None
-        )
+
+        parse_inpaint = lambda x: "/".join(x).split("/") if x is not None else None
+        self.inpaint_seq = parse_inpaint(inpaint_seq)
+        self.inpaint_str = parse_inpaint(inpaint_str)
+
+        self.inpaint_str_helix=parse_inpaint(inpaint_str_helix)
+        self.inpaint_str_strand=parse_inpaint(inpaint_str_strand)
+        self.inpaint_str_loop=parse_inpaint(inpaint_str_loop)
+
         self.inpaint_seq_tensor = inpaint_seq_tensor
         self.inpaint_str_tensor = inpaint_str_tensor
         self.parsed_pdb = parsed_pdb
@@ -125,6 +131,39 @@ def __init__(
                 else:
                     self.inpaint_seq[int(i)] = True
 
+        """
+        We have now added the ability to specify the secondary structure of provided sequence.
+        This is described in Liu et al., 2024
+        https://www.biorxiv.org/content/10.1101/2024.07.16.603789v1
+        This is for the case that e.g. you have a sequence, but don't know the structure (like an IDR), but
+        want to specify the secondary structure of this sequence.
+        Making this compatible with the contigmap object allows all the variable length stuff to be handled.
+        
+        The logic:
+        Secondary structure is provided at the command line, using the following three flags:
+            inpaint_str_helix
+            inpaint_str_strand
+            inpaint_str_loop
+        
+        These are so named because they pertain to the region of the input pdb that you have applied inpaint_str to
+        In other words, any part of the input protein you are masking the structure of, you can specify the secondary structure of.
+        However, you can't specify the secondary structure of a region you're not applying inpaint_str to, as this doesn't make sense.
+        """
+
+        if any(x is not None for x in (inpaint_str_helix, inpaint_str_strand, inpaint_str_loop)):
+            self.ss_spec={}
+            order=['helix','strand','loop']
+            for idx, i in enumerate([inpaint_str_helix, inpaint_str_strand, inpaint_str_loop]):
+                if i is not None:
+                    self.ss_spec[order[idx]] = ~self.get_inpaint_seq_str(i, ss=True)
+                else:
+                    self.ss_spec[order[idx]] = np.zeros(len(self.inpaint_seq), dtype=bool)
+            # some sensible checks
+            for key, mask in self.ss_spec.items():
+                assert sum(mask*self.inpaint_str) == 0, f"You've specified {key} residues that are not structure-masked with inpaint_str. This doesn't really make sense."
+            stack=np.vstack([mask for mask in self.ss_spec.values()])
+            assert np.max(np.sum(stack, axis=0)) == 1, "You've given multiple secondary structure assignations to an input residue. This doesn't make sense."
+
     def get_sampled_mask(self):
         """
         Function to get a sampled mask from a contig.
@@ -306,11 +345,14 @@ def expand_sampled_mask(self):
             inpaint_rf.tolist(),
         )
 
-    def get_inpaint_seq_str(self, inpaint_s):
-        """
+    def get_inpaint_seq_str(self, inpaint_s, ss=False):
+        '''
         function to generate inpaint_str or inpaint_seq masks specific to this contig
-        """
-        s_mask = np.copy(self.mask_1d)
+        '''
+        if not ss:
+            s_mask = np.copy(self.mask_1d)
+        else:
+            s_mask= np.ones(len(self.mask_1d), dtype=bool)
         inpaint_s_list = []
         for i in inpaint_s:
             if "-" in i:

rfdiffusion/inference/model_runners.py

@@ -741,7 +741,15 @@ def __init__(self, conf: DictConfig):
         """
         super().__init__(conf)
         # initialize BlockAdjacency sampling class
-        self.blockadjacency = iu.BlockAdjacency(conf, conf.inference.num_designs)
+        if conf.scaffoldguided.scaffold_dir is None:
+            assert any(x is not None for x in (conf.contigmap.inpaint_str_helix, conf.contigmap.inpaint_str_strand, conf.contigmap.inpaint_str_loop))
+            if conf.contigmap.inpaint_str_loop is not None:
+                assert conf.scaffoldguided.mask_loops == False, "You shouldn't be masking loops if you're specifying loop secondary structure"
+        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)
+
 
         #################################################
         ### Initialize target, if doing binder design ###
@@ -773,8 +781,11 @@ def sample_init(self):
         ##########################
         ### Process Fold Input ###
         ##########################
-        self.L, self.ss, self.adj = self.blockadjacency.get_scaffold()
-        self.adj = nn.one_hot(self.adj.long(), num_classes=3)
+        if hasattr(self, 'blockadjacency'):
+            self.L, self.ss, self.adj = self.blockadjacency.get_scaffold()
+            self.adj = nn.one_hot(self.adj.long(), num_classes=3)
+        else:
+            self.L=100 # shim. Get's overwritten
 
         ##############################
         ### Auto-contig generation ###
@@ -848,7 +859,7 @@ def sample_init(self):
             seq_T=torch.full((L_mapped,),21)
             seq_T[contig_map.hal_idx0] = seq_orig[contig_map.ref_idx0]
             seq_T[~self.mask_seq.squeeze()] = 21
-            assert L_mapped==self.adj.shape[0]
+
             diffusion_mask = self.mask_str
             self.diffusion_mask = diffusion_mask
 
@@ -857,7 +868,13 @@ def sample_init(self):
             xT = get_init_xyz(xT).squeeze()
             atom_mask = torch.full((L_mapped, 27), False)
             atom_mask[contig_map.hal_idx0] = mask_27[contig_map.ref_idx0]
- 
+
+            if hasattr(self.contig_map, 'ss_spec'):
+                self.adj=torch.full((L_mapped, L_mapped),2) # masked
+                self.adj=nn.one_hot(self.adj.long(), num_classes=3)
+                self.ss=iu.ss_from_contig(self.contig_map.ss_spec)
+            assert L_mapped==self.adj.shape[0]
+            
         ####################
         ### Get hotspots ###
         ####################

rfdiffusion/inference/utils.py

@@ -1001,3 +1001,15 @@ def contig_crop(self, contig_crop, residue_offset=200) -> None:
 
     def get_target(self):
         return self.pdb
+
+def ss_from_contig(ss_masks: dict):
+    """  
+    Function for taking 1D masks for each of the ss types, and outputting a secondary structure input
+    """
+    L=len(ss_masks['helix'])
+    ss=torch.zeros((L, 4)).long()
+    ss[:,3] = 1 #mask
+    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