diff --git a/pyproject.toml b/pyproject.toml index a4e5f52..9ffb556 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,8 @@ train = [ "pybaselines", "tifffile", "pandas", + "kneed", + "threadpoolctl", ] [project.scripts] @@ -94,4 +96,7 @@ train = [ "pybaselines", "tifffile", "pandas", + + "kneed", + "threadpoolctl", ] diff --git a/src/tokeye/training/big_tf_unet_2/README.md b/src/tokeye/training/big_tf_unet_2/README.md new file mode 100644 index 0000000..b6002dd --- /dev/null +++ b/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="")` (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//deploy_manifest.yaml`. diff --git a/src/tokeye/training/big_tf_unet_2/__init__.py b/src/tokeye/training/big_tf_unet_2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/tokeye/training/big_tf_unet_2/auto_resolve.py b/src/tokeye/training/big_tf_unet_2/auto_resolve.py new file mode 100644 index 0000000..c4df06f --- /dev/null +++ b/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) + ) diff --git a/src/tokeye/training/big_tf_unet_2/clearing.py b/src/tokeye/training/big_tf_unet_2/clearing.py new file mode 100644 index 0000000..0e705f1 --- /dev/null +++ b/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) diff --git a/src/tokeye/training/big_tf_unet_2/config/defaults.yaml b/src/tokeye/training/big_tf_unet_2/config/defaults.yaml new file mode 100644 index 0000000..3dae17f --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/config/defaults.yaml @@ -0,0 +1,121 @@ +# Complete defaults for the big_tf_unet_2 single-scale pipeline. +# A run's run.yaml is MERGED ON TOP of this file (OmegaConf.merge), so run.yaml +# only ever needs the keys being overridden. Values of "auto" are resolved once +# per run and recorded (with their resolved values) in resolved_params.yaml. + +run: + nfft: 512 + hop: 128 + allow_custom_scale: false # true permits (nfft, hop) outside the canonical grid + +modalities: + ece: { input_key: ece, channels: [8, 12, 16, 20, 24, 28, 32, 36] } + mhr: { input_key: mhr, channels: [3, 4, 5, 6] } + bes: { input_key: bes, channels: [26, 28, 30, 32, 34, 36, 38, 40] } + # co2 disabled: the preserved co2 top-up data is all-zero for every shot/chord. + # Re-enable once co2 is re-fetched with real signal. + # co2: { input_key: co2, channels: [0, 1, 2, 3] } + +extraction: + subseq_len: 66000 # fixed physical window (132 ms @ 500 kHz) at every scale + preemphasis_coeff: 0.99 + fs_khz: 500 + target_rate_khz: 500 # intake resamples every modality to this + ip_threshold: 0.1 + max_windows_per_shot_precap: 60 + +window_filter: + enabled: true + weights: model/big_tf_unet_251210_weights.pt + max_windows_per_shot: 25 + activity_threshold: 0.5 # sigmoid cutoff for an "active" pixel + min_activity: 0.0005 # min active-pixel fraction; drops near-empty windows + mean: auto # per-modality log-magnitude stats (never global) + std: auto + +baseline: + method: fabc + lam: auto # auto = 1e6 * (F/513)^4 — same physical smoothing per scale + edge_method: energy # energy-level edge-bin detection (or "gradient" legacy) + edge_k: 2.0 # mask contiguous edge bins with |log E(f) - log M| > log k + edge_max_fraction: 0.15 # cap masked bins per edge + gradient_threshold: 0.5 # only used when edge_method == "gradient" + +denoise: + representation: complex # complex | magnitude + normalization: robust_asinh # robust_asinh | zscore (legacy) + a: 1.0 # N_a compression scale at the denoiser input + first_layer_size: 32 + num_layers: auto # auto = min(5, floor(log2(min(F, T) / 4))) + batch_size: auto # auto = base_batch_size * (513*516) / (F_pad*T_pad) + base_batch_size: 8 + precision: "32-true" + max_epochs: 30 + tv_patience: 3 + num_workers: 8 + +labels: + knee_sensitivity: 1.0 # kneed S: larger = later/more conservative knee + delta: 0.0 # additive threshold offset, in robust-sigma units + fallback_frac: 0.02 # survival quantile used when no knee is found + min_size: auto + min_size_fraction: 0.0002 + remove_bottom_rows: auto + remove_top_rows: auto + row_removal_fraction_bottom: 0.01 + row_removal_fraction_top: 0.004 + +dataset: + a: 3.0 # N_a compression scale at the segmenter input + stats_windows: 100 # windows sampled for per-modality robust stats + +refine: + model_trust: 0.5 # lambda in q = (1-lambda)*y0 + lambda*p_oof; 0 = no refine + n_folds: 5 + first_layer_size: 32 + num_layers: auto + batch_size: auto + base_batch_size: 16 + precision: "32-true" + max_epochs: 200 + loss_type: symmetric_bce_dice + num_workers: 8 + +final: + first_layer_size: 32 + num_layers: auto + batch_size: auto + base_batch_size: 28 + precision: "32-true" + max_epochs: 100 + loss_type: bce # plain BCE — the soft target q is already smoothed + gamma: 2.0 # only used by focal-family loss_types + num_workers: 8 + +eval: + dataset_dir: data/eval/TJII2021 + n_thresholds: 50 # threshold sweep resolution for the F1-optimal point + +paths: + shots_path: data/autoprocess/settings/shots.txt + foundation_dir: data/autoprocess/foundation + +slurm: + account: null + cpu_partition: null # null = scheduler default routing (CPU work, no gres) + gpu_partition: gpu + cpu_cpus: 42 + cpu_mem: 43G + cpu_time: "08:00:00" + gpu_cpus: 16 + gpu_mem: 48G + gpu_time: "18:00:00" + +smoke: + enabled: false + n_shots: 2 + max_windows_per_shot: 2 + n_folds: 2 + max_epochs: 1 + refine_max_epochs: 1 + final_max_epochs: 1 diff --git a/src/tokeye/training/big_tf_unet_2/config/run_template.yaml b/src/tokeye/training/big_tf_unet_2/config/run_template.yaml new file mode 100644 index 0000000..b7c249c --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/config/run_template.yaml @@ -0,0 +1,36 @@ +# Run configuration — {{RUN_ID}} +# +# This is the ONLY file you edit. Everything not listed here uses the pipeline +# defaults (src/tokeye/training/big_tf_unet_2/config/defaults.yaml). +# +# Workflow for every step: run it, look at the gallery in the notebook, and if +# the pictures look wrong, change the relevant knob below, clear the step, and +# run it again. "auto" means the pipeline computes a suggested value from your +# data — the notebook's suggest() cell shows you what it picked. + +run: + nfft: {{NFFT}} # STFT window size (LOCKED once step_1 has run) + hop: {{HOP}} # STFT hop length (LOCKED once step_1 has run) + +# --- step_2 baseline removal ------------------------------------------------- +baseline: + lam: auto # baseline stiffness; bigger = smoother baseline + edge_k: 2.0 # edge-bin cut aggressiveness; bigger = fewer bins cut + +# --- step_3 denoiser --------------------------------------------------------- +denoise: + max_epochs: 30 + +# --- step_4 labels (the mask thresholds) -------------------------------------- +labels: + knee_sensitivity: 1.0 # bigger = higher threshold = fewer labeled pixels + delta: 0.0 # extra threshold offset in noise-sigma units + min_size_fraction: 0.0002 # smallest object kept, as fraction of image area + +# --- step_6 refine ----------------------------------------------------------- +refine: + model_trust: 0.5 # 0 = trust the step_4 labels, 1 = trust the model + +# --- smoke mode: tiny 2-shot run to sanity-check the whole pipeline ----------- +smoke: + enabled: false diff --git a/src/tokeye/training/big_tf_unet_2/gallery.py b/src/tokeye/training/big_tf_unet_2/gallery.py new file mode 100644 index 0000000..f8737a9 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/gallery.py @@ -0,0 +1,248 @@ +"""Visual checkpoints for every step boundary — the intern's decision tool. + +All plots lazy-load at most ``n`` samples via per-sample HDF5 reads; nothing +here ever loads a whole step file. ``show(step, ...)`` dispatches to the +right view for each step. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import h5py +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import yaml + +from .utils.hdf5_io import read_sample + +if TYPE_CHECKING: + from .paths import RunPaths + +_CMAP = "gist_heat" + + +def _mag(arr: np.ndarray) -> np.ndarray: + """(C,F,T,2) -> (C,F,T) magnitude; passthrough otherwise.""" + if arr.ndim == 4 and arr.shape[-1] == 2: + return np.sqrt(arr[..., 0] ** 2 + arr[..., 1] ** 2) + return np.abs(arr) + + +def _sample_keys(h5_path, n: int) -> list[int]: + with h5py.File(h5_path, "r") as f: + keys = sorted((int(k) for k in f["samples"]), key=int) + if len(keys) <= n: + return keys + idx = np.linspace(0, len(keys) - 1, n).astype(int) + return [keys[i] for i in idx] + + +def _imshow(ax, img: np.ndarray, title: str, log: bool = False) -> None: + data = np.log1p(np.abs(img)) if log else img + lo, hi = np.quantile(data, [0.01, 0.99]) + ax.imshow(data, aspect="auto", origin="lower", cmap=_CMAP, vmin=lo, vmax=hi) + ax.set_title(title, fontsize=8) + ax.set_xticks([]) + ax.set_yticks([]) + + +def timeseries_grid(paths: RunPaths, modality: str, n: int = 4) -> None: + """step_0: raw windowed signals, one channel per row.""" + h5 = paths.step_h5("step_0", modality) + keys = _sample_keys(h5, n) + fig, axes = plt.subplots(len(keys), 1, figsize=(10, 2 * len(keys)), squeeze=False) + for ax, k in zip(axes[:, 0], keys): + arr = read_sample(h5, k) + ax.plot(arr[0], lw=0.3) + ax.set_title(f"{modality} window {k} ch0", fontsize=8) + fig.tight_layout() + plt.show() + + +def spectrogram_grid( + paths: RunPaths, step: str, modality: str, n: int = 6, log: bool = True +) -> None: + """Any spectrogram-shaped step: grid of channel-0 fields.""" + h5 = paths.step_h5(step, modality) + keys = _sample_keys(h5, n) + cols = min(3, len(keys)) + rows = int(np.ceil(len(keys) / cols)) + fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3 * rows), squeeze=False) + for ax, k in zip(axes.flat, keys): + arr = _mag(read_sample(h5, k)) + _imshow(ax, arr[0], f"{modality} {step} sample {k} ch0", log=log) + for ax in axes.flat[len(keys) :]: + ax.axis("off") + fig.tight_layout() + plt.show() + + +def pair_grid( + paths: RunPaths, + step_a: str, + step_b: str, + modality: str, + n: int = 4, + log_a: bool = True, + log_b: bool = False, +) -> None: + """Before/after pairs (e.g. step_1 vs step_2, step_2 vs step_3).""" + h5_a, h5_b = paths.step_h5(step_a, modality), paths.step_h5(step_b, modality) + keys = _sample_keys(h5_b, n) + fig, axes = plt.subplots(len(keys), 2, figsize=(9, 3 * len(keys)), squeeze=False) + for row, k in zip(axes, keys): + _imshow(row[0], _mag(read_sample(h5_a, k))[0], f"{step_a} s{k}", log=log_a) + _imshow(row[1], _mag(read_sample(h5_b, k))[0], f"{step_b} s{k}", log=log_b) + fig.tight_layout() + plt.show() + + +def mask_overlay(paths: RunPaths, modality: str, n: int = 4) -> None: + """step_4: denoised field with coherent/transient mask contours.""" + h5_den = paths.step_h5("step_3", modality) + h5_mask = paths.step_h5("step_4", modality) + keys = _sample_keys(h5_mask, n) + fig, axes = plt.subplots(len(keys), 2, figsize=(9, 3 * len(keys)), squeeze=False) + for row, k in zip(axes, keys): + den = _mag(read_sample(h5_den, k))[0] + mask = read_sample(h5_mask, k) # (C, 2, F, T) + for ax, ch, name in ((row[0], 0, "coherent"), (row[1], 1, "transient")): + _imshow(ax, den, f"s{k} {name}") + m = mask[0, ch].astype(float) + if m.any(): + ax.contour(m, levels=[0.5], colors="cyan", linewidths=0.6) + ax.set_title(f"s{k} {name} ({m.mean():.2%} px)", fontsize=8) + fig.tight_layout() + plt.show() + + +def knee_plot(paths: RunPaths, modality: str, max_curves: int = 8) -> None: + """step_4: the thresholds.csv knee decisions, one curve per (shot, target).""" + df = pd.read_csv(paths.thresholds_csv(modality)) + fig, ax = plt.subplots(figsize=(8, 4)) + shown = df.head(max_curves) + ax.scatter( + range(len(shown)), + shown["threshold"], + c=["tab:blue" if t == "coherent" else "tab:orange" for t in shown["target"]], + ) + for i, (_, r) in enumerate(shown.iterrows()): + marker = " (fallback)" if r["used_fallback"] else "" + ax.annotate( + f"{r['shotn']} ch{r['channel']} {r['target']}{marker}\n" + f"{r['positive_fraction']:.2%} px", + (i, r["threshold"]), + fontsize=6, + rotation=45, + ) + ax.set_ylabel("threshold (robust sigma)") + ax.set_xticks([]) + ax.set_title(f"{modality}: knee thresholds (blue=coherent, orange=transient)") + fig.tight_layout() + plt.show() + n_fb = int(df["used_fallback"].sum()) + if n_fb: + print(f"note: {n_fb}/{len(df)} thresholds used the quantile fallback") + + +def dataset_grid(paths: RunPaths, n: int = 6) -> None: + """step_5: normalized training images with their masks.""" + h5 = paths.step_h5("step_5") + with h5py.File(h5, "r") as f: + total = int(f.attrs["n_samples"]) + idx = np.linspace(0, total - 1, min(n, total)).astype(int) + fig, axes = plt.subplots( + len(idx), 2, figsize=(9, 3 * len(idx)), squeeze=False + ) + for row, i in zip(axes, idx): + img = np.asarray(f["images"][str(i)])[0] + mask = np.asarray(f["masks"][str(i)]) + mod = f["prov_modality"][i].decode() + _imshow(row[0], img, f"{mod} sample {i} (N_a image)") + _imshow(row[1], mask[0] + 2 * mask[1], "mask (coh=1, tra=2)") + fig.tight_layout() + plt.show() + + +def refine_triptych(paths: RunPaths, n: int = 4) -> None: + """step_6: image / knee label / OOF probability / disagreement.""" + h5_data = paths.step_h5("step_5") + h5_ref = paths.step_h5("step_6") + with h5py.File(h5_data, "r") as fd, h5py.File(h5_ref, "r") as fr: + total = fr["p_oof"].shape[0] + # Most-disagreeing samples are the ones worth human eyes. + means = [float(np.mean(fr["disagreement"][i])) for i in range(total)] + idx = np.argsort(means)[::-1][:n] + fig, axes = plt.subplots( + len(idx), 4, figsize=(14, 3 * len(idx)), squeeze=False + ) + for row, i in zip(axes, idx): + img = np.asarray(fd["images"][str(i)])[0] + y0 = np.asarray(fd["masks"][str(i)])[0] + p = fr["p_oof"][i][0] + d = fr["disagreement"][i][0] + _imshow(row[0], img, f"s{i} image") + _imshow(row[1], y0.astype(float), "knee label y0") + _imshow(row[2], p, "OOF prob") + _imshow(row[3], d, f"disagreement ({means[i]:.3f})") + fig.tight_layout() + plt.show() + + +def eval_curves(paths: RunPaths) -> None: + """step_8: the TJ-II threshold sweep + best point vs deployed anchor.""" + df = pd.read_csv(paths.eval_csv) + best = df.iloc[df["f1"].idxmax()] + fig, ax = plt.subplots(figsize=(7, 4)) + ax.plot(df["threshold"], df["f1"], label="F1") + ax.plot(df["threshold"], df["iou_per_image_mean"], label="IoU (per-image)") + ax.axvline(best["threshold"], ls="--", lw=0.8, color="gray") + ax.axhline(0.26, ls=":", lw=0.8, color="green", label="deployed IoU anchor 0.26") + ax.set_xlabel("threshold") + ax.legend() + ax.set_title( + f"TJ-II: IoU={best['iou_per_image_mean']:.3f}, F1={best['f1']:.3f} " + f"@ {best['threshold']:.2f}" + ) + fig.tight_layout() + plt.show() + + +def show( + paths: RunPaths, + step: str, + modality: str | None = None, + modalities: list[str] | None = None, + n: int = 6, +) -> None: + """Dispatch to the right view for a step. Per-modality steps show every + modality unless one is named.""" + mods = [modality] if modality else (modalities or []) + if step == "step_0": + for m in mods: + timeseries_grid(paths, m, n=min(n, 4)) + elif step == "step_1": + for m in mods: + spectrogram_grid(paths, "step_1", m, n=n) + elif step == "step_2": + for m in mods: + pair_grid(paths, "step_1", "step_2", m, n=min(n, 4)) + elif step == "step_3": + for m in mods: + pair_grid(paths, "step_2", "step_3", m, n=min(n, 4), log_a=False) + elif step == "step_4": + for m in mods: + mask_overlay(paths, m, n=min(n, 4)) + knee_plot(paths, m) + elif step == "step_5": + dataset_grid(paths, n=n) + elif step == "step_6": + refine_triptych(paths, n=min(n, 4)) + elif step == "step_7": + print(yaml.safe_dump(yaml.safe_load(paths.deploy_manifest.read_text()))) + elif step == "step_8": + eval_curves(paths) + else: + raise KeyError(f"No gallery view for {step}") diff --git a/src/tokeye/training/big_tf_unet_2/notebook_api.py b/src/tokeye/training/big_tf_unet_2/notebook_api.py new file mode 100644 index 0000000..87f86ba --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebook_api.py @@ -0,0 +1,193 @@ +"""The ``Run`` facade — everything the notebooks call, nothing else. + +Every notebook cell is 1-3 calls into this class. The loop per step: + + run.suggest("step_2") # what the autos picked + which knobs matter + run.run("step_2") # or run.submit(...) for GPU/heavy steps + run.status() # table; rerun the cell to refresh + run.gallery("step_2") # look at the pictures + run.accept("step_2") # sign off — or edit run.yaml, run.clear, redo + +Guardrails: inline execution refuses GPU steps and steps whose upstream +isn't complete+accepted; every error surfaces as a one-line message naming +what to do next. +""" + +from __future__ import annotations + +import subprocess +import time + +import pandas as pd + +from . import auto_resolve, gallery, slurm +from .clearing import clear_run, clear_step +from .paths import STEP_ORDER, RunPaths, get_step +from .run_config import check_scale_lock, load_run_config +from .runner import run_step +from .scaffold import scaffold_run +from .task_matrix import RunTaskMatrix + + +class Run: + """Handle on one pipeline run, addressed by its run_id.""" + + def __init__(self, run_id: str) -> None: + self.run_id = run_id + self.paths = RunPaths(run_id) + if not self.paths.run_yaml.exists(): + raise FileNotFoundError( + f"No run.yaml at {self.paths.run_yaml} — create the run first: " + f"python -m tokeye.training.big_tf_unet_2.scaffold " + f"--nfft --hop " + ) + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + + @classmethod + def open(cls, run_id: str) -> Run: + return cls(run_id) + + @classmethod + def create(cls, nfft: int, hop: int, suffix: str = "") -> Run: + paths = scaffold_run(nfft, hop, suffix) + return cls(paths.run_id) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + @property + def cfg(self): + cfg = load_run_config(self.paths.run_yaml) + check_scale_lock(cfg, self.paths.run_meta) + return cfg + + @property + def _matrix(self) -> RunTaskMatrix: + return RunTaskMatrix(self.paths.task_matrix_path) + + def _upstream_unaccepted(self, step: str) -> list[str]: + mods = self.cfg.modality_names + matrix = self._matrix + return [ + name + for name in STEP_ORDER[: STEP_ORDER.index(step)] + if not matrix.is_accepted(name, mods) + ] + + # ------------------------------------------------------------------ + # Status + # ------------------------------------------------------------------ + + def status(self) -> pd.DataFrame: + """Progress table; SLURM state is refreshed for submitted jobs.""" + rows = self._matrix.to_rows(self.cfg.modality_names) + for row in rows: + if row["status"] in ("submitted", "running") and row["job_id"]: + row["slurm"] = slurm.job_state(str(row["job_id"])) + else: + row["slurm"] = "" + return pd.DataFrame(rows) + + def suggest(self, step: str) -> None: + """Show auto-resolved values + where the step's knobs live.""" + cfg = self.cfg + spec = get_step(step) + mods = cfg.modality_names if spec.per_modality else [None] + for mod in mods: + values = auto_resolve.suggest(cfg, step, mod, self.paths) + label = f"{step}" + (f" [{mod}]" if mod else "") + if values: + print(f"{label} auto suggestions:") + for k, v in values.items(): + print(f" {k} = {v}") + else: + print(f"{label}: no auto values (or upstream not run yet)") + self.edit_hint(step) + + def edit_hint(self, step: str) -> None: + section = get_step(step).knob_section + print(f"knobs: section '{section}:' in {self.paths.run_yaml}") + + # ------------------------------------------------------------------ + # Execution + # ------------------------------------------------------------------ + + def run(self, step: str, modality: str | None = None, force: bool = False): + """Run a step inline. Refuses GPU steps and unaccepted upstream.""" + spec = get_step(step) + if spec.exec_mode == "sbatch_gpu" and not force: + print(f"{step} is a GPU step — use run.submit({step!r}) instead") + return + unaccepted = [] if force else self._upstream_unaccepted(step) + if unaccepted: + print( + f"{step} needs your sign-off on {', '.join(unaccepted)} first " + f"(run.accept(...) after checking the gallery), or force=True" + ) + return + run_step(self.paths.run_yaml, step, modality, force) + + def submit(self, steps: str | list[str], modality: str | None = None) -> str: + """Submit step(s) as one SLURM job; returns the job id.""" + first = steps.split(",")[0] if isinstance(steps, str) else steps[0] + unaccepted = self._upstream_unaccepted(first.strip()) + if unaccepted: + print( + f"heads-up: upstream not signed off yet: {', '.join(unaccepted)}" + ) + job_id = slurm.submit_step(self.paths.run_yaml, steps, modality) + print(f"submitted job {job_id} — run.status() / run.log({job_id!r})") + return job_id + + def wait(self, job_id: str, poll: int = 30) -> str: + """Block until a SLURM job leaves the queue; returns the final state.""" + while True: + state = slurm.job_state(job_id) + if state in ("PENDING", "RUNNING", "CONFIGURING", "COMPLETING"): + time.sleep(poll) + continue + print(f"job {job_id}: {state}") + return state + + def log(self, job_id: str, n: int = 40) -> None: + print(slurm.tail_log(self.paths, str(job_id), n)) + + # ------------------------------------------------------------------ + # Inspection + sign-off + # ------------------------------------------------------------------ + + def gallery(self, step: str, modality: str | None = None, n: int = 6) -> None: + gallery.show( + self.paths, step, modality, modalities=self.cfg.modality_names, n=n + ) + + def accept(self, step: str) -> None: + try: + self._matrix.accept(step, self.cfg.modality_names) + print(f"{step} accepted — next step unlocked") + except ValueError as err: + print(err) + + def clear(self, step: str, modality: str | None = None) -> None: + removed = clear_step(self.paths, self.cfg, step, modality) + print( + f"cleared {len(removed)} artifact(s) for {step}" + + (f" [{modality}]" if modality else "") + + "; downstream marked stale" + ) + + def clear_all(self, confirm: str) -> None: + """Delete every cache/model artifact. confirm must be the run_id.""" + clear_run(self.paths, confirm) + print(f"run {self.run_id} cache cleared (run.yaml + notebooks kept)") + + def jobstats(self, job_id: str) -> None: + """Cluster efficiency report for a finished/running job.""" + result = subprocess.run( + ["jobstats", str(job_id)], capture_output=True, text=True + ) + print(result.stdout or result.stderr) diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/00_setup.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/00_setup.ipynb new file mode 100644 index 0000000..420d63b --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/00_setup.ipynb @@ -0,0 +1,62 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 00 \u2014 Setup\n", + "\n", + "This notebook opens your run and shows its status. Your only editable file is `run.yaml` in this folder \u2014 every knob lives there.\n", + "\n", + "The loop for every step is the same:\n", + "\n", + "1. **Suggest** \u2014 see what values the pipeline picked automatically.\n", + "2. **Run** (light steps) or **Submit** (cluster steps) \u2014 then re-run the status cell until it says `complete`.\n", + "3. **Gallery** \u2014 look at the pictures. Do the spectrograms/masks look right?\n", + "4. **Decide** \u2014 if yes: `run.accept(...)`. If no: edit the knob in `run.yaml`, `run.clear(...)`, and run it again." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Re-run the cell above any time to refresh the table.\n", + "\n", + "- `status` column: pending \u2192 submitted \u2192 running \u2192 complete (or **failed** / **stale**)\n", + "- **stale** means you changed a knob upstream \u2014 that step must be rerun before training\n", + "- `accepted` is your sign-off from the gallery" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Danger zone: wipe every computed artifact of this run (keeps run.yaml + notebooks).\n", + "# To use it, type the run id yourself:\n", + "# run.clear_all(confirm=\"...\")" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/01_spectrograms.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/01_spectrograms.ipynb new file mode 100644 index 0000000..277a636 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/01_spectrograms.ipynb @@ -0,0 +1,166 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 01 \u2014 Signals \u2192 Spectrograms \u2192 Baseline removal (steps 0, 1, 2)\n", + "\n", + "**step_0** loads the raw signals and cuts them into windows. **step_1** computes the spectrograms and keeps only the most active windows. **step_2** removes the smooth background (baseline) so fluctuating modes stand out.\n", + "\n", + "Knobs you might touch here (in `run.yaml`): `baseline.lam` (bigger = smoother background), `baseline.edge_k` (bigger = fewer edge frequency bins cut)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_0 \u2014 intake (runs here, a few minutes)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.run(\"step_0\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_0\", n=2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_0\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_1 \u2014 spectrograms + window filter (runs here)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.suggest(\"step_1\")\n", + "run.run(\"step_1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_2 \u2014 baseline removal (heavy: submits to a CPU node)\n", + "\n", + "After submitting, re-run the status/log cell until complete, then look at the gallery." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.suggest(\"step_2\")\n", + "job = run.submit(\"step_2\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.log(job)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_2\") # left: before, right: after baseline removal" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_2\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Not happy?** Edit `baseline:` in run.yaml, then `run.clear(\"step_2\")` and resubmit." + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/02_denoise.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/02_denoise.ipynb new file mode 100644 index 0000000..eb5f73f --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/02_denoise.ipynb @@ -0,0 +1,90 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 02 \u2014 Denoiser (step_3, GPU)\n", + "\n", + "A small network learns to predict each channel from its neighbors \u2014 what it can predict is coherent signal, what it can't is noise. Submits to a GPU node; training takes a while.\n", + "\n", + "Knob: `denoise.max_epochs` (more = cleaner, slower)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.suggest(\"step_3\")\n", + "job = run.submit(\"step_3\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.log(job)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.jobstats(job) # after it finishes: was the GPU actually busy?" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_3\") # left: input residual, right: denoised" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_3\")" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/03_labels.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/03_labels.ipynb new file mode 100644 index 0000000..df18b06 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/03_labels.ipynb @@ -0,0 +1,76 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 03 \u2014 Label masks (step_4)\n", + "\n", + "Thresholds the denoised fields into coherent/transient masks. The threshold is found automatically at the 'knee' of the intensity distribution \u2014 your knobs shift it:\n", + "\n", + "- `labels.knee_sensitivity` \u2014 bigger = higher threshold = fewer labeled pixels\n", + "- `labels.delta` \u2014 direct offset in noise-sigma units (e.g. 0.5 = half a sigma stricter)\n", + "- `labels.min_size_fraction` \u2014 bigger = small specks removed more aggressively\n", + "\n", + "This step runs in seconds \u2014 it is the tightest tweak-and-look loop in the pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.run(\"step_4\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_4\") # masks over the denoised field + the thresholds picked" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Not happy? Edit labels: in run.yaml, then:\n", + "# run.clear(\"step_4\"); run.run(\"step_4\"); run.gallery(\"step_4\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_4\")" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/04_train.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/04_train.ipynb new file mode 100644 index 0000000..9dd5873 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/04_train.ipynb @@ -0,0 +1,173 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 04 \u2014 Dataset \u2192 Refine \u2192 Final model (steps 5, 6, 7)\n", + "\n", + "**step_5** packs all modalities into one training set (runs here). **step_6** trains 5 cross-validation models and records where they disagree with your masks (GPU). **step_7** trains the final model on the blend (GPU).\n", + "\n", + "Knob: `refine.model_trust` \u2014 0 = trust the step_4 masks, 1 = trust the models. Changing it only re-runs step_7 (cheap), not step_6." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_5 \u2014 dataset (runs here)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.run(\"step_5\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_5\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_5\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_6 \u2014 cross-validation refine (GPU, the long one)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "job = run.submit(\"step_6\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.log(job)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_6\") # shows the samples where model and masks disagree MOST" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_6\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## step_7 \u2014 final model (GPU)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "job = run.submit(\"step_7\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.log(job)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_7\") # the deploy manifest" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_7\")" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/notebooks/05_eval.ipynb b/src/tokeye/training/big_tf_unet_2/notebooks/05_eval.ipynb new file mode 100644 index 0000000..ced2e72 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/notebooks/05_eval.ipynb @@ -0,0 +1,60 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 05 \u2014 TJ-II benchmark (step_8)\n", + "\n", + "Runs the exported model on the public TJ-II set and reports IoU/F1 at the best threshold. The deployed reference model scores IoU 0.260 / F1 0.478 \u2014 but TJ-II images are rendered at one fixed resolution, so for a single-scale teacher treat this as a **relative tracker across your runs**, not a target." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from tokeye.training.big_tf_unet_2.notebook_api import Run\n", + "\n", + "RUN_ID = \"__RUN_ID__\"\n", + "run = Run.open(RUN_ID)\n", + "run.status()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.run(\"step_8\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.gallery(\"step_8\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run.accept(\"step_8\")" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/tokeye/training/big_tf_unet_2/paths.py b/src/tokeye/training/big_tf_unet_2/paths.py new file mode 100644 index 0000000..d68f971 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/paths.py @@ -0,0 +1,203 @@ +"""Single source of truth for run identity, step registry, and artifact paths. + +One pipeline run = one (nfft, hop) scale. Every artifact of a run lives under +``data/cache/big_tf_unet_2//`` (pipeline caches) and +``model/big_tf_unet_2//`` (trained weights). The intern workspace +(run.yaml + notebooks) lives in ``dev/training//`` (gitignored). + +All paths are repo-root relative — run every entry point from the repo root, +matching the convention of the other training pipelines. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +# The canonical scale grid (migrated from big_tf_unet/step_7b_train_multiscale). +SCALE_CONFIGS: list[tuple[int, int]] = [ + (128, 64), + (256, 64), + (256, 128), + (512, 128), + (512, 256), + (1024, 128), + (1024, 256), + (1024, 512), + (2048, 256), + (2048, 512), +] + + +def run_id_for(nfft: int, hop: int, suffix: str = "") -> str: + base = f"nfft{nfft}_hop{hop}" + return f"{base}_{suffix}" if suffix else base + + +def repo_root() -> Path: + """Locate the repo root robustly, independent of the current directory. + + Notebooks run with their own directory as cwd, so the root cannot be + ``Path.cwd()``. Walk up from cwd for a ``pyproject.toml`` (handles running + from anywhere inside a checkout); fall back to this module's fixed location + (``/src/tokeye/training/big_tf_unet_2/paths.py``). + """ + for candidate in (Path.cwd(), *Path.cwd().parents): + if (candidate / "pyproject.toml").exists(): + return candidate + return Path(__file__).resolve().parents[4] + + +@dataclass(frozen=True) +class StepSpec: + """Static description of one pipeline step.""" + + name: str # e.g. "step_2" + title: str # short human label for status tables / notebooks + module: str # module basename inside this package + per_modality: bool # True: runs once per modality, artifacts under / + exec_mode: str # "inline" | "sbatch_cpu" | "sbatch_gpu" + knob_section: str # run.yaml section holding this step's knobs + + +STEPS: list[StepSpec] = [ + StepSpec("step_0", "intake", "step_0_intake", True, "inline", "extraction"), + StepSpec( + "step_1", "spectrogram", "step_1_spectrogram", True, "inline", "window_filter" + ), + StepSpec("step_2", "baseline", "step_2_baseline", True, "sbatch_cpu", "baseline"), + StepSpec("step_3", "denoise", "step_3_denoise", True, "sbatch_gpu", "denoise"), + StepSpec("step_4", "labels", "step_4_labels", True, "inline", "labels"), + StepSpec("step_5", "dataset", "step_5_dataset", False, "inline", "dataset"), + StepSpec("step_6", "refine", "step_6_refine", False, "sbatch_gpu", "refine"), + StepSpec("step_7", "final", "step_7_final", False, "sbatch_gpu", "final"), + StepSpec("step_8", "eval", "step_8_eval", False, "inline", "eval"), +] + +STEP_ORDER: list[str] = [s.name for s in STEPS] +_STEP_BY_NAME: dict[str, StepSpec] = {s.name: s for s in STEPS} + + +def get_step(name: str) -> StepSpec: + try: + return _STEP_BY_NAME[name] + except KeyError: + raise KeyError( + f"Unknown step {name!r}. Valid steps: {', '.join(STEP_ORDER)}" + ) from None + + +def steps_after(name: str) -> list[StepSpec]: + """Steps strictly downstream of ``name`` in pipeline order.""" + idx = STEP_ORDER.index(get_step(name).name) + return STEPS[idx + 1 :] + + +@dataclass(frozen=True) +class RunPaths: + """Every artifact path of one run, derived from its run_id.""" + + run_id: str + root: Path = field(default_factory=repo_root) + + # ------------------------------------------------------------------ + # Roots + # ------------------------------------------------------------------ + + @property + def cache_root(self) -> Path: + return self.root / "data" / "cache" / "big_tf_unet_2" / self.run_id + + @property + def model_dir(self) -> Path: + return self.root / "model" / "big_tf_unet_2" / self.run_id + + @property + def workspace(self) -> Path: + return self.root / "dev" / "training" / self.run_id + + # ------------------------------------------------------------------ + # Run-level files + # ------------------------------------------------------------------ + + @property + def run_yaml(self) -> Path: + return self.workspace / "run.yaml" + + @property + def run_meta(self) -> Path: + return self.cache_root / "run_meta.json" + + @property + def task_matrix_path(self) -> Path: + return self.cache_root / "task_matrix.json" + + @property + def resolved_params_path(self) -> Path: + return self.cache_root / "resolved_params.yaml" + + @property + def slurm_dir(self) -> Path: + return self.cache_root / "slurm" + + # ------------------------------------------------------------------ + # Step artifacts + # ------------------------------------------------------------------ + + def mod_dir(self, modality: str) -> Path: + return self.cache_root / modality + + def step_h5(self, step: str, modality: str | None = None) -> Path: + spec = get_step(step) + if spec.per_modality: + if modality is None: + raise ValueError(f"{step} is per-modality; a modality is required") + return self.mod_dir(modality) / f"{step}.h5" + return self.cache_root / f"{step}.h5" + + def baseline_h5(self, modality: str) -> Path: + return self.mod_dir(modality) / "step_2_baseline.h5" + + def frame_info(self, modality: str, raw: bool = False) -> Path: + name = "frame_info_raw.csv" if raw else "frame_info.csv" + return self.mod_dir(modality) / name + + def thresholds_csv(self, modality: str) -> Path: + return self.mod_dir(modality) / "thresholds.csv" + + @property + def eval_csv(self) -> Path: + return self.cache_root / "eval_tjii.csv" + + @property + def deploy_manifest(self) -> Path: + return self.model_dir / "deploy_manifest.yaml" + + def artifacts(self, step: str, modalities: list[str]) -> list[Path]: + """Every file/dir a step writes — the exact clear list for that step.""" + spec = get_step(step) + out: list[Path] = [] + if spec.per_modality: + for mod in modalities: + out.append(self.step_h5(step, mod)) + if step == "step_0": + out.append(self.frame_info(mod, raw=True)) + elif step == "step_1": + out.append(self.frame_info(mod)) + elif step == "step_2": + out.append(self.baseline_h5(mod)) + elif step == "step_4": + out.append(self.thresholds_csv(mod)) + elif step == "step_7": + out.append(self.model_dir) + elif step == "step_8": + out.append(self.eval_csv) + else: + out.append(self.step_h5(step)) + return out + + +NOTEBOOK_TEMPLATE_DIR = Path(__file__).parent / "notebooks" +CONFIG_DIR = Path(__file__).parent / "config" +DEFAULTS_YAML = CONFIG_DIR / "defaults.yaml" +RUN_TEMPLATE_YAML = CONFIG_DIR / "run_template.yaml" diff --git a/src/tokeye/training/big_tf_unet_2/run_config.py b/src/tokeye/training/big_tf_unet_2/run_config.py new file mode 100644 index 0000000..e4b2bd4 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/run_config.py @@ -0,0 +1,241 @@ +"""Typed, validated run configuration. + +``load_run_config`` merges the intern's ``run.yaml`` ON TOP of the bundled +``config/defaults.yaml`` (so run.yaml only carries overrides — the +replace-not-merge behavior of the legacy ``load_settings`` is retired) and +validates the result into a pydantic model. Bad knob values die here, before +any compute, with one-line messages naming the field and the allowed range. +Unknown keys (typos) are rejected too. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Literal + +from omegaconf import OmegaConf +from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator + +from .paths import DEFAULTS_YAML, SCALE_CONFIGS + +Auto = Literal["auto"] + + +class ConfigError(ValueError): + """A run.yaml problem, formatted for humans (one line per issue).""" + + +class _Section(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class RunSection(_Section): + nfft: int = Field(gt=0) + hop: int = Field(gt=0) + allow_custom_scale: bool = False + + @model_validator(mode="after") + def _check_scale(self) -> RunSection: + if self.hop > self.nfft: + raise ValueError(f"hop ({self.hop}) must be <= nfft ({self.nfft})") + if not self.allow_custom_scale and (self.nfft, self.hop) not in SCALE_CONFIGS: + grid = ", ".join(f"{n}/{h}" for n, h in SCALE_CONFIGS) + raise ValueError( + f"(nfft={self.nfft}, hop={self.hop}) is not in the scale grid " + f"[{grid}]. Set run.allow_custom_scale: true to use it anyway." + ) + return self + + +class ModalityConfig(_Section): + input_key: str + channels: list[int] = Field(min_length=1) + + +class ExtractionSection(_Section): + subseq_len: int = Field(gt=0) + preemphasis_coeff: float = Field(ge=0.0, le=1.0) + fs_khz: float = Field(gt=0) + target_rate_khz: float = Field(gt=0) + ip_threshold: float = Field(ge=0.0) + max_windows_per_shot_precap: int = Field(gt=0) + + +class WindowFilterSection(_Section): + enabled: bool + weights: str + max_windows_per_shot: int = Field(gt=0) + activity_threshold: float = Field(gt=0.0, lt=1.0) + min_activity: float = Field(ge=0.0, lt=1.0) + mean: float | Auto + std: float | Auto + + +class BaselineSection(_Section): + method: str + lam: float | Auto + edge_method: Literal["energy", "gradient"] + edge_k: float = Field(gt=1.0) + edge_max_fraction: float = Field(gt=0.0, le=0.5) + gradient_threshold: float = Field(gt=0.0) + + +class DenoiseSection(_Section): + representation: Literal["complex", "magnitude"] + normalization: Literal["robust_asinh", "zscore"] + a: float = Field(gt=0.0) + first_layer_size: int = Field(gt=0) + num_layers: int | Auto + batch_size: int | Auto + base_batch_size: int = Field(gt=0) + precision: str + max_epochs: int = Field(gt=0) + tv_patience: int = Field(ge=0) + num_workers: int = Field(ge=0) + + +class LabelsSection(_Section): + knee_sensitivity: float = Field(gt=0.0) + delta: float + fallback_frac: float = Field(gt=0.0, lt=0.5) + min_size: int | Auto + min_size_fraction: float = Field(gt=0.0, lt=0.1) + remove_bottom_rows: int | Auto + remove_top_rows: int | Auto + row_removal_fraction_bottom: float = Field(ge=0.0, lt=0.2) + row_removal_fraction_top: float = Field(ge=0.0, lt=0.2) + + +class DatasetSection(_Section): + a: float = Field(gt=0.0) + stats_windows: int = Field(gt=0) + + +class RefineSection(_Section): + model_trust: float = Field(ge=0.0, le=1.0) + n_folds: int = Field(ge=2) + first_layer_size: int = Field(gt=0) + num_layers: int | Auto + batch_size: int | Auto + base_batch_size: int = Field(gt=0) + precision: str + max_epochs: int = Field(gt=0) + loss_type: str + num_workers: int = Field(ge=0) + + +class FinalSection(_Section): + first_layer_size: int = Field(gt=0) + num_layers: int | Auto + batch_size: int | Auto + base_batch_size: int = Field(gt=0) + precision: str + max_epochs: int = Field(gt=0) + loss_type: str + gamma: float = Field(gt=0.0) + num_workers: int = Field(ge=0) + + +class EvalSection(_Section): + dataset_dir: str + n_thresholds: int = Field(ge=2) + + +class PathsSection(_Section): + shots_path: str + foundation_dir: str + + +class SlurmSection(_Section): + account: str | None + cpu_partition: str | None + gpu_partition: str | None + cpu_cpus: int = Field(gt=0) + cpu_mem: str + cpu_time: str + gpu_cpus: int = Field(gt=0) + gpu_mem: str + gpu_time: str + + +class SmokeSection(_Section): + enabled: bool + n_shots: int = Field(gt=0) + max_windows_per_shot: int = Field(gt=0) + n_folds: int = Field(ge=2) + max_epochs: int = Field(gt=0) + refine_max_epochs: int = Field(gt=0) + final_max_epochs: int = Field(gt=0) + + +class RunConfig(_Section): + run: RunSection + modalities: dict[str, ModalityConfig] = Field(min_length=1) + extraction: ExtractionSection + window_filter: WindowFilterSection + baseline: BaselineSection + denoise: DenoiseSection + labels: LabelsSection + dataset: DatasetSection + refine: RefineSection + final: FinalSection + eval: EvalSection + paths: PathsSection + slurm: SlurmSection + smoke: SmokeSection + + @property + def modality_names(self) -> list[str]: + return list(self.modalities) + + @property + def n_freq(self) -> int: + return self.run.nfft // 2 + 1 + + @property + def n_time(self) -> int: + return self.extraction.subseq_len // self.run.hop + 1 + + +def _format_validation_error(err: ValidationError) -> str: + lines = [] + for issue in err.errors(): + loc = ".".join(str(p) for p in issue["loc"]) + lines.append(f" {loc}: {issue['msg']}") + n = len(lines) + plural = "s" if n != 1 else "" + return f"{n} problem{plural} in run.yaml:\n" + "\n".join(lines) + + +def load_run_config(run_yaml: str | Path) -> RunConfig: + """Merge run.yaml over the bundled defaults and validate.""" + run_yaml = Path(run_yaml) + if not run_yaml.exists(): + raise ConfigError(f"run.yaml not found: {run_yaml}") + merged = OmegaConf.merge(OmegaConf.load(DEFAULTS_YAML), OmegaConf.load(run_yaml)) + raw = OmegaConf.to_container(merged, resolve=True) + try: + return RunConfig.model_validate(raw) + except ValidationError as err: + raise ConfigError(_format_validation_error(err)) from None + + +def check_scale_lock(cfg: RunConfig, run_meta_path: Path) -> None: + """The scale in run.yaml must match the scale this run was created with. + + A different scale is a different run (different run_id, cache root, and + workspace) — changing nfft/hop in an existing run.yaml would silently mix + artifacts, so it is refused here. + """ + if not run_meta_path.exists(): + return + meta = json.loads(run_meta_path.read_text()) + locked = (meta.get("nfft"), meta.get("hop")) + current = (cfg.run.nfft, cfg.run.hop) + if locked != current: + raise ConfigError( + f"run.yaml scale nfft={current[0]}/hop={current[1]} does not match " + f"this run's locked scale nfft={locked[0]}/hop={locked[1]}. " + f"To train a different scale, scaffold a new run instead." + ) diff --git a/src/tokeye/training/big_tf_unet_2/runner.py b/src/tokeye/training/big_tf_unet_2/runner.py new file mode 100644 index 0000000..f06d7a1 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/runner.py @@ -0,0 +1,320 @@ +"""Step orchestration: settings builders, gating, and the CLI entry point. + +The same entry point serves notebook inline cells and sbatch payloads: + + python -m tokeye.training.big_tf_unet_2.runner \ + --run-config dev/training/nfft512_hop128/run.yaml --steps step_0,step_1 + +Contract with step modules: each ``step_N_*.py`` exposes +``main(settings: dict) -> dict | None``. The runner builds the settings dict +(config values + resolved autos + artifact paths), clears the step's previous +artifacts, tracks status in the task matrix, and records any in-step resolved +values the step returns (e.g. per-modality stats) into the ledger. +""" + +from __future__ import annotations + +import argparse +import importlib +import logging +from pathlib import Path +from typing import Any + +from . import auto_resolve +from .clearing import clear_step +from .paths import STEP_ORDER, RunPaths, get_step +from .run_config import RunConfig, check_scale_lock, load_run_config +from .task_matrix import RunTaskMatrix, params_hash + +logger = logging.getLogger(__name__) + +# Keys excluded from the params hash: they change resources, never outputs. +_VOLATILE_KEYS = {"n_workers", "num_workers"} + + +# --------------------------------------------------------------------------- +# Settings builders +# --------------------------------------------------------------------------- + +def build_step_settings( + cfg: RunConfig, step: str, modality: str | None, paths: RunPaths +) -> dict[str, Any]: + """Assemble one step's settings dict (config + autos + paths).""" + smoke = cfg.smoke.enabled + resolved = auto_resolve.resolve_step_autos(cfg, step, modality, paths) + + def auto(section_value: Any, key: str) -> Any: + return resolved[key] if section_value == "auto" else section_value + + s: dict[str, Any] = {"run_id": paths.run_id, "smoke": smoke} + + def rooted(rel: str) -> Path: + """Resolve a config path against the repo root (cwd-independent).""" + p = Path(rel) + return p if p.is_absolute() else paths.root / p + + if step == "step_0": + mod = cfg.modalities[modality] + ext = cfg.extraction + s.update( + shots_path=rooted(cfg.paths.shots_path), + foundation_dir=rooted(cfg.paths.foundation_dir), + modality=modality, + input_key=mod.input_key, + channels=list(mod.channels), + subseq_len=ext.subseq_len, + preemphasis_coeff=ext.preemphasis_coeff, + fs_khz=ext.fs_khz, + target_rate_khz=ext.target_rate_khz, + ip_threshold=ext.ip_threshold, + max_windows_per_shot=( + cfg.smoke.max_windows_per_shot + if smoke + else ext.max_windows_per_shot_precap + ), + n_shots=cfg.smoke.n_shots if smoke else None, + out_h5=paths.step_h5("step_0", modality), + frame_info_csv=paths.frame_info(modality, raw=True), + ) + + elif step == "step_1": + wf = cfg.window_filter + s.update( + in_h5=paths.step_h5("step_0", modality), + out_h5=paths.step_h5("step_1", modality), + frame_info_in=paths.frame_info(modality, raw=True), + frame_info_out=paths.frame_info(modality), + nfft=cfg.run.nfft, + hop=cfg.run.hop, + filter_enabled=wf.enabled, + weights=rooted(wf.weights), + max_windows_per_shot=( + cfg.smoke.max_windows_per_shot if smoke else wf.max_windows_per_shot + ), + activity_threshold=wf.activity_threshold, + min_activity=wf.min_activity, + mean=wf.mean, # "auto" -> step computes per-modality logmag stats + std=wf.std, + ) + + elif step == "step_2": + b = cfg.baseline + s.update( + in_h5=paths.step_h5("step_1", modality), + out_h5=paths.step_h5("step_2", modality), + baseline_h5=paths.baseline_h5(modality), + method=b.method, + lam=auto(b.lam, "lam"), + edge_bins_lower=resolved.get("edge_bins_lower", 1), + edge_bins_upper=resolved.get("edge_bins_upper", 1), + n_workers=None, # step resolves: setting > SLURM_CPUS_PER_TASK > cpus + ) + + elif step == "step_3": + d = cfg.denoise + n_channels = len(cfg.modalities[modality].channels) + s.update( + in_h5=paths.step_h5("step_2", modality), + out_h5=paths.step_h5("step_3", modality), + representation=d.representation, + normalization=d.normalization, + a=d.a, + total_channels=n_channels, + adjacent_channels=max(1, n_channels // 2), + first_layer_size=d.first_layer_size, + num_layers=auto(d.num_layers, "num_layers"), + batch_size=auto(d.batch_size, "batch_size"), + precision=d.precision, + max_epochs=cfg.smoke.max_epochs if smoke else d.max_epochs, + tv_patience=d.tv_patience, + num_workers=d.num_workers, + edge_bins_lower=resolved.get("edge_bins_lower", 1), + edge_bins_upper=resolved.get("edge_bins_upper", 1), + ) + + elif step == "step_4": + la = cfg.labels + s.update( + denoised_h5=paths.step_h5("step_3", modality), + baseline_h5=paths.baseline_h5(modality), + out_h5=paths.step_h5("step_4", modality), + frame_info=paths.frame_info(modality), + thresholds_csv=paths.thresholds_csv(modality), + knee_sensitivity=la.knee_sensitivity, + delta=la.delta, + fallback_frac=la.fallback_frac, + min_size=auto(la.min_size, "min_size"), + remove_bottom_rows=auto(la.remove_bottom_rows, "remove_bottom_rows"), + remove_top_rows=auto(la.remove_top_rows, "remove_top_rows"), + ) + + elif step == "step_5": + s.update( + inputs={ + mod: { + "img_h5": paths.step_h5("step_1", mod), + "mask_h5": paths.step_h5("step_4", mod), + "frame_info": paths.frame_info(mod), + } + for mod in cfg.modality_names + }, + out_h5=paths.step_h5("step_5"), + a=cfg.dataset.a, + stats_windows=cfg.dataset.stats_windows, + ) + + elif step == "step_6": + r = cfg.refine + s.update( + in_h5=paths.step_h5("step_5"), + out_h5=paths.step_h5("step_6"), + ckpt_dir=paths.cache_root / "step_6_ckpts", + n_folds=cfg.smoke.n_folds if smoke else r.n_folds, + first_layer_size=r.first_layer_size, + num_layers=auto(r.num_layers, "num_layers"), + batch_size=auto(r.batch_size, "batch_size"), + precision=r.precision, + max_epochs=cfg.smoke.refine_max_epochs if smoke else r.max_epochs, + loss_type=r.loss_type, + num_workers=r.num_workers, + ) + + elif step == "step_7": + f = cfg.final + s.update( + dataset_h5=paths.step_h5("step_5"), + refine_h5=paths.step_h5("step_6"), + # model_trust applies HERE (not in step_6): tuning the lambda knob + # re-trains one final model instead of five folds. + model_trust=cfg.refine.model_trust, + model_dir=paths.model_dir, + deploy_manifest=paths.deploy_manifest, + resolved_params=paths.resolved_params_path, + nfft=cfg.run.nfft, + hop=cfg.run.hop, + first_layer_size=f.first_layer_size, + num_layers=auto(f.num_layers, "num_layers"), + batch_size=auto(f.batch_size, "batch_size"), + precision=f.precision, + max_epochs=cfg.smoke.final_max_epochs if smoke else f.max_epochs, + loss_type=f.loss_type, + gamma=f.gamma, + num_workers=f.num_workers, + ) + + elif step == "step_8": + s.update( + model_dir=paths.model_dir, + dataset_dir=rooted(cfg.eval.dataset_dir), + n_thresholds=cfg.eval.n_thresholds, + out_csv=paths.eval_csv, + ) + + else: + raise KeyError(f"No settings builder for {step}") + + return s + + +# --------------------------------------------------------------------------- +# Gating + execution +# --------------------------------------------------------------------------- + +def _check_upstream( + matrix: RunTaskMatrix, cfg: RunConfig, step: str, modality: str | None +) -> None: + target = get_step(step) + for name in STEP_ORDER[: STEP_ORDER.index(step)]: + spec = get_step(name) + if spec.per_modality: + mods = [modality] if (target.per_modality and modality) else ( + cfg.modality_names + ) + missing = [m for m in mods if not matrix.is_complete(name, m)] + if missing: + raise RuntimeError( + f"{step} needs {name} complete for {', '.join(missing)} " + f"first (rerun with --force to override)" + ) + elif not matrix.is_complete(name): + raise RuntimeError( + f"{step} needs {name} complete first " + f"(rerun with --force to override)" + ) + + +def run_step( + run_yaml: str | Path, + step: str, + modality: str | None = None, + force: bool = False, +) -> None: + """Run one step (one modality for per-modality steps) end to end.""" + run_yaml = Path(run_yaml).resolve() + cfg = load_run_config(run_yaml) + paths = RunPaths(run_yaml.parent.name) + check_scale_lock(cfg, paths.run_meta) + + spec = get_step(step) + if spec.per_modality and modality is None: + for mod in cfg.modality_names: + run_step(run_yaml, step, mod, force) + return + + matrix = RunTaskMatrix(paths.task_matrix_path) + if not force: + _check_upstream(matrix, cfg, step, modality) + + settings = build_step_settings(cfg, step, modality, paths) + pre_resolved = auto_resolve.resolve_step_autos(cfg, step, modality, paths) + auto_resolve.record(paths, step, modality, pre_resolved, source="auto") + + # Idempotent rerun: previous artifacts go away before the step starts, and + # the matrix (not file existence) is the source of truth for completion. + clear_step(paths, cfg, step, modality) + for target in paths.artifacts(step, [modality] if modality else cfg.modality_names): + target.parent.mkdir(parents=True, exist_ok=True) + + hashable = {k: v for k, v in settings.items() if k not in _VOLATILE_KEYS} + matrix.mark_running(step, modality) + label = f"{step}" + (f":{modality}" if modality else "") + logger.info(f"[{paths.run_id}] running {label}") + try: + module = importlib.import_module(f".{spec.module}", package=__package__) + in_step = module.main(settings) + except Exception: + matrix.mark_failed(step, modality) + raise + if in_step: + auto_resolve.record(paths, step, modality, in_step, source="in-step") + matrix.mark_complete(step, modality, params_hash(hashable), cfg.modality_names) + logger.info(f"[{paths.run_id}] {label} complete") + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(name)s %(message)s" + ) + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-config", required=True) + parser.add_argument( + "--steps", required=True, help="comma-separated, e.g. step_0,step_1" + ) + parser.add_argument( + "--modalities", default=None, help="comma-separated subset (default: all)" + ) + parser.add_argument("--force", action="store_true") + args = parser.parse_args() + + mods = args.modalities.split(",") if args.modalities else [None] + for step in args.steps.split(","): + step = step.strip() + if get_step(step).per_modality and args.modalities: + for mod in mods: + run_step(args.run_config, step, mod, args.force) + else: + run_step(args.run_config, step, None, args.force) + + +if __name__ == "__main__": + main() diff --git a/src/tokeye/training/big_tf_unet_2/scaffold.py b/src/tokeye/training/big_tf_unet_2/scaffold.py new file mode 100644 index 0000000..8544ee1 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/scaffold.py @@ -0,0 +1,95 @@ +"""Create a new run: workspace (run.yaml + notebooks), cache dirs, scale lock. + +Usage (from the repo root): + + python -m tokeye.training.big_tf_unet_2.scaffold --nfft 512 --hop 128 + python -m tokeye.training.big_tf_unet_2.scaffold --nfft 512 --hop 128 --suffix v2 + +Refuses to overwrite an existing run.yaml — a run's knob file is its +experiment record. +""" + +from __future__ import annotations + +import argparse +import json +from datetime import datetime + +from .paths import ( + NOTEBOOK_TEMPLATE_DIR, + RUN_TEMPLATE_YAML, + RunPaths, + run_id_for, +) +from .run_config import ConfigError, RunSection + + +def scaffold_run( + nfft: int, hop: int, suffix: str = "", allow_custom_scale: bool = False +) -> RunPaths: + # Validate the scale with the same rules run.yaml will be held to later. + try: + RunSection(nfft=nfft, hop=hop, allow_custom_scale=allow_custom_scale) + except Exception as err: # pydantic ValidationError + raise ConfigError(str(err)) from None + + run_id = run_id_for(nfft, hop, suffix) + paths = RunPaths(run_id) + + if paths.run_yaml.exists(): + raise FileExistsError( + f"{paths.run_yaml} already exists — refusing to overwrite. " + f"Use --suffix for a fresh variant of this scale." + ) + + paths.workspace.mkdir(parents=True, exist_ok=True) + paths.cache_root.mkdir(parents=True, exist_ok=True) + paths.slurm_dir.mkdir(parents=True, exist_ok=True) + + yaml_text = ( + RUN_TEMPLATE_YAML.read_text() + .replace("{{RUN_ID}}", run_id) + .replace("{{NFFT}}", str(nfft)) + .replace("{{HOP}}", str(hop)) + ) + paths.run_yaml.write_text(yaml_text) + + notebooks = sorted(NOTEBOOK_TEMPLATE_DIR.glob("*.ipynb")) + for nb in notebooks: + target = paths.workspace / nb.name + target.write_text(nb.read_text().replace("__RUN_ID__", run_id)) + + paths.run_meta.write_text( + json.dumps( + { + "run_id": run_id, + "nfft": nfft, + "hop": hop, + "created": datetime.now().isoformat(), + }, + indent=2, + ) + ) + + print(f"Run {run_id} created.") + print(f" knobs: {paths.run_yaml}") + print(f" notebooks: {paths.workspace}/") + if not notebooks: + print(" (no notebook templates found — package notebooks/ is empty)") + print(f" cache: {paths.cache_root}") + print(f"Start with {paths.workspace}/00_setup.ipynb") + return paths + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--nfft", type=int, required=True) + parser.add_argument("--hop", type=int, required=True) + parser.add_argument("--suffix", type=str, default="") + parser.add_argument("--allow-custom-scale", action="store_true") + args = parser.parse_args() + scaffold_run(args.nfft, args.hop, args.suffix, args.allow_custom_scale) + + +if __name__ == "__main__": + main() diff --git a/src/tokeye/training/big_tf_unet_2/slurm.py b/src/tokeye/training/big_tf_unet_2/slurm.py new file mode 100644 index 0000000..002bee0 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/slurm.py @@ -0,0 +1,132 @@ +"""SLURM integration: render, submit, and monitor per-step batch jobs. + +GPU steps go to the GPU partition with a gres; CPU-heavy steps go to a CPU +node with NO gres (CPU work must never hold a GPU). The sbatch payload is the +same runner CLI the notebooks use inline, so behavior is identical either +way. Scripts + logs live under the run's ``slurm/`` cache dir. +""" + +from __future__ import annotations + +import subprocess +from typing import TYPE_CHECKING + +from .paths import RunPaths, get_step +from .run_config import load_run_config +from .task_matrix import RunTaskMatrix + +if TYPE_CHECKING: + from pathlib import Path + + from .run_config import RunConfig + + +def _needs_gpu(steps: list[str]) -> bool: + return any(get_step(s).exec_mode == "sbatch_gpu" for s in steps) + + +def _render_script( + cfg: RunConfig, + paths: RunPaths, + run_yaml: Path, + steps: list[str], + modality: str | None, + tag: str, +) -> str: + gpu = _needs_gpu(steps) + slurm = cfg.slurm + lines = [ + "#!/bin/bash", + f"#SBATCH --job-name=bt2_{tag}", + f"#SBATCH --output={paths.slurm_dir.resolve()}/%j_{tag}.out", + f"#SBATCH --time={slurm.gpu_time if gpu else slurm.cpu_time}", + f"#SBATCH --mem={slurm.gpu_mem if gpu else slurm.cpu_mem}", + f"#SBATCH --cpus-per-task={slurm.gpu_cpus if gpu else slurm.cpu_cpus}", + ] + if gpu: + lines.append("#SBATCH --gres=gpu:1") + if slurm.gpu_partition: + lines.append(f"#SBATCH --partition={slurm.gpu_partition}") + elif slurm.cpu_partition: + lines.append(f"#SBATCH --partition={slurm.cpu_partition}") + if slurm.account: + lines.append(f"#SBATCH --account={slurm.account}") + + cmd = ( + "srun python -m tokeye.training.big_tf_unet_2.runner " + f"--run-config {run_yaml.resolve()} --steps {','.join(steps)}" + ) + if modality: + cmd += f" --modalities {modality}" + lines += [ + "", + f'cd "{paths.root.resolve()}" && source .venv/bin/activate', + "export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True", + cmd, + "", + ] + return "\n".join(lines) + + +def submit_step( + run_yaml: str | Path, + steps: list[str] | str, + modality: str | None = None, +) -> str: + """Render + sbatch one job running the given steps; returns the job id.""" + from pathlib import Path as _Path + + run_yaml = _Path(run_yaml) + if isinstance(steps, str): + steps = [s.strip() for s in steps.split(",")] + cfg = load_run_config(run_yaml) + paths = RunPaths(run_yaml.parent.name) + paths.slurm_dir.mkdir(parents=True, exist_ok=True) + + tag = "_".join(s.removeprefix("step_") for s in steps) + ( + f"_{modality}" if modality else "" + ) + script = paths.slurm_dir / f"{tag}.sh" + script.write_text(_render_script(cfg, paths, run_yaml, steps, modality, tag)) + + result = subprocess.run( + ["sbatch", "--parsable", str(script)], + capture_output=True, + text=True, + check=True, + ) + job_id = result.stdout.strip().split(";")[0] + + matrix = RunTaskMatrix(paths.task_matrix_path) + for step in steps: + mods = ( + [modality] if modality else cfg.modality_names + ) if get_step(step).per_modality else [None] + for mod in mods: + matrix.record_job(step, mod, job_id) + return job_id + + +def job_state(job_id: str) -> str: + """Live state via squeue, falling back to sacct for finished jobs.""" + result = subprocess.run( + ["squeue", "-j", job_id, "-h", "-o", "%T"], capture_output=True, text=True + ) + state = result.stdout.strip() + if state: + return state + result = subprocess.run( + ["sacct", "-j", job_id, "--format=State", "-n", "-P", "-X"], + capture_output=True, + text=True, + ) + lines = [ln.strip() for ln in result.stdout.splitlines() if ln.strip()] + return lines[0] if lines else "UNKNOWN" + + +def tail_log(paths: RunPaths, job_id: str, n: int = 40) -> str: + """Last n lines of a submitted job's log file.""" + logs = sorted(paths.slurm_dir.glob(f"{job_id}_*.out")) + if not logs: + return f"(no log yet for job {job_id} — still queued?)" + return "\n".join(logs[-1].read_text().splitlines()[-n:]) diff --git a/src/tokeye/training/big_tf_unet_2/step_0_intake.py b/src/tokeye/training/big_tf_unet_2/step_0_intake.py new file mode 100644 index 0000000..2c2dce0 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_0_intake.py @@ -0,0 +1,176 @@ +"""Step 0: intake — load, resample, pre-emphasize, and window shot data. + +Ported from ``big_tf_unet_ablation/step_0f_foundation.py`` to the +``big_tf_unet_2`` runner contract: ``main(settings)`` receives a fully +resolved settings dict (paths, channels, rates already picked by the +runner) and returns ``None`` — there are no in-step resolved params for +this step. + +Each shot's ``{shot}_processed.h5`` holds one HDF5 group per modality, with +``xdata (N,)`` (time, in seconds) and ``ydata (C, N)`` (channels x samples). +Native sample rates differ per modality -- and even per shot, since they +come from real hardware timestamps rather than a nominal rate -- so the +native rate is derived from the ``xdata`` spacing and each modality is +resampled to ``target_rate_khz`` with an anti-aliased polyphase filter +before windowing (mirrors the pattern in +``big_tf_unet_ablation/step_0g_raw_fast.py``). + +Processing per shot: select channels -> resample to the target rate -> +pre-emphasis -> consecutive non-overlapping windows of ``subseq_len`` +samples, capped at ``max_windows_per_shot``. +""" + +from __future__ import annotations + +import logging +from fractions import Fraction +from typing import TYPE_CHECKING + +import h5py +import numpy as np +import pandas as pd +import torch +from scipy.signal import resample_poly + +from .utils.hdf5_io import create_step_file, write_sample +from .utils.signal_filters import Preemphasis + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +# Number of leading samples read from ``xdata`` to estimate the native rate -- +# spacing is uniform enough that a small prefix is plenty, and avoids reading +# a multi-million-sample time axis just to compute one number. +_RATE_PROBE_SAMPLES = 1024 + + +def _read_shots(shots_path: Path, n_shots: int | None) -> list[int]: + shots = [int(s) for s in shots_path.read_text().split()] + return shots[:n_shots] if n_shots is not None else shots + + +def _native_rate_khz(xdata: np.ndarray) -> float: + """Native sample rate (kHz) implied by ``xdata`` spacing (seconds).""" + dt_s = float(np.median(np.diff(xdata))) + return 1.0 / (dt_s * 1000.0) + + +def _resample_to(sig: np.ndarray, src_khz: float, target_khz: float) -> np.ndarray: + """Resample ``(C, N)`` from ``src_khz`` to ``target_khz`` (anti-aliased).""" + frac = Fraction(int(round(target_khz)), int(round(src_khz))).limit_denominator( + 1000 + ) + up, down = frac.numerator, frac.denominator + if (up, down) == (1, 1): + return sig + return resample_poly(sig, up, down, axis=1) + + +def main(settings: dict) -> dict | None: + shots_path: Path = settings["shots_path"] + foundation_dir: Path = settings["foundation_dir"] + modality: str = settings["modality"] + input_key: str = settings["input_key"] + channels: list[int] = list(settings["channels"]) + subseq_len: int = int(settings["subseq_len"]) + coeff: float = float(settings["preemphasis_coeff"]) + target_khz: float = float(settings["target_rate_khz"]) + max_windows_per_shot: int = int(settings["max_windows_per_shot"]) + n_shots: int | None = settings["n_shots"] + out_h5: Path = settings["out_h5"] + frame_info_csv: Path = settings["frame_info_csv"] + run_id: str = settings["run_id"] + smoke: bool = bool(settings["smoke"]) + + shots = _read_shots(shots_path, n_shots) + preemph = Preemphasis(coeff) + + h5 = create_step_file( + out_h5, + metadata={ + "modality": modality, + "input_key": input_key, + "num_channels": len(channels), + "subseq_len": subseq_len, + "target_rate_khz": target_khz, + "run_id": run_id, + }, + ) + + rows: list[dict[str, int]] = [] + index = 0 + n_shots_used = 0 + try: + for shot in shots: + path = foundation_dir / f"{shot}_processed.h5" + if not path.exists(): + logger.warning(f"{modality}: shot {shot} absent in foundation_dir") + continue + + with h5py.File(path, "r") as fh: + if input_key not in fh or "ydata" not in fh[input_key]: + logger.warning( + f"{modality}: shot {shot} has no {input_key!r} group" + ) + continue + group = fh[input_key] + ydata = group["ydata"] + xdata = group["xdata"] + n = ydata.shape[1] + if n < subseq_len: + logger.warning( + f"{modality}: shot {shot} degenerate (N={n} < " + f"subseq_len={subseq_len}); skip" + ) + continue + if max(channels) >= ydata.shape[0]: + logger.warning( + f"{modality}: shot {shot} has {ydata.shape[0]} channels " + f"< requested max {max(channels)}; skip" + ) + continue + if xdata.shape[0] != n: + logger.warning( + f"{modality}: shot {shot} xdata({xdata.shape[0]}) != " + f"ydata N({n}); skip" + ) + continue + sig = np.asarray(ydata[channels, :], dtype=np.float64) + x_probe = np.asarray( + xdata[: min(_RATE_PROBE_SAMPLES, n)], dtype=np.float64 + ) + + if not np.any(sig): + logger.warning(f"{modality}: shot {shot} is all-zero; skip") + continue + + src_khz = _native_rate_khz(x_probe) + sig = _resample_to(sig, src_khz, target_khz) + sig_t = preemph(torch.from_numpy(sig).float()).numpy() + + n_windows = min(sig_t.shape[1] // subseq_len, max_windows_per_shot) + for w in range(n_windows): + s0, s1 = w * subseq_len, (w + 1) * subseq_len + write_sample(h5, index, sig_t[:, s0:s1]) + rows.append({"index": index, "shotn": shot, "window_start": s0}) + index += 1 + n_shots_used += 1 + logger.info( + f"{modality}: shot {shot} -> {n_windows} windows " + f"({src_khz:.2f} kHz -> {target_khz:.0f} kHz)" + ) + finally: + h5.close() + + pd.DataFrame(rows, columns=["index", "shotn", "window_start"]).to_csv( + frame_info_csv, index=False + ) + + smoke_tag = " (smoke)" if smoke else "" + logger.info( + f"[{run_id}] step_0_intake [{modality}]{smoke_tag}: {index} windows " + f"from {n_shots_used}/{len(shots)} shots -> {out_h5}" + ) + return None diff --git a/src/tokeye/training/big_tf_unet_2/step_1_spectrogram.py b/src/tokeye/training/big_tf_unet_2/step_1_spectrogram.py new file mode 100644 index 0000000..37dbee3 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_1_spectrogram.py @@ -0,0 +1,165 @@ +"""Step 1: STFT spectrograms + optional model-based window filtering. + +Pass 1 computes an STFT for every step_0 window (``torch.stft`` with a hann +window; real/imag stacked into float32 ``(C, F, T, 2)``). With the window +filter disabled the full set goes straight to ``out_h5`` and the frame info +is copied through unchanged. With it enabled, the full set goes to a +temporary ``*.full.h5``, pass 2 scores every window with the surrogate +U-Net, keeps the top windows per shot, and re-indexes the survivors 0..N-1 +into ``out_h5`` (with a matching filtered frame info). The temporary file is +always removed. + +Merged port of ``big_tf_unet_ablation`` steps 2a (STFT) and 2f (window +filtering). +""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path +from typing import TYPE_CHECKING + +import pandas as pd +import torch + +from .utils.hdf5_io import ( + create_step_file, + get_sample_count, + iter_samples, + read_sample, + write_sample, +) +from .utils.window_filter import ( + compute_logmag_stats, + load_filter_model, + score_window, + select_window_indices, +) + +if TYPE_CHECKING: + import numpy as np + +logger = logging.getLogger(__name__) + + +def _compute_stft(data: np.ndarray, nfft: int, hop: int) -> np.ndarray: + """STFT of a ``(C, T)`` window -> float32 ``(C, F, T_spec, 2)`` (real/imag).""" + x = torch.from_numpy(data).float() + window = torch.hann_window(nfft) + sxx = torch.stft( + x, n_fft=nfft, hop_length=hop, window=window, return_complex=True + ) + return torch.stack([sxx.real, sxx.imag], dim=-1).numpy() + + +def _write_spectrograms(in_h5: Path, out_h5: Path, nfft: int, hop: int) -> int: + """Pass 1: STFT every sample of ``in_h5`` into a new step file at ``out_h5``.""" + n = get_sample_count(in_h5) + logger.info(f"computing STFT (nfft={nfft}, hop={hop}) for {n} windows") + h5 = create_step_file( + out_h5, metadata={"nfft": nfft, "hop": hop, "num_samples": n} + ) + try: + for idx, data in iter_samples(in_h5): + write_sample(h5, idx, _compute_stft(data, nfft, hop)) + finally: + h5.close() + return n + + +def _filter_windows( + full_h5: Path, + out_h5: Path, + frame_info_in: Path, + frame_info_out: Path, + nfft: int, + hop: int, + settings: dict, +) -> dict | None: + """Pass 2: score every window, keep the most active per shot, re-index.""" + device = "cpu" # inline step (login node): never grab a GPU + fi = pd.read_csv(frame_info_in) # row order aligns with sample idx + + # Per-modality normalization: each diagnostic has its own intensity scale + # and 1/f structure, so "auto" computes this modality's own log-magnitude + # stats from the freshly computed spectrograms. + mean, std = settings["mean"], settings["std"] + stats_computed = mean == "auto" or std == "auto" + if stats_computed: + mean, std = compute_logmag_stats(full_h5) + mean, std = float(mean), float(std) + logger.info( + f"window-filter normalization (per modality): mean={mean:.4f} std={std:.4f}" + ) + + model = load_filter_model(settings["weights"], device=device) + threshold = float(settings["activity_threshold"]) + n = get_sample_count(full_h5) + with torch.no_grad(): + scores = { + idx: score_window( + model, read_sample(full_h5, idx), mean, std, threshold, device + ) + for idx in range(n) + } + + max_windows = int(settings["max_windows_per_shot"]) + min_activity = float(settings["min_activity"]) + kept: list[int] = [] + for shotn, grp in fi.groupby("shotn"): + shot_scores = {i: scores[i] for i in grp.index if i in scores} + selected = select_window_indices(shot_scores, max_windows, min_activity) + kept.extend(selected) + logger.info(f"shot {shotn}: kept {len(selected)}/{len(grp)} windows") + kept.sort() + logger.info(f"window filter: kept {len(kept)}/{n} windows total") + + h5 = create_step_file( + out_h5, + metadata={"nfft": nfft, "hop": hop, "filtered": True, "kept": len(kept)}, + ) + try: + for new_idx, old_idx in enumerate(kept): + write_sample(h5, new_idx, read_sample(full_h5, old_idx)) + finally: + h5.close() + + # Filtered frame info: new contiguous index, original sample idx kept as + # orig_index, all original columns (shotn, window_start, ...) carried over. + filtered = fi.iloc[kept].reset_index(drop=True) + filtered["orig_index"] = kept + filtered["index"] = filtered.index + lead = ["index", "orig_index"] + cols = lead + [c for c in filtered.columns if c not in lead] + filtered[cols].to_csv(frame_info_out, index=False) + logger.info(f"wrote {out_h5} and {frame_info_out}") + + if stats_computed: + return {"mean": mean, "std": std} + return None + + +def main(settings: dict) -> dict | None: + in_h5 = Path(settings["in_h5"]) + out_h5 = Path(settings["out_h5"]) + frame_info_in = Path(settings["frame_info_in"]) + frame_info_out = Path(settings["frame_info_out"]) + nfft = int(settings["nfft"]) + hop = int(settings["hop"]) + + if not settings["filter_enabled"]: + _write_spectrograms(in_h5, out_h5, nfft, hop) + shutil.copyfile(frame_info_in, frame_info_out) + logger.info("window filter disabled: wrote unfiltered spectrograms") + return None + + full_h5 = out_h5.with_suffix(".full.h5") + try: + _write_spectrograms(in_h5, full_h5, nfft, hop) + return _filter_windows( + full_h5, out_h5, frame_info_in, frame_info_out, nfft, hop, settings + ) + finally: + # Not a registered artifact -- it must never survive the step. + full_h5.unlink(missing_ok=True) diff --git a/src/tokeye/training/big_tf_unet_2/step_2_baseline.py b/src/tokeye/training/big_tf_unet_2/step_2_baseline.py new file mode 100644 index 0000000..d7b5065 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_2_baseline.py @@ -0,0 +1,225 @@ +"""Step 2: baseline-correct spectrograms (FABC residual). + +Reads ``(C, F, T, 2)`` real/imag samples from the step_1 HDF5, applies 2D +baseline fitting (pybaselines FABC along the frequency axis) to the log1p +magnitude of each component, and writes the relative residual to ``out_h5`` +and the fitted baseline to ``baseline_h5``. + +Edge frequency bins (``edge_bins_lower`` bottom rows, ``edge_bins_upper`` top +rows) are filled with the sample's post-log1p median before the fit and zeroed +in the residual after correction. +""" + +from __future__ import annotations + +import logging +import multiprocessing +import os +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait +from pathlib import Path + +import numpy as np +from pybaselines import Baseline2D + +from .utils.hdf5_io import ( + create_step_file, + get_sample_count, + iter_samples, + write_sample, +) + +logger = logging.getLogger(__name__) + + +def _fit_baseline(data: np.ndarray, method: str, lam: float) -> np.ndarray: + """Fit a 2D baseline along the frequency axis of one ``(F, T)`` field.""" + x = np.arange(data.shape[0]) + y = np.arange(data.shape[1]) + fitter = Baseline2D(x, y) + baseline, _ = fitter.individual_axes( + data, axes=0, method=method, method_kwargs={"lam": lam} + ) + return baseline + + +def _process_component( + sxx: np.ndarray, + edge_lower: int, + edge_upper: int, + method: str, + lam: float, +) -> tuple[np.ndarray, np.ndarray]: + """Baseline-correct one post-log1p ``(F, T)`` component field. + + The masked edge rows are EXCLUDED from the fit rather than filled: FABC + classifies baseline points by noise statistics, and a block of constant + synthetic rows starves the classifier ("not enough baseline points") and + makes the Whittaker system singular. The baseline is fit on the interior + rows only; edge rows get the nearest interior baseline value in the + baseline output and 0 in the residual. + """ + n_freq = sxx.shape[0] + lo = edge_lower + hi = n_freq - edge_upper if edge_upper > 0 else n_freq + + interior_baseline = _fit_baseline(sxx[lo:hi], method, lam) + baseline = np.empty_like(sxx) + baseline[lo:hi] = interior_baseline + if lo > 0: + baseline[:lo] = interior_baseline[0] + if edge_upper > 0: + baseline[hi:] = interior_baseline[-1] + + residual = np.zeros_like(sxx) + residual[lo:hi] = (sxx[lo:hi] - interior_baseline) / (interior_baseline + 1e-6) + return residual, baseline + + +def _process_sample( + data: np.ndarray, + edge_lower: int, + edge_upper: int, + method: str, + lam: float, +) -> tuple[np.ndarray, np.ndarray]: + """Baseline-correct one full sample ``(C, F, T, 2)`` across all channels.""" + out = np.zeros_like(data) + bl = np.zeros_like(data) + log_field = np.log1p(np.abs(data)) + for c in range(data.shape[0]): + for r in range(data.shape[-1]): + out[c, ..., r], bl[c, ..., r] = _process_component( + log_field[c, ..., r], edge_lower, edge_upper, method, lam + ) + return out, bl + + +# Kept alive for the worker's lifetime so the thread limit isn't GC'd away. +_THREAD_LIMITER = None + + +def _init_worker() -> None: + """Cap each worker's native thread pools to 1 thread. + + Baseline fitting is run across many worker processes; without this each + process's BLAS/OpenMP pool would try to use every core, oversubscribing them + and thrashing. ``threadpoolctl`` limits at runtime (after numpy/scipy + import), so it works regardless of start method. + """ + global _THREAD_LIMITER + try: + from threadpoolctl import threadpool_limits + + _THREAD_LIMITER = threadpool_limits(limits=1) + except Exception: # optional dependency; fall back to inherited env vars + _THREAD_LIMITER = None + + +def _worker( + args: tuple[int, np.ndarray, int, int, str, float], +) -> tuple[int, np.ndarray, np.ndarray]: + """ProcessPool task: baseline-correct one sample -> ``(idx, out, baseline)``.""" + idx, data, edge_lower, edge_upper, method, lam = args + out, bl = _process_sample(data, edge_lower, edge_upper, method, lam) + return idx, out, bl + + +def _resolve_n_workers(settings: dict) -> int: + """Worker count: explicit setting > SLURM allocation > CPU count, min 1.""" + n = settings.get("n_workers") + if n is None: + n = int(os.environ.get("SLURM_CPUS_PER_TASK", 0)) or (os.cpu_count() or 1) + return max(1, int(n)) + + +def main(settings: dict) -> dict | None: + in_h5 = Path(settings["in_h5"]) + out_h5 = Path(settings["out_h5"]) + baseline_h5 = Path(settings["baseline_h5"]) + method = str(settings["method"]) + lam = float(settings["lam"]) + edge_lower = int(settings["edge_bins_lower"]) + edge_upper = int(settings["edge_bins_upper"]) + + n_samples = get_sample_count(in_h5) + n_workers = _resolve_n_workers(settings) + logger.info( + f"Baseline-correcting {n_samples} spectrograms " + f"(method={method}, lam={lam:g}, " + f"edges: lower={edge_lower}, upper={edge_upper}; {n_workers} workers)" + ) + + h5_out = create_step_file( + out_h5, + metadata={ + "run_id": settings["run_id"], + "method": method, + "lam": lam, + "edge_bins_lower": edge_lower, + "edge_bins_upper": edge_upper, + "num_samples": n_samples, + }, + ) + h5_bl = create_step_file( + baseline_h5, + metadata={"run_id": settings["run_id"], "num_samples": n_samples}, + ) + + log_every = max(1, n_samples // 20) + n_done = 0 + + def _store(idx: int, out: np.ndarray, bl: np.ndarray) -> None: + nonlocal n_done + write_sample(h5_out, idx, out) + write_sample(h5_bl, idx, bl) + n_done += 1 + if n_done % log_every == 0 or n_done == n_samples: + logger.info(f" baseline-corrected {n_done}/{n_samples}") + + def _arg(idx: int, data: np.ndarray) -> tuple: + return (idx, data, edge_lower, edge_upper, method, lam) + + try: + if n_workers <= 1: + for idx, data in iter_samples(in_h5): + _store( + idx, + *_process_sample(data, edge_lower, edge_upper, method, lam), + ) + else: + # Per-sample baseline fitting is the pipeline bottleneck and is pure CPU + # (each sample = C channels x 2 components x ~T 1D FABC fits, no GPU + # work), so a serial loop pegs one core for hours. Fan out over the + # allocated cores with a bounded in-flight set so the fitting saturates + # every core while peak memory stays flat. A forkserver context starts + # workers from a clean server process: no per-worker re-import of the + # (torch-importing) entrypoint that spawn incurs, and no inherited open + # HDF5 handle/CUDA state that plain fork would carry. Results may arrive + # out of order, but write_sample keys by index so ordering does not + # matter. + samples = iter_samples(in_h5) + max_inflight = max(1, n_workers * 2) + with ProcessPoolExecutor( + max_workers=n_workers, + mp_context=multiprocessing.get_context("forkserver"), + initializer=_init_worker, + ) as ex: + pending = set() + for _ in range(max_inflight): + nxt = next(samples, None) + if nxt is None: + break + pending.add(ex.submit(_worker, _arg(*nxt))) + while pending: + done, pending = wait(pending, return_when=FIRST_COMPLETED) + for fut in done: + _store(*fut.result()) + nxt = next(samples, None) + if nxt is not None: + pending.add(ex.submit(_worker, _arg(*nxt))) + finally: + h5_out.close() + h5_bl.close() + + logger.info(f"Wrote residuals to {out_h5} and baselines to {baseline_h5}") + return None diff --git a/src/tokeye/training/big_tf_unet_2/step_3_denoise.py b/src/tokeye/training/big_tf_unet_2/step_3_denoise.py new file mode 100644 index 0000000..d5efb2b --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_3_denoise.py @@ -0,0 +1,556 @@ +"""step_3 — self-supervised cross-channel denoising (train + per-sample extract). + +Trains a U-Net to predict each channel's residual spectrogram from its +``adjacent_channels`` neighbors: sensor noise is independent across channels +while coherent modes are shared, so the network can only reproduce the shared +signal. After training, prediction runs over the full dataset in order and each +sample's denoised field is written to ``out_h5`` keyed by its original index — +the ablation pipeline's separate step_3b unpack is folded into this step. + +Preprocessing per sample (input ``(C, F, T, 2)`` step_2 residual): + +- ``representation="magnitude"`` collapses (Re, Im) -> |Z| before normalizing. +- ``normalization="robust_asinh"``: N_a(x) = a * asinh((x - med) / (a * 1.4826 + * MAD)) per channel (and per real/imag component); ``a=1`` reproduces the + ablation's robust-asinh exactly. +- ``normalization="zscore"``: plain per-sample per-channel z-score. +- The bottom ``edge_bins_lower`` and top ``edge_bins_upper`` frequency rows are + zeroed AFTER normalization (0 is the post-N_a median). + +Output ``/samples/{idx}`` (float32, same indexing as ``in_h5``): the mean of +the model's two denoised estimates (adjacent-channel prediction and +self-reconstruction) — ``(C, F, T, 2)`` real/imag for the complex +representation, ``(C, F, T)`` for magnitude. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import h5py +import lightning as L +import torch +from lightning.pytorch.callbacks import BasePredictionWriter, Callback +from torch import nn +from torch.utils.data import DataLoader +from torchmetrics.image import TotalVariation + +from tokeye.models.modules.unet import UNet + +from .utils.hdf5_io import ( + HDF5StepDataset, + create_step_file, + worker_init_fn, + write_sample, +) + +# TF32: fp32 storage + tensor-core matmul/conv (10-bit operands, fp32 +# accumulate) -- ~bf16 speed, ~fp32 stability. MUST use the legacy API: +# Lightning calls torch.get_float32_matmul_precision() internally, so the new +# backends.*.fp32_precision API raises "mix of legacy and new APIs" at Trainer +# setup. +torch.set_float32_matmul_precision("high") +torch.backends.cudnn.allow_tf32 = True + +logger = logging.getLogger(__name__) + + +# ------------------------------------------------------------------ +# Normalization +# ------------------------------------------------------------------ + + +def _robust_asinh(data: torch.Tensor, a: float) -> torch.Tensor: + """N_a(x) = a * asinh((x - median) / (a * 1.4826 * MAD)), per channel. + + The robust scale tracks the noise floor (breakdown point 50% vs std's 0%), + so a strong observation cannot inflate the divisor and suppress weaker + modes; asinh soft-bounds the strong tail without discarding it, with ``a`` + setting where compression kicks in (``a=1`` reproduces the ablation's + robust-asinh exactly). Handles ``(C, F, T)`` magnitude and ``(C, F, T, 2)`` + complex, normalizing per channel (and component) over the (F, T) axes. + """ + c = data.shape[0] + if data.dim() == 3: # (C, F, T) + flat = data.reshape(c, -1) + med = flat.median(dim=1, keepdim=True).values + mad = (flat - med).abs().median(dim=1, keepdim=True).values + z = a * torch.asinh((flat - med) / (a * (1.4826 * mad + 1e-6))) + return z.reshape(data.shape) + # (C, F, T, 2): per channel & component + k = data.shape[-1] + flat = data.permute(0, 3, 1, 2).reshape(c, k, -1) # (C, K, F*T) + med = flat.median(dim=2, keepdim=True).values + mad = (flat - med).abs().median(dim=2, keepdim=True).values + z = a * torch.asinh((flat - med) / (a * (1.4826 * mad + 1e-6))) + return z.reshape(c, k, data.shape[1], data.shape[2]).permute(0, 2, 3, 1) + + +class Preprocess: + """Per-sample transform: representation collapse, normalize, zero edges. + + Module-level class (not a closure) so DataLoader workers can pickle it. + """ + + def __init__( + self, + representation: str, + normalization: str, + a: float, + edge_bins_lower: int, + edge_bins_upper: int, + ) -> None: + self.representation = representation + self.normalization = normalization + self.a = a + self.edge_bins_lower = edge_bins_lower + self.edge_bins_upper = edge_bins_upper + + def __call__(self, data: torch.Tensor) -> torch.Tensor: + data = torch.nan_to_num(data) + if self.representation == "magnitude": + # collapse (Re, Im) -> |Z|: (C, F, T, 2) -> (C, F, T) + data = torch.sqrt(data[..., 0] ** 2 + data[..., 1] ** 2) + if self.normalization == "robust_asinh": + data = _robust_asinh(data, self.a) + else: # per-sample per-channel z-score (per component for complex) + mean = data.mean(dim=(1, 2), keepdim=True) + std = data.std(dim=(1, 2), keepdim=True) + data = (data - mean) / (std + 1e-6) + # Edge-bin masking after normalization: 0 is the post-N_a median. + if self.edge_bins_lower > 0: + data[:, : self.edge_bins_lower] = 0.0 + if self.edge_bins_upper > 0: + data[:, -self.edge_bins_upper :] = 0.0 + return data + + +# ------------------------------------------------------------------ +# Callbacks +# ------------------------------------------------------------------ + + +class TotalVariationEarlyStopping(Callback): + """Early stopping based on total variation increase.""" + + def __init__(self, patience: int = 3, min_delta: float = 0.0) -> None: + super().__init__() + self.patience = patience + self.min_delta = min_delta + self.best_tv = float("inf") + self.tv_increase_count = 0 + + def on_train_epoch_end( + self, trainer: L.Trainer, pl_module: L.LightningModule + ) -> None: + current_tv = trainer.callback_metrics.get("train_tv", None) + if current_tv is None: + return + current_tv = float(current_tv) + if current_tv > self.best_tv + self.min_delta: + self.tv_increase_count += 1 + if self.tv_increase_count >= self.patience: + trainer.should_stop = True + else: + self.best_tv = current_tv + self.tv_increase_count = 0 + + +class DenoisedSampleWriter(BasePredictionWriter): + """Streams each denoised sample to the step HDF5 as batches complete. + + Keeps the file open across batches (opened at predict start) and keys + every sample by its original ``in_h5`` index — the predict loader is + sequential, so a running position maps 1:1 onto the sorted sample keys. + Predictions never accumulate in memory. + """ + + def __init__( + self, + out_h5: Path, + sample_keys: list[str], + metadata: dict | None = None, + ) -> None: + super().__init__(write_interval="batch") + self.out_h5 = out_h5 + self.sample_keys = sample_keys + self.metadata = metadata + self._h5: h5py.File | None = None + self._pos = 0 + + def on_predict_start( + self, trainer: L.Trainer, pl_module: L.LightningModule + ) -> None: + self._h5 = create_step_file(self.out_h5, metadata=self.metadata) + self._pos = 0 + + def write_on_batch_end( + self, + trainer: L.Trainer, + pl_module: L.LightningModule, + prediction: torch.Tensor, + batch_indices, + batch, + batch_idx: int, + dataloader_idx: int, + ) -> None: + for sample in prediction.numpy(): + write_sample(self._h5, int(self.sample_keys[self._pos]), sample) + self._pos += 1 + + def close(self) -> None: + if self._h5 is not None: + self._h5.attrs["num_samples"] = self._pos + self._h5.close() + self._h5 = None + + +# ------------------------------------------------------------------ +# Model +# ------------------------------------------------------------------ + + +class BTN(nn.Module): + """Complex-field denoiser: one shared U-Net over the (Re, -Im) components. + + Negating the imaginary part means the net sees the conjugate; paired with + the ``y.flip(-1)`` training target this ties the two zero-mean Gaussian + components together instead of treating them as unrelated planes. + """ + + def __init__( + self, + in_channels: int = 4, + num_layers: int = 5, + first_layer_size: int = 32, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.unet = UNet( + in_channels=in_channels, + out_channels=in_channels, + num_layers=num_layers, + first_layer_size=first_layer_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: # (B, 2A, F, T, 2) + x_real, x_imag = x[..., 0], -x[..., 1] + return torch.stack([self.unet(x_real), self.unet(x_imag)], dim=-1) + + +class BTNMag(nn.Module): + """Magnitude-only denoiser: single-component U-Net (no real/imag split). + + Same width/depth as :class:`BTN` so the contrast between representations + is representation, not capacity. + """ + + def __init__( + self, + in_channels: int = 6, + num_layers: int = 5, + first_layer_size: int = 32, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.unet = UNet( + in_channels=in_channels, + out_channels=in_channels, + num_layers=num_layers, + first_layer_size=first_layer_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: # (B, 2A, F, T) + return self.unet(x) + + +class DenoiseModule(L.LightningModule): + """Self-supervised cross-channel denoiser (L1, TV-tracked). + + Two training paths per batch: predict a random target channel from its + ``adjacent_channels`` neighbors, and reconstruct each neighbor from the + target channel repeated. Prediction averages the two estimates per channel. + """ + + def __init__( + self, + representation: str = "complex", + adjacent_channels: int = 3, + total_channels: int = 8, + num_layers: int = 5, + first_layer_size: int = 32, + compile_model: bool = False, + ) -> None: + super().__init__() + if adjacent_channels >= total_channels: + raise ValueError( + f"adjacent_channels ({adjacent_channels}) must be < " + f"total_channels ({total_channels})" + ) + self.representation = representation + self.adjacent_channels = adjacent_channels + self.in_channels = 2 * adjacent_channels + self.total_channels = total_channels + net_cls = BTNMag if representation == "magnitude" else BTN + self.net = net_cls( + in_channels=self.in_channels, + num_layers=num_layers, + first_layer_size=first_layer_size, + ) + if compile_model: + self.net.compile() + + self.pad_len = self.adjacent_channels + # Pads the channel axis: dim -3 of (B, C, F, T) inputs. + self.padding = nn.ReflectionPad3d((0, 0, 0, 0, self.pad_len, self.pad_len)) + + self.loss_fn = nn.L1Loss() + self.train_tv = TotalVariation() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + # --- channel gather (complex, (B, C, F, T, 2)) --- + + def _load_adjacent_channels(self, x: torch.Tensor, target: int) -> torch.Tensor: + """The 2A channels around ``target`` -> (B, 2A, F, T, 2).""" + channel_idx = target + self.pad_len + x_real, x_imag = x[..., 0], x[..., 1] + xp = torch.stack([self.padding(x_real), self.padding(x_imag)], dim=-1) + front = xp[:, channel_idx - self.pad_len : channel_idx] + back = xp[:, channel_idx + 1 : channel_idx + 1 + self.pad_len] + return torch.cat([front, back], dim=1) + + def _load_target_channels(self, x: torch.Tensor, target: int) -> torch.Tensor: + """``target`` repeated to the net's input width -> (B, 2A, F, T, 2).""" + return x[:, target : target + 1].repeat(1, self.in_channels, 1, 1, 1) + + # --- channel gather (magnitude, (B, C, F, T)) --- + + def _load_adjacent_channels_mag( + self, x: torch.Tensor, target: int + ) -> torch.Tensor: + channel_idx = target + self.pad_len + xp = self.padding(x) + front = xp[:, channel_idx - self.pad_len : channel_idx] + back = xp[:, channel_idx + 1 : channel_idx + 1 + self.pad_len] + return torch.cat([front, back], dim=1) + + def _load_target_channels_mag( + self, x: torch.Tensor, target: int + ) -> torch.Tensor: + return x[:, target : target + 1].repeat(1, self.in_channels, 1, 1) + + # --- losses (complex) --- + + def _single_channel_loss( + self, y_hat: torch.Tensor, y: torch.Tensor + ) -> torch.Tensor: + loss = self.loss_fn(y_hat, y.flip(-1)) + self.train_tv.update(y_hat[..., 0]) + self.train_tv.update(y_hat[..., 1]) + return loss + + def _multichannel_loss( + self, y_hat: torch.Tensor, y: torch.Tensor + ) -> torch.Tensor: + loss = torch.tensor(0.0, device=y_hat.device) + for i in range(self.in_channels): + y_hat_i = y_hat[:, i : i + 1] + y_i = y[:, i : i + 1] + loss = loss + self.loss_fn(y_hat_i, y_i.flip(-1)) + self.train_tv.update(y_hat_i[..., 0] / self.in_channels) + self.train_tv.update(y_hat_i[..., 1] / self.in_channels) + return loss / self.in_channels + + # --- training --- + + def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + if self.representation == "magnitude": + return self._training_step_mag(batch) + target = int(torch.randint(0, self.total_channels, (1,))) + + x1 = self._load_adjacent_channels(batch, target) + y1 = batch[:, target : target + 1] + y_hat1 = self(x1).mean(dim=1, keepdim=True) + loss1 = self._single_channel_loss(y_hat1, y1) + self.log("train_loss1", loss1) + + x2 = self._load_target_channels(batch, target) + loss2 = self._multichannel_loss(self(x2), x1) + self.log("train_loss2", loss2) + + loss = loss1 + loss2 + self.log("train_loss", loss) + return loss + + def _training_step_mag(self, batch: torch.Tensor) -> torch.Tensor: + target = int(torch.randint(0, self.total_channels, (1,))) + + x1 = self._load_adjacent_channels_mag(batch, target) + y1 = batch[:, target : target + 1] + y_hat1 = self(x1).mean(dim=1, keepdim=True) + loss1 = self.loss_fn(y_hat1, y1) + self.train_tv.update(y_hat1) + self.log("train_loss1", loss1) + + x2 = self._load_target_channels_mag(batch, target) + y_hat2 = self(x2) + loss2 = torch.tensor(0.0, device=batch.device) + for i in range(self.in_channels): + y_hat_i = y_hat2[:, i : i + 1] + loss2 = loss2 + self.loss_fn(y_hat_i, x1[:, i : i + 1]) + self.train_tv.update(y_hat_i / self.in_channels) + loss2 = loss2 / self.in_channels + self.log("train_loss2", loss2) + + loss = loss1 + loss2 + self.log("train_loss", loss) + return loss + + def on_train_epoch_end(self) -> None: + self.log("train_tv", self.train_tv.compute(), logger=True) + self.train_tv.reset() + + # --- prediction --- + + def _denoise_batch(self, batch: torch.Tensor) -> torch.Tensor: + """Denoise every channel of a batch -> (B, C, F, T[, 2]).""" + targets = range(self.total_channels) + if self.representation == "magnitude": + xs = [self._load_adjacent_channels_mag(batch, t) for t in targets] + ys = [self._load_target_channels_mag(batch, t) for t in targets] + y_hats1 = torch.cat( + [self(x).mean(dim=1, keepdim=True) for x in xs], dim=1 + ) + y_hats2 = torch.cat( + [self(y).mean(dim=1, keepdim=True) for y in ys], dim=1 + ) + return (y_hats1 + y_hats2) / 2 + xs = [self._load_adjacent_channels(batch, t) for t in targets] + ys = [self._load_target_channels(batch, t) for t in targets] + y_hats1 = torch.cat([self(x).mean(dim=1, keepdim=True) for x in xs], dim=1) + y_hats2 = torch.cat([self(y).mean(dim=1, keepdim=True) for y in ys], dim=1) + # The complex net is trained against y.flip(-1) (conjugate trick), so + # the estimates come out (Im, Re); flip back to (Re, Im). + return ((y_hats1 + y_hats2) / 2).flip(-1) + + def predict_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + return self._denoise_batch(batch).float().cpu() + + def configure_optimizers(self) -> torch.optim.Optimizer: + return torch.optim.AdamW(self.parameters(), lr=1e-4) + + +# ------------------------------------------------------------------ +# Main +# ------------------------------------------------------------------ + + +def _make_loader( + dataset: HDF5StepDataset, + batch_size: int, + num_workers: int, + shuffle: bool, +) -> DataLoader: + """DataLoader with the efficiency flags, guarded when num_workers == 0.""" + kwargs: dict = {} + if num_workers > 0: + kwargs.update( + persistent_workers=True, + prefetch_factor=2, + worker_init_fn=worker_init_fn, + ) + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=True, + **kwargs, + ) + + +def main(settings: dict) -> dict | None: + in_h5 = Path(settings["in_h5"]) + out_h5 = Path(settings["out_h5"]) + representation = str(settings["representation"]) + normalization = str(settings["normalization"]) + a = float(settings["a"]) + edge_lower = int(settings["edge_bins_lower"]) + edge_upper = int(settings["edge_bins_upper"]) + batch_size = int(settings["batch_size"]) + num_workers = int(settings["num_workers"]) + smoke = bool(settings["smoke"]) + + L.seed_everything(42) + + dataset = HDF5StepDataset( + in_h5, + transform=Preprocess( + representation=representation, + normalization=normalization, + a=a, + edge_bins_lower=edge_lower, + edge_bins_upper=edge_upper, + ), + ) + with h5py.File(in_h5, "r") as f: + sample_keys = sorted(f["samples"].keys(), key=int) + logger.info( + f"Denoising {len(dataset)} samples from {in_h5} " + f"(representation={representation}, normalization={normalization}, " + f"a={a:g}, edges: lower={edge_lower}, upper={edge_upper})" + ) + + model = DenoiseModule( + representation=representation, + adjacent_channels=int(settings["adjacent_channels"]), + total_channels=int(settings["total_channels"]), + num_layers=int(settings["num_layers"]), + first_layer_size=int(settings["first_layer_size"]), + compile_model=not smoke and torch.cuda.is_available(), + ) + + writer = DenoisedSampleWriter( + out_h5, + sample_keys, + metadata={ + "run_id": settings["run_id"], + "representation": representation, + "normalization": normalization, + "a": a, + "edge_bins_lower": edge_lower, + "edge_bins_upper": edge_upper, + "first_layer_size": int(settings["first_layer_size"]), + }, + ) + trainer = L.Trainer( + accelerator="auto", + devices=1, + max_epochs=int(settings["max_epochs"]), + precision=settings["precision"], + log_every_n_steps=10, + enable_progress_bar=False, + enable_checkpointing=False, + logger=False, + callbacks=[ + TotalVariationEarlyStopping(patience=int(settings["tv_patience"])), + writer, + ], + ) + + trainer.fit(model, _make_loader(dataset, batch_size, num_workers, shuffle=True)) + + try: + trainer.predict( + model, + dataloaders=_make_loader(dataset, batch_size, num_workers, shuffle=False), + return_predictions=False, + ) + finally: + writer.close() + + logger.info(f"Wrote denoised samples to {out_h5}") + return None diff --git a/src/tokeye/training/big_tf_unet_2/step_4_labels.py b/src/tokeye/training/big_tf_unet_2/step_4_labels.py new file mode 100644 index 0000000..0584c20 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_4_labels.py @@ -0,0 +1,128 @@ +"""step_4 — label masks from knee-point thresholds on robust-z fields. + +Two label channels per window, thresholded per (shot, channel) so a strong +shot cannot set another shot's threshold: + +- **coherent** (channel 0): the denoised field from step_3, +- **transient** (channel 1): |dt B| of the step_2 baseline — broadband bursts + move the baseline level between neighboring time columns. + +Both fields are robust-standardized ((x - median) / (1.4826*MAD), pooled over +the shot's windows per channel) and thresholded at the Kneedle knee of the +positive-z ECDF (``kneed``), replacing the hand-rolled triangle method. The +threshold lives in robust-sigma units, so knobs are comparable across shots, +channels, modalities, and scales. No asinh here: knee location is not +invariant under monotone compression. + +Output ``/samples/{idx}``: uint8 ``(C, 2, F, T)`` (coherent, transient). +Every (shot, channel, target) threshold is logged to ``thresholds.csv``. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import h5py +import numpy as np +import pandas as pd +from skimage.morphology import remove_small_objects + +from .utils.auto_params import knee_threshold, robust_stats +from .utils.hdf5_io import create_step_file, read_sample, write_sample + +if TYPE_CHECKING: + from typing import Any + +logger = logging.getLogger(__name__) + + +def _magnitude(arr: np.ndarray) -> np.ndarray: + """(C, F, T, 2) real/imag -> (C, F, T) magnitude; (C, F, T) passes through.""" + if arr.ndim == 4 and arr.shape[-1] == 2: + return np.sqrt(arr[..., 0] ** 2 + arr[..., 1] ** 2) + return np.abs(arr) + + +def _transient_field(baseline: np.ndarray) -> np.ndarray: + """|dt| of the per-channel mean baseline level, width-preserving.""" + level = baseline.mean(axis=-1) if baseline.ndim == 4 else baseline # (C, F, T) + return np.abs(np.diff(level, axis=-1, prepend=level[..., :1])) + + +def _postprocess(mask: np.ndarray, settings: dict) -> np.ndarray: + """Row removal + small-object removal on one (F, T) bool mask.""" + bottom = settings["remove_bottom_rows"] + top = settings["remove_top_rows"] + if bottom > 0: + mask[:bottom] = False + if top > 0: + mask[-top:] = False + return remove_small_objects(mask, min_size=settings["min_size"]) + + +def _label_shot( + fields: np.ndarray, # (W, C, F, T) one target field for one shot + settings: dict, +) -> tuple[np.ndarray, list[dict[str, Any]]]: + """Threshold one target for one shot. Returns (W, C, F, T) bool + rows.""" + n_windows, n_channels = fields.shape[:2] + masks = np.zeros(fields.shape, dtype=bool) + rows = [] + for c in range(n_channels): + pooled = fields[:, c] # (W, F, T) + med, scale = robust_stats(pooled) + z = (pooled - med) / scale + result = knee_threshold( + z, + sensitivity=settings["knee_sensitivity"], + delta=settings["delta"], + fallback_frac=settings["fallback_frac"], + ) + raw = z > result["threshold"] + for w in range(n_windows): + masks[w, c] = _postprocess(raw[w], settings) + rows.append({"channel": c, **result}) + return masks, rows + + +def main(settings: dict) -> None: + frame_info = pd.read_csv(settings["frame_info"]) + by_shot = frame_info.groupby("shotn")["index"].apply(list) + + threshold_rows: list[dict[str, Any]] = [] + out = create_step_file(settings["out_h5"], metadata={"targets": "coherent,transient"}) + try: + with ( + h5py.File(settings["denoised_h5"], "r") as f_den, + h5py.File(settings["baseline_h5"], "r") as f_base, + ): + for shot, indices in by_shot.items(): + coherent = np.stack( + [_magnitude(read_sample(f_den, i)) for i in indices] + ) + transient = np.stack( + [_transient_field(read_sample(f_base, i)) for i in indices] + ) + shot_masks = [] + for target, fields in ( + ("coherent", coherent), + ("transient", transient), + ): + masks, rows = _label_shot(fields, settings) + shot_masks.append(masks) + for row in rows: + threshold_rows.append({"shotn": shot, "target": target, **row}) + # (W, C, F, T) x2 -> per window (C, 2, F, T) + stacked = np.stack(shot_masks, axis=2).astype(np.uint8) + for w, idx in enumerate(indices): + write_sample(out, idx, stacked[w]) + logger.info( + f"shot {shot}: {len(indices)} windows labeled " + f"(coherent+transient, {coherent.shape[1]} channels)" + ) + finally: + out.close() + + pd.DataFrame(threshold_rows).to_csv(settings["thresholds_csv"], index=False) + logger.info(f"thresholds -> {settings['thresholds_csv']}") diff --git a/src/tokeye/training/big_tf_unet_2/step_5_dataset.py b/src/tokeye/training/big_tf_unet_2/step_5_dataset.py new file mode 100644 index 0000000..e689b19 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_5_dataset.py @@ -0,0 +1,84 @@ +"""step_5 — merge all modalities into the segmenter training dataset. + +Replaces the TIF-based step_6a: one HDF5 with ``/images/{n}`` float32 +``(1, F, T)`` and ``/masks/{n}`` uint8 ``(2, F, T)`` (each (window, channel) +pair is one training sample), plus provenance datasets so any sample maps +back to (modality, window, channel). + +Image normalization is the pipeline's single convention at the segmenter +site: ``N_a(x) = a*asinh((x - median)/(a*scale))`` of ``log1p(|Z|)`` with +**per-modality global** robust stats (median/1.4826*MAD from sampled +windows) and a=3 — the smooth, invertible generalization of the old +``clip(z, -3, 3)``. The stats are stored in the file attrs and returned to +the runner's ledger; step_7 exports them in the deploy manifest, replacing +the old hardcoded ``NORMALIZATION_CONFIGS`` constants. +""" + +from __future__ import annotations + +import logging + +import h5py +import numpy as np + +from .utils.auto_params import compute_logmag_robust_stats, normalize_asinh +from .utils.hdf5_io import iter_samples, read_sample + +logger = logging.getLogger(__name__) + + +def main(settings: dict) -> dict: + a = settings["a"] + out_path = settings["out_h5"] + out_path.parent.mkdir(parents=True, exist_ok=True) + + in_step: dict[str, float] = {} + prov_modality: list[str] = [] + prov_window: list[int] = [] + prov_channel: list[int] = [] + + with h5py.File(out_path, "w") as out: + out.attrs["a"] = a + images = out.create_group("images") + masks = out.create_group("masks") + n = 0 + for mod, spec in settings["inputs"].items(): + med, scale = compute_logmag_robust_stats( + spec["img_h5"], max_samples=settings["stats_windows"] + ) + out.attrs[f"{mod}_median"] = med + out.attrs[f"{mod}_scale"] = scale + in_step[f"{mod}_median"] = med + in_step[f"{mod}_scale"] = scale + + n_mod = 0 + with h5py.File(spec["mask_h5"], "r") as f_mask: + for idx, img in iter_samples(spec["img_h5"]): + mag = np.sqrt(img[..., 0] ** 2 + img[..., 1] ** 2) # (C, F, T) + norm = normalize_asinh(np.log1p(mag), a, med, scale).astype( + np.float32 + ) + mask = read_sample(f_mask, idx).astype(np.uint8) # (C, 2, F, T) + for c in range(norm.shape[0]): + images.create_dataset( + str(n), data=norm[c][None], compression="lzf" + ) + masks.create_dataset( + str(n), data=mask[c], compression="lzf" + ) + prov_modality.append(mod) + prov_window.append(idx) + prov_channel.append(c) + n += 1 + n_mod += 1 + logger.info(f"{mod}: {n_mod} samples (median={med:.4f}, scale={scale:.4f})") + + out.create_dataset( + "prov_modality", data=np.array(prov_modality, dtype="S8") + ) + out.create_dataset("prov_window", data=np.array(prov_window, dtype=np.int64)) + out.create_dataset("prov_channel", data=np.array(prov_channel, dtype=np.int64)) + out.attrs["n_samples"] = n + logger.info(f"dataset complete: {n} samples -> {out_path}") + + return in_step diff --git a/src/tokeye/training/big_tf_unet_2/step_6_refine.py b/src/tokeye/training/big_tf_unet_2/step_6_refine.py new file mode 100644 index 0000000..315e279 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_6_refine.py @@ -0,0 +1,283 @@ +"""step_6 — leakage-free out-of-fold (OOF) probabilities for label refinement. + +K-fold cross-validation where each fold's model predicts ONLY its held-out +samples, so every sample gets exactly one probability from a model that never +saw it. This replaces the old step_6b/6c refiner, which (a) predicted the +full dataset with every fold and averaged — 4/5 of each "refined" label came +from models that had memorized that sample — and (b) added pixels where the +MC-dropout ensemble was most *uncertain* (entropy > threshold peaks at +p = 0.5). No MC-dropout here: one inference pass per sample. + +Outputs (``step_6.h5``): +- ``p_oof`` (N, 2, F, T) float32 — out-of-fold sigmoid probabilities +- ``disagreement`` (N, 2, F, T) float32 — |p_oof - y0|, the intern's visual + QC map (replaces the deleted manual mask editor) +- ``fold`` (N,) int8 — which fold held each sample out +- attrs: per-fold validation dice/iou + +The soft target q = (1 - model_trust)*y0 + model_trust*p_oof is deliberately +NOT stored here: model_trust is an intern knob, and applying it downstream +(step_7) means tuning it re-trains one final model instead of five folds. +""" + +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING + +import h5py +import lightning as L +import numpy as np +import torch +from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from sklearn.model_selection import KFold +from torch.utils.data import DataLoader, Dataset, Subset + +from tokeye.models.modules.unet import UNet + +from .utils.augmentations import get_augmentation +from .utils.losses import dice_coefficient, get_loss_function, iou_score + +if TYPE_CHECKING: + from pathlib import Path + +# Legacy API on purpose: Lightning still calls get_float32_matmul_precision(), +# and mixing it with the new torch.backends.*.fp32_precision API raises. +torch.set_float32_matmul_precision("high") +torch.backends.cudnn.allow_tf32 = True + +logger = logging.getLogger(__name__) + +# Augmentation + loss knobs the legacy refiner trained with (kept fixed — +# they are not intern knobs). +_TRAIN_DEFAULTS = { + "augmentation": True, + "aug_rotation_degrees": 180, + "aug_prob_flip": 0.5, + "aug_elastic": True, + "aug_elastic_alpha": 50.0, + "aug_elastic_sigma": 5.0, + "aug_scale_range": [0.8, 1.2], + "aug_intensity": True, + "aug_brightness_range": [0.8, 1.2], + "aug_contrast_range": [0.8, 1.2], + "aug_noise_std": 0.05, + "aug_blur_prob": 0.3, + "aug_blur_sigma_range": [0.5, 1.5], + "aug_gamma_range": [0.8, 1.2], + "aug_apply_prob": 0.8, + "specaugment": True, + "specaug_time_warp_W": 20, + "specaug_freq_mask_F": 5, + "specaug_time_mask_T": 5, + "specaug_freq_mask_num": 0, + "specaug_time_mask_num": 0, + "label_smoothing": 0.1, + "symmetric_alpha": 0.1, + "symmetric_beta": 1.0, + "symmetric_weight": 0.5, + "dice_weight": 0.5, + "bce_weight": 0.5, + "focal_alpha": 0.25, + "focal_gamma": 2.0, + "focal_weight": 0.5, + "learning_rate": 1e-4, + "weight_decay": 1e-5, + "lr_factor": 0.5, + "lr_patience": 5, + "lr_min": 1e-6, +} + + +class DatasetH5(Dataset): + """(image, mask) pairs from step_5.h5, lazy per-worker file handle.""" + + def __init__(self, h5_path: Path, transform=None) -> None: + self.h5_path = h5_path + self.transform = transform + with h5py.File(h5_path, "r") as f: + self.n = int(f.attrs["n_samples"]) + self._file: h5py.File | None = None + + def __len__(self) -> int: + return self.n + + def _f(self) -> h5py.File: + if self._file is None: + self._file = h5py.File(self.h5_path, "r", swmr=True) + return self._file + + def __getitem__(self, idx: int): + f = self._f() + x = torch.from_numpy(np.asarray(f["images"][str(idx)])).float() + y = torch.from_numpy(np.asarray(f["masks"][str(idx)])).float() + if self.transform is not None: + x, y = self.transform(x, y) + return x, y + + +class RefineModule(L.LightningModule): + def __init__(self, settings: dict) -> None: + super().__init__() + self.unet = UNet( + in_channels=1, + out_channels=2, + num_layers=settings["num_layers"], + first_layer_size=settings["first_layer_size"], + ) + if not settings["smoke"] and torch.cuda.is_available(): + self.unet.compile() + loss_settings = {**_TRAIN_DEFAULTS, "loss_type": settings["loss_type"]} + self.loss_fn = get_loss_function(loss_settings) + self.lr_settings = loss_settings + + def forward(self, x): + return self.unet(x) + + def training_step(self, batch, _): + x, y = batch + loss = self.loss_fn(self(x), y) + self.log("train_loss", loss) + return loss + + def validation_step(self, batch, _): + x, y = batch + y_hat = self(x) + loss = self.loss_fn(y_hat, y) + probs = torch.sigmoid(y_hat) + self.log("val_loss", loss, prog_bar=True) + self.log("val_dice", dice_coefficient(probs, y)) + self.log("val_iou", iou_score(probs, y)) + return loss + + def predict_step(self, batch, _): + x, _y = batch + return torch.sigmoid(self(x)) + + def configure_optimizers(self): + s = self.lr_settings + opt = torch.optim.Adam( + self.parameters(), lr=s["learning_rate"], weight_decay=s["weight_decay"] + ) + sched = torch.optim.lr_scheduler.ReduceLROnPlateau( + opt, factor=s["lr_factor"], patience=s["lr_patience"], min_lr=s["lr_min"] + ) + return { + "optimizer": opt, + "lr_scheduler": {"scheduler": sched, "monitor": "val_loss"}, + } + + +def _loader(dataset, settings: dict, shuffle: bool) -> DataLoader: + workers = settings["num_workers"] + return DataLoader( + dataset, + batch_size=settings["batch_size"], + shuffle=shuffle, + num_workers=workers, + pin_memory=torch.cuda.is_available(), + persistent_workers=workers > 0, + prefetch_factor=2 if workers > 0 else None, + ) + + +def main(settings: dict) -> None: + L.seed_everything(42) + in_h5 = settings["in_h5"] + ckpt_dir = settings["ckpt_dir"] + ckpt_dir.mkdir(parents=True, exist_ok=True) + + train_full = DatasetH5(in_h5, transform=get_augmentation(_TRAIN_DEFAULTS)) + eval_full = DatasetH5(in_h5, transform=None) + n = len(eval_full) + sample_shape = eval_full[0][1].shape # (2, F, T) + + with h5py.File(settings["out_h5"], "w") as out: + p_oof = out.create_dataset( + "p_oof", + shape=(n, *sample_shape), + dtype=np.float32, + chunks=(1, *sample_shape), + compression="lzf", + ) + disagreement = out.create_dataset( + "disagreement", + shape=(n, *sample_shape), + dtype=np.float32, + chunks=(1, *sample_shape), + compression="lzf", + ) + fold_of = out.create_dataset("fold", shape=(n,), dtype=np.int8) + + kfold = KFold(n_splits=settings["n_folds"], shuffle=True, random_state=42) + fold_metrics = [] + for fold, (train_idx, val_idx) in enumerate(kfold.split(np.arange(n))): + logger.info( + f"fold {fold}: train={len(train_idx)}, held-out={len(val_idx)}" + ) + module = RefineModule(settings) + ckpt = ModelCheckpoint( + dirpath=ckpt_dir / f"fold_{fold}", + monitor="val_loss", + save_top_k=1, + ) + trainer = L.Trainer( + max_epochs=settings["max_epochs"], + precision=settings["precision"], + accelerator="auto", + devices=1, + callbacks=[ + ckpt, + EarlyStopping(monitor="val_loss", patience=3), + ], + enable_progress_bar=False, + log_every_n_steps=10, + logger=False, + ) + trainer.fit( + module, + train_dataloaders=_loader( + Subset(train_full, train_idx), settings, shuffle=True + ), + val_dataloaders=_loader( + Subset(eval_full, val_idx), settings, shuffle=False + ), + ) + fold_metrics.append( + { + "fold": fold, + "val_loss": float(trainer.callback_metrics["val_loss"]), + "val_dice": float(trainer.callback_metrics["val_dice"]), + "val_iou": float(trainer.callback_metrics["val_iou"]), + } + ) + + # Predict ONLY the held-out fold — this is the leakage fix. + preds = trainer.predict( + module, + dataloaders=_loader( + Subset(eval_full, val_idx), settings, shuffle=False + ), + ckpt_path=ckpt.best_model_path or None, + ) + probs = torch.cat(preds).float().cpu().numpy() + for j, global_idx in enumerate(val_idx): + y0 = np.asarray(eval_full._f()["masks"][str(int(global_idx))]) + p_oof[global_idx] = probs[j] + disagreement[global_idx] = np.abs(probs[j] - y0) + fold_of[global_idx] = fold + logger.info(f"fold {fold} OOF written ({len(val_idx)} samples)") + + out.attrs["n_folds"] = settings["n_folds"] + out.attrs["fold_metrics"] = json.dumps(fold_metrics) + + # Every sample must have been predicted by exactly one fold. + counts = np.zeros(n, dtype=int) + for _, val_idx in KFold( + n_splits=settings["n_folds"], shuffle=True, random_state=42 + ).split(np.arange(n)): + counts[val_idx] += 1 + if not np.all(counts == 1): + raise RuntimeError("OOF coverage broken: some samples not predicted once") + logger.info(f"step_6 complete: OOF coverage verified for {n} samples") diff --git a/src/tokeye/training/big_tf_unet_2/step_7_final.py b/src/tokeye/training/big_tf_unet_2/step_7_final.py new file mode 100644 index 0000000..6be9f20 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_7_final.py @@ -0,0 +1,182 @@ +"""step_7 — final segmenter trained on soft targets, exported for deployment. + +The training target is the shrinkage blend of the knee labels and the +out-of-fold model evidence from step_6: + + q = (1 - model_trust) * y0 + model_trust * p_oof + +trained with BCE (soft bootstrapping): BCE's minimizer is E[q], so the model +learns the smoothed consensus. model_trust = 0 reproduces training directly +on the knee labels (no refine); no extra label smoothing is applied (q is +already soft). Unlike the old never-drop-positives heuristic this is +symmetric — confident model disagreement can remove knee false positives. + +Exports to ``model_dir``: best checkpoint, ``final.torchscript.pt`` (traced +from a fresh uncompiled UNet), and ``deploy_manifest.yaml`` carrying +everything inference needs: scale (nfft/hop), per-modality robust stats and +``a`` from step_5, and the edge bins from the resolved-params ledger. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import h5py +import lightning as L +import numpy as np +import torch +import yaml +from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from torch.utils.data import Dataset, Subset + +from tokeye.models.modules.unet import UNet + +from .step_6_refine import _TRAIN_DEFAULTS, RefineModule, _loader +from .utils.augmentations import get_augmentation + +if TYPE_CHECKING: + from pathlib import Path + +torch.set_float32_matmul_precision("high") +torch.backends.cudnn.allow_tf32 = True + +logger = logging.getLogger(__name__) + + +class SoftTargetDataset(Dataset): + """(image, q) pairs: images from step_5, q blended on the fly.""" + + def __init__( + self, + dataset_h5: Path, + refine_h5: Path, + model_trust: float, + transform=None, + ) -> None: + self.dataset_h5 = dataset_h5 + self.refine_h5 = refine_h5 + self.model_trust = model_trust + self.transform = transform + with h5py.File(dataset_h5, "r") as f: + self.n = int(f.attrs["n_samples"]) + self._data: h5py.File | None = None + self._refine: h5py.File | None = None + + def __len__(self) -> int: + return self.n + + def _files(self) -> tuple[h5py.File, h5py.File]: + if self._data is None: + self._data = h5py.File(self.dataset_h5, "r", swmr=True) + self._refine = h5py.File(self.refine_h5, "r", swmr=True) + return self._data, self._refine + + def __getitem__(self, idx: int): + data, refine = self._files() + x = torch.from_numpy(np.asarray(data["images"][str(idx)])).float() + y0 = torch.from_numpy(np.asarray(data["masks"][str(idx)])).float() + lam = self.model_trust + if lam > 0: + p = torch.from_numpy(np.asarray(refine["p_oof"][idx])).float() + q = (1 - lam) * y0 + lam * p + else: + q = y0 + if self.transform is not None: + x, q = self.transform(x, q) + return x, q + + +def main(settings: dict) -> None: + L.seed_everything(42) + model_dir = settings["model_dir"] + model_dir.mkdir(parents=True, exist_ok=True) + + train_settings = {**settings, "loss_type": settings["loss_type"]} + train_ds = SoftTargetDataset( + settings["dataset_h5"], + settings["refine_h5"], + settings["model_trust"], + transform=get_augmentation(_TRAIN_DEFAULTS), + ) + val_ds = SoftTargetDataset( + settings["dataset_h5"], + settings["refine_h5"], + settings["model_trust"], + transform=None, + ) + n = len(train_ds) + rng = np.random.default_rng(42) + perm = rng.permutation(n) + n_val = max(1, int(0.1 * n)) + val_idx, train_idx = perm[:n_val], perm[n_val:] + + module = RefineModule(train_settings) + ckpt = ModelCheckpoint(dirpath=model_dir, monitor="val_loss", save_top_k=1) + trainer = L.Trainer( + max_epochs=settings["max_epochs"], + precision=settings["precision"], + accelerator="auto", + devices=1, + callbacks=[ckpt, EarlyStopping(monitor="val_loss", patience=5)], + enable_progress_bar=False, + log_every_n_steps=10, + logger=False, + ) + trainer.fit( + module, + train_dataloaders=_loader(Subset(train_ds, train_idx), settings, shuffle=True), + val_dataloaders=_loader(Subset(val_ds, val_idx), settings, shuffle=False), + ) + logger.info(f"best checkpoint: {ckpt.best_model_path}") + + # TorchScript export: trace a FRESH uncompiled UNet (compiled modules + # cannot be traced), loading the best weights. + best = RefineModule.load_from_checkpoint( + ckpt.best_model_path, settings={**train_settings, "smoke": True} + ) + export_net = UNet( + in_channels=1, + out_channels=2, + num_layers=settings["num_layers"], + first_layer_size=settings["first_layer_size"], + ) + export_net.load_state_dict(best.unet.state_dict()) + export_net.eval() + sample_x, _ = val_ds[0] + example = sample_x[None] # (1, 1, F, T) + with torch.no_grad(): + traced = torch.jit.trace(export_net, example) + ts_path = model_dir / "final.torchscript.pt" + traced.save(str(ts_path)) + logger.info(f"TorchScript -> {ts_path}") + + # Deploy manifest: everything inference needs to reproduce the input + # normalization at this scale. + with h5py.File(settings["dataset_h5"], "r") as f: + stats = { + k: float(v) + for k, v in f.attrs.items() + if k.endswith(("_median", "_scale")) + } + a = float(f.attrs["a"]) + resolved = {} + if settings["resolved_params"].exists(): + resolved = yaml.safe_load(settings["resolved_params"].read_text()) or {} + edge_bins = { + mod: {k: entry[k]["value"] for k in entry if k.startswith("edge_bins")} + for mod, entry in resolved.get("step_2", {}).items() + } + manifest = { + "run_id": settings["run_id"], + "nfft": settings["nfft"], + "hop": settings["hop"], + "normalization": {"form": "a*asinh((log1p|Z| - median)/(a*scale))", "a": a}, + "modality_stats": stats, + "edge_bins": edge_bins, + "model_trust": settings["model_trust"], + "checkpoint": str(ckpt.best_model_path), + "torchscript": str(ts_path), + } + settings["deploy_manifest"].write_text(yaml.safe_dump(manifest, sort_keys=False)) + logger.info(f"manifest -> {settings['deploy_manifest']}") diff --git a/src/tokeye/training/big_tf_unet_2/step_8_eval.py b/src/tokeye/training/big_tf_unet_2/step_8_eval.py new file mode 100644 index 0000000..84479dc --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/step_8_eval.py @@ -0,0 +1,127 @@ +"""step_8 — quick TJ-II benchmark: one meaningful number per run. + +Runs the exported TorchScript model on the public TJ-II 2021 set +(``data/eval/TJII2021``) with the SAME preprocessing as the paper eval +(vertical flip, the TJ-II pixel mean/std, reflect-pad to /32), sweeps the +sigmoid threshold, and reports per-image IoU + F1 at the F1-optimal +threshold. The deployed big_tf_unet anchor on this metric is IoU 0.260 / +F1 0.478 (F1-opt threshold 0.80). + +Caveat printed with the result: the TJ-II images are fixed-resolution +pre-rendered spectrograms, so a teacher trained at a different (nfft, hop) +sees a mild domain shift here — treat the number as a relative tracker +across runs, not an absolute target. +""" + +from __future__ import annotations + +import logging + +import numpy as np +import pandas as pd +import torch +from PIL import Image + +logger = logging.getLogger(__name__) + +# TJ-II pixel statistics used by the paper eval (dev/paper/eval/TJII2021*.py). +_TJII_MEAN = 17.84620821169868 +_TJII_STD = 25.016818830630463 +_ANCHOR = {"iou": 0.260, "f1": 0.478, "threshold": 0.80} + + +def _load_images(dataset_dir): + """Yield (spec (1,1,H,W) float, gt (H,W) bool) per TJ-II shot.""" + for spec_path in sorted((dataset_dir / "input").glob("*.png")): + gt_path = dataset_dir / "gt" / spec_path.name + if not gt_path.exists(): + continue + spec = np.array(Image.open(spec_path).convert("L")) + ann = np.array(Image.open(gt_path).convert("L")) + if spec.shape != ann.shape: # one malformed shot in the public set + logger.warning(f"skipping {spec_path.name}: size mismatch") + continue + spec = np.flip(spec, axis=0).copy() + ann = np.flip(ann, axis=0).copy() + spec = (spec - _TJII_MEAN) / _TJII_STD + yield ( + torch.from_numpy(spec).float()[None, None], + (ann // 255).astype(bool), + ) + + +def _pad32(x: torch.Tensor) -> torch.Tensor: + h, w = x.shape[-2:] + ph, pw = (-h) % 32, (-w) % 32 + if ph or pw: + x = torch.nn.functional.pad(x, (0, pw, 0, ph), mode="reflect") + return x + + +def main(settings: dict) -> None: + ts_path = settings["model_dir"] / "final.torchscript.pt" + if not ts_path.exists(): + raise FileNotFoundError(f"no exported model at {ts_path} — run step_7 first") + device = "cuda" if torch.cuda.is_available() else "cpu" + model = torch.jit.load(str(ts_path), map_location=device).eval() + + thresholds = np.linspace(0.02, 0.98, settings["n_thresholds"]) + tp = np.zeros(len(thresholds)) + fp = np.zeros(len(thresholds)) + fn = np.zeros(len(thresholds)) + iou_per_image: list[np.ndarray] = [] + + n_images = 0 + with torch.no_grad(): + for spec, gt in _load_images(settings["dataset_dir"]): + h, w = spec.shape[-2:] + out = model(_pad32(spec).to(device)) + if isinstance(out, (tuple, list)): + out = out[0] + prob = torch.sigmoid(out[:, 0:1])[..., :h, :w].cpu().numpy()[0, 0] + gt_sum = gt.sum() + image_iou = np.zeros(len(thresholds)) + for i, t in enumerate(thresholds): + pred = prob > t + inter = np.logical_and(pred, gt).sum() + pred_sum = pred.sum() + tp[i] += inter + fp[i] += pred_sum - inter + fn[i] += gt_sum - inter + union = pred_sum + gt_sum - inter + image_iou[i] = inter / union if union > 0 else 1.0 + iou_per_image.append(image_iou) + n_images += 1 + + if n_images == 0: + raise RuntimeError(f"no evaluable images in {settings['dataset_dir']}") + + iou_matrix = np.stack(iou_per_image) # (n_images, n_thresholds) + precision = tp / np.maximum(tp + fp, 1) + recall = tp / np.maximum(tp + fn, 1) + f1 = 2 * precision * recall / np.maximum(precision + recall, 1e-12) + best = int(np.argmax(f1)) + + rows = pd.DataFrame( + { + "threshold": thresholds, + "precision": precision, + "recall": recall, + "f1": f1, + "iou_per_image_mean": iou_matrix.mean(axis=0), + } + ) + rows.to_csv(settings["out_csv"], index=False) + + iou_opt = float(iou_matrix.mean(axis=0)[best]) + logger.info( + f"TJ-II ({n_images} images): IoU={iou_opt:.3f}, F1={f1[best]:.3f} " + f"at threshold {thresholds[best]:.2f} " + f"(deployed anchor: IoU={_ANCHOR['iou']}, F1={_ANCHOR['f1']} " + f"at {_ANCHOR['threshold']})" + ) + logger.info( + "note: TJ-II images are fixed-resolution renders — for a per-scale " + "teacher this number is a relative tracker across runs, not a target" + ) + logger.info(f"sweep -> {settings['out_csv']}") diff --git a/src/tokeye/training/big_tf_unet_2/task_matrix.py b/src/tokeye/training/big_tf_unet_2/task_matrix.py new file mode 100644 index 0000000..d898337 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/task_matrix.py @@ -0,0 +1,229 @@ +"""Persistent per-run progress tracker with human sign-off and staleness. + +JSON-backed (atomic write + flock, safe for concurrent SLURM jobs), one file +per run at ``/task_matrix.json``. Entries are keyed ``step_N`` for +combined steps and ``step_N:`` for per-modality steps. + +Beyond the ablation tracker this adds: + +- a status enum (``pending``/``running``/``complete``/``failed``/``stale``), +- ``accepted`` — the intern's explicit visual sign-off per step, +- ``params_hash`` — hash of the resolved settings that produced the artifact; + rerunning a step marks everything downstream ``stale`` so nothing can train + on outputs that no longer match their inputs, +- ``job_id`` — the SLURM job that is producing (or produced) the artifact. +""" + +from __future__ import annotations + +import fcntl +import hashlib +import json +from datetime import datetime +from pathlib import Path +from typing import Any + +from .paths import get_step, steps_after + +STATUSES = ("pending", "submitted", "running", "complete", "failed", "stale") + + +def params_hash(settings: dict[str, Any]) -> str: + """Stable hash of a step's resolved settings dict.""" + blob = json.dumps(settings, sort_keys=True, default=str) + return hashlib.sha1(blob.encode()).hexdigest()[:12] + + +class RunTaskMatrix: + """JSON-backed progress tracker for one pipeline run.""" + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._data: dict = self._load() + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + def _load(self) -> dict: + if self.path.exists(): + return json.loads(self.path.read_text()) + return {"entries": {}} + + def _save(self) -> None: + tmp = self.path.with_suffix(".tmp") + content = json.dumps(self._data, indent=2) + with tmp.open("w") as f: + fcntl.flock(f, fcntl.LOCK_EX) + f.write(content) + f.flush() + fcntl.flock(f, fcntl.LOCK_UN) + tmp.rename(self.path) + + def reload(self) -> None: + self._data = self._load() + + # ------------------------------------------------------------------ + # Keys + # ------------------------------------------------------------------ + + @staticmethod + def key(step: str, modality: str | None = None) -> str: + spec = get_step(step) + if spec.per_modality: + if modality is None: + raise ValueError(f"{step} is per-modality; a modality is required") + return f"{step}:{modality}" + return step + + def keys_for(self, step: str, modalities: list[str]) -> list[str]: + spec = get_step(step) + if spec.per_modality: + return [f"{step}:{m}" for m in modalities] + return [step] + + def _entry(self, key: str) -> dict: + return self._data["entries"].get(key, {}) + + # ------------------------------------------------------------------ + # State transitions + # ------------------------------------------------------------------ + + def _set(self, key: str, **fields: Any) -> None: + entry = self._data["entries"].setdefault(key, {}) + entry.update(fields, timestamp=datetime.now().isoformat()) + + def mark_running( + self, step: str, modality: str | None = None, job_id: str | None = None + ) -> None: + self.reload() + entry = self._entry(self.key(step, modality)) + job_id = job_id or entry.get("job_id") + self._set( + self.key(step, modality), status="running", job_id=job_id, accepted=False + ) + self._save() + + def record_job( + self, step: str, modality: str | None, job_id: str + ) -> None: + """Mark a step as submitted to SLURM (queued, not yet running).""" + self.reload() + self._set( + self.key(step, modality), status="submitted", job_id=job_id, + accepted=False, + ) + self._save() + + def mark_complete( + self, + step: str, + modality: str | None, + step_params_hash: str, + modalities: list[str], + ) -> None: + """Record completion and mark all downstream complete work stale.""" + self.reload() + self._set( + self.key(step, modality), + status="complete", + params_hash=step_params_hash, + accepted=False, + ) + self._invalidate_downstream(step, modality, modalities) + self._save() + + def mark_failed(self, step: str, modality: str | None = None) -> None: + self.reload() + self._set(self.key(step, modality), status="failed", accepted=False) + self._save() + + def mark_pending(self, step: str, modalities: list[str]) -> None: + """Forget a step entirely (used by clearing) and stale its downstream.""" + self.reload() + for key in self.keys_for(step, modalities): + self._data["entries"].pop(key, None) + self._invalidate_downstream(step, None, modalities) + self._save() + + def _invalidate_downstream( + self, step: str, modality: str | None, modalities: list[str] + ) -> None: + """Stale every downstream entry that currently claims completion. + + A per-modality change only stales that modality's own downstream chain + plus every combined step; other modalities' per-modality work survives. + """ + for spec in steps_after(step): + if spec.per_modality: + mods = [modality] if modality is not None else modalities + keys = [f"{spec.name}:{m}" for m in mods] + else: + keys = [spec.name] + for key in keys: + if self._entry(key).get("status") == "complete": + self._set(key, status="stale", accepted=False) + + # ------------------------------------------------------------------ + # Acceptance (human sign-off) + # ------------------------------------------------------------------ + + def accept(self, step: str, modalities: list[str]) -> None: + self.reload() + keys = self.keys_for(step, modalities) + not_done = [k for k in keys if self._entry(k).get("status") != "complete"] + if not_done: + raise ValueError( + f"Cannot accept {step}: not complete for {', '.join(not_done)}" + ) + for key in keys: + self._set(key, accepted=True) + self._save() + + def is_accepted(self, step: str, modalities: list[str]) -> bool: + return all( + self._entry(k).get("accepted", False) + for k in self.keys_for(step, modalities) + ) + + # ------------------------------------------------------------------ + # Queries + # ------------------------------------------------------------------ + + def status(self, step: str, modality: str | None = None) -> str: + return self._entry(self.key(step, modality)).get("status", "pending") + + def is_complete(self, step: str, modality: str | None = None) -> bool: + return self.status(step, modality) == "complete" + + def all_complete(self, step: str, modalities: list[str]) -> bool: + spec = get_step(step) + if spec.per_modality: + return all(self.is_complete(step, m) for m in modalities) + return self.is_complete(step) + + def job_id(self, step: str, modality: str | None = None) -> str | None: + return self._entry(self.key(step, modality)).get("job_id") + + def to_rows(self, modalities: list[str]) -> list[dict[str, Any]]: + """One row per (step, modality) for status display.""" + from .paths import STEPS + + rows = [] + for spec in STEPS: + mods = modalities if spec.per_modality else [None] + for mod in mods: + entry = self._entry(self.key(spec.name, mod)) + rows.append( + { + "step": spec.name, + "title": spec.title, + "modality": mod or "-", + "status": entry.get("status", "pending"), + "accepted": entry.get("accepted", False), + "job_id": entry.get("job_id"), + "updated": entry.get("timestamp"), + } + ) + return rows diff --git a/src/tokeye/training/big_tf_unet_2/utils/augmentations.py b/src/tokeye/training/big_tf_unet_2/utils/augmentations.py new file mode 100644 index 0000000..a051ad3 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/augmentations.py @@ -0,0 +1,528 @@ +""" +Data augmentation transforms for robust segmentation training with noisy labels. +Implements geometric and intensity transforms that apply consistently to image-mask pairs. +Also includes SpecAugment for spectrogram-specific augmentations. +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from scipy.ndimage import gaussian_filter, map_coordinates + +# Module-level random generator for reproducibility +_rng = np.random.default_rng() + + +class SegmentationAugmentation: + """ + Compose multiple augmentation transforms for segmentation. + Applies transforms consistently to both image and mask. + """ + + def __init__( + self, + rotation_degrees: float = 180, + prob_flip_h: float = 0.5, + prob_flip_v: float = 0.5, + elastic: bool = True, + elastic_alpha: float = 50.0, + elastic_sigma: float = 5.0, + scale_range: tuple[float, float] = (0.8, 1.2), + intensity_transforms: bool = True, + brightness_range: tuple[float, float] = (0.8, 1.2), + contrast_range: tuple[float, float] = (0.8, 1.2), + noise_std: float = 0.05, + blur_prob: float = 0.3, + blur_sigma_range: tuple[float, float] = (0.5, 1.5), + gamma_range: tuple[float, float] = (0.8, 1.2), + apply_prob: float = 0.8, + specaugment: SpecAugment | None = None, + rng: np.random.Generator | None = None, + ): + """ + Args: + rotation_degrees: Maximum rotation angle in degrees (±) + prob_flip_h: Probability of horizontal flip + prob_flip_v: Probability of vertical flip + elastic: Enable elastic deformation + elastic_alpha: Elastic deformation alpha parameter + elastic_sigma: Elastic deformation sigma parameter + scale_range: Random scaling range (min, max) + intensity_transforms: Enable intensity augmentations + brightness_range: Brightness multiplier range (min, max) + contrast_range: Contrast multiplier range (min, max) + noise_std: Standard deviation of Gaussian noise + blur_prob: Probability of applying Gaussian blur + blur_sigma_range: Gaussian blur sigma range (min, max) + gamma_range: Gamma correction range (min, max) + apply_prob: Probability of applying augmentation at all + specaugment: Optional SpecAugment instance for spectrogram-specific augmentations + rng: Optional random number generator for reproducibility + """ + self.rotation_degrees = rotation_degrees + self.prob_flip_h = prob_flip_h + self.prob_flip_v = prob_flip_v + self.elastic = elastic + self.elastic_alpha = elastic_alpha + self.elastic_sigma = elastic_sigma + self.scale_range = scale_range + self.intensity_transforms = intensity_transforms + self.brightness_range = brightness_range + self.contrast_range = contrast_range + self.noise_std = noise_std + self.blur_prob = blur_prob + self.blur_sigma_range = blur_sigma_range + self.gamma_range = gamma_range + self.apply_prob = apply_prob + self.specaugment = specaugment + self._rng = rng if rng is not None else _rng + + def __call__( + self, image: torch.Tensor, mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply augmentation to image and mask consistently. + + Args: + image: Tensor of shape (C, H, W) or (H, W) + mask: Tensor of shape (C, H, W) or (H, W) + + Returns: + Augmented (image, mask) pair + """ + if self._rng.random() > self.apply_prob: + return image, mask + + # Ensure proper shape + if image.ndim == 2: + image = image.unsqueeze(0) + if mask.ndim == 2: + mask = mask.unsqueeze(0) + + # Geometric transforms (apply to both image and mask) + image, mask = self._apply_geometric_transforms(image, mask) + + # Intensity transforms (apply only to image) + if self.intensity_transforms: + image = self._apply_intensity_transforms(image) + + return image, mask + + def _apply_geometric_transforms( + self, image: torch.Tensor, mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply geometric transforms consistently to image and mask.""" + + # Random rotation + if self.rotation_degrees > 0: + angle = self._rng.uniform(-self.rotation_degrees, self.rotation_degrees) + image = self._rotate(image, angle, mode="bilinear") + mask = self._rotate(mask, angle, mode="nearest") + + # Random horizontal flip + if self._rng.random() < self.prob_flip_h: + image = torch.flip(image, dims=[-1]) + mask = torch.flip(mask, dims=[-1]) + + # Random vertical flip + if self._rng.random() < self.prob_flip_v: + image = torch.flip(image, dims=[-2]) + mask = torch.flip(mask, dims=[-2]) + + # Random scaling + if self.scale_range[0] != 1.0 or self.scale_range[1] != 1.0: + scale = self._rng.uniform(self.scale_range[0], self.scale_range[1]) + image = self._scale(image, scale, mode="bilinear") + mask = self._scale(mask, scale, mode="nearest") + + # Elastic deformation + if self.elastic and self._rng.random() < 0.5: + image, mask = self._elastic_transform(image, mask) + + return image, mask + + def _apply_intensity_transforms(self, image: torch.Tensor) -> torch.Tensor: + """Apply intensity transforms to image only.""" + + # Apply SpecAugment if provided (spectrogram-specific augmentation) + if self.specaugment is not None: + image = self.specaugment(image) + + # Random brightness + if self.brightness_range[0] != 1.0 or self.brightness_range[1] != 1.0: + brightness = self._rng.uniform( + self.brightness_range[0], self.brightness_range[1] + ) + image = image * brightness + + # Random contrast + if self.contrast_range[0] != 1.0 or self.contrast_range[1] != 1.0: + contrast = self._rng.uniform(self.contrast_range[0], self.contrast_range[1]) + mean = image.mean() + image = (image - mean) * contrast + mean + + # Gaussian noise + if self.noise_std > 0 and self._rng.random() < 0.5: + noise = torch.randn_like(image) * self.noise_std + image = image + noise + + # Gaussian blur + if self._rng.random() < self.blur_prob: + sigma = self._rng.uniform( + self.blur_sigma_range[0], self.blur_sigma_range[1] + ) + image = self._gaussian_blur(image, sigma) + + # Gamma correction + if ( + self.gamma_range[0] != 1.0 or self.gamma_range[1] != 1.0 + ) and self._rng.random() < 0.5: + gamma = self._rng.uniform(self.gamma_range[0], self.gamma_range[1]) + # Normalize to [0, 1], apply gamma, then denormalize + img_min, img_max = image.min(), image.max() + if img_max > img_min: + image_norm = (image - img_min) / (img_max - img_min) + image_norm = torch.pow(image_norm, gamma) + image = image_norm * (img_max - img_min) + img_min + + return image + + @staticmethod + def _rotate( + tensor: torch.Tensor, angle: float, mode: str = "bilinear" + ) -> torch.Tensor: + """Rotate tensor by angle in degrees.""" + angle_rad = angle * np.pi / 180.0 + cos_a, sin_a = np.cos(angle_rad), np.sin(angle_rad) + + # Rotation matrix + theta = torch.tensor( + [[cos_a, -sin_a, 0], [sin_a, cos_a, 0]], + dtype=tensor.dtype, + device=tensor.device, + ).unsqueeze(0) + + # Add batch dimension if needed + needs_squeeze = False + if tensor.ndim == 3: + tensor = tensor.unsqueeze(0) + needs_squeeze = True + + grid = F.affine_grid(theta, tensor.size(), align_corners=False) + rotated = F.grid_sample(tensor, grid, mode=mode, align_corners=False) + + if needs_squeeze: + rotated = rotated.squeeze(0) + + return rotated + + @staticmethod + def _scale( + tensor: torch.Tensor, scale: float, mode: str = "bilinear" + ) -> torch.Tensor: + """Scale tensor by scale factor.""" + theta = torch.tensor( + [[scale, 0, 0], [0, scale, 0]], dtype=tensor.dtype, device=tensor.device + ).unsqueeze(0) + + # Add batch dimension if needed + needs_squeeze = False + if tensor.ndim == 3: + tensor = tensor.unsqueeze(0) + needs_squeeze = True + + grid = F.affine_grid(theta, tensor.size(), align_corners=False) + scaled = F.grid_sample(tensor, grid, mode=mode, align_corners=False) + + if needs_squeeze: + scaled = scaled.squeeze(0) + + return scaled + + def _elastic_transform( + self, image: torch.Tensor, mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply elastic deformation to image and mask consistently. + Based on Simard et al. "Best Practices for Convolutional Neural Networks" + """ + # Convert to numpy for scipy operations + shape = image.shape[-2:] + + # Generate random displacement fields + dx = ( + gaussian_filter( + (self._rng.random(shape) * 2 - 1), + self.elastic_sigma, + mode="constant", + cval=0, + ) + * self.elastic_alpha + ) + dy = ( + gaussian_filter( + (self._rng.random(shape) * 2 - 1), + self.elastic_sigma, + mode="constant", + cval=0, + ) + * self.elastic_alpha + ) + + # Create meshgrid + x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0])) + indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1)) + + # Apply to image + image_np = image.cpu().numpy() + image_warped = np.empty_like(image_np) + for c in range(image_np.shape[0]): + image_warped[c] = map_coordinates( + image_np[c], indices, order=1, mode="reflect" + ).reshape(shape) + + # Apply to mask + mask_np = mask.cpu().numpy() + mask_warped = np.empty_like(mask_np) + for c in range(mask_np.shape[0]): + mask_warped[c] = map_coordinates( + mask_np[c], indices, order=0, mode="reflect" + ).reshape(shape) + + return torch.from_numpy(image_warped), torch.from_numpy(mask_warped) + + @staticmethod + def _gaussian_blur(tensor: torch.Tensor, sigma: float) -> torch.Tensor: + """Apply Gaussian blur to tensor.""" + tensor_np = tensor.cpu().numpy() + blurred = np.empty_like(tensor_np) + for c in range(tensor_np.shape[0]): + blurred[c] = gaussian_filter(tensor_np[c], sigma=sigma) + return torch.from_numpy(blurred).to(tensor.device) + + +class SpecAugment: + """ + SpecAugment for spectrograms: Time Warping, Frequency Masking, Time Masking. + + Reference: "SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition" + Park et al., Interspeech 2019 + + Note: This augmentation is applied only to images, not to masks. + """ + + def __init__( + self, + time_warp_W: int = 40, + freq_mask_F: int = 15, + time_mask_T: int = 20, + num_freq_masks: int = 0, + num_time_masks: int = 0, + apply_prob: float = 0.8, + rng: np.random.Generator | None = None, + ): + """ + Args: + time_warp_W: Time warp parameter W (max displacement along time axis) + freq_mask_F: Maximum width of frequency mask + time_mask_T: Maximum width of time mask + num_freq_masks: Number of frequency masks to apply (0 to disable) + num_time_masks: Number of time masks to apply (0 to disable) + apply_prob: Probability of applying SpecAugment + rng: Optional random number generator for reproducibility + """ + self.time_warp_W = time_warp_W + self.freq_mask_F = freq_mask_F + self.time_mask_T = time_mask_T + self.num_freq_masks = num_freq_masks + self.num_time_masks = num_time_masks + self.apply_prob = apply_prob + self._rng = rng if rng is not None else _rng + + def __call__(self, image: torch.Tensor) -> torch.Tensor: + """ + Apply SpecAugment to image only (not mask). + + Args: + image: Tensor of shape (C, H, W) - H=frequency, W=time + + Returns: + Augmented image + """ + if self._rng.random() > self.apply_prob: + return image + + # Ensure proper shape + if image.ndim == 2: + image = image.unsqueeze(0) + + # Apply time warping + if self.time_warp_W > 0: + image = self._time_warp(image) + + # Apply frequency masking + for _ in range(self.num_freq_masks): + image = self._freq_mask(image) + + # Apply time masking + for _ in range(self.num_time_masks): + image = self._time_mask(image) + + return image + + def _time_warp(self, image: torch.Tensor) -> torch.Tensor: + """ + Apply time warping along the time axis (horizontal). + Warps the spectrogram by displacing a random point along the time axis. + """ + C, H, W = image.shape + + if 2 * self.time_warp_W >= W: + # Image too small for warping + return image + + # Choose a random center point in the middle of time axis + center = self._rng.integers(self.time_warp_W, W - self.time_warp_W) + + # Random warp displacement + warp = self._rng.integers(-self.time_warp_W, self.time_warp_W + 1) + + # Create warping grid + # Generate base grid + y_grid, x_grid = torch.meshgrid( + torch.arange(H, dtype=torch.float32, device=image.device), + torch.arange(W, dtype=torch.float32, device=image.device), + indexing="ij", + ) + + # Apply warping to x coordinates around the center + if warp != 0: + # Linear warping from center + left_dist = torch.clamp(center - x_grid, min=0) + right_dist = torch.clamp(x_grid - center, min=0) + + # Warp left side + warp_left = (left_dist / center) * warp if center > 0 else 0 + # Warp right side + warp_right = (right_dist / (W - center)) * warp if center < W else 0 + + x_grid = x_grid + torch.where(x_grid < center, warp_left, -warp_right) + + # Normalize to [-1, 1] for grid_sample + x_grid = 2.0 * x_grid / (W - 1) - 1.0 + y_grid = 2.0 * y_grid / (H - 1) - 1.0 + + # Stack to (H, W, 2) and add batch dimension + grid = torch.stack([x_grid, y_grid], dim=-1).unsqueeze(0) + + # Apply warping + image_batch = image.unsqueeze(0) + warped = F.grid_sample( + image_batch, + grid, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + return warped.squeeze(0) + + def _freq_mask(self, image: torch.Tensor) -> torch.Tensor: + """ + Apply frequency masking (vertical stripes). + Masks consecutive frequency bins. + """ + C, H, W = image.shape + + # Random mask width + f = self._rng.integers(0, self.freq_mask_F + 1) + + if f == 0 or f >= H: + return image + + # Random starting frequency + f0 = self._rng.integers(0, H - f + 1) + + # Create masked image + masked = image.clone() + masked[:, f0 : f0 + f, :] = 0 + + return masked + + def _time_mask(self, image: torch.Tensor) -> torch.Tensor: + """ + Apply time masking (horizontal stripes). + Masks consecutive time frames. + """ + C, H, W = image.shape + + # Random mask width + t = self._rng.integers(0, self.time_mask_T + 1) + + if t == 0 or t >= W: + return image + + # Random starting time + t0 = self._rng.integers(0, W - t + 1) + + # Create masked image + masked = image.clone() + masked[:, :, t0 : t0 + t] = 0 + + return masked + + +class NoAugmentation: + """Identity transform - no augmentation applied.""" + + def __call__( + self, image: torch.Tensor, mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return image, mask + + +def get_augmentation(settings: dict) -> SegmentationAugmentation | NoAugmentation: + """ + Create augmentation pipeline from settings. + + Args: + settings: Dictionary with augmentation parameters + + Returns: + Augmentation callable + """ + if not settings.get("augmentation", False): + return NoAugmentation() + + # Create SpecAugment if enabled + specaugment = None + if settings.get("specaugment", False): + specaugment = SpecAugment( + time_warp_W=settings.get("specaug_time_warp_W", 40), + freq_mask_F=settings.get("specaug_freq_mask_F", 15), + time_mask_T=settings.get("specaug_time_mask_T", 20), + num_freq_masks=settings.get("specaug_freq_mask_num", 0), + num_time_masks=settings.get("specaug_time_mask_num", 0), + apply_prob=settings.get("aug_apply_prob", 0.8), + ) + + return SegmentationAugmentation( + rotation_degrees=settings.get("aug_rotation_degrees", 180), + prob_flip_h=settings.get("aug_prob_flip", 0.5), + prob_flip_v=settings.get("aug_prob_flip", 0.5), + elastic=settings.get("aug_elastic", True), + elastic_alpha=settings.get("aug_elastic_alpha", 50.0), + elastic_sigma=settings.get("aug_elastic_sigma", 5.0), + scale_range=tuple(settings.get("aug_scale_range", [0.8, 1.2])), + intensity_transforms=settings.get("aug_intensity", True), + brightness_range=tuple(settings.get("aug_brightness_range", [0.8, 1.2])), + contrast_range=tuple(settings.get("aug_contrast_range", [0.8, 1.2])), + noise_std=settings.get("aug_noise_std", 0.05), + blur_prob=settings.get("aug_blur_prob", 0.3), + blur_sigma_range=tuple(settings.get("aug_blur_sigma_range", [0.5, 1.5])), + gamma_range=tuple(settings.get("aug_gamma_range", [0.8, 1.2])), + apply_prob=settings.get("aug_apply_prob", 0.8), + specaugment=specaugment, + ) diff --git a/src/tokeye/training/big_tf_unet_2/utils/auto_params.py b/src/tokeye/training/big_tf_unet_2/utils/auto_params.py new file mode 100644 index 0000000..3423f7f --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/auto_params.py @@ -0,0 +1,355 @@ +"""Automatic parameter derivation for the single-scale pipeline. + +All functions are computed **once per run** from a sample of the data, +producing stable values that are recorded in ``resolved_params.yaml`` and +reused across all samples in that run. + +Ported from ``big_tf_unet_ablation`` and extended with: + +- robust (median/MAD) statistics and the ``N_a`` asinh normalization used at + every model-facing site, +- an energy-level edge-bin detector (the gradient detector masks only 1 of + bes's ~70-bin low-frequency plateau), +- a Kneedle knee-point threshold via the ``kneed`` package (replacing the + hand-rolled triangle method), +- scale-covariant auto values (baseline ``lam``, UNet ``num_layers``, + ``batch_size``) so one config works at every (nfft, hop). +""" + +from __future__ import annotations + +import logging +import math +from typing import TYPE_CHECKING + +import h5py +import numpy as np + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +# Reference geometry: nfft=1024/hop=128 with subseq_len=66000 -> (513, 516). +_REF_FREQ_BINS = 513 +_REF_TIME_BINS = 516 + + +# --------------------------------------------------------------------------- +# Sampling helper +# --------------------------------------------------------------------------- + +def _sample_arrays(h5_path: Path, max_samples: int) -> list[np.ndarray]: + """Read up to *max_samples* random sample arrays from a step HDF5 file.""" + arrays: list[np.ndarray] = [] + with h5py.File(h5_path, "r") as f: + grp = f["samples"] + keys = list(grp.keys()) + rng = np.random.default_rng(42) + indices = rng.choice( + len(keys), size=min(max_samples, len(keys)), replace=False + ) + for idx in indices: + arrays.append(np.asarray(grp[keys[idx]])) + return arrays + + +# --------------------------------------------------------------------------- +# Robust statistics + N_a normalization +# --------------------------------------------------------------------------- + +def robust_stats( + x: np.ndarray, axis: int | tuple[int, ...] | None = None, keepdims: bool = False +) -> tuple[np.ndarray, np.ndarray]: + """Median and scaled MAD (1.4826*MAD ~ sigma for Gaussian data). + + Median/MAD have a 50% breakdown point, so a strong mode cannot inflate the + scale and suppress weak modes the way per-sample mean/std does. + """ + med = np.median(x, axis=axis, keepdims=True) + mad = np.median(np.abs(x - med), axis=axis, keepdims=True) + scale = 1.4826 * mad + 1e-6 + if not keepdims: + med = np.squeeze(med, axis=axis) if axis is not None else med.squeeze() + scale = np.squeeze(scale, axis=axis) if axis is not None else scale.squeeze() + return med, scale + + +def normalize_asinh( + x: np.ndarray, + a: float, + median: np.ndarray | float, + scale: np.ndarray | float, +) -> np.ndarray: + """``N_a(x) = a * asinh((x - median) / (a * scale))``. + + Invertible smooth compression: unit slope through the bulk (|z| << a it IS + the robust z-score), logarithmic in the tails. ``a`` sets where compression + starts — a=3 is the smooth generalization of the old hard clip(z, -3, 3); + unlike clip (or tanh, which saturates), extremes stay ordered and + gradients never die. + """ + return a * np.arcsinh((x - median) / (a * scale)) + + +# --------------------------------------------------------------------------- +# Automatic clamp_range (legacy zscore path of the denoiser) +# --------------------------------------------------------------------------- + +def compute_clamp_range( + h5_path: Path, + percentiles: tuple[float, float] = (1.0, 99.0), + max_samples: int = 100, +) -> tuple[float, float]: + """Global percentile clamp, used only by the legacy zscore normalization.""" + values = [a.ravel() for a in _sample_arrays(h5_path, max_samples)] + all_vals = np.concatenate(values) + lo, hi = np.percentile(all_vals, percentiles) + logger.info(f"Auto clamp_range: ({lo:.4f}, {hi:.4f}) from {len(values)} samples") + return float(lo), float(hi) + + +# --------------------------------------------------------------------------- +# Edge-bin detection (step_2, step_3) +# --------------------------------------------------------------------------- + +def _energy_profile(arrays: list[np.ndarray]) -> np.ndarray: + """Median per-frequency-bin energy profile across sampled windows.""" + profiles = [] + for arr in arrays: + if arr.ndim == 4: # (C, F, T, Z) + energy = np.sqrt(np.mean(arr**2, axis=(0, 2, 3))) + elif arr.ndim == 3: # (C, F, T) + energy = np.mean(np.abs(arr), axis=(0, 2)) + else: # (F, T) + energy = np.mean(np.abs(arr), axis=-1) + profiles.append(energy) + return np.median(np.stack(profiles), axis=0) + + +def detect_edge_bins( + h5_path: Path, + gradient_threshold: float = 0.5, + max_samples: int = 100, +) -> tuple[int, int]: + """Legacy gradient-based edge detector (kept for comparison). + + Known weakness: a broad edge plateau (bes: ~70 bins) has a small interior + gradient, so only ~1 bin gets masked. Prefer ``detect_edge_bins_energy``. + """ + median_energy = _energy_profile(_sample_arrays(h5_path, max_samples)) + grad = np.gradient(median_energy) + grad_norm = np.abs(grad) / (np.max(np.abs(grad)) + 1e-8) + + lower_idx = 0 + for i in range(len(grad_norm)): + if grad_norm[i] < gradient_threshold: + lower_idx = i + break + upper_idx = 0 + for i in range(len(grad_norm) - 1, -1, -1): + if grad_norm[i] < gradient_threshold: + upper_idx = len(grad_norm) - 1 - i + break + + lower_idx = max(1, lower_idx) + upper_idx = max(1, upper_idx) + logger.info(f"Auto edge bins (gradient): lower={lower_idx}, upper={upper_idx}") + return lower_idx, upper_idx + + +def detect_edge_bins_energy( + h5_path: Path, + k: float = 2.0, + max_fraction: float = 0.15, + max_samples: int = 100, +) -> tuple[int, int]: + """Energy-level edge-bin detection. + + From the median per-bin energy profile E(f): the interior band (excluding + a ``max_fraction`` margin at each end) gives a reference median M; from + each edge, the contiguous run of bins with ``|log E(f) - log M| > log k`` + is masked, capped at ``max_fraction`` of the bins per edge. + + Index-based on the run's own profile, so it adapts to any (nfft, hop) and + catches broad plateaus (bes) that the gradient criterion misses. + """ + profile = _energy_profile(_sample_arrays(h5_path, max_samples)) + n = len(profile) + margin = max(1, int(n * max_fraction)) + interior = profile[margin : n - margin] + log_m = np.log(np.median(interior) + 1e-12) + dev = np.abs(np.log(profile + 1e-12) - log_m) + log_k = np.log(k) + + cap = max(1, int(n * max_fraction)) + lower_idx = 0 + while lower_idx < cap and dev[lower_idx] > log_k: + lower_idx += 1 + upper_idx = 0 + while upper_idx < cap and dev[n - 1 - upper_idx] > log_k: + upper_idx += 1 + + lower_idx = max(1, lower_idx) + upper_idx = max(1, upper_idx) + logger.info( + f"Auto edge bins (energy): lower={lower_idx}, upper={upper_idx} " + f"(k={k}, cap={cap} of {n} bins/edge)" + ) + return lower_idx, upper_idx + + +# --------------------------------------------------------------------------- +# Knee-point threshold (step_4) +# --------------------------------------------------------------------------- + +def knee_threshold( + z: np.ndarray, + sensitivity: float = 1.0, + delta: float = 0.0, + fallback_frac: float = 0.02, + grid_points: int = 512, +) -> dict: + """Kneedle threshold on the ECDF of positive robust-z values. + + ``z`` must already be robust-standardized ((x - median) / (1.4826*MAD)), + NOT asinh-compressed — knee location is not invariant under monotone + compression. Modes are positive excess by construction, so only ``z > 0`` + enters the curve (equivalent to the old clamp-below-median). + + Returns a dict with ``threshold`` (in robust-sigma units), ``knee``, + ``used_fallback``, and ``positive_fraction`` (fraction of z above the + threshold) for logging into thresholds.csv. + """ + from kneed import KneeLocator + + zp = np.sort(z[z > 0].ravel()) + if zp.size < 100: + t = float("inf") + return { + "threshold": t, + "knee": None, + "used_fallback": True, + "positive_fraction": 0.0, + } + + x = np.linspace(0.0, float(np.quantile(zp, 0.999)), grid_points) + y = np.searchsorted(zp, x, side="right") / zp.size # ECDF (no data clipped) + + knee = None + try: + kl = KneeLocator( + x, y, curve="concave", direction="increasing", S=sensitivity + ) + knee = float(kl.knee) if kl.knee is not None else None + except Exception: # noqa: BLE001 — kneed can fail on degenerate curves + knee = None + + used_fallback = knee is None + base = knee if knee is not None else float(np.quantile(zp, 1 - fallback_frac)) + if used_fallback: + logger.warning( + f"No knee found; falling back to the {1 - fallback_frac:.3f} quantile" + ) + t = base + delta + positive_fraction = float(np.mean(z > t)) + return { + "threshold": t, + "knee": knee, + "used_fallback": used_fallback, + "positive_fraction": positive_fraction, + } + + +# --------------------------------------------------------------------------- +# Scale-covariant autos +# --------------------------------------------------------------------------- + +def compute_lam( + freq_bins: int, base_lam: float = 1.0e6, base_freq_bins: int = _REF_FREQ_BINS +) -> float: + """Baseline stiffness with constant *physical* smoothing across scales. + + The Whittaker-type smoother's halfwidth scales ~lam^(1/4) in bin units, so + lam must scale with (freq_bins/base)^4 to smooth the same physical + frequency span at every nfft. + """ + lam = base_lam * (freq_bins / base_freq_bins) ** 4 + logger.info(f"Auto lam: {lam:.3e} (freq_bins={freq_bins})") + return float(lam) + + +def compute_num_layers(freq_bins: int, time_bins: int, cap: int = 5) -> int: + """UNet depth that keeps the bottleneck at least ~4 px in each dim.""" + depth = min(cap, int(math.floor(math.log2(min(freq_bins, time_bins) / 4)))) + depth = max(1, depth) + logger.info(f"Auto num_layers: {depth} (F={freq_bins}, T={time_bins})") + return depth + + +def pad_to_multiple(n: int, multiple: int) -> int: + return int(math.ceil(n / multiple) * multiple) + + +def compute_batch_size( + base_batch_size: int, freq_bins: int, time_bins: int, num_layers: int +) -> int: + """Scale the reference batch size by the per-sample activation footprint.""" + mult = 2**num_layers + f_pad = pad_to_multiple(freq_bins, mult) + t_pad = pad_to_multiple(time_bins, mult) + ref = _REF_FREQ_BINS * _REF_TIME_BINS + batch = max(1, int(base_batch_size * ref / (f_pad * t_pad))) + logger.info(f"Auto batch_size: {batch} (pad {f_pad}x{t_pad}, base {base_batch_size})") + return batch + + +# --------------------------------------------------------------------------- +# Row removal / min size (step_4) +# --------------------------------------------------------------------------- + +def compute_row_removal( + freq_bins: int, + bottom_fraction: float = 0.01, + top_fraction: float = 0.004, +) -> tuple[int, int]: + """Rows to strip from the mask edges, as a fraction of frequency bins. + + For nfft=1024 (freq_bins=513): bottom=5, top=2 — the original defaults. + """ + remove_bottom = max(1, int(freq_bins * bottom_fraction)) + remove_top = max(1, int(freq_bins * top_fraction)) + logger.info(f"Auto row removal: bottom={remove_bottom}, top={remove_top}") + return remove_bottom, remove_top + + +def compute_min_size(height: int, width: int, fraction: float = 0.0002) -> int: + """Minimum connected-component size as a fraction of image area.""" + min_size = max(1, int(height * width * fraction)) + logger.info(f"Auto min_size: {min_size} (image {height}x{width})") + return min_size + + +# --------------------------------------------------------------------------- +# Per-modality robust log-magnitude stats (step_1 window filter, step_5) +# --------------------------------------------------------------------------- + +def compute_logmag_robust_stats( + h5_path: Path, max_samples: int = 100 +) -> tuple[float, float]: + """Per-modality median/scale of log1p(|Z|) from sampled spectrograms. + + Each diagnostic has its own intensity scale and 1/f structure, so these + stats must come from that modality's own data — never a global constant. + """ + values = [] + for arr in _sample_arrays(h5_path, max_samples): + if arr.ndim == 4: # (C, F, T, 2) real/imag + mag = np.sqrt(arr[..., 0] ** 2 + arr[..., 1] ** 2) + else: + mag = np.abs(arr) + values.append(np.log1p(mag).ravel()) + all_vals = np.concatenate(values) + med, scale = robust_stats(all_vals) + logger.info(f"Auto logmag stats: median={med:.4f}, scale={scale:.4f}") + return float(med), float(scale) diff --git a/src/tokeye/training/big_tf_unet_2/utils/hdf5_io.py b/src/tokeye/training/big_tf_unet_2/utils/hdf5_io.py new file mode 100644 index 0000000..16a38e8 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/hdf5_io.py @@ -0,0 +1,386 @@ +""" +Utility functions for reading/writing HDF5 files. + +Provides: +- Generic step-level I/O (create, write, read, iterate) for the multiscale + pipeline's intermediate HDF5 files. +- ``HDF5StepDataset`` — a PyTorch Dataset backed by an HDF5 step file. +- Original step_6b/6c prediction helpers (unchanged). +""" + +from __future__ import annotations + +import contextlib +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import h5py +import numpy as np +import torch +from torch.utils.data import Dataset + +if TYPE_CHECKING: + from collections.abc import Iterator + +# ===================================================================== +# Generic step-level HDF5 I/O +# ===================================================================== + + +def create_step_file( + path: str | Path, + metadata: dict[str, Any] | None = None, +) -> h5py.File: + """Create a new HDF5 step file with a ``/samples`` group and metadata. + + Returns the open file handle — caller must close it (or use as context + manager). + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + f = h5py.File(path, "w") + f.create_group("samples") + if metadata: + for k, v in metadata.items(): + if v is not None: + f.attrs[k] = v + return f + + +def write_sample( + h5_file: h5py.File, + idx: int, + data: np.ndarray, + group: str = "samples", +) -> None: + """Write one sample to an HDF5 file with LZF compression.""" + grp = h5_file.require_group(group) + ds_name = str(idx) + if ds_name in grp: + del grp[ds_name] + grp.create_dataset(ds_name, data=data, compression="lzf") + + +def read_sample( + h5_file: h5py.File | str | Path, + idx: int, + group: str = "samples", +) -> np.ndarray: + """Read one sample from an HDF5 file.""" + if isinstance(h5_file, (str, Path)): + with h5py.File(h5_file, "r") as f: + return np.asarray(f[group][str(idx)]) + return np.asarray(h5_file[group][str(idx)]) + + +def get_sample_count(h5_path: str | Path, group: str = "samples") -> int: + """Return the number of samples in an HDF5 step file.""" + with h5py.File(h5_path, "r") as f: + if group not in f: + return 0 + return len(f[group]) + + +def iter_samples( + h5_path: str | Path, + group: str = "samples", +) -> Iterator[tuple[int, np.ndarray]]: + """Lazily iterate ``(idx, array)`` pairs from an HDF5 step file.""" + with h5py.File(h5_path, "r") as f: + grp = f[group] + for key in sorted(grp.keys(), key=int): + yield int(key), np.asarray(grp[key]) + + +def get_step_metadata(h5_path: str | Path) -> dict[str, Any]: + """Read top-level metadata attrs from an HDF5 step file.""" + with h5py.File(h5_path, "r") as f: + return dict(f.attrs) + + +# ===================================================================== +# PyTorch Dataset backed by an HDF5 step file +# ===================================================================== + + +class HDF5StepDataset(Dataset): + """PyTorch Dataset that reads samples from an HDF5 step file. + + Opens the file lazily on first access (or in ``worker_init``) and + reads one sample per ``__getitem__`` call — no full-file load. + """ + + def __init__( + self, + h5_path: str | Path, + group: str = "samples", + transform=None, + ) -> None: + self.h5_path = str(h5_path) + self.group = group + self.transform = transform + + # Determine length and sorted keys up-front + with h5py.File(self.h5_path, "r") as f: + grp = f[self.group] + self._keys = sorted(grp.keys(), key=int) + self._h5: h5py.File | None = None + + def _open(self) -> None: + if self._h5 is None: + self._h5 = h5py.File(self.h5_path, "r", swmr=True) + + def worker_init(self) -> None: + """Call from DataLoader ``worker_init_fn`` to open per-worker handle.""" + self._open() + + def __len__(self) -> int: + return len(self._keys) + + def __getitem__(self, idx: int) -> torch.Tensor | Any: + self._open() + data = np.asarray(self._h5[self.group][self._keys[idx]]) + data = torch.from_numpy(data).float() + if self.transform is not None: + data = self.transform(data) + return data + + def __del__(self) -> None: + if self._h5 is not None: + self._h5.close() + + +def worker_init_fn(worker_id: int) -> None: + """Generic ``worker_init_fn`` for DataLoaders using HDF5StepDataset.""" + import torch.utils.data as data_utils + + dataset = data_utils.get_worker_info().dataset + if hasattr(dataset, "worker_init"): + dataset.worker_init() + + +# ===================================================================== +# Original step_6b/6c prediction helpers (unchanged API) +# ===================================================================== + + +def read_fold_prediction( + hdf5_path: Path, fold_idx: int, sample_idx: int, data_type: str = "pred" +) -> np.ndarray | None: + """ + Read a single prediction from a fold HDF5 file. + + Args: + hdf5_path: Path to HDF5 file (e.g., fold_0_predictions.h5) + fold_idx: Fold index + sample_idx: Sample index + data_type: Type of data ('pred', 'std', 'entropy') + + Returns: + Array of shape (2, H, W) where channel 0 is normal, channel 1 is baseline + Returns None if not found + """ + if not hdf5_path.exists(): + return None + + try: + with h5py.File(hdf5_path, "r") as f: + fold_group = f[f"fold_{fold_idx}"] + dataset_name = f"sample_{sample_idx}_{data_type}" + if dataset_name in fold_group: + return fold_group[dataset_name][:] + return None + except (KeyError, OSError): + return None + + +def read_ensemble_prediction( + hdf5_path: Path, sample_idx: int, data_type: str = "ensemble_mean" +) -> np.ndarray | None: + """ + Read an ensemble prediction from the ensemble HDF5 file. + + Args: + hdf5_path: Path to ensemble HDF5 file (e.g., ensemble.h5) + sample_idx: Sample index + data_type: Type of data ('ensemble_mean', 'ensemble_std', 'ensemble_entropy', + 'mc_std_mean', 'mc_entropy_mean') + + Returns: + Array of shape (2, H, W) where channel 0 is normal, channel 1 is baseline + Returns None if not found + """ + if not hdf5_path.exists(): + return None + + try: + with h5py.File(hdf5_path, "r") as f: + dataset_name = f"sample_{sample_idx}_{data_type}" + if dataset_name in f: + return f[dataset_name][:] + return None + except (KeyError, OSError): + return None + + +def get_channel(data: np.ndarray, channel: str = "normal") -> np.ndarray: + """ + Extract a specific channel from 2-channel prediction data. + + Args: + data: Array of shape (2, H, W) + channel: 'normal' (channel 0) or 'baseline' (channel 1) + + Returns: + Array of shape (H, W) for the specified channel + """ + channel_idx = 0 if channel == "normal" else 1 + return data[channel_idx] + + +def list_samples(hdf5_path: Path, fold_idx: int | None = None) -> list[int]: + """ + List all sample indices in an HDF5 file. + + Args: + hdf5_path: Path to HDF5 file + fold_idx: Fold index (for step_6b output), or None for step_6c output + + Returns: + Sorted list of sample indices + """ + if not hdf5_path.exists(): + return [] + + try: + with h5py.File(hdf5_path, "r") as f: + if fold_idx is not None: + # Step 6b output + fold_group = f[f"fold_{fold_idx}"] + keys = fold_group.keys() + else: + # Step 6c output + keys = f.keys() + + # Extract sample indices from dataset names + sample_indices = set() + for key in keys: + if "sample_" in key: + idx_str = key.split("_")[1] + with contextlib.suppress(ValueError): + sample_indices.add(int(idx_str)) + + return sorted(sample_indices) + except (KeyError, OSError): + return [] + + +def get_metadata(hdf5_path: Path, fold_idx: int | None = None) -> dict: + """ + Get metadata from an HDF5 file. + + Args: + hdf5_path: Path to HDF5 file + fold_idx: Fold index (for step_6b output), or None for step_6c output + + Returns: + Dictionary of metadata attributes + """ + if not hdf5_path.exists(): + return {} + + try: + with h5py.File(hdf5_path, "r") as f: + if fold_idx is not None: + # Step 6b output + fold_group = f[f"fold_{fold_idx}"] + return dict(fold_group.attrs) + # Step 6c output + return dict(f.attrs) + except (KeyError, OSError): + return {} + + +def read_all_samples( + hdf5_path: Path, fold_idx: int | None = None, data_type: str = "pred" +) -> tuple[list[int], np.ndarray]: + """ + Read all samples from an HDF5 file. + + Args: + hdf5_path: Path to HDF5 file + fold_idx: Fold index (for step_6b output), or None for step_6c output + data_type: Type of data to load + + Returns: + Tuple of (sample_indices, data_array) + data_array has shape (N, 2, H, W) + """ + sample_indices = list_samples(hdf5_path, fold_idx) + + if not sample_indices: + return [], np.array([]) + + # Load all samples + samples = [] + valid_indices = [] + + for idx in sample_indices: + if fold_idx is not None: + # Step 6b output + data = read_fold_prediction(hdf5_path, fold_idx, idx, data_type) + else: + # Step 6c output + data = read_ensemble_prediction(hdf5_path, idx, data_type) + + if data is not None: + samples.append(data) + valid_indices.append(idx) + + if samples: + return valid_indices, np.stack(samples, axis=0) + return [], np.array([]) + + +# Example usage functions +def example_read_fold_predictions(): + """Example: Read predictions from step_6b output.""" + from pathlib import Path + + hdf5_path = Path("data/cache/step_6b_segmenter/fold_0_predictions.h5") + fold_idx = 0 + sample_idx = 0 + + # Read prediction for a specific sample + pred = read_fold_prediction(hdf5_path, fold_idx, sample_idx, "pred") + + if pred is not None: + print(f"Prediction shape: {pred.shape}") # (2, H, W) + + # Extract individual channels + normal = get_channel(pred, "normal") # (H, W) + baseline = get_channel(pred, "baseline") # (H, W) + + print(f"Normal channel shape: {normal.shape}") + print(f"Baseline channel shape: {baseline.shape}") + + +def example_read_ensemble(): + """Example: Read ensemble predictions from step_6c output.""" + from pathlib import Path + + hdf5_path = Path("data/cache/step_6c_ensemble/ensemble.h5") + sample_idx = 0 + + # Read ensemble mean + ensemble_mean = read_ensemble_prediction(hdf5_path, sample_idx, "ensemble_mean") + + if ensemble_mean is not None: + print(f"Ensemble mean shape: {ensemble_mean.shape}") # (2, H, W) + + # Extract individual channels + normal_mean = get_channel(ensemble_mean, "normal") + baseline_mean = get_channel(ensemble_mean, "baseline") + + print(f"Normal channel mean shape: {normal_mean.shape}") + print(f"Baseline channel mean shape: {baseline_mean.shape}") diff --git a/src/tokeye/training/big_tf_unet_2/utils/losses.py b/src/tokeye/training/big_tf_unet_2/utils/losses.py new file mode 100644 index 0000000..fc618a0 --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/losses.py @@ -0,0 +1,511 @@ +""" +Noise-robust loss functions for semantic segmentation with noisy labels. +Includes label smoothing, symmetric losses, and Dice-based losses. +Also includes pixel-wise supervised contrastive learning for segmentation. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class LabelSmoothingBCELoss(nn.Module): + """ + Binary Cross-Entropy with Label Smoothing. + Smooths hard 0/1 labels to reduce overconfidence on noisy labels. + """ + + def __init__(self, smoothing: float = 0.1): + """ + Args: + smoothing: Label smoothing factor (0 = no smoothing, 1 = maximum smoothing) + Labels are transformed: 0 -> smoothing, 1 -> 1-smoothing + """ + super().__init__() + self.smoothing = smoothing + assert 0 <= smoothing < 0.5, "Smoothing should be in [0, 0.5)" + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Args: + logits: Model predictions (before sigmoid), shape (B, C, H, W) or (B, H, W) + targets: Ground truth labels (0/1), same shape as logits + + Returns: + Scalar loss value + """ + # Apply label smoothing + targets_smooth = targets * (1 - self.smoothing) + self.smoothing + + # Compute BCE with logits + return F.binary_cross_entropy_with_logits(logits, targets_smooth) + + +class SymmetricCrossEntropyLoss(nn.Module): + """ + Symmetric Cross-Entropy Loss for robust learning with noisy labels. + Combines standard CE with reverse CE to make loss symmetric and robust. + + Reference: "Symmetric Cross Entropy for Robust Learning with Noisy Labels" + Wang et al., ICCV 2019 + """ + + def __init__(self, alpha: float = 0.1, beta: float = 1.0): + """ + Args: + alpha: Weight for reverse cross-entropy term + beta: Weight for standard cross-entropy term + """ + super().__init__() + self.alpha = alpha + self.beta = beta + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + + Returns: + Scalar loss value + """ + # Get predictions + probs = torch.sigmoid(logits) + + # Standard cross-entropy: -y*log(p) - (1-y)*log(1-p) + ce = F.binary_cross_entropy_with_logits(logits, targets) + + # Reverse cross-entropy: -p*log(y) - (1-p)*log(1-y) + # Add small epsilon to targets to avoid log(0) + targets_eps = torch.clamp(targets, min=1e-7, max=1 - 1e-7) + rce = -probs * torch.log(targets_eps) - (1 - probs) * torch.log(1 - targets_eps) + rce = rce.mean() + + # Combine + return self.alpha * rce + self.beta * ce + + +class DiceLoss(nn.Module): + """ + Dice Loss for semantic segmentation. + More robust to class imbalance and boundary errors than BCE. + """ + + def __init__(self, smooth: float = 1.0): + """ + Args: + smooth: Smoothing factor to avoid division by zero + """ + super().__init__() + self.smooth = smooth + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + + Returns: + Scalar loss value (1 - Dice coefficient) + """ + # Get predictions + probs = torch.sigmoid(logits) + + # Flatten + probs_flat = probs.reshape(-1) + targets_flat = targets.reshape(-1) + + # Dice coefficient + intersection = (probs_flat * targets_flat).sum() + union = probs_flat.sum() + targets_flat.sum() + + dice = (2.0 * intersection + self.smooth) / (union + self.smooth) + + # Return loss (1 - Dice) + return 1 - dice + + +class CombinedLoss(nn.Module): + """ + Combines multiple loss functions with configurable weights. + """ + + def __init__(self, losses: list, weights: list): + """ + Args: + losses: List of loss modules + weights: List of weights for each loss (should sum to 1.0) + """ + super().__init__() + assert len(losses) == len(weights), "Number of losses and weights must match" + self.losses = nn.ModuleList(losses) + self.weights = weights + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Compute weighted combination of all losses. + """ + total_loss = 0.0 + for loss_fn, weight in zip(self.losses, self.weights, strict=False): + total_loss += weight * loss_fn(logits, targets) + return total_loss + + +class FocalLoss(nn.Module): + """ + Focal Loss for addressing class imbalance. + Focuses on hard examples by down-weighting easy ones. + + Reference: "Focal Loss for Dense Object Detection" Lin et al., ICCV 2017 + """ + + def __init__(self, alpha: float = 0.25, gamma: float = 2.0): + """ + Args: + alpha: Weighting factor for class balance + gamma: Focusing parameter (higher = more focus on hard examples) + """ + super().__init__() + self.alpha = alpha + self.gamma = gamma + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + + Returns: + Scalar loss value + """ + # Get predictions + probs = torch.sigmoid(logits) + + # BCE loss + bce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none") + + # Focal weight + pt = torch.where(targets == 1, probs, 1 - probs) + focal_weight = (1 - pt) ** self.gamma + + # Alpha weighting + alpha_weight = torch.where(targets == 1, self.alpha, 1 - self.alpha) + + # Combine + loss = alpha_weight * focal_weight * bce + return loss.mean() + + +class IoULoss(nn.Module): + """ + Intersection over Union (IoU) Loss. + Directly optimizes IoU metric. + """ + + def __init__(self, smooth: float = 1.0): + """ + Args: + smooth: Smoothing factor to avoid division by zero + """ + super().__init__() + self.smooth = smooth + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + + Returns: + Scalar loss value (1 - IoU) + """ + # Get predictions + probs = torch.sigmoid(logits) + + # Flatten + probs_flat = probs.reshape(-1) + targets_flat = targets.reshape(-1) + + # IoU + intersection = (probs_flat * targets_flat).sum() + union = probs_flat.sum() + targets_flat.sum() - intersection + + iou = (intersection + self.smooth) / (union + self.smooth) + + # Return loss (1 - IoU) + return 1 - iou + + +class PixelContrastiveLoss(nn.Module): + """ + Pixel-wise Supervised Contrastive Loss for semantic segmentation. + + Learns discriminative pixel embeddings by pulling together pixels of the same class + and pushing apart pixels of different classes. Supports multi-class segmentation + with background class. + + Reference: "Supervised Contrastive Learning" (Khosla et al., NeurIPS 2020) + Adapted for dense pixel-wise prediction with efficient sampling. + """ + + def __init__( + self, + temperature: float = 0.07, + num_samples: int = 512, + dim: int = 128, + ): + """ + Args: + temperature: Temperature parameter for scaling similarities + num_samples_per_class: Number of pixels to sample per class per batch + embedding_dim: Expected embedding dimension (for validation) + """ + super().__init__() + self.temperature = temperature + self.num_samples = num_samples + self.dim = dim + + def forward(self, embeddings: torch.Tensor, masks: torch.Tensor) -> torch.Tensor: + """ + Compute pixel-wise contrastive loss. + + Args: + embeddings: Pixel embeddings of shape (B, D, H_emb, W_emb) where D is embedding dimension + masks: Ground truth masks of shape (B, C, H_mask, W_mask) where C=2 (normal, baseline channels) + Background is defined as pixels where both channels are 0 + Will be downsampled to match embedding spatial resolution + + Returns: + Scalar contrastive loss + """ + B, D, H, W = embeddings.shape + _, C, H_mask, W_mask = masks.shape + + # Downsample masks to match embedding spatial resolution + if H_mask != H or W_mask != W: + masks = F.interpolate(masks, size=(H, W), mode="nearest") + + # Flatten spatial dimensions: (B, D, H*W) -> (B*H*W, D) + embeddings_flat = embeddings.permute(0, 2, 3, 1).reshape(-1, D) + + # Create class labels from masks: (B, C, H, W) -> (B*H*W,) + # Class 0: background (both channels are 0) + # Class 1: normal (channel 0 is 1) + # Class 2: baseline (channel 1 is 1) + masks_flat = masks.permute(0, 2, 3, 1).reshape(-1, C) # (B*H*W, C) + + # Assign class labels based on mask channels + # Priority: baseline > normal > background (in case of overlap) + class_labels = torch.zeros( + B * H * W, dtype=torch.long, device=embeddings.device + ) + # First assign normal (class 1) + class_labels = torch.where( + masks_flat[:, 0] > 0.5, torch.ones_like(class_labels), class_labels + ) + # Then assign baseline (class 2) - overrides normal if overlap + class_labels = torch.where( + masks_flat[:, 1] > 0.5, torch.full_like(class_labels, 2), class_labels + ) + + # Sample pixels from each class to make computation tractable + sampled_indices = [] + sampled_labels = [] + + for class_id in range(3): # 0: background, 1: normal, 2: baseline + class_mask = class_labels == class_id + class_indices = torch.where(class_mask)[0] + + if len(class_indices) == 0: + continue + + # Sample pixels from this class + num_samples = min(self.num_samples, len(class_indices)) + if num_samples > 0: + sampled_idx = class_indices[ + torch.randperm(len(class_indices), device=embeddings.device)[ + :num_samples + ] + ] + sampled_indices.append(sampled_idx) + sampled_labels.append( + torch.full( + (num_samples,), + class_id, + dtype=torch.long, + device=embeddings.device, + ) + ) + + if len(sampled_indices) == 0: + return torch.tensor(0.0, device=embeddings.device, requires_grad=True) + + sampled_indices = torch.cat(sampled_indices) + sampled_labels = torch.cat(sampled_labels) + sampled_embeddings = embeddings_flat[sampled_indices] + + sampled_embeddings = F.normalize(sampled_embeddings, p=2, dim=1) + + similarity_matrix = ( + torch.matmul(sampled_embeddings, sampled_embeddings.t()) / self.temperature + ) + + label_mask = ( + sampled_labels.unsqueeze(0) == sampled_labels.unsqueeze(1) + ).float() + + logits_mask = torch.ones_like(label_mask) - torch.eye( + len(sampled_labels), device=embeddings.device + ) + label_mask = label_mask * logits_mask + + similarity_matrix = ( + similarity_matrix - similarity_matrix.max(dim=1, keepdim=True)[0].detach() + ) + + exp_logits = torch.exp(similarity_matrix) * logits_mask + log_prob = similarity_matrix - torch.log( + exp_logits.sum(dim=1, keepdim=True) + 1e-7 + ) + + num_positives_per_row = label_mask.sum(dim=1) + valid_rows = num_positives_per_row > 0 + + if valid_rows.sum() == 0: + return torch.tensor(0.0, device=embeddings.device, requires_grad=True) + + mean_log_prob_pos = (label_mask * log_prob).sum(dim=1) / ( + num_positives_per_row + 1e-7 + ) + mean_log_prob_pos = mean_log_prob_pos[valid_rows] + + return -mean_log_prob_pos.mean() + + + +def get_loss_function(settings: dict) -> nn.Module: + """ + Create loss function from settings. + + Args: + settings: Dictionary with loss configuration + + Returns: + Loss module + """ + loss_type = settings.get("loss_type", "bce") + + if loss_type == "bce": + return nn.BCEWithLogitsLoss() + + if loss_type == "label_smooth_bce": + smoothing = settings.get("label_smoothing", 0.1) + return LabelSmoothingBCELoss(smoothing=smoothing) + + if loss_type == "symmetric_bce": + alpha = settings.get("symmetric_alpha", 0.1) + beta = settings.get("symmetric_beta", 1.0) + return SymmetricCrossEntropyLoss(alpha=alpha, beta=beta) + + if loss_type == "dice": + return DiceLoss() + + if loss_type == "dice_bce": + dice_weight = settings.get("dice_weight", 0.5) + bce_weight = settings.get("bce_weight", 0.5) + return CombinedLoss( + losses=[DiceLoss(), nn.BCEWithLogitsLoss()], + weights=[dice_weight, bce_weight], + ) + + if loss_type == "symmetric_bce_dice": + # Default: symmetric BCE + Dice for maximum robustness + symmetric_weight = settings.get("symmetric_weight", 0.5) + dice_weight = settings.get("dice_weight", 0.5) + alpha = settings.get("symmetric_alpha", 0.1) + beta = settings.get("symmetric_beta", 1.0) + return CombinedLoss( + losses=[SymmetricCrossEntropyLoss(alpha=alpha, beta=beta), DiceLoss()], + weights=[symmetric_weight, dice_weight], + ) + + if loss_type == "focal": + alpha = settings.get("focal_alpha", 0.25) + gamma = settings.get("focal_gamma", 2.0) + return FocalLoss(alpha=alpha, gamma=gamma) + + if loss_type == "focal_dice": + focal_weight = settings.get("focal_weight", 0.5) + dice_weight = settings.get("dice_weight", 0.5) + alpha = settings.get("focal_alpha", 0.25) + gamma = settings.get("focal_gamma", 2.0) + return CombinedLoss( + losses=[FocalLoss(alpha=alpha, gamma=gamma), DiceLoss()], + weights=[focal_weight, dice_weight], + ) + + if loss_type == "iou": + return IoULoss() + + if loss_type == "mse": + return nn.MSELoss() + + raise ValueError(f"Unknown loss type: {loss_type}") + + +def dice_coefficient( + logits: torch.Tensor, targets: torch.Tensor, threshold: float = 0.5 +) -> torch.Tensor: + """ + Compute Dice coefficient metric. + + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + threshold: Threshold for binarizing predictions + + Returns: + Dice coefficient (scalar) + """ + probs = torch.sigmoid(logits) + preds = (probs > threshold).float() + + preds_flat = preds.reshape(-1) + targets_flat = targets.reshape(-1) + + intersection = (preds_flat * targets_flat).sum() + union = preds_flat.sum() + targets_flat.sum() + + if union == 0: + return torch.tensor(1.0, device=logits.device) + + return (2.0 * intersection) / union + + +def iou_score( + logits: torch.Tensor, targets: torch.Tensor, threshold: float = 0.5 +) -> torch.Tensor: + """ + Compute IoU (Jaccard) score. + + Args: + logits: Model predictions (before sigmoid) + targets: Ground truth labels (0/1) + threshold: Threshold for binarizing predictions + + Returns: + IoU score (scalar) + """ + probs = torch.sigmoid(logits) + preds = (probs > threshold).float() + + preds_flat = preds.reshape(-1) + targets_flat = targets.reshape(-1) + + intersection = (preds_flat * targets_flat).sum() + union = preds_flat.sum() + targets_flat.sum() - intersection + + if union == 0: + return torch.tensor(1.0, device=logits.device) + + return intersection / union diff --git a/src/tokeye/training/big_tf_unet_2/utils/signal_filters.py b/src/tokeye/training/big_tf_unet_2/utils/signal_filters.py new file mode 100644 index 0000000..ff8cc8f --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/signal_filters.py @@ -0,0 +1,13 @@ +"""Signal-processing filters shared by the intake steps. + +``Preemphasis`` used to be re-exported from ``torchaudio.transforms`` by each +step module that needed it (``step_0b_filter_faithdata.py`` here and in the +ablation pipeline, ``step_0f_foundation.py``). It now lives in one place so +steps depend on this module instead of on each other. +""" + +from __future__ import annotations + +from torchaudio.transforms import Preemphasis + +__all__ = ["Preemphasis"] diff --git a/src/tokeye/training/big_tf_unet_2/utils/window_filter.py b/src/tokeye/training/big_tf_unet_2/utils/window_filter.py new file mode 100644 index 0000000..4443c5b --- /dev/null +++ b/src/tokeye/training/big_tf_unet_2/utils/window_filter.py @@ -0,0 +1,137 @@ +"""Model-based window activity filter. + +Runs the full-recipe TokEye surrogate on each window's per-channel +magnitude spectrogram and keeps the most-active windows per shot. Applied +identically to every run, so it cannot bias comparisons -- it only selects +*which* windows to train on, not the labels. + +Ported from ``big_tf_unet_ablation/window_filter.py``. The surrogate is the +shared U-Net building block (``tokeye.models.modules.unet``). +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + + +def _remap_legacy_state_dict(sd: dict) -> dict: + """Map legacy ``double_conv``/``maxpool_conv`` checkpoint keys to the U-Net.""" + idx_map = {"0": "0", "1": "1", "4": "3", "5": "4"} + out = {} + for k, v in sd.items(): + nk = k.replace(".double_conv.", ".conv.").replace( + ".maxpool_conv.1.", ".down.1." + ) + parts = nk.split(".") + for i, p in enumerate(parts): + if p == "conv" and i + 1 < len(parts) and parts[i + 1] in idx_map: + parts[i + 1] = idx_map[parts[i + 1]] + break + out[".".join(parts)] = v + return out + + +def load_filter_model( + weights: str | Path, + fallback: str | Path | None = None, + device: str = "cpu", +) -> torch.nn.Module: + """Load the surrogate U-Net used for window activity scoring.""" + from tokeye.models.modules.unet import UNet + + wp = Path(weights) + if not wp.exists() and fallback is not None: + wp = Path(fallback) + model = UNet( + in_channels=1, + out_channels=2, + num_layers=5, + first_layer_size=32, + dropout_rate=0.0, + ) + sd = _remap_legacy_state_dict(torch.load(wp, weights_only=True, map_location="cpu")) + model.load_state_dict(sd, strict=False) + model.to(device).eval() + logger.info(f"filter model loaded from {wp}") + return model + + +def _pad_to_multiple(x: torch.Tensor, m: int = 32) -> torch.Tensor: + """Pad (1,1,H,W) so H,W are multiples of m (reflect padding).""" + h, w = x.shape[-2:] + ph, pw = (-h) % m, (-w) % m + if ph or pw: + x = torch.nn.functional.pad(x, (0, pw, 0, ph), mode="reflect") + return x + + +def activity_score_from_sigmoid(sig: np.ndarray, threshold: float) -> float: + """sig: (2, H, W) sigmoid outputs -> max over channels of active-pixel fraction.""" + active = (sig > threshold).mean(axis=(1, 2)) # per channel fraction + return float(active.max()) + + +@torch.no_grad() +def score_window( + model: torch.nn.Module, + complex_window: np.ndarray, + mean: float, + std: float, + threshold: float, + device: str, +) -> float: + """complex_window: (C, F, T, 2). Returns max activity over channels.""" + n_channels = complex_window.shape[0] + best = 0.0 + for c in range(n_channels): + mag = np.sqrt(complex_window[c, ..., 0] ** 2 + complex_window[c, ..., 1] ** 2) + mag = np.log1p(mag).astype(np.float32) + mag = (mag - mean) / std + x = torch.from_numpy(mag).float().unsqueeze(0).unsqueeze(0).to(device) + x = _pad_to_multiple(x, 32) + out = model(x) # (1, 2, H, W) + sig = torch.sigmoid(out[0]).cpu().numpy() # (2, H, W) + best = max(best, activity_score_from_sigmoid(sig, threshold)) + return best + + +def compute_logmag_stats( + h5_path: str | Path, max_samples: int = 60 +) -> tuple[float, float]: + """Per-modality log-magnitude mean/std from a sample of spectrogram windows. + + Each diagnostic has its own intensity scale and 1/f structure, so the filter + model must see inputs normalized with that modality's own statistics (the + surrogate was trained with per-modality normalization). Computed + automatically here rather than hard-coded per diagnostic. + """ + from .hdf5_io import get_sample_count, read_sample + + n = get_sample_count(h5_path) + if n == 0: + return 0.0, 1.0 + step = max(1, n // max_samples) + vals = [] + for idx in range(0, n, step): + data = read_sample(h5_path, idx) # (C, F, T, 2) + mag = np.log1p(np.sqrt(data[..., 0] ** 2 + data[..., 1] ** 2)) + vals.append(mag.reshape(-1)) + allv = np.concatenate(vals) + mean = float(allv.mean()) + std = float(allv.std()) + return mean, (std if std > 1e-6 else 1.0) + + +def select_window_indices( + scores: dict[int, float], max_windows: int, min_activity: float +) -> list[int]: + """Keep indices above floor, top-`max_windows` by score, sorted by score desc.""" + above = [(i, s) for i, s in scores.items() if s >= min_activity] + above.sort(key=lambda t: t[1], reverse=True) + return [i for i, _ in above[:max_windows]] diff --git a/tests/test_big_tf_unet_2.py b/tests/test_big_tf_unet_2.py new file mode 100644 index 0000000..2790173 --- /dev/null +++ b/tests/test_big_tf_unet_2.py @@ -0,0 +1,241 @@ +"""Unit tests for the big_tf_unet_2 workflow layer + math core. + +Collection-safe under `uv sync --dev` (no train deps): modules that pull +h5py/kneed are guarded with importorskip. +""" + +from __future__ import annotations + +import pytest + +from tokeye.training.big_tf_unet_2.paths import ( + STEP_ORDER, + RunPaths, + get_step, + run_id_for, + steps_after, +) +from tokeye.training.big_tf_unet_2.run_config import ( + ConfigError, + check_scale_lock, + load_run_config, +) +from tokeye.training.big_tf_unet_2.task_matrix import RunTaskMatrix, params_hash + +MODS = ["ece", "mhr", "bes"] + + +def _write_run_yaml(path, body="run: {nfft: 512, hop: 128}\n"): + path.write_text(body) + return path + + +# --------------------------------------------------------------------------- +# paths +# --------------------------------------------------------------------------- + + +def test_run_id_and_registry(): + assert run_id_for(512, 128) == "nfft512_hop128" + assert run_id_for(512, 128, "v2") == "nfft512_hop128_v2" + assert STEP_ORDER[0] == "step_0" and STEP_ORDER[-1] == "step_8" + assert get_step("step_3").exec_mode == "sbatch_gpu" + assert not get_step("step_5").per_modality + assert [s.name for s in steps_after("step_6")] == ["step_7", "step_8"] + with pytest.raises(KeyError): + get_step("step_99") + + +def test_artifacts_are_registered_per_step(tmp_path): + paths = RunPaths("t", root=tmp_path) + per_mod = paths.artifacts("step_2", ["ece"]) + assert paths.step_h5("step_2", "ece") in per_mod + assert paths.baseline_h5("ece") in per_mod + combined = paths.artifacts("step_5", MODS) + assert combined == [paths.step_h5("step_5")] + assert paths.artifacts("step_7", MODS) == [paths.model_dir] + + +# --------------------------------------------------------------------------- +# run_config +# --------------------------------------------------------------------------- + + +def test_valid_config_loads(tmp_path): + cfg = load_run_config(_write_run_yaml(tmp_path / "run.yaml")) + assert cfg.run.nfft == 512 + assert cfg.n_freq == 257 + assert cfg.modality_names == MODS + + +@pytest.mark.parametrize( + "body", + [ + "run: {nfft: 512, hop: 1024}", # hop > nfft + "run: {nfft: 500, hop: 100}", # off-grid + "run: {nfft: 512, hop: 128}\nrefine: {model_trust: 2.0}", # out of range + "run: {nfft: 512, hop: 128}\nlabels: {knee_sensitivty: 1.0}", # typo + ], +) +def test_bad_configs_raise_config_error(tmp_path, body): + with pytest.raises(ConfigError): + load_run_config(_write_run_yaml(tmp_path / "run.yaml", body)) + + +def test_custom_scale_needs_opt_in(tmp_path): + body = "run: {nfft: 500, hop: 100, allow_custom_scale: true}" + cfg = load_run_config(_write_run_yaml(tmp_path / "run.yaml", body)) + assert cfg.run.nfft == 500 + + +def test_scale_lock(tmp_path): + cfg = load_run_config(_write_run_yaml(tmp_path / "run.yaml")) + meta = tmp_path / "run_meta.json" + meta.write_text('{"nfft": 1024, "hop": 256}') + with pytest.raises(ConfigError): + check_scale_lock(cfg, meta) + meta.write_text('{"nfft": 512, "hop": 128}') + check_scale_lock(cfg, meta) # no raise + + +# --------------------------------------------------------------------------- +# task matrix +# --------------------------------------------------------------------------- + + +def test_staleness_propagates_per_modality(tmp_path): + m = RunTaskMatrix(tmp_path / "tm.json") + h = params_hash({"x": 1}) + m.mark_complete("step_0", "ece", h, MODS) + m.mark_complete("step_1", "ece", h, MODS) + m.mark_complete("step_1", "mhr", h, MODS) + m.mark_complete("step_5", None, h, MODS) + # rerun of ece step_0 stales ece's chain + combined steps, NOT mhr's + m.mark_complete("step_0", "ece", params_hash({"x": 2}), MODS) + assert m.status("step_1", "ece") == "stale" + assert m.status("step_1", "mhr") == "complete" + assert m.status("step_5") == "stale" + + +def test_accept_requires_complete(tmp_path): + m = RunTaskMatrix(tmp_path / "tm.json") + h = params_hash({}) + m.mark_complete("step_0", "ece", h, MODS) + with pytest.raises(ValueError, match="not complete"): + m.accept("step_0", MODS) + for mod in MODS[1:]: + m.mark_complete("step_0", mod, h, MODS) + m.accept("step_0", MODS) + assert m.is_accepted("step_0", MODS) + # rerunning revokes acceptance + m.mark_complete("step_0", "ece", h, MODS) + assert not m.is_accepted("step_0", MODS) + + +def test_clear_resets_and_stales(tmp_path): + m = RunTaskMatrix(tmp_path / "tm.json") + h = params_hash({}) + for mod in MODS: + m.mark_complete("step_0", mod, h, MODS) + m.mark_complete("step_1", mod, h, MODS) + m.mark_pending("step_0", MODS) + assert m.status("step_0", "ece") == "pending" + assert m.status("step_1", "ece") == "stale" + + +# --------------------------------------------------------------------------- +# clearing (fence) +# --------------------------------------------------------------------------- + + +def test_clear_fence_refuses_outside_roots(tmp_path): + from tokeye.training.big_tf_unet_2.clearing import _assert_fenced + + paths = RunPaths("t", root=tmp_path) + inside = paths.cache_root / "ece" / "step_0.h5" + _assert_fenced(inside, [paths.cache_root]) # no raise + with pytest.raises(RuntimeError, match="Refusing"): + _assert_fenced(tmp_path / "outside.txt", [paths.cache_root]) + + +# --------------------------------------------------------------------------- +# math core (train deps required) +# --------------------------------------------------------------------------- + + +def test_normalize_asinh_properties(): + ap = pytest.importorskip("tokeye.training.big_tf_unet_2.utils.auto_params") + np = pytest.importorskip("numpy") + x = np.random.default_rng(0).normal(0, 1, 50_000) + med, scale = ap.robust_stats(x) + n3 = ap.normalize_asinh(x, 3.0, med, scale) + bulk = np.abs(x) < 0.5 + assert np.allclose(n3[bulk], (x[bulk] - med) / scale, atol=0.02) + grid = ap.normalize_asinh(np.linspace(-100, 100, 1000), 3.0, 0.0, 1.0) + assert np.all(np.diff(grid) > 0) # strictly monotone (invertible) + + +def test_knee_threshold_synthetic(): + pytest.importorskip("kneed") + ap = pytest.importorskip("tokeye.training.big_tf_unet_2.utils.auto_params") + np = pytest.importorskip("numpy") + rng = np.random.default_rng(1) + z = rng.normal(0, 1, 200_000) + z[:2000] = rng.normal(8, 0.5, 2000) # 1% strong signal + r = ap.knee_threshold(z) + assert not r["used_fallback"] + assert 1.0 < r["threshold"] < 8.0 + r2 = ap.knee_threshold(z, delta=1.5) + assert r2["threshold"] == pytest.approx(r["threshold"] + 1.5) + degenerate = ap.knee_threshold(np.full(10, -1.0)) + assert degenerate["used_fallback"] + + +def test_edge_bins_energy_catches_plateau(tmp_path): + h5py = pytest.importorskip("h5py") + ap = pytest.importorskip("tokeye.training.big_tf_unet_2.utils.auto_params") + np = pytest.importorskip("numpy") + n_freq = 257 + profile = np.ones(n_freq) + profile[:35] = 40.0 # broad low-frequency plateau (the bes failure mode) + profile[-3:] = 30.0 + path = tmp_path / "t.h5" + with h5py.File(path, "w") as f: + grp = f.create_group("samples") + for i in range(4): + noise = np.random.default_rng(i).normal(1, 0.05, (2, n_freq, 64, 2)) + grp.create_dataset(str(i), data=(profile[None, :, None, None] * noise)) + lower, upper = ap.detect_edge_bins_energy(path, k=2.0) + assert 33 <= lower <= 38 # full plateau; gradient method finds ~1 + assert 2 <= upper <= 5 + + +def test_scale_covariant_autos(): + ap = pytest.importorskip("tokeye.training.big_tf_unet_2.utils.auto_params") + assert ap.compute_lam(513) == pytest.approx(1.0e6) + assert 0.05e6 < ap.compute_lam(257) < 0.07e6 # (257/513)^4 + assert ap.compute_num_layers(257, 516) == 5 + assert ap.compute_num_layers(65, 1032) == 4 # nfft=128 shrinks the net + assert ap.compute_batch_size(8, 257, 516, 5) >= 8 + + +# --------------------------------------------------------------------------- +# scaffold (isolated cwd) +# --------------------------------------------------------------------------- + + +def test_scaffold_creates_and_refuses_overwrite(tmp_path, monkeypatch): + from tokeye.training.big_tf_unet_2.scaffold import scaffold_run + + # repo_root() walks up for pyproject.toml; give the sandbox one so the + # scaffold lands inside tmp_path instead of the real checkout. + (tmp_path / "pyproject.toml").write_text("[project]\nname='sandbox'\n") + monkeypatch.chdir(tmp_path) + paths = scaffold_run(512, 128) + assert paths.root == tmp_path + assert paths.run_yaml.exists() + assert paths.run_meta.exists() + cfg = load_run_config(paths.run_yaml) + check_scale_lock(cfg, paths.run_meta) + with pytest.raises(FileExistsError): + scaffold_run(512, 128) diff --git a/uv.lock b/uv.lock index 0c19110..416ba68 100644 --- a/uv.lock +++ b/uv.lock @@ -1178,6 +1178,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/80/be/3578e8afd18c88cdf9cb4cffde75a96d2be38c5a903f1ed0ceec061bd09e/kiwisolver-1.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:4a48a2ce79d65d363597ef7b567ce3d14d68783d2b2263d98db3d9477805ba32", size = 70260, upload-time = "2025-08-10T21:27:36.606Z" }, ] +[[package]] +name = "kneed" +version = "0.8.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/64/4bb8f8a7a4627b585a66d5bec0c9b30ae5b39a4caea1775c8bfb3fb3f4cf/kneed-0.8.6.tar.gz", hash = "sha256:65b22727c623661701f15edf057f2e6c73e2b1ad4e68cd9ca4291675c318b5ef", size = 13161, upload-time = "2026-03-20T21:01:51.966Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/cd/23c89d53c36028bccb39f55aa5dd24c4bdaab76c4d556ad43dc8cf026918/kneed-0.8.6-py3-none-any.whl", hash = "sha256:3412e7b70bce07717386d24fab37f0f985968d1b85ea0c749a6b98caccaf65ec", size = 10797, upload-time = "2026-03-20T21:01:50.87Z" }, +] + [[package]] name = "lazy-loader" version = "0.4" @@ -2735,11 +2748,13 @@ train = [ { name = "ipykernel" }, { name = "ipywidgets" }, { name = "joblib" }, + { name = "kneed" }, { name = "lightning" }, { name = "pandas" }, { name = "pybaselines" }, { name = "scikit-image" }, { name = "scikit-learn" }, + { name = "threadpoolctl" }, { name = "tifffile" }, { name = "torchaudio" }, { name = "torchmetrics" }, @@ -2760,11 +2775,13 @@ train = [ { name = "ipykernel" }, { name = "ipywidgets" }, { name = "joblib" }, + { name = "kneed" }, { name = "lightning" }, { name = "pandas" }, { name = "pybaselines" }, { name = "scikit-image" }, { name = "scikit-learn" }, + { name = "threadpoolctl" }, { name = "tifffile" }, { name = "torchaudio" }, { name = "torchmetrics" }, @@ -2779,6 +2796,7 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'train'" }, { name = "ipywidgets", marker = "extra == 'train'" }, { name = "joblib", marker = "extra == 'train'" }, + { name = "kneed", marker = "extra == 'train'" }, { name = "lightning", marker = "extra == 'train'" }, { name = "matplotlib" }, { name = "numpy" }, @@ -2792,6 +2810,7 @@ requires-dist = [ { name = "scikit-learn", marker = "extra == 'train'" }, { name = "scipy" }, { name = "tables" }, + { name = "threadpoolctl", marker = "extra == 'train'" }, { name = "tifffile", marker = "extra == 'train'" }, { name = "torch" }, { name = "torchaudio", marker = "extra == 'train'" }, @@ -2818,11 +2837,13 @@ train = [ { name = "ipykernel" }, { name = "ipywidgets" }, { name = "joblib" }, + { name = "kneed" }, { name = "lightning" }, { name = "pandas" }, { name = "pybaselines" }, { name = "scikit-image" }, { name = "scikit-learn" }, + { name = "threadpoolctl" }, { name = "tifffile" }, { name = "torchaudio" }, { name = "torchmetrics" },