evaluation/README.md

# Model evaluation scripts

These scripts score sequences from a FASTA file using checkpoints bundled with this repository under `weights/`. ESM pipelines also read **fair-esm** `.pt` caches: one dict per sequence, named `{fasta_header}.pt` (see `extract_embeddings/README.md`). Headers must match the FASTA IDs you evaluate.

## `weights/` layout (repo-relative defaults)

Training outputs for the shipped models mirror the trainers in `models/` and live here so that a fresh Hub download resolves paths without edits:

| Subdirectory | Trained by | Used by evaluator |
|--------------|------------|---------------------|
| `weights/nn_one_hot/` | `models/nn_one_hot.py` | `evaluate_nn_one_hot.py` |
| `weights/nn_mean_pertoken_esm/` | `models/nn_mean_pertoken_esm.py` | `evaluate_nn_mean_pertoken_esm.py` |
| `weights/rf_one_hot/` | `models/rf_one_hot.py` | `evaluate_rf_one_hot.py` |
| `weights/rf_mean_pertoken_esm/` | `models/rf_mean_pertoken_esm.py` | `evaluate_rf_mean_pertoken_esm.py` |

**Composition logistic regression** is **not** pre-bundled under `weights/`. Train with `models/lr_sequence_composition_baseline.py`, then pass the training output directory via `--model-dir` to `evaluate_lr_sequence_composition.py` (or symlink it to `weights/composition_lr/`).

### What each bundle contains

Shipped checkpoints are trimmed to **inference essentials** plus reproducibility metadata. Full per-file inventory: [`weights/README.md`](../weights/README.md).

**Required for all evaluators:**

- Serialized model weights (`.h5`, `.joblib`, or `model_weights/` for the one-hot NN)
- `*_validation_temperature.csv` — temperature scaling factor
- `*_validation_validation_diagnostics_summary.csv` — F1-optimal threshold

**Also shipped for reproducibility:** `model_parameters.json`, `random_seed.txt`, `best_hyperparameters.csv` (one-hot models), `*_hyperparameters.csv` (ESM models).

Training-time performance diagnostics (ROC/PR/calibration CSVs, confusion matrices, Optuna trial tables, PNG plots) are **not** in the repo.

**One-hot neural net (`nn_one_hot`):** Keras checkpoints under `model_weights/` (`one_hot_model_full_model.h5` and/or `one_hot_model_weights.h5`, `one_hot_model_architecture.json`), plus calibration CSVs at the run root.

**ESM Keras heads (`nn_mean_pertoken_esm`):** `mean_embeddings_model_weights.h5`, `per_token_embeddings_model_weights.h5`, matching `*_scaler.npy` and `per_token_embeddings_pca.npy`, plus validation CSVs in `weights/nn_mean_pertoken_esm/`.

**ESM embedding caches** for inference are **not** included in `weights/`. After extraction (`extract_embeddings/`), place or symlink:

- `weights/nn_mean_pertoken_esm/mean_embeddings_pt/` — `{header}.pt` with `mean_representations`
- `weights/nn_mean_pertoken_esm/per_token_embeddings_pt/` — `{header}.pt` with `representations[layer]`

Same layout applies to `weights/rf_mean_pertoken_esm/`. Pass `--emb-mean` / `--emb-token` if your `.pt` trees live elsewhere.

**One-hot Random Forest (`rf_one_hot`):** `one_hot_rf_model.joblib` plus calibration CSVs.

**ESM Random Forest (`rf_mean_pertoken_esm`):** `mean_embeddings_rf_model.joblib`, `per_token_embeddings_rf_model.joblib`, scalers/PCA `.joblib` when applicable, plus calibration CSVs.

**Composition logistic regression:** expects `composition_lr_model.joblib` (clf + scaler bundle), `composition_lr_validation_temperature.csv`, and `composition_lr_validation_diagnostics_summary.csv` from training. No embeddings or fixed sequence length — features are computed from the full amino-acid sequence (Biopython required).

## Overrides

Paths are built from `_REPO_ROOT` (repository root: parent of `evaluation/`). **`--model-dir`** overrides **`MODEL_CHECKPOINT_DIR`**. **`--emb-mean`** / **`--emb-token`** override the default `.pt` roots for ESM evaluators only.

## Install

From the repository root:

```bash
pip install -r evaluation/requirements.txt
```

