Skip to main content

Fast population-genetics PCA from scratch

Make a Patterson PCA over raw, unindexed VCF text as fast as possible in Python with NumPy and SciPy only, scored by accuracy-gated, memory-weighted speedup over a full scan on four cohorts the agent has not seen.

Final score per graded run
051015202530
  • Claude Opus 5
  • GPT-5.6 Sol
  • Gemini 3.7 Flash
  • Kimi K3
  • Grok 4.6

Principal component analysis is a routine early step in population genetics: it summarises the ancestry structure of a cohort in a few axes, and those axes are used everywhere from quality control to correcting association studies. On a real cohort the genotypes arrive as a multi-gigabyte plain-text VCF with no index, and reading every byte, decoding every genotype and standardising every marker is correct and slow. This task hands the agent a correct, naive implementation, eight CPUs, no GPU, NumPy and SciPy, and four hours, and asks how much faster the same computation can be made without changing what it computes.

A correct full scan of the file is the anchor and scores 1.0. Anything above that has to come from reading less of the file, decoding it in bulk, choosing the right linear-algebra orientation for the cohort's shape, and using BLAS well, while staying accurate on structure that is subtle, files that are messy, and a cohort whose sample count makes the obvious matrix not fit in memory.

The task

The deliverable is one program, pca <vcf> <k> <out.tsv>, that reads a plain-text multi-sample VCF and writes the leading kk principal component scores for every sample. The object is the Patterson HWE-standardised PCA over every eligible marker. Writing xijx_{ij} for the genotype dosage of sample ii at marker jj and pjp_j for the marker's alternate-allele frequency,

zij=xij2pj2pj(1pj),z_{ij} = \frac{x_{ij} - 2p_j}{\sqrt{2p_j(1-p_j)}},

with missing calls imputed to the marker mean and no marker filtering of any kind. The denominator is what makes this a population-genetics PCA rather than a plain covariance PCA, and one of the grader's probes checks for exactly it.

A PCA is only defined up to rotation, sign and scale within its subspace, and that is how it is scored: the grader compares the span of the submission's columns with the span of the full scan's, never the entries. Only the standard library, NumPy, SciPy and the program's own modules may be imported, with no native code, compilers or child executables; check applies the same allowlist the grader applies. The starting point in /app/src/pca is a correct, naive full scan. It scores 0.04: it runs out of time on the largest fold and out of memory on the widest one.

Environment

The agent works in an isolated container with no general internet access, on eight CPUs with 32 GB of memory and four hours on the clock. /app/src is a git repo and the tree at HEAD is what gets submitted, so only committed work counts. The public package the grader runs is in the environment in full: the Patterson math, the full-scan anchor, the accuracy metric, the cohort generator, the fold catalogue, the identity probes and the import allowlist. The fast reference is not, and its speed is the number to aspire to.

The graded folds are four synthetic cohorts with planted population structure, regenerated from a fresh seed on every grade. The files differ from grade to grade; the shapes, the structure and the requested rank do not, and the agent can generate any of them locally at full or reduced size.

foldshapewhat it tests
continental500 samples, 500k markers, six populations, about 1 GBstrong structure in a large file: reading a fraction of it wins
subtle500 samples, 150k markers, small Fststructure that needs wide coverage: undersampling loses axes
messy500 samples, 120k markers, variable-width FORMAT, missing calls, messy records, CRLFthe parser
biobank64,000 samples, 2,000 markersthe sample Gram does not fit; the marker-space solve does

Three tools do the work. check runs the allowlist, the output contract and the accuracy metric on small generated cohorts of each kind, free and unlimited. bench runs the grader's own rule on development folds at a quarter of the graded size and prints a dev_score, and it appends every attempt to results.tsv and the experiment journal. submit seals the tree at HEAD with a code snapshot and sends it for grading. Local runs and scored submissions are both unlimited; a grade takes about ten minutes.

