From f5b2251208191e1028accb41afabc4cbdd8c0567 Mon Sep 17 00:00:00 2001 From: SonOfAnton Date: Mon, 7 Sep 2026 14:36:15 -0700 Subject: [PATCH 1/4] FAISS ANN search backend for proteogram similarity Adds an optional approximate-nearest-neighbour search path as an alternative to brute-force cosine similarity, for corpora where the O(N^2) scan dominates runtime. - proteogram/v2/faiss_search.py: FaissIndex wrapper (IVFFlat, and IVF-PQ for very large corpora). Embeddings are L2-normalised before indexing so inner-product search is true cosine similarity. - Img2Vec gains four thin delegators (build/save/load/similarities_faiss) so scripts keep a single entry point. - measure_similarity_v2.py: --faiss, --faiss_pq, --faiss_index_file. Brute-force remains the default; the index is cached to disk and reused unless --overwrite is passed. - Adds faiss-cpu to dependencies and relocks. Both search paths share the same preprocessing function so ANN results are directly comparable to the brute-force baseline. --- docs/improvements_v2.md | 1366 +++++++++++++++++++++++++++ proteogram/v2/__init__.py | 4 +- proteogram/v2/faiss_search.py | 327 +++++++ proteogram/v2/image_similarity.py | 100 ++ pyproject.toml | 1 + scripts/v2/measure_similarity_v2.py | 66 +- uv.lock | 314 +++--- 7 files changed, 1981 insertions(+), 197 deletions(-) create mode 100644 docs/improvements_v2.md create mode 100644 proteogram/v2/faiss_search.py diff --git a/docs/improvements_v2.md b/docs/improvements_v2.md new file mode 100644 index 0000000..2e9478f --- /dev/null +++ b/docs/improvements_v2.md @@ -0,0 +1,1366 @@ +# Proteogram v2 — Three Targeted Improvements + +## Overview + +This document covers the design, implementation, and validation plan for three independent but complementary improvements to the Proteogram v2 pipeline: + +| # | Name | File(s) Affected | Impact | +|---|------|-----------------|--------| +| [1](#1-faiss-approximate-nearest-neighbour-search) | FAISS Approximate Nearest Neighbour Search | `proteogram/v2/image_similarity.py`, `scripts/v2/measure_similarity_v2.py` | Scales corpus search from minutes to milliseconds | +| [2](#2-global-percentile-normalisation) | Global Percentile Normalisation | `proteogram/v2/proteogram.py`, `scripts/v2/create_v2_proteograms.py` | Preserves physically meaningful inter-protein scale | +| [3](#3-grad-cam-explainability) | Grad-CAM Explainability | `proteogram/v2/image_similarity.py` (new method), new `scripts/v2/explain_similarity.py` | Residue-pair attribution for any similar pair | + +Each section follows the same structure: motivation → design decisions → full implementation → validation steps. + +--- + +## Operational Notes (May 2026): Environment Setup + Long v2 Runs + +This section documents practical lessons from running the current v2 pipeline on Linux with mixed toolchains (`uv`, local Conda bootstrap, OpenMM). + +### A. Why `create_v2_proteograms.py` may be slow + +`scripts/v2/create_v2_proteograms.py` runs a full MD pipeline per protein (minimization + NPT + NVT + production) before image export. If OpenMM CUDA is unavailable, this falls back to CPU and runtime increases significantly. + +At default MD lengths, CPU runtime can be many minutes per protein; with 2,008 proteins this can become multi-day if not accelerated. + +### B. Distinguish PyTorch CUDA vs OpenMM CUDA + +It is common to have: + +- `torch.cuda.is_available() == True` +- OpenMM platforms = `['Reference', 'CPU', 'OpenCL']` + +In this case, similarity scripts can use GPU (PyTorch), but MD in v2 proteogram generation still runs without CUDA. + +Check OpenMM platforms directly: + +```bash +python - <<'PY' +from openmm import Platform +names = [Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())] +print('OpenMM platforms:', names) +print('CUDA available:', 'CUDA' in names) +PY +``` + +### C. Conda bootstrap pitfall encountered + +In this run, `conda` was pointing to a local bootstrap install under: + +`scripts/v2/exit/bin/conda` + +This caused solver and lock issues (e.g., sqlite lock, libmamba plugin mismatch), and prevented reliable environment creation. + +Recommended safeguards: + +1. Confirm which conda is active (`which conda`, `conda info --base`). +2. Prefer a stable system conda/mamba/micromamba install for OpenMM-CUDA env creation. +3. If needed, force classic solver when libmamba plugin is unavailable. + +### D. Minimal reliable runbook (CUDA-capable OpenMM env) + +```bash +# 1) create env with python 3.11 (recommended for this project stack) +conda create -n proteogram-openmm-cuda -c conda-forge python=3.11 openmm pdbfixer -y + +# 2) activate env +conda activate proteogram-openmm-cuda + +# 3) verify OpenMM CUDA platform visibility +python - <<'PY' +from openmm import Platform +print([Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())]) +PY + +# 4) install project in editable mode +cd /path/to/proteogram +pip install -e . + +# 5) run v2 proteogram creation +cd scripts/v2 +python create_v2_proteograms.py --overwrite +``` + +### E. Monitoring a long run + +During `create_v2_proteograms.py`, JPG outputs are written incrementally (not only at end). To watch output growth: + +```bash +# from repo root +watch -n 5 'find data/scope2.08_all_proteograms_v2 -maxdepth 1 -name "*.jpg" | wc -l' + +# or from scripts/v2 +watch -n 5 'find ../data/scope2.08_all_proteograms_v2 -maxdepth 1 -name "*.jpg" | wc -l' +``` + +### F. Current observed status + +- Run is progressing through structures and skipping out-of-range chains (`sequence length outside [20, 200]`) as designed. +- Output directory image count should rise continuously as proteins complete. +- If OpenMM CUDA remains unavailable, OpenCL/CPU execution is expected and slower than CUDA. + +--- + +## 1. FAISS Approximate Nearest Neighbour Search + +### 1.1 Motivation + +The current `Img2Vec.similarities()` method computes cosine similarity between every pair of embeddings in the corpus: + +```python +# proteogram/v2/image_similarity.py — existing inner loop (O(N²)) +for image_path_i, embedding_i in tqdm(self.dataset.items()): + for image_path_j, embedding_j in self.dataset.items(): + sim = cosine(embedding_i, embedding_j)[0].item() +``` + +For a corpus of N proteins this is O(N²) in both time and sequential memory access. At the current SCOPe 2.08 scale (~100 K domains) this already takes tens of minutes. The AlphaFold Database (AFDB) contains ~200 million predicted structures — brute-force search there would take weeks. + +FAISS (Facebook AI Similarity Search) replaces this with an **Inverted File index with Product Quantisation (IVF-PQ)** that gives sub-linear query time with controllable recall-accuracy trade-offs. + +### 1.2 Design Decisions + +#### Index type: `IndexIVFFlat` for small corpora, `IndexIVFPQ` for large + +| Corpus size | Recommended index | Reason | +|-------------|------------------|--------| +| ≤ 100 K | `IndexIVFFlat` | Exact L2/IP; no quantisation error; fast enough | +| 100 K – 10 M | `IndexIVFPQ` | 8–32× memory reduction; ~1% recall loss | +| > 10 M | `IndexIVFPQ` + `OPQ` pre-rotation | Best recall at extreme scale | + +The `Img2Vec` class will default to `IndexIVFFlat` and let callers opt into `IndexIVFPQ`. + +#### Inner-product (IP) vs. L2 + +FAISS supports both. Because the rest of the codebase uses **cosine similarity**, we L2-normalise embeddings before indexing and use **inner product** — on unit vectors, inner product equals cosine similarity exactly. This avoids any change to how scores are interpreted. + +#### `nlist` (number of Voronoi cells) + +A rule of thumb is `nlist = sqrt(N)`. For 100 K vectors: `nlist = 316`. For 10 M vectors: `nlist = 3162`. These values will be set automatically if not specified. + +#### `nprobe` (cells searched at query time) + +Higher `nprobe` → better recall, slower query. Default: `nprobe = max(1, nlist // 10)`. Can be tuned by the caller. + +#### Backward compatibility + +The existing `similarities()` method signature must not change. FAISS is added as an optional code path activated by passing `use_faiss=True`. The `.sim_dict` output format stays identical so `measure_similarity_v2.py` and `evaluate_methods_v2.py` require no changes. + +### 1.3 New Dependency + +```toml +# pyproject.toml — add to [project] dependencies +"faiss-cpu>=1.8; extra != 'cuda12'", +"faiss-gpu>=1.8; extra == 'cuda12'", +``` + +Or install manually: +```bash +# CPU +uv add faiss-cpu + +# GPU (CUDA 12) +uv add faiss-gpu +``` + +### 1.4 Implementation + +#### 1.4.1 New method: `Img2Vec.build_faiss_index()` + +Add to `proteogram/v2/image_similarity.py` inside the `Img2Vec` class: + +```python +def build_faiss_index(self, + use_pq: bool = False, + nlist: int = None, + nprobe: int = None, + pq_m: int = 8, + pq_nbits: int = 8) -> None: + """Build a FAISS index from the currently loaded embedding dataset. + + Embeddings are L2-normalised before indexing so that inner-product + search is equivalent to cosine similarity. + + Args: + use_pq: If True, use IVF-PQ (compressed) index. Recommended for + corpora > 100 K. Defaults to False (IVFFlat, exact). + nlist: Number of Voronoi cells. Defaults to sqrt(N). + nprobe: Number of cells to search at query time. Higher = better + recall, slower query. Defaults to nlist // 10. + pq_m: Number of PQ sub-quantisers (IVF-PQ only). Must divide + the embedding dimension evenly. + pq_nbits: Bits per sub-quantiser (IVF-PQ only). 8 is standard. + """ + try: + import faiss + except ImportError: + raise ImportError( + "faiss is required for build_faiss_index(). " + "Install with: uv add faiss-cpu (or faiss-gpu for GPU builds)." + ) + + if not self.dataset: + raise RuntimeError("embed_dataset() must be called before build_faiss_index().") + + # Stack embeddings and keys in a consistent order + keys = list(self.dataset.keys()) + vecs = torch.cat([self.dataset[k].cpu() for k in keys]).float() # (N, d) + + # L2-normalise so inner product == cosine similarity + faiss.normalize_L2(vecs.numpy()) + + N, d = vecs.shape + _nlist = nlist if nlist is not None else max(1, int(N ** 0.5)) + _nprobe = nprobe if nprobe is not None else max(1, _nlist // 10) + + # Build quantiser (flat inner-product) + quantiser = faiss.IndexFlatIP(d) + + if use_pq: + # Ensure pq_m divides d evenly + while d % pq_m != 0 and pq_m > 1: + pq_m -= 1 + index = faiss.IndexIVFPQ(quantiser, d, _nlist, pq_m, pq_nbits, + faiss.METRIC_INNER_PRODUCT) + else: + index = faiss.IndexIVFFlat(quantiser, d, _nlist, + faiss.METRIC_INNER_PRODUCT) + + index.train(vecs.numpy()) + index.add(vecs.numpy()) + index.nprobe = _nprobe + + # Store on instance for re-use across queries + self._faiss_index = index + self._faiss_keys = keys # maps integer index → filename key + self._faiss_vecs_norm = vecs # keep L2-normalised vecs for query normalisation + + print(f"FAISS index built: {index.ntotal} vectors | d={d} | " + f"nlist={_nlist} | nprobe={_nprobe} | " + f"type={'IVF-PQ' if use_pq else 'IVFFlat'}") +``` + +#### 1.4.2 New method: `Img2Vec.similarities_faiss()` + +Add directly below `build_faiss_index()`: + +```python +def similarities_faiss(self, + n: int = 10, + save_result_images_dir: str = None, + pad_fn=None) -> float: + """Compute top-N similar images for every entry in the corpus using FAISS. + + Populates self.sim_dict with the same format as similarities(), so all + downstream scripts (evaluate_methods_v2.py, measure_similarity_v2.py) + work without modification. + + Call build_faiss_index() first. + + Args: + n: Top-N results per query (self-hit included + at rank 0; callers should request n+1 and + strip the self-hit themselves if needed). + save_result_images_dir: Optional directory to write result images. + pad_fn: Optional padding callable passed to save_images(). + + Returns: + float: Wall-clock seconds spent in FAISS search (excludes image saving). + """ + try: + import faiss + except ImportError: + raise ImportError("faiss not installed. Run: uv add faiss-cpu") + + if not hasattr(self, '_faiss_index'): + raise RuntimeError("Call build_faiss_index() before similarities_faiss().") + + keys = self._faiss_keys + vecs = self._faiss_vecs_norm.numpy() # already L2-normalised + + start = time() + # Batch query: search all N vectors at once — single FAISS call + scores_matrix, indices_matrix = self._faiss_index.search(vecs, n + 1) + elapsed = time() - start + + # Build sim_dict in the same format as similarities() + self.sim_dict = {} + for i, key in enumerate(keys): + hits = [] + for rank in range(n + 1): + j = indices_matrix[i, rank] + if j < 0: # FAISS pads with -1 when fewer results exist + continue + target_key = keys[j] + score = float(scores_matrix[i, rank]) + hits.append((target_key, score)) + self.sim_dict[key] = hits # includes self-hit at rank 0 + + if save_result_images_dir: + for image_path in self.sim_dict: + self.save_images(os.path.join(self.files[0].rsplit('/', 1)[0], image_path), + save_result_images_dir, pad_fn=pad_fn) + + return elapsed +``` + +#### 1.4.3 FAISS index persistence + +Add two methods for saving and loading the built index: + +```python +def save_faiss_index(self, index_path: str) -> None: + """Persist the FAISS index and key mapping to disk. + + Args: + index_path: File path for the index (e.g. 'corpus.faiss'). + A companion '.keys.pkl' file is written + alongside for the key mapping. + """ + import faiss, pickle + faiss.write_index(self._faiss_index, index_path) + keys_path = index_path + '.keys.pkl' + with open(keys_path, 'wb') as f: + pickle.dump(self._faiss_keys, f) + print(f"Saved FAISS index → {index_path}") + print(f"Saved key mapping → {keys_path}") + + +def load_faiss_index(self, index_path: str) -> None: + """Load a previously saved FAISS index and key mapping. + + Args: + index_path: Path to the '.faiss' index file. + """ + import faiss, pickle + self._faiss_index = faiss.read_index(index_path) + keys_path = index_path + '.keys.pkl' + with open(keys_path, 'rb') as f: + self._faiss_keys = pickle.load(f) + # Reconstruct normalised vecs for future queries (needed for single-query search) + keys = self._faiss_keys + vecs = torch.cat([self.dataset[k].cpu() for k in keys]).float() + faiss.normalize_L2(vecs.numpy()) + self._faiss_vecs_norm = vecs + print(f"Loaded FAISS index from {index_path} " + f"({self._faiss_index.ntotal} vectors)") +``` + +#### 1.4.4 Update `measure_similarity_v2.py` + +Add `--faiss` flag and wire it up: + +```python +# Add to the argparse block +parser.add_argument('--faiss', action='store_true', + help='Use FAISS ANN index for similarity search instead of ' + 'brute-force cosine similarity. Much faster for large corpora.') +parser.add_argument('--faiss_index_file', type=str, default=None, + help='Path to save/load the FAISS index. Defaults to ' + 'embed_file with .faiss extension.') +parser.add_argument('--faiss_pq', action='store_true', + help='Use IVF-PQ compressed index (recommended for > 100K proteins). ' + 'Slightly lower recall but much lower memory.') + +# Replace the similarities() call block with: +if args.faiss: + faiss_index_file = args.faiss_index_file or embed_file.replace('.pkl', '.faiss') + if os.path.exists(faiss_index_file) and not args.overwrite: + print(f'Loading existing FAISS index from {faiss_index_file}') + img_sim.load_faiss_index(faiss_index_file) + else: + print('Building FAISS index ...') + img_sim.build_faiss_index(use_pq=args.faiss_pq) + img_sim.save_faiss_index(faiss_index_file) + sim_time = img_sim.similarities_faiss( + n=n_results, + save_result_images_dir=None, + pad_fn=pad_to_size) +else: + sim_time = img_sim.similarities(n=n_results, + save_result_images_dir=None, + pad_fn=pad_to_size) +``` + +#### 1.4.5 Single-protein query update in `query_similar_proteins.py` + +```python +# Replace the inner loop in query_similar_proteins.py with: +def query_with_faiss(img_sim, query_embedding, top_k, corpus_dir): + """Query a built FAISS index with a single new embedding.""" + import faiss + import numpy as np + + query_vec = query_embedding.cpu().float().numpy() # (1, d) + faiss.normalize_L2(query_vec) + scores, indices = img_sim._faiss_index.search(query_vec, top_k + 1) + + results = [] + for rank in range(top_k + 1): + j = indices[0, rank] + if j < 0: + continue + key = img_sim._faiss_keys[j] + score = float(scores[0, rank]) + if key != os.path.basename(query_path): # skip self-hit if present + results.append((key, score)) + if len(results) >= top_k: + break + return results +``` + +### 1.5 Validation Steps + +#### Step 1 — Recall parity test (automated) + +Run both methods on the eval set and assert that FAISS Recall@K ≥ 0.99 × brute-force Recall@K at every SCOPe level: + +```python +# scripts/v2/tests/test_faiss_recall.py +import pytest +from proteogram.v2 import Img2Vec +import torch, pickle, os + +EMBED_FILE = os.environ.get('EMBED_FILE', 'corpus_embeddings.pkl') + +@pytest.fixture(scope='module') +def img_sim(): + sim = Img2Vec('resnet_ft', dataset_dir=[], device='cpu') + with open(EMBED_FILE, 'rb') as f: + sim.dataset = pickle.load(f) + return sim + +def test_faiss_topk_recall_at_5(img_sim): + """FAISS top-5 results should overlap ≥99% with brute-force top-5.""" + TOP_K = 5 + # Brute-force + img_sim.similarities(n=TOP_K) + bf_dict = {k: set(t for t, _ in v[:TOP_K]) for k, v in img_sim.sim_dict.items()} + + # FAISS IVFFlat + img_sim.build_faiss_index(use_pq=False) + img_sim.similarities_faiss(n=TOP_K) + faiss_dict = {k: set(t for t, _ in v[1:TOP_K+1]) for k, v in img_sim.sim_dict.items()} + + overlaps = [] + for key in bf_dict: + if key in faiss_dict: + overlap = len(bf_dict[key] & faiss_dict[key]) / TOP_K + overlaps.append(overlap) + + mean_recall = sum(overlaps) / len(overlaps) + print(f'Mean FAISS/BF overlap at top-{TOP_K}: {mean_recall:.4f}') + assert mean_recall >= 0.99, f'FAISS recall too low: {mean_recall:.4f}' +``` + +Run: +```bash +EMBED_FILE=/path/to/corpus_embeddings.pkl pytest scripts/v2/tests/test_faiss_recall.py -v +``` + +#### Step 2 — Timing benchmark + +```bash +# Brute-force +time python measure_similarity_v2.py --no-embed + +# FAISS IVFFlat +time python measure_similarity_v2.py --no-embed --faiss + +# FAISS IVF-PQ (large corpus) +time python measure_similarity_v2.py --no-embed --faiss --faiss_pq +``` + +Expected results on ~10 K eval set: + +| Method | Expected time | +|--------|--------------| +| Brute-force | ~2–5 min | +| FAISS IVFFlat | < 5 sec | +| FAISS IVF-PQ | < 2 sec | + +#### Step 3 — MAP@K parity + +Run `evaluate_methods_v2.py` on outputs from both methods. FAISS MAP@K should be within ±0.005 of brute-force MAP@K at all SCOPe levels. Any larger gap indicates the `nprobe` needs increasing. + +#### Step 4 — Index round-trip test + +```python +# Verify save/load produces identical results +img_sim.build_faiss_index() +img_sim.save_faiss_index('/tmp/test.faiss') + +img_sim2 = Img2Vec(model_file, dataset_dir=[], device='cpu') +img_sim2.dataset = img_sim.dataset +img_sim2.load_faiss_index('/tmp/test.faiss') + +img_sim.similarities_faiss(n=5) +img_sim2.similarities_faiss(n=5) + +for key in img_sim.sim_dict: + assert img_sim.sim_dict[key] == img_sim2.sim_dict[key], f"Mismatch at {key}" +print("Round-trip test passed.") +``` + +--- + +## 2. Global Percentile Normalisation + +> **Update:** the code snippets below reflect the original design. The shipped +> `--save_npy_matrices` implementation initially had a bug that fed +> already-normalised pixel data (not raw physical-unit energies) into +> `compute_norm_stats.py`, defeating the purpose of this feature. See +> [`percentile_normalisation_bug_fix_and_validation.md`](percentile_normalisation_bug_fix_and_validation.md) +> for the bug, the fix, and the new `validate_normalisation.py` tool for +> measuring before/after impact. + +### 2.1 Motivation + +The current `ProteogramV2.normalize_map()` applies **per-protein min-max normalisation** independently to each energy channel: + +```python +# proteogram/v2/proteogram.py — current implementation +arr = ((arr - arr.min()) * (1 / (arr.max() - arr.min()) * 255)).astype('uint8') +``` + +This has a critical flaw: every protein's energy map is stretched to fill the full [0, 255] dynamic range, regardless of the actual energy magnitudes. A small, weakly-interacting loop region and a tightly-packed hydrophobic core will produce identical grey levels after normalisation. The model never sees the absolute energy scale — only the relative rank order within each protein. + +**Concrete example**: Protein A has VdW attractive energies in [-50, -5] kJ/mol and Protein B has VdW attractive energies in [-200, -20] kJ/mol. After per-protein normalisation, both are mapped to [0, 255]. A CNN comparing the two images cannot tell that Protein B has 4× stronger packing. + +Global percentile normalisation computes bounds from the entire training corpus once, then applies those fixed bounds to every protein — preserving inter-protein energy scale in the pixel values. + +### 2.2 Design Decisions + +#### Percentile instead of global min/max + +Extreme outlier structures (e.g., very short peptides, structures with unusual post-translational modifications) would dominate a global min/max and compress most proteins into a narrow band. Using the **1st and 99th percentiles** clips ~2% of values but gives a robust, representative range. + +#### Separate bounds per channel + +Each of the 6 channels (VdW attractive, VdW repulsive, ES attractive, ES repulsive, distance, hydrophobicity) has a different physical unit and magnitude range. Bounds must be computed and stored independently per channel. + +#### Where bounds are stored + +A single JSON file `norm_stats.json` is written alongside the proteogram corpus. It is read at proteogram creation time when global normalisation is enabled. This keeps the bounds portable and version-controlled. + +#### Backward compatibility flag + +The new behaviour is opt-in via `--global_norm` flag in `create_v2_proteograms.py`. Per-protein normalisation remains the default so existing proteogram datasets are unaffected. + +### 2.3 New File: `scripts/v2/compute_norm_stats.py` + +This one-time script samples up to `--max_samples` existing `.npy` energy matrices (or re-runs the MD pipeline for a random subset) to compute global percentile bounds. + +```python +#!/usr/bin/env python +"""Compute global percentile normalisation statistics from a corpus of energy matrices. + +Run this ONCE after generating a representative sample of proteograms with +--save_npy_matrices (a new flag added to create_v2_proteograms.py). Outputs +norm_stats.json which is read by create_v2_proteograms.py --global_norm. + +Usage: + python compute_norm_stats.py \\ + --npy_dir /path/to/energy_matrices \\ + --out_file /path/to/norm_stats.json \\ + --low_pct 1.0 \\ + --high_pct 99.0 \\ + --max_samples 5000 +""" +import argparse +import json +import glob +import os +import numpy as np +from tqdm import tqdm + +CHANNEL_NAMES = [ + 'vdw_attractive', + 'vdw_repulsive', + 'es_attractive', + 'es_repulsive', + 'distance', + 'hydrophobicity', +] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--npy_dir', required=True, + help='Directory of .npy files, one per channel per protein ' + '(naming: _.npy).') + parser.add_argument('--out_file', required=True, + help='Output JSON path for normalisation bounds.') + parser.add_argument('--low_pct', type=float, default=1.0, + help='Lower percentile bound (default: 1.0).') + parser.add_argument('--high_pct', type=float, default=99.0, + help='Upper percentile bound (default: 99.0).') + parser.add_argument('--max_samples', type=int, default=5000, + help='Maximum number of energy matrices to sample per channel ' + '(default: 5000). More samples = more accurate statistics.') + args = parser.parse_args() + + # Collect all values per channel across the sampled corpus + channel_values = {ch: [] for ch in CHANNEL_NAMES} + + for channel in CHANNEL_NAMES: + files = sorted(glob.glob(os.path.join(args.npy_dir, f'*_{channel}.npy'))) + if not files: + print(f'WARNING: No .npy files found for channel "{channel}" in {args.npy_dir}') + continue + + # Random subsample if corpus is large + if len(files) > args.max_samples: + rng = np.random.default_rng(seed=42) + files = list(rng.choice(files, size=args.max_samples, replace=False)) + + print(f'Channel {channel}: sampling {len(files)} matrices ...') + for fpath in tqdm(files, desc=channel): + arr = np.load(fpath) + # Only include non-zero upper-triangle values (lower triangle is 0) + vals = arr[arr != 0].ravel() + channel_values[channel].append(vals) + + # Compute and store bounds + stats = {} + for channel in CHANNEL_NAMES: + if not channel_values[channel]: + stats[channel] = {'p_low': 0.0, 'p_high': 255.0} + continue + all_vals = np.concatenate(channel_values[channel]) + p_low = float(np.percentile(all_vals, args.low_pct)) + p_high = float(np.percentile(all_vals, args.high_pct)) + stats[channel] = {'p_low': p_low, 'p_high': p_high} + print(f' {channel}: p{args.low_pct}={p_low:.4f} p{args.high_pct}={p_high:.4f} ' + f'N={len(all_vals):,}') + + stats['_meta'] = { + 'low_pct': args.low_pct, + 'high_pct': args.high_pct, + 'n_files_per_channel': args.max_samples, + 'npy_dir': args.npy_dir, + } + + os.makedirs(os.path.dirname(os.path.abspath(args.out_file)), exist_ok=True) + with open(args.out_file, 'w') as f: + json.dump(stats, f, indent=2) + print(f'\nSaved normalisation stats → {args.out_file}') + + +if __name__ == '__main__': + main() +``` + +### 2.4 Implementation: Changes to `proteogram/v2/proteogram.py` + +#### 2.4.1 New static method: `normalize_map_global()` + +Add alongside the existing `normalize_map()`: + +```python +@staticmethod +def normalize_map_global(arr: np.ndarray, + p_low: float, + p_high: float) -> tuple[np.ndarray, str]: + """Normalise an energy/property map to [0, 255] using corpus-level percentile bounds. + + Unlike normalize_map(), which uses per-protein min/max, this method + applies fixed bounds derived from the full training corpus so that + inter-protein energy scale is preserved in pixel values. + + Zero values (unfilled lower-triangle entries) are mapped to 128 (mid-grey) + to distinguish them visually from true low-energy interactions (which map + near 0) — matching the existing gray padding convention in the training code. + + Args: + arr: Input energy matrix (upper triangle populated; lower = 0). + p_low: Lower percentile bound in physical units (kJ/mol or Å). + p_high: Upper percentile bound in physical units. + + Returns: + Tuple of (normalised uint8 array, error string or ''). + """ + err = '' + try: + scale = p_high - p_low + if scale == 0: + return np.full_like(arr, 128, dtype='uint8'), 'zero scale range' + + # Clip to [p_low, p_high] then scale to [0, 255] + clipped = np.clip(arr, p_low, p_high) + normalised = ((clipped - p_low) / scale * 255).astype('uint8') + + # Remap structural zeros (unfilled lower triangle) to mid-grey (128) + # so they don't contaminate the 0-end of the energy scale + normalised[arr == 0] = 128 + + except Exception as e: + err = f'Problem in normalize_map_global: {e}' + normalised = np.full_like(arr, 128, dtype='uint8') + return normalised, err +``` + +#### 2.4.2 Update `calculate_proteogram()` to accept `norm_stats` + +Modify the method signature and normalisation block: + +```python +def calculate_proteogram(self, + return_simulated_pdb: bool = False, + debug: bool = False, + subtract_solvent_energies: bool = True, + memory_efficient: bool = False, + norm_stats: dict = None): # <-- NEW parameter + """ + ... (existing docstring) ... + + Args: + ... + norm_stats: Optional dict loaded from norm_stats.json. When supplied, + normalize_map_global() is used for all 6 channels instead + of per-protein min-max. Keys: 'vdw_attractive', 'vdw_repulsive', + 'es_attractive', 'es_repulsive', 'distance', 'hydrophobicity'. + Each value is a dict with 'p_low' and 'p_high'. + """ + # ... existing MD pipeline code unchanged ... + + # ---- Replace the normalisation block ---- + def _norm(arr, channel_name): + if norm_stats and channel_name in norm_stats: + s = norm_stats[channel_name] + return self.normalize_map_global(arr, s['p_low'], s['p_high']) + return self.normalize_map(arr) + + norm_disto_map, disto_err = _norm(disto_map, 'distance') + norm_hydro_map, hydro_err = _norm(hydro_map, 'hydrophobicity') + norm_vdw_att_map, vdw_att_err = _norm(vdw_e_att, 'vdw_attractive') + norm_vdw_rep_map, vdw_rep_err = _norm(vdw_e_rep, 'vdw_repulsive') + norm_es_att_map, es_att_err = _norm(es_e_att, 'es_attractive') + norm_es_rep_map, es_rep_err = _norm(es_e_rep, 'es_repulsive') + # ... rest of stacking unchanged ... +``` + +#### 2.4.3 Update `create_v2_proteograms.py` + +```python +# Add to argparse +parser.add_argument('--global_norm', action='store_true', + help='Use global percentile normalisation bounds from norm_stats.json ' + 'instead of per-protein min-max. Requires --norm_stats_file.') +parser.add_argument('--norm_stats_file', type=str, default=None, + help='Path to norm_stats.json produced by compute_norm_stats.py.') +parser.add_argument('--save_npy_matrices', action='store_true', + help='Save raw energy matrices as .npy files alongside proteogram JPGs. ' + 'Required input for compute_norm_stats.py.') + +# Load norm_stats once before the proteogram creation loop +norm_stats = None +if args.global_norm: + if not args.norm_stats_file or not os.path.exists(args.norm_stats_file): + raise ValueError('--global_norm requires --norm_stats_file pointing to norm_stats.json') + import json + with open(args.norm_stats_file) as f: + norm_stats = json.load(f) + print(f'Loaded global norm stats from {args.norm_stats_file}') + +# Pass norm_stats into the ProteogramV2 call inside the creation loop +proteogram_data, errors = prot.calculate_proteogram( + subtract_solvent_energies=True, + memory_efficient=args.memory_efficient, + norm_stats=norm_stats, # <-- new +) +``` + +### 2.5 End-to-End Workflow + +```bash +# Step 1: Generate proteograms with raw .npy matrix saving (first pass or subset) +python create_v2_proteograms.py --save_npy_matrices + +# Step 2: Compute global bounds from saved matrices +python compute_norm_stats.py \ + --npy_dir /path/to/proteograms/energy_matrices \ + --out_file /path/to/norm_stats.json \ + --max_samples 5000 + +# Step 3: Re-generate proteograms using global normalisation +python create_v2_proteograms.py \ + --global_norm \ + --norm_stats_file /path/to/norm_stats.json \ + --overwrite +``` + +### 2.6 Validation Steps + +> Steps 1 and 4 below (pixel-distribution and clipping-rate checks) are now +> implemented as a single runnable tool, `scripts/v2/validate_normalisation.py`, +> which also adds an inter-protein variance ratio and a correlation-with-raw-scale +> check that these steps didn't originally include. See +> [`percentile_normalisation_bug_fix_and_validation.md`](percentile_normalisation_bug_fix_and_validation.md) +> for exact commands and how to read the output. Steps 2 and 3 (visual inspection, +> downstream MAP@K) remain manual/expensive as described below. + +#### Step 1 — Sanity check: pixel distribution + +For a random sample of 100 proteograms, compare the pixel value histograms between per-protein and global normalisation: + +```python +import numpy as np +import matplotlib.pyplot as plt +from PIL import Image +import glob + +per_protein_files = glob.glob('/path/to/proteograms_per_protein/*.jpg')[:100] +global_files = glob.glob('/path/to/proteograms_global/*.jpg')[:100] + +for label, files in [('per-protein', per_protein_files), ('global', global_files)]: + pixels = np.concatenate([np.array(Image.open(f)).ravel() for f in files]) + plt.hist(pixels, bins=50, alpha=0.6, label=label) + +plt.legend() +plt.xlabel('Pixel value') +plt.title('Pixel distribution: per-protein vs. global normalisation') +plt.savefig('norm_comparison.png', dpi=150) +``` + +Expected result: global normalisation produces a wider, less clipped distribution with meaningful variation near 0 and 255. Per-protein should look nearly uniform (every image uses the full range). + +#### Step 2 — Visual inspection + +Side-by-side comparison of the same protein normalised both ways: + +```python +from PIL import Image +import matplotlib.pyplot as plt + +pdb_id = 'd3kfda_' +per_protein = Image.open(f'/path/per_protein/{pdb_id}.jpg') +global_norm = Image.open(f'/path/global/{pdb_id}.jpg') + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) +ax1.imshow(per_protein); ax1.set_title('Per-protein normalisation'); ax1.axis('off') +ax2.imshow(global_norm); ax2.set_title('Global normalisation'); ax2.axis('off') +plt.savefig(f'{pdb_id}_norm_comparison.png', dpi=150) +``` + +Proteins with strong hydrophobic cores (e.g., globins, immunoglobulins) should appear noticeably brighter in the VdW channels under global normalisation compared to proteins with weak packing. + +#### Step 3 — Downstream MAP@K comparison + +Retrain the ResNet18 model on global-normalised proteograms and compare MAP@K on the eval set against the baseline model trained on per-protein normalised proteograms. Expected outcome: global normalisation improves MAP@K at the fold and superfamily levels (where energy magnitude differences are most discriminative), with neutral or marginal effect at the class level. + +#### Step 4 — Robustness check: unseen proteins + +Compute the fraction of pixel values clipped to 0 or 255 for 50 randomly selected held-out proteins (not used in `compute_norm_stats.py`). If > 5% of non-zero pixels are clipped, the percentile bounds are too tight and `--low_pct`/`--high_pct` should be widened (e.g., 0.5 and 99.5). + +--- + +## 3. Grad-CAM Explainability + +### 3.1 Motivation + +When Proteogram reports that two proteins share 87% cosine similarity, a structural biologist naturally asks: *which residue-residue interactions drove that score?* Currently there is no answer — the model is a black box. + +Because proteogram pixels directly encode pairwise residue interactions (pixel at row `i`, column `j` represents the interaction between residue `i` and residue `j`), a saliency heatmap over the input image is directly interpretable as a **residue-pair importance map**. This is a unique property of the proteogram representation that does not exist for most CV tasks. + +Grad-CAM (Gradient-weighted Class Activation Mapping) computes a heatmap by backpropagating the gradient of a target score through the last convolutional layer of the CNN. High activation regions in the heatmap indicate which spatial features (and thus which residue pairs) most influenced the model's output. + +### 3.2 Design Decisions + +#### Target layer selection + +For ResNet18, the natural target is the output of `layer4` (the last residual block), which has spatial resolution 7×7 for 224px input or 13×13 for 200px padded proteograms. This gives meaningful spatial resolution after upsampling back to the full NxN image. + +For the custom ConvNet, the target is `block4` (after the 4th MaxPool, spatial resolution ≈ 12×12 for 200px input). + +#### Score to differentiate + +Standard Grad-CAM differentiates with respect to the **class logit** for a classification task. For a *retrieval* task we instead differentiate with respect to the **cosine similarity score** between a query and a target embedding. This gives a "similarity-attribution" heatmap: *which parts of the query proteogram, when activated, push the cosine similarity with the target higher?* + +Formally, if `f_q` and `f_t` are the embedding vectors for query and target: + +``` +S = cos(f_q, f_t) = (f_q · f_t) / (||f_q|| · ||f_t||) +``` + +We compute `∂S / ∂A_k` for each activation map `A_k` in the target convolutional layer. + +#### Output format + +The Grad-CAM heatmap is: +- An NxN float32 array in [0, 1] — matching the proteogram dimensions +- Saved as both a matplotlib figure (with residue axis labels) and a raw `.npy` file +- Overlaid as a semi-transparent colour map on top of the original proteogram image + +### 3.3 Implementation + +#### 3.3.1 New method: `Img2Vec.gradcam_similarity()` + +Add to `proteogram/v2/image_similarity.py`: + +```python +def gradcam_similarity(self, + query_image_path: str, + target_image_path: str, + output_dir: str, + query_sequence: str = None, + target_sequence: str = None) -> np.ndarray: + """Compute a Grad-CAM residue-pair importance map for a query→target similarity. + + The heatmap shows which residue-pair interactions in the QUERY proteogram + most influence the cosine similarity with the TARGET proteogram. + + The model must be a ResNet18 fine-tuned with train_multiple_models.py + (--model resnet18) or the from-scratch ConvNet (--model cnn). + + Args: + query_image_path: Path to the query proteogram JPG. + target_image_path: Path to the target proteogram JPG. + output_dir: Directory to save the heatmap figure and .npy file. + query_sequence: Optional 1-letter amino acid sequence for axis labels. + target_sequence: Optional 1-letter amino acid sequence (unused currently, + reserved for cross-proteogram attribution in future). + + Returns: + np.ndarray: Upsampled Grad-CAM heatmap, shape (H, W), values in [0, 1]. + """ + import torch.nn.functional as F + + os.makedirs(output_dir, exist_ok=True) + + # ------------------------------------------------------------------ # + # 1. Identify the target convolutional layer # + # ------------------------------------------------------------------ # + target_layer = self._get_gradcam_target_layer() + + # ------------------------------------------------------------------ # + # 2. Register forward/backward hooks # + # ------------------------------------------------------------------ # + activations = {} + gradients = {} + + def _save_activation(module, input, output): + activations['value'] = output.detach() + + def _save_gradient(module, grad_input, grad_output): + gradients['value'] = grad_output[0].detach() + + fwd_hook = target_layer.register_forward_hook(_save_activation) + bwd_hook = target_layer.register_full_backward_hook(_save_gradient) + + try: + # ------------------------------------------------------------------ # + # 3. Forward pass for both query and target # + # ------------------------------------------------------------------ # + query_tensor = self._load_and_preprocess(query_image_path) # (1, 3, H, W) + target_tensor = self._load_and_preprocess(target_image_path) # (1, 3, H, W) + + # Embeddings from the penultimate layer + # Switch to full model (not self.embed which stripped the head) + self.model.eval() + query_tensor = query_tensor.to(self.device).requires_grad_(True) + target_tensor = target_tensor.to(self.device) + + # Get embedding for query (triggers forward hook and saves activations) + query_feat = self.embed(query_tensor) # (1, d) + target_feat = self.embed(target_tensor).detach() # (1, d) + + # ------------------------------------------------------------------ # + # 4. Compute cosine similarity and differentiate # + # ------------------------------------------------------------------ # + # Manually compute cosine similarity (not through nn.CosineSimilarity + # so we can call backward on the scalar) + q_norm = F.normalize(query_feat, dim=1) + t_norm = F.normalize(target_feat, dim=1) + cos_sim = (q_norm * t_norm).sum() # scalar + + self.model.zero_grad() + cos_sim.backward() + + # ------------------------------------------------------------------ # + # 5. Compute Grad-CAM weights # + # ------------------------------------------------------------------ # + grads = gradients['value'] # (1, C, h, w) + acts = activations['value'] # (1, C, h, w) + + # Global-average-pool the gradients over the spatial dims → (1, C, 1, 1) + weights = grads.mean(dim=(2, 3), keepdim=True) + + # Weighted combination of activation maps → (1, 1, h, w) + cam = (weights * acts).sum(dim=1, keepdim=True) + cam = F.relu(cam) # keep only positive contributions + + # Normalise to [0, 1] + cam_min, cam_max = cam.min(), cam.max() + if cam_max > cam_min: + cam = (cam - cam_min) / (cam_max - cam_min) + + # ------------------------------------------------------------------ # + # 6. Upsample to input image size # + # ------------------------------------------------------------------ # + input_h = query_tensor.shape[2] + input_w = query_tensor.shape[3] + cam_upsampled = F.interpolate(cam, + size=(input_h, input_w), + mode='bilinear', + align_corners=False) + cam_np = cam_upsampled.squeeze().cpu().numpy() # (H, W) + + finally: + fwd_hook.remove() + bwd_hook.remove() + + # ------------------------------------------------------------------ # + # 7. Save outputs # + # ------------------------------------------------------------------ # + query_name = os.path.splitext(os.path.basename(query_image_path))[0] + target_name = os.path.splitext(os.path.basename(target_image_path))[0] + stem = f'{query_name}_vs_{target_name}' + + # Save raw heatmap + npy_path = os.path.join(output_dir, f'{stem}_gradcam.npy') + np.save(npy_path, cam_np) + + # Save overlay figure + query_img = np.array(Image.open(query_image_path).convert('RGB')) + self._save_gradcam_figure( + query_img=query_img, + cam=cam_np, + cos_sim=cos_sim.item(), + query_name=query_name, + target_name=target_name, + output_dir=output_dir, + query_sequence=query_sequence, + ) + + print(f'Grad-CAM saved → {output_dir}/{stem}_gradcam.png') + return cam_np + + +def _get_gradcam_target_layer(self): + """Return the last convolutional layer for Grad-CAM based on architecture.""" + children = list(self.model.children()) + # ResNet18: children order is conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4, avgpool, fc + # Find the last nn.Sequential that contains Conv2d layers + target = None + for child in children: + if isinstance(child, nn.Sequential): + for submodule in child.modules(): + if isinstance(submodule, nn.Conv2d): + target = child + if target is None: + # Fallback: use the last Conv2d found anywhere in the model + for module in self.model.modules(): + if isinstance(module, nn.Conv2d): + target = module + return target + + +def _load_and_preprocess(self, image_path: str) -> torch.Tensor: + """Load and preprocess a single proteogram image matching training transforms.""" + img = Image.open(image_path).convert('RGB') + # Apply the same pad-to-200 + ImageNet normalisation used in training + from torchvision import transforms as T + import numpy as np + arr = np.array(img) + H, W = arr.shape[:2] + target = 200 + + def get_pad(curr, tgt): + d = tgt - curr + if d <= 0: + return (0, 0) + p1 = d // 2 + return (p1, d - p1) + + padding = (get_pad(H, target), get_pad(W, target), (0, 0)) + arr = np.pad(arr, padding, constant_values=128) + arr = arr[:target, :target, :] + img_padded = Image.fromarray(arr.astype('uint8')) + + transform = T.Compose([ + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + return transform(img_padded).unsqueeze(0) # (1, 3, H, W) + + +def _save_gradcam_figure(self, + query_img: np.ndarray, + cam: np.ndarray, + cos_sim: float, + query_name: str, + target_name: str, + output_dir: str, + query_sequence: str = None) -> None: + """Save a 3-panel Grad-CAM figure: original, heatmap, overlay.""" + import matplotlib.pyplot as plt + import matplotlib.cm as cm + from matplotlib.colors import Normalize + + fig, axes = plt.subplots(1, 3, figsize=(18, 6)) + fig.suptitle( + f'Grad-CAM: {query_name} → {target_name} (cosine similarity = {cos_sim:.4f})', + fontsize=13, y=1.01 + ) + + # Panel 1: original query proteogram + axes[0].imshow(query_img) + axes[0].set_title('Query proteogram', fontsize=11) + axes[0].axis('off') + + # Panel 2: Grad-CAM heatmap alone + im = axes[1].imshow(cam, cmap='hot', vmin=0, vmax=1) + axes[1].set_title('Grad-CAM heatmap\n(high = important residue pairs)', fontsize=11) + axes[1].axis('off') + plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04) + + # Panel 3: overlay (proteogram + semi-transparent heatmap) + axes[2].imshow(query_img) + overlay = axes[2].imshow(cam, cmap='hot', alpha=0.55, vmin=0, vmax=1) + axes[2].set_title('Overlay', fontsize=11) + axes[2].axis('off') + plt.colorbar(overlay, ax=axes[2], fraction=0.046, pad=0.04) + + # Optional: add residue index tick labels if sequence is provided + if query_sequence and len(query_sequence) <= 200: + step = max(1, len(query_sequence) // 20) # show ~20 tick labels + ticks = list(range(0, len(query_sequence), step)) + tick_labels = [f'{i}\n{query_sequence[i]}' for i in ticks] + for ax in axes: + ax.set_xticks(ticks); ax.set_xticklabels(tick_labels, fontsize=6) + ax.set_yticks(ticks); ax.set_yticklabels(tick_labels, fontsize=6) + ax.tick_params(axis='both', length=2) + + plt.tight_layout() + stem = f'{query_name}_vs_{target_name}' + fig_path = os.path.join(output_dir, f'{stem}_gradcam.png') + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + plt.close(fig) +``` + +#### 3.3.2 New script: `scripts/v2/explain_similarity.py` + +```python +#!/usr/bin/env python +"""Generate Grad-CAM residue-pair importance maps for similar protein pairs. + +For each query in the eval set, explains the similarity to its top-1 hit +(or a user-specified target) using Grad-CAM over the last convolutional layer. + +Usage — explain top-1 hit for all eval proteograms: + python explain_similarity.py + +Usage — explain a specific query→target pair: + python explain_similarity.py \\ + --query /path/to/d3kfda_.jpg \\ + --target /path/to/d1yl4r1.jpg \\ + --query_seq ACDEFGHIKLMNPQRSTVWY... \\ + --output_dir gradcam_results/ + +Usage — explain top-K hits for the 50 queries with lowest MAP@K (worst cases): + python explain_similarity.py --explain_worst 50 --top_k 3 +""" +import argparse +import os +import pickle +import torch +import pandas as pd + +from proteogram.v2 import Img2Vec +from proteogram.common import read_yaml + + +def main(): + parser = argparse.ArgumentParser(description='Grad-CAM explainability for Proteogram.') + parser.add_argument('--query', '-q', type=str, default=None, + help='Path to a single query proteogram JPG.') + parser.add_argument('--target', '-t', type=str, default=None, + help='Path to a single target proteogram JPG.') + parser.add_argument('--query_seq', type=str, default=None, + help='1-letter amino acid sequence of the query protein ' + '(optional, for residue axis labels).') + parser.add_argument('--output_dir', '-o', type=str, default='gradcam_output', + help='Directory to save Grad-CAM figures and .npy files.') + parser.add_argument('--explain_worst', type=int, default=None, + help='Explain the top-1 hit for the N queries with the lowest ' + 'MAP@K score (most informative failures). ' + 'Requires proteogram_sim_results in config.yml.') + parser.add_argument('--top_k', type=int, default=1, + help='Number of top hits to explain per query (default: 1).') + args = parser.parse_args() + + config = read_yaml('config.yml') + model_file = config['model_file'] + embed_file = config['embed_file'] + corpus_dir = config['proteograms_for_sim_dir'] + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + img_sim = Img2Vec(model_file, dataset_dir=[], device=device) + + # Load corpus embeddings + with open(embed_file, 'rb') as f: + img_sim.dataset = pickle.load(f) + + os.makedirs(args.output_dir, exist_ok=True) + + if args.query and args.target: + # Single pair mode + img_sim.gradcam_similarity( + query_image_path=args.query, + target_image_path=args.target, + output_dir=args.output_dir, + query_sequence=args.query_seq, + ) + + elif args.explain_worst: + # Explain worst-performing queries from similarity results + results_file = config.get('proteogram_sim_results') + if not results_file or not os.path.exists(results_file): + raise FileNotFoundError( + f'proteogram_sim_results not found: {results_file}. ' + 'Run measure_similarity_v2.py first.') + + results_df = pd.read_csv(results_file, sep='\t') + # Sort by MAP@K score (ascending = worst first) if the column exists, + # otherwise just take the last N rows as a proxy + n = args.explain_worst + queries_to_explain = results_df.head(n) + + for _, row in queries_to_explain.iterrows(): + query_path = row.iloc[0] + top_hits = [row.iloc[k+1].split(',')[0] # filename part of 'filename,score' + for k in range(min(args.top_k, len(row) - 1))] + for target_stem in top_hits: + target_path = os.path.join(corpus_dir, target_stem + '.jpg') + if not os.path.exists(target_path): + print(f'Target not found, skipping: {target_path}') + continue + img_sim.gradcam_similarity( + query_image_path=query_path, + target_image_path=target_path, + output_dir=args.output_dir, + ) + else: + parser.print_help() + + +if __name__ == '__main__': + main() +``` + +### 3.4 Validation Steps + +#### Step 1 — Sanity check: heatmap is not uniform + +For 10 random query-target pairs, verify the heatmap has meaningful spatial variation (std > 0.05): + +```python +import numpy as np +import glob + +npy_files = glob.glob('gradcam_output/*_gradcam.npy') +for f in npy_files: + cam = np.load(f) + std = cam.std() + max_val = cam.max() + print(f'{f}: std={std:.4f} max={max_val:.4f}') + assert std > 0.05, f'Heatmap appears uniform for {f} — check hook registration' +``` + +#### Step 2 — Biological sanity check + +Run Grad-CAM on a well-studied protein pair from the same SCOPe superfamily (e.g., two globins: haemoglobin α-chain vs. myoglobin). Expect high activation at: +- The haem-binding pocket region (residues ~60–90 and ~130–150 in the sequence) +- The conserved F-helix contacts +- The hydrophobic core residues + +Cross-reference high-activation residue pairs against the known structural alignment from US-align or GTalign. If the top-10 residue pairs by Grad-CAM score overlap significantly with the US-align-identified structurally equivalent residue pairs, the explainer is working correctly. + +#### Step 3 — Negative control + +Run Grad-CAM on a query-target pair from *different* SCOPe classes (e.g., an all-alpha vs. an all-beta protein with low cosine similarity score ~0.3). The heatmap should be diffuse and low-magnitude — no clear hotspot — because no specific structural motif is driving the (low) similarity. + +```python +# Confirm: mean activation for negative pairs should be < mean activation for positive pairs +import numpy as np + +positive_cams = [np.load(f) for f in glob.glob('gradcam_output/same_class_*.npy')] +negative_cams = [np.load(f) for f in glob.glob('gradcam_output/diff_class_*.npy')] + +pos_mean = np.mean([c.max() for c in positive_cams]) +neg_mean = np.mean([c.max() for c in negative_cams]) +print(f'Positive pairs max activation: {pos_mean:.4f}') +print(f'Negative pairs max activation: {neg_mean:.4f}') +assert pos_mean > neg_mean, 'Grad-CAM not discriminating positive/negative pairs' +``` + +#### Step 4 — Hook cleanup test + +Verify hooks are always removed even when an exception occurs mid-computation (hooks left dangling slow down subsequent forward passes and may accumulate memory): + +```python +# Deliberately pass a corrupted image path and confirm no hook leakage +from proteogram.v2 import Img2Vec +import torch + +img_sim = Img2Vec(model_file, dataset_dir=[], device='cpu') +img_sim.dataset = {} # minimal setup + +hook_count_before = len(list(img_sim.model._forward_hooks.values())) +try: + img_sim.gradcam_similarity('/nonexistent/query.jpg', '/nonexistent/target.jpg', '/tmp') +except Exception: + pass +hook_count_after = len(list(img_sim.model._forward_hooks.values())) +assert hook_count_before == hook_count_after, 'Forward hooks leaked after exception!' +print('Hook cleanup test passed.') +``` + +--- + +## 4. Combined Integration Checklist + +Before merging all three changes, run through this checklist end-to-end on the eval set: + +``` +[ ] pip install faiss-cpu (or faiss-gpu) added to pyproject.toml +[ ] compute_norm_stats.py runs without error on 500 random energy matrices +[ ] norm_stats.json is committed to the repository alongside the dataset +[ ] create_v2_proteograms.py --global_norm produces visually distinct images + vs. per-protein normalised equivalents (visual inspection on 5 proteins) +[ ] measure_similarity_v2.py --faiss produces sim_dict identical in format to + existing brute-force output (evaluate_methods_v2.py accepts it unchanged) +[ ] FAISS Recall@5 ≥ 0.99 × brute-force Recall@5 (automated test passes) +[ ] MAP@K from FAISS results within ±0.005 of MAP@K from brute-force results +[ ] explain_similarity.py runs on a single pair and produces a 3-panel PNG +[ ] Grad-CAM heatmap std > 0.05 for same-class pairs (sanity check passes) +[ ] No memory leaks: all forward/backward hooks removed after gradcam_similarity() +[ ] All existing tests in scripts/v2/tests/ still pass +[ ] README.md updated with: + - New --global_norm / --norm_stats_file flags in Step 1 + - New --faiss / --faiss_pq flags in Step 4 + - New explain_similarity.py in the scripts reference table +``` + +--- + +## 5. Config additions (`scripts/v2/config.example.yml`) + +Add these keys to the example config for discoverability: + +```yaml +# ── Global normalisation (Improvement 2) ────────────────────────────── +# Path to norm_stats.json produced by compute_norm_stats.py. +# Required when running create_v2_proteograms.py --global_norm. +norm_stats_file: /path/to/norm_stats.json + +# ── FAISS index (Improvement 1) ─────────────────────────────────────── +# Path to save/load the FAISS index (auto-derived from embed_file if omitted). +faiss_index_file: /path/to/corpus_embeddings.faiss + +# ── Grad-CAM output (Improvement 3) ─────────────────────────────────── +# Directory to write Grad-CAM figures and .npy heatmap files. +gradcam_output_dir: /path/to/gradcam_output +``` diff --git a/proteogram/v2/__init__.py b/proteogram/v2/__init__.py index 1686c46..2a2845e 100644 --- a/proteogram/v2/__init__.py +++ b/proteogram/v2/__init__.py @@ -3,7 +3,9 @@ from .atomistic_nonbonded_forces import AtomisticNonBondedForceModel from .martini_nonbonded_forces import MartiniNonBondedForceModel from .losses import HierarchicalTripletLoss, HierarchicalPKSampler, SCOPE_LEVELS +from .faiss_search import FaissIndex __all__ = ['ProteogramV2', 'Img2Vec', 'AtomisticNonBondedForceModel', 'MartiniNonBondedForceModel', - 'HierarchicalTripletLoss', 'HierarchicalPKSampler', 'SCOPE_LEVELS'] \ No newline at end of file + 'HierarchicalTripletLoss', 'HierarchicalPKSampler', 'SCOPE_LEVELS', + 'FaissIndex'] diff --git a/proteogram/v2/faiss_search.py b/proteogram/v2/faiss_search.py new file mode 100644 index 0000000..8d68052 --- /dev/null +++ b/proteogram/v2/faiss_search.py @@ -0,0 +1,327 @@ +"""FAISS-based Approximate Nearest Neighbour index for proteogram embeddings. + +This module is fully self-contained and has no dependency on Img2Vec or any other +proteogram class. It operates on plain numpy float32 arrays and stores a key list +(filename strings) so integer FAISS indices can be mapped back to protein IDs. + +Typical usage +------------- +>>> from proteogram.v2.faiss_search import FaissIndex +>>> import numpy as np + +>>> # Build from a dict of {filename: embedding_tensor} (same format as Img2Vec.dataset) +>>> index = FaissIndex.from_dataset(img2vec.dataset) + +>>> # All-vs-all search (returns same format as Img2Vec.sim_dict) +>>> sim_dict = index.search_all(top_k=5) + +>>> # Single-query search +>>> hits = index.search_one(query_vec, top_k=5) + +>>> # Persistence +>>> index.save("/path/to/corpus.faiss") +>>> index2 = FaissIndex.load("/path/to/corpus.faiss") +""" + +from __future__ import annotations + +import os +import pickle +from typing import Dict, List, Tuple + +import numpy as np +import torch + + +# --------------------------------------------------------------------------- +# Public constants +# --------------------------------------------------------------------------- + +#: Default lower percentile for nlist auto-selection. +_NLIST_SQRT_FACTOR: float = 1.0 + + +# --------------------------------------------------------------------------- +# Helper +# --------------------------------------------------------------------------- + +def _l2_normalise(mat: np.ndarray) -> np.ndarray: + """Return an L2-normalised copy of *mat* (shape N×d, float32). + + Does NOT modify the input array in-place so the caller's embeddings stay intact. + """ + mat = mat.copy().astype(np.float32) + norms = np.linalg.norm(mat, axis=1, keepdims=True) + norms = np.where(norms == 0, 1.0, norms) # avoid div-by-zero for zero vectors + mat /= norms + return mat + + +def _stack_dataset(dataset: Dict[str, torch.Tensor]) -> Tuple[List[str], np.ndarray]: + """Convert an Img2Vec-style dataset dict to an ordered (keys, matrix) pair. + + Args: + dataset: Mapping of filename → 1-D or 1×d embedding tensor. + + Returns: + keys: List of filenames in the same row order as the matrix. + matrix: float32 numpy array of shape (N, d). + """ + keys = list(dataset.keys()) + vecs = torch.cat([dataset[k].cpu().reshape(1, -1) for k in keys]).float().numpy() + return keys, vecs + + +# --------------------------------------------------------------------------- +# FaissIndex +# --------------------------------------------------------------------------- + +class FaissIndex: + """Wraps a FAISS IVFFlat or IVF-PQ index with a key mapping. + + Parameters + ---------- + keys: + Ordered list of protein filenames. ``keys[i]`` is the protein + corresponding to FAISS integer index ``i``. + vecs_norm: + L2-normalised embedding matrix, shape (N, d), float32. Stored so + single-query searches can normalise the query the same way. + index: + A trained and populated FAISS index (inner-product metric). + """ + + def __init__(self, + keys: List[str], + vecs_norm: np.ndarray, + index) -> None: + self.keys = keys + self.vecs_norm = vecs_norm + self._index = index + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + + @classmethod + def from_dataset(cls, + dataset: Dict[str, torch.Tensor], + use_pq: bool = False, + nlist: int = None, + nprobe: int = None, + pq_m: int = 8, + pq_nbits: int = 8) -> "FaissIndex": + """Build a FAISS index from an Img2Vec-style embedding dataset. + + Embeddings are L2-normalised before indexing so inner-product search + is equivalent to cosine similarity. + + Args: + dataset: ``{filename: embedding_tensor}`` dict (same as + ``Img2Vec.dataset``). + use_pq: Use IVF-PQ compressed index. Recommended for corpora + larger than 100 K proteins. Slightly lower recall but + 4–32× lower memory. Defaults to ``False`` (IVFFlat, + exact). + nlist: Number of Voronoi cells. Defaults to + ``max(1, int(sqrt(N)))``. + nprobe: Cells searched per query. Higher → better recall, + slower. Defaults to ``max(1, nlist // 10)``. + pq_m: Sub-quantiser count for IVF-PQ. Must divide ``d`` + evenly. Auto-reduced if necessary. + pq_nbits: Bits per sub-quantiser (IVF-PQ only). 8 is standard. + + Returns: + A trained and populated ``FaissIndex`` ready for search. + + Raises: + ImportError: If ``faiss`` is not installed. + ValueError: If ``dataset`` is empty. + """ + try: + import faiss + except ImportError as exc: + raise ImportError( + "faiss is required. Install with:\n" + " uv add faiss-cpu # CPU\n" + " uv add faiss-gpu # GPU / CUDA build" + ) from exc + + if not dataset: + raise ValueError("dataset is empty — embed_dataset() must be called first.") + + keys, vecs = _stack_dataset(dataset) + vecs_norm = _l2_normalise(vecs) + + N, d = vecs_norm.shape + _nlist = nlist if nlist is not None else max(1, int(N ** _NLIST_SQRT_FACTOR ** 0.5)) + + # Cannot have more cells than vectors during training + _nlist = min(_nlist, N) + _nprobe = nprobe if nprobe is not None else max(1, _nlist // 10) + + quantiser = faiss.IndexFlatIP(d) + + if use_pq and N >= 256: + # pq_m must divide d evenly + while d % pq_m != 0 and pq_m > 1: + pq_m -= 1 + index = faiss.IndexIVFPQ( + quantiser, d, _nlist, pq_m, pq_nbits, + faiss.METRIC_INNER_PRODUCT, + ) + index_type = "IVF-PQ" + else: + if use_pq and N < 256: + print("WARNING: corpus too small for IVF-PQ (N<256) — falling back to IVFFlat.") + index = faiss.IndexIVFFlat(quantiser, d, _nlist, faiss.METRIC_INNER_PRODUCT) + index_type = "IVFFlat" + + index.train(vecs_norm) + index.add(vecs_norm) + index.nprobe = _nprobe + + print( + f"FAISS index built: {index.ntotal} vectors | d={d} | " + f"nlist={_nlist} | nprobe={_nprobe} | type={index_type}" + ) + return cls(keys=keys, vecs_norm=vecs_norm, index=index) + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + def search_all(self, top_k: int = 10) -> Dict[str, List[Tuple[str, float]]]: + """Batch cosine-similarity search for every vector in the index. + + Returns a dict with the same structure as ``Img2Vec.sim_dict``: + ``{filename: [(target_filename, score), ...]}``. + + Self-hits (rank 0, score ≈ 1.0) are included so callers can + decide whether to strip them. + + Args: + top_k: Number of results to return per query (including self-hit). + + Returns: + Similarity dict keyed by query filename. + """ + scores_mat, idx_mat = self._index.search(self.vecs_norm, top_k + 1) + + sim_dict: Dict[str, List[Tuple[str, float]]] = {} + for i, key in enumerate(self.keys): + hits: List[Tuple[str, float]] = [] + for rank in range(top_k + 1): + j = int(idx_mat[i, rank]) + if j < 0: # FAISS pads with -1 when fewer results exist + continue + hits.append((self.keys[j], float(scores_mat[i, rank]))) + if len(hits) >= top_k: + break + sim_dict[key] = hits + return sim_dict + + def search_one(self, + query_vec: np.ndarray, + top_k: int = 10, + exclude_self_key: str = None) -> List[Tuple[str, float]]: + """Search for the ``top_k`` most similar proteins to a single query. + + Args: + query_vec: 1-D float embedding (not necessarily normalised). + top_k: Number of results to return. + exclude_self_key: If provided, any hit matching this key is skipped + (useful when the query is already in the corpus). + + Returns: + List of ``(filename, cosine_score)`` tuples, descending by score. + """ + qvec = query_vec.copy().reshape(1, -1).astype(np.float32) + norm = np.linalg.norm(qvec) + if norm > 0: + qvec /= norm + + scores, indices = self._index.search(qvec, top_k + 1) + results: List[Tuple[str, float]] = [] + for rank in range(top_k + 1): + j = int(indices[0, rank]) + if j < 0: + continue + key = self.keys[j] + if exclude_self_key and key == exclude_self_key: + continue + results.append((key, float(scores[0, rank]))) + if len(results) >= top_k: + break + return results + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + def save(self, index_path: str) -> None: + """Save the FAISS index and key mapping to disk. + + Two files are written: + - ``index_path`` — the FAISS binary index + - ``index_path + '.keys.pkl'`` — the ordered key list + + Args: + index_path: Destination path, e.g. ``/data/corpus.faiss``. + """ + try: + import faiss + except ImportError as exc: + raise ImportError("faiss required for save().") from exc + + os.makedirs(os.path.dirname(os.path.abspath(index_path)), exist_ok=True) + faiss.write_index(self._index, index_path) + keys_path = index_path + ".keys.pkl" + with open(keys_path, "wb") as fh: + pickle.dump({"keys": self.keys, "vecs_norm": self.vecs_norm}, fh) + print(f"Saved FAISS index → {index_path}") + print(f"Saved key mapping → {keys_path}") + + @classmethod + def load(cls, index_path: str) -> "FaissIndex": + """Load a previously saved FAISS index and key mapping. + + Args: + index_path: Path to the ``.faiss`` file written by ``save()``. + + Returns: + A ready-to-use ``FaissIndex`` instance. + """ + try: + import faiss + except ImportError as exc: + raise ImportError("faiss required for load().") from exc + + index = faiss.read_index(index_path) + keys_path = index_path + ".keys.pkl" + with open(keys_path, "rb") as fh: + data = pickle.load(fh) + keys = data["keys"] + vecs_norm = data["vecs_norm"] + print(f"Loaded FAISS index ← {index_path} ({index.ntotal} vectors)") + return cls(keys=keys, vecs_norm=vecs_norm, index=index) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def n_vectors(self) -> int: + """Number of vectors stored in the index.""" + return self._index.ntotal + + @property + def dim(self) -> int: + """Embedding dimension.""" + return self._index.d + + def __repr__(self) -> str: + return ( + f"FaissIndex(n={self.n_vectors}, d={self.dim}, " + f"nprobe={getattr(self._index, 'nprobe', 'N/A')})" + ) diff --git a/proteogram/v2/image_similarity.py b/proteogram/v2/image_similarity.py index bdda7a9..4d4678b 100644 --- a/proteogram/v2/image_similarity.py +++ b/proteogram/v2/image_similarity.py @@ -757,3 +757,103 @@ def cluster_dataset(self, nclusters, dist="euclidean", display=False): self.display_clusters() return + + # ------------------------------------------------------------------ + # FAISS wrapper methods + # These are thin delegators to proteogram.v2.faiss_search.FaissIndex. + # The FaissIndex instance is stored as self._faiss so scripts can also + # access it directly if needed. + # ------------------------------------------------------------------ + + def build_faiss_index(self, + use_pq: bool = False, + nlist: int = None, + nprobe: int = None, + pq_m: int = 8, + pq_nbits: int = 8) -> None: + """Build a FAISS ANN index from the currently loaded embedding dataset. + + Delegates to :class:`~proteogram.v2.faiss_search.FaissIndex`. + After calling this, ``similarities_faiss()`` can be used as a fast + drop-in replacement for ``similarities()``. + + Args: + use_pq: Use IVF-PQ compressed index (recommended for > 100 K + proteins). Defaults to ``False`` (IVFFlat, exact). + nlist: Voronoi cell count. Defaults to ``sqrt(N)``. + nprobe: Cells searched per query. Defaults to ``nlist // 10``. + pq_m: Sub-quantiser count (IVF-PQ only). + pq_nbits: Bits per sub-quantiser (IVF-PQ only). + """ + from .faiss_search import FaissIndex + self._faiss = FaissIndex.from_dataset( + self.dataset, + use_pq=use_pq, + nlist=nlist, + nprobe=nprobe, + pq_m=pq_m, + pq_nbits=pq_nbits, + ) + + def similarities_faiss(self, + n: int = 10, + save_result_images_dir: str = None, + pad_fn=None) -> float: + """ANN similarity search via FAISS — drop-in replacement for ``similarities()``. + + Populates ``self.sim_dict`` with the same ``{filename: [(target, score)]}`` + format so all downstream scripts work without modification. + + Call ``build_faiss_index()`` (or ``load_faiss_index()``) first. + + Args: + n: Top-N results per query (self-hit included + at rank 0). + save_result_images_dir: Optional directory to write result images. + pad_fn: Padding callable passed to ``save_images()``. + + Returns: + Wall-clock seconds spent in FAISS search. + """ + if not hasattr(self, '_faiss'): + raise RuntimeError( + "Call build_faiss_index() or load_faiss_index() before " + "similarities_faiss()." + ) + start = time() + self.sim_dict = self._faiss.search_all(top_k=n) + elapsed = time() - start + + if save_result_images_dir: + for image_path in self.sim_dict: + full_path = os.path.join( + os.path.dirname(self.files[0]) if self.files else '', + image_path, + ) + self.save_images(full_path, save_result_images_dir, + scores_n_arr=self.sim_dict[image_path], + pad_fn=pad_fn) + return elapsed + + def save_faiss_index(self, index_path: str) -> None: + """Persist the FAISS index to disk. + + Args: + index_path: Destination file path (e.g. ``corpus.faiss``). + A companion ``.keys.pkl`` is written + alongside. + """ + if not hasattr(self, '_faiss'): + raise RuntimeError("No FAISS index to save. Call build_faiss_index() first.") + self._faiss.save(index_path) + + def load_faiss_index(self, index_path: str) -> None: + """Load a previously saved FAISS index. + + Args: + index_path: Path to the ``.faiss`` file written by + ``save_faiss_index()``. + """ + from .faiss_search import FaissIndex + self._faiss = FaissIndex.load(index_path) + diff --git a/pyproject.toml b/pyproject.toml index 20e222f..5880d87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "goatools>=1.4", "pyrotein @ git+https://github.com/carbonscott/pyrotein.git@main", "rcsb-api>=1.7.3", + "faiss-cpu>=1.13.2", ] [project.optional-dependencies] diff --git a/scripts/v2/measure_similarity_v2.py b/scripts/v2/measure_similarity_v2.py index 1220669..9fed008 100644 --- a/scripts/v2/measure_similarity_v2.py +++ b/scripts/v2/measure_similarity_v2.py @@ -50,6 +50,25 @@ def get_pad(curr, tgt): parser.add_argument('--embed', action=argparse.BooleanOptionalAction, default=True, help='Recompute and save embeddings (default: True). ' 'Use --no-embed to load from embed_file instead.') + # ── FAISS options ──────────────────────────────────────────────────────── + parser.add_argument('--faiss', action='store_true', + help=( + 'Use FAISS ANN index for similarity search instead of ' + 'brute-force cosine similarity. Much faster for large ' + 'corpora (> 10 K proteins). Requires faiss-cpu or ' + 'faiss-gpu to be installed.' + )) + parser.add_argument('--faiss_pq', action='store_true', + help=( + 'Use IVF-PQ compressed FAISS index (recommended for ' + '> 100 K proteins). Slightly lower recall but 4-32x ' + 'lower memory than IVFFlat.' + )) + parser.add_argument('--faiss_index_file', type=str, default=None, + help=( + 'Path to save / load the FAISS index. Defaults to ' + 'embed_file with a .faiss extension.' + )) args = parser.parse_args() # Run embedding vs loading saved embeddings @@ -114,6 +133,14 @@ def _confirm_overwrite(path, label, is_dir=False): if os.path.splitext(os.path.basename(f))[0] not in excluded_sids] print(f'Excluded {before - len(prot_files)} proteograms from class(es): ' + ', '.join(sorted(excluded))) + + if not prot_files: + raise ValueError( + 'No proteogram .jpg files found for similarity search. '\ + f'Checked dataset_dir={dataset_dir!r}. '\ + 'If you are running from scripts/v2/, ensure config paths are correct '\ + 'relative to that working directory.' + ) device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f'Using device: {device}') @@ -184,6 +211,11 @@ def _prep_fn(img): with torch.no_grad(): if embed: img_sim.embed_dataset() + if not img_sim.dataset: + raise ValueError( + 'Embedding dataset is empty after embed_dataset(). '\ + 'Verify input proteogram files are readable and preprocessing succeeded.' + ) # Save embeddings with open(embed_file, 'wb') as pklout: pickle.dump(img_sim.dataset, pklout) @@ -192,6 +224,11 @@ def _prep_fn(img): if embed_file: with open(embed_file, 'rb') as pklin: img_sim.dataset = pickle.load(pklin) + if not img_sim.dataset: + raise ValueError( + 'Loaded embedding dataset is empty. '\ + f'Check embed_file={embed_file!r} or rerun with --embed.' + ) # Search to find similar images using cosine-similarity amongst embeddings. # Save all corpus results (including self-hit) so Recall@K can be computed at @@ -199,9 +236,29 @@ def _prep_fn(img): # Image saving is done separately at top_k to avoid PIL's 65500px dimension limit. start = time() n_results = len(prot_files) # all including self-hit - sim_time = img_sim.similarities(n=n_results, - save_result_images_dir=None, - pad_fn=_prep_fn) + + if args.faiss: + # ── FAISS ANN search ───────────────────────────────────────────── + if args.faiss_index_file: + faiss_index_file = args.faiss_index_file + else: + base, _ = os.path.splitext(embed_file) + faiss_index_file = base + '.faiss' + if os.path.exists(faiss_index_file) and not args.overwrite: + print(f'Loading existing FAISS index from {faiss_index_file}') + img_sim.load_faiss_index(faiss_index_file) + else: + print(f'Building FAISS index (use_pq={args.faiss_pq}) ...') + img_sim.build_faiss_index(use_pq=args.faiss_pq) + img_sim.save_faiss_index(faiss_index_file) + sim_time = img_sim.similarities_faiss(n=n_results, + save_result_images_dir=None, + pad_fn=_prep_fn) + else: + # ── Brute-force cosine search (original) ───────────────────────── + sim_time = img_sim.similarities(n=n_results, + save_result_images_dir=None, + pad_fn=_prep_fn) # Save top-k result images with padding full_sim_dict = {k: list(v) for k, v in img_sim.sim_dict.items()} @@ -222,7 +279,8 @@ def _prep_fn(img): for i, image_path in enumerate(prot_files): try: scores = img_sim.sim_dict[os.path.basename(image_path)] - df_res.iloc[i, :n_results] = [f'{a},{b}' for (a, b) in scores] + row_vals = [f'{a},{b}' for (a, b) in scores[:n_results]] + df_res.iloc[i, :len(row_vals)] = row_vals except KeyError as e: print(f'Key error for {e}') # Reorder cols diff --git a/uv.lock b/uv.lock index 39ac717..d2db451 100644 --- a/uv.lock +++ b/uv.lock @@ -651,30 +651,6 @@ toml = [ { name = "tomli", marker = "python_full_version <= '3.11'" }, ] -[[package]] -name = "cuda-bindings" -version = "12.9.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, - { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, - { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" }, - { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" }, - { url = "https://files.pythonhosted.org/packages/d1/af/6dfd8f2ed90b1d4719bc053ff8940e494640fe4212dc3dd72f383e4992da/cuda_bindings-12.9.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8b72ee72a9cc1b531db31eebaaee5c69a8ec3500e32c6933f2d3b15297b53686", size = 11922703, upload-time = "2025-10-21T14:52:03.585Z" }, - { url = "https://files.pythonhosted.org/packages/6c/19/90ac264acc00f6df8a49378eedec9fd2db3061bf9263bf9f39fd3d8377c3/cuda_bindings-12.9.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d80bffc357df9988dca279734bc9674c3934a654cab10cadeed27ce17d8635ee", size = 11924658, upload-time = "2025-10-21T14:52:10.411Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.4.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/60/d8f1dbfb7f06b94c662e98c95189e6f39b817da638bc8fcea0d003f89e5d/cuda_pathfinder-1.4.0-py3-none-any.whl", hash = "sha256:437079ca59e7b61ae439ecc501d69ed87b3accc34d58153ef1e54815e2c2e118", size = 38406, upload-time = "2026-02-25T22:13:00.807Z" }, -] - [[package]] name = "cycler" version = "0.12.1" @@ -764,6 +740,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" }, ] +[[package]] +name = "faiss-cpu" +version = "1.15.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "packaging" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/68/20e91694ad9a8b2bb48af956899e52b645cb1501e7e2ec31cb733da4d4c5/faiss_cpu-1.15.0-cp310-abi3-macosx_14_0_arm64.whl", hash = "sha256:50ea471ef1f4f3580eda8ab0ec9727d4bf65fd71c444bf306ce7cdbba8a42b21", size = 4904897, upload-time = "2026-08-03T17:49:37.003Z" }, + { url = "https://files.pythonhosted.org/packages/d2/cd/ef4cf498977c4a84af7a8920bc97ca49fc19060c8464c63fab58847b4692/faiss_cpu-1.15.0-cp310-abi3-macosx_15_0_x86_64.whl", hash = "sha256:dd383bb1ce06fabcff5785f998f253aa88f88dcbe1fe36c922417cd6666dd896", size = 7087977, upload-time = "2026-08-03T17:49:38.947Z" }, + { url = "https://files.pythonhosted.org/packages/94/c8/88b072bf55714405d0d7e11c12349510f15a69ae56033b1cd894fb2be7d6/faiss_cpu-1.15.0-cp310-abi3-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5d0a2d5d33fe023e263d0d355a837f20db67578e3be27fc5f4012a273274abf6", size = 9835009, upload-time = "2026-08-03T17:49:40.8Z" }, + { url = "https://files.pythonhosted.org/packages/c8/3b/8878dbfc78a0084bbd408b34827a58b530be98132fcf620b7e15f9191614/faiss_cpu-1.15.0-cp310-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ec9b29aae29e428c085c2d49dbb02e4673cdea75db418d420f9e60e0b4184498", size = 18764625, upload-time = "2026-08-03T17:49:43.676Z" }, + { url = "https://files.pythonhosted.org/packages/db/2a/654116e6ee2808562a6b2a11c396bdb46d45689e3bf7206ee99400589cab/faiss_cpu-1.15.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:30da3029952f0de69f16ce31946fd63fc3e292c867749bbcd2c0a0f09fd06f65", size = 11413863, upload-time = "2026-08-03T17:49:46.471Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8c/0a0f09659c1972aa83b9820cd3dd7f68f6678cfcfebde542e1c23d7d8663/faiss_cpu-1.15.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:88fbe1acac6978869063cb2f9477f85718da596a6e0a17751618f9c756bce255", size = 19470092, upload-time = "2026-08-03T17:49:50.253Z" }, + { url = "https://files.pythonhosted.org/packages/5b/b1/47967207659650ad74c1b06c42671e6beb4f7d798fe6eb2d53ba5e77ad90/faiss_cpu-1.15.0-cp311-cp311-win_amd64.whl", hash = "sha256:90169515a95ea58a9a95d419e518907927a8ef54c46788396365ec5902c9c8df", size = 16247194, upload-time = "2026-08-03T17:49:56.38Z" }, + { url = "https://files.pythonhosted.org/packages/b2/76/021398ec5608314124b554bb025878a86f129bcf3576c293826352d9a783/faiss_cpu-1.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:5b940897b317febaa761088513a3db164fad3ac71a5e1ed7be9a052c9bf1a447", size = 16251530, upload-time = "2026-08-03T17:50:00.166Z" }, + { url = "https://files.pythonhosted.org/packages/96/74/4a70395a6e07036628a1bd0b3f709101a6aecfa6a746db13b6e7921cf291/faiss_cpu-1.15.0-cp313-cp313-win_amd64.whl", hash = "sha256:22dddb013e764aad66dac6cd15b49c7598d60339e0591b73b5e081629419c21b", size = 16251914, upload-time = "2026-08-03T17:50:03.293Z" }, + { url = "https://files.pythonhosted.org/packages/ec/13/0a021b9df16963f839a3f325657656b70f23b5a6dbeb422eaa187d0121b3/faiss_cpu-1.15.0-cp314-cp314-win_amd64.whl", hash = "sha256:37170d5e9ead4b6bfd9c314afc39e17e92064068a0c5a4063dd3f39568c2667e", size = 16535739, upload-time = "2026-08-03T17:50:06.714Z" }, +] + [[package]] name = "fastjsonschema" version = "2.21.2" @@ -1052,6 +1049,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "intel-openmp" +version = "2021.4.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/18/527f247d673ff84c38e0b353b6901539b99e83066cd505be42ad341ab16d/intel_openmp-2021.4.0-py2.py3-none-win32.whl", hash = "sha256:6e863d8fd3d7e8ef389d52cf97a50fe2afe1a19247e8c0d168ce021546f96fc9", size = 1860605, upload-time = "2021-09-28T17:03:44.748Z" }, + { url = "https://files.pythonhosted.org/packages/6f/21/b590c0cc3888b24f2ac9898c41d852d7454a1695fbad34bee85dba6dc408/intel_openmp-2021.4.0-py2.py3-none-win_amd64.whl", hash = "sha256:eef4c8bcc8acefd7f5cd3b9384dbf73d59e2c99fc56545712ded913f43c4a94f", size = 3516906, upload-time = "2021-09-28T17:03:50.453Z" }, +] + [[package]] name = "ipykernel" version = "7.2.0" @@ -1782,6 +1788,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/f7/4a5e785ec9fbd65146a27b6b70b6cdc161a66f2024e4b04ac06a67f5578b/mistune-3.2.0-py3-none-any.whl", hash = "sha256:febdc629a3c78616b94393c6580551e0e34cc289987ec6c35ed3f4be42d0eee1", size = 53598, upload-time = "2025-12-23T11:36:33.211Z" }, ] +[[package]] +name = "mkl" +version = "2021.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "intel-openmp", marker = "sys_platform == 'win32'" }, + { name = "tbb", marker = "sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/c6/892fe3bc91e811b78e4f85653864f2d92541d5e5c306b0cb3c2311e9ca64/mkl-2021.4.0-py2.py3-none-win32.whl", hash = "sha256:439c640b269a5668134e3dcbcea4350459c4a8bc46469669b2d67e07e3d330e8", size = 129048357, upload-time = "2021-09-28T17:08:58.256Z" }, + { url = "https://files.pythonhosted.org/packages/fe/1c/5f6dbf18e8b73e0a5472466f0ea8d48ce9efae39bd2ff38cebf8dce61259/mkl-2021.4.0-py2.py3-none-win_amd64.whl", hash = "sha256:ceef3cafce4c009dd25f65d7ad0d833a0fbadc3d8903991ec92351fe5de1e718", size = 228499609, upload-time = "2021-09-28T17:09:19.683Z" }, +] + [[package]] name = "mmtf-python" version = "1.1.3" @@ -2092,15 +2111,15 @@ wheels = [ [[package]] name = "nvidia-cublas-cu12" -version = "12.8.4.1" +version = "12.1.3.1" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, + { url = "https://files.pythonhosted.org/packages/37/6d/121efd7382d5b0284239f4ab1fc1590d86d34ed4a4a2fdb13b30ca8e5740/nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl", hash = "sha256:ee53ccca76a6fc08fb9701aa95b6ceb242cdaab118c3bb152af4e579af792728", size = 410594774, upload-time = "2023-04-19T15:50:03.519Z" }, ] [[package]] name = "nvidia-cuda-cupti-cu12" -version = "12.8.90" +version = "12.1.105" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", @@ -2109,8 +2128,7 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/1f/b3bd73445e5cb342727fd24fe1f7b748f690b460acadc27ea22f904502c8/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed", size = 9533318, upload-time = "2025-03-07T01:40:10.421Z" }, - { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, + { url = "https://files.pythonhosted.org/packages/7e/00/6b218edd739ecfc60524e585ba8e6b00554dd908de2c9c66c1af3e44e18d/nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:e54fde3983165c624cb79254ae9818a456eb6e87a7fd4d56a2352c24ee542d7e", size = 14109015, upload-time = "2023-04-19T15:47:32.502Z" }, ] [[package]] @@ -2137,54 +2155,25 @@ wheels = [ [[package]] name = "nvidia-cuda-nvcc-cu12" -version = "12.9.86" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, - { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, - { url = "https://files.pythonhosted.org/packages/d2/9e/c71c53655a65d7531c89421c282359e2f626838762f1ce6180ea0bbebd29/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:8ed7f0b17dea662755395be029376db3b94fed5cbb17c2d35cc866c5b1b84099", size = 34669845, upload-time = "2025-06-05T20:11:56.308Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.8.93" +version = "12.1.105" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", -] wheels = [ - { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, - { url = "https://files.pythonhosted.org/packages/eb/d1/e50d0acaab360482034b84b6e27ee83c6738f7d32182b987f9c7a4e32962/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8", size = 43106076, upload-time = "2025-03-07T01:41:59.817Z" }, + { url = "https://files.pythonhosted.org/packages/ac/12/1d881feaf81b0b56e11ee7fdde9a688ffd4e09208136ea6cfeea7b05edb9/nvidia_cuda_nvcc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:a0712a05ffd57d9f017580ffcc808b207f7f4ef7a5ca8ac1312318b153739a18", size = 19996355, upload-time = "2023-04-19T15:48:02.05Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d8/c32d44063f58a38643392a737b281f22d04543410147d2c9bcd6a68f236c/nvidia_cuda_nvcc_cu12-12.1.105-py3-none-win_amd64.whl", hash = "sha256:87079d62e4507c7f60d355752edaceeaa2de762f69d8c914634655b6115f20af", size = 15992470, upload-time = "2023-04-19T15:54:03.342Z" }, ] [[package]] name = "nvidia-cuda-nvrtc-cu12" -version = "12.9.86" +version = "12.1.105" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.14' and platform_machine == 'ARM64' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine != 'ARM64' and sys_platform == 'win32'", - "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine == 'ARM64' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine != 'ARM64' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and platform_machine == 'ARM64' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and platform_machine != 'ARM64' and sys_platform == 'win32'", - "python_full_version < '3.12' and platform_machine == 'ARM64' and sys_platform == 'win32'", - "python_full_version < '3.12' and platform_machine != 'ARM64' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version < '3.12' and sys_platform == 'emscripten'", -] wheels = [ - { url = "https://files.pythonhosted.org/packages/52/de/823919be3b9d0ccbf1f784035423c5f18f4267fb0123558d58b813c6ec86/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:72972ebdcf504d69462d3bcd67e7b81edd25d0fb85a2c46d3ea3517666636349", size = 76408187, upload-time = "2025-06-05T20:12:27.819Z" }, + { url = "https://files.pythonhosted.org/packages/b6/9f/c64c03f49d6fbc56196664d05dba14e3a561038a81a638eeb47f4d4cfd48/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:339b385f50c309763ca65456ec75e17bbefcbbf2893f462cb8b90584cd27a1c2", size = 23671734, upload-time = "2023-04-19T15:48:32.42Z" }, + { url = "https://files.pythonhosted.org/packages/ad/1d/f76987c4f454eb86e0b9a0e4f57c3bf1ac1d13ad13cd1a4da4eb0e0c0ce9/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-win_amd64.whl", hash = "sha256:0a98a522d9ff138b96c010a65e145dc1b4850e9ecb75a0172371793752fd46ed", size = 19331863, upload-time = "2023-04-19T15:54:34.603Z" }, ] [[package]] name = "nvidia-cuda-runtime-cu12" -version = "12.8.90" +version = "12.1.105" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", @@ -2193,8 +2182,7 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265, upload-time = "2025-03-07T01:39:43.533Z" }, - { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, + { url = "https://files.pythonhosted.org/packages/eb/d5/c68b1d2cdfcc59e72e8a5949a37ddb22ae6cade80cd4a57a84d4c8b55472/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40", size = 823596, upload-time = "2023-04-19T15:47:22.471Z" }, ] [[package]] @@ -2221,18 +2209,18 @@ wheels = [ [[package]] name = "nvidia-cudnn-cu12" -version = "9.10.2.21" +version = "8.9.2.26" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, + { url = "https://files.pythonhosted.org/packages/ff/74/a2e2be7fb83aaedec84f391f082cf765dfb635e7caa9b49065f73e4835d8/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl", hash = "sha256:5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9", size = 731725872, upload-time = "2023-06-01T19:24:57.328Z" }, ] [[package]] name = "nvidia-cufft-cu12" -version = "11.3.3.83" +version = "11.0.2.54" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", @@ -2240,12 +2228,8 @@ resolution-markers = [ "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] -dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" }, - { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, + { url = "https://files.pythonhosted.org/packages/86/94/eb540db023ce1d162e7bea9f8f5aa781d57c65aed513c33ee9a5123ead4d/nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl", hash = "sha256:794e3948a1aa71fd817c3775866943936774d1c14e7628c74f6f7417224cdf56", size = 121635161, upload-time = "2023-04-19T15:50:46Z" }, ] [[package]] @@ -2273,25 +2257,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/ee/29955203338515b940bd4f60ffdbc073428f25ef9bfbce44c9a066aedc5c/nvidia_cufft_cu12-11.4.1.4-py3-none-win_amd64.whl", hash = "sha256:8e5bfaac795e93f80611f807d42844e8e27e340e0cde270dcb6c65386d795b80", size = 200067309, upload-time = "2025-06-05T20:13:59.762Z" }, ] -[[package]] -name = "nvidia-cufile-cu12" -version = "1.13.1.3" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, -] - [[package]] name = "nvidia-curand-cu12" -version = "10.3.9.90" +version = "10.3.2.106" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, + { url = "https://files.pythonhosted.org/packages/44/31/4890b1c9abc496303412947fc7dcea3d14861720642b49e8ceed89636705/nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl", hash = "sha256:9d264c5036dde4e64f1de8c50ae753237c12e0b1348738169cd0f8a536c0e1e0", size = 56467784, upload-time = "2023-04-19T15:51:04.804Z" }, ] [[package]] name = "nvidia-cusolver-cu12" -version = "11.7.3.90" +version = "11.4.5.107" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, @@ -2299,34 +2275,26 @@ dependencies = [ { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, + { url = "https://files.pythonhosted.org/packages/bc/1d/8de1e5c67099015c834315e333911273a8c6aaba78923dd1d1e25fc5f217/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl", hash = "sha256:8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd", size = 124161928, upload-time = "2023-04-19T15:51:25.781Z" }, ] [[package]] name = "nvidia-cusparse-cu12" -version = "12.5.8.93" +version = "12.1.0.106" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, -] - -[[package]] -name = "nvidia-cusparselt-cu12" -version = "0.7.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, + { url = "https://files.pythonhosted.org/packages/65/5b/cfaeebf25cd9fdec14338ccb16f6b2c4c7fa9163aefcf057d86b9cc248bb/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl", hash = "sha256:f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c", size = 195958278, upload-time = "2023-04-19T15:51:49.939Z" }, ] [[package]] name = "nvidia-nccl-cu12" -version = "2.27.5" +version = "2.20.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, + { url = "https://files.pythonhosted.org/packages/4b/2a/0a131f572aa09f741c30ccd45a8e56316e8be8dfc7bc19bf0ab7cfef7b19/nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:057f6bf9685f75215d0c53bf3ac4a10b3e6578351de307abad9e18a99182af56", size = 176249402, upload-time = "2024-03-06T04:30:20.663Z" }, ] [[package]] @@ -2341,7 +2309,6 @@ resolution-markers = [ ] wheels = [ { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, - { url = "https://files.pythonhosted.org/packages/2a/a2/8cee5da30d13430e87bf99bb33455d2724d0a4a9cb5d7926d80ccb96d008/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7", size = 38386204, upload-time = "2025-03-07T01:49:43.612Z" }, ] [[package]] @@ -2366,20 +2333,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dd/7e/2eecb277d8a98184d881fb98a738363fd4f14577a4d2d7f8264266e82623/nvidia_nvjitlink_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:cc6fcec260ca843c10e34c936921a1c426b351753587fdd638e8cff7b16bb9db", size = 35584936, upload-time = "2025-06-05T20:16:08.525Z" }, ] -[[package]] -name = "nvidia-nvshmem-cu12" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, -] - [[package]] name = "nvidia-nvtx-cu12" -version = "12.8.90" +version = "12.1.105" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, + { url = "https://files.pythonhosted.org/packages/da/d3/8057f0587683ed2fcd4dbfbdfdfa807b9160b809976099d36b8f60d08f03/nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:dc21cf308ca5691e7c04d962e213f8a4aa9bbfa23d95412f452254c2caeb09e5", size = 99138, upload-time = "2023-04-19T15:48:43.556Z" }, ] [[package]] @@ -2422,14 +2381,13 @@ name = "openmm-cuda-12" version = "8.4.0.post2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-cupti-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cuda-cupti-cu12", version = "12.1.105", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, { name = "nvidia-cuda-cupti-cu12", version = "12.9.79", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, { name = "nvidia-cuda-nvcc-cu12" }, - { name = "nvidia-cuda-nvrtc-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cuda-nvrtc-cu12", version = "12.9.86", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, - { name = "nvidia-cuda-runtime-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cuda-nvrtc-cu12" }, + { name = "nvidia-cuda-runtime-cu12", version = "12.1.105", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, { name = "nvidia-cuda-runtime-cu12", version = "12.9.79", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, - { name = "nvidia-cufft-cu12", version = "11.3.3.83", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cufft-cu12", version = "11.0.2.54", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, { name = "nvidia-cufft-cu12", version = "11.4.1.4", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, ] wheels = [ @@ -2745,6 +2703,7 @@ version = "0.0.4" source = { editable = "." } dependencies = [ { name = "biopython" }, + { name = "faiss-cpu" }, { name = "goatools" }, { name = "kmeans-pytorch" }, { name = "matplotlib" }, @@ -2767,8 +2726,12 @@ dependencies = [ [package.optional-dependencies] cuda12 = [ + { name = "nvidia-cuda-nvcc-cu12" }, + { name = "nvidia-cuda-nvrtc-cu12" }, { name = "openmm" }, { name = "openmm-cuda-12" }, + { name = "torch" }, + { name = "torchvision" }, ] notebook = [ { name = "jupyterlab" }, @@ -2782,6 +2745,7 @@ test = [ [package.metadata] requires-dist = [ { name = "biopython", specifier = ">=1.8" }, + { name = "faiss-cpu", specifier = ">=1.13.2" }, { name = "goatools", specifier = ">=1.4" }, { name = "jupyterlab", marker = "extra == 'notebook'", specifier = ">=4.2.5" }, { name = "kmeans-pytorch", specifier = ">=0.3" }, @@ -2789,6 +2753,8 @@ requires-dist = [ { name = "mdanalysis", extras = ["analysis", "extra-formats", "parallel"], specifier = ">=2.10.0" }, { name = "nglview", marker = "extra == 'notebook'", specifier = ">=3.1.4" }, { name = "numpy", specifier = ">=1.26" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "extra == 'cuda12'", specifier = "==12.1.105" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "extra == 'cuda12'", specifier = "==12.1.105" }, { name = "objgraph", specifier = ">=3.6.2" }, { name = "openmm", specifier = ">=8.4" }, { name = "openmm", marker = "extra == 'cuda12'", specifier = "==8.4.0" }, @@ -2803,8 +2769,10 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "rcsb-api", specifier = ">=1.7.3" }, { name = "torch", specifier = ">=2.2,<2.11" }, + { name = "torch", marker = "extra == 'cuda12'", specifier = "==2.3.0" }, { name = "torchsummary", specifier = ">=1.5" }, { name = "torchvision", specifier = ">=0.17,<0.27" }, + { name = "torchvision", marker = "extra == 'cuda12'", specifier = "==0.18.0" }, { name = "tqdm", specifier = ">=4.67" }, ] provides-extras = ["cuda12", "test", "notebook"] @@ -3608,6 +3576,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "tbb" +version = "2021.13.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/8a/5062b00c378c051e26507e5eca8d3b5c91ed63f8a2139f6f0f422be84b02/tbb-2021.13.1-py3-none-win32.whl", hash = "sha256:00f5e5a70051650ddd0ab6247c0549521968339ec21002e475cd23b1cbf46d66", size = 248994, upload-time = "2024-08-07T15:10:08.934Z" }, + { url = "https://files.pythonhosted.org/packages/9b/24/84ce997e8ae6296168a74d0d9c4dde572d90fb23fd7c0b219c30ff71e00e/tbb-2021.13.1-py3-none-win_amd64.whl", hash = "sha256:cbf024b2463fdab3ebe3fa6ff453026358e6b903839c80d647e08ad6d0796ee9", size = 286908, upload-time = "2024-08-07T15:09:05.677Z" }, +] + [[package]] name = "terminado" version = "0.18.1" @@ -3720,68 +3697,38 @@ wheels = [ [[package]] name = "torch" -version = "2.10.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, + { name = "mkl", marker = "sys_platform == 'win32'" }, { name = "networkx" }, { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-cupti-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-runtime-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", version = "12.1.105", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", version = "12.1.105", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufft-cu12", version = "11.3.3.83", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", version = "11.0.2.54", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, { name = "sympy" }, - { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "triton", marker = "python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" }, - { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, - { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, - { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" }, - { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, - { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, - { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, - { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" }, - { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" }, - { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" }, - { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" }, - { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" }, - { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" }, - { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, - { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, - { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, - { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, - { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, - { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, - { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" }, - { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" }, - { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" }, - { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" }, - { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" }, - { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" }, - { url = "https://files.pythonhosted.org/packages/4f/93/716b5ac0155f1be70ed81bacc21269c3ece8dba0c249b9994094110bfc51/torch-2.10.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:bf0d9ff448b0218e0433aeb198805192346c4fd659c852370d5cc245f602a06a", size = 79464992, upload-time = "2026-01-21T16:23:05.162Z" }, - { url = "https://files.pythonhosted.org/packages/69/2b/51e663ff190c9d16d4a8271203b71bc73a16aa7619b9f271a69b9d4a936b/torch-2.10.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:233aed0659a2503b831d8a67e9da66a62c996204c0bba4f4c442ccc0c68a3f60", size = 146018567, upload-time = "2026-01-21T16:22:23.393Z" }, - { url = "https://files.pythonhosted.org/packages/5e/cd/4b95ef7f293b927c283db0b136c42be91c8ec6845c44de0238c8c23bdc80/torch-2.10.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:682497e16bdfa6efeec8cde66531bc8d1fbbbb4d8788ec6173c089ed3cc2bfe5", size = 915721646, upload-time = "2026-01-21T16:21:16.983Z" }, - { url = "https://files.pythonhosted.org/packages/56/97/078a007208f8056d88ae43198833469e61a0a355abc0b070edd2c085eb9a/torch-2.10.0-cp314-cp314-win_amd64.whl", hash = "sha256:6528f13d2a8593a1a412ea07a99812495bec07e9224c28b2a25c0a30c7da025c", size = 113752373, upload-time = "2026-01-21T16:22:13.471Z" }, - { url = "https://files.pythonhosted.org/packages/d8/94/71994e7d0d5238393df9732fdab607e37e2b56d26a746cb59fdb415f8966/torch-2.10.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f5ab4ba32383061be0fb74bda772d470140a12c1c3b58a0cfbf3dae94d164c28", size = 79850324, upload-time = "2026-01-21T16:22:09.494Z" }, - { url = "https://files.pythonhosted.org/packages/e2/65/1a05346b418ea8ccd10360eef4b3e0ce688fba544e76edec26913a8d0ee0/torch-2.10.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:716b01a176c2a5659c98f6b01bf868244abdd896526f1c692712ab36dbaf9b63", size = 146006482, upload-time = "2026-01-21T16:22:18.42Z" }, - { url = "https://files.pythonhosted.org/packages/1d/b9/5f6f9d9e859fc3235f60578fa64f52c9c6e9b4327f0fe0defb6de5c0de31/torch-2.10.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d8f5912ba938233f86361e891789595ff35ca4b4e2ac8fe3670895e5976731d6", size = 915613050, upload-time = "2026-01-21T16:20:49.035Z" }, - { url = "https://files.pythonhosted.org/packages/66/4d/35352043ee0eaffdeff154fad67cd4a31dbed7ff8e3be1cc4549717d6d51/torch-2.10.0-cp314-cp314t-win_amd64.whl", hash = "sha256:71283a373f0ee2c89e0f0d5f446039bdabe8dbc3c9ccf35f0f784908b0acd185", size = 113995816, upload-time = "2026-01-21T16:22:05.312Z" }, + { url = "https://files.pythonhosted.org/packages/35/3a/a39f354fa3119785be87e2f94ffa2620f8a270c8560f7356358ee62fb4c5/torch-2.3.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:493d54ee2f9df100b5ce1d18c96dbb8d14908721f76351e908c9d2622773a788", size = 779160265, upload-time = "2024-04-24T15:46:28.108Z" }, + { url = "https://files.pythonhosted.org/packages/91/3c/98a9b901b40278b40a9ff5b9c6be387b20997269f929f2182d8845c94085/torch-2.3.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:bce43af735c3da16cc14c7de2be7ad038e2fbf75654c2e274e575c6c05772ace", size = 88536251, upload-time = "2024-04-24T15:47:25.229Z" }, + { url = "https://files.pythonhosted.org/packages/2a/b7/a3cf5fd40334b9785cc83ee0c96b50603026eb3aa70210a33729018e7029/torch-2.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:729804e97b7cf19ae9ab4181f91f5e612af07956f35c8b2c8e9d9f3596a8e877", size = 159803952, upload-time = "2024-04-24T15:47:01.3Z" }, + { url = "https://files.pythonhosted.org/packages/ad/08/c5e41eb22323db4a52260607598a207a2e1918916ae8201aa7a8ae005fcd/torch-2.3.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:d24e328226d8e2af7cf80fcb1d2f1d108e0de32777fab4aaa2b37b9765d8be73", size = 60998957, upload-time = "2024-04-24T15:47:12.528Z" }, + { url = "https://files.pythonhosted.org/packages/51/03/1abad10990c76bee3703857b1617563b241f87d297ee466dbad922b0c308/torch-2.3.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:b0de2bdc0486ea7b14fc47ff805172df44e421a7318b7c4d92ef589a75d27410", size = 779062531, upload-time = "2024-04-24T15:45:26.461Z" }, + { url = "https://files.pythonhosted.org/packages/f1/9d/dfe273e19b7165148208bd182fac87488c5a0f7a3318d1646f5f37af872f/torch-2.3.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:a306c87a3eead1ed47457822c01dfbd459fe2920f2d38cbdf90de18f23f72542", size = 88437165, upload-time = "2024-04-24T15:47:07.772Z" }, + { url = "https://files.pythonhosted.org/packages/37/04/a5cd83baccbf2d4329990ec06b8abf3a644e1559a7b1f764f42d2cb77d51/torch-2.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:f9b98bf1a3c8af2d4c41f0bf1433920900896c446d1ddc128290ff146d1eb4bd", size = 159749140, upload-time = "2024-04-24T15:46:47.303Z" }, + { url = "https://files.pythonhosted.org/packages/55/51/4bdee83e6fa9cca8e3a6cdf81a2695ede9d3fd7148e4fd4188dff142d7b0/torch-2.3.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:dca986214267b34065a79000cee54232e62b41dff1ec2cab9abc3fc8b3dee0ad", size = 60968873, upload-time = "2024-04-24T15:46:55.552Z" }, ] [[package]] @@ -3795,7 +3742,7 @@ wheels = [ [[package]] name = "torchvision" -version = "0.25.0" +version = "0.18.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, @@ -3803,30 +3750,14 @@ dependencies = [ { name = "torch" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/3e/be/c704bceaf11c4f6b19d64337a34a877fcdfe3bd68160a8c9ae9bea4a35a3/torchvision-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:db74a551946b75d19f9996c419a799ffdf6a223ecf17c656f90da011f1d75b20", size = 1874923, upload-time = "2026-01-21T16:27:46.574Z" }, - { url = "https://files.pythonhosted.org/packages/ae/e9/f143cd71232430de1f547ceab840f68c55e127d72558b1061a71d0b193cd/torchvision-0.25.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f49964f96644dbac2506dffe1a0a7ec0f2bf8cf7a588c3319fed26e6329ffdf3", size = 2344808, upload-time = "2026-01-21T16:27:43.191Z" }, - { url = "https://files.pythonhosted.org/packages/43/ae/ad5d6165797de234c9658752acb4fce65b78a6a18d82efdf8367c940d8da/torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:153c0d2cbc34b7cf2da19d73450f24ba36d2b75ec9211b9962b5022fb9e4ecee", size = 8070752, upload-time = "2026-01-21T16:27:33.748Z" }, - { url = "https://files.pythonhosted.org/packages/23/19/55b28aecdc7f38df57b8eb55eb0b14a62b470ed8efeb22cdc74224df1d6a/torchvision-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:ea580ffd6094cc01914ad32f8c8118174f18974629af905cea08cb6d5d48c7b7", size = 4038722, upload-time = "2026-01-21T16:27:41.355Z" }, - { url = "https://files.pythonhosted.org/packages/56/3a/6ea0d73f49a9bef38a1b3a92e8dd455cea58470985d25635beab93841748/torchvision-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2abe430c90b1d5e552680037d68da4eb80a5852ebb1c811b2b89d299b10573b", size = 1874920, upload-time = "2026-01-21T16:27:45.348Z" }, - { url = "https://files.pythonhosted.org/packages/51/f8/c0e1ef27c66e15406fece94930e7d6feee4cb6374bbc02d945a630d6426e/torchvision-0.25.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b75deafa2dfea3e2c2a525559b04783515e3463f6e830cb71de0fb7ea36fe233", size = 2344556, upload-time = "2026-01-21T16:27:40.125Z" }, - { url = "https://files.pythonhosted.org/packages/68/2f/f24b039169db474e8688f649377de082a965fbf85daf4e46c44412f1d15a/torchvision-0.25.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:f25aa9e380865b11ea6e9d99d84df86b9cc959f1a007cd966fc6f1ab2ed0e248", size = 8072351, upload-time = "2026-01-21T16:27:21.074Z" }, - { url = "https://files.pythonhosted.org/packages/ad/16/8f650c2e288977cf0f8f85184b90ee56ed170a4919347fc74ee99286ed6f/torchvision-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:f9c55ae8d673ab493325d1267cbd285bb94d56f99626c00ac4644de32a59ede3", size = 4303059, upload-time = "2026-01-21T16:27:11.08Z" }, - { url = "https://files.pythonhosted.org/packages/f5/5b/1562a04a6a5a4cf8cf40016a0cdeda91ede75d6962cff7f809a85ae966a5/torchvision-0.25.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:24e11199e4d84ba9c5ee7825ebdf1cd37ce8deec225117f10243cae984ced3ec", size = 1874918, upload-time = "2026-01-21T16:27:39.02Z" }, - { url = "https://files.pythonhosted.org/packages/36/b1/3d6c42f62c272ce34fcce609bb8939bdf873dab5f1b798fd4e880255f129/torchvision-0.25.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f271136d2d2c0b7a24c5671795c6e4fd8da4e0ea98aeb1041f62bc04c4370ef", size = 2309106, upload-time = "2026-01-21T16:27:30.624Z" }, - { url = "https://files.pythonhosted.org/packages/c7/60/59bb9c8b67cce356daeed4cb96a717caa4f69c9822f72e223a0eae7a9bd9/torchvision-0.25.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:855c0dc6d37f462482da7531c6788518baedca1e0847f3df42a911713acdfe52", size = 8071522, upload-time = "2026-01-21T16:27:29.392Z" }, - { url = "https://files.pythonhosted.org/packages/32/a5/9a9b1de0720f884ea50dbf9acb22cbe5312e51d7b8c4ac6ba9b51efd9bba/torchvision-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:cef0196be31be421f6f462d1e9da1101be7332d91984caa6f8022e6c78a5877f", size = 4321911, upload-time = "2026-01-21T16:27:35.195Z" }, - { url = "https://files.pythonhosted.org/packages/52/99/dca81ed21ebaeff2b67cc9f815a20fdaa418b69f5f9ea4c6ed71721470db/torchvision-0.25.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a8f8061284395ce31bcd460f2169013382ccf411148ceb2ee38e718e9860f5a7", size = 1896209, upload-time = "2026-01-21T16:27:32.159Z" }, - { url = "https://files.pythonhosted.org/packages/28/cc/2103149761fdb4eaed58a53e8437b2d716d48f05174fab1d9fcf1e2a2244/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:146d02c9876858420adf41f3189fe90e3d6a409cbfa65454c09f25fb33bf7266", size = 2310735, upload-time = "2026-01-21T16:27:22.327Z" }, - { url = "https://files.pythonhosted.org/packages/76/ad/f4c985ad52ddd3b22711c588501be1b330adaeaf6850317f66751711b78c/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c4d395cb2c4a2712f6eb93a34476cdf7aae74bb6ea2ea1917f858e96344b00aa", size = 8089557, upload-time = "2026-01-21T16:27:27.666Z" }, - { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" }, - { url = "https://files.pythonhosted.org/packages/9e/1f/fa839532660e2602b7e704d65010787c5bb296258b44fa8b9c1cd6175e7d/torchvision-0.25.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:620a236288d594dcec7634c754484542dc0a5c1b0e0b83a34bda5e91e9b7c3a1", size = 1896193, upload-time = "2026-01-21T16:27:24.785Z" }, - { url = "https://files.pythonhosted.org/packages/80/ed/d51889da7ceaf5ff7a0574fb28f9b6b223df19667265395891f81b364ab3/torchvision-0.25.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:0b5e7f50002a8145a98c5694a018e738c50e2972608310c7e88e1bd4c058f6ce", size = 2309331, upload-time = "2026-01-21T16:27:19.97Z" }, - { url = "https://files.pythonhosted.org/packages/90/a5/f93fcffaddd8f12f9e812256830ec9c9ca65abbf1bc369379f9c364d1ff4/torchvision-0.25.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:632db02300e83793812eee4f61ae6a2686dab10b4cfd628b620dc47747aa9d03", size = 8088713, upload-time = "2026-01-21T16:27:15.281Z" }, - { url = "https://files.pythonhosted.org/packages/1f/eb/d0096eed5690d962853213f2ee00d91478dfcb586b62dbbb449fb8abc3a6/torchvision-0.25.0-cp314-cp314-win_amd64.whl", hash = "sha256:d1abd5ed030c708f5dbf4812ad5f6fbe9384b63c40d6bd79f8df41a4a759a917", size = 4325058, upload-time = "2026-01-21T16:27:26.165Z" }, - { url = "https://files.pythonhosted.org/packages/97/36/96374a4c7ab50dea9787ce987815614ccfe988a42e10ac1a2e3e5b60319a/torchvision-0.25.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ad9a8a5877782944d99186e4502a614770fe906626d76e9cd32446a0ac3075f2", size = 1896207, upload-time = "2026-01-21T16:27:23.383Z" }, - { url = "https://files.pythonhosted.org/packages/b5/e2/7abb10a867db79b226b41da419b63b69c0bd5b82438c4a4ed50e084c552f/torchvision-0.25.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:40a122c3cf4d14b651f095e0f672b688dde78632783fc5cd3d4d5e4f6a828563", size = 2310741, upload-time = "2026-01-21T16:27:18.712Z" }, - { url = "https://files.pythonhosted.org/packages/08/e6/0927784e6ffc340b6676befde1c60260bd51641c9c574b9298d791a9cda4/torchvision-0.25.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:846890161b825b38aa85fc37fb3ba5eea74e7091ff28bab378287111483b6443", size = 8089772, upload-time = "2026-01-21T16:27:14.048Z" }, - { url = "https://files.pythonhosted.org/packages/b6/37/e7ca4ec820d434c0f23f824eb29f0676a0c3e7a118f1514f5b949c3356da/torchvision-0.25.0-cp314-cp314t-win_amd64.whl", hash = "sha256:f07f01d27375ad89d72aa2b3f2180f07da95dd9d2e4c758e015c0acb2da72977", size = 4425879, upload-time = "2026-01-21T16:27:12.579Z" }, + { url = "https://files.pythonhosted.org/packages/b5/14/c05da13c98f528ba5fd99897320a7684df5dd136ec6faa6a5766f25e4a7e/torchvision-0.18.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6896a52168befe1105fb3c9335287390ed227e71d1e4ec4d68b62e8a3099fc09", size = 1554225, upload-time = "2024-04-24T15:48:28.866Z" }, + { url = "https://files.pythonhosted.org/packages/6e/7d/bc67ec2d1011226e75cdcc45a2c85d97b8eaac32a7d648b71c432d584367/torchvision-0.18.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:3d7955398d4ceaad77c487c2c44f6f7813112402c9bab8cd906d346005891048", size = 6955003, upload-time = "2024-04-24T15:48:25.034Z" }, + { url = "https://files.pythonhosted.org/packages/70/1d/107894816bf5ebecbc5a8556743fd89c0a1dd6da82b1c562d6becd5a108a/torchvision-0.18.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:e5a24d620cea14a4bb89f24aa2b506230c0a16a3ada57fc53ad80cfd256a2128", size = 13995855, upload-time = "2024-04-24T15:48:11.523Z" }, + { url = "https://files.pythonhosted.org/packages/12/c2/7c89c62f2b0a606070aa7cdb8af8af0c905562942778ebdd77600642c3b9/torchvision-0.18.0-cp311-cp311-win_amd64.whl", hash = "sha256:6ad70ddfa879bda5ed886b2518fe562640e0059787cbd65cb2bffa7674541410", size = 1174125, upload-time = "2024-04-24T15:48:21.598Z" }, + { url = "https://files.pythonhosted.org/packages/7c/12/49d4fd4ae7a48c6d33babb01a523594aca38365d378518320a10f9a5baa1/torchvision-0.18.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:eb9d83c0e1dbb54ecb0fb04c87f786333e3a6fb8b9c400aca7c31081f9aa5707", size = 1554232, upload-time = "2024-04-24T15:48:26.578Z" }, + { url = "https://files.pythonhosted.org/packages/6d/fe/729256fec03403b0bfdc229d4350936e29020d618f1740009c6a4f995b06/torchvision-0.18.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:b657d052d146f24cb3b2a78219bfc82ae70a9706671c50f632528907d10cccec", size = 6955047, upload-time = "2024-04-24T15:47:57.381Z" }, + { url = "https://files.pythonhosted.org/packages/51/2d/30883e9c6734546f9e7e0c429b76bd2b651aa25f6ced087c9c11cbf2ef41/torchvision-0.18.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:a964afbc7ddf50a46b941477f6c35729b416deedd139756befd488245e2e226d", size = 13996688, upload-time = "2024-04-24T15:47:54.406Z" }, + { url = "https://files.pythonhosted.org/packages/53/8a/864c3969af219a95213a5065d453313a96598e7c744b859e99b6ac134e16/torchvision-0.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:7c770f0f748e0b17f57c0297508d7254f686cdf03fc2e2949f422b20574f4c0f", size = 1174123, upload-time = "2024-04-24T15:48:20.241Z" }, ] [[package]] @@ -3871,15 +3802,14 @@ wheels = [ [[package]] name = "triton" -version = "3.6.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock", marker = "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, - { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, - { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, - { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, - { url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" }, - { url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" }, + { url = "https://files.pythonhosted.org/packages/3c/00/84e0006f2025260fa111ddfc66194bd1af731b3ee18e2fd611a00f290b5e/triton-2.3.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3c3d9607f85103afdb279938fc1dd2a66e4f5999a58eb48a346bd42738f986dd", size = 168079300, upload-time = "2024-04-05T02:47:52.164Z" }, + { url = "https://files.pythonhosted.org/packages/90/2f/3e8f0ea4ef7ac6d8720a48b9b9700fd32c9909ee83b2eb1f25209ace0767/triton-2.3.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:218d742e67480d9581bafb73ed598416cc8a56f6316152e5562ee65e33de01c0", size = 168091361, upload-time = "2024-04-05T02:48:01.044Z" }, ] [[package]] From a1012733aed24eca3aef33027aaa40c7386a2b17 Mon Sep 17 00:00:00 2001 From: SonOfAnton Date: Fri, 11 Sep 2026 22:52:40 -0700 Subject: [PATCH 2/4] Fix FAISS index defaults, result truncation, and make faiss optional Review fixes on top of the initial FAISS backend: nlist defaulted to N instead of sqrt(N). The expression `N ** _NLIST_SQRT_FACTOR ** 0.5` parses as `N ** (1.0 ** 0.5)` because ** is right-associative, so every vector got its own Voronoi cell. FAISS warned on every build and training cost scaled quadratically (3.7s vs 0.5s at 20k x 512). search_all() silently truncated deep rankings. measure_similarity_v2.py asks for the full corpus ordering so Recall@K works at any K, but an IVF search only returns vectors in the cells it probes: 200 of 2008 requested results came back per query. The blank trailing CSV cells read back as NaN and crashed evaluate_methods_v2.py in target.split(','). search_all() now widens nprobe to reach the requested depth, restoring it afterwards so one deep search does not leave the index scanning exhaustively, and the CSV is sized to the shortest ranking actually produced. Added --faiss_top_k, because a full-corpus ranking forces an exhaustive scan and is 15x slower than brute force at the current corpus size (1.67s vs 0.11s at N=2008). Capping the depth is what makes the ANN path worth using. faiss-cpu moves to an optional "search" extra, matching the lazy import and install hint already in faiss_search.py. Adds tests/test_faiss_search.py covering the defaults, truncation, nprobe state, persistence and PQ fallbacks. Drops docs/improvements_v2.md: it plans three improvements, only one of which is implemented here, and its snippets no longer match the shipped API. --- docs/improvements_v2.md | 1366 --------------------------- proteogram/v2/faiss_search.py | 377 ++++---- proteogram/v2/image_similarity.py | 133 +-- pyproject.toml | 7 +- scripts/v2/measure_similarity_v2.py | 84 +- tests/test_faiss_search.py | 175 ++++ uv.lock | 8 +- 7 files changed, 473 insertions(+), 1677 deletions(-) delete mode 100644 docs/improvements_v2.md create mode 100644 tests/test_faiss_search.py diff --git a/docs/improvements_v2.md b/docs/improvements_v2.md deleted file mode 100644 index 2e9478f..0000000 --- a/docs/improvements_v2.md +++ /dev/null @@ -1,1366 +0,0 @@ -# Proteogram v2 — Three Targeted Improvements - -## Overview - -This document covers the design, implementation, and validation plan for three independent but complementary improvements to the Proteogram v2 pipeline: - -| # | Name | File(s) Affected | Impact | -|---|------|-----------------|--------| -| [1](#1-faiss-approximate-nearest-neighbour-search) | FAISS Approximate Nearest Neighbour Search | `proteogram/v2/image_similarity.py`, `scripts/v2/measure_similarity_v2.py` | Scales corpus search from minutes to milliseconds | -| [2](#2-global-percentile-normalisation) | Global Percentile Normalisation | `proteogram/v2/proteogram.py`, `scripts/v2/create_v2_proteograms.py` | Preserves physically meaningful inter-protein scale | -| [3](#3-grad-cam-explainability) | Grad-CAM Explainability | `proteogram/v2/image_similarity.py` (new method), new `scripts/v2/explain_similarity.py` | Residue-pair attribution for any similar pair | - -Each section follows the same structure: motivation → design decisions → full implementation → validation steps. - ---- - -## Operational Notes (May 2026): Environment Setup + Long v2 Runs - -This section documents practical lessons from running the current v2 pipeline on Linux with mixed toolchains (`uv`, local Conda bootstrap, OpenMM). - -### A. Why `create_v2_proteograms.py` may be slow - -`scripts/v2/create_v2_proteograms.py` runs a full MD pipeline per protein (minimization + NPT + NVT + production) before image export. If OpenMM CUDA is unavailable, this falls back to CPU and runtime increases significantly. - -At default MD lengths, CPU runtime can be many minutes per protein; with 2,008 proteins this can become multi-day if not accelerated. - -### B. Distinguish PyTorch CUDA vs OpenMM CUDA - -It is common to have: - -- `torch.cuda.is_available() == True` -- OpenMM platforms = `['Reference', 'CPU', 'OpenCL']` - -In this case, similarity scripts can use GPU (PyTorch), but MD in v2 proteogram generation still runs without CUDA. - -Check OpenMM platforms directly: - -```bash -python - <<'PY' -from openmm import Platform -names = [Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())] -print('OpenMM platforms:', names) -print('CUDA available:', 'CUDA' in names) -PY -``` - -### C. Conda bootstrap pitfall encountered - -In this run, `conda` was pointing to a local bootstrap install under: - -`scripts/v2/exit/bin/conda` - -This caused solver and lock issues (e.g., sqlite lock, libmamba plugin mismatch), and prevented reliable environment creation. - -Recommended safeguards: - -1. Confirm which conda is active (`which conda`, `conda info --base`). -2. Prefer a stable system conda/mamba/micromamba install for OpenMM-CUDA env creation. -3. If needed, force classic solver when libmamba plugin is unavailable. - -### D. Minimal reliable runbook (CUDA-capable OpenMM env) - -```bash -# 1) create env with python 3.11 (recommended for this project stack) -conda create -n proteogram-openmm-cuda -c conda-forge python=3.11 openmm pdbfixer -y - -# 2) activate env -conda activate proteogram-openmm-cuda - -# 3) verify OpenMM CUDA platform visibility -python - <<'PY' -from openmm import Platform -print([Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())]) -PY - -# 4) install project in editable mode -cd /path/to/proteogram -pip install -e . - -# 5) run v2 proteogram creation -cd scripts/v2 -python create_v2_proteograms.py --overwrite -``` - -### E. Monitoring a long run - -During `create_v2_proteograms.py`, JPG outputs are written incrementally (not only at end). To watch output growth: - -```bash -# from repo root -watch -n 5 'find data/scope2.08_all_proteograms_v2 -maxdepth 1 -name "*.jpg" | wc -l' - -# or from scripts/v2 -watch -n 5 'find ../data/scope2.08_all_proteograms_v2 -maxdepth 1 -name "*.jpg" | wc -l' -``` - -### F. Current observed status - -- Run is progressing through structures and skipping out-of-range chains (`sequence length outside [20, 200]`) as designed. -- Output directory image count should rise continuously as proteins complete. -- If OpenMM CUDA remains unavailable, OpenCL/CPU execution is expected and slower than CUDA. - ---- - -## 1. FAISS Approximate Nearest Neighbour Search - -### 1.1 Motivation - -The current `Img2Vec.similarities()` method computes cosine similarity between every pair of embeddings in the corpus: - -```python -# proteogram/v2/image_similarity.py — existing inner loop (O(N²)) -for image_path_i, embedding_i in tqdm(self.dataset.items()): - for image_path_j, embedding_j in self.dataset.items(): - sim = cosine(embedding_i, embedding_j)[0].item() -``` - -For a corpus of N proteins this is O(N²) in both time and sequential memory access. At the current SCOPe 2.08 scale (~100 K domains) this already takes tens of minutes. The AlphaFold Database (AFDB) contains ~200 million predicted structures — brute-force search there would take weeks. - -FAISS (Facebook AI Similarity Search) replaces this with an **Inverted File index with Product Quantisation (IVF-PQ)** that gives sub-linear query time with controllable recall-accuracy trade-offs. - -### 1.2 Design Decisions - -#### Index type: `IndexIVFFlat` for small corpora, `IndexIVFPQ` for large - -| Corpus size | Recommended index | Reason | -|-------------|------------------|--------| -| ≤ 100 K | `IndexIVFFlat` | Exact L2/IP; no quantisation error; fast enough | -| 100 K – 10 M | `IndexIVFPQ` | 8–32× memory reduction; ~1% recall loss | -| > 10 M | `IndexIVFPQ` + `OPQ` pre-rotation | Best recall at extreme scale | - -The `Img2Vec` class will default to `IndexIVFFlat` and let callers opt into `IndexIVFPQ`. - -#### Inner-product (IP) vs. L2 - -FAISS supports both. Because the rest of the codebase uses **cosine similarity**, we L2-normalise embeddings before indexing and use **inner product** — on unit vectors, inner product equals cosine similarity exactly. This avoids any change to how scores are interpreted. - -#### `nlist` (number of Voronoi cells) - -A rule of thumb is `nlist = sqrt(N)`. For 100 K vectors: `nlist = 316`. For 10 M vectors: `nlist = 3162`. These values will be set automatically if not specified. - -#### `nprobe` (cells searched at query time) - -Higher `nprobe` → better recall, slower query. Default: `nprobe = max(1, nlist // 10)`. Can be tuned by the caller. - -#### Backward compatibility - -The existing `similarities()` method signature must not change. FAISS is added as an optional code path activated by passing `use_faiss=True`. The `.sim_dict` output format stays identical so `measure_similarity_v2.py` and `evaluate_methods_v2.py` require no changes. - -### 1.3 New Dependency - -```toml -# pyproject.toml — add to [project] dependencies -"faiss-cpu>=1.8; extra != 'cuda12'", -"faiss-gpu>=1.8; extra == 'cuda12'", -``` - -Or install manually: -```bash -# CPU -uv add faiss-cpu - -# GPU (CUDA 12) -uv add faiss-gpu -``` - -### 1.4 Implementation - -#### 1.4.1 New method: `Img2Vec.build_faiss_index()` - -Add to `proteogram/v2/image_similarity.py` inside the `Img2Vec` class: - -```python -def build_faiss_index(self, - use_pq: bool = False, - nlist: int = None, - nprobe: int = None, - pq_m: int = 8, - pq_nbits: int = 8) -> None: - """Build a FAISS index from the currently loaded embedding dataset. - - Embeddings are L2-normalised before indexing so that inner-product - search is equivalent to cosine similarity. - - Args: - use_pq: If True, use IVF-PQ (compressed) index. Recommended for - corpora > 100 K. Defaults to False (IVFFlat, exact). - nlist: Number of Voronoi cells. Defaults to sqrt(N). - nprobe: Number of cells to search at query time. Higher = better - recall, slower query. Defaults to nlist // 10. - pq_m: Number of PQ sub-quantisers (IVF-PQ only). Must divide - the embedding dimension evenly. - pq_nbits: Bits per sub-quantiser (IVF-PQ only). 8 is standard. - """ - try: - import faiss - except ImportError: - raise ImportError( - "faiss is required for build_faiss_index(). " - "Install with: uv add faiss-cpu (or faiss-gpu for GPU builds)." - ) - - if not self.dataset: - raise RuntimeError("embed_dataset() must be called before build_faiss_index().") - - # Stack embeddings and keys in a consistent order - keys = list(self.dataset.keys()) - vecs = torch.cat([self.dataset[k].cpu() for k in keys]).float() # (N, d) - - # L2-normalise so inner product == cosine similarity - faiss.normalize_L2(vecs.numpy()) - - N, d = vecs.shape - _nlist = nlist if nlist is not None else max(1, int(N ** 0.5)) - _nprobe = nprobe if nprobe is not None else max(1, _nlist // 10) - - # Build quantiser (flat inner-product) - quantiser = faiss.IndexFlatIP(d) - - if use_pq: - # Ensure pq_m divides d evenly - while d % pq_m != 0 and pq_m > 1: - pq_m -= 1 - index = faiss.IndexIVFPQ(quantiser, d, _nlist, pq_m, pq_nbits, - faiss.METRIC_INNER_PRODUCT) - else: - index = faiss.IndexIVFFlat(quantiser, d, _nlist, - faiss.METRIC_INNER_PRODUCT) - - index.train(vecs.numpy()) - index.add(vecs.numpy()) - index.nprobe = _nprobe - - # Store on instance for re-use across queries - self._faiss_index = index - self._faiss_keys = keys # maps integer index → filename key - self._faiss_vecs_norm = vecs # keep L2-normalised vecs for query normalisation - - print(f"FAISS index built: {index.ntotal} vectors | d={d} | " - f"nlist={_nlist} | nprobe={_nprobe} | " - f"type={'IVF-PQ' if use_pq else 'IVFFlat'}") -``` - -#### 1.4.2 New method: `Img2Vec.similarities_faiss()` - -Add directly below `build_faiss_index()`: - -```python -def similarities_faiss(self, - n: int = 10, - save_result_images_dir: str = None, - pad_fn=None) -> float: - """Compute top-N similar images for every entry in the corpus using FAISS. - - Populates self.sim_dict with the same format as similarities(), so all - downstream scripts (evaluate_methods_v2.py, measure_similarity_v2.py) - work without modification. - - Call build_faiss_index() first. - - Args: - n: Top-N results per query (self-hit included - at rank 0; callers should request n+1 and - strip the self-hit themselves if needed). - save_result_images_dir: Optional directory to write result images. - pad_fn: Optional padding callable passed to save_images(). - - Returns: - float: Wall-clock seconds spent in FAISS search (excludes image saving). - """ - try: - import faiss - except ImportError: - raise ImportError("faiss not installed. Run: uv add faiss-cpu") - - if not hasattr(self, '_faiss_index'): - raise RuntimeError("Call build_faiss_index() before similarities_faiss().") - - keys = self._faiss_keys - vecs = self._faiss_vecs_norm.numpy() # already L2-normalised - - start = time() - # Batch query: search all N vectors at once — single FAISS call - scores_matrix, indices_matrix = self._faiss_index.search(vecs, n + 1) - elapsed = time() - start - - # Build sim_dict in the same format as similarities() - self.sim_dict = {} - for i, key in enumerate(keys): - hits = [] - for rank in range(n + 1): - j = indices_matrix[i, rank] - if j < 0: # FAISS pads with -1 when fewer results exist - continue - target_key = keys[j] - score = float(scores_matrix[i, rank]) - hits.append((target_key, score)) - self.sim_dict[key] = hits # includes self-hit at rank 0 - - if save_result_images_dir: - for image_path in self.sim_dict: - self.save_images(os.path.join(self.files[0].rsplit('/', 1)[0], image_path), - save_result_images_dir, pad_fn=pad_fn) - - return elapsed -``` - -#### 1.4.3 FAISS index persistence - -Add two methods for saving and loading the built index: - -```python -def save_faiss_index(self, index_path: str) -> None: - """Persist the FAISS index and key mapping to disk. - - Args: - index_path: File path for the index (e.g. 'corpus.faiss'). - A companion '.keys.pkl' file is written - alongside for the key mapping. - """ - import faiss, pickle - faiss.write_index(self._faiss_index, index_path) - keys_path = index_path + '.keys.pkl' - with open(keys_path, 'wb') as f: - pickle.dump(self._faiss_keys, f) - print(f"Saved FAISS index → {index_path}") - print(f"Saved key mapping → {keys_path}") - - -def load_faiss_index(self, index_path: str) -> None: - """Load a previously saved FAISS index and key mapping. - - Args: - index_path: Path to the '.faiss' index file. - """ - import faiss, pickle - self._faiss_index = faiss.read_index(index_path) - keys_path = index_path + '.keys.pkl' - with open(keys_path, 'rb') as f: - self._faiss_keys = pickle.load(f) - # Reconstruct normalised vecs for future queries (needed for single-query search) - keys = self._faiss_keys - vecs = torch.cat([self.dataset[k].cpu() for k in keys]).float() - faiss.normalize_L2(vecs.numpy()) - self._faiss_vecs_norm = vecs - print(f"Loaded FAISS index from {index_path} " - f"({self._faiss_index.ntotal} vectors)") -``` - -#### 1.4.4 Update `measure_similarity_v2.py` - -Add `--faiss` flag and wire it up: - -```python -# Add to the argparse block -parser.add_argument('--faiss', action='store_true', - help='Use FAISS ANN index for similarity search instead of ' - 'brute-force cosine similarity. Much faster for large corpora.') -parser.add_argument('--faiss_index_file', type=str, default=None, - help='Path to save/load the FAISS index. Defaults to ' - 'embed_file with .faiss extension.') -parser.add_argument('--faiss_pq', action='store_true', - help='Use IVF-PQ compressed index (recommended for > 100K proteins). ' - 'Slightly lower recall but much lower memory.') - -# Replace the similarities() call block with: -if args.faiss: - faiss_index_file = args.faiss_index_file or embed_file.replace('.pkl', '.faiss') - if os.path.exists(faiss_index_file) and not args.overwrite: - print(f'Loading existing FAISS index from {faiss_index_file}') - img_sim.load_faiss_index(faiss_index_file) - else: - print('Building FAISS index ...') - img_sim.build_faiss_index(use_pq=args.faiss_pq) - img_sim.save_faiss_index(faiss_index_file) - sim_time = img_sim.similarities_faiss( - n=n_results, - save_result_images_dir=None, - pad_fn=pad_to_size) -else: - sim_time = img_sim.similarities(n=n_results, - save_result_images_dir=None, - pad_fn=pad_to_size) -``` - -#### 1.4.5 Single-protein query update in `query_similar_proteins.py` - -```python -# Replace the inner loop in query_similar_proteins.py with: -def query_with_faiss(img_sim, query_embedding, top_k, corpus_dir): - """Query a built FAISS index with a single new embedding.""" - import faiss - import numpy as np - - query_vec = query_embedding.cpu().float().numpy() # (1, d) - faiss.normalize_L2(query_vec) - scores, indices = img_sim._faiss_index.search(query_vec, top_k + 1) - - results = [] - for rank in range(top_k + 1): - j = indices[0, rank] - if j < 0: - continue - key = img_sim._faiss_keys[j] - score = float(scores[0, rank]) - if key != os.path.basename(query_path): # skip self-hit if present - results.append((key, score)) - if len(results) >= top_k: - break - return results -``` - -### 1.5 Validation Steps - -#### Step 1 — Recall parity test (automated) - -Run both methods on the eval set and assert that FAISS Recall@K ≥ 0.99 × brute-force Recall@K at every SCOPe level: - -```python -# scripts/v2/tests/test_faiss_recall.py -import pytest -from proteogram.v2 import Img2Vec -import torch, pickle, os - -EMBED_FILE = os.environ.get('EMBED_FILE', 'corpus_embeddings.pkl') - -@pytest.fixture(scope='module') -def img_sim(): - sim = Img2Vec('resnet_ft', dataset_dir=[], device='cpu') - with open(EMBED_FILE, 'rb') as f: - sim.dataset = pickle.load(f) - return sim - -def test_faiss_topk_recall_at_5(img_sim): - """FAISS top-5 results should overlap ≥99% with brute-force top-5.""" - TOP_K = 5 - # Brute-force - img_sim.similarities(n=TOP_K) - bf_dict = {k: set(t for t, _ in v[:TOP_K]) for k, v in img_sim.sim_dict.items()} - - # FAISS IVFFlat - img_sim.build_faiss_index(use_pq=False) - img_sim.similarities_faiss(n=TOP_K) - faiss_dict = {k: set(t for t, _ in v[1:TOP_K+1]) for k, v in img_sim.sim_dict.items()} - - overlaps = [] - for key in bf_dict: - if key in faiss_dict: - overlap = len(bf_dict[key] & faiss_dict[key]) / TOP_K - overlaps.append(overlap) - - mean_recall = sum(overlaps) / len(overlaps) - print(f'Mean FAISS/BF overlap at top-{TOP_K}: {mean_recall:.4f}') - assert mean_recall >= 0.99, f'FAISS recall too low: {mean_recall:.4f}' -``` - -Run: -```bash -EMBED_FILE=/path/to/corpus_embeddings.pkl pytest scripts/v2/tests/test_faiss_recall.py -v -``` - -#### Step 2 — Timing benchmark - -```bash -# Brute-force -time python measure_similarity_v2.py --no-embed - -# FAISS IVFFlat -time python measure_similarity_v2.py --no-embed --faiss - -# FAISS IVF-PQ (large corpus) -time python measure_similarity_v2.py --no-embed --faiss --faiss_pq -``` - -Expected results on ~10 K eval set: - -| Method | Expected time | -|--------|--------------| -| Brute-force | ~2–5 min | -| FAISS IVFFlat | < 5 sec | -| FAISS IVF-PQ | < 2 sec | - -#### Step 3 — MAP@K parity - -Run `evaluate_methods_v2.py` on outputs from both methods. FAISS MAP@K should be within ±0.005 of brute-force MAP@K at all SCOPe levels. Any larger gap indicates the `nprobe` needs increasing. - -#### Step 4 — Index round-trip test - -```python -# Verify save/load produces identical results -img_sim.build_faiss_index() -img_sim.save_faiss_index('/tmp/test.faiss') - -img_sim2 = Img2Vec(model_file, dataset_dir=[], device='cpu') -img_sim2.dataset = img_sim.dataset -img_sim2.load_faiss_index('/tmp/test.faiss') - -img_sim.similarities_faiss(n=5) -img_sim2.similarities_faiss(n=5) - -for key in img_sim.sim_dict: - assert img_sim.sim_dict[key] == img_sim2.sim_dict[key], f"Mismatch at {key}" -print("Round-trip test passed.") -``` - ---- - -## 2. Global Percentile Normalisation - -> **Update:** the code snippets below reflect the original design. The shipped -> `--save_npy_matrices` implementation initially had a bug that fed -> already-normalised pixel data (not raw physical-unit energies) into -> `compute_norm_stats.py`, defeating the purpose of this feature. See -> [`percentile_normalisation_bug_fix_and_validation.md`](percentile_normalisation_bug_fix_and_validation.md) -> for the bug, the fix, and the new `validate_normalisation.py` tool for -> measuring before/after impact. - -### 2.1 Motivation - -The current `ProteogramV2.normalize_map()` applies **per-protein min-max normalisation** independently to each energy channel: - -```python -# proteogram/v2/proteogram.py — current implementation -arr = ((arr - arr.min()) * (1 / (arr.max() - arr.min()) * 255)).astype('uint8') -``` - -This has a critical flaw: every protein's energy map is stretched to fill the full [0, 255] dynamic range, regardless of the actual energy magnitudes. A small, weakly-interacting loop region and a tightly-packed hydrophobic core will produce identical grey levels after normalisation. The model never sees the absolute energy scale — only the relative rank order within each protein. - -**Concrete example**: Protein A has VdW attractive energies in [-50, -5] kJ/mol and Protein B has VdW attractive energies in [-200, -20] kJ/mol. After per-protein normalisation, both are mapped to [0, 255]. A CNN comparing the two images cannot tell that Protein B has 4× stronger packing. - -Global percentile normalisation computes bounds from the entire training corpus once, then applies those fixed bounds to every protein — preserving inter-protein energy scale in the pixel values. - -### 2.2 Design Decisions - -#### Percentile instead of global min/max - -Extreme outlier structures (e.g., very short peptides, structures with unusual post-translational modifications) would dominate a global min/max and compress most proteins into a narrow band. Using the **1st and 99th percentiles** clips ~2% of values but gives a robust, representative range. - -#### Separate bounds per channel - -Each of the 6 channels (VdW attractive, VdW repulsive, ES attractive, ES repulsive, distance, hydrophobicity) has a different physical unit and magnitude range. Bounds must be computed and stored independently per channel. - -#### Where bounds are stored - -A single JSON file `norm_stats.json` is written alongside the proteogram corpus. It is read at proteogram creation time when global normalisation is enabled. This keeps the bounds portable and version-controlled. - -#### Backward compatibility flag - -The new behaviour is opt-in via `--global_norm` flag in `create_v2_proteograms.py`. Per-protein normalisation remains the default so existing proteogram datasets are unaffected. - -### 2.3 New File: `scripts/v2/compute_norm_stats.py` - -This one-time script samples up to `--max_samples` existing `.npy` energy matrices (or re-runs the MD pipeline for a random subset) to compute global percentile bounds. - -```python -#!/usr/bin/env python -"""Compute global percentile normalisation statistics from a corpus of energy matrices. - -Run this ONCE after generating a representative sample of proteograms with ---save_npy_matrices (a new flag added to create_v2_proteograms.py). Outputs -norm_stats.json which is read by create_v2_proteograms.py --global_norm. - -Usage: - python compute_norm_stats.py \\ - --npy_dir /path/to/energy_matrices \\ - --out_file /path/to/norm_stats.json \\ - --low_pct 1.0 \\ - --high_pct 99.0 \\ - --max_samples 5000 -""" -import argparse -import json -import glob -import os -import numpy as np -from tqdm import tqdm - -CHANNEL_NAMES = [ - 'vdw_attractive', - 'vdw_repulsive', - 'es_attractive', - 'es_repulsive', - 'distance', - 'hydrophobicity', -] - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--npy_dir', required=True, - help='Directory of .npy files, one per channel per protein ' - '(naming: _.npy).') - parser.add_argument('--out_file', required=True, - help='Output JSON path for normalisation bounds.') - parser.add_argument('--low_pct', type=float, default=1.0, - help='Lower percentile bound (default: 1.0).') - parser.add_argument('--high_pct', type=float, default=99.0, - help='Upper percentile bound (default: 99.0).') - parser.add_argument('--max_samples', type=int, default=5000, - help='Maximum number of energy matrices to sample per channel ' - '(default: 5000). More samples = more accurate statistics.') - args = parser.parse_args() - - # Collect all values per channel across the sampled corpus - channel_values = {ch: [] for ch in CHANNEL_NAMES} - - for channel in CHANNEL_NAMES: - files = sorted(glob.glob(os.path.join(args.npy_dir, f'*_{channel}.npy'))) - if not files: - print(f'WARNING: No .npy files found for channel "{channel}" in {args.npy_dir}') - continue - - # Random subsample if corpus is large - if len(files) > args.max_samples: - rng = np.random.default_rng(seed=42) - files = list(rng.choice(files, size=args.max_samples, replace=False)) - - print(f'Channel {channel}: sampling {len(files)} matrices ...') - for fpath in tqdm(files, desc=channel): - arr = np.load(fpath) - # Only include non-zero upper-triangle values (lower triangle is 0) - vals = arr[arr != 0].ravel() - channel_values[channel].append(vals) - - # Compute and store bounds - stats = {} - for channel in CHANNEL_NAMES: - if not channel_values[channel]: - stats[channel] = {'p_low': 0.0, 'p_high': 255.0} - continue - all_vals = np.concatenate(channel_values[channel]) - p_low = float(np.percentile(all_vals, args.low_pct)) - p_high = float(np.percentile(all_vals, args.high_pct)) - stats[channel] = {'p_low': p_low, 'p_high': p_high} - print(f' {channel}: p{args.low_pct}={p_low:.4f} p{args.high_pct}={p_high:.4f} ' - f'N={len(all_vals):,}') - - stats['_meta'] = { - 'low_pct': args.low_pct, - 'high_pct': args.high_pct, - 'n_files_per_channel': args.max_samples, - 'npy_dir': args.npy_dir, - } - - os.makedirs(os.path.dirname(os.path.abspath(args.out_file)), exist_ok=True) - with open(args.out_file, 'w') as f: - json.dump(stats, f, indent=2) - print(f'\nSaved normalisation stats → {args.out_file}') - - -if __name__ == '__main__': - main() -``` - -### 2.4 Implementation: Changes to `proteogram/v2/proteogram.py` - -#### 2.4.1 New static method: `normalize_map_global()` - -Add alongside the existing `normalize_map()`: - -```python -@staticmethod -def normalize_map_global(arr: np.ndarray, - p_low: float, - p_high: float) -> tuple[np.ndarray, str]: - """Normalise an energy/property map to [0, 255] using corpus-level percentile bounds. - - Unlike normalize_map(), which uses per-protein min/max, this method - applies fixed bounds derived from the full training corpus so that - inter-protein energy scale is preserved in pixel values. - - Zero values (unfilled lower-triangle entries) are mapped to 128 (mid-grey) - to distinguish them visually from true low-energy interactions (which map - near 0) — matching the existing gray padding convention in the training code. - - Args: - arr: Input energy matrix (upper triangle populated; lower = 0). - p_low: Lower percentile bound in physical units (kJ/mol or Å). - p_high: Upper percentile bound in physical units. - - Returns: - Tuple of (normalised uint8 array, error string or ''). - """ - err = '' - try: - scale = p_high - p_low - if scale == 0: - return np.full_like(arr, 128, dtype='uint8'), 'zero scale range' - - # Clip to [p_low, p_high] then scale to [0, 255] - clipped = np.clip(arr, p_low, p_high) - normalised = ((clipped - p_low) / scale * 255).astype('uint8') - - # Remap structural zeros (unfilled lower triangle) to mid-grey (128) - # so they don't contaminate the 0-end of the energy scale - normalised[arr == 0] = 128 - - except Exception as e: - err = f'Problem in normalize_map_global: {e}' - normalised = np.full_like(arr, 128, dtype='uint8') - return normalised, err -``` - -#### 2.4.2 Update `calculate_proteogram()` to accept `norm_stats` - -Modify the method signature and normalisation block: - -```python -def calculate_proteogram(self, - return_simulated_pdb: bool = False, - debug: bool = False, - subtract_solvent_energies: bool = True, - memory_efficient: bool = False, - norm_stats: dict = None): # <-- NEW parameter - """ - ... (existing docstring) ... - - Args: - ... - norm_stats: Optional dict loaded from norm_stats.json. When supplied, - normalize_map_global() is used for all 6 channels instead - of per-protein min-max. Keys: 'vdw_attractive', 'vdw_repulsive', - 'es_attractive', 'es_repulsive', 'distance', 'hydrophobicity'. - Each value is a dict with 'p_low' and 'p_high'. - """ - # ... existing MD pipeline code unchanged ... - - # ---- Replace the normalisation block ---- - def _norm(arr, channel_name): - if norm_stats and channel_name in norm_stats: - s = norm_stats[channel_name] - return self.normalize_map_global(arr, s['p_low'], s['p_high']) - return self.normalize_map(arr) - - norm_disto_map, disto_err = _norm(disto_map, 'distance') - norm_hydro_map, hydro_err = _norm(hydro_map, 'hydrophobicity') - norm_vdw_att_map, vdw_att_err = _norm(vdw_e_att, 'vdw_attractive') - norm_vdw_rep_map, vdw_rep_err = _norm(vdw_e_rep, 'vdw_repulsive') - norm_es_att_map, es_att_err = _norm(es_e_att, 'es_attractive') - norm_es_rep_map, es_rep_err = _norm(es_e_rep, 'es_repulsive') - # ... rest of stacking unchanged ... -``` - -#### 2.4.3 Update `create_v2_proteograms.py` - -```python -# Add to argparse -parser.add_argument('--global_norm', action='store_true', - help='Use global percentile normalisation bounds from norm_stats.json ' - 'instead of per-protein min-max. Requires --norm_stats_file.') -parser.add_argument('--norm_stats_file', type=str, default=None, - help='Path to norm_stats.json produced by compute_norm_stats.py.') -parser.add_argument('--save_npy_matrices', action='store_true', - help='Save raw energy matrices as .npy files alongside proteogram JPGs. ' - 'Required input for compute_norm_stats.py.') - -# Load norm_stats once before the proteogram creation loop -norm_stats = None -if args.global_norm: - if not args.norm_stats_file or not os.path.exists(args.norm_stats_file): - raise ValueError('--global_norm requires --norm_stats_file pointing to norm_stats.json') - import json - with open(args.norm_stats_file) as f: - norm_stats = json.load(f) - print(f'Loaded global norm stats from {args.norm_stats_file}') - -# Pass norm_stats into the ProteogramV2 call inside the creation loop -proteogram_data, errors = prot.calculate_proteogram( - subtract_solvent_energies=True, - memory_efficient=args.memory_efficient, - norm_stats=norm_stats, # <-- new -) -``` - -### 2.5 End-to-End Workflow - -```bash -# Step 1: Generate proteograms with raw .npy matrix saving (first pass or subset) -python create_v2_proteograms.py --save_npy_matrices - -# Step 2: Compute global bounds from saved matrices -python compute_norm_stats.py \ - --npy_dir /path/to/proteograms/energy_matrices \ - --out_file /path/to/norm_stats.json \ - --max_samples 5000 - -# Step 3: Re-generate proteograms using global normalisation -python create_v2_proteograms.py \ - --global_norm \ - --norm_stats_file /path/to/norm_stats.json \ - --overwrite -``` - -### 2.6 Validation Steps - -> Steps 1 and 4 below (pixel-distribution and clipping-rate checks) are now -> implemented as a single runnable tool, `scripts/v2/validate_normalisation.py`, -> which also adds an inter-protein variance ratio and a correlation-with-raw-scale -> check that these steps didn't originally include. See -> [`percentile_normalisation_bug_fix_and_validation.md`](percentile_normalisation_bug_fix_and_validation.md) -> for exact commands and how to read the output. Steps 2 and 3 (visual inspection, -> downstream MAP@K) remain manual/expensive as described below. - -#### Step 1 — Sanity check: pixel distribution - -For a random sample of 100 proteograms, compare the pixel value histograms between per-protein and global normalisation: - -```python -import numpy as np -import matplotlib.pyplot as plt -from PIL import Image -import glob - -per_protein_files = glob.glob('/path/to/proteograms_per_protein/*.jpg')[:100] -global_files = glob.glob('/path/to/proteograms_global/*.jpg')[:100] - -for label, files in [('per-protein', per_protein_files), ('global', global_files)]: - pixels = np.concatenate([np.array(Image.open(f)).ravel() for f in files]) - plt.hist(pixels, bins=50, alpha=0.6, label=label) - -plt.legend() -plt.xlabel('Pixel value') -plt.title('Pixel distribution: per-protein vs. global normalisation') -plt.savefig('norm_comparison.png', dpi=150) -``` - -Expected result: global normalisation produces a wider, less clipped distribution with meaningful variation near 0 and 255. Per-protein should look nearly uniform (every image uses the full range). - -#### Step 2 — Visual inspection - -Side-by-side comparison of the same protein normalised both ways: - -```python -from PIL import Image -import matplotlib.pyplot as plt - -pdb_id = 'd3kfda_' -per_protein = Image.open(f'/path/per_protein/{pdb_id}.jpg') -global_norm = Image.open(f'/path/global/{pdb_id}.jpg') - -fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) -ax1.imshow(per_protein); ax1.set_title('Per-protein normalisation'); ax1.axis('off') -ax2.imshow(global_norm); ax2.set_title('Global normalisation'); ax2.axis('off') -plt.savefig(f'{pdb_id}_norm_comparison.png', dpi=150) -``` - -Proteins with strong hydrophobic cores (e.g., globins, immunoglobulins) should appear noticeably brighter in the VdW channels under global normalisation compared to proteins with weak packing. - -#### Step 3 — Downstream MAP@K comparison - -Retrain the ResNet18 model on global-normalised proteograms and compare MAP@K on the eval set against the baseline model trained on per-protein normalised proteograms. Expected outcome: global normalisation improves MAP@K at the fold and superfamily levels (where energy magnitude differences are most discriminative), with neutral or marginal effect at the class level. - -#### Step 4 — Robustness check: unseen proteins - -Compute the fraction of pixel values clipped to 0 or 255 for 50 randomly selected held-out proteins (not used in `compute_norm_stats.py`). If > 5% of non-zero pixels are clipped, the percentile bounds are too tight and `--low_pct`/`--high_pct` should be widened (e.g., 0.5 and 99.5). - ---- - -## 3. Grad-CAM Explainability - -### 3.1 Motivation - -When Proteogram reports that two proteins share 87% cosine similarity, a structural biologist naturally asks: *which residue-residue interactions drove that score?* Currently there is no answer — the model is a black box. - -Because proteogram pixels directly encode pairwise residue interactions (pixel at row `i`, column `j` represents the interaction between residue `i` and residue `j`), a saliency heatmap over the input image is directly interpretable as a **residue-pair importance map**. This is a unique property of the proteogram representation that does not exist for most CV tasks. - -Grad-CAM (Gradient-weighted Class Activation Mapping) computes a heatmap by backpropagating the gradient of a target score through the last convolutional layer of the CNN. High activation regions in the heatmap indicate which spatial features (and thus which residue pairs) most influenced the model's output. - -### 3.2 Design Decisions - -#### Target layer selection - -For ResNet18, the natural target is the output of `layer4` (the last residual block), which has spatial resolution 7×7 for 224px input or 13×13 for 200px padded proteograms. This gives meaningful spatial resolution after upsampling back to the full NxN image. - -For the custom ConvNet, the target is `block4` (after the 4th MaxPool, spatial resolution ≈ 12×12 for 200px input). - -#### Score to differentiate - -Standard Grad-CAM differentiates with respect to the **class logit** for a classification task. For a *retrieval* task we instead differentiate with respect to the **cosine similarity score** between a query and a target embedding. This gives a "similarity-attribution" heatmap: *which parts of the query proteogram, when activated, push the cosine similarity with the target higher?* - -Formally, if `f_q` and `f_t` are the embedding vectors for query and target: - -``` -S = cos(f_q, f_t) = (f_q · f_t) / (||f_q|| · ||f_t||) -``` - -We compute `∂S / ∂A_k` for each activation map `A_k` in the target convolutional layer. - -#### Output format - -The Grad-CAM heatmap is: -- An NxN float32 array in [0, 1] — matching the proteogram dimensions -- Saved as both a matplotlib figure (with residue axis labels) and a raw `.npy` file -- Overlaid as a semi-transparent colour map on top of the original proteogram image - -### 3.3 Implementation - -#### 3.3.1 New method: `Img2Vec.gradcam_similarity()` - -Add to `proteogram/v2/image_similarity.py`: - -```python -def gradcam_similarity(self, - query_image_path: str, - target_image_path: str, - output_dir: str, - query_sequence: str = None, - target_sequence: str = None) -> np.ndarray: - """Compute a Grad-CAM residue-pair importance map for a query→target similarity. - - The heatmap shows which residue-pair interactions in the QUERY proteogram - most influence the cosine similarity with the TARGET proteogram. - - The model must be a ResNet18 fine-tuned with train_multiple_models.py - (--model resnet18) or the from-scratch ConvNet (--model cnn). - - Args: - query_image_path: Path to the query proteogram JPG. - target_image_path: Path to the target proteogram JPG. - output_dir: Directory to save the heatmap figure and .npy file. - query_sequence: Optional 1-letter amino acid sequence for axis labels. - target_sequence: Optional 1-letter amino acid sequence (unused currently, - reserved for cross-proteogram attribution in future). - - Returns: - np.ndarray: Upsampled Grad-CAM heatmap, shape (H, W), values in [0, 1]. - """ - import torch.nn.functional as F - - os.makedirs(output_dir, exist_ok=True) - - # ------------------------------------------------------------------ # - # 1. Identify the target convolutional layer # - # ------------------------------------------------------------------ # - target_layer = self._get_gradcam_target_layer() - - # ------------------------------------------------------------------ # - # 2. Register forward/backward hooks # - # ------------------------------------------------------------------ # - activations = {} - gradients = {} - - def _save_activation(module, input, output): - activations['value'] = output.detach() - - def _save_gradient(module, grad_input, grad_output): - gradients['value'] = grad_output[0].detach() - - fwd_hook = target_layer.register_forward_hook(_save_activation) - bwd_hook = target_layer.register_full_backward_hook(_save_gradient) - - try: - # ------------------------------------------------------------------ # - # 3. Forward pass for both query and target # - # ------------------------------------------------------------------ # - query_tensor = self._load_and_preprocess(query_image_path) # (1, 3, H, W) - target_tensor = self._load_and_preprocess(target_image_path) # (1, 3, H, W) - - # Embeddings from the penultimate layer - # Switch to full model (not self.embed which stripped the head) - self.model.eval() - query_tensor = query_tensor.to(self.device).requires_grad_(True) - target_tensor = target_tensor.to(self.device) - - # Get embedding for query (triggers forward hook and saves activations) - query_feat = self.embed(query_tensor) # (1, d) - target_feat = self.embed(target_tensor).detach() # (1, d) - - # ------------------------------------------------------------------ # - # 4. Compute cosine similarity and differentiate # - # ------------------------------------------------------------------ # - # Manually compute cosine similarity (not through nn.CosineSimilarity - # so we can call backward on the scalar) - q_norm = F.normalize(query_feat, dim=1) - t_norm = F.normalize(target_feat, dim=1) - cos_sim = (q_norm * t_norm).sum() # scalar - - self.model.zero_grad() - cos_sim.backward() - - # ------------------------------------------------------------------ # - # 5. Compute Grad-CAM weights # - # ------------------------------------------------------------------ # - grads = gradients['value'] # (1, C, h, w) - acts = activations['value'] # (1, C, h, w) - - # Global-average-pool the gradients over the spatial dims → (1, C, 1, 1) - weights = grads.mean(dim=(2, 3), keepdim=True) - - # Weighted combination of activation maps → (1, 1, h, w) - cam = (weights * acts).sum(dim=1, keepdim=True) - cam = F.relu(cam) # keep only positive contributions - - # Normalise to [0, 1] - cam_min, cam_max = cam.min(), cam.max() - if cam_max > cam_min: - cam = (cam - cam_min) / (cam_max - cam_min) - - # ------------------------------------------------------------------ # - # 6. Upsample to input image size # - # ------------------------------------------------------------------ # - input_h = query_tensor.shape[2] - input_w = query_tensor.shape[3] - cam_upsampled = F.interpolate(cam, - size=(input_h, input_w), - mode='bilinear', - align_corners=False) - cam_np = cam_upsampled.squeeze().cpu().numpy() # (H, W) - - finally: - fwd_hook.remove() - bwd_hook.remove() - - # ------------------------------------------------------------------ # - # 7. Save outputs # - # ------------------------------------------------------------------ # - query_name = os.path.splitext(os.path.basename(query_image_path))[0] - target_name = os.path.splitext(os.path.basename(target_image_path))[0] - stem = f'{query_name}_vs_{target_name}' - - # Save raw heatmap - npy_path = os.path.join(output_dir, f'{stem}_gradcam.npy') - np.save(npy_path, cam_np) - - # Save overlay figure - query_img = np.array(Image.open(query_image_path).convert('RGB')) - self._save_gradcam_figure( - query_img=query_img, - cam=cam_np, - cos_sim=cos_sim.item(), - query_name=query_name, - target_name=target_name, - output_dir=output_dir, - query_sequence=query_sequence, - ) - - print(f'Grad-CAM saved → {output_dir}/{stem}_gradcam.png') - return cam_np - - -def _get_gradcam_target_layer(self): - """Return the last convolutional layer for Grad-CAM based on architecture.""" - children = list(self.model.children()) - # ResNet18: children order is conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4, avgpool, fc - # Find the last nn.Sequential that contains Conv2d layers - target = None - for child in children: - if isinstance(child, nn.Sequential): - for submodule in child.modules(): - if isinstance(submodule, nn.Conv2d): - target = child - if target is None: - # Fallback: use the last Conv2d found anywhere in the model - for module in self.model.modules(): - if isinstance(module, nn.Conv2d): - target = module - return target - - -def _load_and_preprocess(self, image_path: str) -> torch.Tensor: - """Load and preprocess a single proteogram image matching training transforms.""" - img = Image.open(image_path).convert('RGB') - # Apply the same pad-to-200 + ImageNet normalisation used in training - from torchvision import transforms as T - import numpy as np - arr = np.array(img) - H, W = arr.shape[:2] - target = 200 - - def get_pad(curr, tgt): - d = tgt - curr - if d <= 0: - return (0, 0) - p1 = d // 2 - return (p1, d - p1) - - padding = (get_pad(H, target), get_pad(W, target), (0, 0)) - arr = np.pad(arr, padding, constant_values=128) - arr = arr[:target, :target, :] - img_padded = Image.fromarray(arr.astype('uint8')) - - transform = T.Compose([ - T.ToTensor(), - T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ]) - return transform(img_padded).unsqueeze(0) # (1, 3, H, W) - - -def _save_gradcam_figure(self, - query_img: np.ndarray, - cam: np.ndarray, - cos_sim: float, - query_name: str, - target_name: str, - output_dir: str, - query_sequence: str = None) -> None: - """Save a 3-panel Grad-CAM figure: original, heatmap, overlay.""" - import matplotlib.pyplot as plt - import matplotlib.cm as cm - from matplotlib.colors import Normalize - - fig, axes = plt.subplots(1, 3, figsize=(18, 6)) - fig.suptitle( - f'Grad-CAM: {query_name} → {target_name} (cosine similarity = {cos_sim:.4f})', - fontsize=13, y=1.01 - ) - - # Panel 1: original query proteogram - axes[0].imshow(query_img) - axes[0].set_title('Query proteogram', fontsize=11) - axes[0].axis('off') - - # Panel 2: Grad-CAM heatmap alone - im = axes[1].imshow(cam, cmap='hot', vmin=0, vmax=1) - axes[1].set_title('Grad-CAM heatmap\n(high = important residue pairs)', fontsize=11) - axes[1].axis('off') - plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04) - - # Panel 3: overlay (proteogram + semi-transparent heatmap) - axes[2].imshow(query_img) - overlay = axes[2].imshow(cam, cmap='hot', alpha=0.55, vmin=0, vmax=1) - axes[2].set_title('Overlay', fontsize=11) - axes[2].axis('off') - plt.colorbar(overlay, ax=axes[2], fraction=0.046, pad=0.04) - - # Optional: add residue index tick labels if sequence is provided - if query_sequence and len(query_sequence) <= 200: - step = max(1, len(query_sequence) // 20) # show ~20 tick labels - ticks = list(range(0, len(query_sequence), step)) - tick_labels = [f'{i}\n{query_sequence[i]}' for i in ticks] - for ax in axes: - ax.set_xticks(ticks); ax.set_xticklabels(tick_labels, fontsize=6) - ax.set_yticks(ticks); ax.set_yticklabels(tick_labels, fontsize=6) - ax.tick_params(axis='both', length=2) - - plt.tight_layout() - stem = f'{query_name}_vs_{target_name}' - fig_path = os.path.join(output_dir, f'{stem}_gradcam.png') - plt.savefig(fig_path, dpi=150, bbox_inches='tight') - plt.close(fig) -``` - -#### 3.3.2 New script: `scripts/v2/explain_similarity.py` - -```python -#!/usr/bin/env python -"""Generate Grad-CAM residue-pair importance maps for similar protein pairs. - -For each query in the eval set, explains the similarity to its top-1 hit -(or a user-specified target) using Grad-CAM over the last convolutional layer. - -Usage — explain top-1 hit for all eval proteograms: - python explain_similarity.py - -Usage — explain a specific query→target pair: - python explain_similarity.py \\ - --query /path/to/d3kfda_.jpg \\ - --target /path/to/d1yl4r1.jpg \\ - --query_seq ACDEFGHIKLMNPQRSTVWY... \\ - --output_dir gradcam_results/ - -Usage — explain top-K hits for the 50 queries with lowest MAP@K (worst cases): - python explain_similarity.py --explain_worst 50 --top_k 3 -""" -import argparse -import os -import pickle -import torch -import pandas as pd - -from proteogram.v2 import Img2Vec -from proteogram.common import read_yaml - - -def main(): - parser = argparse.ArgumentParser(description='Grad-CAM explainability for Proteogram.') - parser.add_argument('--query', '-q', type=str, default=None, - help='Path to a single query proteogram JPG.') - parser.add_argument('--target', '-t', type=str, default=None, - help='Path to a single target proteogram JPG.') - parser.add_argument('--query_seq', type=str, default=None, - help='1-letter amino acid sequence of the query protein ' - '(optional, for residue axis labels).') - parser.add_argument('--output_dir', '-o', type=str, default='gradcam_output', - help='Directory to save Grad-CAM figures and .npy files.') - parser.add_argument('--explain_worst', type=int, default=None, - help='Explain the top-1 hit for the N queries with the lowest ' - 'MAP@K score (most informative failures). ' - 'Requires proteogram_sim_results in config.yml.') - parser.add_argument('--top_k', type=int, default=1, - help='Number of top hits to explain per query (default: 1).') - args = parser.parse_args() - - config = read_yaml('config.yml') - model_file = config['model_file'] - embed_file = config['embed_file'] - corpus_dir = config['proteograms_for_sim_dir'] - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - img_sim = Img2Vec(model_file, dataset_dir=[], device=device) - - # Load corpus embeddings - with open(embed_file, 'rb') as f: - img_sim.dataset = pickle.load(f) - - os.makedirs(args.output_dir, exist_ok=True) - - if args.query and args.target: - # Single pair mode - img_sim.gradcam_similarity( - query_image_path=args.query, - target_image_path=args.target, - output_dir=args.output_dir, - query_sequence=args.query_seq, - ) - - elif args.explain_worst: - # Explain worst-performing queries from similarity results - results_file = config.get('proteogram_sim_results') - if not results_file or not os.path.exists(results_file): - raise FileNotFoundError( - f'proteogram_sim_results not found: {results_file}. ' - 'Run measure_similarity_v2.py first.') - - results_df = pd.read_csv(results_file, sep='\t') - # Sort by MAP@K score (ascending = worst first) if the column exists, - # otherwise just take the last N rows as a proxy - n = args.explain_worst - queries_to_explain = results_df.head(n) - - for _, row in queries_to_explain.iterrows(): - query_path = row.iloc[0] - top_hits = [row.iloc[k+1].split(',')[0] # filename part of 'filename,score' - for k in range(min(args.top_k, len(row) - 1))] - for target_stem in top_hits: - target_path = os.path.join(corpus_dir, target_stem + '.jpg') - if not os.path.exists(target_path): - print(f'Target not found, skipping: {target_path}') - continue - img_sim.gradcam_similarity( - query_image_path=query_path, - target_image_path=target_path, - output_dir=args.output_dir, - ) - else: - parser.print_help() - - -if __name__ == '__main__': - main() -``` - -### 3.4 Validation Steps - -#### Step 1 — Sanity check: heatmap is not uniform - -For 10 random query-target pairs, verify the heatmap has meaningful spatial variation (std > 0.05): - -```python -import numpy as np -import glob - -npy_files = glob.glob('gradcam_output/*_gradcam.npy') -for f in npy_files: - cam = np.load(f) - std = cam.std() - max_val = cam.max() - print(f'{f}: std={std:.4f} max={max_val:.4f}') - assert std > 0.05, f'Heatmap appears uniform for {f} — check hook registration' -``` - -#### Step 2 — Biological sanity check - -Run Grad-CAM on a well-studied protein pair from the same SCOPe superfamily (e.g., two globins: haemoglobin α-chain vs. myoglobin). Expect high activation at: -- The haem-binding pocket region (residues ~60–90 and ~130–150 in the sequence) -- The conserved F-helix contacts -- The hydrophobic core residues - -Cross-reference high-activation residue pairs against the known structural alignment from US-align or GTalign. If the top-10 residue pairs by Grad-CAM score overlap significantly with the US-align-identified structurally equivalent residue pairs, the explainer is working correctly. - -#### Step 3 — Negative control - -Run Grad-CAM on a query-target pair from *different* SCOPe classes (e.g., an all-alpha vs. an all-beta protein with low cosine similarity score ~0.3). The heatmap should be diffuse and low-magnitude — no clear hotspot — because no specific structural motif is driving the (low) similarity. - -```python -# Confirm: mean activation for negative pairs should be < mean activation for positive pairs -import numpy as np - -positive_cams = [np.load(f) for f in glob.glob('gradcam_output/same_class_*.npy')] -negative_cams = [np.load(f) for f in glob.glob('gradcam_output/diff_class_*.npy')] - -pos_mean = np.mean([c.max() for c in positive_cams]) -neg_mean = np.mean([c.max() for c in negative_cams]) -print(f'Positive pairs max activation: {pos_mean:.4f}') -print(f'Negative pairs max activation: {neg_mean:.4f}') -assert pos_mean > neg_mean, 'Grad-CAM not discriminating positive/negative pairs' -``` - -#### Step 4 — Hook cleanup test - -Verify hooks are always removed even when an exception occurs mid-computation (hooks left dangling slow down subsequent forward passes and may accumulate memory): - -```python -# Deliberately pass a corrupted image path and confirm no hook leakage -from proteogram.v2 import Img2Vec -import torch - -img_sim = Img2Vec(model_file, dataset_dir=[], device='cpu') -img_sim.dataset = {} # minimal setup - -hook_count_before = len(list(img_sim.model._forward_hooks.values())) -try: - img_sim.gradcam_similarity('/nonexistent/query.jpg', '/nonexistent/target.jpg', '/tmp') -except Exception: - pass -hook_count_after = len(list(img_sim.model._forward_hooks.values())) -assert hook_count_before == hook_count_after, 'Forward hooks leaked after exception!' -print('Hook cleanup test passed.') -``` - ---- - -## 4. Combined Integration Checklist - -Before merging all three changes, run through this checklist end-to-end on the eval set: - -``` -[ ] pip install faiss-cpu (or faiss-gpu) added to pyproject.toml -[ ] compute_norm_stats.py runs without error on 500 random energy matrices -[ ] norm_stats.json is committed to the repository alongside the dataset -[ ] create_v2_proteograms.py --global_norm produces visually distinct images - vs. per-protein normalised equivalents (visual inspection on 5 proteins) -[ ] measure_similarity_v2.py --faiss produces sim_dict identical in format to - existing brute-force output (evaluate_methods_v2.py accepts it unchanged) -[ ] FAISS Recall@5 ≥ 0.99 × brute-force Recall@5 (automated test passes) -[ ] MAP@K from FAISS results within ±0.005 of MAP@K from brute-force results -[ ] explain_similarity.py runs on a single pair and produces a 3-panel PNG -[ ] Grad-CAM heatmap std > 0.05 for same-class pairs (sanity check passes) -[ ] No memory leaks: all forward/backward hooks removed after gradcam_similarity() -[ ] All existing tests in scripts/v2/tests/ still pass -[ ] README.md updated with: - - New --global_norm / --norm_stats_file flags in Step 1 - - New --faiss / --faiss_pq flags in Step 4 - - New explain_similarity.py in the scripts reference table -``` - ---- - -## 5. Config additions (`scripts/v2/config.example.yml`) - -Add these keys to the example config for discoverability: - -```yaml -# ── Global normalisation (Improvement 2) ────────────────────────────── -# Path to norm_stats.json produced by compute_norm_stats.py. -# Required when running create_v2_proteograms.py --global_norm. -norm_stats_file: /path/to/norm_stats.json - -# ── FAISS index (Improvement 1) ─────────────────────────────────────── -# Path to save/load the FAISS index (auto-derived from embed_file if omitted). -faiss_index_file: /path/to/corpus_embeddings.faiss - -# ── Grad-CAM output (Improvement 3) ─────────────────────────────────── -# Directory to write Grad-CAM figures and .npy heatmap files. -gradcam_output_dir: /path/to/gradcam_output -``` diff --git a/proteogram/v2/faiss_search.py b/proteogram/v2/faiss_search.py index 8d68052..dd36e63 100644 --- a/proteogram/v2/faiss_search.py +++ b/proteogram/v2/faiss_search.py @@ -1,108 +1,110 @@ -"""FAISS-based Approximate Nearest Neighbour index for proteogram embeddings. +"""FAISS approximate nearest neighbour search over Proteogram embeddings. -This module is fully self-contained and has no dependency on Img2Vec or any other -proteogram class. It operates on plain numpy float32 arrays and stores a key list -(filename strings) so integer FAISS indices can be mapped back to protein IDs. +Img2Vec.similarities() scores every query against every corpus vector, which +is O(N^2) in both time and memory and becomes the bottleneck well before the +embedding step does. This module wraps a FAISS IVF index over the same +embeddings so search cost scales with nprobe rather than corpus size. -Typical usage -------------- ->>> from proteogram.v2.faiss_search import FaissIndex ->>> import numpy as np +Nothing here imports Img2Vec. The index operates on plain float32 numpy +arrays plus an ordered key list, so the integer ids FAISS returns can be +mapped back to proteogram filenames. faiss itself is imported lazily inside +the methods that need it, so importing proteogram.v2 stays cheap. ->>> # Build from a dict of {filename: embedding_tensor} (same format as Img2Vec.dataset) ->>> index = FaissIndex.from_dataset(img2vec.dataset) +Embeddings are L2-normalised before indexing, so the inner-product metric +FAISS searches with is equivalent to the cosine similarity Img2Vec reports. ->>> # All-vs-all search (returns same format as Img2Vec.sim_dict) ->>> sim_dict = index.search_all(top_k=5) +Example: +----------- ->>> # Single-query search ->>> hits = index.search_one(query_vec, top_k=5) + from proteogram.v2.faiss_search import FaissIndex ->>> # Persistence ->>> index.save("/path/to/corpus.faiss") ->>> index2 = FaissIndex.load("/path/to/corpus.faiss") -""" + index = FaissIndex.from_dataset(img_sim.dataset) + sim_dict = index.search_all(top_k=5) # same shape as Img2Vec.sim_dict + hits = index.search_one(query_vec, top_k=5) + index.save('corpus.faiss') + index = FaissIndex.load('corpus.faiss') +""" from __future__ import annotations import os import pickle +from contextlib import contextmanager from typing import Dict, List, Tuple import numpy as np import torch -# --------------------------------------------------------------------------- -# Public constants -# --------------------------------------------------------------------------- +_FAISS_INSTALL_HINT = ( + "faiss is required for ANN search. Install with 'uv add faiss-cpu', or " + "'uv add faiss-gpu' for a CUDA build." +) -#: Default lower percentile for nlist auto-selection. -_NLIST_SQRT_FACTOR: float = 1.0 +def _import_faiss(): + """Import faiss, re-raising with an install hint if it is missing.""" + try: + import faiss + except ImportError as exc: + raise ImportError(_FAISS_INSTALL_HINT) from exc + return faiss -# --------------------------------------------------------------------------- -# Helper -# --------------------------------------------------------------------------- def _l2_normalise(mat: np.ndarray) -> np.ndarray: - """Return an L2-normalised copy of *mat* (shape N×d, float32). + """Return an L2-normalised float32 copy of mat (shape N x d). - Does NOT modify the input array in-place so the caller's embeddings stay intact. + Normalising a copy rather than in place matters because the caller's + embeddings are the same arrays Img2Vec keeps in self.dataset. """ - mat = mat.copy().astype(np.float32) + mat = mat.astype(np.float32, copy=True) norms = np.linalg.norm(mat, axis=1, keepdims=True) - norms = np.where(norms == 0, 1.0, norms) # avoid div-by-zero for zero vectors + # A zero vector has no direction; leave it as-is instead of dividing by 0. + norms[norms == 0] = 1.0 mat /= norms return mat def _stack_dataset(dataset: Dict[str, torch.Tensor]) -> Tuple[List[str], np.ndarray]: - """Convert an Img2Vec-style dataset dict to an ordered (keys, matrix) pair. + """Flatten an Img2Vec-style dataset dict into an ordered (keys, matrix) pair. Args: - dataset: Mapping of filename → 1-D or 1×d embedding tensor. + dataset: mapping of filename to a 1-D or 1 x d embedding tensor. Returns: - keys: List of filenames in the same row order as the matrix. - matrix: float32 numpy array of shape (N, d). + keys: filenames in the same row order as the matrix. + matrix: float32 array of shape (N, d). """ keys = list(dataset.keys()) vecs = torch.cat([dataset[k].cpu().reshape(1, -1) for k in keys]).float().numpy() return keys, vecs -# --------------------------------------------------------------------------- -# FaissIndex -# --------------------------------------------------------------------------- - class FaissIndex: - """Wraps a FAISS IVFFlat or IVF-PQ index with a key mapping. - - Parameters - ---------- - keys: - Ordered list of protein filenames. ``keys[i]`` is the protein - corresponding to FAISS integer index ``i``. - vecs_norm: - L2-normalised embedding matrix, shape (N, d), float32. Stored so - single-query searches can normalise the query the same way. - index: - A trained and populated FAISS index (inner-product metric). + """A FAISS IVFFlat or IVF-PQ index paired with its proteogram key mapping. + + Parameters: + ----------- + keys: ordered proteogram filenames. keys[i] is the proteogram stored at + FAISS integer id i. + vecs_norm: the L2-normalised embedding matrix, shape (N, d), float32. Kept + so search_all() can re-query the corpus and so queries can be + normalised the same way the corpus was. + index: a trained, populated FAISS index using METRIC_INNER_PRODUCT. + + See also: + ----------- + FaissIndex.from_dataset(): build an index from Img2Vec.dataset + FaissIndex.search_all(): all-vs-all search, returns an Img2Vec.sim_dict + FaissIndex.search_one(): search a single query vector + FaissIndex.save()/load(): persist and restore an index """ - def __init__(self, - keys: List[str], - vecs_norm: np.ndarray, - index) -> None: + def __init__(self, keys: List[str], vecs_norm: np.ndarray, index) -> None: self.keys = keys self.vecs_norm = vecs_norm self._index = index - # ------------------------------------------------------------------ - # Construction - # ------------------------------------------------------------------ - @classmethod def from_dataset(cls, dataset: Dict[str, torch.Tensor], @@ -111,113 +113,135 @@ def from_dataset(cls, nprobe: int = None, pq_m: int = 8, pq_nbits: int = 8) -> "FaissIndex": - """Build a FAISS index from an Img2Vec-style embedding dataset. - - Embeddings are L2-normalised before indexing so inner-product search - is equivalent to cosine similarity. + """Build an index from an Img2Vec-style embedding dataset. Args: - dataset: ``{filename: embedding_tensor}`` dict (same as - ``Img2Vec.dataset``). - use_pq: Use IVF-PQ compressed index. Recommended for corpora - larger than 100 K proteins. Slightly lower recall but - 4–32× lower memory. Defaults to ``False`` (IVFFlat, - exact). - nlist: Number of Voronoi cells. Defaults to - ``max(1, int(sqrt(N)))``. - nprobe: Cells searched per query. Higher → better recall, - slower. Defaults to ``max(1, nlist // 10)``. - pq_m: Sub-quantiser count for IVF-PQ. Must divide ``d`` - evenly. Auto-reduced if necessary. - pq_nbits: Bits per sub-quantiser (IVF-PQ only). 8 is standard. + dataset: {filename: embedding_tensor}, i.e. Img2Vec.dataset. + use_pq: use a product-quantised (IVF-PQ) index. Worth it above + roughly 100k proteograms, where the uncompressed vectors stop + fitting comfortably in memory; costs some recall. Ignored for + corpora under 256 vectors, which are too small to train PQ. + nlist: number of Voronoi cells. Defaults to sqrt(N), the usual + FAISS starting point, capped so cells never outnumber vectors. + nprobe: cells visited per query. Higher is more accurate and + slower. Defaults to nlist // 10. + pq_m: number of PQ sub-quantisers. Must divide d evenly, and is + reduced automatically until it does. + pq_nbits: bits per sub-quantiser. 8 is the standard choice. Returns: - A trained and populated ``FaissIndex`` ready for search. + A trained, populated FaissIndex ready to search. Raises: - ImportError: If ``faiss`` is not installed. - ValueError: If ``dataset`` is empty. + ImportError: faiss is not installed. + ValueError: dataset is empty. """ - try: - import faiss - except ImportError as exc: - raise ImportError( - "faiss is required. Install with:\n" - " uv add faiss-cpu # CPU\n" - " uv add faiss-gpu # GPU / CUDA build" - ) from exc + faiss = _import_faiss() if not dataset: - raise ValueError("dataset is empty — embed_dataset() must be called first.") + raise ValueError('dataset is empty, call embed_dataset() first.') keys, vecs = _stack_dataset(dataset) vecs_norm = _l2_normalise(vecs) - - N, d = vecs_norm.shape - _nlist = nlist if nlist is not None else max(1, int(N ** _NLIST_SQRT_FACTOR ** 0.5)) - - # Cannot have more cells than vectors during training - _nlist = min(_nlist, N) - _nprobe = nprobe if nprobe is not None else max(1, _nlist // 10) - - quantiser = faiss.IndexFlatIP(d) - - if use_pq and N >= 256: - # pq_m must divide d evenly - while d % pq_m != 0 and pq_m > 1: + n_vecs, dim = vecs_norm.shape + + _nlist = nlist if nlist is not None else int(n_vecs ** 0.5) + # Training clusters the corpus itself, so it cannot ask for more + # centroids than there are vectors to cluster. + _nlist = max(1, min(_nlist, n_vecs)) + _nprobe = nprobe if nprobe is not None else _nlist // 10 + _nprobe = max(1, min(_nprobe, _nlist)) + + quantiser = faiss.IndexFlatIP(dim) + if use_pq and n_vecs >= 256: + while dim % pq_m != 0 and pq_m > 1: pq_m -= 1 - index = faiss.IndexIVFPQ( - quantiser, d, _nlist, pq_m, pq_nbits, - faiss.METRIC_INNER_PRODUCT, - ) - index_type = "IVF-PQ" + index = faiss.IndexIVFPQ(quantiser, dim, _nlist, pq_m, pq_nbits, + faiss.METRIC_INNER_PRODUCT) + index_type = 'IVF-PQ' else: - if use_pq and N < 256: - print("WARNING: corpus too small for IVF-PQ (N<256) — falling back to IVFFlat.") - index = faiss.IndexIVFFlat(quantiser, d, _nlist, faiss.METRIC_INNER_PRODUCT) - index_type = "IVFFlat" + if use_pq: + print('Corpus too small to train IVF-PQ (N<256), using IVFFlat instead.') + index = faiss.IndexIVFFlat(quantiser, dim, _nlist, + faiss.METRIC_INNER_PRODUCT) + index_type = 'IVFFlat' index.train(vecs_norm) index.add(vecs_norm) index.nprobe = _nprobe - print( - f"FAISS index built: {index.ntotal} vectors | d={d} | " - f"nlist={_nlist} | nprobe={_nprobe} | type={index_type}" - ) + print(f'Built {index_type} index: {index.ntotal} vectors, d={dim}, ' + f'nlist={_nlist}, nprobe={_nprobe}') return cls(keys=keys, vecs_norm=vecs_norm, index=index) - # ------------------------------------------------------------------ - # Search - # ------------------------------------------------------------------ + @property + def nprobe(self) -> int: + """Number of cells visited per query.""" + return self._index.nprobe + + @nprobe.setter + def nprobe(self, value: int) -> None: + self._index.nprobe = max(1, min(int(value), self._index.nlist)) + + def _reachable(self) -> int: + """Estimate how many hits a query can return at the current nprobe. + + An IVF search only ever sees the vectors inside the cells it probes, + so asking for more than that yields -1 padding no matter how large + top_k is. Cells are not evenly filled, so treat this as a guide. + """ + nlist = self._index.nlist + return max(1, int(round(self.n_vectors * self._index.nprobe / nlist))) + + @contextmanager + def _nprobe_for(self, top_k: int): + """Temporarily raise nprobe so top_k results are actually reachable. + + Without this a deep request (measure_similarity_v2.py asks for the + whole corpus ordering) silently comes back short, padded with -1, + because the cells probed simply do not hold that many vectors. The + original nprobe is restored afterwards so one deep search does not + leave the index scanning exhaustively for every later query. + """ + nlist = self._index.nlist + original = self._index.nprobe + while self._index.nprobe < nlist and self._reachable() < top_k: + self._index.nprobe = min(nlist, self._index.nprobe * 2) + if self._index.nprobe != original: + print(f'Raised nprobe {original} -> {self._index.nprobe} of {nlist} cells ' + f'to reach {top_k} results' + + (' (exhaustive, no ANN speedup at this depth)' + if self._index.nprobe >= nlist else '')) + try: + yield + finally: + self._index.nprobe = original def search_all(self, top_k: int = 10) -> Dict[str, List[Tuple[str, float]]]: - """Batch cosine-similarity search for every vector in the index. + """Search the whole corpus against itself. - Returns a dict with the same structure as ``Img2Vec.sim_dict``: - ``{filename: [(target_filename, score), ...]}``. + Self-hits are kept at rank 0 (score ~1.0) so the result matches what + Img2Vec.similarities() produces and callers can strip them or not. - Self-hits (rank 0, score ≈ 1.0) are included so callers can - decide whether to strip them. + nprobe is widened automatically if top_k is deeper than the current + setting can reach, so the returned ranking is never silently truncated. Args: - top_k: Number of results to return per query (including self-hit). + top_k: results per query, self-hit included. Returns: - Similarity dict keyed by query filename. + {filename: [(target_filename, score), ...]}, the same structure as + Img2Vec.sim_dict. """ - scores_mat, idx_mat = self._index.search(self.vecs_norm, top_k + 1) + top_k = min(top_k, self.n_vectors) + with self._nprobe_for(top_k): + scores_mat, idx_mat = self._index.search(self.vecs_norm, top_k) sim_dict: Dict[str, List[Tuple[str, float]]] = {} for i, key in enumerate(self.keys): - hits: List[Tuple[str, float]] = [] - for rank in range(top_k + 1): - j = int(idx_mat[i, rank]) - if j < 0: # FAISS pads with -1 when fewer results exist - continue - hits.append((self.keys[j], float(scores_mat[i, rank]))) - if len(hits) >= top_k: - break + hits = [(self.keys[j], float(scores_mat[i, rank])) + for rank, j in enumerate(idx_mat[i]) + if j >= 0] # FAISS pads short result rows with -1 sim_dict[key] = hits return sim_dict @@ -225,26 +249,30 @@ def search_one(self, query_vec: np.ndarray, top_k: int = 10, exclude_self_key: str = None) -> List[Tuple[str, float]]: - """Search for the ``top_k`` most similar proteins to a single query. + """Search the corpus for the top_k proteograms closest to one query. Args: - query_vec: 1-D float embedding (not necessarily normalised). - top_k: Number of results to return. - exclude_self_key: If provided, any hit matching this key is skipped - (useful when the query is already in the corpus). + query_vec: 1-D embedding, normalised here so it need not be. + top_k: number of results to return. + exclude_self_key: drop any hit with this key, for when the query + is itself part of the indexed corpus. Returns: - List of ``(filename, cosine_score)`` tuples, descending by score. + [(filename, cosine_score), ...], highest score first. """ - qvec = query_vec.copy().reshape(1, -1).astype(np.float32) + # Search one deeper than asked so dropping the self-hit still leaves + # top_k results. + n_request = min(top_k + 1, self.n_vectors) + + qvec = np.asarray(query_vec, dtype=np.float32).reshape(1, -1).copy() norm = np.linalg.norm(qvec) if norm > 0: qvec /= norm - scores, indices = self._index.search(qvec, top_k + 1) + with self._nprobe_for(n_request): + scores, indices = self._index.search(qvec, n_request) results: List[Tuple[str, float]] = [] - for rank in range(top_k + 1): - j = int(indices[0, rank]) + for rank, j in enumerate(indices[0]): if j < 0: continue key = self.keys[j] @@ -255,64 +283,49 @@ def search_one(self, break return results - # ------------------------------------------------------------------ - # Persistence - # ------------------------------------------------------------------ - def save(self, index_path: str) -> None: - """Save the FAISS index and key mapping to disk. + """Write the index and its key mapping to disk. - Two files are written: - - ``index_path`` — the FAISS binary index - - ``index_path + '.keys.pkl'`` — the ordered key list + Two files are produced: index_path holds the FAISS binary index, and + index_path + '.keys.pkl' holds the key list and normalised vectors + needed to reconstruct the wrapper. Args: - index_path: Destination path, e.g. ``/data/corpus.faiss``. + index_path: destination path, e.g. 'corpus.faiss'. """ - try: - import faiss - except ImportError as exc: - raise ImportError("faiss required for save().") from exc + faiss = _import_faiss() os.makedirs(os.path.dirname(os.path.abspath(index_path)), exist_ok=True) faiss.write_index(self._index, index_path) - keys_path = index_path + ".keys.pkl" - with open(keys_path, "wb") as fh: - pickle.dump({"keys": self.keys, "vecs_norm": self.vecs_norm}, fh) - print(f"Saved FAISS index → {index_path}") - print(f"Saved key mapping → {keys_path}") + keys_path = index_path + '.keys.pkl' + with open(keys_path, 'wb') as pklout: + pickle.dump({'keys': self.keys, 'vecs_norm': self.vecs_norm}, pklout) + print(f'Saved FAISS index to {index_path}') + print(f'Saved key mapping to {keys_path}') @classmethod def load(cls, index_path: str) -> "FaissIndex": - """Load a previously saved FAISS index and key mapping. + """Restore an index written by save(). Args: - index_path: Path to the ``.faiss`` file written by ``save()``. + index_path: path to the .faiss file, with its .keys.pkl companion + alongside it. Returns: - A ready-to-use ``FaissIndex`` instance. + A ready-to-search FaissIndex. """ - try: - import faiss - except ImportError as exc: - raise ImportError("faiss required for load().") from exc + faiss = _import_faiss() index = faiss.read_index(index_path) - keys_path = index_path + ".keys.pkl" - with open(keys_path, "rb") as fh: - data = pickle.load(fh) - keys = data["keys"] - vecs_norm = data["vecs_norm"] - print(f"Loaded FAISS index ← {index_path} ({index.ntotal} vectors)") - return cls(keys=keys, vecs_norm=vecs_norm, index=index) - - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ + keys_path = index_path + '.keys.pkl' + with open(keys_path, 'rb') as pklin: + data = pickle.load(pklin) + print(f'Loaded FAISS index from {index_path} ({index.ntotal} vectors)') + return cls(keys=data['keys'], vecs_norm=data['vecs_norm'], index=index) @property def n_vectors(self) -> int: - """Number of vectors stored in the index.""" + """Number of vectors held in the index.""" return self._index.ntotal @property @@ -321,7 +334,5 @@ def dim(self) -> int: return self._index.d def __repr__(self) -> str: - return ( - f"FaissIndex(n={self.n_vectors}, d={self.dim}, " - f"nprobe={getattr(self._index, 'nprobe', 'N/A')})" - ) + return (f'FaissIndex(n={self.n_vectors}, d={self.dim}, ' + f'nprobe={self.nprobe})') diff --git a/proteogram/v2/image_similarity.py b/proteogram/v2/image_similarity.py index 4d4678b..f6dc122 100644 --- a/proteogram/v2/image_similarity.py +++ b/proteogram/v2/image_similarity.py @@ -20,6 +20,8 @@ from kmeans_pytorch import kmeans from PIL import Image, ImageDraw, ImageFont +from .faiss_search import FaissIndex + class Img2Vec: """ @@ -110,6 +112,9 @@ def __init__(self, model_name_or_path, dataset_dir, embed_file=None, weights="DE self.image_clusters = {} self.cluster_centers = {} self.sim_dict = {} + # Built on demand by build_faiss_index()/load_faiss_index(); only + # similarities_faiss() needs it, similarities() ignores it entirely. + self.faiss_index = None self.files = self.validate_source(dataset_dir) def validate_model(self, model_name_or_path): @@ -758,102 +763,66 @@ def cluster_dataset(self, nclusters, dist="euclidean", display=False): return - # ------------------------------------------------------------------ - # FAISS wrapper methods - # These are thin delegators to proteogram.v2.faiss_search.FaissIndex. - # The FaissIndex instance is stored as self._faiss so scripts can also - # access it directly if needed. - # ------------------------------------------------------------------ - - def build_faiss_index(self, - use_pq: bool = False, - nlist: int = None, - nprobe: int = None, - pq_m: int = 8, - pq_nbits: int = 8) -> None: - """Build a FAISS ANN index from the currently loaded embedding dataset. - - Delegates to :class:`~proteogram.v2.faiss_search.FaissIndex`. - After calling this, ``similarities_faiss()`` can be used as a fast - drop-in replacement for ``similarities()``. - - Args: - use_pq: Use IVF-PQ compressed index (recommended for > 100 K - proteins). Defaults to ``False`` (IVFFlat, exact). - nlist: Voronoi cell count. Defaults to ``sqrt(N)``. - nprobe: Cells searched per query. Defaults to ``nlist // 10``. - pq_m: Sub-quantiser count (IVF-PQ only). - pq_nbits: Bits per sub-quantiser (IVF-PQ only). + def build_faiss_index(self, use_pq=False, nlist=None, nprobe=None, + pq_m=8, pq_nbits=8): + """Build a FAISS ANN index over the currently loaded embedding dataset. + + The index is kept on self.faiss_index, so callers that need more than + similarities_faiss() offers can reach through to it directly. See + proteogram.v2.faiss_search.FaissIndex for what the arguments mean. + + Parameters: + ----------- + use_pq: use a product-quantised index, for corpora large enough that + the uncompressed vectors are a memory problem (roughly >100k). + nlist: number of Voronoi cells, defaults to sqrt(N). + nprobe: cells visited per query, defaults to nlist // 10. + pq_m: number of PQ sub-quantisers, IVF-PQ only. + pq_nbits: bits per sub-quantiser, IVF-PQ only. """ - from .faiss_search import FaissIndex - self._faiss = FaissIndex.from_dataset( - self.dataset, - use_pq=use_pq, - nlist=nlist, - nprobe=nprobe, - pq_m=pq_m, - pq_nbits=pq_nbits, - ) + self.faiss_index = FaissIndex.from_dataset(self.dataset, use_pq=use_pq, + nlist=nlist, nprobe=nprobe, + pq_m=pq_m, pq_nbits=pq_nbits) - def similarities_faiss(self, - n: int = 10, - save_result_images_dir: str = None, - pad_fn=None) -> float: - """ANN similarity search via FAISS — drop-in replacement for ``similarities()``. + def similarities_faiss(self, n=10, save_result_images_dir=None, pad_fn=None): + """ANN equivalent of similarities(), backed by FAISS. - Populates ``self.sim_dict`` with the same ``{filename: [(target, score)]}`` - format so all downstream scripts work without modification. + Fills self.sim_dict with the same {filename: [(target, score), ...]} + structure, self-hit included at rank 0, so downstream scripts do not + need to know which search path produced it. - Call ``build_faiss_index()`` (or ``load_faiss_index()``) first. + Call build_faiss_index() or load_faiss_index() first. - Args: - n: Top-N results per query (self-hit included - at rank 0). - save_result_images_dir: Optional directory to write result images. - pad_fn: Padding callable passed to ``save_images()``. + Parameters: + ----------- + n: number of results per query, including the self-hit. + save_result_images_dir: directory to write result images to, or None. + pad_fn: padding callable handed to save_images(). - Returns: - Wall-clock seconds spent in FAISS search. + Returns the seconds spent in the FAISS search. """ - if not hasattr(self, '_faiss'): - raise RuntimeError( - "Call build_faiss_index() or load_faiss_index() before " - "similarities_faiss()." - ) + if self.faiss_index is None: + raise RuntimeError('No FAISS index, call build_faiss_index() or ' + 'load_faiss_index() before similarities_faiss().') start = time() - self.sim_dict = self._faiss.search_all(top_k=n) + self.sim_dict = self.faiss_index.search_all(top_k=n) elapsed = time() - start if save_result_images_dir: + corpus_dir = os.path.dirname(self.files[0]) if self.files else '' for image_path in self.sim_dict: - full_path = os.path.join( - os.path.dirname(self.files[0]) if self.files else '', - image_path, - ) - self.save_images(full_path, save_result_images_dir, + self.save_images(os.path.join(corpus_dir, image_path), + save_result_images_dir, scores_n_arr=self.sim_dict[image_path], pad_fn=pad_fn) return elapsed - def save_faiss_index(self, index_path: str) -> None: - """Persist the FAISS index to disk. - - Args: - index_path: Destination file path (e.g. ``corpus.faiss``). - A companion ``.keys.pkl`` is written - alongside. - """ - if not hasattr(self, '_faiss'): - raise RuntimeError("No FAISS index to save. Call build_faiss_index() first.") - self._faiss.save(index_path) - - def load_faiss_index(self, index_path: str) -> None: - """Load a previously saved FAISS index. - - Args: - index_path: Path to the ``.faiss`` file written by - ``save_faiss_index()``. - """ - from .faiss_search import FaissIndex - self._faiss = FaissIndex.load(index_path) + def save_faiss_index(self, index_path): + """Write the FAISS index to index_path, plus a .keys.pkl companion.""" + if self.faiss_index is None: + raise RuntimeError('No FAISS index to save, call build_faiss_index() first.') + self.faiss_index.save(index_path) + def load_faiss_index(self, index_path): + """Load a FAISS index previously written by save_faiss_index().""" + self.faiss_index = FaissIndex.load(index_path) diff --git a/pyproject.toml b/pyproject.toml index 5880d87..e8cc44d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,6 @@ dependencies = [ "goatools>=1.4", "pyrotein @ git+https://github.com/carbonscott/pyrotein.git@main", "rcsb-api>=1.7.3", - "faiss-cpu>=1.13.2", ] [project.optional-dependencies] @@ -41,6 +40,10 @@ cuda12 = [ "torch==2.3.0", "torchvision==0.18.0", ] +# ANN search backend for Img2Vec.similarities_faiss(). Optional because +# the brute-force similarities() path needs nothing extra. Swap faiss-cpu +# for faiss-gpu on a CUDA box. +search = ["faiss-cpu>=1.13.2"] test = ["pytest", "pytest-cov"] notebook = [ "jupyterlab>=4.2.5", @@ -49,7 +52,7 @@ notebook = [ [tool.setuptools.packages.find] # Exclude patterns should match the full package name -exclude = ["assets", "docs", "paper", "notebooks", "scripts", "tmp"] +exclude = ["assets", "docs", "paper", "notebooks", "scripts", "tests", "tmp"] [build-system] requires = ["setuptools>=61.0"] diff --git a/scripts/v2/measure_similarity_v2.py b/scripts/v2/measure_similarity_v2.py index 9fed008..bb50dcc 100644 --- a/scripts/v2/measure_similarity_v2.py +++ b/scripts/v2/measure_similarity_v2.py @@ -50,25 +50,24 @@ def get_pad(curr, tgt): parser.add_argument('--embed', action=argparse.BooleanOptionalAction, default=True, help='Recompute and save embeddings (default: True). ' 'Use --no-embed to load from embed_file instead.') - # ── FAISS options ──────────────────────────────────────────────────────── + # FAISS search options parser.add_argument('--faiss', action='store_true', - help=( - 'Use FAISS ANN index for similarity search instead of ' - 'brute-force cosine similarity. Much faster for large ' - 'corpora (> 10 K proteins). Requires faiss-cpu or ' - 'faiss-gpu to be installed.' - )) + help='Search with a FAISS ANN index instead of brute-force ' + 'cosine similarity. Pays off on large corpora; on a few ' + 'thousand proteograms the brute-force path is usually ' + 'faster. Needs faiss-cpu or faiss-gpu installed.') parser.add_argument('--faiss_pq', action='store_true', - help=( - 'Use IVF-PQ compressed FAISS index (recommended for ' - '> 100 K proteins). Slightly lower recall but 4-32x ' - 'lower memory than IVFFlat.' - )) + help='Use a product-quantised (IVF-PQ) FAISS index, which ' + 'trades some recall for much lower memory. Worth it ' + 'above roughly 100k proteograms.') + parser.add_argument('--faiss_top_k', type=int, default=None, + help='Depth of the FAISS ranking to write out. Defaults to the ' + 'whole corpus, which forces an exhaustive search and gives ' + 'up the ANN speedup; set it to the largest K you evaluate ' + 'at to keep the search approximate.') parser.add_argument('--faiss_index_file', type=str, default=None, - help=( - 'Path to save / load the FAISS index. Defaults to ' - 'embed_file with a .faiss extension.' - )) + help='Where to save or load the FAISS index. Defaults to ' + 'embed_file with a .faiss extension.') args = parser.parse_args() # Run embedding vs loading saved embeddings @@ -135,13 +134,10 @@ def _confirm_overwrite(path, label, is_dir=False): + ', '.join(sorted(excluded))) if not prot_files: - raise ValueError( - 'No proteogram .jpg files found for similarity search. '\ - f'Checked dataset_dir={dataset_dir!r}. '\ - 'If you are running from scripts/v2/, ensure config paths are correct '\ - 'relative to that working directory.' - ) - + raise ValueError(f'No proteogram .jpg files found under {dataset_dir!r}. ' + 'Config paths are resolved against the working directory, ' + 'so check them if running from inside scripts/v2/.') + device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f'Using device: {device}') @@ -212,10 +208,8 @@ def _prep_fn(img): if embed: img_sim.embed_dataset() if not img_sim.dataset: - raise ValueError( - 'Embedding dataset is empty after embed_dataset(). '\ - 'Verify input proteogram files are readable and preprocessing succeeded.' - ) + raise ValueError('embed_dataset() produced no embeddings. Check that the ' + 'proteogram files are readable and preprocessing worked.') # Save embeddings with open(embed_file, 'wb') as pklout: pickle.dump(img_sim.dataset, pklout) @@ -225,10 +219,8 @@ def _prep_fn(img): with open(embed_file, 'rb') as pklin: img_sim.dataset = pickle.load(pklin) if not img_sim.dataset: - raise ValueError( - 'Loaded embedding dataset is empty. '\ - f'Check embed_file={embed_file!r} or rerun with --embed.' - ) + raise ValueError(f'No embeddings in {embed_file!r}, rerun with --embed ' + 'to recompute them.') # Search to find similar images using cosine-similarity amongst embeddings. # Save all corpus results (including self-hit) so Recall@K can be computed at @@ -238,24 +230,28 @@ def _prep_fn(img): n_results = len(prot_files) # all including self-hit if args.faiss: - # ── FAISS ANN search ───────────────────────────────────────────── if args.faiss_index_file: faiss_index_file = args.faiss_index_file else: - base, _ = os.path.splitext(embed_file) - faiss_index_file = base + '.faiss' + faiss_index_file = os.path.splitext(embed_file)[0] + '.faiss' if os.path.exists(faiss_index_file) and not args.overwrite: print(f'Loading existing FAISS index from {faiss_index_file}') img_sim.load_faiss_index(faiss_index_file) else: - print(f'Building FAISS index (use_pq={args.faiss_pq}) ...') + print(f'Building FAISS index (use_pq={args.faiss_pq})') img_sim.build_faiss_index(use_pq=args.faiss_pq) img_sim.save_faiss_index(faiss_index_file) - sim_time = img_sim.similarities_faiss(n=n_results, + # An IVF search only returns what sits in the cells it probes, so + # asking for the full corpus ranking makes it scan every cell. + # Ranking less deeply is what keeps the search approximate, and fast. + faiss_top_k = min(args.faiss_top_k or n_results, n_results) + sim_time = img_sim.similarities_faiss(n=faiss_top_k, save_result_images_dir=None, pad_fn=_prep_fn) + if faiss_top_k < n_results: + print(f'Ranked the top {faiss_top_k} of {n_results} results per query; ' + f'metrics beyond K={faiss_top_k} cannot be computed from this run.') else: - # ── Brute-force cosine search (original) ───────────────────────── sim_time = img_sim.similarities(n=n_results, save_result_images_dir=None, pad_fn=_prep_fn) @@ -273,16 +269,22 @@ def _prep_fn(img): print(f'Took {time()-start} seconds overall (including optional image result saving).') # Create dataframe of results - scores_tmp = [[''] * n_results] * len(prot_files) - df_res = pd.DataFrame(scores_tmp, columns=[str(i) for i in range(n_results)]) + # Width the table to the deepest ranking actually produced. The FAISS + # path can return fewer than n_results per query, and blank trailing + # cells read back as NaN, which evaluate_methods_v2.py cannot parse. + n_cols = min((len(v) for v in img_sim.sim_dict.values()), default=n_results) + scores_tmp = [[''] * n_cols] * len(prot_files) + df_res = pd.DataFrame(scores_tmp, columns=[str(i) for i in range(n_cols)]) df_res['query_image'] = prot_files for i, image_path in enumerate(prot_files): try: scores = img_sim.sim_dict[os.path.basename(image_path)] - row_vals = [f'{a},{b}' for (a, b) in scores[:n_results]] - df_res.iloc[i, :len(row_vals)] = row_vals + df_res.iloc[i, :n_cols] = [f'{a},{b}' for (a, b) in scores[:n_cols]] except KeyError as e: print(f'Key error for {e}') + if n_cols < n_results: + print(f'Wrote {n_cols} of {n_results} possible result columns, limited by ' + f'the query with the fewest hits.') # Reorder cols df_res.drop('query_image', inplace=True, axis=1) df_res.insert(0, 'query_image', prot_files) diff --git a/tests/test_faiss_search.py b/tests/test_faiss_search.py new file mode 100644 index 0000000..4c2e144 --- /dev/null +++ b/tests/test_faiss_search.py @@ -0,0 +1,175 @@ +"""Tests for the FAISS ANN search backend. + +Run with: uv run --extra search --extra test pytest tests/test_faiss_search.py + +The corpora here are random gaussian vectors, which is close to the worst case +for an IVF index: there is no cluster structure for the coarse quantiser to +exploit, so approximate recall is low. That is deliberate. These tests check +the index's mechanics (defaults, truncation, state, persistence), not the +recall the real proteogram embeddings achieve, which only a run against the +actual corpus can tell you. +""" +import numpy as np +import pytest +import torch + +faiss = pytest.importorskip('faiss', reason='install the "search" extra to run these') + +from proteogram.v2.faiss_search import FaissIndex + + +N_VECS = 600 +DIM = 128 + + +@pytest.fixture(scope='module') +def vectors(): + rng = np.random.default_rng(0) + return rng.normal(size=(N_VECS, DIM)).astype(np.float32) + + +@pytest.fixture +def dataset(vectors): + return {f'p{i}.jpg': torch.from_numpy(vectors[i].copy()) for i in range(N_VECS)} + + +@pytest.fixture +def index(dataset): + return FaissIndex.from_dataset(dataset) + + +def brute_force_top_k(vectors, k): + """Exact cosine ranking, as Img2Vec.similarities() would produce it.""" + unit = vectors / np.linalg.norm(vectors, axis=1, keepdims=True) + return np.argsort(-(unit @ unit.T), axis=1)[:, :k] + + +def test_defaults_follow_sqrt_n(index): + # nlist defaulting to N rather than sqrt(N) gives every vector its own cell, + # which trains slowly and is not an approximate index in any useful sense. + assert index._index.nlist == int(N_VECS ** 0.5) + assert index.nprobe == index._index.nlist // 10 + assert index.n_vectors == N_VECS + assert index.dim == DIM + + +def test_empty_dataset_rejected(): + with pytest.raises(ValueError, match='empty'): + FaissIndex.from_dataset({}) + + +def test_source_embeddings_are_not_normalised_in_place(dataset, vectors): + FaissIndex.from_dataset(dataset) + after = np.stack([dataset[f'p{i}.jpg'].numpy() for i in range(N_VECS)]) + np.testing.assert_allclose(after, vectors) + + +def test_deep_ranking_is_not_truncated(index): + # measure_similarity_v2.py asks for the whole corpus ordering so Recall@K + # can be computed at any K afterwards. An IVF index only returns what sits + # in the cells it probes, so this has to widen nprobe rather than pad -1. + sim_dict = index.search_all(top_k=N_VECS) + assert {len(hits) for hits in sim_dict.values()} == {N_VECS} + + +def test_deep_search_does_not_leave_nprobe_raised(index): + before = index.nprobe + index.search_all(top_k=N_VECS) + assert index.nprobe == before + index.search_one(np.zeros(DIM, dtype=np.float32), top_k=N_VECS) + assert index.nprobe == before + + +def test_exhaustive_ranking_matches_brute_force(index, vectors): + sim_dict = index.search_all(top_k=N_VECS) + expected = brute_force_top_k(vectors, 20) + for i in range(N_VECS): + got = [key for key, _ in sim_dict[f'p{i}.jpg'][:20]] + assert got == [f'p{j}.jpg' for j in expected[i]] + + +def test_self_hit_leads_each_ranking(index): + sim_dict = index.search_all(top_k=5) + for i in range(N_VECS): + key, score = sim_dict[f'p{i}.jpg'][0] + assert key == f'p{i}.jpg' + assert score == pytest.approx(1.0, abs=1e-4) + + +def test_raising_nprobe_improves_recall(index, vectors): + expected = [{f'p{j}.jpg' for j in row} for row in brute_force_top_k(vectors, 20)] + + def recall_at_20(): + sim_dict = index.search_all(top_k=20) + found = sum(len({k for k, _ in sim_dict[f'p{i}.jpg']} & expected[i]) + for i in range(N_VECS)) + return found / (N_VECS * 20) + + default = recall_at_20() + index.nprobe = index._index.nlist + assert recall_at_20() > default + + +def test_nprobe_setter_clamps_to_valid_range(index): + index.nprobe = 10 ** 6 + assert index.nprobe == index._index.nlist + index.nprobe = -5 + assert index.nprobe == 1 + + +def test_search_one_can_drop_the_query_itself(index, vectors): + hits = index.search_one(vectors[7], top_k=5, exclude_self_key='p7.jpg') + assert len(hits) == 5 + assert all(key != 'p7.jpg' for key, _ in hits) + # Scores must still come back descending after the self-hit is removed. + assert [s for _, s in hits] == sorted((s for _, s in hits), reverse=True) + + +def test_search_one_normalises_the_query(index, vectors): + scaled = index.search_one(vectors[3] * 17.0, top_k=5) + plain = index.search_one(vectors[3], top_k=5) + assert [key for key, _ in scaled] == [key for key, _ in plain] + # Scaling then normalising in float32 costs a few ulps, so scores match + # approximately rather than bit for bit. + for (_, a), (_, b) in zip(scaled, plain): + assert a == pytest.approx(b, abs=1e-5) + + +def test_zero_vector_does_not_divide_by_zero(dataset): + dataset['zero.jpg'] = torch.zeros(DIM) + idx = FaissIndex.from_dataset(dataset) + hits = idx.search_one(np.zeros(DIM, dtype=np.float32), top_k=3) + assert len(hits) == 3 + + +def test_save_load_round_trip(index, tmp_path): + path = str(tmp_path / 'nested' / 'corpus.faiss') # parent dirs are created + index.save(path) + restored = FaissIndex.load(path) + assert restored.keys == index.keys + assert restored.search_all(top_k=10) == index.search_all(top_k=10) + + +def test_pq_falls_back_on_a_corpus_too_small_to_train(vectors): + small = {f'p{i}.jpg': torch.from_numpy(vectors[i].copy()) for i in range(100)} + idx = FaissIndex.from_dataset(small, use_pq=True) + assert isinstance(idx._index, faiss.IndexIVFFlat) + + +def test_pq_index_is_searchable(dataset): + idx = FaissIndex.from_dataset(dataset, use_pq=True) + assert isinstance(idx._index, faiss.IndexIVFPQ) + sim_dict = idx.search_all(top_k=5) + assert len(sim_dict) == N_VECS + + +def test_pq_sub_quantiser_count_is_reduced_to_divide_the_dimension(dataset): + # 7 does not divide 128, so from_dataset() has to walk pq_m down to 4. + idx = FaissIndex.from_dataset(dataset, use_pq=True, pq_m=7) + assert DIM % idx._index.pq.M == 0 + + +def test_nlist_cannot_exceed_the_corpus_size(vectors): + tiny = {f'p{i}.jpg': torch.from_numpy(vectors[i].copy()) for i in range(4)} + idx = FaissIndex.from_dataset(tiny, nlist=500) + assert idx._index.nlist == 4 diff --git a/uv.lock b/uv.lock index d2db451..85c73f4 100644 --- a/uv.lock +++ b/uv.lock @@ -2703,7 +2703,6 @@ version = "0.0.4" source = { editable = "." } dependencies = [ { name = "biopython" }, - { name = "faiss-cpu" }, { name = "goatools" }, { name = "kmeans-pytorch" }, { name = "matplotlib" }, @@ -2737,6 +2736,9 @@ notebook = [ { name = "jupyterlab" }, { name = "nglview" }, ] +search = [ + { name = "faiss-cpu" }, +] test = [ { name = "pytest" }, { name = "pytest-cov" }, @@ -2745,7 +2747,7 @@ test = [ [package.metadata] requires-dist = [ { name = "biopython", specifier = ">=1.8" }, - { name = "faiss-cpu", specifier = ">=1.13.2" }, + { name = "faiss-cpu", marker = "extra == 'search'", specifier = ">=1.13.2" }, { name = "goatools", specifier = ">=1.4" }, { name = "jupyterlab", marker = "extra == 'notebook'", specifier = ">=4.2.5" }, { name = "kmeans-pytorch", specifier = ">=0.3" }, @@ -2775,7 +2777,7 @@ requires-dist = [ { name = "torchvision", marker = "extra == 'cuda12'", specifier = "==0.18.0" }, { name = "tqdm", specifier = ">=4.67" }, ] -provides-extras = ["cuda12", "test", "notebook"] +provides-extras = ["cuda12", "search", "test", "notebook"] [[package]] name = "psutil" From a5d3371792bf71e3ccf2defb3eefda0f8e735c7e Mon Sep 17 00:00:00 2001 From: SonOfAnton Date: Fri, 11 Sep 2026 23:23:33 -0700 Subject: [PATCH 3/4] Document the FAISS search path Adds docs/faiss_search.md and wires it into the README's Step 4 and the scripts reference table. The thing the doc leads with is that --faiss on its own is slower than the brute-force path it replaces. measure_similarity_v2.py ranks the whole corpus per query so Recall@K works at any K, and an exhaustive ranking gives an IVF index nothing to skip, so it does brute-force work plus indexing overhead. Measured on random 512-d vectors: 1.46s vs 0.11s at N=2008, and 77.37s vs 4.28s at N=13503, the released demo corpus size. Capping the depth with --faiss_top_k 100 turns that into 1.60s vs 4.28s. Also records what has not been established: recall on real proteogram embeddings. On random vectors, which have no cluster structure for the coarse quantiser to exploit, Recall@20 is 0.26 at the default nprobe and 0.71 at nlist//2. Real embeddings are clustered by fold and superfamily so this probably understates it, but it should be measured against the brute-force path before any --faiss numbers are used in a GTalign/USalign/Foldseek comparison. --- README.md | 10 ++- docs/faiss_search.md | 158 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 docs/faiss_search.md diff --git a/README.md b/README.md index c31b628..7a8101e 100644 --- a/README.md +++ b/README.md @@ -350,6 +350,14 @@ python measure_similarity_v2.py Key optional flags: - `--no-embed`: Skip embedding and load from `embed_file` (faster if embeddings already exist) - `--exclude_classes h,i,j,k,l`: Exclude classes from the search corpus +- `--faiss`: Use an approximate nearest neighbour index instead of brute-force cosine search. Only worth it in combination with `--faiss_top_k` — see [Approximate search with FAISS](docs/faiss_search.md) + +By default the script ranks the **entire** corpus for every query, so that +Recall@K can be computed afterwards at any K. That is the right default for +benchmarking, and it is why plain `--faiss` is *slower* than brute force: an +exhaustive ranking gives an ANN index nothing to skip. Cap the ranking depth +with `--faiss_top_k` to get the speedup, at the cost of not being able to +evaluate beyond that depth. --- @@ -535,7 +543,7 @@ The `v1` and `v2` subfolders have their own `config.yml` (copy from the correspo |--------|---------|--------------------------------|------------------------| | `v2/create_v2_proteograms.py` | Create proteograms using MD-based nonbonded energy calculations, distances, and hydrophobicity deltas. `cg_method: martini` selects the fast coarse-grained path; null/empty uses the all-atom path. Note: `calpha_atom_distance_cutoff=10` Å is **hard-coded** in the script (not config/CLI) | `limit_file`, `scope_structures_dir`, `all_proteograms_dir`, `cg_method` | `--max_workers/-w`, `--overwrite`, `--verbose`, `--debug`, `--memory-efficient`, `--save_simulated_pdb`, `--sequence_len_lower_cutoff` (default 20), `--sequence_len_upper_cutoff` (default 200) | | `v2/query_similar_proteins.py` | Create a proteogram for a single query PDB and find the top-K most similar proteins from a pre-computed corpus | `top_k`, `model_file`, `embed_file`, `cg_method`, `proteograms_for_sim_dir` (optional — parent or root directory containing corpus `.jpg` files, searched recursively, needed for result image) | `--pdb_file/-p`, `--chain_id/-c`, `--cg_method` (`martini`\|`atomistic`), `--output_dir/-o`, `--top_k/-k`, `--model_file`, `--embed_file`, `--annot_file` (optional agreement report), `--target_size`, `--resize/--no-resize`, `--sequence_len_lower_cutoff`, `--sequence_len_upper_cutoff`. Preprocessing (grid/resize/cutoff) defaults to the checkpoint meta | -| `v2/measure_similarity_v2.py` | Batch similarity search across all proteograms | `top_k`, `model_file`, `embed_file`, `proteogram_sim_results`, `proteograms_for_sim_dir`, `search_images_dir` | `--exclude_classes/-x`, `--overwrite`, `--embed/--no-embed` | +| `v2/measure_similarity_v2.py` | Batch similarity search across all proteograms. Brute-force cosine by default; `--faiss` switches to an ANN index (see [docs/faiss_search.md](docs/faiss_search.md)) | `top_k`, `model_file`, `embed_file`, `proteogram_sim_results`, `proteograms_for_sim_dir`, `search_images_dir` | `--exclude_classes/-x`, `--overwrite`, `--embed/--no-embed`, `--faiss`, `--faiss_top_k`, `--faiss_pq`, `--faiss_index_file` | | `v2/train_multiple_models_randomized_eval.py` | **Current trainer.** Train ResNet18/ConvNet/ViT-B/16 with a reproducible seeded train/val/**test** split from a single directory; writes a self-describing checkpoint (grid/resize/cutoff in `meta`). Supports classification (`ce`/`focal`) and retrieval (`triplet_hierarchy`) losses | `training_data_dir`, `model_file_prefix`, `pretrained` | `--data_dir/-d`, `--model/-m` (`cnn`\|`resnet18`\|`vit`), `--epochs/-e`, `--batch_size/-b`, `--lr/-l`, `--pretrained/--no-pretrained`, `--seed`, `--max_image_size`, `--input_size`, `--resize`, `--level`, `--min_class_size`, `--loss` (`ce`\|`focal`\|`triplet_hierarchy`), `--patience`, `--val_size`, `--test_size`, `--save_test_list/--no-save_test_list`, `--save_train_list`, `--test_list`, `--exclude_classes/-x`, `--tsv_file/-t`, `--overwrite/-o`, `--verbose/-v` (plus `--triplet_*`, `--focal_gamma`, `--embed_dim` for the respective losses) | | `v2/train_multiple_models.py` | Legacy trainer (manual `train/`/`eval/` split, no self-describing meta). Prefer the randomized-eval trainer above for new models | `training_data_dir`, `num_epochs`, `learning_rate`, `batch_size`, `scope_level`, `model_file_prefix` | `--data_dir/-d` (overrides `training_data_dir`), `--epochs/-e`, `--batch_size/-b`, `--lr/-l`, `--model/-m` (`cnn`\|`resnet18`), `--level` (`class`\|`fold`\|`superfamily`\|`family`, default: `class`), `--tsv_file/-t`, `--patience`, `--val_size`, `--exclude_classes/-x`, `--overwrite/-o`, `--resize`, `--verbose/-v` | | `v2/evaluate_methods_v2.py` | Evaluate proteogram approach vs GTalign, USalign, and Foldseek | `top_k`, `scope_eval_set`, `proteogram_sim_results`, `gtalign_results_dir`, `usalign_results`, `foldseek_results` (optional), `search_images_dir`, `save_bad_searches_dir`, `save_good_searches_dir`, `scope_cla_file`, `scope_des_file`, `scope_hie_file` | `--overwrite`, `--exclude_classes/-x`, `--bootstrap`, `--n_boot` (default 10000), `--boot_seed` (default 0) | diff --git a/docs/faiss_search.md b/docs/faiss_search.md new file mode 100644 index 0000000..5795f37 --- /dev/null +++ b/docs/faiss_search.md @@ -0,0 +1,158 @@ +# Approximate Nearest Neighbour Search with FAISS + +*Applies to `scripts/v2/measure_similarity_v2.py` and `proteogram.v2.FaissIndex`.* + +> **Read this before turning on `--faiss`.** Used with its default settings it +> is *slower* than the brute-force path it replaces. The speedup only exists +> if you also cap the ranking depth with `--faiss_top_k`. + +## Why + +`Img2Vec.similarities()` scores every query against every corpus vector. That +is O(N^2) in time and memory, and it becomes the bottleneck long before the +embedding step does. A FAISS IVF index instead partitions the corpus into +Voronoi cells and visits only a few of them per query, so search cost scales +with `nprobe` rather than with corpus size. + +Embeddings are L2-normalised before indexing, so the inner-product metric +FAISS searches with is exactly the cosine similarity the brute-force path +reports. Both paths fill `Img2Vec.sim_dict` with the same +`{filename: [(target, score), ...]}` structure, self-hit at rank 0, so +`evaluate_methods_v2.py` does not care which one produced the results. + +## The catch: ranking depth + +`measure_similarity_v2.py` ranks the **whole corpus** for every query by +default, so Recall@K can be computed afterwards at any K. This is a sensible +benchmarking default, but it is the worst possible case for an ANN index. An +IVF search only ever returns vectors that live in the cells it probes, so +asking for all N results forces it to probe every cell. At that point it is +doing the same work as brute force plus the indexing overhead. + +`FaissIndex.search_all()` detects this and widens `nprobe` until the requested +depth is actually reachable, printing what it did. It never silently returns a +short ranking. (An earlier revision did, returning 200 of 2008 requested +results per query and writing blank CSV cells that crashed +`evaluate_methods_v2.py` when it tried to parse them as `target,score`.) + +So: **cap the depth.** Set `--faiss_top_k` to the largest K you actually +evaluate at. The script prints a warning that metrics beyond that K are not +computable from the run. + +## Measured cost + +Random 512-dimensional vectors, single CPU. `N=2008` is the size of the SCOPe +eval set used in Step 4; `N=13503` is the released demo corpus. + +| corpus | brute force, full ranking | `--faiss`, full ranking | `--faiss --faiss_top_k 100` | +|---|---|---|---| +| N = 2,008 | 0.11 s | 1.46 s | 0.07 s | +| N = 13,503 | 4.28 s | 77.37 s | 1.60 s | + +Index build time is small and is not the problem: 0.10 s at N=2,008 and 0.16 s +at N=13,503. + +Two things to take from this. Plain `--faiss` costs you roughly **18x** at the +demo corpus size and gets worse as the corpus grows, because the full ranking +is exhaustive either way. With the depth capped at 100, FAISS is **~2.7x +faster** than brute force at N=13,503, and the gap widens with N since the +brute-force path stays O(N^2) while the capped ANN path does not. + +Note that these are random vectors, which have no cluster structure for the +coarse quantiser to exploit. Real proteogram embeddings are clustered by fold +and superfamily, so recall at a given `nprobe` should be better than what +random data suggests -- but measure it on your own corpus rather than assuming. + +## Recall + +ANN trades recall for speed, and the default `nprobe = nlist // 10` is +aggressive. On random 512-d vectors, Recall@20 against the exact ranking was +**0.26** at the default and **0.71** at `nprobe = nlist // 2`. Again, random +vectors are the worst case -- but the shape of the tradeoff is real. + +Before trusting `--faiss` numbers in a comparison against GTalign, USalign or +Foldseek, run both paths on the same corpus and confirm the retrieval metrics +agree. If they do not, raise `nprobe`: + +```python +from proteogram.v2 import FaissIndex + +index = FaissIndex.from_dataset(img_sim.dataset) +index.nprobe = index._index.nlist // 2 # clamped to [1, nlist] +``` + +## Usage + +Install the extra (FAISS is optional; the brute-force path needs nothing): + +```bash +uv sync --extra search +``` + +Then from `scripts/v2/`: + +```bash +# Rank the top 100 per query with an ANN index +python measure_similarity_v2.py --no-embed --faiss --faiss_top_k 100 + +# Compressed index, for corpora large enough that the raw vectors are a +# memory problem (roughly >100k proteograms) +python measure_similarity_v2.py --no-embed --faiss --faiss_top_k 100 --faiss_pq +``` + +The index is saved next to `embed_file` with a `.faiss` extension and reused +on later runs; pass `--overwrite` to rebuild it, or `--faiss_index_file` to put +it somewhere else. A companion `.keys.pkl` holds the filename mapping. + +### Flags + +| Flag | Meaning | +|---|---| +| `--faiss` | Use the ANN index instead of brute-force cosine search | +| `--faiss_top_k N` | Rank only the top N per query. **This is what makes it fast.** Defaults to the whole corpus | +| `--faiss_pq` | Use a product-quantised (IVF-PQ) index: much lower memory, some recall lost. Ignored below 256 vectors, which is too few to train | +| `--faiss_index_file` | Where to save/load the index. Defaults to `embed_file` with a `.faiss` extension | + +## Library API + +`FaissIndex` is independent of `Img2Vec` -- it operates on plain float32 numpy +arrays plus an ordered key list, and imports `faiss` lazily, so importing +`proteogram.v2` works without the extra installed. + +```python +from proteogram.v2 import FaissIndex + +index = FaissIndex.from_dataset(img_sim.dataset) # {filename: tensor} +sim_dict = index.search_all(top_k=100) # same shape as Img2Vec.sim_dict +hits = index.search_one(query_vec, top_k=10, exclude_self_key='d1abca_.jpg') + +index.save('corpus.faiss') +index = FaissIndex.load('corpus.faiss') +``` + +`Img2Vec` also wraps this: `build_faiss_index()`, `similarities_faiss()`, +`save_faiss_index()`, `load_faiss_index()`. + +### Tuning + +| Parameter | Default | Effect | +|---|---|---| +| `nlist` | `sqrt(N)` | Voronoi cells. More cells means finer partitioning and slower training; capped at N | +| `nprobe` | `nlist // 10` | Cells visited per query. **The main recall/speed dial** | +| `pq_m` | 8 | IVF-PQ sub-quantisers. Must divide the embedding dimension; reduced automatically until it does | +| `pq_nbits` | 8 | Bits per sub-quantiser | + +`search_all()` and `search_one()` raise `nprobe` temporarily when the requested +depth needs it, then put it back, so one deep query does not leave the index +scanning exhaustively for everything afterwards. + +## Tests + +```bash +uv run --extra search --extra test pytest tests/test_faiss_search.py +``` + +Covers the `nlist`/`nprobe` defaults, that deep rankings are not truncated, +that `nprobe` is restored after a deep search, self-hit ordering, exactness +against brute force when the search is exhaustive, the recall/`nprobe` +relationship, save/load round trips, and the IVF-PQ fallbacks. From 085634278f7ecd93595ae868d3a33e5d3dd6cafb Mon Sep 17 00:00:00 2001 From: SonOfAnton Date: Fri, 11 Sep 2026 23:37:27 -0700 Subject: [PATCH 4/4] Correct the FAISS benchmarks using the real corpus embeddings The performance numbers in the previous commit were measured on random gaussian vectors and were wrong in both directions. Replaced with a sweep over the released 13,503-proteogram corpus embeddings. Two errors. The claimed 2.7x speedup at top-100 compared FAISS against brute force doing a full argsort of the entire NxN matrix, when only the top K was needed; against a chunked argpartition top-K the same comparison is 1.1x. And the recall figures came from random vectors, which understate real embeddings by about 3x at the same scanned fraction (Recall@10 of 0.31 vs 0.87 at 8.7% of the corpus) because the model clusters structures by fold and superfamily, giving the coarse quantiser real structure to exploit. On the real corpus the feature looks considerably better than the random-vector numbers suggested. Brute-force top-10 is 1.82s; the index reaches 4.8x at Recall@10 0.962 (nprobe=4, 3.4% of the corpus) and 12.6x at 0.841 (nprobe=1). The shipped default nprobe = nlist//10 sits at the conservative end, 1.9x at 0.991. What has not changed is that a full-corpus ranking is the wrong way to use an ANN index: at nprobe = nlist it scans everything and is 5x slower than brute force, so --faiss_top_k is required rather than optional. Documents the random-vs-real comparison so the mistake is not repeated, and notes that the sweep should be re-run on substantially larger corpora since nlist = sqrt(N) shifts the useful nprobe range. --- README.md | 14 ++++-- docs/faiss_search.md | 100 ++++++++++++++++++++++++++++--------------- 2 files changed, 75 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 7a8101e..737c198 100644 --- a/README.md +++ b/README.md @@ -354,10 +354,16 @@ Key optional flags: By default the script ranks the **entire** corpus for every query, so that Recall@K can be computed afterwards at any K. That is the right default for -benchmarking, and it is why plain `--faiss` is *slower* than brute force: an -exhaustive ranking gives an ANN index nothing to skip. Cap the ranking depth -with `--faiss_top_k` to get the speedup, at the cost of not being able to -evaluate beyond that depth. +benchmarking, and it is why plain `--faiss` is *5x slower* than brute force: an +exhaustive ranking gives an ANN index nothing to skip. Always pair `--faiss` +with `--faiss_top_k`, at the cost of not being able to evaluate beyond that +depth. + +On the released 13,503-proteogram corpus, a depth-capped index is **~5x faster +than brute force while retaining 96% of the exact top-10**, and up to 12.6x if +you accept 84%. The shipped `nprobe` default is deliberately conservative +(1.9x, 99.1% recall); [docs/faiss_search.md](docs/faiss_search.md) has the full +speed/recall sweep and how to tune it. --- diff --git a/docs/faiss_search.md b/docs/faiss_search.md index 5795f37..33a4e68 100644 --- a/docs/faiss_search.md +++ b/docs/faiss_search.md @@ -2,9 +2,12 @@ *Applies to `scripts/v2/measure_similarity_v2.py` and `proteogram.v2.FaissIndex`.* -> **Read this before turning on `--faiss`.** Used with its default settings it -> is *slower* than the brute-force path it replaces. The speedup only exists -> if you also cap the ranking depth with `--faiss_top_k`. +> **Read this before turning on `--faiss`.** Plain `--faiss` ranks the whole +> corpus, which makes it *5x slower* than brute force. Always pair it with +> `--faiss_top_k`. On the released 13,503-proteogram corpus a tuned index is +> ~5x faster than brute force at 96% Recall@10; the shipped default is +> deliberately conservative at 1.9x and 99.1%. See +> [Measured cost and recall](#measured-cost-and-recall). ## Why @@ -39,47 +42,74 @@ So: **cap the depth.** Set `--faiss_top_k` to the largest K you actually evaluate at. The script prints a warning that metrics beyond that K are not computable from the run. -## Measured cost +## Measured cost and recall -Random 512-dimensional vectors, single CPU. `N=2008` is the size of the SCOPe -eval set used in Step 4; `N=13503` is the released demo corpus. +Measured on the **released 13,503-proteogram corpus embeddings** (the ResNet18 +superfamily CE checkpoint, 512-d), retrieving top-10, single machine, 4 FAISS +threads against a 4-thread OpenBLAS brute-force reference. Brute-force top-10 +over this corpus takes 1.82 s. -| corpus | brute force, full ranking | `--faiss`, full ranking | `--faiss --faiss_top_k 100` | -|---|---|---|---| -| N = 2,008 | 0.11 s | 1.46 s | 0.07 s | -| N = 13,503 | 4.28 s | 77.37 s | 1.60 s | +`nlist` defaults to `sqrt(N) = 116` here, so `nprobe` is the fraction of the +corpus scanned. Recall@10 is measured against the exact cosine ranking: -Index build time is small and is not the problem: 0.10 s at N=2,008 and 0.16 s -at N=13,503. +| `nprobe` | % of corpus scanned | time | vs brute force | Recall@10 | +|---|---|---|---|---| +| 1 | 0.9% | 0.14 s | **12.6x** | 0.841 | +| 2 | 1.7% | 0.25 s | **7.3x** | 0.916 | +| 3 | 2.6% | 0.34 s | 5.4x | 0.946 | +| 4 | 3.4% | 0.38 s | 4.8x | 0.962 | +| 6 | 5.2% | 0.52 s | 3.5x | 0.977 | +| 8 | 6.9% | 0.69 s | 2.6x | 0.984 | +| 11 *(default)* | 9.5% | 0.94 s | 1.9x | 0.991 | +| 16 | 13.8% | 1.65 s | 1.1x | 0.997 | +| 24 | 20.7% | 2.97 s | 0.6x | 0.999 | +| 116 | 100% | 9.76 s | 0.2x | 1.000 | -Two things to take from this. Plain `--faiss` costs you roughly **18x** at the -demo corpus size and gets worse as the corpus grows, because the full ranking -is exhaustive either way. With the depth capped at 100, FAISS is **~2.7x -faster** than brute force at N=13,503, and the gap widens with N since the -brute-force path stays O(N^2) while the capped ANN path does not. +There is a broad useful range here. Around `nprobe = 4` (3.4% of the corpus) +the index is roughly **5x faster than brute force while keeping 96% of the +exact top-10**. Pushing to `nprobe = 1` buys 12.6x at 84% recall. -Note that these are random vectors, which have no cluster structure for the -coarse quantiser to exploit. Real proteogram embeddings are clustered by fold -and superfamily, so recall at a given `nprobe` should be better than what -random data suggests -- but measure it on your own corpus rather than assuming. +The shipped default of `nprobe = nlist // 10` sits at the conservative end: +1.9x faster, 99.1% recall. That is a defensible default -- it barely perturbs +the ranking -- but if search time matters, lowering it is where the gains are. -## Recall +Two settings to avoid: -ANN trades recall for speed, and the default `nprobe = nlist // 10` is -aggressive. On random 512-d vectors, Recall@20 against the exact ranking was -**0.26** at the default and **0.71** at `nprobe = nlist // 2`. Again, random -vectors are the worst case -- but the shape of the tradeoff is real. +- **A full-corpus ranking.** At `nprobe = nlist` the index scans everything and + is **5x slower** than brute force (9.76 s vs 1.82 s), since it does the same + work plus indexing overhead. This is what plain `--faiss` does by default, + which is why `--faiss_top_k` matters. +- **`nprobe` above ~15% of the corpus.** Past that the index is slower than + brute force for recall gains in the third decimal place. -Before trusting `--faiss` numbers in a comparison against GTalign, USalign or -Foldseek, run both paths on the same corpus and confirm the retrieval metrics -agree. If they do not, raise `nprobe`: +Index build time is negligible and is not part of this tradeoff: 0.16 s. -```python -from proteogram.v2 import FaissIndex +### Why random test vectors are not a proxy -index = FaissIndex.from_dataset(img_sim.dataset) -index.nprobe = index._index.nlist // 2 # clamped to [1, nlist] -``` +Earlier revisions of this document quoted figures measured on random gaussian +vectors. Those understated real performance by a wide margin and have been +removed. At matched N and dimension, Recall@10 on real proteogram embeddings +versus random vectors: + +| % of corpus scanned | real embeddings | random gaussian | +|---|---|---| +| 4.3% | 0.702 | 0.225 | +| 8.7% | 0.869 | 0.307 | +| 17.4% | 0.968 | 0.437 | +| 34.8% | 0.998 | 0.633 | + +Roughly a 3x difference at the same scanned fraction. This is expected: the +model is trained to cluster structures by fold and superfamily, so the coarse +quantiser has genuine structure to exploit, whereas isotropic gaussian vectors +have none. Benchmark against real embeddings. + +### Scaling + +The numbers above are for one corpus size. Brute-force cost grows as O(N^2) +while the IVF path at fixed `nprobe/nlist` does not, so the advantage should +widen with N -- but that has not been measured beyond 13,503 on real data, and +the useful `nprobe` may shift as `nlist = sqrt(N)` grows. Re-run the sweep if +you move to a substantially larger corpus. ## Usage @@ -138,7 +168,7 @@ index = FaissIndex.load('corpus.faiss') | Parameter | Default | Effect | |---|---|---| | `nlist` | `sqrt(N)` | Voronoi cells. More cells means finer partitioning and slower training; capped at N | -| `nprobe` | `nlist // 10` | Cells visited per query. **The main recall/speed dial** | +| `nprobe` | `nlist // 10` | Cells visited per query. **The main recall/speed dial.** The default is conservative; `nlist // 30` gave ~5x at 96% Recall@10 on the released corpus | | `pq_m` | 8 | IVF-PQ sub-quantisers. Must divide the embedding dimension; reduced automatically until it does | | `pq_nbits` | 8 | Bits per sub-quantiser |