rc-foundry 0.1.1__py3-none-any.whl

rf3/loss/loss.py

@@ -0,0 +1,179 @@
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def calc_ddihedralmse_dxyz(a, b, c, d, true_dih, eps=1e-6):
+    """
+    Calculates the gradient of the dihedral angle with respect to the xyz coordinates using the closed form derivative.
+    a, b, c, and d are atoms participating in the chiral center. true_dih is the true dihedral angle.
+
+    Unlike the original implementation, this does NOT use autograd.
+    """
+    # I need to reshape this from n_symm, batch, n, 3 to n_symm * batch * n, 3)
+    og_shape = a.shape
+    # Expand the dihedral by the batch dimension to match n_atoms*batchs
+    true_dih = true_dih.unsqueeze(0).repeat(a.shape[0], 1)
+    a = a.view(-1, 3)
+    b = b.view(-1, 3)
+    c = c.view(-1, 3)
+    d = d.view(-1, 3)
+    true_dih = true_dih.view(-1)
+
+    batch_size = a.shape[0]  # Support for batch size
+    I = (
+        torch.eye(3).unsqueeze(0).repeat(batch_size, 1, 1).to(a.device)
+    )  # Make batch-aware identity matrix
+
+    # Compute b0, b1, b2
+    b0 = a - b
+    b1 = c - b
+    b2 = d - c
+
+    # Normalize b1
+    b1_norm = torch.norm(b1, dim=-1, keepdim=True)
+    b1n = b1 / (b1_norm + eps)
+
+    # Compute orthogonal components v and w
+    v = b0 - torch.sum(b0 * b1n, dim=-1, keepdim=True) * b1n
+    w = b2 - torch.sum(b2 * b1n, dim=-1, keepdim=True) * b1n
+
+    # Dihedral components x and y
+    x = torch.sum(v * w, dim=-1)
+    y = torch.sum(torch.cross(b1n, v, dim=-1) * w, dim=-1)
+
+    # Dihedral angle
+    dih = torch.atan2(y + eps, x + eps)
+
+    # Compute MSE loss and manual gradients
+    # mse_loss = torch.mean(torch.square(dih - true_dih))
+    # mse_loss = torch.sum(torch.square(dih - true_dih))
+
+    # Define matrices and gradients, adapted for batch
+    db0_db = -I
+    db1_db = -I
+    db1_dc = I
+    db2_dc = -I
+    db0_da = I
+    db2_dd = I
+    # dmse_ddih = 2 * (dih - true_dih) / batch_size
+    dmse_ddih = 2 * (dih - true_dih)
+    ddih_dx = -y / (x**2 + y**2 + eps)
+    ddih_dy = x / (x**2 + y**2 + eps)
+    dy_dv = -torch.cross(b1n, w, dim=-1)
+    dy_dw = torch.cross(b1n, v, dim=-1)
+    dx_dv = w
+    dx_dw = v
+
+    dw_db1n = -torch.sum(b2 * b1n, dim=-1, keepdim=True).unsqueeze(-1) * I - torch.bmm(
+        b2.unsqueeze(-1), b1n.unsqueeze(1)
+    )
+
+    db1n_db1 = (b1_norm + eps).unsqueeze(-1) * I / (b1_norm**2 + eps).unsqueeze(
+        -1
+    ) - torch.bmm(b1.unsqueeze(-1), b1.unsqueeze(1)) / (b1_norm**2 + eps).unsqueeze(-1)
+
+    dv_db1n = -torch.sum(b0 * b1n, dim=-1, keepdim=True).unsqueeze(-1) * I - torch.bmm(
+        b0.unsqueeze(-1), b1n.unsqueeze(1)
+    )
+    dv_db0 = I - torch.bmm(b1n.unsqueeze(-1), b1n.unsqueeze(1))
+    dw_db2 = I - torch.bmm(b1n.unsqueeze(-1), b1n.unsqueeze(1))
+
+    # Adjust sizes now for efficiency
+    ddih_dx = ddih_dx.view(-1, 1, 1)
+    ddih_dy = ddih_dy.view(-1, 1, 1)
+    dmse_ddih = dmse_ddih.view(-1, 1, 1)
+    dx_dv = dx_dv.unsqueeze(1)
+    dx_dw = dx_dw.unsqueeze(1)
+    dy_dv = dy_dv.unsqueeze(1)
+    dy_dw = dy_dw.unsqueeze(1)
+
+    # Gradient computations
+    # wrt a
+    dv_da = torch.matmul(dv_db0, db0_da)
+    ddih_da = torch.bmm((ddih_dx * dx_dv), dv_da) + torch.bmm((ddih_dy * dy_dv), dv_da)
+    dmse_da = torch.bmm(dmse_ddih, ddih_da)
+
+    # wrt b
+    db1n_db = torch.matmul(db1n_db1, db1_db)
+    dv_db = torch.matmul(dv_db0, db0_db) + torch.matmul(
+        dv_db1n.transpose(-1, -2), db1n_db
+    )
+    dw_db = torch.matmul(dw_db1n.transpose(-1, -2), db1n_db)
+    dx_db = torch.bmm(dx_dv, dv_db) + torch.bmm(dx_dw, dw_db)
+    dy_db = torch.bmm(dy_dv, dv_db) + torch.bmm(dy_dw, dw_db)
+    ddih_db = torch.bmm(ddih_dx, dx_db) + torch.bmm(ddih_dy, dy_db)
+    dmse_db = torch.bmm(dmse_ddih, ddih_db)
+
+    # wrt c
+    db1n_dc = torch.matmul(db1n_db1, db1_dc)
+    dv_dc = torch.matmul(dv_db1n.transpose(-1, -2), db1n_dc)
+    dw_dc = torch.matmul(dw_db2, db2_dc) + torch.matmul(
+        dw_db1n.transpose(-1, -2), db1n_dc
+    )
+    dx_dc = torch.bmm(dx_dv, dv_dc) + torch.bmm(dx_dw, dw_dc)
+    dy_dc = torch.bmm(dy_dv, dv_dc) + torch.bmm(dy_dw, dw_dc)
+    ddih_dc = torch.bmm(ddih_dx, dx_dc) + torch.bmm(ddih_dy, dy_dc)
+    dmse_dc = torch.bmm(dmse_ddih, ddih_dc)
+
+    # wrt d
+    dw_dd = torch.matmul(dw_db2, db2_dd)
+    ddih_dd = torch.bmm((ddih_dx * dx_dw), dw_dd) + torch.bmm((ddih_dy * dy_dw), dw_dd)
+    dmse_dd = torch.bmm(dmse_ddih, ddih_dd)
+
+    # Reshape gradients back to original shape and prep for cat
+    dmse_da = dmse_da.view(og_shape).unsqueeze(-2)
+    dmse_db = dmse_db.view(og_shape).unsqueeze(-2)
+    dmse_dc = dmse_dc.view(og_shape).unsqueeze(-2)
+    dmse_dd = dmse_dd.view(og_shape).unsqueeze(-2)
+
+    grads = torch.cat([dmse_da, dmse_db, dmse_dc, dmse_dd], dim=-2)
+    return grads
+
+
+def calc_chiral_grads_flat_impl(
+    xyz, chiral_centers, chiral_center_dihedral_angles, no_grad_on_chiral_center
+):
+    """
+    Calculates the gradient of the chiral centers with respect to the xyz coordinates using the closed form derivative.
+    Args:
+    xyz: torch.Tensor, shape (batch, n_atoms, 3)
+    chiral_centers: torch.Tensor, shape (long) (n_centers, 4)
+    chiral_center_dihedral_angles: torch.Tensor, shape (float) (n_centers, 1)
+
+    Returns:
+    grads: torch.Tensor, shape (batch, n_atoms, 3)
+    """
+    # (We want to track the gradient of the dihedral angle loss with respect to the xyz coordinates)
+    xyz.requires_grad_(True)
+
+    # Edge case: No chiral centers, return zero gradients
+    if chiral_centers.shape[0] == 0:
+        return torch.zeros(xyz.shape, device=xyz.device)
+
+    # Get the coordinates of the four atoms that make up the chiral center
+    chiral_dih = xyz[:, chiral_centers, :]
+
+    # Calculate the gradient of the dihedral angle loss with respect to the xyz coordinates
+    grads = torch.zeros_like(xyz).to(xyz.device)
+    chiral_grads = calc_ddihedralmse_dxyz(
+        chiral_dih[..., 0, :],
+        chiral_dih[..., 1, :],
+        chiral_dih[..., 2, :],
+        chiral_dih[..., 3, :],
+        chiral_center_dihedral_angles,
+    )  # n_center, 4, 3
+
+    if no_grad_on_chiral_center:
+        chiral_grads[:, :, 0] = 0.0  # no gradient on chiral center
+
+    # back to atom
+    grads.index_add_(
+        1,
+        chiral_centers.flatten(),
+        chiral_grads.flatten(start_dim=1, end_dim=2),
+    )
+
+    return grads