README.md86 lines
# Working in this environment

Everything lives under `/app`:

```
/app/src/pca              your program; the tree at git HEAD is what you submit
/app/task/pcabench/       the public code the grader runs: the Patterson math and the full-scan
                          anchor (full_scan.py), the accuracy metric (subspace.py), the cohort
                          generator (generate.py), the fold catalogue (folds.py), the identity
                          probes (probes.py), the import allowlist (library_scan.py)
/app/task/grade.py        the grader itself, runnable locally on folds you generate
/app/tools/               check, bench, submit, clock
/app/experiments/         your lab notebook (history.jsonl)
/app/submissions/         your submission record, managed by `submit`; never edit by hand
```

Python 3.12 with NumPy and SciPy, eight CPUs, no GPU, no network beyond the leaderboard.

## The program

`pca <vcf_path> <k> <out_path>` reads a plain-text, unindexed, multi-sample VCF and writes a TSV
with the exact header `sample_id<TAB>PC1<TAB>...<TAB>PCk`, one row per sample in `#CHROM` order,
exactly `k` finite columns, exit zero. The object is the leading-k sample score subspace of the
Patterson HWE-standardised genotype matrix over every eligible marker:

- eligible markers: polymorphic biallelic SNVs, REF and ALT each one of A/C/G/T, any case; skip
  indels, multiallelic, symbolic and monomorphic records;
- `GT` is the first FORMAT subfield, `/` or `|` separated; haploid `0` and `1` are dosages 0
  and 2; anything missing, partial or malformed is a missing call;
- `z = (x - 2p) / sqrt(2p(1 - p))` with `p` the alternate-allele frequency over called genotypes
  and missing calls imputed to the marker mean.

A PCA is defined up to rotation, sign and scale within its subspace, and that is exactly how it
is scored: the grader compares the span of your columns with the anchor's, never entries.
The output file exists before you run and is the only writable path besides `TMPDIR`; write it
in place. Only stdlib, NumPy, SciPy and your own modules may be imported; no native
extensions, no compilers, no other interpreters, no child executables (Linux `fork` workers are
fine). `check` applies the grader's allowlist scan.

## The folds

Four synthetic cohorts, regenerated from a fresh seed on every grade (`/app/task/pcabench/folds.py`):

| fold | shape | what it tests |
|---|---|---|
| continental | 500 samples × 500k markers, about 1 GB | strong structure in a large file: reading a fraction of it wins |
| subtle | 500 × 150k, small Fst | structure that needs wide coverage: undersampling loses axes |
| messy | 500 × 120k, variable-width FORMAT, missing calls, messy records, CRLF | the parser |
| biobank | 64,000 × 2,000 | the sample Gram does not fit in memory; the marker-space solve does |

Generate any of them yourself: `python3 -m pcabench.generate continental src/dev/continental.vcf --seed 1`
(from `/app/task`), full size or `--scale 0.25`. The truth sidecar has the population labels.

## Scoring

Per fold: accuracy against the full-scan PCA of the same file must clear 0.9, else the fold is
0; then `score = full_scan_time / your_time × min(1, memory_budget / your_peak_rss)`, medians of
three interleaved timed runs each of you, the full scan and the fast reference, in the same
sandbox with the same limits. The task score is the mean over folds times the gate
factors (allowlist; an HWE-normalisation probe that separates a Patterson fit from a plain
covariance PCA; a coverage probe that a few-hundred-marker fit cannot resolve). The full scan is
1.0 by definition; the fast reference lands well above it; nothing is clipped.

## Local loop

```
/app/tools/check                                   # allowlist + contract + accuracy on small cohorts
/app/tools/bench --label <idea> > run.log 2>&1     # the grader's rule on dev folds, quarter size: dev_score:
/app/tools/submit --name <idea>                    # the tree at git HEAD, graded on the board
/app/tools/submit --list                           # every submission and its public score (refreshes pending grades)
/app/tools/submit --wait <id>                      # block until that grade is in
/app/tools/clock
```

