feat: keep AlphaPulldown's own files in alphapulldown/ and debug/

alphapulldown/folding_backend/alphafold2_backend.py

@@ -25,6 +25,8 @@
 #from alphapulldown.utils.calculate_rmsd import calculate_rmsd_and_superpose
 from alphapulldown.utils.modelling_setup import pad_input_features
 from alphapulldown.prediction import fast_kernels as fast_kernels_policy
+from alphapulldown.prediction.inference_kernels import record_inference_kernels
+from alphapulldown.utils.output_paths import debug_dir
 from alphapulldown.prediction.inference_flags import validate_model_configuration
 from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache
 from alphapulldown.prediction.resume import (
@@ -412,7 +414,8 @@ def _write_processed_msa_debug_artifact(
         ),
         query_sequence=_query_sequence_for_debug(multimeric_object),
     )
-    out_path = os.path.join(output_dir, f"{model_name}_processed_msa.a3m")
+    os.makedirs(debug_dir(output_dir), exist_ok=True)
+    out_path = os.path.join(debug_dir(output_dir), f"{model_name}_processed_msa.a3m")
     with open(out_path, "wt") as handle:
         handle.write(a3m_text)
     logging.info(f"Wrote processed MSA A3M to {out_path}")
@@ -497,8 +500,8 @@ def _write_processed_template_debug_artifacts(
     else:
         chain_index = np.asarray(_normalize_asym_id(processed_feature_dict)["asym_id"]).reshape(-1)
 
-    debug_dir = os.path.join(output_dir, "templates_debug")
-    os.makedirs(debug_dir, exist_ok=True)
+    templates_dir = os.path.join(debug_dir(output_dir), "templates")
+    os.makedirs(templates_dir, exist_ok=True)
 
     domain_names = processed_feature_dict.get("template_domain_names")
     if domain_names is None:
@@ -522,14 +525,14 @@ def _write_processed_template_debug_artifacts(
                 domain_name, f"template_{template_index}"
             )
             debug_path = os.path.join(
-                debug_dir, f"{model_name}_{template_name}_idx{template_index}.pdb"
+                templates_dir, f"{model_name}_{template_name}_idx{template_index}.pdb"
             )
             with open(debug_path, "wt") as handle:
                 handle.write(protein.to_pdb(template_protein))
             logging.info(f"Wrote processed template PDB to {debug_path}")
         except Exception as exc:
             error_path = os.path.join(
-                debug_dir, f"ERROR_{model_name}_template_{template_index}.txt"
+                templates_dir, f"ERROR_{model_name}_template_{template_index}.txt"
             )
             with open(error_path, "wt") as handle:
                 handle.write(f"Error: {exc}\n")
@@ -575,11 +578,9 @@ def _msa_depths(
 
 
 def _record_inference_kernels(output_dir, model_name, model_runner) -> None:
-    """Record in inference_kernels.json whether a model ran ColabFold's fused kernels.
+    """Record whether a model ran ColabFold's fused kernels (``inference_kernels``).
 
-    Fast and stock predictions differ slightly in their numerics, so the record lets
-    ranking and analysis tell them apart. Read back from the runner's own config, so it
-    states what actually ran.
+    Read back from the runner's own config, so it states what actually ran.
     """
     try:
         global_config = model_runner.config.model.global_config
@@ -589,17 +590,7 @@ def _record_inference_kernels(output_dir, model_name, model_runner) -> None:
         "fused_kernels": bool(global_config.get("use_pallas", False)),
         "compute_capability": global_config.get("compute_capability", None),
     }
-    path = os.path.join(output_dir, "inference_kernels.json")
-    existing = {}
-    if os.path.exists(path):
-        try:
-            with open(path) as f:
-                existing = json.load(f)
-        except (OSError, ValueError):
-            existing = {}
-    existing[model_name] = record
-    with open(path, "w") as f:
-        json.dump(existing, f, indent=2, sort_keys=True)
+    record_inference_kernels(output_dir, model_name, record)
 
 
 class AlphaFold2Backend(FoldingBackend):

alphapulldown/folding_backend/alphafold3_backend.py

@@ -45,6 +45,7 @@
 from alphapulldown.prediction.inference_flags import validate_model_configuration
 from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache
 from alphapulldown.prediction.resume import finish_fold_record
+from alphapulldown.utils.output_paths import debug_dir
 from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject
 from alphapulldown.folding_backend.alphafold3_storage import post_prediction_process_af3
 
@@ -134,6 +135,7 @@ def prepare_input(
             features_directory=features_directory,
             debug_templates=debug_templates,
             debug_msas=debug_msas,
+            debug_dir=debug_dir(objects_to_model[0]['output_dir']) if objects_to_model else None,
         )
         for entry in objects_to_model:
             object_to_model = entry['object']

alphapulldown/folding_backend/alphafold3_debug.py

@@ -18,6 +18,7 @@
 
 from alphapulldown.utils.af2_to_af3_msa import Af2ToAf3TranslationResult
 from alphapulldown.utils.msa_encoding import ids_to_a3m_af3
+from alphapulldown.utils.output_paths import debug_dir
 
 
 @dataclasses.dataclass(frozen=True, slots=True)
@@ -43,6 +44,7 @@ def _write_translated_msa_debug_artifacts(
 ) -> None:
     if not chain_records or not translation_results:
         return
+    output_dir = debug_dir(output_dir)
 
     def _extract_paired_species_identifier(description: str) -> str:
         if not description.startswith(("tr|", "sp|")):
@@ -229,11 +231,12 @@ def _write_final_msa_a3m(
     try:
         if output_dir is None:
             return
-        os.makedirs(output_dir, exist_ok=True)
+        output_dir = debug_dir(output_dir)
 
         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)
+            os.makedirs(output_dir, exist_ok=True)
             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)

alphapulldown/folding_backend/alphafold3_inputs.py

@@ -644,10 +644,14 @@ def __init__(
         features_directory: str | None = None,
         debug_templates: bool | None = None,
         debug_msas: bool = False,
+        debug_dir: str | None = None,
     ) -> None:
         self.features_directory = features_directory
         self.debug_templates = debug_templates
         self.debug_msas = debug_msas
+        # The fold's debug/ directory: generated template mmCIFs, and a note on each
+        # template that failed to convert. Without it, no debug files are written.
+        self.debug_dir = debug_dir
         self.chains: list[
             folding_input.ProteinChain
             | folding_input.RnaChain
@@ -861,11 +865,12 @@ def _monomeric_to_chain(
                     # 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:
+                    if self.debug_templates and self.debug_dir is not None:
                         try:
-                            os.makedirs("templates_debug", exist_ok=True)
+                            templates_dir = os.path.join(self.debug_dir, "templates")
+                            os.makedirs(templates_dir, exist_ok=True)
                             debug_fname = os.path.join(
-                                "templates_debug",
+                                templates_dir,
                                 f"generated_template_{pdb_code}_{chain_id_template}_chain{chain_id}_idx{i}.cif",
                             )
                             with open(debug_fname, "w") as f_dbg:
@@ -895,9 +900,12 @@ def _monomeric_to_chain(
                     )
                 except Exception as e:
                     logging.error(f"Error processing template {i} for chain {chain_id}: {e}")
+                    if self.debug_dir is None:
+                        raise
                     try:
-                        os.makedirs("templates_debug", exist_ok=True)
-                        error_filename = os.path.join("templates_debug", f"ERROR_template_{i}_chain_{chain_id}.txt")
+                        templates_dir = os.path.join(self.debug_dir, "templates")
+                        os.makedirs(templates_dir, exist_ok=True)
+                        error_filename = os.path.join(templates_dir, 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")

alphapulldown/folding_backend/alphafold3_runner.py

@@ -11,7 +11,6 @@
 import dataclasses
 import functools
 import importlib
-import json
 import logging
 import os
 import pathlib
@@ -36,6 +35,7 @@
 from alphapulldown.folding_backend.alphafold3_debug import _write_final_msa_a3m
 from alphapulldown.folding_backend.alphafold3_outputs import write_seed_outputs
 from alphapulldown.prediction import af3_fused_triangles
+from alphapulldown.prediction.inference_kernels import record_inference_kernels
 
 
 # -----------------------------------------------------------------------------
@@ -191,13 +191,11 @@ def _record_inference_kernels(
     model_runner: ModelRunner,
     example: features.BatchDict,
 ) -> None:
-    """Record in inference_kernels.json which kernels a seed's triangle layers ran.
+    """Record which kernels a seed's triangle layers ran (``inference_kernels``).
 
-    Fused and original layers differ slightly in their numerics, so the record lets
-    ranking and analysis tell them apart, like AlphaFold 2's. Entries are keyed
-    ``seed-N`` like AF3's other outputs and are read back from the runner's own config,
-    so they state what actually ran. The record never fails a finished prediction: an
-    unreadable file is replaced, and any other failure is logged and skipped.
+    Entries are keyed ``seed-N`` like AF3's other outputs and are read back from the
+    runner's own config, so they state what actually ran. The record never fails a
+    finished prediction: a failure to describe the kernels is logged and skipped.
     """
     fused_triangles = getattr(model_runner, "fused_triangles", None)
     tokens = example.get("token_index", example.get("aatype"))
@@ -211,22 +209,7 @@ def _record_inference_kernels(
     except Exception as exc:  # provenance only; the prediction itself succeeded
         logging.warning("Not recording the inference kernels of %s: %s", key, exc)
         return
-    path = pathlib.Path(output_dir) / "inference_kernels.json"
-    try:
-        records = json.loads(path.read_text()) if path.exists() else {}
-        if not isinstance(records, dict):
-            raise ValueError("not a JSON object")
-    except (OSError, ValueError) as exc:
-        logging.warning("Replacing unreadable %s: %s", path, exc)
-        records = {}
-    records[key] = record
-    try:
-        os.makedirs(output_dir, exist_ok=True)
-        temporary = path.with_suffix(".json.tmp")
-        temporary.write_text(json.dumps(records, indent=2, sort_keys=True) + "\n")
-        temporary.replace(path)
-    except OSError as exc:
-        logging.warning("Not recording the inference kernels of %s: %s", key, exc)
+    record_inference_kernels(output_dir, key, record)
 
 
 def predict_structure(

alphapulldown/prediction/fold_preparation.py

@@ -19,6 +19,7 @@
 from absl import logging
 
 from alphapulldown.objects import ChoppedObject, MonomericObject, MultimericObject
+from alphapulldown.utils.output_paths import FEATURE_METADATA_DIR
 
 
 # Most filesystems reject paths longer than this.
@@ -94,14 +95,16 @@ def _copy_feature_metadata(interactors: List[Any], flags: Any, output_dir: str)
             )
             continue
         latest = max(metadata_files, key=os.path.getmtime)
-        destination = os.path.join(output_dir, os.path.basename(latest))
+        metadata_dir = os.path.join(output_dir, FEATURE_METADATA_DIR)
+        os.makedirs(metadata_dir, exist_ok=True)
+        destination = os.path.join(metadata_dir, os.path.basename(latest))
         if latest.endswith(".json.xz"):
             destination = destination[: -len(".xz")]
             logging.info("Decompressing %s to %s", latest, destination)
             with lzma.open(latest, "rb") as source, open(destination, "wb") as target:
                 target.write(source.read())
         else:
-            logging.info("Copying %s to %s", latest, output_dir)
+            logging.info("Copying %s to %s", latest, metadata_dir)
             shutil.copyfile(latest, destination)
 
 

alphapulldown/prediction/inference_kernels.py

@@ -0,0 +1,62 @@
+"""Which kernels predicted a fold: ``alphapulldown/inference_kernels.json`` in its directory.
+
+Fused and original kernels differ slightly in their numerics, so the record lets ranking
+and analysis tell predictions apart. Each backend adds one entry per unit it predicts,
+keyed ``model_<name>`` by AlphaFold 2 and ``seed-<N>`` by AlphaFold 3. Versions up to
+2.10 kept the file at the top of the fold's directory; the next entry written moves it.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import Any, Mapping
+
+from alphapulldown.utils.file_handling import open_atomic
+from alphapulldown.utils.output_paths import alphapulldown_dir
+
+INFERENCE_KERNELS = "inference_kernels.json"
+
+
+def inference_kernels_path(output_dir) -> str:
+    """Where the record of the fold in ``output_dir`` is kept."""
+    return os.path.join(alphapulldown_dir(output_dir), INFERENCE_KERNELS)
+
+
+def _legacy_path(output_dir) -> str:
+    return os.path.join(os.fspath(output_dir), INFERENCE_KERNELS)
+
+
+def read_inference_kernels(output_dir) -> dict[str, Any]:
+    """The fold's record, wherever this version or an earlier one kept it; empty if
+    there is none or it is unreadable."""
+    for path in (inference_kernels_path(output_dir), _legacy_path(output_dir)):
+        if not os.path.exists(path):
+            continue
+        try:
+            with open(path, encoding="utf-8") as handle:
+                records = json.load(handle)
+            if not isinstance(records, dict):
+                raise ValueError("not a JSON object")
+        except (OSError, ValueError) as exc:
+            logging.warning("Replacing unreadable %s: %s", path, exc)
+            return {}
+        return records
+    return {}
+
+
+def record_inference_kernels(output_dir, key: str, entry: Mapping[str, Any]) -> None:
+    """Add ``entry`` under ``key``, keeping the other entries. Provenance never fails a
+    finished prediction: an unreadable record is replaced, a failed write only logged."""
+    records = read_inference_kernels(output_dir)
+    records[key] = dict(entry)
+    try:
+        os.makedirs(alphapulldown_dir(output_dir), exist_ok=True)
+        with open_atomic(inference_kernels_path(output_dir)) as handle:
+            json.dump(records, handle, indent=2, sort_keys=True)
+            handle.write("\n")
+        if os.path.exists(_legacy_path(output_dir)):
+            os.remove(_legacy_path(output_dir))
+    except OSError as exc:
+        logging.warning("Not recording the inference kernels of %s: %s", key, exc)

alphapulldown/prediction/resume.py

@@ -1,7 +1,7 @@
 """Writing prediction outputs so that an interrupted run can be resumed.
 
 Each fold directory holds a record of the run that predicted it
-(``alphapulldown_prediction.json``): the settings that shape the fold's outputs, and
+(``alphapulldown/prediction.json``): the settings that shape the fold's outputs, and
 whether the run finished. A backend starts the record before predicting and marks it
 complete as the last step of post-processing, so a later run with ``--allow_resume`` can
 skip a finished fold without inspecting, or redoing, any of its outputs.
@@ -16,10 +16,16 @@
 from absl import logging
 
 from alphapulldown.utils.file_handling import open_atomic
+from alphapulldown.utils.output_paths import alphapulldown_dir
 
 PathLike = Union[str, "os.PathLike[str]"]
 
-FOLD_RECORD_NAME = "alphapulldown_prediction.json"
+FOLD_RECORD_NAME = "prediction.json"
+
+
+def fold_record_path(output_dir: PathLike) -> str:
+    """Where the record of the fold in ``output_dir`` is kept."""
+    return os.path.join(alphapulldown_dir(output_dir), FOLD_RECORD_NAME)
 
 
 def _as_recorded(settings: Mapping[str, Any]) -> Dict[str, Any]:
@@ -29,7 +35,7 @@ def _as_recorded(settings: Mapping[str, Any]) -> Dict[str, Any]:
 
 def read_fold_record(output_dir: PathLike) -> Dict[str, Any]:
     """The record of the run that predicted the fold in ``output_dir``; empty if none."""
-    path = os.path.join(os.fspath(output_dir), FOLD_RECORD_NAME)
+    path = fold_record_path(output_dir)
     try:
         with open(path, "r", encoding="utf-8") as handle:
             record = json.load(handle)
@@ -44,8 +50,8 @@ def read_fold_record(output_dir: PathLike) -> Dict[str, Any]:
 def _write_fold_record(output_dir: PathLike, record: Mapping[str, Any]) -> None:
     from alphapulldown import __version__
 
-    os.makedirs(output_dir, exist_ok=True)
-    with open_atomic(os.path.join(os.fspath(output_dir), FOLD_RECORD_NAME)) as handle:
+    os.makedirs(alphapulldown_dir(output_dir), exist_ok=True)
+    with open_atomic(fold_record_path(output_dir)) as handle:
         json.dump(
             {**record, "alphapulldown_version": __version__}, handle, indent=2, sort_keys=True
         )

alphapulldown/scripts/convert_to_modelcif.py

@@ -31,6 +31,7 @@
 
 from alphapulldown.utils.file_handling import iter_seqs
 from alphapulldown.utils.feature_metadata import (
+    copied_feature_metadata_paths,
     decode_metadata_from_description,
     read_feature_metadata,
 )
@@ -805,11 +806,7 @@ def _register_metadata(monomer_name: str, metadata: dict) -> None:
             if new_entry not in fasta_dicts:
                 fasta_dicts.append(new_entry)
 
-    feature_json_files = sorted(
-        set(glob.glob(os.path.join(out_dir, "*_feature_metadata_*.json")))
-        | set(glob.glob(os.path.join(out_dir, "*_feature_metadata_*.json.xz")))
-    )
-    for feature_json in feature_json_files:
+    for feature_json in copied_feature_metadata_paths(out_dir):
         _file_exists_or_exit(
             feature_json, f"No feature metadata file '{feature_json}' found."
         )