Skip to content
ndnngPublic

About

ARCHER: Amortized cross-specimen pose estimation for cryo-electron microscopy

Topics

Resources

Stars

4 stars

Watchers

1 watching

Forks

Latest commit

 

History

4 Commits

Folders and files

Repository files navigation

ARCHER

Amortized cross-specimen pose estimation for cryo-electron microscopy.

Paper: bioRxiv 2026.08.21.746234

Pose assignment is normally solved anew for every dataset: iterative refinement restarts from scratch, and learned estimators store the molecule in their weights. ARCHER learns the operation once, across 3,330 protein structures, by supplying the reference volume as an input rather than absorbing it into the parameters — so a single set of weights poses particles of proteins it has not seen, with no retraining.


Install

git clone https://github.com/ndnng/ARCHER.git
cd ARCHER
pip install -r requirements.txt

PyTorch with CUDA is required — the model is not usable on CPU at realistic particle counts.

The checkpoints are hosted outside git. Fetch them with:

python scripts/download_weights.py --list     # what exists
python scripts/download_weights.py            # the published model
python scripts/download_weights.py --all      # every checkpoint
python scripts/download_weights.py --verify   # checksum what you have

Each download is checked against the sha256 in weights/MANIFEST.json, and a file that fails is deleted rather than kept — a truncated checkpoint loads fine and gives wrong answers, which is worse than no checkpoint.

Quick start

# 1. weights (they do not ship with the code)
python scripts/download_weights.py

# 2. an example dataset -- see docs/DATASETS.md
python scripts/download_data.py list
pip install cesped && python scripts/download_data.py cesped --targets 10409

# 3. check the geometry conventions on your machine BEFORE trusting any number
python scripts/selftest.py

selftest.py downloads one small EMDB map (128³, ~6 MB), runs the checks against it and cleans up; it needs no dataset and no checkpoint. To include the real-particle check, give it a star file and a map of the same specimen: python scripts/selftest.py --star particles.star --particles_dir <dir> --ref_map ref.mrc.

Run this first, and read what it prints. Convention errors — rotation origin, CTF sign, shift sign — are the dominant failure mode in this kind of pipeline, and every one of them is invisible in a test phase. verify runs the checks that cross that boundary.

Inference -- Posing your own particles

scripts/infer.py is the entry point. It takes particles and a reference map, and writes a star file with rlnAngleRot/Tilt/Psi filled in plus a per-particle confidence column.

python scripts/infer.py \
    --star particles.star --particles_dir <dir holding the .mrcs> \
    --ref_map reference.mrc \
    --iters 3 --refine \
    --out poses.star

You need three things:

particles a star file and its .mrcs stacks, with CTF columns and an optics table. Angles are not required — that is the output.
a reference map ARCHER is reference-conditioned: it answers "where does this particle sit relative to this map". Any map of the same molecule works — a published EMDB entry, a map from an earlier run, an ab-initio volume — but past roughly 25 Å it stops carrying usable orientation information. If you have no map at all, build one from a subset you have already processed with scripts/make_reference.py below.
a sensible box the D=64 crop has to land in the 2–10 Å/px range the model was trained on. A 128–256 px box at 1–2 Å/px does; infer.py prints the sampling it ends up with and warns if it falls outside.

