feat: align antgen-antibody with DSW

profold2/model/dsw.py

@@ -0,0 +1,124 @@
+""""Differentiable Soft Smith-Waterman
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DifferentiableSmithWaterman(nn.Module):
+  """
+  Differentiable Soft Smith-Waterman (local alignment).
+
+  Input:
+      S : [B, Lq, Lr] substitution score matrix (higher = more likely to align)
+  Output:
+      score : [B]           soft alignment score
+      P     : [B, Lq, Lr]   soft alignment path matrix (soft correspondence,
+                            row/col sums <= 1 after sinkhorn normalization)
+  """
+
+  def __init__(self, gap_open=0.5, gap_extend=0.1, temperature=1.0, eps=1e-8):
+    super().__init__()
+    # Gap penalties stored as learnable parameters (kept positive via softplus).
+    self.log_gap_open = nn.Parameter(torch.tensor(float(gap_open)))
+    self.log_gap_ext = nn.Parameter(torch.tensor(float(gap_extend)))
+    self.temperature = temperature
+    self.eps = eps
+
+  @property
+  def gap_open(self):
+    return F.softplus(self.log_gap_open)
+
+  @property
+  def gap_extend(self):
+    return F.softplus(self.log_gap_ext)
+
+  def _lse(self, *xs):
+    """Numerically stable log-sum-exp over a list of tensors."""
+    stacked = torch.stack(xs, dim=0)
+    maxv = stacked.max(dim=0).values
+    return maxv + torch.logsumexp(stacked - maxv.unsqueeze(0), dim=0)
+
+  def forward(self, S):
+    """
+    S: [B, Lq, Lr] substitution scores.
+    """
+    B, Lq, Lr = S.shape
+    tau = self.temperature
+    go = self.gap_open
+    ge = self.gap_extend
+
+    # DP matrices:
+    #   M = best score ending with a match/mismatch (diagonal step)
+    #   I = best score ending with a gap on the query side (ref advances)
+    #   D = best score ending with a gap on the ref side (query advances)
+    # Initialize to a large negative value (soft version of -inf).
+    neg_inf = torch.finfo(S.dtype).min / 4  # avoid overflow in exp
+    M = torch.full((B, Lq + 1, Lr + 1), neg_inf, device=S.device, dtype=S.dtype)
+    I = torch.full((B, Lq + 1, Lr + 1), neg_inf, device=S.device, dtype=S.dtype)
+    D = torch.full((B, Lq + 1, Lr + 1), neg_inf, device=S.device, dtype=S.dtype)
+
+    # Local alignment: free start from any position -> 0.
+    M[:, 0, :] = 0.0
+    M[:, :, 0] = 0.0
+    I[:, 0, :] = 0.0
+    D[:, :, 0] = 0.0
+
+    for u in range(1, Lq + 1):
+      for j in range(1, Lr + 1):
+        s = S[:, u - 1, j - 1]  # [B]
+
+        # --- I: gap on query side (advance along ref, j) ---
+        # Either extend an existing I-gap or open a new one from M.
+        i_ext = I[:, u, j - 1] - ge * tau
+        i_open = M[:, u, j - 1] - go * tau
+        I[:, u, j] = self._lse(i_ext, i_open, torch.zeros_like(i_ext))
+
+        # --- D: gap on ref side (advance along query, u) ---
+        d_ext = D[:, u - 1, j] - ge * tau
+        d_open = M[:, u - 1, j] - go * tau
+        D[:, u, j] = self._lse(d_ext, d_open, torch.zeros_like(d_ext))
+
+        # --- M: diagonal match/mismatch step ---
+        m_diag = M[:, u - 1, j - 1] + s * tau
+        m_from_i = I[:, u - 1, j - 1] + s * tau
+        m_from_d = D[:, u - 1, j - 1] + s * tau
+        zero = torch.zeros_like(m_diag)
+
+        # Soft-max over the 4 possible sources -> normalized weights.
+        stack = torch.stack([m_diag, m_from_i, m_from_d, zero], dim=0)  # [4, B]
+        w = F.softmax(stack, dim=0)  # [4, B]
+        # Soft value: weighted combination, rescaled back by temperature.
+        M[:, u, j] = (
+            w[0] * m_diag + w[1] * m_from_i + w[2] * m_from_d + w[3] * zero
+        ) / tau
+
+    # Soft alignment score = soft-max over all M cells (local alignment picks
+    # the best region in a differentiable way).
+    M_flat = M[:, 1:, 1:]  # [B, Lq, Lr]
+    score = torch.logsumexp(M_flat * tau, dim=(1, 2)) / tau  # [B]
+
+    # --- Soft path matrix P ---
+    # Use the normalized M table as a soft-correspondence proxy:
+    # high temperature -> uniform; low temperature -> sharp.
+    logits = M_flat * tau
+    P = F.softmax(logits.reshape(B, -1), dim=-1).reshape(B, Lq, Lr)
+    # A few Sinkhorn rounds push P toward a soft doubly-stochastic matching.
+    for _ in range(3):
+        P = P / (P.sum(dim=2, keepdim=True) + self.eps)
+        P = P / (P.sum(dim=1, keepdim=True) + self.eps)
+
+    return score, P
+
+
+def soft_align_query(A_raw, P):
+  """
+  Smear the raw Query amino-acid distribution onto the Reference coordinate
+  system through the soft alignment matrix P.
+
+  A_raw : [B, Lq, 20]  raw Query sequence (one-hot or soft distribution)
+  P     : [B, Lq, Lr]  soft alignment matrix (query -> ref)
+  returns [B, Lr, 20]  soft amino-acid distribution in Reference coordinates (A_tilde)
+  """
+  # A_tilde[j, d] = sum_u P[u, j] * A_raw[u, d]
+  return torch.einsum("buj,bud->bjd", P, A_raw)