fix(notebook): recompute projections on every generate

apps/protspace/CLAUDE.md

@@ -269,17 +269,18 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
 | `test_interpro_annotation_retriever.py` | InterPro API mocking, parsing |
 | `test_settings_converter.py` | Settings table ↔ visualization state conversion |
 | `test_uniprot_annotation_retriever.py` | UniProt API mocking, inactive entry resolution |
-| `test_pipeline_utils.py` | ReductionPipeline, EmbeddingSet, method parsing, multi-input merging, inline param overrides |
+| `test_pipeline_utils.py` | ReductionPipeline, notebook input/annotation/projection cache identity, EmbeddingSet, method parsing, multi-input merging, inline param overrides |
 | `test_stats.py` | Projection statistics: elbow, annotation-based validity (silhouette/DBI/CH per annotation), auto-cluster ARI/NMI agreement, auto-cluster self-validity (filed under the membership column, gated on it, and equal to driving `AnnotationValidityStatistic` directly so an out-of-band re-score cannot drift), faithfulness (dual continuity + global metrics), cluster-selection (elbow/silhouette/both), subsample determinism/order-invariance, silhouette consistency, `_align` no-id guard, silhouette→elbow fallback |
 | `test_stats_cli.py` | `protspace stats` CLI + `prepare` stats wiring, `--stats-annotation` (auto/list) wiring, `--settings-out` guard, `--cluster-selection` validation |
 | `test_stats_carriage.py` | Routing rows to bundle parts (metadata quality, annotation columns, cluster legend) |
 | `test_stats_bundle.py` | Optional 5th (statistics) bundle part round-trip |
 | `test_annotation_select.py` | Annotation selection: suitability filter (cardinality/numeric/id-like exclusion), `auto` vs explicit-list label building (explicit names bypass the heuristic), missing-value dropping |
 | `test_annotation_validity.py` | `AnnotationValidityStatistic`: silhouette/DBI/CH scored per annotation on `ctx.coords`, embedding vs. projection `space_kind`, missing-value exclusion, single-category no-op, id-canonical subsample determinism |
 | `test_biocentral_embedder.py` | Biocentral API client, embedding flow |
-| `test_backend_switch.py` | Embedding backend switch: `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
+| `test_backend_switch.py` | Embedding backend switch: notebook cache ownership/reuse, `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
 | `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, empty-output guard, esm2_8m end-to-end + resume (slow) |
 | `test_fasta.py` | FASTA parsing, edge cases, CSV annotation loading |
+| `test_query.py` | UniProt query FASTA download validation and atomic cache publication |
 | `test_biocentral_retriever.py` | Biocentral prediction retriever (TMbed parsing, per-sequence) |
 | `test_taxonomy_annotation_retriever.py` | Taxonomy via UniProt Taxonomy API (mocked + integration) |
 | `test_config_validation.py` | DimensionReductionConfig parameter validation |

apps/protspace/docs/cli.md

@@ -348,11 +348,12 @@ With `--keep-tmp` (default), all intermediate results are cached in `{output}/tm
 | ----------- | ---- | -------------- |
 | FASTA sequences | `sequences.fasta` | Skip UniProt query download |
 | Embeddings | `{embedder}.h5` | Skip already-embedded proteins |
-| Annotations | `all_annotations.parquet` | Fetch only missing annotation sources |
+| Annotations | `all_annotations.parquet` | Fetch missing sources when the cache covers every requested identifier; rebuild when requested identifiers are absent |
 | Similarity matrix | `similarity_matrix.npy` | Skip MMseqs2 recomputation |
 | DR projections | `proj_{name}_{method}_{hash}.npz` | Skip dimensionality reduction |
 
 - Annotation cache always includes scores regardless of `--no-scores`
+- An annotation cache may cover more proteins than the current run; those extra rows are filtered later. If any requested identifier is absent, annotations are rebuilt for the current input and the cache is replaced.
 - DR projection caches are keyed by embedding name, method, dimensions, and all parameters — changing any parameter creates a new cache entry
 - Use `--refetch all` to bypass all caches, or `--refetch <stages>` selectively (e.g., `--refetch ted,biocentral`)
 

