flax_model/alphafold3/data/pipeline.py

"""Functions for running the MSA and template tools for the AlphaFold model."""

from concurrent import futures
import dataclasses
import datetime
import functools
import logging
import time

from flax_model.alphafold3.common import folding_input
from flax_model.alphafold3.constants import mmcif_names
from flax_model.alphafold3.data import msa
from flax_model.alphafold3.data import msa_config
from flax_model.alphafold3.data import structure_stores
from flax_model.alphafold3.data import templates as templates_lib



# Cache to avoid re-running template search for the same sequence in homomers.
@functools.cache
def _get_protein_templates(
    sequence: str,
    input_msa_a3m: str,
    run_template_search: bool,
    templates_config: msa_config.TemplatesConfig,
    pdb_database_path: str,
) -> templates_lib.Templates:
  """Searches for templates for a single protein chain."""
  if run_template_search:
    templates_start_time = time.time()
    logging.info('Getting protein templates for sequence %s', sequence)
    protein_templates = templates_lib.Templates.from_seq_and_a3m(
        query_sequence=sequence,
        msa_a3m=input_msa_a3m,
        max_template_date=templates_config.filter_config.max_template_date,
        database_path=templates_config.template_tool_config.database_path,
        hmmsearch_config=templates_config.template_tool_config.hmmsearch_config,
        max_a3m_query_sequences=None,
        chain_poly_type=mmcif_names.PROTEIN_CHAIN,
        structure_store=structure_stores.StructureStore(pdb_database_path),
        filter_config=templates_config.filter_config,
    )
    logging.info(
        'Getting %d protein templates took %.2f seconds for sequence %s',
        protein_templates.num_hits,
        time.time() - templates_start_time,
        sequence,
    )
  else:
    logging.info('Skipping template search for sequence %s', sequence)
    protein_templates = templates_lib.Templates(
        query_sequence=sequence,
        hits=[],
        max_template_date=templates_config.filter_config.max_template_date,
        structure_store=structure_stores.StructureStore(pdb_database_path),
    )
  return protein_templates


# Cache to avoid re-running the MSA tools for the same sequence in homomers.
@functools.cache
def _get_protein_msa_and_templates(
    sequence: str,
    run_template_search: bool,
    uniref90_msa_config: msa_config.RunConfig,
    mgnify_msa_config: msa_config.RunConfig,
    small_bfd_msa_config: msa_config.RunConfig,
    uniprot_msa_config: msa_config.RunConfig,
    templates_config: msa_config.TemplatesConfig,
    pdb_database_path: str,
    model_loader_callback=None,
) -> tuple[msa.Msa, msa.Msa, templates_lib.Templates]:
  """Processes a single protein chain."""
  logging.info('Getting protein MSAs for sequence %s', sequence)
  msa_start_time = time.time()
  # Run various MSA tools in parallel. Use a ThreadPoolExecutor because
  # they're not blocked by the GIL, as they're sub-shelled out.
  with futures.ThreadPoolExecutor(max_workers=4) as executor:
    uniref90_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=uniref90_msa_config,
        chain_poly_type=mmcif_names.PROTEIN_CHAIN,
    )
    mgnify_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=mgnify_msa_config,
        chain_poly_type=mmcif_names.PROTEIN_CHAIN,
    )
    small_bfd_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=small_bfd_msa_config,
        chain_poly_type=mmcif_names.PROTEIN_CHAIN,
    )
    uniprot_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=uniprot_msa_config,
        chain_poly_type=mmcif_names.PROTEIN_CHAIN,
    )
  uniref90_msa = uniref90_msa_future.result()
  mgnify_msa = mgnify_msa_future.result()
  small_bfd_msa = small_bfd_msa_future.result()
  uniprot_msa = uniprot_msa_future.result()
  logging.info(
      'Getting protein MSAs took %.2f seconds for sequence %s',
      time.time() - msa_start_time,
      sequence,
  )
  
  # Load model after protein MSAs are completed
  if model_loader_callback is not None:
    logging.info('Protein MSAs completed. Loading model parameters...')
    model_loader_callback()

  logging.info('Deduplicating MSAs for sequence %s', sequence)
  msa_dedupe_start_time = time.time()
  with futures.ThreadPoolExecutor() as executor:
    unpaired_protein_msa_future = executor.submit(
        msa.Msa.from_multiple_msas,
        msas=[uniref90_msa, small_bfd_msa, mgnify_msa],
        deduplicate=True,
    )
    paired_protein_msa_future = executor.submit(
        msa.Msa.from_multiple_msas, msas=[uniprot_msa], deduplicate=False
    )
  unpaired_protein_msa = unpaired_protein_msa_future.result()
  paired_protein_msa = paired_protein_msa_future.result()
  logging.info(
      'Deduplicating MSAs took %.2f seconds for sequence %s, found %d unpaired'
      ' sequences, %d paired sequences',
      time.time() - msa_dedupe_start_time,
      sequence,
      unpaired_protein_msa.depth,
      paired_protein_msa.depth,
  )

  protein_templates = _get_protein_templates(
      sequence=sequence,
      input_msa_a3m=unpaired_protein_msa.to_a3m(),
      run_template_search=run_template_search,
      templates_config=templates_config,
      pdb_database_path=pdb_database_path,
  )

  return unpaired_protein_msa, paired_protein_msa, protein_templates


