models/README.md

# FRET binary classifiers: model training scripts

This folder contains **command-line training pipelines** for a binary FRET task (positive class **`highFRET`**, negative **`lowFRET`**). Each script fits a classifier on a train/validation split, applies **temperature scaling** and an **F1-optimal threshold** on validation, evaluates on a held-out test set (metrics are **printed to the terminal**), and writes artifacts under `--output`.

**Default (no extra flag):** outputs are **minimal** — serialized model weights (and scalers / PCA arrays for ESM scripts), **`random_seed.txt`** when applicable, **`*_validation_temperature.csv`**, and **`*_validation_diagnostics_summary.csv`** (threshold for downstream evaluation scripts). **Plots (PNG)** and detailed **evaluation CSVs** are **off** unless you pass **`--save-diagnostics`**.

Install dependencies from the repo (see **Requirements** below), then run a script with `python` from your environment.

---

## Requirements

Install packages listed in [`requirements_models_training.txt`](requirements_models_training.txt) (numpy, pandas, scikit-learn, matplotlib, h5py, tqdm, psutil, optuna, joblib, torch, **fair-esm**, tensorflow, biopython).

A **GPU helps** TensorFlow/PyTorch but is not strictly required if datasets are modest in size.

---

## Input data: FASTA and CSV

### FASTA

- Use standard FASTA records; scripts read sequences with **`esm.data.read_fasta`**.
- **The FASTA header (text after `>`)** must equal the identifier used to look up rows in your CSV (**no extra spaces** vs the CSV unless your IDs genuinely include them).

Example:

```fasta
>variant_A001
MSKGEELFTGVVPILVELDGDV...
```

### CSV column names (**important**)

The four scripts **do not all use the same CSV schema**:

| Script | ID column | Label column | Allowed label strings |
|--------|-----------|--------------|------------------------|
| `nn_one_hot.py` | **`sequence_id`** | **`per_token_model_class`** | `highFRET`, `lowFRET` |
| `rf_one_hot.py` | **`variant`** | **`label`** | `highFRET`, `lowFRET` |
| `nn_mean_pertoken_esm.py` | **`variant`** | **`label`** | `highFRET`, `lowFRET` |
| `rf_mean_pertoken_esm.py` | **`variant`** | **`label`** | `highFRET`, `lowFRET` |

If you run **`nn_one_hot.py`** alongside the others, prepare a CSV with **`sequence_id`** and **`per_token_model_class`**, **or** add/duplicate columns so that file matches the script you call.

Rows with labels other than those two strings are skipped.

### ESM embedding files (`*.pt`)

The two **ESM** scripts expect **one mean file and one per-token file per variant**, named after the **`variant`** value:

- **Mean embeddings directory** (`--emb-mean`): for each variant ID `V`, a file **`V.pt`** loaded with PyTorch (`torch.load`). Inside the checkpoint:
  - `mean_representations[emb_layer]` — vector used as the pooled representation.
- **Per-token embeddings directory** (`--emb-token`): **`V.pt`** with:
  - `representations[emb_layer]` — token-level tensor (flattened internally for modeling).

Layer index **`--emb-layer`** defaults to **`33`** (typical **ESM-2** last-layer convention in **fair-esm**); change it only if your `.pt` files were extracted with another layer indexing scheme.

---

## Train / validation / test split (high level)

- **One-hot scripts** (`nn_one_hot.py`, `rf_one_hot.py`): stratified split **80%** train+val / **20%** test; from train, another **20%** is held out as **validation** (calibration / threshold tuning).
- **ESM scripts** (`nn_mean_pertoken_esm.py`, `rf_mean_pertoken_esm.py`): stratified **80%** train / **20%** test; each model internally uses a validation slice from train for tuning (details in script).

Unless you pass **`--seed`**, a seed is sampled and saved to **`random_seed.txt`** in the output directory so runs can be reproduced.

---

## 1. `nn_one_hot.py` — neural network on one-hot sequences

