flax_model/alphafold3/data/parsers.py

"""Functions for parsing various file formats."""

from collections.abc import Iterable, Sequence
from typing import IO, TypeAlias

from flax_model.alphafold3.cpp import fasta_iterator
from flax_model.alphafold3.cpp import msa_conversion


DeletionMatrix: TypeAlias = Sequence[Sequence[int]]


def lazy_parse_fasta_string(fasta_string: str) -> Iterable[tuple[str, str]]:
  """Lazily parses a FASTA/A3M string and yields (sequence, description) tuples.

  This implementation is more memory friendly than `fasta_sequence` while
  offering comparable performance. The underlying implementation is in C++ and
  is therefore faster than a pure Python implementation.

  Use this method when parsing FASTA files where you already have the FASTA
  string, but need to control how far you iterate through its sequences.

  Arguments:
    fasta_string: A string with the contents of FASTA/A3M file.

  Returns:
    Iterator of (sequence, description). In the description, the leading ">" is
    stripped.

  Raises:
    ValueError if the FASTA/A3M file is invalid, e.g. empty.
  """

  # The lifetime of the FastaStringIterator is tied to the lifetime of
  # fasta_string - fasta_string must be kept while the iterator is in use.
  return fasta_iterator.FastaStringIterator(fasta_string)


def parse_fasta(fasta_string: str) -> tuple[Sequence[str], Sequence[str]]:
  """Parses FASTA string and returns list of strings with amino-acid sequences.

  Arguments:
    fasta_string: The string contents of a FASTA file.

  Returns:
    A tuple of two lists:
    * A list of sequences.
    * A list of sequence descriptions taken from the comment lines. In the
      same order as the sequences.
  """
  return fasta_iterator.parse_fasta_include_descriptions(fasta_string)


def convert_a3m_to_stockholm(a3m: str, max_seqs: int | None = None) -> str:
  """Converts MSA in the A3M format to the Stockholm format."""
  sequences, descriptions = parse_fasta(a3m)
  if max_seqs is not None:
    sequences = sequences[:max_seqs]
    descriptions = descriptions[:max_seqs]

  stockholm = ['# STOCKHOLM 1.0', '']

  # Add the Stockholm header with the sequence metadata.
  names = []
  for i, description in enumerate(descriptions):
    name, _, rest = description.replace('\t', ' ').partition(' ')
    # Ensure that the names are unique - stockholm format requires that
    # the sequence names are unique.
    name = f'{name}_{i}'
    names.append(name)
    # Avoid zero-length description due to historic hmmbuild parsing bug.
    desc = rest.strip() or '<EMPTY>'
    stockholm.append(f'#=GS {name.strip()} DE {desc}')
  stockholm.append('')

  # Convert insertions in a sequence into gaps in all other sequences that don't
  # have an insertion in that column as well.
  sequences = msa_conversion.convert_a3m_to_stockholm(sequences)

  # Add the MSA data.
  max_name_width = max(len(name) for name in names)
  for name, sequence in zip(names, sequences, strict=True):
    # Align the names to the left and pad with spaces to the maximum length.
    stockholm.append(f'{name:<{max_name_width}s} {sequence}')

  # Add the reference annotation for the query (the first sequence).
  ref_annotation = ''.join('.' if c == '-' else 'x' for c in sequences[0])
  stockholm.append(f'{"#=GC RF":<{max_name_width}s} {ref_annotation}')
  stockholm.append('//')

  return '\n'.join(stockholm)


