Tom/eDNA_Stream_E01_custom_models
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial release: E1 genus probe package moved out of the E01 workflow

Browse Source 
- src/custom_models torch modules, store, training, metrics, runner and CLI
  (package-relative imports for pip installability)
- test suite (49 tests) + conftest with fast synthetic-store fixtures
- uv-managed dev env (pyproject + CPU torch index), hatchling build
- README: uv for the package, pip install into conda envs for the workflow
This commit is contained in:
Tom Kasper
2026-09-22 13:22:50 +01:00
commit 88e98effad
18 changed files with 5988 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_pipeline.py
+416
View File
@@ -0,0 +1,416 @@
"""Pipeline acceptance tests: store, splits, metrics, training and the CLI."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import polars as pl
import pytest
import torch
from custom_models.cli import main as cli_main
from custom_models.config import (
DataConfig,
ModelConfig,
RunConfig,
TrainConfig,
load_config,
run_id_for,
save_config,
)
from custom_models.data import make_splits, permute_read_labels
from custom_models.metrics import choose_tau_msp, closed_set_report, msp
from custom_models.runner import evaluate_test, evaluate_val, train_run
from custom_models.seed import seed_everything
from custom_models.store import SignalStore, synthetic_store
from custom_models.train import select_device
# ---------------------------------------------------------------- store
def test_synthetic_store_roundtrip(tmp_path: Path) -> None:
out = tmp_path / "s"
synthetic_store(
out_dir=out,
n_genera=3,
reads_per_genus=4,
window_samples=1200,
shard_windows=4,
seed=0,
)
store = SignalStore(out)
assert len(store) == 3 * 4
assert store.window_len() == 1200
meta = store.meta()
assert meta["synthetic"] is True
assert meta["n_genera"] == 3
assert meta["n_windows"] == len(store)
windows = [store.get(i) for i in range(len(store))]
assert all(w.shape == (1200,) and w.dtype == np.float32 for w in windows)
manifest = store.manifest()
assert set(manifest.columns) >= {
"window_id",
"read_id",
"genus",
"species",
"source_run",
"window_index",
"n_samples",
}
assert manifest["window_id"].to_list() == list(range(len(store)))
genera = manifest["genus"].unique().sort().to_list()
assert genera == [f"genus_{i:02d}" for i in range(3)]
def test_store_labels_and_split_determinism(fast_store: Path) -> None:
store = SignalStore(fast_store)
manifest = store.manifest()
selected = ("genus_00", "genus_01", "genus_02")
cfg = DataConfig(genera=selected, min_test_windows_per_genus=0, seed=0)
a = make_splits(manifest, cfg)
b = make_splits(manifest, cfg)
assert all(
np.array_equal(x, y) for x, y in zip((a.train, a.val, a.test), (b.train, b.val, b.test))
)
assert a.genera == list(selected)
# read-level purity: every window of a read shares one split
read_arr = manifest["read_id"].to_numpy()
wid_arr = manifest["window_id"].to_numpy()
split_of = {int(w): 0 for w in a.train}
split_of.update({int(w): 1 for w in a.val})
split_of.update({int(w): 2 for w in a.test})
by_read: dict[str, set[int]] = {}
for wid, read_id in zip(wid_arr, read_arr):
if int(wid) in split_of: # windows of unselected genera have no split
by_read.setdefault(str(read_id), set()).add(split_of[int(wid)])
assert all(len(v) == 1 for v in by_read.values())
# a different data seed re-shuffles the read assignment
c = make_splits(manifest, DataConfig(genera=selected, min_test_windows_per_genus=0, seed=1))
assert not np.array_equal(a.train, c.train)
# each read keeps its genus label; class ids map onto splits.genera
assert len(a.labels) == len(manifest)
assert (a.labels[a.train] >= 0).all()
# ---------------------------------------------------------------- splits
def test_make_splits_min_test_windows_guard(trap_store: Path) -> None:
store = SignalStore(trap_store)
with pytest.raises(ValueError, match="min_test_windows_per_genus"):
make_splits(
store.manifest(),
DataConfig(genera=("trap_00", "trap_01"), min_test_windows_per_genus=10**6),
)
def test_make_splits_holdout_runs_guard(fast_store: Path) -> None:
store = SignalStore(fast_store)
manifest = store.manifest()
# two source_runs, alternating per read: holdout picks news_run_*1 for test
runs = [
"news_run_0" if i % 2 == 0 else "news_run_1"
for i, read_id in enumerate(manifest["read_id"].unique().sort().to_list())
]
split_map = dict(zip(manifest["read_id"].unique().sort().to_list(), runs))
two_runs = manifest.with_columns(
pl.col("read_id").replace_strict(split_map).alias("source_run")
)
assert (
len(make_splits(two_runs, DataConfig(holdout_runs=True, min_test_windows_per_genus=1)).test)
> 0
)
single_run = manifest.with_columns(pl.lit("only_run").alias("source_run"))
with pytest.raises(ValueError, match="holdout_runs"):
make_splits(single_run, DataConfig(holdout_runs=True, min_test_windows_per_genus=0))
def test_permute_read_labels_preserves_counts(fast_store: Path) -> None:
store = SignalStore(fast_store)
manifest = store.manifest()
pd = permute_read_labels(manifest, (), 3)
counts_before = manifest["genus"].value_counts().sort("genus")
counts_after = pd["genus"].value_counts().sort("genus")
assert counts_before.equals(counts_after)
# windows keep their read's label: read -> single genus before and after
for frame in (manifest, pd):
per_read = frame.group_by("read_id").agg(pl.col("genus").n_unique().alias("k"))
assert per_read["k"].max() == 1
pd2 = permute_read_labels(manifest, (), 3)
assert pd.equals(pd2)
with pytest.raises(ValueError, match="at least two reads"):
permute_read_labels(manifest.filter(pl.col("read_id") == manifest["read_id"][0]), (), 0)
# ---------------------------------------------------------------- configs
def test_config_roundtrip(tmp_path: Path) -> None:
cfg = RunConfig(
run_id="cnn_1.0M_s4_g11",
stage="arch_ladder",
store_dir=Path("/data/store"),
data=DataConfig(genera=("A", "B"), val_fraction=0.2, seed=7),
model=ModelConfig(
arch="cnn",
param_budget=1_000_000,
stride=4,
d_model=64,
n_layers=8,
params_realised=912_345,
),
train=TrainConfig(lr=1e-3, seed=3),
out_dir=Path("/runs"),
genera=("A", "B"),
notes="hello",
)
path = tmp_path / "config.json"
save_config(cfg, path)
assert load_config(path) == cfg
def test_run_id_format() -> None:
assert run_id_for("cnn", 1_000_000, 4, 11, "arch_ladder") == "cnn_1.0M_s4_g11"
assert run_id_for("linatt", 5_000_000, 8, 6, "windows_ablation") == "linatt_5.0M_s8_g6"
assert run_id_for("cnn", 1_000_000, 4, 11, "control") == "cnn_1.0M_s4_g11-shuf"
# ---------------------------------------------------------------- metrics
def test_msp_and_choose_tau() -> None:
logits = np.array([[10.0, 0.0], [0.0, 10.0], [1.0, 1.0]])
values = msp(logits)
assert np.all(values >= 0.5) and values[2] == pytest.approx(0.5)
with pytest.raises(ValueError):
msp(np.array([1.0, 0.0]))
with pytest.raises(ValueError):
choose_tau_msp(np.array([]))
arr = np.linspace(0.1, 0.9, 9)
tau = choose_tau_msp(arr, target_known_recall=0.95)
assert tau == pytest.approx(np.quantile(arr, 0.05))
def test_closed_set_report_perfect_predictions() -> None:
y = np.array([0, 0, 1, 1, 2])
report = closed_set_report(y, y.copy(), ["a", "b", "c"], n_boot=50, seed=0)
assert report.recall_macro == pytest.approx(1.0)
assert report.recall_micro == pytest.approx(1.0)
assert report.f1_macro == pytest.approx(1.0)
assert (report.confusion == np.diag(np.diag(report.confusion))).all()
assert report.ci["recall_macro"] == (1.0, 1.0)
# ---------------------------------------------------------------- training
def test_trainability_duty_cycle(tmp_path: Path) -> None:
seed_everything(0)
store_dir = tmp_path / "s"
synthetic_store(
store_dir, n_genera=2, reads_per_genus=12, window_samples=1200, shard_windows=12
)
cfg = RunConfig(
run_id="",
stage="arch_ladder",
store_dir=store_dir,
data=DataConfig(
genera=(), val_fraction=0.25, test_fraction=0.25, min_test_windows_per_genus=1, seed=0
),
model=ModelConfig(arch="cnn", param_budget=100_000, stride=4),
train=TrainConfig(
max_epochs=25,
patience=25,
num_workers=0,
amp="off",
device="cpu",
batch_size=8,
),
out_dir=tmp_path / "runs",
)
realised, _result, _splits = train_run(cfg)
run_dir = realised.out_dir / realised.run_id
assert (run_dir / "history.parquet").is_file()
val_payload = evaluate_val(run_dir, n_boot=0)
assert val_payload["closed"]["recall_macro"] >= 0.9
def test_control_run_lands_at_chance(tmp_path: Path) -> None:
seed_everything(1)
store_dir = tmp_path / "s"
synthetic_store(
store_dir, n_genera=3, reads_per_genus=12, window_samples=1200, shard_windows=12
)
cfg = RunConfig(
run_id="",
stage="control",
store_dir=store_dir,
data=DataConfig(genera=(), val_fraction=0.25, min_test_windows_per_genus=1, seed=0),
model=ModelConfig(arch="cnn", param_budget=100_000, stride=4),
train=TrainConfig(
max_epochs=4, patience=4, num_workers=0, amp="off", device="cpu", batch_size=8
),
out_dir=tmp_path / "runs",
)
realised, _result, _splits = train_run(cfg)
assert realised.run_id.endswith("-shuf")
payload = evaluate_test(realised.out_dir / realised.run_id, n_boot=0)
chance = 1.0 / payload["n_classes"]
assert 0.0 <= payload["closed"]["recall_macro"] <= 1.5 * chance
# ---------------------------------------------------------------- CLI
def _cli_store_cmd(store_dir: Path) -> int:
return cli_main(
[
"store",
"--synthetic",
"--out",
str(store_dir),
"--n-genera",
"4",
"--reads-per-genus",
"8",
"--window-samples",
"1200",
"--shard-windows",
"64",
"--seed",
"0",
]
)
def test_cli_end_to_end(tmp_path: Path) -> None:
store_dir = tmp_path / "cli_store"
assert _cli_store_cmd(store_dir) == 0
out_dir = tmp_path / "runs"
assert (
cli_main(
[
"train",
"--store",
str(store_dir),
"--out-dir",
str(out_dir),
"--arch",
"cnn",
"--budget",
"100000",
"--batch",
"8",
"--epochs",
"3",
"--patience",
"3",
"--workers",
"0",
"--min-test-windows",
"1",
"--amp",
"off",
"--device",
"cpu",
"--stage",
"arch_ladder",
]
)
== 0
)
run_dir = out_dir / "cnn_0.1M_s4_g4"
assert (run_dir / "config.json").is_file()
assert (run_dir / "ckpt.pt").is_file()
assert cli_main(["validate", "--run-dir", str(run_dir), "--n-boot", "50"]) == 0
assert (run_dir / "val_report.json").is_file()
assert cli_main(["test", "--run-dir", str(run_dir), "--n-boot", "50"]) == 0
assert (run_dir / "report.json").is_file()
# a control run for the gate, then the verdict
assert (
cli_main(
[
"train",
"--store",
str(store_dir),
"--out-dir",
str(out_dir),
"--arch",
"cnn",
"--budget",
"100000",
"--batch",
"8",
"--epochs",
"3",
"--patience",
"3",
"--workers",
"0",
"--min-test-windows",
"1",
"--amp",
"off",
"--device",
"cpu",
"--control",
]
)
== 0
)
control_dir = out_dir / "cnn_0.1M_s4_g4-shuf"
assert cli_main(["test", "--run-dir", str(control_dir), "--n-boot", "50"]) == 0
assert cli_main(["gate", "--runs", str(run_dir), "--control", str(control_dir)]) == 0
assert (out_dir / "gate_verdict.json").is_file()
def test_cli_test_with_trap_store(tmp_path: Path, trap_store: Path) -> None:
store_dir = tmp_path / "s"
synthetic_store(
store_dir, n_genera=4, reads_per_genus=10, window_samples=1200, shard_windows=64
)
out_dir = tmp_path / "runs"
assert cli_main(["train", "--store", str(store_dir), "--out-dir", str(out_dir),
"--arch", "cnn", "--budget", "100000", "--batch", "8", "--epochs", "2",
"--min-test-windows", "1", "--workers", "0", "--amp", "off",
"--device", "cpu"]) == 0
run_dir = out_dir / "cnn_0.1M_s4_g4"
assert (
cli_main(
["test", "--run-dir", str(run_dir), "--trap-store", str(trap_store), "--n-boot", "50"]
)
== 0
)
report = json.loads((run_dir / "report.json").read_text())
assert report["trap"] is not None
assert "trap_fpr_at_tau" in report["trap"]
assert (run_dir / "trap_report.json").is_file()
def test_build_subcommand_prints_realised(capsys: pytest.CaptureFixture) -> None:
assert cli_main(["build", "--arch", "cnn", "--budget", "1000000", "--n-classes", "11"]) == 0
payload = json.loads(capsys.readouterr().out)
assert payload["params_realised"] <= int(1.1 * 1_000_000)
assert payload["params_realised"] == payload["count_params"]
assert payload["n_classes"] == 11
# ---------------------------------------------------------------- device/AMP
def test_select_device_and_amp_policies() -> None:
assert select_device("cpu").type == "cpu"
assert select_device("").type == "cpu"
with pytest.raises(ValueError):
select_device("cuda")
with pytest.raises(ValueError):
select_device("not-a-device")
device = torch.device("cpu")
from custom_models.train import resolve_amp
assert resolve_amp("off", device) == "fp32"
assert resolve_amp("bf16", device) == "fp32"
assert resolve_amp("fp32", device) == "fp32"
assert resolve_amp("fp16", device) == "fp32"