README.md

# mgxLens

A retrieval-based metagenomic taxonomic profiler. Given short DNA reads (~150 bp) from a
sequencing run, it reports the estimated **relative abundance** of microbial taxa at the
**genus** and **species** level.

**Reference coverage.** The index covers **13,399 distinct species-level SGBs across 3,297 genera**
— the full MetaPhlAn vJan25 SGB reference. Every indexed SGB resolves to a named Linnaean species
(no unnamed placeholder bins among the *indexed* organisms; the naming rate over indexed SGBs is
100%). This means "closed-world" is about organisms *outside MetaPhlAn's reference entirely* — not
gaps within it. An organism that is in MetaPhlAn's reference is representable; one that is not will
be misassigned to its nearest reference genus (see Caveats → closed-world).

---

## 1. What's in the bundle

| File | What it is |
|---|---|
| `encoder.pt` | The neural encoder (seqLens-89M backbone + attention pooling → 256-d embedding). Turns a DNA string into a vector. |
| `index.faiss` | A FAISS vector index of ~2.32M reference marker embeddings. Cosine search via inner product. |
| `index.clades.npy` | Row-aligned to the FAISS index: the SGB (clade) id for each indexed vector. |
| `index.markers.npy` | Row-aligned marker source ids (provenance; not needed for inference). |
| `index.config.json` | Index build parameters (window=150, nwin=5, emb_dim=256, max_length=128). |
| `lineage.tsv` | `clade_id → genus, species, display_name`. Turns a retrieved SGB into a taxon name (display_name is MetaPhlAn-style). |
| `mgx_encoder.py` | The encoder class + `load_encoder()` / `embed()` helpers. Import this; don't reimplement. |
| `MANIFEST.json` | Machine-readable summary + file checksums + caveats. |

**Encoder and index are from the same training run** (`index.config.json.source_ckpt` ==
`encoder.pt`). Do not mix this encoder with any other index or vice-versa — the vectors would be
meaningless.

---

## 2. Environment

```bash
pip install torch transformers faiss-cpu numpy
# GPU optional: faiss-gpu + a CUDA torch build speed up large batches, but faiss-cpu is fine for
# moderate throughput. The encoder runs on CPU or GPU; GPU is ~10-50x faster for embedding.
```

The encoder downloads its base model (`omicseye/seqLens_4096_512_89M-at-base-multi`) from
HuggingFace on first load via `transformers`, so the serving host needs network access on first
run (or pre-cache the HF model). `trust_remote_code=True` is required for that base model.

---

## 3. How it works (the inference recipe)

For each input read:

1. **Window** the read to 150 bp (if longer, take a centered 150 bp window; if ≤150, use as-is).
2. **Embed** it with the encoder → a 256-d L2-normalized vector.
3. **Search** the FAISS index for the top-k nearest reference vectors (k=25). Because vectors are
   L2-normalized and the index is inner-product, the score **is** cosine similarity ∈ [-1, 1].
4. **Gate**: if the top-1 cosine is below a threshold τ (start with **0.66**), the read is
   *abstained* (dropped, counted as "unclassified"). Otherwise it's assigned.
5. **Assign**: the read's taxon is the clade (`index.clades.npy`) of its top-1 hit. Map that clade
   to genus/species via `lineage.tsv`.
6. **Aggregate** across all assigned reads: count reads per genus (and per species), normalize to
   fractions → the abundance profile. Report the unclassified fraction separately.

That's the whole pipeline. It's nearest-neighbor retrieval, not a trained classifier.

---

## 4. Reference implementation (copy-paste starting point)

```python
import numpy as np, faiss, torch
from collections import defaultdict
from mgx_encoder import load_encoder, embed   # ships in this bundle

BUNDLE = "."                       # path to the bundle dir
K = 25                             # neighbors per read
GATE = 0.66                        # top-1 cosine threshold; below -> abstain
WINDOW = 150

# ---- load once at startup ----
device = "cuda" if torch.cuda.is_available() else "cpu"
model, tok, cfg = load_encoder(f"{BUNDLE}/encoder.pt", device)
index = faiss.read_index(f"{BUNDLE}/index.faiss")
clades = np.load(f"{BUNDLE}/index.clades.npy", allow_pickle=True)   # (N,) SGB id per vector
lineage = {}                       # clade_id -> (genus, species, display_name)
with open(f"{BUNDLE}/lineage.tsv") as fh:
    next(fh)
    for line in fh:
        cid, g, s, disp = line.rstrip("\n").split("\t")
        lineage[cid] = (g, s, disp)

def center_window(seq, w=WINDOW):
    if len(seq) <= w:
        return seq
    s = (len(seq) - w) // 2
    return seq[s:s+w]

def profile(reads, batch_size=4096):
    """reads: list of DNA strings. Returns genus/species abundance + unclassified fraction.
    Species keys are MetaPhlAn-style display names ('Escherichia coli', or
    'Escherichia sp. (SGB123)' for unnamed genome bins)."""
    genus_counts, species_counts = defaultdict(float), defaultdict(float)
    n_total = len(reads); n_assigned = 0
    for i in range(0, n_total, batch_size):
        batch = [center_window(r) for r in reads[i:i+batch_size]]
        ml = cfg["max_length"] if "max_length" in cfg else 128
        Z = embed(model, tok, batch, device, max_length=ml).astype("float32")
        D, I = index.search(Z, K)              # D=cosine sims, I=index rows
        for r in range(len(batch)):
            if D[r, 0] < GATE:                 # gate: abstain
                continue
            clade = str(clades[I[r, 0]])       # top-1 hit's clade
            g, s, disp = lineage.get(clade, ("", "", ""))
            if not g:
                continue
            n_assigned += 1
            genus_counts[g] += 1
            if disp:
                species_counts[disp] += 1      # display name, not raw s__ / SGB id
    def norm(d):
        tot = sum(d.values()) or 1.0
        return {k: v/tot for k, v in sorted(d.items(), key=lambda x: -x[1])}
    return {
        "genus": norm(genus_counts),
        "species": norm(species_counts),
        "n_reads": n_total,
        "n_assigned": n_assigned,
        "unclassified_fraction": 1.0 - (n_assigned / n_total if n_total else 0.0),
    }

# ---- example ----
if __name__ == "__main__":
    reads = ["ACGT..."]            # your reads (from a FASTQ)
    print(profile(reads))
```

