rc-foundry 0.1.1__py3-none-any.whl

rf3/metrics/metric_utils.py

@@ -0,0 +1,192 @@
+from itertools import combinations
+from typing import Union
+
+import numpy as np
+import torch
+from jaxtyping import Bool, Float
+from numpy.typing import NDArray
+
+
+def find_bin_midpoints(
+    max_distance: float, num_bins: int, device: Union[str, torch.device] = "cpu"
+) -> Float[torch.Tensor, "num_bins"]:
+    """
+    Find the bin midpoints for a given binning scheme. Used to find expectation of values when converting binned
+    predictions to unbinned predictions. Assumes the minimum of the schema is 0.
+    Args:
+        max_distance: float, maximum distance
+        num_bins: int, number of bins
+        device: device to run on
+    Returns:
+        pae_midpoints: [num_bins], bin midpoints
+    """
+    bin_size = max_distance / num_bins
+    bins = torch.linspace(
+        bin_size, max_distance - bin_size, num_bins - 1, device=device
+    )
+    midpoints = (bins[1:] + bins[:-1]) / 2
+    midpoints = torch.cat(
+        [(bins[0] - bin_size / 2)[None], midpoints, bins[-1:] + bin_size / 2]
+    )
+
+    return midpoints
+
+
+def unbin_logits(
+    logits: Float[torch.Tensor, "B num_bins L X"], max_distance: float, num_bins: int
+) -> Float[torch.Tensor, "B L L"]:
+    """
+    Unbin the logits to get the matrix
+    Args:
+        logits: [B, num_bins, L, X], binned logits  where X is 23 for plddt and L for pae and pde
+        max_distance: float, maximum distance
+        num_bins: int, number of bins
+    Returns:
+        unbinned: [B, L, L], unbinned matrix
+    """
+    midpoints = find_bin_midpoints(max_distance, num_bins, device=logits.device)
+    probabilities = torch.nn.Softmax(dim=1)(logits).detach().float()
+    unbinned = (probabilities * midpoints[None, :, None, None]).sum(dim=1)
+    return unbinned
+
+
+def create_chainwise_masks_1d(
+    ch_label: NDArray[np.str_], device: Union[str, torch.device] = "cpu"
+) -> dict[str, Bool[torch.Tensor, "L"]]:
+    """
+    Create 1D chainwise masks for a set of chain labels
+    Args:
+        ch_label: np.ndarray [L], chain labels
+        device: torch.device, device to run on
+    Returns:
+        ch_masks: dict, chain maps chain letter to which elements to score for that chain
+    """
+    unique_chains = np.unique(ch_label)
+    ch_masks = {}
+    for chain in unique_chains:
+        indices = torch.from_numpy((ch_label == chain)).to(
+            dtype=torch.bool, device=device
+        )
+        ch_masks[chain] = indices
+    return ch_masks
+
+
+def create_chainwise_masks_2d(
+    ch_label: NDArray[np.str_], device: Union[str, torch.device] = "cpu"
+) -> dict[str, Bool[torch.Tensor, "L L"]]:
+    """
+    Create 2D chainwise masks for a set of chain labels
+    Args:
+        ch_label: np.ndarray [L], chain labels
+        device: torch.device, device to run on
+    Returns:
+        ch_masks: dict, chain maps chain letter to which elements to score for that chain
+    """
+    unique_chains = np.unique(ch_label)
+    ch_masks = {}
+    for chain in unique_chains:
+        indices = torch.from_numpy((ch_label == chain))
+        mask = torch.outer(indices, indices).to(dtype=torch.bool, device=device)
+        ch_masks[chain] = mask
+    return ch_masks
+
+
+def create_interface_masks_2d(
+    ch_label: NDArray[np.str_], device: Union[str, torch.device] = "cpu"
+) -> dict[tuple[str, str], Bool[torch.Tensor, "L L"]]:
+    """
+    Create interface masks for a set of chain labels
+    Args:
+        ch_label: np.ndarray [L], chain labels
+        device: torch.device, device to run on
+    Returns:
+        pairs_to_score: dict mapping chain pairs to boolean masks
+    """
+    unique_chains = np.unique(ch_label)
+    pairs_to_score = {}
+    for chain_i, chain_j in combinations(unique_chains, 2):
+        chain_i_indices = torch.from_numpy((ch_label == chain_i))
+        chain_j_indices = torch.from_numpy((ch_label == chain_j))
+        to_be_scored = torch.outer(chain_i_indices, chain_j_indices).to(
+            dtype=torch.bool, device=device
+        ) + torch.outer(chain_j_indices, chain_i_indices).to(
+            dtype=torch.bool, device=device
+        )
+        pairs_to_score[(chain_i, chain_j)] = to_be_scored
+    return pairs_to_score
+
+
+def compute_mean_over_subsampled_pairs(
+    matrix_to_mean: Float[torch.Tensor, "B L M"],
+    pairs_to_score: Bool[torch.Tensor, "L M"],
+    eps: float = 1e-6,
+) -> Float[torch.Tensor, "B"]:
+    """
+    Compute the mean over a subsample of pairs in a 2d matrix. Returns a tensor with an element for each batch
+    Args:
+        matrix_to_mean: tensor of shape (batch, L, L)
+        pairs_to_score: 2d tensor of shape (L, L) with 1s where pairs should be scored and 0s elsewhere
+        eps: small epsilon value to avoid division by zero
+    Returns:
+        1d tensor of shape (batch,) with the mean over the subsampled pairs for each batch
+    """
+    B, L, M = matrix_to_mean.shape
+    assert matrix_to_mean.shape == (
+        B,
+        L,
+        M,
+    ), "Matrix to mean should be of shape (batch, L, M)"
+    assert pairs_to_score.shape == (L, M), "Pairs to score should be of shape (L, M)"
+    batch = (matrix_to_mean * pairs_to_score).sum(dim=(-1, -2)) / (
+        pairs_to_score.sum() + eps
+    )
+    assert batch.shape == (B,), "Batch should be of shape (batch,)"
+    return batch
+
+
+def compute_min_over_subsampled_pairs(
+    matrix_to_min: Float[torch.Tensor, "B L M"],
+    pairs_to_score: Bool[torch.Tensor, "L M"],
+) -> Float[torch.Tensor, "B"]:
+    """
+    Compute the min over a subsample of pairs in a 2d matrix. Returns a tensor with an element for each batch
+    Args:
+        matrix_to_min: tensor of shape (batch, L, L)
+        pairs_to_score: 2d tensor of shape (L, L) with 1s where pairs should be scored and 0s elsewhere
+    Returns:
+        1d tensor of shape (batch,) with the min over the subsampled pairs for each batch
+    """
+    B, L, M = matrix_to_min.shape
+    assert matrix_to_min.shape == (
+        B,
+        L,
+        M,
+    ), "Matrix to min should be of shape (batch, L, M)"
+    assert pairs_to_score.shape == (L, M), "Pairs to score should be of shape (L, M)"
+    # Use torch.where to efficiently mask without cloning the entire matrix
+    # This broadcasts pairs_to_score across the batch dimension
+    masked_matrix = torch.where(
+        pairs_to_score.bool(),  # condition (L, M) -> broadcasts to (B, L, M)
+        matrix_to_min,  # if True: use original values (B, L, M)
+        torch.tensor(
+            float("inf"), device=matrix_to_min.device, dtype=matrix_to_min.dtype
+        ),  # if False: use inf
+    )
+
+    # Flatten the last two dimensions and compute min across them
+    batch = masked_matrix.view(B, -1).min(dim=-1)[0]
+
+    assert batch.shape == (B,), "Batch should be of shape (batch,)"
+    return batch
+
+
+def spread_batch_into_dictionary(batch: Float[torch.Tensor, "B"]) -> dict[int, float]:
+    """
+    Given a batch of data, create a dictionary with keys as the batch index and value as the corresponding data
+    Args:
+        batch: 1D tensor of shape (B,)
+    Returns:
+        Dictionary mapping batch indices to float values
+    """
+    assert len(batch.shape) == 1, f"Batch should be a 1d tensor, {batch}"
+    return {i: data.item() for i, data in enumerate(batch)}