apps/protspace/notebooks/ProtSpace_Preparation.ipynb

@@ -57,7 +57,10 @@
     "    PipelineConfig,\n",
     "    ReducerParams,\n",
     "    ReductionPipeline,\n",
+    "    _embedding_cache_path,\n",
+    "    _input_cache_dir,\n",
     "    parse_methods_arg,\n",
+    "    _query_fasta_cache_path,\n",
     ")"
    ]
   },
@@ -759,8 +762,8 @@
     "\n",
     "        out_dir = Path(\"output\")\n",
     "        out_dir.mkdir(exist_ok=True)\n",
-    "        cache_dir = out_dir / \"tmp\"\n",
-    "        cache_dir.mkdir(exist_ok=True)\n",
+    "        cache_root = out_dir / \"tmp\"\n",
+    "        cache_root.mkdir(exist_ok=True)\n",
     "        output_path = out_dir / \"data.parquetbundle\"\n",
     "\n",
     "        step_html = HTML(value=\"<b>Step 1/4: Loading embeddings...</b>\")\n",
@@ -776,7 +779,7 @@
     "                    print(\"Select at least one embedder.\")\n",
     "                    return\n",
     "                step_html.value = \"<b>Step 1/6: Downloading FASTA...</b>\"\n",
-    "                fasta_cache = cache_dir / \"sequences.fasta\"\n",
+    "                fasta_cache = _query_fasta_cache_path(cache_root, inp[\"query\"])\n",
     "                if fasta_cache.exists() and fasta_cache.stat().st_size > 0:\n",
     "                    from protspace.data.loaders.query import (\n",
     "                        extract_identifiers_from_fasta,\n",
@@ -792,6 +795,8 @@
     "                if not headers:\n",
     "                    print(f\"No sequences found for query: {inp['query']}\")\n",
     "                    return\n",
+    "                cache_dir = _input_cache_dir(cache_root, fasta_path)\n",
+    "                cache_dir.mkdir(parents=True, exist_ok=True)\n",
     "                backend, _emb_cfg = _resolve_backend_and_config()\n",
     "                embs = _drop_incompatible(embs, backend)\n",
     "                if not embs:\n",
@@ -804,7 +809,7 @@
     "                        emb_name,\n",
     "                        backend=backend,\n",
     "                        embed_config=_emb_cfg,\n",
-    "                        embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
+    "                        embedding_cache=_embedding_cache_path(cache_dir, emb_name, backend),\n",
     "                    )\n",
     "                    emb_set.fasta_path = fasta_path\n",
     "                    embedding_sets.append(emb_set)\n",
@@ -813,6 +818,9 @@
     "                if not embs:\n",
     "                    print(\"Select at least one embedder.\")\n",
     "                    return\n",
+    "                fasta_path = Path(inp[\"path\"])\n",
+    "                cache_dir = _input_cache_dir(cache_root, fasta_path)\n",
+    "                cache_dir.mkdir(parents=True, exist_ok=True)\n",
     "                backend, _emb_cfg = _resolve_backend_and_config()\n",
     "                embs = _drop_incompatible(embs, backend)\n",
     "                if not embs:\n",
@@ -821,23 +829,25 @@
     "                for emb_name in embs:\n",
     "                    step_html.value = f\"<b>Step 1/5: Computing {emb_name} embeddings ({backend})...</b>\"\n",
     "                    emb_set = embed_fasta(\n",
-    "                        Path(inp[\"path\"]),\n",
+    "                        fasta_path,\n",
     "                        emb_name,\n",
     "                        backend=backend,\n",
     "                        embed_config=_emb_cfg,\n",
-    "                        embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
+    "                        embedding_cache=_embedding_cache_path(cache_dir, emb_name, backend),\n",
     "                    )\n",
-    "                    emb_set.fasta_path = Path(inp[\"path\"])\n",
+    "                    emb_set.fasta_path = fasta_path\n",
     "                    embedding_sets.append(emb_set)\n",
     "            else:\n",
     "                h5_path = Path(inp[\"path\"])\n",