**Idea:** Each sequence is encoded as **L × 20** one-hot positions (canonical amino acids **ACDEFGHIKLMNPQRSTVWY**; unknown residues skipped; truncate/pad to `--target-length`).

**Model:** Keras **MLP** (flatten → dense layers → 2-class softmax). Optional **Optuna** search over width, learning rate, batch size, and epochs.

**Pipeline highlights:** validation **temperature scaling**, **F1-optimal threshold** on calibrated probabilities, then **test** metrics printed to the console; **`--save-diagnostics`** restores full ROC/PR/calibration CSVs and PNGs.

### Required arguments

```bash
python models/nn_one_hot.py \
  --fasta /path/to/sequences.fasta \
  --csv /path/to/labels.csv \
  --output /path/to/run_output/
```

### Useful options

| Flag | Default | Meaning |
|------|---------|--------|
| `--target-length` | 120 | Max sequence length for encoding; may be **auto-increased** to `max(seq_len) + 4` if needed |
| `--n-trials` | 50 | Optuna trials (when enabled) |
| `--no-optuna` | — | Fixed default MLP (64→32 units, 50 epochs, batch 32) |
| `--no-weights-export` | — | Skip extra weight export formats; still saves `one_hot_model_weights.h5` |
| `--seed` | random | Integer seed for NumPy/TensorFlow |
| `--save-diagnostics` | off | PNG plots plus full evaluation CSVs (ROC/PR, calibration files, Optuna tables, training history, `model_parameters.json`) |

### Outputs (selection)

- **Always:** **`one_hot_model_weights.h5`** (and optional **`model_weights/`** exports unless `--no-weights-export`), **`one_hot_nn_validation_temperature.csv`**, **`one_hot_nn_validation_validation_diagnostics_summary.csv`**, **`random_seed.txt`** if `--seed` was omitted.
- **With `--save-diagnostics`:** validation/test diagnostic CSV and PNG prefixes **`one_hot_nn_*`**, **`one_hot_training_history.csv`**, Optuna CSVs when Optuna ran, **`model_parameters.json`**.

---

## 2. `rf_one_hot.py` — random forest on one-hot sequences

**Idea:** Same one-hot tensor as above, **flattened** to a 2D feature matrix for **sklearn `RandomForestClassifier`**.

### Required arguments

```bash
python models/rf_one_hot.py \
  --fasta /path/to/sequences.fasta \
  --csv /path/to/labels_with_variant_label.csv \
  --output /path/to/run_output/
```

### Useful options

| Flag | Default | Meaning |
|------|---------|--------|
| `--target-length` | 120 | Same semantics as NN script (`max(seq_len)+4` bump possible) |
| `--n-trials` | 50 | Optuna trials over RF hyperparameters |
| `--no-optuna` | — | Default RandomForest hyperparameters |
| `--seed` | random | Random state for splitting and forests |
| `--save-diagnostics` | off | PNG plots, test/validation CSVs, **`model_parameters.json`**, Optuna CSVs when used |

### Outputs (selection)

- **Always:** **`one_hot_rf_model.joblib`**, **`one_hot_rf_validation_temperature.csv`**, **`one_hot_rf_validation_validation_diagnostics_summary.csv`**, **`random_seed.txt`** if `--seed` was omitted.
- **With `--save-diagnostics`:** **`one_hot_rf_*`** diagnostics (calibration, ROC, PR, etc.), **`model_parameters.json`**, Optuna CSVs when Optuna ran.

---

## 3. `nn_mean_pertoken_esm.py` — Keras L1 logistic regression on ESM embeddings

**Idea:** Trains **two** separate binary heads: one on **mean** embeddings (**1280**-dim typical for ESM-2), one on **flattened per-token** embeddings (optionally reduced with **PCA**). Each head is essentially **multinomial logistic regression** implemented as a single **Dense(2, softmax)** Keras layer with **L1** regularization and **Adam**; **Optuna** tunes learning rate and L1 strength (**20 trials** per model in code).

