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 (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:

>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

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

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

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

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

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:

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