**Reading FASTQ**: reads come from `.fastq`/`.fastq.gz`. Use `pysam`, `Bio.SeqIO`, or a
2nd/4th-line parser to get the sequence strings, then pass the list to `profile()`. Paired-end
R1/R2 can both be fed as independent reads.

---

## 5. Serving it (FastAPI sketch)

```python
from fastapi import FastAPI, UploadFile
app = FastAPI()
# model/index/lineage loaded once at module import (see section 4)

@app.post("/profile")
async def profile_endpoint(file: UploadFile):
    reads = parse_fastq(await file.read())     # your FASTQ parser -> list[str]
    return profile(reads)                      # JSON: {genus:{...}, species:{...}, ...}
```

Load the model, index, and lineage **once at startup** (they're large — encoder 342 MB, index
2.3 GB). Never reload per request. The index holds ~2.3 GB in RAM; size the host accordingly
(≥8 GB RAM recommended, ≥16 GB comfortable).

---

## 6. Input / output contract

**Input**: DNA reads as strings (from FASTQ), ~150 bp each. Non-ACGT characters are tolerated by
the tokenizer but degrade the embedding; upstream QC/trimming (e.g. fastp) is assumed.

**Output** (JSON):
```json
{
  "genus":   {"escherichia": 0.29, "pseudomonas": 0.24, "...": 0.0},
  "species": {"Escherichia coli": 0.21, "Pseudomonas aeruginosa": 0.18, "...": 0.0},
  "n_reads": 200000,
  "n_assigned": 94459,
  "unclassified_fraction": 0.528
}
```
Abundances are fractions summing to ~1.0 within each level (over assigned reads). All 13,399
indexed SGBs have real species names, so species keys are Linnaean names (the `Genus sp. (SGBxxxx)`
fallback only appears if the index is rebuilt to include unnamed bins).

---

## 7. Tunable knobs

| Knob | Default | Effect |
|---|---|---|
| `GATE` (top-1 cosine) | 0.66 | Higher = stricter, more reads abstained, fewer false assignments (but see Caveats — it does **not** reliably reject novel organisms). |
| `K` (neighbors) | 25 | Only top-1 is used for assignment here; k>1 matters if you switch to k-NN vote aggregation. |
| aggregation | top-1 vote | Simplest. Alternatives (k-NN vote, similarity-weighted) exist in the research code but top-1 is the documented default. |

---

## 8. Caveats — read before building product on this

- **Proof of concept.** Not accuracy-tuned. The goal of this handoff is to enable infrastructure
  and integration work, not to ship a validated diagnostic.
- **Genus is the reliable level; species is approximate.** At 150 bp, reads from different species
  in the same genus are near-indistinguishable to the encoder — the model reliably gets the
  *genus* right but frequently assigns the wrong *species within that genus*. Report species
  output, but do not treat it as trustworthy.
- **Species naming.** All 13,399 indexed SGBs have real Linnaean species names, so normal output
  shows names like "Escherichia coli", not bin ids. (The `display_name` logic still falls back to
  `Genus sp. (SGBxxxx)` for unnamed bins — this only matters if the index is ever rebuilt to
  include the unnamed SGBs that exist elsewhere in the MetaPhlAn taxonomy; the shipped index has
  none.)
- **Closed-world — this is the big one.** The model can only assign a read to a taxon that exists
  in its reference index. An organism that is **not** in the reference will still be confidently
  assigned to its nearest reference genus and reported as present. There is effectively **no
  novelty rejection** (abstain rate ~0% in evaluation), and this behavior was **never quantified**.
  **Do not deploy this against arbitrary/open real-world samples** and present the output as
  complete or trustworthy without adding a novelty-rejection layer. It is validated only in the
  closed-world setting where the sample's organisms are known to be in the reference.
- **Evaluated on one mock community (Zymo D6300).** Broader validation was not done.

---

## 9. Provenance

- Encoder: `enc_attention_bg1.best_genus.pt` — seqLens-89M backbone, attention pooling, 256-d
  projection, trained with contrastive InfoNCE + background negatives, checkpoint selected on
  real-Zymo genus accuracy.
- Index: `index_bakeoff_mw_notest` — leakage-fixed (held-out test markers excluded).
- Base model pulled at runtime: `omicseye/seqLens_4096_512_89M-at-base-multi`.

See `MANIFEST.json` for file checksums and the machine-readable summary.