# Cache to avoid re-running the Nhmmer for the same sequence in homomers.
@functools.cache
def _get_rna_msa(
    sequence: str,
    nt_rna_msa_config: msa_config.NhmmerConfig,
    rfam_msa_config: msa_config.NhmmerConfig,
    rnacentral_msa_config: msa_config.NhmmerConfig,
) -> msa.Msa:
  """Processes a single RNA chain."""
  logging.info('Getting RNA MSAs for sequence %s', sequence)
  rna_msa_start_time = time.time()
  # Run various MSA tools in parallel. Use a ThreadPoolExecutor because
  # they're not blocked by the GIL, as they're sub-shelled out.
  with futures.ThreadPoolExecutor() as executor:
    nt_rna_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=nt_rna_msa_config,
        chain_poly_type=mmcif_names.RNA_CHAIN,
    )
    rfam_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=rfam_msa_config,
        chain_poly_type=mmcif_names.RNA_CHAIN,
    )
    rnacentral_msa_future = executor.submit(
        msa.get_msa,
        target_sequence=sequence,
        run_config=rnacentral_msa_config,
        chain_poly_type=mmcif_names.RNA_CHAIN,
    )
  nt_rna_msa = nt_rna_msa_future.result()
  rfam_msa = rfam_msa_future.result()
  rnacentral_msa = rnacentral_msa_future.result()
  rna_msa = msa.Msa.from_multiple_msas(
      msas=[rfam_msa, rnacentral_msa, nt_rna_msa],
      deduplicate=True,
  )
  logging.info(
      'Getting RNA MSAs took %.2f seconds for sequence %s, found %d unpaired'
      ' sequences',
      time.time() - rna_msa_start_time,
      sequence,
      rna_msa.depth,
  )
  return rna_msa