rf3/metrics/chiral.py

@@ -0,0 +1,179 @@
+import torch
+from atomworks.io.transforms.atom_array import ensure_atom_array_stack
+from atomworks.ml.transforms.af3_reference_molecule import (
+    get_af3_reference_molecule_features,
+)
+from atomworks.ml.transforms.chirals import add_af3_chiral_features
+from atomworks.ml.transforms.rdkit_utils import get_rdkit_chiral_centers
+from beartype.typing import Any
+from biotite.structure import AtomArray, AtomArrayStack
+from jaxtyping import Bool, Float
+from rf3.kinematics import get_dih
+
+from foundry.metrics.metric import Metric
+
+
+def calc_chiral_metrics_masked(
+    pred: Float[torch.Tensor, "B L ... 3"],
+    chirals: Float[torch.Tensor, "n_chiral 5"],
+    mask: Bool[torch.Tensor, "I"],
+):
+    """Calculate metrics for chiral centers, including:
+        - n_chiral_centers (B): number of chiral centers in the structure
+        - chiral_loss_mean (B): mean of the squared errors of chiral angles
+        - percent_correct_chirality (B): percentage of correctly predicted chiral centers
+
+    Args:
+        pred: predicted coords (B, L, :, 3)
+        chirals: True coords (nchiral, 5); skip if 0 chiral sites. 5 dimension are indices for 4 atoms that make dihedral and the ideal angle they should form
+        mask: Boolean mask of shape (I) indicating valid positions (e.g., non-NaN coordinates, desired residue type)
+
+    Returns:
+        chiral_loss_sum: sum of squared errors of chiral angles (B)
+        n_chiral_centers: number of chiral centers in the structure
+        percent_correct_chirality: percentage of correctly predicted chiral centers (B)
+    """
+    if not chirals.numel() or not mask.sum():
+        # ... no chiral centers; exit
+        return {}
+
+    # ... get the coordinates of all four atoms involved in each chiral center
+    chiral_dih = pred[
+        :, chirals[..., :-1].long()
+    ]  # (n_chiral 5) -> (B, n_chiral, 4, 3)
+
+    # ... for each chiral center, compute the dihedral angle
+    pred_dih = get_dih(
+        chiral_dih[..., 0, :],
+        chiral_dih[..., 1, :],
+        chiral_dih[..., 2, :],
+        chiral_dih[..., 3, :],
+    )  # [B, n_chiral]
+
+    # ... total chiral loss (sum of squared errors)
+    diff = pred_dih - chirals[..., -1]  # [B, n_chiral]
+    is_correct_chirality = torch.sign(pred_dih) == torch.sign(
+        chirals[..., -1]
+    )  # [B, n_chiral]
+
+    # To avoid over-counting chirals, we should only keep one "row" for each chiral center (rather than enumerating all orderings)
+    inf_tensor = torch.tensor(
+        [-float("inf")], device=chirals.device, dtype=chirals.dtype
+    )
+    shifted = torch.cat([inf_tensor, chirals[:-1, 0]], dim=0)  # Shape [24]
+    first_occurence_mask = chirals[:, 0] != shifted
+
+    is_valid_chiral_center = mask[chirals[..., :-1].long()].all(
+        dim=-1
+    )  # [L] -> [n_chiral] (a chiral center is valid iff ALL atoms are included)
+    # ... and only keep the first occurrence of each chiral center
+    is_valid_chiral_center = is_valid_chiral_center & first_occurence_mask
+
+    percent_correct_chirality = (is_correct_chirality[:, is_valid_chiral_center]).sum(
+        dim=-1
+    ) / is_valid_chiral_center.sum(dim=-1)  # [B]
+
+    l = torch.square(diff[:, is_valid_chiral_center]).sum(dim=-1)  # [B]
+
+    return {
+        "chiral_loss_mean": l / mask.sum(),  # [B]
+        "n_chiral_centers": is_valid_chiral_center.sum(dim=-1),  # [B]
+        "percent_correct_chirality": percent_correct_chirality,  # [B]
+    }
+
+
+def compute_chiral_metrics(
+    predicted_atom_array_stack: AtomArrayStack | AtomArray,
+    ground_truth_atom_array_stack: AtomArrayStack | AtomArray,
+    chiral_feats: Float[torch.Tensor, "n_chiral 5"] | None = None,
+):
+    """Compute chiral metrics from the predicted and ground truth atom arrays.
+
+    If chiral features are not directly provided, they will be re-computed from the AtomArrays.
+
+    Returns:
+        dict: Dictionary containing chiral metrics, separated for polymers and non-polymers. The metrics are:
+            - n_chiral_centers: number of chiral centers in the structure
+            - chiral_loss_mean: mean of the squared errors of chiral angles
+            - percent_correct_chirality: percentage of correctly predicted chiral centers
+    """
+    predicted_atom_array_stack = ensure_atom_array_stack(predicted_atom_array_stack)
+    ground_truth_atom_array_stack = ensure_atom_array_stack(
+        ground_truth_atom_array_stack
+    )
+
+    chiral_metrics = {}
+    # (Choose the first model  - chirality does not depend on our data augmentation)
+    ground_truth_atom_array = ground_truth_atom_array_stack[0]
+
+    if chiral_feats is None:
+        # Generate chiral features if not provided
+        _, rdkit_mols = get_af3_reference_molecule_features(ground_truth_atom_array)
+        chiral_centers = get_rdkit_chiral_centers(rdkit_mols)
+        chiral_feats = add_af3_chiral_features(
+            ground_truth_atom_array, chiral_centers, rdkit_mols
+        )
+
+    X_L = torch.from_numpy(predicted_atom_array_stack.coord).to(
+        device=chiral_feats.device
+    )
+
+    categories = ["polymer", "non_polymer"]
+    _polymer_mask = torch.from_numpy(ground_truth_atom_array.is_polymer).to(
+        device=chiral_feats.device
+    )
+    # (Only consider non-NaN coordinates in the ground truth, since otherwise we can't compare dihedral angles)
+    _valid_coord_mask = ~torch.isnan(
+        torch.from_numpy(ground_truth_atom_array.coord)
+    ).any(dim=1).to(device=chiral_feats.device)
+    masks = [_polymer_mask, ~_polymer_mask]
+
+    for category, mask in zip(categories, masks):
+        # ... compute the chiral loss, given the mask
+        result = calc_chiral_metrics_masked(
+            X_L,
+            chiral_feats,
+            mask=mask & _valid_coord_mask,
+        )
+
+        if not result:
+            # No chiral centers - skip
+            continue
+
+        # ... store the metric results, meaned over the diffusion batch
+        if result["n_chiral_centers"] > 0:
+            chiral_metrics[f"{category}_n_chiral_centers"] = result[
+                "n_chiral_centers"
+            ].item()
+            chiral_metrics[f"{category}_chiral_loss_mean"] = (
+                result["chiral_loss_mean"].mean().item()
+            )
+            chiral_metrics[f"{category}_percent_correct_chirality"] = (
+                result["percent_correct_chirality"].mean().item()
+            )
+
+    return chiral_metrics
+
+
+class ChiralLoss(Metric):
+    @property
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
+        return {
+            "predicted_atom_array_stack": "predicted_atom_array_stack",
+            "ground_truth_atom_array_stack": "ground_truth_atom_array_stack",
+            "chiral_feats": ("network_input", "f", "chiral_feats"),
+        }
+
+    def compute(
+        self,
+        predicted_atom_array_stack: AtomArrayStack | AtomArray,
+        ground_truth_atom_array_stack: AtomArrayStack | AtomArray,
+        chiral_feats: Float[torch.Tensor, "n_chiral 5"] = None,
+    ):
+        chiral_metrics = compute_chiral_metrics(
+            predicted_atom_array_stack,
+            ground_truth_atom_array_stack,
+            chiral_feats=chiral_feats,
+        )
+
+        return chiral_metrics

