File size: 1,800 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
"""Interface and implementations for fetching MSA data."""
from collections.abc import Sequence
from typing import Protocol, TypeAlias
from flax_model.alphafold3.data import msa
from flax_model.alphafold3.data import msa_config
MsaErrors: TypeAlias = Sequence[tuple[msa_config.RunConfig, str]]
class MsaProvider(Protocol):
"""Interface for providing Multiple Sequence Alignments."""
def __call__(
self,
query_sequence: str,
chain_polymer_type: str,
) -> tuple[msa.Msa, MsaErrors]:
"""Retrieve MSA for the given polymer query_sequence.
Args:
query_sequence: The residue sequence of the polymer to search for.
chain_polymer_type: The polymer type of the query_sequence. This must
match the chain_polymer_type of the provider.
Returns:
A tuple containing the MSA and MsaErrors. MsaErrors is a Sequence
containing a tuple for each msa_query that failed. Each tuple contains
the failing query and the associated error message.
"""
class EmptyMsaProvider:
"""MSA provider that returns just the query sequence, useful for testing."""
def __init__(self, chain_polymer_type: str):
self._chain_polymer_type = chain_polymer_type
def __call__(
self, query_sequence: str, chain_polymer_type: str
) -> tuple[msa.Msa, MsaErrors]:
"""Returns an MSA containing just the query sequence, never errors."""
if chain_polymer_type != self._chain_polymer_type:
raise ValueError(
f'EmptyMsaProvider of type {self._chain_polymer_type} called with '
f'sequence of {chain_polymer_type=}, {query_sequence=}.'
)
return (
msa.Msa.from_empty(
query_sequence=query_sequence,
chain_poly_type=self._chain_polymer_type,
),
(),
)
|