rf3/metrics/predicted_error.py

@@ -0,0 +1,134 @@
+from typing import Any
+
+import torch
+from rf3.metrics.metric_utils import find_bin_midpoints
+
+from foundry.metrics.metric import Metric
+
+
+def compute_ptm(
+    pae: torch.Tensor,
+    to_calculate: torch.Tensor | None,
+    max_distance: float = 32,
+    bin_count: int = 64,
+):
+    """Compute the predicted TM-score (PTM) from the predicted aligned error (PAE).
+
+    Args:
+        pae: Predicted aligned error tensor.
+        to_calculate: Tensor indicating which residues to calculate PTM for.
+
+    Returns:
+        ptm: Computed predicted TM-score.
+    """
+    D, I = pae.shape[0], pae.shape[1]
+    if to_calculate is None:
+        to_calculate = torch.ones((I, I), dtype=torch.bool, device=pae.device)
+
+    bin_centers = find_bin_midpoints(
+        max_distance, bin_count, device=pae.device
+    )  # TODO: get this from config
+    pae = torch.nn.Softmax(dim=-1)(pae).detach().float()
+    normalization_factor = 1.24 * (max(I, 19) - 15.0) ** (1 / 3) - 1.8
+    denominator = 1 / (1 + (bin_centers / (normalization_factor)) ** 2)
+    pae = pae * denominator[None, None, None, :]  # Broadcast to match pae shape
+
+    pae = pae.sum(dim=-1)  # Sum over the last dimension
+    pae = (pae * to_calculate[None]).sum(dim=-1) / (to_calculate.sum(dim=-1) + 1e-6)
+    ptm = pae.max(dim=-1).values
+    assert ptm.shape == (D,)
+    return ptm
+
+
+class ComputePTM(Metric):
+    @property
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
+        return {
+            "pae": ("network_output", "pae"),
+            "asym_id": ("network_input", "f", "asym_id"),
+        }
+
+    def compute(
+        self,
+        pae: torch.Tensor,
+        asym_id: torch.Tensor,
+    ) -> dict[str, float]:
+        """Compute the predicted TM-score (PTM) from the predicted aligned error (PAE).
+        Args:
+            pae: Predicted aligned error tensor.
+            asym_id: AtomArrayStack containing the predicted structure.
+        Returns:
+            ptm: Computed predicted TM-score.
+        """
+        ptm = compute_ptm(pae, None)
+        # split the batch dimension into separate keys in the output dictionary
+        ptm = ptm.cpu().numpy()
+        ptm = {f"ptm_{i}": ptm[i] for i in range(len(ptm))}
+        return ptm
+
+
+class ComputeIPTM(Metric):
+    @property
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
+        return {
+            "pae": ("network_output", "pae"),
+            "asym_id": ("network_input", "f", "asym_id"),
+            "is_ligand": ("network_input", "f", "is_ligand"),
+        }
+
+    def compute(
+        self,
+        pae: torch.Tensor,
+        asym_id: torch.Tensor,
+        is_ligand: torch.Tensor,
+    ) -> dict[str, float]:
+        """Compute the predicted interface TM-score (iPTM) from the predicted aligned error (PAE).
+        Args:
+            pae: Predicted aligned error tensor.
+            predicted_atom_array_stack: AtomArrayStack containing the predicted structure.
+        Returns:
+            iptm: Computed interface TM-score.
+        """
+        unique, counts = torch.unique(asym_id, return_counts=True)
+        to_calculate = asym_id[None, :] != asym_id[:, None]
+        iptm = compute_ptm(pae, to_calculate)
+
+        # make a protein - ligand mask
+        protein_mask = is_ligand == 0
+        ligand_mask = is_ligand == 1
+        # calculate iptm for protein-protein, protein-ligand, and ligand-ligand interfaces
+        protein_protein_mask = (
+            protein_mask[None, :] & protein_mask[:, None] * to_calculate
+        )
+        protein_ligand_mask = (
+            (protein_mask[None, :] & ligand_mask[:, None])
+            | (ligand_mask[None, :] & protein_mask[:, None])
+        ) * to_calculate
+        ligand_ligand_mask = ligand_mask[None, :] & ligand_mask[:, None] * to_calculate
+        # calculate iptm for each interface type
+        iptm_protein_protein = compute_ptm(pae, protein_protein_mask)
+        iptm_protein_ligand = compute_ptm(pae, protein_ligand_mask)
+        iptm_ligand_ligand = compute_ptm(pae, ligand_ligand_mask)
+
+        # split the batch dimension into separate keys in the output dictionary
+        iptm = iptm.cpu().numpy()
+        iptm = {f"iptm_{i}": iptm[i] for i in range(len(iptm))}
+        iptm_protein_protein = iptm_protein_protein.cpu().numpy()
+        iptm_protein_protein = {
+            f"iptm_protein_protein_{i}": iptm_protein_protein[i]
+            for i in range(len(iptm_protein_protein))
+        }
+        iptm_protein_ligand = iptm_protein_ligand.cpu().numpy()
+        iptm_protein_ligand = {
+            f"iptm_protein_ligand_{i}": iptm_protein_ligand[i]
+            for i in range(len(iptm_protein_ligand))
+        }
+        iptm_ligand_ligand = iptm_ligand_ligand.cpu().numpy()
+        iptm_ligand_ligand = {
+            f"iptm_ligand_ligand_{i}": iptm_ligand_ligand[i]
+            for i in range(len(iptm_ligand_ligand))
+        }
+        iptm.update(iptm_protein_protein)
+        iptm.update(iptm_protein_ligand)
+        iptm.update(iptm_ligand_ligand)
+        return iptm

