diff --git a/.gitignore b/.gitignore index 8932772..addefab 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ bench/figures/ # Generated example stores (virtual-reference Icechunk repos, rebuilt from committed URI lists) examples/hubble/hubble_store/ +examples/sdss/sdss_store/ # MkDocs build output (docs/figures/ is committed for the published site) site/ diff --git a/examples/README.md b/examples/README.md index 98cd1d1..e12d591 100644 --- a/examples/README.md +++ b/examples/README.md @@ -121,9 +121,9 @@ uv sync --extra torch 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 \ - python -m examples.hubble.train_torch --source hubble --build +# `astronomy` build-time stack), then streams and trains: +uv sync --extra torch --extra astronomy +uv run 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 @@ -140,6 +140,58 @@ 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. +## sdss/ — reconstructing SDSS spectra streamed in place (no reshard) + +The **decode-amortization** companion to Hubble, and an [astroML](https://github.com/astroML/astroML) +mirror. astroML's spectral-PCA workflow (`fetch_sdss_corrected_spectra` → `compute_sdss_pca`) first +downloads the raw archive and resamples every spectrum onto a common grid into one `spec4000.npz` +file — the exact download-and-reshard step insitubatch argues against. [`sdss/`](sdss/data.py) +instead indexes a real **`spPlate`** frame on `data.sdss.org` (over HTTPS) as virtual references: +one plate holds all ~640 fibers on a *common* log-wavelength grid as a single `(fiber, wave)` image, +which [`_rechunk_fibers`](sdss/data.py) splits along the fiber axis into contiguous virtual chunks +by **pure byte arithmetic** — no pixels moved. `sample_axis=0` makes each fiber one sample, and +because a chunk is many fibers this lands in the **O(chunks) decode-amortization** regime (contrast +Hubble's one image per chunk). + +```bash +uv sync --extra torch +uv run python -m examples.sdss.train_torch # offline synthetic spectra (default) + +# real SDSS spPlate over HTTPS -- indexes it into a virtual-reference store first (needs the +# `astronomy` build-time stack), then streams and trains: +uv sync --extra torch --extra astronomy +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, aligned +``` + +One plate (~640 spectra, ~10 MB) fits in memory, so `--plates` scales the real story. The two +modes trade off along the FITS byte layout, and **both move no pixels — only the chunk manifest is +rewritten:** a **single** plate is re-chunked along the fiber axis at full wave width (many fibers +per chunk — the *decode-amortization* regime); **several** plates are cropped to their shared +log-wavelength window and folded into one flat fiber axis (one fiber per chunk — the *streaming* +regime). The crop is exact — every SDSS plate uses the same `dloglam` and their `COEFF0` start +wavelengths differ by whole bins, so the windows align onto one grid with no resampling — which is +the only no-reshard way to concatenate the ragged-width plates. That is the corpus-scale path: the +full archive is thousands of plates (tens of GB) you would otherwise download and reshard. + +The task mirrors astroML's **spectral reconstruction**: recover the clean spectrum through a +low-dimensional bottleneck. The baseline is **PCA** at the same latent dim — the optimal *linear* +reconstruction; a small **1-D convolutional autoencoder**, trained by SGD over the streamed +mini-batches, beats it because galaxy spectra are translation-structured (varying redshift shifts +the lines, which a fixed-dim linear basis reconstructs poorly). A per-fiber robust normalization +(`normalize`) runs vectorized on the decode pool; the per-sample reconstruction noise lives in the +`Corrupt` batch stage, per the transform-cost contract. `--source synthetic` (the default) needs no +network or FITS stack — redshift-shifted synthetic spectra on which the conv AE clearly beats PCA — +and backs the drift test in `tests/test_sdss.py`. The real run is a *same-pipeline, real-archive* +demonstration (one plate = ~640 spectra); the synthetic default carries the beats-PCA claim. + +> Multi-plate is one-fiber-per-chunk by construction: cropping to the common window breaks the +> multi-fiber contiguity the single-plate mode relies on, so decode-amortization and corpus scale +> are mutually exclusive here — an honest boundary of the FITS byte layout. SDSS spectra also ship +> as FITS **binary tables** (structured big-endian dtypes); `tests/test_fits_bintable.py` validates +> that path end-to-end through insitubatch and pins the VirtualiZarr fix it needs +> ([#1037](https://github.com/zarr-developers/VirtualiZarr/pull/1037)). + ## 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/data.py b/examples/hubble/data.py index 08a7a71..470fa82 100644 --- a/examples/hubble/data.py +++ b/examples/hubble/data.py @@ -109,10 +109,10 @@ def build_store( import icechunk import xarray as xr + from obspec_utils.registry import ObjectStoreRegistry 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) @@ -138,7 +138,7 @@ def build_store( ), ) session = repo.writable_session("main") - combined.virtualize.to_icechunk(session.store) + combined.vz.to_icechunk(session.store) session.commit(f"index {len(uris)} Hubble WFC3/IR frames") return store_path diff --git a/examples/sdss/__init__.py b/examples/sdss/__init__.py new file mode 100644 index 0000000..af1645a --- /dev/null +++ b/examples/sdss/__init__.py @@ -0,0 +1,8 @@ +"""SDSS spectra example: reconstruct galaxy spectra streamed in place from spPlate FITS. + +Mirrors astroML's SDSS spectral-PCA reconstruction (``compute_sdss_pca`` / +``fetch_sdss_corrected_spectra``), but where astroML first downloads the raw archive and +resamples every spectrum onto a common grid into a single ``spec4000.npz`` file, this streams +per-plate ``spPlate`` frames -- already on a common log-wavelength grid -- straight from the +SDSS archive as virtual references, with no reshard. See :mod:`examples.sdss.data`. +""" diff --git a/examples/sdss/data.py b/examples/sdss/data.py new file mode 100644 index 0000000..4fcdced --- /dev/null +++ b/examples/sdss/data.py @@ -0,0 +1,464 @@ +"""Data layer for the SDSS spectral-reconstruction example: spectra streamed in place. + +Pipeline +-------- +1. **Index (build time, once).** Each SDSS ``spPlate-PLATE-MJD.fits`` holds all ~640 fibers of + one plate observation as a single ``(fiber, wave) float32`` image on a *common* log-wavelength + grid (unlike the per-object ``spec-lite`` files, which are trimmed to ragged lengths). + :func:`build_store` opens each plate's ``PRIMARY`` flux HDU as a *virtual* dataset + (VirtualiZarr -> ``kerchunk.fits``), concatenates them along the ``fiber`` axis, and commits + the byte-range references to a local Icechunk repo. No flux is copied or resampled -- the store + is a few kB of references pointing at the original FITS objects on ``data.sdss.org``. + +2. **Stream (train time).** :func:`open_store` reopens the Icechunk repo (resolving the virtual + chunks straight over HTTPS), and :func:`reconstruct_dataset` builds an + :class:`~insitubatch.InSituDataset` with ``sample_axis=0`` (one fiber = one sample). Because a + whole plate is one chunk, each decoded chunk yields ~640 samples -- the O(chunks) + decode-amortization regime (contrast the Hubble example, one image per chunk). A per-fiber + robust normalization (:func:`normalize`) runs vectorized on the decode pool; the per-sample + reconstruction noise lives in the :class:`Corrupt` batch stage, per the transform-cost contract. + +The task mirrors astroML's ``compute_sdss_pca`` (spectral reconstruction / eigenspectra): recover +the clean spectrum through a low-dimensional bottleneck. The baseline is **PCA** at the same latent +dimension -- the optimal *linear* reconstruction; a small autoencoder, trained by SGD over the +streamed mini-batches, beats it when the spectra lie on a nonlinear manifold. astroML fits its PCA +after downloading + resampling the archive into one ``spec4000.npz``; here we stream the common-grid +plates in place, no reshard. +""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import xarray as xr +import zarr +from zarr.abc.store import Store + +if TYPE_CHECKING: + from virtualizarr.manifests import ManifestArray + +from insitubatch import ( + Batch, + ChunkTransform, + DecodedChunk, + InSituDataset, + ensure_local_dir, + obstore_store, + open_geometries, + split_by_chunk, +) + +FLUX_VAR = "flux" +SDSS_HOST = "https://data.sdss.org" +REDUX = "/sas/dr17/sdss/spectro/redux/26" # legacy (DR8-era) reduction; spPlate lives per-plate +LATENT_DIM = 16 # bottleneck / PCA components -- the reconstruction budget both methods share +NOISE_SIGMA = 0.3 # per-sample input noise (normalized units), added in the batch stage + +HERE = Path(__file__).resolve().parent +DEFAULT_URIS = HERE / "dr17_plate_uris.json" +DEFAULT_STORE = HERE / "sdss_store" + + +# --------------------------------------------------------------------------- offline synthetic +def make_synthetic_store( + url: str, + *, + n_plates: int = 12, + fibers_per_plate: int = 128, + n_wave: int = 512, + intrinsic_dim: int = 8, + seed: int = 0, +) -> None: + """Write a synthetic ``flux (fiber, wave)`` zarr mimicking the spPlate-derived store's geometry. + + For offline runs (``--source synthetic``) and drift tests -- no network, no FITS/VZ stack. + One plate per chunk (``chunks=(fibers_per_plate, n_wave)``), like the real spPlate layout, so + each chunk carries many fiber samples (the decode-amortization regime). + + Each spectrum is a gentle continuum plus a set of emission lines shifted by a per-fiber + redshift, with per-fiber line amplitudes -- a low-dimensional but strongly *nonlinear* manifold + (a moving line is not a linear combination of a few fixed templates). This mirrors real galaxy + spectra, where varying redshift is exactly what makes a fixed-dimension linear PCA reconstruct + poorly: at ``LATENT_DIM`` components PCA leaves the shifted lines smeared, while a nonlinear + autoencoder of the same bottleneck learns the shift -- so the trained model beats the baseline. + ``intrinsic_dim`` is accepted for API stability but the redshift+amplitude latent sets the true + dimension. + """ + del intrinsic_dim # latent dimension is set by the redshift + per-line amplitudes below + rng = np.random.default_rng(seed) + n_fiber = n_plates * fibers_per_plate + wave = np.linspace(0.0, 1.0, n_wave) + + rest_lines = np.linspace(0.12, 0.88, 5) # rest-frame line centers + width = 0.012 + zshift = rng.uniform(-0.12, 0.12, size=(n_fiber, 1)) # dominant nonlinear (redshift) parameter + slope = rng.uniform(-1.0, 1.0, size=(n_fiber, 1)) + flux = 8.0 + 1.5 * slope * (wave[None, :] - 0.5) # gentle continuum, e-/s-like pedestal + for c in rest_lines: + centers = c + zshift # (n_fiber, 1): the line sweeps across the grid with redshift + amp = rng.uniform(0.5, 3.0, size=(n_fiber, 1)) + flux = flux + amp * np.exp(-((wave[None, :] - centers) ** 2) / (2 * width**2)) + + flux = flux.astype("f4") + flux += rng.normal(0.0, 0.05, size=flux.shape).astype("f4") # small intrinsic scatter + bad = (rng.integers(0, n_fiber, 3 * n_plates), rng.integers(0, n_wave, 3 * n_plates)) + flux[bad] = np.nan # a few bad pixels, as real coadds carry + + ensure_local_dir(url) + group = zarr.open_group(store=obstore_store(url, read_only=False), mode="w") + arr = group.create_array( + FLUX_VAR, + shape=flux.shape, + chunks=(fibers_per_plate, n_wave), + dtype="f4", + dimension_names=("fiber", "wave"), + ) + arr[:] = flux + + +# --------------------------------------------------------------------------- build time +# +# A single plate is ~640 fibers x ~3864 wave (~10 MB) -- it fits in memory, so one plate is a demo, +# not a streaming workload. The two build modes trade off along the FITS byte layout (no reshard in +# either -- both only rewrite the *chunk manifest*, never the pixels): +# +# one plate -> _single_plate_fiber_chunks: full wave width, many fibers per chunk +# (the O(chunks) decode-amortization regime; toy scale). +# many plates -> _many_plates_common_grid: plates cropped to a shared wavelength window and +# folded into one flat fiber axis, one fiber per chunk (the streaming regime; the +# only no-reshard way to concat ragged-width plates -- scales to the archive). +# +FIBERS_PER_CHUNK = 64 # one plate's 640 fibers -> 10 contiguous chunks + + +def _chunk_ref(array: ManifestArray) -> tuple[str, int]: + """The ``(path, byte offset)`` of the single virtual chunk kerchunk emits for a plate HDU.""" + manifest = array.manifest + return str(manifest._paths[0, 0]), int(manifest._offsets[0, 0]) + + +def _grid_start(plate: xr.Dataset) -> float: + return float(plate[FLUX_VAR].attrs["COEFF0"]) # log10(wavelength) of bin 0 + + +def _grid_step(plate: xr.Dataset) -> float: + return float(plate[FLUX_VAR].attrs["COEFF1"]) # dloglam per bin (identical across SDSS plates) + + +def _virtual_flux( + template: xr.Dataset, + shape: tuple[int, int], + chunk_shape: tuple[int, int], + entries: dict[str, dict[str, object]], +) -> xr.Dataset: + """Assemble a virtual ``flux (fiber, wave)`` dataset from explicit byte-range chunk refs. + + Reuses ``template``'s dtype/codec metadata (from the opened plate) and overrides only the shape + and chunk grid -- the chunks point back into the original FITS bytes, so nothing is copied. + """ + from virtualizarr.manifests import ChunkManifest, ManifestArray + + metadata = template[FLUX_VAR].data.metadata + grid = replace(metadata.chunk_grid, chunk_shape=chunk_shape) + array = ManifestArray( + metadata=replace(metadata, shape=shape, chunk_grid=grid), + chunkmanifest=ChunkManifest(entries=entries), + ) + return xr.Dataset({FLUX_VAR: xr.Variable(("fiber", "wave"), array)}) + + +def _single_plate_fiber_chunks(plate: xr.Dataset, fibers_per_chunk: int) -> xr.Dataset: + """One plate -> contiguous fiber-block chunks at full wave width (many fibers per chunk). + + kerchunk maps the ``(fiber, wave)`` HDU to ONE virtual chunk: the byte range + ``[offset, offset + n_fiber*row_bytes)``. The fiber axis is the outer (row) axis, so fiber block + ``i`` is the contiguous sub-range ``offset + i*fibers_per_chunk*row_bytes`` -- pure byte + arithmetic. Each chunk still holds many fibers: the decode-amortization regime. + """ + array = plate[FLUX_VAR].data + path, base = _chunk_ref(array) + n_fiber, n_wave = array.shape + row_bytes = n_wave * array.dtype.itemsize + if n_fiber % fibers_per_chunk: + raise ValueError(f"{n_fiber} fibers not divisible by fibers_per_chunk={fibers_per_chunk}") + + block_bytes = fibers_per_chunk * row_bytes + entries: dict[str, dict[str, object]] = { + f"{i}.0": {"path": path, "offset": base + i * block_bytes, "length": block_bytes} + for i in range(n_fiber // fibers_per_chunk) + } + return _virtual_flux(plate, (n_fiber, n_wave), (fibers_per_chunk, n_wave), entries) + + +def _many_plates_common_grid(plates: list[xr.Dataset]) -> xr.Dataset: + """N plates -> one flat fiber axis on a shared wavelength window, one fiber per chunk. + + Plates cover slightly different wavelength ranges, so they cannot share a rectangular ``wave`` + axis at full width. But every plate uses the same ``dloglam`` (COEFF1) and their start + wavelengths (COEFF0) differ by whole bins, so cropping each plate to the common overlap window + lands them on ONE grid *exactly* -- no resampling. After the crop each fiber's window is still a + contiguous byte sub-range, so a fiber is one virtual chunk and the plates concatenate into a + flat ``(total_fiber, width)`` sample axis. One fiber per chunk is the streaming regime, but it + scales to the whole archive (thousands of plates) with no reshard and no download. + """ + step = _grid_step(plates[0]) + starts = [_grid_start(p) for p in plates] + window_lo = max(starts) + window_hi = min( + start + p[FLUX_VAR].data.shape[1] * step for p, start in zip(plates, starts, strict=True) + ) + width = round((window_hi - window_lo) / step) + + entries: dict[str, dict[str, object]] = {} + fiber = 0 + for plate, start in zip(plates, starts, strict=True): + array = plate[FLUX_VAR].data + path, base = _chunk_ref(array) + n_fiber, n_wave = array.shape + itemsize = array.dtype.itemsize + row_bytes = n_wave * itemsize + offset_bins = round((window_lo - start) / step) + if abs((window_lo - start) / step - offset_bins) > 1e-3: + raise ValueError("plate wavelength grids are not bin-aligned; cannot crop losslessly") + window_bytes = width * itemsize + for f in range(n_fiber): + byte0 = base + f * row_bytes + offset_bins * itemsize + entries[f"{fiber}.0"] = {"path": path, "offset": byte0, "length": window_bytes} + fiber += 1 + return _virtual_flux(plates[0], (fiber, width), (1, width), entries) + + +def build_store( + plate_urls: list[str], + store_path: str | os.PathLike[str] = DEFAULT_STORE, + *, + fibers_per_chunk: int = FIBERS_PER_CHUNK, +) -> str: + """Index one or more SDSS ``spPlate`` FITS (HTTPS) into a local Icechunk repo of virtual refs. + + Idempotent: rebuilds ``store_path`` from scratch. Requires the ``astronomy`` extra + (``virtualizarr``, ``kerchunk``, ``astropy``, ``icechunk``). A **single** plate becomes + full-width fiber-block chunks (many fibers per chunk -- decode-amortization); **several** plates + are cropped to a shared wavelength window and folded into one flat fiber axis (one fiber per + chunk -- streaming at archive scale). Both modes move no pixels -- see + :func:`_single_plate_fiber_chunks` and :func:`_many_plates_common_grid`. + """ + import shutil + + import icechunk + from obspec_utils.registry import ObjectStoreRegistry + from obstore.store import HTTPStore + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + + store_path = str(store_path) + shutil.rmtree(store_path, ignore_errors=True) + + registry = ObjectStoreRegistry({SDSS_HOST: HTTPStore.from_url(SDSS_HOST)}) + plates = [ + open_virtual_dataset(url=u, registry=registry, parser=FITSParser()).rename( + {"PRIMARY": FLUX_VAR} + ) + for u in plate_urls + ] + virtual = ( + _single_plate_fiber_chunks(plates[0], fibers_per_chunk) + if len(plates) == 1 + else _many_plates_common_grid(plates) + ) + + ice_prefix = SDSS_HOST + "/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.http_store()) + ) + repo = icechunk.Repository.create( + icechunk.local_filesystem_storage(store_path), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), + ) + session = repo.writable_session("main") + virtual.vz.to_icechunk(session.store) + session.commit(f"index {len(plate_urls)} SDSS spPlate frame(s)") + return store_path + + +def open_store(store_path: str | os.PathLike[str] = DEFAULT_STORE) -> Store: + """Reopen the Icechunk repo built by :func:`build_store` as a read-only zarr ``Store``. + + Resolves virtual chunks straight over HTTPS from ``data.sdss.org``. Needs ``icechunk`` (a Rust + read path -- not ``kerchunk``): the build-time index libraries are absent from here. + """ + import icechunk + + ice_prefix = SDSS_HOST + "/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.http_store()) + ) + repo = icechunk.Repository.open( + icechunk.local_filesystem_storage(str(store_path)), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), + ) + return repo.readonly_session("main").store + + +# --------------------------------------------------------------------------- transforms +def normalize(chunk: DecodedChunk) -> DecodedChunk: + """chunk_transform: NaN-scrub + per-fiber robust standardization, vectorized. + + Coadded spectra span a wide flux range with bad-pixel NaNs. Subtract the per-fiber median and + divide by a MAD-based scale, then clip -- so a reconstruction MSE is well-conditioned across + fibers. Per-(variable, chunk), deterministic, pure numpy: the cacheable chunk stage, and it + releases the GIL on the decode pool. + """ + if chunk.read.array != FLUX_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, keepdims=True) + mad = np.median(np.abs(x - med), axis=1, 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 Corrupt: + """batch_transform: split each spectrum into a ``(noisy, clean)`` pair with fresh input noise. + + The reconstruction target is the clean normalized spectrum; the model sees a noisy copy. The + noise is per-sample random, so it belongs in the (uncached) batch stage, not the chunk stage. + ``clean``/``noisy`` are ``(B, W)`` float32. + """ + + sigma: float = NOISE_SIGMA + seed: int | None = None + + def __call__(self, batch: Batch) -> Batch: + clean = batch.arrays.pop(FLUX_VAR).astype(np.float32) # (B, 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 reconstruct_dataset( + store: Store, + *, + batch_size: int = 64, + shuffle: bool = True, + sigma: float = NOISE_SIGMA, + cache_dir: str | None = None, + max_inflight: int | None = None, +) -> InSituDataset: + """One reconstruction dataset over the SDSS store: iterate ``.train`` / ``.val`` / ``.test``. + + Each batch carries ``noisy`` and ``clean`` ``(B, W)`` spectra. Split is by chunk (plate), so no + fiber leaks across train/val/test. ``max_inflight`` is unset by default (``data.sdss.org`` over + HTTPS is not an anonymous-S3 throttler like MAST); raise/cap it to tune read-ahead. + """ + geoms = open_geometries(store, variables=[FLUX_VAR], sample_axis=0) + chunk_transforms: list[ChunkTransform] = [normalize] + return InSituDataset( + store, + split_by_chunk(geoms[FLUX_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=[Corrupt(sigma=sigma)], + ) + + +# --------------------------------------------------------------------------- metrics / baseline +def collect(ds: InSituDataset, split: str) -> tuple[np.ndarray, np.ndarray]: + """Materialize a split into ``(noisy, clean)`` matrices ``(N, W)``. + + Used to fit/score the PCA baseline -- which, like astroML's ``spec4000`` workflow (and unlike + the streamed autoencoder), needs every spectrum resident at once. + """ + noisy, clean = [], [] + for b in getattr(ds, split): + noisy.append(b.arrays["noisy"]) + clean.append(b.arrays["clean"]) + return np.concatenate(noisy), np.concatenate(clean) + + +def fit_pca(clean: np.ndarray, k: int = LATENT_DIM) -> tuple[np.ndarray, np.ndarray]: + """Fit a ``k``-component PCA on clean training spectra: return ``(mean, components)``.""" + mean = clean.mean(axis=0) + _, _, vt = np.linalg.svd(clean - mean, full_matrices=False) + return mean.astype(np.float32), vt[:k].astype(np.float32) + + +def pca_reconstruct(noisy: np.ndarray, mean: np.ndarray, components: np.ndarray) -> np.ndarray: + """Reconstruct spectra by projecting onto the top-``k`` PCA subspace (the linear baseline).""" + centered = noisy - mean + return (mean + (centered @ components.T) @ components).astype(np.float32) + + +def recon_mse(pred: np.ndarray, clean: np.ndarray) -> float: + """Mean squared reconstruction error (lower is better).""" + return float(np.mean((pred.astype(np.float64) - clean.astype(np.float64)) ** 2)) + + +# --------------------------------------------------------------------------- CLI glue +def load_uris(path: str | os.PathLike[str] = DEFAULT_URIS) -> list[str]: + """Load the cached list of SDSS spPlate URLs.""" + return json.loads(Path(path).read_text()) + + +def build_datasets(args: argparse.Namespace) -> InSituDataset: + """One reconstruction dataset from CLI args: offline ``synthetic`` spectra (written fresh to a + temp store) or the real ``sdss`` Icechunk store (built from URIs first if ``--build``).""" + if args.source == "synthetic": + import tempfile + + url = f"file://{tempfile.mkdtemp()}/sdss_synth.zarr" + make_synthetic_store(url, n_plates=args.n_plates, fibers_per_plate=args.fibers) + store = obstore_store(url) + else: + if args.build: + plate_urls = load_uris(args.uris)[: args.plates] + print(f"building store from {len(plate_urls)} plate(s): {plate_urls} ...") + build_store(plate_urls, args.store) + store = open_store(args.store) + return reconstruct_dataset(store, batch_size=args.batch_size, sigma=args.sigma) + + +def cli(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description="SDSS spectral reconstruction -- train in place.") + p.add_argument( + "--source", + choices=("synthetic", "sdss"), + default="synthetic", + help="offline synthetic spectra | real SDSS spPlate frames (needs --build once)", + ) + p.add_argument("--uris", default=str(DEFAULT_URIS), help="JSON list of spPlate FITS URLs") + p.add_argument("--store", default=str(DEFAULT_STORE), help="local Icechunk repo path") + p.add_argument("--build", action="store_true", help="(re)build the sdss store first") + p.add_argument( + "--plates", + type=int, + default=1, + help="real spPlate count: 1 = full-width many-fibers/chunk (decode-amortization); " + ">1 = common-window 1-fiber/chunk (streaming at archive scale)", + ) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--epochs", type=int, default=15) + p.add_argument("--latent-dim", type=int, default=LATENT_DIM) + p.add_argument("--sigma", type=float, default=NOISE_SIGMA) + p.add_argument("--n-plates", type=int, default=12, help="synthetic plate count") + p.add_argument("--fibers", type=int, default=128, help="synthetic fibers per plate") + return p.parse_args(argv) diff --git a/examples/sdss/dr17_plate_uris.json b/examples/sdss/dr17_plate_uris.json new file mode 100644 index 0000000..d8d6510 --- /dev/null +++ b/examples/sdss/dr17_plate_uris.json @@ -0,0 +1,8 @@ +[ + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0266/spPlate-0266-51602.fits", + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0267/spPlate-0267-51608.fits", + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0268/spPlate-0268-51633.fits", + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0269/spPlate-0269-51581.fits", + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0270/spPlate-0270-51909.fits", + "https://data.sdss.org/sas/dr17/sdss/spectro/redux/26/0271/spPlate-0271-51883.fits" +] diff --git a/examples/sdss/train_torch.py b/examples/sdss/train_torch.py new file mode 100644 index 0000000..6db25c3 --- /dev/null +++ b/examples/sdss/train_torch.py @@ -0,0 +1,144 @@ +"""Train a small autoencoder on real SDSS spectra, 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.sdss.train_torch --source sdss --build --epochs 15 + +The store indexes real SDSS ``spPlate`` frames on ``data.sdss.org`` as virtual references -- no +flux is resampled or downloaded ahead of time. The task mirrors astroML's spectral-PCA +reconstruction; the streamed autoencoder is compared against a PCA baseline at the same latent dim. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.nn import functional as F + +from insitubatch import InSituDataset, to_torch + +from .data import ( + LATENT_DIM, + build_datasets, + cli, + collect, + fit_pca, + pca_reconstruct, + recon_mse, +) + + +class AutoEncoder(nn.Module): + """A small 1-D convolutional autoencoder: compress a spectrum to ``latent_dim``, reconstruct it. + + Convolutional on purpose: galaxy spectra are translation-structured (varying redshift shifts the + lines), which a *linear* PCA of a fixed latent dim reconstructs poorly -- a moving line is not a + low-rank combination of fixed templates. The conv encoder learns the shift, so at the same + bottleneck it beats the PCA baseline (as on real spectra). + """ + + def __init__(self, n_wave: int, latent_dim: int = LATENT_DIM, channels: int = 16) -> None: + super().__init__() + self.n_wave = n_wave + self.padded = ((n_wave + 3) // 4) * 4 # conv path halves length twice + bottleneck_len = self.padded // 4 + c = channels + self.enc_conv = nn.Sequential( + nn.Conv1d(1, c, 7, padding=3), + nn.ReLU(inplace=True), + nn.MaxPool1d(2), + nn.Conv1d(c, 2 * c, 7, padding=3), + nn.ReLU(inplace=True), + nn.MaxPool1d(2), + ) + self.enc_lin = nn.Linear(2 * c * bottleneck_len, latent_dim) + self.dec_lin = nn.Linear(latent_dim, 2 * c * bottleneck_len) + self.unflatten = nn.Unflatten(1, (2 * c, bottleneck_len)) + self.dec_conv = nn.Sequential( + nn.Upsample(scale_factor=2), + nn.Conv1d(2 * c, c, 7, padding=3), + nn.ReLU(inplace=True), + nn.Upsample(scale_factor=2), + nn.Conv1d(c, 1, 7, padding=3), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + xp = F.pad(x, (0, self.padded - self.n_wave)).unsqueeze(1) # (B, 1, padded) + z = self.enc_lin(self.enc_conv(xp).flatten(1)) + out = self.dec_conv(self.unflatten(self.dec_lin(z))).squeeze(1) # (B, padded) + return out[:, : self.n_wave] + + +def evaluate(model: nn.Module, ds: InSituDataset, device: str = "cpu") -> float: + """Mean reconstruction MSE of the model over the val split (lower is better).""" + model.eval() + total, n = 0.0, 0 + with torch.no_grad(): + for batch in ds.val: + d = to_torch(batch) + pred = model(d["noisy"].to(device)) + clean = d["clean"].to(device) + total += float(((pred - clean) ** 2).mean()) * clean.shape[0] + n += clean.shape[0] + return total / max(n, 1) + + +def pca_baseline_mse(ds: InSituDataset, latent_dim: int = LATENT_DIM) -> float: + """Mean reconstruction MSE of the PCA baseline (fit on train, scored on val). + + Like astroML's ``spec4000`` workflow, PCA needs every training spectrum resident to fit -- the + opposite of the streamed autoencoder. That contrast is the point of the example. + """ + _, train_clean = collect(ds, "train") + val_noisy, val_clean = collect(ds, "val") + mean, components = fit_pca(train_clean, k=latent_dim) + return recon_mse(pca_reconstruct(val_noisy, mean, components), val_clean) + + +def train(ds: InSituDataset, *, n_wave: int, epochs: int, device: str = "cpu") -> nn.Module: + model = AutoEncoder(n_wave).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_mse {val:.4f}") + return model + + +def spectrum_width(ds: InSituDataset) -> int: + """Wavelength-bin count, read from the first val batch (the clean spectrum width).""" + for batch in ds.val: + return int(batch.arrays["clean"].shape[-1]) + raise RuntimeError("empty dataset: cannot infer spectrum width") + + +def main(argv: list[str] | None = None) -> None: + args = cli(argv) + ds = build_datasets(args) + n_wave = spectrum_width(ds) + + base = pca_baseline_mse(ds, latent_dim=args.latent_dim) + print(f"PCA baseline val MSE: {base:.4f} (latent dim {args.latent_dim}, needs all spectra)\n") + model = train(ds, n_wave=n_wave, epochs=args.epochs) + print(f"\nautoencoder val MSE: {evaluate(model, ds):.4f}") + print(f"PCA baseline val MSE: {base:.4f}") + ds.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index baa22fe..65d2024 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,6 +102,25 @@ gcsfs = [ arraylake = [ "arraylake>=0.14", ] +# Build-time FITS index stack for the astronomy examples (Hubble images, SDSS +# spectra): kerchunk.fits parses each FITS header, VirtualiZarr exposes the +# byte-range references, and they are committed to an Icechunk store. Needed only +# to BUILD the reference store -- the training hot path is icechunk + numpy, never +# kerchunk/astropy. astropy also backs the ground-truth read in tests. The SDSS +# (binary-table / structured-dtype) path additionally requires an unreleased +# upstream fix, pinned to git in [tool.uv.sources] below. +astronomy = [ + "virtualizarr>=2.6", + # ObjectStoreRegistry, which maps a URL prefix to an obstore Store for the FITS + # parser. VirtualiZarr re-exported it until v2.7.1 deprecated the shim in favour of + # this package (its own dependency), so we import from the source and declare it + # rather than lean on a re-export that is scheduled for removal. + "obspec-utils>=0.9", + "kerchunk>=0.2.10", + "astropy>=8.0", + "s3fs>=2024.6", + "icechunk>=0.2", +] [dependency-groups] dev = [ @@ -116,6 +135,20 @@ dev = [ requires = ["uv_build>=0.11.21,<0.12.0"] build-backend = "uv_build" +# Temporary git pin for FITS binary-table (structured-dtype) support, exercised by +# tests/test_fits_bintable.py. This entry is dev-environment only (uv strips it from the +# published wheel, so downstream requirements are unchanged). Drop it once a release ships. +# - virtualizarr: kerchunk struct-dtype refs for FITS binary tables +# (zarr-developers/VirtualiZarr#1037, merged 2026-07-15). Released VirtualiZarr raises on +# the list-valued field specs; the fix ships in v2.7.1, so this stays pinned to a main +# commit until that release, then becomes a plain `virtualizarr>=2.7.1` in [astronomy]. +# NOTE: the related zarr struct-dtype byte-order fix (zarr-developers/zarr-python#4142) is NOT +# pinned -- it fixes the direct `zarr.create_array('>f4')` path, but VirtualiZarr/icechunk store +# the struct fields with native-order labels, so the codec fix never engages: consumers reinterpret +# each field with `.view('>f4')` at the numpy boundary regardless of the zarr version. +[tool.uv.sources] +virtualizarr = { git = "https://github.com/zarr-developers/VirtualiZarr.git", rev = "f46902d2d3346bff31c88bf7212a14a2a50934bc" } + [tool.ruff] line-length = 100 target-version = "py312" diff --git a/tests/test_fits_bintable.py b/tests/test_fits_bintable.py new file mode 100644 index 0000000..f6c795b --- /dev/null +++ b/tests/test_fits_bintable.py @@ -0,0 +1,149 @@ +"""Validate FITS *binary-table* (structured big-endian dtype) support end-to-end via insitubatch. + +The SDSS spectra example streams spPlate *images*, but the archival format also carries binary +tables (per-object spectra, catalogs) as big-endian structured dtypes. Reading those as virtual +references exercises a fragile path: ``kerchunk.fits`` maps the table to one structured-dtype zarr +array, and VirtualiZarr must translate the kerchunk refs into a zarr-v3 structured dtype. That +translation needs the fix in ``zarr-developers/VirtualiZarr#1037`` (pinned in ``[tool.uv.sources]`` +via the ``astronomy`` extra): without it, ``from_kerchunk_refs`` raises on the list-valued field +specs and base64 fill value. + +This offline test synthesizes a big-endian ``BinTable`` (no network, no SDSS), runs it through the +whole chain -- kerchunk -> VirtualiZarr -> Icechunk -> insitubatch -- and checks the delivered field +values against the astropy ground truth. It is the drift guard that keeps FITS binary-table support +(and the pinned VZ fix) exercised. + +Endianness note: VirtualiZarr/icechunk store the struct fields with native (little-endian) *labels* +while the referenced FITS bytes stay big-endian, so a consumer reinterprets each field with +``.view('>f4')`` at the numpy boundary before handing native-order data to a framework (DLPack +rejects both structured dtypes and non-native byte order). That projection is the documented, +vectorized boundary step -- it is required independently of the zarr codec, so this test does not +depend on the zarr byte-order pin. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The FITS binary-table chain needs the whole build-time stack (the `astronomy` extra). +pytest.importorskip("virtualizarr") +pytest.importorskip("kerchunk") +pytest.importorskip("astropy") +pytest.importorskip("icechunk") + +from insitubatch import InSituDataset, open_geometries, split_by_chunk # noqa: E402 + + +def _build_bintable_store(tmp_dir: str) -> tuple[str, np.ndarray]: + """Write a big-endian FITS BinTable and index it into a local Icechunk repo of virtual refs. + + Returns ``(store_path, ground_truth_flux)``. + """ + import icechunk + from astropy.io import fits + from obspec_utils.registry import ObjectStoreRegistry + from obstore.store import LocalStore + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + + n = 64 + flux = (np.sin(np.linspace(0.0, 6.0, n)) * 10.0).astype(">f4") # big-endian, as in real FITS + ivar = np.linspace(1.0, 2.0, n).astype(">f4") + cols = fits.ColDefs( + [ + fits.Column(name="flux", format="E", array=flux), + fits.Column(name="ivar", format="E", array=ivar), + ] + ) + fits.HDUList([fits.PrimaryHDU(), fits.BinTableHDU.from_columns(cols, name="COADD")]).writeto( + f"{tmp_dir}/spec.fits" + ) + + prefix = f"file://{tmp_dir}" + registry = ObjectStoreRegistry({prefix: LocalStore(prefix=tmp_dir)}) + vds = open_virtual_dataset(url=f"{prefix}/spec.fits", registry=registry, parser=FITSParser()) + + ice_prefix = prefix + "/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.local_filesystem_store(tmp_dir)) + ) + repo = icechunk.Repository.create( + icechunk.local_filesystem_storage(f"{tmp_dir}/repo"), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), + ) + session = repo.writable_session("main") + vds.vz.to_icechunk(session.store) + session.commit("index bintable") + return f"{tmp_dir}/repo", flux + + +def test_fits_bintable_roundtrips_through_insitubatch(tmp_path) -> None: + import icechunk + + store_path, flux_truth = _build_bintable_store(str(tmp_path)) + + ice_prefix = f"file://{tmp_path}/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.local_filesystem_store(str(tmp_path))) + ) + repo = icechunk.Repository.open( + icechunk.local_filesystem_storage(store_path), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), + ) + store = repo.readonly_session("main").store + + geoms = open_geometries(store, sample_axis=0) + (var,) = geoms # kerchunk names the HDU variable (here "COADD") + geom = geoms[var] + assert geom.n_samples == flux_truth.size + assert geom.dtype.names == ("flux", "ivar") # a structured (record) dtype survived the chain + + ds = InSituDataset( + store, + split_by_chunk(geom, fractions=(1.0, 0.0, 0.0)), + geometries=geoms, + batch_size=64, + shuffle=False, + ) + got: dict[int, np.ndarray] = {} + for batch in ds.train: + rows = batch.arrays[var] + for k, idx in enumerate(batch.sample_indices): + got[int(idx)] = rows[k] + ds.close() + + assert len(got) == flux_truth.size + delivered = np.array([got[i] for i in range(flux_truth.size)]) + # The referenced bytes are big-endian; reinterpret the field at the numpy boundary. + flux = delivered["flux"].view(">f4").astype(np.float32) + np.testing.assert_allclose(flux, flux_truth.astype(np.float32)) + + +def test_fits_bintable_endianness_note_is_load_bearing(tmp_path) -> None: + # Guard the documented endianness behaviour: the native-label field reads as garbage, and the + # ``.view('>f4')`` reinterpret is what recovers the values. If a future zarr/VZ release delivers + # correctly-swapped native data, THIS test flips -- the signal to drop the workaround. + import icechunk + + store_path, flux_truth = _build_bintable_store(str(tmp_path)) + ice_prefix = f"file://{tmp_path}/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(ice_prefix, icechunk.local_filesystem_store(str(tmp_path))) + ) + repo = icechunk.Repository.open( + icechunk.local_filesystem_storage(store_path), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), + ) + import zarr + + (var,) = open_geometries(repo.readonly_session("main").store, sample_axis=0) + raw = zarr.open_array(repo.readonly_session("main").store, path=var, mode="r")[:] + assert not np.allclose(raw["flux"].astype("f8"), flux_truth.astype("f8")) # native label: wrong + assert np.allclose(raw["flux"].view(">f4").astype("f8"), flux_truth.astype("f8")) # reinterpret diff --git a/tests/test_sdss.py b/tests/test_sdss.py new file mode 100644 index 0000000..d7125ef --- /dev/null +++ b/tests/test_sdss.py @@ -0,0 +1,94 @@ +"""SDSS spectral-reconstruction example: offline synthetic spectra guard the stream+train path. + +The real example indexes an SDSS ``spPlate`` FITS over HTTPS as virtual references (build-time +VirtualiZarr/kerchunk/astropy + network). These tests instead drive a synthetic ``flux +(fiber, wave)`` zarr with the *same* geometry (many fibers 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, the PCA baseline, 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.sdss.data import ( + FLUX_VAR, + fit_pca, + make_synthetic_store, + pca_reconstruct, + recon_mse, + reconstruct_dataset, +) +from insitubatch import obstore_store, open_geometries + + +@pytest.fixture +def synth_store(tmp_path) -> str: + """A small synthetic spPlate-like store: 8 plates x 96 fibers, NaN pixels, redshift lines.""" + url = f"file://{tmp_path}/sdss.zarr" + make_synthetic_store(url, n_plates=8, fibers_per_plate=96, n_wave=256, seed=0) + return url + + +def test_many_fibers_per_chunk(synth_store) -> None: + # The spPlate-derived geometry: one plate is one chunk, and a chunk holds many fiber samples + # (the decode-amortization regime, unlike Hubble's one image per chunk). + geoms = open_geometries(obstore_store(synth_store), variables=[FLUX_VAR], sample_axis=0) + geom = geoms[FLUX_VAR] + assert geom.n_samples == 8 * 96 + assert geom.sample_chunk_size == 96 + assert geom.n_chunks == 8 + + +def test_reconstruct_batch_shapes_and_scrubbing(synth_store) -> None: + ds = reconstruct_dataset(obstore_store(synth_store), batch_size=32, shuffle=False) + ds.set_epoch(0) + batch = next(iter(ds.train)) + + assert set(batch.arrays) == {"noisy", "clean"} + assert batch.arrays["clean"].shape == (32, 256) + assert batch.arrays["noisy"].shape == batch.arrays["clean"].shape + # 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 + # Corrupt actually perturbs the spectrum (the batch stage ran). + assert not np.allclose(batch.arrays["noisy"], batch.arrays["clean"]) + + +def test_split_is_by_chunk(synth_store) -> None: + # No fiber leaks across train/val/test: splits partition whole plates (chunks). + ds = reconstruct_dataset(obstore_store(synth_store), batch_size=64, shuffle=False) + ds.set_epoch(0) + seen = {s: set() for s in ("train", "val", "test")} + for split in seen: + for b in getattr(ds, split): + seen[split].update(int(i) for i in b.sample_indices) + assert seen["train"] and seen["val"] and seen["test"] + assert not (seen["train"] & seen["val"]) + assert not (seen["train"] & seen["test"]) + assert not (seen["val"] & seen["test"]) + + +def test_pca_baseline_reconstructs(synth_store) -> None: + # The no-training baseline must beat the trivial mean-spectrum reconstruction. + ds = reconstruct_dataset(obstore_store(synth_store), batch_size=64, shuffle=False) + ds.set_epoch(0) + clean = np.concatenate([b.arrays["clean"] for b in ds.train]) + noisy = np.concatenate([b.arrays["noisy"] for b in ds.train]) + mean, comps = fit_pca(clean, k=16) + trivial = np.broadcast_to(mean, clean.shape) + assert recon_mse(pca_reconstruct(noisy, mean, comps), clean) < recon_mse(trivial, clean) + + +def test_torch_beats_baseline(synth_store) -> None: + pytest.importorskip("torch") + from examples.sdss.train_torch import evaluate, pca_baseline_mse, spectrum_width, train + + ds = reconstruct_dataset(obstore_store(synth_store), batch_size=32) + base = pca_baseline_mse(ds, latent_dim=16) + model = train(ds, n_wave=spectrum_width(ds), epochs=25) + # A nonlinear conv autoencoder beats linear PCA at the same latent dim on the redshift manifold. + assert evaluate(model, ds) < base diff --git a/tests/test_sdss_build.py b/tests/test_sdss_build.py new file mode 100644 index 0000000..af8a2a8 --- /dev/null +++ b/tests/test_sdss_build.py @@ -0,0 +1,144 @@ +"""Guard the SDSS spPlate virtual-reference build modes -- offline, with synthetic FITS images. + +The real example indexes SDSS ``spPlate`` frames over HTTPS; :func:`build_store` re-chunks them by +rewriting the *chunk manifest* (byte arithmetic), never the pixels. Two modes: + +* :func:`_single_plate_fiber_chunks` -- one plate -> contiguous fiber-block chunks (full width). +* :func:`_many_plates_common_grid` -- N plates cropped to a shared wavelength window -> one flat + fiber axis, one fiber per chunk. + +These tests write tiny big-endian FITS images with SDSS-style ``COEFF0``/``COEFF1`` grid headers +(no network), run each mode through Icechunk + insitubatch, and check the delivered spectra are +byte-identical to the source -- so the manifest byte-offset math can't drift. Need the ``astronomy`` +extra. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("virtualizarr") +pytest.importorskip("kerchunk") +pytest.importorskip("astropy") +pytest.importorskip("icechunk") + +from examples.sdss.data import ( # noqa: E402 + FLUX_VAR, + _many_plates_common_grid, + _single_plate_fiber_chunks, +) +from insitubatch import InSituDataset, open_geometries, split_by_chunk # noqa: E402 + +DLOGLAM = 1e-4 + + +def _write_plate( + tmp_dir: str, name: str, *, n_fiber: int, n_wave: int, coeff0: float, seed: int +) -> np.ndarray: + """Write a synthetic spPlate-like FITS: a big-endian ``(fiber, wave)`` image on a log grid.""" + from astropy.io import fits + + rng = np.random.default_rng(seed) + data = (10.0 + rng.normal(0.0, 1.0, size=(n_fiber, n_wave))).astype(">f4") + hdu = fits.PrimaryHDU(data) + hdu.header["COEFF0"] = coeff0 + hdu.header["COEFF1"] = DLOGLAM + hdu.writeto(f"{tmp_dir}/{name}") + return data.astype(np.float32) + + +def _open_plate(tmp_dir: str, name: str): # -> xr.Dataset + from obspec_utils.registry import ObjectStoreRegistry + from obstore.store import LocalStore + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + + prefix = f"file://{tmp_dir}" + registry = ObjectStoreRegistry({prefix: LocalStore(prefix=tmp_dir)}) + vds = open_virtual_dataset(url=f"{prefix}/{name}", registry=registry, parser=FITSParser()) + return vds.rename({"PRIMARY": FLUX_VAR}) + + +def _commit_and_open(virtual, tmp_dir: str): # -> zarr Store + import icechunk + + prefix = f"file://{tmp_dir}/" + config = icechunk.RepositoryConfig.default() + config.set_virtual_chunk_container( + icechunk.VirtualChunkContainer(prefix, icechunk.local_filesystem_store(tmp_dir)) + ) + repo = icechunk.Repository.create( + icechunk.local_filesystem_storage(f"{tmp_dir}/repo"), + config=config, + authorize_virtual_chunk_access=icechunk.containers_credentials({prefix: None}), + ) + session = repo.writable_session("main") + virtual.vz.to_icechunk(session.store) + session.commit("build") + return repo.readonly_session("main").store + + +def _read_all(store) -> dict[int, np.ndarray]: + geoms = open_geometries(store, variables=[FLUX_VAR], sample_axis=0) + ds = InSituDataset( + store, + split_by_chunk(geoms[FLUX_VAR], fractions=(1.0, 0.0, 0.0)), + geometries=geoms, + batch_size=256, + shuffle=False, + ) + got: dict[int, np.ndarray] = {} + for batch in ds.train: + for k, idx in enumerate(batch.sample_indices): + got[int(idx)] = batch.arrays[FLUX_VAR][k] + ds.close() + return got + + +def test_single_plate_fiber_chunks(tmp_path) -> None: + truth = _write_plate(str(tmp_path), name="p.fits", n_fiber=12, n_wave=20, coeff0=3.58, seed=0) + plate = _open_plate(str(tmp_path), "p.fits") + + store = _commit_and_open(_single_plate_fiber_chunks(plate, fibers_per_chunk=4), str(tmp_path)) + geom = open_geometries(store, sample_axis=0)[FLUX_VAR] + assert geom.n_samples == 12 + assert geom.sample_chunk_size == 4 # many fibers per chunk (decode-amortization) + assert geom.n_chunks == 3 + + got = _read_all(store) + for i in range(12): + np.testing.assert_array_equal(got[i], truth[i]) + + +def test_many_plates_common_grid_aligns_and_concats(tmp_path) -> None: + # Two plates offset by whole bins with different widths: cropping to the common window must + # align them onto one grid exactly and concat into a flat fiber axis. + a = _write_plate(str(tmp_path), name="a.fits", n_fiber=6, n_wave=24, coeff0=3.5800, seed=1) + b = _write_plate(str(tmp_path), name="b.fits", n_fiber=6, n_wave=22, coeff0=3.5803, seed=2) + plates = [_open_plate(str(tmp_path), "a.fits"), _open_plate(str(tmp_path), "b.fits")] + + virtual = _many_plates_common_grid(plates) + store = _commit_and_open(virtual, str(tmp_path)) + geom = open_geometries(store, sample_axis=0)[FLUX_VAR] + assert geom.n_samples == 12 # 6 + 6 fibers, folded into one axis + assert geom.sample_chunk_size == 1 # one fiber per chunk (streaming) + + # common window: lo = max start = 3.5803; a starts 3 bins earlier, b at bin 0. + off_a, off_b = 3, 0 + width = geom.inner_shape[0] + got = _read_all(store) + for f in range(6): + np.testing.assert_array_equal(got[f], a[f, off_a : off_a + width]) # plate a fibers + np.testing.assert_array_equal(got[6 + f], b[f, off_b : off_b + width]) # plate b fibers + # both plates now share bin 0 == loglam 3.5803, so the grids are aligned, not merely truncated. + assert width == min(24 - off_a, 22 - off_b) + + +def test_many_plates_rejects_unaligned_grids(tmp_path) -> None: + # A half-bin COEFF0 offset cannot be cropped losslessly -> fail fast, do not silently misalign. + _write_plate(str(tmp_path), name="a.fits", n_fiber=4, n_wave=16, coeff0=3.5800, seed=3) + _write_plate(str(tmp_path), name="b.fits", n_fiber=4, n_wave=16, coeff0=3.58005, seed=4) + plates = [_open_plate(str(tmp_path), "a.fits"), _open_plate(str(tmp_path), "b.fits")] + with pytest.raises(ValueError, match="not bin-aligned"): + _many_plates_common_grid(plates) diff --git a/uv.lock b/uv.lock index 5057472..df4f652 100644 --- a/uv.lock +++ b/uv.lock @@ -25,6 +25,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" }, ] +[[package]] +name = "aiobotocore" +version = "3.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "aioitertools" }, + { name = "botocore" }, + { name = "jmespath" }, + { name = "multidict" }, + { name = "python-dateutil" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e7/75/42cce839c2ec263ff74b10b650fe36b066fbb124cbee6f247eac0983e1ab/aiobotocore-3.7.0.tar.gz", hash = "sha256:c64d871ed5491a6571948dd48eabd185b46c6c23b64e3afd0c059fc7593ada30", size = 127054, upload-time = "2026-05-09T10:02:52.332Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/5f/85535dfb3cfd6442d66d1df1694062c5d6df02f895329e7e120b2a3d2b8b/aiobotocore-3.7.0-py3-none-any.whl", hash = "sha256:680bde7c64679a821a9312641b759d9497f790ba8b2e88c6959e6273ee765b8e", size = 89539, upload-time = "2026-05-09T10:02:50.389Z" }, +] + [[package]] name = "aiofiles" version = "25.1.0" @@ -143,6 +161,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/71/6e22be134a4061ada85a92951b842f2657f17d926b727f3f94c56ae963d6/aiohttp-3.14.1-cp314-cp314t-win_arm64.whl", hash = "sha256:90d53f1609c29ccc2193945ef732428382a28f78d0456ae4d3daf0d48b74f0f6", size = 469640, upload-time = "2026-06-07T21:09:33.028Z" }, ] +[[package]] +name = "aioitertools" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/3c/53c4a17a05fb9ea2313ee1777ff53f5e001aefd5cc85aa2f4c2d982e1e38/aioitertools-0.13.0.tar.gz", hash = "sha256:620bd241acc0bbb9ec819f1ab215866871b4bbd1f73836a55f799200ee86950c", size = 19322, upload-time = "2025-11-06T22:17:07.609Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl", hash = "sha256:0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be", size = 24182, upload-time = "2025-11-06T22:17:06.502Z" }, +] + [[package]] name = "aiosignal" version = "1.4.0" @@ -251,6 +278,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/19/cc8bd127d28a43da249aa955cfd164cf8fd534e79e42cea96c4854d72fd0/ast_serialize-0.5.0-cp39-abi3-win_arm64.whl", hash = "sha256:92a31c9c20d25a076edaeec76b128a3535d74a24f340b9a8a7e96c9b86dc9642", size = 1081181, upload-time = "2026-05-17T17:48:28.122Z" }, ] +[[package]] +name = "astropy" +version = "8.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "astropy-iers-data" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyerfa" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/c4/21be4313ddfde5f60e0607fd307f367b9e0f0bf153a89b10cbd036dd8cfd/astropy-8.0.1.tar.gz", hash = "sha256:45ca31d5b91fa294cd590a4791a32db94de7f9c8a343155f4d5877baa82351da", size = 7152500, upload-time = "2026-07-05T07:24:48.482Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/80/69a84af0b35d55a859cde38e16b71553840a156bdf7dc4b767ec2b2d2829/astropy-8.0.1-cp311-abi3-macosx_10_9_x86_64.whl", hash = "sha256:24813c764d5ea111e7f716127121e22a0cfe233ca1eac2ed23c4a3848390610e", size = 6635649, upload-time = "2026-07-05T07:24:31.957Z" }, + { url = "https://files.pythonhosted.org/packages/a6/34/074f367d5699a1008b24c50acc1539f05f2cfc881b1901f4b110d57eae20/astropy-8.0.1-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:222f0b9837e79fd3d101e6d4e4579e3aa98c490000bb01e35cd32fd3e3c12bf5", size = 6607950, upload-time = "2026-07-05T07:24:34.148Z" }, + { url = "https://files.pythonhosted.org/packages/56/3b/d29035da0a08e3b0b03d63abe9ce69b900d8ef7863dc557a2e34ebf56a4e/astropy-8.0.1-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9a4044396c15969bb029b648521f7c25d294e982dbd23e7c2c7dd328178b8b98", size = 10414634, upload-time = "2026-07-05T07:24:36.129Z" }, + { url = "https://files.pythonhosted.org/packages/fa/14/6f4427419f8a02e1d0a885fb64d8d4454c5d443627aa7cb1a3a36e16abb7/astropy-8.0.1-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fa11d56855e10107ea2231a6b6a33dbf1edbea6890adf34634c1f1d8f25c5a5a", size = 10452095, upload-time = "2026-07-05T07:24:38.424Z" }, + { url = "https://files.pythonhosted.org/packages/b2/ef/a5cf36a402c4511a776405f12d0b35196f55f4312ce407012a2fbcf1e4e0/astropy-8.0.1-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5b482bc6c57c966e6c6234410a1d9afbcf92bcac858cf287812f8c99ddc3fafc", size = 10407732, upload-time = "2026-07-05T07:24:40.774Z" }, + { url = "https://files.pythonhosted.org/packages/58/6d/59e7fb6f9ade53889c9553afff89991b8076c021581ba66eb51d913ca6c9/astropy-8.0.1-cp311-abi3-win32.whl", hash = "sha256:13a6ef59347a15c55406bb630e0d25d0248a179af84fae06a7ddb85a229065b2", size = 6394540, upload-time = "2026-07-05T07:24:43.146Z" }, + { url = "https://files.pythonhosted.org/packages/ec/6a/8534fd98e2d3eb51cbaa261e6c6b59e4ad969a1787d72bb1208499080dbd/astropy-8.0.1-cp311-abi3-win_amd64.whl", hash = "sha256:ee9024f2237243ebc98726f8eff2efbf201ea4e0bba1d21f6e2dee0cb499fbdf", size = 6529976, upload-time = "2026-07-05T07:24:44.693Z" }, + { url = "https://files.pythonhosted.org/packages/8a/df/96b8bf4ec218d03969a8b904ade5fc5b1a63f4a7f02c4e107b6264bcf152/astropy-8.0.1-cp311-abi3-win_arm64.whl", hash = "sha256:435e784f2cd9a4c31a03102e6ad27a51b901a8d5587bc78afe8d28e41eb44824", size = 6395252, upload-time = "2026-07-05T07:24:46.56Z" }, +] + +[[package]] +name = "astropy-iers-data" +version = "0.2026.7.13.0.54.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/f2/8e52b70264edc5664c069939224405c6e9085fa3388403b8313303a62b82/astropy_iers_data-0.2026.7.13.0.54.2.tar.gz", hash = "sha256:d86e32e95e98a86f83b5b073627442924a0f45cf61a7828da4f565a17d726c2f", size = 1941259, upload-time = "2026-07-13T00:54:51.205Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/54/121240c06fff63d358f2a96187e927d05628fdf933a84246151954e95c7a/astropy_iers_data-0.2026.7.13.0.54.2-py3-none-any.whl", hash = "sha256:0f0b22f43d0917d78382f35ef13acd72600752d65289a3f91e39ea67653264cd", size = 1996278, upload-time = "2026-07-13T00:54:49.386Z" }, +] + [[package]] name = "astunparse" version = "1.6.3" @@ -295,6 +354,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/5c/fb93d3092640a24dfb7bd7727a24016d7c01774ca013e60efd3f683c8002/backrefs-7.0-py314-none-any.whl", hash = "sha256:a6448b28180e3ca01134c9cf09dcebafad8531072e09903c5451748a05f24bc9", size = 412349, upload-time = "2026-04-28T16:28:02.412Z" }, ] +[[package]] +name = "botocore" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/28/79/2f4be1896db3db7ccf44504253a175d56b6bd6b669619edc5147d1aa21ea/botocore-1.43.0.tar.gz", hash = "sha256:e933b31a2d644253e1d029d7d39e99ba41b87e29300534f189744cc438cdf928", size = 15286817, upload-time = "2026-04-29T22:07:31.723Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bf/4b/afc1fef8a43bafb139f57f73bbd70df82807af5934321e8112ae50668827/botocore-1.43.0-py3-none-any.whl", hash = "sha256:cc5b15eaec3c6eac05d8012cb5ef17ebe891beb88a16ca13c374bfaece1241e6", size = 14970102, upload-time = "2026-04-29T22:07:27Z" }, +] + [[package]] name = "certifi" version = "2026.5.20" @@ -1263,6 +1336,14 @@ dependencies = [ arraylake = [ { name = "arraylake" }, ] +astronomy = [ + { name = "astropy" }, + { name = "icechunk" }, + { name = "kerchunk" }, + { name = "obspec-utils" }, + { name = "s3fs" }, + { name = "virtualizarr" }, +] bench = [ { name = "gcsfs" }, { name = "plotly" }, @@ -1316,32 +1397,38 @@ dev = [ [package.metadata] requires-dist = [ { name = "arraylake", marker = "extra == 'arraylake'", specifier = ">=0.14" }, + { name = "astropy", marker = "extra == 'astronomy'", specifier = ">=8.0" }, { name = "cloudpickle", marker = "extra == 'cache'", specifier = ">=3" }, { name = "cupy-cuda12x", marker = "extra == 'gpu'", specifier = ">=13.0" }, { name = "flax", marker = "extra == 'jax'", specifier = ">=0.8" }, { name = "gcsfs", marker = "extra == 'bench'", specifier = ">=2024.6" }, { name = "gcsfs", marker = "extra == 'gcsfs'", specifier = ">=2024.6" }, + { name = "icechunk", marker = "extra == 'astronomy'", specifier = ">=0.2" }, { name = "icechunk", marker = "extra == 'icechunk'", specifier = ">=0.2" }, { name = "jax", marker = "extra == 'jax'", specifier = ">=0.4.30" }, + { name = "kerchunk", marker = "extra == 'astronomy'", specifier = ">=0.2.10" }, { name = "kvikio-cu12", marker = "extra == 'gpu'", specifier = ">=24.10" }, { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.5" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.26" }, { name = "numpy", specifier = ">=1.26" }, + { name = "obspec-utils", marker = "extra == 'astronomy'", specifier = ">=0.9" }, { name = "obstore", specifier = ">=0.3" }, { name = "optax", marker = "extra == 'jax'", specifier = ">=0.2" }, { name = "plotly", marker = "extra == 'bench'", specifier = ">=5.20" }, { name = "psutil", marker = "extra == 'bench'", specifier = ">=5.9" }, { name = "py-spy", marker = "extra == 'bench'", specifier = ">=0.4" }, + { name = "s3fs", marker = "extra == 'astronomy'", specifier = ">=2024.6" }, { name = "scikit-learn", marker = "extra == 'bench'", specifier = ">=1.4" }, { name = "tensorflow", marker = "extra == 'tf'", specifier = ">=2.16" }, { name = "torch", marker = "extra == 'torch'", specifier = ">=2.2" }, { name = "torchdata", marker = "extra == 'torch'", specifier = ">=0.10" }, - { name = "virtualizarr", marker = "extra == 'virtual'", specifier = ">=1.2" }, + { name = "virtualizarr", marker = "extra == 'astronomy'", git = "https://github.com/zarr-developers/VirtualiZarr.git?rev=f46902d2d3346bff31c88bf7212a14a2a50934bc" }, + { name = "virtualizarr", marker = "extra == 'virtual'", git = "https://github.com/zarr-developers/VirtualiZarr.git?rev=f46902d2d3346bff31c88bf7212a14a2a50934bc" }, { name = "xarray", specifier = ">=2024.9" }, { name = "xbatcher", marker = "extra == 'bench'", specifier = ">=0.3" }, { name = "zarr", specifier = ">=3.0" }, ] -provides-extras = ["torch", "jax", "tf", "gpu", "virtual", "cache", "bench", "docs", "icechunk", "gcsfs", "arraylake"] +provides-extras = ["torch", "jax", "tf", "gpu", "virtual", "cache", "bench", "docs", "icechunk", "gcsfs", "arraylake", "astronomy"] [package.metadata.requires-dev] dev = [ @@ -1410,6 +1497,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joblib" version = "1.5.3" @@ -1438,6 +1534,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/03/184267c1d09783dd070f1ddfd0d4beb7503139dfc7bd75b422867cf282fd/keras-3.14.1-py3-none-any.whl", hash = "sha256:ebd2c14d2af3c9de18083604d408483996407fc7d2f9ebd1d565961f96608c29", size = 1628606, upload-time = "2026-05-07T21:43:32.737Z" }, ] +[[package]] +name = "kerchunk" +version = "0.2.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fsspec" }, + { name = "numcodecs" }, + { name = "numpy" }, + { name = "ujson" }, + { name = "zarr" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/a0/ebeca522912e68f360117404a6b2f740d4b0a343d98e9586377a2fd21567/kerchunk-0.2.10.tar.gz", hash = "sha256:aae63c0fe4ca2e97025f026578a0577545011fe6679751392c815e0d1d6bf954", size = 716526, upload-time = "2026-03-30T13:49:05.303Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/e4/3c356a9ea448a48caa5e44cd51293f7e896cd606f2ef86da96f5d61cc427/kerchunk-0.2.10-py3-none-any.whl", hash = "sha256:7fdaa77dae25c75d3ec9402c49208f37ae51d346ab082724e3e32608438f8c66", size = 68490, upload-time = "2026-03-30T13:49:03.855Z" }, +] + [[package]] name = "kvikio-cu12" version = "26.6.0" @@ -2956,6 +3068,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, ] +[[package]] +name = "pyerfa" +version = "2.0.1.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/39/63cc8291b0cf324ae710df41527faf7d331bce573899199d926b3e492260/pyerfa-2.0.1.5.tar.gz", hash = "sha256:17d6b24fe4846c65d5e7d8c362dcb08199dc63b30a236aedd73875cc83e1f6c0", size = 818430, upload-time = "2024-11-11T15:22:30.852Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/d9/3448a57cb5bd19950de6d6ab08bd8fbb3df60baa71726de91d73d76c481b/pyerfa-2.0.1.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:b282d7c60c4c47cf629c484c17ac504fcb04abd7b3f4dfcf53ee042afc3a5944", size = 341818, upload-time = "2024-11-11T15:22:16.467Z" }, + { url = "https://files.pythonhosted.org/packages/11/4a/31a363370478b63c6289a34743f2ba2d3ae1bd8223e004d18ab28fb92385/pyerfa-2.0.1.5-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:be1aeb70390dd03a34faf96749d5cabc58437410b4aab7213c512323932427df", size = 329370, upload-time = "2024-11-11T15:22:17.829Z" }, + { url = "https://files.pythonhosted.org/packages/cb/96/b6210fc624123c8ae13e1eecb68fb75e3f3adff216d95eee1c7b05843e3e/pyerfa-2.0.1.5-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b0603e8e1b839327d586c8a627cdc634b795e18b007d84f0cda5500a0908254e", size = 692794, upload-time = "2024-11-11T15:22:19.429Z" }, + { url = "https://files.pythonhosted.org/packages/e5/e0/050018d855d26d3c0b4a7d1b2ed692be758ce276d8289e2a2b44ba1014a5/pyerfa-2.0.1.5-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0e43c7194e3242083f2350b46c09fd4bf8ba1bcc0ebd1460b98fc47fe2389906", size = 738711, upload-time = "2024-11-11T15:22:20.661Z" }, + { url = "https://files.pythonhosted.org/packages/b9/f5/ff91ee77308793ae32fa1e1de95e9edd4551456dd888b4e87c5938657ca5/pyerfa-2.0.1.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:07b80cd70701f5d066b1ac8cce406682cfcd667a1186ec7d7ade597239a6021d", size = 722966, upload-time = "2024-11-11T15:22:21.905Z" }, + { url = "https://files.pythonhosted.org/packages/2c/56/b22b35c8551d2228ff8d445e63787112927ca13f6dc9e2c04f69d742c95b/pyerfa-2.0.1.5-cp39-abi3-win32.whl", hash = "sha256:d30b9b0df588ed5467e529d851ea324a67239096dd44703125072fd11b351ea2", size = 339955, upload-time = "2024-11-11T15:22:23.087Z" }, + { url = "https://files.pythonhosted.org/packages/b4/11/97233cf23ad5411ac6f13b1d6ee3888f90ace4f974d9bf9db887aa428912/pyerfa-2.0.1.5-cp39-abi3-win_amd64.whl", hash = "sha256:66292d437dcf75925b694977aa06eb697126e7b86553e620371ed3e48b5e0ad0", size = 349410, upload-time = "2024-11-11T15:22:24.817Z" }, +] + [[package]] name = "pygments" version = "2.20.0" @@ -3174,6 +3304,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" }, ] +[[package]] +name = "s3fs" +version = "2026.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiobotocore" }, + { name = "aiohttp" }, + { name = "fsspec" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/99/00/6677343dc919d6c072bb04d80210afdd22c16838a8d16b3315c122dc728f/s3fs-2026.6.0.tar.gz", hash = "sha256:b28de7082d0a4f72392884bdc497e34a4a1582f675d214c7da0acf6e950a0083", size = 87358, upload-time = "2026-06-16T02:05:48.719Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/0b/f68a968b49876eae0f2a515387093cebb2eb9451380a96741cc20efac0d0/s3fs-2026.6.0-py3-none-any.whl", hash = "sha256:60576e31bb31193c1f643f32b4c6439548720ea6918ac702e21cd757c80b5db8", size = 32573, upload-time = "2026-06-16T02:05:47.608Z" }, +] + [[package]] name = "scikit-learn" version = "1.9.0" @@ -3717,8 +3861,8 @@ wheels = [ [[package]] name = "virtualizarr" -version = "2.6.2" -source = { registry = "https://pypi.org/simple" } +version = "2.7.1.dev14+gf46902d2d" +source = { git = "https://github.com/zarr-developers/VirtualiZarr.git?rev=f46902d2d3346bff31c88bf7212a14a2a50934bc#f46902d2d3346bff31c88bf7212a14a2a50934bc" } dependencies = [ { name = "numcodecs" }, { name = "numpy" }, @@ -3730,10 +3874,6 @@ dependencies = [ { name = "xarray" }, { name = "zarr" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/85/82/6a97e5319836c5cd8c8c24cb1ac087f6a57373da54643e200417bb40d453/virtualizarr-2.6.2.tar.gz", hash = "sha256:cc6df21fa0acdeae50a9f7d170ea4fd9efe4fb8ee7061a70d9ddaef82a8b443e", size = 281960, upload-time = "2026-05-18T22:01:51.575Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/83/be/b2a5639c2bdd3aa2e63aafe2a93cf2fb2ae738707579decb3047457dde99/virtualizarr-2.6.2-py3-none-any.whl", hash = "sha256:8a128f4c2ea71eb8a7ef7e0c9ccfa173c9c313089834adc641999d89f75d5132", size = 239323, upload-time = "2026-05-18T22:01:49.84Z" }, -] [[package]] name = "watchdog"