feat: big_tf_unet_2 — intern-driven single-scale teacher pipeline

pyproject.toml

@@ -48,6 +48,8 @@ train = [
     "pybaselines",
     "tifffile",
     "pandas",
+    "kneed",
+    "threadpoolctl",
 ]
 
 [project.scripts]
@@ -94,4 +96,7 @@ train = [
     "pybaselines",
     "tifffile",
     "pandas",
+
+    "kneed",
+    "threadpoolctl",
 ]

src/tokeye/training/big_tf_unet_2/README.md

@@ -0,0 +1,107 @@
+# big_tf_unet_2 — single-scale teacher training, human-in-the-loop
+
+This pipeline trains **one model per STFT resolution** (an `(nfft, hop)` pair).
+You run it once per scale; the resulting per-scale "teacher" models are later
+distilled into one multiscale student (not part of this pipeline).
+
+You drive the pipeline from Jupyter notebooks. You never edit code — every
+adjustable value (a "knob") lives in one YAML file, and every step shows you
+pictures so you can decide whether to keep its output or tweak a knob and
+rerun it.
+
+## One-time setup
+
+```bash
+cd /scratch/gpfs/nc1514/tokeye
+uv sync --group train          # installs everything the pipeline needs
+```
+
+Open JupyterLab through the cluster's OnDemand portal (or however you usually
+start notebooks), with this repository as the working directory. The kernel
+must use the project environment (`.venv`).
+
+## Starting a run
+
+From the repository root:
+
+```bash
+uv run python -m tokeye.training.big_tf_unet_2.scaffold --nfft 512 --hop 128
+```
+
+This creates `dev/training/nfft512_hop128/` containing:
+
+- **`run.yaml`** — the only file you ever edit. It starts with sensible
+  defaults; most knobs say `auto`, meaning the pipeline computes a suggested
+  value from your data and records what it picked.
+- **`00_setup.ipynb` … `05_eval.ipynb`** — the six notebooks you work
+  through, in order.
+
+To try a different resolution, scaffold again with different numbers — each
+scale is a completely separate run with its own folder, cache, and model.
+Nothing you do in one run can affect another.
+
+## The loop (same for every step)
+
+1. **Suggest** — `run.suggest("step_2")` shows the auto-chosen values and
+   tells you which section of `run.yaml` holds this step's knobs.
+2. **Run** — light steps run right in the notebook
+   (`run.run("step_4")`); heavy steps are submitted to the cluster
+   (`job = run.submit("step_3")`). Re-run `run.status()` until the step says
+   `complete`; `run.log(job)` shows the job's output.
+3. **Look** — `run.gallery("step_2")` shows the step's results as images.
+4. **Decide** — if it looks right, `run.accept("step_2")` unlocks the next
+   step. If not: edit the knob in `run.yaml`, `run.clear("step_2")`, and run
+   it again. Clearing a step automatically marks everything after it
+   `stale`, so you can't accidentally train on outdated intermediate data.
+
+## What the steps do
+
+| Step | What it does | Where it runs |
+|---|---|---|
+| step_0 | load raw signals, cut into windows | notebook |
+| step_1 | spectrograms + keep the most active windows | notebook |
+| step_2 | remove the smooth background (baseline) | cluster (CPU) |
+| step_3 | self-supervised denoiser | cluster (GPU) |
+| step_4 | threshold into coherent/transient masks | notebook (seconds) |
+| step_5 | pack all diagnostics into a training set | notebook |
+| step_6 | 5-fold cross-validation "second opinion" on the masks | cluster (GPU) |
+| step_7 | train + export the final model | cluster (GPU) |
+| step_8 | quick TJ-II benchmark number | notebook |
+
+## Knob glossary (the ones you'll actually touch)
+
+| Knob (run.yaml) | Meaning |
+|---|---|
+| `baseline.lam` | Background smoothness. Bigger = smoother background estimate. `auto` scales it to your resolution. |
+| `baseline.edge_k` | How aggressively noisy frequency bins at the spectrogram's top/bottom edges are cut. Bigger = fewer bins cut. |
+| `labels.knee_sensitivity` | Mask threshold strictness. Bigger = higher threshold = fewer labeled pixels. |
+| `labels.delta` | Direct threshold offset in noise-sigma units (0.5 = half a sigma stricter). |
+| `labels.min_size_fraction` | Smallest object kept in a mask, as a fraction of the image. Bigger = more small specks removed. |
+| `refine.model_trust` | 0 = trust your step_4 masks completely; 1 = trust the cross-validation models completely. Changing it only re-runs step_7 (cheap). |
+| `denoise.max_epochs` | Denoiser training length. More = cleaner, slower. |
+
+Bad values are caught the moment a step starts, with a message naming the
+field and its allowed range — you cannot break anything by mistyping a knob.
+
+## If something fails
+
+- `run.status()` shows `failed` — `run.log(job)` (or the printed error for
+  notebook steps) says why. Most failures are fixed by a knob change +
+  `run.clear(step)` + rerun.
+- A cluster job seems slow or stuck — `run.jobstats(job)` shows whether it's
+  actually using its CPUs/GPU.
+- You want to start a scale completely over —
+  `run.clear_all(confirm="<run_id>")` (it makes you type the run id so it
+  can't happen by accident). Your `run.yaml` and notebooks are kept.
+
+## For developers
+
+Steps expose `main(settings: dict)`; the settings are built centrally in
+`runner.py` from the merged config (`config/defaults.yaml` ← `run.yaml`,
+validated by `run_config.py`). `"auto"` knobs resolve in `auto_resolve.py`
+and every resolved value is recorded with its source in the run's
+`resolved_params.yaml`. Progress, human sign-off, and staleness live in
+`task_matrix.json` (`task_matrix.py`). `paths.py` is the single source of
+truth for step registry and artifact locations. Deployment inputs
+(normalization stats, edge bins, scale) are exported per run in
+`model/big_tf_unet_2/<run_id>/deploy_manifest.yaml`.

src/tokeye/training/big_tf_unet_2/auto_resolve.py

@@ -0,0 +1,142 @@
+"""Resolution of ``"auto"`` knob values + the resolved-parameter ledger.
+
+Every knob that supports ``"auto"`` is resolved once per run (from the run's
+own data or geometry) and recorded in ``resolved_params.yaml`` together with
+its source (``auto`` / ``user`` / ``in-step``). That file is the experiment
+record the notebooks display, and the source of the deployment manifest.
+
+Two kinds of resolution:
+
+- **pre-run** (this module): values computable from the config geometry or an
+  upstream artifact before the step starts (lam, edge bins, num_layers, ...).
+- **in-step**: values that only exist mid-step (per-modality stats computed
+  from data the step itself produces). Step ``main()`` returns them and the
+  runner records them here.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import yaml
+
+from .utils import auto_params
+
+if TYPE_CHECKING:
+    from .paths import RunPaths
+    from .run_config import RunConfig
+
+
+def _pad_geometry(cfg: RunConfig, num_layers: int) -> tuple[int, int]:
+    mult = 2**num_layers
+    return (
+        auto_params.pad_to_multiple(cfg.n_freq, mult),
+        auto_params.pad_to_multiple(cfg.n_time, mult),
+    )
+
+
+def _resolve_model_geometry(cfg: RunConfig, section: Any) -> dict[str, Any]:
+    """num_layers + batch_size for one model section (denoise/refine/final)."""
+    out: dict[str, Any] = {}
+    num_layers = section.num_layers
+    if num_layers == "auto":
+        num_layers = auto_params.compute_num_layers(cfg.n_freq, cfg.n_time)
+        out["num_layers"] = num_layers
+    if section.batch_size == "auto":
+        out["batch_size"] = auto_params.compute_batch_size(
+            section.base_batch_size, cfg.n_freq, cfg.n_time, num_layers
+        )
+    return out
+
+
+def resolve_step_autos(
+    cfg: RunConfig, step: str, modality: str | None, paths: RunPaths
+) -> dict[str, Any]:
+    """Pre-run auto resolutions for one step (empty dict if none apply)."""
+    resolved: dict[str, Any] = {}
+
+    if step == "step_2":
+        if cfg.baseline.lam == "auto":
+            resolved["lam"] = auto_params.compute_lam(cfg.n_freq)
+        in_h5 = paths.step_h5("step_1", modality)
+        if in_h5.exists():
+            if cfg.baseline.edge_method == "energy":
+                lower, upper = auto_params.detect_edge_bins_energy(
+                    in_h5,
+                    k=cfg.baseline.edge_k,
+                    max_fraction=cfg.baseline.edge_max_fraction,
+                )
+            else:
+                lower, upper = auto_params.detect_edge_bins(
+                    in_h5, gradient_threshold=cfg.baseline.gradient_threshold
+                )
+            resolved["edge_bins_lower"] = lower
+            resolved["edge_bins_upper"] = upper
+
+    elif step == "step_3":
+        resolved.update(_resolve_model_geometry(cfg, cfg.denoise))
+        # Reuse step_2's edge bins for input masking (consistency over
+        # recomputation) — fall back to recomputing only if absent.
+        ledger = read_ledger(paths)
+        prior = ledger.get("step_2", {}).get(modality or "-", {})
+        for key in ("edge_bins_lower", "edge_bins_upper"):
+            if key in prior:
+                resolved[key] = prior[key]["value"]
+
+    elif step == "step_4":
+        labels = cfg.labels
+        if labels.min_size == "auto":
+            resolved["min_size"] = auto_params.compute_min_size(
+                cfg.n_freq, cfg.n_time, labels.min_size_fraction
+            )
+        if labels.remove_bottom_rows == "auto" or labels.remove_top_rows == "auto":
+            bottom, top = auto_params.compute_row_removal(
+                cfg.n_freq,
+                labels.row_removal_fraction_bottom,
+                labels.row_removal_fraction_top,
+            )
+            if labels.remove_bottom_rows == "auto":
+                resolved["remove_bottom_rows"] = bottom
+            if labels.remove_top_rows == "auto":
+                resolved["remove_top_rows"] = top
+
+    elif step in ("step_6", "step_7"):
+        section = cfg.refine if step == "step_6" else cfg.final
+        resolved.update(_resolve_model_geometry(cfg, section))
+
+    return resolved
+
+
+def suggest(cfg: RunConfig, step: str, modality: str | None, paths: RunPaths) -> dict:
+    """Dry-run resolution for notebook display; never writes the ledger."""
+    return resolve_step_autos(cfg, step, modality, paths)
+
+
+# ---------------------------------------------------------------------------
+# Ledger
+# ---------------------------------------------------------------------------
+
+def read_ledger(paths: RunPaths) -> dict:
+    if paths.resolved_params_path.exists():
+        return yaml.safe_load(paths.resolved_params_path.read_text()) or {}
+    return {}
+
+
+def record(
+    paths: RunPaths,
+    step: str,
+    modality: str | None,
+    values: dict[str, Any],
+    source: str,
+) -> None:
+    """Merge resolved values for (step, modality) into resolved_params.yaml."""
+    if not values:
+        return
+    ledger = read_ledger(paths)
+    slot = ledger.setdefault(step, {}).setdefault(modality or "-", {})
+    for key, value in values.items():
+        slot[key] = {"value": value, "source": source}
+    paths.resolved_params_path.parent.mkdir(parents=True, exist_ok=True)
+    paths.resolved_params_path.write_text(
+        yaml.safe_dump(ledger, sort_keys=False, default_flow_style=False)
+    )

src/tokeye/training/big_tf_unet_2/clearing.py

@@ -0,0 +1,79 @@
+"""Safe, file-granular clearing of step artifacts.
+
+Replaces the legacy ``setup_directory(overwrite=True)`` pattern (which
+``shutil.rmtree``'d whole directories and could wipe sibling modalities).
+Every deletion here is resolved against an explicit artifact list from
+``RunPaths.artifacts`` and fenced to the run's own cache/model roots —
+clearing can never touch another run or anything outside the pipeline dirs.
+"""
+
+from __future__ import annotations
+
+import shutil
+from typing import TYPE_CHECKING
+
+from .paths import RunPaths, get_step
+from .task_matrix import RunTaskMatrix
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from .run_config import RunConfig
+
+
+def _assert_fenced(target: Path, roots: list[Path]) -> None:
+    resolved = target.resolve()
+    for root in roots:
+        if resolved.is_relative_to(root.resolve()):
+            return
+    raise RuntimeError(
+        f"Refusing to delete {target}: outside the run's cache/model roots"
+    )
+
+
+def clear_step(
+    paths: RunPaths,
+    cfg: RunConfig,
+    step: str,
+    modality: str | None = None,
+) -> list[Path]:
+    """Delete one step's artifacts (optionally one modality) and mark it
+    pending; everything downstream that was complete becomes stale.
+
+    Returns the paths that were removed.
+    """
+    spec = get_step(step)
+    if modality is not None and not spec.per_modality:
+        raise ValueError(f"{step} is not per-modality; drop the modality argument")
+    modalities = [modality] if modality is not None else cfg.modality_names
+
+    roots = [paths.cache_root, paths.model_dir]
+    removed: list[Path] = []
+    for target in paths.artifacts(step, modalities):
+        _assert_fenced(target, roots)
+        if target.is_dir():
+            shutil.rmtree(target)
+            removed.append(target)
+        elif target.exists():
+            target.unlink()
+            removed.append(target)
+
+    matrix = RunTaskMatrix(paths.task_matrix_path)
+    matrix.mark_pending(step, modalities)
+    return removed
+
+
+def clear_run(paths: RunPaths, confirm: str) -> None:
+    """Delete ALL cache and model artifacts of a run. Keeps the workspace
+    (run.yaml + notebooks). ``confirm`` must equal the run_id, typed out.
+    """
+    if confirm != paths.run_id:
+        raise ValueError(
+            f"clear_run needs confirm={paths.run_id!r} (got {confirm!r}). "
+            f"Type the run id exactly to confirm you mean it."
+        )
+    for root in (paths.cache_root, paths.model_dir):
+        if root.exists():
+            _assert_fenced(root, [root])
+            shutil.rmtree(root)
+    paths.cache_root.mkdir(parents=True, exist_ok=True)