rf3/metrics/clashing_chains.py

@@ -0,0 +1,68 @@
+import itertools
+from typing import Any
+
+import torch
+from biotite.structure import AtomArrayStack
+
+from foundry.metrics.metric import Metric
+
+
+class CountClashingChains(Metric):
+    @property
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
+        return {
+            "X_L": ("network_output", "X_L"),
+            "predicted_atom_array_stack": ("predicted_atom_array_stack",),
+        }
+
+    def compute(
+        self,
+        X_L: torch.Tensor,
+        predicted_atom_array_stack: AtomArrayStack,
+    ) -> dict[str, float]:
+        """Compute the predicted interface TM-score (IPTM) from the predicted aligned error (PAE).
+        Args:
+            X_L: Predicted aligned error tensor.
+            predicted_atom_array_stack: AtomArrayStack containing the predicted structure.
+        Returns:
+            clash_count: Computed clashing chains count.
+        """
+        D = X_L.shape[0]
+        MIN_CLASH_DISTANCE = 1.1  # Minimum distance to consider a clash
+        # Count the number of clashing chains
+        pn_units = set(predicted_atom_array_stack.pn_unit_id)
+
+        has_clash = torch.zeros((D), dtype=torch.bool)
+        for chain_i, chain_j in itertools.combinations(pn_units, 2):
+            # check to make sure they are both polymer chains
+
+            chain_i_atoms = predicted_atom_array_stack[
+                :, predicted_atom_array_stack.pn_unit_id == chain_i
+            ]
+            chain_j_atoms = predicted_atom_array_stack[
+                :, predicted_atom_array_stack.pn_unit_id == chain_j
+            ]
+            if not chain_i_atoms[0, 0].is_polymer or not chain_j_atoms[0, 0].is_polymer:
+                continue
+
+            distances = torch.cdist(
+                torch.from_numpy(chain_i_atoms.coord),
+                torch.from_numpy(chain_j_atoms.coord),
+            )
+            num_clashes = (distances < MIN_CLASH_DISTANCE).sum(dim=-1).sum(dim=-1)
+            has_clash_pair = (num_clashes > 100) | (
+                num_clashes
+                / (
+                    max(chain_i_atoms.coord.shape[0], chain_j_atoms.coord.shape[0])
+                    + 1e-6
+                )
+                > 0.5
+            )
+            has_clash = torch.logical_or(has_clash, has_clash_pair)
+        assert has_clash.shape == (D,)
+        # unpack the batch dimension into separate keys in the output dictionary
+
+        clash_count_per_batch = {
+            f"has_clash_{i}": int(has_clash[i]) for i in range(len(has_clash))
+        }
+        return clash_count_per_batch