fastplms/models/esmfold2/esmfold2_protein_complex.py

"""Protein-complex data, assembly expansion, and geometry for ESMFold2."""

from __future__ import annotations

import io
import itertools
import random
import re
import warnings
from collections.abc import Iterable, Sequence
from dataclasses import asdict, dataclass, replace
from functools import cached_property
from pathlib import Path
from subprocess import check_output
from tempfile import TemporaryDirectory
from typing import Any

import biotite.structure as bs
import brotli
import msgpack
import msgpack_numpy
import numpy as np
import torch
from biotite.database import rcsb
from biotite.file import InvalidFileError
from biotite.structure.io.pdb import PDBFile
from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile
from biotite.structure.io.pdbx import set_structure as set_structure_pdbx
from biotite.structure.io.pdbx.convert import _get_transformations, get_structure
from biotite.structure.util import matrix_rotate
from scipy.spatial import KDTree

from . import esmfold2_residue_constants as residue_constants
from .esmfold2_affine3d import Affine3D
from .esmfold2_aligner import Aligner
from .esmfold2_atom_indexer import AtomIndexer
from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca
from .esmfold2_misc import slice_python_object_as_numpy
from .esmfold2_mmcif_parsing import (
    MmcifWrapper,
    NoProteinError,
    round_mmcif_columns,
)
from .esmfold2_protein_chain import (
    ProteinChain,
    _str_key_to_int_key,
    chain_to_ndarray,
    index_by_atom_name,
    infer_cb,
)
from .esmfold2_utils_types import PathOrBuffer

SINGLE_LETTER_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"


def _parse_operation_expression(expression: str) -> list[tuple[str, ...]]:
    """Expand an mmCIF operation expression in application order."""

    def expand_group(group: str) -> list[str]:
        operation_ids: list[str] = []
        for term in group.split(","):
            if "-" not in term:
                operation_ids.append(term)
                continue
            first, last = (int(value) for value in term.split("-"))
            operation_ids.extend(str(value) for value in range(first, last + 1))
        return operation_ids

    groups = [group for group in expression.replace(")", "").split("(") if group]
    groups.reverse()
    return list(itertools.product(*(expand_group(group) for group in groups)))


def _apply_transformations_fast(chains, transformation_dict, operations):
    """Return transformed copies of each affected protein chain."""
    transformed_chains = []
    for chain in chains:
        for operation in operations:
            coordinates = chain.atom37_positions.copy()
            for op_step in operation:
                transform = transformation_dict[op_step]
                coordinates = matrix_rotate(coordinates, transform.rotation)
                coordinates += transform.target_translation
            transformed_chains.append(replace(chain, atom37_positions=coordinates))
    return transformed_chains


@dataclass
class ProteinComplexMetadata:
    entity_lookup: dict[int, int | str]
    chain_lookup: dict[int, str]
    mmcif: MmcifWrapper | None = None
    # This is a dictionary that maps assembly ids to the list of unique chains
    # in that assembly. Allows for usage of `switch_assembly`.
    assembly_composition: dict[str, list[str]] | None = None


@dataclass
class DockQSingleScore:
    native_chains: tuple[str, str]
    DockQ: float
    interface_rms: float
    ligand_rms: float
    fnat: float
    fnonnat: float
    clashes: float
    F1: float
    DockQ_F1: float


@dataclass
class DockQResult:
    total_dockq: float
    native_interfaces: int
    chain_mapping: dict[str, str]
    interfaces: dict[tuple[str, str], DockQSingleScore]
    # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing
    # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling
    aligned: ProteinComplex
    aligned_rmsd: float


