Geometric Sketching for Protein Sequences

Select a small, maximally diverse subset from a large set of protein sequences using geometric sketching (Hie et al. 2019, Cell Systems 8, 483–493).

Two embedding modes are supported:

1. K-mer mode — fast, self-contained C pipeline: read a FASTA file, compute k-mer count-sketch embeddings, and sketch.
2. Matrix mode — bring your own embeddings (e.g. from a protein language model): provide a .npy matrix and an IDs file.

In both cases the output is a text file of selected sequence IDs, one per line, in the order they appeared in the input.

---

Quick start

# Build (requires a C11 compiler — gcc, clang, etc.)
cd c/
make

# Select 100 diverse sequences from a FASTA
./bin/geosketch --fasta seqs.fa --M 100 --out sketch.txt

# Check the result
wc -l sketch.txt    # 100
head sketch.txt     # sequence IDs

---

Build

cd c/
make                       # produces bin/geosketch and bin/clustersketch
make clustersketch         # build only the clustering baseline tool
make clean                 # remove binaries

clustersketch links against easel (hmmer-3.4/easel), so ensure HMMER is built and libeasel.a exists. You can override EASEL_DIR if needed:

cd c/
make clustersketch EASEL_DIR=/path/to/hmmer-3.4/easel

The geosketch binary has no external dependencies beyond a standard C11 compiler and the math library (-lm), both of which are available on any Linux / macOS system.

On the FASRC cluster the default gcc module works out of the box:

module load gcc     # if not already loaded
cd c/ && make

---

Usage

FASTA mode (k-mer embedding + geometric sketching)

./bin/geosketch \
    --fasta input.fa \
    --M 100 \
    --out sketch.txt \
    --seed 0

This reads input.fa, computes a k-mer count-sketch embedding for each sequence (default: 4-mers hashed into 512 dimensions), runs geometric sketching to select 100 diverse sequences, and writes their IDs to sketch.txt.

Matrix mode (precomputed embeddings)

./bin/geosketch \
    --matrix embeddings.npy \
    --ids ids.txt \
    --M 100 \
    --out sketch.txt \
    --seed 0

embeddings.npy is a 2D NumPy array of shape (N, D) with dtype float32 or float64. ids.txt has one sequence ID per line (same order as the matrix rows). This mode is intended for protein language model embeddings produced by an external Python script.

Score-based within-cell selection

By default, when a grid cell is sampled, a random point within that cell is chosen. An optional scores file changes this to pick the highest-scoring (or lowest-scoring) point instead:

./bin/geosketch \
    --fasta input.fa \
    --M 100 \
    --out sketch.txt \
    --scores bitscores.txt \
    --score-policy max

bitscores.txt has one floating-point value per line, in the same order as the FASTA sequences.

---

Full option reference

Usage:
  geosketch --fasta <path> --M <int> --out <path> [options]
  geosketch --matrix <path> --ids <path> --M <int> --out <path> [options]

Modes:
  --fasta <path>        Input FASTA file (k-mer embedding mode)
  --matrix <path>       Precomputed embedding matrix (.npy, float32/64)
  --ids <path>          Sequence IDs file (required with --matrix)

Required:
  --M <int>             Sketch size (number of sequences to select)
  --out <path>          Output file (selected IDs, one per line)

K-mer options (--fasta mode only):
  --kmer-k <int>        K-mer length                    [default: 4]
  --dim <int>           Embedding dimension (hash buckets) [default: 512]

Geosketch options:
  --seed <uint64>       Random seed                     [default: 0]
  --alpha <float>       Binary-search tolerance          [default: 0.1]
  --max-iter <int>      Max binary-search iterations     [default: 200]
  --k auto|<int>        Target covering-box count        [default: auto = M]

Score-based within-cell selection:
  --scores <path>       Scores file (one float per line)
  --score-policy <p>    uniform | max | min | max_ties | min_ties
                                                         [default: uniform]
Other:
  --verbose <int>       Verbosity level (0 = silent)     [default: 0]
  --help                Show help message

---

Examples

1. Basic sketch from FASTA

./bin/geosketch --fasta proteins.fa --M 200 --out sketch200.txt --seed 42

2. Reproducibility check

# Same seed → identical output
./bin/geosketch --fasta proteins.fa --M 100 --out a.txt --seed 0
./bin/geosketch --fasta proteins.fa --M 100 --out b.txt --seed 0
diff a.txt b.txt    # no differences

3. Tune k-mer parameters

# 3-mers into 256 dimensions (faster, lower resolution)
./bin/geosketch --fasta proteins.fa --M 100 --out sketch.txt \
    --kmer-k 3 --dim 256

# 5-mers into 1024 dimensions (higher resolution, more memory)
./bin/geosketch --fasta proteins.fa --M 100 --out sketch.txt \
    --kmer-k 5 --dim 1024

