ESMRank

A learn-to-rank approach for protein variant effect prediction

ESMRank is a pipeline for ranking protein variants according to their predicted functional impact.
The method frames variant effect prediction as a learning-to-rank problem, leveraging protein language models and sequence-derived features to prioritize mutations based on fitness, activity, or stability.

Installation

Requirements

• Python 3.10+
• Linux / macOS
• A CUDA GPU (the ESM-2 feature step needs one)

Install with pip

Install a torch build that matches your CUDA setup first (see https://pytorch.org), then:

pip install git+https://github.com/gambalab/ESMRank.git

This provides the esmrank command (and python -m esmrank).

Download the pretrained models

The package does not bundle the models (~830 MB, hosted on Zenodo: https://zenodo.org/records/18773439). Fetch them once with:

esmrank download-models

They are stored in ~/.cache/esmrank/model (or --dest DIR). If the installation fails the downloaded archive is kept, so re-running does not download it again; an archive you already have can be used with esmrank download-models --archive ESMRank.tar.gz. At run time the model directory is looked up in this order: --model-dir, $ESMRANK_MODEL_DIR, ./model, then the cache directory above.

Install from source (development)

git clone https://github.com/gambalab/ESMRank.git
cd ESMRank
python -m venv ESMRank_venv && source ESMRank_venv/bin/activate
pip install -e ".[test]"
esmrank download-models
pytest

requirements.txt still lists the exact versions the pipeline was developed against; pyproject.toml declares compatible ranges.

Usage

From Python

import pandas as pd
import esmrank

# make sure the pretrained files are available (downloads them only if missing)
model_dir = esmrank.ensure_models()

# a protein sequence instead of a FASTA file
seq = "MSEQNNTEMTFQIQRIYTKDISFEAPNAPHVFQKDWQPEVKLDLDTASSQLADDVYEVVLRVTVTASLG"
df = esmrank.rank_variants(seq, sub=True)                    # all missense variants
df = esmrank.rank_variants(sequence=seq, alanines=True)      # explicit form, any length

# a DataFrame instead of a CSV file (columns hgvsp, seq_wt, seq_mu; extra columns are kept)
variants = pd.DataFrame({"hgvsp": ["p.Q12V"], "seq_wt": [seq], "seq_mu": [seq[:11] + "V" + seq[12:]]})
df = esmrank.rank_variants(variants)

# files work as well
df = esmrank.rank_variants("protein.fasta", sub=True, model_dir=model_dir)
df = esmrank.rank_variants("variants.csv", csv=True)

# download the models first if they are not there yet
df = esmrank.rank_variants(seq, alanines=True, auto_download=True)

rank_variants returns a DataFrame with columns ids and prediction.

How the input is interpreted:

A sequence needs at least one of sub, indel, alanines, and may only contain the 20 standard amino acids. Any other string that is not a file raises FileNotFoundError, so a mistyped file name is never read as a sequence (use sequence= for short ones). Results are still written to <outdir>/<name>/: name defaults to the file name, the FASTA header, or a short hash of the sequence / DataFrame, and can be set with name="...". The command line only accepts files.

rank_variants and run_pipeline take the same options as the CLI: csv, sub, indel, alanines, outdir, model_dir, esm_model, batch_size, gpu_ids, plus sequence, name and auto_download. On the command line the equivalent of auto_download is esmrank run --download ....

From command line

esmrank run --input INPUT [-csv] [-sub] [-indel] [-alanines]
            [--outdir DIR] [--model-dir DIR] [--download] [--gpus 0,1] [--batch-size N] [--esm-model NAME]

Results are written to <outdir>/<input file name>/ (default output/<name>/): predictions.csv, plus profile.tsv and profile.png in -sub / -alanines mode.

Input modes

ESMRank supports two input modes: CSV mode and FASTA mode.

From the command line the input is always a file. From Python the input can also be given in memory: a pandas DataFrame instead of the CSV file, and a protein sequence string instead of the FASTA file (see From Python for all the details).

CSV mode

In CSV mode, the input must be a CSV file (or, from Python, a pandas DataFrame) containing three mandatory columns:

Example:

hgvsp,seq_wt,seq_mu
p.A123V,MSEQNNTEMTFQIQRIYTKDISFEAPNAPHVFQ...,MSEQNNTEMTFQIQRIYTKDISFEVPNAPHVFQ...

To run in CSV mode:

esmrank run --input variants.csv -csv

From Python you can pass a DataFrame directly, without writing a CSV (extra columns are kept):

import pandas as pd
import esmrank

variants = pd.read_csv("variants.csv")      # or any DataFrame with the three columns above
df = esmrank.rank_variants(variants)        # DataFrame with columns `ids` and `prediction`

FASTA mode

In FASTA mode, the input must be a single wild-type protein sequence in FASTA format (or, from Python, the sequence itself as a string).

The pipeline will automatically generate variants based on the selected mutation types:

• -sub : generate all possible missense substitutions
• -indel : generate all possible insertions and deletions
• -alanines: mutually exclusive to sub and indel, generate only alanines substitutions for each aminoacid

Example FASTA:

>protein_X
MSEQNNTEMTFQIQRIYTKDISFEAPNAPHVFQ...

Example execution:

generating all missenses variants:

esmrank run --input protein.fasta -sub

generating all single single indel

esmrank run --input protein.fasta -indel

generating all missenses and single indel

esmrank run --input protein.fasta -sub -indel

generate alanine scanning ( alanines substitution only )

esmrank run --input protein.fasta -alanines

From Python you can pass the sequence as a string instead of a FASTA file; the same sub, indel and alanines options apply:

import esmrank

seq = "MSEQNNTEMTFQIQRIYTKDISFEAPNAPHVFQKDWQPEVKLDLDTASSQLADDVYEVVLRVTVTASLG"
df = esmrank.rank_variants(seq, sub=True)                  # all missense variants
df = esmrank.rank_variants(sequence=seq, alanines=True)    # alanine scanning

Command-line arguments

The command line only accepts files. Sequence strings and DataFrames are available from Python, see From Python.

The environment variables ESM_GPU_IDS, ESM_WORKER_BATCH_SIZE, ESM_MODEL_NAME and ESM_CHECKPOINT_DIR are still honoured.

Notes

• CSV and FASTA modes are mutually exclusive.
• In FASTA mode, at least one of -sub, -indel or -alanines must be specified.
• Runtime and memory usage scale with protein length and number of generated variants.

License

• Source code: Apache License 2.0, see LICENSE.
• Pretrained model parameters (downloaded with esmrank download-models): ESMRank Model Parameters Non-Commercial Research License, see LICENSE_WEIGHTS.md. They may be used for non-commercial research and educational purposes only; commercial use requires written permission from the authors.
• Third-party resources (Pfam-A HMMs, ESM-2 weights) keep their own licenses.

Citation

If you use ESMRank in your research, please cite:

@article{esmrank,
  title={ESMRank reveals a transferable axis of protein mutational constraint from overlapping variant effect assays},
  author={Riccardo Arnese, Gennaro Gambardella},
  journal={TBD},
  year={TBD}
}