From 4b45657d901e82ecb9ab9bdcf4e8df26c18d54fe Mon Sep 17 00:00:00 2001 From: Johan Mathe Date: Thu, 16 Jul 2026 22:49:28 +0000 Subject: [PATCH 1/3] Unify experiment comparison protocol and build real-data grid figure pipeline Addresses figure feedback: consistent "Aug. CNN" (R-trained) baselines and OOD y-labels across rows, a single power-spectrum entry whose capacity curve covers the bispectrum budget, norm-pool (power-spectrum analog) baseline for PCam, full method coverage in every panel, and a caption justifying the Cohen S2CNN published-reference line. - pcam: new norm_pool model variant (GroupNormPool as final invariant pool), sweep scripts retrained per protocol (standard=R, invariants=C, rotation eval always on) - spherical_mnist: new run_capacity_sweep.sh (accuracy-vs-params curves); matched power-spectrum concept removed - organ3d: all 4 models in data-efficiency and capacity sweeps, 3 seeds; redundant run_tier1_sweep.sh removed - experiments/make_grid_figure.py: single figure entry point with results.json loaders + --mock layout preview, replaces make_mock_grid_panels.py; emits caption.txt - experiments/README.md: run matrix, tmux commands, rsync instructions for external GPU machines Co-authored-by: Cursor --- .gitignore | 4 + experiments/README.md | 85 ++ experiments/make_grid_figure.py | 1148 +++++++++++++++++ experiments/organ3d/README.md | 4 +- experiments/organ3d/analyze_results.py | 56 +- experiments/organ3d/run_dataeff_multiseed.sh | 38 +- experiments/organ3d/run_sweep.sh | 35 +- experiments/organ3d/run_tier1_sweep.sh | 97 -- experiments/organ3d/run_wider_multiseed.sh | 41 +- experiments/pcam/README.md | 5 +- experiments/pcam/analyze_data_pareto.py | 23 +- experiments/pcam/analyze_pareto.py | 43 +- experiments/pcam/model.py | 39 +- experiments/pcam/run_data_pareto_sweep.sh | 62 +- experiments/pcam/run_matched_sweep.sh | 77 +- experiments/pcam/train.py | 40 +- experiments/spherical_mnist/.gitignore | 2 +- experiments/spherical_mnist/README.md | 4 +- .../spherical_mnist/analyze_results.py | 57 +- .../spherical_mnist/run_capacity_sweep.sh | 64 + .../spherical_mnist/run_data_efficiency.sh | 36 +- 21 files changed, 1614 insertions(+), 346 deletions(-) create mode 100644 experiments/README.md create mode 100644 experiments/make_grid_figure.py delete mode 100755 experiments/organ3d/run_tier1_sweep.sh create mode 100755 experiments/spherical_mnist/run_capacity_sweep.sh diff --git a/.gitignore b/.gitignore index f531328..c407a26 100644 --- a/.gitignore +++ b/.gitignore @@ -97,3 +97,7 @@ uv.lock *.snm *.toc *.vrb + +# Experiment figure outputs +experiments/grid_mockup/ +experiments/grid_figure/ diff --git a/experiments/README.md b/experiments/README.md new file mode 100644 index 0000000..e524861 --- /dev/null +++ b/experiments/README.md @@ -0,0 +1,85 @@ +# Experiments + +Three benchmark experiments comparing G-bispectrum pooling against baseline invariant/equivariant pooling strategies, plus the reconstruction demo. Results feed the 3x3 grid figure built by `make_grid_figure.py`. + +## Comparison protocol (all experiments) + +- Invariant models are trained on canonical (non-augmented) data (`train_mode C`). +- The non-equivariant CNN baseline is trained with G-augmentation (`train_mode R`) and labeled "Aug. CNN" in the figures. +- Every run records both the canonical test metric and the rotated (OOD) test metric; line plots report the rotated metric, mean ± std over seeds 42/123/456. +- Cohen et al. (2018) S²CNN appears in the Spherical MNIST row as a published-reference dashed line (cited, not re-run). + +## Setup (on the GPU machine) + +```bash +git clone bispectrum && cd bispectrum +git checkout +python -m venv .venv && source .venv/bin/activate +pip install -e ".[dev,experiments]" +``` + +Datasets download automatically on first use (PCam from Zenodo ≈ 8 GB, OrganMNIST3D via `medmnist`, MNIST via `torchvision`). + +## Run matrix + +Each script is resumable: completed runs (existing `results.json`) are skipped, so re-running after an interruption is safe. + +| # | Experiment | Script | Runs | Feeds | +|---|------------|--------|------|-------| +| 1 | PCam | `pcam/run_matched_sweep.sh` (then `--phase-b`) | 5 CNN configs (R) + 5x5 equivariant (C) + 7 so2_disk, x3 seeds | row 1 params | +| 2 | PCam | `pcam/run_data_pareto_sweep.sh` (then `--phase-b`) | 6 models x 5 sizes x 3 seeds at ~100K params | row 1 bars + data | +| 3 | Organ3D | `organ3d/run_sweep.sh` | 4 models x 3 seeds at channels (4,8) | row 2 bars + curve anchors | +| 4 | Organ3D | `organ3d/run_wider_multiseed.sh` | 4 models x 2 wider channel configs x 3 seeds | row 2 params | +| 5 | Organ3D | `organ3d/run_dataeff_multiseed.sh` | 4 models x 4 sizes x 3 seeds | row 2 data | +| 6 | SMNIST | `spherical_mnist/run_sweep.sh` | 3 models x 2 modes x 3 seeds | row 3 bars + Cohen table | +| 7 | SMNIST | `spherical_mnist/run_capacity_sweep.sh` | 2 models x 5 widths x 3 seeds (C) | row 3 params | +| 8 | SMNIST | `spherical_mnist/run_data_efficiency.sh` | 3 models x 4 sizes x 3 seeds | row 3 data | + +Expected results directories (created next to each script): + +``` +pcam/pcam_results_pareto/ # 1 +pcam/pcam_results_data_pareto/ # 2 (n_100/ ... n_full/ subdirs) +organ3d/organ3d_results/ # 3, 4, 5 (shared) +spherical_mnist/smnist_results/ # 6, 8 (shared) +spherical_mnist/smnist_results_capacity/ # 7 +``` + +## Running on the GPU machine (tmux) + +One detached session per experiment family; each logs to a file. The three families are independent — run them on separate GPUs/machines if available (`CUDA_VISIBLE_DEVICES=` before `bash` to pin a GPU). + +```bash +cd ~/bispectrum/experiments + +tmux new-session -d -s pcam 'cd pcam && { bash run_matched_sweep.sh && bash run_matched_sweep.sh --phase-b && bash run_data_pareto_sweep.sh && bash run_data_pareto_sweep.sh --phase-b; } 2>&1 | tee pcam_sweeps.log' + +tmux new-session -d -s organ3d 'cd organ3d && { bash run_sweep.sh && bash run_wider_multiseed.sh && bash run_dataeff_multiseed.sh; } 2>&1 | tee organ3d_sweeps.log' + +tmux new-session -d -s smnist 'cd spherical_mnist && { bash run_sweep.sh && bash run_capacity_sweep.sh && bash run_data_efficiency.sh; } 2>&1 | tee smnist_sweeps.log' +``` + +Monitor with `tmux attach -t pcam` (detach: `Ctrl-b d`) or `tail -f /_sweeps.log`. + +## Pulling results back and building the figure + +From the analysis machine: + +```bash +REMOTE=user@gpu-machine:~/bispectrum/experiments +rsync -avz --include='*/' --include='results.json' --exclude='*' \ + "$REMOTE/pcam/pcam_results_pareto/" experiments/pcam/pcam_results_pareto/ +rsync -avz --include='*/' --include='results.json' --exclude='*' \ + "$REMOTE/pcam/pcam_results_data_pareto/" experiments/pcam/pcam_results_data_pareto/ +rsync -avz --include='*/' --include='results.json' --exclude='*' \ + "$REMOTE/organ3d/organ3d_results/" experiments/organ3d/organ3d_results/ +rsync -avz --include='*/' --include='results.json' --exclude='*' \ + "$REMOTE/spherical_mnist/smnist_results/" experiments/spherical_mnist/smnist_results/ +rsync -avz --include='*/' --include='results.json' --exclude='*' \ + "$REMOTE/spherical_mnist/smnist_results_capacity/" experiments/spherical_mnist/smnist_results_capacity/ + +python experiments/make_grid_figure.py # real data -> experiments/grid_figure/ +python experiments/make_grid_figure.py --mock # layout preview -> experiments/grid_mockup/ +``` + +Outputs: per-panel PDFs (`grid_r{row}c{col}_*.pdf`), per-row legends, an assembled contact sheet (PNG + PDF), and `caption.txt` with the figure caption including the S²CNN justification. diff --git a/experiments/make_grid_figure.py b/experiments/make_grid_figure.py new file mode 100644 index 0000000..0429980 --- /dev/null +++ b/experiments/make_grid_figure.py @@ -0,0 +1,1148 @@ +#!/usr/bin/env python3 +"""Build the 3x3 experiment grid figure (bars + param/data efficiency curves). + +Rows: PCam (C8), OrganMNIST3D (octahedral O), Spherical MNIST (SO(3)). +Columns: + 1. OOD rotation robustness — canonical vs rotated test, bar pairs. + 2. Parameter efficiency — rotated-test metric vs trainable params. + 3. Data efficiency — rotated-test metric vs training examples. + +Protocol (consistent across rows): + - "Aug. CNN" baselines are non-equivariant CNNs trained with + G-augmentation (train_mode R); invariant models are C-trained. + - All line plots report the rotated (OOD) test metric, mean +/- std + over seeds. + - Cohen et al. (2018) S2CNN is shown as a published-reference dashed + line in the Spherical MNIST row (cited, not re-run). + +Usage: + # Real data (after syncing results dirs from the GPU machines): + python make_grid_figure.py + + # Layout preview with deterministic mock data: + python make_grid_figure.py --mock +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import matplotlib + +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.axes import Axes +from matplotlib.lines import Line2D +from matplotlib.patches import Patch +from matplotlib.ticker import FuncFormatter +from numpy.typing import NDArray + +FloatArray = NDArray[np.float64] +LegendHandle = Line2D | Patch + +EXPERIMENTS_DIR: Final = Path(__file__).parent +PANEL_SIZE: Final = (4.0, 3.0) + +COHEN_S2CNN_NRR: Final = 0.94 # Cohen et al. (2018), NR/R accuracy, published. + + +# -------------------------------------------------------------------------- +# Data model +# -------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class MethodStyle: + label: str + color: str + marker: str + + +@dataclass(frozen=True) +class ReferenceLine: + value: float + label: str + linestyle: str = '--' + + +@dataclass(frozen=True) +class BarPanel: + methods: tuple[str, ...] + original: FloatArray + ood: FloatArray + original_error: FloatArray + ood_error: FloatArray + ylabel: str + ylim: tuple[float, float] + reference: ReferenceLine | None = None + + +@dataclass(frozen=True) +class CurveSeries: + method: str + x: FloatArray + y: FloatArray + error: FloatArray + + +@dataclass(frozen=True) +class CurvePanel: + series: tuple[CurveSeries, ...] + ylabel: str + ylim: tuple[float, float] + xticks: tuple[float, ...] | None = None + xticklabels: tuple[str, ...] | None = None + reference: ReferenceLine | None = None + + +@dataclass(frozen=True) +class RowData: + bar: BarPanel + params: CurvePanel + data: CurvePanel + + +GridData = tuple[RowData, RowData, RowData] + + +METHOD_STYLE: Final[dict[str, MethodStyle]] = { + # Non-equivariant CNN baselines, trained with G-augmentation. + 'standard': MethodStyle('Aug. CNN', '#888888', 's'), + 'standard_3d': MethodStyle('Aug. 3D CNN', '#888888', 's'), + 'standard_s2': MethodStyle('Aug. CNN', '#888888', 's'), + # Equivariant-pooling baselines. + 'norm': MethodStyle('NormReLU', '#2196F3', 'o'), + 'gate': MethodStyle('Gated', '#FF9800', '^'), + 'fourier_elu': MethodStyle('Fourier-ELU', '#9C27B0', 'D'), + 'max_pool': MethodStyle('Max Pool', '#00897B', 'P'), + # Incomplete second-order invariants (shared color family across rows). + 'norm_pool': MethodStyle('Norm pool', '#5C6BC0', 'v'), + 'power_spectrum': MethodStyle('Power spectrum', '#5C6BC0', 'v'), + # Complete invariant (ours). + 'bispectrum': MethodStyle('Bispectrum', '#E53935', '*'), +} + +ROW_METHODS: Final[tuple[tuple[str, ...], ...]] = ( + ('standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum'), + ('standard_3d', 'norm_pool', 'max_pool', 'bispectrum'), + ('standard_s2', 'power_spectrum', 'bispectrum'), +) + +ROW_LABELS: Final = ( + r'PCam — $C_8$ rotations', + r'OrganMNIST3D — octahedral $O$', + r'Spherical MNIST — $\mathrm{SO}(3)$', +) + +COLUMN_LABELS: Final = ( + 'OOD rotation robustness', + 'Parameter efficiency', + 'Data efficiency', +) + +PANEL_FILENAMES: Final = ( + ('grid_r1c1_pcam_ood.pdf', 'grid_r1c2_pcam_params.pdf', 'grid_r1c3_pcam_data.pdf'), + ('grid_r2c1_organ3d_ood.pdf', 'grid_r2c2_organ3d_params.pdf', 'grid_r2c3_organ3d_data.pdf'), + ('grid_r3c1_smnist_ood.pdf', 'grid_r3c2_smnist_params.pdf', 'grid_r3c3_smnist_data.pdf'), +) + +COHEN_REFERENCE: Final = ReferenceLine(COHEN_S2CNN_NRR, r'Cohen S$^2$CNN (published)') + +CAPTION_TEXT: Final = """\ +Consistent comparison of invariant pooling strategies across three +group-structured benchmarks: PatchCamelyon (planar C8 rotations, test AUC), +OrganMNIST3D (octahedral rotations, test accuracy), and Spherical MNIST +(SO(3) rotations, test accuracy). Columns: (left) canonical vs rotated +(OOD) test performance as paired bars; (middle) parameter efficiency; +(right) data efficiency. Line plots report the rotated-test metric, +mean +/- std over 3 seeds. "Aug. CNN" denotes the non-equivariant CNN +baseline trained with G-augmentation; all invariant models are trained on +canonical (non-augmented) data. "Norm pool" (finite groups) and "Power +spectrum" (SO(3)) are incomplete second-order invariant baselines with +matched backbones; the bispectrum is the complete invariant. The dashed +line marks the published Spherical CNN result of Cohen et al. (2018) +(NR/R = 0.94), cited rather than re-run. We include equivariant +architectures in every row: the NormReLU, Gated, Fourier-ELU, Max-Pool, +and Norm-Pool variants are G-equivariant networks differing only in +their invariant map. The S2CNN reference appears only in the Spherical +MNIST row because it is the canonical published architecture designed +specifically for spherical images; rows 1-2 have no comparably canonical +external baseline, and their equivariant variants already fill that role. +""" + + +# -------------------------------------------------------------------------- +# Results loading helpers +# -------------------------------------------------------------------------- + + +def _load_runs(root: Path) -> list[dict]: + """Recursively load every results.json under *root*.""" + runs: list[dict] = [] + if not root.exists(): + print(f'WARNING: results dir not found: {root}') + return runs + for p in sorted(root.rglob('results.json')): + with open(p) as f: + runs.append(json.load(f)) + return runs + + +def _mean_std(values: list[float]) -> tuple[float, float]: + arr = np.asarray(values, dtype=np.float64) + return float(arr.mean()), float(arr.std()) + + +def _metric(run: dict, key: str, metric: str) -> float | None: + """Extract run[key][metric], returning None when absent or empty.""" + block = run.get(key) + if not block: + return None + value = block.get(metric) + return float(value) if value is not None else None + + +def _expected_mode(method: str) -> str: + """Training protocol per method: Aug. CNN baselines R, invariant models C.""" + return 'R' if method.startswith('standard') else 'C' + + +def _canonical_mode(mode: str) -> str: + return 'C' if mode in ('C', 'NR') else 'R' + + +@dataclass(frozen=True) +class Aggregate: + canonical_mean: float + canonical_std: float + ood_mean: float + ood_std: float + n_params: float + n_seeds: int + + +def _aggregate( + runs: list[dict], + canonical_key: str, + metric: str, +) -> Aggregate | None: + """Aggregate canonical and rotated metrics over a list of seed runs.""" + canon = [v for r in runs if (v := _metric(r, canonical_key, metric)) is not None] + ood = [v for r in runs if (v := _metric(r, 'test_r', metric)) is not None] + if not canon or not ood: + return None + c_mean, c_std = _mean_std(canon) + o_mean, o_std = _mean_std(ood) + return Aggregate( + canonical_mean=c_mean, + canonical_std=c_std, + ood_mean=o_mean, + ood_std=o_std, + n_params=float(np.mean([r['n_params'] for r in runs])), + n_seeds=len(runs), + ) + + +def _build_bar_panel( + methods: tuple[str, ...], + aggregates: dict[str, Aggregate], + ylabel: str, + ylim: tuple[float, float], + reference: ReferenceLine | None = None, +) -> BarPanel: + present = tuple(m for m in methods if m in aggregates) + missing = [m for m in methods if m not in aggregates] + if missing: + print(f'WARNING: bar panel missing methods {missing}') + return BarPanel( + methods=present, + original=np.asarray([aggregates[m].canonical_mean for m in present]), + ood=np.asarray([aggregates[m].ood_mean for m in present]), + original_error=np.asarray([aggregates[m].canonical_std for m in present]), + ood_error=np.asarray([aggregates[m].ood_std for m in present]), + ylabel=ylabel, + ylim=ylim, + reference=reference, + ) + + +def _build_curve_panel( + methods: tuple[str, ...], + points: dict[str, list[tuple[float, float, float]]], + ylabel: str, + ylim: tuple[float, float], + xticks: tuple[float, ...] | None = None, + xticklabels: tuple[str, ...] | None = None, + reference: ReferenceLine | None = None, +) -> CurvePanel: + """Build a curve panel from per-method (x, y_mean, y_std) points.""" + series: list[CurveSeries] = [] + for method in methods: + pts = sorted(points.get(method, [])) + if not pts: + print(f'WARNING: curve panel missing method {method}') + continue + xs, ys, es = zip(*pts, strict=True) + series.append( + CurveSeries( + method, + np.asarray(xs, dtype=np.float64), + np.asarray(ys, dtype=np.float64), + np.asarray(es, dtype=np.float64), + ) + ) + return CurvePanel( + series=tuple(series), + ylabel=ylabel, + ylim=ylim, + xticks=xticks, + xticklabels=xticklabels, + reference=reference, + ) + + +def _group_runs( + runs: list[dict], + methods: tuple[str, ...], + model_key: str = 'model', +) -> dict[str, list[dict]]: + """Group runs by method, keeping only protocol-conforming train modes.""" + grouped: dict[str, list[dict]] = defaultdict(list) + for r in runs: + method = r.get(model_key) + if method not in methods: + continue + if _canonical_mode(r.get('train_mode', 'C')) != _expected_mode(method): + continue + grouped[method].append(r) + return grouped + + +# -------------------------------------------------------------------------- +# Experiment loaders +# -------------------------------------------------------------------------- + +PCAM_METHODS: Final = ('standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum') +ORGAN3D_METHODS: Final = ('standard', 'max_pool', 'norm_pool', 'bispectrum') +SMNIST_METHODS: Final = ('standard', 'power_spectrum', 'bispectrum') + + +def _is_full_train(run: dict) -> bool: + size = run.get('train_size') + return size is None or size <= 0 + + +def load_pcam(pareto_dir: Path, data_pareto_dir: Path) -> RowData: + """Row 1: PCam. Bars from matched-budget full-data runs, curves from sweeps.""" + pareto_runs = _group_runs(_load_runs(pareto_dir), PCAM_METHODS) + data_runs = _group_runs(_load_runs(data_pareto_dir), PCAM_METHODS) + + # Bars: matched ~100K-param configs trained on the full training set. + aggregates: dict[str, Aggregate] = {} + for method, runs in data_runs.items(): + full = [r for r in runs if _is_full_train(r)] + by_seed = defaultdict(list) + for r in full: + by_seed[r['seed']].append(r) + agg = _aggregate([rs[0] for rs in by_seed.values()], 'test_c', 'auc') + if agg is not None: + aggregates[method] = agg + + bar = _build_bar_panel( + PCAM_METHODS, + aggregates, + ylabel='Test AUC', + ylim=(0.82, 0.97), + ) + + # Param curves: pareto sweep grouped by growth rate. + param_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + for method, runs in pareto_runs.items(): + by_gr = defaultdict(list) + for r in runs: + by_gr[r.get('growth_rate')].append(r) + for gr_runs in by_gr.values(): + agg = _aggregate(gr_runs, 'test_c', 'auc') + if agg is not None: + param_points[method].append((agg.n_params, agg.ood_mean, agg.ood_std)) + + params = _build_curve_panel( + PCAM_METHODS, + param_points, + ylabel='Rotated test AUC', + ylim=(0.85, 0.97), + ) + + # Data curves: data-pareto sweep grouped by training-set size. + data_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + for method, runs in data_runs.items(): + by_size = defaultdict(list) + for r in runs: + by_size[r.get('train_examples')].append(r) + for size, size_runs in by_size.items(): + agg = _aggregate(size_runs, 'test_c', 'auc') + if agg is not None: + data_points[method].append((float(size), agg.ood_mean, agg.ood_std)) + + data = _build_curve_panel( + PCAM_METHODS, + data_points, + ylabel='Rotated test AUC', + ylim=(0.6, 0.98), + ) + return RowData(bar=bar, params=params, data=data) + + +def load_organ3d(results_dir: Path) -> RowData: + """Row 2: OrganMNIST3D. Bars from ch(4,8) full-data runs, curves from sweeps.""" + grouped = _group_runs(_load_runs(results_dir), ORGAN3D_METHODS) + + aggregates: dict[str, Aggregate] = {} + param_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + data_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + + for method, runs in grouped.items(): + full = [r for r in runs if _is_full_train(r)] + subset = [r for r in runs if not _is_full_train(r)] + + # Bars: default (4, 8) channel config, full training set. + base = [r for r in full if tuple(r.get('channels', ())) == (4, 8)] + agg = _aggregate(base, 'test_c', 'accuracy') + if agg is not None: + aggregates[method] = agg + + # Param curve: one point per channel config (full training set). + by_channels = defaultdict(list) + for r in full: + by_channels[tuple(r.get('channels', ()))].append(r) + for ch_runs in by_channels.values(): + ch_agg = _aggregate(ch_runs, 'test_c', 'accuracy') + if ch_agg is not None: + param_points[method].append((ch_agg.n_params, ch_agg.ood_mean, ch_agg.ood_std)) + + # Data curve: (4, 8) channels across training-set sizes + full point. + by_size = defaultdict(list) + for r in subset: + if tuple(r.get('channels', ())) == (4, 8): + by_size[r.get('train_examples')].append(r) + if base: + by_size[base[0].get('train_examples')] = base + for size, size_runs in by_size.items(): + size_agg = _aggregate(size_runs, 'test_c', 'accuracy') + if size_agg is not None: + data_points[method].append((float(size), size_agg.ood_mean, size_agg.ood_std)) + + remap = {'standard': 'standard_3d'} + aggregates = {remap.get(k, k): v for k, v in aggregates.items()} + param_points = {remap.get(k, k): v for k, v in param_points.items()} + data_points = {remap.get(k, k): v for k, v in data_points.items()} + + bar = _build_bar_panel( + ROW_METHODS[1], + aggregates, + ylabel='Test accuracy', + ylim=(0.0, 0.88), + ) + params = _build_curve_panel( + ROW_METHODS[1], + param_points, + ylabel='Rotated test accuracy', + ylim=(0.3, 0.86), + ) + data = _build_curve_panel( + ROW_METHODS[1], + data_points, + ylabel='Rotated test accuracy', + ylim=(0.05, 0.82), + ) + return RowData(bar=bar, params=params, data=data) + + +def load_smnist(results_dir: Path, capacity_dir: Path) -> RowData: + """Row 3: Spherical MNIST. Bars + data from main sweep, params from capacity.""" + main_runs = _load_runs(results_dir) + capacity_runs = _load_runs(capacity_dir) + + # Capacity runs use run_label 'model_h{width}'; recover the base model. + grouped_main = _group_runs(main_runs, SMNIST_METHODS, model_key='base_model') + for r in main_runs: # older results may lack base_model + if 'base_model' not in r and r.get('model') in SMNIST_METHODS: + method = r['model'] + if _canonical_mode(r.get('train_mode', 'C')) == _expected_mode(method): + grouped_main[method].append(r) + + aggregates: dict[str, Aggregate] = {} + data_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + + for method, runs in grouped_main.items(): + full = [r for r in runs if _is_full_train(r)] + subset = [r for r in runs if not _is_full_train(r)] + + agg = _aggregate(full, 'test_nr', 'accuracy') + if agg is not None: + aggregates[method] = agg + + by_size = defaultdict(list) + for r in subset: + by_size[r.get('train_examples')].append(r) + if full: + by_size[full[0].get('train_examples')] = full + for size, size_runs in by_size.items(): + size_agg = _aggregate(size_runs, 'test_nr', 'accuracy') + if size_agg is not None: + data_points[method].append((float(size), size_agg.ood_mean, size_agg.ood_std)) + + # Param curve: capacity sweep grouped by run label (one label per width). + param_points: dict[str, list[tuple[float, float, float]]] = defaultdict(list) + by_label = defaultdict(list) + for r in capacity_runs: + base = r.get('base_model') + if base in SMNIST_METHODS and _canonical_mode(r.get('train_mode', 'C')) == 'C': + by_label[(base, r.get('model'))].append(r) + for (base, _label), label_runs in by_label.items(): + agg = _aggregate(label_runs, 'test_nr', 'accuracy') + if agg is not None: + param_points[base].append((agg.n_params, agg.ood_mean, agg.ood_std)) + + # The Aug. CNN has a fixed architecture: single point from the main sweep. + if 'standard' in aggregates: + std = aggregates['standard'] + param_points['standard'].append((std.n_params, std.ood_mean, std.ood_std)) + + remap = {'standard': 'standard_s2'} + aggregates = {remap.get(k, k): v for k, v in aggregates.items()} + param_points = {remap.get(k, k): v for k, v in param_points.items()} + data_points = {remap.get(k, k): v for k, v in data_points.items()} + + bar = _build_bar_panel( + ROW_METHODS[2], + aggregates, + ylabel='Test accuracy', + ylim=(0.0, 1.08), + reference=COHEN_REFERENCE, + ) + params = _build_curve_panel( + ROW_METHODS[2], + param_points, + ylabel='Rotated test accuracy', + ylim=(0.1, 1.0), + reference=COHEN_REFERENCE, + ) + data = _build_curve_panel( + ROW_METHODS[2], + data_points, + ylabel='Rotated test accuracy', + ylim=(0.05, 1.0), + reference=COHEN_REFERENCE, + ) + return RowData(bar=bar, params=params, data=data) + + +# -------------------------------------------------------------------------- +# Mock data (layout preview only) +# -------------------------------------------------------------------------- + + +def _array(values: list[float]) -> FloatArray: + return np.asarray(values, dtype=np.float64) + + +def mock_grid_data() -> GridData: + """Deterministic mock numbers exercising the exact real-data layout.""" + pcam = RowData( + bar=BarPanel( + methods=ROW_METHODS[0], + original=_array([0.896, 0.942, 0.941, 0.945, 0.930, 0.941]), + ood=_array([0.861, 0.941, 0.940, 0.944, 0.929, 0.941]), + original_error=_array([0.010, 0.004, 0.009, 0.004, 0.007, 0.004]), + ood_error=_array([0.013, 0.004, 0.008, 0.004, 0.007, 0.004]), + ylabel='Test AUC', + ylim=(0.82, 0.97), + ), + params=CurvePanel( + series=( + CurveSeries( + 'standard', + _array([30_000, 102_000, 267_000, 582_000, 786_000]), + _array([0.880, 0.861, 0.870, 0.858, 0.862]), + _array([0.012, 0.013, 0.011, 0.013, 0.012]), + ), + CurveSeries( + 'norm', + _array([69_000, 110_000, 222_000, 372_000, 791_000]), + _array([0.927, 0.941, 0.924, 0.911, 0.901]), + _array([0.007, 0.004, 0.008, 0.011, 0.012]), + ), + CurveSeries( + 'gate', + _array([136_000, 218_000, 440_000, 741_000, 1_580_000]), + _array([0.940, 0.938, 0.935, 0.930, 0.924]), + _array([0.009, 0.008, 0.007, 0.009, 0.010]), + ), + CurveSeries( + 'fourier_elu', + _array([69_000, 110_000, 222_000, 372_000, 790_000]), + _array([0.935, 0.944, 0.943, 0.937, 0.926]), + _array([0.005, 0.004, 0.004, 0.006, 0.008]), + ), + CurveSeries( + 'norm_pool', + _array([69_000, 110_000, 222_000, 372_000, 790_000]), + _array([0.921, 0.929, 0.928, 0.922, 0.915]), + _array([0.008, 0.007, 0.006, 0.008, 0.009]), + ), + CurveSeries( + 'bispectrum', + _array([80_000, 128_000, 258_000, 433_000, 920_000]), + _array([0.934, 0.941, 0.943, 0.941, 0.943]), + _array([0.005, 0.004, 0.004, 0.004, 0.004]), + ), + ), + ylabel='Rotated test AUC', + ylim=(0.85, 0.96), + ), + data=CurvePanel( + series=( + CurveSeries( + 'standard', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.660, 0.735, 0.805, 0.861, 0.930]), + _array([0.032, 0.024, 0.019, 0.013, 0.007]), + ), + CurveSeries( + 'norm', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.720, 0.805, 0.873, 0.941, 0.953]), + _array([0.028, 0.020, 0.018, 0.004, 0.004]), + ), + CurveSeries( + 'gate', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.735, 0.820, 0.884, 0.940, 0.952]), + _array([0.027, 0.018, 0.015, 0.009, 0.005]), + ), + CurveSeries( + 'fourier_elu', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.715, 0.810, 0.876, 0.944, 0.954]), + _array([0.031, 0.019, 0.016, 0.004, 0.004]), + ), + CurveSeries( + 'norm_pool', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.700, 0.790, 0.860, 0.928, 0.940]), + _array([0.030, 0.021, 0.017, 0.006, 0.005]), + ), + CurveSeries( + 'bispectrum', + _array([100, 500, 2_500, 12_500, 262_144]), + _array([0.770, 0.850, 0.912, 0.941, 0.956]), + _array([0.025, 0.017, 0.025, 0.004, 0.004]), + ), + ), + ylabel='Rotated test AUC', + ylim=(0.6, 0.98), + xticks=(100, 500, 2_500, 12_500, 262_144), + xticklabels=('100', '500', '2.5K', '12.5K', 'full'), + ), + ) + + organ3d = RowData( + bar=BarPanel( + methods=ROW_METHODS[1], + original=_array([0.601, 0.568, 0.730, 0.726]), + ood=_array([0.576, 0.568, 0.730, 0.726]), + original_error=_array([0.017, 0.099, 0.033, 0.027]), + ood_error=_array([0.021, 0.099, 0.033, 0.027]), + ylabel='Test accuracy', + ylim=(0.0, 0.88), + ), + params=CurvePanel( + series=( + CurveSeries( + 'standard_3d', + _array([16_000, 60_000, 230_000]), + _array([0.576, 0.610, 0.640]), + _array([0.021, 0.020, 0.022]), + ), + CurveSeries( + 'norm_pool', + _array([375_000, 1_520_000, 6_100_000]), + _array([0.568, 0.590, 0.600]), + _array([0.099, 0.080, 0.075]), + ), + CurveSeries( + 'max_pool', + _array([374_000, 1_500_000, 6_000_000]), + _array([0.730, 0.743, 0.785]), + _array([0.033, 0.015, 0.032]), + ), + CurveSeries( + 'bispectrum', + _array([463_000, 1_700_000, 6_300_000]), + _array([0.726, 0.745, 0.685]), + _array([0.027, 0.006, 0.039]), + ), + ), + ylabel='Rotated test accuracy', + ylim=(0.3, 0.86), + ), + data=CurvePanel( + series=( + CurveSeries( + 'standard_3d', + _array([50, 100, 250, 500, 971]), + _array([0.115, 0.155, 0.320, 0.465, 0.576]), + _array([0.026, 0.031, 0.036, 0.031, 0.021]), + ), + CurveSeries( + 'norm_pool', + _array([50, 100, 250, 500, 971]), + _array([0.150, 0.230, 0.390, 0.500, 0.568]), + _array([0.060, 0.070, 0.080, 0.090, 0.099]), + ), + CurveSeries( + 'max_pool', + _array([50, 100, 250, 500, 971]), + _array([0.189, 0.280, 0.500, 0.640, 0.730]), + _array([0.035, 0.040, 0.045, 0.038, 0.033]), + ), + CurveSeries( + 'bispectrum', + _array([50, 100, 250, 500, 971]), + _array([0.332, 0.430, 0.580, 0.680, 0.726]), + _array([0.040, 0.038, 0.035, 0.030, 0.027]), + ), + ), + ylabel='Rotated test accuracy', + ylim=(0.05, 0.82), + xticks=(50, 100, 250, 500, 971), + xticklabels=('50', '100', '250', '500', 'full'), + ), + ) + + smnist = RowData( + bar=BarPanel( + methods=ROW_METHODS[2], + original=_array([0.460, 0.792, 0.950]), + ood=_array([0.230, 0.790, 0.951]), + original_error=_array([0.012, 0.010, 0.001]), + ood_error=_array([0.015, 0.010, 0.001]), + ylabel='Test accuracy', + ylim=(0.0, 1.08), + reference=COHEN_REFERENCE, + ), + params=CurvePanel( + series=( + CurveSeries('standard_s2', _array([185_000]), _array([0.230]), _array([0.015])), + CurveSeries( + 'power_spectrum', + _array([3_700, 11_500, 39_000, 144_000, 550_000]), + _array([0.740, 0.765, 0.777, 0.786, 0.792]), + _array([0.012, 0.008, 0.001, 0.007, 0.010]), + ), + CurveSeries( + 'bispectrum', + _array([25_000, 52_000, 108_000, 232_000, 529_000]), + _array([0.915, 0.932, 0.944, 0.951, 0.951]), + _array([0.007, 0.005, 0.003, 0.001, 0.002]), + ), + ), + ylabel='Rotated test accuracy', + ylim=(0.1, 1.0), + reference=COHEN_REFERENCE, + ), + data=CurvePanel( + series=( + CurveSeries( + 'standard_s2', + _array([100, 500, 2_500, 12_500, 60_000]), + _array([0.130, 0.155, 0.185, 0.212, 0.230]), + _array([0.008, 0.008, 0.009, 0.010, 0.015]), + ), + CurveSeries( + 'power_spectrum', + _array([100, 500, 2_500, 12_500, 60_000]), + _array([0.420, 0.580, 0.690, 0.750, 0.777]), + _array([0.018, 0.015, 0.012, 0.006, 0.001]), + ), + CurveSeries( + 'bispectrum', + _array([100, 500, 2_500, 12_500, 60_000]), + _array([0.670, 0.840, 0.920, 0.946, 0.951]), + _array([0.016, 0.012, 0.006, 0.002, 0.001]), + ), + ), + ylabel='Rotated test accuracy', + ylim=(0.05, 1.0), + xticks=(100, 500, 2_500, 12_500, 60_000), + xticklabels=('100', '500', '2.5K', '12.5K', 'full'), + reference=COHEN_REFERENCE, + ), + ) + + return (pcam, organ3d, smnist) + + +# -------------------------------------------------------------------------- +# Rendering +# -------------------------------------------------------------------------- + + +def configure_style() -> None: + """Apply the shared Illustrator-friendly matplotlib style.""" + plt.rcParams.update( + { + 'pdf.fonttype': 42, + 'ps.fonttype': 42, + 'font.family': 'sans-serif', + 'font.sans-serif': ['Helvetica', 'Arial', 'DejaVu Sans'], + 'font.size': 9, + 'axes.labelsize': 10, + 'axes.linewidth': 0.65, + 'xtick.labelsize': 8, + 'ytick.labelsize': 8, + 'xtick.major.width': 0.6, + 'ytick.major.width': 0.6, + 'xtick.major.size': 3, + 'ytick.major.size': 3, + 'lines.linewidth': 1.6, + 'figure.dpi': 150, + 'savefig.dpi': 300, + 'savefig.bbox': 'tight', + 'savefig.pad_inches': 0.04, + } + ) + + +def style_axes(ax: Axes) -> None: + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.set_axisbelow(True) + ax.yaxis.grid(True, color='#D9D9D9', linewidth=0.5, alpha=0.8) + ax.tick_params(axis='x', pad=2) + ax.tick_params(axis='y', pad=2) + + +def add_reference(ax: Axes, reference: ReferenceLine | None) -> None: + if reference is None: + return + ax.axhline( + reference.value, + color='#666666', + linestyle=reference.linestyle, + linewidth=1.0, + zorder=1, + ) + + +def draw_bar_panel(ax: Axes, panel: BarPanel) -> None: + """Draw one canonical-vs-rotated grouped bar panel.""" + x = np.arange(len(panel.methods), dtype=np.float64) + width = 0.35 + + for index, method in enumerate(panel.methods): + style = METHOD_STYLE[method] + ax.bar( + x[index] - width / 2, + panel.original[index], + width, + yerr=panel.original_error[index], + color=style.color, + edgecolor=style.color, + linewidth=0.8, + capsize=2, + error_kw={'elinewidth': 0.8, 'capthick': 0.8}, + zorder=3, + ) + ax.bar( + x[index] + width / 2, + panel.ood[index], + width, + yerr=panel.ood_error[index], + color=style.color, + edgecolor=style.color, + linewidth=0.8, + hatch='////', + alpha=0.42, + capsize=2, + error_kw={'elinewidth': 0.8, 'capthick': 0.8}, + zorder=3, + ) + + ax.set_xticks(x) + ax.set_xticklabels([METHOD_STYLE[method].label for method in panel.methods], rotation=24, ha='right') + ax.set_ylabel(panel.ylabel) + ax.set_ylim(panel.ylim) + if 'accuracy' in panel.ylabel.lower(): + ax.yaxis.set_major_formatter(FuncFormatter(lambda value, _position: f'{value:.0%}')) + add_reference(ax, panel.reference) + style_axes(ax) + + +def draw_curve_panel(ax: Axes, panel: CurvePanel, *, xlabel: str) -> None: + """Draw one parameter- or data-efficiency curve panel.""" + for series in panel.series: + style = METHOD_STYLE[series.method] + line_style = '-' if len(series.x) > 1 else 'none' + ax.errorbar( + series.x, + series.y, + yerr=series.error, + color=style.color, + marker=style.marker, + linestyle=line_style, + markersize=6.5 if style.marker == '*' else 4.5, + markeredgewidth=0.7, + markeredgecolor='white', + capsize=2, + elinewidth=0.7, + zorder=3, + ) + + ax.set_xscale('log') + ax.set_xlabel(xlabel) + ax.set_ylabel(panel.ylabel) + ax.set_ylim(panel.ylim) + if panel.xticks is not None: + ax.set_xticks(panel.xticks) + if panel.xticklabels is not None: + ax.set_xticklabels(panel.xticklabels) + if 'accuracy' in panel.ylabel.lower(): + ax.yaxis.set_major_formatter(FuncFormatter(lambda value, _position: f'{value:.0%}')) + add_reference(ax, panel.reference) + style_axes(ax) + + +def draw_panel(ax: Axes, row: RowData, column: int) -> None: + if column == 0: + draw_bar_panel(ax, row.bar) + elif column == 1: + draw_curve_panel(ax, row.params, xlabel='Trainable parameters') + else: + draw_curve_panel(ax, row.data, xlabel='Training examples') + + +def method_legend_handle(method: str) -> Line2D: + style = METHOD_STYLE[method] + return Line2D( + [0], + [0], + color=style.color, + marker=style.marker, + linewidth=1.6, + markersize=7 if style.marker == '*' else 5, + markeredgecolor='white', + markeredgewidth=0.7, + label=style.label, + ) + + +def row_legend_handles(row_index: int) -> list[LegendHandle]: + handles: list[LegendHandle] = [method_legend_handle(method) for method in ROW_METHODS[row_index]] + handles.extend( + [ + Patch(facecolor='#777777', edgecolor='#777777', label='Canonical test'), + Patch( + facecolor='#D8D8D8', + edgecolor='#777777', + hatch='////', + linewidth=0.8, + label='Rotated test (OOD)', + ), + ] + ) + if row_index == 2: + handles.append( + Line2D( + [0], + [0], + color='#666666', + linestyle='--', + linewidth=1.0, + label=COHEN_REFERENCE.label, + ) + ) + return handles + + +def save_panels(grid: GridData, output_dir: Path) -> None: + """Save title-free, legend-free vector panels plus per-row legends.""" + for row_index, row in enumerate(grid): + for column in range(3): + fig, ax = plt.subplots(figsize=PANEL_SIZE, constrained_layout=True) + draw_panel(ax, row, column) + fig.savefig(output_dir / PANEL_FILENAMES[row_index][column]) + plt.close(fig) + + handles = row_legend_handles(row_index) + height = 0.30 * len(handles) + 0.25 + fig = plt.figure(figsize=(2.45, height)) + fig.legend( + handles=handles, + labels=[handle.get_label() for handle in handles], + loc='center left', + bbox_to_anchor=(0.02, 0.5), + frameon=False, + handlelength=2.1, + handletextpad=0.8, + borderaxespad=0.0, + labelspacing=0.7, + fontsize=9, + ) + fig.savefig(output_dir / f'legend_row{row_index + 1}.pdf') + plt.close(fig) + + +def save_contact_sheet(grid: GridData, output_dir: Path, *, mock: bool) -> None: + """Save an assembled 3x3 sheet with column headers and side legends.""" + fig = plt.figure(figsize=(15.2, 9.6), constrained_layout=False) + gridspec = fig.add_gridspec( + 3, + 4, + width_ratios=(1.0, 1.0, 1.0, 0.58), + left=0.075, + right=0.985, + bottom=0.07, + top=0.91, + wspace=0.34, + hspace=0.42, + ) + + axes: list[list[Axes]] = [] + for row_index, row in enumerate(grid): + row_axes: list[Axes] = [] + for column in range(3): + ax = fig.add_subplot(gridspec[row_index, column]) + draw_panel(ax, row, column) + row_axes.append(ax) + axes.append(row_axes) + + legend_ax = fig.add_subplot(gridspec[row_index, 3]) + legend_ax.axis('off') + handles = row_legend_handles(row_index) + legend_ax.legend( + handles=handles, + labels=[handle.get_label() for handle in handles], + loc='center left', + frameon=False, + handlelength=2.0, + handletextpad=0.7, + labelspacing=0.55, + fontsize=8.5, + ) + + for column, label in enumerate(COLUMN_LABELS): + position = axes[0][column].get_position() + fig.text( + (position.x0 + position.x1) / 2, + 0.945, + label, + ha='center', + va='center', + fontsize=13, + fontweight='medium', + ) + + for row_index, label in enumerate(ROW_LABELS): + position = axes[row_index][0].get_position() + fig.text( + 0.018, + (position.y0 + position.y1) / 2, + label, + ha='center', + va='center', + rotation=90, + fontsize=11, + fontweight='medium', + ) + + if mock: + fig.text( + 0.985, + 0.015, + 'MOCK DATA — layout preview only', + ha='right', + va='bottom', + fontsize=8, + color='#777777', + ) + fig.savefig(output_dir / 'grid_contact_sheet.png', dpi=300, bbox_inches='tight') + fig.savefig(output_dir / 'grid_contact_sheet.pdf', bbox_inches='tight') + plt.close(fig) + + +# -------------------------------------------------------------------------- +# Entry point +# -------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description='Build the 3x3 experiment grid figure') + parser.add_argument( + '--mock', + action='store_true', + help='Render deterministic mock data instead of loading results.', + ) + parser.add_argument( + '--pcam_pareto_dir', + type=Path, + default=EXPERIMENTS_DIR / 'pcam' / 'pcam_results_pareto', + ) + parser.add_argument( + '--pcam_data_pareto_dir', + type=Path, + default=EXPERIMENTS_DIR / 'pcam' / 'pcam_results_data_pareto', + ) + parser.add_argument( + '--organ3d_dir', + type=Path, + default=EXPERIMENTS_DIR / 'organ3d' / 'organ3d_results', + ) + parser.add_argument( + '--smnist_dir', + type=Path, + default=EXPERIMENTS_DIR / 'spherical_mnist' / 'smnist_results', + ) + parser.add_argument( + '--smnist_capacity_dir', + type=Path, + default=EXPERIMENTS_DIR / 'spherical_mnist' / 'smnist_results_capacity', + ) + parser.add_argument( + '--output_dir', + type=Path, + default=None, + help='Defaults to ./grid_mockup for --mock, ./grid_figure otherwise.', + ) + args = parser.parse_args() + + output_dir: Path = args.output_dir or (EXPERIMENTS_DIR / ('grid_mockup' if args.mock else 'grid_figure')) + output_dir.mkdir(parents=True, exist_ok=True) + + configure_style() + + if args.mock: + grid = mock_grid_data() + else: + grid = ( + load_pcam(args.pcam_pareto_dir, args.pcam_data_pareto_dir), + load_organ3d(args.organ3d_dir), + load_smnist(args.smnist_dir, args.smnist_capacity_dir), + ) + + save_panels(grid, output_dir) + save_contact_sheet(grid, output_dir, mock=args.mock) + (output_dir / 'caption.txt').write_text(CAPTION_TEXT) + + print(f'Wrote grid assets to {output_dir}') + + +if __name__ == '__main__': + main() diff --git a/experiments/organ3d/README.md b/experiments/organ3d/README.md index a81ddbc..ee5dd9d 100644 --- a/experiments/organ3d/README.md +++ b/experiments/organ3d/README.md @@ -5,5 +5,7 @@ ```bash pip install -e "../../[dev]" python train.py --model bispectrum --data_dir ./organ3d_data -./run_sweep.sh # full sweep: 4 models x 3 seeds +./run_sweep.sh # main sweep: 4 models x 3 seeds (channels 4,8) +./run_wider_multiseed.sh # accuracy-vs-params curves (channels 8,16 and 16,32) +./run_dataeff_multiseed.sh # accuracy-vs-train-size curves ``` diff --git a/experiments/organ3d/analyze_results.py b/experiments/organ3d/analyze_results.py index afcad12..d77a734 100644 --- a/experiments/organ3d/analyze_results.py +++ b/experiments/organ3d/analyze_results.py @@ -23,7 +23,7 @@ MODEL_ORDER = ['standard', 'norm_pool', 'max_pool', 'bispectrum'] MODEL_LABELS = { - 'standard': 'Standard 3D CNN', + 'standard': 'Aug. 3D CNN', 'max_pool': 'O-Equiv + Max Pool', 'norm_pool': 'O-Equiv + Norm Pool', 'bispectrum': 'O-Equiv + Bispectrum', @@ -44,6 +44,11 @@ def _result_train_mode(record: dict) -> str: return _canonical_train_mode(record.get('train_mode', DEFAULT_TRAIN_MODE)) +def _expected_train_mode(model: str) -> str: + """Protocol mode per model: Aug. CNN baseline is R-trained, invariants C.""" + return 'R' if model == 'standard' else 'C' + + def _test_c(record: dict) -> dict: return record.get('test_c') or record.get('test') or {} @@ -147,9 +152,7 @@ def _baseline_runs( f = [ r for r in runs - if r.get('channels', [4, 8]) == [4, 8] - and _is_full_train(r) - and _result_train_mode(r) == train_mode + if r.get('channels', [4, 8]) == [4, 8] and _is_full_train(r) and _result_train_mode(r) == train_mode ] if f: filtered[model] = f @@ -162,10 +165,7 @@ def print_our_results(grouped: dict[str, list[dict]]): print('\n' + '=' * 90) print('OUR RESULTS (controlled ablation — same backbone, different pooling, C-trained)') print('=' * 90) - header = ( - f'{"Model":<25} {"Params":>8} {"Test ACC":>12} {"Test AUC":>12} ' - f'{"Rot ACC":>12} {"Rot \u03c3_ACC":>10}' - ) + header = f'{"Model":<25} {"Params":>8} {"Test ACC":>12} {"Test AUC":>12} {"Rot ACC":>12} {"Rot \u03c3_ACC":>10}' print(header) print('-' * 90) @@ -215,9 +215,7 @@ def print_published_baselines(): print('-' * 80) for b in PUBLISHED_BASELINES: auc_str = f'{b["auc"]:.3f}' if b['auc'] is not None else '—' - print( - f'{b["method"]:<30} {b["venue"]:<25} {b["params"]:>8} {b["acc"]:.3f} {auc_str:>8}' - ) + print(f'{b["method"]:<30} {b["venue"]:<25} {b["params"]:>8} {b["acc"]:.3f} {auc_str:>8}') def plot_rotation_comparison(grouped: dict[str, list[dict]], output_path: str): @@ -246,7 +244,7 @@ def plot_rotation_comparison(grouped: dict[str, list[dict]], output_path: str): ) PLOT_LABELS = { - 'standard': 'Standard', + 'standard': 'Aug. 3D CNN', 'max_pool': 'Max Pool', 'norm_pool': 'Norm Pool', 'bispectrum': 'Bispectrum', @@ -381,10 +379,20 @@ def plot_data_efficiency(grouped: dict[str, list[dict]], output_path: str): organ_full = 971 sizes = [100, 500, organ_full] - models_to_plot = ['standard', 'max_pool', 'bispectrum'] - colors = {'standard': '#888888', 'max_pool': '#2D6A9F', 'bispectrum': '#C44E52'} - markers = {'standard': 's', 'max_pool': 'D', 'bispectrum': 'o'} - labels = {'standard': 'Standard CNN', 'max_pool': 'Max Pool', 'bispectrum': 'Bispectrum'} + models_to_plot = ['standard', 'norm_pool', 'max_pool', 'bispectrum'] + colors = { + 'standard': '#888888', + 'norm_pool': '#5C6BC0', + 'max_pool': '#2D6A9F', + 'bispectrum': '#C44E52', + } + markers = {'standard': 's', 'norm_pool': 'v', 'max_pool': 'D', 'bispectrum': 'o'} + labels = { + 'standard': 'Aug. 3D CNN', + 'norm_pool': 'Norm Pool', + 'max_pool': 'Max Pool', + 'bispectrum': 'Bispectrum', + } fig, ax = plt.subplots(figsize=(5.5, 3.5)) @@ -398,7 +406,7 @@ def plot_data_efficiency(grouped: dict[str, list[dict]], output_path: str): for r in grouped.get(model, []) if _train_size_value(r) == target and r.get('channels', [4, 8]) == [4, 8] - and _result_train_mode(r) == 'C' + and _result_train_mode(r) == _expected_train_mode(model) ] if runs: vals = [_test_c(r).get('accuracy', 0.0) for r in runs] @@ -472,9 +480,7 @@ def print_tier1_summary(grouped: dict[str, list[dict]]): print('\n' + '=' * 90) print('DATA EFFICIENCY (channels [4,8], C-trained)') print('=' * 90) - print( - f'{"Model":<15} {"N":>6} {"Params":>10} {"Test ACC":>18} {"Test AUC":>18} {"Rot ACC":>10}' - ) + print(f'{"Model":<15} {"N":>6} {"Params":>10} {"Test ACC":>18} {"Test AUC":>18} {"Rot ACC":>10}') print('-' * 85) for model in models: @@ -485,7 +491,7 @@ def print_tier1_summary(grouped: dict[str, list[dict]]): for r in grouped.get(model, []) if _train_size_value(r) == target and r.get('channels', [4, 8]) == [4, 8] - and _result_train_mode(r) == 'C' + and _result_train_mode(r) == _expected_train_mode(model) ] if runs: accs = [_test_c(r).get('accuracy', 0.0) for r in runs] @@ -502,9 +508,7 @@ def print_tier1_summary(grouped: dict[str, list[dict]]): print('\n' + '=' * 90) print('WIDER CHANNELS (full training set)') print('=' * 90) - print( - f'{"Model":<15} {"Channels":>10} {"Params":>10} {"Test ACC":>18} {"Test AUC":>18} {"Rot ACC":>10}' - ) + print(f'{"Model":<15} {"Channels":>10} {"Params":>10} {"Test ACC":>18} {"Test AUC":>18} {"Rot ACC":>10}') print('-' * 85) for ch in [[4, 8], [8, 16], [16, 32]]: @@ -512,9 +516,7 @@ def print_tier1_summary(grouped: dict[str, list[dict]]): runs = [ r for r in grouped.get(model, []) - if r.get('channels', [4, 8]) == ch - and _is_full_train(r) - and _result_train_mode(r) == 'C' + if r.get('channels', [4, 8]) == ch and _is_full_train(r) and _result_train_mode(r) == 'C' ] if runs: accs = [_test_c(r).get('accuracy', 0.0) for r in runs] diff --git a/experiments/organ3d/run_dataeff_multiseed.sh b/experiments/organ3d/run_dataeff_multiseed.sh index 4632087..88781b9 100755 --- a/experiments/organ3d/run_dataeff_multiseed.sh +++ b/experiments/organ3d/run_dataeff_multiseed.sh @@ -1,6 +1,17 @@ #!/bin/bash -# Data efficiency: 3 seeds x 4 sample-count steps x 3 models -# Seed 42 already done — this runs seeds 123, 456 +# Data efficiency: 3 seeds x 4 sample-count steps x 4 models at channels (4,8). +# Full-training-set points come from run_sweep.sh. +# +# Protocol (consistent-comparison): +# - Invariant models (max_pool, norm_pool, bispectrum) are trained +# canonical (C). +# - The standard 3D CNN is trained with octahedral augmentation (R) — it is +# the "Aug. 3D CNN" baseline in the figures. +# - Rotation (OOD) evaluation is always on: the figure curves plot rotated +# test accuracy. +# +# Usage (run in tmux): +# ./run_dataeff_multiseed.sh set -euo pipefail @@ -15,6 +26,15 @@ export PYTHONUNBUFFERED=1 OUTPUT_DIR="./organ3d_results" COMMON="--patience 15 --epochs 100 --data_dir ./organ3d_data" +train_mode_for() { + local model=$1 + if [[ "$model" == "standard" ]]; then + echo "R" + else + echo "C" + fi +} + batch_size_for() { local model=$1 if [[ "$model" == "bispectrum" ]]; then echo 16 @@ -25,11 +45,13 @@ batch_size_for() { run_single() { local model=$1 seed=$2 size=$3 local channels="4 8" + local mode + mode=$(train_mode_for "$model") local size_tag="_n${size}" - local out_dir="${OUTPUT_DIR}/${model}_ch4_8_seed${seed}${size_tag}" + local out_dir="${OUTPUT_DIR}/${model}_${mode}_ch4_8_seed${seed}${size_tag}" if [[ -f "${out_dir}/results.json" ]]; then - echo "SKIP (already done): model=$model seed=$seed size=$size" + echo "SKIP (already done): model=$model mode=$mode seed=$seed size=$size" return 0 fi @@ -38,16 +60,16 @@ run_single() { echo "" echo "============================================================" - echo " model=$model seed=$seed size=$size bs=$bs $(date)" + echo " model=$model mode=$mode seed=$seed size=$size bs=$bs $(date)" echo "============================================================" - python train.py --model "$model" --channels $channels \ + python train.py --model "$model" --channels $channels --train_mode "$mode" \ --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" \ --train_size "$size" $COMMON } -for seed in 123 456; do +for seed in 42 123 456; do for size in 50 100 250 500; do - for model in standard max_pool bispectrum; do + for model in standard max_pool norm_pool bispectrum; do run_single "$model" "$seed" "$size" done done diff --git a/experiments/organ3d/run_sweep.sh b/experiments/organ3d/run_sweep.sh index 6b37186..5eff150 100755 --- a/experiments/organ3d/run_sweep.sh +++ b/experiments/organ3d/run_sweep.sh @@ -1,5 +1,13 @@ #!/bin/bash -# OrganMNIST3D sweep: run all 4 model variants x 3 seeds. +# OrganMNIST3D main sweep: all 4 model variants x 3 seeds at channels (4,8). +# +# Protocol (consistent-comparison): +# - Invariant models (max_pool, norm_pool, bispectrum) are trained +# canonical (C). +# - The standard 3D CNN is trained with random octahedral augmentation (R) +# — it is the "Aug. 3D CNN" baseline in the figures. +# - Rotation (OOD) evaluation is always on: the figures plot rotated test +# accuracy next to canonical test accuracy. # # Param counts (default channels 4,8): # standard: ~16K @@ -7,8 +15,6 @@ # norm_pool: ~375K # bispectrum: ~463K # -# Full sweep: 4 models x 3 seeds with rotation eval (~2 hours) -# # Usage (run in tmux): # ./run_sweep.sh @@ -26,6 +32,15 @@ MODELS=(standard max_pool norm_pool bispectrum) OUTPUT_DIR="./organ3d_results" COMMON="--patience 15 --epochs 100 --data_dir ./organ3d_data" +train_mode_for() { + local model=$1 + if [[ "$model" == "standard" ]]; then + echo "R" + else + echo "C" + fi +} + batch_size_for() { local model=$1 if [[ "$model" == "bispectrum" ]]; then @@ -38,21 +53,23 @@ batch_size_for() { } run_single() { - local model=$1 seed=$2 extra=${3:-} + local model=$1 seed=$2 local channels="4 8" - local out_dir="${OUTPUT_DIR}/${model}_ch4_8_seed${seed}" + local mode + mode=$(train_mode_for "$model") + local out_dir="${OUTPUT_DIR}/${model}_${mode}_ch4_8_seed${seed}" if [[ -f "${out_dir}/results.json" ]]; then - echo "SKIP (already done): model=$model seed=$seed" + echo "SKIP (already done): model=$model mode=$mode seed=$seed" return 0 fi local bs bs=$(batch_size_for "$model") echo "" echo "============================================================" - echo " model=$model seed=$seed bs=$bs $(date)" + echo " model=$model mode=$mode seed=$seed bs=$bs $(date)" echo "============================================================" - python train.py --model "$model" --channels $channels \ - --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" $COMMON $extra + python train.py --model "$model" --channels $channels --train_mode "$mode" \ + --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" $COMMON } echo "=== Full sweep: all seeds with rotation eval ===" diff --git a/experiments/organ3d/run_tier1_sweep.sh b/experiments/organ3d/run_tier1_sweep.sh deleted file mode 100755 index 880b539..0000000 --- a/experiments/organ3d/run_tier1_sweep.sh +++ /dev/null @@ -1,97 +0,0 @@ -#!/bin/bash -# Tier 1 experiments: data efficiency + wider channels -# -# Part A: Data efficiency — train at N ∈ {50, 100, 250, 500} examples -# (full set already done in the main sweep, ~971 samples) -# Models: standard, max_pool, bispectrum (skip norm_pool — it's unstable) -# Single seed (42) for speed; rotation eval on all. -# -# Part B: Wider channels — (8,16) and (16,32) -# Models: max_pool, bispectrum (the comparison that matters) -# Single seed (42); rotation eval on all. -# -# Usage (run in tmux): -# ./run_tier1_sweep.sh - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR" - -REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -source "$REPO_ROOT/.venv/bin/activate" -export PYTHONPATH="$REPO_ROOT/src:${PYTHONPATH:-}" -export PYTHONUNBUFFERED=1 - -OUTPUT_DIR="./organ3d_results" -COMMON="--patience 15 --epochs 100 --data_dir ./organ3d_data" - -batch_size_for() { - local model=$1 channels=$2 - local c1 - c1=$(echo "$channels" | awk '{print $NF}') - if [[ "$model" == "bispectrum" ]]; then - if (( c1 >= 32 )); then echo 4 - elif (( c1 >= 16 )); then echo 8 - else echo 16; fi - elif [[ "$model" == "standard" ]]; then - echo 64 - else - if (( c1 >= 32 )); then echo 8 - elif (( c1 >= 16 )); then echo 16 - else echo 32; fi - fi -} - -run_single() { - local model=$1 seed=$2 channels=$3 size=$4 - local ch_tag="${channels// /_}" - local size_arg=() - local size_tag="" - if [[ "$size" != "full" ]]; then - size_arg=(--train_size "$size") - size_tag="_n${size}" - fi - local out_dir="${OUTPUT_DIR}/${model}_ch${ch_tag}_seed${seed}${size_tag}" - - if [[ -f "${out_dir}/results.json" ]]; then - echo "SKIP (already done): model=$model seed=$seed channels=$channels size=$size" - return 0 - fi - - local bs - bs=$(batch_size_for "$model" "$channels") - - echo "" - echo "============================================================" - echo " model=$model seed=$seed channels=$channels size=$size bs=$bs $(date)" - echo "============================================================" - python train.py --model "$model" --channels $channels \ - --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" \ - "${size_arg[@]}" $COMMON -} - -echo "============================================================" -echo " PART A: Data efficiency sweep" -echo "============================================================" -for size in 50 100 250 500; do - for model in standard max_pool bispectrum; do - run_single "$model" 42 "4 8" "$size" - done -done - -echo "" -echo "============================================================" -echo " PART B: Wider channels sweep" -echo "============================================================" -for channels in "8 16" "16 32"; do - for model in max_pool bispectrum; do - run_single "$model" 42 "$channels" "full" - done -done - -echo "" -echo "============================================================" -echo " ALL DONE — $(date)" -echo " Results in $OUTPUT_DIR" -echo "============================================================" diff --git a/experiments/organ3d/run_wider_multiseed.sh b/experiments/organ3d/run_wider_multiseed.sh index cd162e5..2cb60cf 100755 --- a/experiments/organ3d/run_wider_multiseed.sh +++ b/experiments/organ3d/run_wider_multiseed.sh @@ -1,6 +1,18 @@ #!/bin/bash -# Wider channels multi-seed: seeds 123, 456 for (8,16) and (16,32). -# Seed 42 already done in run_tier1_sweep.sh Part B. +# Capacity (parameter-efficiency) sweep: 3 seeds x 3 channel widths x 4 models. +# Together with run_sweep.sh (channels 4,8) this builds the +# accuracy-vs-params curves for the grid figure. +# +# Protocol (consistent-comparison): +# - Invariant models (max_pool, norm_pool, bispectrum) are trained +# canonical (C). +# - The standard 3D CNN is trained with octahedral augmentation (R) — it is +# the "Aug. 3D CNN" baseline in the figures. +# - Rotation (OOD) evaluation is always on: the figure curves plot rotated +# test accuracy. +# +# Usage (run in tmux): +# ./run_wider_multiseed.sh set -euo pipefail @@ -15,6 +27,15 @@ export PYTHONUNBUFFERED=1 OUTPUT_DIR="./organ3d_results" COMMON="--patience 15 --epochs 100 --data_dir ./organ3d_data" +train_mode_for() { + local model=$1 + if [[ "$model" == "standard" ]]; then + echo "R" + else + echo "C" + fi +} + batch_size_for() { local model=$1 channels=$2 local c1 @@ -23,6 +44,8 @@ batch_size_for() { if (( c1 >= 32 )); then echo 4 elif (( c1 >= 16 )); then echo 8 else echo 16; fi + elif [[ "$model" == "standard" ]]; then + echo 64 else if (( c1 >= 32 )); then echo 8 elif (( c1 >= 16 )); then echo 16 @@ -33,10 +56,12 @@ batch_size_for() { run_single() { local model=$1 seed=$2 channels=$3 local ch_tag="${channels// /_}" - local out_dir="${OUTPUT_DIR}/${model}_ch${ch_tag}_seed${seed}" + local mode + mode=$(train_mode_for "$model") + local out_dir="${OUTPUT_DIR}/${model}_${mode}_ch${ch_tag}_seed${seed}" if [[ -f "${out_dir}/results.json" ]]; then - echo "SKIP (already done): model=$model seed=$seed channels=$channels" + echo "SKIP (already done): model=$model mode=$mode seed=$seed channels=$channels" return 0 fi @@ -45,16 +70,16 @@ run_single() { echo "" echo "============================================================" - echo " model=$model seed=$seed channels=$channels bs=$bs $(date)" + echo " model=$model mode=$mode seed=$seed channels=$channels bs=$bs $(date)" echo "============================================================" - python train.py --model "$model" --channels $channels \ + python train.py --model "$model" --channels $channels --train_mode "$mode" \ --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" \ $COMMON } -for seed in 123 456; do +for seed in 42 123 456; do for channels in "8 16" "16 32"; do - for model in max_pool bispectrum; do + for model in standard max_pool norm_pool bispectrum; do run_single "$model" "$seed" "$channels" done done diff --git a/experiments/pcam/README.md b/experiments/pcam/README.md index 001958d..22b4c12 100644 --- a/experiments/pcam/README.md +++ b/experiments/pcam/README.md @@ -1,9 +1,10 @@ # PatchCamelyon Classification -Binary classification on PatchCamelyon (96x96 histopathology patches) comparing bispectral pooling against norm, gate, and FourierELU nonlinearities in equivariant DenseNets. +Binary classification on PatchCamelyon (96x96 histopathology patches) comparing bispectral pooling against norm, gate, FourierELU, and norm-pool (power-spectrum-like) baselines in equivariant DenseNets. ```bash pip install -e "../../[dev]" python train.py --model bispectrum --group c8 --data_dir ./pcam_data -./run_matched_sweep.sh # full Pareto sweep: 5 models x 5 growth rates x 3 seeds +./run_matched_sweep.sh # Pareto sweep: 6 models x 5 growth rates x 3 seeds +./run_data_pareto_sweep.sh # AUC-vs-train-size curves at matched ~100K params ``` diff --git a/experiments/pcam/analyze_data_pareto.py b/experiments/pcam/analyze_data_pareto.py index 49e78fb..de74731 100644 --- a/experiments/pcam/analyze_data_pareto.py +++ b/experiments/pcam/analyze_data_pareto.py @@ -14,13 +14,14 @@ PCAM_FULL_TRAIN = 262_144 MODEL_STYLE: dict[str, dict] = { - 'standard': {'color': '#888888', 'marker': 's', 'label': 'Standard'}, + 'standard': {'color': '#888888', 'marker': 's', 'label': 'Aug. CNN'}, 'norm': {'color': '#2196F3', 'marker': 'o', 'label': 'NormReLU'}, 'gate': {'color': '#FF9800', 'marker': '^', 'label': 'Gated'}, 'fourier_elu': {'color': '#9C27B0', 'marker': 'D', 'label': 'Fourier-ELU'}, + 'norm_pool': {'color': '#5C6BC0', 'marker': 'v', 'label': 'Norm pool'}, 'bispectrum': {'color': '#E53935', 'marker': '*', 'label': 'Bispectrum'}, } -MODEL_ORDER = ['standard', 'norm', 'gate', 'fourier_elu', 'bispectrum'] +MODEL_ORDER = ['standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum'] DEFAULT_TRAIN_MODE = 'R' @@ -37,6 +38,11 @@ def _result_train_mode(record: dict) -> str: return _canonical_train_mode(record.get('train_mode', DEFAULT_TRAIN_MODE)) +def _expected_train_mode(model: str) -> str: + """Protocol mode per model: Aug. CNN baseline is R-trained, invariants C.""" + return 'R' if model == 'standard' else 'C' + + def _result_test_metrics(record: dict) -> dict: return record.get('test_c') or record.get('test') or {} @@ -67,12 +73,14 @@ def load_all_results(base_dir: Path) -> list[dict]: def aggregate( results: list[dict], - train_mode: str = DEFAULT_TRAIN_MODE, ) -> dict[str, dict[int, dict[str, float]]]: - """Group by (model, train_size) for *train_mode*; report mean/std test AUC.""" + """Group by (model, train_size); report mean/std test AUC. + + Only protocol-conforming runs are kept (Aug. CNN: R, invariants: C). + """ grouped: dict[tuple[str, int], list[float]] = defaultdict(list) for r in results: - if _result_train_mode(r) != train_mode: + if _result_train_mode(r) != _expected_train_mode(r.get('model', '')): continue key = (r['model'], _train_size_value(r)) grouped[key].append(_result_test_metrics(r).get('auc', 0.0)) @@ -210,10 +218,7 @@ def main() -> None: deltas = [agg[model][n]['mean'] - standard_data[n]['mean'] for n in sizes] if deltas: best_n = sizes[int(np.argmax(deltas))] - print( - f' {MODEL_STYLE[model]["label"]:25s} ' - f'best advantage at N={best_n}: {max(deltas):+.4f}' - ) + print(f' {MODEL_STYLE[model]["label"]:25s} best advantage at N={best_n}: {max(deltas):+.4f}') if __name__ == '__main__': diff --git a/experiments/pcam/analyze_pareto.py b/experiments/pcam/analyze_pareto.py index 1920231..1c2b7c4 100644 --- a/experiments/pcam/analyze_pareto.py +++ b/experiments/pcam/analyze_pareto.py @@ -15,13 +15,14 @@ DATA_PARETO_DIR = Path(__file__).parent / 'pcam_results_data_pareto' MODEL_STYLE = { - 'standard': {'color': '#888888', 'marker': 's', 'label': 'Standard'}, + 'standard': {'color': '#888888', 'marker': 's', 'label': 'Aug. CNN'}, 'norm': {'color': '#2196F3', 'marker': 'o', 'label': 'NormReLU'}, 'gate': {'color': '#FF9800', 'marker': '^', 'label': 'Gated'}, 'fourier_elu': {'color': '#9C27B0', 'marker': 'D', 'label': 'Fourier-ELU'}, + 'norm_pool': {'color': '#5C6BC0', 'marker': 'v', 'label': 'Norm pool'}, 'bispectrum': {'color': '#E53935', 'marker': '*', 'label': 'Bispectrum'}, } -MODEL_ORDER = ['standard', 'norm', 'gate', 'fourier_elu', 'bispectrum'] +MODEL_ORDER = ['standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum'] DEFAULT_TRAIN_MODE = 'R' @@ -42,6 +43,15 @@ def _result_train_mode(record: dict) -> str: return _canonical_train_mode(record.get('train_mode', DEFAULT_TRAIN_MODE)) +def _expected_train_mode(model: str) -> str: + """Protocol mode per model: Aug. CNN baseline is R-trained, invariants C.""" + return 'R' if model == 'standard' else 'C' + + +def _matches_protocol(record: dict) -> bool: + return _result_train_mode(record) == _expected_train_mode(record.get('model', '')) + + def _result_test_metrics(record: dict) -> dict: return record.get('test_c') or record.get('test') or {} @@ -58,13 +68,13 @@ def load_results(directory: Path) -> list[dict]: return results -def aggregate_1pct(results: list[dict], train_mode: str = DEFAULT_TRAIN_MODE) -> dict[str, dict]: - """Group 1% results by model (filtered to *train_mode*), compute mean +/- std.""" +def aggregate_1pct(results: list[dict]) -> dict[str, dict]: + """Group 1% results by model (protocol-conforming runs only), mean +/- std.""" from collections import defaultdict grouped = defaultdict(list) for r in results: - if _result_train_mode(r) != train_mode: + if not _matches_protocol(r): continue grouped[r['model']].append(r) agg = {} @@ -82,17 +92,17 @@ def aggregate_1pct(results: list[dict], train_mode: str = DEFAULT_TRAIN_MODE) -> def aggregate_pareto_multiseed( results: list[dict], - train_mode: str = DEFAULT_TRAIN_MODE, ) -> dict[str, list[dict]]: - """Group pareto results by (model, n_params) for *train_mode*, mean +/- std AUC. + """Group pareto results by (model, n_params), mean +/- std AUC. + Only protocol-conforming runs are kept (Aug. CNN: R, invariants: C). Returns model -> sorted list of {n_params, mean_auc, std_auc, n_seeds}. """ from collections import defaultdict grouped: dict[tuple[str, int], list[float]] = defaultdict(list) for r in results: - if _result_train_mode(r) != train_mode: + if not _matches_protocol(r): continue grouped[(r['model'], r['n_params'])].append(_result_test_metrics(r).get('auc', 0.0)) @@ -129,12 +139,11 @@ def _train_size_value(record: dict) -> int: return n if n > 0 else PCAM_FULL_TRAIN -def load_data_pareto( - train_mode: str = DEFAULT_TRAIN_MODE, -) -> dict[str, dict[int, dict[str, float]]]: +def load_data_pareto() -> dict[str, dict[int, dict[str, float]]]: """Load data-efficiency results with multi-seed aggregation. - Filters by *train_mode* (default `R`). Returns model -> {train_size: stats}. + Only protocol-conforming runs are kept (Aug. CNN: R, invariants: C). + Returns model -> {train_size: stats}. """ from collections import defaultdict @@ -144,7 +153,7 @@ def load_data_pareto( for run_dir in sorted(DATA_PARETO_DIR.rglob('results.json')): with open(run_dir) as f: r = json.load(f) - if _result_train_mode(r) != train_mode: + if not _matches_protocol(r): continue size = _train_size_value(r) raw[(r['model'], size)].append(_result_test_metrics(r).get('auc', 0.0)) @@ -230,9 +239,7 @@ def main() -> None: ax1.set_title('(a) AUC vs. parameter count (10% data)') _style_ax(ax1) ax1.yaxis.grid(True, linewidth=0.4, alpha=0.5, zorder=0) - ax1.legend( - frameon=True, fancybox=False, edgecolor='#cccccc', framealpha=0.95, loc='lower right' - ) + ax1.legend(frameon=True, fancybox=False, edgecolor='#cccccc', framealpha=0.95, loc='lower right') matched_data = [] for model in MODEL_ORDER: @@ -240,9 +247,7 @@ def main() -> None: if not pts: continue best_at_100k = min(pts, key=lambda x: abs(x['n_params'] - 100_000)) - matched_data.append( - (model, best_at_100k['n_params'], best_at_100k['mean_auc'], best_at_100k['std_auc']) - ) + matched_data.append((model, best_at_100k['n_params'], best_at_100k['mean_auc'], best_at_100k['std_auc'])) matched_data.sort(key=lambda x: x[2]) y_pos = np.arange(len(matched_data)) diff --git a/experiments/pcam/model.py b/experiments/pcam/model.py index 1b2ad33..45acb72 100644 --- a/experiments/pcam/model.py +++ b/experiments/pcam/model.py @@ -3,14 +3,15 @@ No escnn dependency — group-equivariant convolutions are implemented from scratch following Cohen & Welling (2016). -Six model variants: +Seven model variants: 1. ``standard`` — vanilla DenseNet + data augmentation (no group structure) 2. ``norm`` — equivariant DenseNet + NormReLU nonlinearity 3. ``gate`` — equivariant DenseNet + gated nonlinearity 4. ``fourier_elu``— equivariant DenseNet + FFT→ELU→IFFT nonlinearity - 5. ``bispectrum`` — equivariant DenseNet + bispectral invariant pooling - 6. ``so2_disk`` — SO2onDisk disk bispectrum on raw patches + MLP (no backbone) + 5. ``norm_pool`` — equivariant DenseNet + norm (power-spectrum-like) pooling + 6. ``bispectrum`` — equivariant DenseNet + bispectral invariant pooling + 7. ``so2_disk`` — SO2onDisk disk bispectrum on raw patches + MLP (no backbone) """ from __future__ import annotations @@ -125,9 +126,7 @@ def __init__( self.register_buffer('inverses', inverses) self.group_order = elements.shape[0] - self.weight = nn.Parameter( - torch.empty(out_channels, in_channels, kernel_size, kernel_size) - ) + self.weight = nn.Parameter(torch.empty(out_channels, in_channels, kernel_size, kernel_size)) nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) self.padding = padding @@ -174,9 +173,7 @@ def __init__( self.out_channels = out_channels self.kernel_size = kernel_size - self.weight = nn.Parameter( - torch.empty(out_channels, in_channels, self.group_order, kernel_size, kernel_size) - ) + self.weight = nn.Parameter(torch.empty(out_channels, in_channels, self.group_order, kernel_size, kernel_size)) nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) self.padding = padding @@ -398,9 +395,7 @@ def __init__( self.bn2 = EquivBatchNorm(inter) self.conv2 = GroupConv2d(inter, growth_rate * gate_factor, 3, group, padding='same') - self.nonlin2 = _make_nonlinearity( - nonlin_type, growth_rate * gate_factor, group_order, growth_rate - ) + self.nonlin2 = _make_nonlinearity(nonlin_type, growth_rate * gate_factor, group_order, growth_rate) else: self.bn1 = nn.BatchNorm2d(in_channels) self.conv1 = nn.Conv2d(in_channels, inter, 1, bias=False) @@ -428,9 +423,7 @@ def __init__( super().__init__() layers = [] for i in range(num_layers): - layers.append( - _DenseLayer(in_channels + i * growth_rate, growth_rate, group, nonlin_type) - ) + layers.append(_DenseLayer(in_channels + i * growth_rate, growth_rate, group, nonlin_type)) self.layers = nn.ModuleList(layers) def forward(self, x): @@ -457,9 +450,7 @@ def __init__( gate_factor = 2 if nonlin_type == 'gate' else 1 self.bn = EquivBatchNorm(in_channels) self.conv = GroupConv2d(in_channels, out_channels * gate_factor, 1, group, padding=0) - self.nonlin = _make_nonlinearity( - nonlin_type, out_channels * gate_factor, group_order, out_channels - ) + self.nonlin = _make_nonlinearity(nonlin_type, out_channels * gate_factor, group_order, out_channels) else: self.bn = nn.BatchNorm2d(in_channels) self.conv = nn.Conv2d(in_channels, out_channels, 1, bias=False) @@ -487,9 +478,11 @@ def _make_nonlinearity( return GatedNonlinearity(out_channels or in_channels // 2) elif nonlin_type == 'fourier_elu': return FourierELU(group_order) - elif nonlin_type == 'bispectrum': + elif nonlin_type in ('bispectrum', 'norm_pool'): # For intermediate layers, use ReLU (equivariant for regular repr). - # Bispectrum is applied only as the final invariant pool. + # The invariant map (bispectrum or group norm) is applied only as the + # final pool, making the two variants a matched-architecture ablation: + # complete (bispectrum) vs incomplete (norm/power-spectrum) invariant. return _RegularReLU() else: raise ValueError(f'Unknown nonlinearity: {nonlin_type}') @@ -506,7 +499,7 @@ def forward(self, x): return F.relu(x) -NonlinType = Literal['standard', 'norm', 'gate', 'fourier_elu', 'bispectrum'] +NonlinType = Literal['standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum'] class PCamDenseNet(nn.Module): @@ -514,7 +507,7 @@ class PCamDenseNet(nn.Module): Args: nonlin_type: One of ``"standard"``, ``"norm"``, ``"gate"``, - ``"fourier_elu"``, ``"bispectrum"``. + ``"fourier_elu"``, ``"norm_pool"``, ``"bispectrum"``. group: ``"c8"`` or ``"d4"``. Ignored when ``nonlin_type="standard"``. growth_rate: DenseNet growth rate *k*. block_config: Number of layers in each dense block. @@ -563,6 +556,8 @@ def __init__( if nonlin_type == 'bispectrum': self.invariant_pool = BispectrumPool(group, channels) + elif nonlin_type == 'norm_pool': + self.invariant_pool = GroupNormPool() elif self.equivariant: self.invariant_pool = GroupMaxPool() else: diff --git a/experiments/pcam/run_data_pareto_sweep.sh b/experiments/pcam/run_data_pareto_sweep.sh index e747204..3876212 100755 --- a/experiments/pcam/run_data_pareto_sweep.sh +++ b/experiments/pcam/run_data_pareto_sweep.sh @@ -2,20 +2,25 @@ # Data-efficiency Pareto sweep: run each model at matched ~100K params across # multiple training set fractions to build AUC-vs-data curves. # +# Protocol (consistent-comparison): +# - Invariant models are trained in canonical mode (C). +# - The standard CNN is trained with rotation augmentation (R) — it is the +# "Aug. CNN" baseline in the figures. +# - Rotation (OOD) evaluation is always on: the figure curves plot rotated +# test AUC. +# # Matched growth rates (from find_growth_rates.py, target ~100K params): # standard: gr=12 → 102K # norm: gr=4 → 110K # gate: gr=3 → 136K # fourier_elu: gr=4 → 110K +# norm_pool: gr=4 → ~110K (same backbone as fourier_elu, paramless pool) # bispectrum: gr=4 → 128K # so2_disk: bl=30 → ~100K (MLP auto-sized) # -# Phase A (36 runs, single seed, skip rotation): ~4-6 hours -# Phase B (72 runs, 2 more seeds, with rotation): ~10-15 hours -# # Usage (run in tmux): -# ./run_data_pareto_sweep.sh # Phase A -# ./run_data_pareto_sweep.sh --phase-b # Phase B +# ./run_data_pareto_sweep.sh # Phase A (seed 42) +# ./run_data_pareto_sweep.sh --phase-b # Phase B (seeds 123, 456) set -euo pipefail @@ -32,16 +37,26 @@ MODEL_GR[standard]=12 MODEL_GR[norm]=4 MODEL_GR[gate]=3 MODEL_GR[fourier_elu]=4 +MODEL_GR[norm_pool]=4 MODEL_GR[bispectrum]=4 -MODELS=(standard norm gate fourier_elu bispectrum) +MODELS=(standard norm gate fourier_elu norm_pool bispectrum) # Absolute training-set sizes; "full" maps to PCam's 262144 examples. SIZES=(100 500 2500 12500 full) -SO2_DISK_BL=10 +SO2_DISK_BL=30 BASE_OUTPUT_DIR="./pcam_results_data_pareto" COMMON="--patience 10 --epochs 50" +train_mode_for() { + local model=$1 + if [[ "$model" == "standard" ]]; then + echo "R" + else + echo "C" + fi +} + size_args() { local size=$1 if [[ "$size" == "full" ]]; then @@ -61,20 +76,23 @@ size_tag() { } batch_size_for() { - local model=$1 gr=$2 + local model=$1 case "$model" in standard) echo 1024 ;; norm) echo 128 ;; gate) echo 128 ;; fourier_elu) echo 64 ;; + norm_pool) echo 128 ;; bispectrum) echo 128 ;; *) echo 128 ;; esac } run_single() { - local model=$1 size=$2 seed=$3 extra=${4:-} + local model=$1 size=$2 seed=$3 local gr=${MODEL_GR[$model]} + local mode + mode=$(train_mode_for "$model") local tag tag=$(size_tag "$size") local output_dir="${BASE_OUTPUT_DIR}/${tag}" @@ -82,26 +100,26 @@ run_single() { if [[ "$size" != "full" ]]; then suffix="_n${size}" fi - local out_dir="${output_dir}/${model}_c8_gr${gr}_seed${seed}${suffix}" + local out_dir="${output_dir}/${model}_c8_gr${gr}_${mode}_seed${seed}${suffix}" if [[ -f "${out_dir}/results.json" ]]; then echo "SKIP (already done): model=$model size=$size seed=$seed" return 0 fi local bs - bs=$(batch_size_for "$model" "$gr") + bs=$(batch_size_for "$model") local size_arg size_arg=$(size_args "$size") echo "" echo "============================================================" - echo " model=$model gr=$gr size=$size seed=$seed bs=$bs $(date)" + echo " model=$model gr=$gr mode=$mode size=$size seed=$seed bs=$bs $(date)" echo "============================================================" - python train.py --model "$model" --growth_rate "$gr" \ + python train.py --model "$model" --growth_rate "$gr" --train_mode "$mode" \ --output_dir "$output_dir" --seed "$seed" --batch_size "$bs" \ - $size_arg $COMMON $extra + $size_arg $COMMON } run_so2_disk() { - local size=$1 seed=$2 extra=${3:-} + local size=$1 seed=$2 local tag tag=$(size_tag "$size") local output_dir="${BASE_OUTPUT_DIR}/${tag}" @@ -109,7 +127,7 @@ run_so2_disk() { if [[ "$size" != "full" ]]; then suffix="_n${size}" fi - local out_dir="${output_dir}/so2_disk_bl${SO2_DISK_BL}_seed${seed}${suffix}" + local out_dir="${output_dir}/so2_disk_bl${SO2_DISK_BL}_C_seed${seed}${suffix}" if [[ -f "${out_dir}/results.json" ]]; then echo "SKIP (already done): model=so2_disk size=$size seed=$seed" return 0 @@ -120,13 +138,13 @@ run_so2_disk() { echo "============================================================" echo " model=so2_disk bl=$SO2_DISK_BL size=$size seed=$seed $(date)" echo "============================================================" - python train.py --model so2_disk --bandlimit "$SO2_DISK_BL" \ + python train.py --model so2_disk --bandlimit "$SO2_DISK_BL" --train_mode C \ --output_dir "$output_dir" --seed "$seed" --batch_size 256 \ - $size_arg $COMMON $extra + $size_arg $COMMON } if [[ "${1:-}" == "--phase-b" ]]; then - echo "=== PHASE B: remaining seeds (123, 456) with rotation eval ===" + echo "=== PHASE B: remaining seeds (123, 456) ===" for size in "${SIZES[@]}"; do for seed in 123 456; do for model in "${MODELS[@]}"; do @@ -136,12 +154,12 @@ if [[ "${1:-}" == "--phase-b" ]]; then done done else - echo "=== PHASE A: single seed (42), skip rotation ===" + echo "=== PHASE A: seed 42 ===" for size in "${SIZES[@]}"; do for model in "${MODELS[@]}"; do - run_single "$model" "$size" 42 "--skip_rotation" + run_single "$model" "$size" 42 done - run_so2_disk "$size" 42 "--skip_rotation" + run_so2_disk "$size" 42 done fi diff --git a/experiments/pcam/run_matched_sweep.sh b/experiments/pcam/run_matched_sweep.sh index 4719a15..16333fe 100755 --- a/experiments/pcam/run_matched_sweep.sh +++ b/experiments/pcam/run_matched_sweep.sh @@ -1,17 +1,26 @@ #!/bin/bash # Pareto sweep: run each model at multiple growth_rates to build AUC-vs-params curves. # +# Protocol (consistent-comparison): +# - Invariant models (norm, gate, fourier_elu, norm_pool, bispectrum, so2_disk) +# are trained in canonical mode (C). +# - The standard CNN is trained with rotation augmentation (R) — it is the +# "Aug. CNN" baseline in the figures. +# - Rotation (OOD) evaluation is always on: the figure curves plot rotated +# test AUC, so every run must record test_r / rotation_robustness. +# # Param counts (from find_growth_rates.py): # # standard: gr=6→30K, gr=12→102K, gr=20→267K, gr=30→582K, gr=35→786K # norm: gr=3→69K, gr=4→110K, gr=6→222K, gr=8→372K, gr=12→791K # gate: gr=3→136K, gr=4→218K, gr=6→440K, gr=8→741K, gr=12→1.58M # fourier_elu: gr=3→69K, gr=4→110K, gr=6→222K, gr=8→372K, gr=12→790K +# norm_pool: ~same as fourier_elu (identical backbone, paramless pool) # bispectrum: gr=3→80K, gr=4→128K, gr=6→258K, gr=8→433K, gr=12→920K # so2_disk: all ~100K (MLP auto-sized), bandlimit controls feature quality # -# Phase A (25 CNN + N so2_disk runs, single seed, skip rotation): ~4-6 hours -# Phase B (remaining seeds with rotation): ~10-15 hours +# Phase A (seed 42, all configs): first full curves +# Phase B (seeds 123, 456): error bars # # Usage (run in tmux): # ./run_matched_sweep.sh # Phase A @@ -30,13 +39,16 @@ export PYTHONUNBUFFERED=1 STANDARD_GRS=(6 12 20 30 35) EQUIVARIANT_GRS=(3 4 6 8 12) SO2_DISK_BLS=(10 15 20 25 30 40 50) +EQUIVARIANT_MODELS=(norm gate fourier_elu norm_pool bispectrum) OUTPUT_DIR="./pcam_results_pareto" -COMMON="--train_size 12500 --patience 10 --epochs 50" +TRAIN_SIZE=12500 +COMMON="--train_size ${TRAIN_SIZE} --patience 10 --epochs 50" +N_TAG="_n${TRAIN_SIZE}" batch_size_for() { local model=$1 gr=$2 - if [[ "$model" == "fourier_elu" || "$model" == "bispectrum" || "$model" == "norm" ]]; then + if [[ "$model" == "fourier_elu" || "$model" == "bispectrum" || "$model" == "norm" || "$model" == "norm_pool" ]]; then if (( gr >= 8 )); then echo 64 else echo 128 fi @@ -48,25 +60,25 @@ batch_size_for() { } run_single() { - local model=$1 gr=$2 seed=$3 extra=${4:-} - local out_dir="${OUTPUT_DIR}/${model}_c8_gr${gr}_seed${seed}" + local model=$1 gr=$2 seed=$3 mode=${4:-C} + local out_dir="${OUTPUT_DIR}/${model}_c8_gr${gr}_${mode}_seed${seed}${N_TAG}" if [[ -f "${out_dir}/results.json" ]]; then - echo "SKIP (already done): model=$model gr=$gr seed=$seed" + echo "SKIP (already done): model=$model gr=$gr mode=$mode seed=$seed" return 0 fi local bs bs=$(batch_size_for "$model" "$gr") echo "" echo "============================================================" - echo " model=$model gr=$gr seed=$seed bs=$bs $(date)" + echo " model=$model gr=$gr mode=$mode seed=$seed bs=$bs $(date)" echo "============================================================" - python train.py --model "$model" --growth_rate "$gr" \ - --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" $COMMON $extra + python train.py --model "$model" --growth_rate "$gr" --train_mode "$mode" \ + --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size "$bs" $COMMON } run_so2_disk() { - local bl=$1 seed=$2 extra=${3:-} - local out_dir="${OUTPUT_DIR}/so2_disk_bl${bl}_seed${seed}" + local bl=$1 seed=$2 + local out_dir="${OUTPUT_DIR}/so2_disk_bl${bl}_C_seed${seed}${N_TAG}" if [[ -f "${out_dir}/results.json" ]]; then echo "SKIP (already done): model=so2_disk bl=$bl seed=$seed" return 0 @@ -75,38 +87,33 @@ run_so2_disk() { echo "============================================================" echo " model=so2_disk bl=$bl seed=$seed $(date)" echo "============================================================" - python train.py --model so2_disk --bandlimit "$bl" \ - --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size 256 $COMMON $extra + python train.py --model so2_disk --bandlimit "$bl" --train_mode C \ + --output_dir "$OUTPUT_DIR" --seed "$seed" --batch_size 256 $COMMON } -if [[ "${1:-}" == "--phase-b" ]]; then - echo "=== PHASE B: remaining seeds (123, 456) with rotation eval ===" - for seed in 123 456; do - for gr in "${STANDARD_GRS[@]}"; do - run_single standard "$gr" "$seed" - done - for model in norm gate fourier_elu bispectrum; do - for gr in "${EQUIVARIANT_GRS[@]}"; do - run_single "$model" "$gr" "$seed" - done - done - for bl in "${SO2_DISK_BLS[@]}"; do - run_so2_disk "$bl" "$seed" - done - done -else - echo "=== PHASE A: single seed (42), skip rotation ===" +run_all_for_seed() { + local seed=$1 for gr in "${STANDARD_GRS[@]}"; do - run_single standard "$gr" 42 "--skip_rotation" + run_single standard "$gr" "$seed" R done - for model in norm gate fourier_elu bispectrum; do + for model in "${EQUIVARIANT_MODELS[@]}"; do for gr in "${EQUIVARIANT_GRS[@]}"; do - run_single "$model" "$gr" 42 "--skip_rotation" + run_single "$model" "$gr" "$seed" C done done for bl in "${SO2_DISK_BLS[@]}"; do - run_so2_disk "$bl" 42 "--skip_rotation" + run_so2_disk "$bl" "$seed" done +} + +if [[ "${1:-}" == "--phase-b" ]]; then + echo "=== PHASE B: remaining seeds (123, 456) ===" + for seed in 123 456; do + run_all_for_seed "$seed" + done +else + echo "=== PHASE A: seed 42, all configs ===" + run_all_for_seed 42 fi echo "" diff --git a/experiments/pcam/train.py b/experiments/pcam/train.py index da67b95..31be0ca 100644 --- a/experiments/pcam/train.py +++ b/experiments/pcam/train.py @@ -18,7 +18,7 @@ # SO2onDisk disk bispectrum baseline (no backbone) python train.py --model so2_disk --bandlimit 30 --data_dir ./pcam_data - # Run all 6 baselines × 3 seeds (full sweep) + # Run all 6 model variants × 2 train modes × 3 seeds (full sweep) python train.py --sweep --data_dir ./pcam_data """ @@ -177,9 +177,7 @@ def evaluate_rotation_robustness( results['mean_auc'] = sum(aucs) / len(aucs) results['std_auc'] = (sum((a - results['mean_auc']) ** 2 for a in aucs) / len(aucs)) ** 0.5 results['mean_accuracy'] = sum(accs) / len(accs) - results['std_accuracy'] = ( - sum((a - results['mean_accuracy']) ** 2 for a in accs) / len(accs) - ) ** 0.5 + results['std_accuracy'] = (sum((a - results['mean_accuracy']) ** 2 for a in accs) / len(accs)) ** 0.5 return results @@ -378,9 +376,7 @@ def train(args: argparse.Namespace) -> dict: if args.model == 'so2_disk': print(f'Model: {args.model} (bandlimit={args.bandlimit}), {n_params:,} params') else: - print( - f'Model: {args.model} (group={args.group}, gr={args.growth_rate}), {n_params:,} params' - ) + print(f'Model: {args.model} (group={args.group}, gr={args.growth_rate}), {n_params:,} params') if args.dry_run: info: dict = { @@ -408,10 +404,7 @@ def train(args: argparse.Namespace) -> dict: seed=args.seed, subset_dir=args.subset_dir, ) - print( - f'Data: train={len(train_loader.dataset)}, ' - f'val={len(val_loader.dataset)}, test={len(test_loader.dataset)}' - ) + print(f'Data: train={len(train_loader.dataset)}, val={len(val_loader.dataset)}, test={len(test_loader.dataset)}') if args.compile: model = torch.compile(model) @@ -428,8 +421,7 @@ def train(args: argparse.Namespace) -> dict: n_tag = f'_n{args.train_size}' if args.train_size and args.train_size > 0 else '' if args.model == 'so2_disk': out_dir = ( - Path(args.output_dir) - / f'{args.model}_bl{args.bandlimit:.0f}_{args.train_mode}_seed{args.seed}{n_tag}' + Path(args.output_dir) / f'{args.model}_bl{args.bandlimit:.0f}_{args.train_mode}_seed{args.seed}{n_tag}' ) else: out_dir = ( @@ -486,10 +478,7 @@ def train(args: argparse.Namespace) -> dict: if not args.skip_rotation: rot_metrics = evaluate_rotation_robustness(model, test_loader, device, geometry_group) - print( - f'Rotation robustness: mean_auc={rot_metrics["mean_auc"]:.4f}, ' - f'std_auc={rot_metrics["std_auc"]:.4f}' - ) + print(f'Rotation robustness: mean_auc={rot_metrics["mean_auc"]:.4f}, std_auc={rot_metrics["std_auc"]:.4f}') test_r = _mean_group_metrics(rot_metrics) else: rot_metrics = {} @@ -524,7 +513,7 @@ def train(args: argparse.Namespace) -> dict: def run_sweep(args: argparse.Namespace): """Run the full experimental sweep.""" - models = ['standard', 'norm', 'gate', 'fourier_elu', 'bispectrum'] + models = ['standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum'] train_modes = ['C', 'R'] seeds = [42, 123, 456] all_results = [] @@ -536,10 +525,7 @@ def run_sweep(args: argparse.Namespace): args.train_mode = train_mode args.seed = seed print(f'\n{"=" * 60}') - print( - f'Running: model={model_name}, group={args.group}, ' - f'train_mode={train_mode}, seed={seed}' - ) + print(f'Running: model={model_name}, group={args.group}, train_mode={train_mode}, seed={seed}') print(f'{"=" * 60}') results = train(args) all_results.append(results) @@ -560,11 +546,7 @@ def run_sweep(args: argparse.Namespace): runs = grouped[(model_name, train_mode)] aucs = [r['test_c']['auc'] for r in runs] accs = [r['test_c']['accuracy'] for r in runs] - rot_stds = [ - r['rotation_robustness'].get('std_auc', 0.0) - for r in runs - if r.get('rotation_robustness') - ] + rot_stds = [r['rotation_robustness'].get('std_auc', 0.0) for r in runs if r.get('rotation_robustness')] n_params = runs[0]['n_params'] mean_auc = sum(aucs) / len(aucs) @@ -593,7 +575,7 @@ def main(): ) parser.add_argument( '--model', - choices=['standard', 'norm', 'gate', 'fourier_elu', 'bispectrum', 'so2_disk'], + choices=['standard', 'norm', 'gate', 'fourier_elu', 'norm_pool', 'bispectrum', 'so2_disk'], default='bispectrum', help='Nonlinearity / model variant.', ) @@ -648,7 +630,7 @@ def main(): parser.add_argument( '--sweep', action='store_true', - help='Run all 6 baselines × 3 seeds.', + help='Run all 6 model variants × 2 train modes × 3 seeds.', ) parser.add_argument( '--dry_run', diff --git a/experiments/spherical_mnist/.gitignore b/experiments/spherical_mnist/.gitignore index ec61eb5..0d4fbc0 100644 --- a/experiments/spherical_mnist/.gitignore +++ b/experiments/spherical_mnist/.gitignore @@ -1,6 +1,6 @@ smnist_data/ smnist_results/ -smnist_results_matched/ +smnist_results_capacity/ *.png *.pdf __pycache__/ diff --git a/experiments/spherical_mnist/README.md b/experiments/spherical_mnist/README.md index 87ec0ad..35a947b 100644 --- a/experiments/spherical_mnist/README.md +++ b/experiments/spherical_mnist/README.md @@ -5,5 +5,7 @@ MNIST digits projected onto S2, classified using SO(3)-invariant bispectral feat ```bash pip install -e "../../[dev]" python train.py --model bispectrum --train_mode NR -./run_sweep.sh # full sweep: 3 models x 2 modes x 3 seeds +./run_sweep.sh # full sweep: 3 models x 2 modes x 3 seeds +./run_capacity_sweep.sh # accuracy-vs-params curves (power_spectrum, bispectrum) +./run_data_efficiency.sh # accuracy-vs-train-size curves ``` diff --git a/experiments/spherical_mnist/analyze_results.py b/experiments/spherical_mnist/analyze_results.py index a2d4572..d6662db 100644 --- a/experiments/spherical_mnist/analyze_results.py +++ b/experiments/spherical_mnist/analyze_results.py @@ -10,7 +10,6 @@ Usage: python analyze_results.py [--results_dir ./smnist_results] - [--matched_dir ./smnist_results_matched] """ from __future__ import annotations @@ -27,18 +26,16 @@ import matplotlib.ticker as mticker import numpy as np -MODEL_ORDER = ['standard', 'power_spectrum', 'power_spectrum_matched', 'bispectrum'] +MODEL_ORDER = ['standard', 'power_spectrum', 'bispectrum'] MODEL_LABELS = { - 'standard': 'Std. CNN', + 'standard': 'Aug. CNN', 'power_spectrum': 'PowSpec', - 'power_spectrum_matched': 'PowSpec (matched)', 'bispectrum': 'Bispectrum', } MODEL_LABELS_FULL = { - 'standard': 'Standard CNN', + 'standard': 'Aug. CNN', 'power_spectrum': 'Power Spectrum', - 'power_spectrum_matched': 'Power Spectrum (matched)', 'bispectrum': 'Bispectrum (ours)', } @@ -128,8 +125,7 @@ def _is_full_train(record: dict) -> bool: COLORS = { 'standard': '#7f7f7f', - 'power_spectrum': '#aec7e8', - 'power_spectrum_matched': '#1f77b4', + 'power_spectrum': '#1f77b4', 'bispectrum': '#d62728', } @@ -142,13 +138,8 @@ def _is_full_train(record: dict) -> bool: def load_results( results_dir: str, - matched_dir: str | None = None, ) -> dict[tuple[str, str], list[dict]]: - """Load all results.json files, keyed by (model, train_mode). - - If *matched_dir* is provided, loads matched-param power spectrum runs - from that directory under the key ``power_spectrum_matched``. - """ + """Load all results.json files, keyed by (model, train_mode).""" grouped: dict[tuple[str, str], list[dict]] = defaultdict(list) results_path = Path(results_dir) if not results_path.exists(): @@ -162,15 +153,6 @@ def load_results( mode = _canonical_train_mode(r['train_mode']) grouped[(r['model'], mode)].append(r) - if matched_dir: - matched_path = Path(matched_dir) - if matched_path.exists(): - for p in sorted(matched_path.glob('*/results.json')): - with open(p) as f: - r = json.load(f) - mode = _canonical_train_mode(r['train_mode']) - grouped[('power_spectrum_matched', mode)].append(r) - return grouped @@ -200,9 +182,7 @@ def print_cohen_table(grouped: dict[tuple[str, str], list[dict]]): print('ACCURACY TABLE (Cohen et al. 2018 protocol, canonical/random)') print('Train/Test: X/Y = trained on X, evaluated on Y. Cohen reference uses NR for canonical.') print('=' * 100) - hdr = ( - f'{"Method":<32} {"Params":>10} {"C/C":>12} {"R/R":>12} {"C/R":>12} {"Rot \u03c3":>10}' - ) + hdr = f'{"Method":<32} {"Params":>10} {"C/C":>12} {"R/R":>12} {"C/R":>12} {"Rot \u03c3":>10}' print(hdr) print('-' * 100) @@ -234,11 +214,7 @@ def print_cohen_table(grouped: dict[tuple[str, str], list[dict]]): rot_str = f'{np.mean(rot_stds):.5f}' if rot_stds else '\u2014' label = _full_label(model, n_params) - print( - f'{label:<32} {n_params:>10,} ' - f'{_fmt(c_c):>12} {_fmt(r_r):>12} {_fmt(c_r):>12} ' - f'{rot_str:>10}' - ) + print(f'{label:<32} {n_params:>10,} {_fmt(c_c):>12} {_fmt(r_r):>12} {_fmt(c_r):>12} {rot_str:>10}') def print_latex_table(grouped: dict[tuple[str, str], list[dict]]): @@ -248,11 +224,7 @@ def print_latex_table(grouped: dict[tuple[str, str], list[dict]]): print(r'% Model & Params & C/C & R/R & C/R & Rot.\ $\sigma$ \\') for method, accs in COHEN_RESULTS.items(): - print( - f' {method} & --- ' - f'& ${accs["C/C"]:.2f}$ & ${accs["R/R"]:.2f}$ ' - f'& ${accs["C/R"]:.2f}$ & --- \\\\' - ) + print(f' {method} & --- & ${accs["C/C"]:.2f}$ & ${accs["R/R"]:.2f}$ & ${accs["C/R"]:.2f}$ & --- \\\\') print(r' \midrule') for model in MODEL_ORDER: @@ -424,16 +396,10 @@ def print_data_efficiency_table( c_accs = [_test_c(r)['accuracy'] for r in runs] r_accs = [_test_r(r)['accuracy'] for r in runs if _test_r(r)] n = len(runs) - c_str = ( - f'{np.mean(c_accs):.4f}\u00b1{np.std(c_accs):.4f}' - if n > 1 - else f'{c_accs[0]:.4f}' - ) + c_str = f'{np.mean(c_accs):.4f}\u00b1{np.std(c_accs):.4f}' if n > 1 else f'{c_accs[0]:.4f}' if r_accs: r_str = ( - f'{np.mean(r_accs):.4f}\u00b1{np.std(r_accs):.4f}' - if len(r_accs) > 1 - else f'{r_accs[0]:.4f}' + f'{np.mean(r_accs):.4f}\u00b1{np.std(r_accs):.4f}' if len(r_accs) > 1 else f'{r_accs[0]:.4f}' ) else: r_str = '\u2014' @@ -695,7 +661,6 @@ def main(): description='Analyze Spherical MNIST results', ) parser.add_argument('--results_dir', type=str, default='./smnist_results') - parser.add_argument('--matched_dir', type=str, default='./smnist_results_matched') parser.add_argument( '--output', type=str, @@ -703,7 +668,7 @@ def main(): ) args = parser.parse_args() - grouped = load_results(args.results_dir, args.matched_dir) + grouped = load_results(args.results_dir) if not grouped: print('No results found. Run the sweep first.') diff --git a/experiments/spherical_mnist/run_capacity_sweep.sh b/experiments/spherical_mnist/run_capacity_sweep.sh new file mode 100755 index 0000000..699e966 --- /dev/null +++ b/experiments/spherical_mnist/run_capacity_sweep.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# Capacity (parameter-efficiency) sweep: power_spectrum and bispectrum at +# multiple MLP widths x 3 seeds, C-trained, always evaluated on the rotated +# test set. Builds the accuracy-vs-params curves for the grid figure. +# +# Param counts (lmax=15): +# power_spectrum: h=128->3.7K h=256->11K h=512->39K h=1024->144K h=2048->550K +# bispectrum: h=32->25K h=64->52K h=128->108K h=256->232K h=512->529K +# +# The largest power-spectrum width (550K) matches the largest bispectrum +# budget (529K), so the incomplete-vs-complete invariant comparison is +# capacity-controlled along the whole curve. The standard CNN has a fixed +# architecture (185K) and appears as a single point from the main sweep. +# +# Usage (run in tmux): +# ./run_capacity_sweep.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +source "$REPO_ROOT/.venv/bin/activate" +export PYTHONPATH="$REPO_ROOT/src:${PYTHONPATH:-}" +export PYTHONUNBUFFERED=1 + +OUTPUT_DIR="./smnist_results_capacity" +COMMON="--epochs 50 --patience 10 --output_dir ${OUTPUT_DIR}" + +PS_HIDDENS=(128 256 512 1024 2048) +BSP_HIDDENS=(32 64 128 256 512) +SEEDS=(42 123 456) + +run_single() { + local model=$1 hidden=$2 seed=$3 + local label="${model}_h${hidden}" + local out_dir="${OUTPUT_DIR}/${label}_C_seed${seed}" + if [[ -f "${out_dir}/results.json" ]]; then + echo "SKIP (already done): model=$model hidden=$hidden seed=$seed" + return 0 + fi + echo "" + echo "============================================================" + echo " model=$model hidden=$hidden seed=$seed $(date)" + echo "============================================================" + python train.py --model "$model" --hidden "$hidden" \ + --run_label "$label" --train_mode C --seed "$seed" $COMMON +} + +for seed in "${SEEDS[@]}"; do + for hidden in "${PS_HIDDENS[@]}"; do + run_single power_spectrum "$hidden" "$seed" + done + for hidden in "${BSP_HIDDENS[@]}"; do + run_single bispectrum "$hidden" "$seed" + done +done + +echo "" +echo "============================================================" +echo " ALL DONE — $(date)" +echo " Results in $OUTPUT_DIR" +echo "============================================================" diff --git a/experiments/spherical_mnist/run_data_efficiency.sh b/experiments/spherical_mnist/run_data_efficiency.sh index a954c9e..9d8eee1 100755 --- a/experiments/spherical_mnist/run_data_efficiency.sh +++ b/experiments/spherical_mnist/run_data_efficiency.sh @@ -1,6 +1,15 @@ #!/bin/bash -# Data efficiency sweep: 3 models x 2 modes x 3 seeds x 2 sample-count steps = 36 runs. -# Full-set already done in run_sweep.sh. +# Data efficiency sweep for the grid figure's data-efficiency curves. +# +# Protocol (consistent-comparison): +# - Invariant models (power_spectrum, bispectrum) are trained canonical (C). +# - The standard CNN is trained on SO(3)-rotated data (R) — it is the +# "Aug. CNN" baseline in the figures. +# - test_r (rotated test set) is always recorded; the figure curves plot +# rotated test accuracy. --skip_rotation only skips the extra +# per-rotation robustness stats, which the curves do not need. +# +# Full-training-set points come from run_sweep.sh. # # Usage (run in tmux): # ./run_data_efficiency.sh @@ -15,14 +24,23 @@ source "$REPO_ROOT/.venv/bin/activate" export PYTHONPATH="$REPO_ROOT/src:${PYTHONPATH:-}" export PYTHONUNBUFFERED=1 -MODELS=(standard power_spectrum bispectrum) -TRAIN_MODES=(NR R) -SIZES=(500 6000) +SIZES=(100 500 2500 12500) OUTPUT_DIR="./smnist_results" COMMON="--patience 10 --epochs 50" +train_mode_for() { + local model=$1 + if [[ "$model" == "standard" ]]; then + echo "R" + else + echo "C" + fi +} + run_single() { - local model=$1 mode=$2 seed=$3 size=$4 + local model=$1 seed=$2 size=$3 + local mode + mode=$(train_mode_for "$model") local out_dir="${OUTPUT_DIR}/${model}_${mode}_seed${seed}_n${size}" if [[ -f "${out_dir}/results.json" ]]; then echo "SKIP (already done): model=$model mode=$mode seed=$seed size=$size" @@ -41,10 +59,8 @@ run_single() { for seed in 42 123 456; do for size in "${SIZES[@]}"; do - for mode in "${TRAIN_MODES[@]}"; do - for model in "${MODELS[@]}"; do - run_single "$model" "$mode" "$seed" "$size" - done + for model in standard power_spectrum bispectrum; do + run_single "$model" "$seed" "$size" done done done From f7a2cad074909baead521a4602e1dd6441785ef9 Mon Sep 17 00:00:00 2001 From: Johan Mathe Date: Thu, 16 Jul 2026 22:58:54 +0000 Subject: [PATCH 2/3] Document Python 3.12 + uv setup for GPU machines (old system pip breaks hatchling builds) Co-authored-by: Cursor --- experiments/README.md | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/experiments/README.md b/experiments/README.md index e524861..344cef8 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -11,10 +11,24 @@ Three benchmark experiments comparing G-bispectrum pooling against baseline inva ## Setup (on the GPU machine) +Requires Python 3.12 (`requires-python = ">=3.12,<3.13"`). The simplest path is `uv`, which downloads 3.12 if the machine only has an older system Python: + ```bash +curl -LsSf https://astral.sh/uv/install.sh | sh # skip if uv is installed +source ~/.local/bin/env + git clone bispectrum && cd bispectrum git checkout -python -m venv .venv && source .venv/bin/activate +uv venv --python 3.12 +source .venv/bin/activate +uv pip install -e ".[dev,experiments]" +``` + +Without uv: create the venv with a real `python3.12` binary and upgrade pip **before** installing (stock Ubuntu pip leaks an old `packaging` into the build env and fails with `No module named 'packaging.licenses'`): + +```bash +python3.12 -m venv .venv && source .venv/bin/activate +pip install -U pip pip install -e ".[dev,experiments]" ``` From a9dd55111c63d7b7b3b8ed284ae6236abaddb832 Mon Sep 17 00:00:00 2001 From: Johan Mathe Date: Fri, 17 Jul 2026 01:02:04 +0000 Subject: [PATCH 3/3] Document CUDA wheel/driver mismatch fix (cu130 wheels need r580+ drivers) Co-authored-by: Cursor --- experiments/README.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/experiments/README.md b/experiments/README.md index 344cef8..b0288f0 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -32,6 +32,15 @@ pip install -U pip pip install -e ".[dev,experiments]" ``` +**CUDA check.** PyPI's default torch wheels are built against CUDA 13.0, which requires NVIDIA driver r580+. On older drivers (e.g. 575.x = CUDA 12.9) torch prints a "driver too old" warning, reports `cuda.is_available() == False`, and the sweeps silently run on CPU. Fix by swapping in the CUDA 12.8 wheels (driver ≥ 570, same torch version so the `torch-harmonics` ABI pin holds): + +```bash +uv pip install --force-reinstall "torch==2.11.0" torchvision --index-url https://download.pytorch.org/whl/cu128 +python -c "import torch; assert torch.cuda.is_available(), 'CUDA not available'; print(torch.cuda.get_device_name(0))" +``` + +Run the assert line before launching tmux — do not start sweeps on a machine where it fails. + Datasets download automatically on first use (PCam from Zenodo ≈ 8 GB, OrganMNIST3D via `medmnist`, MNIST via `torchvision`). ## Run matrix