rf3/metrics/rasa.py

@@ -0,0 +1,108 @@
+import numpy as np
+from atomworks.ml.transforms.sasa import calculate_atomwise_rasa
+from beartype.typing import Any
+from biotite.structure import AtomArrayStack
+
+from foundry.metrics.metric import Metric
+
+
+class UnresolvedRegionRASA(Metric):
+    """
+    This metric computes the RASA score for unresolved regions in a protein structure.
+    The RASA score is defined as the ratio of the solvent-accessible surface area (SASA)
+    of a residue in a protein structure to the SASA of the same residue in an extended conformation.
+    """
+
+    def __init__(
+        self,
+        probe_radius: float = 1.4,
+        atom_radii: str | np.ndarray = "ProtOr",
+        point_number: int = 100,
+        include_resolved: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.probe_radius = probe_radius
+        self.atom_radii = atom_radii
+        self.point_number = point_number
+        self.include_resolved = include_resolved
+
+    @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",),
+        }
+
+    def compute(
+        self,
+        predicted_atom_array_stack: AtomArrayStack,
+        ground_truth_atom_array_stack: AtomArrayStack,
+    ) -> dict[str, Any]:
+        """Compute the RASA score for unresolved regions in a protein structure.
+
+        Args:
+            predicted_atom_array (AtomArray): The input atom array representing the  predicted protein structure.
+            ground_truth_atom_array (AtomArray): The input atom array representing the ground truth protein structure.
+            probe_radius (float, optional): Van-der-Waals radius of the probe in Angstrom. Defaults to 1.4 (for water).
+            atom_radii (str | np.ndarray, optional): Atom radii set to use for calculation. Defaults to "ProtOr".
+            point_number (int, optional): Number of points in the Shrake-Rupley algorithm to sample for calculating SASA. Defaults to 100.
+            include_resolved (bool, optional): Whether to include resolved regions in the RASA score. Defaults to False.
+
+        Returns:
+            dict: A dictionary containing the RASA score and other relevant information.
+        """
+
+        # find unresolved regions
+        # (polymer atoms with occupancy 0.0)
+        atoms_to_score_unresolved = ground_truth_atom_array_stack.is_polymer & (
+            ground_truth_atom_array_stack.occupancy == 0.0
+        )
+
+        # find resolved regions (polymer atoms with occupancy > 0.0)
+        atoms_to_score_resolved = ground_truth_atom_array_stack.is_polymer & (
+            ground_truth_atom_array_stack.occupancy > 0.0
+        )
+
+        unresolved_rasas = []
+        resolved_rasas = []
+
+        # Calculate RASA
+        for atom_array in predicted_atom_array_stack:
+            try:
+                rasa = calculate_atomwise_rasa(
+                    atom_array=atom_array,
+                    probe_radius=self.probe_radius,
+                    atom_radii=self.atom_radii,
+                    point_number=self.point_number,
+                )
+                unresolved_rasas.append(rasa[atoms_to_score_unresolved].mean())
+                if self.include_resolved:
+                    resolved_rasas.append(rasa[atoms_to_score_resolved].mean())
+            except KeyError:
+                unresolved_rasas.append(np.nan)
+                if self.include_resolved:
+                    resolved_rasas.append(np.nan)
+
+        # Calculate the mean RASA scores
+        # Pattern-match other metrics by appending "_i" to the metric name to represent multiple batches
+        # (e.g., "unresolved_polymer_rasa_0", "unresolved_polymer_rasa_1", etc.)
+        unresolved_rasa = np.nanmean(unresolved_rasas)
+        output_dictionary = {
+            f"unresolved_polymer_rasa_{i}": rasa
+            for i, rasa in enumerate(unresolved_rasas)
+        }
+        output_dictionary["mean_unresolved_polymer_rasa"] = unresolved_rasa
+
+        # ...  and add resolved region RASA scores if flag is enabled
+        if self.include_resolved:
+            resolved_rasa = np.nanmean(resolved_rasas)
+            output_dictionary.update(
+                {
+                    f"resolved_polymer_rasa_{i}": rasa
+                    for i, rasa in enumerate(resolved_rasas)
+                }
+            )
+            output_dictionary["mean_resolved_polymer_rasa"] = resolved_rasa
+
+        return output_dictionary

