flax_model/alphafold3/structure/table.py

"""Table module for atom/residue/chain tables in Structure.

Tables are intended to be lightweight collections of columns, loosely based
on a pandas dataframe, for use in the Structure class.
"""

import abc
from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence
import dataclasses
import functools
import graphlib
import typing
from typing import Any, Protocol, Self, TypeAlias, TypeVar, overload

from flax_model.alphafold3.cpp import string_array
import numpy as np


TableEntry: TypeAlias = str | int | float | None
FilterPredicate: TypeAlias = (
    TableEntry
    | Iterable[Any]  # Workaround for b/326384670. Tighten once fixed.
    | Callable[[Any], bool]  # Workaround for b/326384670. Tighten once fixed.
    | Callable[[np.ndarray], bool]
)


class RowLookup(Protocol):

  def get_row_by_key(
      self,
      key: int,
      column_name_map: Mapping[str, str] | None = None,
  ) -> Mapping[str, Any]:
    ...


@dataclasses.dataclass(frozen=True, kw_only=True)
class Table:
  """Parent class for structure tables.

  A table is a collection of columns of equal length, where one column is the
  key. The key uniquely identifies each row in the table.

  A table can refer to other tables by including a foreign key column, whose
  values are key values from the other table's key column. These column can have
  arbitrary names and are treated like any other integer-valued column.

  See the `Database` class in this module for utilities for handing sets of
  tables that are related via foreign keys.

  NB: This does not correspond to an mmCIF table.
  """

  key: np.ndarray

  def __post_init__(self):
    for col_name in self.columns:
      if (col_len := self.get_column(col_name).shape[-1]) != self.size:
        raise ValueError(
            f'All columns should have length {self.size} but got "{col_name}"'
            f' with length {col_len}.'
        )
      self.get_column(col_name).flags.writeable = False  # Make col immutable.
    if self.key.size and self.key.min() < 0:
      raise ValueError(
          'Key values must be non-negative. Got negative values:'
          f' {set(self.key[self.key < 0])}'
      )
    self.key.flags.writeable = False  # Make key immutable.

  def __getstate__(self) -> dict[str, Any]:
    """Returns members with cached properties removed for pickling."""
    cached_props = {
        k
        for k, v in self.__class__.__dict__.items()
        if isinstance(v, functools.cached_property)
    }
    return {k: v for k, v in self.__dict__.items() if k not in cached_props}

  @functools.cached_property
  def index_by_key(self) -> np.ndarray:
    """Mapping from key values to their index in the column arrays.

    i.e.: self.key[index_by_key[k]] == k
    """
    if not self.key.size:
      return np.array([], dtype=np.int64)
    else:
      index_by_key = np.zeros(np.max(self.key) + 1, dtype=np.int64)
      index_by_key[self.key] = np.arange(self.size)
      return index_by_key

  @functools.cached_property
  def columns(self) -> tuple[str, ...]:
    """The names of the columns in the table, including the key column."""
    return tuple(field.name for field in dataclasses.fields(self))

  @functools.cached_property
  def items(self) -> Mapping[str, np.ndarray]:
    """Returns the mapping from column names to column values."""
    return {col: getattr(self, col) for col in self.columns}

  @functools.cached_property
  def size(self) -> int:
    """The number of rows in the table."""
    return self.key.shape[-1]

  def __len__(self) -> int:
    return self.size

  def get_column(self, column_name: str) -> np.ndarray:
    """Gets a column by name."""
    # Performance optimisation: use the cached columns, instead of getattr.
    return self.items[column_name]

  def apply_array(self, arr: np.ndarray) -> Self:
    """Returns a sliced table using a key (!= index) array or a boolean mask."""
    if arr.dtype == bool and np.all(arr):
      return self  # Shortcut: No-op, so just return.

    return self.copy_and_update(**{
        column_name: self.apply_array_to_column(column_name, arr)
        for column_name in self.columns
    })

  def apply_index(self, index_arr: np.ndarray) -> Self:
    """Returns a sliced table using an index (!= key) array."""
    if index_arr.dtype == bool:
      raise ValueError('The index array must not be a boolean mask.')

    return self.copy_and_update(
        **{col: self.get_column(col)[..., index_arr] for col in self.columns}
    )

  def apply_array_to_column(
      self,
      column_name: str,
      arr: np.ndarray,
  ) -> np.ndarray:
    """Returns a sliced column array using a key array or a boolean mask."""
    if arr.dtype == bool:
      return self.get_column(column_name)[..., arr]
    else:
      return self.get_column(column_name)[..., self.index_by_key[arr]]

  def get_value_by_index(self, column_name: str, index: int) -> Any:
    return self.get_column(column_name)[index]

  def get_value_by_key(
      self,
      column_name: str,
      key: int | np.integer,
  ) -> TableEntry:
    """Gets the value of a column at the row with specified key value."""
    return self.get_value_by_index(column_name, self.index_by_key[key])

  @overload
  def __getitem__(self, key: str) -> np.ndarray:
    ...

  @overload
  def __getitem__(self, key: np.ndarray) -> 'Table':
    ...

  @overload
  def __getitem__(self, key: tuple[str, int | np.integer]) -> TableEntry:
    ...

  @overload
  def __getitem__(self, key: tuple[str, np.ndarray]) -> np.ndarray:
    ...

  def __getitem__(self, key):
    match key:
      case str():
        return self.get_column(key)
      case np.ndarray() as key_arr_or_mask:
        return self.apply_array(key_arr_or_mask)
      case str() as col, int() | np.integer() as key_val:
        return self.get_value_by_key(col, key_val)
      case str() as col, np.ndarray() as key_arr_or_mask:
        return self.apply_array_to_column(col, key_arr_or_mask)
      case _:
        if isinstance(key, tuple):
          err_msg = f'{key}, type: tuple({[type(v) for v in key]})'
        else:
          err_msg = f'{key}, type: {type(key)}'
        raise KeyError(err_msg)

  def get_row_by_key(
      self,
      key: int,
      column_name_map: Mapping[str, str] | None = None,
  ) -> dict[str, Any]:
    """Gets the row with specified key value."""
    return self.get_row_by_index(
        self.index_by_key[key], column_name_map=column_name_map
    )

  def get_row_by_index(
      self,
      index: int,
      column_name_map: Mapping[str, str] | None = None,
  ) -> dict[str, Any]:
    """Gets the row at the specified index."""
    if column_name_map is not None:
      return {
          renamed_col: self.get_value_by_index(col, index)
          for renamed_col, col in column_name_map.items()
      }
    else:
      return {col: self.get_value_by_index(col, index) for col in self.columns}

  def iterrows(
      self,
      *,
      row_keys: np.ndarray | None = None,
      column_name_map: Mapping[str, str] | None = None,
      **table_by_foreign_key_col: RowLookup,
  ) -> Iterator[Mapping[str, Any]]:
    """Yields rows from the table.

    This can be used to easily convert a table to a Pandas dataframe:

    ```py
    df = pd.DataFrame(table.iterrows())
    ```

    Args:
      row_keys: An optional array of keys of rows to yield. If None, all rows
        will be yielded.
      column_name_map: An optional mapping from desired keys in the row dicts to
        the names of the columns they correspond to.
      **table_by_foreign_key_col: An optional mapping from column names in this
        table, which are expected to be columns of foreign keys, to the table
        that the foreign keys point into. If provided, then the yielded rows
        will include data from the foreign tables at the appropriate key.
    """
    if row_keys is not None:
      row_indices = self.index_by_key[row_keys]
    else:
      row_indices = range(self.size)
    for i in row_indices:
      row = self.get_row_by_index(i, column_name_map=column_name_map)
      for key_col, table in table_by_foreign_key_col.items():
        foreign_key = self[key_col][i]
        foreign_row = table.get_row_by_key(foreign_key)
        row.update(foreign_row)
      yield row

  def with_column_names(
      self, column_name_map: Mapping[str, str]
  ) -> 'RenamedTableView':
    """Returns a view of this table with mapped column names."""
    return RenamedTableView(self, column_name_map=column_name_map)

  def make_filter_mask(
      self,
      mask: np.ndarray | None = None,
      *,
      apply_per_element: bool = False,
      **predicate_by_col: FilterPredicate,
  ) -> np.ndarray | None:
    """Returns a boolean array of rows to keep, or None if all can be kept.

    Args:
      mask: See `Table.filter`.
      apply_per_element: See `Table.filter`.
      **predicate_by_col: See `Table.filter`.

    Returns:
      Either a boolean NumPy array of length `(self.size,)` denoting which rows
      should be kept according to the input mask and predicates, or None. None
      implies there is no filtering required, and is used where possible
      instead of an all-True array to save time and space.
    """
    if mask is None:
      if not predicate_by_col:
        return None
      else:
        mask = np.ones((self.size,), dtype=bool)
    else:
      if mask.shape != (self.size,):
        raise ValueError(
            f'mask must have shape ({self.size},). Got: {mask.shape}.'
        )
      if mask.dtype != bool:
        raise ValueError(f'mask must have dtype bool. Got: {mask.dtype}.')

    for col, predicate in predicate_by_col.items():
      if self[col].ndim > 1:
        raise ValueError(
            f'Cannot filter by column {col} with more than 1 dimension.'
        )

      callable_predicates = []
      if not callable(predicate):
        if isinstance(predicate, Iterable) and not isinstance(predicate, str):
          target_vals = predicate
        else:
          target_vals = [predicate]
        for target_val in target_vals:
          callable_predicates.append(lambda x, target=target_val: x == target)
      else:
        callable_predicates.append(predicate)

      field_mask = np.zeros_like(mask)
      for callable_predicate in callable_predicates:
        if not apply_per_element:
          callable_predicate = typing.cast(
              Callable[[np.ndarray], bool], callable_predicate
          )
          predicate_result = callable_predicate(self.get_column(col))
        else:
          predicate_result = np.array(
              [callable_predicate(elem) for elem in self.get_column(col)]
          )
        np.logical_or(field_mask, predicate_result, out=field_mask)
      np.logical_and(mask, field_mask, out=mask)  # Update in-place.
    return mask

  def filter(
      self,
      mask: np.ndarray | None = None,
      *,
      apply_per_element: bool = False,
      invert: bool = False,
      **predicate_by_col: FilterPredicate,
  ) -> Self:
    """Filters the table using mask and/or predicates and returns a new table.

    Predicates can be either:
      1. A constant value, e.g. `'CA'`. In this case then only rows that match
        this value for the given column are retained.
      2. A (non-string) iterable e.g. `('A', 'B')`. In this
        case then rows are retained if they match any of the provided values for
        the given column.
      3. A boolean function e.g. `lambda b_fac: b_fac < 100.0`.
        In this case then only rows that evaluate to `True` are retained. By
        default this function's parameter is expected to be an array, unless
        `apply_per_element=True`.

    Args:
      mask: An optional boolean NumPy array with length equal to the table size.
        If provided then this will be combined with the other predicates so that
        a row is included if it is masked-in *and* matches all the predicates.
      apply_per_element: Whether apply predicates to each element in the column
        individually, or to pass the whole column array to the predicate.
      invert: If True then the returned table will contain exactly those rows
        that would be removed if this was `False`.
      **predicate_by_col: A mapping from column name to a predicate. Filtered
        columns must be 1D arrays. If multiple columns are provided as keyword
        arguments then each predicate is applied and the results are combined
        using a boolean AND operation, so an atom is only retained if it passes
        all predicates.

    Returns:
      A new table with the desired rows retained (or filtered out if
      `invert=True`).

    Raises:
      ValueError: If mask is provided and is not a bool array with shape
        `(num_atoms,)`.
    """
    filter_mask = self.make_filter_mask(
        mask, apply_per_element=apply_per_element, **predicate_by_col
    )
    if filter_mask is None:
      # No mask or predicate was specified, so we can return early.
      if not invert:
        return self
      else:
        return self[np.array((), dtype=np.int64)]
    else:
      return self[~filter_mask if invert else filter_mask]

  def _validate_keys_are_column_names(self, keys: Collection[str]) -> None:
    """Raises an error if any of the keys are not column names."""
    if mismatches := set(keys) - set(self.columns):
      raise ValueError(f'Invalid column names: {sorted(mismatches)}.')

  def copy_and_update(self, **new_column_by_column_name: np.ndarray) -> Self:
    """Returns a copy of this table with the specified changes applied.

    Args:
      **new_column_by_column_name: New values for the specified columns.

    Raises:
      ValueError: If a specified column name is not a column in this table.
    """
    self._validate_keys_are_column_names(new_column_by_column_name)
    return dataclasses.replace(self, **new_column_by_column_name)

  def copy_and_remap(
      self, **mapping_by_col: Mapping[TableEntry, TableEntry]
  ) -> Self:
    """Returns a copy of the table with the specified columns remapped.

    Args:
      **mapping_by_col: Each kwarg key should be the name of one of this table's
        columns, and each value should be a mapping. The values in the column
        will be looked up in the mapping and replaced with the result if one is
        found.

    Raises:
      ValueError: If a specified column name is not a column in this table.
    """
    self._validate_keys_are_column_names(mapping_by_col)
    if not self.size:
      return self
    remapped_cols = {}
    for column_name, mapping in mapping_by_col.items():
      col_arr = self.get_column(column_name)
      if col_arr.dtype == object:
        remapped = string_array.remap(col_arr, mapping)
      else:
        remapped = np.vectorize(lambda x: mapping.get(x, x))(col_arr)  # pylint: disable=cell-var-from-loop
      remapped_cols[column_name] = remapped
    return self.copy_and_update(**remapped_cols)