## Shared options

- **`--csv`**: Optional spreadsheet with an ID column (`variant`, `sequence_id`, or `id`) and a **`label`** or **`target`** column with values `highFRET` / `lowFRET`. When present, metrics are computed on the intersection of FASTA IDs and CSV rows.
- **`--no-detailed-metrics`**: Use with `--csv`. Skips writing diagnostic CSVs (ROC points, reliability bins, per-threshold tables, etc.) and the long classification report. A short summary is still printed.

### Variable sequence length

| Evaluator | Behaviour |
|-----------|-----------|
| **One-hot NN / RF** | Shorter sequences are **zero-padded** to the model `target_length`. **RF:** length is taken from `one_hot_rf_model.joblib` (`n_features / 20`), not from a stale `model_parameters.json`. Longer sequences use the **N-terminal** `target_length` residues by default (`--long-sequences truncate`). Use `--long-sequences skip` to drop long sequences (legacy behaviour). Predictions CSV includes `original_length`, `truncated`. |
| **ESM mean head** | Length-invariant (mean-pooled vectors); no extra padding in the eval script. |
| **ESM per-token head** | Token matrices are **padded or truncated** to the token count implied by the trained PCA input size (or `--target-tokens`). With **early PCA** training, inference uses PCA input dim ÷ 1280 (not scaler dim). Predictions include `per_token_n_tokens`, `per_token_truncated`. |
| **Composition LR** | Uses the **full** sequence (canonical AAs only for feature math). Rows with undefined composition features are skipped (`skipped_sequences.csv`). Optional `--include-features` adds gravy / pI / charge / complexity columns to the predictions table. |

## Scripts

| Script | Model type | Default checkpoint bundle |
|--------|------------|---------------------------|
| `evaluate_nn_one_hot.py` | One-hot → Keras MLP | `weights/nn_one_hot/` |
| `evaluate_rf_one_hot.py` | One-hot flattened → sklearn RF | `weights/rf_one_hot/` |
| `evaluate_nn_mean_pertoken_esm.py` | ESM `.pt` → Keras softmax head(s) | `weights/nn_mean_pertoken_esm/` + embedding subdirs |
| `evaluate_rf_mean_pertoken_esm.py` | ESM `.pt` → sklearn RF | `weights/rf_mean_pertoken_esm/` + embedding subdirs |
| `evaluate_lr_sequence_composition.py` | Sequence composition → sklearn LR | training output dir via `--model-dir` |

### ESM scripts: `--embedding-type`

Both ESM-backed evaluators accept:

- **`mean`**: Uses **`EMBEDDING_PT_MEAN_DIR`** (or `--emb-mean`).
- **`per_token`**: Uses **`EMBEDDING_PT_TOKEN_DIR`** (or `--emb-token`).
- **`both`** (default): Both heads — both embedding directories apply.

Examples:

```bash
# One-hot NN (defaults to weights/nn_one_hot/)
python evaluation/evaluate_nn_one_hot.py --fasta seqs.fa --output out/

# Composition LR (pass training output directory)
python evaluation/evaluate_lr_sequence_composition.py \
  --fasta seqs.fa --output out/ --model-dir /path/to/composition_lr_baseline/

# ESM-Keras mean head only
python evaluation/evaluate_nn_mean_pertoken_esm.py \
  --fasta seqs.fa --output out/ --embedding-type mean

# Per-token with explicit embedding roots
python evaluation/evaluate_nn_mean_pertoken_esm.py \
  --fasta seqs.fa --output out/ --embedding-type per_token \
  --emb-token /path/to/per_token/pt_dir
```

Outputs are written under `--output`: prediction tables (`one_hot_predictions.csv`, `predictions.csv`, etc.) and, unless `--no-detailed-metrics` is set, metric artifacts named like `{mean_embeddings|per_token_embeddings}_*` when labels are supplied.

## Layout

- `eval_metrics.py` — shared calibration/curve helpers and supervised metric reporting used by all entry points.
- `eval_metrics.py` also provides one-hot pad/truncate helpers (`encode_aa_sequences`, `prepare_sequences_from_fasta`) for `evaluate_*_one_hot.py`. `one_hot_encode.py` re-exports those for backward compatibility.
- `esm_embed_util.py` — per-token ESM pad/truncate to match trained feature size.