diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index f4a518b3..4570ad8f 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -9,14 +9,12 @@ import dataclasses import functools import importlib -import inspect import json import logging import os import pathlib import time import typing -from collections import Counter from collections.abc import Sequence from typing import Any, List, Dict, Union, overload @@ -29,16 +27,21 @@ from alphafold3.common import folding_input from alphafold3.constants import chemical_components from alphafold3.data import featurisation -from alphafold3.data import parsers as af3_parsers from alphafold3.model import features, params from alphafold3.model import model from alphafold3.model.components import utils -from alphafold3.structure import mmcif as af3_mmcif from jax import numpy as jnp -from alphafold.common import residue_constants -from alphafold.common.protein import Protein, to_mmcif from alphapulldown.folding_backend.folding_backend import FoldingBackend +from alphapulldown.folding_backend.alphafold3_debug import ( + _write_final_msa_a3m, + _write_translated_msa_debug_artifacts, +) +from alphapulldown.folding_backend.alphafold3_inputs import ( + FoldInputBuilder, + PreparedFold, + _build_output_job_name, +) from alphapulldown.folding_backend.alphafold3_outputs import ( _modelcif_metadata, _start_or_resume_fold, @@ -49,27 +52,7 @@ from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache from alphapulldown.prediction.resume import finish_fold_record from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject -from alphapulldown.utils.af2_to_af3_msa import ( - Af2ToAf3TranslationResult, - msa_rows_and_deletions_to_a3m, - translate_af2_individual_chain_features_to_af3_msas_with_stats, - translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats, -) -from alphapulldown.utils.msa_encoding import ids_to_a3m_af3 from alphapulldown.utils.post_modelling import post_prediction_process_af3 -from alphapulldown.utils.file_handling import read_maybe_xz, resolve_af3_input -from alphapulldown.utils.output_paths import ( - _collapse_repeated_name_fragments, - _compact_existing_compound_name, - _compact_output_job_name, - _json_input_basename, - _normalise_output_name_fragment, - _regions_to_name_fragment, -) -from alphapulldown.utils.feature_metadata import ( - encode_metadata_in_description, - find_feature_metadata, -) # ----------------------------------------------------------------------------- @@ -190,43 +173,6 @@ def extract_structures( ) -def _object_name_fragment(obj: typing.Any) -> str: - """Builds a deterministic output-name fragment for one modelling object.""" - if isinstance(obj, dict) and "json_input" in obj: - fragment = _json_input_basename(obj["json_input"]) - regions = obj.get("regions") - if isinstance(regions, Sequence) and regions: - fragment = f"{fragment}__{_regions_to_name_fragment(regions)}" - return _normalise_output_name_fragment(fragment) - - if isinstance(obj, MultimericObject): - return _compact_existing_compound_name(obj.description or "multimer") - - if isinstance(obj, (MonomericObject, ChoppedObject)): - return _normalise_output_name_fragment(obj.description or "monomer") - - if isinstance(obj, folding_input.Input): - return _compact_existing_compound_name(obj.name) - - return _normalise_output_name_fragment(type(obj).__name__) - - -def _build_output_job_name(objects_to_model: Sequence[dict]) -> str: - """Builds a readable AF3 job name from the requested modelling objects.""" - fragments: list[str] = [] - for entry in objects_to_model: - object_to_model = entry["object"] - if isinstance(object_to_model, list): - fragments.extend(_object_name_fragment(obj) for obj in object_to_model) - else: - fragments.append(_object_name_fragment(object_to_model)) - fragments = [fragment for fragment in fragments if fragment] - if not fragments: - return "ranked_0" - readable_name = "_and_".join(_collapse_repeated_name_fragments(fragments)) - return _compact_output_job_name(readable_name) - - def _record_inference_kernels( output_dir: os.PathLike[str] | str, seed: int, @@ -301,37 +247,15 @@ def predict_structure( all_inference_start_time = time.time() all_inference_results = [] - # Utility: write final complex MSA (after pairing/dedup/merge) to A3M - def _write_final_msa_a3m(example_batch, seed_value: int): - try: - if output_dir is None: - return - os.makedirs(output_dir, exist_ok=True) - - def write_from_array(rows: np.ndarray, suffix: str): - # AF3 uses a different integer alphabet for MSA arrays. - a3m_text = ids_to_a3m_af3(rows) - a3m_path = os.path.join(output_dir, f"{fold_input.sanitised_name()}_seed-{seed_value}_{suffix}.a3m") - with open(a3m_path, 'wt') as f: - f.write(a3m_text) - logging.info(f"Wrote {suffix} A3M to {a3m_path}") - - # Write merged msa - msa_rows = example_batch.get('msa', None) - if msa_rows is not None: - num_alignments = int(example_batch.get('num_alignments', 0)) if 'num_alignments' in example_batch else msa_rows.shape[0] - if num_alignments > 0: - write_from_array(np.asarray(msa_rows)[:num_alignments], 'final_complex_msa') - - except Exception as e: - logging.error(f"Failed to write final complex MSA A3M: {e}") for seed, example in zip(fold_input.rng_seeds, featurised_examples): logging.info(f'Running model inference for seed {seed}...') inference_start_time = time.time() # If requested, dump the post-featurisation merged complex MSA for inspection. if debug_msas: - _write_final_msa_a3m(example, seed) + _write_final_msa_a3m( + example, seed, output_dir=output_dir, job_name=fold_input.sanitised_name() + ) rng_key = jax.random.PRNGKey(seed) result = model_runner.run_inference(example, rng_key) @@ -577,1229 +501,57 @@ def prepare_input( features_directory: str = None, debug_templates: bool | None = None, debug_msas: bool = False, - ) -> list: + ) -> list[PreparedFold]: """Prepare input for AlphaFold3 prediction.""" logging.info(f"prepare_input called with {len(objects_to_model)} objects to model") for i, entry in enumerate(objects_to_model): logging.info(f"Object {i}: {entry}") - @functools.lru_cache(maxsize=None) - def _description_with_af2_feature_metadata( - output_description: str, - source_description: str, - ) -> str: - try: - metadata = find_feature_metadata( - source_description, features_directory - ) - except (OSError, ValueError) as exc: - logging.warning( - "Could not read feature metadata for %s: %s", - source_description, - exc, - ) - return output_description - if metadata is None: - return output_description - return encode_metadata_in_description(output_description, metadata) - - @dataclasses.dataclass(frozen=True, slots=True) - class _TranslatedMsaDebugRecord: - chain_id: str - chain_description: str - chain_length: int - paired_msa: str - unpaired_msa: str - paired_msa_row_count: int - unpaired_msa_row_count: int - paired_species_identifier_count: int - paired_rows_without_species_identifier_count: int - paired_rows_with_generated_accession_count: int - - @functools.lru_cache(maxsize=None) - def _supported_chain_kwargs(chain_cls: type) -> frozenset[str]: - return frozenset(inspect.signature(chain_cls).parameters) - - def _construct_chain(chain_cls: type, **kwargs): - supported_kwargs = _supported_chain_kwargs(chain_cls) - filtered_kwargs = { - key: value - for key, value in kwargs.items() - if key in supported_kwargs - } - return chain_cls(**filtered_kwargs) - - def _chain_input_sequence( - chain: ( - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - ), - ) -> str: - canonical_sequence = getattr(chain, "_sequence", None) - if isinstance(canonical_sequence, str): - return canonical_sequence - return chain.sequence - - def _clone_chain_with_id(chain, new_id: str): - description = getattr(chain, "description", None) - if isinstance(chain, folding_input.ProteinChain): - return _construct_chain( - folding_input.ProteinChain, - id=new_id, - sequence=_chain_input_sequence(chain), - description=description, - residue_ids=getattr(chain, "residue_ids", None), - ptms=chain.ptms, - paired_msa=chain.paired_msa, - unpaired_msa=chain.unpaired_msa, - templates=chain.templates, - ) - if isinstance(chain, folding_input.RnaChain): - return _construct_chain( - folding_input.RnaChain, - id=new_id, - sequence=_chain_input_sequence(chain), - description=description, - residue_ids=getattr(chain, "residue_ids", None), - modifications=chain.modifications, - unpaired_msa=chain.unpaired_msa, - ) - if isinstance(chain, folding_input.DnaChain): - return _construct_chain( - folding_input.DnaChain, - id=new_id, - sequence=_chain_input_sequence(chain), - description=description, - residue_ids=getattr(chain, "residue_ids", None), - modifications=chain.modifications(), - ) - if isinstance(chain, folding_input.Ligand): - return _construct_chain( - folding_input.Ligand, - id=new_id, - ccd_ids=chain.ccd_ids, - smiles=chain.smiles, - description=description, - ) - raise TypeError(f"Unsupported chain type: {type(chain)}") - - def _adjacent_duplicate_keep_indices( - residue_ids: Sequence[int] | None, - ) -> np.ndarray | None: - if residue_ids is None: - return None - residue_ids_array = np.asarray(residue_ids, dtype=np.int32).reshape(-1) - if residue_ids_array.size <= 1: - return None - keep_mask = np.ones(residue_ids_array.shape[0], dtype=bool) - keep_mask[1:] = residue_ids_array[1:] != residue_ids_array[:-1] - if np.all(keep_mask): - return None - return np.flatnonzero(keep_mask) - - def _slice_sequence_like( - value: str | bytes | bytearray, - keep_indices: np.ndarray, - *, - tolerate_shorter_input: bool = False, - context: str = "sequence", - ) -> str | bytes: - as_text = ( - value.decode("utf-8") - if isinstance(value, (bytes, bytearray)) - else str(value) - ) - if keep_indices.size: - max_index = int(np.max(keep_indices)) - if max_index >= len(as_text): - if tolerate_shorter_input: - logging.warning( - "Leaving %s unsliced because length %d is shorter " - "than requested keep index %d", - context, - len(as_text), - max_index, - ) - return value - raise IndexError( - f"{context} length {len(as_text)} is shorter than keep index {max_index}" - ) - sliced = "".join(as_text[int(index)] for index in keep_indices.tolist()) - if isinstance(value, (bytes, bytearray)): - return sliced.encode("utf-8") - return sliced - - def _normalise_monomeric_inputs_for_af3( - mono_obj: Union[MonomericObject, ChoppedObject], - ) -> tuple[str, Dict[str, Any], list[int] | None, bool]: - sequence = mono_obj.sequence - feature_dict = mono_obj.feature_dict - residue_index = feature_dict.get("residue_index") - if residue_index is None: - return sequence, feature_dict, None, False - - residue_ids = ( - np.asarray(residue_index, dtype=np.int32).reshape(-1) + 1 - ).astype(int).tolist() - keep_indices = _adjacent_duplicate_keep_indices(residue_ids) - if keep_indices is None: - return sequence, feature_dict, residue_ids, False - - if len(sequence) != len(residue_ids): - logging.warning( - "Skipping adjacent-duplicate AF3 normalization for %s because " - "sequence length %d != residue ID count %d", - mono_obj.description, - len(sequence), - len(residue_ids), - ) - return sequence, feature_dict, residue_ids, False - - normalized_feature_dict = dict(feature_dict) - normalized_sequence = _slice_sequence_like(sequence, keep_indices) - if not isinstance(normalized_sequence, str): - normalized_sequence = normalized_sequence.decode("utf-8") - normalized_residue_ids = [residue_ids[int(index)] for index in keep_indices.tolist()] - - for key in ("aatype", "between_segment_residues", "residue_index"): - if key in normalized_feature_dict: - normalized_feature_dict[key] = np.asarray( - normalized_feature_dict[key] - )[keep_indices] - - for key in ( - "msa", - "deletion_matrix_int", - "deletion_matrix", - "msa_all_seq", - "deletion_matrix_int_all_seq", - "deletion_matrix_all_seq", - ): - if key in normalized_feature_dict: - normalized_feature_dict[key] = np.asarray( - normalized_feature_dict[key] - )[:, keep_indices] - - for key in ( - "template_aatype", - "template_all_atom_masks", - "template_confidence_scores", - ): - if key in normalized_feature_dict: - normalized_feature_dict[key] = np.asarray( - normalized_feature_dict[key] - )[:, keep_indices] - - if "template_all_atom_positions" in normalized_feature_dict: - normalized_feature_dict["template_all_atom_positions"] = np.asarray( - normalized_feature_dict["template_all_atom_positions"] - )[:, keep_indices, :, :] - - if "template_sequence" in normalized_feature_dict: - normalized_feature_dict["template_sequence"] = np.array( - [ - _slice_sequence_like( - template_sequence, - keep_indices, - tolerate_shorter_input=True, - context="template_sequence", - ) - for template_sequence in normalized_feature_dict["template_sequence"] - ], - dtype=object, - ) - - if "sequence" in normalized_feature_dict: - normalized_feature_dict["sequence"] = np.array( - [normalized_sequence.encode("utf-8")] - ) - if "seq_length" in normalized_feature_dict: - normalized_feature_dict["seq_length"] = np.full( - len(normalized_sequence), - len(normalized_sequence), - dtype=np.int32, - ) - if "num_alignments" in normalized_feature_dict: - num_alignments = int( - np.asarray(normalized_feature_dict["num_alignments"]).reshape(-1)[0] - ) - normalized_feature_dict["num_alignments"] = np.full( - len(normalized_sequence), - num_alignments, - dtype=np.int32, - ) - - logging.info( - "Collapsed %d adjacent duplicate residue position(s) for %s to " - "satisfy AF3 input constraints", - len(residue_ids) - len(normalized_residue_ids), - mono_obj.description, - ) - return ( - normalized_sequence, - normalized_feature_dict, - normalized_residue_ids, - True, - ) - - @functools.lru_cache(maxsize=None) - def _validate_json_template_mmcif( - mmcif_string: str, - ) -> tuple[bool, str | None]: - try: - af3_mmcif.from_string(mmcif_string) - except Exception as exc: - return False, str(exc) - return True, None - - def _sanitize_json_input_chain_templates( - chain: ( - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - | folding_input.Ligand - ), - *, - json_path: str, - ): - if not isinstance(chain, folding_input.ProteinChain): - return chain - if chain.templates is None: - return chain - - valid_templates = [] - dropped_template_count = 0 - for template_index, template in enumerate(chain.templates): - is_valid, error_message = _validate_json_template_mmcif( - template.mmcif - ) - if is_valid: - valid_templates.append(template) - continue - - dropped_template_count += 1 - logging.warning( - "Skipping invalid template %d for JSON input %s chain %s: %s", - template_index, - json_path, - chain.id, - error_message, - ) - - if dropped_template_count == 0: - return chain - - logging.info( - "Kept %d/%d template(s) for JSON input %s chain %s after validation", - len(valid_templates), - len(chain.templates), - json_path, - chain.id, - ) - return _construct_chain( - folding_input.ProteinChain, - id=chain.id, - sequence=_chain_input_sequence(chain), - description=getattr(chain, "description", None), - residue_ids=getattr(chain, "residue_ids", None), - ptms=chain.ptms, - paired_msa=chain.paired_msa, - unpaired_msa=chain.unpaired_msa, - templates=valid_templates, - ) - - def _slice_a3m_row_to_region(a3m_row: str, start: int, end: int) -> str: - query_position = 0 - sliced_chars: list[str] = [] - for char in a3m_row: - if char.islower(): - # Lowercase A3M letters are insertions after the current - # query position, so keep only those between retained - # residues and drop insertions before `start` or after `end`. - if start <= query_position < end: - sliced_chars.append(char) - continue - - query_position += 1 - if query_position < start: - continue - if query_position > end: - break - sliced_chars.append(char) - return "".join(sliced_chars) - - def _slice_a3m_to_regions( - a3m_text: str | None, - regions: Sequence[tuple[int, int]], - ) -> str | None: - if a3m_text in (None, "") or not regions: - return a3m_text - - sequences, descriptions = af3_parsers.parse_fasta(a3m_text) - sliced_sequences = [ - "".join( - _slice_a3m_row_to_region(sequence, start, end) - for start, end in regions - ) - for sequence in sequences - ] - return ( - "\n".join( - f">{description}\n{sequence}" - for description, sequence in zip( - descriptions, - sliced_sequences, - strict=True, - ) - ) - + "\n" - ) - - def _slice_positioned_modifications_to_regions( - modifications: Sequence[tuple[str, int]], - regions: Sequence[tuple[int, int]], - ) -> list[tuple[str, int]]: - sliced_modifications: list[tuple[str, int]] = [] - offset = 0 - for start, end in regions: - sliced_modifications.extend( - ( - modification_name, - offset + modification_position - start + 1, - ) - for modification_name, modification_position in modifications - if start <= modification_position <= end - ) - offset += end - start + 1 - return sliced_modifications - - def _slice_templates_to_regions( - templates: Sequence[folding_input.Template] | None, - regions: Sequence[tuple[int, int]], - ) -> Sequence[folding_input.Template] | None: - if templates is None: - return None - - sliced_templates = [] - for template in templates: - remapped_indices = {} - offset = 0 - for start, end in regions: - start_index = start - 1 - remapped_indices.update({ - offset + (query_index - start_index): template_index - for query_index, template_index in template.query_to_template_map.items() - if start_index <= query_index < end - }) - offset += end - start + 1 - if remapped_indices: - sliced_templates.append( - folding_input.Template( - mmcif=template.mmcif, - query_to_template_map=remapped_indices, - ) - ) - return sliced_templates - - def _chain_residue_ids( - chain: ( - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - | folding_input.Ligand - ), - ) -> list[int] | None: - residue_ids = getattr(chain, "residue_ids", None) - if residue_ids is not None: - return [int(residue_id) for residue_id in residue_ids] - if isinstance( - chain, - ( - folding_input.ProteinChain, - folding_input.RnaChain, - folding_input.DnaChain, - ), - ): - return list(range(1, len(_chain_input_sequence(chain)) + 1)) - return None - - def _slice_af3_chain_to_regions( - chain: ( - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - | folding_input.Ligand - ), - *, - regions: Sequence[tuple[int, int]], - json_path: str, - ): - if isinstance(chain, folding_input.Ligand): - raise ValueError( - f"Region ranges are not supported for ligand AF3 JSON inputs: {json_path}" - ) - - chain_sequence = _chain_input_sequence(chain) - sequence_length = len(chain_sequence) - for start, end in regions: - if not 1 <= start <= end <= sequence_length: - raise ValueError( - f"Requested region {start}-{end} is outside the sequence " - f"length {sequence_length} for AF3 JSON input {json_path}." - ) - - sliced_sequence = "".join( - chain_sequence[start - 1:end] - for start, end in regions - ) - sliced_residue_ids = [ - residue_id - for start, end in regions - for residue_id in (_chain_residue_ids(chain) or [])[start - 1:end] - ] - if isinstance(chain, folding_input.ProteinChain): - return _construct_chain( - folding_input.ProteinChain, - id=chain.id, - sequence=sliced_sequence, - description=getattr(chain, "description", None), - ptms=_slice_positioned_modifications_to_regions( - chain.ptms, - regions, - ), - residue_ids=sliced_residue_ids, - unpaired_msa=_slice_a3m_to_regions( - chain.unpaired_msa, - regions, - ), - paired_msa=_slice_a3m_to_regions( - chain.paired_msa, - regions, - ), - templates=_slice_templates_to_regions( - chain.templates, - regions, - ), - ) - - if isinstance(chain, folding_input.RnaChain): - return _construct_chain( - folding_input.RnaChain, - id=chain.id, - sequence=sliced_sequence, - description=getattr(chain, "description", None), - modifications=_slice_positioned_modifications_to_regions( - chain.modifications, - regions, - ), - residue_ids=sliced_residue_ids, - unpaired_msa=_slice_a3m_to_regions( - chain.unpaired_msa, - regions, - ), - ) - - if isinstance(chain, folding_input.DnaChain): - return _construct_chain( - folding_input.DnaChain, - id=chain.id, - sequence=sliced_sequence, - description=getattr(chain, "description", None), - modifications=_slice_positioned_modifications_to_regions( - chain.modifications(), - regions, - ), - residue_ids=sliced_residue_ids, - ) - - raise TypeError(f"Unsupported chain type for AF3 JSON slicing: {type(chain)}") - - def _expand_json_input_chains( - *, - chains: Sequence[ - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - | folding_input.Ligand - ], - json_path: str, - regions: list[tuple[int, int]] | None, - ) -> list[ - folding_input.ProteinChain - | folding_input.RnaChain - | folding_input.DnaChain - | folding_input.Ligand - ]: - sanitized_chains = [ - _sanitize_json_input_chain_templates(chain, json_path=json_path) - for chain in chains - ] - if not regions: - return list(sanitized_chains) - - if len(sanitized_chains) != 1: - raise ValueError( - "Region ranges for AF3 JSON feature inputs require exactly " - f"one chain per file, but {json_path} contains {len(sanitized_chains)} chains." - ) - - merged_chain = _slice_af3_chain_to_regions( - sanitized_chains[0], - regions=regions, - json_path=json_path, - ) - logging.info( - "Collapsed AF3 JSON input %s into one gapped chain for regions %s", - json_path, - regions, - ) - return [merged_chain] - - def get_chain_id(index: int) -> str: - if index < 26: - return chr(ord('A') + index) - else: - first_char = chr(ord('A') + (index // 26) - 1) - second_char = chr(ord('A') + (index % 26)) - return first_char + second_char - - def get_next_available_chain_id(used_chain_ids: set, chain_id_counter_ref) -> str: - """Get the next available chain ID that's not already used.""" - while True: - chain_id = get_chain_id(chain_id_counter_ref[0]) - if chain_id not in used_chain_ids: - chain_id_counter_ref[0] += 1 - return chain_id - chain_id_counter_ref[0] += 1 - - def _write_translated_msa_debug_artifacts( - *, - job_name: str, - output_dir: str, - chain_records: list[_TranslatedMsaDebugRecord], - translation_results: list[Af2ToAf3TranslationResult], - ) -> None: - if not chain_records or not translation_results: - return - - def _extract_paired_species_identifier(description: str) -> str: - if not description.startswith(("tr|", "sp|")): - return "" - description_tail = description.split("|")[-1] - if "_" not in description_tail: - return "" - return description_tail.rsplit("_", maxsplit=1)[-1].strip() - - def _summarise_effective_af3_pairing( - records: list[_TranslatedMsaDebugRecord], - ) -> dict[str, object]: - per_chain_species_counts: dict[str, Counter[str]] = {} - effective_non_gap_rows_by_chain: dict[str, int] = {} - effective_gap_rows_by_chain: dict[str, int] = {} - - for record in records: - _, descriptions = af3_parsers.parse_fasta(record.paired_msa) - species_counts: Counter[str] = Counter() - for description in descriptions[1:]: - species_identifier = _extract_paired_species_identifier(description) - if species_identifier: - species_counts[species_identifier] += 1 - per_chain_species_counts[record.chain_id] = species_counts - effective_non_gap_rows_by_chain[record.chain_id] = 0 - effective_gap_rows_by_chain[record.chain_id] = 0 - - effective_paired_row_count = 0 - effective_paired_row_histogram_by_num_chains: Counter[str] = Counter() - all_species = sorted( - { - species_identifier - for species_counts in per_chain_species_counts.values() - for species_identifier in species_counts - } - ) - for species_identifier in all_species: - present_chain_ids = [ - chain_id - for chain_id, species_counts in per_chain_species_counts.items() - if species_counts.get(species_identifier, 0) > 0 - ] - if len(present_chain_ids) <= 1: - continue - - kept_rows = min( - per_chain_species_counts[chain_id][species_identifier] - for chain_id in present_chain_ids - ) - effective_paired_row_count += kept_rows - effective_paired_row_histogram_by_num_chains[ - str(len(present_chain_ids)) - ] += kept_rows - for chain_id in per_chain_species_counts: - if chain_id in present_chain_ids: - effective_non_gap_rows_by_chain[chain_id] += kept_rows - else: - effective_gap_rows_by_chain[chain_id] += kept_rows - - return { - "effective_paired_row_count": int(effective_paired_row_count), - "effective_paired_row_histogram_by_num_chains": { - key: int(value) - for key, value in sorted( - effective_paired_row_histogram_by_num_chains.items(), - key=lambda item: int(item[0]), - ) - }, - "effective_non_gap_rows_by_chain": { - key: int(value) - for key, value in sorted(effective_non_gap_rows_by_chain.items()) - }, - "effective_gap_rows_by_chain": { - key: int(value) - for key, value in sorted(effective_gap_rows_by_chain.items()) - }, - } - - os.makedirs(output_dir, exist_ok=True) - effective_pairing_summary = _summarise_effective_af3_pairing(chain_records) - summary = { - "job_name": job_name, - "translation_modes": sorted( - {result.translation_mode for result in translation_results} - ), - "total_rows_considered": int( - sum(result.total_rows_considered for result in translation_results) - ), - "occupancy_histogram": { - "0": int( - sum(result.occupancy_histogram.get("0", 0) for result in translation_results) - ), - "1": int( - sum(result.occupancy_histogram.get("1", 0) for result in translation_results) - ), - "ge_2": int( - sum(result.occupancy_histogram.get("ge_2", 0) for result in translation_results) - ), - }, - "paired_row_count": int( - effective_pairing_summary["effective_paired_row_count"] - ), - "translated_paired_input_row_count": int( - sum(result.paired_row_count for result in translation_results) - ), - "effective_paired_row_count": int( - effective_pairing_summary["effective_paired_row_count"] - ), - "effective_paired_row_histogram_by_num_chains": dict( - effective_pairing_summary["effective_paired_row_histogram_by_num_chains"] - ), - "invalid_paired_rows": int( - sum(result.invalid_paired_rows for result in translation_results) - ), - "invalid_unpaired_rows": int( - sum(result.invalid_unpaired_rows for result in translation_results) - ), - "paired_rows_valid": all( - result.invalid_paired_rows == 0 for result in translation_results - ), - "unpaired_rows_valid": all( - result.invalid_unpaired_rows == 0 for result in translation_results - ), - "chains": [], - } - - for record in chain_records: - chain_id = record.chain_id - paired_path = os.path.join( - output_dir, f"{job_name}_chain-{chain_id}_paired_input.a3m" - ) - unpaired_path = os.path.join( - output_dir, f"{job_name}_chain-{chain_id}_unpaired_input.a3m" - ) - with open(paired_path, "wt") as handle: - handle.write(record.paired_msa) - with open(unpaired_path, "wt") as handle: - handle.write(record.unpaired_msa) - - summary["chains"].append( - { - "chain_id": chain_id, - "chain_description": record.chain_description, - "chain_length": int(record.chain_length), - "paired_msa_row_count": int(record.paired_msa_row_count), - "unpaired_msa_row_count": int(record.unpaired_msa_row_count), - "paired_species_identifier_count": int( - record.paired_species_identifier_count - ), - "paired_rows_without_species_identifier_count": int( - record.paired_rows_without_species_identifier_count - ), - "paired_rows_with_generated_accession_count": int( - record.paired_rows_with_generated_accession_count - ), - "effective_paired_msa_row_count": int( - effective_pairing_summary["effective_non_gap_rows_by_chain"].get( - chain_id, 0 - ) - ), - "effective_paired_gap_row_count": int( - effective_pairing_summary["effective_gap_rows_by_chain"].get( - chain_id, 0 - ) - ), - } - ) - - summary_path = os.path.join( - output_dir, f"{job_name}_af2_to_af3_translation_summary.json" - ) - with open(summary_path, "wt") as handle: - json.dump(summary, handle, indent=2, sort_keys=True) - - def insert_release_date_into_mmcif( - mmcif_string: str, revision_date: str = '2100-01-01' - ) -> str: - lines = mmcif_string.splitlines() - insert_index = None - for i, line in enumerate(lines): - if line.startswith('data_'): - insert_index = i + 1 - break - if insert_index is None: - return mmcif_string - revision_lines = [ - "", - "loop_", - "_pdbx_audit_revision_history.ordinal", - "_pdbx_audit_revision_history.data_content_type", - "_pdbx_audit_revision_history.major_revision", - "_pdbx_audit_revision_history.minor_revision", - "_pdbx_audit_revision_history.revision_date", - "1 'Structure model' 1 0 " + revision_date, - "" - ] - lines[insert_index:insert_index] = revision_lines - return '\n'.join(lines) - - def strip_entity_poly_seq_block(mmcif_string: str) -> str: - """Remove _entity_poly_seq loop to let AF3 reconstruct scheme from _atom_site. - - AF3's parser can infer _pdbx_poly_seq_scheme directly from _atom_site and - avoids mismatches when our synthetic _entity_poly_seq has UNK entries. - """ - lines = mmcif_string.splitlines() - out_lines = [] - i = 0 - while i < len(lines): - line = lines[i] - if line.strip() == 'loop_': - # Peek ahead to see if this loop is for _entity_poly_seq - j = i + 1 - is_entity_poly_seq = False - while j < len(lines) and lines[j].strip().startswith('_'): - if lines[j].strip().startswith('_entity_poly_seq.'): - is_entity_poly_seq = True - # Stop header scan when we reach first data row (no leading underscore) - j += 1 - if j < len(lines) and not lines[j].strip().startswith('_'): - break - if is_entity_poly_seq: - # Skip until next blank line or next loop_/category header - k = j - while k < len(lines): - if lines[k].strip() == '' or lines[k].strip() == '#': - k += 1 - break - if lines[k].strip() == 'loop_': - break - k += 1 - i = k - continue - out_lines.append(line) - i += 1 - return '\n'.join(out_lines) - - def _monomeric_to_chain( - mono_obj: Union[MonomericObject, ChoppedObject], - chain_id: str - ) -> folding_input.ProteinChain: - sequence, feature_dict, residue_ids, _ = _normalise_monomeric_inputs_for_af3( - mono_obj - ) - msa_array = feature_dict.get('msa') - deletion_matrix = feature_dict.get('deletion_matrix_int') - if deletion_matrix is None: - deletion_matrix = feature_dict.get('deletion_matrix') - # Standalone AlphaPulldown monomer objects carry only a single custom MSA. - # MultimericObject instances override this with an AF2->AF3 translation. - unpaired_msa = ( - msa_rows_and_deletions_to_a3m( - msa_rows=np.asarray(msa_array), - deletion_rows=None if deletion_matrix is None else np.asarray(deletion_matrix), - query_sequence=sequence, - ) - if msa_array is not None - else "" - ) - paired_msa = "" - templates = [] - if 'template_aatype' in feature_dict: - num_templates = feature_dict['template_aatype'].shape[0] - for i in range(num_templates): - try: - pdb_code_chain = feature_dict["template_domain_names"][i].decode('utf-8') - if '_' in pdb_code_chain: - pdb_code, chain_id_template = pdb_code_chain.split('_') - else: - pdb_code = pdb_code_chain - chain_id_template = 'A' - template_sequence = feature_dict["template_sequence"][i] - if isinstance(template_sequence, bytes): - template_sequence = template_sequence.decode('utf-8') - template_mask = feature_dict["template_all_atom_masks"][i] - hh_ids = np.argmax(feature_dict["template_aatype"][i], axis=-1) - if np.sum(template_mask) == 0: - logging.info(f"Skipping template {i} ({pdb_code_chain}) - no atoms in region") - continue - tmpl_aatype = np.array([ - residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE[j] - for j in hh_ids - ], dtype=np.int32) - tmpl_aatype[tmpl_aatype == 21] = 20 - if np.any(tmpl_aatype < 0) or np.any(tmpl_aatype >= len(residue_constants.resnames)): - raise ValueError(f"Invalid aatype: {tmpl_aatype}") - protein = Protein( - atom_positions=feature_dict["template_all_atom_positions"][i], - atom_mask=feature_dict["template_all_atom_masks"][i], - aatype=tmpl_aatype, - residue_index=feature_dict["residue_index"], - chain_index=np.zeros_like(feature_dict["residue_index"], dtype=int), - b_factors=np.zeros(template_mask.shape, dtype=float), - ) - mmcif_string = to_mmcif( - prot=protein, file_id=pdb_code, model_type='Monomer' - ) - new_mmcif_string = insert_release_date_into_mmcif(mmcif_string) - # Remove _entity_poly_seq to avoid mismatches during AF3 parsing - new_mmcif_string = strip_entity_poly_seq_block(new_mmcif_string) - # Optional debug dump of generated template mmCIF - if debug_templates: - try: - os.makedirs("templates_debug", exist_ok=True) - debug_fname = os.path.join( - "templates_debug", - f"generated_template_{pdb_code}_{chain_id_template}_chain{chain_id}_idx{i}.cif", - ) - with open(debug_fname, "w") as f_dbg: - f_dbg.write(new_mmcif_string) - except Exception: - # Best-effort debug; ignore failures - pass - # Build query->template mapping only for residues that have any atoms - # present in the template (avoid gaps causing index errors downstream). - try: - per_res_atom_mask = feature_dict["template_all_atom_masks"][i] - # Boolean mask: True if any atom present for this residue - present_res_mask = np.any(per_res_atom_mask > 0, axis=-1) - # Map compact template indices (0..N_present-1) to original residue indices - residue_indices_present = [idx for idx, present in enumerate(present_res_mask) if present] - compact_index_by_residue = {res_idx: compact_idx for compact_idx, res_idx in enumerate(residue_indices_present)} - # Query indices are positions in the template sequence aligned to the query - query_to_template_map = {res_idx: compact_index_by_residue[res_idx] for res_idx in residue_indices_present} - except Exception: - # Fallback to identity if masks missing, though this may error later - query_to_template_map = {j: j for j in range(len(template_sequence))} - templates.append( - folding_input.Template( - mmcif=new_mmcif_string, - query_to_template_map=query_to_template_map, - ) - ) - except Exception as e: - logging.error(f"Error processing template {i} for chain {chain_id}: {e}") - try: - os.makedirs("templates_debug", exist_ok=True) - error_filename = os.path.join("templates_debug", f"ERROR_template_{i}_chain_{chain_id}.txt") - with open(error_filename, "w") as error_file: - error_file.write(f"Error: {e}\n") - error_file.write(f"Template index: {i}\n") - error_file.write(f"Chain ID: {chain_id}\n") - if 'template_domain_names' in feature_dict: - error_file.write(f"Domain name: {feature_dict['template_domain_names'][i]}\n") - if 'template_sequence' in feature_dict: - error_file.write(f"Template sequence: {feature_dict['template_sequence'][i]}\n") - except Exception as save_error: - logging.error(f"Failed to save error info: {save_error}") - raise - source_description = getattr( - mono_obj, "monomeric_description", mono_obj.description - ) - chain_description = _description_with_af2_feature_metadata( - mono_obj.description, - source_description, - ) - chain = _construct_chain( - folding_input.ProteinChain, - id=chain_id, - sequence=sequence, - ptms=[], - residue_ids=residue_ids, - description=chain_description, - unpaired_msa=unpaired_msa, - paired_msa=paired_msa, - templates=templates, - ) - return chain - - def _expand_monomeric_object( - mono_obj: Union[MonomericObject, ChoppedObject], - ) -> list[Union[MonomericObject, ChoppedObject]]: - return [mono_obj] - - def _process_single_object(obj, chain_id_counter_ref, used_chain_ids: set): - nonlocal all_chains - if isinstance(obj, dict) and 'json_input' in obj: - json_path = obj['json_input'] - json_regions = obj.get('regions') - logging.info(f"Processing JSON file: {json_path}") - try: - # Accept a compressed feature file: the published database - # ships AF3 inputs as .json.xz, mirroring AF2's .pkl.xz. - resolved = resolve_af3_input(json_path) - if resolved is None: - raise FileNotFoundError(json_path) - if resolved != json_path: - logging.info(f"Using compressed feature file: {resolved}") - json_str = read_maybe_xz(resolved) - input_obj = folding_input.Input.from_json(json_str) - chains_to_add = _expand_json_input_chains( - chains=input_obj.chains, - json_path=json_path, - regions=( - json_regions if isinstance(json_regions, list) and json_regions else None - ), - ) - # Track the chain IDs from the JSON file - logging.info( - "JSON file %s contributes chains with IDs: %s", - json_path, - [chain.id for chain in chains_to_add], - ) - - # Check for duplicate chain IDs and modify them if necessary - modified_chains = [] - for chain in chains_to_add: - original_id = chain.id - new_id = original_id - - # If this chain ID is already used, generate a new one - if original_id in used_chain_ids: - new_id = get_next_available_chain_id(used_chain_ids, chain_id_counter_ref) - logging.info(f"Chain ID '{original_id}' already used, changing to '{new_id}'") - - used_chain_ids.add(new_id) - logging.info(f"Added chain ID '{new_id}' from JSON file {json_path}") - - modified_chains.append(_clone_chain_with_id(chain, new_id)) - - all_chains.extend(modified_chains) - except Exception as e: - logging.error(f"Failed to parse JSON file {json_path}: {e}") - raise - elif isinstance(obj, (MonomericObject, ChoppedObject)): - for expanded_obj in _expand_monomeric_object(obj): - chain_id = get_next_available_chain_id(used_chain_ids, chain_id_counter_ref) - used_chain_ids.add(chain_id) - logging.info( - f"Added chain ID '{chain_id}' for AlphaPulldown object " - f"{expanded_obj.description}" - ) - all_chains.append(_monomeric_to_chain(expanded_obj, chain_id)) - elif isinstance(obj, MultimericObject): - chains = [] - translated_result = None - combined_msa = obj.feature_dict.get('msa') - expanded_interactors = list(obj.interactors) - normalized_interactor_inputs = [ - _normalise_monomeric_inputs_for_af3(interactor) - for interactor in expanded_interactors - ] - has_adjacent_duplicate_residue_ids = any( - normalized_input[3] for normalized_input in normalized_interactor_inputs - ) - - translated_result = ( - translate_af2_individual_chain_features_to_af3_msas_with_stats( - chain_feature_dicts=[ - normalized_input[1] - for normalized_input in normalized_interactor_inputs - ], - chain_sequences=[ - normalized_input[0] - for normalized_input in normalized_interactor_inputs - ], - ) - ) - num_pairable_chains = sum( - chain_stats.paired_species_identifier_count > 0 - for chain_stats in translated_result.chain_stats - ) - if ( - has_adjacent_duplicate_residue_ids - and num_pairable_chains < 2 - and combined_msa is not None - ): - logging.info( - "Skipping AF2 merged-MSA fallback for multimeric object %s " - "because adjacent duplicate residue IDs were collapsed for AF3 input compatibility.", - obj.description, - ) - elif num_pairable_chains < 2 and combined_msa is not None: - # Fall back to the merged AF2 multimer MSA transport path when the - # individual `_all_seq` features do not carry usable species IDs. - translated_result = ( - translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats( - merged_msa=np.asarray(combined_msa), - chain_sequences=[ - interactor.sequence for interactor in obj.interactors - ], - num_alignments=obj.feature_dict.get('num_alignments'), - deletion_matrix=( - obj.feature_dict.get('deletion_matrix_int') - if obj.feature_dict.get('deletion_matrix_int') is not None - else obj.feature_dict.get('deletion_matrix') - ), - asym_id=obj.feature_dict.get('asym_id'), - ) - ) - expanded_interactors = list(obj.interactors) - - for chain_index, interactor in enumerate(expanded_interactors): - chain_id = get_next_available_chain_id(used_chain_ids, chain_id_counter_ref) - used_chain_ids.add(chain_id) - logging.info(f"Added chain ID '{chain_id}' for multimeric interactor") - base_chain = _monomeric_to_chain(interactor, chain_id) - if translated_result is not None: - chain_msas = translated_result.chain_msas[chain_index] - chain_stats = translated_result.chain_stats[chain_index] - base_chain = _construct_chain( - folding_input.ProteinChain, - id=base_chain.id, - sequence=_chain_input_sequence(base_chain), - ptms=base_chain.ptms, - residue_ids=getattr(base_chain, "residue_ids", None), - description=getattr(base_chain, "description", None), - unpaired_msa=chain_msas.unpaired_msa, - paired_msa=chain_msas.paired_msa, - templates=base_chain.templates, - ) - af2_translated_msa_chain_ids.add(base_chain.id) - if debug_msas: - translation_debug_chain_records.append( - _TranslatedMsaDebugRecord( - chain_id=base_chain.id, - chain_description=interactor.description, - chain_length=len(interactor.sequence), - paired_msa=chain_msas.paired_msa, - unpaired_msa=chain_msas.unpaired_msa, - paired_msa_row_count=chain_stats.paired_msa_row_count, - unpaired_msa_row_count=chain_stats.unpaired_msa_row_count, - paired_species_identifier_count=( - chain_stats.paired_species_identifier_count - ), - paired_rows_without_species_identifier_count=( - chain_stats.paired_rows_without_species_identifier_count - ), - paired_rows_with_generated_accession_count=( - chain_stats.paired_rows_with_generated_accession_count - ), - ) - ) - chains.append(base_chain) - if debug_msas and translated_result is not None: - translation_debug_results.append(translated_result) - all_chains.extend(chains) - else: - raise TypeError(f"Unsupported object type for folding input conversion: {type(obj)}") - - prepared_inputs = [] - chain_id_counter = [0] # Use a list to allow pass-by-reference - used_chain_ids = set() # Track used chain IDs - all_chains = [] - # Track chains whose MSAs were translated from AF2 features; they must - # not be rewritten by the promotion heuristic below. - af2_translated_msa_chain_ids: set[str] = set() - translation_debug_chain_records: list[_TranslatedMsaDebugRecord] = [] - translation_debug_results: list[Af2ToAf3TranslationResult] = [] - + # A fresh builder per call: its chain IDs and caches belong to this preparation. + builder = FoldInputBuilder( + features_directory=features_directory, + debug_templates=debug_templates, + debug_msas=debug_msas, + ) for entry in objects_to_model: object_to_model = entry['object'] - output_dir = entry['output_dir'] - # If object_to_model is a list, process each element if isinstance(object_to_model, list): for sub_obj in object_to_model: - _process_single_object(sub_obj, chain_id_counter, used_chain_ids) + builder.add(sub_obj) continue # Done with this entry # Otherwise, process as a single object - _process_single_object(object_to_model, chain_id_counter, used_chain_ids) + builder.add(object_to_model) # Debug: Print all chain IDs before creating the combined input - logging.info(f"All chain IDs before creating combined input: {[chain.id for chain in all_chains]}") - logging.info(f"Used chain IDs set: {used_chain_ids}") + logging.info(f"All chain IDs before creating combined input: {[chain.id for chain in builder.chains]}") + logging.info(f"Used chain IDs set: {builder.used_chain_ids}") # Create a single combined input with all chains - if all_chains: - # Promote unpaired->paired only for chains not constructed from AF2 multimer features. - promoted_chains: list[folding_input.ProteinChain | folding_input.RnaChain | folding_input.DnaChain | folding_input.Ligand] = [] - for ch in all_chains: - if ( - isinstance(ch, folding_input.ProteinChain) - and ch.id not in af2_translated_msa_chain_ids - ): - try: - has_empty_paired = (getattr(ch, 'paired_msa', None) in (None, '')) - has_unpaired = bool(getattr(ch, 'unpaired_msa', '')) - if has_empty_paired and has_unpaired: - new_chain = _construct_chain( - folding_input.ProteinChain, - id=ch.id, - sequence=_chain_input_sequence(ch), - ptms=ch.ptms, - residue_ids=getattr(ch, "residue_ids", None), - description=getattr(ch, "description", None), - paired_msa=ch.unpaired_msa, - unpaired_msa='', - templates=ch.templates, - ) - promoted_chains.append(new_chain) - continue - except Exception: - pass - promoted_chains.append(ch) - - all_chains = promoted_chains - job_name = _build_output_job_name(objects_to_model) - combined_input = folding_input.Input( - name=job_name, - rng_seeds=[random_seed], - chains=all_chains, + if not builder.chains: + return [] + combined_input = builder.build( + job_name=_build_output_job_name(objects_to_model), random_seed=random_seed + ) + # Use the output directory from the first object + first_output_dir = objects_to_model[0]['output_dir'] + if debug_msas: + _write_translated_msa_debug_artifacts( + job_name=combined_input.sanitised_name(), + output_dir=first_output_dir, + chain_records=builder.translation_debug_chain_records, + translation_results=builder.translation_debug_results, ) - # Use the output directory from the first object - first_output_dir = objects_to_model[0]['output_dir'] - if debug_msas: - _write_translated_msa_debug_artifacts( - job_name=combined_input.sanitised_name(), - output_dir=first_output_dir, - chain_records=translation_debug_chain_records, - translation_results=translation_debug_results, - ) - # Disable overlap resolution when we provide translated AF2 multimer MSAs. - disable_overlaps = len(af2_translated_msa_chain_ids) > 0 - prepared_inputs.append({combined_input: (first_output_dir, not disable_overlaps)}) - - return prepared_inputs + # Disable overlap resolution when we provide translated AF2 multimer MSAs. + disable_overlaps = len(builder.af2_translated_msa_chain_ids) > 0 + return [ + PreparedFold( + fold_input=combined_input, + output_dir=first_output_dir, + resolve_msa_overlaps=not disable_overlaps, + ) + ] @staticmethod def predict( @@ -1837,19 +589,10 @@ def predict( debug_msas=kwargs.get("debug_msas", False), ) # Run predictions - for mapping in prepared_inputs: - logging.info(f"Processing mapping: {mapping}") - if len(mapping) != 1: - raise ValueError(f"Expected exactly one item in mapping, got {len(mapping)}: {mapping}") - fold_input_obj, mapping_value = next(iter(mapping.items())) - if isinstance(mapping_value, tuple): - output_dir, resolve_overlaps_flag = mapping_value - else: - output_dir = mapping_value - resolve_overlaps_flag = True - + for prepared in prepared_inputs: + output_dir = prepared.output_dir started = _start_or_resume_fold( - fold_input_obj, output_dir, random_seed, allow_resume, kwargs, postprocess_flags + prepared.fold_input, output_dir, random_seed, allow_resume, kwargs, postprocess_flags ) if started is None: continue @@ -1864,7 +607,7 @@ def predict( model_runner=model_runner, output_dir=output_dir, buckets=buckets, - resolve_msa_overlaps=resolve_overlaps_flag, + resolve_msa_overlaps=prepared.resolve_msa_overlaps, debug_msas=kwargs.get('debug_msas', False), seeds_done=seeds_done, modelcif_metadata=( diff --git a/alphapulldown/folding_backend/alphafold3_debug.py b/alphapulldown/folding_backend/alphafold3_debug.py new file mode 100644 index 00000000..911d2821 --- /dev/null +++ b/alphapulldown/folding_backend/alphafold3_debug.py @@ -0,0 +1,250 @@ +""" +Debug artifacts of AlphaFold 3 input preparation and inference: the MSAs translated from +AF2 features with their summary, and each seed's final complex MSA. + +Copyright (c) 2024 European Molecular Biology Laboratory + +Author: Dmitry Molodenskiy +""" + +import dataclasses +import json +import logging +import os +from collections import Counter + +import numpy as np +from alphafold3.data import parsers as af3_parsers + +from alphapulldown.utils.af2_to_af3_msa import Af2ToAf3TranslationResult +from alphapulldown.utils.msa_encoding import ids_to_a3m_af3 + + +@dataclasses.dataclass(frozen=True, slots=True) +class _TranslatedMsaDebugRecord: + chain_id: str + chain_description: str + chain_length: int + paired_msa: str + unpaired_msa: str + paired_msa_row_count: int + unpaired_msa_row_count: int + paired_species_identifier_count: int + paired_rows_without_species_identifier_count: int + paired_rows_with_generated_accession_count: int + + +def _write_translated_msa_debug_artifacts( + *, + job_name: str, + output_dir: str, + chain_records: list[_TranslatedMsaDebugRecord], + translation_results: list[Af2ToAf3TranslationResult], +) -> None: + if not chain_records or not translation_results: + return + + def _extract_paired_species_identifier(description: str) -> str: + if not description.startswith(("tr|", "sp|")): + return "" + description_tail = description.split("|")[-1] + if "_" not in description_tail: + return "" + return description_tail.rsplit("_", maxsplit=1)[-1].strip() + + def _summarise_effective_af3_pairing( + records: list[_TranslatedMsaDebugRecord], + ) -> dict[str, object]: + per_chain_species_counts: dict[str, Counter[str]] = {} + effective_non_gap_rows_by_chain: dict[str, int] = {} + effective_gap_rows_by_chain: dict[str, int] = {} + + for record in records: + _, descriptions = af3_parsers.parse_fasta(record.paired_msa) + species_counts: Counter[str] = Counter() + for description in descriptions[1:]: + species_identifier = _extract_paired_species_identifier(description) + if species_identifier: + species_counts[species_identifier] += 1 + per_chain_species_counts[record.chain_id] = species_counts + effective_non_gap_rows_by_chain[record.chain_id] = 0 + effective_gap_rows_by_chain[record.chain_id] = 0 + + effective_paired_row_count = 0 + effective_paired_row_histogram_by_num_chains: Counter[str] = Counter() + all_species = sorted( + { + species_identifier + for species_counts in per_chain_species_counts.values() + for species_identifier in species_counts + } + ) + for species_identifier in all_species: + present_chain_ids = [ + chain_id + for chain_id, species_counts in per_chain_species_counts.items() + if species_counts.get(species_identifier, 0) > 0 + ] + if len(present_chain_ids) <= 1: + continue + + kept_rows = min( + per_chain_species_counts[chain_id][species_identifier] + for chain_id in present_chain_ids + ) + effective_paired_row_count += kept_rows + effective_paired_row_histogram_by_num_chains[ + str(len(present_chain_ids)) + ] += kept_rows + for chain_id in per_chain_species_counts: + if chain_id in present_chain_ids: + effective_non_gap_rows_by_chain[chain_id] += kept_rows + else: + effective_gap_rows_by_chain[chain_id] += kept_rows + + return { + "effective_paired_row_count": int(effective_paired_row_count), + "effective_paired_row_histogram_by_num_chains": { + key: int(value) + for key, value in sorted( + effective_paired_row_histogram_by_num_chains.items(), + key=lambda item: int(item[0]), + ) + }, + "effective_non_gap_rows_by_chain": { + key: int(value) + for key, value in sorted(effective_non_gap_rows_by_chain.items()) + }, + "effective_gap_rows_by_chain": { + key: int(value) + for key, value in sorted(effective_gap_rows_by_chain.items()) + }, + } + + os.makedirs(output_dir, exist_ok=True) + effective_pairing_summary = _summarise_effective_af3_pairing(chain_records) + summary = { + "job_name": job_name, + "translation_modes": sorted( + {result.translation_mode for result in translation_results} + ), + "total_rows_considered": int( + sum(result.total_rows_considered for result in translation_results) + ), + "occupancy_histogram": { + "0": int( + sum(result.occupancy_histogram.get("0", 0) for result in translation_results) + ), + "1": int( + sum(result.occupancy_histogram.get("1", 0) for result in translation_results) + ), + "ge_2": int( + sum(result.occupancy_histogram.get("ge_2", 0) for result in translation_results) + ), + }, + "paired_row_count": int( + effective_pairing_summary["effective_paired_row_count"] + ), + "translated_paired_input_row_count": int( + sum(result.paired_row_count for result in translation_results) + ), + "effective_paired_row_count": int( + effective_pairing_summary["effective_paired_row_count"] + ), + "effective_paired_row_histogram_by_num_chains": dict( + effective_pairing_summary["effective_paired_row_histogram_by_num_chains"] + ), + "invalid_paired_rows": int( + sum(result.invalid_paired_rows for result in translation_results) + ), + "invalid_unpaired_rows": int( + sum(result.invalid_unpaired_rows for result in translation_results) + ), + "paired_rows_valid": all( + result.invalid_paired_rows == 0 for result in translation_results + ), + "unpaired_rows_valid": all( + result.invalid_unpaired_rows == 0 for result in translation_results + ), + "chains": [], + } + + for record in chain_records: + chain_id = record.chain_id + paired_path = os.path.join( + output_dir, f"{job_name}_chain-{chain_id}_paired_input.a3m" + ) + unpaired_path = os.path.join( + output_dir, f"{job_name}_chain-{chain_id}_unpaired_input.a3m" + ) + with open(paired_path, "wt") as handle: + handle.write(record.paired_msa) + with open(unpaired_path, "wt") as handle: + handle.write(record.unpaired_msa) + + summary["chains"].append( + { + "chain_id": chain_id, + "chain_description": record.chain_description, + "chain_length": int(record.chain_length), + "paired_msa_row_count": int(record.paired_msa_row_count), + "unpaired_msa_row_count": int(record.unpaired_msa_row_count), + "paired_species_identifier_count": int( + record.paired_species_identifier_count + ), + "paired_rows_without_species_identifier_count": int( + record.paired_rows_without_species_identifier_count + ), + "paired_rows_with_generated_accession_count": int( + record.paired_rows_with_generated_accession_count + ), + "effective_paired_msa_row_count": int( + effective_pairing_summary["effective_non_gap_rows_by_chain"].get( + chain_id, 0 + ) + ), + "effective_paired_gap_row_count": int( + effective_pairing_summary["effective_gap_rows_by_chain"].get( + chain_id, 0 + ) + ), + } + ) + + summary_path = os.path.join( + output_dir, f"{job_name}_af2_to_af3_translation_summary.json" + ) + with open(summary_path, "wt") as handle: + json.dump(summary, handle, indent=2, sort_keys=True) + + +# Utility: write final complex MSA (after pairing/dedup/merge) to A3M +def _write_final_msa_a3m( + example_batch, + seed_value: int, + *, + output_dir: os.PathLike[str] | str | None, + job_name: str, +) -> None: + try: + if output_dir is None: + return + os.makedirs(output_dir, exist_ok=True) + + def write_from_array(rows: np.ndarray, suffix: str): + # AF3 uses a different integer alphabet for MSA arrays. + a3m_text = ids_to_a3m_af3(rows) + a3m_path = os.path.join(output_dir, f"{job_name}_seed-{seed_value}_{suffix}.a3m") + with open(a3m_path, 'wt') as f: + f.write(a3m_text) + logging.info(f"Wrote {suffix} A3M to {a3m_path}") + + # Write merged msa + msa_rows = example_batch.get('msa', None) + if msa_rows is not None: + num_alignments = int(example_batch.get('num_alignments', 0)) if 'num_alignments' in example_batch else msa_rows.shape[0] + if num_alignments > 0: + write_from_array(np.asarray(msa_rows)[:num_alignments], 'final_complex_msa') + + except Exception as e: + logging.error(f"Failed to write final complex MSA A3M: {e}") diff --git a/alphapulldown/folding_backend/alphafold3_inputs.py b/alphapulldown/folding_backend/alphafold3_inputs.py new file mode 100644 index 00000000..35414916 --- /dev/null +++ b/alphapulldown/folding_backend/alphafold3_inputs.py @@ -0,0 +1,1133 @@ +""" +Translates AlphaPulldown objects and AlphaFold 3 JSON inputs into one AlphaFold 3 fold +input: chain construction, monomer normalisation, region slicing of JSON chains, template +mmCIF fixes, chain IDs and the job name. + +Copyright (c) 2024 European Molecular Biology Laboratory + +Author: Dmitry Molodenskiy +""" + +import dataclasses +import functools +import inspect +import logging +import os +import typing +from collections.abc import Sequence +from typing import Any, Dict, Union + +import numpy as np +from alphafold3.common import folding_input +from alphafold3.data import parsers as af3_parsers +from alphafold3.structure import mmcif as af3_mmcif + +from alphafold.common import residue_constants +from alphafold.common.protein import Protein, to_mmcif +from alphapulldown.folding_backend.alphafold3_debug import _TranslatedMsaDebugRecord +from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject +from alphapulldown.utils.af2_to_af3_msa import ( + Af2ToAf3TranslationResult, + msa_rows_and_deletions_to_a3m, + translate_af2_individual_chain_features_to_af3_msas_with_stats, + translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats, +) +from alphapulldown.utils.file_handling import read_maybe_xz, resolve_af3_input +from alphapulldown.utils.output_paths import ( + _collapse_repeated_name_fragments, + _compact_existing_compound_name, + _compact_output_job_name, + _json_input_basename, + _normalise_output_name_fragment, + _regions_to_name_fragment, +) +from alphapulldown.utils.feature_metadata import ( + encode_metadata_in_description, + find_feature_metadata, +) + + +def _object_name_fragment(obj: typing.Any) -> str: + """Builds a deterministic output-name fragment for one modelling object.""" + if isinstance(obj, dict) and "json_input" in obj: + fragment = _json_input_basename(obj["json_input"]) + regions = obj.get("regions") + if isinstance(regions, Sequence) and regions: + fragment = f"{fragment}__{_regions_to_name_fragment(regions)}" + return _normalise_output_name_fragment(fragment) + + if isinstance(obj, MultimericObject): + return _compact_existing_compound_name(obj.description or "multimer") + + if isinstance(obj, (MonomericObject, ChoppedObject)): + return _normalise_output_name_fragment(obj.description or "monomer") + + if isinstance(obj, folding_input.Input): + return _compact_existing_compound_name(obj.name) + + return _normalise_output_name_fragment(type(obj).__name__) + + +def _build_output_job_name(objects_to_model: Sequence[dict]) -> str: + """Builds a readable AF3 job name from the requested modelling objects.""" + fragments: list[str] = [] + for entry in objects_to_model: + object_to_model = entry["object"] + if isinstance(object_to_model, list): + fragments.extend(_object_name_fragment(obj) for obj in object_to_model) + else: + fragments.append(_object_name_fragment(object_to_model)) + fragments = [fragment for fragment in fragments if fragment] + if not fragments: + return "ranked_0" + readable_name = "_and_".join(_collapse_repeated_name_fragments(fragments)) + return _compact_output_job_name(readable_name) + + +def _supported_chain_kwargs(chain_cls: type) -> frozenset[str]: + return frozenset(inspect.signature(chain_cls).parameters) + + +def _construct_chain(chain_cls: type, **kwargs): + supported_kwargs = _supported_chain_kwargs(chain_cls) + filtered_kwargs = { + key: value + for key, value in kwargs.items() + if key in supported_kwargs + } + return chain_cls(**filtered_kwargs) + + +def _chain_input_sequence( + chain: ( + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + ), +) -> str: + canonical_sequence = getattr(chain, "_sequence", None) + if isinstance(canonical_sequence, str): + return canonical_sequence + return chain.sequence + + +def _clone_chain_with_id(chain, new_id: str): + description = getattr(chain, "description", None) + if isinstance(chain, folding_input.ProteinChain): + return _construct_chain( + folding_input.ProteinChain, + id=new_id, + sequence=_chain_input_sequence(chain), + description=description, + residue_ids=getattr(chain, "residue_ids", None), + ptms=chain.ptms, + paired_msa=chain.paired_msa, + unpaired_msa=chain.unpaired_msa, + templates=chain.templates, + ) + if isinstance(chain, folding_input.RnaChain): + return _construct_chain( + folding_input.RnaChain, + id=new_id, + sequence=_chain_input_sequence(chain), + description=description, + residue_ids=getattr(chain, "residue_ids", None), + modifications=chain.modifications, + unpaired_msa=chain.unpaired_msa, + ) + if isinstance(chain, folding_input.DnaChain): + return _construct_chain( + folding_input.DnaChain, + id=new_id, + sequence=_chain_input_sequence(chain), + description=description, + residue_ids=getattr(chain, "residue_ids", None), + modifications=chain.modifications(), + ) + if isinstance(chain, folding_input.Ligand): + return _construct_chain( + folding_input.Ligand, + id=new_id, + ccd_ids=chain.ccd_ids, + smiles=chain.smiles, + description=description, + ) + raise TypeError(f"Unsupported chain type: {type(chain)}") + + +def _adjacent_duplicate_keep_indices( + residue_ids: Sequence[int] | None, +) -> np.ndarray | None: + if residue_ids is None: + return None + residue_ids_array = np.asarray(residue_ids, dtype=np.int32).reshape(-1) + if residue_ids_array.size <= 1: + return None + keep_mask = np.ones(residue_ids_array.shape[0], dtype=bool) + keep_mask[1:] = residue_ids_array[1:] != residue_ids_array[:-1] + if np.all(keep_mask): + return None + return np.flatnonzero(keep_mask) + + +def _slice_sequence_like( + value: str | bytes | bytearray, + keep_indices: np.ndarray, + *, + tolerate_shorter_input: bool = False, + context: str = "sequence", +) -> str | bytes: + as_text = ( + value.decode("utf-8") + if isinstance(value, (bytes, bytearray)) + else str(value) + ) + if keep_indices.size: + max_index = int(np.max(keep_indices)) + if max_index >= len(as_text): + if tolerate_shorter_input: + logging.warning( + "Leaving %s unsliced because length %d is shorter " + "than requested keep index %d", + context, + len(as_text), + max_index, + ) + return value + raise IndexError( + f"{context} length {len(as_text)} is shorter than keep index {max_index}" + ) + sliced = "".join(as_text[int(index)] for index in keep_indices.tolist()) + if isinstance(value, (bytes, bytearray)): + return sliced.encode("utf-8") + return sliced + + +def _normalise_monomeric_inputs_for_af3( + mono_obj: Union[MonomericObject, ChoppedObject], +) -> tuple[str, Dict[str, Any], list[int] | None, bool]: + sequence = mono_obj.sequence + feature_dict = mono_obj.feature_dict + residue_index = feature_dict.get("residue_index") + if residue_index is None: + return sequence, feature_dict, None, False + + residue_ids = ( + np.asarray(residue_index, dtype=np.int32).reshape(-1) + 1 + ).astype(int).tolist() + keep_indices = _adjacent_duplicate_keep_indices(residue_ids) + if keep_indices is None: + return sequence, feature_dict, residue_ids, False + + if len(sequence) != len(residue_ids): + logging.warning( + "Skipping adjacent-duplicate AF3 normalization for %s because " + "sequence length %d != residue ID count %d", + mono_obj.description, + len(sequence), + len(residue_ids), + ) + return sequence, feature_dict, residue_ids, False + + normalized_feature_dict = dict(feature_dict) + normalized_sequence = _slice_sequence_like(sequence, keep_indices) + if not isinstance(normalized_sequence, str): + normalized_sequence = normalized_sequence.decode("utf-8") + normalized_residue_ids = [residue_ids[int(index)] for index in keep_indices.tolist()] + + for key in ("aatype", "between_segment_residues", "residue_index"): + if key in normalized_feature_dict: + normalized_feature_dict[key] = np.asarray( + normalized_feature_dict[key] + )[keep_indices] + + for key in ( + "msa", + "deletion_matrix_int", + "deletion_matrix", + "msa_all_seq", + "deletion_matrix_int_all_seq", + "deletion_matrix_all_seq", + ): + if key in normalized_feature_dict: + normalized_feature_dict[key] = np.asarray( + normalized_feature_dict[key] + )[:, keep_indices] + + for key in ( + "template_aatype", + "template_all_atom_masks", + "template_confidence_scores", + ): + if key in normalized_feature_dict: + normalized_feature_dict[key] = np.asarray( + normalized_feature_dict[key] + )[:, keep_indices] + + if "template_all_atom_positions" in normalized_feature_dict: + normalized_feature_dict["template_all_atom_positions"] = np.asarray( + normalized_feature_dict["template_all_atom_positions"] + )[:, keep_indices, :, :] + + if "template_sequence" in normalized_feature_dict: + normalized_feature_dict["template_sequence"] = np.array( + [ + _slice_sequence_like( + template_sequence, + keep_indices, + tolerate_shorter_input=True, + context="template_sequence", + ) + for template_sequence in normalized_feature_dict["template_sequence"] + ], + dtype=object, + ) + + if "sequence" in normalized_feature_dict: + normalized_feature_dict["sequence"] = np.array( + [normalized_sequence.encode("utf-8")] + ) + if "seq_length" in normalized_feature_dict: + normalized_feature_dict["seq_length"] = np.full( + len(normalized_sequence), + len(normalized_sequence), + dtype=np.int32, + ) + if "num_alignments" in normalized_feature_dict: + num_alignments = int( + np.asarray(normalized_feature_dict["num_alignments"]).reshape(-1)[0] + ) + normalized_feature_dict["num_alignments"] = np.full( + len(normalized_sequence), + num_alignments, + dtype=np.int32, + ) + + logging.info( + "Collapsed %d adjacent duplicate residue position(s) for %s to " + "satisfy AF3 input constraints", + len(residue_ids) - len(normalized_residue_ids), + mono_obj.description, + ) + return ( + normalized_sequence, + normalized_feature_dict, + normalized_residue_ids, + True, + ) + + +def _validate_json_template_mmcif( + mmcif_string: str, +) -> tuple[bool, str | None]: + try: + af3_mmcif.from_string(mmcif_string) + except Exception as exc: + return False, str(exc) + return True, None + + +def _slice_a3m_row_to_region(a3m_row: str, start: int, end: int) -> str: + query_position = 0 + sliced_chars: list[str] = [] + for char in a3m_row: + if char.islower(): + # Lowercase A3M letters are insertions after the current + # query position, so keep only those between retained + # residues and drop insertions before `start` or after `end`. + if start <= query_position < end: + sliced_chars.append(char) + continue + + query_position += 1 + if query_position < start: + continue + if query_position > end: + break + sliced_chars.append(char) + return "".join(sliced_chars) + + +def _slice_a3m_to_regions( + a3m_text: str | None, + regions: Sequence[tuple[int, int]], +) -> str | None: + if a3m_text in (None, "") or not regions: + return a3m_text + + sequences, descriptions = af3_parsers.parse_fasta(a3m_text) + sliced_sequences = [ + "".join( + _slice_a3m_row_to_region(sequence, start, end) + for start, end in regions + ) + for sequence in sequences + ] + return ( + "\n".join( + f">{description}\n{sequence}" + for description, sequence in zip( + descriptions, + sliced_sequences, + strict=True, + ) + ) + + "\n" + ) + + +def _slice_positioned_modifications_to_regions( + modifications: Sequence[tuple[str, int]], + regions: Sequence[tuple[int, int]], +) -> list[tuple[str, int]]: + sliced_modifications: list[tuple[str, int]] = [] + offset = 0 + for start, end in regions: + sliced_modifications.extend( + ( + modification_name, + offset + modification_position - start + 1, + ) + for modification_name, modification_position in modifications + if start <= modification_position <= end + ) + offset += end - start + 1 + return sliced_modifications + + +def _slice_templates_to_regions( + templates: Sequence[folding_input.Template] | None, + regions: Sequence[tuple[int, int]], +) -> Sequence[folding_input.Template] | None: + if templates is None: + return None + + sliced_templates = [] + for template in templates: + remapped_indices = {} + offset = 0 + for start, end in regions: + start_index = start - 1 + remapped_indices.update({ + offset + (query_index - start_index): template_index + for query_index, template_index in template.query_to_template_map.items() + if start_index <= query_index < end + }) + offset += end - start + 1 + if remapped_indices: + sliced_templates.append( + folding_input.Template( + mmcif=template.mmcif, + query_to_template_map=remapped_indices, + ) + ) + return sliced_templates + + +def _chain_residue_ids( + chain: ( + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ), +) -> list[int] | None: + residue_ids = getattr(chain, "residue_ids", None) + if residue_ids is not None: + return [int(residue_id) for residue_id in residue_ids] + if isinstance( + chain, + ( + folding_input.ProteinChain, + folding_input.RnaChain, + folding_input.DnaChain, + ), + ): + return list(range(1, len(_chain_input_sequence(chain)) + 1)) + return None + + +def _slice_af3_chain_to_regions( + chain: ( + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ), + *, + regions: Sequence[tuple[int, int]], + json_path: str, +): + if isinstance(chain, folding_input.Ligand): + raise ValueError( + f"Region ranges are not supported for ligand AF3 JSON inputs: {json_path}" + ) + + chain_sequence = _chain_input_sequence(chain) + sequence_length = len(chain_sequence) + for start, end in regions: + if not 1 <= start <= end <= sequence_length: + raise ValueError( + f"Requested region {start}-{end} is outside the sequence " + f"length {sequence_length} for AF3 JSON input {json_path}." + ) + + sliced_sequence = "".join( + chain_sequence[start - 1:end] + for start, end in regions + ) + sliced_residue_ids = [ + residue_id + for start, end in regions + for residue_id in (_chain_residue_ids(chain) or [])[start - 1:end] + ] + if isinstance(chain, folding_input.ProteinChain): + return _construct_chain( + folding_input.ProteinChain, + id=chain.id, + sequence=sliced_sequence, + description=getattr(chain, "description", None), + ptms=_slice_positioned_modifications_to_regions( + chain.ptms, + regions, + ), + residue_ids=sliced_residue_ids, + unpaired_msa=_slice_a3m_to_regions( + chain.unpaired_msa, + regions, + ), + paired_msa=_slice_a3m_to_regions( + chain.paired_msa, + regions, + ), + templates=_slice_templates_to_regions( + chain.templates, + regions, + ), + ) + + if isinstance(chain, folding_input.RnaChain): + return _construct_chain( + folding_input.RnaChain, + id=chain.id, + sequence=sliced_sequence, + description=getattr(chain, "description", None), + modifications=_slice_positioned_modifications_to_regions( + chain.modifications, + regions, + ), + residue_ids=sliced_residue_ids, + unpaired_msa=_slice_a3m_to_regions( + chain.unpaired_msa, + regions, + ), + ) + + if isinstance(chain, folding_input.DnaChain): + return _construct_chain( + folding_input.DnaChain, + id=chain.id, + sequence=sliced_sequence, + description=getattr(chain, "description", None), + modifications=_slice_positioned_modifications_to_regions( + chain.modifications(), + regions, + ), + residue_ids=sliced_residue_ids, + ) + + raise TypeError(f"Unsupported chain type for AF3 JSON slicing: {type(chain)}") + + +def get_chain_id(index: int) -> str: + if index < 26: + return chr(ord('A') + index) + else: + first_char = chr(ord('A') + (index // 26) - 1) + second_char = chr(ord('A') + (index % 26)) + return first_char + second_char + + +def insert_release_date_into_mmcif( + mmcif_string: str, revision_date: str = '2100-01-01' +) -> str: + lines = mmcif_string.splitlines() + insert_index = None + for i, line in enumerate(lines): + if line.startswith('data_'): + insert_index = i + 1 + break + if insert_index is None: + return mmcif_string + revision_lines = [ + "", + "loop_", + "_pdbx_audit_revision_history.ordinal", + "_pdbx_audit_revision_history.data_content_type", + "_pdbx_audit_revision_history.major_revision", + "_pdbx_audit_revision_history.minor_revision", + "_pdbx_audit_revision_history.revision_date", + "1 'Structure model' 1 0 " + revision_date, + "" + ] + lines[insert_index:insert_index] = revision_lines + return '\n'.join(lines) + + +def strip_entity_poly_seq_block(mmcif_string: str) -> str: + """Remove _entity_poly_seq loop to let AF3 reconstruct scheme from _atom_site. + + AF3's parser can infer _pdbx_poly_seq_scheme directly from _atom_site and + avoids mismatches when our synthetic _entity_poly_seq has UNK entries. + """ + lines = mmcif_string.splitlines() + out_lines = [] + i = 0 + while i < len(lines): + line = lines[i] + if line.strip() == 'loop_': + # Peek ahead to see if this loop is for _entity_poly_seq + j = i + 1 + is_entity_poly_seq = False + while j < len(lines) and lines[j].strip().startswith('_'): + if lines[j].strip().startswith('_entity_poly_seq.'): + is_entity_poly_seq = True + # Stop header scan when we reach first data row (no leading underscore) + j += 1 + if j < len(lines) and not lines[j].strip().startswith('_'): + break + if is_entity_poly_seq: + # Skip until next blank line or next loop_/category header + k = j + while k < len(lines): + if lines[k].strip() == '' or lines[k].strip() == '#': + k += 1 + break + if lines[k].strip() == 'loop_': + break + k += 1 + i = k + continue + out_lines.append(line) + i += 1 + return '\n'.join(out_lines) + + +def _expand_monomeric_object( + mono_obj: Union[MonomericObject, ChoppedObject], +) -> list[Union[MonomericObject, ChoppedObject]]: + return [mono_obj] + + +@dataclasses.dataclass(frozen=True) +class PreparedFold: + """A fold input ready to predict: where its outputs go, and whether AF3 may resolve + MSA overlaps (not when its MSAs were translated from AF2 multimer features).""" + + fold_input: folding_input.Input + output_dir: str + resolve_msa_overlaps: bool = True + + +class FoldInputBuilder: + """Translates one preparation's AlphaPulldown objects and AF3 JSON inputs into AF3 chains. + + It holds what translating one object needs from the others: the chains so far, the + chain IDs in use, which chains carry MSAs translated from AF2 features, the debug + records, and caches. ``prepare_input`` makes a fresh builder per call, so none of + them outlives the preparation. + """ + + def __init__( + self, + *, + features_directory: str | None = None, + debug_templates: bool | None = None, + debug_msas: bool = False, + ) -> None: + self.features_directory = features_directory + self.debug_templates = debug_templates + self.debug_msas = debug_msas + self.chains: list[ + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ] = [] + self.used_chain_ids: set[str] = set() + self._chain_id_counter = 0 + # Track chains whose MSAs were translated from AF2 features; they must + # not be rewritten by the promotion heuristic in build(). + self.af2_translated_msa_chain_ids: set[str] = set() + self.translation_debug_chain_records: list[_TranslatedMsaDebugRecord] = [] + self.translation_debug_results: list[Af2ToAf3TranslationResult] = [] + # Per-preparation caches, as lru_cache wrappers on this instance. + self._description_with_af2_feature_metadata = functools.lru_cache(maxsize=None)( + self._description_with_af2_feature_metadata + ) + self._validate_json_template_mmcif = functools.lru_cache(maxsize=None)( + _validate_json_template_mmcif + ) + + def _next_chain_id(self) -> str: + """Get the next available chain ID that's not already used.""" + while True: + chain_id = get_chain_id(self._chain_id_counter) + self._chain_id_counter += 1 + if chain_id not in self.used_chain_ids: + return chain_id + + def _description_with_af2_feature_metadata( + self, + output_description: str, + source_description: str, + ) -> str: + try: + metadata = find_feature_metadata( + source_description, self.features_directory + ) + except (OSError, ValueError) as exc: + logging.warning( + "Could not read feature metadata for %s: %s", + source_description, + exc, + ) + return output_description + if metadata is None: + return output_description + return encode_metadata_in_description(output_description, metadata) + + def _sanitize_json_input_chain_templates( + self, + chain: ( + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ), + *, + json_path: str, + ): + if not isinstance(chain, folding_input.ProteinChain): + return chain + if chain.templates is None: + return chain + + valid_templates = [] + dropped_template_count = 0 + for template_index, template in enumerate(chain.templates): + is_valid, error_message = self._validate_json_template_mmcif( + template.mmcif + ) + if is_valid: + valid_templates.append(template) + continue + + dropped_template_count += 1 + logging.warning( + "Skipping invalid template %d for JSON input %s chain %s: %s", + template_index, + json_path, + chain.id, + error_message, + ) + + if dropped_template_count == 0: + return chain + + logging.info( + "Kept %d/%d template(s) for JSON input %s chain %s after validation", + len(valid_templates), + len(chain.templates), + json_path, + chain.id, + ) + return _construct_chain( + folding_input.ProteinChain, + id=chain.id, + sequence=_chain_input_sequence(chain), + description=getattr(chain, "description", None), + residue_ids=getattr(chain, "residue_ids", None), + ptms=chain.ptms, + paired_msa=chain.paired_msa, + unpaired_msa=chain.unpaired_msa, + templates=valid_templates, + ) + + def _expand_json_input_chains( + self, + *, + chains: Sequence[ + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ], + json_path: str, + regions: list[tuple[int, int]] | None, + ) -> list[ + folding_input.ProteinChain + | folding_input.RnaChain + | folding_input.DnaChain + | folding_input.Ligand + ]: + sanitized_chains = [ + self._sanitize_json_input_chain_templates(chain, json_path=json_path) + for chain in chains + ] + if not regions: + return list(sanitized_chains) + + if len(sanitized_chains) != 1: + raise ValueError( + "Region ranges for AF3 JSON feature inputs require exactly " + f"one chain per file, but {json_path} contains {len(sanitized_chains)} chains." + ) + + merged_chain = _slice_af3_chain_to_regions( + sanitized_chains[0], + regions=regions, + json_path=json_path, + ) + logging.info( + "Collapsed AF3 JSON input %s into one gapped chain for regions %s", + json_path, + regions, + ) + return [merged_chain] + + def _monomeric_to_chain( + self, + mono_obj: Union[MonomericObject, ChoppedObject], + chain_id: str + ) -> folding_input.ProteinChain: + sequence, feature_dict, residue_ids, _ = _normalise_monomeric_inputs_for_af3( + mono_obj + ) + msa_array = feature_dict.get('msa') + deletion_matrix = feature_dict.get('deletion_matrix_int') + if deletion_matrix is None: + deletion_matrix = feature_dict.get('deletion_matrix') + # Standalone AlphaPulldown monomer objects carry only a single custom MSA. + # MultimericObject instances override this with an AF2->AF3 translation. + unpaired_msa = ( + msa_rows_and_deletions_to_a3m( + msa_rows=np.asarray(msa_array), + deletion_rows=None if deletion_matrix is None else np.asarray(deletion_matrix), + query_sequence=sequence, + ) + if msa_array is not None + else "" + ) + paired_msa = "" + templates = [] + if 'template_aatype' in feature_dict: + num_templates = feature_dict['template_aatype'].shape[0] + for i in range(num_templates): + try: + pdb_code_chain = feature_dict["template_domain_names"][i].decode('utf-8') + if '_' in pdb_code_chain: + pdb_code, chain_id_template = pdb_code_chain.split('_') + else: + pdb_code = pdb_code_chain + chain_id_template = 'A' + template_sequence = feature_dict["template_sequence"][i] + if isinstance(template_sequence, bytes): + template_sequence = template_sequence.decode('utf-8') + template_mask = feature_dict["template_all_atom_masks"][i] + hh_ids = np.argmax(feature_dict["template_aatype"][i], axis=-1) + if np.sum(template_mask) == 0: + logging.info(f"Skipping template {i} ({pdb_code_chain}) - no atoms in region") + continue + tmpl_aatype = np.array([ + residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE[j] + for j in hh_ids + ], dtype=np.int32) + tmpl_aatype[tmpl_aatype == 21] = 20 + if np.any(tmpl_aatype < 0) or np.any(tmpl_aatype >= len(residue_constants.resnames)): + raise ValueError(f"Invalid aatype: {tmpl_aatype}") + protein = Protein( + atom_positions=feature_dict["template_all_atom_positions"][i], + atom_mask=feature_dict["template_all_atom_masks"][i], + aatype=tmpl_aatype, + residue_index=feature_dict["residue_index"], + chain_index=np.zeros_like(feature_dict["residue_index"], dtype=int), + b_factors=np.zeros(template_mask.shape, dtype=float), + ) + mmcif_string = to_mmcif( + prot=protein, file_id=pdb_code, model_type='Monomer' + ) + new_mmcif_string = insert_release_date_into_mmcif(mmcif_string) + # Remove _entity_poly_seq to avoid mismatches during AF3 parsing + new_mmcif_string = strip_entity_poly_seq_block(new_mmcif_string) + # Optional debug dump of generated template mmCIF + if self.debug_templates: + try: + os.makedirs("templates_debug", exist_ok=True) + debug_fname = os.path.join( + "templates_debug", + f"generated_template_{pdb_code}_{chain_id_template}_chain{chain_id}_idx{i}.cif", + ) + with open(debug_fname, "w") as f_dbg: + f_dbg.write(new_mmcif_string) + except Exception: + # Best-effort debug; ignore failures + pass + # Build query->template mapping only for residues that have any atoms + # present in the template (avoid gaps causing index errors downstream). + try: + per_res_atom_mask = feature_dict["template_all_atom_masks"][i] + # Boolean mask: True if any atom present for this residue + present_res_mask = np.any(per_res_atom_mask > 0, axis=-1) + # Map compact template indices (0..N_present-1) to original residue indices + residue_indices_present = [idx for idx, present in enumerate(present_res_mask) if present] + compact_index_by_residue = {res_idx: compact_idx for compact_idx, res_idx in enumerate(residue_indices_present)} + # Query indices are positions in the template sequence aligned to the query + query_to_template_map = {res_idx: compact_index_by_residue[res_idx] for res_idx in residue_indices_present} + except Exception: + # Fallback to identity if masks missing, though this may error later + query_to_template_map = {j: j for j in range(len(template_sequence))} + templates.append( + folding_input.Template( + mmcif=new_mmcif_string, + query_to_template_map=query_to_template_map, + ) + ) + except Exception as e: + logging.error(f"Error processing template {i} for chain {chain_id}: {e}") + try: + os.makedirs("templates_debug", exist_ok=True) + error_filename = os.path.join("templates_debug", f"ERROR_template_{i}_chain_{chain_id}.txt") + with open(error_filename, "w") as error_file: + error_file.write(f"Error: {e}\n") + error_file.write(f"Template index: {i}\n") + error_file.write(f"Chain ID: {chain_id}\n") + if 'template_domain_names' in feature_dict: + error_file.write(f"Domain name: {feature_dict['template_domain_names'][i]}\n") + if 'template_sequence' in feature_dict: + error_file.write(f"Template sequence: {feature_dict['template_sequence'][i]}\n") + except Exception as save_error: + logging.error(f"Failed to save error info: {save_error}") + raise + source_description = getattr( + mono_obj, "monomeric_description", mono_obj.description + ) + chain_description = self._description_with_af2_feature_metadata( + mono_obj.description, + source_description, + ) + chain = _construct_chain( + folding_input.ProteinChain, + id=chain_id, + sequence=sequence, + ptms=[], + residue_ids=residue_ids, + description=chain_description, + unpaired_msa=unpaired_msa, + paired_msa=paired_msa, + templates=templates, + ) + return chain + + def add(self, obj) -> None: + """Translate one AlphaPulldown object or AF3 JSON input into chains, with chain IDs + unique across this preparation.""" + if isinstance(obj, dict) and 'json_input' in obj: + json_path = obj['json_input'] + json_regions = obj.get('regions') + logging.info(f"Processing JSON file: {json_path}") + try: + # Accept a compressed feature file: the published database + # ships AF3 inputs as .json.xz, mirroring AF2's .pkl.xz. + resolved = resolve_af3_input(json_path) + if resolved is None: + raise FileNotFoundError(json_path) + if resolved != json_path: + logging.info(f"Using compressed feature file: {resolved}") + json_str = read_maybe_xz(resolved) + input_obj = folding_input.Input.from_json(json_str) + chains_to_add = self._expand_json_input_chains( + chains=input_obj.chains, + json_path=json_path, + regions=( + json_regions if isinstance(json_regions, list) and json_regions else None + ), + ) + # Track the chain IDs from the JSON file + logging.info( + "JSON file %s contributes chains with IDs: %s", + json_path, + [chain.id for chain in chains_to_add], + ) + + # Check for duplicate chain IDs and modify them if necessary + modified_chains = [] + for chain in chains_to_add: + original_id = chain.id + new_id = original_id + + # If this chain ID is already used, generate a new one + if original_id in self.used_chain_ids: + new_id = self._next_chain_id() + logging.info(f"Chain ID '{original_id}' already used, changing to '{new_id}'") + + self.used_chain_ids.add(new_id) + logging.info(f"Added chain ID '{new_id}' from JSON file {json_path}") + + modified_chains.append(_clone_chain_with_id(chain, new_id)) + + self.chains.extend(modified_chains) + except Exception as e: + logging.error(f"Failed to parse JSON file {json_path}: {e}") + raise + elif isinstance(obj, (MonomericObject, ChoppedObject)): + for expanded_obj in _expand_monomeric_object(obj): + chain_id = self._next_chain_id() + self.used_chain_ids.add(chain_id) + logging.info( + f"Added chain ID '{chain_id}' for AlphaPulldown object " + f"{expanded_obj.description}" + ) + self.chains.append(self._monomeric_to_chain(expanded_obj, chain_id)) + elif isinstance(obj, MultimericObject): + chains = [] + translated_result = None + combined_msa = obj.feature_dict.get('msa') + expanded_interactors = list(obj.interactors) + normalized_interactor_inputs = [ + _normalise_monomeric_inputs_for_af3(interactor) + for interactor in expanded_interactors + ] + has_adjacent_duplicate_residue_ids = any( + normalized_input[3] for normalized_input in normalized_interactor_inputs + ) + + translated_result = ( + translate_af2_individual_chain_features_to_af3_msas_with_stats( + chain_feature_dicts=[ + normalized_input[1] + for normalized_input in normalized_interactor_inputs + ], + chain_sequences=[ + normalized_input[0] + for normalized_input in normalized_interactor_inputs + ], + ) + ) + num_pairable_chains = sum( + chain_stats.paired_species_identifier_count > 0 + for chain_stats in translated_result.chain_stats + ) + if ( + has_adjacent_duplicate_residue_ids + and num_pairable_chains < 2 + and combined_msa is not None + ): + logging.info( + "Skipping AF2 merged-MSA fallback for multimeric object %s " + "because adjacent duplicate residue IDs were collapsed for AF3 input compatibility.", + obj.description, + ) + elif num_pairable_chains < 2 and combined_msa is not None: + # Fall back to the merged AF2 multimer MSA transport path when the + # individual `_all_seq` features do not carry usable species IDs. + translated_result = ( + translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats( + merged_msa=np.asarray(combined_msa), + chain_sequences=[ + interactor.sequence for interactor in obj.interactors + ], + num_alignments=obj.feature_dict.get('num_alignments'), + deletion_matrix=( + obj.feature_dict.get('deletion_matrix_int') + if obj.feature_dict.get('deletion_matrix_int') is not None + else obj.feature_dict.get('deletion_matrix') + ), + asym_id=obj.feature_dict.get('asym_id'), + ) + ) + expanded_interactors = list(obj.interactors) + + for chain_index, interactor in enumerate(expanded_interactors): + chain_id = self._next_chain_id() + self.used_chain_ids.add(chain_id) + logging.info(f"Added chain ID '{chain_id}' for multimeric interactor") + base_chain = self._monomeric_to_chain(interactor, chain_id) + if translated_result is not None: + chain_msas = translated_result.chain_msas[chain_index] + chain_stats = translated_result.chain_stats[chain_index] + base_chain = _construct_chain( + folding_input.ProteinChain, + id=base_chain.id, + sequence=_chain_input_sequence(base_chain), + ptms=base_chain.ptms, + residue_ids=getattr(base_chain, "residue_ids", None), + description=getattr(base_chain, "description", None), + unpaired_msa=chain_msas.unpaired_msa, + paired_msa=chain_msas.paired_msa, + templates=base_chain.templates, + ) + self.af2_translated_msa_chain_ids.add(base_chain.id) + if self.debug_msas: + self.translation_debug_chain_records.append( + _TranslatedMsaDebugRecord( + chain_id=base_chain.id, + chain_description=interactor.description, + chain_length=len(interactor.sequence), + paired_msa=chain_msas.paired_msa, + unpaired_msa=chain_msas.unpaired_msa, + paired_msa_row_count=chain_stats.paired_msa_row_count, + unpaired_msa_row_count=chain_stats.unpaired_msa_row_count, + paired_species_identifier_count=( + chain_stats.paired_species_identifier_count + ), + paired_rows_without_species_identifier_count=( + chain_stats.paired_rows_without_species_identifier_count + ), + paired_rows_with_generated_accession_count=( + chain_stats.paired_rows_with_generated_accession_count + ), + ) + ) + chains.append(base_chain) + if self.debug_msas and translated_result is not None: + self.translation_debug_results.append(translated_result) + self.chains.extend(chains) + else: + raise TypeError(f"Unsupported object type for folding input conversion: {type(obj)}") + + def build(self, *, job_name: str, random_seed: int) -> folding_input.Input: + """The fold input of all chains added, with unpaired MSAs promoted to paired.""" + # Promote unpaired->paired only for chains not constructed from AF2 multimer features. + promoted_chains: list[folding_input.ProteinChain | folding_input.RnaChain | folding_input.DnaChain | folding_input.Ligand] = [] + for ch in self.chains: + if ( + isinstance(ch, folding_input.ProteinChain) + and ch.id not in self.af2_translated_msa_chain_ids + ): + try: + has_empty_paired = (getattr(ch, 'paired_msa', None) in (None, '')) + has_unpaired = bool(getattr(ch, 'unpaired_msa', '')) + if has_empty_paired and has_unpaired: + new_chain = _construct_chain( + folding_input.ProteinChain, + id=ch.id, + sequence=_chain_input_sequence(ch), + ptms=ch.ptms, + residue_ids=getattr(ch, "residue_ids", None), + description=getattr(ch, "description", None), + paired_msa=ch.unpaired_msa, + unpaired_msa='', + templates=ch.templates, + ) + promoted_chains.append(new_chain) + continue + except Exception: + pass + promoted_chains.append(ch) + + return folding_input.Input( + name=job_name, + rng_seeds=[random_seed], + chains=promoted_chains, + ) diff --git a/test/af3_fixtures.py b/test/af3_fixtures.py index aac6688b..e51b4372 100644 --- a/test/af3_fixtures.py +++ b/test/af3_fixtures.py @@ -278,5 +278,5 @@ def _prepare_fold_input( debug_msas=debug_msas, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input return fold_input_obj diff --git a/test/integration/test_af3_input_preparation.py b/test/integration/test_af3_input_preparation.py index 1080bd73..13bfa7d4 100644 --- a/test/integration/test_af3_input_preparation.py +++ b/test/integration/test_af3_input_preparation.py @@ -393,10 +393,9 @@ def test_af3_real_fixture_pipeline_tolerates_mixed_missing_accession_ids(self): debug_msas=True, ) self.assertLen(mappings, 1) - fold_input_obj, ( - prepared_output_dir, - resolve_msa_overlaps, - ) = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input + prepared_output_dir = mappings[0].output_dir + resolve_msa_overlaps = mappings[0].resolve_msa_overlaps process_fold_input( fold_input=fold_input_obj, @@ -456,7 +455,7 @@ def test_af3_on_the_fly_pairing_from_json_features(self): # Use backend to prepare the combined input mappings = AlphaFold3Backend.prepare_input(objects_to_model=objects_to_model, random_seed=42) self.assertEqual(len(mappings), 1) - fold_input_obj, out_dir = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input # Ask the backend helper to write *_data.json without inference res = process_fold_input( @@ -661,7 +660,7 @@ def test_af3_output_job_name_compacts_long_homomer_names(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input self.assertEqual( fold_input_obj.sanitised_name(), @@ -683,7 +682,7 @@ def test_af3_output_job_name_compacts_long_homomer_names(self): def test_af3_output_job_name_hashes_overlong_unique_compound_names(self): """AF3 job names should fall back to a deterministic hash suffix when needed.""" - from alphapulldown.folding_backend.alphafold3_backend import ( + from alphapulldown.folding_backend.alphafold3_inputs import ( _build_output_job_name, ) @@ -743,7 +742,7 @@ def test_af3_prepare_input_accepts_monomer_plus_ligand_json(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input self.assertEqual([chain.id for chain in fold_input_obj.chains], ["A", "L"]) self.assertIsInstance(fold_input_obj.chains[0], folding_input.ProteinChain) @@ -797,7 +796,7 @@ def test_af3_prepare_input_skips_invalid_json_templates_for_ptm_input(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input self.assertEqual([chain.id for chain in fold_input_obj.chains], ["P"]) self.assertLen(fold_input_obj.chains, 1) @@ -858,7 +857,7 @@ def test_af3_prepare_input_keeps_valid_json_templates(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input self.assertEqual([chain.id for chain in fold_input_obj.chains], ["A"]) self.assertLen(fold_input_obj.chains, 1) @@ -1115,7 +1114,7 @@ def test_af3_keeps_discontinuous_chopped_regions_in_one_gapped_chain(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input chopped_region_sequences = self._get_region_sequences( "A0A075B6L2", @@ -1201,7 +1200,7 @@ def test_af3_keeps_two_out_of_order_gapped_copies_as_two_chains(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input expected_regions = [(8, 10), (2, 5)] expected_sequence = "".join( @@ -1280,7 +1279,7 @@ def test_af3_json_feature_ranges_collapse_into_one_gapped_chain(self): random_seed=42, ) self.assertLen(mappings, 1) - fold_input_obj, _ = next(iter(mappings[0].items())) + fold_input_obj = mappings[0].fold_input json_sequences = self._get_sequence_from_json( "af3_features/protein/A0A024R1R8_af3_input.json" diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index d154c795..49ee7439 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -11,6 +11,7 @@ import pytest from alphapulldown.objects import MonomericObject +from alphapulldown.utils import output_paths from alphapulldown.utils.feature_metadata import ( decode_metadata_from_description, encode_metadata_in_description, @@ -400,6 +401,10 @@ class Config: structure_pkg.mmcif = mmcif_mod +# The modules split out of alphafold3_backend.py; each imports AF3 itself. +_EXTRACTED_MODULES = ("alphafold3_debug", "alphafold3_inputs", "alphafold3_outputs") + + @pytest.fixture(scope="module") def af3_backend_module(tmp_path_factory): tmp_path = tmp_path_factory.mktemp("af3_backend_stubs") @@ -410,7 +415,8 @@ def af3_backend_module(tmp_path_factory): sys.modules.pop("alphapulldown.folding_backend.alphafold3_backend", None) # Load the extracted modules afresh too, against the stubs: a copy imported # earlier would hold the real AF3. - sys.modules.pop("alphapulldown.folding_backend.alphafold3_outputs", None) + for extracted in _EXTRACTED_MODULES: + sys.modules.pop(f"alphapulldown.folding_backend.{extracted}", None) spec = importlib.util.spec_from_file_location( "alphapulldown.folding_backend.alphafold3_backend", MODULE_PATH, @@ -429,6 +435,18 @@ def af3_outputs_module(af3_backend_module): return sys.modules["alphapulldown.folding_backend.alphafold3_outputs"] +@pytest.fixture(scope="module") +def af3_debug_module(af3_backend_module): + """The debug-artifact module the stubbed backend imported.""" + return sys.modules["alphapulldown.folding_backend.alphafold3_debug"] + + +@pytest.fixture(scope="module") +def af3_inputs_module(af3_backend_module): + """The input-preparation module the stubbed backend imported.""" + return sys.modules["alphapulldown.folding_backend.alphafold3_inputs"] + + class FakeChainsTable: def __init__(self, chain_ids): self._chain_ids = list(chain_ids) @@ -459,38 +477,38 @@ def copy_and_update_residues(self, **kwargs): return SimpleNamespace(updated_residues=kwargs) -def test_output_name_helpers_compact_and_normalise_fragments(af3_backend_module): - assert af3_backend_module._normalise_output_name_fragment(" A/B : C ") == "A_B___C" - assert af3_backend_module._collapse_repeated_name_fragments( +def test_output_name_helpers_compact_and_normalise_fragments(): + assert output_paths._normalise_output_name_fragment(" A/B : C ") == "A_B___C" + assert output_paths._collapse_repeated_name_fragments( ["a", "a", "b", "b", "b", "c"] ) == ["a__x2", "b__x3", "c"] - compacted = af3_backend_module._compact_output_job_name("x" * 260, max_chars=40) + compacted = output_paths._compact_output_job_name("x" * 260, max_chars=40) assert compacted.startswith("x") assert len(compacted) <= 40 -def test_compound_and_region_name_helpers_handle_existing_suffixes(af3_backend_module): - assert af3_backend_module._compact_existing_compound_name( +def test_compound_and_region_name_helpers_handle_existing_suffixes(): + assert output_paths._compact_existing_compound_name( "protA_and_protA_and_protB" ) == "protA__x2_and_protB" - assert af3_backend_module._regions_to_name_fragment([(1, 10), (20, 25)]) == "1-10_20-25" - assert af3_backend_module._json_input_basename("/tmp/A0ABD7FQG0_af3_input.json") == "A0ABD7FQG0" - assert af3_backend_module._json_input_basename("/tmp/job_input.json") == "job" + assert output_paths._regions_to_name_fragment([(1, 10), (20, 25)]) == "1-10_20-25" + assert output_paths._json_input_basename("/tmp/A0ABD7FQG0_af3_input.json") == "A0ABD7FQG0" + assert output_paths._json_input_basename("/tmp/job_input.json") == "job" def test_object_and_output_job_name_helpers_cover_json_monomer_and_multimer_inputs( - af3_backend_module, + af3_inputs_module, ): monomer = MonomericObject("protA", "ACDE") - multimer = object.__new__(af3_backend_module.MultimericObject) + multimer = object.__new__(af3_inputs_module.MultimericObject) multimer.description = "protA_and_protA_and_protB" json_entry = {"json_input": "/tmp/sample_af3_input.json", "regions": [(5, 9)]} - assert af3_backend_module._object_name_fragment(json_entry) == "sample__5-9" - assert af3_backend_module._object_name_fragment(monomer) == "protA" - assert af3_backend_module._object_name_fragment(multimer) == "protA__x2_and_protB" + assert af3_inputs_module._object_name_fragment(json_entry) == "sample__5-9" + assert af3_inputs_module._object_name_fragment(monomer) == "protA" + assert af3_inputs_module._object_name_fragment(multimer) == "protA__x2_and_protB" - job_name = af3_backend_module._build_output_job_name( + job_name = af3_inputs_module._build_output_job_name( [{"object": [json_entry, monomer, monomer, multimer]}] ) assert job_name == "sample__5-9_and_protA__x2_and_protA__x2_and_protB" @@ -868,6 +886,7 @@ def test_process_fold_input_predicts_only_the_seeds_not_yet_written( def test_predict_structure_writes_final_msa_and_collects_optional_outputs( af3_backend_module, + af3_debug_module, monkeypatch, tmp_path, ): @@ -891,7 +910,7 @@ def test_predict_structure_writes_final_msa_and_collects_optional_outputs( lambda **kwargs: [example], ) monkeypatch.setattr( - af3_backend_module, + af3_debug_module, "ids_to_a3m_af3", lambda rows: ">query\nABC\n", ) @@ -1276,7 +1295,9 @@ def test_af3_predict_expands_num_seeds_and_calls_process_fold_input( monkeypatch.setattr( af3_backend_module.AlphaFold3Backend, "prepare_input", - staticmethod(lambda **kwargs: [{fake_input: (tmp_path, False)}]), + staticmethod( + lambda **kwargs: [af3_backend_module.PreparedFold(fake_input, tmp_path, False)] + ), ) monkeypatch.setattr( af3_backend_module, @@ -1316,11 +1337,12 @@ def _af3_run(module, monkeypatch, output_dir, *, finish=True, random_seed=None, monkeypatch.setattr( backend, "prepare_input", - staticmethod(lambda **kwargs: [{ + staticmethod(lambda **kwargs: [module.PreparedFold( module.folding_input.Input( name="job", chains=("A",), rng_seeds=(kwargs["random_seed"],) - ): (output_dir, True) - }]), + ), + output_dir, + )]), ) monkeypatch.setattr( module, @@ -1475,7 +1497,7 @@ def _modelcif_metadata_given_to_seed_outputs(module, monkeypatch, output_dir, fo monkeypatch.setattr( module.AlphaFold3Backend, "prepare_input", - staticmethod(lambda **kwargs: [{fold_input: (output_dir, True)}]), + staticmethod(lambda **kwargs: [module.PreparedFold(fold_input, output_dir)]), ) calls = [] monkeypatch.setattr( @@ -1623,9 +1645,9 @@ def test_prepare_input_merges_json_chain_and_rewrites_duplicate_ids( ) assert len(prepared_inputs) == 1 - fold_input, (output_dir, resolve_msa_overlaps) = next( - iter(prepared_inputs[0].items()) - ) + fold_input = prepared_inputs[0].fold_input + output_dir = prepared_inputs[0].output_dir + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert output_dir == str(tmp_path) assert resolve_msa_overlaps is True assert [chain.id for chain in fold_input.chains] == ["A", "B"] @@ -1635,6 +1657,117 @@ def test_prepare_input_merges_json_chain_and_rewrites_duplicate_ids( assert fold_input.chains[1].unpaired_msa == "" +def test_fold_input_builder_numbers_chains_past_z(af3_inputs_module): + builder = af3_inputs_module.FoldInputBuilder() + for index in range(28): + builder.add(MonomericObject(f"prot{index}", "ACDE")) + + assert [chain.id for chain in builder.chains][-3:] == ["Z", "AA", "AB"] + + +def test_fold_input_builder_never_gives_two_chains_one_id(af3_inputs_module, tmp_path): + first = _write_stub_af3_json( + tmp_path / "first.json", + { + "name": "first", + "chains": [ + {"type": "protein", "id": "A", "sequence": "GG"}, + {"type": "protein", "id": "C", "sequence": "HH"}, + ], + }, + ) + second = _write_stub_af3_json( + tmp_path / "second.json", + {"name": "second", "chains": [{"type": "protein", "id": "A", "sequence": "KK"}]}, + ) + builder = af3_inputs_module.FoldInputBuilder() + builder.add({"json_input": str(first)}) + builder.add(MonomericObject("protA", "ACDE")) + builder.add(MonomericObject("protB", "ACDE")) + builder.add({"json_input": str(second)}) + + # A JSON chain keeps its ID while it is free, generated IDs skip the taken ones, + # and a JSON chain whose ID is taken gets the next free one. + assert [chain.id for chain in builder.chains] == ["A", "C", "B", "D", "E"] + assert [chain.sequence for chain in builder.chains] == ["GG", "HH", "ACDE", "ACDE", "KK"] + + +def test_successive_preparations_share_no_state(af3_backend_module, tmp_path): + """Each preparation numbers its chains afresh and reads feature metadata from its own + features directory, even for a description an earlier preparation already read.""" + fold_inputs = [] + for version in ("2.5.0", "2.6.0"): + feature_dir = tmp_path / version + feature_dir.mkdir() + metadata = { + "software": {"AlphaPulldown": {"version": version}}, + "databases": {}, + "other": {"data_pipeline": "alphafold2"}, + } + (feature_dir / "protA_feature_metadata_2026-08-19.json").write_text( + json.dumps(metadata), encoding="utf-8" + ) + prepared_inputs = af3_backend_module.AlphaFold3Backend.prepare_input( + objects_to_model=[ + {"object": MonomericObject("protA", "ACDE"), "output_dir": str(tmp_path)} + ], + random_seed=7, + features_directory=[str(feature_dir)], + ) + fold_input = prepared_inputs[0].fold_input + fold_inputs.append(fold_input) + + assert [[chain.id for chain in fold_input.chains] for fold_input in fold_inputs] == [ + ["A"], + ["A"], + ] + versions = [ + decode_metadata_from_description(fold_input.chains[0].description)[1]["software"][ + "AlphaPulldown" + ]["version"] + for fold_input in fold_inputs + ] + assert versions == ["2.5.0", "2.6.0"] + + +def test_a3m_region_slice_keeps_only_the_insertions_inside_it(af3_inputs_module): + # Lowercase letters are insertions after the query residue before them. + row = "AbCDeF" + + assert af3_inputs_module._slice_a3m_row_to_region(row, 1, 3) == "AbCD" + assert af3_inputs_module._slice_a3m_row_to_region(row, 2, 3) == "CD" + + +def test_json_chain_regions_slice_sequence_residue_ids_ptms_and_msa(af3_inputs_module): + chain = af3_inputs_module.folding_input.ProteinChain( + id="A", + sequence="ACDEFG", + ptms=[("SEP", 4)], + unpaired_msa=">query\nACDEFG\n>hit\nAC-EFG\n", + ) + + sliced = af3_inputs_module._slice_af3_chain_to_regions( + chain, regions=[(1, 2), (4, 5)], json_path="input.json" + ) + + assert sliced.sequence == "ACEF" + assert sliced.residue_ids == [1, 2, 4, 5] + assert sliced.ptms == [("SEP", 3)] + assert sliced.unpaired_msa == ">query\nACEF\n>hit\nACEF\n" + + +def test_json_regions_need_exactly_one_chain(af3_inputs_module): + protein = af3_inputs_module.folding_input.ProteinChain + builder = af3_inputs_module.FoldInputBuilder() + + with pytest.raises(ValueError, match="exactly one chain per file"): + builder._expand_json_input_chains( + chains=[protein(id="A", sequence="AC"), protein(id="B", sequence="DE")], + json_path="input.json", + regions=[(1, 1)], + ) + + def test_prepare_input_transports_af2_sidecar_metadata_into_af3_description( af3_backend_module, tmp_path, @@ -1660,7 +1793,7 @@ def test_prepare_input_transports_af2_sidecar_metadata_into_af3_description( features_directory=[str(feature_dir)], ) - fold_input, _ = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input clean_description, recovered = decode_metadata_from_description( fold_input.chains[0].description ) @@ -1718,7 +1851,7 @@ def test_prepare_input_preserves_mixed_af2_and_af3_metadata( features_directory=[str(feature_dir)], ) - fold_input, _ = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input assert extract_metadata_from_fold_input(fold_input) == [af2_metadata, af3_metadata] @@ -1796,7 +1929,8 @@ def test_prepare_input_slices_single_chain_json_regions_and_promotes_msa( random_seed=13, ) - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps chain = fold_input.chains[0] assert resolve_msa_overlaps is True @@ -1843,7 +1977,7 @@ def test_prepare_input_drops_invalid_json_templates_but_keeps_valid_ones( random_seed=17, ) - fold_input, _ = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input chain = fold_input.chains[0] assert len(chain.templates) == 1 assert chain.templates[0].mmcif == "data_valid" @@ -1852,6 +1986,7 @@ def test_prepare_input_drops_invalid_json_templates_but_keeps_valid_ones( def test_prepare_input_normalises_adjacent_duplicate_residues_for_monomers( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -1880,7 +2015,7 @@ def _fake_msa_rows_and_deletions_to_a3m( return ">query\nACD\n" monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "msa_rows_and_deletions_to_a3m", _fake_msa_rows_and_deletions_to_a3m, ) @@ -1895,7 +2030,8 @@ def _fake_msa_rows_and_deletions_to_a3m( random_seed=19, ) - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps chain = fold_input.chains[0] assert resolve_msa_overlaps is True @@ -1915,7 +2051,7 @@ def _fake_msa_rows_and_deletions_to_a3m( def _make_af3_translation_result( - af3_backend_module, + af3_inputs_module, *, chain_msas, chain_stats, @@ -1925,7 +2061,7 @@ def _make_af3_translation_result( invalid_paired_rows=0, invalid_unpaired_rows=0, ): - return af3_backend_module.Af2ToAf3TranslationResult( + return af3_inputs_module.Af2ToAf3TranslationResult( chain_msas=tuple(chain_msas), chain_stats=tuple(chain_stats), translation_mode=translation_mode, @@ -2030,6 +2166,7 @@ def _make_stub_multimeric_object(af3_backend_module, *, description: str, intera def test_prepare_input_multimer_writes_translated_msa_debug_artifacts( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -2044,7 +2181,7 @@ def test_prepare_input_multimer_writes_translated_msa_debug_artifacts( multimer.feature_dict = {"msa": np.asarray([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=np.int32)} translation_result = _make_af3_translation_result( - af3_backend_module, + af3_inputs_module, chain_msas=[ _translation_msas( paired_msa=( @@ -2085,13 +2222,13 @@ def test_prepare_input_multimer_writes_translated_msa_debug_artifacts( ) monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_individual_chain_features_to_af3_msas_with_stats", lambda **kwargs: translation_result, ) fallback_calls = [] monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats", lambda **kwargs: fallback_calls.append(kwargs), ) @@ -2108,7 +2245,9 @@ def test_prepare_input_multimer_writes_translated_msa_debug_artifacts( ) assert fallback_calls == [] - fold_input, (output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + output_dir = prepared_inputs[0].output_dir + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert output_dir == str(tmp_path) assert resolve_msa_overlaps is False assert [chain.id for chain in fold_input.chains] == ["A", "B"] @@ -2175,6 +2314,7 @@ def test_prepare_input_multimer_writes_translated_msa_debug_artifacts( def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavailable( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -2196,7 +2336,7 @@ def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavai individual_calls = [] fallback_calls = [] individual_result = _make_af3_translation_result( - af3_backend_module, + af3_inputs_module, chain_msas=[ _translation_msas(paired_msa="", unpaired_msa=">indA\nACDE\n"), _translation_msas(paired_msa="", unpaired_msa=">indB\nFGHI\n"), @@ -2216,7 +2356,7 @@ def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavai translation_mode="af3_species_pairing_from_af2_individual_msas", ) fallback_result = _make_af3_translation_result( - af3_backend_module, + af3_inputs_module, chain_msas=[ _translation_msas(paired_msa="", unpaired_msa=">fallbackA\nACDE\n"), _translation_msas(paired_msa="", unpaired_msa=">fallbackB\nFGHI\n"), @@ -2235,12 +2375,12 @@ def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavai ) monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_individual_chain_features_to_af3_msas_with_stats", lambda **kwargs: individual_calls.append(kwargs) or individual_result, ) monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats", lambda **kwargs: fallback_calls.append(kwargs) or fallback_result, ) @@ -2263,7 +2403,8 @@ def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavai np.asarray([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=np.int32), ) - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert resolve_msa_overlaps is False assert [chain.unpaired_msa for chain in fold_input.chains] == [ ">fallbackA\nACDE\n", @@ -2274,6 +2415,7 @@ def test_prepare_input_multimer_falls_back_to_complex_msa_when_pairing_is_unavai def test_prepare_input_multimer_skips_complex_fallback_after_duplicate_residue_normalisation( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -2288,7 +2430,7 @@ def test_prepare_input_multimer_skips_complex_fallback_after_duplicate_residue_n multimer.feature_dict = {"msa": np.asarray([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=np.int32)} individual_result = _make_af3_translation_result( - af3_backend_module, + af3_inputs_module, chain_msas=[ _translation_msas(paired_msa="", unpaired_msa=">indA\nACD\n"), _translation_msas(paired_msa="", unpaired_msa=">indB\nFGHI\n"), @@ -2310,12 +2452,12 @@ def test_prepare_input_multimer_skips_complex_fallback_after_duplicate_residue_n fallback_calls = [] monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_individual_chain_features_to_af3_msas_with_stats", lambda **kwargs: individual_result, ) monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats", lambda **kwargs: fallback_calls.append(kwargs), ) @@ -2331,7 +2473,8 @@ def test_prepare_input_multimer_skips_complex_fallback_after_duplicate_residue_n ) assert fallback_calls == [] - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert resolve_msa_overlaps is False assert [chain.sequence for chain in fold_input.chains] == ["ACD", "FGHI"] assert [chain.unpaired_msa for chain in fold_input.chains] == [ @@ -2342,6 +2485,7 @@ def test_prepare_input_multimer_skips_complex_fallback_after_duplicate_residue_n def test_prepare_input_multimer_trimer_preserves_sparse_middle_chain_pairing( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -2374,7 +2518,7 @@ def _unexpected_fallback(**kwargs): raise AssertionError(f"Unexpected AF2 merged-MSA fallback: {kwargs}") monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats", _unexpected_fallback, ) @@ -2390,7 +2534,8 @@ def _unexpected_fallback(**kwargs): debug_msas=True, ) - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert resolve_msa_overlaps is False assert [chain.id for chain in fold_input.chains] == ["A", "B", "C"] assert fold_input.chains[0].paired_msa @@ -2419,6 +2564,7 @@ def _unexpected_fallback(**kwargs): def test_prepare_input_multimer_trimer_reports_effective_min_count_pairing( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): @@ -2451,7 +2597,7 @@ def _unexpected_fallback(**kwargs): raise AssertionError(f"Unexpected AF2 merged-MSA fallback: {kwargs}") monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "translate_af2_complex_msa_to_af3_unpaired_chain_msas_with_stats", _unexpected_fallback, ) @@ -2467,7 +2613,8 @@ def _unexpected_fallback(**kwargs): debug_msas=True, ) - fold_input, (_output_dir, resolve_msa_overlaps) = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input + resolve_msa_overlaps = prepared_inputs[0].resolve_msa_overlaps assert resolve_msa_overlaps is False assert all(chain.paired_msa for chain in fold_input.chains) @@ -2493,13 +2640,14 @@ def _unexpected_fallback(**kwargs): def test_prepare_input_builds_template_mmcif_and_skips_zero_atom_templates( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): monomer = MonomericObject("templated", "ACDE") - atom_type_num = af3_backend_module.residue_constants.atom_type_num + atom_type_num = af3_inputs_module.residue_constants.atom_type_num hhblits_alphabet_size = len( - af3_backend_module.residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE + af3_inputs_module.residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE ) template_aatype = np.zeros((2, 4, hhblits_alphabet_size), dtype=np.float32) template_aatype[:, np.arange(4), [0, 4, 3, 6]] = 1.0 @@ -2535,7 +2683,7 @@ def _fake_to_mmcif(*, prot, file_id, model_type): "ATOM\n" ) - monkeypatch.setattr(af3_backend_module, "to_mmcif", _fake_to_mmcif) + monkeypatch.setattr(af3_inputs_module, "to_mmcif", _fake_to_mmcif) monkeypatch.chdir(tmp_path) prepared_inputs = af3_backend_module.AlphaFold3Backend.prepare_input( @@ -2550,7 +2698,7 @@ def _fake_to_mmcif(*, prot, file_id, model_type): ) assert len(mmcif_calls) == 1 - fold_input, _ = next(iter(prepared_inputs[0].items())) + fold_input = prepared_inputs[0].fold_input chain = fold_input.chains[0] assert len(chain.templates) == 1 assert chain.templates[0].query_to_template_map == {0: 0, 2: 1} @@ -2566,13 +2714,14 @@ def _fake_to_mmcif(*, prot, file_id, model_type): def test_prepare_input_writes_template_error_artifacts_on_mmcif_generation_failure( af3_backend_module, + af3_inputs_module, monkeypatch, tmp_path, ): monomer = MonomericObject("templated", "ACDE") - atom_type_num = af3_backend_module.residue_constants.atom_type_num + atom_type_num = af3_inputs_module.residue_constants.atom_type_num hhblits_alphabet_size = len( - af3_backend_module.residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE + af3_inputs_module.residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE ) template_aatype = np.zeros((1, 4, hhblits_alphabet_size), dtype=np.float32) template_aatype[0, np.arange(4), [0, 4, 3, 6]] = 1.0 @@ -2590,7 +2739,7 @@ def test_prepare_input_writes_template_error_artifacts_on_mmcif_generation_failu } monkeypatch.setattr( - af3_backend_module, + af3_inputs_module, "to_mmcif", lambda **kwargs: (_ for _ in ()).throw(RuntimeError("mmcif boom")), )