Skip to content

Examples

The sample axis is a role, not a fixed dimension — so the same engine trains on weather over time, segments microscopy volumes over Z, and streams telescope frames out of an archival format that was never zarr. Each example below is a different geometry on a real public store, with no reshard anywhere.

All of them are runnable from a checkout (they are not shipped in the wheel), and every one has an offline synthetic --source so you can run it with no network or cloud credentials.

Example Domain & store What is different about the data What it proves
advection/ Weather — WeatherBench2 ERA5 (gs://, anonymous), Arraylake/Icechunk, or synthetic Input at t, target at t+24 h as offset views of one array Windowed multi-offset sampling; one dataset → torch, JAX and TF
microscopy/ Bio-imaging — IDR OME-NGFF (T,C,Z,Y,X) on s3://idr Samples a middle axis (sample_axis=2); two co-registered variables chunked 1 vs 30 planes deep Arbitrary sample axis + per-variable chunk size — the engine is not weather-specific
hubble/ Astronomy — Hubble WFC3/IR frames of M16 on MAST's s3://stpubdata FITS, not zarr — indexed as virtual byte-range references; one frame is one chunk Training in place over an archival format; streaming value without decode amortization
sdss/ Astronomy — SDSS DR17 spPlate spectra, or our public reference stores FITS binary-image tables re-chunked by byte arithmetic; two layouts from the same bytes Both regimes from one dataset: many-fibers-per-chunk decode amortization vs one-fiber-per-chunk archive-scale streaming
WB2 pair Weather — WeatherBench2 ERA5 The same task on two engines The cold-start / memory trade-off vs an xbatcher worker stack
transforms.py · fit_scaler.py Any — tiny offline store Why there are two transform stages, and how to fit a scaler over the loader

Full flags, data-source tables and design notes for each live in examples/README.md.

advection — a 24-hour forecast in three frameworks

One InSituDataset reads three fields at time t (temperature t2m and the 10 m wind u10, v10) and the target t2m 24 hours later via g.shift(horizon). Input and target are offset views of the same in-place array — the windowing unlock — and nothing is resharded. The resulting numpy Batch then trains the same tiny CNN in three frameworks through the DLPack adapters; the three files differ only in framework calls.

uv sync --extra bench --extra torch               # PyTorch   (torch.nn)
uv run python -m examples.advection.train_torch

uv sync --extra bench --extra jax                 # JAX       (flax + optax)
uv run python -m examples.advection.train_jax

uv sync --extra bench --extra tf                  # TensorFlow (Keras)
uv run python -m examples.advection.train_tf

Each run prints 24-hour forecast skill on held-out data — a model that reads the wind to predict advection, versus the persistence baseline. --source wb2 runs the same code against real ERA5 in the cloud; there the claim is "same pipeline, real data", not SOTA skill (24 h temperature persistence is a strong baseline).

Install one framework at a time

Having torch, JAX and TensorFlow present in the same uv venv can segfault. Sync the extra for the one you are running.

microscopy — cell segmentation over Z

The cross-domain showcase: same engine, different geometry. Where advection samples the outer time axis, microscopy/ samples a middle axis — one Z-plane of an OME-NGFF (T,C,Z,Y,X) confocal stack — and gathers two co-registered variables per anchor: a 2-channel raw image chunked one plane deep on Z, and its mask label chunked 30 planes deep and tiled in Y/X. Different physical chunking, different channel counts, one sample grid, no reshard.

uv sync --extra torch
uv run python -m examples.microscopy.train_torch                 # synthetic cells (offline)
uv run python -m examples.microscopy.train_torch --source idr    # the real IDR image (streamed)

The task is per-plane foreground segmentation and the baseline is a global Otsu threshold — the segmentation analogue of persistence. Otsu reads each pixel's intensity alone, so a smooth autofluorescence haze gradient defeats it; a tiny CNN that reads the neighbourhood beats it. Each run prints held-out foreground IoU, model vs Otsu.

hubble — denoising real telescope frames from FITS

The archival-format showcase: this data never was zarr. hubble/ indexes real Hubble WFC3/IR frames of M16 (the Eagle Nebula) on MAST's public AWS bucket as virtual referencesVirtualiZarr parses each _flt.fits header and commits byte-range references to a local Icechunk repo. No pixels are copied: the store is a few kB pointing at the original FITS objects, and insitubatch streams frames straight from S3 with sample_axis=0. The indexing libraries are build-time only — the training hot path is icechunk + numpy.

uv sync --extra torch
uv run python -m examples.hubble.train_torch      # offline synthetic frames (default)

# real Hubble frames on S3 — build the virtual-reference store, then stream and train:
uv sync --extra torch --extra astronomy
uv run python -m examples.hubble.train_torch --source hubble --build

Because a FITS image is one chunk, this demonstrates the streaming-in-place value (no reshard over a giant archive) rather than the many-samples-per-chunk decode amortization the chunked-zarr examples show — the honest boundary of the thesis, kept visible on purpose.

sdss — reconstructing galaxy spectra streamed in place

The companion to Hubble, and an astroML mirror. astroML's spectral-PCA workflow downloads the raw archive and resamples every spectrum onto a common grid into one spec4000.npz — the download-and-reshard step this project argues against. sdss/ instead indexes SDSS DR17 spPlate frames as virtual references and trains on them where they lie.

The trick is that a plate already holds ~640 fibers on one common log-wavelength grid, stored fiber-major — so a block of fibers is a contiguous byte range, and re-chunking is arithmetic on the manifest, not a copy. The same bytes give two layouts:

  • one plate → 64 fibers per chunk. One read decodes 64 spectra: the decode-amortization regime.
  • many plates → one fiber per chunk. Plates cover slightly different wavelength ranges, so joining them means cropping each to the shared window, which breaks fiber contiguity. You get archive-scale concatenation (extending to the ~2800-plate archive) and pay one read per spectrum: the streaming regime.

The crop is exact, not interpolated: every plate uses the same dloglam step and their start wavelengths differ by whole bins, so the windows land on one grid. The two regimes are mutually exclusive here — an honest property of the FITS byte layout, not a tuning choice.

Streaming without building anything

The stores are published read-anonymous, so the scan is already done. This needs icechunk only — no FITS stack, no credentials, no build step:

uv sync --extra torch --extra astronomy
uv run python -m examples.sdss.train_torch --source published                      # 6 plates
uv run python -m examples.sdss.train_torch --source published --published 1plate   # 1 plate

--published picks the layout: 1plate (640 spectra, 64 per chunk), 6plate (3840 spectra, one per chunk), or 6plate-mirror (same as 6plate, references pointing at our mirror of the FITS files — the default, and the one to prefer). Bucket layout and a plain-zarr recipe are in the dataset README.

Or index the archive yourself

uv run python -m examples.sdss.train_torch                                   # offline synthetic
uv run python -m examples.sdss.train_torch --source sdss --build             # 1 plate
uv run python -m examples.sdss.train_torch --source sdss --build --plates 8  # 8 plates

build_store reads the spPlate headers with VirtualiZarr and writes an Icechunk repo of byte-range references — kilobytes, whatever the archive size. Point it at more plate URLs to index a larger slice.

The task mirrors astroML's spectral reconstruction: recover the spectrum through a low-dimensional bottleneck. The baseline is PCA at the same latent dim — the optimal linear reconstruction — and a small 1-D convolutional autoencoder trained over the streamed batches beats it, because varying redshift shifts the lines and a fixed linear basis reconstructs a shifted spectrum poorly. --source synthetic (the default) needs no network and carries that claim; the real run is the same pipeline on the real archive.

The WeatherBench2 cold-start pair

The same task two ways, so you can see the trade-off and pick per workload: wb2_dataloader.py is the insitu single-event-loop loader (with --backend fsspec for the gcsfs A/B), and wb2_xbatcher.py is the xbatcher + torch DataLoader worker stack, following Earthmover's dataloader-demo, focused on cold-start latency and how forkserver-preload cuts it.

uv run python -m examples.wb2_dataloader          # tiny synthetic data, no network

The WeatherBench2 walkthrough narrates this pair end to end, and Benchmarks has the measured numbers.

Transforms and normalization

transforms.py puts the two user transform stages side by side on a tiny offline store — a Kelvin→Celsius chunk_transform (per chunk, one variable, cached) and a cross-variable windspeed batch_transform (needs the assembled batch, uncached). It is the clearest illustration of why there are two; see Transforms for the placement model.

fit_scaler.py fits a StandardScaler over the loader with sklearn partial_fit — the recommended pattern (it warms the cache) versus caching scaled chunks.

uv run python -m examples.transforms              # no network