model/boltzgen/data/crop/multimer.py

from dataclasses import replace
from typing import List, Optional, Set

import numpy as np
from scipy.spatial.distance import cdist

from boltzgen.data import const
from boltzgen.data.crop.cropper import Cropper
from boltzgen.data.data import Input, Token, TokenBond, Tokenized


def pick_random_token(
    tokens: np.ndarray,
    random: np.random.Generator,
) -> np.ndarray:
    """Pick a random token from the data.

    Parameters
    ----------
    tokens : np.ndarray
        The token data.
    random : np.random.Generator
        The random state for reproducibility.

    Returns
    -------
    np.ndarray
        The selected token.

    """
    return tokens[random.integers(len(tokens))]


def pick_chain_token(
    tokens: np.ndarray,
    chain_id: int,
    random: np.random.Generator,
) -> np.ndarray:
    """Pick a random token from a chain.

    Parameters
    ----------
    tokens : np.ndarray
        The token data.
    chain_id : int
        The chain ID.
    random : np.ndarray
        The random state for reproducibility.

    Returns
    -------
    np.ndarray
        The selected token.

    """
    # Filter to chain
    chain_tokens = tokens[tokens["asym_id"] == chain_id]

    # Pick from chain, fallback to all tokens
    if chain_tokens.size:
        query = pick_random_token(chain_tokens, random)
    else:
        query = pick_random_token(tokens, random)

    return query


def pick_interface_token(
    tokens: np.ndarray,
    interface: np.ndarray,
    random: np.random.Generator,
) -> np.ndarray:
    """Pick a random token from an interface.

    Parameters
    ----------
    tokens : np.ndarray
        The token data.
    interface : int
        The interface ID.
    random : np.random.Generator
        The random state for reproducibility.

    Returns
    -------
    np.ndarray
        The selected token.

    """
    # Sample random interface
    chain_1 = int(interface["chain_1"])
    chain_2 = int(interface["chain_2"])

    tokens_1 = tokens[tokens["asym_id"] == chain_1]
    tokens_2 = tokens[tokens["asym_id"] == chain_2]

    # If no interface, pick from the chains
    if tokens_1.size and (not tokens_2.size):
        query = pick_random_token(tokens_1, random)
    elif tokens_2.size and (not tokens_1.size):
        query = pick_random_token(tokens_2, random)
    elif (not tokens_1.size) and (not tokens_2.size):
        query = pick_random_token(tokens, random)
    else:
        # If we have tokens, compute distances
        tokens_1_coords = tokens_1["center_coords"]
        tokens_2_coords = tokens_2["center_coords"]

        dists = cdist(tokens_1_coords, tokens_2_coords)
        cuttoff = dists < const.interface_cutoff

        # In rare cases, the interface cuttoff is slightly
        # too small, then we slightly expand it if it happens
        if not np.any(cuttoff):
            cuttoff = dists < (const.interface_cutoff + 5.0)

        tokens_1 = tokens_1[np.any(cuttoff, axis=1)]
        tokens_2 = tokens_2[np.any(cuttoff, axis=0)]

        # Select random token
        candidates = np.concatenate([tokens_1, tokens_2])
        query = pick_random_token(candidates, random)

    return query


def pick_initial_crop_token(
    tokens: np.ndarray,
    initial_crop: List[int],
    random: np.random.Generator,
) -> np.ndarray:
    """Pick a random token from the initial crop.

    Parameters
    ----------
    tokens : np.ndarray
        The token data.
    initial_crop : List[int]
        The initial crop.
    random : np.random.Generator
        The random state for reproducibility.

    Returns
    -------
    np.ndarray

    """
    # Compute crop centroid
    crop_centroid = np.mean(tokens[initial_crop]["center_coords"], axis=0)

    # Compute distances to all tokens
    dists = cdist(tokens["center_coords"], crop_centroid[None])

    # Pick the closest token
    return tokens[np.argmin(dists[:, 0])]