A grade takes about ten minutes and only graded submissions count at the end: submit your last
candidate at least fifteen minutes before the wall, and confirm it with `submit --list`.

`bench` at `--scale 1.0` is the full-size measurement and takes about as long as a grade.
A run that exceeds fifteen times the full scan's own time on a fold (never under three minutes) scores 0 on it; the other folds still count.

## Experiment discipline

`/app/src` is a git repo; only what is committed is submitted. One idea per commit, the checked-
out commit is your best-so-far, `git reset --hard HEAD~1` rejects a failed idea. `bench` appends
to `/app/src/results.tsv` and `/app/experiments/history.jsonl`.
Figure 1. The agent's workspace at /app, with the files it is given. src/pca starts as the naive full scan in a fresh git repo, and task/pcabench is the public package the grader runs. The fast reference and the grader's own fold draws exist nowhere in the environment. The task instruction is delivered as the agent's prompt rather than as a file, and is on the Prompt tab.

Sampling markers is the intended lever, and the instruction says so. What it also says is that sampling correlated with allele frequency, missingness, record width or position changes the object, and the probes are built to catch it.

Evaluation

A submission is graded on a fresh draw of the four folds. On each fold the full scan computes the truth and its own time, and the submission's output is compared with it. That accuracy is a subspace agreement rather than an entrywise error: for each reference PC that stands above the noise in the anchor's spectrum, the grader measures how much of that direction lies inside the span of the submission's first kk columns, takes the geometric mean over those structured PCs, and rescales so that a random subspace reads 0 and one containing every structured direction reads 1. Dropping any real axis collapses the product. The gate is 0.9, the worst a genuine Patterson PCA scores; below it the fold is 0.

Past the gate the submission, the full scan and the fast reference are each timed three times, interleaved, under the same limits. Writing tanchort_{\mathrm{anchor}} and tsubt_{\mathrm{sub}} for the median wall times, MM for the submission's peak memory and BB for the fold's budget,

sfold=tanchortsubmin ⁣(1,BM).s_{\mathrm{fold}} = \frac{t_{\mathrm{anchor}}}{t_{\mathrm{sub}}} \cdot \min\!\left(1, \frac{B}{M}\right).

The budgets are 2.0, 1.0, 2.5 and 4.0 GB for the four folds, about twice the reference's measured peaks. A run that exceeds fifteen times the full scan's own time on a fold, never under three minutes, or crashes, or fails the output contract, scores 0 on that fold and the other folds still count.

Two identity probes follow, each a small plain cohort with planted structure. The HWE probe tells a Patterson fit from a plain covariance fit; the coverage probe plants structure so subtle that only thousands of markers resolve it. Each is collapsed to a factor that reads 1 for a genuine fit. The task score is

S=14foldssfoldgatesg,S = \frac{1}{4}\sum_{\mathrm{folds}} s_{\mathrm{fold}} \cdot \prod_{\mathrm{gates}} g,

the mean over folds rather than the geometric mean, so that a submission which fails one fold still registers its progress on the others. The full scan is 1.0 by definition and nothing is clipped. Public and private are the same number: every grade is its own fresh draw of the folds, so there is no visible half to overfit.

For orientation, at calibration the full scan took 12.4, 3.8, 14.3 and 28.2 seconds on the four folds. The fast reference ran 3.9, 1.2, 0.66 and 7.9 times faster, for a task score of 3.41. It is slower than the full scan on the messy fold, where its sampling path has to stream variable-width records, and that is the fold where a submission can beat it outright. An independent, correct, unsampled full scan scores 0.24, because the anchor is itself a fast full scan with a batched decoder. Every cheat in the calibration set, from a plain covariance PCA to a 200-marker fit to a forbidden import, scores 0 on the axis it cheats.