+    "                cache_dir = _input_cache_dir(cache_root, h5_path)\n",
+    "                cache_dir.mkdir(parents=True, exist_ok=True)\n",
     "                name_override = inp.get(\"name\")\n",
     "                emb_set = load_h5([h5_path], name_override=name_override)\n",
     "                embedding_sets.append(emb_set)\n",
     "\n",
     "            n_proteins = len(embedding_sets[0].headers)\n",
     "\n",
-    "            # Build pipeline with caching enabled\n",
+    "            # Build pipeline with non-projection caching enabled\n",
     "            reducer_params = ReducerParams(\n",
     "                n_neighbors=pw[\"n_neighbors\"].value,\n",
     "                min_dist=pw[\"min_dist\"].value,\n",
@@ -852,6 +862,7 @@
     "                bundled=True,\n",
     "                keep_tmp=True,\n",
     "                intermediate_dir=cache_dir,\n",
+    "                refetch_stages=frozenset({\"projections\"}),\n",
     "                annotations=ann,\n",
     "                reducer_params=reducer_params,\n",
     "                stats=compute_stats_cb.value,\n",
@@ -860,7 +871,9 @@
     "\n",
     "            # Step 2: Annotations (cached after first run)\n",
     "            step_html.value = \"<b>Step 2/4: Fetching annotations...</b>\"\n",
-    "            metadata = pipeline._fetch_annotations(embedding_sets[0].headers)\n",
+    "            metadata = pipeline._fetch_annotations(\n",
+    "                embedding_sets[0].headers, embedding_sets\n",
+    "            )\n",
     "\n",
     "            # Step 3: Dimensionality reduction\n",
     "            step_html.value = \"<b>Step 3/4: Reducing dimensions...</b>\"\n",

apps/protspace/src/protspace/data/loaders/query.py

@@ -6,6 +6,7 @@
 
 import gzip
 import logging
+import os
 import tempfile
 from pathlib import Path
 
@@ -33,34 +34,48 @@ def query_uniprot(
 
     base_url = "https://rest.uniprot.org/uniprotkb/stream"
     params = {"compressed": "true", "format": "fasta", "query": query}
+    temp_gz_file: Path | None = None
+    staged_path: Path | None = None
+    extracted_path: Path | None = None
+    completed = False
 
     try:
         response = requests.get(base_url, params=params, stream=True)
         response.raise_for_status()
 
         # Download to temporary compressed file
-        temp_file = tempfile.NamedTemporaryFile(
-            mode="wb", suffix=".fasta.gz", delete=False
-        )
-        temp_gz_file = Path(temp_file.name)
-
         total_size = int(response.headers.get("content-length", 0))
-        with tqdm(
-            total=total_size, unit="B", unit_scale=True, desc="Downloading FASTA"
-        ) as pbar:
-            for chunk in response.iter_content(chunk_size=8192):
-                if chunk:
-                    temp_file.write(chunk)
-                    pbar.update(len(chunk))
-        temp_file.close()
+        with tempfile.NamedTemporaryFile(
+            mode="wb", suffix=".fasta.gz", delete=False
+        ) as temp_file:
+            temp_gz_file = Path(temp_file.name)
+            with tqdm(
+                total=total_size,
+                unit="B",
+                unit_scale=True,
+                desc="Downloading FASTA",
+            ) as pbar:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        temp_file.write(chunk)
+                        pbar.update(len(chunk))
 
         # Extract identifiers from compressed FASTA
         identifiers = _extract_identifiers_gz(temp_gz_file)
 
         # Extract FASTA to final location
-        if save_to:
-            extracted_path = save_to
-            extracted_path.parent.mkdir(parents=True, exist_ok=True)
+        if save_to is not None:
+            save_to = Path(save_to)
+            save_to.parent.mkdir(parents=True, exist_ok=True)
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                prefix=f".{save_to.name}.",
+                suffix=".tmp",
+                dir=save_to.parent,
+                delete=False,
+            ) as staged_file:
+                staged_path = Path(staged_file.name)
+            extracted_path = staged_path
         else:
             extracted_path = temp_gz_file.with_suffix("")
 