@dataclass(frozen=True)
class ProteinComplex:
    """Dataclass with atom37 representation of an entire protein complex."""

    id: str
    sequence: str
    entity_id: np.ndarray  # entities map to unique sequences
    chain_id: np.ndarray  # multiple chains might share an entity id
    sym_id: np.ndarray  # complexes might be copies of the same chain
    residue_index: np.ndarray
    insertion_code: np.ndarray
    atom37_positions: np.ndarray
    atom37_mask: np.ndarray
    confidence: np.ndarray
    # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort.
    metadata: ProteinComplexMetadata
    atom37_confidence: np.ndarray | None = None  # P has shape (l, 37).

    # Coordinate completion, concatenation, and comparison
    def infer_oxygen(self) -> ProteinComplex:
        """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()

        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
        N = torch.roll(N, -3)
        N[..., -1, :] = torch.nan

        # Get the frame defined by the CA-C-N atom
        frames = Affine3D.from_graham_schmidt(CA, C, N)
        oxygen_coordinates = frames.apply(O_vector)
        atom37_positions = self.atom37_positions.copy()
        atom37_mask = self.atom37_mask.copy()

        atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[
            O_missing_indices
        ].numpy()
        atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
            atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
        ).any(-1)
        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
        return new_chain

    def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex:
        """Return a new chain with inferred CB atoms at all residues except GLY.

        Args:
            infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine
                residues, even though that residue doesn't have one.  Default off.

                NOTE(rverkuil): The reason for having this switch in the first place
                is that sometimes we want a (inferred) CB coordinate for every residue,
                for example for making a pairwise distance matrix, or doing an RMSD
                calculation between two designs for a given structural template, w/
                CB atoms.
        """
        atom37_positions = self.atom37_positions.copy()
        atom37_mask = self.atom37_mask.copy()

        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
        # See usage in trDesign codebase.
        # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
        inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143)
        if not infer_cbeta_for_glycine:
            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan

        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions
        atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
            atom37_positions[:, residue_constants.atom_order["CB"]]
        ).any(-1)
        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
        return new_chain

    @classmethod
    def from_open_source(cls, pc: ProteinComplex):
        # TODO(@zeming): deprecated, should delete
        return pc

    @classmethod
    def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex:
        pdb_ids = [obj.id for obj in objs]
        if len(set(pdb_ids)) > 1:
            raise RuntimeError(
                "Concatention of protein complexes across different PDB ids is unsupported"
            )
        return ProteinComplex.from_chains(
            list(itertools.chain.from_iterable(obj.chain_iter() for obj in objs))
        )

    def _sanity_check_complexes_are_comparable(self, other: ProteinComplex):
        if len(self) != len(other):
            raise ValueError("Protein complexes must have the same length")
        if len(list(self.chain_iter())) != len(list(other.chain_iter())):
            raise ValueError("Protein complexes must have the same number of chains")

    def rmsd(
        self,
        target: ProteinComplex,
        also_check_reflection: bool = False,
        mobile_inds: list[int] | np.ndarray | None = None,
        target_inds: list[int] | np.ndarray | None = None,
        only_compute_backbone_rmsd: bool = False,
        compute_chain_assignment: bool = True,
    ):
        """
        Compute the RMSD between this protein chain and another.

        Args:
            target (ProteinComplex): The target (other) protein complex to compare to.
            also_check_reflection: Compare the reflected mobile coordinates too.
            mobile_inds: Mobile atom indices, not residue indices.
            target_inds: Target atom indices, not residue indices.
            only_compute_backbone_rmsd: Restrict the score to backbone atoms.
        """
        aligned = self.dockq(target).aligned if compute_chain_assignment else self

        aligner = Aligner(
            aligned if mobile_inds is None else aligned[mobile_inds],
            target if target_inds is None else target[target_inds],
            only_compute_backbone_rmsd,
        )
        avg_rmsd = aligner.rmsd

        if not also_check_reflection:
            return avg_rmsd

        aligner = Aligner(
            aligned if mobile_inds is None else aligned[mobile_inds],
            target if target_inds is None else target[target_inds],
            only_compute_backbone_rmsd,
            use_reflection=True,
        )
        avg_rmsd_neg = aligner.rmsd

        return min(avg_rmsd, avg_rmsd_neg)

    def lddt_ca(
        self,
        target: ProteinComplex,
        mobile_inds: list[int] | np.ndarray | None = None,
        target_inds: list[int] | np.ndarray | None = None,
        compute_chain_assignment: bool = True,
        **kwargs,
    ) -> float | np.ndarray:
        """Compute the LDDT between this protein complex and another.

        Arguments:
            target (ProteinComplex): The other protein complex to compare to.
            mobile_inds: Mobile atom indices, not residue indices.
            target_inds: Target atom indices, not residue indices.

        Returns:
            float | np.ndarray: The LDDT score between the two protein chains, either
                a single float or per-residue LDDT scores if `per_residue` is True.
        """
        aligned = self.dockq(target).aligned if compute_chain_assignment else self
        lddt = compute_lddt_ca(
            torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0),
            torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0),
            torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0),
            **kwargs,
        )
        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()

    def gdt_ts(
        self,
        target: ProteinComplex,
        mobile_inds: list[int] | np.ndarray | None = None,
        target_inds: list[int] | np.ndarray | None = None,
        compute_chain_assignment: bool = True,
        **kwargs,
    ) -> float | np.ndarray:
        """Compute the GDT_TS between this protein complex and another.

        Arguments:
            target (ProteinComplex): The other protein complex to compare to.
            mobile_inds: Mobile atom indices, not residue indices.
            target_inds: Target atom indices, not residue indices.

        Returns:
            float: The GDT_TS score between the two protein chains.
        """
        aligned = self.dockq(target).aligned if compute_chain_assignment else self
        gdt_ts = compute_gdt_ts(
            mobile=torch.tensor(
                index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"),
                dtype=torch.float32,
            ).unsqueeze(0),
            target=torch.tensor(
                index_by_atom_name(target.atom37_positions[target_inds], "CA"),
                dtype=torch.float32,
            ).unsqueeze(0),
            atom_exists_mask=torch.tensor(
                index_by_atom_name(aligned.atom37_mask[mobile_inds], "CA", dim=-1)
                & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
            ).unsqueeze(0),
            **kwargs,
        )
        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()

    def dockq(self, native: ProteinComplex):
        # This function uses dockqv2 to compute the DockQ score. Because it does a mapping
        # over all possible chains, it's quite slow. Be careful not to use this in an inference loop
        # or something that requires fast scoring. It defaults to 8 CPUs.
        #
        # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif,
        # if your protein has multi-letter or repeated chain IDs, this will fail. Please call
        # Normalize chain IDs before DockQ when IDs repeat or use multiple letters.

        try:
            pass
        except BaseException:
            raise RuntimeError("DockQ is not installed. Please update your environment.") from None
        self._sanity_check_complexes_are_comparable(native)

        def sanity_check_chain_ids(pc: ProteinComplex):
            ids = []
            for i, chain in enumerate(pc.chain_iter()):
                if i > len(SINGLE_LETTER_CHAIN_IDS):
                    raise ValueError("Too many chains to write to PDB file")
                if len(chain.chain_id) > 1:
                    raise ValueError("We only supports single letter chain IDs for DockQ")
                ids.append(chain.chain_id)
            if len(set(ids)) != len(ids):
                raise ValueError(f"Duplicate chain IDs in protein complex: {ids}")
            return ids

        sanity_check_chain_ids(self)
        sanity_check_chain_ids(native)

        with TemporaryDirectory() as tdir:
            dir = Path(tdir)
            self.to_pdb(dir / "self.pdb")
            native.to_pdb(dir / "native.pdb")

            output = check_output(["DockQ", dir / "self.pdb", dir / "native.pdb"])
        lines = output.decode().split("\n")

        # Remove the header comments
        start_index = next(i for i, line in enumerate(lines) if line.startswith("Model"))
        lines = lines[start_index:]

        result = {}
        interfaces = []
        current_interface: dict = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if line.startswith(("Model  :", "Native :")):
                pass  # Tmp pdb file location, it's useless...
            elif line.startswith("Total DockQ"):
                total_dockq_match = re.search(
                    r"Total DockQ over (\d+) native interfaces: ([\d.]+) with "
                    r"(.*) model:native mapping",
                    line,
                )
                if total_dockq_match:
                    result["value"] = float(total_dockq_match.group(2))
                    result["native interfaces"] = int(total_dockq_match.group(1))
                    native_chains, self_chains = total_dockq_match.group(3).split(":")
                    result["mapping"] = dict(zip(native_chains, self_chains, strict=False))
                else:
                    raise RuntimeError(
                        "Failed to parse DockQ output, maybe your DockQ version is wrong?"
                    )
            elif line.startswith("Native chains:"):
                if current_interface:
                    interfaces.append(current_interface)
                current_interface = {"Native chains": line.split(":")[1].strip().split(", ")}
            elif line.startswith("Model chains:"):
                current_interface["Model chains"] = line.split(":")[1].strip().split(", ")
            elif ":" in line:
                key, value = line.split(":", 1)
                current_interface[key.strip()] = float(value.strip())

        if current_interface:
            interfaces.append(current_interface)

        def parse_dict(d: dict[str, Any]) -> DockQSingleScore:
            return DockQSingleScore(
                native_chains=tuple(d["Native chains"]),  # type: ignore
                DockQ=float(d["DockQ"]),
                interface_rms=float(d["irms"]),
                ligand_rms=float(d["Lrms"]),  # Note the capitalization difference
                fnat=float(d["fnat"]),
                fnonnat=float(d["fnonnat"]),
                clashes=float(d["clashes"]),
                F1=float(d["F1"]),
                DockQ_F1=float(d["DockQ_F1"]),
            )

        inv_mapping = {v: k for k, v in result["mapping"].items()}

        self_chain_map = {c.chain_id: c for c in self.chain_iter()}
        realigned = []
        for chain in native.chain_iter():
            realigned.append(self_chain_map[inv_mapping[chain.chain_id]])

        realigned = ProteinComplex.from_chains(realigned)
        aligner = Aligner(realigned, native)
        realigned = aligner.apply(realigned)

        result = DockQResult(
            total_dockq=result["value"],
            native_interfaces=result["native interfaces"],
            chain_mapping=result["mapping"],
            interfaces={
                (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) for i in interfaces
            },
            aligned=realigned,
            aligned_rmsd=aligner.rmsd,
        )

        return result

    # Object invariants, slicing, and chain views
    def __post_init__(self):
        if not isinstance(self.sequence, str):
            raise TypeError("sequence must be a string.")
        sequence_length = len(self.sequence)
        aligned = {
            "atom37_positions": self.atom37_positions,
            "atom37_mask": self.atom37_mask,
            "residue_index": self.residue_index,
            "insertion_code": self.insertion_code,
            "confidence": self.confidence,
            "entity_id": self.entity_id,
            "chain_id": self.chain_id,
            "sym_id": self.sym_id,
        }
        for name, values in aligned.items():
            if not isinstance(values, np.ndarray):
                raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.")
            if values.ndim == 0 or values.shape[0] != sequence_length:
                raise ValueError(
                    f"{name} shape {values.shape} does not align with "
                    f"sequence length {sequence_length}."
                )
        if self.atom37_positions.shape != (sequence_length, 37, 3):
            raise ValueError(
                "atom37_positions must have shape "
                f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}."
            )
        if self.atom37_mask.shape != (sequence_length, 37):
            raise ValueError(
                "atom37_mask must have shape "
                f"({sequence_length}, 37), got {self.atom37_mask.shape}."
            )
        if self.atom37_mask.dtype != bool:
            raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.")
        if not np.issubdtype(self.atom37_positions.dtype, np.number):
            raise TypeError("atom37_positions must use a numeric dtype.")
        for name, values in (
            ("residue_index", self.residue_index),
            ("insertion_code", self.insertion_code),
            ("confidence", self.confidence),
            ("entity_id", self.entity_id),
            ("chain_id", self.chain_id),
            ("sym_id", self.sym_id),
        ):
            if values.shape != (sequence_length,):
                raise ValueError(
                    f"{name} must have shape ({sequence_length},), got {values.shape}."
                )
        if not np.issubdtype(self.confidence.dtype, np.number):
            raise TypeError("confidence must use a numeric dtype.")
        atom37_confidence = self.atom37_confidence
        if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray):
            raise TypeError("atom37_confidence must be a NumPy array when provided.")
        if (
            isinstance(atom37_confidence, np.ndarray)
            and atom37_confidence.shape != self.atom37_mask.shape
        ):
            raise ValueError(
                "atom37_confidence shape must match atom37_mask: "
                f"{atom37_confidence.shape} != {self.atom37_mask.shape}."
            )

    def __getitem__(self, idx: int | list[int] | slice | np.ndarray):
        """This function slices protein complexes without consideration of chain breaks
        NOTE: When slicing with a boolean mask, it's possible that the output array won't
        be the expected length. This is because we do our best to preserve chainbreak tokens.
        """

        if isinstance(idx, int):
            idx = [idx]
        if isinstance(idx, list):
            raise ValueError("ProteinComplex doesn't supports indexing with lists of indices")

        if isinstance(idx, np.ndarray):
            is_chainbreak = np.asarray([s == "|" for s in self.sequence])
            idx = idx.astype(bool) | is_chainbreak

        complex = self._unsafe_slice(idx)
        if len(complex) == 0:
            return complex

        # detect runs of chainbreaks by searching for instances of '||' in complex.sequence
        chainbreak_runs = np.asarray(
            [complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1)]
            + [complex.sequence[-1] == "|"]
        )
        # We should remove as many chainbreaks as possible from the start of the sequence
        for i in range(len(chainbreak_runs)):
            if complex.sequence[i] == "|":
                chainbreak_runs[i] = True
            else:
                break
        complex = complex._unsafe_slice(~chainbreak_runs)
        return complex

    def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray):
        sequence = slice_python_object_as_numpy(self.sequence, idx)
        return replace(
            self,
            sequence=sequence,
            entity_id=self.entity_id[..., idx],
            chain_id=self.chain_id[..., idx],
            sym_id=self.sym_id[..., idx],
            residue_index=self.residue_index[..., idx],
            insertion_code=self.insertion_code[..., idx],
            atom37_positions=self.atom37_positions[..., idx, :, :],
            atom37_mask=self.atom37_mask[..., idx, :],
            confidence=self.confidence[..., idx],
            atom37_confidence=self.atom37_confidence[..., idx, :]
            if self.atom37_confidence is not None
            else None,
        )

    def __len__(self):
        return len(self.sequence)

    @property
    def num_chains(self):
        return len(self.chain_boundaries)

    @cached_property
    def atoms(self) -> AtomIndexer:
        return AtomIndexer(self, property="atom37_positions", dim=-2)

    @cached_property
    def atom_mask(self) -> AtomIndexer:
        return AtomIndexer(self, property="atom37_mask", dim=-1)

    @cached_property
    def chain_lengths(self) -> np.ndarray:
        return np.diff(self.chain_boundaries, axis=1).flatten()

    @cached_property
    def chain_boundaries(self) -> list[tuple[int, int]]:
        cb = [-1]
        for i, s in enumerate(self.sequence):
            if s == "|":
                cb.append(i)
        cb.append(len(self))
        return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)]

    def get_chain_by_index(self, index: int) -> ProteinChain:
        try:
            start, end = self.chain_boundaries[index]
            return self[start:end].as_chain()
        except IndexError:
            raise IndexError(f"Chain index {index} out of bounds") from None

    def get_chain_by_id(
        self, chain_id: str, sample_chain_if_duplicate: bool = True
    ) -> ProteinChain:
        valid_indices = [
            index
            for index, id_of_index in self.metadata.chain_lookup.items()
            if id_of_index == chain_id
        ]
        if not valid_indices:
            raise KeyError(f"Chain ID {chain_id} not found")
        if sample_chain_if_duplicate:
            index_to_return = random.choice(valid_indices)
            return self.get_chain_by_index(index_to_return)
        else:
            if len(valid_indices) > 1:
                raise ValueError(f"Multiple chains with chain ID {chain_id} found")
            return self.get_chain_by_index(valid_indices[0])

    def chain_iter(self) -> Iterable[ProteinChain]:
        for start, end in self.chain_boundaries:
            c = self[start:end]
            yield c.as_chain()

    def as_chain(self, force_conversion: bool = False) -> ProteinChain:
        """Convert the ProteinComplex to a ProteinChain.

        Args:
            force_conversion: Flatten multiple chains to access chain-only utilities.

        """
        if not force_conversion:
            if len(np.unique(self.chain_id)) != 1:
                raise ValueError(
                    f"Protein complex {self.id!r} has multiple chains; "
                    "pass force_conversion=True to flatten it."
                )
            if len(np.unique(self.entity_id)) != 1:
                raise ValueError(
                    f"Protein complex {self.id!r} has multiple entities; "
                    "pass force_conversion=True to flatten it."
                )
            if self.chain_id[0] not in self.metadata.chain_lookup:
                warnings.warn(
                    "Chain ID not found in metadata, using 'A' as default",
                    stacklevel=2,
                )
            if self.entity_id[0] not in self.metadata.entity_lookup:
                warnings.warn(
                    "Entity ID not found in metadata, using None as default",
                    stacklevel=2,
                )
            chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A")
            entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None)
        else:
            chain_id = "A"
            entity_id = None

        return ProteinChain(
            id=self.id,
            sequence=self.sequence,
            chain_id=chain_id,
            entity_id=entity_id,
            atom37_positions=self.atom37_positions,
            atom37_mask=self.atom37_mask,
            residue_index=self.residue_index,
            insertion_code=self.insertion_code,
            confidence=self.confidence,
            mmcif=self.metadata.mmcif,
            atom37_confidence=self.atom37_confidence,
        )

    # Contact topology and mmCIF export
    @cached_property
    def per_chain_kd_trees(self):
        # Iterate over chains, build KDTree for each chain
        kdtrees = []

        CA = self.atoms["CA"]

        for start, end in self.chain_boundaries:
            chain_CA = CA[start:end]
            chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)]
            kdtrees.append(KDTree(chain_CA))

        return kdtrees

    def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray:
        # Compute adjacency matrix for protein complex
        num_chains = self.num_chains
        adjacency = np.zeros((num_chains, num_chains), dtype=bool)
        for (i, kdtree), (j, kdtree2) in itertools.combinations(
            enumerate(self.per_chain_kd_trees), 2
        ):
            adj = kdtree.query_ball_tree(kdtree2, cutoff)
            any_is_adjacent = any(len(a) > 0 for a in adj)
            adjacency[i, j] = any_is_adjacent
            adjacency[j, i] = any_is_adjacent
        return adjacency

    def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray:
        num_chains = len(self.chain_boundaries)
        adjacency = np.zeros(num_chains, dtype=bool)
        for i, kdtree in enumerate(self.per_chain_kd_trees):
            if i == index:
                continue
            adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff)
            adjacency[i] = any(len(a) > 0 for a in adj)
        return adjacency

    def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex:
        """Rename all chains in the complex with a given prefix.

        Args:
            prefix (str): The prefix to use for the new chain IDs. Each chain will be
                named as "{prefix}_{chain_id}".

        Returns:
            ProteinComplex: A new protein complex with renamed chains.
        """
        new_chains = []
        for chain in self.chain_iter():
            # Create new chain with updated chain_id
            new_chain = replace(chain, chain_id=f"{prefix}_{chain.chain_id}")
            new_chains.append(new_chain)
        return ProteinComplex.from_chains(new_chains)

    def sasa(self, by_residue: bool = True):
        chain = self.as_chain(force_conversion=True)
        return chain.sasa(by_residue=by_residue)

    def to_mmcif_string(self) -> str:
        """Convert the ProteinComplex to mmCIF format.

        Returns:
            str: The mmCIF content as a string.
        """
        # Convert the ProteinComplex to a biotite AtomArray
        # Collect all atoms from all chains
        all_atoms = []
        for chain in self.chain_iter():
            chain_atom_array = chain.atom_array
            # Convert AtomArray to list of atoms and add to collection
            all_atoms.extend(chain_atom_array)

        # Create combined AtomArray from all atoms
        if not all_atoms:
            raise ValueError("No atoms found in protein complex")

        atom_array = bs.array(all_atoms)

        # Create CIF file
        f = CIFFile()
        set_structure_pdbx(f, atom_array, data_block=self.id)

        # Add entity information for proper mmCIF structure
        self._add_entity_information(f)
        round_mmcif_columns(f)

        # Write to string
        output = io.StringIO()
        f.write(output)
        return output.getvalue()

    def _add_entity_information(self, cif_file: CIFFile) -> None:
        """Add entity, entity_poly, and struct_asym sections to CIF file."""

        # Group chains by sequence to create unique entities
        entity_map = {}  # sequence -> entity_id
        chain_to_entity = {}  # chain_id -> entity_id
        entity_sequences = {}  # entity_id -> sequence
        entity_id_counter = 1

        for chain in self.chain_iter():
            sequence = chain.sequence
            if sequence not in entity_map:
                entity_map[sequence] = entity_id_counter
                entity_sequences[entity_id_counter] = sequence
                entity_id_counter += 1
            chain_to_entity[chain.chain_id] = entity_map[sequence]

        # Create _entity section
        entity_ids = []
        entity_types = []
        entity_descriptions = []

        for entity_id in sorted(entity_sequences.keys()):
            entity_ids.append(str(entity_id))
            entity_types.append("polymer")
            entity_descriptions.append(f"Protein chain (entity {entity_id})")

        cif_file.block["entity"] = CIFCategory(
            name="entity",
            columns={
                "id": CIFColumn(data=CIFData(array=np.array(entity_ids), dtype=np.str_)),
                "type": CIFColumn(data=CIFData(array=np.array(entity_types), dtype=np.str_)),
                "pdbx_description": CIFColumn(
                    data=CIFData(array=np.array(entity_descriptions), dtype=np.str_)
                ),
            },
        )

        # Create _entity_poly section
        poly_entity_ids = []
        poly_types = []
        poly_nstd_linkages = []
        poly_sequences = []

        for entity_id in sorted(entity_sequences.keys()):
            poly_entity_ids.append(str(entity_id))
            poly_types.append("polypeptide(L)")
            poly_nstd_linkages.append("no")
            poly_sequences.append(entity_sequences[entity_id])

        cif_file.block["entity_poly"] = CIFCategory(
            name="entity_poly",
            columns={
                "entity_id": CIFColumn(
                    data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_)
                ),
                "type": CIFColumn(data=CIFData(array=np.array(poly_types), dtype=np.str_)),
                "nstd_linkage": CIFColumn(
                    data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_)
                ),
                "pdbx_seq_one_letter_code": CIFColumn(
                    data=CIFData(array=np.array(poly_sequences), dtype=np.str_)
                ),
            },
        )

        # Create _struct_asym section
        asym_ids = []
        asym_entity_ids = []
        asym_details = []

        for chain in self.chain_iter():
            asym_ids.append(chain.chain_id)
            asym_entity_ids.append(str(chain_to_entity[chain.chain_id]))
            asym_details.append("")

        cif_file.block["struct_asym"] = CIFCategory(
            name="struct_asym",
            columns={
                "id": CIFColumn(data=CIFData(array=np.array(asym_ids), dtype=np.str_)),
                "entity_id": CIFColumn(
                    data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_)
                ),
                "details": CIFColumn(data=CIFData(array=np.array(asym_details), dtype=np.str_)),
            },
        )

    # Construction, PDB interchange, and compact storage
    @classmethod
    def from_pdb(
        cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False
    ) -> ProteinComplex:
        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])

        chains = []
        for chain in bs.chain_iter(atom_array):
            chain = chain[~chain.hetero]
            if len(chain) == 0:
                continue
            chains.append(ProteinChain.from_atomarray(chain, id, is_predicted))
        return ProteinComplex.from_chains(chains)

    def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
        atom_array = None
        for chain in self.chain_iter():
            carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions
            atom_array = carr if atom_array is None else atom_array + carr
        f = PDBFile()
        f.set_structure(atom_array)
        f.write(path)

    def to_pdb_string(self, include_insertions: bool = True) -> str:
        buf = io.StringIO()
        self.to_pdb(buf, include_insertions=include_insertions)
        buf.seek(0)
        return buf.read()

    def normalize_chain_ids_for_pdb(self):
        # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index,
        # we can normalize it instead which might be necessary for DockQ and to_pdb.
        ids = SINGLE_LETTER_CHAIN_IDS
        chains = []
        for i, chain in enumerate(self.chain_iter()):
            chain = replace(chain, chain_id=ids[i])
            if i > len(ids):
                raise RuntimeError("Too many chains to write to PDB file")
            chains.append(chain)

        return ProteinComplex.from_chains(chains)

    def find_assembly_ids_with_chain(self, id: str) -> list[str]:
        good_chains = []
        if (comp := self.metadata.assembly_composition) is not None:
            for assembly_id, chain_ids in comp.items():
                if id in chain_ids:
                    good_chains.append(assembly_id)
        else:
            raise ValueError(
                "Cannot switch assemblies on this ProteinComplex; construct it from "
                "mmCIF to retain assembly metadata"
            )
        return good_chains

    def switch_assembly(self, id: str):
        if self.metadata.mmcif is None:
            raise ValueError(
                "Cannot switch assemblies without retained mmCIF source metadata."
            )
        return get_assembly_fast(self.metadata.mmcif, assembly_id=id)

    def state_dict(self, backbone_only=False, json_serializable=False):
        """This state dict is optimized for storage, so it turns things to fp16 whenever
        possible. Note that we also only support int32 residue indices, I'm hoping we don't
        need more than 2**32 residues..."""
        dct = {k: v for k, v in vars(self).items()}
        if backbone_only:
            # Frozen dataclasses do not make their NumPy members immutable.  Work on a
            # private mask so requesting a compact backbone payload cannot clear the
            # caller's side-chain atoms in-place.
            atom37_mask = dct["atom37_mask"].copy()
            atom37_mask[:, 3:] = False
            dct["atom37_mask"] = atom37_mask
        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
        if dct.get("atom37_confidence") is not None:
            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
        else:
            dct.pop("atom37_confidence", None)
        for k, v in dct.items():
            if isinstance(v, np.ndarray):
                match v.dtype:
                    case np.int64:
                        dct[k] = v.astype(np.int32)
                    case np.float64 | np.float32:
                        dct[k] = v.astype(np.float16)
                    case _:
                        pass
                if json_serializable:
                    dct[k] = v.tolist()
            elif isinstance(v, ProteinComplexMetadata):
                dct[k] = asdict(v)
        dct["metadata"]["mmcif"] = None
        # These can be populated with non-serializable objects and are not needed for reconstruction
        dct.pop("atoms", None)
        dct.pop("atom_mask", None)
        dct.pop("per_chain_kd_trees", None)
        return dct

    def to_blob(self, backbone_only=False) -> bytes:
        payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode)
        return brotli.compress(payload, quality=5)

    @classmethod
    def from_state_dict(cls, dct):
        # Note: assembly_composition is *supposed* to have string keys.
        dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"])

        for k, v in dct.items():
            if isinstance(v, list):
                dct[k] = np.array(v)

        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
        atom37[dct["atom37_mask"]] = dct["atom37_positions"]
        dct["atom37_positions"] = atom37
        if "atom37_confidence" in dct:
            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
            dct["atom37_confidence"] = atom37_conf
        dct = {
            k: (
                v.astype(np.float32)
                if k in ["atom37_positions", "confidence", "atom37_confidence"]
                else v
            )
            for k, v in dct.items()
        }
        if "chain_boundaries" in dct:
            del dct["chain_boundaries"]
        if "chain_boundaries" in dct["metadata"]:
            del dct["metadata"]["chain_boundaries"]
        dct["metadata"] = ProteinComplexMetadata(**dct["metadata"])
        return cls(**dct)

    @classmethod
    def from_blob(cls, input: Path | str | io.BytesIO | bytes):
        """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory
        of chains from 52G/1M chains to 20G/1M chains, I think this is a good first
        shot at compressing and dumping chains to disk. I'm sure there's better ways."""
        match input:
            case Path() | str():
                bytes = Path(input).read_bytes()
            case io.BytesIO():
                bytes = input.getvalue()
            case _:
                bytes = input
        state = msgpack.loads(
            brotli.decompress(bytes),
            object_hook=msgpack_numpy.decode,
            strict_map_key=False,
        )
        return cls.from_state_dict(state)

    @classmethod
    def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex:
        f: io.StringIO = rcsb.fetch(pdb_id, "cif")  # type: ignore
        return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False)

    @classmethod
    def from_mmcif(
        cls,
        path: PathOrBuffer,
        id: str | None = None,
        assembly_id: str | None = None,
        is_predicted: bool = False,
        keep_source: bool = False,
    ):
        """Return a ProteinComplex object from an mmcif file.
        TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation,
        we only consider the first defined complex!

        Args:
            path: Uncompressed mmCIF path or text buffer.
            id: Optional structure identifier.
            is_predicted (bool): If True, reads b factor as the confidence readout. Default: False.
            chain_id (str, optional): Select a chain corresponding to (author) chain id.
        """
        mmcif = MmcifWrapper.read(path, id)
        return get_assembly_fast(mmcif, assembly_id=assembly_id)

    @classmethod
    def from_chains(
        cls,
        chains: Sequence[ProteinChain],
        mmcif: MmcifWrapper | None = None,
        all_assembly_metadata_dictionary: dict[str, list[str]] | None = None,
    ):
        if not chains:
            raise ValueError("Cannot create a ProteinComplex from an empty list of chains")

        # TODO(roshan): Make a proper protein complex class
        def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
            full_array = []
            for array in arrays:
                full_array.append(array)
                full_array.append(sep)
            full_array = full_array[:-1]
            return np.concatenate(full_array, 0)

        sep_tokens = {
            "residue_index": np.array([-1]),
            "insertion_code": np.array([""]),
            "atom37_positions": np.full([1, 37, 3], np.nan),
            "atom37_mask": np.zeros([1, 37], dtype=bool),
            "confidence": np.array([0]),
        }

        any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains)
        if any_has_atom37_conf:
            sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32)

        def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray:
            val = getattr(chain, name)
            if val is None and name == "atom37_confidence":
                return np.full([len(chain), 37], np.nan, dtype=np.float32)
            return val

        array_args: dict[str, np.ndarray] = {
            name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep)
            for name, sep in sep_tokens.items()
        }

        multimer_arrays = []
        chain2num_max = -1
        chain2num = {}
        ent2num_max = -1
        ent2num = {}
        total_index = 0
        for i, c in enumerate(chains):
            num_res = c.residue_index.shape[0]
            if c.chain_id not in chain2num:
                chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1)
            chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64)

            if c.entity_id is None:
                entity_num = (ent2num_max := ent2num_max + 1)
            else:
                if c.entity_id not in ent2num:
                    ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1)
                entity_num = ent2num[c.entity_id]
            entity_id_array = np.full([num_res], entity_num, dtype=np.int64)

            sym_id_array = np.full([num_res], i, dtype=np.int64)

            multimer_arrays.append(
                {
                    "chain_id": chain_id_array,
                    "entity_id": entity_id_array,
                    "sym_id": sym_id_array,
                }
            )

            total_index += num_res + 1

        sep = np.array([-1])
        update = {
            name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep)
            for name in ["chain_id", "entity_id", "sym_id"]
        }
        array_args.update(update)

        metadata = ProteinComplexMetadata(
            mmcif=mmcif,
            chain_lookup={v: k for k, v in chain2num.items()},
            entity_lookup={v: k for k, v in ent2num.items()},
            assembly_composition=all_assembly_metadata_dictionary,
        )

        return cls(
            id=chains[0].id,
            sequence=residue_constants.CHAIN_BREAK_TOKEN.join(chain.sequence for chain in chains),
            metadata=metadata,
            **array_args,
        )


# Biological-assembly expansion
def get_assembly_fast(
    mmcif: MmcifWrapper,
    assembly_id=None,
    model=None,
    data_block=None,
    altloc="first",
    use_author_fields=True,
):
    pdbx_file = mmcif.raw
    if pdbx_file is None:
        raise InvalidFileError("No mmCIF data loaded")
    assembly_gen_category = pdbx_file.block["pdbx_struct_assembly_gen"]
    if assembly_gen_category is None:
        raise InvalidFileError("File has no 'pdbx_struct_assembly_gen' category")

    struct_oper_category = pdbx_file.block["pdbx_struct_oper_list"]
    if struct_oper_category is None:
        raise InvalidFileError("File has no 'pdbx_struct_oper_list' category")

    if assembly_id is None:
        assembly_id = assembly_gen_category["assembly_id"].data.array[0]
    elif assembly_id not in assembly_gen_category["assembly_id"].data.array:
        raise KeyError(f"File has no Assembly ID '{assembly_id}'")

    ### Calculate all possible transformations
    transformations = _get_transformations(struct_oper_category)

    ### Get structure according to additional parameters
    structure = get_structure(
        pdbx_file, model, data_block, altloc, ["label_asym_id"], use_author_fields
    )[0]  # type: ignore
    # TODO(@zeming) This line will remove all non-protein structural elements,
    # we should remove this when we want to parse these too.
    structure: bs.AtomArray = structure[
        bs.filter_amino_acids(structure) & ~structure.hetero  # type: ignore
    ]
    if len(structure) == 0:
        raise NoProteinError
    unique_asym_ids = np.unique(structure.label_asym_id)  # type: ignore
    asym2chain = {}
    asym2auth = {}
    for asym_id in unique_asym_ids:
        sub_structure: bs.AtomArray = structure[structure.label_asym_id == asym_id]  # type: ignore
        chain_id: str = sub_structure[0].chain_id  # type: ignore
        (
            sequence,
            atom_positions,
            atom_mask,
            residue_index,
            insertion_code,
            confidence,
            entity_id,
        ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False)

        asym2chain[asym_id] = ProteinChain(
            id=mmcif.id or "unknown",
            sequence=sequence,
            chain_id=chain_id,
            entity_id=entity_id,
            atom37_positions=atom_positions,
            atom37_mask=atom_mask,
            residue_index=residue_index,
            insertion_code=insertion_code,
            confidence=confidence,
            mmcif=None,
        )
        asym2auth[asym_id] = chain_id

    ### Get transformations and apply them to the affected asym IDs
    assembly = []
    assembly_id_dict: dict[str, list[str]] = {}

    # Process the target assembly ID
    for aid, op_expr, asym_id_expr in zip(
        assembly_gen_category["assembly_id"].data.array,
        assembly_gen_category["oper_expression"].data.array,
        assembly_gen_category["asym_id_list"].data.array,
        strict=False,
    ):
        if aid == assembly_id:
            # Parse operations and asym IDs for this specific entry
            operations = _parse_operation_expression(op_expr)
            asym_ids = asym_id_expr.split(",")

            # Filter affected asym IDs to only protein chains, preserving order
            sub_structures = [asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain]

            # Apply transformations
            sub_assembly = _apply_transformations_fast(sub_structures, transformations, operations)
            assembly.extend(sub_assembly)

            # Build assembly_id_dict for this entry
            assembly_id_dict[aid] = assembly_id_dict.get(aid, []) + [
                asym2auth[id_] for id_ in asym_ids if id_ in asym2auth
            ]

    if len(assembly) == 0:
        raise NoProteinError
    return ProteinComplex.from_chains(assembly, mmcif, assembly_id_dict)


def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex:
    if "|" not in chain.sequence:
        return ProteinComplex.from_chains([chain])
    chain_breaks = np.array(list(chain.sequence)) == "|"
    chain_break_inds = np.where(chain_breaks)[0]
    chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]])
    chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds)))
    complex_chains = []
    for start, end in chain_break_inds:
        if start != 0:
            start += 1
        complex_chains.append(chain[start:end])
    complex_chains = [
        ProteinChain.from_atom37(
            chain.atom37_positions,
            sequence=chain.sequence,
            chain_id=SINGLE_LETTER_CHAIN_IDS[i],
            entity_id=i,
        )
        for i, chain in enumerate(complex_chains)
    ]
    return ProteinComplex.from_chains(complex_chains)