add mask for binder loss

Author: sokrypton

colabdesign/af/alphafold/common/confidence.py

@@ -14,74 +14,79 @@
 
 """Functions for processing confidence metrics."""
 
-from typing import Dict, Optional, Tuple
+import jax.numpy as jnp
+import jax
 import numpy as np
+from colabdesign.af.alphafold.common import residue_constants
 import scipy.special
 
+def compute_tol(prev_pos, current_pos, mask, use_jnp=False):
+    # Early stopping criteria based on criteria used in
+    # AF2Complex: https://www.nature.com/articles/s41467-022-29394-2    
+    _np = jnp if use_jnp else np
+    dist = lambda x:_np.sqrt(((x[:,None] - x[None,:])**2).sum(-1))
+    ca_idx = residue_constants.atom_order['CA']
+    sq_diff = _np.square(dist(prev_pos[:,ca_idx])-dist(current_pos[:,ca_idx]))
+    mask_2d = mask[:,None] * mask[None,:]
+    return _np.sqrt((sq_diff * mask_2d).sum()/mask_2d.sum() + 1e-8)
 
-def compute_plddt(logits: np.ndarray) -> np.ndarray:
-  """Computes per-residue pLDDT from logits.
 
+def compute_plddt(logits, use_jnp=False):
+  """Computes per-residue pLDDT from logits.
   Args:
     logits: [num_res, num_bins] output from the PredictedLDDTHead.
-
   Returns:
     plddt: [num_res] per-residue pLDDT.
   """
+  if use_jnp:
+    _np, _softmax = jnp, jax.nn.softmax
+  else:
+    _np, _softmax = np, scipy.special.softmax
+  
   num_bins = logits.shape[-1]
   bin_width = 1.0 / num_bins
-  bin_centers = np.arange(start=0.5 * bin_width, stop=1.0, step=bin_width)
-  probs = scipy.special.softmax(logits, axis=-1)
-  predicted_lddt_ca = np.sum(probs * bin_centers[None, :], axis=-1)
+  bin_centers = _np.arange(start=0.5 * bin_width, stop=1.0, step=bin_width)
+  probs = _softmax(logits, axis=-1)
+  predicted_lddt_ca = (probs * bin_centers[None, :]).sum(-1)
   return predicted_lddt_ca * 100
 
-
-def _calculate_bin_centers(breaks: np.ndarray):
+def _calculate_bin_centers(breaks, use_jnp=False):
   """Gets the bin centers from the bin edges.
-
   Args:
     breaks: [num_bins - 1] the error bin edges.
-
   Returns:
     bin_centers: [num_bins] the error bin centers.
   """
-  step = (breaks[1] - breaks[0])
+  _np = jnp if use_jnp else np
+  step = breaks[1] - breaks[0]
 
   # Add half-step to get the center
   bin_centers = breaks + step / 2
-  # Add a catch-all bin at the end.
-  bin_centers = np.concatenate([bin_centers, [bin_centers[-1] + step]],
-                               axis=0)
-  return bin_centers
 
+  # Add a catch-all bin at the end.
+  return _np.append(bin_centers, bin_centers[-1] + step)
 
 def _calculate_expected_aligned_error(
-    alignment_confidence_breaks: np.ndarray,
-    aligned_distance_error_probs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+  alignment_confidence_breaks,
+  aligned_distance_error_probs,
+  use_jnp=False):
   """Calculates expected aligned distance errors for every pair of residues.
-
   Args:
     alignment_confidence_breaks: [num_bins - 1] the error bin edges.
     aligned_distance_error_probs: [num_res, num_res, num_bins] the predicted
       probs for each error bin, for each pair of residues.
-
   Returns:
     predicted_aligned_error: [num_res, num_res] the expected aligned distance
       error for each pair of residues.
     max_predicted_aligned_error: The maximum predicted error possible.
   """
-  bin_centers = _calculate_bin_centers(alignment_confidence_breaks)
-
+  bin_centers = _calculate_bin_centers(alignment_confidence_breaks, use_jnp=use_jnp)
   # Tuple of expected aligned distance error and max possible error.
-  return (np.sum(aligned_distance_error_probs * bin_centers, axis=-1),
-          np.asarray(bin_centers[-1]))
-
+  pae = (aligned_distance_error_probs * bin_centers).sum(-1)
+  return (pae, bin_centers[-1])
 
-def compute_predicted_aligned_error(
-    logits: np.ndarray,
-    breaks: np.ndarray) -> Dict[str, np.ndarray]:
+def compute_predicted_aligned_error(logits, breaks, use_jnp=False):
   """Computes aligned confidence metrics from logits.