4. Score-guided selection (e.g. by HMMER bitscore)

# Pick the highest-bitscore sequence from each grid cell
./bin/geosketch --fasta hits.fa --M 100 --out sketch.txt \
    --scores bitscores.txt --score-policy max

# Break ties randomly among equally-scoring sequences
./bin/geosketch --fasta hits.fa --M 100 --out sketch.txt \
    --scores bitscores.txt --score-policy max_ties --seed 7

5. Use PLM embeddings (matrix mode)

# Python side: save embeddings + IDs
# python save_plm_embeddings.py --fasta proteins.fa \
#     --out-matrix embeddings.npy --out-ids ids.txt

# C side: sketch from the precomputed matrix
./bin/geosketch --matrix embeddings.npy --ids ids.txt \
    --M 100 --out sketch.txt

6. Fewer covering boxes than sketch size (overflow)

# Force 30 boxes but select 100 sequences.
# First round: 1 sample per box (30).  Overflow rounds fill the
# remaining 70 by prioritising cells with fewer sequences left.
./bin/geosketch --fasta proteins.fa --M 100 --k 30 --out sketch.txt

7. Verbose output for debugging

./bin/geosketch --fasta proteins.fa --M 100 --out sketch.txt \
    --verbose 2 2>geosketch.log
# geosketch.log contains binary-search iterations and cell counts

---

How it works

1. Embed — each protein sequence is mapped to a D-dimensional vector. In FASTA mode this is a CountSketch of k-mer frequencies, L2-normalised so sequence length does not dominate the geometry.

2. Translate & scale — the embedding matrix is shifted so all values are non-negative, then divided by the global maximum so everything lies in [0, 1].

3. Plaid covering — a binary search finds a box side length unit such that the number of non-empty axis-aligned grid cells is approximately equal to the target k (default: M).

4. Sample — one cell is drawn uniformly at random, then one point is drawn from that cell (uniformly, or by score). The cell is removed from the round; once all cells have been visited, a new round begins.

5. Overflow — if M exceeds the number of cells, additional samples are drawn by sorting the remaining non-empty cells by size (ascending) and sampling one from each, smallest first. This maximises diversity by avoiding over-sampling from large cells.

---

Memory

Peak memory is roughly N × D × 24 bytes (the embedding matrix in double, plus argsort indices and the grid table in int32).

N (sequences)	D	Approx. peak
100 K	512	~1.2 GB
1 M	512	~12 GB
1 M	256	~6 GB
1 M	1024	~24 GB

These numbers are comfortable on HPC nodes. To reduce memory, lower D or use a high-memory node.

---

Low-compute baseline (jackhmmer + sampling)

Scripts under scripts/ generate baseline sketch lists from a jackhmmer run (run dir must contain hits/hits.fa, hits/hits_scores.txt, hits/hits.sto).

Generate list files (four methods: lowest E-value, random, stratified E-random, stratified E-minE) then build .fa, .sto, .hmm, and *_alistat.txt per list using the shared build step below.

Submit as one SLURM job (recommended):

cd /path/to/Sketching
mkdir -p logs
# Optional: use a different run dir
export RUN=experiments/globins
sbatch scripts/run_baseline_pipeline.sh

The job runs from Sketching/; default run dir is experiments/globins. Logs go to logs/run_baseline.out and logs/run_baseline.err. This only generates sketch list files (no build, no tables).

Or run the same steps by hand (no sbatch):

# From Sketching/; set RUN for your run directory (can use a different run later)
RUN=experiments/globins

for M in 1000 500 100 50; do
  ./scripts/low_compute_baseline.sh "$RUN" $M
done

# Build all sketch dirs and write tables
./scripts/run_build_and_tables.sh

Method 1 writes one list file per M. Methods 2–4 write 100 list files per M (different seeds; stratifiedEminE breaks ties by seed).

---

High-compute baseline (single-linkage MSA clustering)

Cluster hits/hits.sto by percent identity using easel’s esl_msacluster_SingleLinkage() and sample sequences to size M with geosketch-style overflow sampling.

SLURM job (recommended):

cd /path/to/Sketching
mkdir -p logs
RUN=experiments/globins sbatch scripts/run_cluster_pipeline.sh 1000 500 100 50

This writes:

• Random replicates: experiments/<exp>/msacluster_<M>/sketch1.txt ... sketch100.txt
• minE replicates: experiments/<exp>/msacluster_minE_<M>/sketch1.txt ... sketch100.txt

To change score sampling policy or parameters:

RUN=experiments/<exp> sbatch scripts/run_cluster_pipeline.sh 100 \
  --score-policy uniform --replicates 100 --alpha 0.1 --max-iter 200

---

Experiment workflow (recommended)

cd /path/to/Sketching
mkdir -p logs

# 1) Search (put query.fa in experiments/<exp>/ first)
sbatch scripts/search.sh experiments/<exp>

