diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 00000000..38c5e873 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,152 @@ +name: benchmark + +on: + pull_request: + branches: ["main"] + workflow_dispatch: + +permissions: + contents: read + pull-requests: write + +concurrency: + group: benchmark-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + benchmark: + runs-on: ubuntu-latest + timeout-minutes: 60 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: "pip" + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e . + + # CI runs only the real CT/MRI cases. The 400^3 / 256^3 synthetic cases + # dominate wall time (~30s per case per run × 4 runs) and are the noisiest + # single knob in the workflow; they stay available locally by omitting + # --cases. --repeats 5 buys back a little statistical stability now that + # we are not paying for synth. + - name: Benchmark HEAD (speed) + run: python benchmarks/nii_poi/bench_speed.py --cases ct_3d,ct_2d,mri_3d,mri_2d --repeats 5 --json head.json + + - name: Benchmark HEAD (memory) + run: python benchmarks/nii_poi/bench_mem.py --cases ct_3d,ct_2d,mri_3d,mri_2d --repeats 5 --json head_mem.json + + - name: Benchmark baseline (base commit's TPTBox, HEAD's harness) + if: github.event_name == 'pull_request' + shell: bash + run: | + # Only the TPTBox/ tree is swapped. benchmarks/ stays at HEAD so the + # workloads, repeats, warmup and JSON schema are identical on both + # sides — otherwise the comparison measures the harness, not the code. + cp -a TPTBox /tmp/TPTBox-head + # Nuke first, then repopulate from base. A plain `git checkout base --` + # leaves behind files that are new in HEAD, and the resulting mix of + # old and new modules fails with TypeError on unknown kwargs. + rm -rf TPTBox + git checkout ${{ github.event.pull_request.base.sha }} -- TPTBox + python benchmarks/nii_poi/bench_speed.py --cases ct_3d,ct_2d,mri_3d,mri_2d --repeats 5 --json baseline.json + python benchmarks/nii_poi/bench_mem.py --cases ct_3d,ct_2d,mri_3d,mri_2d --repeats 5 --json baseline_mem.json + rm -rf TPTBox + cp -a /tmp/TPTBox-head TPTBox + + - name: Compare against on-runner baseline + id: compare + if: github.event_name == 'pull_request' + shell: bash + run: | + set +e + python benchmarks/nii_poi/compare.py baseline.json head.json \ + --fail-on-regression-pct 50 > speed.md + speed_status=$? + python benchmarks/nii_poi/compare_mem.py baseline_mem.json head_mem.json \ + --fail-on-regression-pct 50 > mem.md + mem_status=$? + + { + echo '' + echo '# Benchmark comparison' + echo '' + cat speed.md + echo '' + cat mem.md + } > comparison.md + + cat comparison.md >> "$GITHUB_STEP_SUMMARY" + + status=0 + if [ "$speed_status" -ne 0 ] || [ "$mem_status" -ne 0 ]; then + status=1 + fi + echo "compare_status=$status" >> "$GITHUB_OUTPUT" + exit 0 + + - name: Post/update PR comment + if: github.event_name == 'pull_request' && always() + continue-on-error: true + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + let body; + try { + body = fs.readFileSync('comparison.md', 'utf8'); + } catch (e) { + core.warning(`comparison.md missing (${e.message}); skipping PR comment.`); + return; + } + const marker = ''; + const { owner, repo } = context.repo; + const issue_number = context.payload.pull_request.number; + const { data: comments } = await github.rest.issues.listComments({ + owner, repo, issue_number, + }); + const existing = comments.find(c => c.body && c.body.includes(marker)); + if (existing) { + await github.rest.issues.updateComment({ + owner, repo, comment_id: existing.id, body, + }); + } else { + await github.rest.issues.createComment({ + owner, repo, issue_number, body, + }); + } + + - name: Fail on regression + if: github.event_name == 'pull_request' + shell: bash + run: exit "${{ steps.compare.outputs.compare_status }}" + + - name: Head-only summary (workflow_dispatch) + if: github.event_name != 'pull_request' + shell: bash + run: | + { + echo '# Benchmark (head only)' + echo '' + echo 'No baseline to compare against on a manual run; raw numbers are in the artifacts.' + } >> "$GITHUB_STEP_SUMMARY" + + - name: Upload artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: benchmark-results + path: | + head.json + head_mem.json + baseline.json + baseline_mem.json + comparison.md + if-no-files-found: ignore diff --git a/benchmarks/nii_poi/README.md b/benchmarks/nii_poi/README.md new file mode 100644 index 00000000..8435a6a5 --- /dev/null +++ b/benchmarks/nii_poi/README.md @@ -0,0 +1,182 @@ +# NII / POI speed + memory benchmark + +Regression detection for the two classes almost every TPTBox workflow runs +through: `NII` (`TPTBox/core/nii_wrapper.py`) and `POI` (`TPTBox/core/poi.py`, +`TPTBox/core/poi_fun/`). + +`.github/workflows/benchmark.yml` runs this on every PR to `main`: once on the +PR head, once on the PR base commit (same runner, same harness), then posts a +comparison table as a PR comment and **fails the job** if anything regressed. + +This is unrelated to `benchmarks/benchmark_nnunet_inference.py`, which is a +GPU-only inference harness. + +## Why it exists + +`TPTBox/tests/speedtests/` compares *candidate implementations against each +other*, by hand, locally. This compares *HEAD against main*, automatically, and +covers memory as well as time. + +## Workloads + +Six cases, spanning a small/large × 2D/3D grid, so a change that only hurts +large 3D volumes shows up as such instead of being averaged away: + +| case | source | shape | +|---|---|---| +| `ct_3d` | `TPTBox/tests/sample_ct` | `(73, 47, 73)` | +| `ct_2d` | centre slice of `ct_3d` | `(73, 47, 1)` | +| `mri_3d` | `TPTBox/tests/sample_mri` | `(68, 52, 67)` | +| `mri_2d` | centre slice of `mri_3d` | `(68, 52, 1)` | +| `synth_3d` | generated | `(400, 400, 400)` | +| `synth_2d` | generated | `(400, 400, 1)` | + +**2D means a singleton third axis, not a 2-element shape.** `NII` has no 2D code +path — orientation, `reorient` and `rescale` all assume three axes — so a real +2-D `Nifti1Image` would crash. A `(X, Y, 1)` volume is geometrically 2D and goes +through the normal machinery. The slice is taken with `NII.apply_crop`, which +updates the affine along with the data. + +The synthetic volume is hollow labelled cuboids on a regular grid: hollow so +`fill_holes` has work to do, gridded so they never touch and the result is +byte-identical on every machine, ~20% foreground so `use_crop=True` cannot +trivialise the morphology measurements. + +## What is measured + +~33 measurements per case, defined once in `measurements.py` so the speed and +memory runs always cover exactly the same keys. + +- **NII IO** — `load` (image and segmentation), `save` +- **NII arrays** — `get_array`, `set_dtype` +- **NII geometry** — `reorient`, `rescale` (image and segmentation), + `resample_from_to`, `compute_crop`, `apply_crop`, `pad_to` +- **NII statistics** — `unique`, `volumes`, `center_of_masses` +- **NII labels** — `extract_label`, `map_labels` +- **NII morphology / cc3d** — `dilate_msk`, `erode_msk`, `fill_holes`, + `get_connected_components`, `filter_connected_components` +- **POI** — `calc_centroids` (and the legacy `_crop=False` per-label path as a + standing A/B), `reorient`, `rescale`, `to_global`, `local_to_global_arr`, + `resample_from_to`, `map_labels`, `save`, `load` +- **POI pipeline** — `calc_poi_from_subreg_vert`, on the real CT/MRI 3D cases + only; it is meaningless on a single slice or on synthetic cuboids + +Two measurements are skipped where they would not be informative: +`poi_calc_centroids_nocrop` on the large cases (it is a per-label scipy loop and +would dominate the entire run) and `poi_calc_poi_from_subreg_vert` on anything +that is not real spine data. + +Every call is wrapped in try/except. If an op raises, or does not exist on the +older baseline commit, the key is dropped and listed on a `skipped:` line rather +than aborting the run — that is what lets one harness measure two commits. + +`metric_*` rows (voxel count, label count, foreground %) are context, not +measurements, and never gate. + +## Running it + +```bash +# defaults: 10 repeats + 1 warmup, 400^3 synthetic case (~25 min) +python benchmarks/nii_poi/bench_speed.py --json head.json +python benchmarks/nii_poi/bench_mem.py --json head_mem.json + +# what CI runs: real CT/MRI cases only, 5 repeats + 1 warmup (~1 min each) +python benchmarks/nii_poi/bench_speed.py --cases ct_3d,ct_2d,mri_3d,mri_2d --repeats 5 --json head.json + +# quick local run including the 256^3 synthetic case (~2 min each) +python benchmarks/nii_poi/bench_speed.py --quick --json head.json + +# one case, fast iteration +python benchmarks/nii_poi/bench_speed.py --cases ct_3d,ct_2d --repeats 3 + +# compare two runs +python benchmarks/nii_poi/compare.py baseline.json head.json +python benchmarks/nii_poi/compare_mem.py baseline_mem.json head_mem.json +``` + +`--quick` lowers the synthetic edge length to 256 as well as the repeat count. A +full pass over 400³ is ~2 minutes; 256³ is ~4× cheaper and still a large 3D +volume. Both sides of any comparison always use the same value, so this changes +sensitivity, never correctness. + +**CI skips the synthetic cases entirely** and only runs the real CT/MRI 3D/2D +cases: the synthetic 400³/256³ workloads dominate wall time (~30s per case per +run, four runs per PR) and are also the noisiest single knob in the workflow, +since they exercise code paths whose runtime scales with volume rather than with +what a spine workflow actually looks like. Regressions specific to very large +volumes are still catchable locally by omitting `--cases`. + +To reproduce the baseline swap locally: + +```bash +cp -a TPTBox /tmp/TPTBox-head +rm -rf TPTBox && git checkout main -- TPTBox +python benchmarks/nii_poi/bench_speed.py --quick --json baseline.json +rm -rf TPTBox && cp -a /tmp/TPTBox-head TPTBox +``` + +## Reading the output + +Both comparison tools emit one table per case: + +``` +| Measurement | baseline ms (median ±½·range) | head ms (median ±½·range) | Δ % | p | +| `nii_dilate_msk` | 12.30 ±0.45 | 15.60 ±0.32 | +26.8% | 0.002 | +``` + +`±` is half the min→p90 range, i.e. how noisy that measurement was, not a +standard deviation. 🔴 marks a gated regression, 🟢 a gated improvement, and +`(noise)` flags a row whose baseline is below the display noise floor (3 ms / +3 MiB): its Δ% is dominated by shared-runner jitter, not by any real change. + +Each per-case table shows only the **five most-changed rows**; the rest fold +into a `
` block so the PR comment stays scannable. Rank order is by +`|Δ%|`, with `(noise)` rows always sorted after real changes so a large swing +on a sub-ms measurement cannot push a real regression out of the headline. + +### The gate + +A measurement fails the build only when **all three** hold: + +1. the baseline is at least **1 ms** / **1 MiB** and the key is not `metric_*` — + sub-unit jitter on a shared runner must never block a merge; +2. the median grew by at least `--fail-on-regression-pct` (CI uses **50%**); +3. Welch's t-test over the two samples gives `p < --alpha` (default 0.05). With + fewer than two samples per side there is no t-test, so it falls back to + "head's best run is still worse than baseline's p90". + +### Memory numbers specifically + +The number is **peak RSS growth per call, in MiB**, not allocated bytes. + +Each sampled iteration runs in a forked child. CPython's heap does not shrink, +so a second in-process iteration of the same call reuses the arena freed by the +first and reports ≈0 MiB; a fresh address space per iteration removes that. The +baseline RSS is read *inside* the child after the fork, so copy-on-write pages +of the input volume are already accounted for. + +There is a constant **floor of ~1 MiB** per measurement (the sampler thread, the +gc pass, copy-on-write faults from refcount writes on inherited objects). It is +measured explicitly at startup and reported as `floor_mib` in the JSON and in +the comparison header. It cancels out between the two sides, and it sits at the +1 MiB gating threshold, so floor-level rows are excluded from the gate anyway. + +Warmup for the memory benchmark runs in the **parent**, deliberately: a forked +child inherits whatever the parent has already imported, so the job of the +warmup here is to pull in lazily-imported modules (`scipy.ndimage`, `cc3d`, …) +before the fork, rather than making every child pay for them. + +On a platform without `fork` or without `/proc`, the harness degrades to +in-process measurement and/or `getrusage` and records which, as `isolation` and +`sampler` in the JSON. + +## Files + +- `workloads.py` — case construction. Reads the sample data **by path**, not via + `TPTBox.tests.test_utils`: the workflow swaps the whole `TPTBox/` tree to + produce the baseline, so anything imported from inside it would silently + become the old version. +- `measurements.py` — the shared measurement registry. +- `bench_speed.py` / `bench_mem.py` — the two harnesses. +- `compare.py` / `compare_mem.py` — thin CLI wrappers over `_compare_core.py`. +- `_common.py` — statistics, JSON schema helpers, shared CLI flags. diff --git a/benchmarks/nii_poi/_common.py b/benchmarks/nii_poi/_common.py new file mode 100644 index 00000000..08da2ca0 --- /dev/null +++ b/benchmarks/nii_poi/_common.py @@ -0,0 +1,105 @@ +"""Statistics, JSON schema helpers and CLI plumbing shared by all four scripts.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +DEFAULT_REPEATS = 10 +DEFAULT_WARMUP = 1 +DEFAULT_LARGE_SIZE = 400 +# --quick is what CI runs. A full pass over the 400^3 case is ~2 min, which twice +# over (head + baseline) and twice again (speed + memory) is far too slow to gate +# every PR on; 256^3 is ~4x cheaper and still a large 3D volume. Both sides of a +# comparison always use the same value, so this only changes sensitivity, never +# correctness. +QUICK_REPEATS = 3 +QUICK_WARMUP = 1 +QUICK_LARGE_SIZE = 256 + + +def summarize(samples: list[float]) -> dict[str, float]: + """Reduce raw per-iteration samples to the schema both comparison tools read.""" + a = np.asarray(samples, dtype=float) + return { + "min": float(a.min()), + "median": float(np.median(a)), + "p90": float(np.percentile(a, 90)), + "mean": float(a.mean()), + "stddev": float(a.std(ddof=1)) if a.size > 1 else 0.0, + "n": int(a.size), + } + + +def commit_hash() -> str: + try: + out = subprocess.run( # noqa: S603 + ["git", "rev-parse", "--short", "HEAD"], # noqa: S607 + capture_output=True, + text=True, + check=True, + cwd=Path(__file__).resolve().parents[2], + ) + except Exception: + return "unknown" + return out.stdout.strip() or "unknown" + + +def python_version() -> str: + return ".".join(str(v) for v in sys.version_info[:3]) + + +def add_common_args(p: argparse.ArgumentParser) -> None: + p.add_argument("--quick", action="store_true", help=f"fewer repeats ({QUICK_REPEATS}) for CI") + p.add_argument("--repeats", type=int, default=None, help=f"sampled iterations per measurement (default {DEFAULT_REPEATS})") + p.add_argument("--warmup", type=int, default=None, help=f"discarded iterations per measurement (default {DEFAULT_WARMUP})") + p.add_argument( + "--large-size", + type=int, + default=None, + help=f"edge length of the synthetic cases (default {DEFAULT_LARGE_SIZE}, {QUICK_LARGE_SIZE} with --quick)", + ) + p.add_argument("--cases", type=str, default=None, help="comma-separated subset of case names") + p.add_argument("--json", type=str, default=None, help="write machine-readable results here") + + +def resolve_repeats(args: argparse.Namespace) -> tuple[int, int]: + repeats = args.repeats if args.repeats is not None else (QUICK_REPEATS if args.quick else DEFAULT_REPEATS) + warmup = args.warmup if args.warmup is not None else (QUICK_WARMUP if args.quick else DEFAULT_WARMUP) + return max(repeats, 1), max(warmup, 0) + + +def resolve_large_size(args: argparse.Namespace) -> int: + if args.large_size is not None: + return args.large_size + return QUICK_LARGE_SIZE if args.quick else DEFAULT_LARGE_SIZE + + +def write_json(path: str | None, doc: dict[str, Any]) -> None: + if path is None: + return + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(doc, indent=2), encoding="utf-8") + print(f"\nwrote {out}") + + +def load_json(path: str | Path) -> dict[str, Any]: + return json.loads(Path(path).read_text(encoding="utf-8")) + + +def print_case_table(case_name: str, shape, summaries: dict[str, dict[str, float]], unit: str) -> None: + """Human-readable per-case table, slowest/largest first.""" + print(f"\n=== {case_name} shape={tuple(shape)} ===") + if not summaries: + print(" (nothing measured)") + return + width = max(len(k) for k in summaries) + for key, s in sorted(summaries.items(), key=lambda kv: -kv[1]["median"]): + print(f" {key:<{width}} {s['median']:10.3f} {unit} (min {s['min']:.3f}, p90 {s['p90']:.3f}, n={s['n']})") diff --git a/benchmarks/nii_poi/_compare_core.py b/benchmarks/nii_poi/_compare_core.py new file mode 100644 index 00000000..a55983f3 --- /dev/null +++ b/benchmarks/nii_poi/_compare_core.py @@ -0,0 +1,244 @@ +"""Shared baseline-vs-head comparison: markdown table plus the regression gate. + +``compare.py`` and ``compare_mem.py`` are thin wrappers around this; the only +differences are the measurements key (``measurements_ms`` / ``measurements_mb``) +and the unit label. + +The gate is deliberately hard to trip. A measurement fails the build only when +*all three* hold: + +1. the baseline is at least ``MIN_GATED`` (1 ms / 1 MiB) and the key is not a + non-timing ``metric_*`` value — sub-unit jitter on a shared runner must never + block a merge; +2. the regression is practically significant: ``|delta| >= --fail-on-regression-pct``; +3. it is statistically significant: Welch's t-test over the two samples gives + ``p < --alpha``. With fewer than two samples per side there is no t-test, so + it falls back to "head's best run is still worse than baseline's p90". + +The rendered report shows only the ``TOP_N`` most-changed measurements per case +in the visible table; the rest go in a collapsible ``
`` block so the PR +comment stays scannable. Rows whose baseline is below ``NOISE_FLOOR`` get a +``(noise)`` tag: with only three repeats on a shared runner, a sub-unit +baseline can swing ±30% between runs without anything in the code having +changed, and that must be visually obvious. +""" + +from __future__ import annotations + +import argparse +import math +from dataclasses import dataclass +from typing import Any + +MIN_GATED = 1.0 +#: Baselines below this are shown with a ``(noise)`` tag; sub-unit measurements +#: on a shared runner jitter by tens of percent between repeats. +NOISE_FLOOR = 3.0 +#: How many rows to show up-front per case. The rest go into a
block. +TOP_N = 5 + + +def _stats(entry: Any) -> dict[str, float] | None: + """Normalise a measurement entry; a bare number is treated as min=median=p90.""" + if isinstance(entry, (int, float)): + v = float(entry) + return {"min": v, "median": v, "p90": v, "mean": v, "stddev": 0.0, "n": 1} + if isinstance(entry, dict) and "median" in entry: + out = {k: float(entry.get(k, entry["median"])) for k in ("min", "median", "p90", "mean")} + out["stddev"] = float(entry.get("stddev", 0.0)) + out["n"] = int(entry.get("n", 1)) + return out + return None + + +def _welch_p(b: dict[str, float], h: dict[str, float]) -> float | None: + if b["n"] < 2 or h["n"] < 2 or (b["stddev"] == 0.0 and h["stddev"] == 0.0): + return None + try: + from scipy.stats import ttest_ind_from_stats + except ImportError: + return None + res = ttest_ind_from_stats( + mean1=b["mean"], std1=b["stddev"], nobs1=b["n"], mean2=h["mean"], std2=h["stddev"], nobs2=h["n"], equal_var=False + ) + p = float(res.pvalue) + return None if math.isnan(p) else p + + +@dataclass +class Row: + key: str + base: dict[str, float] + head: dict[str, float] + + @property + def delta(self) -> float: + return self.head["median"] - self.base["median"] + + @property + def pct(self) -> float | None: + return (self.delta / self.base["median"] * 100.0) if self.base["median"] > 0 else None + + @property + def p_value(self) -> float | None: + return _welch_p(self.base, self.head) + + @property + def gateable(self) -> bool: + return self.base["median"] >= MIN_GATED and not self.key.startswith("metric_") + + @property + def is_noisy(self) -> bool: + """True for rows whose baseline sits in the sub-unit / near-floor range. + + Δ% on these is dominated by runner jitter and should be read as such, + not as a real change. + """ + return not self.key.startswith("metric_") and self.base["median"] < NOISE_FLOOR + + def regressed(self, threshold_pct: float, alpha: float) -> bool: + pct = self.pct + if not self.gateable or pct is None or pct < threshold_pct: + return False + p = self.p_value + if p is None: + return self.head["min"] > self.base["p90"] + return p < alpha + + +def _fmt(s: dict[str, float]) -> str: + spread = max(s["p90"] - s["min"], 0.0) + return f"{s['median']:.2f} ±{spread / 2:.2f}" + + +def _fmt_pct(row: Row, threshold_pct: float, alpha: float) -> str: + pct = row.pct + if pct is None: + return "n/a" + if row.regressed(threshold_pct, alpha): + mark = " 🔴" + elif row.gateable and pct <= -threshold_pct: + mark = " 🟢" + elif row.is_noisy: + mark = " (noise)" + else: + mark = "" + return f"{pct:+.1f}%{mark}" + + +def _rank_key(row: Row) -> tuple[int, float]: + """Sort order for the visible top-N: real changes first, noise last, then |Δ%|. + + Noisy rows are pushed to the bottom of the ranking so a big-percentage swing + on a sub-ms row does not push a real 15% regression on a 200 ms row into + the collapsed block. + """ + pct = row.pct if row.pct is not None else 0.0 + return (1 if row.is_noisy else 0, -abs(pct)) + + +def _render_rows(rows: list[Row], unit: str, threshold_pct: float, alpha: float) -> list[str]: + lines = [ + f"| Measurement | baseline {unit} (median ±½·range) | head {unit} (median ±½·range) | Δ % | p |", + "| --- | ---: | ---: | :--- | ---: |", + ] + for r in rows: + p = r.p_value + lines.append( + f"| `{r.key}` | {_fmt(r.base)} | {_fmt(r.head)} | {_fmt_pct(r, threshold_pct, alpha)} | {'—' if p is None else f'{p:.3f}'} |" + ) + return lines + + +def compare( + baseline: dict[str, Any], + head: dict[str, Any], + *, + measurements_key: str, + unit: str, + threshold_pct: float, + alpha: float, + title: str, +) -> tuple[str, bool]: + """Render the markdown report and report whether anything regressed.""" + base_cases = {c["name"]: c for c in baseline.get("cases", [])} + head_cases = {c["name"]: c for c in head.get("cases", [])} + + lines: list[str] = [f"## {title}", ""] + lines.append( + f"baseline `{baseline.get('commit', '?')}` vs head `{head.get('commit', '?')}` · " + f"python {head.get('python', '?')} · {head.get('repeats', '?')} repeats + {head.get('warmup', '?')} warmup" + + (f" · sampler `{head.get('sampler')}`, isolation `{head.get('isolation')}`" if head.get("sampler") else "") + + (f" · measurement floor ≈ {head['floor_mib']:.2f} MiB" if head.get("floor_mib") else "") + ) + lines.append("") + lines.append( + f"_Showing the {TOP_N} most-changed measurements per case; the rest are collapsed. " + f"Rows with a baseline below {NOISE_FLOOR:g} {unit} are tagged `(noise)` — runner jitter " + f"on those swamps any real change._" + ) + lines.append("") + + any_regression = False + for name, head_case in head_cases.items(): + shape = tuple(head_case.get("shape", [])) + lines.append(f"### `{name}` — shape {shape}") + lines.append("") + base_case = base_cases.get(name) + if base_case is None: + lines.append("_new case, no baseline to compare against._") + lines.append("") + continue + + b_meas = base_case.get(measurements_key, {}) + h_meas = head_case.get(measurements_key, {}) + rows: list[Row] = [] + new_keys: list[str] = [] + for key, entry in h_meas.items(): + h = _stats(entry) + b = _stats(b_meas.get(key)) if key in b_meas else None + if h is None: + continue + if b is None: + new_keys.append(key) + else: + rows.append(Row(key, b, h)) + removed_keys = [k for k in b_meas if k not in h_meas] + + rows.sort(key=_rank_key) + for r in rows: + any_regression = any_regression or r.regressed(threshold_pct, alpha) + + top = rows[:TOP_N] + rest = rows[TOP_N:] + lines.extend(_render_rows(top, unit, threshold_pct, alpha)) + lines.append("") + if rest: + lines.append(f"
… {len(rest)} more measurements") + lines.append("") + lines.extend(_render_rows(rest, unit, threshold_pct, alpha)) + lines.append("") + lines.append("
") + lines.append("") + if new_keys: + lines.append(f"_new (no baseline): {', '.join(f'`{k}`' for k in new_keys)}_") + lines.append("") + if removed_keys: + lines.append(f"_removed (baseline only): {', '.join(f'`{k}`' for k in removed_keys)}_") + lines.append("") + + lines.append( + f"Gate: a measurement fails when the baseline is ≥ {MIN_GATED:g} {unit}, the median grows by " + f"≥ {threshold_pct:g}%, **and** Welch's t-test gives p < {alpha:g}. `metric_*` rows are context only." + ) + lines.append("") + return "\n".join(lines), any_regression + + +def build_parser(description: str, default_pct: float) -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=description) + p.add_argument("baseline", help="baseline JSON produced by the bench script") + p.add_argument("head", help="head JSON produced by the bench script") + p.add_argument("--fail-on-regression-pct", type=float, default=default_pct, help=f"regression threshold (default {default_pct:g}%%)") + p.add_argument("--alpha", type=float, default=0.05, help="Welch's t-test significance level (default 0.05)") + return p diff --git a/benchmarks/nii_poi/bench_mem.py b/benchmarks/nii_poi/bench_mem.py new file mode 100644 index 00000000..433da097 --- /dev/null +++ b/benchmarks/nii_poi/bench_mem.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python +"""Memory benchmark for the NII and POI hot paths. + +Same workloads and same measurement registry as ``bench_speed.py``, but the +number reported is **peak RSS growth per call, in MiB**. + +Two details make the numbers mean something: + +* **Sampling.** A daemon thread polls ``VmRSS`` from ``/proc/self/status`` every + 5 ms while the call runs, and ``resource.getrusage`` is read afterwards as a + belt-and-braces high-water mark in case a short-lived spike fell between two + polls. On a platform without ``/proc`` only the latter is used. +* **Isolation.** Each sampled iteration runs in a forked child. CPython's heap + does not shrink, so a second in-process iteration of the same call typically + reuses the arena freed by the first and reports ~0 MiB. A fresh address space + per iteration removes that. The baseline RSS is taken *inside* the child after + the fork, so the copy-on-write pages of the input volume are already accounted + for and do not inflate the result. + + python benchmarks/nii_poi/bench_mem.py # 10 repeats + 1 warmup + python benchmarks/nii_poi/bench_mem.py --quick --json head_mem.json +""" + +from __future__ import annotations + +import argparse +import gc +import multiprocessing as mp +import os +import resource +import sys +import tempfile +import threading +import time +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from _common import ( # noqa: E402 + add_common_args, + commit_hash, + print_case_table, + python_version, + resolve_large_size, + resolve_repeats, + summarize, + write_json, +) +from measurements import build_measurements, silenced # noqa: E402 +from workloads import build_case, parse_case_names # noqa: E402 + +_STATUS = Path("/proc/self/status") +_HAS_PROC = _STATUS.exists() +_POLL_SECONDS = 0.005 +_MIB = 1024.0 * 1024.0 + + +def _rss_bytes() -> float: + """Current resident set size in bytes, or 0.0 where /proc is unavailable.""" + try: + for line in _STATUS.read_text(encoding="utf-8").splitlines(): + if line.startswith("VmRSS:"): + return float(line.split()[1]) * 1024.0 + except OSError: + pass + return 0.0 + + +def _peak_rusage_bytes() -> float: + """High-water RSS of this process. + + Linux does *not* reset ``ru_maxrss`` at fork — a child starts out carrying the + parent's high-water mark — but it does track the child's own growth on top of + it. Since the baseline is read inside the child, the difference is still the + child's own peak, which is all this is used for. + """ + maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # Linux reports KiB, macOS reports bytes. + return float(maxrss) * (1.0 if sys.platform == "darwin" else 1024.0) + + +class _PeakSampler(threading.Thread): + """Poll VmRSS in the background and remember the maximum seen.""" + + def __init__(self) -> None: + super().__init__(daemon=True) + # NB: not `_stop` — threading.Thread already uses that name for an + # internal method, and shadowing it breaks join(). + self._done = threading.Event() + self.peak = 0.0 + + def run(self) -> None: + while not self._done.is_set(): + self.peak = max(self.peak, _rss_bytes()) + time.sleep(_POLL_SECONDS) + self.peak = max(self.peak, _rss_bytes()) + + def stop(self) -> float: + self._done.set() + self.join(timeout=1.0) + return self.peak + + +def _measure_in_process(call) -> float: + """Peak RSS growth of one call, in MiB. Used inside the forked child, or as fallback.""" + gc.collect() + base = max(_rss_bytes(), _peak_rusage_bytes()) + sampler = _PeakSampler() if _HAS_PROC else None + if sampler is not None: + sampler.start() + try: + with silenced(): + call() + finally: + peak = sampler.stop() if sampler is not None else 0.0 + peak = max(peak, _rss_bytes(), _peak_rusage_bytes()) + return max(peak - base, 0.0) / _MIB + + +def _child(conn, call) -> None: + try: + conn.send(("ok", _measure_in_process(call))) + except Exception as exc: # noqa: BLE001 -- surfaced to the parent as a skip + conn.send(("err", f"{type(exc).__name__}: {exc}")) + finally: + conn.close() + + +def measure(call, ctx) -> float: + """One isolated sample. Falls back to in-process measurement without fork().""" + if ctx is None: + return _measure_in_process(call) + parent, child = ctx.Pipe(duplex=False) + proc = ctx.Process(target=_child, args=(child, call)) + proc.start() + child.close() + try: + status, payload = parent.recv() + except EOFError as exc: + proc.join() + raise RuntimeError(f"measurement child died (exit {proc.exitcode})") from exc + finally: + parent.close() + proc.join() + if status != "ok": + raise RuntimeError(payload) + return float(payload) + + +def _fork_context(): + if os.name != "posix" or "fork" not in mp.get_all_start_methods(): + return None + return mp.get_context("fork") + + +def measure_floor(ctx, n: int = 5) -> float: + """Peak RSS attributed to an empty call. + + A forked child pays a small constant cost of its own (the sampler thread, the + gc pass, copy-on-write faults from refcount writes on inherited objects). + It is the same on both sides of a comparison so it cancels out, but it is + worth reporting so nobody reads a 1 MiB row as 1 MiB of real allocation. + """ + try: + return float(np.median([measure(lambda: None, ctx) for _ in range(n)])) + except Exception: # noqa: BLE001 -- diagnostic only + return 0.0 + + +def run_case(name: str, large_size: int, repeats: int, warmup: int, tmp: Path, ctx) -> dict: + case = build_case(name, large_size=large_size) + measurements, metrics, setup_notes = build_measurements(case, tmp) + + summaries: dict[str, dict[str, float]] = {} + skipped: list[str] = [] + for meas in measurements: + try: + # Warmup runs in the *parent*, on purpose. A forked child inherits + # whatever the parent has already imported, so the point of a warmup + # here is to pull in lazily-imported modules (scipy.ndimage, cc3d, …) + # before the fork — otherwise every child pays for them and every + # sample is inflated by the same import cost. + with silenced(): + for _ in range(warmup): + meas.call() + summaries[meas.key] = summarize([measure(meas.call, ctx) for _ in range(repeats)]) + except Exception as exc: # noqa: BLE001 -- a missing/broken op must not abort the run + skipped.append(f"{meas.key} ({type(exc).__name__}: {exc})") + + print_case_table(case.name, case.shape, summaries, "MiB") + if setup_notes: + print(f" setup failed: {'; '.join(setup_notes)}") + if skipped: + print(f" skipped: {'; '.join(skipped)}") + + summaries.update(metrics) + return {"name": case.name, "shape": list(case.shape), "voxels": case.voxels, "measurements_mb": summaries} + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_common_args(p) + args = p.parse_args(argv) + repeats, warmup = resolve_repeats(args) + large_size = resolve_large_size(args) + names = parse_case_names(args.cases) + + ctx = _fork_context() + sampler = "proc-status" if _HAS_PROC else "getrusage" + isolation = "fork" if ctx is not None else "in-process" + floor = measure_floor(ctx) + print( + f"memory benchmark | python {python_version()} | commit {commit_hash()} | " + f"repeats={repeats} warmup={warmup} large_size={large_size} | " + f"sampler={sampler} isolation={isolation} floor={floor:.2f} MiB" + ) + started = time.perf_counter() + with tempfile.TemporaryDirectory(prefix="tptbox-benchmem-") as td: + tmp = Path(td) + cases = [run_case(n, large_size, repeats, warmup, tmp, ctx) for n in names] + print(f"\ntotal wall time: {time.perf_counter() - started:.1f}s") + + write_json( + args.json, + { + "python": python_version(), + "commit": commit_hash(), + "repeats": repeats, + "warmup": warmup, + "large_size": large_size, + "sampler": sampler, + "isolation": isolation, + "floor_mib": round(floor, 3), + "cases": cases, + }, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/nii_poi/bench_speed.py b/benchmarks/nii_poi/bench_speed.py new file mode 100644 index 00000000..01cc5668 --- /dev/null +++ b/benchmarks/nii_poi/bench_speed.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python +"""Speed benchmark for the NII and POI hot paths. + +Times every measurement in ``measurements.py`` across the six workloads in +``workloads.py`` and writes a JSON document that ``compare.py`` can diff against +a second run (typically the PR base commit). + + python benchmarks/nii_poi/bench_speed.py # 10 repeats + 1 warmup + python benchmarks/nii_poi/bench_speed.py --quick --json head.json + python benchmarks/nii_poi/bench_speed.py --cases ct_3d,ct_2d --repeats 3 +""" + +from __future__ import annotations + +import argparse +import gc +import sys +import tempfile +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from _common import ( # noqa: E402 + add_common_args, + commit_hash, + print_case_table, + python_version, + resolve_large_size, + resolve_repeats, + summarize, + write_json, +) +from measurements import build_measurements, silenced # noqa: E402 +from workloads import build_case, parse_case_names # noqa: E402 + + +def time_measurement(call, repeats: int, warmup: int) -> list[float]: + """Return per-iteration wall times in milliseconds.""" + with silenced(): + for _ in range(warmup): + call() + samples: list[float] = [] + for _ in range(repeats): + gc.collect() + with silenced(): + t0 = time.perf_counter_ns() + call() + t1 = time.perf_counter_ns() + samples.append((t1 - t0) / 1e6) + return samples + + +def run_case(name: str, large_size: int, repeats: int, warmup: int, tmp: Path) -> dict: + case = build_case(name, large_size=large_size) + measurements, metrics, setup_notes = build_measurements(case, tmp) + + summaries: dict[str, dict[str, float]] = {} + skipped: list[str] = [] + for meas in measurements: + try: + summaries[meas.key] = summarize(time_measurement(meas.call, repeats, warmup)) + except Exception as exc: # noqa: BLE001 -- a missing/broken op must not abort the run + skipped.append(f"{meas.key} ({type(exc).__name__}: {exc})") + + print_case_table(case.name, case.shape, summaries, "ms") + if setup_notes: + print(f" setup failed: {'; '.join(setup_notes)}") + if skipped: + print(f" skipped: {'; '.join(skipped)}") + + summaries.update(metrics) + return {"name": case.name, "shape": list(case.shape), "voxels": case.voxels, "measurements_ms": summaries} + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_common_args(p) + args = p.parse_args(argv) + repeats, warmup = resolve_repeats(args) + large_size = resolve_large_size(args) + names = parse_case_names(args.cases) + + print( + f"speed benchmark | python {python_version()} | commit {commit_hash()} | repeats={repeats} warmup={warmup} large_size={large_size}" + ) + started = time.perf_counter() + with tempfile.TemporaryDirectory(prefix="tptbox-bench-") as td: + tmp = Path(td) + cases = [run_case(n, large_size, repeats, warmup, tmp) for n in names] + print(f"\ntotal wall time: {time.perf_counter() - started:.1f}s") + + write_json( + args.json, + { + "python": python_version(), + "commit": commit_hash(), + "repeats": repeats, + "warmup": warmup, + "large_size": large_size, + "cases": cases, + }, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/nii_poi/compare.py b/benchmarks/nii_poi/compare.py new file mode 100644 index 00000000..cdb168e9 --- /dev/null +++ b/benchmarks/nii_poi/compare.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python +"""Compare two ``bench_speed.py`` JSON files and emit a markdown report. + +Writes the report to stdout and exits 1 if any measurement regressed past the +threshold (see ``_compare_core`` for the exact gate). + + python benchmarks/nii_poi/compare.py baseline.json head.json --fail-on-regression-pct 50 +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from _common import load_json # noqa: E402 +from _compare_core import build_parser, compare # noqa: E402 + + +def main(argv: list[str] | None = None) -> int: + args = build_parser(__doc__, default_pct=50.0).parse_args(argv) + report, regressed = compare( + load_json(args.baseline), + load_json(args.head), + measurements_key="measurements_ms", + unit="ms", + threshold_pct=args.fail_on_regression_pct, + alpha=args.alpha, + title="Speed (wall time per call)", + ) + print(report) + return 1 if regressed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/nii_poi/compare_mem.py b/benchmarks/nii_poi/compare_mem.py new file mode 100644 index 00000000..f9d9365e --- /dev/null +++ b/benchmarks/nii_poi/compare_mem.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python +"""Compare two ``bench_mem.py`` JSON files and emit a markdown report. + +Writes the report to stdout and exits 1 if any measurement regressed past the +threshold (see ``_compare_core`` for the exact gate). + + python benchmarks/nii_poi/compare_mem.py baseline_mem.json head_mem.json --fail-on-regression-pct 50 +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from _common import load_json # noqa: E402 +from _compare_core import build_parser, compare # noqa: E402 + + +def main(argv: list[str] | None = None) -> int: + args = build_parser(__doc__, default_pct=50.0).parse_args(argv) + report, regressed = compare( + load_json(args.baseline), + load_json(args.head), + measurements_key="measurements_mb", + unit="MiB", + threshold_pct=args.fail_on_regression_pct, + alpha=args.alpha, + title="Memory (peak RSS growth per call)", + ) + print(report) + return 1 if regressed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/nii_poi/measurements.py b/benchmarks/nii_poi/measurements.py new file mode 100644 index 00000000..9da91c45 --- /dev/null +++ b/benchmarks/nii_poi/measurements.py @@ -0,0 +1,256 @@ +"""The measurement registry shared by the speed and the memory benchmark. + +Both benches call :func:`build_measurements` so the two JSON files always cover +exactly the same keys, and so a comparison never accidentally lines up different +work under the same name. + +Everything here has to survive being pointed at **two different versions of +TPTBox**: the workflow runs this same file against the PR head and against the +PR base commit. Three things make that work. + +* :func:`bind` pre-filters keyword arguments against the signature actually + present, so a kwarg added in the PR is simply dropped on the baseline side + instead of raising ``TypeError``. (Same trick as + ``benchmarks/benchmark_nnunet_inference.py``.) +* :func:`accepts` lets a measurement opt *out* entirely when the feature it + exists to measure is absent, rather than silently measuring something else. +* Setup and calls are individually guarded. A failure drops the affected keys + and is reported on a ``skipped:`` / ``setup:`` line; it never aborts the run. + The comparison tools list keys present on only one side as ``new`` / + ``removed`` and never gate on them. +""" + +from __future__ import annotations + +import contextlib +import functools +import inspect +import io +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + +# This directory is added to sys.path so `workloads` resolves whether the bench +# scripts are run as files or this module is imported directly. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from workloads import Case # noqa: E402 + +from TPTBox.core import poi as poi_module # noqa: E402 +from TPTBox.core.nii_wrapper import NII # noqa: E402 +from TPTBox.core.poi import POI # noqa: E402 + +# Resolved by name rather than imported: a module-level `from ... import foo` +# would raise ImportError on a baseline commit that predates `foo`, taking the +# whole harness down instead of dropping one measurement. +calc_centroids = getattr(poi_module, "calc_centroids", None) +calc_poi_from_subreg_vert = getattr(poi_module, "calc_poi_from_subreg_vert", None) + +#: Number of points fed to POI.local_to_global_arr. +_ARR_POINTS = 10_000 + + +@dataclass +class Measurement: + """A single measured call, already bound to its arguments.""" + + key: str + call: Callable[[], Any] + + +@contextlib.contextmanager +def silenced(): + """Swallow the chatty ``print``-based progress output of NII/POI operations. + + Entered *outside* the measured region so it costs nothing in the numbers. + """ + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + yield + + +def _parameters(fn) -> dict[str, inspect.Parameter] | None: + try: + return dict(inspect.signature(fn).parameters) + except (TypeError, ValueError): + return None + + +def accepts(fn, name: str) -> bool: + """True if ``fn`` in *this* version of TPTBox takes a ``name`` keyword.""" + params = _parameters(fn) + return bool(params) and name in params + + +def bind(fn, *args, **kwargs) -> Callable[[], Any]: + """Pre-bind a zero-argument call, dropping kwargs this version does not accept. + + Resolved once here rather than at call time, so the signature introspection + never lands inside a measured region. + """ + params = _parameters(fn) + if params is not None and not any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + kwargs = {k: v for k, v in kwargs.items() if k in params} + return functools.partial(fn, *args, **kwargs) + + +def _other_orientation(nii: NII): + return ("R", "A", "S") if tuple(nii.orientation) == ("P", "I", "R") else ("P", "I", "R") + + +def _pad_target(case: Case) -> tuple[int, ...]: + """Grow the in-plane axes by 8 voxels, leaving a singleton axis singleton.""" + return tuple(s + 8 if s > 1 else s for s in case.shape) + + +def build_measurements(case: Case, tmp: Path) -> tuple[list[Measurement], dict[str, float], list[str]]: + """Return the measurements for ``case``, its ``metric_*`` context values, and setup notes. + + ``tmp`` is a scratch directory for the save/load measurements; it is written + to once here (untimed) so the load measurements have something real to read. + """ + img, vert, subreg = case.img, case.vert, case.subreg + notes: list[str] = [] + ok = object() + + def resolve(label: str, target): + """Turn ``(obj, "method")`` into a bound callable, or None if absent here. + + Methods are referred to by *name* rather than attribute access so that a + method which does not exist on the baseline commit drops one measurement + instead of raising AttributeError while the registry is being built. + """ + if isinstance(target, tuple): + obj, name = target + fn = getattr(obj, name, None) + if fn is None: + notes.append(f"{label}: {type(obj).__name__}.{name} does not exist in this version") + return fn + if target is None: + notes.append(f"{label}: not available in this version") + return target + + def setup(label: str, target, *args, **kwargs): + """Run one piece of untimed setup, recording rather than raising on failure. + + Returns the call's result, or ``None`` if it was unavailable or failed. + Calls that return nothing on success (the ``save`` ones) get the ``ok`` + sentinel instead, so success stays distinguishable from failure. + """ + fn = resolve(label, target) + if fn is None: + return None + try: + with silenced(): + result = bind(fn, *args, **kwargs)() + except Exception as exc: # noqa: BLE001 -- an absent API must not abort the run + notes.append(f"{label} ({type(exc).__name__}: {exc})") + return None + return ok if result is None else result + + orient_to = _other_orientation(vert) + labels = setup("unique", (vert, "unique")) + labels = list(labels) if isinstance(labels, (list, tuple, np.ndarray)) else [] + first_label = labels[0] if labels else 1 + grid_2mm = setup("rescale-grid", (img, "rescale"), (2, 2, 2), verbose=False) + crop = setup("compute-crop", (vert, "compute_crop"), raise_error=False) + poi = setup("calc_centroids", calc_centroids, vert) + region_map = None + if poi is not None: + regions = setup("poi-keys_region", (poi, "keys_region")) + if regions is not None: + region_map = {int(k): int(k) + 100 for k in regions} + + img_file = tmp / f"{case.name}_img.nii.gz" + seg_file = tmp / f"{case.name}_seg.nii.gz" + poi_file = tmp / f"{case.name}_poi.json" + out_img = tmp / f"{case.name}_roundtrip.nii.gz" + have_img = setup("write-img", (img, "save"), img_file, verbose=False) is not None + have_seg = setup("write-seg", (vert, "save"), seg_file, verbose=False) is not None + have_poi_file = poi is not None and setup("write-poi", (poi, "save"), poi_file, verbose=False) is not None + + label_map = {int(a): int(a) + 100 for a in labels} + coords = np.asarray([[i % 97, (i * 7) % 89, (i * 13) % 83] for i in range(_ARR_POINTS)], dtype=float) + pad_to = _pad_target(case) + + m: list[Measurement] = [] + + def add(key: str, target, *args, **kwargs) -> None: + fn = resolve(key, target) + if fn is not None: + m.append(Measurement(key, bind(fn, *args, **kwargs))) + + # --- NII: IO ----------------------------------------------------------- + if have_img: + add("nii_load_img", lambda: NII.load(img_file, seg=False).get_array()) + if have_seg: + add("nii_load_seg", lambda: NII.load(seg_file, seg=True).get_seg_array()) + add("nii_save", (img, "save"), out_img, verbose=False) + + # --- NII: array access / dtype ---------------------------------------- + add("nii_get_array", (img, "get_array")) + add("nii_set_dtype", (img, "set_dtype"), np.float32) + + # --- NII: geometry ----------------------------------------------------- + add("nii_reorient", (img, "reorient"), orient_to) + add("nii_rescale", (img, "rescale"), (2, 2, 2), verbose=False) + add("nii_rescale_seg", (vert, "rescale"), (2, 2, 2), verbose=False) + if grid_2mm is not None: + add("nii_resample_from_to", (img, "resample_from_to"), grid_2mm, verbose=False) + add("nii_compute_crop", (vert, "compute_crop"), raise_error=False) + if crop is not None: + add("nii_apply_crop", (vert, "apply_crop"), crop) + add("nii_pad_to", (img, "pad_to"), pad_to) + + # --- NII: statistics --------------------------------------------------- + add("nii_unique", (vert, "unique")) + add("nii_volumes", (vert, "volumes")) + add("nii_center_of_masses", (vert, "center_of_masses")) + + # --- NII: label ops ---------------------------------------------------- + add("nii_extract_label", (vert, "extract_label"), first_label) + add("nii_map_labels", (vert, "map_labels"), label_map, verbose=False) + + # --- NII: morphology / connected components ---------------------------- + add("nii_dilate_msk", (vert, "dilate_msk"), 2, verbose=False) + add("nii_erode_msk", (vert, "erode_msk"), 1, verbose=False) + add("nii_fill_holes", (vert, "fill_holes")) + add("nii_get_connected_components", (vert, "get_connected_components"), connectivity=3) + add("nii_filter_connected_components", (vert, "filter_connected_components"), min_volume=8, connectivity=3) + + # --- POI --------------------------------------------------------------- + add("poi_calc_centroids", calc_centroids, vert) + if not case.is_heavy and calc_centroids is not None and accepts(calc_centroids, "_crop"): + # The _crop=False path is a per-label scipy loop kept around as an A/B + # against the single-pass np_center_of_mass. On a 400^3 volume it is + # minutes per repeat and would dominate the whole run, so large cases + # skip it. Where the kwarg does not exist there is nothing to compare, + # and measuring it anyway would just duplicate poi_calc_centroids. + add("poi_calc_centroids_nocrop", calc_centroids, vert, _crop=False) + if poi is not None: + add("poi_reorient", (poi, "reorient"), orient_to) + add("poi_rescale", (poi, "rescale"), (2, 2, 2), verbose=False) + add("poi_to_global", (poi, "to_global")) + add("poi_local_to_global_arr", (poi, "local_to_global_arr"), coords) + add("poi_save", (poi, "save"), poi_file, verbose=False) + if region_map is not None: + add("poi_map_labels", (poi, "map_labels"), label_map_region=region_map) + if grid_2mm is not None: + add("poi_resample_from_to", (poi, "resample_from_to"), grid_2mm) + if have_poi_file: + add("poi_load", (POI, "load"), poi_file) + if case.anatomical: + # The heavy strategy pipeline. Only meaningful on real spine data, and + # nonsense on a single slice or on synthetic cuboids. + add("poi_calc_poi_from_subreg_vert", calc_poi_from_subreg_vert, vert, subreg, verbose=False) + + metrics = { + "metric_voxels": float(case.voxels), + "metric_labels": float(len(labels)), + "metric_foreground_pct": float(round(100.0 * float(np.count_nonzero(vert.get_seg_array())) / max(case.voxels, 1), 3)), + } + return m, metrics, notes diff --git a/benchmarks/nii_poi/workloads.py b/benchmarks/nii_poi/workloads.py new file mode 100644 index 00000000..881daed7 --- /dev/null +++ b/benchmarks/nii_poi/workloads.py @@ -0,0 +1,204 @@ +"""Workload construction for the NII/POI benchmarks. + +Six cases spanning a small/large x 2D/3D grid, so a change that only hurts (say) +large 3D volumes shows up as such instead of being averaged away: + + ct_3d sample CT (73, 47, 73) small 3D + ct_2d centre slice of ct_3d small 2D + mri_3d sample MRI (68, 52, 67) small 3D + mri_2d centre slice of mri_3d small 2D + synth_3d generated (400, 400, 400) large 3D + synth_2d generated (400, 400, 1) large 2D + +"2D" is a singleton third axis, not a 2-element shape: ``NII`` has no 2D code +path (orientation/reorient/rescale all assume three axes), so a genuine 2-D +``Nifti1Image`` would crash. A ``(X, Y, 1)`` volume is geometrically 2D and goes +through the normal machinery unharmed. + +The sample data is read by path rather than through ``TPTBox.tests.test_utils``. +The benchmark workflow swaps the whole ``TPTBox/`` tree to the PR base commit to +produce the baseline numbers, so anything imported from inside that tree would +silently become the *old* version; this module lives outside it and stays fixed. +""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import nibabel as nib # noqa: E402 + +from TPTBox.core.nii_wrapper import NII # noqa: E402 + +SAMPLE_CT = REPO_ROOT / "TPTBox" / "tests" / "sample_ct" +SAMPLE_MRI = REPO_ROOT / "TPTBox" / "tests" / "sample_mri" + +CASE_NAMES = ("ct_3d", "ct_2d", "mri_3d", "mri_2d", "synth_3d", "synth_2d") +DEFAULT_LARGE_SIZE = 400 +#: Cases above this voxel count skip measurements marked as quadratic/per-label. +HEAVY_VOXEL_LIMIT = 4_000_000 + + +@dataclass +class Case: + """One benchmark workload: an intensity image plus its two segmentations.""" + + name: str + img: NII + vert: NII + subreg: NII + is_2d: bool + #: True for the real spine samples, where the full POI pipeline is meaningful. + anatomical: bool + img_path: Path | None = None + vert_path: Path | None = None + + @property + def shape(self) -> tuple[int, ...]: + return tuple(int(s) for s in self.vert.shape) + + @property + def voxels(self) -> int: + return int(np.prod(self.shape)) + + @property + def is_heavy(self) -> bool: + return self.voxels > HEAVY_VOXEL_LIMIT + + +def _center_slice(nii: NII) -> NII: + """Take the centre slice along the raw third array axis, keeping it singleton. + + Uses ``NII.apply_crop`` so the affine is updated with the slice, which a plain + numpy index would not do. Indexing the raw axis (rather than an anatomical + one) keeps the result reproducible regardless of the file's orientation. + """ + c = nii.shape[2] // 2 + return nii.apply_crop((slice(None), slice(None), slice(c, c + 1))) + + +def _load_sample(folder: Path, stem_img: str, stem_subreg: str, stem_vert: str): + img = NII.load(folder / stem_img, seg=False) + subreg = NII.load(folder / stem_subreg, seg=True) + vert = NII.load(folder / stem_vert, seg=True) + return img, subreg, vert + + +def _sample_ct(): + return _load_sample( + SAMPLE_CT, + "sub-ct_label-22_ct.nii.gz", + "sub-ct_seg-subreg_label-22_msk.nii.gz", + "sub-ct_seg-vert_label-22_msk.nii.gz", + ) + + +def _sample_mri(): + return _load_sample( + SAMPLE_MRI, + "sub-mri_label-6_T2w.nii.gz", + "sub-mri_seg-subreg_label-6_msk.nii.gz", + "sub-mri_seg-vert_label-6_msk.nii.gz", + ) + + +def _synthetic_arrays(size: int, flat: bool): + """Build a deterministic labelled volume plus a matching intensity volume. + + Labels are hollow cuboids laid out on a regular grid: hollow so ``fill_holes`` + has real work to do, on a grid so they never touch (one connected component + per label) and the result is identical on every machine. Roughly 20% of the + volume is foreground, which is enough that ``use_crop=True`` cannot trivialise + the morphology measurements. + + The intensity volume is a smooth ramp plus the labels rather than noise, so + that gzip-compressed save/load stays a sane thing to measure. + """ + shape = (size, size, 1) if flat else (size, size, size) + seg = np.zeros(shape, dtype=np.uint16) + + grid = 5 if flat else 3 + axes = (0, 1) if flat else (0, 1, 2) + cell = [shape[a] // grid for a in axes] + label = 1 + steps = [range(grid) for _ in axes] + for i in steps[0]: + for j in steps[1]: + for k in steps[2] if len(steps) > 2 else [0]: + idx = (i, j, k)[: len(axes)] + lo = [idx[a] * cell[a] + cell[a] // 8 for a in range(len(axes))] + hi = [lo[a] + (cell[a] * 5) // 8 for a in range(len(axes))] + sl = [slice(lo[a], hi[a]) for a in range(len(axes))] + if flat: + sl.append(slice(None)) + seg[tuple(sl)] = label + # carve a cavity so fill_holes is not a no-op + inner = [] + for a in range(len(axes)): + pad = max((hi[a] - lo[a]) // 4, 1) + inner.append(slice(lo[a] + pad, hi[a] - pad)) + if flat: + inner.append(slice(None)) + seg[tuple(inner)] = 0 + label += 1 + + ramp = np.linspace(-500, 500, shape[0], dtype=np.int16)[:, None, None] + np.linspace(-200, 200, shape[1], dtype=np.int16)[None, :, None] + img = (seg.astype(np.int16) * 40 + ramp).astype(np.int16) + return seg, img + + +def _synthetic_case(name: str, size: int, flat: bool) -> Case: + seg, img = _synthetic_arrays(size, flat) + affine = np.eye(4) + affine[0, 0] = affine[1, 1] = affine[2, 2] = 1.5 + affine[:3, 3] = (-100.0, -50.0, -25.0) + vert = NII(nib.nifti1.Nifti1Image(seg, affine), seg=True) + # A second label space so calc_poi_from_two_segs-style calls have two masks; + # the synthetic subreg is the vertebra mask remapped into the 50-range. + subreg = NII(nib.nifti1.Nifti1Image((seg > 0).astype(np.uint16) * 50, affine), seg=True) + image = NII(nib.nifti1.Nifti1Image(img, affine), seg=False) + return Case(name=name, img=image, vert=vert, subreg=subreg, is_2d=flat, anatomical=False) + + +def build_case(name: str, large_size: int = DEFAULT_LARGE_SIZE) -> Case: + """Construct a single case by name. Built one at a time so the large volumes are not all resident at once.""" + if name in ("ct_3d", "ct_2d"): + img, subreg, vert = _sample_ct() + paths = (SAMPLE_CT / "sub-ct_label-22_ct.nii.gz", SAMPLE_CT / "sub-ct_seg-vert_label-22_msk.nii.gz") + elif name in ("mri_3d", "mri_2d"): + img, subreg, vert = _sample_mri() + paths = (SAMPLE_MRI / "sub-mri_label-6_T2w.nii.gz", SAMPLE_MRI / "sub-mri_seg-vert_label-6_msk.nii.gz") + elif name == "synth_3d": + return _synthetic_case(name, large_size, flat=False) + elif name == "synth_2d": + return _synthetic_case(name, large_size, flat=True) + else: + raise ValueError(f"unknown case {name!r}; known: {', '.join(CASE_NAMES)}") + + if name.endswith("_2d"): + return Case( + name=name, + img=_center_slice(img), + vert=_center_slice(vert), + subreg=_center_slice(subreg), + is_2d=True, + anatomical=False, # a single slice is not a spine + ) + return Case(name=name, img=img, vert=vert, subreg=subreg, is_2d=False, anatomical=True, img_path=paths[0], vert_path=paths[1]) + + +def parse_case_names(raw: str | None) -> list[str]: + if not raw: + return list(CASE_NAMES) + out = [c.strip() for c in raw.split(",") if c.strip()] + unknown = [c for c in out if c not in CASE_NAMES] + if unknown: + raise ValueError(f"unknown case(s) {unknown}; known: {', '.join(CASE_NAMES)}") + return out