fix(examples): refuse a split too small to score at setup, naming the size that works

CHANGELOG.md

@@ -2,6 +2,26 @@
 
 ## Unreleased
 
+- **Shrinking an example until a split came up empty gave an error that named neither the
+  cause nor the knob to turn. Resolves #87.** Tuning an example down is usually the first thing
+  a new user does, and three paths still failed the way #71 fixed for `examples/advection`:
+  - `examples.sdss.train_torch --n-plates 3` (one plate is one chunk, so the 0.15 val
+    fraction rounds to none) raised `empty dataset: cannot infer spectrum width` from a width
+    probe.
+  - `bench.advection_sweep --n-steps 192` passed its own precheck, which claimed to catch
+    exactly this, and then died in the child with a `CalledProcessError` stacked on top of the
+    real message. The precheck's `3 * sample_chunk` bound ignored the 24-step horizon; the true
+    minimum is 384.
+  - `examples.microscopy` with fewer than 6 Z-planes (`--n-planes`, or a narrow
+    `--sample-range` on the IDR image) trained a full run and only then refused in `evaluate`.
+
+  #71's solver now lives in one place (`examples/_splits.py`). All three examples refuse at
+  setup and name the smallest size that works, and the sweep applies the child's exact rule
+  at argparse, before any run starts. The move also fixes a latent bug in #71's solver: its
+  stand-in geometry lengthened axis 0 rather than the sample axis, so on a middle-axis layout
+  like the IDR mask (Z-chunk 30 behind a T axis chunked 1) it named 30 samples where 180 are
+  needed. No example hit this yet, because only advection (axis 0) called it.
+
 - **The microscopy example's `--source idr` would have stopped working when EMBL-EBI retires
   its `uk1s3.embassy.ebi.ac.uk` S3 endpoint (#97).** The IDR image has moved to
   `livingobjects.ebi.ac.uk`: same `idr` bucket, same key layout, same path-style anonymous

bench/advection_sweep.py

@@ -89,7 +89,11 @@
 from pathlib import Path
 from typing import Any
 
-from insitubatch import InSituDataset
+import numpy as np
+
+from examples._splits import require_drawable_splits
+from examples.advection.data import FRACTIONS, SYNTH_HORIZON
+from insitubatch import ArrayGeometry, InSituDataset
 
 DEFAULT_OUT = Path(__file__).parent / "results" / "advection_sweep.jsonl"
 
@@ -187,6 +191,23 @@ def _configs(
             yield {**cfg, "sweep": "payload", "batch_size": bs, "geom": f"{cfg['store']}b{bs}"}
 
 
+def require_scorable(cfg: dict[str, Any], n_steps: int) -> None:
+    """Raise ``ValueError`` if a synthetic config's child would refuse ``n_steps`` at setup.
+
+    The child's own rule, applied to a stand-in geometry of the store it would write -- the
+    same split fractions and forecast horizon, so the parent cannot pass a trajectory the
+    child then rejects an hour into the sweep. Real-store configs are the child's to check.
+    """
+    if cfg["source"] != "synthetic":
+        return
+    size, sc = cfg["size"], cfg["sample_chunk"]
+    ic = cfg["inner_chunk"] or size
+    geom = ArrayGeometry(
+        path="t2m", shape=(n_steps, size, size), chunks=(sc, ic, ic), dtype=np.dtype("f4")
+    )
+    require_drawable_splits(geom, FRACTIONS, knob="--n-steps", offsets=(0, SYNTH_HORIZON))
+
+
 def store_key(cfg: dict[str, Any]) -> str:
     """Which zarr a config reads -- the identity both the URL and the persistence check use.
 
@@ -444,14 +465,13 @@ def main() -> None:
     configs = list(
         _configs(sweeps, wb2_range=args.wb2_range, payload_batch_sizes=args.payload_batch_sizes)
     )
-    # A 0.8/0.1/0.1 split needs >=3 sample-axis chunks, so the synthetic trajectory must be a
-    # few chunks long; catch a too-short --n-steps here rather than as an empty-val crash later.
+    # Catch a too-short --n-steps here, with the child's own rule, rather than as a refused
+    # child (and a CalledProcessError on top of its message) partway through the sweep.
     for cfg in configs:
-        if cfg["source"] == "synthetic" and args.n_steps < 3 * cfg["sample_chunk"]:
-            p.error(
-                f"--n-steps {args.n_steps} too short for sample_chunk {cfg['sample_chunk']} "
-                f"(need >= 3 chunks to split; use --n-steps >= {3 * cfg['sample_chunk']})"
-            )
+        try:
+            require_scorable(cfg, args.n_steps)
+        except ValueError as e:
+            p.error(f"{cfg['sweep']} {cfg['geom']} (sample_chunk {cfg['sample_chunk']}): {e}")
         # A batch never crosses a shuffle block, so an oversized batch_size is silently clipped
         # to the block -- which on the payload sweep would collapse two payload points onto one
         # while still labelling them apart. Fail instead.

examples/_splits.py

@@ -0,0 +1,94 @@
+"""Refuse a run whose train or val split has nothing to draw -- at setup, naming the fix.
+
+Both facts are fixed by the arguments, so neither is worth an epoch to discover: a run with an
+empty val split would train, validate, and only then fail, keeping none of it. A split can come
+up empty two ways that look identical from outside -- it rounds to zero chunks, or it gets
+chunks whose every anchor a windowed view drops off the end of the array -- and
+:func:`~insitubatch.summary.drawable_samples` is the one number that separates them.
+
+The remedy is *solved* against the same :func:`~insitubatch.split_by_chunk` and window
+arithmetic the run uses rather than quoted as a constant, so the advice cannot drift from the
+check. It is pure integer work on a stand-in geometry: no store is opened, which is also what
+lets a driver (``bench.advection_sweep``) apply the child's exact rule before launching it.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import replace
+
+from insitubatch import ArrayGeometry, split_by_chunk
+from insitubatch.summary import drawable_samples
+
+SCORED = ("train", "val")  # the splits a run trains on and is scored against
+
+
+def _drawable(
+    geom: ArrayGeometry,
+    fractions: tuple[float, float, float],
+    offsets: Sequence[int],
+    sample_range: tuple[int, int] | None = None,
+) -> tuple[dict[str, int], dict[str, int]]:
+    """``(chunks, drawable samples)`` per split, for ``geom`` read at each of ``offsets``."""
+    manifest = split_by_chunk(geom, fractions=fractions, sample_range=sample_range)
+    views = {f"offset{k}": geom.shift(k) for k in offsets}
+    chunks = {name: len(ids) for name, ids in manifest.chunks.items()}
+    return chunks, drawable_samples(views, manifest)
+
+
+def minimum_samples(
+    geom: ArrayGeometry, fractions: tuple[float, float, float], offsets: Sequence[int] = (0,)
+) -> int:
+    """The smallest sample-axis length, in whole chunks, for which every scored split draws.
+
+    Returns 0 when nothing up to 64 chunks works; the caller then reports what it has.
+    """
+    step = geom.sample_chunk_size
+    for n in range(step, step * 64 + 1, step):
+        shape = list(geom.shape)
+        shape[geom.sample_axis] = n
+        _, drawable = _drawable(replace(geom, shape=tuple(shape)), fractions, offsets)
+        if all(drawable[name] > 0 for name in SCORED):
+            return n
+    return 0
+
+
+def require_drawable_splits(
+    geom: ArrayGeometry,
+    fractions: tuple[float, float, float],
+    *,
+    knob: str,
+    offsets: Sequence[int] = (0,),
+    sample_range: tuple[int, int] | None = None,
+) -> None:
+    """Raise ``ValueError`` if the train or val split of ``geom`` has nothing to draw.
+
+    ``knob`` names what the reader changes to lengthen the sample axis (e.g. ``--n-steps``);
+    ``offsets`` are the windowed views the run reads (``(0, horizon)`` for a forecast target).
+    """
+    chunks, drawable = _drawable(geom, fractions, offsets, sample_range)
+    empty = [name for name in SCORED if not drawable[name]]
+    if not empty:
+        return
+    counts = ", ".join(
+        f"{name}={chunks[name]} chunk(s)/{drawable[name]} drawable samples" for name in SCORED
+    )
+    need = minimum_samples(geom, fractions, offsets)
+    remedy = (
+        f"Use at least {need} samples on the sample axis ({need // geom.sample_chunk_size} "
+        f"chunks of {geom.sample_chunk_size}) -- {knob}; you have {geom.n_samples}."
+        if need
+        else f"Widen the sample axis ({knob}) or change the split fractions."
+    )
+    ahead = max(offsets)
+    window = (
+        f" A view read {ahead} samples ahead also drops the last {ahead} anchors, which can "
+        f"empty a split that did get chunks."
+        if ahead
+        else ""
+    )
+    raise ValueError(
+        f"the {' and '.join(empty)} split has nothing to draw, so this run could train but "
+        f"never be scored. {geom.n_samples} samples over chunks of {geom.sample_chunk_size} is "
+        f"{geom.n_chunks} chunk(s); split {fractions} gives {counts}.{window} {remedy}"
+    )

examples/advection/data.py

@@ -29,7 +29,6 @@
 from zarr.abc.store import Store
 
 from insitubatch import (
-    ArrayGeometry,
     Batch,
     InSituDataset,
     arraylake_store,
@@ -38,9 +37,9 @@
     open_geometries,
     split_by_chunk,
 )
-from insitubatch.summary import drawable_samples
 
 from .._logging import add_log_level, configure_logging
+from .._splits import require_drawable_splits
 
 # The public WeatherBench2 ERA5 (6-hourly) ARCO store the wb2_* examples use.
 WB2_URL = (
@@ -55,6 +54,7 @@
 
 SYNTH_VARS = ("t2m", "u10", "v10")
 SYNTH_HORIZON = 24  # 24 steps x 1 h = 24 h
+FRACTIONS = (0.8, 0.1, 0.1)  # train / val / test, by chunk
 
 LABELS = ("t2m", "u10", "v10")  # canonical input labels (store-independent)
 
@@ -215,64 +215,6 @@ def _synthetic_ready(
         return False
 
 
-def _minimum_samples(
-    geom: ArrayGeometry, fractions: tuple[float, float, float], horizon: int
-) -> int:
-    """The smallest sample-axis length for which train *and* val can both be drawn.
-
-    Solved against the same `split_by_chunk` and `drawable_samples` the run will use, rather
-    than quoted as a constant, so it stays true if either changes. Pure integer work on a
-    stand-in geometry -- no store is opened.
-    """
-    step = geom.sample_chunk_size
-    for n in range(step, step * 64 + 1, step):
-        trial = ArrayGeometry(
-            path=geom.path, shape=(n, *geom.shape[1:]), chunks=geom.chunks, dtype=geom.dtype
-        )
-        views = {"now": trial, "target": trial.shift(horizon)}
-        drawable = drawable_samples(views, split_by_chunk(trial, fractions=fractions))
-        if all(drawable[name] > 0 for name in ("train", "val")):
-            return n
-    return 0  # nothing in range works -- the caller reports what it has instead
-
-
-def _require_drawable_splits(
-    ds: InSituDataset, geom: ArrayGeometry, horizon: int, fractions: tuple[float, float, float]
-) -> None:
-    """Refuse a geometry whose train or val view is empty, before any training runs.
-
-    Both are known from the arguments, so neither is worth an epoch to discover: the run
-    would train, validate, and only then reach `evaluate` with no batches to score.
-
-    The counts come from :func:`~insitubatch.summary.drawable_samples`, which is what the
-    engine itself reports -- a split can hold chunks and still draw nothing, because a
-    windowed target drops the anchors whose horizon runs off the array.
-    """
-    report = ds.describe()["config"]
-    drawable, chunks = report["drawable_samples"], report["split_chunks"]
-    empty = [name for name in ("train", "val") if not drawable[name]]
-    if not empty:
-        return
-    counts = ", ".join(
-        f"{name}={chunks[name]} chunk(s)/{drawable[name]} drawable samples"
-        for name in ("train", "val")
-    )
-    need = _minimum_samples(geom, fractions, horizon)
-    remedy = (
-        f"Use at least {need} samples on the sample axis (--n-steps for the synthetic store, "
-        f"--sample-range for a real one); you have {geom.n_samples}."
-        if need
-        else "Widen the sample axis (--n-steps / --sample-range) or change the split fractions."
-    )
-    raise ValueError(
-        f"the {' and '.join(empty)} split has nothing to draw, so this run could train but "
-        f"never be scored. {geom.n_samples} samples over chunks of {geom.sample_chunk_size} is "
-        f"{geom.n_chunks} chunk(s); split {fractions} gives {counts}. A {horizon}-step target "
-        f"also drops the last {horizon} anchors, which can empty a split that did get chunks. "
-        f"{remedy}"
-    )
-
-
 def forecast_dataset(
     store: Store,
     *,
@@ -300,9 +242,15 @@ def forecast_dataset(
     opened = open_geometries(store, variables=list(variables))
     geoms = {label: opened[var] for label, var in zip(LABELS, variables, strict=True)}
     geoms["target"] = opened[variables[0]].shift(horizon)
-    fractions = (0.8, 0.1, 0.1)
-    manifest = split_by_chunk(opened[variables[0]], fractions=fractions, sample_range=sample_range)
-    ds = InSituDataset(
+    require_drawable_splits(
+        opened[variables[0]],
+        FRACTIONS,
+        knob="--n-steps for the synthetic store, --sample-range for a real one",
+        offsets=(0, horizon),
+        sample_range=sample_range,
+    )
+    manifest = split_by_chunk(opened[variables[0]], fractions=FRACTIONS, sample_range=sample_range)
+    return InSituDataset(
         store,
         manifest,
         geometries=geoms,
@@ -311,8 +259,6 @@ def forecast_dataset(
         cache_dir=cache_dir,
         max_inflight=max_inflight,
     )
-    _require_drawable_splits(ds, opened[variables[0]], horizon, fractions)
-    return ds
 
 
 def inputs_and_targets(batch: Batch) -> tuple[np.ndarray, np.ndarray, np.ndarray]:

examples/microscopy/data.py

@@ -40,6 +40,7 @@
 )
 
 from .._logging import add_log_level, configure_logging
+from .._splits import require_drawable_splits
 
 # A public OME-NGFF (zarr v0.1) image in the EMBL-EBI Image Data Repository: a 3D two-channel
 # confocal stack with an expert instance-segmentation label. Read anonymously off the IDR S3.
@@ -52,6 +53,7 @@
 # (leading T=1) so one code path serves both sources.
 SAMPLE_AXIS = 2
 LABELS = ("raw", "mask")  # canonical labels (store-independent)
+FRACTIONS = (0.8, 0.1, 0.1)  # train / val / test, by raw chunk (one Z-plane)
 
 
 def _idr_store() -> Store:
@@ -204,13 +206,20 @@ def segmentation_dataset(
     store; the dataset's labels are always ``raw, mask`` so the model code is store-agnostic.
     Both share the Z (sample) axis *length* but may chunk it differently -- the engine maps the
     single sample grid (the ``raw`` chunking, one plane per anchor) onto each variable's own
-    chunks. ``sample_range`` restricts the split to a finite Z window. Iterate the returned
-    dataset's ``.train`` / ``.val`` views.
+    chunks. ``sample_range`` restricts the split to a finite Z window; a window too shallow to
+    give train and val a plane each is refused here, naming the depth that works. Iterate the
+    returned dataset's ``.train`` / ``.val`` views.
     """
     opened = open_geometries(store, variables=[raw_var, mask_var], sample_axis=sample_axis)
     geoms = {"raw": opened[raw_var], "mask": opened[mask_var]}
     # The raw (one plane per chunk) is the reference sample grid: one Z-plane = one sample.
-    manifest = split_by_chunk(opened[raw_var], fractions=(0.8, 0.1, 0.1), sample_range=sample_range)
+    require_drawable_splits(
+        opened[raw_var],
+        FRACTIONS,
+        knob="--n-planes for the synthetic store, --sample-range for the real one",
+        sample_range=sample_range,
+    )
+    manifest = split_by_chunk(opened[raw_var], fractions=FRACTIONS, sample_range=sample_range)
     return InSituDataset(
         store,
         manifest,

examples/sdss/data.py

@@ -54,11 +54,14 @@
     split_by_chunk,
 )
 
+from .._splits import require_drawable_splits
+
 FLUX_VAR = "flux"
 SDSS_HOST = "https://data.sdss.org"
 REDUX = "/sas/dr17/sdss/spectro/redux/26"  # legacy (DR8-era) reduction; spPlate lives per-plate
 LATENT_DIM = 16  # bottleneck / PCA components -- the reconstruction budget both methods share
 NOISE_SIGMA = 0.3  # per-sample input noise (normalized units), added in the batch stage
+FRACTIONS = (0.7, 0.15, 0.15)  # train / val / test, by chunk (a plate, synthetically)
 
 HERE = Path(__file__).resolve().parent
 DEFAULT_URIS = HERE / "dr17_plate_uris.json"
@@ -436,14 +439,21 @@ def reconstruct_dataset(
     """One reconstruction dataset over the SDSS store: iterate ``.train`` / ``.val`` / ``.test``.
 
     Each batch carries ``noisy`` and ``clean`` ``(B, W)`` spectra. Split is by chunk (plate), so no
-    fiber leaks across train/val/test. ``max_inflight`` is unset by default (``data.sdss.org`` over
-    HTTPS is not an anonymous-S3 throttler like MAST); raise/cap it to tune read-ahead.
+    fiber leaks across train/val/test; a store too small to give train and val a chunk each is
+    refused here, naming the size that works. ``max_inflight`` is unset by default
+    (``data.sdss.org`` over HTTPS is not an anonymous-S3 throttler like MAST); raise/cap it to
+    tune read-ahead.
     """
     geoms = open_geometries(store, variables=[FLUX_VAR], sample_axis=0)
+    require_drawable_splits(
+        geoms[FLUX_VAR],
+        FRACTIONS,
+        knob="--n-plates for the synthetic store (one plate is one chunk), --plates for a real one",
+    )
     chunk_transforms: list[ChunkTransform] = [normalize]
     return InSituDataset(
         store,
-        split_by_chunk(geoms[FLUX_VAR], fractions=(0.7, 0.15, 0.15)),
+        split_by_chunk(geoms[FLUX_VAR], fractions=FRACTIONS),
         geometries=geoms,
         batch_size=batch_size,
         shuffle=shuffle,