boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/model/loss/diffusion.py

@@ -0,0 +1,171 @@
+# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
+
+from einops import einsum
+import torch
+import torch.nn.functional as F
+
+
+def weighted_rigid_align(
+    true_coords,
+    pred_coords,
+    weights,
+    mask,
+):
+    """Compute weighted alignment.
+
+    Parameters
+    ----------
+    true_coords: torch.Tensor
+        The ground truth atom coordinates
+    pred_coords: torch.Tensor
+        The predicted atom coordinates
+    weights: torch.Tensor
+        The weights for alignment
+    mask: torch.Tensor
+        The atoms mask
+
+    Returns
+    -------
+    torch.Tensor
+        Aligned coordinates
+
+    """
+
+    batch_size, num_points, dim = true_coords.shape
+    weights = (mask * weights).unsqueeze(-1)
+
+    # Compute weighted centroids
+    true_centroid = (true_coords * weights).sum(dim=1, keepdim=True) / weights.sum(
+        dim=1, keepdim=True
+    )
+    pred_centroid = (pred_coords * weights).sum(dim=1, keepdim=True) / weights.sum(
+        dim=1, keepdim=True
+    )
+
+    # Center the coordinates
+    true_coords_centered = true_coords - true_centroid
+    pred_coords_centered = pred_coords - pred_centroid
+
+    if num_points < (dim + 1):
+        print(
+            "Warning: The size of one of the point clouds is <= dim+1. "
+            + "`WeightedRigidAlign` cannot return a unique rotation."
+        )
+
+    # Compute the weighted covariance matrix
+    cov_matrix = einsum(
+        weights * pred_coords_centered, true_coords_centered, "b n i, b n j -> b i j"
+    )
+
+    # Compute the SVD of the covariance matrix, required float32 for svd and determinant
+    original_dtype = cov_matrix.dtype
+    cov_matrix_32 = cov_matrix.to(dtype=torch.float32)
+    U, S, V = torch.linalg.svd(
+        cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
+    )
+    V = V.mH
+
+    # Catch ambiguous rotation by checking the magnitude of singular values
+    if (S.abs() <= 1e-15).any() and not (num_points < (dim + 1)):
+        print(
+            "Warning: Excessively low rank of "
+            + "cross-correlation between aligned point clouds. "
+            + "`WeightedRigidAlign` cannot return a unique rotation."
+        )
+
+    # Compute the rotation matrix
+    rot_matrix = torch.einsum("b i j, b k j -> b i k", U, V).to(dtype=torch.float32)
+
+    # Ensure proper rotation matrix with determinant 1
+    F = torch.eye(dim, dtype=cov_matrix_32.dtype, device=cov_matrix.device)[
+        None
+    ].repeat(batch_size, 1, 1)
+    F[:, -1, -1] = torch.det(rot_matrix)
+    rot_matrix = einsum(U, F, V, "b i j, b j k, b l k -> b i l")
+    rot_matrix = rot_matrix.to(dtype=original_dtype)
+
+    # Apply the rotation and translation
+    aligned_coords = (
+        einsum(true_coords_centered, rot_matrix, "b n i, b j i -> b n j")
+        + pred_centroid
+    )
+    aligned_coords.detach_()
+
+    return aligned_coords
+
+
+def smooth_lddt_loss(
+    pred_coords,
+    true_coords,
+    is_nucleotide,
+    coords_mask,
+    nucleic_acid_cutoff: float = 30.0,
+    other_cutoff: float = 15.0,
+    multiplicity: int = 1,
+):
+    """Compute weighted alignment.
+
+    Parameters
+    ----------
+    pred_coords: torch.Tensor
+        The predicted atom coordinates
+    true_coords: torch.Tensor
+        The ground truth atom coordinates
+    is_nucleotide: torch.Tensor
+        The weights for alignment
+    coords_mask: torch.Tensor
+        The atoms mask
+    nucleic_acid_cutoff: float
+        The nucleic acid cutoff
+    other_cutoff: float
+        The non nucleic acid cutoff
+    multiplicity: int
+        The multiplicity
+    Returns
+    -------
+    torch.Tensor
+        Aligned coordinates
+
+    """
+    B, N, _ = true_coords.shape
+    true_dists = torch.cdist(true_coords, true_coords)
+    is_nucleotide = is_nucleotide.repeat_interleave(multiplicity, 0)
+
+    coords_mask = coords_mask.repeat_interleave(multiplicity, 0)
+    is_nucleotide_pair = is_nucleotide.unsqueeze(-1).expand(
+        -1, -1, is_nucleotide.shape[-1]
+    )
+
+    mask = (
+        is_nucleotide_pair * (true_dists < nucleic_acid_cutoff).float()
+        + (1 - is_nucleotide_pair) * (true_dists < other_cutoff).float()
+    )
+    mask = mask * (1 - torch.eye(pred_coords.shape[1], device=pred_coords.device))
+    mask = mask * (coords_mask.unsqueeze(-1) * coords_mask.unsqueeze(-2))
+
+    # Compute distances between all pairs of atoms
+    pred_dists = torch.cdist(pred_coords, pred_coords)
+    dist_diff = torch.abs(true_dists - pred_dists)
+
+    # Compute epsilon values
+    eps = (
+        (
+            (
+                F.sigmoid(0.5 - dist_diff)
+                + F.sigmoid(1.0 - dist_diff)
+                + F.sigmoid(2.0 - dist_diff)
+                + F.sigmoid(4.0 - dist_diff)
+            )
+            / 4.0
+        )
+        .view(multiplicity, B // multiplicity, N, N)
+        .mean(dim=0)
+    )
+
+    # Calculate masked averaging
+    eps = eps.repeat_interleave(multiplicity, 0)
+    num = (eps * mask).sum(dim=(-1, -2))
+    den = mask.sum(dim=(-1, -2)).clamp(min=1)
+    lddt = num / den
+
+    return 1.0 - lddt.mean()

boltz/model/loss/diffusionv2.py

@@ -0,0 +1,134 @@
+# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
+
+import einx
+import torch
+import torch.nn.functional as F
+from einops import einsum, rearrange
+
+
+def weighted_rigid_align(
+    true_coords,  # Float['b n 3'],       #  true coordinates
+    pred_coords,  # Float['b n 3'],       # predicted coordinates
+    weights,  # Float['b n'],             # weights for each atom
+    mask,  # Bool['b n'] | None = None    # mask for variable lengths
+):  # -> Float['b n 3']:
+    """Algorithm 28 : note there is a problem with the pseudocode in the paper where predicted and
+    GT are swapped in algorithm 28, but correct in equation (2)."""
+
+    batch_size, num_points, dim = true_coords.shape
+    weights = (mask * weights).unsqueeze(-1)
+
+    # Compute weighted centroids
+    true_centroid = (true_coords * weights).sum(dim=1, keepdim=True) / weights.sum(
+        dim=1, keepdim=True
+    )
+    pred_centroid = (pred_coords * weights).sum(dim=1, keepdim=True) / weights.sum(
+        dim=1, keepdim=True
+    )
+
+    # Center the coordinates
+    true_coords_centered = true_coords - true_centroid
+    pred_coords_centered = pred_coords - pred_centroid
+
+    if torch.any(mask.sum(dim=-1) < (dim + 1)):
+        print(
+            "Warning: The size of one of the point clouds is <= dim+1. "
+            + "`WeightedRigidAlign` cannot return a unique rotation."
+        )
+
+    # Compute the weighted covariance matrix
+    cov_matrix = einsum(
+        weights * pred_coords_centered, true_coords_centered, "b n i, b n j -> b i j"
+    )
+
+    # Compute the SVD of the covariance matrix, required float32 for svd and determinant
+    original_dtype = cov_matrix.dtype
+    cov_matrix_32 = cov_matrix.to(dtype=torch.float32)
+
+    U, S, V = torch.linalg.svd(
+        cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
+    )
+    V = V.mH
+
+    # Catch ambiguous rotation by checking the magnitude of singular values
+    if (S.abs() <= 1e-15).any() and not (num_points < (dim + 1)):
+        print(
+            "Warning: Excessively low rank of "
+            + "cross-correlation between aligned point clouds. "
+            + "`WeightedRigidAlign` cannot return a unique rotation."
+        )
+
+    # Compute the rotation matrix
+    rot_matrix = torch.einsum("b i j, b k j -> b i k", U, V).to(dtype=torch.float32)
+
+    # Ensure proper rotation matrix with determinant 1
+    F = torch.eye(dim, dtype=cov_matrix_32.dtype, device=cov_matrix.device)[
+        None
+    ].repeat(batch_size, 1, 1)
+    F[:, -1, -1] = torch.det(rot_matrix)
+    rot_matrix = einsum(U, F, V, "b i j, b j k, b l k -> b i l")
+    rot_matrix = rot_matrix.to(dtype=original_dtype)
+
+    # Apply the rotation and translation
+    aligned_coords = (
+        einsum(true_coords_centered, rot_matrix, "b n i, b j i -> b n j")
+        + pred_centroid
+    )
+    aligned_coords.detach_()
+
+    return aligned_coords
+
+
+def smooth_lddt_loss(
+    pred_coords,  # Float['b n 3'],
+    true_coords,  # Float['b n 3'],
+    is_nucleotide,  # Bool['b n'],
+    coords_mask,  # Bool['b n'] | None = None,
+    nucleic_acid_cutoff: float = 30.0,
+    other_cutoff: float = 15.0,
+    multiplicity: int = 1,
+):  # -> Float['']:
+    """Algorithm 27
+    pred_coords: predicted coordinates
+    true_coords: true coordinates
+    Note: for efficiency pred_coords is the only one with the multiplicity expanded
+    TODO: add weighing which overweight the smooth lddt contribution close to t=0 (not present in the paper)
+    """
+    lddt = []
+    for i in range(true_coords.shape[0]):
+        true_dists = torch.cdist(true_coords[i], true_coords[i])
+
+        is_nucleotide_i = is_nucleotide[i // multiplicity]
+        coords_mask_i = coords_mask[i // multiplicity]
+
+        is_nucleotide_pair = is_nucleotide_i.unsqueeze(-1).expand(
+            -1, is_nucleotide_i.shape[-1]
+        )
+
+        mask = is_nucleotide_pair * (true_dists < nucleic_acid_cutoff).float()
+        mask += (1 - is_nucleotide_pair) * (true_dists < other_cutoff).float()
+        mask *= 1 - torch.eye(pred_coords.shape[1], device=pred_coords.device)
+        mask *= coords_mask_i.unsqueeze(-1)
+        mask *= coords_mask_i.unsqueeze(-2)
+
+        valid_pairs = mask.nonzero()
+        true_dists_i = true_dists[valid_pairs[:, 0], valid_pairs[:, 1]]
+
+        pred_coords_i1 = pred_coords[i, valid_pairs[:, 0]]
+        pred_coords_i2 = pred_coords[i, valid_pairs[:, 1]]
+        pred_dists_i = F.pairwise_distance(pred_coords_i1, pred_coords_i2)
+
+        dist_diff_i = torch.abs(true_dists_i - pred_dists_i)
+
+        eps_i = (
+            F.sigmoid(0.5 - dist_diff_i)
+            + F.sigmoid(1.0 - dist_diff_i)
+            + F.sigmoid(2.0 - dist_diff_i)
+            + F.sigmoid(4.0 - dist_diff_i)
+        ) / 4.0
+
+        lddt_i = eps_i.sum() / (valid_pairs.shape[0] + 1e-5)
+        lddt.append(lddt_i)
+
+    # average over batch & multiplicity
+    return 1.0 - torch.stack(lddt, dim=0).mean(dim=0)

boltz/model/loss/distogram.py

@@ -0,0 +1,48 @@
+import torch
+from torch import Tensor
+
+
+def distogram_loss(
+    output: dict[str, Tensor],
+    feats: dict[str, Tensor],
+) -> tuple[Tensor, Tensor]:
+    """Compute the  distogram loss.
+
+    Parameters
+    ----------
+    output : Dict[str, Tensor]
+        Output of the model
+    feats : Dict[str, Tensor]
+        Input features
+
+    Returns
+    -------
+    Tensor
+        The globally averaged loss.
+    Tensor
+        Per example loss.
+
+    """
+    # Get predicted distograms
+    pred = output["pdistogram"]
+
+    # Compute target distogram
+    target = feats["disto_target"]
+
+    # Combine target mask and padding mask
+    mask = feats["token_disto_mask"]
+    mask = mask[:, None, :] * mask[:, :, None]
+    mask = mask * (1 - torch.eye(mask.shape[1])[None]).to(pred)
+
+    # Compute the distogram loss
+    errors = -1 * torch.sum(
+        target * torch.nn.functional.log_softmax(pred, dim=-1),
+        dim=-1,
+    )
+    denom = 1e-5 + torch.sum(mask, dim=(-1, -2))
+    mean = errors * mask
+    mean = torch.sum(mean, dim=-1)
+    mean = mean / denom[..., None]
+    batch_loss = torch.sum(mean, dim=-1)
+    global_loss = torch.mean(batch_loss)
+    return global_loss, batch_loss

boltz/model/loss/distogramv2.py

@@ -0,0 +1,105 @@
+import torch
+from torch import Tensor
+
+
+def distogram_loss(
+    output: dict[str, Tensor],
+    feats: dict[str, Tensor],
+    aggregate_distogram: bool = True,
+) -> tuple[Tensor, Tensor]:
+    """Compute the  distogram loss.
+
+    Parameters
+    ----------
+    output : Dict[str, Tensor]
+        Output of the model
+    feats : Dict[str, Tensor]
+        Input features
+
+    Returns
+    -------
+    Tensor
+        The globally averaged loss.
+    Tensor
+        Per example loss.
+
+    """
+    with torch.autocast("cuda", enabled=False):
+        # Get predicted distograms
+        pred = output["pdistogram"].float()  # (B, L, L, num_distograms, disto_bins)
+        D = pred.shape[3]  # num_distograms  # noqa: N806
+        assert len(pred.shape) == 5  # noqa: PLR2004
+
+        # Compute target distogram
+        target = feats["disto_target"]  # (B, L, L, K, disto_bins)
+        assert len(target.shape) == 5  # noqa: PLR2004
+
+        if aggregate_distogram:
+            msg = "Cannot aggregate GT distogram when num_distograms > 1"
+            assert pred.shape[3] == 1, msg
+
+            pred = pred.squeeze(3)  # (B, L, L, disto_bins)
+
+            # Aggregate distogram over K conformers
+            target = target.sum(dim=3)  # (B, L, L, disto_bins)
+
+            # Normalize distogram
+            P = target / target.sum(-1)[..., None].clamp(min=1)  # noqa: N806
+
+            # Combine target mask and padding mask
+            mask = feats["token_disto_mask"]
+            mask = mask[:, None, :] * mask[:, :, None]
+            mask = mask * (1 - torch.eye(mask.shape[1])[None]).to(pred)
+
+            # Compute the distogram loss
+            log_Q = torch.nn.functional.log_softmax(pred, dim=-1)  # noqa: N806
+            errors = -1 * torch.sum(
+                P * log_Q,
+                dim=-1,
+            )
+            denom = 1e-5 + torch.sum(mask, dim=(-1, -2))
+            mean = errors * mask
+            mean = torch.sum(mean, dim=-1)
+            mean = mean / denom[..., None]
+            batch_loss = torch.sum(mean, dim=-1)
+            global_loss = torch.mean(batch_loss)
+        else:
+            # We want to compute the loss for each pair of conformer K and predicted
+            # distogram
+
+            # Loop through conformers and compute the loss
+            batch_loss = []
+            for k in range(target.shape[3]):
+                # Get the target distogram for conformer k
+                # (B, L, L, K, disto_bins) -> (B, L, L, D, disto_bins)
+                P_k = target[:, :, :, k : k + 1, :].repeat_interleave(D, dim=3)  # noqa: N806
+
+                # Compute the distogram loss to all predicted distograms
+                log_Q = torch.nn.functional.log_softmax(pred, dim=-1)  # noqa: N806
+                errors = -1 * torch.sum(
+                    P_k * log_Q,
+                    dim=-1,
+                )  # (B, L, L, D)
+
+                # Compute mask
+                mask = feats["token_disto_mask"]
+                mask = mask[:, None, :] * mask[:, :, None]
+                mask = mask * (1 - torch.eye(mask.shape[1])[None]).to(pred)
+                mask = mask.unsqueeze(-1).repeat_interleave(D, -1)  # (B, L, L, D)
+
+                denom = 1e-5 + torch.sum(mask, dim=(-2, -3))  # (B, D)
+                mean = errors * mask
+                mean = torch.sum(mean, dim=-2)  # (B, L, D)
+                mean = mean / denom[..., None, :]
+                b_loss = torch.sum(mean, dim=-2)  # (B, D)
+
+                batch_loss.append(b_loss)
+
+            batch_loss = torch.stack(batch_loss, dim=1)  # (B, K, D)
+
+            # Compute the batch loss by taking the min over the predicted distograms
+            # and the average across conformers
+            batch_loss = torch.min(batch_loss, dim=-1).values.mean(dim=1)
+            global_loss = torch.mean(batch_loss)
+
+        return global_loss, batch_loss