From 70de7a36162be24b321953c72baeaaefbb2e44ac Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:30:45 +0000 Subject: [PATCH 1/4] build: pin unreleased zarr/VirtualiZarr fixes + add astronomy extra The SDSS/spectra astronomy example reads FITS binary tables (structured dtypes) end-to-end, which needs two upstream fixes not yet in a release: - zarr struct-dtype byte-order handling in the bytes codec (zarr-developers/zarr-python#4142, merged, past v3.2.1) - VirtualiZarr kerchunk struct-dtype refs for FITS binary tables (zarr-developers/VirtualiZarr#1037, open on the emfdavid fork) Pin both to git commits via [tool.uv.sources]. These are dev-environment only -- uv strips them from the published wheel, so the abstract zarr>=3.0 requirement downstream is unchanged. Drop each pin once a release ships. Add an `astronomy` optional-dependency extra (virtualizarr/kerchunk/astropy/ s3fs/icechunk) that formalizes the FITS build-time index stack for both the Hubble and SDSS examples, replacing the ad-hoc `uv run --with ...` install; point the Hubble README at it. Verified: uv sync resolves, the struct-dtype endianness repro passes on the pinned env, and the full suite is green (168 passed / 10 skipped) with the git-pinned dev-zarr. Co-Authored-By: Claude Opus 4.8 --- examples/README.md | 6 +- pyproject.toml | 28 ++++++++ uv.lock | 164 ++++++++++++++++++++++++++++++++++++++++----- 3 files changed, 180 insertions(+), 18 deletions(-) diff --git a/examples/README.md b/examples/README.md index 98cd1d1..2cec948 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 diff --git a/pyproject.toml b/pyproject.toml index baa22fe..9a76454 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,6 +102,20 @@ 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 two unreleased +# upstream fixes, pinned to git in [tool.uv.sources] below. +astronomy = [ + "virtualizarr>=2.6", + "kerchunk>=0.2.10", + "astropy>=8.0", + "s3fs>=2024.6", + "icechunk>=0.2", +] [dependency-groups] dev = [ @@ -116,6 +130,20 @@ dev = [ requires = ["uv_build>=0.11.21,<0.12.0"] build-backend = "uv_build" +# Temporary git pins for the SDSS/spectra astronomy example, which reads FITS binary +# tables (structured dtypes) end-to-end. Both fixes are merged/in-flight upstream but +# not yet in a release; these entries are dev-environment only (uv strips them from the +# published wheel, so the abstract `zarr>=3.0` requirement downstream is unchanged). +# Drop each pin once a release carrying the fix ships. +# - zarr: struct-dtype byte-order handling in the bytes codec +# (zarr-developers/zarr-python#4142, merged; not in a release past v3.2.1). +# - virtualizarr: kerchunk struct-dtype refs for FITS binary tables +# (zarr-developers/VirtualiZarr#1037, open on the emfdavid fork; on merge, repoint +# to a zarr-developers/VirtualiZarr main commit, then to the release). +[tool.uv.sources] +zarr = { git = "https://github.com/zarr-developers/zarr-python.git", rev = "bbe27d7ce9713d50e34df2cda5108f4dcf410b92" } +virtualizarr = { git = "https://github.com/emfdavid/VirtualiZarr.git", rev = "8a1563f236355534d06362034c1a05ab27667eba" } + [tool.ruff] line-length = 100 target-version = "py312" diff --git a/uv.lock b/uv.lock index 5057472..c40f071 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,13 @@ dependencies = [ arraylake = [ { name = "arraylake" }, ] +astronomy = [ + { name = "astropy" }, + { name = "icechunk" }, + { name = "kerchunk" }, + { name = "s3fs" }, + { name = "virtualizarr" }, +] bench = [ { name = "gcsfs" }, { name = "plotly" }, @@ -1316,13 +1396,16 @@ 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" }, @@ -1332,16 +1415,18 @@ requires-dist = [ { 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/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba" }, + { name = "virtualizarr", marker = "extra == 'virtual'", git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba" }, { name = "xarray", specifier = ">=2024.9" }, { name = "xbatcher", marker = "extra == 'bench'", specifier = ">=0.3" }, - { name = "zarr", specifier = ">=3.0" }, + { name = "zarr", git = "https://github.com/zarr-developers/zarr-python.git?rev=bbe27d7ce9713d50e34df2cda5108f4dcf410b92" }, ] -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 +1495,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 +1532,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 +3066,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 +3302,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 +3859,8 @@ wheels = [ [[package]] name = "virtualizarr" -version = "2.6.2" -source = { registry = "https://pypi.org/simple" } +version = "10000.dev608+g8a1563f23" +source = { git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba#8a1563f236355534d06362034c1a05ab27667eba" } dependencies = [ { name = "numcodecs" }, { name = "numpy" }, @@ -3730,10 +3872,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" @@ -3947,8 +4085,8 @@ wheels = [ [[package]] name = "zarr" -version = "3.2.1" -source = { registry = "https://pypi.org/simple" } +version = "3.2.2.dev106+gbbe27d7c" +source = { git = "https://github.com/zarr-developers/zarr-python.git?rev=bbe27d7ce9713d50e34df2cda5108f4dcf410b92#bbe27d7ce9713d50e34df2cda5108f4dcf410b92" } dependencies = [ { name = "donfig" }, { name = "google-crc32c" }, @@ -3957,10 +4095,6 @@ dependencies = [ { name = "packaging" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/93/8d/aeb164004f87543b06ef54f885d02c342c31ceb274e2bbec470a98927621/zarr-3.2.1.tar.gz", hash = "sha256:71565b738a0e7e8ed226f0516eba8c6bb53440ad7669a8c48ebb3534a161d035", size = 675161, upload-time = "2026-05-05T12:37:22.383Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/0a/469e2bd01be1490336e6c8707386845655d59261543315778a3ccc7e8019/zarr-3.2.1-py3-none-any.whl", hash = "sha256:f78cdd3d9687ad0e9f9cba2c5683b64f0c52589c19f685eeabe872e93cc0d2c7", size = 319617, upload-time = "2026-05-05T12:37:20.66Z" }, -] [[package]] name = "zipp" From cf82347ea10847e6f8e962c23f61a3edecac4dd7 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Wed, 15 Jul 2026 03:16:42 +0000 Subject: [PATCH 2/4] examples: SDSS spectral reconstruction (spPlate, train in place, no reshard) An astroML cross-domain companion to the Hubble example, and the decode- amortization counterpart to it. astroML's spectral-PCA workflow first downloads the raw archive and resamples every spectrum onto a common grid into one spec4000.npz -- the download-and-reshard step insitubatch argues against. examples/sdss indexes a real SDSS spPlate frame on data.sdss.org (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 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 is the O(chunks) decode-amortization regime (contrast Hubble's one image per chunk). Verified byte-identical to astropy. The task mirrors astroML's spectral reconstruction: a 1-D conv autoencoder trained by SGD over the streamed mini-batches beats a PCA baseline at the same latent dim, because galaxy spectra are translation-structured (redshift shifts the lines, which a fixed-dim linear basis reconstructs poorly). The offline synthetic default (redshift-shifted spectra) carries the beats-PCA claim and backs tests/test_sdss.py; the real run is a same-pipeline, real-archive demo. Only one plate per store: plates cover slightly different wavelength ranges, so concatenating them would require resampling -- the reshard this example avoids. 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 (kerchunk -> VirtualiZarr -> Icechunk -> insitubatch) against astropy, and pins the VirtualiZarr fix it needs (zarr-developers/VirtualiZarr#1037) -- released VirtualiZarr raises on the list-valued kerchunk field specs. Endianness is dropped in the VZ/icechunk struct metadata, so consumers reinterpret each field with .view('>f4') at the numpy boundary (the zarr #4142 codec fix does not engage in this path, so it is not pinned). Co-Authored-By: Claude Opus 4.8 --- .gitignore | 1 + examples/README.md | 40 +++ examples/sdss/__init__.py | 8 + examples/sdss/data.py | 375 +++++++++++++++++++++++++++++ examples/sdss/dr17_plate_uris.json | 8 + examples/sdss/train_torch.py | 144 +++++++++++ pyproject.toml | 20 +- tests/test_fits_bintable.py | 149 ++++++++++++ tests/test_sdss.py | 94 ++++++++ uv.lock | 10 +- 10 files changed, 836 insertions(+), 13 deletions(-) create mode 100644 examples/sdss/__init__.py create mode 100644 examples/sdss/data.py create mode 100644 examples/sdss/dr17_plate_uris.json create mode 100644 examples/sdss/train_torch.py create mode 100644 tests/test_fits_bintable.py create mode 100644 tests/test_sdss.py 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 2cec948..ddae181 100644 --- a/examples/README.md +++ b/examples/README.md @@ -140,6 +140,46 @@ 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 +``` + +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. + +> Only one plate per store: different plates cover slightly different wavelength ranges, so they +> cannot share a rectangular `wave` axis without resampling — exactly the reshard this example +> avoids. 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/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..0c2a470 --- /dev/null +++ b/examples/sdss/data.py @@ -0,0 +1,375 @@ +"""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 +from pathlib import Path + +import numpy as np +import xarray as xr +import zarr +from zarr.abc.store import Store + +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 +FIBERS_PER_CHUNK = 64 # split one plate's fibers into contiguous virtual chunks (640 -> 10 chunks) + + +def _rechunk_fibers(plate: xr.Dataset, fibers_per_chunk: int) -> xr.Dataset: + """Split a plate's whole-image virtual chunk into contiguous fiber-block chunks -- no reshard. + + ``kerchunk.fits`` maps the ``(fiber, wave)`` HDU to a single virtual chunk: one byte range + ``[offset, offset + n_fiber*n_wave*itemsize)``. Because the fiber axis is the *outer* axis of + the row-major image, fiber block ``i`` occupies a *contiguous* sub-range + ``offset + i*fibers_per_chunk*n_wave*itemsize`` of length ``fibers_per_chunk*n_wave*itemsize``. + So we rebuild the manifest as ``n_fiber // fibers_per_chunk`` sub-refs -- pure byte arithmetic, + no pixels moved -- giving insitubatch multiple sample-axis chunks (each still many fibers, the + decode-amortization regime) that :func:`~insitubatch.split_by_chunk` can partition. + """ + import dataclasses + + from virtualizarr.manifests import ChunkManifest, ManifestArray + + ma = plate[FLUX_VAR].data + manifest = ma.manifest + path = str(manifest._paths[0, 0]) + base_offset = int(manifest._offsets[0, 0]) + n_fiber, n_wave = ma.shape + row_bytes = n_wave * ma.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 = { + f"{i}.0": {"path": path, "offset": base_offset + i * block_bytes, "length": block_bytes} + for i in range(n_fiber // fibers_per_chunk) + } + grid = dataclasses.replace(ma.metadata.chunk_grid, chunk_shape=(fibers_per_chunk, n_wave)) + rechunked = ManifestArray( + metadata=dataclasses.replace(ma.metadata, chunk_grid=grid), + chunkmanifest=ChunkManifest(entries=entries), + ) + return xr.Dataset({FLUX_VAR: xr.Variable(("fiber", "wave"), rechunked)}) + + +def build_store( + plate_url: str, + store_path: str | os.PathLike[str] = DEFAULT_STORE, + *, + fibers_per_chunk: int = FIBERS_PER_CHUNK, +) -> str: + """Index one SDSS ``spPlate`` FITS (over HTTPS) into a local Icechunk repo of virtual refs. + + Idempotent: rebuilds ``store_path`` from scratch. Requires the ``astronomy`` extra + (``virtualizarr``, ``kerchunk``, ``astropy``, ``icechunk``). The plate's ``PRIMARY`` flux HDU + -- all ~640 fibers on one common log-wavelength grid -- is virtually re-chunked along the fiber + axis (:func:`_rechunk_fibers`) so each fiber is one sample and each chunk is many fibers. (Only + one plate: different plates cover slightly different wavelength ranges, so they cannot share a + rectangular ``wave`` axis without resampling -- exactly the reshard this example avoids.) + """ + import shutil + + import icechunk + from obstore.store import HTTPStore + 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({SDSS_HOST: HTTPStore.from_url(SDSS_HOST)}) + plate = open_virtual_dataset(url=plate_url, registry=registry, parser=FITSParser()).rename( + {"PRIMARY": FLUX_VAR} + ) + rechunked = _rechunk_fibers(plate, fibers_per_chunk) + + 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") + rechunked.virtualize.to_icechunk(session.store) + session.commit(f"index SDSS spPlate {plate_url.rsplit('/', 1)[-1]}") + 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_url = load_uris(args.uris)[0] + print(f"building store from {plate_url} ...") + build_store(plate_url, 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("--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 9a76454..35535f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -130,18 +130,18 @@ dev = [ requires = ["uv_build>=0.11.21,<0.12.0"] build-backend = "uv_build" -# Temporary git pins for the SDSS/spectra astronomy example, which reads FITS binary -# tables (structured dtypes) end-to-end. Both fixes are merged/in-flight upstream but -# not yet in a release; these entries are dev-environment only (uv strips them from the -# published wheel, so the abstract `zarr>=3.0` requirement downstream is unchanged). -# Drop each pin once a release carrying the fix ships. -# - zarr: struct-dtype byte-order handling in the bytes codec -# (zarr-developers/zarr-python#4142, merged; not in a release past v3.2.1). +# 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, open on the emfdavid fork; on merge, repoint -# to a zarr-developers/VirtualiZarr main commit, then to the release). +# (zarr-developers/VirtualiZarr#1037, open on the emfdavid fork). Released VirtualiZarr +# raises on the list-valued field specs; on merge, repoint to a zarr-developers/VirtualiZarr +# main commit, then to the release. +# 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] -zarr = { git = "https://github.com/zarr-developers/zarr-python.git", rev = "bbe27d7ce9713d50e34df2cda5108f4dcf410b92" } virtualizarr = { git = "https://github.com/emfdavid/VirtualiZarr.git", rev = "8a1563f236355534d06362034c1a05ab27667eba" } [tool.ruff] diff --git a/tests/test_fits_bintable.py b/tests/test_fits_bintable.py new file mode 100644 index 0000000..d13efe2 --- /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 obstore.store import LocalStore + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + from virtualizarr.registry import ObjectStoreRegistry + + 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.virtualize.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/uv.lock b/uv.lock index c40f071..d8aa3e8 100644 --- a/uv.lock +++ b/uv.lock @@ -1424,7 +1424,7 @@ requires-dist = [ { name = "virtualizarr", marker = "extra == 'virtual'", git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba" }, { name = "xarray", specifier = ">=2024.9" }, { name = "xbatcher", marker = "extra == 'bench'", specifier = ">=0.3" }, - { name = "zarr", git = "https://github.com/zarr-developers/zarr-python.git?rev=bbe27d7ce9713d50e34df2cda5108f4dcf410b92" }, + { name = "zarr", specifier = ">=3.0" }, ] provides-extras = ["torch", "jax", "tf", "gpu", "virtual", "cache", "bench", "docs", "icechunk", "gcsfs", "arraylake", "astronomy"] @@ -4085,8 +4085,8 @@ wheels = [ [[package]] name = "zarr" -version = "3.2.2.dev106+gbbe27d7c" -source = { git = "https://github.com/zarr-developers/zarr-python.git?rev=bbe27d7ce9713d50e34df2cda5108f4dcf410b92#bbe27d7ce9713d50e34df2cda5108f4dcf410b92" } +version = "3.2.1" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "donfig" }, { name = "google-crc32c" }, @@ -4095,6 +4095,10 @@ dependencies = [ { name = "packaging" }, { name = "typing-extensions" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/93/8d/aeb164004f87543b06ef54f885d02c342c31ceb274e2bbec470a98927621/zarr-3.2.1.tar.gz", hash = "sha256:71565b738a0e7e8ed226f0516eba8c6bb53440ad7669a8c48ebb3534a161d035", size = 675161, upload-time = "2026-05-05T12:37:22.383Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/0a/469e2bd01be1490336e6c8707386845655d59261543315778a3ccc7e8019/zarr-3.2.1-py3-none-any.whl", hash = "sha256:f78cdd3d9687ad0e9f9cba2c5683b64f0c52589c19f685eeabe872e93cc0d2c7", size = 319617, upload-time = "2026-05-05T12:37:20.66Z" }, +] [[package]] name = "zipp" From aa6c3042550551ea6b320f2e5467bcc270f602c1 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Wed, 15 Jul 2026 03:38:28 +0000 Subject: [PATCH 3/4] examples/sdss: add multi-plate common-grid concat (streaming at archive scale) One spPlate (~640 spectra, ~10 MB) fits in memory, so a single-plate store is a demo, not a streaming workload. build_store now takes a list of plates and picks one of two modes, both of which rewrite only the chunk manifest (no pixels moved): * one plate -> _single_plate_fiber_chunks: full wave width, many fibers per chunk (the decode-amortization regime). * many plates -> _many_plates_common_grid: every plate uses the same dloglam and their COEFF0 start wavelengths differ by whole bins, so cropping each to the shared window aligns them onto one grid exactly (no resampling) and folds (plate, fiber) into one flat sample axis, one fiber per chunk -- the streaming regime, and the only no-reshard way to concat the ragged-width plates. Scales to the full archive (thousands of plates, tens of GB). The two modes are mutually exclusive by construction: cropping to the common window breaks the multi-fiber contiguity the single-plate mode rides on, so decode-amortization and corpus scale cannot both hold -- an honest boundary of the FITS byte layout. A `--plates N` CLI flag selects the count. tests/test_sdss_build.py guards both modes offline with synthetic big-endian FITS images (byte-identical reads, exact cross-plate alignment, and a fail-fast on non-bin-aligned grids). Verified end-to-end on real plates: 3 plates -> 1920 aligned spectra, byte-identical to astropy. Co-Authored-By: Claude Opus 4.8 --- examples/README.md | 24 ++++-- examples/sdss/data.py | 171 +++++++++++++++++++++++++++++---------- tests/test_sdss_build.py | 144 +++++++++++++++++++++++++++++++++ 3 files changed, 292 insertions(+), 47 deletions(-) create mode 100644 tests/test_sdss_build.py diff --git a/examples/README.md b/examples/README.md index ddae181..e12d591 100644 --- a/examples/README.md +++ b/examples/README.md @@ -160,9 +160,20 @@ uv run python -m examples.sdss.train_torch # offline synth # 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 +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 @@ -174,11 +185,12 @@ network or FITS stack — redshift-shifted synthetic spectra on which the conv A 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. -> Only one plate per store: different plates cover slightly different wavelength ranges, so they -> cannot share a rectangular `wave` axis without resampling — exactly the reshard this example -> avoids. 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)). +> 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) diff --git a/examples/sdss/data.py b/examples/sdss/data.py index 0c2a470..ec64afa 100644 --- a/examples/sdss/data.py +++ b/examples/sdss/data.py @@ -31,14 +31,18 @@ import argparse import json import os -from dataclasses import dataclass +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, @@ -119,60 +123,131 @@ def make_synthetic_store( # --------------------------------------------------------------------------- build time -FIBERS_PER_CHUNK = 64 # split one plate's fibers into contiguous virtual chunks (640 -> 10 chunks) +# +# 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 _rechunk_fibers(plate: xr.Dataset, fibers_per_chunk: int) -> xr.Dataset: - """Split a plate's whole-image virtual chunk into contiguous fiber-block chunks -- no reshard. +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) - ``kerchunk.fits`` maps the ``(fiber, wave)`` HDU to a single virtual chunk: one byte range - ``[offset, offset + n_fiber*n_wave*itemsize)``. Because the fiber axis is the *outer* axis of - the row-major image, fiber block ``i`` occupies a *contiguous* sub-range - ``offset + i*fibers_per_chunk*n_wave*itemsize`` of length ``fibers_per_chunk*n_wave*itemsize``. - So we rebuild the manifest as ``n_fiber // fibers_per_chunk`` sub-refs -- pure byte arithmetic, - no pixels moved -- giving insitubatch multiple sample-axis chunks (each still many fibers, the - decode-amortization regime) that :func:`~insitubatch.split_by_chunk` can partition. - """ - import dataclasses +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 - ma = plate[FLUX_VAR].data - manifest = ma.manifest - path = str(manifest._paths[0, 0]) - base_offset = int(manifest._offsets[0, 0]) - n_fiber, n_wave = ma.shape - row_bytes = n_wave * ma.dtype.itemsize + 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 = { - f"{i}.0": {"path": path, "offset": base_offset + i * block_bytes, "length": block_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) } - grid = dataclasses.replace(ma.metadata.chunk_grid, chunk_shape=(fibers_per_chunk, n_wave)) - rechunked = ManifestArray( - metadata=dataclasses.replace(ma.metadata, chunk_grid=grid), - chunkmanifest=ChunkManifest(entries=entries), + 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) ) - return xr.Dataset({FLUX_VAR: xr.Variable(("fiber", "wave"), rechunked)}) + 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_url: str, + plate_urls: list[str], store_path: str | os.PathLike[str] = DEFAULT_STORE, *, fibers_per_chunk: int = FIBERS_PER_CHUNK, ) -> str: - """Index one SDSS ``spPlate`` FITS (over HTTPS) into a local Icechunk repo of virtual refs. + """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``). The plate's ``PRIMARY`` flux HDU - -- all ~640 fibers on one common log-wavelength grid -- is virtually re-chunked along the fiber - axis (:func:`_rechunk_fibers`) so each fiber is one sample and each chunk is many fibers. (Only - one plate: different plates cover slightly different wavelength ranges, so they cannot share a - rectangular ``wave`` axis without resampling -- exactly the reshard this example avoids.) + (``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 @@ -186,10 +261,17 @@ def build_store( shutil.rmtree(store_path, ignore_errors=True) registry = ObjectStoreRegistry({SDSS_HOST: HTTPStore.from_url(SDSS_HOST)}) - plate = open_virtual_dataset(url=plate_url, registry=registry, parser=FITSParser()).rename( - {"PRIMARY": FLUX_VAR} + 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) ) - rechunked = _rechunk_fibers(plate, fibers_per_chunk) ice_prefix = SDSS_HOST + "/" config = icechunk.RepositoryConfig.default() @@ -202,8 +284,8 @@ def build_store( authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), ) session = repo.writable_session("main") - rechunked.virtualize.to_icechunk(session.store) - session.commit(f"index SDSS spPlate {plate_url.rsplit('/', 1)[-1]}") + virtual.virtualize.to_icechunk(session.store) + session.commit(f"index {len(plate_urls)} SDSS spPlate frame(s)") return store_path @@ -348,9 +430,9 @@ def build_datasets(args: argparse.Namespace) -> InSituDataset: store = obstore_store(url) else: if args.build: - plate_url = load_uris(args.uris)[0] - print(f"building store from {plate_url} ...") - build_store(plate_url, args.store) + 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) @@ -366,6 +448,13 @@ def cli(argv: list[str] | None = None) -> argparse.Namespace: 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) diff --git a/tests/test_sdss_build.py b/tests/test_sdss_build.py new file mode 100644 index 0000000..fb4dfa4 --- /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 obstore.store import LocalStore + from virtualizarr import open_virtual_dataset + from virtualizarr.parsers import FITSParser + from virtualizarr.registry import ObjectStoreRegistry + + 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.virtualize.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) From ec43f84c8b0207e6eb97b54163aacbf02db03c82 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:00:59 +0000 Subject: [PATCH 4/4] build: repoint virtualizarr pin upstream, migrate off deprecated APIs The kerchunk struct-dtype fix for FITS binary tables merged upstream as zarr-developers/VirtualiZarr#1037, so drop the emfdavid fork pin. That rev (8a1563f) was never pushed to the fork remote and only ever resolved from uv's local cache -- it was unresolvable from a clean checkout. Pin to zarr-developers main (f46902d), verified to carry the merged fix as an ancestor. The fix ships in v2.7.1, so a git pin is still required; it becomes a plain virtualizarr>=2.7.1 once that releases. Moving to current upstream deprecated two APIs the astronomy build-time path uses: - .virtualize accessor -> .vz - virtualizarr.registry.ObjectStoreRegistry -> obspec_utils.registry Both are confined to the four FITS index-build sites (two tests, two examples); nothing in src/ or on the hot path. obspec-utils was already present as a VirtualiZarr dependency, so declaring it in [astronomy] adds no package (uv still resolves 202) -- it just stops us importing through a re-export scheduled for removal. Astronomy tests pass under -W error::DeprecationWarning. Co-Authored-By: Claude Opus 4.8 --- examples/hubble/data.py | 4 ++-- examples/sdss/data.py | 4 ++-- pyproject.toml | 17 +++++++++++------ tests/test_fits_bintable.py | 4 ++-- tests/test_sdss_build.py | 4 ++-- uv.lock | 10 ++++++---- 6 files changed, 25 insertions(+), 18 deletions(-) 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/data.py b/examples/sdss/data.py index ec64afa..4fcdced 100644 --- a/examples/sdss/data.py +++ b/examples/sdss/data.py @@ -252,10 +252,10 @@ def build_store( 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 - from virtualizarr.registry import ObjectStoreRegistry store_path = str(store_path) shutil.rmtree(store_path, ignore_errors=True) @@ -284,7 +284,7 @@ def build_store( authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), ) session = repo.writable_session("main") - virtual.virtualize.to_icechunk(session.store) + virtual.vz.to_icechunk(session.store) session.commit(f"index {len(plate_urls)} SDSS spPlate frame(s)") return store_path diff --git a/pyproject.toml b/pyproject.toml index 35535f5..65d2024 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,10 +107,15 @@ arraylake = [ # 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 two unreleased -# upstream fixes, pinned to git in [tool.uv.sources] below. +# (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", @@ -134,15 +139,15 @@ build-backend = "uv_build" # 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, open on the emfdavid fork). Released VirtualiZarr -# raises on the list-valued field specs; on merge, repoint to a zarr-developers/VirtualiZarr -# main commit, then to the release. +# (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/emfdavid/VirtualiZarr.git", rev = "8a1563f236355534d06362034c1a05ab27667eba" } +virtualizarr = { git = "https://github.com/zarr-developers/VirtualiZarr.git", rev = "f46902d2d3346bff31c88bf7212a14a2a50934bc" } [tool.ruff] line-length = 100 diff --git a/tests/test_fits_bintable.py b/tests/test_fits_bintable.py index d13efe2..f6c795b 100644 --- a/tests/test_fits_bintable.py +++ b/tests/test_fits_bintable.py @@ -42,10 +42,10 @@ def _build_bintable_store(tmp_dir: str) -> tuple[str, np.ndarray]: """ 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 - from virtualizarr.registry import ObjectStoreRegistry n = 64 flux = (np.sin(np.linspace(0.0, 6.0, n)) * 10.0).astype(">f4") # big-endian, as in real FITS @@ -75,7 +75,7 @@ def _build_bintable_store(tmp_dir: str) -> tuple[str, np.ndarray]: authorize_virtual_chunk_access=icechunk.containers_credentials({ice_prefix: None}), ) session = repo.writable_session("main") - vds.virtualize.to_icechunk(session.store) + vds.vz.to_icechunk(session.store) session.commit("index bintable") return f"{tmp_dir}/repo", flux diff --git a/tests/test_sdss_build.py b/tests/test_sdss_build.py index fb4dfa4..af8a2a8 100644 --- a/tests/test_sdss_build.py +++ b/tests/test_sdss_build.py @@ -49,10 +49,10 @@ def _write_plate( 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 - from virtualizarr.registry import ObjectStoreRegistry prefix = f"file://{tmp_dir}" registry = ObjectStoreRegistry({prefix: LocalStore(prefix=tmp_dir)}) @@ -74,7 +74,7 @@ def _commit_and_open(virtual, tmp_dir: str): # -> zarr Store authorize_virtual_chunk_access=icechunk.containers_credentials({prefix: None}), ) session = repo.writable_session("main") - virtual.virtualize.to_icechunk(session.store) + virtual.vz.to_icechunk(session.store) session.commit("build") return repo.readonly_session("main").store diff --git a/uv.lock b/uv.lock index d8aa3e8..df4f652 100644 --- a/uv.lock +++ b/uv.lock @@ -1340,6 +1340,7 @@ astronomy = [ { name = "astropy" }, { name = "icechunk" }, { name = "kerchunk" }, + { name = "obspec-utils" }, { name = "s3fs" }, { name = "virtualizarr" }, ] @@ -1410,6 +1411,7 @@ requires-dist = [ { 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" }, @@ -1420,8 +1422,8 @@ requires-dist = [ { 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 == 'astronomy'", git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba" }, - { name = "virtualizarr", marker = "extra == 'virtual'", git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba" }, + { 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" }, @@ -3859,8 +3861,8 @@ wheels = [ [[package]] name = "virtualizarr" -version = "10000.dev608+g8a1563f23" -source = { git = "https://github.com/emfdavid/VirtualiZarr.git?rev=8a1563f236355534d06362034c1a05ab27667eba#8a1563f236355534d06362034c1a05ab27667eba" } +version = "2.7.1.dev14+gf46902d2d" +source = { git = "https://github.com/zarr-developers/VirtualiZarr.git?rev=f46902d2d3346bff31c88bf7212a14a2a50934bc#f46902d2d3346bff31c88bf7212a14a2a50934bc" } dependencies = [ { name = "numcodecs" }, { name = "numpy" },