protenix/data/filter.py

# Copyright 2024 ByteDance and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import biotite.structure as struc
import numpy as np
from biotite.structure import AtomArray, get_molecule_indices
from scipy.spatial.distance import cdist

from protenix.data.constants import CRYSTALLIZATION_AIDS


class Filter(object):
    """
    Ref: AlphaFold3 SI Chapter 2.5.4
    """

    @staticmethod
    def remove_hydrogens(atom_array: AtomArray) -> AtomArray:
        """remove hydrogens and deuteriums"""
        return atom_array[~np.isin(atom_array.element, ["H", "D"])]

    @staticmethod
    def remove_water(atom_array: AtomArray) -> AtomArray:
        """remove water (HOH) and deuterated water (DOD)"""
        return atom_array[~np.isin(atom_array.res_name, ["HOH", "DOD"])]

    @staticmethod
    def remove_element_X(atom_array: AtomArray) -> AtomArray:
        """
        remove element X
        following residues have element X:
        - UNX: unknown one atom or ion
        - UNL: unknown ligand, some atoms are marked as X
        - ASX: ASP/ASN ambiguous, two ambiguous atoms are marked as X, 6 entries in the PDB
        - GLX: GLU/GLN ambiguous, two ambiguous atoms are marked as X, 5 entries in the PDB
        """
        X_mask = np.zeros(len(atom_array), dtype=bool)
        starts = struc.get_residue_starts(atom_array, add_exclusive_stop=True)
        for start, stop in zip(starts[:-1], starts[1:]):
            res_name = atom_array.res_name[start]
            if res_name in ["UNX", "UNL"]:
                X_mask[start:stop] = True
        atom_array = atom_array[~X_mask]

        # map ASX to ASP, as ASP is more symmetric than ASN
        mask = atom_array.res_name == "ASX"
        atom_array.res_name[mask] = "ASP"
        atom_array.atom_name[mask & (atom_array.atom_name == "XD1")] = "OD1"
        atom_array.atom_name[mask & (atom_array.atom_name == "XD2")] = "OD2"
        atom_array.element[mask & (atom_array.element == "X")] = "O"

        # map GLX to GLU, as GLU is more symmetric than GLN
        mask = atom_array.res_name == "GLX"
        atom_array.res_name[mask] = "GLU"
        atom_array.atom_name[mask & (atom_array.atom_name == "XE1")] = "OE1"
        atom_array.atom_name[mask & (atom_array.atom_name == "XE2")] = "OE2"
        atom_array.element[mask & (atom_array.element == "X")] = "O"
        return atom_array

    @staticmethod
    def remove_crystallization_aids(
        atom_array: AtomArray, entity_poly_type: dict
    ) -> AtomArray:
        """remove crystallization aids, eg: SO4, GOL, etc.

        Only remove crystallization aids if the chain is not polymer.

        Ref: AlphaFold3 SI Chapter 2.5.4
        """
        non_aids_mask = ~np.isin(atom_array.res_name, CRYSTALLIZATION_AIDS)
        poly_mask = np.isin(atom_array.label_entity_id, list(entity_poly_type.keys()))
        return atom_array[poly_mask | non_aids_mask]

    @staticmethod
    def _get_clashing_chains(
        atom_array: AtomArray, chain_ids: list[str]
    ) -> tuple[np.ndarray, list[int]]:
        """
        Calculate the number of atoms clashing with other chains for each chain
        and return a matrix that records the count of clashing atoms.

        Note: if two chains are covalent, they are not considered as clashing.

        Args:
            atom_array (AtomArray): All atoms, including those not resolved.
            chain_ids (list[str]): Unique chain indices of resolved atoms.

        Returns:
            tuple:
                clash_records (numpy.ndarray): Matrix of clashing atom num.
                                               (i, j) means the ratio of i's atom clashed with j's atoms.
                                               Note: (i, j) != (j, i).
                chain_resolved_atom_nums (list[int]): The number of resolved atoms corresponding to each chain ID.
        """
        is_resolved_centre_atom = (
            atom_array.centre_atom_mask == 1
        ) & atom_array.is_resolved
        cell_list = struc.CellList(
            atom_array, cell_size=1.7, selection=is_resolved_centre_atom
        )

        # (i, j) means the ratio of i's atom clashed with j's atoms
        clash_records = np.zeros((len(chain_ids), len(chain_ids)))

        # record the number of resolved atoms for each chain
        chain_resolved_atom_nums = []

        # record covalent relationship between chains
        chains_covalent_dict = {}
        for idx, chain_id_i in enumerate(chain_ids):
            for chain_id_j in chain_ids[idx + 1 :]:
                mol_indices = get_molecule_indices(
                    atom_array[np.isin(atom_array.chain_id, [chain_id_i, chain_id_j])]
                )
                if len(mol_indices) == 1:
                    covalent = 1
                else:
                    covalent = 0
                chains_covalent_dict[(chain_id_i, chain_id_j)] = covalent
                chains_covalent_dict[(chain_id_j, chain_id_i)] = covalent

        for i, chain_id in enumerate(chain_ids):
            coords = atom_array.coord[
                (atom_array.chain_id == chain_id) & is_resolved_centre_atom
            ]
            chain_resolved_atom_nums.append(len(coords))
            chain_atom_ids = np.where(atom_array.chain_id == chain_id)[0]
            chain_atom_ids_set = set(chain_atom_ids) | {-1}

            # Get atom indices from the current cell and the eight surrounding cells.
            neighbors_ids_2d = cell_list.get_atoms_in_cells(coords, cell_radius=1)
            neighbors_ids = np.unique(neighbors_ids_2d)

            # Remove the atom indices of the current chain.
            other_chain_atom_ids = list(set(neighbors_ids) - chain_atom_ids_set)

            if not other_chain_atom_ids:
                continue
            else:
                # Calculate the distance matrix with neighboring atoms.
                other_chain_atom_coords = atom_array.coord[other_chain_atom_ids]
                dist_mat = cdist(coords, other_chain_atom_coords, metric="euclidean")
                clash_mat = dist_mat < 1.6  # change 1.7 to 1.6 for more compatibility
                if np.any(clash_mat):
                    clashed_other_chain_ids = atom_array.chain_id[other_chain_atom_ids]

                    for other_chain_id in set(clashed_other_chain_ids):

                        # two chains covalent with each other
                        if chains_covalent_dict[(chain_id, other_chain_id)]:
                            continue

                        cols = np.where(clashed_other_chain_ids == other_chain_id)[0]

                        # how many i's atoms clashed with j
                        any_atom_clashed = np.any(
                            clash_mat[:, cols].astype(int), axis=1
                        )
                        clashed_atom_num = np.sum(any_atom_clashed.astype(int))

                        if clashed_atom_num > 0:
                            j = chain_ids.index(other_chain_id)
                            clash_records[i][j] += clashed_atom_num
        return clash_records, chain_resolved_atom_nums

    @staticmethod
    def _get_removed_clash_chain_ids(
        clash_records: np.ndarray,
        chain_ids: list[str],
        chain_resolved_atom_nums: list[int],
        core_chain_id: np.ndarray = [],
    ) -> list[str]:
        """
        Perform pairwise comparisons on the chains, and select the chain IDs
        to be deleted according to the clahsing chain rules.

        Args:
            clash_records (numpy.ndarray): Matrix of clashing atom num.
                                           (i, j) means the ratio of i's atom clashed with j's atoms.
                                           Note: (i, j) != (j, i).
            chain_ids (list[str]): Unique chain indices of resolved atoms.
            chain_resolved_atom_nums (list[int]): The number of resolved atoms corresponding to each chain ID.
            core_chain_id (np.ndarray): The chain ID of the core chain.

        Returns:
            list[str]: A list of chain IDs that have been determined for deletion.
        """
        removed_chain_ids = []
        for i in range(len(chain_ids)):
            atom_num_i = chain_resolved_atom_nums[i]
            chain_idx_i = chain_ids[i]

            if chain_idx_i in removed_chain_ids:
                continue

            for j in range(i + 1, len(chain_ids)):
                atom_num_j = chain_resolved_atom_nums[j]
                chain_idx_j = chain_ids[j]

                if chain_idx_j in removed_chain_ids:
                    continue

                clash_num_ij, clash_num_ji = (
                    clash_records[i][j],
                    clash_records[j][i],
                )

                clash_ratio_ij = clash_num_ij / atom_num_i
                clash_ratio_ji = clash_num_ji / atom_num_j

                if clash_ratio_ij <= 0.3 and clash_ratio_ji <= 0.3:
                    # not reaches the threshold
                    continue
                else:
                    # clashing chains
                    if (
                        chain_idx_i in core_chain_id
                        and chain_idx_j not in core_chain_id
                    ):
                        removed_chain_idx = chain_idx_j
                    elif (
                        chain_idx_i not in core_chain_id
                        and chain_idx_j in core_chain_id
                    ):
                        removed_chain_idx = chain_idx_i

                    elif clash_ratio_ij > clash_ratio_ji:
                        removed_chain_idx = chain_idx_i
                    elif clash_ratio_ij < clash_ratio_ji:
                        removed_chain_idx = chain_idx_j
                    else:
                        if atom_num_i < atom_num_j:
                            removed_chain_idx = chain_idx_i
                        elif atom_num_i > atom_num_j:
                            removed_chain_idx = chain_idx_j
                        else:
                            removed_chain_idx = sorted([chain_idx_i, chain_idx_j])[1]

                    removed_chain_ids.append(removed_chain_idx)

                    if removed_chain_idx == chain_idx_i:
                        # chain i already removed
                        break
        return removed_chain_ids

    @staticmethod
    def remove_polymer_chains_all_residues_unknown(
        atom_array: AtomArray,
        entity_poly_type: dict,
    ) -> AtomArray:
        """remove chains with all residues unknown"""
        chain_starts = struc.get_chain_starts(atom_array, add_exclusive_stop=True)
        invalid_chains = []  # list of [start, end)
        for index in range(len(chain_starts) - 1):
            start, end = chain_starts[index], chain_starts[index + 1]
            entity_id = atom_array[start].label_entity_id
            if (
                entity_poly_type.get(entity_id, "non-poly") == "polypeptide(L)"
                and np.all(atom_array.res_name[start:end] == "UNK")
            ) or (
                entity_poly_type.get(entity_id, "non-poly")
                in (
                    "polyribonucleotide",
                    "polydeoxyribonucleotide",
                )
                and np.all(atom_array.res_name[start:end] == "N")
            ):
                invalid_chains.append((start, end))
        mask = np.ones(len(atom_array), dtype=bool)
        for start, end in invalid_chains:
            mask[start:end] = False
        atom_array = atom_array[mask]
        return atom_array

    @staticmethod
    def remove_polymer_chains_too_short(
        atom_array: AtomArray, entity_poly_type: dict
    ) -> AtomArray:
        chain_starts = struc.get_chain_starts(atom_array, add_exclusive_stop=True)
        invalid_chains = []  # list of [start, end)
        for index in range(len(chain_starts) - 1):
            start, end = chain_starts[index], chain_starts[index + 1]
            entity_id = atom_array[start].label_entity_id
            num_residue_ids = len(set(atom_array.label_seq_id[start:end]))
            if (
                entity_poly_type.get(entity_id, "non-poly")
                in (
                    "polypeptide(L)",  # TODO: how to handle polypeptide(D)?
                    "polyribonucleotide",
                    "polydeoxyribonucleotide",
                )
                and num_residue_ids < 4
            ):
                invalid_chains.append((start, end))
        mask = np.ones(len(atom_array), dtype=bool)
        for start, end in invalid_chains:
            mask[start:end] = False
        atom_array = atom_array[mask]
        return atom_array

    @staticmethod
    def remove_polymer_chains_with_consecutive_c_alpha_too_far_away(
        atom_array: AtomArray, entity_poly_type: dict, max_distance: float = 10.0
    ) -> AtomArray:
        chain_starts = struc.get_chain_starts(atom_array, add_exclusive_stop=True)
        invalid_chains = []  # list of [start, end)
        for index in range(len(chain_starts) - 1):
            start, end = chain_starts[index], chain_starts[index + 1]
            entity_id = atom_array.label_entity_id[start]
            if entity_poly_type.get(entity_id, "non-poly") == "polypeptide(L)":
                peptide_atoms = atom_array[start:end]
                ca_atoms = peptide_atoms[peptide_atoms.atom_name == "CA"]
                seq_ids = ca_atoms.label_seq_id
                seq_ids[seq_ids == "."] = "-100"
                seq_ids = seq_ids.astype(np.int64)
                dist_square = np.sum(
                    (ca_atoms[:-1].coord - ca_atoms[1:].coord) ** 2, axis=-1
                )
                invalid_neighbor_mask = (dist_square > max_distance**2) & (
                    seq_ids[:-1] + 1 == seq_ids[1:]
                )
                if np.any(invalid_neighbor_mask):
                    invalid_chains.append((start, end))
        mask = np.ones(len(atom_array), dtype=bool)
        for start, end in invalid_chains:
            mask[start:end] = False
        atom_array = atom_array[mask]
        return atom_array

    @staticmethod
    def too_many_chains_filter(
        atom_array: AtomArray,
        interface_radius: int = 15,
        max_chains_num: int = 20,
        core_indices: list[int] = None,
        max_tokens_num: int = None,
    ) -> tuple[AtomArray, int]:
        """
        Ref: AlphaFold3 SI Chapter 2.5.4

        For bioassemblies with greater than 20 chains, we select a random interface token
        (with a centre atom <15 Å to the centre atom of a token in another chain)
        and select the closest 20 chains to this token based on
        minimum distance between any tokens centre atom.

        Note: due to the presence of covalent small molecules,
        treat the covalent small molecule and the polymer it is attached to
        as a single chain to avoid inadvertently removing the covalent small molecules.
        Use the mol_id added to the AtomArray to differentiate between the various
        parts of the structure composed of covalent bonds.

        Args:
            atom_array (AtomArray): Biotite AtomArray Object of a Bioassembly.
            interface_radius (int, optional): Atoms within this distance of the central atom are considered interface atoms.
                                            Defaults to 15.
            max_chains_num (int, optional): The maximum number of chains permitted in a bioassembly.
                                            Filtration will be applied if exceeds this value. Defaults to 20.
            core_indices (list[int], optional): A list of indices to be used as chose the central atom.
                                                     And corresponding chains in the list will be selected proriority.
                                                     If None, a random index from whole AtomArray will be selected. Defaults to None.
            max_tokens_num (int, optional): The maximum number of tokens permitted in a bioassembly.
                                            If not None,  after more than max_chains_num, if the max_tokens_num is not reached,
                                            it will continue to append the chains.

        Returns:
            tuple:
                - atom_array (AtomArray): An AtomArray that has been processed through this filter.
                - input_chains_num (int): The number of chain in the input AtomArray.
                                          This is to log whether the filter has been utilized.
        """
        # each mol is a so called "chain" in the context of this filter.
        input_chains_num = len(np.unique(atom_array.mol_id))
        if input_chains_num <= max_chains_num:
            # no change
            return atom_array, input_chains_num

        is_resolved_centre_atom = (
            atom_array.centre_atom_mask == 1
        ) & atom_array.is_resolved

        cell_list = struc.CellList(
            atom_array, cell_size=interface_radius, selection=is_resolved_centre_atom
        )
        resolved_centre_atom = atom_array[is_resolved_centre_atom]

        assert resolved_centre_atom, "There is no resolved central atom."

        # random pick centre atom
        if core_indices is None:
            index_shuf = np.random.default_rng(seed=42).permutation(
                len(resolved_centre_atom)
            )
        else:
            index_shuf = np.array(core_indices)
            resolved_centre_atom_indices = np.nonzero(is_resolved_centre_atom)[0]

            # get indices of resolved_centre_atom
            index_shuf = np.array(
                [
                    np.where(resolved_centre_atom_indices == idx)[0][0]
                    for idx in index_shuf
                    if idx in resolved_centre_atom_indices
                ]
            )
            np.random.default_rng(seed=42).shuffle(index_shuf)

        chosen_centre_atom = None
        for idx in index_shuf:
            centre_atom = resolved_centre_atom[idx]
            neighbors_indices = cell_list.get_atoms(
                centre_atom.coord, radius=interface_radius
            )
            neighbors_indices = neighbors_indices[neighbors_indices != -1]

            neighbors_chain_ids = np.unique(atom_array.mol_id[neighbors_indices])
            # neighbors include centre atom itself
            if len(neighbors_chain_ids) > 1:
                chosen_centre_atom = centre_atom
                break

        # The distance between the central atoms in any two chains is greater than 15 angstroms.
        if chosen_centre_atom is None:
            return None, input_chains_num

        dist_mat = cdist(centre_atom.coord.reshape((1, -1)), resolved_centre_atom.coord)
        sorted_chain_id = np.array(
            [
                chain_id
                for chain_id, _dist in sorted(
                    zip(resolved_centre_atom.mol_id, dist_mat[0]),
                    key=lambda pair: pair[1],
                )
            ]
        )

        if core_indices is not None:
            # select core proriority
            core_mol_id = np.unique(atom_array.mol_id[core_indices])
            in_core_mask = np.isin(sorted_chain_id, core_mol_id)
            sorted_chain_id = np.concatenate(
                (sorted_chain_id[in_core_mask], sorted_chain_id[~in_core_mask])
            )

        closest_chain_id = set()
        chain_ids_to_token_num = {}
        if max_tokens_num is None:
            max_tokens_num = 0

        tokens = 0
        for chain_id in sorted_chain_id:
            # get token num
            if chain_id not in chain_ids_to_token_num:
                chain_ids_to_token_num[chain_id] = atom_array.centre_atom_mask[
                    atom_array.mol_id == chain_id
                ].sum()
            chain_token_num = chain_ids_to_token_num[chain_id]

            if len(closest_chain_id) >= max_chains_num:
                if tokens + chain_token_num > max_tokens_num:
                    break

            closest_chain_id.add(chain_id)
            tokens += chain_token_num

        atom_array = atom_array[np.isin(atom_array.mol_id, list(closest_chain_id))]
        output_chains_num = len(np.unique(atom_array.mol_id))
        assert (
            output_chains_num == max_chains_num
            or atom_array.centre_atom_mask.sum() <= max_tokens_num
        )
        return atom_array, input_chains_num

    @staticmethod
    def remove_clashing_chains(
        atom_array: AtomArray,
        core_indices: list[int] = None,
    ) -> AtomArray:
        """
        Ref: AlphaFold3 SI Chapter 2.5.4

        Clashing chains are removed.
        Clashing chains are defined as those with >30% of atoms within 1.7 Å of an atom in another chain.
        If two chains are clashing with each other, the chain with the greater percentage of clashing atoms will be removed.
        If the same fraction of atoms are clashing, the chain with fewer total atoms is removed.
        If the chains have the same number of atoms, then the chain with the larger chain id is removed.

        Note: if two chains are covalent, they are not considered as clashing.

        Args:
            atom_array (AtomArray): Biotite AtomArray Object of a Bioassembly.
            core_indices (list[int]): A list of indices for core structures,
                                      where these indices correspond to structures that will be preferentially
                                      retained when pairwise clash chain assessments are performed.

        Returns:
            atom_array (AtomArray): An AtomArray that has been processed through this filter.
            removed_chain_ids (list[str]): A list of chain IDs that have been determined for deletion.
                                           This is to log whether the filter has been utilized.
        """
        chain_ids = np.unique(atom_array.chain_id[atom_array.is_resolved]).tolist()

        if core_indices is not None:
            core_chain_id = np.unique(atom_array.chain_id[core_indices])
        else:
            core_chain_id = np.array([])

        clash_records, chain_resolved_atom_nums = Filter._get_clashing_chains(
            atom_array, chain_ids
        )
        removed_chain_ids = Filter._get_removed_clash_chain_ids(
            clash_records,
            chain_ids,
            chain_resolved_atom_nums,
            core_chain_id=core_chain_id,
        )

        atom_array = atom_array[~np.isin(atom_array.chain_id, removed_chain_ids)]
        return atom_array, removed_chain_ids

    @staticmethod
    def remove_unresolved_mols(atom_array: AtomArray) -> AtomArray:
        """
        Remove molecules from a bioassembly object which all atoms are not resolved.

        Args:
            atom_array (AtomArray): Biotite AtomArray Object of a bioassembly.

        Returns:
            AtomArray: An AtomArray object with unresolved molecules removed.
        """
        valid_mol_id = []
        for mol_id in np.unique(atom_array.mol_id):
            resolved = atom_array.is_resolved[atom_array.mol_id == mol_id]
            if np.any(resolved):
                valid_mol_id.append(mol_id)

        atom_array = atom_array[np.isin(atom_array.mol_id, valid_mol_id)]
        return atom_array

    @staticmethod
    def remove_asymmetric_polymer_ligand_bonds(
        atom_array: AtomArray, entity_poly_type: dict
    ) -> AtomArray:
        """remove asymmetric polymer ligand bonds (including protein-protein bond, like disulfide bond).

        AF3 SI 5.1 Structure filters
        Bonds for structures with homomeric subcomplexes lacking the corresponding homomeric symmetry are also removed
        - e.g. if a certain bonded ligand only exists for some of the symmetric copies, but not for all,
        we remove the corresponding bond information from the input.
        In consequence the model has to learn to infer these bonds by itself.

        Args:
            atom_array (AtomArray): input atom array

        Returns:
            AtomArray: output atom array with asymmetric polymer ligand bonds removed.
        """
        # get inter chain bonds
        inter_chain_bonds = set()
        for i, j, b in atom_array.bonds.as_array():
            if atom_array.chain_id[i] != atom_array.chain_id[j]:
                inter_chain_bonds.add((i, j))

        # get asymmetric polymer ligand bonds
        asymmetric_bonds = set()
        chain_starts = struc.get_chain_starts(atom_array, add_exclusive_stop=False)
        for bond in inter_chain_bonds:

            if bond in asymmetric_bonds:
                continue

            i, j = bond
            atom_i = atom_array[i]
            atom_j = atom_array[j]
            i_is_polymer = atom_i.label_entity_id in entity_poly_type
            j_is_polymer = atom_j.label_entity_id in entity_poly_type
            if i_is_polymer:
                pass
            elif j_is_polymer:
                i, j = j, i
                atom_i, atom_j = atom_j, atom_i
                i_is_polymer, j_is_polymer = j_is_polymer, i_is_polymer
            else:
                # both entity is not polymer
                continue

            # get atom i mask from all entity i copies
            entity_mask_i = atom_array.label_entity_id == atom_i.label_entity_id
            num_copies = np.isin(chain_starts, np.flatnonzero(entity_mask_i)).sum()
            mask_i = (
                entity_mask_i
                & (atom_array.res_id == atom_i.res_id)
                & (atom_array.atom_name == atom_i.atom_name)
            )
            indices_i = np.flatnonzero(mask_i)

            if len(indices_i) != num_copies:
                # not every copy of entity i has atom i.
                asymmetric_bonds.add(bond)
                continue

            # check all atom i in entity i bond to an atom j in entity j.
            target_bonds = []
            for ii in indices_i:
                ii_bonds = [b for b in inter_chain_bonds if ii in b]
                for bond in ii_bonds:
                    jj = bond[1] if ii == bond[0] else bond[0]
                    atom_jj = atom_array[jj]
                    if atom_jj.label_entity_id != atom_j.label_entity_id:
                        continue
                    if atom_jj.res_name != atom_j.res_name:
                        continue
                    if atom_jj.atom_name != atom_j.atom_name:
                        continue
                    if j_is_polymer and atom_jj.res_id != atom_j.res_id:
                        # only for polymer, check res_id
                        continue
                    # found bond (ii, jj) with same enity_id, res_name, atom_name to bond (i,j)
                    target_bonds.append((min(ii, jj), max(ii, jj)))
                    break
            if len(target_bonds) != num_copies:
                asymmetric_bonds |= set(target_bonds)

        for bond in asymmetric_bonds:
            atom_array.bonds.remove_bond(bond[0], bond[1])
        return atom_array