class MultimerCropper(Cropper):
    """Interpolate between contiguous and spatial crops."""

    def __init__(
        self, neighborhood_sizes: List[int], dna_double_helix: bool = False
    ) -> None:
        """Initialize the cropper.

        Parameters
        ----------
        neighborhood_sizes : List[int]
            Modulates the type of cropping to be performed.
            Smaller neighborhoods result in more spatial
            cropping. Larger neighborhoods result in more
            continuous cropping. A mix can be achieved by
            providing a list of sizes from which to sample.

        """
        self.neighborhood_sizes = neighborhood_sizes
        self.dna_double_helix = dna_double_helix

    def crop(  # noqa: PLR0915
        self,
        data: Tokenized,
        max_tokens: int,
        random: np.random.Generator,
        chain_id: Optional[int] = None,
        interface_id: Optional[int] = None,
        max_atoms: Optional[int] = None,
        return_indices: bool = False,
        initial_crop: Optional[List[int]] = None,
        prefer_protein_queries: Optional[bool] = False,
    ) -> Tokenized:
        """Crop the data to a maximum number of tokens.

        Parameters
        ----------
        data : Tokenized
            The tokenized data.
        max_tokens : int
            The maximum number of tokens to crop.
        random : np.random.Generator
            The random state for reproducibility.
        max_atoms : Optional[int]
            The maximum number of atoms to consider.

        Returns
        -------
        Tokenized
            The cropped data.

        """
        # Check inputs
        if chain_id is not None and interface_id is not None:
            msg = "Only one of chain_id or interface_id can be provided."
            raise ValueError(msg)

        # randomly select a neighborhood size
        neighborhood_size = random.choice(self.neighborhood_sizes)

        # Get token data
        token_data = data.tokens
        token_to_res = data.token_to_res
        token_bonds = data.bonds
        mask = data.structure.mask
        chains = data.structure.chains
        interfaces = data.structure.interfaces

        # Filter to valid chains
        valid_chains = chains[mask]

        # Filter to valid interfaces
        valid_interfaces = interfaces
        valid_interfaces = valid_interfaces[mask[valid_interfaces["chain_1"]]]
        valid_interfaces = valid_interfaces[mask[valid_interfaces["chain_2"]]]

        # Filter to resolved tokens
        valid_tokens = token_data[token_data["resolved_mask"]]

        # Check if we have any valid tokens
        if not valid_tokens.size:
            msg = "No valid tokens in structure"
            raise ValueError(msg)

        is_protein_mask = valid_tokens["mol_type"] == const.chain_type_ids["PROTEIN"]
        if prefer_protein_queries and is_protein_mask.sum() > 0:
            query_choice_tokens = valid_tokens[is_protein_mask]
        else:
            query_choice_tokens = valid_tokens

        # Pick a random token, chain, or interface
        if initial_crop is not None:
            query = pick_initial_crop_token(token_data, initial_crop, random)
        elif chain_id is not None:
            query = pick_chain_token(query_choice_tokens, chain_id, random)
        elif interface_id is not None:
            interface = interfaces[interface_id]
            query = pick_interface_token(query_choice_tokens, interface, random)
        elif valid_interfaces.size:
            idx = random.integers(len(valid_interfaces))
            interface = valid_interfaces[idx]
            query = pick_interface_token(query_choice_tokens, interface, random)
        else:
            idx = random.integers(len(valid_chains))
            chain_id = valid_chains[idx]["asym_id"]
            query = pick_chain_token(query_choice_tokens, chain_id, random)

        cropped = self.select_cropped_indices(
            token_data,
            valid_tokens,
            query,
            neighborhood_size,
            max_atoms,
            max_tokens,
            initial_crop,
        )

        # Get the cropped tokens sorted by index
        token_data = token_data[cropped]
        token_to_res = token_to_res[cropped]

        # Only keep bonds within the cropped tokens and reindex token_idx
        old_indices = token_data["token_idx"]
        old_to_new = {old: new for new, old in enumerate(old_indices)}
        token_bonds = np.array(
            [
                (old_to_new[bond["token_1"]], old_to_new[bond["token_2"]], bond["type"])
                for bond in token_bonds
                if bond["token_1"] in old_indices and bond["token_2"] in old_indices
            ],
            dtype=TokenBond,
        )
        token_data["token_idx"] = np.arange(len(token_data))

        # Return the cropped tokens
        if return_indices:
            return replace(
                data, tokens=token_data, bonds=token_bonds, token_to_res=token_to_res
            ), cropped
        else:
            return replace(
                data, tokens=token_data, bonds=token_bonds, token_to_res=token_to_res
            )

    def select_cropped_indices(
        self,
        tokens: np.ndarray,
        valid_tokens: np.ndarray,
        query: np.ndarray,
        neighborhood_size: int,
        max_atoms: Optional[int] = None,
        max_tokens: Optional[int] = None,
        initial_crop: Optional[np.ndarray] = None,
    ):
        # Sort all tokens by distance to query_coords
        dists = valid_tokens["center_coords"] - query["center_coords"]
        indices = np.argsort(np.linalg.norm(dists, axis=1))

        # Select cropped indices
        cropped: Set[int] = set()
        total_atoms = 0

        if initial_crop is not None:
            cropped.update(initial_crop)
            total_atoms = sum(tokens[idx]["atom_num"] for idx in initial_crop)

        for idx in indices:
            # Get the token
            token = valid_tokens[idx]

            # Get all tokens from this chain
            chain_tokens = tokens[tokens["asym_id"] == token["asym_id"]]

            # Pick the whole chain if possible, otherwise select
            # a contiguous subset centered at the query token
            if len(chain_tokens) <= neighborhood_size:
                new_tokens = chain_tokens
            else:
                # First limit to the maximum set of tokens, with the
                # neighboorhood on both sides to handle edges. This
                # is mostly for efficiency with the while loop below.

                min_idx = token["res_idx"] - neighborhood_size
                max_idx = token["res_idx"] + neighborhood_size

                max_token_set = chain_tokens
                max_token_set = max_token_set[max_token_set["res_idx"] >= min_idx]
                max_token_set = max_token_set[max_token_set["res_idx"] <= max_idx]
                if len(max_token_set) < neighborhood_size:
                    max_token_set = chain_tokens

                if len(max_token_set) < neighborhood_size:
                    print(
                        "WARNING! len(max_token_set) < neighborhood_size in cropper in selector"
                    )
                    new_tokens = max_token_set
                else:
                    # Start by adding just the query token
                    new_tokens = max_token_set[
                        max_token_set["res_idx"] == token["res_idx"]
                    ]

                    # Expand the neighborhood until we have enough tokens, one
                    # by one to handle some edge cases with non-standard chains.
                    # We switch to the res_idx instead of the token_idx to always
                    # include all tokens from modified residues or from ligands.
                    min_idx = max_idx = token["res_idx"]
                    counter = 0
                    while new_tokens.size < neighborhood_size:
                        counter += 1
                        min_idx = min_idx - 1
                        max_idx = max_idx + 1
                        new_tokens = max_token_set
                        new_tokens = new_tokens[new_tokens["res_idx"] >= min_idx]
                        new_tokens = new_tokens[new_tokens["res_idx"] <= max_idx]
                        if counter > 1000:
                            raise Exception("Infinite loop in cropper while loop")

            # Compute new tokens and new atoms
            new_token_mask = np.isin(tokens["token_idx"], new_tokens["token_idx"])
            new_indices = np.arange(len(tokens), dtype=tokens["token_idx"].dtype)[
                new_token_mask
            ]
            new_indices = set(new_indices) - cropped
            new_tokens = tokens[list(new_indices)]
            new_atoms = np.sum(new_tokens["atom_num"])

            # Stop if we exceed the max number of tokens or atoms
            if (len(new_indices) > (max_tokens - len(cropped))) or (
                (max_atoms is not None) and ((total_atoms + new_atoms) > max_atoms)
            ):
                break

            # Add new indices
            cropped.update(new_indices)
            total_atoms += new_atoms
        return np.array(sorted(cropped))

    def crop_indices(  # noqa: PLR0915
        self,
        data: Tokenized,
        cropped_indices: List[int],
    ) -> Tokenized:
        token_data = data.tokens
        token_ids_mol = token_data[token_data["mol_type"] == 3]["token_idx"].tolist()  # noqa: PLR2004
        cropped_indices = sorted({*token_ids_mol, *cropped_indices})
        token_data = token_data[cropped_indices]
        indices = token_data["token_idx"]
        token_bonds = data.bonds
        token_bonds = token_bonds[np.isin(token_bonds["token_1"], indices)]
        token_bonds = token_bonds[np.isin(token_bonds["token_2"], indices)]

        return replace(data, tokens=token_data, bonds=token_bonds)