@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
class DataPipelineConfig:
  """The configuration for the data pipeline.

  Attributes:
    jackhmmer_binary_path: Jackhmmer binary path, used for protein MSA search.
    nhmmer_binary_path: Nhmmer binary path, used for RNA MSA search.
    hmmalign_binary_path: Hmmalign binary path, used to align hits to the query
      profile.
    hmmsearch_binary_path: Hmmsearch binary path, used for template search.
    hmmbuild_binary_path: Hmmbuild binary path, used to build HMM profile from
      raw MSA in template search.
    small_bfd_database_path: Small BFD database path, used for protein MSA
      search.
    small_bfd_z_value: The Z-value representing the database size in number of
      sequences for E-value calculation. Must be set for sharded databases.
    mgnify_database_path: Mgnify database path, used for protein MSA search.
    mgnify_z_value: The Z-value representing the database size in number of
      sequences for E-value calculation. Must be set for sharded databases.
    uniprot_cluster_annot_database_path: Uniprot database path, used for protein
      paired MSA search.
    uniprot_cluster_annot_z_value: The Z-value representing the database size in
      number of sequences for E-value calculation. Must be set for sharded
      databases.
    uniref90_database_path: UniRef90 database path, used for MSA search, and the
      MSA obtained by searching it is used to construct the profile for template
      search.
    uniref90_z_value: The Z-value representing the database size in number of
      sequences for E-value calculation. Must be set for sharded databases.
    ntrna_database_path: NT-RNA database path, used for RNA MSA search.
    ntrna_z_value: The Z-value representing the database size in megabases for
      E-value calculation. Must be set for sharded databases.
    rfam_database_path: Rfam database path, used for RNA MSA search.
    rfam_z_value: The Z-value representing the database size in megabases for
      E-value calculation. Must be set for sharded databases.
    rna_central_database_path: RNAcentral database path, used for RNA MSA
      search.
    rna_central_z_value: The Z-value representing the database size in megabases
      for E-value calculation. Must be set for sharded databases.
    seqres_database_path: PDB sequence database path, used for template search.
    pdb_database_path: PDB database directory with mmCIF files path, used for
      template search.
    jackhmmer_n_cpu: Number of CPUs to use for Jackhmmer.
    jackhmmer_max_parallel_shards: Maximum number of shards to search against in
      parallel. If None, one Jackhmmer instance will be run per shard. Only
      applicable if the database is sharded.
    nhmmer_n_cpu: Number of CPUs to use for Nhmmer.
    nhmmer_max_parallel_shards: Maximum number of shards to search against in
      parallel. If None, one Nhmmer instance will be run per shard. Only
      applicable if the database is sharded.
    max_template_date: The latest date of templates to use.
  """

  # Binary paths.
  jackhmmer_binary_path: str
  nhmmer_binary_path: str
  hmmalign_binary_path: str
  hmmsearch_binary_path: str
  hmmbuild_binary_path: str
  mmseqs_binary_path: str

  # Jackhmmer databases.
  small_bfd_database_path: str
  small_bfd_z_value: int | None = None
  mgnify_database_path: str
  mgnify_z_value: int | None = None
  uniprot_cluster_annot_database_path: str
  uniprot_cluster_annot_z_value: int | None = None
  uniref90_database_path: str
  uniref90_z_value: int | None = None
  # Nhmmer databases.
  ntrna_database_path: str
  ntrna_z_value: int | None = None
  rfam_database_path: str
  rfam_z_value: int | None = None
  rna_central_database_path: str
  rna_central_z_value: int | None = None
  # Template search databases.
  seqres_database_path: str
  pdb_database_path: str

  # Optional configuration for MSA tools.
  jackhmmer_n_cpu: int = 8
  jackhmmer_max_parallel_shards: int | None = None
  jackhmmer_max_threads: int | None = None
  nhmmer_n_cpu: int = 8
  nhmmer_max_parallel_shards: int | None = None
  nhmmer_max_threads: int | None = None
  mmseqs_n_cpu: int = 8

  max_template_date: datetime.date

  use_mmseqs: bool = False
  use_mmseqs_gpu: bool = False
  mmseqs_options: str
  result2msa_options: str

