flax_model/alphafold3/structure/bonds.py

"""Bond representation for structure module."""

import collections
from collections.abc import Mapping, Sequence
import dataclasses
import typing
from typing import Self

from flax_model.alphafold3.structure import table
import numpy as np


@dataclasses.dataclass(frozen=True, kw_only=True)
class Bonds(table.Table):
  """Table of atomic bonds."""

  # mmCIF column: _struct_conn.conn_type_id
  # mmCIF desc: This data item is a pointer to _struct_conn_type.id in the
  #             STRUCT_CONN_TYPE category.
  # E.g.: "covale", "disulf", "hydrog", "metalc".
  type: np.ndarray

  # mmCIF column: _struct_conn.pdbx_role
  # mmCIF desc: The chemical or structural role of the interaction.
  # E.g.: "N-Glycosylation", "O-Glycosylation".
  role: np.ndarray

  # mmCIF columns: _struct_conn.ptnr1_*
  from_atom_key: np.ndarray

  # mmCIF columns: _struct_conn.ptnr2_*
  dest_atom_key: np.ndarray

  @classmethod
  def make_empty(cls) -> Self:
    return cls(
        key=np.empty((0,), dtype=np.int64),
        from_atom_key=np.empty((0,), dtype=np.int64),
        dest_atom_key=np.empty((0,), dtype=np.int64),
        type=np.empty((0,), dtype=object),
        role=np.empty((0,), dtype=object),
    )

  def get_atom_indices(
      self,
      atom_key: np.ndarray,
  ) -> tuple[np.ndarray, np.ndarray]:
    """Returns the indices of the from/dest atoms in the atom_key array."""
    from_atom_missing = ~np.isin(self.from_atom_key, atom_key)
    dest_atom_missing = ~np.isin(self.dest_atom_key, atom_key)
    if np.any(from_atom_missing):
      raise ValueError(
          f'No atoms for from_atom_key {self.from_atom_key[from_atom_missing]}'
      )
    if np.any(dest_atom_missing):
      raise ValueError(
          f'No atoms for dest_atom_key {self.dest_atom_key[dest_atom_missing]}'
      )
    sort_indices = np.argsort(atom_key)
    from_indices_sorted = np.searchsorted(
        atom_key, self.from_atom_key, sorter=sort_indices
    )
    dest_indices_sorted = np.searchsorted(
        atom_key, self.dest_atom_key, sorter=sort_indices
    )
    from_indices = sort_indices[from_indices_sorted]
    dest_indices = sort_indices[dest_indices_sorted]
    return from_indices, dest_indices

  def restrict_to_atoms(self, atom_key: np.ndarray) -> Self:
    if not self.size:  # Early-out for empty table.
      return self
    from_atom_mask = np.isin(self.from_atom_key, atom_key)
    dest_atom_mask = np.isin(self.dest_atom_key, atom_key)
    mask = np.logical_and(from_atom_mask, dest_atom_mask)
    return typing.cast(Bonds, self.filter(mask=mask))

  def to_mmcif_dict_from_atom_arrays(
      self,
      atom_key: np.ndarray,
      chain_id: np.ndarray,
      res_id: np.ndarray,
      res_name: np.ndarray,
      atom_name: np.ndarray,
      auth_asym_id: np.ndarray,
      auth_seq_id: np.ndarray,
      insertion_code: np.ndarray,
  ) -> Mapping[str, Sequence[str] | np.ndarray]:
    """Returns a dict suitable for building a CifDict, representing bonds.

    Args:
      atom_key: A (num_atom,) integer array of atom_keys.
      chain_id: A (num_atom,) array of label_asym_id strings.
      res_id: A (num_atom,) array of label_seq_id strings.
      res_name: A (num_atom,) array of label_comp_id strings.
      atom_name: A (num_atom,) array of label_atom_id strings.
      auth_asym_id: A (num_atom,) array of auth_asym_id strings.
      auth_seq_id: A (num_atom,) array of auth_seq_id strings.
      insertion_code: A (num_atom,) array of insertion code strings.
    """
    mmcif_dict = collections.defaultdict(list)
    ptnr1_indices, ptnr2_indices = self.get_atom_indices(atom_key)

    mmcif_dict['_struct_conn.ptnr1_label_asym_id'] = chain_id[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_label_asym_id'] = chain_id[ptnr2_indices]
    mmcif_dict['_struct_conn.ptnr1_label_comp_id'] = res_name[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_label_comp_id'] = res_name[ptnr2_indices]
    mmcif_dict['_struct_conn.ptnr1_label_seq_id'] = res_id[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_label_seq_id'] = res_id[ptnr2_indices]
    mmcif_dict['_struct_conn.ptnr1_label_atom_id'] = atom_name[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_label_atom_id'] = atom_name[ptnr2_indices]

    mmcif_dict['_struct_conn.ptnr1_auth_asym_id'] = auth_asym_id[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_auth_asym_id'] = auth_asym_id[ptnr2_indices]
    mmcif_dict['_struct_conn.ptnr1_auth_seq_id'] = auth_seq_id[ptnr1_indices]
    mmcif_dict['_struct_conn.ptnr2_auth_seq_id'] = auth_seq_id[ptnr2_indices]
    mmcif_dict['_struct_conn.pdbx_ptnr1_PDB_ins_code'] = insertion_code[
        ptnr1_indices
    ]
    mmcif_dict['_struct_conn.pdbx_ptnr2_PDB_ins_code'] = insertion_code[
        ptnr2_indices
    ]

    label_alt_id = ['?'] * self.size
    mmcif_dict['_struct_conn.pdbx_ptnr1_label_alt_id'] = label_alt_id
    mmcif_dict['_struct_conn.pdbx_ptnr2_label_alt_id'] = label_alt_id

    # We need to set this to make visualisation work in NGL/PyMOL.
    mmcif_dict['_struct_conn.pdbx_value_order'] = ['?'] * self.size

    # We use a symmetry of 1_555 which is the no-op transformation. Other
    # values are used when bonds involve atoms that only exist after expanding
    # the bioassembly, but we don't support this kind of bond at the moment.
    symmetry = ['1_555'] * self.size
    mmcif_dict['_struct_conn.ptnr1_symmetry'] = symmetry
    mmcif_dict['_struct_conn.ptnr2_symmetry'] = symmetry
    bond_type_counter = collections.Counter()
    for bond_row in self.iterrows():
      bond_type = bond_row['type']
      bond_type_counter[bond_type] += 1
      mmcif_dict['_struct_conn.id'].append(
          f'{bond_type}{bond_type_counter[bond_type]}'
      )
      mmcif_dict['_struct_conn.pdbx_role'].append(bond_row['role'])
      mmcif_dict['_struct_conn.conn_type_id'].append(bond_type)

    bond_types = np.unique(self.type)
    mmcif_dict['_struct_conn_type.id'] = bond_types
    unknown = ['?'] * len(bond_types)
    mmcif_dict['_struct_conn_type.criteria'] = unknown
    mmcif_dict['_struct_conn_type.reference'] = unknown

    return dict(mmcif_dict)


def concat_with_atom_keys(
    bonds_tables: Sequence[Bonds | None],
    atom_key_arrays: Sequence[np.ndarray],
) -> tuple[Bonds | None, np.ndarray]:
  """Concatenates bonds tables and atom keys simultaneously.

  Args:
    bonds_tables: A sequence of `Bonds` instances to concatenate. If any are
      None then these are skipped.
    atom_key_arrays: A sequence of integer `atom_key` arrays, where the n-th
      bonds_table referrs to the atoms in the n-th atom_key array. These must
      all be non-None.

  Returns:
    A pair of (bonds, atom_key) where atom_key is a unique atom_key array with
    length equal to the sum of the input atom array sizes, and the bonds table
    contains all the bonds from the individual bonds table inputs.
  """
  if not bonds_tables or not atom_key_arrays:
    if bonds_tables or atom_key_arrays:
      raise ValueError(
          'bonds_tables and atom_keys must have same length but got'
          f' {len(bonds_tables)=} and {len(atom_key_arrays)=}'
      )
    return None, np.array([], dtype=np.int64)
  max_key = -1
  atom_keys_to_concat = []
  types_to_concat = []
  roles_to_concat = []
  from_atom_keys_to_concat = []
  dest_atom_keys_to_concat = []
  for bonds, atom_key in zip(bonds_tables, atom_key_arrays, strict=True):
    if not atom_key.size:
      assert bonds is None or bonds.size == 0
      continue
    assert np.min(atom_key, initial=0) >= 0  # Should always be non-negative!
    offset = max_key + 1
    offset_atom_key = atom_key + offset
    atom_keys_to_concat.append(offset_atom_key)
    max_key = np.max(offset_atom_key)
    if bonds is not None:
      types_to_concat.append(bonds.type)
      roles_to_concat.append(bonds.role)
      from_atom_keys_to_concat.append(bonds.from_atom_key + offset)
      dest_atom_keys_to_concat.append(bonds.dest_atom_key + offset)

  if atom_keys_to_concat:
    concatted_atom_keys = np.concatenate(atom_keys_to_concat, axis=0)
  else:
    concatted_atom_keys = np.array([], dtype=np.int64)

  if types_to_concat:
    assert (
        len(types_to_concat)
        == len(roles_to_concat)
        == len(from_atom_keys_to_concat)
        == len(dest_atom_keys_to_concat)
    )
    num_bonds = sum(b.size for b in bonds_tables if b is not None)
    concatted_bonds = Bonds(
        key=np.arange(num_bonds, dtype=np.int64),
        type=np.concatenate(types_to_concat, axis=0),
        role=np.concatenate(roles_to_concat, axis=0),
        from_atom_key=np.concatenate(from_atom_keys_to_concat, axis=0),
        dest_atom_key=np.concatenate(dest_atom_keys_to_concat, axis=0),
    )
  else:
    concatted_bonds = None

  return concatted_bonds, concatted_atom_keys