flax_model/alphafold3/model/features.py

"""Data-side of the input features processing."""

import dataclasses
import datetime
import itertools
from typing import Any, Self, TypeAlias

from absl import logging
from flax_model.alphafold3 import structure
from flax_model.alphafold3.common import folding_input
from flax_model.alphafold3.constants import chemical_components
from flax_model.alphafold3.constants import mmcif_names
from flax_model.alphafold3.constants import periodic_table
from flax_model.alphafold3.constants import residue_names
from flax_model.alphafold3.cpp import cif_dict
from flax_model.alphafold3.data import msa as msa_module
from flax_model.alphafold3.data import templates
from flax_model.alphafold3.data.tools import rdkit_utils
from flax_model.alphafold3.model import data3
from flax_model.alphafold3.model import data_constants
from flax_model.alphafold3.model import merging_features
from flax_model.alphafold3.model import msa_pairing
from flax_model.alphafold3.model.atom_layout import atom_layout
from flax_model.alphafold3.structure import chemical_components as struc_chem_comps
import jax
import jax.numpy as jnp
import numpy as np
from rdkit import Chem


xnp_ndarray: TypeAlias = np.ndarray | jnp.ndarray  # pylint: disable=invalid-name
BatchDict: TypeAlias = dict[str, xnp_ndarray]

_STANDARD_RESIDUES = frozenset({
    *residue_names.PROTEIN_TYPES_WITH_UNKNOWN,
    *residue_names.NUCLEIC_TYPES_WITH_2_UNKS,
})


@dataclasses.dataclass(frozen=True)
class PaddingShapes:
  num_tokens: int
  msa_size: int
  num_chains: int
  num_templates: int
  num_atoms: int


def _pad_to(
    arr: np.ndarray, shape: tuple[int | None, ...], **kwargs
) -> np.ndarray:
  """Pads an array to a given shape. Wrapper around np.pad().

  Args:
    arr: numpy array to pad
    shape: target shape, use None for axes that should stay the same
    **kwargs: additional args for np.pad, e.g. constant_values=-1

  Returns:
    the padded array

  Raises:
    ValueError if arr and shape have a different number of axes.
  """
  if arr.ndim != len(shape):
    raise ValueError(
        f'arr and shape have different number of axes. {arr.shape=}, {shape=}'
    )

  num_pad = []
  for axis, width in enumerate(shape):
    if width is None:
      num_pad.append((0, 0))
    else:
      if width >= arr.shape[axis]:
        num_pad.append((0, width - arr.shape[axis]))
      else:
        raise ValueError(
            f'Can not pad to a smaller shape. {arr.shape=}, {shape=}'
        )
  padded_arr = np.pad(arr, pad_width=num_pad, **kwargs)
  return padded_arr


def _unwrap(obj):
  """Unwrap an object from a zero-dim np.ndarray."""
  if isinstance(obj, np.ndarray) and obj.ndim == 0:
    return obj.item()
  else:
    return obj


@dataclasses.dataclass(frozen=True)
class Chains:
  chain_id: np.ndarray
  asym_id: np.ndarray
  entity_id: np.ndarray
  sym_id: np.ndarray


jax.tree_util.register_dataclass(
    Chains,
    data_fields=[f.name for f in dataclasses.fields(Chains)],
    meta_fields=[],
)


def _compute_asym_entity_and_sym_id(
    all_tokens: atom_layout.AtomLayout,
) -> Chains:
  """Compute asym_id, entity_id and sym_id.

  Args:
    all_tokens: atom layout containing a representative atom for each token.

  Returns:
    A Chains object
  """

  # Find identical sequences and assign entity_id and sym_id to every chain.
  seq_to_entity_id_sym_id = {}
  seen_chain_ids = set()
  chain_ids = []
  asym_ids = []
  entity_ids = []
  sym_ids = []
  for chain_id in all_tokens.chain_id:
    if chain_id not in seen_chain_ids:
      asym_id = len(seen_chain_ids) + 1
      seen_chain_ids.add(chain_id)
      seq = ','.join(all_tokens.res_name[all_tokens.chain_id == chain_id])
      if seq not in seq_to_entity_id_sym_id:
        entity_id = len(seq_to_entity_id_sym_id) + 1
        sym_id = 1
      else:
        entity_id, sym_id = seq_to_entity_id_sym_id[seq]
        sym_id += 1
      seq_to_entity_id_sym_id[seq] = (entity_id, sym_id)

      chain_ids.append(chain_id)
      asym_ids.append(asym_id)
      entity_ids.append(entity_id)
      sym_ids.append(sym_id)

  return Chains(
      chain_id=np.array(chain_ids),
      asym_id=np.array(asym_ids),
      entity_id=np.array(entity_ids),
      sym_id=np.array(sym_ids),
  )


def tokenizer(
    flat_output_layout: atom_layout.AtomLayout,
    ccd: chemical_components.Ccd,
    max_atoms_per_token: int,
    flatten_non_standard_residues: bool,
    logging_name: str,
) -> tuple[atom_layout.AtomLayout, atom_layout.AtomLayout, np.ndarray]:
  """Maps a flat atom layout to tokens for evoformer.

  Creates the evoformer tokens as one token per polymer residue and one token
  per ligand atom. The tokens are represented as AtomLayouts all_tokens
  (1 representative atom per token) atoms per residue, and
  all_token_atoms_layout (num_tokens, max_atoms_per_token). The atoms in a
  residue token use the layout of the corresponding CCD entry

  Args:
    flat_output_layout: flat AtomLayout containing all atoms that the model
      wants to predict.
    ccd: The chemical components dictionary.
    max_atoms_per_token: number of slots per token.
    flatten_non_standard_residues: whether to flatten non-standard residues,
      i.e. whether to use one token per atom for non-standard residues.
    logging_name: logging name for debugging (usually the mmcif_id).

  Returns:
    A tuple (all_tokens, all_tokens_atoms_layout) with
      all_tokens: AtomLayout shape (num_tokens,) containing one representative
        atom per token.
      all_token_atoms_layout: AtomLayout with shape
        (num_tokens, max_atoms_per_token) containing all atoms per token.
      standard_token_idxs: The token index that each token would have if not
        flattening non standard resiudes.
  """
  # Select  the representative atom for each token.
  token_idxs = []
  single_atom_token = []
  standard_token_idxs = []
  current_standard_token_id = 0
  # Iterate over residues, and provide a group_iter over the atoms of each
  # residue.
  for key, group_iter in itertools.groupby(
      zip(
          flat_output_layout.chain_type,
          flat_output_layout.chain_id,
          flat_output_layout.res_id,
          flat_output_layout.res_name,
          flat_output_layout.atom_name,
          np.arange(flat_output_layout.shape[0]),
      ),
      key=lambda x: x[:3],
  ):

    # Get chain type and chain id of this residue
    chain_type, chain_id, _ = key

    # Get names and global idxs for all atoms of this residue
    _, _, _, res_names, atom_names, idxs = zip(*group_iter)

    # As of March 2023, all OTHER CHAINs in pdb are artificial nucleics.
    is_nucleic_backbone = (
        chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES
        or chain_type == mmcif_names.OTHER_CHAIN
    )
    if chain_type in mmcif_names.PEPTIDE_CHAIN_TYPES:
      res_name = res_names[0]
      if (
          flatten_non_standard_residues
          and res_name not in residue_names.PROTEIN_TYPES_WITH_UNKNOWN
          and res_name != residue_names.MSE
      ):
        # For non-standard protein residues take all atoms.
        # NOTE: This may get very large if we include hydrogens.
        token_idxs.extend(idxs)
        single_atom_token += [True] * len(idxs)
        standard_token_idxs.extend([current_standard_token_id] * len(idxs))
      else:
        # For standard protein residues take 'CA' if it exists, else first atom.
        if 'CA' in atom_names:
          token_idxs.append(idxs[atom_names.index('CA')])
        else:
          token_idxs.append(idxs[0])
        single_atom_token += [False]
        standard_token_idxs.append(current_standard_token_id)
      current_standard_token_id += 1
    elif is_nucleic_backbone:
      res_name = res_names[0]
      if (
          flatten_non_standard_residues
          and res_name not in residue_names.NUCLEIC_TYPES_WITH_2_UNKS
      ):
        # For non-standard nucleic residues take all atoms.
        token_idxs.extend(idxs)
        single_atom_token += [True] * len(idxs)
        standard_token_idxs.extend([current_standard_token_id] * len(idxs))
      else:
        # For standard nucleic residues take C1' if it exists, else first atom.
        if "C1'" in atom_names:
          token_idxs.append(idxs[atom_names.index("C1'")])
        else:
          token_idxs.append(idxs[0])
        single_atom_token += [False]
        standard_token_idxs.append(current_standard_token_id)
      current_standard_token_id += 1
    elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
      # For non-polymers take all atoms
      token_idxs.extend(idxs)
      single_atom_token += [True] * len(idxs)
      standard_token_idxs.extend([current_standard_token_id] * len(idxs))
      current_standard_token_id += len(idxs)
    else:
      # Chain type that we don't handle yet.
      logging.warning(
          '%s: ignoring chain %s with chain type %s.',
          logging_name,
          chain_id,
          chain_type,
      )

  assert len(token_idxs) == len(single_atom_token)
  assert len(token_idxs) == len(standard_token_idxs)
  standard_token_idxs = np.array(standard_token_idxs, dtype=np.int32)

  # Create the list of all tokens, represented as a flat AtomLayout with 1
  # representative atom per token.
  all_tokens = flat_output_layout[token_idxs]

  # Create the 2D atoms_per_token layout
  num_tokens = all_tokens.shape[0]

  # Target lists.
  target_atom_names = []
  target_atom_elements = []
  target_res_ids = []
  target_res_names = []
  target_chain_ids = []
  target_chain_types = []

  # uids of all atoms in the flat layout, to check whether the dense atoms
  # exist -- This is necessary for terminal atoms (e.g. 'OP3' or 'OXT')
  all_atoms_uids = set(
      zip(
          flat_output_layout.chain_id,
          flat_output_layout.res_id,
          flat_output_layout.atom_name,
      )
  )

  for idx, single_atom in enumerate(single_atom_token):
    if not single_atom:
      # Standard protein and nucleic residues have many atoms per token
      chain_id = all_tokens.chain_id[idx]
      res_id = all_tokens.res_id[idx]
      res_name = all_tokens.res_name[idx]
      atom_names = []
      atom_elements = []

      res_atoms = struc_chem_comps.get_all_atoms_in_entry(
          ccd=ccd, res_name=res_name
      )
      atom_names_elements = list(
          zip(
              res_atoms['_chem_comp_atom.atom_id'],
              res_atoms['_chem_comp_atom.type_symbol'],
              strict=True,
          )
      )

      for atom_name, atom_element in atom_names_elements:
        # Remove hydrogens if they are not in flat layout.
        if atom_element in ['H', 'D'] and (
            (chain_id, res_id, atom_name) not in all_atoms_uids
        ):
          continue
        elif (chain_id, res_id, atom_name) in all_atoms_uids:
          atom_names.append(atom_name)
          atom_elements.append(atom_element)
        # Leave spaces for OXT etc.
        else:
          atom_names.append('')
          atom_elements.append('')

      if len(atom_names) > max_atoms_per_token:
        logging.warning(
            'Atom list for chain %s '
            'residue %s %s is too long and will be truncated: '
            '%s to the max atoms limit %s. Dropped atoms: %s',
            chain_id,
            res_id,
            res_name,
            len(atom_names),
            max_atoms_per_token,
            list(
                zip(
                    atom_names[max_atoms_per_token:],
                    atom_elements[max_atoms_per_token:],
                    strict=True,
                )
            ),
        )
        atom_names = atom_names[:max_atoms_per_token]
        atom_elements = atom_elements[:max_atoms_per_token]

      num_pad = max_atoms_per_token - len(atom_names)
      atom_names.extend([''] * num_pad)
      atom_elements.extend([''] * num_pad)

    else:
      # ligands have only 1 atom per token
      padding = [''] * (max_atoms_per_token - 1)
      atom_names = [all_tokens.atom_name[idx]] + padding
      atom_elements = [all_tokens.atom_element[idx]] + padding

    # Append the atoms to the target lists.
    target_atom_names.append(atom_names)
    target_atom_elements.append(atom_elements)
    target_res_names.append([all_tokens.res_name[idx]] * max_atoms_per_token)
    target_res_ids.append([all_tokens.res_id[idx]] * max_atoms_per_token)
    target_chain_ids.append([all_tokens.chain_id[idx]] * max_atoms_per_token)
    target_chain_types.append(
        [all_tokens.chain_type[idx]] * max_atoms_per_token
    )

  # Make sure to get the right shape also for 0 tokens
  trg_shape = (num_tokens, max_atoms_per_token)
  all_token_atoms_layout = atom_layout.AtomLayout(
      atom_name=np.array(target_atom_names, dtype=object).reshape(trg_shape),
      atom_element=np.array(target_atom_elements, dtype=object).reshape(
          trg_shape
      ),
      res_name=np.array(target_res_names, dtype=object).reshape(trg_shape),
      res_id=np.array(target_res_ids, dtype=int).reshape(trg_shape),
      chain_id=np.array(target_chain_ids, dtype=object).reshape(trg_shape),
      chain_type=np.array(target_chain_types, dtype=object).reshape(trg_shape),
  )

  return all_tokens, all_token_atoms_layout, standard_token_idxs


