compute potentials from either partially denoised or predicted structures

Author: geraseva

rfdiffusion/inference/utils.py

@@ -360,7 +360,7 @@ def rmsd(V, W, eps=0):
         px0_[~atom_mask] = float("nan")
         return torch.Tensor(px0_)
 
-    def get_potential_gradients(self, xyz, diffusion_mask):
+    def get_potential_gradients(self, xyz, diffusion_mask, t, predicted=False):
         """
         This could be moved into potential manager if desired - NRB
 
@@ -375,7 +375,8 @@ def get_potential_gradients(self, xyz, diffusion_mask):
             Ca_grads (torch.tensor): [L,3] The gradient at each Ca atom
         """
 
-        if self.potential_manager == None or self.potential_manager.is_empty():
+        if (self.potential_manager == None or 
+            sum([potential.predicted==predicted for potential in self.potential_manager.potentials_to_apply])<1):
             return torch.zeros(xyz.shape[0], 3)
 
         use_Cb = False
@@ -386,7 +387,7 @@ def get_potential_gradients(self, xyz, diffusion_mask):
         if not xyz.grad is None:
             xyz.grad.zero_()
 
-        current_potential = self.potential_manager.compute_all_potentials(xyz)
+        current_potential = self.potential_manager.compute_all_potentials(xyz, predicted)
         current_potential.backward()
 
         # Since we are not moving frames, Cb grads are same as Ca grads
@@ -395,6 +396,10 @@ def get_potential_gradients(self, xyz, diffusion_mask):
 
         if not diffusion_mask == None:
             Ca_grads[diffusion_mask, :] = 0
+        
+        # clamp gradients
+        dist=(Ca_grads**2).sum(-1, keepdim=True).sqrt()
+        Ca_grads=torch.where(dist<=1/self.alphabar_schedule[t-1], Ca_grads, Ca_grads/dist/self.alphabar_schedule[t-1]).detach()
 
         # check for NaN's
         if torch.isnan(Ca_grads).any():