Three levels, in increasing cost: the default is one forward pass (the network's argmax over the SO(3) grid); --refine adds the curvature step that moves poses off the grid; --iters N alternates reconstructing from the current poses and re-scoring against that reconstruction. The iterations need no labels and close most of the gap to the classical ceiling.

Getting a reference map

If your particles already carry poses, scripts/make_reference.py backprojects them directly:

python scripts/make_reference.py \
    --star particles.star --particles_dir <dir> --n 20000 --out reference.mrc

If you would rather use RELION, either of these produces a usable reference.mrc:

# A. from particles that already have poses -- one reconstruction, no refinement
relion_reconstruct --i particles.star --o reference.mrc --ctf --sym C1 --pad 2

# B. from scratch -- 3D initial model, then use its output as the reference
relion_refine --i particles.star --o initial/run --denovo_3dref \
    --ctf --K 1 --sym C1 --iter 100 --tau2_fudge 4 \
    --particle_diameter <diameter in A> --zero_mask --flatten_solvent \
    --pool 10 --j 8
cp initial/run_it100_class001.mrc reference.mrc

relion_reconstruct needs rlnAngleRot/Tilt/Psi in the star file; relion_refine --denovo_3dref does not, which is the option to use when you have no poses at all. Either output can go straight into --ref_map. Only the low-resolution content matters here — ARCHER stops reading a reference past roughly 25 Å — so an initial model is enough; it does not need to be refined.

A worked example, end to end

EMPIAR-10409 with nothing but what download_data.py gives you. Half 1 builds the reference, half 0 gets posed against it, so no particle is scored against a map it helped build:

CESPED=$ARCHER_CESPED_ROOT/10409

# a reference from half 1's consensus poses
python scripts/make_reference.py \
    --star $CESPED/particles_1.star --particles_dir $CESPED \
    --n 20000 --out ref_half1.mrc

# pose half 0 against it
python scripts/infer.py \
    --star $CESPED/particles_0.star --particles_dir $CESPED \
    --ref_map ref_half1.mrc \
    --n 4000 --iters 3 --refine --shuffled_control \
    --reconstruct recon.mrc --out poses.star

Because this star file happens to carry consensus angles, infer.py reports how far its predictions land from them — an agreement number, not an accuracy, since those angles are themselves an estimate. On your own data without angles that line is simply absent.

Always pass --shuffled_control on a new dataset. It repeats the run against a phase-scrambled reference, which has the same power spectrum and no structure. If your real run does not beat it, the model is not reading your images and nothing downstream means anything. It is the cheapest way to catch a convention mismatch in your own files.

What is in here

archer/
  model.py           the reference-conditioned contrastive classifier
  slicing.py         central-slice extraction and Wiener backprojection, one convention
  conventions.py     geometry, CTF and shift conventions -- the single source of truth
  so3_grid.py        HEALPix SO(3) grid, nearest-grid assignment, measured spacing
  episodes.py        cross-protein episodes: clean projections + real CTF + real-PSD noise
  refine.py          SO(3) geometry and the curvature-based local refinement step
  scoring.py         per-shell normalised correlation against a volume
  symmetry.py        symmetry-aware geodesic error and point-group handling
  augment.py         appearance and pose-equivariant geometric augmentation
  volume_ops.py      per-volume normalisation, exact 24-fold cube rotations
  diffusion_head.py  optional denoising channel (see "the diffusion channel" below)
  score_net.py       U-Net score network used only by that channel
  cryodrgn_io.py     read cryoDRGN-format particles into these conventions

scripts/
  infer.py           POSE YOUR OWN PARTICLES -- the entry point for new data
  make_reference.py  build a reference map from particles that already have poses
  train.py           training
  download_weights.py / download_data.py    fetch checkpoints and datasets
  eval_arms_cesped.py / cesped_official_eval.py / classical_global_search.py
                     benchmark reproduction: our arms, CESPED's metric, the classical baseline
  junk_rejection.py / noise_sweep_denoiser.py / benchmark_*.py
                     the paper's ablations and throughput measurements

benchmarks/          how every external comparison in the paper was run -- cryoPARES,
                     CryoFastAR, cryoDRGN, cryoFM, the classical matched filter, and
                     CESPED's own evaluator. Scripts only, no results.

weights/             MANIFEST.json (the checkpoints themselves live outside git)
data/                measured noise models (small; the datasets are not here)
assets/              benchmark metadata
docs/                dataset preparation
pipeline.py          single entry point: verify / train / pose / refine

Released weights

These are the models to use — the longest-trained checkpoint of each arm:

file input steps
weights/archer_base_diffusion_3M.pt 3-channel (raw + phase-flipped + denoised) 3,542,232
weights/archer_base_ctrl_4M.pt 2-channel (raw + phase-flipped) 4,162,845

scripts/download_weights.py with no arguments fetches the two-channel model; add --only archer_base_diffusion_3M.pt for the three-channel one, which is the stronger model in the low-SNR regime real particles occupy (see below).

Two earlier checkpoints stay available via --all or --only — archer_base_2M.pt (2,000,000 steps) and archer_base_diffusion.pt (1,348,797). They are the same recipe stopped sooner, and the numbers in the paper were measured with archer_base_2M.pt.

All four are inference checkpoints: the Adam moments are stripped, which is two thirds of the file. Re-export your own with scripts/export_checkpoint.py.

They are not in this repository — scripts/download_weights.py fetches them and verifies the checksums recorded in weights/MANIFEST.json. They are hosted in a public Drive folder: https://drive.google.com/drive/folders/1B8LFFAOSahmVM2OvHBeqbNCqe7Gum4KF, which you can also browse by hand; drop a file into weights/ yourself and --verify will checksum it.

The diffusion channel

--input_mode raw_flip_diff adds a jointly-trained denoiser whose output image becomes a third input channel. Denoising enhances low SNR, which is dominant in experimental particles.

Inference is the same command — the denoiser is restored from the checkpoint automatically, so only --ckpt changes:

python scripts/download_weights.py --only archer_base_diffusion_3M.pt

python scripts/infer.py \
    --ckpt weights/archer_base_diffusion_3M.pt \
    --star particles.star --particles_dir <dir> \
    --ref_map reference.mrc \
    --iters 3 --refine \
    --out poses.star

To train one:

python scripts/train.py \
    --noise_psd data/noise_model_D64.pt \
    --ctf_star <a star file supplying real defocus parameters> \
    --outdir runs/my_diff_run \
    --size base --input_mode raw_flip_diff \
    --diffusion_head --diff_head_base_ch 64 --diff_head_wiener 0.15 \
    --hp_order 3 --n_particles 96 --steps 2000000 --lr 1e-4 \
    --snr_min 0.005 --snr_max 0.075

--diff_head_null shuffles the denoised channel across the batch, keeping its statistics and its compute but destroying its correspondence to the particle — the control that says whether a gain comes from denoising or from the extra capacity.

Reproducing the evaluation

The external comparisons — cryoPARES, CryoFastAR, cryoDRGN, cryoFM, the classical baseline and CESPED's official metric — each have their run scripts under benchmarks/, with a README that says plainly which comparisons favour which method and why. bash benchmarks/smoke.sh checks they all run without the external tools installed.

Our own arms:

# our learned and classical arms on one target
python scripts/eval_arms_cesped.py --target 10409 --outdir runs/eval_10409

# the matched-filter baseline, on the same particles and reference
python scripts/classical_global_search.py \
    --star $ARCHER_CESPED_ROOT/10409/particles_0.star \
    --particles_dir $ARCHER_CESPED_ROOT/10409 \
    --ref_map ref_half1.mrc --out runs/classical_10409.json

# confidence vs constructed junk (AUROC)
python scripts/junk_rejection.py --target 10409 --out runs/junk_10409.json

# pooled FSC across targets, once the per-target runs exist
python scripts/compute_pooled_fsc.py --runs_dir runs --out runs/pooled_fsc.json

# CESPED's own metric -- needs a RELION install and a validated mask, see --help
python scripts/cesped_official_eval.py --target 10409 --relion_bin <path> --mask_cache <path>

Training

The released model's exact configuration, which the defaults do not reproduce (they run a short 20k-step job on a coarser grid):

python scripts/train.py \
    --noise_psd data/noise_model_D64.pt \
    --ctf_star  <a star file supplying real defocus parameters> \
    --outdir    runs/my_run \
    --size base --input_mode raw_flip \
    --hp_order 3 --n_particles 96 --n_train_templates 512 \
    --n_heldout 100 --steps 2000000 --lr 1e-4 \
    --snr_min 0.005 --snr_max 0.075

Every released checkpoint carries the arguments it was trained with, so you can always read the truth off the file rather than trusting this block:

python -c "import torch; print(torch.load('weights/archer_base_2M.pt', \
    map_location='cpu', weights_only=False)['args'])"

Needs the D=64 corpus (docs/DATASETS.md). The split is by structure ID, not by particle: a model that has seen any particle of a protein has seen its frame.

Capacity is not the binding constraint — a 69M-parameter model scores 0.701× the 19M one over 16 matched checkpoints — so --size base is the reported configuration.

Citation

@article{archer2026,
  title   = {ARCHER: Amortized cross-specimen pose estimation for cryo-electron microscopy},
  author  = {Nguyen, Nhan and Pham, Bao},
  journal = {bioRxiv},
  year    = {2026},
  doi     = {10.64898/2026.08.21.746234},
  url     = {https://doi.org/10.64898/2026.08.21.746234}
}

License

MIT, see LICENSE. The datasets are not covered by it: EMPIAR and EMDB entries carry their own terms, per entry, and none are redistributed here.

About

ARCHER: Amortized cross-specimen pose estimation for cryo-electron microscopy

Topics

Resources

Stars

4 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages