"""API for retrieving and manipulating template search results.""" from collections.abc import Iterable, Iterator, Mapping, Sequence import dataclasses import datetime import functools import os import re from typing import Any, Final, Self, TypeAlias from absl import logging from flax_model.alphafold3 import structure from flax_model.alphafold3.common import resources from flax_model.alphafold3.constants import atom_types from flax_model.alphafold3.constants import mmcif_names from flax_model.alphafold3.constants import residue_names from flax_model.alphafold3.data import msa_config from flax_model.alphafold3.data import parsers from flax_model.alphafold3.data import structure_stores from flax_model.alphafold3.data import template_realign from flax_model.alphafold3.data.tools import hmmsearch from flax_model.alphafold3.structure import mmcif import numpy as np _POLYMER_FEATURES: Final[Mapping[str, np.float64 | np.int32 | object]] = { 'template_aatype': np.int32, 'template_all_atom_masks': np.float64, 'template_all_atom_positions': np.float64, 'template_domain_names': object, 'template_release_date': object, 'template_sequence': object, } _LIGAND_FEATURES: Final[Mapping[str, Any]] = { 'ligand_features': Mapping[str, Any] } TemplateFeatures: TypeAlias = Mapping[ str, np.ndarray | bytes | Mapping[str, np.ndarray | bytes] ] _REQUIRED_METADATA_COLUMNS: Final[Sequence[str]] = ( 'seq_release_date', 'seq_unresolved_res_num', 'seq_author_chain_id', 'seq_sequence', ) @dataclasses.dataclass(frozen=True, kw_only=True, slots=True) class _Polymer: """Container for alphabet specific (dna, rna, protein) atom information.""" min_atoms: int num_atom_types: int atom_order: Mapping[str, int] _POLYMERS = { mmcif_names.PROTEIN_CHAIN: _Polymer( min_atoms=5, num_atom_types=atom_types.ATOM37_NUM, atom_order=atom_types.ATOM37_ORDER, ), mmcif_names.DNA_CHAIN: _Polymer( min_atoms=21, num_atom_types=atom_types.ATOM29_NUM, atom_order=atom_types.ATOM29_ORDER, ), mmcif_names.RNA_CHAIN: _Polymer( min_atoms=20, num_atom_types=atom_types.ATOM29_NUM, atom_order=atom_types.ATOM29_ORDER, ), } def _encode_restype( chain_poly_type: str, sequence: str, ) -> Sequence[int]: """Encodes a sequence of residue names as a sequence of ints. Args: chain_poly_type: Polymer chain type to determine sequence encoding. sequence: Polymer residues. Protein encoded by single letters. RNA and DNA encoded by multi-letter CCD codes. Returns: A sequence of integers encoding amino acid types for the given chain type. """ if chain_poly_type == mmcif_names.PROTEIN_CHAIN: return [ residue_names.PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP_TO_INT[ _STANDARDIZED_AA.get(res, res) ] for res in sequence ] unk_nucleic = residue_names.UNK_NUCLEIC_ONE_LETTER unk_nucleic_idx = residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP[ unk_nucleic ] if chain_poly_type == mmcif_names.RNA_CHAIN: return [ residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP.get( res, unk_nucleic_idx ) for res in sequence ] elif chain_poly_type == mmcif_names.DNA_CHAIN: # Map UNK DNA to the generic nucleic UNK (N), which happens to also be the # same as the RNA UNK. return [ residue_names.POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP.get( residue_names.DNA_COMMON_ONE_TO_TWO.get(res, unk_nucleic), unk_nucleic_idx, ) for res in sequence ] raise NotImplementedError(f'"{chain_poly_type}" unsupported.') _DAYS_BEFORE_QUERY_DATE: Final[int] = 60 _HIT_DESCRIPTION_REGEX = re.compile( r'(?P[a-z0-9]{4,})_(?P\w+)/(?P\d+)-(?P\d+) ' r'.* length:(?P\d+)\b.*' ) _STANDARDIZED_AA = {'B': 'D', 'J': 'X', 'O': 'X', 'U': 'C', 'Z': 'E'} class Error(Exception): """Base class for exceptions.""" class HitDateError(Error): """An error indicating that invalid release date was detected.""" class InvalidTemplateError(Error): """An error indicating that template is invalid.""" @dataclasses.dataclass(frozen=True, kw_only=True) class Hit: """Template hit metrics derived from the MSA for filtering and featurising. Attributes: pdb_id: The PDB ID of the hit. auth_chain_id: The author chain ID of the hit. hmmsearch_sequence: Hit sequence as given in hmmsearch a3m output. structure_sequence: Hit sequence as given in PDB structure. unresolved_res_indices: Indices of unresolved residues in the structure sequence. 0-based. query_sequence: The query nucleotide/amino acid sequence. start_index: The start index of the sequence relative to the full PDB seqres sequence. Inclusive and uses 0-based indexing. end_index: The end index of the sequence relative to the full PDB seqres sequence. Exclusive and uses 0-based indexing. full_length: Length of the full PDB seqres sequence. This can be different from the length from the actual sequence we get from the mmCIF and we use this to detect whether we need to realign or not. release_date: The release date of the PDB corresponding to this hit. chain_poly_type: The polymer type of the selected hit structure. """ pdb_id: str auth_chain_id: str hmmsearch_sequence: str structure_sequence: str unresolved_res_indices: Sequence[int] | None query_sequence: str start_index: int end_index: int full_length: int release_date: datetime.date chain_poly_type: str @functools.cached_property def query_to_hit_mapping(self) -> Mapping[int, int]: """0-based query index to hit index mapping.""" query_to_hit_mapping = {} hit_index = 0 query_index = 0 for residue in self.hmmsearch_sequence: # Gap inserted in the template if residue == '-': query_index += 1 # Deleted residue in the template (would be a gap in the query). elif residue.islower(): hit_index += 1 # Normal aligned residue, in both query and template. Add to mapping. elif residue.isupper(): query_to_hit_mapping[query_index] = hit_index query_index += 1 hit_index += 1 structure_subseq = self.structure_sequence[ self.start_index : self.end_index ] if self.matching_sequence != structure_subseq: # The seqres sequence doesn't match the structure sequence. Two cases: # 1. The sequences have the same length. The sequences are different # because our 3->1 residue code mapping is different from the one PDB # uses. We don't do anything in this case as both sequences have the # same length, so the original query to hit mapping stays valid. # 2. The sequences don't have the same length, the one in structure is # shorter. In this case we change the mapping to match the actual # structure sequence using a simple realignment algorithm. # This procedure was validated on all PDB seqres (2023_01_12) sequences # and handles all cases that can happen. if self.full_length != len(self.structure_sequence): return template_realign.realign_hit_to_structure( hit_sequence=self.matching_sequence, hit_start_index=self.start_index, hit_end_index=self.end_index, full_length=self.full_length, structure_sequence=self.structure_sequence, query_to_hit_mapping=query_to_hit_mapping, ) # Hmmsearch returns a subsequence and so far indices have been relative to # the subsequence. Add an offset to index relative to the full structure # sequence. return {q: h + self.start_index for q, h in query_to_hit_mapping.items()} @property def matching_sequence(self) -> str: """Returns the matching hit sequence including insertions. Make deleted residues uppercase and remove gaps ("-"). """ return self.hmmsearch_sequence.upper().replace('-', '') @functools.cached_property def output_templates_sequence(self) -> str: """Returns the final template sequence.""" result_seq = ['-'] * len(self.query_sequence) for query_index, template_index in self.query_to_hit_mapping.items(): result_seq[query_index] = self.structure_sequence[template_index] return ''.join(result_seq) @property def length_ratio(self) -> float: """Ratio of the length of the hit sequence to the query.""" return len(self.matching_sequence) / len(self.query_sequence) @property def align_ratio(self) -> float: """Ratio of the number of aligned residues to the query length.""" return len(self.query_to_hit_mapping) / len(self.query_sequence) @functools.cached_property def is_valid(self) -> bool: """Whether hit can be used as a template.""" if self.unresolved_res_indices is None: return False return bool( set(self.query_to_hit_mapping.values()) - set(self.unresolved_res_indices) ) @property def full_name(self) -> str: """A full name of the hit.""" return f'{self.pdb_id}_{self.auth_chain_id}' def __post_init__(self): if not self.pdb_id.islower() and not self.pdb_id.isdigit(): raise ValueError(f'pdb_id must be lowercase {self.pdb_id}') if not (0 <= self.start_index <= self.end_index): raise ValueError( 'Start must be non-negative and less than or equal to end index. ' f'Range: {self.start_index}-{self.end_index}' ) if len(self.matching_sequence) != (self.end_index - self.start_index): raise ValueError( 'Sequence length must be equal to end_index - start_index. ' f'{len(self.matching_sequence)} != {self.end_index} - ' f'{self.start_index}' ) if self.full_length < 0: raise ValueError(f'Full length must be non-negative: {self.full_length}') def keep( self, *, release_date_cutoff: datetime.date | None, max_subsequence_ratio: float | None, min_hit_length: int | None, min_align_ratio: float | None, ) -> bool: """Returns whether the hit should be kept. In addition to filtering on all of the provided parameters, this method also excludes hits with unresolved residues. Args: release_date_cutoff: Maximum release date of the template. max_subsequence_ratio: If set, excludes hits which are an exact subsequence of the query sequence, and longer than this ratio. Useful to avoid ground truth leakage. min_hit_length: If set, excludes hits which have fewer residues than this. min_align_ratio: If set, excludes hits where the number of residues aligned to the query is less than this proportion of the template length. """ # Exclude hits which are too recent. if ( release_date_cutoff is not None and self.release_date > release_date_cutoff ): return False # Exclude hits which are large duplicates of the query_sequence. if ( max_subsequence_ratio is not None and self.length_ratio > max_subsequence_ratio ): if self.matching_sequence in self.query_sequence: return False # Exclude hits which are too short. if ( min_hit_length is not None and len(self.matching_sequence) < min_hit_length ): return False # Exclude hits with unresolved residues. if not self.is_valid: return False # Exclude hits with too few alignments. try: if min_align_ratio is not None and self.align_ratio <= min_align_ratio: return False except template_realign.AlignmentError as e: logging.warning('Failed to align %s: %s', self, str(e)) return False return True def _filter_hits( hits: Iterable[Hit], release_date_cutoff: datetime.date, max_subsequence_ratio: float | None, min_align_ratio: float | None, min_hit_length: int | None, deduplicate_sequences: bool, max_hits: int | None, ) -> Sequence[Hit]: """Filters hits based on the filter config.""" filtered_hits = [] seen_before = set() for hit in hits: if not hit.keep( max_subsequence_ratio=max_subsequence_ratio, min_align_ratio=min_align_ratio, min_hit_length=min_hit_length, release_date_cutoff=release_date_cutoff, ): continue # Remove duplicate templates, keeping the first. if deduplicate_sequences: if hit.output_templates_sequence in seen_before: continue seen_before.add(hit.output_templates_sequence) filtered_hits.append(hit) if max_hits and len(filtered_hits) == max_hits: break return filtered_hits @dataclasses.dataclass(init=False) class Templates: """A container for templates that were found for the given query sequence. The structure_store is constructed from the config by default. Callers can optionally supply a structure_store to the constructor to avoid the cost of construction and metadata loading. """ def __init__( self, *, query_sequence: str, hits: Sequence[Hit], max_template_date: datetime.date, structure_store: structure_stores.StructureStore, query_release_date: datetime.date | None = None, ): self._query_sequence = query_sequence self._hits = tuple(hits) self._max_template_date = max_template_date self._query_release_date = query_release_date self._hit_structures = {} self._structure_store = structure_store if any(h.query_sequence != self._query_sequence for h in self.hits): raise ValueError('All hits must match the query sequence.') if self._hits: chain_poly_type = self._hits[0].chain_poly_type if any(h.chain_poly_type != chain_poly_type for h in self.hits): raise ValueError('All hits must have the same chain_poly_type.') @classmethod def from_seq_and_a3m( cls, *, query_sequence: str, msa_a3m: str, max_template_date: datetime.date, database_path: os.PathLike[str] | str, hmmsearch_config: msa_config.HmmsearchConfig, max_a3m_query_sequences: int | None, structure_store: structure_stores.StructureStore, filter_config: msa_config.TemplateFilterConfig | None = None, query_release_date: datetime.date | None = None, chain_poly_type: str = mmcif_names.PROTEIN_CHAIN, ) -> Self: """Creates templates from a run of hmmsearch tool against a custom a3m. Args: query_sequence: The polymer sequence of the target query. msa_a3m: An a3m of related polymers aligned to the query sequence, this is used to create an HMM for the hmmsearch run. max_template_date: This is used to filter templates for training, ensuring that they do not leak ground truth information used in testing sets. database_path: A path to the sequence database to search for templates. hmmsearch_config: Config with Hmmsearch settings. max_a3m_query_sequences: The maximum number of input MSA sequences to use to construct the profile which is then used to search for templates. structure_store: Structure store to fetch template structures from. filter_config: Optional config that controls which and how many hits to keep. More performant than constructing and then filtering. If not provided, no filtering is done. query_release_date: The release_date of the template query, this is used to filter templates for training, ensuring that they do not leak structure information from the future. chain_poly_type: The polymer type of the templates. Returns: Templates object containing a list of Hits initialised from the structure_store metadata and a3m alignments. """ hmmsearch_a3m = run_hmmsearch_with_a3m( database_path=database_path, hmmsearch_config=hmmsearch_config, max_a3m_query_sequences=max_a3m_query_sequences, a3m=msa_a3m, ) return cls.from_hmmsearch_a3m( query_sequence=query_sequence, a3m=hmmsearch_a3m, max_template_date=max_template_date, query_release_date=query_release_date, chain_poly_type=chain_poly_type, structure_store=structure_store, filter_config=filter_config, ) @classmethod def from_hmmsearch_a3m( cls, *, query_sequence: str, a3m: str, max_template_date: datetime.date, structure_store: structure_stores.StructureStore, filter_config: msa_config.TemplateFilterConfig | None = None, query_release_date: datetime.date | None = None, chain_poly_type: str = mmcif_names.PROTEIN_CHAIN, ) -> Self: """Creates Templates from a Hmmsearch A3M. Args: query_sequence: The polymer sequence of the target query. a3m: Results of Hmmsearch in A3M format. This provides a list of potential template alignments and pdb codes. max_template_date: This is used to filter templates for training, ensuring that they do not leak ground truth information used in testing sets. structure_store: Structure store to fetch template structures from. filter_config: Optional config that controls which and how many hits to keep. More performant than constructing and then filtering. If not provided, no filtering is done. query_release_date: The release_date of the template query, this is used to filter templates for training, ensuring that they do not leak structure information from the future. chain_poly_type: The polymer type of the templates. Returns: Templates object containing a list of Hits initialised from the structure_store metadata and a3m alignments. """ def hit_generator(a3m: str): if not a3m: return # Hmmsearch could return an empty string if there are no hits. for hit_seq, hit_desc in parsers.lazy_parse_fasta_string(a3m): pdb_id, auth_chain_id, start, end, full_length = _parse_hit_description( hit_desc ) release_date, sequence, unresolved_res_ids = _parse_hit_metadata( structure_store, pdb_id, auth_chain_id ) if unresolved_res_ids is None: continue # seq_unresolved_res_num are 1-based, setting to 0-based indices. unresolved_indices = [i - 1 for i in unresolved_res_ids] yield Hit( pdb_id=pdb_id, auth_chain_id=auth_chain_id, hmmsearch_sequence=hit_seq, structure_sequence=sequence, query_sequence=query_sequence, unresolved_res_indices=unresolved_indices, start_index=start - 1, # Raw value is residue number, not index. end_index=end, full_length=full_length, release_date=datetime.date.fromisoformat(release_date), chain_poly_type=chain_poly_type, ) if filter_config is None: hits = tuple(hit_generator(a3m)) else: hits = _filter_hits( hit_generator(a3m), release_date_cutoff=filter_config.max_template_date, max_subsequence_ratio=filter_config.max_subsequence_ratio, min_align_ratio=filter_config.min_align_ratio, min_hit_length=filter_config.min_hit_length, deduplicate_sequences=filter_config.deduplicate_sequences, max_hits=filter_config.max_hits, ) return Templates( query_sequence=query_sequence, query_release_date=query_release_date, hits=hits, max_template_date=max_template_date, structure_store=structure_store, ) @property def query_sequence(self) -> str: return self._query_sequence @property def hits(self) -> tuple[Hit, ...]: return self._hits @property def query_release_date(self) -> datetime.date | None: return self._query_release_date @property def num_hits(self) -> int: return len(self._hits) @functools.cached_property def release_date_cutoff(self) -> datetime.date: if self.query_release_date is None: return self._max_template_date return min( self._max_template_date, self.query_release_date - datetime.timedelta(days=_DAYS_BEFORE_QUERY_DATE), ) def __repr__(self) -> str: return f'Templates({self.num_hits} hits)' def filter( self, *, max_subsequence_ratio: float | None, min_align_ratio: float | None, min_hit_length: int | None, deduplicate_sequences: bool, max_hits: int | None, ) -> Self: """Returns a new Templates object with only the hits that pass all filters. This also filters on query_release_date and max_template_date. Args: max_subsequence_ratio: If set, excludes hits which are an exact subsequence of the query sequence, and longer than this ratio. Useful to avoid ground truth leakage. min_align_ratio: If set, excludes hits where the number of residues aligned to the query is less than this proportion of the template length. min_hit_length: If set, excludes hits which have fewer residues than this. deduplicate_sequences: Whether to exclude duplicate template sequences, keeping only the first. This can be useful in increasing the diversity of hits especially in the case of homomer hits. max_hits: If set, excludes any hits which exceed this count. """ filtered_hits = _filter_hits( hits=self._hits, release_date_cutoff=self.release_date_cutoff, max_subsequence_ratio=max_subsequence_ratio, min_align_ratio=min_align_ratio, min_hit_length=min_hit_length, deduplicate_sequences=deduplicate_sequences, max_hits=max_hits, ) return Templates( query_sequence=self.query_sequence, query_release_date=self.query_release_date, hits=filtered_hits, max_template_date=self._max_template_date, structure_store=self._structure_store, ) def get_hits_with_structures( self, ) -> Sequence[tuple[Hit, structure.Structure]]: """Returns hits + Structures, Structures filtered to the hit's chain.""" results = [] structures = {struc.name.lower(): struc for struc in self.structures} for hit in self.hits: if not hit.is_valid: raise InvalidTemplateError( 'Hits must be filtered before calling get_hits_with_structures.' ) struc = structures[hit.pdb_id] label_chain_id = struc.polymer_auth_asym_id_to_label_asym_id().get( hit.auth_chain_id ) results.append((hit, struc.filter(chain_id=label_chain_id))) return results def featurize( self, include_ligand_features: bool = True, ) -> TemplateFeatures: """Featurises the templates and returns a map of feature names to features. NB: If you don't do any prefiltering, this method might be slow to run as it has to fetch many CIFs and featurize them all. Args: include_ligand_features: Whether to compute ligand features. Returns: Template features: A mapping of template feature labels to features, which may be numpy arrays, bytes objects, or for the special case of label `ligand_features` (if `include_ligand_features` is True), a nested feature map of labels to numpy arrays. Raises: InvalidTemplateError: If hits haven't been filtered before featurization. """ hits_by_pdb_id = {} for idx, hit in enumerate(self.hits): if not hit.is_valid: raise InvalidTemplateError( f'Hits must be filtered before featurizing, got unprocessed {hit=}' ) hits_by_pdb_id.setdefault(hit.pdb_id, []).append((idx, hit)) unsorted_features = [] for struc in self.structures: pdb_id = str(struc.name).lower() for idx, hit in hits_by_pdb_id[pdb_id]: try: label_chain_id = struc.polymer_auth_asym_id_to_label_asym_id()[ hit.auth_chain_id ] hit_features = { **get_polymer_features( chain=struc.filter(chain_id=label_chain_id), chain_poly_type=hit.chain_poly_type, query_sequence_length=len(hit.query_sequence), query_to_hit_mapping=hit.query_to_hit_mapping, ), } if include_ligand_features: hit_features['ligand_features'] = _get_ligand_features(struc) unsorted_features.append((idx, hit_features)) except Error as e: raise type(e)(f'Failed to featurise {hit=}') from e sorted_features = sorted(unsorted_features, key=lambda x: x[0]) sorted_features = [feat for _, feat in sorted_features] return package_template_features( hit_features=sorted_features, include_ligand_features=include_ligand_features, ) @property def structures(self) -> Iterator[structure.Structure]: """Yields template structures for each unique PDB ID among hits. If there are multiple hits in the same Structure, the Structure will be included only once by this method. Yields: A Structure object for each unique PDB ID among hits. Raises: HitDateError: If template's release date exceeds max cutoff date. """ for hit in self.hits: if hit.release_date > self.release_date_cutoff: # pylint: disable=comparison-with-callable raise HitDateError( f'Invalid release date for hit {hit.pdb_id=}, when release date ' f'cutoff is {self.release_date_cutoff}.' ) # Get the set of pdbs to load. In particular, remove duplicate PDB IDs. targets_to_load = tuple({hit.pdb_id for hit in self.hits}) for target_name in targets_to_load: yield structure.from_mmcif( mmcif_string=self._structure_store.get_mmcif_str(target_name), fix_mse_residues=True, fix_arginines=True, include_water=False, include_bonds=False, include_other=True, # For non-standard polymer chains. ) def _parse_hit_description(description: str) -> tuple[str, str, int, int, int]: """Parses the hmmsearch A3M sequence description line.""" # Example lines (protein, nucleic, no description): # >4pqx_A/2-217 [subseq from] mol:protein length:217 Free text # >4pqx_A/2-217 [subseq from] mol:na length:217 Free text # >5g3r_A/1-55 [subseq from] mol:protein length:352 if match := re.fullmatch(_HIT_DESCRIPTION_REGEX, description): return ( match['pdb_id'], match['chain_id'], int(match['start']), int(match['end']), int(match['length']), ) else: raise ValueError(f'Could not parse description "{description}"') def _parse_hit_metadata( structure_store: structure_stores.StructureStore, pdb_id: str, auth_chain_id: str, ) -> tuple[Any, str | None, Sequence[int] | None]: """Parse hit metadata by parsing mmCIF from structure store.""" try: cif = mmcif.from_string(structure_store.get_mmcif_str(pdb_id)) except structure_stores.NotFoundError: logging.warning( 'Failed to get mmCIF for %s (author chain %s).', pdb_id, auth_chain_id ) return None, None, None release_date = mmcif.get_release_date(cif) try: struc = structure.from_parsed_mmcif( cif, model_id=structure.ModelID.ALL, include_water=True, include_other=True, include_bonds=False, ) except ValueError: struc = structure.from_parsed_mmcif( cif, model_id=structure.ModelID.FIRST, include_water=True, include_other=True, include_bonds=False, ) sequence = struc.polymer_author_chain_single_letter_sequence( include_missing_residues=True, protein=True, dna=True, rna=True, other=True, )[auth_chain_id] unresolved_res_ids = struc.filter( chain_auth_asym_id=auth_chain_id ).unresolved_residues.id return release_date, sequence, unresolved_res_ids def get_polymer_features( *, chain: structure.Structure, chain_poly_type: str, query_sequence_length: int, query_to_hit_mapping: Mapping[int, int], ) -> Mapping[str, Any]: """Returns features for this polymer chain. Args: chain: Structure object representing the template. Must be already filtered to a single chain. chain_poly_type: The chain polymer type (protein, DNA, RNA). query_sequence_length: The length of the query sequence. query_to_hit_mapping: 0-based query index to hit index mapping. Returns: A dictionary with polymer features for template_chain_id in the struc. Raises: ValueError: If the input structure contains more than just a single chain. """ # if len(chain.polymer_auth_asym_id_to_label_asym_id()) != 1: # raise ValueError('The structure must be filtered to a single chain.') if chain.name is None: raise ValueError('Template structure must have a name.') if chain.release_date is None: raise ValueError( f'Template structure {chain.name} must have a release date. You can do' ' this by setting "_pdbx_audit_revision_history.revision_date" in the' ' template mmCIF to a date in the ISO-8601 format (e.g. 1989-11-17).' ) num_polymer_chains = len(chain.polymer_auth_asym_id_to_label_asym_id()) if num_polymer_chains != 1: raise ValueError( f'Template structure {chain.name} must be filtered to a single polymer' f' chain but got a structure with {num_polymer_chains} polymer chains.' ) auth_chain_id, label_chain_id = next( iter(chain.polymer_auth_asym_id_to_label_asym_id().items()) ) chain_sequence = chain.chain_single_letter_sequence()[label_chain_id] polymer = _POLYMERS[chain_poly_type] res_arrays = chain.to_res_arrays( include_missing_residues=True, atom_order=polymer.atom_order ) positions = res_arrays.atom_positions positions_mask = res_arrays.atom_mask template_all_atom_positions = np.zeros( (query_sequence_length, polymer.num_atom_types, 3), dtype=np.float64 ) template_all_atom_masks = np.zeros( (query_sequence_length, polymer.num_atom_types), dtype=np.int64 ) template_sequence = ['-'] * query_sequence_length for query_index, template_index in query_to_hit_mapping.items(): template_all_atom_positions[query_index] = positions[template_index] template_all_atom_masks[query_index] = positions_mask[template_index] template_sequence[query_index] = chain_sequence[template_index] template_sequence = ''.join(template_sequence) template_aatype = _encode_restype(chain_poly_type, template_sequence) template_name = f'{chain.name.lower()}_{auth_chain_id}' release_date = chain.release_date.strftime('%Y-%m-%d') return { 'template_all_atom_positions': template_all_atom_positions, 'template_all_atom_masks': template_all_atom_masks, 'template_sequence': template_sequence.encode(), 'template_aatype': np.array(template_aatype, dtype=np.int32), 'template_domain_names': np.array(template_name.encode(), dtype=object), 'template_release_date': np.array(release_date.encode(), dtype=object), } def _get_ligand_features( struc: structure.Structure, ) -> Mapping[str, Mapping[str, np.ndarray | bytes]]: """Returns features for the ligands in this structure.""" ligand_struc = struc.filter_to_entity_type(ligand=True) assert ligand_struc.coords is not None assert ligand_struc.atom_name is not None assert ligand_struc.atom_occupancy is not None ligand_features = {} for ligand_chain_id in ligand_struc.chains: idxs = np.where(ligand_struc.chain_id == ligand_chain_id)[0] if idxs.shape[0]: ligand_features[ligand_chain_id] = { 'ligand_atom_positions': ( ligand_struc.coords[idxs, :].astype(np.float32) ), 'ligand_atom_names': ligand_struc.atom_name[idxs].astype(object), 'ligand_atom_occupancies': ( ligand_struc.atom_occupancy[idxs].astype(np.float32) ), 'ccd_id': ligand_struc.res_name[idxs][0].encode(), } return ligand_features def package_template_features( *, hit_features: Sequence[Mapping[str, Any]], include_ligand_features: bool, ) -> Mapping[str, Any]: """Stacks polymer features, adds empty and keeps ligand features unstacked.""" features_to_include = set(_POLYMER_FEATURES) if include_ligand_features: features_to_include.update(_LIGAND_FEATURES) features = { feat: [single_hit_features[feat] for single_hit_features in hit_features] for feat in features_to_include } stacked_features = {} for k, v in features.items(): if k in _POLYMER_FEATURES: v = np.stack(v, axis=0) if v else np.array([], dtype=_POLYMER_FEATURES[k]) stacked_features[k] = v return stacked_features def _resolve_path(path: os.PathLike[str] | str) -> str: """Resolves path for data dep paths, stringifies otherwise.""" # Data dependency paths: db baked into the binary. resolved_path = resources.filename(path) if os.path.exists(resolved_path): return resolved_path else: # Other paths, e.g. local. return str(path) def run_hmmsearch_with_a3m( *, database_path: os.PathLike[str] | str, hmmsearch_config: msa_config.HmmsearchConfig, max_a3m_query_sequences: int | None, a3m: str | None, ) -> str: """Runs Hmmsearch to get a3m string of hits.""" searcher = hmmsearch.Hmmsearch( binary_path=hmmsearch_config.hmmsearch_binary_path, hmmbuild_binary_path=hmmsearch_config.hmmbuild_binary_path, database_path=_resolve_path(database_path), e_value=hmmsearch_config.e_value, inc_e=hmmsearch_config.inc_e, dom_e=hmmsearch_config.dom_e, incdom_e=hmmsearch_config.incdom_e, alphabet=hmmsearch_config.alphabet, filter_f1=hmmsearch_config.filter_f1, filter_f2=hmmsearch_config.filter_f2, filter_f3=hmmsearch_config.filter_f3, filter_max=hmmsearch_config.filter_max, ) # STO enables us to annotate query non-gap columns as reference columns. sto = parsers.convert_a3m_to_stockholm(a3m, max_a3m_query_sequences) return searcher.query_with_sto(sto, model_construction='hand')