Single-owner dataset statistics: pure stats functions, unified transform init, on-disk stats cache

[stefaanhessmann]

Follow-up to the dataset refactor (#781), executing the statistics refactor plan (REVIEW_FIXES.md items S1/D3). Statistics become a pure function of (base dataset, train indices) with a single owner and one query interface, and are persisted next to the split file.

What changed

1. Pure-function statistics — calculate_stats / estimate_atomrefs take an optional indices parameter and build their own raw (transform-free, split-free) view internally. "Statistics are computed on raw data" is now the functions' contract, enforced in one place; the shallow-copy trick in the provider is deleted.
2. One initialization path — transforms are initialized through a single initialize(stats) taking any object with get_stats/get_atomrefs (the provider or the datamodule). The Transform.datamodule() hook is hard-deleted; AtomisticModel.initialize_transforms passes the datamodule as the stats source. The strict no-estimation atomref case is an explicit estimate=False mode of the atomref query, replacing the side-channel atomrefs= argument.
3. Disk cache next to the split file — computed stats and estimated atomrefs are persisted to <split>_stats.npz, keyed by a fingerprint of the train partition, read-modify-written while holding the same inter-process lock as split creation (held across miss→compute→write, so concurrent DDP ranks compute each entry at most once). On by default; stats_file=None opts out. Dataset-provided atomrefs are never persisted. Unreadable cache files fall back to recompute-and-overwrite.

The cornerstone test (commit 1) ports the review's end-to-end S1 check into the suite: RemoveOffsets and a late-initialized AddOffsets postprocessor must hold the identical, independently recomputed raw training mean. It passes on the base branch and on every commit of this PR.

Commits

Ten planned commits (tests-first, each leaving the suite green — verified by running the full suite at every commit), plus black formatting, two hardening fixes from a pre-PR review pass (lock held across compute; corrupt-cache tolerance), and two commits from a post-implementation audit: a stats-file extension fix (a user-supplied stats_file without .npz was written by np.savez under the appended extension but read back at the verbatim path, so the cache silently never hit) and a literal regenerated-split invalidation test.

Known limitations (deliberate)

• The splitting/stats lock file remains cwd-relative (inherited from the existing split lock); processes with different working directories sharing one stats file take different locks.
• The fingerprint covers (dataset length, train partition) — regenerating the database in place with identical length would not invalidate persisted stats, same as it would not invalidate split.npz. Delete the stats file or set stats_file=None in that case.
• Persisted stats entries are scalar per (property, flags) — matching calculate_stats, which reduces each property to scalar mean/std.

Test plan

• pytest tests/: 86 passed, 3 skipped (base branch: 60 passed, 3 skipped); suite verified green at each individual commit.
• black --check . clean with the pinned black==24.4.2.
• New tests: e2e offset consistency incl. late postprocessor init (S1 regression), raw-view guarantee of the stats functions, provider pure interface, strict atomref mode, fingerprint determinism/invalidation, disk cache read-instead-of-recompute, split-regeneration invalidation, opt-out, never-persist dataset atomrefs, corrupt-cache recovery, non-.npz stats-file round-trip.

🤖 Generated with Claude Code

[stefaanhessmann]

Review summary (post-implementation audit)

Audited the branch against REFACTOR_PLAN_STATS.md (10 planned commits → 13 actual, mapped 1:1 plus black + two hardening fixes) and against the base tip 9ca9186.

Verified:

• Plan fidelity — indices contract with the raw view built inside calculate_stats/estimate_atomrefs; provider constructed from (base dataset, train indices) with the shallow-copy trick gone; strict estimate=False atomref mode; single initialize(stats) path with the datamodule() hook fully deleted (no leftovers in transform/base.py, model/base.py, or docs); fingerprint over (dataset length, sorted train indices); persistence keyed by query triple + fingerprint under the shared SPLITTING_LOCK; stats_file=None opt-out; dataset-provided atomrefs never persisted.
• Concurrency — no nested acquisition of the inter-process lock on any call path (_load_partitions releases it before the provider exists), so the POSIX release-drops-all hazard is not reachable; torn np.savez writes are only observable after a crash and are handled by the corrupt-file fallback; compute-under-lock gives at-most-once semantics per entry across ranks. The cwd-relative lock path remains a documented limitation.
• Tests — every Testing Decisions bullet is covered, all asserting external behavior only (no test touches internal caches). Full suite green at each of the 13 commits (63 → 84 tests, now 86 with the audit commits); black --check clean with the pinned 24.4.2.

Found & fixed (pushed to this branch):

• 8e5fa6e's parent e4eac8e — a user-supplied stats_file without .npz was written by np.savez under the appended extension but read back at the verbatim path: the cache silently never hit. Normalized in the provider + regression test.
• 8e5fa6e — the plan's "regenerated split → fingerprint mismatch → recompute" decision is now tested literally at the datamodule level (delete split.npz, new partition, assert recompute).

Remaining review notes are style-level judgement calls (mirror-image RemoveOffsets/AddOffsets initialize bodies — pre-existing shape, net reduced; provider combines stats + persistence concerns) — none blocking.

Verdict: ready to undraft.

🤖 Generated with Claude Code