class RenamedTableView:
  """View of a table with renamed column names."""

  def __init__(self, table: Table, column_name_map: Mapping[str, str]):
    self._table = table
    self._column_name_map = column_name_map

  def get_row_by_key(
      self,
      key: int,
      column_name_map: Mapping[str, str] | None = None,
  ) -> Mapping[str, Any]:
    del column_name_map
    return self._table.get_row_by_key(
        key, column_name_map=self._column_name_map
    )


_DatabaseT = TypeVar('_DatabaseT', bound='Database')


class Database(abc.ABC):
  """Relational database base class."""

  @property
  @abc.abstractmethod
  def tables(self) -> Collection[str]:
    """The names of the tables in this database."""

  @abc.abstractmethod
  def get_table(self, table_name: str) -> Table:
    """Gets the table with the given name."""

  @property
  @abc.abstractmethod
  def foreign_keys(self) -> Mapping[str, Collection[tuple[str, str]]]:
    """Describes the relationship between keys in the database.

    Returns:
      A map from table names to pairs of `(column_name, foreign_table_name)`
      where `column_name` is a column containing foreign keys in the table named
      by the key, and the `foreign_table_name` is the name of the table that
      those foreign keys refer to.
    """

  @abc.abstractmethod
  def copy_and_update(
      self: _DatabaseT,
      **new_field_by_field_name: ...,
  ) -> _DatabaseT:
    """Returns a copy of this database with the specified changes applied."""


