From c6ee1a4d882fe2a40b4144f5c3873793f3334197 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 13 Jul 2026 04:31:55 +0000 Subject: [PATCH 1/2] examples: Hubble FITS denoising (train in place, no reshard) A cross-domain example on an archival format that never was zarr. Real Hubble WFC3/IR frames of M16 on MAST's public S3 bucket are indexed as virtual references (VirtualiZarr -> kerchunk.fits -> Icechunk) and streamed by insitubatch with no reshard: sample_axis=0 makes each frame one sample, a NaN-scrubbing + robust-normalizing chunk stage and a block-mean Coarsen run on the decode pool, and per-sample Gaussian noise lives in the batch stage. A tiny residual CNN denoiser beats a median-filter baseline (~26.4 vs ~24.7 dB PSNR). The build-time index stack (virtualizarr/kerchunk/astropy) is needed only to build the store; the train hot path is icechunk + numpy. Ships an offline synthetic source (sharp point sources a median filter can't keep) as the default -- no network or FITS stack -- which also backs tests/test_hubble.py, a drift guard on the insitubatch-facing wiring. Co-Authored-By: Claude Opus 4.8 --- .gitignore | 3 + examples/README.md | 35 +++ examples/hubble/__init__.py | 6 + examples/hubble/data.py | 305 +++++++++++++++++++++++++++ examples/hubble/train_torch.py | 93 ++++++++ examples/hubble/wfc3ir_m16_uris.json | 14 ++ tests/test_hubble.py | 73 +++++++ 7 files changed, 529 insertions(+) create mode 100644 examples/hubble/__init__.py create mode 100644 examples/hubble/data.py create mode 100644 examples/hubble/train_torch.py create mode 100644 examples/hubble/wfc3ir_m16_uris.json create mode 100644 tests/test_hubble.py diff --git a/.gitignore b/.gitignore index af54dfc..8932772 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,9 @@ wheels/ bench/results/ bench/figures/ +# Generated example stores (virtual-reference Icechunk repos, rebuilt from committed URI lists) +examples/hubble/hubble_store/ + # MkDocs build output (docs/figures/ is committed for the published site) site/ diff --git a/examples/README.md b/examples/README.md index 939a34e..47754e4 100644 --- a/examples/README.md +++ b/examples/README.md @@ -105,6 +105,41 @@ prints the held-out foreground IoU, model vs Otsu. Only `train_torch.py` ships here — framework-neutrality is the advection example's job; this example's job is to prove the *geometry* generalizes. +## hubble/ — denoising real telescope frames from FITS (no reshard) + +The **archival-format** showcase: the data never was zarr. [`hubble/`](hubble/data.py) indexes +real Hubble WFC3/IR frames of **M16 (the Eagle Nebula)** on MAST's public AWS bucket +(`s3://stpubdata`, anonymous) as **virtual references** — [VirtualiZarr](https://github.com/zarr-developers/VirtualiZarr) +parses each `_flt.fits` header (`kerchunk.fits`) and commits byte-range references to a local +Icechunk repo. No pixels are copied or resharded; the store is a few kB pointing at the original +FITS objects. insitubatch then streams the frames straight from S3, `sample_axis=0` making each +frame one sample. The build-time index libraries (`virtualizarr`/`kerchunk`/`astropy`) are needed +*only* to build the store — the training hot path is `icechunk` + numpy, never kerchunk. + +```bash +uv sync --extra torch +uv run --with scipy python -m examples.hubble.train_torch # offline synthetic frames (default) + +# real Hubble frames on S3 -- indexes them into a virtual-reference store first (needs the +# build-time stack), then streams and trains: +uv run --with virtualizarr --with kerchunk --with astropy --with icechunk --with s3fs --with scipy \ + python -m examples.hubble.train_torch --source hubble --build +``` + +The task is per-frame **Gaussian-noise removal** (a deliberately didactic stand-in — the point is +training on the real archive, not a SOTA denoiser), and the baseline is a **median filter** (the +no-training reference). Two chunk stages run vectorized on the decode pool — a robust +per-frame normalization (`clean_normalize`) then a block-mean `Coarsen` (1014→253, keeping the CPU +demo light while still reading whole frames) — and the per-sample noise lives in the `AddNoise` +batch stage, per the transform-cost contract. A short run beats the baseline (real ≈26.4 vs ≈24.7 dB +PSNR; synthetic sharp-star frames similar). Because a FITS image is one chunk, this is the +*streaming-in-place* value (no reshard over a giant archive), not the many-samples-per-chunk +decode-amortization of the chunked-zarr examples. `--source synthetic` (the default) needs no +network or FITS stack — it also backs the drift test in `tests/test_hubble.py`. + +> Note: MAST's *anonymous* bucket throttles (HTTP 503) under heavy concurrent read-ahead, so +> `max_inflight` is capped low by default; an authenticated/retrying path on AWS would raise it. + ## The WeatherBench2 cold-start pair (with xbatcher) The same task two ways — complementary engines for the same ndim-batch problem — so you diff --git a/examples/hubble/__init__.py b/examples/hubble/__init__.py new file mode 100644 index 0000000..02f1ae8 --- /dev/null +++ b/examples/hubble/__init__.py @@ -0,0 +1,6 @@ +"""Train-in-place denoising over real Hubble (WFC3/IR) frames. + +A worked cross-domain example: FITS images on MAST's public AWS bucket, indexed as +virtual references (VirtualiZarr -> Icechunk) and streamed by insitubatch with no +reshard. See ``data.py`` for the pipeline and ``train_torch.py`` for the model. +""" diff --git a/examples/hubble/data.py b/examples/hubble/data.py new file mode 100644 index 0000000..2dbc5d0 --- /dev/null +++ b/examples/hubble/data.py @@ -0,0 +1,305 @@ +"""Data layer for the Hubble denoising example: real WFC3/IR frames, streamed in place. + +Pipeline +-------- +1. **Index (build time, once).** Each Hubble ``_flt.fits`` on MAST's public S3 bucket + holds one ``SCI`` image ``(1014, 1014) float32``. :func:`build_store` opens each as a + *virtual* dataset (VirtualiZarr -> ``kerchunk.fits``), concatenates them along a new + ``frame`` axis, and commits the byte-range references to a local Icechunk repo. No pixels + are copied -- the store is a few kB of references pointing at the original FITS objects. + ``kerchunk``/``virtualizarr``/``astropy`` are needed *only here*, never at train time. + +2. **Stream (train time).** :func:`open_store` reopens the Icechunk repo (resolving the + virtual chunks straight from S3), and :func:`denoise_dataset` builds an + :class:`~insitubatch.InSituDataset`: ``sample_axis=0`` makes each frame one sample (one + image = one chunk). Two chunk stages -- :func:`clean_normalize` then + :class:`~examples.transforms.Coarsen` -- run vectorized on the decode pool; the per-sample + random noise lives in the :class:`AddNoise` batch stage, per the transform-cost contract. + +The ML task is deliberately simple (Gaussian-noise removal, a didactic stand-in) -- the point +is that we train on the real archive with no reshard, not that this is a SOTA denoiser. +""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import zarr +from zarr.abc.store import Store + +from examples.transforms import Coarsen +from insitubatch import ( + Batch, + ChunkTransform, + DecodedChunk, + InSituDataset, + ensure_local_dir, + obstore_store, + open_geometries, + split_by_chunk, +) + +SCI_VAR = "SCI" +BUCKET = "stpubdata" +S3_PREFIX = "s3://stpubdata" +REGION = "us-east-1" +COARSEN = 4 # 1014 -> 253 on each axis: keeps the CPU demo light, still whole-frame reads +NOISE_SIGMA = 0.6 # in normalized (robust-standardized) units + +HERE = Path(__file__).resolve().parent +DEFAULT_URIS = HERE / "wfc3ir_m16_uris.json" +DEFAULT_STORE = HERE / "hubble_store" + + +# --------------------------------------------------------------------------- offline synthetic +def make_synthetic_store(url: str, *, n_frames: int = 12, size: int = 128, seed: int = 0) -> None: + """Write a synthetic ``SCI (frame, y, x)`` zarr mimicking the FITS-derived store's geometry. + + For offline runs (``--source synthetic``) and drift tests -- no network, no FITS/VZ stack. + One frame per chunk (``chunks=(1, size, size)``), like the real ``_flt.fits`` layout. Each + frame holds sharp point sources (stars) on a bright pedestal (a wide e-/s-like dynamic range) + with a handful of NaN bad pixels, so :func:`clean_normalize`'s NaN-scrub and robust scaling are + exercised as on real WFC3/IR data. The sources are sharp on purpose: a median filter blurs + them, so a CNN that preserves point sources beats the baseline (as it does on real stars). + """ + rng = np.random.default_rng(seed) + gy, gx = np.mgrid[0:size, 0:size].astype(np.float64) + frames = np.empty((n_frames, size, size), dtype="f4") + for i in range(n_frames): + img = np.full((size, size), rng.uniform(20.0, 80.0)) # bright background pedestal + for _ in range(20): # sharp stars: a 3x3 median blurs/erases them, the CNN keeps them + cy, cx = rng.uniform(0, size, size=2) + r = rng.uniform(1.0, 3.0) + amp = rng.uniform(80.0, 600.0) + img += amp * np.exp(-((gy - cy) ** 2 + (gx - cx) ** 2) / (2 * r * r)) + bad = (rng.integers(0, size, 5), rng.integers(0, size, 5)) + img[bad] = np.nan # a few bad pixels, as WFC3/IR flt frames carry + frames[i] = img.astype("f4") + + ensure_local_dir(url) + group = zarr.open_group(store=obstore_store(url, read_only=False), mode="w") + arr = group.create_array( + SCI_VAR, + shape=frames.shape, + chunks=(1, size, size), + dtype="f4", + dimension_names=("frame", "y", "x"), + ) + arr[:] = frames + + +# --------------------------------------------------------------------------- build time +def build_store( + uris: list[str], + store_path: str | os.PathLike[str] = DEFAULT_STORE, + *, + region: str = REGION, +) -> str: + """Index ``uris`` (Hubble ``_flt.fits`` on S3) into a local Icechunk repo of virtual refs. + + Idempotent: rebuilds ``store_path`` from scratch. Requires the build-time stack + (``virtualizarr``, ``kerchunk``, ``astropy``, ``icechunk``, ``s3fs`` for anonymous S3). + """ + import shutil + + import icechunk + import xarray as xr + from obstore.store import S3Store + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + from virtualizarr.registry import ObjectStoreRegistry + + store_path = str(store_path) + shutil.rmtree(store_path, ignore_errors=True) + + registry = ObjectStoreRegistry({S3_PREFIX: S3Store(BUCKET, region=region, skip_signature=True)}) + parser = FITSParser(reader_options={"storage_options": {"anon": True}}) + frames = [ + open_virtual_dataset(url=u, registry=registry, parser=parser).expand_dims("frame") + for u in uris + ] + combined = xr.concat(frames, dim="frame") # (frame, 1014, 1014), one image per chunk + + ice_prefix = S3_PREFIX + "/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.s3_store(region=region, anonymous=True)) + ) + repo = icechunk.Repository.create( + icechunk.local_filesystem_storage(store_path), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials( + {ice_prefix: icechunk.s3_anonymous_credentials()} + ), + ) + session = repo.writable_session("main") + combined.virtualize.to_icechunk(session.store) + session.commit(f"index {len(uris)} Hubble WFC3/IR frames") + return store_path + + +def open_store( + store_path: str | os.PathLike[str] = DEFAULT_STORE, *, region: str = REGION +) -> Store: + """Reopen the Icechunk repo built by :func:`build_store` as a read-only zarr ``Store``. + + Resolves virtual chunks straight from the public S3 objects. Needs ``icechunk`` (a Rust + read path -- not ``kerchunk``): the build-time index libraries are absent from here. + """ + import icechunk + + ice_prefix = S3_PREFIX + "/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.s3_store(region=region, anonymous=True)) + ) + repo = icechunk.Repository.open( + icechunk.local_filesystem_storage(str(store_path)), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials( + {ice_prefix: icechunk.s3_anonymous_credentials()} + ), + ) + return repo.readonly_session("main").store + + +# --------------------------------------------------------------------------- transforms +def clean_normalize(chunk: DecodedChunk) -> DecodedChunk: + """chunk_transform: NaN-scrub + per-frame robust standardization, vectorized. + + WFC3/IR ``SCI`` frames are in e-/s with bad-pixel NaNs and a huge dynamic range. Subtract + the per-frame median and divide by a MAD-based scale, then clip -- so MSE denoising is + well-conditioned. Per-(variable, chunk), deterministic, pure numpy: the cacheable chunk + stage, and it releases the GIL on the decode pool. + """ + if chunk.read.array != SCI_VAR: + return chunk + x = np.nan_to_num(chunk.data.astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0) + med = np.median(x, axis=(1, 2), keepdims=True) + mad = np.median(np.abs(x - med), axis=(1, 2), keepdims=True) + x = (x - med) / (1.4826 * mad + 1e-6) + chunk.data = np.clip(x, -5.0, 5.0).astype(np.float32) + return chunk + + +@dataclass +class AddNoise: + """batch_transform: split each frame into a ``(noisy, clean)`` pair with fresh Gaussian noise. + + Per-sample random, so it belongs in the (uncached) batch stage, not the chunk stage. + Adds a channel axis: ``clean``/``noisy`` are ``(B, 1, H, W)`` float32. + """ + + sigma: float = NOISE_SIGMA + seed: int | None = None + + def __call__(self, batch: Batch) -> Batch: + clean = batch.arrays.pop(SCI_VAR)[:, None, :, :].astype(np.float32) # (B, 1, H, W) + rng = np.random.default_rng(self.seed) + noise = rng.normal(0.0, self.sigma, size=clean.shape).astype(np.float32) + batch.arrays["clean"] = clean + batch.arrays["noisy"] = clean + noise + return batch + + +# --------------------------------------------------------------------------- dataset +def denoise_dataset( + store: Store, + *, + batch_size: int = 4, + shuffle: bool = True, + coarsen: int = COARSEN, + sigma: float = NOISE_SIGMA, + cache_dir: str | None = None, + max_inflight: int | None = 4, +) -> InSituDataset: + """One denoising dataset over the Hubble store: iterate ``.train`` / ``.val`` / ``.test``. + + Each batch carries ``noisy`` and ``clean`` ``(B, 1, H, W)`` frames. ``coarsen`` block-means + the frame (``1014 -> 1014//coarsen``) to keep the CPU demo light while still reading whole + frames from S3 (the no-reshard, train-in-place stance). + + ``max_inflight`` is capped low by default: MAST's *anonymous* public bucket throttles + (HTTP 503 SlowDown) under heavy concurrent read-ahead, and the virtual-reference fetch does + not retry. A benchmark on AWS would use authenticated/retrying access and raise this. + """ + geoms = open_geometries(store, variables=[SCI_VAR], sample_axis=0) + chunk_transforms: list[ChunkTransform] = [clean_normalize, Coarsen(factor=coarsen)] + return InSituDataset( + store, + split_by_chunk(geoms[SCI_VAR], fractions=(0.7, 0.15, 0.15)), + geometries=geoms, + batch_size=batch_size, + shuffle=shuffle, + cache_dir=cache_dir, + max_inflight=max_inflight, + chunk_transforms=chunk_transforms, + batch_transforms=[AddNoise(sigma=sigma)], + ) + + +# --------------------------------------------------------------------------- metrics / baseline +def psnr(pred: np.ndarray, clean: np.ndarray, data_range: float = 10.0) -> float: + """Peak signal-to-noise ratio (dB). ``data_range`` is the clipped span (-5..5 -> 10).""" + mse = float(np.mean((pred.astype(np.float64) - clean.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else 10.0 * np.log10(data_range**2 / mse) + + +def median_baseline(noisy: np.ndarray, size: int = 3) -> np.ndarray: + """Naive denoiser: a per-frame median filter -- the no-training reference to beat.""" + from scipy.ndimage import median_filter + + out = np.empty_like(noisy) + for i in range(noisy.shape[0]): # (B, 1, H, W) + out[i, 0] = median_filter(noisy[i, 0], size=size) + return out + + +# --------------------------------------------------------------------------- CLI glue +def load_uris(path: str | os.PathLike[str] = DEFAULT_URIS) -> list[str]: + """Load the cached list of Hubble frame S3 URIs (curated from MAST via astroquery).""" + return json.loads(Path(path).read_text()) + + +def build_datasets(args: argparse.Namespace) -> InSituDataset: + """One denoising dataset from CLI args: offline ``synthetic`` frames (written fresh to a temp + store) or the real ``hubble`` Icechunk store (built from URIs first if ``--build``).""" + if args.source == "synthetic": + import tempfile + + url = f"file://{tempfile.mkdtemp()}/hubble_synth.zarr" + make_synthetic_store(url, n_frames=args.n_frames, size=args.size) + store = obstore_store(url) + else: + if args.build: + print(f"building store from {args.uris} ...") + build_store(load_uris(args.uris), args.store) + store = open_store(args.store) + return denoise_dataset( + store, batch_size=args.batch_size, coarsen=args.coarsen, sigma=args.sigma + ) + + +def cli(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description="Hubble WFC3/IR denoising -- train in place.") + p.add_argument( + "--source", + choices=("synthetic", "hubble"), + default="synthetic", + help="offline synthetic frames | real Hubble WFC3/IR frames (needs --build once)", + ) + p.add_argument("--uris", default=str(DEFAULT_URIS), help="JSON list of _flt.fits S3 URIs") + p.add_argument("--store", default=str(DEFAULT_STORE), help="local Icechunk repo path") + p.add_argument("--build", action="store_true", help="(re)build the hubble store first") + p.add_argument("--batch-size", type=int, default=4) + p.add_argument("--epochs", type=int, default=8) + p.add_argument("--coarsen", type=int, default=COARSEN) + p.add_argument("--sigma", type=float, default=NOISE_SIGMA) + p.add_argument("--n-frames", type=int, default=12, help="synthetic frame count") + p.add_argument("--size", type=int, default=128, help="synthetic frame size (Y=X)") + return p.parse_args(argv) diff --git a/examples/hubble/train_torch.py b/examples/hubble/train_torch.py new file mode 100644 index 0000000..4311806 --- /dev/null +++ b/examples/hubble/train_torch.py @@ -0,0 +1,93 @@ +"""Train a small denoiser on real Hubble frames, streamed in place by insitubatch. + +Only this file touches torch: ``to_torch`` is a zero-copy DLPack hand-off, and the loop moves +tensors to the device itself. The data layer (``data.py``) is framework-free numpy. + +Run (builds the reference store on first use):: + + python -m examples.hubble.train_torch --build --epochs 5 + +The store indexes real WFC3/IR frames of M16 (the Eagle Nebula) on MAST's public S3 bucket +as virtual references -- no pixels are resharded or downloaded ahead of time. +""" + +from __future__ import annotations + +import numpy as np +import torch +from torch import nn + +from insitubatch import InSituDataset, to_torch + +from .data import build_datasets, cli, median_baseline, psnr + + +class Denoiser(nn.Module): + """A tiny residual CNN (DnCNN-lite): predict the noise, subtract it from the input.""" + + def __init__(self, channels: int = 32, depth: int = 5) -> None: + super().__init__() + layers: list[nn.Module] = [nn.Conv2d(1, channels, 3, padding=1), nn.ReLU(inplace=True)] + for _ in range(depth - 2): + layers += [nn.Conv2d(channels, channels, 3, padding=1), nn.ReLU(inplace=True)] + layers += [nn.Conv2d(channels, 1, 3, padding=1)] + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x - self.net(x) # residual denoising + + +def evaluate(model: nn.Module, ds: InSituDataset, device: str = "cpu") -> float: + """Mean PSNR (dB) of the model's denoised output vs the clean frame, over the val split.""" + model.eval() + scores: list[float] = [] + with torch.no_grad(): + for batch in ds.val: + d = to_torch(batch) + pred = model(d["noisy"].to(device)).cpu().numpy() + scores.append(psnr(pred, batch.arrays["clean"])) + return float(np.mean(scores)) if scores else float("nan") + + +def baseline_psnr(ds: InSituDataset) -> float: + """Mean PSNR (dB) of the no-training median-filter baseline over the val split.""" + scores = [psnr(median_baseline(b.arrays["noisy"]), b.arrays["clean"]) for b in ds.val] + return float(np.mean(scores)) if scores else float("nan") + + +def train(ds: InSituDataset, *, epochs: int, device: str = "cpu") -> nn.Module: + model = Denoiser().to(device) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + loss_fn = nn.MSELoss() + for epoch in range(epochs): + ds.set_epoch(epoch) + model.train() + total, n = 0.0, 0 + for batch in ds.train: + d = to_torch(batch) + noisy, clean = d["noisy"].to(device), d["clean"].to(device) + opt.zero_grad() + loss = loss_fn(model(noisy), clean) + loss.backward() + opt.step() + total += loss.item() * noisy.shape[0] + n += noisy.shape[0] + val = evaluate(model, ds, device) + print(f"epoch {epoch}: train_mse {total / max(n, 1):.4f} val_psnr {val:.2f} dB") + return model + + +def main(argv: list[str] | None = None) -> None: + args = cli(argv) + ds = build_datasets(args) + + base = baseline_psnr(ds) + print(f"median-filter baseline val PSNR: {base:.2f} dB (no training)\n") + model = train(ds, epochs=args.epochs) + print(f"\ntrained denoiser val PSNR: {evaluate(model, ds):.2f} dB") + print(f"median baseline val PSNR: {base:.2f} dB") + ds.close() + + +if __name__ == "__main__": + main() diff --git a/examples/hubble/wfc3ir_m16_uris.json b/examples/hubble/wfc3ir_m16_uris.json new file mode 100644 index 0000000..49ef03c --- /dev/null +++ b/examples/hubble/wfc3ir_m16_uris.json @@ -0,0 +1,14 @@ +[ + "s3://stpubdata/hst/public/icg9/icg919k0q/icg919k0q_flt.fits", + "s3://stpubdata/hst/public/icg9/icg919k1q/icg919k1q_flt.fits", + "s3://stpubdata/hst/public/ick9/ick901hxq/ick901hxq_flt.fits", + "s3://stpubdata/hst/public/ick9/ick901hzq/ick901hzq_flt.fits", + "s3://stpubdata/hst/public/ick9/ick901i4q/ick901i4q_flt.fits", + "s3://stpubdata/hst/public/ick9/ick901i7q/ick901i7q_flt.fits", + "s3://stpubdata/hst/public/ick9/ick902n7q/ick902n7q_flt.fits", + "s3://stpubdata/hst/public/ick9/ick902n9q/ick902n9q_flt.fits", + "s3://stpubdata/hst/public/ick9/ick902nbq/ick902nbq_flt.fits", + "s3://stpubdata/hst/public/ick9/ick902neq/ick902neq_flt.fits", + "s3://stpubdata/hst/public/idnm/idnm20cmq/idnm20cmq_flt.fits", + "s3://stpubdata/hst/public/idnm/idnm20cnq/idnm20cnq_flt.fits" +] \ No newline at end of file diff --git a/tests/test_hubble.py b/tests/test_hubble.py new file mode 100644 index 0000000..7991a4d --- /dev/null +++ b/tests/test_hubble.py @@ -0,0 +1,73 @@ +"""Hubble FITS denoising example: offline synthetic frames guard the stream+train path. + +The real example indexes Hubble ``_flt.fits`` on S3 as virtual references, which needs the +build-time stack (VirtualiZarr/kerchunk/astropy) and the network. These tests instead drive a +synthetic ``SCI (frame, y, x)`` zarr with the *same* geometry (one frame per chunk over +``sample_axis=0``), so the insitubatch-facing wiring -- the NaN-scrubbing chunk stage, the +per-sample-noise batch stage, the ``(noisy, clean)`` labels, and the torch loop -- can't drift +silently when insitubatch changes. No network, no FITS stack. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from examples.hubble.data import ( + SCI_VAR, + denoise_dataset, + make_synthetic_store, + median_baseline, + psnr, +) +from insitubatch import obstore_store, open_geometries + + +@pytest.fixture +def synth_store(tmp_path) -> str: + """A small synthetic Hubble-like store: 16 frames, NaN bad-pixels, wide dynamic range.""" + url = f"file://{tmp_path}/hubble.zarr" + make_synthetic_store(url, n_frames=16, size=128, seed=0) + return url + + +def test_one_frame_per_chunk(synth_store) -> None: + # The FITS-derived geometry: each frame is one chunk along the sample axis. + geom = open_geometries(obstore_store(synth_store), variables=[SCI_VAR], sample_axis=0)[SCI_VAR] + assert geom.n_samples == 16 + assert geom.sample_chunk_size == 1 + + +def test_denoise_batch_shapes_and_scrubbing(synth_store) -> None: + ds = denoise_dataset(obstore_store(synth_store), batch_size=4, shuffle=False, coarsen=4) + ds.set_epoch(0) + batch = next(iter(ds.train)) + + assert set(batch.arrays) == {"noisy", "clean"} + assert batch.arrays["clean"].shape[1:] == (1, 32, 32) # 128 // coarsen=4, channel axis added + assert batch.arrays["noisy"].shape == batch.arrays["clean"].shape + # clean_normalize scrubs the NaN bad-pixels and clips to the robust-standardized range. + assert np.isfinite(batch.arrays["clean"]).all() + assert np.isfinite(batch.arrays["noisy"]).all() + assert batch.arrays["clean"].min() >= -5.0 and batch.arrays["clean"].max() <= 5.0 + # AddNoise actually perturbs the frame (the batch stage ran). + assert not np.allclose(batch.arrays["noisy"], batch.arrays["clean"]) + + +def test_median_baseline_denoises(synth_store) -> None: + # The no-training reference must at least remove noise: higher PSNR than the noisy input. + ds = denoise_dataset(obstore_store(synth_store), batch_size=8, shuffle=False) + ds.set_epoch(0) + batch = next(iter(ds.train)) + noisy, clean = batch.arrays["noisy"], batch.arrays["clean"] + assert psnr(median_baseline(noisy), clean) > psnr(noisy, clean) + + +def test_torch_beats_baseline(synth_store) -> None: + pytest.importorskip("torch") + from examples.hubble.train_torch import baseline_psnr, evaluate, train + + ds = denoise_dataset(obstore_store(synth_store), batch_size=4) + base = baseline_psnr(ds) + model = train(ds, epochs=20) + assert evaluate(model, ds) > base From 321e4e2749e86508f319e9f749f0d0a33e0728bb Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 13 Jul 2026 05:08:25 +0000 Subject: [PATCH 2/2] examples/hubble: drop scipy, use a pure-numpy median baseline The median-filter baseline reached for scipy.ndimage, but scipy is only a transitive (scikit-learn/bench) dep -- so tests/test_hubble.py would fail on CI matrix jobs that don't install --extra bench (jax, tf, free-threaded). Replace it with a vectorized numpy median (edge-padded sliding windows), matching the pure-numpy baselines in the microscopy/advection examples, so the example and its test need nothing beyond insitubatch core (+ torch, already guarded). Co-Authored-By: Claude Opus 4.8 --- examples/README.md | 4 ++-- examples/hubble/data.py | 17 +++++++++++------ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/examples/README.md b/examples/README.md index 47754e4..98cd1d1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -118,11 +118,11 @@ frame one sample. The build-time index libraries (`virtualizarr`/`kerchunk`/`ast ```bash uv sync --extra torch -uv run --with scipy python -m examples.hubble.train_torch # offline synthetic frames (default) +uv run python -m examples.hubble.train_torch # offline synthetic frames (default) # real Hubble frames on S3 -- indexes them into a virtual-reference store first (needs the # build-time stack), then streams and trains: -uv run --with virtualizarr --with kerchunk --with astropy --with icechunk --with s3fs --with scipy \ +uv run --with virtualizarr --with kerchunk --with astropy --with icechunk --with s3fs \ python -m examples.hubble.train_torch --source hubble --build ``` diff --git a/examples/hubble/data.py b/examples/hubble/data.py index 2dbc5d0..08a7a71 100644 --- a/examples/hubble/data.py +++ b/examples/hubble/data.py @@ -251,13 +251,18 @@ def psnr(pred: np.ndarray, clean: np.ndarray, data_range: float = 10.0) -> float def median_baseline(noisy: np.ndarray, size: int = 3) -> np.ndarray: - """Naive denoiser: a per-frame median filter -- the no-training reference to beat.""" - from scipy.ndimage import median_filter + """Naive denoiser: a per-frame ``size``x``size`` median filter -- the no-training reference. - out = np.empty_like(noisy) - for i in range(noisy.shape[0]): # (B, 1, H, W) - out[i, 0] = median_filter(noisy[i, 0], size=size) - return out + Pure vectorized numpy (edge-padded sliding windows, median over the window stack) so the + example needs no scipy, matching the numpy baselines in the other examples. ``(B, 1, H, W)``. + """ + pad = size // 2 + padded = np.pad(noisy, ((0, 0), (0, 0), (pad, pad), (pad, pad)), mode="edge") + h, w = noisy.shape[2], noisy.shape[3] + windows = np.stack( + [padded[:, :, i : i + h, j : j + w] for i in range(size) for j in range(size)], axis=0 + ) + return np.median(windows, axis=0).astype(noisy.dtype) # --------------------------------------------------------------------------- CLI glue