evoPLM-Tree is a prefix-LM protein language model for SARS-CoV-2 spike/RBD evolution, trained with an optional tree-aware conditioning signal (phylogenetic features encoded as discrete tokens).
This repository contains the full pipeline: raw genome download and QC, phylogenetic-tree feature extraction, ancestor-descendant sequence pairs (trajectories) construction, model training, sequence generation, deep-mutational-scanning (DMS) validation, and figure generation.
configs/ Tokenizer config (protein_tree_tokenizer.json)
src/evoplm/ Installable package: tokenizer (ProteinTreeTokenizer)
scripts/
preprocessing/ 01-12 Python steps: NCBI download -> QC -> Nextclade -> dedup ->
RBD extraction -> split1 -> UShER tree features -> trajectory pairs ->
filter trivial pairs (launchers live under cluster/slurm/preprocessing/)
train_prefix_lm.py Prefix-LM training (rbd/spike x seq-only/tree-aware)
generate_sequences.py Stratified-sample generation + mutation-recovery metrics
diagnose_conditioning.py Masked-vs-visible-prompt teacher-forced loss ratio
dms/ DMS validation: table prep, mutation-probability inference,
enrichment/correlation metrics, plots
figures/ Generation-metric, hotspot, and structure-mapping figures
cluster/slurm/
preprocessing/ Launchers for scripts/preprocessing/01-12 (same numbering)
train/ train_{rbd,spike}_{seqonly,treeaware}.sh (GPU)
generate/ generate_{rbd,spike}_{seqonly,treeaware}.sh (GPU; sampled generation)
diagnose/ diagnose_conditioning_{rbd,spike}_{seqonly,treeaware}.sh (CPU;
masked/visible loss-ratio conditioning-dependence eval)
dms/ prepare_dms_table.sh, infer_mutation_probabilities.sh, evaluate_dms.sh,
plot_dms_comparison.sh, plot_weighted_enrichment.sh
(see "DMS validation" below)
figures/ plot_generation_metrics.sh, plot_mutation_hotspots.sh,
plot_conditioning_context.sh, plot_observed_vs_predicted_frequency.sh,
map_hotspots_to_structure_{rbd,spike}_{pred,true}.sh, render_structures.sh
data/
dms/ Raw DMS assay scores (from MaveDB)
pdb/ Reference structure (6VSB) for hotspot-to-structure mapping (from PDB)
raw/ NOT INCLUDED (regenerated by steps 01-05) — see data/raw/README.md
processed/ NOT INCLUDED (regenerated by steps 06-11) — see data/processed/README.md
usher/ NOT INCLUDED (external public data) — see step 10 below
trajectories/ NOT INCLUDED — see "Data" below
results/
generation/split1/ Sampled generation outputs + hotspot-to-structure mappings
context_dependence/ diagnose_conditioning.py stdout logs (masked/visible loss ratios) +
rendered plots
dms/ DMS validation tables/metrics
models/ NOT INCLUDED — see "Pretrained models" below
This repo covers only split1 (first three Omicron clades -> later Omicron descendants).
Two conda environments are needed.
Main environment (training, generation, evaluation, DMS validation, most figures):
conda env create -f environment.yml
conda activate evoplmThis installs the evoplm package itself (editable, from src/) via the -e . entry in
environment.yml / pyproject.toml.
Tree-feature environment (evoplm-bte, used only by
cluster/slurm/preprocessing/10_add_usher_tree_to_splits.sh and
cluster/slurm/preprocessing/11_gen_trajectories_tree_aware.sh): requires UShER / bte
(big-tree-explorer) Python bindings for mutation-annotated-tree (MAT) parsing.
conda env create -f environment_bte.ymlRun in order (01 -> 12). submit_nextclade_pipeline.sh orchestrates the sharded Nextclade step
(03) and its merge (04); run the others directly with sbatch. Each numbered step reads
the previous step's output (01 -> 02 -> ... -> 12); see the script bodies for exact filenames.
| # | Script | What it does |
|---|---|---|
| 01 | download_ncbi.sh |
Download SARS-CoV-2 genomes from NCBI Datasets |
| 02 | prefilter_split.sh |
Keep near-complete genomes by length, shard for Nextclade |
| 03 | nextclade_array.sh |
Run Nextclade per shard (submit via submit_nextclade_pipeline.sh) |
| 04 | merge_nextclade.sh |
Merge per-shard Nextclade TSVs and spike FASTAs |
| 05 | augment_nextclade.{sh,py} |
Join Nextclade output with NCBI accession/date/geography metadata |
| 06 | filter_metadata.{sh,py} |
QC filters; writes spike.parquet / .tsv / .fasta |
| 07 | rm_duplicates.{sh,py} |
Collapse duplicates, add dup_count |
| 08 | extract_rbd.{sh,py} |
Derive RBD sequence from the deduplicated spike set |
| 09 | splits_with_cutoff.{sh,py} |
Build the split1 train/test partition, exclude recombinants |
| 10 | add_usher_tree_to_splits.{sh,py} |
Map accessions to UShER MAT leaves, compute tree features |
| 11 | gen_trajectories_tree_aware.{sh,py} |
Build ancestor-descendant sequence pairs for train/test |
| 12 | filter_trivial_trajectories.{sh,py} |
Drop pairs with no sequence change / no nt substitutions, per task |
sbatch cluster/slurm/train/train_rbd_treeaware.sh # and rbd_seqonly / spike_seqonly / spike_treeawareEach launcher calls scripts/train_prefix_lm.py --task {rbd,spike} --split split1 [--use_tree_tokens] ...
and writes to models/{seqonly,treeaware}/{task}/{split}/maxT.../final.
sbatch cluster/slurm/generate/generate_rbd_treeaware.sh # sampled generation + mutation-recovery metrics
sbatch cluster/slurm/diagnose/diagnose_conditioning_rbd_treeaware.sh # masked/visible loss-ratio evalgenerate_*.sh calls scripts/generate_sequences.py, producing
results/generation/split1/{rbd,spike}_{seqonly,treeaware}_sample2000.{csv,json}.
diagnose_conditioning_{rbd,spike}_{seqonly,treeaware}.sh calls scripts/diagnose_conditioning.py,
which computes teacher-forced loss twice per example — once with the conditioning prompt visible,
once with its attention-mask zeroed — and reports the ratio (ratio masked/normal). Each launcher tees its full stdout to
results/context_dependence/{rbd,spike}_{seqonly,treeaware}.txt; plot_conditioning_context.py
(see "Figures" below) parses the ratio masked/normal line from those four logs directly.
Pipeline: prepare_dms_table.py (parse raw scores from MaveDB into results/dms/dms_scores/<dataset>_table.csv,
launcher: prepare_dms_table.sh, run once per included MaveDB dataset ID: urn_mavedb_00000670-a-1 and urn_mavedb_00000671-a-1)
-> infer_mutation_probabilities.py (per-site mutation probabilities from a trained checkpoint,
launcher: infer_mutation_probabilities.sh) -> evaluate_dms.py (merge + correlation/enrichment
metrics, launcher: evaluate_dms.sh) -> plot_dms_comparison.sh / plot_weighted_enrichment.sh.
evaluate_dms.sh reads every *_table.csv file from results/dms/dms_scores/ in a single run and writes one subdirectory per dataset directly under results/dms/
(results/dms/urn_mavedb_00000670-a-1/, ...000671-a-1/), each with its own
{seqonly,treeaware}_merged_probs_dms.csv / _top3_mutations.csv / metrics_summary.csv, plus a
single combined results/dms/metrics_summary_all_datasets.csv written at the end of that same run
(concatenating every dataset it just processed).
plot_dms_comparison.sh calls plot_dms_comparison.py once per dataset subdirectory.
plot_generation_metrics.py, plot_mutation_hotspots.py, plot_observed_vs_predicted_frequency.py
plot from results/generation/split1/*sample2000.csv (launchers of the same name).
plot_conditioning_context.py reads the masked/visible loss ratios from
results/context_dependence/{rbd,spike}_{seqonly,treeaware}.txt (produced by
diagnose_conditioning_*.sh, see "Generation & evaluation" above; launcher of the same name;
writes to results/context_dependence/plots). map_hotspots_to_structure.py projects
hotspot frequencies onto data/pdb/6VSB.pdb
(writes a mapping CSV + a PyMOL .pml script) — one launcher per {rbd,spike} x {pred,true}
combination (map_hotspots_to_structure_{rbd,spike}_{pred,true}.sh). render_structures.sh
(self-contained) renders the resulting .pml
files to PNGs via PyMOL.
Included: data/dms/ (raw DMS scores from MaveDB), data/pdb/6VSB.pdb (reference structure from PDB).
Not included: data/trajectories/*.parquet (train/test ancestor-descendant pairs for both
rbd and spike, see data/trajectories/README.md) — regenerate via the preprocessing pipeline
above (step 12), or download from Zenodo. data/usher/ (external UShER MAT
snapshot needed for steps 10-11, see data/usher/README.md).
see models/README.md for the expected directory layout. Can be downloaded from Zenodo.
MIT — see LICENSE.