# 2) Baseline list generation
RUN=experiments/<exp> sbatch scripts/run_baseline_pipeline.sh

# 3) Geosketch list generation (array job)
RUN=experiments/<exp> sbatch scripts/run_geosketch_pipeline.sh

# 4) PLM geosketch list generation (array job)
# Requires merged model NPZ files in hits/ESM2_<size>_<dim>.npz.
# If part shards exist (hits/ESM2_<size>_<dim>.part-*.npz), they are merged first.
# Length sorting is enabled by default; this creates/reuses:
#   experiments/<exp>/hits/hits_scores_sorted.txt
RUN=experiments/<exp> sbatch scripts/run_plm_geosketch_pipeline.sh

# 5) High-compute clustering baseline (optional)
RUN=experiments/<exp> sbatch scripts/run_cluster_pipeline.sh 1000 500 100 50

# 6) Build all sketches + write tables + merge timing
RUN=experiments/<exp> sbatch scripts/run_build_and_tables.sh

# 7) Plots
python3 scripts/plot_baseline_results.py --run-dir experiments/<exp>

Outputs are organized under experiments/<exp>/ (hits/, tables/, plots/, times/).

One-command DAG submission (search -> baseline/kmer/plm/cluster -> build)

Use the orchestrator when you want a single command to submit the full workflow and wait for successful completion of the final build/tables step.

cd /path/to/Sketching
mkdir -p logs

# Uses k=3, dims 64/128/256/512/1024 for kmer geosketch
# Uses merged PLM NPZs (hits/ESM2_<size>_<dim>.npz); merged from part shards when needed
RUN=experiments/<exp> M_VALS="1000 500 100 50" REPLICATES=100 \
./scripts/run_full_pipeline.sh

# Submit only (do not block/wait):
RUN=experiments/<exp> M_VALS="1000 500 100 50" REPLICATES=100 \
./scripts/run_full_pipeline.sh --no-wait

PLM embedding script

scripts/plm_embed.py is a config-free, high-throughput embedding writer adapted from PSALM-2/src/sketch_embed.py and intended for ESM2 HF IDs only.

Supported model IDs:

• facebook/esm2_t6_8M_UR50D (320)
• facebook/esm2_t12_35M_UR50D (480)
• facebook/esm2_t30_150M_UR50D (640)
• facebook/esm2_t33_650M_UR50D (1280)

Example:

source activate /n/eddy_lab/Lab/protein_annotation_dl/PSALM-2/psalm2
python3 scripts/plm_embed.py \
  --fasta experiments/<exp>/hits/hits.fa \
  --output experiments/<exp>/hits/ESM2_650M_1280 \
  --model facebook/esm2_t33_650M_UR50D \
  --scores experiments/<exp>/hits/hits_scores.txt \
  --sorted-scores-out experiments/<exp>/hits/hits_scores_sorted.txt

By default, embeddings are length-sorted for throughput. When sorting is enabled and --scores is provided, hits_scores_sorted.txt is created/reused so minE selection stays aligned with embedding ID order.

---

Experiment layout

experiments/<exp>/
├── hits/
│   ├── hits.fa
│   ├── hits.sto
│   ├── hits.hmm
│   ├── hits_alistat.txt
│   ├── hits_scores.txt
│   ├── hits_scores_sorted.txt           # created/reused when PLM embedding sorts by length
│   ├── hits.domtblout
│   ├── hits_clan.txt                    # seq_id<TAB>family clan assignments
│   ├── hits_fe.txt                      # family counts + X + entropy (bits)
│   ├── hits_fe_bins.txt                 # fixed FE categories used by all sketch FE files
│   ├── <model>_<size>_<dim>.part-*.npz  # raw embedding shards from plm_embed.py
│   └── <model>_<size>_<dim>.npz         # merged embedding per model for sketching
├── minE_<M>/ ...                        # includes sketch*.txt and sketch*_fe.txt
├── randomsample_<M>/ ...                # includes sketch*.txt and sketch*_fe.txt
├── stratifiedErandom_<M>/ ...           # includes sketch*.txt and sketch*_fe.txt
├── stratifiedEminE_<M>/ ...             # includes sketch*.txt and sketch*_fe.txt
├── geosketch_kmer_<k>_<dim>_<M>/ ...    # includes sketch*.txt and sketch*_fe.txt
├── geosketch_kmer_<k>_<dim>_minE_<M>/ ...
├── geosketch_plm_<size>_<dim>_<M>/ ...
├── geosketch_plm_<size>_<dim>_minE_<M>/ ...
├── plots/
├── tables/
└── times/

---

References

• Hie, B., Cho, H., DeMeo, B., Bryson, B. & Berger, B. Geometric Sketching Compactly Summarizes the Single-Cell Transcriptomic Landscape. Cell Systems 8, 483–493 (2019). https://doi.org/10.1016/j.cels.2019.05.003