Designed for **large** per-token tensors: streaming batch loads, optional **HDF5 disk cache**, **`--train-separately`** to peak less memory, **`--max-samples`** for debugging.

### Required arguments

```bash
python models/nn_mean_pertoken_esm.py \
  --fasta /path/to/sequences.fasta \
  --csv /path/to/labels.csv \
  --emb-mean /path/to/mean_emb_dir/ \
  --emb-token /path/to/per_token_emb_dir/ \
  --output /path/to/run_output/
```

### Useful options

| Flag | Default | Meaning |
|------|---------|--------|
| `--emb-layer` | 33 | ESM representation layer index inside `.pt` |
| `--batch-size` | 500 | Embeddings streaming batch |
| `--max-samples` | all | Cap number of variants (testing) |
| `--use-pca-early` | off | PCA per-token rows during loading |
| `--pca-components` | 1280 | PCA size for per-token when early PCA is used |
| `--use-disk-cache` | off | Write batches to HDF5 (+ temp pickle for IDs) |
| `--cache-dir` | temp | Where to write cache files |
| `--train-separately` | off | Train mean model, free memory, then per-token |
| `--seed` | random | Reproducibility |
| `--save-diagnostics` | off | **`neural_net_comparison_results.csv`**, **`model_parameters.json`**, training history / hyperparameter CSVs, per-model diagnostic CSVs + PNGs on validation and test |

### Outputs (selection)

- **Always:** per-head **`*_embeddings_model_weights.h5`**, **`*_embeddings_scaler.npy`** (and **`*_embeddings_pca.npy`** when used), **`mean_embeddings_validation_temperature.csv`** / **`per_token_embeddings_validation_temperature.csv`**, matching **`mean_embeddings_validation_validation_diagnostics_summary.csv`** and **`per_token_embeddings_validation_validation_diagnostics_summary.csv`**, **`random_seed.txt`** when applicable.
- **With `--save-diagnostics`:** **`neural_net_comparison_results.csv`**, **`model_parameters.json`**, **`mean_embeddings_training_history.csv`**, **`per_token_embeddings_training_history.csv`**, **`*_embeddings_hyperparameters.csv`**, and expanded **`mean_embeddings_*` / `per_token_embeddings_*`** diagnostic files.

---

## 4. `rf_mean_pertoken_esm.py` — random forests on ESM embeddings

**Idea:** Same embedding loading alignment as script (3): **two** RandomForest models (**mean** and **per-token**), optional **PCA** on per-token side, **`StandardScaler`**, Optuna tuning of RF hyperparameters (`--n-trials`, `--no-optuna`).

### Required arguments

```bash
python models/rf_mean_pertoken_esm.py \
  --fasta /path/to/sequences.fasta \
  --csv /path/to/labels.csv \
  --emb-mean /path/to/mean_emb_dir/ \
  --emb-token /path/to/per_token_emb_dir/ \
  --output /path/to/run_output/
```

CLI flags **`--emb-layer`**, **`--batch-size`**, **`--max-samples`**, **`--use-pca-early`**, **`--pca-components`**, **`--use-disk-cache`**, **`--cache-dir`**, **`--train-separately`**, **`--seed`**, **`--save-diagnostics`**, **`--n-trials`**, **`--no-optuna`** mirror the NN ESM script where applicable.

### Outputs (selection)

- **Always:** **`mean_embeddings_rf_model.joblib`** & **`per_token_embeddings_rf_model.joblib`** (plus **`*_embeddings_scaler.joblib`**, **`_*_embeddings_pca.joblib`** when PCA is used), **`mean_embeddings_validation_temperature.csv`**, **`per_token_embeddings_validation_temperature.csv`**, **`mean_embeddings_validation_validation_diagnostics_summary.csv`**, **`per_token_embeddings_validation_validation_diagnostics_summary.csv`**, **`random_seed.txt`** when applicable.
- **With `--save-diagnostics`:** **`rf_esm_comparison_results.csv`**, **`model_parameters.json`**, **`mean_embeddings_hyperparameters.csv`** / **`per_token_embeddings_hyperparameters.csv`**, ROC/PR/confusion/feature-importance CSVs + PNGs.