@@ -452,16 +457,7 @@ def get_next_pose(
         # Now done with diffusion mask. if fix motif is False, just set diffusion mask to be all True, and all coordinates can diffuse
         if not fix_motif:
             diffusion_mask[:] = False
-
-        grad_ca = self.get_potential_gradients(
-            px0.clone(), diffusion_mask=diffusion_mask
-        )
-        # clamp gradients
-        dist=(grad_ca**2).sum(-1, keepdim=True).sqrt()
-        grad_ca=torch.where(dist<=1/self.alphabar_schedule[t-1], grad_ca, grad_ca/dist/self.alphabar_schedule[t-1]).detach()
-
-        px0 += self.potential_manager.get_guide_scale(t) * grad_ca[:, None, :]
-
+        
         # get the next set of CA coordinates
         noise_scale_ca = self.noise_schedule_ca(t)
         _, ca_deltas = get_next_ca(
@@ -489,16 +485,16 @@ def get_next_pose(
 
         # Apply gradient step from guiding potentials
         # This can be moved to below where the full atom representation is calculated to allow for potentials involving sidechains
-        '''
+        
         grad_ca = self.get_potential_gradients(
-            xt.clone(), diffusion_mask=diffusion_mask
+            xt.clone(), diffusion_mask=diffusion_mask, t=t
+        ) + self.get_potential_gradients(
+            px0.clone(), diffusion_mask=diffusion_mask,
+            t=t, predicted=True
         )
-        # clamp gradients
-        dist=(grad_ca**2).sum(-1, keepdim=True).sqrt()
-        grad_ca=torch.where(dist<=1/self.alphabar_schedule[t-1], grad_ca, grad_ca/dist/self.alphabar_schedule[t-1]).detach()
-
+        
         ca_deltas += self.potential_manager.get_guide_scale(t) * grad_ca
-        '''
+        
         # add the delta to the new frames
         frames_next = torch.from_numpy(frames_next) + ca_deltas[:, None, :]  # translate
 

rfdiffusion/potentials/manager.py

@@ -167,12 +167,12 @@ def initialize_all_potentials(self, setting_list):
 
         return to_apply
 
-    def compute_all_potentials(self, xyz):
+    def compute_all_potentials(self, xyz, predicted=False):
         '''
             This is the money call. Take the current sequence and structure information and get the sum of all of the potentials that are being used
         '''
 
-        potential_list = [potential.compute(xyz) for potential in self.potentials_to_apply]
+        potential_list = [potential.compute(xyz) for potential in self.potentials_to_apply if potential.predicted==predicted]
         potential_stack = torch.stack(potential_list, dim=0)
 
         return torch.sum(potential_stack, dim=0)

rfdiffusion/potentials/potentials.py

@@ -8,6 +8,9 @@ class Potential:
     '''
         Interface class that defines the functions a potential must implement
     '''
+    def __init__(self):
+
+        self.predicted=False
 
     def compute(self, xyz):
         '''
@@ -31,7 +34,7 @@ class monomer_ROG(Potential):
     '''
 
     def __init__(self, weight=1, min_dist=15):
-
+        super().__init__()
         self.weight   = weight
         self.min_dist = min_dist
 
@@ -57,6 +60,7 @@ class binder_ROG(Potential):
 
     def __init__(self, binderlen, weight=1, min_dist=15):
 
+        super().__init__()
         self.binderlen = binderlen
         self.min_dist  = min_dist
         self.weight    = weight
@@ -87,6 +91,7 @@ class dimer_ROG(Potential):
 
     def __init__(self, binderlen, weight=1, min_dist=15):
 
+        super().__init__()
         self.binderlen = binderlen
         self.min_dist  = min_dist
         self.weight    = weight
@@ -127,6 +132,7 @@ class binder_ncontacts(Potential):
 
     def __init__(self, binderlen, weight=1, r_0=8, d_0=4):
 
+        super().__init__()
         self.binderlen = binderlen
         self.r_0       = r_0
         self.weight    = weight
@@ -161,6 +167,7 @@ class interface_ncontacts(Potential):
 
     def __init__(self, binderlen, weight=1, r_0=8, d_0=6):
 
+        super().__init__()
         self.binderlen = binderlen
         self.r_0       = r_0
         self.weight    = weight
@@ -200,6 +207,7 @@ class monomer_contacts(Potential):
 
     def __init__(self, weight=1, r_0=8, d_0=2, eps=1e-6):
 
+        super().__init__()
         self.r_0       = r_0
         self.weight    = weight
         self.d_0       = d_0
@@ -245,6 +253,7 @@ def __init__(self,
 
             weight (int/float, optional): Scaling/weighting factor
         """
+        super().__init__()
         self.contact_matrix = contact_matrix
         self.weight_intra = weight_intra 
         self.weight_inter = weight_inter 
@@ -354,6 +363,7 @@ class substrate_contacts(Potential):
 
     def __init__(self, weight=1, r_0=8, d_0=2, s=1, eps=1e-6, rep_r_0=5, rep_s=2, rep_r_min=1):
 
+        super().__init__()
         self.r_0       = r_0
         self.weight    = weight
         self.d_0       = d_0
@@ -464,13 +474,16 @@ class dmasif_interactions(Potential):
 
     def __init__(self, binderlen, int_weight=1, non_int_weight=1, threshold=3, seq_model_type='protein_mpnn'):
 
+        super().__init__()
+        self.predicted=True
+
         submodule_path='/'.join(__file__.split('/')[:-4])
         import sys
         sys.path.append(submodule_path)
 
-        from masif_martini.rfdiff_potential import Potential_from_bb
+        from masif_martini.potential import RFdiff_potential_from_bb
 
-        self.potential=Potential_from_bb(binderlen=binderlen, 
+        self.potential=RFdiff_potential_from_bb(binderlen=binderlen, 
                                           int_weight=int_weight, 
                                           non_int_weight=non_int_weight, 
                                           threshold=threshold,