def table_dependency_order(db: Database) -> Iterable[str]:
  """Yields the names of the tables in the database in dependency order.

  This order guarantees that a table appears after all other tables that
  it refers to using foreign keys. Specifically A < B implies that A contains
  no column that refers to B.key as a foreign key.

  Args:
    db: The database that defines the table names and foreign keys.
  """
  connections: dict[str, set[str]] = {}
  for table_name in db.tables:
    connection_set = set()
    for _, foreign_table in db.foreign_keys.get(table_name, ()):
      connection_set.add(foreign_table)
    connections[table_name] = connection_set
  yield from graphlib.TopologicalSorter(connections).static_order()


def concat_databases(dbs: Sequence[_DatabaseT]) -> _DatabaseT:
  """Concatenates the tables across a sequence of databases.

  Args:
    dbs: A non-empty sequence of database instances of the same type.

  Returns:
    A new database containing the concatenated tables from the input databases.

  Raises:
    ValueError: If `dbs` is empty or `dbs` contains different Database
      types.
  """
  if not dbs:
    raise ValueError('Need at least one value to concatenate.')
  distinct_db_types = {type(db) for db in dbs}
  if len(distinct_db_types) > 1:
    raise ValueError(
        f'All `dbs` must be of the same type, got: {distinct_db_types}'
    )

  first_db, *other_dbs = dbs
  concatted_tables: dict[str, Table] = {}
  key_offsets: dict[str, list[int]] = {}
  for table_name in table_dependency_order(first_db):
    first_table = first_db.get_table(table_name)
    columns: dict[str, list[np.ndarray]] = {
        column_name: [first_table.get_column(column_name)]
        for column_name in first_table.columns
    }
    key_offsets[table_name] = [
        first_table.key.max() + 1 if first_table.size else 0
    ]

    for prev_index, db in enumerate(other_dbs):
      table = db.get_table(table_name)
      for col_name in table.columns:
        columns[col_name].append(table.get_column(col_name))
      key_offset = key_offsets[table_name][prev_index]
      offset_key = table.key + key_offset
      columns['key'][-1] = offset_key
      if table.size:
        key_offsets[table_name].append(offset_key.max() + 1)
      else:
        key_offsets[table_name].append(key_offsets[table_name][prev_index])
      for fkey_col_name, foreign_table_name in first_db.foreign_keys.get(
          table_name, []
      ):
        fkey_columns = columns[fkey_col_name]
        fkey_columns[-1] = (
            fkey_columns[-1] + key_offsets[foreign_table_name][prev_index]
        )

    concatted_columns = {
        column_name: np.concatenate(values, axis=-1)
        for column_name, values in columns.items()
    }
    concatted_tables[table_name] = (type(first_table))(**concatted_columns)
  return first_db.copy_and_update(**concatted_tables)