@@ -69,7 +84,19 @@ def query_uniprot(
             with open(extracted_path, "w") as out:
                 out.write(content)
 
-        temp_gz_file.unlink(missing_ok=True)
+        extracted_identifiers = extract_identifiers_from_fasta(extracted_path)
+        if extracted_identifiers != identifiers:
+            raise ValueError("Extracted FASTA identifiers do not match the download")
+
+        if save_to is not None:
+            current_umask = os.umask(0)
+            os.umask(current_umask)
+            staged_path.chmod(0o666 & ~current_umask)
+            staged_path.replace(save_to)
+            staged_path = None
+            extracted_path = save_to
+
+        completed = True
         logger.info(f"Downloaded and extracted {len(identifiers)} sequences")
 
         return identifiers, extracted_path
@@ -80,6 +107,13 @@ def query_uniprot(
     except Exception as e:
         logger.error(f"Error processing FASTA: {e}")
         raise
+    finally:
+        if temp_gz_file is not None:
+            temp_gz_file.unlink(missing_ok=True)
+        if staged_path is not None:
+            staged_path.unlink(missing_ok=True)
+        if not completed and save_to is None and extracted_path is not None:
+            extracted_path.unlink(missing_ok=True)
 
 
 def extract_identifiers_from_fasta(fasta_path: Path) -> list[str]:

apps/protspace/src/protspace/data/processors/pipeline.py

@@ -82,6 +82,26 @@ class PipelineConfig:
     reducer_params: ReducerParams = field(default_factory=ReducerParams)
 
 
+def _query_fasta_cache_path(cache_root: Path, query: str) -> Path:
+    """Return the retained FASTA path owned by one exact UniProt query."""
+    digest = hashlib.sha256(query.encode()).hexdigest()[:12]
+    return cache_root / "queries" / f"{digest}.fasta"
+
+
+def _input_cache_dir(cache_root: Path, input_path: Path) -> Path:
+    """Return the retained intermediate directory owned by one input file."""
+    digest = hashlib.sha256()
+    with input_path.open("rb") as source:
+        for chunk in iter(lambda: source.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return cache_root / "inputs" / digest.hexdigest()[:12]
+
+
+def _embedding_cache_path(cache_dir: Path, embedder: str, backend: str) -> Path:
+    """Return the H5 path owned by one input, model, and producing backend."""
+    return cache_dir / f"{backend}-{embedder}.h5"
+
+
 # Valid override parameter names (from ReducerParams fields)
 _VALID_OVERRIDE_KEYS = {f.name for f in fields(ReducerParams)}
 # Field types for coercion
@@ -394,8 +414,33 @@ def _fetch_annotations(
             intermediate_dir.mkdir(parents=True, exist_ok=True)
             cache_path = intermediate_dir / "all_annotations.parquet"
 
+            def fetch_current_annotations() -> pd.DataFrame:
+                api_df = ProteinAnnotationManager(
+                    headers=headers,
+                    annotations=annotations_list,
+                    output_path=cache_path,
+                    sequences=sequences,
+                ).to_pd()
+                return self._merge_csv(api_df, csv_df)
+
             if cache_path.exists():
                 cached_df = pd.read_parquet(cache_path)
+                cached_identifiers = (
+                    Counter(cached_df["identifier"].astype(str))
+                    if "identifier" in cached_df.columns
+                    else Counter()
+                )
+                requested_identifiers = Counter(map(str, headers))
+                missing_identifiers = requested_identifiers - cached_identifiers
+
+                if missing_identifiers:
+                    logger.warning(
+                        "Annotation cache is missing %d requested identifier(s); "
+                        "fetching annotations for the current identifiers",
+                        sum(missing_identifiers.values()),
+                    )
+                    return fetch_current_annotations()
+
                 cached_annotations = set(cached_df.columns) - {"identifier"}
 
                 if annotations_list is None:
@@ -476,13 +521,7 @@ def _fetch_annotations(
                 ).to_pd()
                 return self._merge_csv(api_df, csv_df)
             else:
-                api_df = ProteinAnnotationManager(
-                    headers=headers,
-                    annotations=annotations_list,
-                    output_path=cache_path,
-                    sequences=sequences,
-                ).to_pd()
-                return self._merge_csv(api_df, csv_df)
+                return fetch_current_annotations()
         else:
             api_df = ProteinAnnotationManager(
                 headers=headers,