@dataclasses.dataclass(frozen=True)
class MSA:
  """Dataclass containing MSA."""

  rows: xnp_ndarray
  mask: xnp_ndarray
  deletion_matrix: xnp_ndarray
  # Occurrence of each residue type along the sequence, averaged over MSA rows.
  profile: xnp_ndarray
  # Occurrence of deletions along the sequence, averaged over MSA rows.
  deletion_mean: xnp_ndarray
  # Number of MSA alignments.
  num_alignments: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      *,
      all_tokens: atom_layout.AtomLayout,
      standard_token_idxs: np.ndarray,
      padding_shapes: PaddingShapes,
      fold_input: folding_input.Input,
      logging_name: str,
      max_paired_sequence_per_species: int,
      resolve_msa_overlaps: bool = True,
  ) -> Self:
    """Compute the msa features."""
    seen_entities = {}

    substruct = atom_layout.make_structure(
        flat_layout=all_tokens,
        atom_coords=np.zeros(all_tokens.shape + (3,)),
        name=logging_name,
    )
    prot = substruct.filter_to_entity_type(protein=True)
    num_unique_chains = len(set(prot.chain_single_letter_sequence().values()))
    need_msa_pairing = num_unique_chains > 1

    np_chains_list = []
    input_chains_by_id = {chain.id: chain for chain in fold_input.chains}
    nonempty_chain_ids = set(all_tokens.chain_id)
    for asym_id, chain_info in enumerate(substruct.iter_chains(), start=1):
      b_chain_id = chain_info['chain_id']
      chain_type = chain_info['chain_type']
      chain = input_chains_by_id[b_chain_id]

      # Generalised "sequence" for ligands (can't trust residue name)
      chain_tokens = all_tokens[all_tokens.chain_id == b_chain_id]
      assert chain_tokens.res_name is not None
      three_letter_sequence = ','.join(chain_tokens.res_name.tolist())
      chain_num_tokens = len(chain_tokens.atom_name)
      if chain_type in mmcif_names.POLYMER_CHAIN_TYPES:
        sequence = substruct.chain_single_letter_sequence()[b_chain_id]
        if chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES:
          # Only allow nucleic residue types for nucleic chains (can have some
          # protein residues in e.g. tRNA, but that causes MSA search failures).
          # Replace non nucleic residue types by UNK_NUCLEIC.
          nucleic_types_one_letter = (
              residue_names.DNA_TYPES_ONE_LETTER
              + residue_names.RNA_TYPES_ONE_LETTER_WITH_UNKNOWN
          )
          sequence = ''.join([
              base
              if base in nucleic_types_one_letter
              else residue_names.UNK_NUCLEIC_ONE_LETTER
              for base in sequence
          ])
      else:
        sequence = 'X' * chain_num_tokens

      skip_chain = (
          chain_type not in mmcif_names.STANDARD_POLYMER_CHAIN_TYPES
          or len(sequence) <= 4
          or b_chain_id not in nonempty_chain_ids
      )
      if three_letter_sequence in seen_entities:
        entity_id = seen_entities[three_letter_sequence]
      else:
        entity_id = len(seen_entities) + 1

      if chain_type in mmcif_names.STANDARD_POLYMER_CHAIN_TYPES:
        unpaired_a3m = ''
        paired_a3m = ''
        if not skip_chain:
          if need_msa_pairing and isinstance(chain, folding_input.ProteinChain):
            paired_a3m = chain.paired_msa
          if isinstance(
              chain, folding_input.RnaChain | folding_input.ProteinChain
          ):
            unpaired_a3m = chain.unpaired_msa
        # If we generated the MSA ourselves, it is already deduplicated. If it
        # is user-provided, keep it as is to prevent destroying desired pairing.
        unpaired_msa = msa_module.Msa.from_a3m(
            query_sequence=sequence,
            chain_poly_type=chain_type,
            a3m=unpaired_a3m,
            deduplicate=False,
        )

        paired_msa = msa_module.Msa.from_a3m(
            query_sequence=sequence,
            chain_poly_type=chain_type,
            a3m=paired_a3m,
            deduplicate=False,
        )
      else:
        unpaired_msa = msa_module.Msa.from_empty(
            query_sequence='-' * len(sequence),
            chain_poly_type=mmcif_names.PROTEIN_CHAIN,
        )
        paired_msa = msa_module.Msa.from_empty(
            query_sequence='-' * len(sequence),
            chain_poly_type=mmcif_names.PROTEIN_CHAIN,
        )

      msa_features = unpaired_msa.featurize()
      all_seqs_msa_features = paired_msa.featurize()

      # msa_features = data3.fix_features(msa_features)
      # all_seqs_msa_features = data3.fix_features(all_seqs_msa_features)

      msa_features = msa_features | {
          f'{k}_all_seq': v for k, v in all_seqs_msa_features.items()
      }
      feats = msa_features
      feats['chain_id'] = b_chain_id
      feats['asym_id'] = np.full(chain_num_tokens, asym_id)
      feats['entity_id'] = entity_id
      np_chains_list.append(feats)

    # Add profile features to each chain.
    for chain in np_chains_list:
      chain.update(
          data3.get_profile_features(chain['msa'], chain['deletion_matrix'])
      )

    # Allow 50% of the MSA to come from MSA pairing.
    max_paired_sequences = padding_shapes.msa_size // 2
    if need_msa_pairing:
      np_chains_list = list(map(dict, np_chains_list))
      np_chains_list = msa_pairing.create_paired_features(
          np_chains_list,
          max_paired_sequences=max_paired_sequences,
          nonempty_chain_ids=nonempty_chain_ids,
          max_hits_per_species=max_paired_sequence_per_species,
      )
      if resolve_msa_overlaps:
        np_chains_list = msa_pairing.deduplicate_unpaired_sequences(
            np_chains_list
        )

    # Remove all gapped rows from all seqs.
    nonempty_asym_ids = []
    for chain in np_chains_list:
      if chain['chain_id'] in nonempty_chain_ids:
        nonempty_asym_ids.append(chain['asym_id'][0])
    if 'msa_all_seq' in np_chains_list[0]:
      np_chains_list = msa_pairing.remove_all_gapped_rows_from_all_seqs(
          np_chains_list, asym_ids=nonempty_asym_ids
      )

    # Crop MSA rows.
    cropped_chains_list = []
    for chain in np_chains_list:
      unpaired_msa_size, paired_msa_size = (
          msa_pairing.choose_paired_unpaired_msa_crop_sizes(
              unpaired_msa=chain['msa'],
              paired_msa=chain.get('msa_all_seq'),
              total_msa_crop_size=padding_shapes.msa_size,
              max_paired_sequences=max_paired_sequences,
          )
      )
      cropped_chain = {
          'asym_id': chain['asym_id'],
          'chain_id': chain['chain_id'],
          'profile': chain['profile'],
          'deletion_mean': chain['deletion_mean'],
      }
      for feat in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
        if feat in chain:
          cropped_chain[feat] = chain[feat][:unpaired_msa_size]
        if feat + '_all_seq' in chain:
          cropped_chain[feat + '_all_seq'] = chain[feat + '_all_seq'][
              :paired_msa_size
          ]
      cropped_chains_list.append(cropped_chain)

    # Merge Chains.
    # Make sure the chain order is unaltered before slicing with tokens.
    curr_chain_order = [chain['chain_id'] for chain in cropped_chains_list]
    orig_chain_order = [chain['chain_id'] for chain in substruct.iter_chains()]
    assert curr_chain_order == orig_chain_order
    np_example = {
        'asym_id': np.concatenate(
            [c['asym_id'] for c in cropped_chains_list], axis=0
        ),
    }
    for feature in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
      for feat in [feature, feature + '_all_seq']:
        if feat in cropped_chains_list[0]:
          np_example[feat] = merging_features.merge_msa_features(
              feat, cropped_chains_list
          )
    for feature in ['profile', 'deletion_mean']:
      feature_list = [c[feature] for c in cropped_chains_list]
      np_example[feature] = np.concatenate(feature_list, axis=0)

    # Crop MSA rows to maximum size given by chains participating in the crop.
    max_allowed_unpaired = max([
        len(chain['msa'])
        for chain in cropped_chains_list
        if chain['asym_id'][0] in nonempty_asym_ids
    ])
    np_example['msa'] = np_example['msa'][:max_allowed_unpaired]
    if 'msa_all_seq' in np_example:
      max_allowed_paired = max([
          len(chain['msa_all_seq'])
          for chain in cropped_chains_list
          if chain['asym_id'][0] in nonempty_asym_ids
      ])
      np_example['msa_all_seq'] = np_example['msa_all_seq'][:max_allowed_paired]

    np_example = merging_features.merge_paired_and_unpaired_msa(np_example)

    # Crop MSA residues. Need to use the standard token indices, since msa does
    # not expand non-standard residues. This means that for expanded residues,
    # we get repeated msa columns.
    new_cropping_idxs = standard_token_idxs
    for feature in data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES:
      if feature in np_example:
        np_example[feature] = np_example[feature][:, new_cropping_idxs].copy()
    for feature in ['profile', 'deletion_mean']:
      np_example[feature] = np_example[feature][new_cropping_idxs]

    # Make MSA mask.
    np_example['msa_mask'] = np.ones_like(np_example['msa'], dtype=np.float32)

    # Count MSA size before padding.
    num_alignments = np_example['msa'].shape[0]

    # Pad:
    msa_size, num_tokens = padding_shapes.msa_size, padding_shapes.num_tokens

    def safe_cast_int8(x):
      return np.clip(x, np.iinfo(np.int8).min, np.iinfo(np.int8).max).astype(
          np.int8
      )

    return MSA(
        rows=_pad_to(safe_cast_int8(np_example['msa']), (msa_size, num_tokens)),
        mask=_pad_to(
            np_example['msa_mask'].astype(bool), (msa_size, num_tokens)
        ),
        # deletion_matrix may be out of int8 range, but we mostly care about
        # small values since we arctan it in the model.
        deletion_matrix=_pad_to(
            safe_cast_int8(np_example['deletion_matrix']),
            (msa_size, num_tokens),
        ),
        profile=_pad_to(np_example['profile'], (num_tokens, None)),
        deletion_mean=_pad_to(np_example['deletion_mean'], (num_tokens,)),
        num_alignments=np.array(num_alignments, dtype=np.int32),
    )

  def index_msa_rows(self, indices: xnp_ndarray) -> Self:
    assert indices.ndim == 1

    return MSA(
        rows=self.rows[indices, :],
        mask=self.mask[indices, :],
        deletion_matrix=self.deletion_matrix[indices, :],
        profile=self.profile,
        deletion_mean=self.deletion_mean,
        num_alignments=self.num_alignments,
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    output = cls(
        rows=batch['msa'],
        mask=batch['msa_mask'],
        deletion_matrix=batch['deletion_matrix'],
        profile=batch['profile'],
        deletion_mean=batch['deletion_mean'],
        num_alignments=batch['num_alignments'],
    )
    return output

  def as_data_dict(self) -> BatchDict:
    return {
        'msa': self.rows,
        'msa_mask': self.mask,
        'deletion_matrix': self.deletion_matrix,
        'profile': self.profile,
        'deletion_mean': self.deletion_mean,
        'num_alignments': self.num_alignments,
    }


jax.tree_util.register_dataclass(
    MSA,
    data_fields=[f.name for f in dataclasses.fields(MSA)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class Templates:
  """Dataclass containing templates."""

  # aatype of templates, int32 w shape [num_templates, num_res]
  aatype: xnp_ndarray
  # atom positions of templates, float32 w shape [num_templates, num_res, 24, 3]
  atom_positions: xnp_ndarray
  # atom mask of templates, bool w shape [num_templates, num_res, 24]
  atom_mask: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      standard_token_idxs: np.ndarray,
      padding_shapes: PaddingShapes,
      fold_input: folding_input.Input,
      max_templates: int,
      logging_name: str,
  ) -> Self:
    """Compute the template features."""

    seen_entities = {}
    polymer_entity_features = {True: {}, False: {}}

    substruct = atom_layout.make_structure(
        flat_layout=all_tokens,
        atom_coords=np.zeros(all_tokens.shape + (3,)),
        name=logging_name,
    )
    np_chains_list = []

    input_chains_by_id = {chain.id: chain for chain in fold_input.chains}

    nonempty_chain_ids = set(all_tokens.chain_id)
    for chain_info in substruct.iter_chains():
      chain_id = chain_info['chain_id']
      chain_type = chain_info['chain_type']
      chain = input_chains_by_id[chain_id]

      # Generalised "sequence" for ligands (can't trust residue name)
      chain_tokens = all_tokens[all_tokens.chain_id == chain_id]
      assert chain_tokens.res_name is not None
      three_letter_sequence = ','.join(chain_tokens.res_name.tolist())
      chain_num_tokens = len(chain_tokens.atom_name)

      # Don't compute features for chains not included in the crop, or ligands.
      skip_chain = (
          chain_type != mmcif_names.PROTEIN_CHAIN
          or chain_num_tokens <= 4  # not cache filled
          or chain_id not in nonempty_chain_ids
      )

      if three_letter_sequence in seen_entities:
        entity_id = seen_entities[three_letter_sequence]
      else:
        entity_id = len(seen_entities) + 1

      if entity_id not in polymer_entity_features[skip_chain]:
        if skip_chain:
          template_features = data3.empty_template_features(chain_num_tokens)
        else:
          assert isinstance(chain, folding_input.ProteinChain)

          sorted_features = []
          for template in chain.templates:
            struc = structure.from_mmcif(
                template.mmcif,
                fix_mse_residues=True,
                fix_arginines=True,
                include_bonds=False,
                include_water=False,
                include_other=True,  # For non-standard polymer chains.
            )
            hit_features = templates.get_polymer_features(
                chain=struc,
                chain_poly_type=mmcif_names.PROTEIN_CHAIN,
                query_sequence_length=len(chain.sequence),
                query_to_hit_mapping=dict(template.query_to_template_map),
            )
            sorted_features.append(hit_features)

          template_features = templates.package_template_features(
              hit_features=sorted_features,
              include_ligand_features=False,
          )

          template_features = data3.fix_template_features(
              template_features=template_features, num_res=len(chain.sequence)
          )

        template_features = _reduce_template_features(
            template_features, max_templates
        )
        polymer_entity_features[skip_chain][entity_id] = template_features

      seen_entities[three_letter_sequence] = entity_id
      feats = polymer_entity_features[skip_chain][entity_id].copy()
      feats['chain_id'] = chain_id
      np_chains_list.append(feats)

    # We pad the num_templates dimension before merging, so that different
    # chains can be concatenated on the num_res dimension.  Masking will be
    # applied so that each chains templates can't see each other.
    for chain in np_chains_list:
      chain['template_aatype'] = _pad_to(
          chain['template_aatype'], (max_templates, None)
      )
      chain['template_atom_positions'] = _pad_to(
          chain['template_atom_positions'], (max_templates, None, None, None)
      )
      chain['template_atom_mask'] = _pad_to(
          chain['template_atom_mask'], (max_templates, None, None)
      )

    # Merge on token dimension.
    np_example = {
        ft: np.concatenate([c[ft] for c in np_chains_list], axis=1)
        for ft in np_chains_list[0]
        if ft in data_constants.TEMPLATE_FEATURES
    }

    # Crop template data. Need to use the standard token indices, since msa does
    # not expand non-standard residues. This means that for expanded residues,
    # we get repeated template information.
    for feature_name, v in np_example.items():
      np_example[feature_name] = v[:max_templates, standard_token_idxs, ...]

    # Pad along the token dimension.
    templates_features = Templates(
        aatype=_pad_to(
            np_example['template_aatype'], (None, padding_shapes.num_tokens)
        ),
        atom_positions=_pad_to(
            np_example['template_atom_positions'],
            (None, padding_shapes.num_tokens, None, None),
        ),
        atom_mask=_pad_to(
            np_example['template_atom_mask'].astype(bool),
            (None, padding_shapes.num_tokens, None),
        ),
    )
    return templates_features

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    """Make Template from batch dictionary."""
    return cls(
        aatype=batch['template_aatype'],
        atom_positions=batch['template_atom_positions'],
        atom_mask=batch['template_atom_mask'],
    )

  def as_data_dict(self) -> BatchDict:
    return {
        'template_aatype': self.aatype,
        'template_atom_positions': self.atom_positions,
        'template_atom_mask': self.atom_mask,
    }


jax.tree_util.register_dataclass(
    Templates,
    data_fields=[f.name for f in dataclasses.fields(Templates)],
    meta_fields=[],
)


def _reduce_template_features(
    template_features: data3.FeatureDict,
    max_templates: int,
) -> data3.FeatureDict:
  """Reduces template features to max num templates and defined feature set."""
  num_templates = template_features['template_aatype'].shape[0]
  template_keep_mask = np.arange(num_templates) < max_templates
  template_fields = data_constants.TEMPLATE_FEATURES + (
      'template_release_timestamp',
  )
  template_features = {
      k: v[template_keep_mask]
      for k, v in template_features.items()
      if k in template_fields
  }
  return template_features


@dataclasses.dataclass(frozen=True)
class TokenFeatures:
  """Dataclass containing features for tokens."""

  residue_index: xnp_ndarray
  token_index: xnp_ndarray
  aatype: xnp_ndarray
  mask: xnp_ndarray
  seq_length: xnp_ndarray

  # Chain symmetry identifiers
  # for an A3B2 stoichiometry the meaning of these features is as follows:
  # asym_id:    1 2 3 4 5
  # entity_id:  1 1 1 2 2
  # sym_id:     1 2 3 1 2
  asym_id: xnp_ndarray
  entity_id: xnp_ndarray
  sym_id: xnp_ndarray

  # token type features
  is_protein: xnp_ndarray
  is_rna: xnp_ndarray
  is_dna: xnp_ndarray
  is_ligand: xnp_ndarray
  is_nonstandard_polymer_chain: xnp_ndarray
  is_water: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Compute the per-token features."""

    residue_index = all_tokens.res_id.astype(np.int32)

    token_index = np.arange(1, len(all_tokens.atom_name) + 1).astype(np.int32)

    aatype = []
    for res_name, chain_type in zip(all_tokens.res_name, all_tokens.chain_type):
      if chain_type in mmcif_names.POLYMER_CHAIN_TYPES:
        res_name = mmcif_names.fix_non_standard_polymer_res(
            res_name=res_name, chain_type=chain_type
        )
        if (
            chain_type == mmcif_names.DNA_CHAIN
            and res_name == residue_names.UNK_DNA
        ):
          res_name = residue_names.UNK_NUCLEIC_ONE_LETTER
      elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
        res_name = residue_names.UNK
      else:
        raise ValueError(f'Chain type {chain_type} not polymer or ligand.')
      aa = residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP[res_name]
      aatype.append(aa)
    aatype = np.array(aatype, dtype=np.int32)

    mask = np.ones(all_tokens.shape[0], dtype=bool)
    chains = _compute_asym_entity_and_sym_id(all_tokens)
    m = dict(zip(chains.chain_id, chains.asym_id))
    asym_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)

    m = dict(zip(chains.chain_id, chains.entity_id))
    entity_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)

    m = dict(zip(chains.chain_id, chains.sym_id))
    sym_id = np.array([m[c] for c in all_tokens.chain_id], dtype=np.int32)

    seq_length = np.array(all_tokens.shape[0], dtype=np.int32)

    is_protein = all_tokens.chain_type == mmcif_names.PROTEIN_CHAIN
    is_rna = all_tokens.chain_type == mmcif_names.RNA_CHAIN
    is_dna = all_tokens.chain_type == mmcif_names.DNA_CHAIN
    is_ligand = np.isin(
        all_tokens.chain_type, list(mmcif_names.LIGAND_CHAIN_TYPES)
    )
    standard_polymer_chain = list(mmcif_names.NON_POLYMER_CHAIN_TYPES) + list(
        mmcif_names.STANDARD_POLYMER_CHAIN_TYPES
    )
    is_nonstandard_polymer_chain = np.isin(
        all_tokens.chain_type, standard_polymer_chain, invert=True
    )
    is_water = all_tokens.chain_type == mmcif_names.WATER

    return TokenFeatures(
        residue_index=_pad_to(residue_index, (padding_shapes.num_tokens,)),
        token_index=_pad_to(token_index, (padding_shapes.num_tokens,)),
        aatype=_pad_to(aatype, (padding_shapes.num_tokens,)),
        mask=_pad_to(mask, (padding_shapes.num_tokens,)),
        asym_id=_pad_to(asym_id, (padding_shapes.num_tokens,)),
        entity_id=_pad_to(entity_id, (padding_shapes.num_tokens,)),
        sym_id=_pad_to(sym_id, (padding_shapes.num_tokens,)),
        seq_length=seq_length,
        is_protein=_pad_to(is_protein, (padding_shapes.num_tokens,)),
        is_rna=_pad_to(is_rna, (padding_shapes.num_tokens,)),
        is_dna=_pad_to(is_dna, (padding_shapes.num_tokens,)),
        is_ligand=_pad_to(is_ligand, (padding_shapes.num_tokens,)),
        is_nonstandard_polymer_chain=_pad_to(
            is_nonstandard_polymer_chain, (padding_shapes.num_tokens,)
        ),
        is_water=_pad_to(is_water, (padding_shapes.num_tokens,)),
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        residue_index=batch['residue_index'],
        token_index=batch['token_index'],
        aatype=batch['aatype'],
        mask=batch['seq_mask'],
        entity_id=batch['entity_id'],
        asym_id=batch['asym_id'],
        sym_id=batch['sym_id'],
        seq_length=batch['seq_length'],
        is_protein=batch['is_protein'],
        is_rna=batch['is_rna'],
        is_dna=batch['is_dna'],
        is_ligand=batch['is_ligand'],
        is_nonstandard_polymer_chain=batch['is_nonstandard_polymer_chain'],
        is_water=batch['is_water'],
    )

  def as_data_dict(self) -> BatchDict:
    return {
        'residue_index': self.residue_index,
        'token_index': self.token_index,
        'aatype': self.aatype,
        'seq_mask': self.mask,
        'entity_id': self.entity_id,
        'asym_id': self.asym_id,
        'sym_id': self.sym_id,
        'seq_length': self.seq_length,
        'is_protein': self.is_protein,
        'is_rna': self.is_rna,
        'is_dna': self.is_dna,
        'is_ligand': self.is_ligand,
        'is_nonstandard_polymer_chain': self.is_nonstandard_polymer_chain,
        'is_water': self.is_water,
    }


jax.tree_util.register_dataclass(
    TokenFeatures,
    data_fields=[f.name for f in dataclasses.fields(TokenFeatures)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class PredictedStructureInfo:
  """Contains information necessary to work with predicted structure."""

  atom_mask: xnp_ndarray
  residue_center_index: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      all_token_atoms_layout: atom_layout.AtomLayout,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Compute the PredictedStructureInfo features.

    Args:
      all_tokens: flat AtomLayout with 1 representative atom per token, shape
        (num_tokens,)
      all_token_atoms_layout: AtomLayout for all atoms per token, shape
        (num_tokens, max_atoms_per_token)
      padding_shapes: padding shapes.

    Returns:
      A PredictedStructureInfo object.
    """
    atom_mask = _pad_to(
        all_token_atoms_layout.atom_name.astype(bool),
        (padding_shapes.num_tokens, None),
    )
    residue_center_index = np.zeros(padding_shapes.num_tokens, dtype=np.int32)
    for idx in range(all_tokens.shape[0]):
      repr_atom = all_tokens.atom_name[idx]
      atoms = list(all_token_atoms_layout.atom_name[idx, :])
      if repr_atom in atoms:
        residue_center_index[idx] = atoms.index(repr_atom)
      else:
        # Representative atoms can be missing if cropping the number of atoms
        # per residue.
        logging.warning(
            'The representative atom in all_tokens (%s) is not in '
            'all_token_atoms_layout (%s)',
            all_tokens[idx : idx + 1],
            all_token_atoms_layout[idx, :],
        )
        residue_center_index[idx] = 0
    return cls(atom_mask=atom_mask, residue_center_index=residue_center_index)

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        atom_mask=batch['pred_dense_atom_mask'],
        residue_center_index=batch['residue_center_index'],
    )

  def as_data_dict(self) -> BatchDict:
    return {
        'pred_dense_atom_mask': self.atom_mask,
        'residue_center_index': self.residue_center_index,
    }


jax.tree_util.register_dataclass(
    PredictedStructureInfo,
    data_fields=[f.name for f in dataclasses.fields(PredictedStructureInfo)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class PolymerLigandBondInfo:
  """Contains information about polymer-ligand bonds."""

  tokens_to_polymer_ligand_bonds: atom_layout.GatherInfo
  # Gather indices to convert from cropped dense atom layout to bonds layout
  # (num_tokens, 2)
  token_atoms_to_bonds: atom_layout.GatherInfo

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      all_token_atoms_layout: atom_layout.AtomLayout,
      bond_layout: atom_layout.AtomLayout | None,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Computes the InterChainBondInfo features.

    Args:
      all_tokens: AtomLayout for tokens; shape (num_tokens,).
      all_token_atoms_layout: Atom Layout for all atoms (num_tokens,
        max_atoms_per_token)
      bond_layout: Bond layout for polymer-ligand bonds.
      padding_shapes: Padding shapes.

    Returns:
      A PolymerLigandBondInfo object.
    """

    if bond_layout is not None:
      # Must convert to list before calling np.isin, will not work raw.
      peptide_types = list(mmcif_names.PEPTIDE_CHAIN_TYPES)
      nucleic_types = list(mmcif_names.NUCLEIC_ACID_CHAIN_TYPES) + [
          mmcif_names.OTHER_CHAIN
      ]
      # These atom renames are so that we can use the atom layout code with
      # all_tokens, which only has a single atom per token.
      atom_names = bond_layout.atom_name.copy()
      atom_names[np.isin(bond_layout.chain_type, peptide_types)] = 'CA'
      atom_names[np.isin(bond_layout.chain_type, nucleic_types)] = "C1'"
      adjusted_bond_layout = atom_layout.AtomLayout(
          atom_name=atom_names,
          res_id=bond_layout.res_id,
          chain_id=bond_layout.chain_id,
          chain_type=bond_layout.chain_type,
      )
      # Remove bonds that are not in the crop.
      cropped_tokens_to_bonds = atom_layout.compute_gather_idxs(
          source_layout=all_tokens, target_layout=adjusted_bond_layout
      )
      bond_is_in_crop = np.all(
          cropped_tokens_to_bonds.gather_mask, axis=1
      ).astype(bool)
      adjusted_bond_layout = adjusted_bond_layout[bond_is_in_crop, :]
    else:
      # Create layout with correct shape when bond_layout is None.
      s = (0, 2)
      adjusted_bond_layout = atom_layout.AtomLayout(
          atom_name=np.array([], dtype=object).reshape(s),
          res_id=np.array([], dtype=int).reshape(s),
          chain_id=np.array([], dtype=object).reshape(s),
      )
    adjusted_bond_layout = adjusted_bond_layout.copy_and_pad_to(
        (padding_shapes.num_tokens, 2)
    )
    tokens_to_polymer_ligand_bonds = atom_layout.compute_gather_idxs(
        source_layout=all_tokens, target_layout=adjusted_bond_layout
    )

    # Stuff for computing the bond loss.
    if bond_layout is not None:
      # Pad to num_tokens (hoping that there are never more bonds than tokens).
      padded_bond_layout = bond_layout.copy_and_pad_to(
          (padding_shapes.num_tokens, 2)
      )
      token_atoms_to_bonds = atom_layout.compute_gather_idxs(
          source_layout=all_token_atoms_layout, target_layout=padded_bond_layout
      )
    else:
      token_atoms_to_bonds = atom_layout.GatherInfo(
          gather_idxs=np.zeros((padding_shapes.num_tokens, 2), dtype=int),
          gather_mask=np.zeros((padding_shapes.num_tokens, 2), dtype=bool),
          input_shape=np.array((
              padding_shapes.num_tokens,
              all_token_atoms_layout.shape[1],
          )),
      )

    return cls(
        tokens_to_polymer_ligand_bonds=tokens_to_polymer_ligand_bonds,
        token_atoms_to_bonds=token_atoms_to_bonds,
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        tokens_to_polymer_ligand_bonds=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='tokens_to_polymer_ligand_bonds'
        ),
        token_atoms_to_bonds=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='token_atoms_to_polymer_ligand_bonds'
        ),
    )

  def as_data_dict(self) -> BatchDict:
    return {
        **self.tokens_to_polymer_ligand_bonds.as_dict(
            key_prefix='tokens_to_polymer_ligand_bonds'
        ),
        **self.token_atoms_to_bonds.as_dict(
            key_prefix='token_atoms_to_polymer_ligand_bonds'
        ),
    }


jax.tree_util.register_dataclass(
    PolymerLigandBondInfo,
    data_fields=[f.name for f in dataclasses.fields(PolymerLigandBondInfo)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class LigandLigandBondInfo:
  """Contains information about the location of ligand-ligand bonds."""

  tokens_to_ligand_ligand_bonds: atom_layout.GatherInfo

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      bond_layout: atom_layout.AtomLayout | None,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Computes the InterChainBondInfo features.

    Args:
      all_tokens: AtomLayout for tokens; shape (num_tokens,).
      bond_layout: Bond layout for ligand-ligand bonds.
      padding_shapes: Padding shapes.

    Returns:
      A LigandLigandBondInfo object.
    """

    if bond_layout is not None:
      # Discard any bonds that do not join to an existing atom.
      keep_mask = []
      all_atom_ids = {
          uid
          for uid in zip(
              all_tokens.chain_id,
              all_tokens.res_id,
              all_tokens.atom_name,
              strict=True,
          )
      }
      for chain_id, res_id, atom_name in zip(
          bond_layout.chain_id,
          bond_layout.res_id,
          bond_layout.atom_name,
          strict=True,
      ):
        atom_a = (chain_id[0], res_id[0], atom_name[0])
        atom_b = (chain_id[1], res_id[1], atom_name[1])
        if atom_a in all_atom_ids and atom_b in all_atom_ids:
          keep_mask.append(True)
        else:
          keep_mask.append(False)
      keep_mask = np.array(keep_mask).astype(bool)
      bond_layout = bond_layout[keep_mask]
      # Remove any bonds to Hydrogen atoms.
      bond_layout = bond_layout[
          ~np.char.startswith(bond_layout.atom_name.astype(str), 'H').any(
              axis=1
          )
      ]
      atom_names = bond_layout.atom_name
      adjusted_bond_layout = atom_layout.AtomLayout(
          atom_name=atom_names,
          res_id=bond_layout.res_id,
          chain_id=bond_layout.chain_id,
          chain_type=bond_layout.chain_type,
      )
    else:
      # Create layout with correct shape when bond_layout is None.
      s = (0, 2)
      adjusted_bond_layout = atom_layout.AtomLayout(
          atom_name=np.array([], dtype=object).reshape(s),
          res_id=np.array([], dtype=int).reshape(s),
          chain_id=np.array([], dtype=object).reshape(s),
      )
    # 10 x num_tokens as max_inter_bonds_ratio + max_intra_bonds_ration = 2.061.
    adjusted_bond_layout = adjusted_bond_layout.copy_and_pad_to(
        (padding_shapes.num_tokens * 10, 2)
    )
    gather_idx = atom_layout.compute_gather_idxs(
        source_layout=all_tokens, target_layout=adjusted_bond_layout
    )
    return cls(tokens_to_ligand_ligand_bonds=gather_idx)

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        tokens_to_ligand_ligand_bonds=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='tokens_to_ligand_ligand_bonds'
        )
    )

  def as_data_dict(self) -> BatchDict:
    return {
        **self.tokens_to_ligand_ligand_bonds.as_dict(
            key_prefix='tokens_to_ligand_ligand_bonds'
        )
    }


jax.tree_util.register_dataclass(
    LigandLigandBondInfo,
    data_fields=[f.name for f in dataclasses.fields(LigandLigandBondInfo)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class PseudoBetaInfo:
  """Contains information for extracting pseudo-beta and equivalent atoms."""

  token_atoms_to_pseudo_beta: atom_layout.GatherInfo

  @classmethod
  def compute_features(
      cls,
      all_token_atoms_layout: atom_layout.AtomLayout,
      ccd: chemical_components.Ccd,
      padding_shapes: PaddingShapes,
      logging_name: str,
  ) -> Self:
    """Compute the PseudoBetaInfo features.

    Args:
      all_token_atoms_layout:  AtomLayout for all atoms per token, shape
        (num_tokens, max_atoms_per_token)
      ccd: The chemical components dictionary.
      padding_shapes: padding shapes.
      logging_name: logging name for debugging (usually the mmcif_id)

    Returns:
      A PseudoBetaInfo object.
    """
    token_idxs = []
    atom_idxs = []
    for token_idx in range(all_token_atoms_layout.shape[0]):
      chain_type = all_token_atoms_layout.chain_type[token_idx, 0]
      atom_names = list(all_token_atoms_layout.atom_name[token_idx, :])
      atom_idx = None
      is_nucleic_backbone = (
          chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES
          or chain_type == mmcif_names.OTHER_CHAIN
      )
      if chain_type == mmcif_names.PROTEIN_CHAIN:
        # Protein chains
        if 'CB' in atom_names:
          atom_idx = atom_names.index('CB')
        elif 'CA' in atom_names:
          atom_idx = atom_names.index('CA')
      elif is_nucleic_backbone:
        # RNA / DNA chains
        res_name = all_token_atoms_layout.res_name[token_idx, 0]
        cifdict = ccd.get(res_name)
        if cifdict:
          parent = cifdict['_chem_comp.mon_nstd_parent_comp_id'][0]
          if parent != '?':
            res_name = parent
        if res_name in {'A', 'G', 'DA', 'DG'}:
          if 'C4' in atom_names:
            atom_idx = atom_names.index('C4')
        else:
          if 'C2' in atom_names:
            atom_idx = atom_names.index('C2')
      elif chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
        # Ligands: there is only one atom per token
        atom_idx = 0
      else:
        logging.warning(
            '%s: Unknown chain type for token %i. (%s)',
            logging_name,
            token_idx,
            all_token_atoms_layout[token_idx : token_idx + 1],
        )
        atom_idx = 0
      if atom_idx is None:
        (valid_atom_idxs,) = np.nonzero(
            all_token_atoms_layout.atom_name[token_idx, :]
        )
        if valid_atom_idxs.shape[0] > 0:
          atom_idx = valid_atom_idxs[0]
        else:
          atom_idx = 0
        logging.warning(
            '%s token %i (%s), does not contain a pseudo-beta atom.'
            'Using first valid atom (%s) instead.',
            logging_name,
            token_idx,
            all_token_atoms_layout[token_idx : token_idx + 1],
            all_token_atoms_layout.atom_name[token_idx, atom_idx],
        )

      token_idxs.append(token_idx)
      atom_idxs.append(atom_idx)

    pseudo_beta_layout = all_token_atoms_layout[token_idxs, atom_idxs]
    pseudo_beta_layout = pseudo_beta_layout.copy_and_pad_to((
        padding_shapes.num_tokens,
    ))
    token_atoms_to_pseudo_beta = atom_layout.compute_gather_idxs(
        source_layout=all_token_atoms_layout, target_layout=pseudo_beta_layout
    )

    return cls(
        token_atoms_to_pseudo_beta=token_atoms_to_pseudo_beta,
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        token_atoms_to_pseudo_beta=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='token_atoms_to_pseudo_beta'
        ),
    )

  def as_data_dict(self) -> BatchDict:
    return {
        **self.token_atoms_to_pseudo_beta.as_dict(
            key_prefix='token_atoms_to_pseudo_beta'
        ),
    }


jax.tree_util.register_dataclass(
    PseudoBetaInfo,
    data_fields=[f.name for f in dataclasses.fields(PseudoBetaInfo)],
    meta_fields=[],
)


_DEFAULT_BLANK_REF = {
    'positions': np.zeros(3),
    'mask': 0,
    'element': 0,
    'charge': 0,
    'atom_name_chars': np.zeros(4),
}


def random_rotation(random_state: np.random.RandomState) -> np.ndarray:
  # Create a random rotation (Gram-Schmidt orthogonalization of two
  # random normal vectors)
  v0, v1 = random_state.normal(size=(2, 3))
  e0 = v0 / np.maximum(1e-10, np.linalg.norm(v0))
  v1 = v1 - e0 * np.dot(v1, e0)
  e1 = v1 / np.maximum(1e-10, np.linalg.norm(v1))
  e2 = np.cross(e0, e1)
  return np.stack([e0, e1, e2])


def random_augmentation(
    positions: np.ndarray,
    random_state: np.random.RandomState,
) -> np.ndarray:
  """Center then apply random translation and rotation."""

  center = np.mean(positions, axis=0)
  rot = random_rotation(random_state)
  positions_target = np.einsum('ij,kj->ki', rot, positions - center)

  translation = random_state.normal(size=(3,))
  positions_target = positions_target + translation
  return positions_target


def _get_reference_positions_from_ccd_cif(
    ccd_cif: cif_dict.CifDict,
    ref_max_modified_date: datetime.date,
    logging_name: str,
) -> np.ndarray:
  """Creates reference positions from a CCD mmcif data block."""
  num_atoms = len(ccd_cif['_chem_comp_atom.atom_id'])
  if '_chem_comp_atom.pdbx_model_Cartn_x_ideal' in ccd_cif:
    atom_x = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_x_ideal']
    atom_y = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_y_ideal']
    atom_z = ccd_cif['_chem_comp_atom.pdbx_model_Cartn_z_ideal']
  else:
    atom_x = np.array(['?'] * num_atoms)
    atom_y = np.array(['?'] * num_atoms)
    atom_z = np.array(['?'] * num_atoms)
  pos = np.array([[x, y, z] for x, y, z in zip(atom_x, atom_y, atom_z)])
  # Unknown reference coordinates are specified by '?' in chem comp dict.
  # Replace unknown reference coords with 0.
  if '?' in pos and '_chem_comp.pdbx_modified_date' in ccd_cif:
    # Use reference coordinates if modifed date is before cutoff.
    modified_dates = [
        datetime.date.fromisoformat(date)
        for date in ccd_cif['_chem_comp.pdbx_modified_date']
    ]
    max_modified_date = max(modified_dates)
    if max_modified_date < ref_max_modified_date:
      atom_x = ccd_cif['_chem_comp_atom.model_Cartn_x']
      atom_y = ccd_cif['_chem_comp_atom.model_Cartn_y']
      atom_z = ccd_cif['_chem_comp_atom.model_Cartn_z']
      pos = np.array([[x, y, z] for x, y, z in zip(atom_x, atom_y, atom_z)])
  if '?' in pos:
    if np.all(pos == '?'):
      logging.warning('All ref positions unknown for: %s', logging_name)
    else:
      logging.warning('Some ref positions unknown for: %s', logging_name)
    pos[pos == '?'] = 0
  return np.array(pos, dtype=np.float32)


def get_reference(
    res_name: str,
    chemical_components_data: struc_chem_comps.ChemicalComponentsData,
    ccd: chemical_components.Ccd,
    random_state: np.random.RandomState,
    ref_max_modified_date: datetime.date,
    conformer_max_iterations: int | None,
) -> tuple[dict[str, Any], Any, Any]:
  """Reference structure for residue from CCD or SMILES.

  Uses CCD entry if available, otherwise uses SMILES from chemical components
  data. Conformer generation is done using RDKit, with a fallback to CCD ideal
  or reference coordinates if RDKit fails and those coordinates are supplied.

  Args:
    res_name: ccd code of the residue.
    chemical_components_data: ChemicalComponentsData for making ref structure.
    ccd: The chemical components dictionary.
    random_state: Numpy RandomState
    ref_max_modified_date: date beyond which reference structures must not be
      modified to be allowed to use reference coordinates.
    conformer_max_iterations: Optional override for maximum number of iterations
      to run for RDKit conformer search.

  Returns:
    Mapping from atom names to features, from_atoms, dest_atoms.
  """

  ccd_cif = ccd.get(res_name)

  mol = None
  if ccd_cif:
    try:
      mol = rdkit_utils.mol_from_ccd_cif(ccd_cif, remove_hydrogens=False)
    except rdkit_utils.MolFromMmcifError:
      logging.warning('Failed to construct mol from ccd_cif for: %s', res_name)
  else:  # No CCD entry, use SMILES from chemical components data.
    if not (
        chemical_components_data.chem_comp
        and res_name in chemical_components_data.chem_comp
        and chemical_components_data.chem_comp[res_name].pdbx_smiles
    ):
      raise ValueError(f'No CCD entry or SMILES for {res_name}.')
    smiles_string = chemical_components_data.chem_comp[res_name].pdbx_smiles
    logging.info('Using SMILES for: %s - %s', res_name, smiles_string)

    mol = Chem.MolFromSmiles(smiles_string)
    if mol is None:
      # In this case the model will not have any information about this molecule
      # and will not be able to predict anything about it.
      raise ValueError(
          f'Failed to construct RDKit Mol for {res_name} from SMILES string: '
          f'{smiles_string} . This is likely due to an issue with the SMILES '
          'string. Note that the userCCD input format provides an alternative '
          'way to define custom molecules directly without RDKit or SMILES.'
      )
    mol = Chem.AddHs(mol)
    # No existing names, we assign them from the graph.
    mol = rdkit_utils.assign_atom_names_from_graph(mol)
    # Temporary CCD cif with just atom and bond information, no coordinates.
    ccd_cif = rdkit_utils.mol_to_ccd_cif(mol, component_id='fake_cif')

  conformer = None
  atom_names = []
  elements = []
  charges = []
  pos = []

  # If mol is not None (must be True for SMILES case), then we try and generate
  # an RDKit conformer.
  if mol is not None:
    conformer_random_seed = int(random_state.randint(1, 1 << 31))
    conformer = rdkit_utils.get_random_conformer(
        mol=mol,
        random_seed=conformer_random_seed,
        max_iterations=conformer_max_iterations,
        logging_name=res_name,
    )
    if conformer:
      for idx, atom in enumerate(mol.GetAtoms()):
        atom_names.append(atom.GetProp('atom_name'))
        elements.append(atom.GetAtomicNum())
        charges.append(atom.GetFormalCharge())
        coords = conformer.GetAtomPosition(idx)
        pos.append([coords.x, coords.y, coords.z])
      pos = np.array(pos, dtype=np.float32)

  # If no mol could be generated (can only happen when using CCD), or no
  # conformer could be generated from the mol (can happen in either case), then
  # use CCD cif instead (which will have zero coordinates for SMILES case).
  if conformer is None:
    atom_names = ccd_cif['_chem_comp_atom.atom_id']
    charges = ccd_cif['_chem_comp_atom.charge']
    type_symbols = ccd_cif['_chem_comp_atom.type_symbol']
    elements = [
        periodic_table.ATOMIC_NUMBER.get(elem_type.capitalize(), 0)
        for elem_type in type_symbols
    ]
    pos = _get_reference_positions_from_ccd_cif(
        ccd_cif=ccd_cif,
        ref_max_modified_date=ref_max_modified_date,
        logging_name=res_name,
    )

  # Augment reference positions.
  pos = random_augmentation(pos, random_state)

  # Extract atom and bond information from CCD cif.
  from_atom = ccd_cif.get('_chem_comp_bond.atom_id_1', None)
  dest_atom = ccd_cif.get('_chem_comp_bond.atom_id_2', None)

  features = {}
  for atom_name in atom_names:
    features[atom_name] = {}
    idx = atom_names.index(atom_name)
    charge = 0 if charges[idx] == '?' else int(charges[idx])
    atom_name_chars = np.array([ord(c) - 32 for c in atom_name], dtype=int)
    atom_name_chars = _pad_to(atom_name_chars, (4,))
    features[atom_name]['positions'] = pos[idx]
    features[atom_name]['mask'] = 1
    features[atom_name]['element'] = elements[idx]
    features[atom_name]['charge'] = charge
    features[atom_name]['atom_name_chars'] = atom_name_chars
  return features, from_atom, dest_atom


@dataclasses.dataclass(frozen=True)
class RefStructure:
  """Contains ref structure information."""

  # Array with positions, float32, shape [num_res, max_atoms_per_token, 3]
  positions: xnp_ndarray
  # Array with masks, bool, shape [num_res, max_atoms_per_token]
  mask: xnp_ndarray
  # Array with elements, int32, shape [num_res, max_atoms_per_token]
  element: xnp_ndarray
  # Array with charges, float32, shape [num_res, max_atoms_per_token]
  charge: xnp_ndarray
  # Array with atom name characters, int32, [num_res, max_atoms_per_token, 4]
  atom_name_chars: xnp_ndarray
  # Array with reference space uids, int32, [num_res, max_atoms_per_token]
  ref_space_uid: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      all_token_atoms_layout: atom_layout.AtomLayout,
      ccd: chemical_components.Ccd,
      padding_shapes: PaddingShapes,
      chemical_components_data: struc_chem_comps.ChemicalComponentsData,
      random_state: np.random.RandomState,
      ref_max_modified_date: datetime.date,
      conformer_max_iterations: int | None,
      ligand_ligand_bonds: atom_layout.AtomLayout | None = None,
  ) -> tuple[Self, Any]:
    """Reference structure information for each residue."""

    # Get features per atom
    padded_shape = (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
    result = {
        'positions': np.zeros((*padded_shape, 3), 'float32'),
        'mask': np.zeros(padded_shape, 'bool'),
        'element': np.zeros(padded_shape, 'int32'),
        'charge': np.zeros(padded_shape, 'float32'),
        'atom_name_chars': np.zeros((*padded_shape, 4), 'int32'),
        'ref_space_uid': np.zeros((*padded_shape,), 'int32'),
    }

    atom_names_all = []
    chain_ids_all = []
    res_ids_all = []

    # Cache reference conformations for each residue.
    conformations = {}
    ref_space_uids = {}
    for idx in np.ndindex(all_token_atoms_layout.shape):
      chain_id = all_token_atoms_layout.chain_id[idx]
      res_id = all_token_atoms_layout.res_id[idx]
      res_name = all_token_atoms_layout.res_name[idx]
      is_non_standard = res_name not in _STANDARD_RESIDUES
      atom_name = all_token_atoms_layout.atom_name[idx]
      if not atom_name:
        ref = _DEFAULT_BLANK_REF
      else:
        if (chain_id, res_id) not in conformations:
          conf, from_atom, dest_atom = get_reference(
              res_name=res_name,
              chemical_components_data=chemical_components_data,
              ccd=ccd,
              random_state=random_state,
              ref_max_modified_date=ref_max_modified_date,
              conformer_max_iterations=conformer_max_iterations,
          )
          conformations[(chain_id, res_id)] = conf

          if (
              is_non_standard
              and (from_atom is not None)
              and (dest_atom is not None)
          ):
            # Add intra-ligand bond graph
            atom_names_ligand = np.stack(
                [from_atom, dest_atom], axis=1, dtype=object
            )
            atom_names_all.append(atom_names_ligand)
            res_ids_all.append(
                np.full_like(atom_names_ligand, res_id, dtype=int)
            )
            chain_ids_all.append(
                np.full_like(atom_names_ligand, chain_id, dtype=object)
            )

        conformation = conformations.get(
            (chain_id, res_id), {atom_name: _DEFAULT_BLANK_REF}
        )
        if atom_name not in conformation:
          logging.warning(
              'Missing atom "%s" for CCD "%s"',
              atom_name,
              all_token_atoms_layout.res_name[idx],
          )
        ref = conformation.get(atom_name, _DEFAULT_BLANK_REF)
      for k in ref:
        result[k][idx] = ref[k]

      # Assign a unique reference space id to each component, to determine which
      # reference positions live in the same reference space.
      space_str_id = (
          all_token_atoms_layout.chain_id[idx],
          all_token_atoms_layout.res_id[idx],
      )
      if space_str_id not in ref_space_uids:
        ref_space_uids[space_str_id] = len(ref_space_uids)
      result['ref_space_uid'][idx] = ref_space_uids[space_str_id]

    if atom_names_all:
      atom_names_all = np.concatenate(atom_names_all, axis=0)
      res_ids_all = np.concatenate(res_ids_all, axis=0)
      chain_ids_all = np.concatenate(chain_ids_all, axis=0)
      if ligand_ligand_bonds is not None:
        adjusted_ligand_ligand_bonds = atom_layout.AtomLayout(
            atom_name=np.concatenate(
                [ligand_ligand_bonds.atom_name, atom_names_all], axis=0
            ),
            chain_id=np.concatenate(
                [ligand_ligand_bonds.chain_id, chain_ids_all], axis=0
            ),
            res_id=np.concatenate(
                [ligand_ligand_bonds.res_id, res_ids_all], axis=0
            ),
        )
      else:
        adjusted_ligand_ligand_bonds = atom_layout.AtomLayout(
            atom_name=atom_names_all,
            chain_id=chain_ids_all,
            res_id=res_ids_all,
        )
    else:
      adjusted_ligand_ligand_bonds = ligand_ligand_bonds

    return cls(**result), adjusted_ligand_ligand_bonds

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        positions=batch['ref_pos'],
        mask=batch['ref_mask'],
        element=batch['ref_element'],
        charge=batch['ref_charge'],
        atom_name_chars=batch['ref_atom_name_chars'],
        ref_space_uid=batch['ref_space_uid'],
    )

  def as_data_dict(self) -> BatchDict:
    return {
        'ref_pos': self.positions,
        'ref_mask': self.mask,
        'ref_element': self.element,
        'ref_charge': self.charge,
        'ref_atom_name_chars': self.atom_name_chars,
        'ref_space_uid': self.ref_space_uid,
    }


jax.tree_util.register_dataclass(
    RefStructure,
    data_fields=[f.name for f in dataclasses.fields(RefStructure)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class ConvertModelOutput:
  """Contains atom layout info."""

  cleaned_struc: structure.Structure
  token_atoms_layout: atom_layout.AtomLayout
  flat_output_layout: atom_layout.AtomLayout
  empty_output_struc: structure.Structure
  polymer_ligand_bonds: atom_layout.AtomLayout
  ligand_ligand_bonds: atom_layout.AtomLayout

  @classmethod
  def compute_features(
      cls,
      all_token_atoms_layout: atom_layout.AtomLayout,
      padding_shapes: PaddingShapes,
      cleaned_struc: structure.Structure,
      flat_output_layout: atom_layout.AtomLayout,
      empty_output_struc: structure.Structure,
      polymer_ligand_bonds: atom_layout.AtomLayout,
      ligand_ligand_bonds: atom_layout.AtomLayout,
  ) -> Self:
    """Pads the all_token_atoms_layout and stores other data."""
    # Crop and pad the all_token_atoms_layout.
    token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
        (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
    )

    return cls(
        cleaned_struc=cleaned_struc,
        token_atoms_layout=token_atoms_layout,
        flat_output_layout=flat_output_layout,
        empty_output_struc=empty_output_struc,
        polymer_ligand_bonds=polymer_ligand_bonds,
        ligand_ligand_bonds=ligand_ligand_bonds,
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    """Construct atom layout object from dictionary."""

    return cls(
        cleaned_struc=_unwrap(batch.get('cleaned_struc', None)),
        token_atoms_layout=_unwrap(batch.get('token_atoms_layout', None)),
        flat_output_layout=_unwrap(batch.get('flat_output_layout', None)),
        empty_output_struc=_unwrap(batch.get('empty_output_struc', None)),
        polymer_ligand_bonds=_unwrap(batch.get('polymer_ligand_bonds', None)),
        ligand_ligand_bonds=_unwrap(batch.get('ligand_ligand_bonds', None)),
    )

  def as_data_dict(self) -> BatchDict:
    return {
        'cleaned_struc': np.array(self.cleaned_struc, object),
        'token_atoms_layout': np.array(self.token_atoms_layout, object),
        'flat_output_layout': np.array(self.flat_output_layout, object),
        'empty_output_struc': np.array(self.empty_output_struc, object),
        'polymer_ligand_bonds': np.array(self.polymer_ligand_bonds, object),
        'ligand_ligand_bonds': np.array(self.ligand_ligand_bonds, object),
    }


jax.tree_util.register_dataclass(
    ConvertModelOutput,
    data_fields=[f.name for f in dataclasses.fields(ConvertModelOutput)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class AtomCrossAtt:
  """Operate on flat atoms."""

  token_atoms_to_queries: atom_layout.GatherInfo
  tokens_to_queries: atom_layout.GatherInfo
  tokens_to_keys: atom_layout.GatherInfo
  queries_to_keys: atom_layout.GatherInfo
  queries_to_token_atoms: atom_layout.GatherInfo

  @classmethod
  def compute_features(
      cls,
      all_token_atoms_layout: atom_layout.AtomLayout,  # (num_tokens, num_dense)
      queries_subset_size: int,
      keys_subset_size: int,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Computes gather indices and meta data to work with a flat atom list."""

    token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
        (padding_shapes.num_tokens, all_token_atoms_layout.shape[1])
    )
    token_atoms_mask = token_atoms_layout.atom_name.astype(bool)
    flat_layout = token_atoms_layout[token_atoms_mask]
    num_atoms = flat_layout.shape[0]

    padded_flat_layout = flat_layout.copy_and_pad_to((
        padding_shapes.num_atoms,
    ))

    # Create the layout for queries
    num_subsets = padding_shapes.num_atoms // queries_subset_size
    lay_arr = padded_flat_layout.to_array()
    queries_layout = atom_layout.AtomLayout.from_array(
        lay_arr.reshape((6, num_subsets, queries_subset_size))
    )

    # Create the layout for the keys (the key subsets are centered around the
    # query subsets)
    # Create initial gather indices (contain out-of-bound indices)
    subset_centers = np.arange(
        queries_subset_size / 2, padding_shapes.num_atoms, queries_subset_size
    )
    flat_to_key_gathers = (
        subset_centers[:, None]
        + np.arange(-keys_subset_size / 2, keys_subset_size / 2)[None, :]
    )
    flat_to_key_gathers = flat_to_key_gathers.astype(int)
    # Shift subsets with out-of-bound indices, such that they are fully within
    # the bounds.
    for row in range(flat_to_key_gathers.shape[0]):
      if flat_to_key_gathers[row, 0] < 0:
        flat_to_key_gathers[row, :] -= flat_to_key_gathers[row, 0]
      elif flat_to_key_gathers[row, -1] > num_atoms - 1:
        overflow = flat_to_key_gathers[row, -1] - (num_atoms - 1)
        flat_to_key_gathers[row, :] -= overflow
    # Create the keys layout.
    keys_layout = padded_flat_layout[flat_to_key_gathers]

    # Create gather indices for conversion between token atoms layout,
    # queries layout and keys layout.
    token_atoms_to_queries = atom_layout.compute_gather_idxs(
        source_layout=token_atoms_layout, target_layout=queries_layout
    )

    token_atoms_to_keys = atom_layout.compute_gather_idxs(
        source_layout=token_atoms_layout, target_layout=keys_layout
    )

    queries_to_keys = atom_layout.compute_gather_idxs(
        source_layout=queries_layout, target_layout=keys_layout
    )

    queries_to_token_atoms = atom_layout.compute_gather_idxs(
        source_layout=queries_layout, target_layout=token_atoms_layout
    )

    # Create gather indices for conversion of tokens layout to
    # queries and keys layout
    token_idxs = np.arange(padding_shapes.num_tokens).astype(np.int64)
    token_idxs = np.broadcast_to(token_idxs[:, None], token_atoms_layout.shape)
    tokens_to_queries = atom_layout.GatherInfo(
        gather_idxs=atom_layout.convert(
            token_atoms_to_queries, token_idxs, layout_axes=(0, 1)
        ),
        gather_mask=atom_layout.convert(
            token_atoms_to_queries, token_atoms_mask, layout_axes=(0, 1)
        ),
        input_shape=np.array((padding_shapes.num_tokens,)),
    )

    tokens_to_keys = atom_layout.GatherInfo(
        gather_idxs=atom_layout.convert(
            token_atoms_to_keys, token_idxs, layout_axes=(0, 1)
        ),
        gather_mask=atom_layout.convert(
            token_atoms_to_keys, token_atoms_mask, layout_axes=(0, 1)
        ),
        input_shape=np.array((padding_shapes.num_tokens,)),
    )

    return cls(
        token_atoms_to_queries=token_atoms_to_queries,
        tokens_to_queries=tokens_to_queries,
        tokens_to_keys=tokens_to_keys,
        queries_to_keys=queries_to_keys,
        queries_to_token_atoms=queries_to_token_atoms,
    )

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(
        token_atoms_to_queries=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='token_atoms_to_queries'
        ),
        tokens_to_queries=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='tokens_to_queries'
        ),
        tokens_to_keys=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='tokens_to_keys'
        ),
        queries_to_keys=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='queries_to_keys'
        ),
        queries_to_token_atoms=atom_layout.GatherInfo.from_dict(
            batch, key_prefix='queries_to_token_atoms'
        ),
    )

  def as_data_dict(self) -> BatchDict:
    return {
        **self.token_atoms_to_queries.as_dict(
            key_prefix='token_atoms_to_queries'
        ),
        **self.tokens_to_queries.as_dict(key_prefix='tokens_to_queries'),
        **self.tokens_to_keys.as_dict(key_prefix='tokens_to_keys'),
        **self.queries_to_keys.as_dict(key_prefix='queries_to_keys'),
        **self.queries_to_token_atoms.as_dict(
            key_prefix='queries_to_token_atoms'
        ),
    }


jax.tree_util.register_dataclass(
    AtomCrossAtt,
    data_fields=[f.name for f in dataclasses.fields(AtomCrossAtt)],
    meta_fields=[],
)


@dataclasses.dataclass(frozen=True)
class Frames:
  """Features for backbone frames."""

  mask: xnp_ndarray

  @classmethod
  def compute_features(
      cls,
      all_tokens: atom_layout.AtomLayout,
      all_token_atoms_layout: atom_layout.AtomLayout,
      ref_structure: RefStructure,
      padding_shapes: PaddingShapes,
  ) -> Self:
    """Computes features for backbone frames."""
    num_tokens = padding_shapes.num_tokens
    all_token_atoms_layout = all_token_atoms_layout.copy_and_pad_to(
        (num_tokens, all_token_atoms_layout.shape[1])
    )

    all_token_atoms_to_all_tokens = atom_layout.compute_gather_idxs(
        source_layout=all_token_atoms_layout, target_layout=all_tokens
    )
    ref_coordinates = atom_layout.convert(
        all_token_atoms_to_all_tokens,
        ref_structure.positions.astype(np.float32),
        layout_axes=(0, 1),
    )
    ref_mask = atom_layout.convert(
        all_token_atoms_to_all_tokens,
        ref_structure.mask.astype(bool),
        layout_axes=(0, 1),
    )
    ref_mask = ref_mask & all_token_atoms_to_all_tokens.gather_mask.astype(bool)

    all_frame_mask = []

    # Iterate over tokens
    for idx, args in enumerate(
        zip(all_tokens.chain_type, all_tokens.chain_id, all_tokens.res_id)
    ):

      chain_type, chain_id, res_id = args

      if chain_type in list(mmcif_names.PEPTIDE_CHAIN_TYPES):
        frame_mask = True
      elif chain_type in list(mmcif_names.NUCLEIC_ACID_CHAIN_TYPES):
        frame_mask = True
      elif chain_type in list(mmcif_names.NON_POLYMER_CHAIN_TYPES):
        # For ligands, build frames from closest atoms from the same molecule.
        (local_token_idxs,) = np.where(
            (all_tokens.chain_type == chain_type)
            & (all_tokens.chain_id == chain_id)
            & (all_tokens.res_id == res_id)
        )

        if len(local_token_idxs) < 3:
          frame_mask = False

        else:
          # [local_tokens]
          local_dist = np.linalg.norm(
              ref_coordinates[idx] - ref_coordinates[local_token_idxs], axis=-1
          )
          local_mask = ref_mask[local_token_idxs]
          cost = local_dist + 1e8 * ~local_mask
          cost = cost + 1e8 * (idx == local_token_idxs)
          # [local_tokens]
          closest_idxs = np.argsort(cost, axis=0)

          # The closest indices index an array of local tokens. Convert this
          # to indices of the full (num_tokens,) array.
          global_closest_idxs = local_token_idxs[closest_idxs]

          # Construct frame by placing the current token at the origin and two
          # nearest atoms on either side.
          global_frame_idxs = np.array(
              (global_closest_idxs[0], idx, global_closest_idxs[1])
          )

          # Check that the frame atoms are not colinear.
          a, b, c = ref_coordinates[global_frame_idxs]
          vec1 = a - b
          vec2 = c - b
          # Reference coordinates can be all zeros, in which case we have
          # to explicitly set colinearity.
          if np.isclose(np.linalg.norm(vec1, axis=-1), 0) or np.isclose(
              np.linalg.norm(vec2, axis=-1), 0
          ):
            is_colinear = True
            logging.info('Found identical coordinates: Assigning as colinear.')
          else:
            vec1 = vec1 / np.linalg.norm(vec1, axis=-1)
            vec2 = vec2 / np.linalg.norm(vec2, axis=-1)
            cos_angle = np.einsum('...k,...k->...', vec1, vec2)
            # <25 degree deviation is considered colinear.
            is_colinear = 1 - np.abs(cos_angle) < 0.0937

          frame_mask = not is_colinear
      else:
        # No frame for other chain types.
        frame_mask = False

      all_frame_mask.append(frame_mask)

    all_frame_mask = np.array(all_frame_mask, dtype=bool)

    mask = _pad_to(all_frame_mask, (padding_shapes.num_tokens,))

    return cls(mask=mask)

  @classmethod
  def from_data_dict(cls, batch: BatchDict) -> Self:
    return cls(mask=batch['frames_mask'])

  def as_data_dict(self) -> BatchDict:
    return {'frames_mask': self.mask}


jax.tree_util.register_dataclass(
    Frames,
    data_fields=[f.name for f in dataclasses.fields(Frames)],
    meta_fields=[],
)