class DataPipeline:
  """Runs the alignment tools and assembles the input features."""

  def __init__(self, data_pipeline_config: DataPipelineConfig, model_loader_callback=None):
    """Initializes the data pipeline with default configurations."""
    if not data_pipeline_config.use_mmseqs:
      self._uniref90_msa_config = msa_config.RunConfig(
          config=msa_config.JackhmmerConfig(
              binary_path=data_pipeline_config.jackhmmer_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='uniref90',
                  path=data_pipeline_config.uniref90_database_path,
              ),
              n_cpu=data_pipeline_config.jackhmmer_n_cpu,
              n_iter=1,
              e_value=1e-4,
              z_value=data_pipeline_config.uniref90_z_value,
              dom_z_value=data_pipeline_config.uniref90_z_value,
              max_sequences=10_000,
              max_parallel_shards=data_pipeline_config.jackhmmer_max_parallel_shards,
              max_threads=data_pipeline_config.jackhmmer_max_threads,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._mgnify_msa_config = msa_config.RunConfig(
          config=msa_config.JackhmmerConfig(
              binary_path=data_pipeline_config.jackhmmer_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='mgnify',
                  path=data_pipeline_config.mgnify_database_path,
              ),
              n_cpu=data_pipeline_config.jackhmmer_n_cpu,
              n_iter=1,
              e_value=1e-4,
              z_value=data_pipeline_config.mgnify_z_value,
              dom_z_value=data_pipeline_config.mgnify_z_value,
              max_sequences=5_000,
              max_parallel_shards=data_pipeline_config.jackhmmer_max_parallel_shards,
              max_threads=data_pipeline_config.jackhmmer_max_threads,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._small_bfd_msa_config = msa_config.RunConfig(
          config=msa_config.JackhmmerConfig(
              binary_path=data_pipeline_config.jackhmmer_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='small_bfd',
                  path=data_pipeline_config.small_bfd_database_path,
              ),
              n_cpu=data_pipeline_config.jackhmmer_n_cpu,
              n_iter=1,
              e_value=1e-4,
              # Set z_value=138_515_945 to match the z_value used in the paper.
              # In practice, this has minimal impact on predicted structures.
              z_value=data_pipeline_config.small_bfd_z_value,
              dom_z_value=data_pipeline_config.small_bfd_z_value,
              max_sequences=5_000,
              max_parallel_shards=data_pipeline_config.jackhmmer_max_parallel_shards,
              max_threads=data_pipeline_config.jackhmmer_max_threads,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._uniprot_msa_config = msa_config.RunConfig(
          config=msa_config.JackhmmerConfig(
              binary_path=data_pipeline_config.jackhmmer_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='uniprot_cluster_annot',
                  path=data_pipeline_config.uniprot_cluster_annot_database_path,
              ),
              n_cpu=data_pipeline_config.jackhmmer_n_cpu,
              n_iter=1,
              e_value=1e-4,
              z_value=data_pipeline_config.uniprot_cluster_annot_z_value,
              dom_z_value=data_pipeline_config.uniprot_cluster_annot_z_value,
              max_sequences=50_000,
              max_parallel_shards=data_pipeline_config.jackhmmer_max_parallel_shards,
              max_threads=data_pipeline_config.jackhmmer_max_threads,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
    else:
      self._uniref90_msa_config = msa_config.RunConfig(
          config=msa_config.MmseqsConfig(
              binary_path=data_pipeline_config.mmseqs_binary_path,
              database_config=msa_config.DatabaseConfig(
                    name='uniref90',
                    path=data_pipeline_config.uniref90_database_path,
              ),
              n_cpu=data_pipeline_config.mmseqs_n_cpu,
              use_gpu=1 if data_pipeline_config.use_mmseqs_gpu else 0,
              mmseqs_options=data_pipeline_config.mmseqs_options,
              result2msa_options=data_pipeline_config.result2msa_options,
              max_sequences=10_000,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._mgnify_msa_config = msa_config.RunConfig(
          config=msa_config.MmseqsConfig(
              binary_path=data_pipeline_config.mmseqs_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='mgnify',
                  path=data_pipeline_config.mgnify_database_path,
              ),
              n_cpu=data_pipeline_config.mmseqs_n_cpu,
              use_gpu=1 if data_pipeline_config.use_mmseqs_gpu else 0,
              mmseqs_options=data_pipeline_config.mmseqs_options,
              result2msa_options=data_pipeline_config.result2msa_options,
              max_sequences=5_000,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._small_bfd_msa_config = msa_config.RunConfig(
          config=msa_config.MmseqsConfig(
              binary_path=data_pipeline_config.mmseqs_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='small_bfd',
                  path=data_pipeline_config.small_bfd_database_path,
              ),
              n_cpu=data_pipeline_config.mmseqs_n_cpu,
              use_gpu=1 if data_pipeline_config.use_mmseqs_gpu else 0,
              mmseqs_options=data_pipeline_config.mmseqs_options,
              result2msa_options=data_pipeline_config.result2msa_options,
              max_sequences=5_000,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
      self._uniprot_msa_config = msa_config.RunConfig(
          config=msa_config.MmseqsConfig(
              binary_path=data_pipeline_config.mmseqs_binary_path,
              database_config=msa_config.DatabaseConfig(
                  name='uniprot_cluster_annot',
                  path=data_pipeline_config.uniprot_cluster_annot_database_path,
              ),
              n_cpu=data_pipeline_config.mmseqs_n_cpu,
              use_gpu=1 if data_pipeline_config.use_mmseqs_gpu else 0,
              mmseqs_options=data_pipeline_config.mmseqs_options,
              result2msa_options=data_pipeline_config.result2msa_options,
              max_sequences=50_000,
          ),
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
          crop_size=None,
      )
    
    self._nt_rna_msa_config = msa_config.RunConfig(
        config=msa_config.NhmmerConfig(
            binary_path=data_pipeline_config.nhmmer_binary_path,
            hmmalign_binary_path=data_pipeline_config.hmmalign_binary_path,
            hmmbuild_binary_path=data_pipeline_config.hmmbuild_binary_path,
            database_config=msa_config.DatabaseConfig(
                name='nt_rna',
                path=data_pipeline_config.ntrna_database_path,
            ),
            n_cpu=data_pipeline_config.nhmmer_n_cpu,
            e_value=1e-3,
            alphabet='rna',
            z_value=data_pipeline_config.ntrna_z_value,
            max_sequences=10_000,
            max_parallel_shards=data_pipeline_config.nhmmer_max_parallel_shards,
            max_threads=data_pipeline_config.nhmmer_max_threads,
        ),
        chain_poly_type=mmcif_names.RNA_CHAIN,
        crop_size=None,
    )
    self._rfam_msa_config = msa_config.RunConfig(
        config=msa_config.NhmmerConfig(
            binary_path=data_pipeline_config.nhmmer_binary_path,
            hmmalign_binary_path=data_pipeline_config.hmmalign_binary_path,
            hmmbuild_binary_path=data_pipeline_config.hmmbuild_binary_path,
            database_config=msa_config.DatabaseConfig(
                name='rfam_rna',
                path=data_pipeline_config.rfam_database_path,
            ),
            n_cpu=data_pipeline_config.nhmmer_n_cpu,
            e_value=1e-3,
            alphabet='rna',
            z_value=data_pipeline_config.rfam_z_value,
            max_sequences=10_000,
            max_parallel_shards=data_pipeline_config.nhmmer_max_parallel_shards,
            max_threads=data_pipeline_config.nhmmer_max_threads,
        ),
        chain_poly_type=mmcif_names.RNA_CHAIN,
        crop_size=None,
    )
    self._rnacentral_msa_config = msa_config.RunConfig(
        config=msa_config.NhmmerConfig(
            binary_path=data_pipeline_config.nhmmer_binary_path,
            hmmalign_binary_path=data_pipeline_config.hmmalign_binary_path,
            hmmbuild_binary_path=data_pipeline_config.hmmbuild_binary_path,
            database_config=msa_config.DatabaseConfig(
                name='rna_central_rna',
                path=data_pipeline_config.rna_central_database_path,
            ),
            n_cpu=data_pipeline_config.nhmmer_n_cpu,
            e_value=1e-3,
            alphabet='rna',
            z_value=data_pipeline_config.rna_central_z_value,
            max_sequences=10_000,
            max_parallel_shards=data_pipeline_config.nhmmer_max_parallel_shards,
            max_threads=data_pipeline_config.nhmmer_max_threads,
        ),
        chain_poly_type=mmcif_names.RNA_CHAIN,
        crop_size=None,
    )

    self._templates_config = msa_config.TemplatesConfig(
        template_tool_config=msa_config.TemplateToolConfig(
            database_path=data_pipeline_config.seqres_database_path,
            chain_poly_type=mmcif_names.PROTEIN_CHAIN,
            hmmsearch_config=msa_config.HmmsearchConfig(
                hmmsearch_binary_path=data_pipeline_config.hmmsearch_binary_path,
                hmmbuild_binary_path=data_pipeline_config.hmmbuild_binary_path,
                filter_f1=0.1,
                filter_f2=0.1,
                filter_f3=0.1,
                e_value=100,
                inc_e=100,
                dom_e=100,
                incdom_e=100,
                alphabet='amino',
            ),
        ),
        filter_config=msa_config.TemplateFilterConfig(
            max_subsequence_ratio=0.95,
            min_align_ratio=0.1,
            min_hit_length=10,
            deduplicate_sequences=True,
            max_hits=4,
            max_template_date=data_pipeline_config.max_template_date,
        ),
    )
    self._pdb_database_path = data_pipeline_config.pdb_database_path
    self._model_loader_callback = model_loader_callback

  def process_protein_chain(
      self, chain: folding_input.ProteinChain
  ) -> folding_input.ProteinChain:
    """Processes a single protein chain."""
    has_unpaired_msa = chain.unpaired_msa is not None
    has_paired_msa = chain.paired_msa is not None
    has_templates = chain.templates is not None

    if not has_unpaired_msa and not has_paired_msa and not chain.templates:
      # MSA None - search. Templates either [] - don't search, or None - search.
      unpaired_msa, paired_msa, template_hits = _get_protein_msa_and_templates(
          sequence=chain.sequence,
          run_template_search=not has_templates,  # Skip template search if [].
          uniref90_msa_config=self._uniref90_msa_config,
          mgnify_msa_config=self._mgnify_msa_config,
          small_bfd_msa_config=self._small_bfd_msa_config,
          uniprot_msa_config=self._uniprot_msa_config,
          templates_config=self._templates_config,
          pdb_database_path=self._pdb_database_path,
          model_loader_callback=self._model_loader_callback,
      )
      unpaired_msa = unpaired_msa.to_a3m()
      paired_msa = paired_msa.to_a3m()
      templates = [
          folding_input.Template(
              mmcif=struc.to_mmcif(),
              query_to_template_map=hit.query_to_hit_mapping,
          )
          for hit, struc in template_hits.get_hits_with_structures()
      ]
    elif has_unpaired_msa and has_paired_msa and not has_templates:
      # Has MSA, but doesn't have templates. Search for templates only.
      empty_msa = msa.Msa.from_empty(
          query_sequence=chain.sequence,
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
      ).to_a3m()
      unpaired_msa = chain.unpaired_msa or empty_msa
      paired_msa = chain.paired_msa or empty_msa
      template_hits = _get_protein_templates(
          sequence=chain.sequence,
          input_msa_a3m=unpaired_msa,
          run_template_search=True,
          templates_config=self._templates_config,
          pdb_database_path=self._pdb_database_path,
      )
      templates = [
          folding_input.Template(
              mmcif=struc.to_mmcif(),
              query_to_template_map=hit.query_to_hit_mapping,
          )
          for hit, struc in template_hits.get_hits_with_structures()
      ]
    else:
      # Has MSA and templates, don't search for anything.
      if not has_unpaired_msa or not has_paired_msa or not has_templates:
        raise ValueError(
            f'Protein chain {chain.id} has unpaired MSA, paired MSA, or'
            ' templates set only partially. If you want to run the pipeline'
            ' with custom MSA/templates, you need to set all of them. You can'
            ' set MSA to empty string and templates to empty list to signify'
            ' that they should not be used and searched for.'
        )
      logging.info(
          'Skipping MSA and template search for protein chain %s because it '
          'already has MSAs and templates.',
          chain.id,
      )
      if not chain.unpaired_msa:
        logging.info('Using empty unpaired MSA for protein chain %s', chain.id)
      if not chain.paired_msa:
        logging.info('Using empty paired MSA for protein chain %s', chain.id)
      if not chain.templates:
        logging.info('Using no templates for protein chain %s', chain.id)
      empty_msa = msa.Msa.from_empty(
          query_sequence=chain.sequence,
          chain_poly_type=mmcif_names.PROTEIN_CHAIN,
      ).to_a3m()
      unpaired_msa = chain.unpaired_msa or empty_msa
      paired_msa = chain.paired_msa or empty_msa
      templates = chain.templates

    return folding_input.ProteinChain(
        id=chain.id,
        sequence=chain.sequence,
        ptms=chain.ptms,
        unpaired_msa=unpaired_msa,
        paired_msa=paired_msa,
        templates=templates,
    )

  def process_rna_chain(
      self, chain: folding_input.RnaChain
  ) -> folding_input.RnaChain:
    """Processes a single RNA chain."""
    if chain.unpaired_msa is not None:
      # Don't run MSA tools if the chain already has an MSA.
      logging.info(
          'Skipping MSA search for RNA chain %s because it already has MSA.',
          chain.id,
      )
      if not chain.unpaired_msa:
        logging.info('Using empty unpaired MSA for RNA chain %s', chain.id)
      empty_msa = msa.Msa.from_empty(
          query_sequence=chain.sequence, chain_poly_type=mmcif_names.RNA_CHAIN
      ).to_a3m()
      unpaired_msa = chain.unpaired_msa or empty_msa
    else:
      unpaired_msa = _get_rna_msa(
          sequence=chain.sequence,
          nt_rna_msa_config=self._nt_rna_msa_config,
          rfam_msa_config=self._rfam_msa_config,
          rnacentral_msa_config=self._rnacentral_msa_config,
      ).to_a3m()
    return folding_input.RnaChain(
        id=chain.id,
        sequence=chain.sequence,
        modifications=chain.modifications,
        unpaired_msa=unpaired_msa,
    )

  def process(self, fold_input: folding_input.Input) -> folding_input.Input:
    """Runs MSA and template tools and returns a new Input with the results."""
    processed_chains = []
    for chain in fold_input.chains:
      print(f'Running data pipeline for chain {chain.id}...')
      process_chain_start_time = time.time()
      match chain:
        case folding_input.ProteinChain():
          processed_chains.append(self.process_protein_chain(chain))
        case folding_input.RnaChain():
          processed_chains.append(self.process_rna_chain(chain))
        case _:
          processed_chains.append(chain)
      print(
          f'Running data pipeline for chain {chain.id} took'
          f' {time.time() - process_chain_start_time:.2f} seconds',
      )

    return dataclasses.replace(fold_input, chains=processed_chains)