flax_model/alphafold3/structure/mmcif.py

"""Low level mmCIF parsing operations and wrappers for nicer C++/Py errors.

Note that the cif_dict.CifDict class has many useful methods to help with data
extraction which are not shown in this file. You can find them in cif_dict.clif
together with docstrings. The cif_dict.CifDict class behaves like an immutable
Python dictionary (some methods are not implemented though).
"""
from collections.abc import Callable, Mapping, Sequence
import functools
import itertools
import re
from typing import ParamSpec, TypeAlias, TypeVar

from flax_model.alphafold3.constants import chemical_components
from flax_model.alphafold3.cpp import cif_dict
from flax_model.alphafold3.cpp import mmcif_atom_site
from flax_model.alphafold3.cpp import mmcif_struct_conn
from flax_model.alphafold3.cpp import string_array
import numpy as np

Mmcif = cif_dict.CifDict


_P = ParamSpec('_P')
_T = TypeVar('_T')
_WappedFn: TypeAlias = Callable[_P, _T]


@functools.lru_cache(maxsize=256)
def int_id_to_str_id(num: int) -> str:
  """Encodes a number as a string, using reverse spreadsheet style naming.

  Args:
    num: A positive integer.

  Returns:
    A string that encodes the positive integer using reverse spreadsheet style,
    naming e.g. 1 = A, 2 = B, ..., 27 = AA, 28 = BA, 29 = CA, ... This is the
    usual way to encode chain IDs in mmCIF files.
  """
  if num <= 0:
    raise ValueError(f'Only positive integers allowed, got {num}.')

  num = num - 1  # 1-based indexing.
  output = []
  while num >= 0:
    output.append(chr(num % 26 + ord('A')))
    num = num // 26 - 1
  return ''.join(output)


@functools.lru_cache(maxsize=256)
def str_id_to_int_id(str_id: str) -> int:
  """Encodes an mmCIF-style string chain ID as an integer.

  The integer IDs are one based so this function is the inverse of
  int_id_to_str_id.

  Args:
    str_id: A string chain ID consisting only of upper case letters A-Z.

  Returns:
    An integer that can be used to order mmCIF chain IDs in the standard
    (reverse spreadsheet style) ordering.
  """
  if not re.match('^[A-Z]+$', str_id):
    raise ValueError(f'String ID must be upper case letters, got {str_id}.')

  offset = ord('A') - 1
  output = 0
  for i, c in enumerate(str_id):
    output += (ord(c) - offset) * int(26**i)
  return output


def from_string(mmcif_string: str | bytes) -> Mmcif:
  return cif_dict.from_string(mmcif_string)


def parse_multi_data_cif(cif_string: str) -> dict[str, Mmcif]:
  """Parses a CIF string with multiple data records.

  For instance, the CIF string:

  ```
  data_001
  _foo bar
  #
  data_002
  _foo baz
  ```

  is parsed as:

  ```
  {'001': Mmcif({'_foo': ['bar']}), '002': Mmcif({'_foo': ['baz']})}
  ```

  Args:
    cif_string: The multi-data CIF string to be parsed.

  Returns:
    A dictionary mapping record names to Mmcif objects with data.
  """
  return cif_dict.parse_multi_data_cif(cif_string)


# def tokenize(mmcif_string: str) -> list[str]:
#   return cif_dict.tokenize(mmcif_string)


# def split_line(line: str) -> list[str]:
#   return cif_dict.split_line(line)


class BondParsingError(Exception):
  """Exception raised by errors when getting bond atom indices."""


def get_bond_atom_indices(
    mmcif: Mmcif,
    model_id: str = '1',
) -> tuple[Sequence[int], Sequence[int]]:
  """Extracts the indices of the atoms that participate in bonds.

  Args:
    mmcif: The mmCIF object to process.
    model_id: The ID of the model that the returned atoms will belong to. This
      should be a value in the mmCIF's _atom_site.pdbx_PDB_model_num column.

  Returns:
    Two lists of atom indices, `from_atoms` and `to_atoms`, each one having
    length num_bonds (as defined by _struct_conn, the bonds table). The bond
    i, defined by the i'th row in _struct_conn, is a bond from atom at index
    from_atoms[i], to the atom at index to_atoms[i]. The indices are simple
    0-based indexes into the columns of the _atom_site table in the input
    mmCIF, and do not necessarily correspond to the values in _atom_site.id,
    or any other column.

  Raises:
    BondParsingError: If any of the required tables or columns are not present
    in
      the mmCIF, or if the _struct_conn table refers to atoms that cannot
      be found in the _atom_site table.
  """
  try:
    return mmcif_struct_conn.get_bond_atom_indices(mmcif, model_id)
  except ValueError as e:
    raise BondParsingError(str(e)) from e


def get_or_infer_type_symbol(
    mmcif: Mmcif, ccd: chemical_components.Ccd | None = None
) -> Sequence[str]:
  """Returns the type symbol (element) for all of the atoms.

  Args:
    mmcif: A parsed mmCIF file in the Mmcif format.
    ccd: The chemical component dictionary. If not provided, defaults to the
      cached CCD.

  If present, returns the _atom_site.type_symbol. If not, infers it using
  _atom_site.label_comp_id (residue name), _atom_site.label_atom_id (atom name)
  and the CCD.
  """
  ccd = ccd or chemical_components.Ccd()
  type_symbol_fn = lambda res_name, atom_name: chemical_components.type_symbol(
      ccd, res_name, atom_name
  )
  return mmcif_atom_site.get_or_infer_type_symbol(mmcif, type_symbol_fn)