-
   Args:
     logits: [num_res, num_res, num_bins] the logits output from
       PredictedAlignedErrorHead.
@@ -94,62 +99,71 @@ def compute_predicted_aligned_error(
       error for each pair of residues.
     max_predicted_aligned_error: The maximum predicted error possible.
   """
-  aligned_confidence_probs = scipy.special.softmax(
-      logits,
-      axis=-1)
-  predicted_aligned_error, max_predicted_aligned_error = (
-      _calculate_expected_aligned_error(
-          alignment_confidence_breaks=breaks,
-          aligned_distance_error_probs=aligned_confidence_probs))
+  _softmax = jax.nn.softmax if use_jnp else scipy.special.softmax
+  aligned_confidence_probs = _softmax(logits,axis=-1)
+  predicted_aligned_error, max_predicted_aligned_error = \
+  _calculate_expected_aligned_error(breaks, aligned_confidence_probs, use_jnp=use_jnp)
+
   return {
       'aligned_confidence_probs': aligned_confidence_probs,
       'predicted_aligned_error': predicted_aligned_error,
       'max_predicted_aligned_error': max_predicted_aligned_error,
   }
 
-
-def predicted_tm_score(
-    logits: np.ndarray,
-    breaks: np.ndarray,
-    residue_weights: Optional[np.ndarray] = None) -> np.ndarray:
-  """Computes predicted TM alignment score.
+def predicted_tm_score(logits, breaks, residue_weights = None,
+    asym_id = None, use_jnp=False):
+  """Computes predicted TM alignment or predicted interface TM alignment score.
 
   Args:
     logits: [num_res, num_res, num_bins] the logits output from
       PredictedAlignedErrorHead.
     breaks: [num_bins] the error bins.
     residue_weights: [num_res] the per residue weights to use for the
       expectation.
+    asym_id: [num_res] the asymmetric unit ID - the chain ID. Only needed for
+      ipTM calculation.
 
   Returns:
-    ptm_score: the predicted TM alignment score.
+    ptm_score: The predicted TM alignment or the predicted iTM score.
   """
+  if use_jnp:
+    _np, _softmax = jnp, jax.nn.softmax
+  else:
+    _np, _softmax = np, scipy.special.softmax
 
   # residue_weights has to be in [0, 1], but can be floating-point, i.e. the
   # exp. resolved head's probability.
   if residue_weights is None:
-    residue_weights = np.ones(logits.shape[0])
+    residue_weights = _np.ones(logits.shape[0])
 
-  bin_centers = _calculate_bin_centers(breaks)
+  bin_centers = _calculate_bin_centers(breaks, use_jnp=use_jnp)
+  num_res = residue_weights.shape[0]
 
-  num_res = np.sum(residue_weights)
   # Clip num_res to avoid negative/undefined d0.
-  clipped_num_res = max(num_res, 19)
+  clipped_num_res = _np.maximum(residue_weights.sum(), 19)
 
-  # Compute d_0(num_res) as defined by TM-score, eqn. (5) in
-  # http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf
-  # Yang & Skolnick "Scoring function for automated
-  # assessment of protein structure template quality" 2004
+  # Compute d_0(num_res) as defined by TM-score, eqn. (5) in Yang & Skolnick
+  # "Scoring function for automated assessment of protein structure template
+  # quality", 2004: http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf
   d0 = 1.24 * (clipped_num_res - 15) ** (1./3) - 1.8
 
-  # Convert logits to probs
-  probs = scipy.special.softmax(logits, axis=-1)
+  # Convert logits to probs.
+  probs = _softmax(logits, axis=-1)
+
+  # TM-Score term for every bin.
+  tm_per_bin = 1. / (1 + _np.square(bin_centers) / _np.square(d0))
+  # E_distances tm(distance).
+  predicted_tm_term = (probs * tm_per_bin).sum(-1)
+
+  if asym_id is None:
+    pair_mask = _np.full((num_res,num_res),True)
+  else:
+    pair_mask = asym_id[:, None] != asym_id[None, :]
+
+  predicted_tm_term *= pair_mask
 
-  # TM-Score term for every bin
-  tm_per_bin = 1. / (1 + np.square(bin_centers) / np.square(d0))
-  # E_distances tm(distance)
-  predicted_tm_term = np.sum(probs * tm_per_bin, axis=-1)
+  pair_residue_weights = pair_mask * (residue_weights[None, :] * residue_weights[:, None])
+  normed_residue_mask = pair_residue_weights / (1e-8 + pair_residue_weights.sum(-1, keepdims=True))
+  per_alignment = (predicted_tm_term * normed_residue_mask).sum(-1)
 
-  normed_residue_mask = residue_weights / (1e-8 + residue_weights.sum())
-  per_alignment = np.sum(predicted_tm_term * normed_residue_mask, axis=-1)
-  return np.asarray(per_alignment[(per_alignment * residue_weights).argmax()])
+  return (per_alignment * residue_weights).max()

colabdesign/af/loss.py

@@ -5,7 +5,7 @@
 from colabdesign.shared.utils import Key, copy_dict
 from colabdesign.shared.protein import jnp_rmsd_w, _np_kabsch, _np_rmsd, _np_get_6D_loss
 from colabdesign.af.alphafold.model import model, folding, all_atom
-from colabdesign.af.alphafold.common import confidence_jax, residue_constants
+from colabdesign.af.alphafold.common import confidence, residue_constants
 
 ####################################################
 # AF_LOSS - setup loss function
@@ -36,12 +36,16 @@ def _loss_binder(self, inputs, outputs, aux):
     '''get losses'''
     opt = inputs["opt"]
     zeros = jnp.zeros(sum(self._lengths))
-    binder_id = zeros.at[self._target_len:].set(1)
+    mask_1d = inputs["seq_mask"]
+    binder_id = zeros.at[-self._binder_len:].set(1)
+    binder_id = jnp.where(mask_1d, binder_id, 0)
     if "hotspot" in opt:
       target_id = zeros.at[opt["hotspot"]].set(1)
+      target_id = jnp.where(mask_1d, target_id, 0)
       i_con_loss = get_con_loss(inputs, outputs, opt["i_con"], mask_1d=target_id, mask_1b=binder_id)
     else:
       target_id = zeros.at[:self._target_len].set(1)
+      target_id = jnp.where(mask_1d, target_id, 0)
       i_con_loss = get_con_loss(inputs, outputs, opt["i_con"], mask_1d=binder_id, mask_1b=target_id)
 
     # unsupervised losses
@@ -68,10 +72,17 @@ def _loss_binder(self, inputs, outputs, aux):
       # compute fape
       fape = get_fape_loss(inputs, outputs, clamp=opt["fape_cutoff"], return_mtx=True)
 
+      mask_1d = inputs["batch"]["all_atom_mask"][:,1]
+      mask_2d = mask_1d[:,None] * mask_1d[None,:]
+      def exclude_target(x):
+        x = x[-self._binder_len:,:]
+        m = mask_2d[-self._binder_len:,:]
+        return (x*m).sum() / (m.sum() + 1e-8)
+
       aux["losses"].update({
         "rmsd":      aln["rmsd"],
-        "dgram_cce": cce[self._target_len:,:].mean(),
-        "fape":      fape[self._target_len:,:].mean()
+        "dgram_cce": exclude_target(cce),
+        "fape":      exclude_target(fape)
       })
 
     else:
@@ -198,10 +209,15 @@ def get_pae(outputs):
   return (prob*bin_centers).sum(-1)
 
 def get_ptm(inputs, outputs, interface=False):
-  pae = outputs["predicted_aligned_error"]
-  if "asym_id" not in pae:
-    pae["asym_id"] = inputs["asym_id"]
-  return confidence_jax.predicted_tm_score_jax(**pae, interface=interface)
+  pae = {"residue_weights":inputs["seq_mask"],
+         **outputs["predicted_aligned_error"]}
+  if interface:
+    if "asym_id" not in pae:
+      pae["asym_id"] = inputs["asym_id"]
+  else:
+    if "asym_id" in pae:
+      pae.pop("asym_id")
+  return confidence.predicted_tm_score(**pae, use_jnp=True)
 
 def get_contact_map(outputs, dist=8.0):
   '''get contact map from distogram'''