def convert_stockholm_to_a3m(
    stockholm: IO[str],
    max_sequences: int | None = None,
    remove_first_row_gaps: bool = True,
    linewidth: int | None = None,
) -> str:
  """Converts MSA in Stockholm format to the A3M format."""
  descriptions = {}
  sequences = {}
  reached_max_sequences = False

  if linewidth is not None and linewidth <= 0:
    raise ValueError('linewidth must be > 0 or None')

  for line in stockholm:
    reached_max_sequences = max_sequences and len(sequences) >= max_sequences
    line = line.strip()
    # Ignore blank lines, markup and end symbols - remainder are alignment
    # sequence parts.
    if not line or line.startswith(('#', '//')):
      continue
    seqname, aligned_seq = line.split(maxsplit=1)
    if seqname not in sequences:
      if reached_max_sequences:
        continue
      sequences[seqname] = ''
    sequences[seqname] += aligned_seq

  if not sequences:
    return ''

  stockholm.seek(0)
  for line in stockholm:
    line = line.strip()
    if line[:4] == '#=GS':
      # Description row - example format is:
      # #=GS UniRef90_Q9H5Z4/4-78            DE [subseq from] cDNA: FLJ22755 ...
      columns = line.split(maxsplit=3)
      seqname, feature = columns[1:3]
      value = columns[3] if len(columns) == 4 else ''
      if feature != 'DE':
        continue
      if reached_max_sequences and seqname not in sequences:
        continue
      descriptions[seqname] = value
      if len(descriptions) == len(sequences):
        break

  assert len(descriptions) <= len(sequences)

  # Convert sto format to a3m line by line
  a3m_sequences = {}
  # query_sequence is assumed to be the first sequence
  query_sequence = next(iter(sequences.values()))
  for seqname, sto_sequence in sequences.items():
    if remove_first_row_gaps:
      a3m_sequences[seqname] = msa_conversion.align_sequence_to_gapless_query(
          sequence=sto_sequence, query_sequence=query_sequence
      ).replace('.', '')
    else:
      a3m_sequences[seqname] = sto_sequence.replace('.', '')

  fasta_chunks = []

  for seqname, seq in a3m_sequences.items():
    fasta_chunks.append(f'>{seqname} {descriptions.get(seqname, "")}')

    if linewidth:
      fasta_chunks.extend(
          seq[i : linewidth + i] for i in range(0, len(seq), linewidth)
      )
    else:
      fasta_chunks.append(seq)

  return '\n'.join(fasta_chunks) + '\n'  # Include terminating newline.


def convert_mmseqs_stockholm_to_a3m(
    stockholm: IO[str],
    max_sequences: int | None = None,
    remove_first_row_gaps: bool = True,
    linewidth: int | None = None,
) -> str:
    """Converts MSA in Stockholm format to the A3M format."""
    from collections import defaultdict

    descriptions = {}
    sequences = {}  # 存储最终序列（自动处理重复键）
    seqname_counter = defaultdict(int)  # 记录原始seqname出现次数
    original_seqnames = {}  # 记录处理后的seqname对应的原始名�?

    if linewidth is not None and linewidth <= 0:
        raise ValueError('linewidth must be > 0 or None')

    # 第一遍：读取所有序列行，处理重复键
    for line in stockholm:
        line = line.strip()
        if not line or line.startswith(('#', '//')):
            continue
        parts = line.split(maxsplit=1)
        if len(parts) < 2:
            continue
        original_seqname, aligned_seq = parts

        # 生成唯一seqname（例�?"_1", "_1_a", "_1_b"�?
        count = seqname_counter[original_seqname]
        if count == 0:
            new_seqname = original_seqname
        else:
            new_seqname = f"{original_seqname}_{chr(97 + count - 1)}"  # 97是ASCII码的'a'
        seqname_counter[original_seqname] += 1

        # 达到最大序列数时跳过后续序�?
        if max_sequences and len(sequences) >= max_sequences:
            continue

        # 累积序列（原逻辑�?
        sequences.setdefault(new_seqname, '')
        sequences[new_seqname] += aligned_seq
        original_seqnames[new_seqname] = original_seqname

    if not sequences:
        return ''

    # 第二遍：读取描述信息
    stockholm.seek(0)
    for line in stockholm:
        line = line.strip()
        if line.startswith('#=GS'):
            columns = line.split(maxsplit=3)
            if len(columns) < 4:
                continue
            seqname, feature = columns[1:3]
            if feature != 'DE':
                continue
            value = columns[3] if len(columns) >= 4 else ''
            descriptions[seqname] = value

    # 转换序列格式
    a3m_sequences = {}
    query_sequence = next(iter(sequences.values())).replace('.', '')  # 假设第一个序列是查询序列

    for seqname, sto_sequence in sequences.items():
        current_seq = sto_sequence.replace('.', '')
        if remove_first_row_gaps:
            aligned = msa_conversion.align_sequence_to_gapless_query(
                sequence=current_seq, query_sequence=query_sequence
            ).replace('.', '')
        else:
            aligned = current_seq
        a3m_sequences[seqname] = aligned

    # 生成FASTA
    fasta_chunks = []
    for seqname, seq in a3m_sequences.items():
        original_seqname = original_seqnames.get(seqname, seqname)
        desc = descriptions.get(original_seqname, "")
        fasta_chunks.append(f'>{seqname} {desc}')

        if linewidth:
            chunks = [seq[i:i+linewidth] for i in range(0, len(seq), linewidth)]
            fasta_chunks.extend(chunks)
        else:
            fasta_chunks.append(seq)

    return '\n'.join(fasta_chunks) + '\n'