Architecture

prima-bench scores model predictions against held-out truth. The library never imports a model framework — predictions come in as EvalData, scores go out as EvalResult.

The code lives under src/prima_bench/. Each subdirectory owns one layer of the pipeline. Read the per-directory README.md for details.

Top-level layout

src/prima_bench/
├── core/         Fundamental types: EvalData, Contract, EvalResult, MetricResult
├── registry/     In-memory registries for metrics, tasks, datasets, benchmarks
├── io/           Hashing, atomic writes, cache paths, downloads, JSON serialize
├── metrics/      Dual-mode Metric base class + device/stats helpers
├── transforms/   Shared numerical helpers (kNN) that metrics call directly
├── genome/       Reference-genome FASTA access + pure sequence-window helpers (VEP)
├── datasets/     Dataset descriptors + bucket artifact URIs (one subdir per dataset)
├── tasks/        Task definitions: contract, metric panel, split preparers (one subdir per task)
├── benchmarks/   Benchmark records (a Benchmark groups Tasks)
├── runner/       The `pb.Evaluation` entry point — feeds a task's metric panel
├── cli/          `prima-bench` click commands (discovery + docs)
├── docs/         Markdown card autogeneration from registry metadata
└── sequences.py  `pb.get_sequences` — consumer helper that builds VEP model inputs

VEP sequence helpers

The variant-effect tasks need reference-genome sequence windows as model inputs. That support lives in two layers, kept off the scoring path:

• genome/ — internal. reference.py downloads/opens a UCSC FASTA (via io/download.py); sequence_extract.py is pure window math (N-padding, allele substitution, ref-base checks).
• sequences.py — the public pb.get_sequences(...) orchestration that turns variant rows into ref/alt sequences for a consumer’s model.

The runner never imports either; they are a consumer-side convenience, not part of pb.Evaluation.

Concepts

A Task is a panel of Metrics that should map to a biological question. All Metrics within a Task must share the same data format. Tasks have a default Dataset that they can be evaluated against. Tasks can be evaluated against multiple Datasets.

A Dataset is versioned set of named Artifacts. A Dataset can be used to evaluate multiple Tasks. An Artifact is an s3:// or gs:// URI to a file or directory that is downloaded from a remote location via a bucket-agnostic boto3 client (see io/); there is no gcloud dependency.

A Dataset exposes one or more Splits, keyed by name. Conventional split names are train, val, test, and solution.

• train exposes model inputs and labels that models are trained or finetuned on.
• val (optional) is a labelled validation split in the same format as train. It is a dataset-level artifact: a dataset may ship it and a recipe may write it, but it need not be wired into a task or exposed to consumers — and a dataset whose upstream defines no validation split should not invent one. (The NT benchmark ships no val: its upstream defines only train/test, and the Pleiades source’s validation is a byte-identical copy of test.)
• test exposes model inputs only
• solution is never exposed to the consumer but has corresponding ground truth outputs for each item in test.

Where censoring lives — two valid patterns. The test/solution split can be derived either at eval time by the SplitPreparer (e.g. cell_type_annotation carves a NaN-censored test out of a golden solution) or recipe-side, where the recipe writes an already-censored test plus a per-item solution and the preparer is a thin mapper. The NT benchmark (sequence_classification) uses the recipe-side pattern: the recipe ships test with its label censored to null and a solution/labels.parquet keyed by a row primary key (read_id); predictions and solution pair on that key. Either way the privacy invariant holds — solution is a registered artifact the preparer fetches internally and pb.fetch_data never returns.

A SplitPreparer is the per-(task, dataset) adapter that composes splits from a Dataset’s artifacts and ships the per-split loaders that turn paths into EvalData. One preparer per (task, dataset) pair — the same dataset can be wired to many tasks by registering more preparers.

Consumer API

import prima_bench as pb

# downloads bytes from remote location to local cache
train_path, test_path = pb.fetch_data(task="perturbation_prediction")

# Optional: invoke the registered loader for this (task, split) .
train = pb.load(train_path, task="perturbation_prediction", split="train")
test  = pb.load(test_path,  task="perturbation_prediction", split="test")

model.fit(train)
predictions = model(test)  # EvalData

run = pb.Evaluation(task="perturbation_prediction")
run.update(predictions=predictions)
result = run.compute()

Invariants to Maintain:

• pb.fetch_data always returns (train_path, test_path) — no per-task polymorphism. Tasks without a separate train file (e.g. unsupervised batch_integration) return the test path twice.

• solution is never reachable from consumer code. The bucket has it, pb.Evaluation loads it internally, no pb.fetch_data call exposes its path.

• How predictions align with solution is a per-task Contract concern. Tasks pick the alignment strategy that fits the modality:

  • perturbation_prediction matches by row position over a shared gene space (censoring held-out rows in the consumer-visible test is the preparer’s job).
  • batch_integration matches by cell_id via metadata.index; predictions can be in any order and the metric joins on the key.

  In both cases, predictions only need to carry what the consumer can produce — labels live on the solution side.

Coordinate conventions

Off-by-one errors are silent correctness bugs in variant effect prediction. Get this right at every boundary.

The rule across prima-bench: a variant position (pos) is 1-based; an index, interval, or array slice is 0-based half-open.

• pb.get_sequences takes 1-based pos (VCF / human-genetics convention). avsec variant_ids (chr3_319804_C_G_hg38) are 1-based. Internally it converts once (pos - 1) before slicing the FASTA.
• variant_index (where the variant lands inside the returned window) is 0-based.
• BED exports (e.g. the TECH-4199 liftover BED) are 0-based half-open (BED standard).

We read FASTA with pyfaidx (not pysam), but both return identical bases for identical coordinates — the library is never the source of an off-by-one; the handoff convention is. Bridging a coordinate from a pysam workflow:

⚠️ .pos means different things on different pysam objects: VariantRecord.pos is 1-based, but AlignedSegment.pos is 0-based. Never assume .pos; check the object type.

When locating a variant inside a read, map read→reference coordinates via the CIGAR (get_reference_positions / get_aligned_pairs) — never add a read-sequence offset to reference_start (soft-clips/indels/splices break it).

get_sequences is variant-centric: it requires ref and alt and raises ValueError unless the FASTA matches one of them. A bare read coordinate is not a variant. The match check is a backstop against wrong-frame coordinates — but in homopolymers a neighbor base can coincidentally match, so convert correctly; do not rely on the guard.

Separation of concerns

Each layer changes for a different reason and depends only on the layer above.

Recipe (external)
    │  writes immutable bytes in a publication-shaped layout
    ▼
Dataset           datasets/<name>/__init__.py
    │  maps  name@version → {artifact: remote_uri}
    ▼
SplitPreparer    tasks/<task>/split_preparers/<dataset>.py
    │  per (task, dataset): artifacts → {train, test, solution} paths
    │                       + per-split loaders (Path → EvalData)
    ▼
Runner           runner/evaluation.py
    │  contract.validate → feed metric panel → EvalResult
    ▼
Consumer         pb.fetch_data / pb.load / pb.Evaluation

Motivation:

• Recipes describe a publication, not an API. They write whatever layout is natural for the source paper. They live outside the library so third-party buckets (e.g. OP-bio) can be wrapped without re-curation.
• Datasets carry only identity and artifact URIs. They don’t know what a “train split” means — that’s a task-level concept.
• SplitPreparers are the (task, dataset) adapter. Same dataset, two tasks? Two preparers, no changes to the dataset descriptor. Where censoring, fold selection, or split derivation lives.
• The Runner never sees bytes or URIs — it operates on EvalData. It does not own domain logic; metric functions do.