---

## 5. `lr_sequence_composition_baseline.py` — logistic regression on composition only (Figure 3 baseline)

**Idea:** No embeddings and no one-hot sequence tensor — only **eight** scalar features per variant computed from the amino-acid sequence:

| Category | Features |
|----------|----------|
| Hydrophobicity | `gravy_mean` (mean Kyte-Doolittle index) |
| Isoelectric point | `isoelectric_point` (Biopython pI) |
| Charge composition | `frac_positive`, `frac_negative`, `frac_charged`, `net_charge_ph7` |
| Sequence complexity | `shannon_entropy_aa`, `complexity_normalized` (entropy / log₂20) |

**Model:** `sklearn.linear_model.LogisticRegression` (L2) on **StandardScaler**-normalized features.

**Split / evaluation:** Same as ESM scripts — stratified **80%** train / **20%** test; validation slice from train for **temperature scaling** and **F1-optimal threshold**; **test ROC-AUC** is the headline metric. The script prints an explicit note if test AUC **≥ 0.85** (composition may explain much of the ESM signal).

### Required arguments

```bash
python models/lr_sequence_composition_baseline.py \
  --fasta /path/to/sequences.fasta \
  --csv /path/to/labels.csv \
  --output /path/to/composition_baseline/
```

Use **`--csv-schema esm`** (default: `variant` + **`label`**, same as ESM/RF embedding scripts). The script **always prefers the `label` column** when it exists (not `per_token_model_class`). Override with `--csv-label-col` / `--csv-id-col` if needed.

### Useful options

| Flag | Default | Meaning |
|------|---------|--------|
| `--c` | 1.0 | L2 inverse regularization for logistic regression |
| `--ablation` | off | Per-group test AUC (hydrophobicity, pI, charge, complexity alone) |
| `--save-features` | off | Write `composition_features_per_variant.csv` |
| `--save-diagnostics` | off | ROC CSV/PNG and related exports |
| `--seed` | random | Reproducible split |
| `--perm-n-repeats` | 30 | Shuffle repeats for permutation importance on the test set |

### Outputs (selection)

- **Always:** `composition_baseline_summary.csv`, `composition_lr_model.joblib` (inference bundle for `evaluation/evaluate_lr_sequence_composition.py`), `composition_lr_auc_score.csv`, `composition_lr_coefficients.csv`, `composition_lr_permutation_importance.csv` (ranked test-set permutation importance), `composition_lr_validation_diagnostics_summary.csv`, `composition_lr_validation_temperature.csv`, `composition_baseline_metadata.json`
- **With `--ablation`:** `composition_baseline_ablation.csv`
- **With `--save-features` or `--save-diagnostics`:** `composition_features_per_variant.csv`

Requires **Biopython** (see `requirements_models_training.txt`).

---

## Quick reference: which script should I run?

| You have… | Consider |
|-----------|----------|
| Only FASTA + CSV (**no embeddings**) | **`lr_sequence_composition_baseline.py`** (composition baseline for Figure 3), **`nn_one_hot.py`**, or **`rf_one_hot.py`** — check CSV column requirements. |
| FASTA + CSV + **`variant`-keyed `.pt` mean & per-token** | **`nn_mean_pertoken_esm.py`** (linear / L1 Keras heads) vs **`rf_mean_pertoken_esm.py`** (nonlinear forests). |
| Prefer interpretable linear weights / smaller models | NN ESM script’s logistic heads. |
| Prefer fewer assumptions about geometry of features | RF ESM or RF one-hot. |

---

## Help

Each script exposes defaults via argparse:

```bash
python models/nn_one_hot.py -h
python models/rf_one_hot.py -h
python models/nn_mean_pertoken_esm.py -h
python models/rf_mean_pertoken_esm.py -h
python models/lr_sequence_composition_baseline.py -h
```