def get_chain_type_by_entity_id(mmcif: Mmcif) -> Mapping[str, str]:
  """Returns mapping from entity ID to its type or polymer type if available.

  If the entity is in the _entity_poly table, returns its polymer chain type.
  If not, returns the type as specified in the _entity table.

  Args:
    mmcif: CifDict holding the mmCIF.
  """
  poly_entity_id = mmcif.get('_entity_poly.entity_id', [])
  poly_type = mmcif.get('_entity_poly.type', [])
  poly_type_by_entity_id = dict(zip(poly_entity_id, poly_type, strict=True))

  chain_type_by_entity_id = {}
  for entity_id, entity_type in zip(
      mmcif.get('_entity.id', []), mmcif.get('_entity.type', []), strict=True
  ):
    chain_type = poly_type_by_entity_id.get(entity_id) or entity_type
    chain_type_by_entity_id[entity_id] = chain_type

  return chain_type_by_entity_id


def get_internal_to_author_chain_id_map(mmcif: Mmcif) -> Mapping[str, str]:
  """Returns a mapping from internal chain ID to the author chain ID.

  Note that this is not a bijection. One author chain ID can map to multiple
  internal chain IDs. For example, a protein chain and a ligand bound to it will
  share the same author chain ID, but they will each have a unique internal
  chain ID).

  Args:
    mmcif: CifDict holding the mmCIF.
  """
  return mmcif_atom_site.get_internal_to_author_chain_id_map(mmcif)


def get_experimental_method(mmcif: Mmcif) -> str | None:
  field = '_exptl.method'
  return ','.join(mmcif[field]).lower() if field in mmcif else None


def get_release_date(mmcif: Mmcif) -> str | None:
  """Returns the oldest revision date."""
  if '_pdbx_audit_revision_history.revision_date' not in mmcif:
    return None

  # Release dates are ISO-8601, hence sort well.
  return min(mmcif['_pdbx_audit_revision_history.revision_date'])


def get_resolution(mmcif: Mmcif) -> float | None:
  """Returns the resolution of the structure.

  More than one resolution can be reported in an mmCIF. This function returns
  the first one (in the order _refine.ls_d_res_high,
  _em_3d_reconstruction.resolution, _reflns.d_resolution_high) that appears
  in the mmCIF as is parseable as a float.

  Args:
    mmcif: An `Mmcif` object.

  Returns:
    The resolution as reported in the mmCIF.
  """
  for res_key in ('_refine.ls_d_res_high',
                  '_em_3d_reconstruction.resolution',
                  '_reflns.d_resolution_high'):
    if res_key in mmcif:
      try:
        raw_resolution = mmcif[res_key][0]
        return float(raw_resolution)
      except ValueError:
        continue
  return None


def parse_oper_expr(oper_expression: str) -> list[tuple[str, ...]]:
  """Determines which transforms to apply based on an MMCIF oper_expression str.

  Args:
    oper_expression: the field oper_expression from MMCIF format data.
      Transform ids may be either numbers or single letters. Hyphens are used to
      denote a numeric range of transforms to apply, and commas are used to
      delimit a sequence of transforms. Where two sets of parentheses are
      adjacent without a comma, the two sets of transforms should be combined as
      a cartesian product, i.e. all possible pairs.
      example 1,2,3 -> generate 3 copies of each chain by applying 1, 2 or 3.
      example (1-3) -> generate 3 copies of each chain by applying 1, 2 or 3.
      example (1-3)(4-6) -> generate 9 copies of each chain by applying one of
        [(1,4), (1,5), (1,6),
         (2,4), (2,5), (2,6),
         (3,4), (3,5), (3,6)]
      example (P) -> apply transform with id P.

  Raises:
    ValueError: Failure to parse oper_expression.

  Returns:
    A list with one element for each chain copy that should be generated.
    Each element is a list of transform ids to apply.
  """
  # Expand ranges, e.g. 1-4 -> 1,2,3,4.
  def range_expander(match):
    return ','.join(
        [str(i) for i in range(int(match.group(1)),
                               int(match.group(2)) + 1)])

  ranges_expanded = re.sub(r'\b(\d+)-(\d+)', range_expander, oper_expression)

  if re.fullmatch(r'(\w+,)*\w+', ranges_expanded):
    # No brackets, just a single range, e.g. "1,2,3".
    return [(t,) for t in ranges_expanded.split(',')]
  elif re.fullmatch(r'\((\w+,)*\w+\)', ranges_expanded):
    # Single range in brackets, e.g. "(1,2,3)".
    return [(t,) for t in ranges_expanded[1:-1].split(',')]
  elif re.fullmatch(r'\((\w+,)*\w+\)\((\w+,)*\w+\)', ranges_expanded):
    # Cartesian product of two ranges, e.g. "(1,2,3)(4,5)".
    part1, part2 = ranges_expanded[1:-1].split(')(')
    return list(itertools.product(part1.split(','), part2.split(',')))
  else:
    raise ValueError(f'Unsupported oper_expression format: {oper_expression}')


def format_float_array(
    values: np.ndarray, num_decimal_places: int) -> Sequence[str]:
  """Converts 1D array to a list of strings with the given number of decimals.

  This function is faster than converting via Python list comprehension, e.g.:
  atoms_x = ['%.3f' % x for x in atoms_x]

  Args:
    values: A numpy array with values to convert. This array is casted to
      float32 before doing the conversion.
    num_decimal_places: The number of decimal points to keep, including trailing
      zeros. E.g. for 1.07 and num_decimal_places=1: 1.1,
      num_decimal_places=2: 1.07, num_decimal_places=3: 1.070.

  Returns:
    A list of formatted strings.
  """
  if values.ndim != 1:
    raise ValueError(f'The given array must be 1D, got {values.ndim}D')

  return string_array.format_float_array(
      values=values.astype(np.float32), num_decimal_places=num_decimal_places
  )