rf3/metrics/selected_distances.py

@@ -0,0 +1,91 @@
+import numpy as np
+from atomworks.ml.utils import nested_dict
+from atomworks.ml.utils.selection import (
+    get_mask_from_atom_selection,
+    parse_selection_string,
+)
+from beartype.typing import Any
+from biotite.structure import AtomArrayStack
+
+from foundry.metrics.metric import Metric
+
+
+class SelectedAtomByAtomDistances(Metric):
+    """Computes all-by-all 2D distances given a list of selection strings"""
+
+    def compute_from_kwargs(self, **kwargs: Any) -> dict[str, Any]:
+        """Override parent class to handle optional selection_strings parameter"""
+        compute_inputs = {
+            "atom_array_stack": nested_dict.getitem(
+                kwargs, key="predicted_atom_array_stack"
+            )
+        }
+
+        # Add selection_strings only if it exists
+        try:
+            compute_inputs["selection_strings"] = nested_dict.getitem(
+                kwargs, key=("extra_info", "selection_strings")
+            )
+        except (KeyError, IndexError, TypeError):
+            pass
+
+        return self.compute(**compute_inputs)
+
+    def compute(
+        self,
+        atom_array_stack: AtomArrayStack,
+        selection_strings: list[str] | None = None,
+    ) -> dict[str, Any]:
+        # Short-circuit if no selection strings are provided
+        if not selection_strings:
+            return {}
+
+        # ... select the specified atoms
+        mask = np.zeros(atom_array_stack.array_length(), dtype=bool)
+        atom_selections = [parse_selection_string(s) for s in selection_strings]
+        for atom_selection in atom_selections:
+            mask |= get_mask_from_atom_selection(atom_array_stack, atom_selection)
+        selected_atom_array_stack = atom_array_stack[:, mask]
+
+        # Create views with added dimensions for broadcasting
+        # coord is (D, L, 3), we want pairwise distances for each D
+        coord_i = selected_atom_array_stack.coord[:, :, np.newaxis, :]  # (D, L, 1, 3)
+        coord_j = selected_atom_array_stack.coord[:, np.newaxis, :, :]  # (D, 1, L, 3)
+
+        # Calculate pairwise differences and distances
+        differences = coord_i - coord_j  # broadcasts to (D, L, L, 3)
+        distances = np.linalg.norm(differences, axis=-1)  # (D, L, L)
+
+        # Compute the mean and standard deviation across the D dimension
+        mean_distances = np.mean(distances, axis=0)  # Shape: (L, L)
+        std_distances = np.std(distances, axis=0)  # Shape: (L, L)
+
+        # Name the features with the chain_id, res_name, res_id, atom_name
+        def _format_atom_id(chain_id, res_name, res_id, atom_name):
+            return f"{chain_id}/{res_name}/{res_id}/{atom_name}"
+
+        vectorized_format = np.vectorize(_format_atom_id)
+        id = vectorized_format(
+            selected_atom_array_stack.chain_id,
+            selected_atom_array_stack.res_name,
+            selected_atom_array_stack.res_id,
+            selected_atom_array_stack.atom_name,
+        )
+
+        # Create a 2x2 numpy arrays of names, where we concatenate the id ...
+        id_i = np.char.add(id, "-")
+        id_II = np.char.add(id_i[:, np.newaxis], id[np.newaxis, :])
+
+        # ... and store the results in a dictionary, naming the columns with the concatenated id
+        results = {}
+        for i in range(len(id)):
+            for j in range(
+                i + 1, len(id)
+            ):  # Only consider j > i to avoid symmetric duplicates
+                col_id = id_II[i, j]
+                mean = mean_distances[i, j]
+                std = std_distances[i, j]
+                results[f"{col_id}_mean"] = mean
+                results[f"{col_id}_std"] = std
+
+        return results