FAISS ANN search backend for proteogram similarity

README.md

@@ -350,6 +350,20 @@ 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 *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.
 
 ---
 
@@ -535,7 +549,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) |

docs/faiss_search.md

@@ -0,0 +1,188 @@
+# Approximate Nearest Neighbour Search with FAISS
+
+*Applies to `scripts/v2/measure_similarity_v2.py` and `proteogram.v2.FaissIndex`.*
+
+> **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
+
+`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 and recall
+
+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.
+
+`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:
+
+| `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 |
+
+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.
+
+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.
+
+Two settings to avoid:
+
+- **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.
+
+Index build time is negligible and is not part of this tradeoff: 0.16 s.
+
+### Why random test vectors are not a proxy
+
+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
+
+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.** 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 |
+
+`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.

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']
+           'HierarchicalTripletLoss', 'HierarchicalPKSampler', 'SCOPE_LEVELS',
+           'FaissIndex']