diff --git a/.binder/environment.yml b/.binder/environment.yml index 98c7a6c36..219c4e48a 100644 --- a/.binder/environment.yml +++ b/.binder/environment.yml @@ -8,11 +8,17 @@ # there is a single source of truth for the notebook dependencies (pyproject.toml). # # Only PyPI/conda-forge resolvable packages appear here: no ``git+`` URLs and no -# spherely fork. The Binder-runnable notebooks read synthetic data, the public -# anonymous source.coop store, or anonymous CMR-STAC granule *metadata* -# (jupyterhub_example, shardmap_viewer) -- all of which use the default HEALPix -# ``mortie`` backend, so the exact-S2 spherely SpatialIndex backend is never on -# their import path. +# spherely fork. The two Binder notebooks (hhdc_viewer, waveform_viewer) READ the +# anonymous public source.coop demo stores and build no catalogs at all, so no +# SpatialIndex backend is on their import path -- exact-S2 spherely least of all. +# +# NOTE (2026-08-26): ``postBuild`` still installs ``[analysis,catalog,viz]``, but +# ``catalog`` (stac-geoparquet) and ``viz`` (ipyleaflet) were carried for the +# jupyterhub_example / shardmap_viewer notebooks, which have been archived off +# main. Neither remaining notebook imports them, so both extras -- and much of +# ``analysis`` -- are now dead weight in the image and could be trimmed to speed +# the build. Left in place rather than trimmed blind: that is a dependency +# change, and the image is not the place to discover a transitive need. name: zagg-binder channels: - conda-forge diff --git a/.binder/postBuild b/.binder/postBuild index 638088dd7..32732b089 100755 --- a/.binder/postBuild +++ b/.binder/postBuild @@ -30,3 +30,29 @@ else HATCH_VCS_PRETEND_VERSION="0.0.0+binder" \ python -m pip install --no-cache-dir ".[analysis,catalog,viz]" fi + +# The reader-only demos (notebooks/hhdc_viewer.ipynb, notebooks/waveform_viewer.ipynb) +# call no zagg PUBLIC API, so their runtime is NOT a zagg extra -- each notebook +# carries its own `%pip install` line and stays runnable outside Binder. They are not +# zagg-free, though: moczarr imports the t-digest algebra from zagg rather than +# vendoring it, and that is exactly what their `moczarr[zagg]` extra carries. +# +# NOTE that extra resolves `zagg>=0.40`, while the tag-less fallback below installs +# this checkout as `0.0.0+binder`, which does not satisfy it. On a Binder build where +# no tag is reachable, the notebooks' own `%pip install` will therefore quietly fetch +# a RELEASED zagg from PyPI and install it over the checkout, in the image's shared +# site-packages -- so the branch under review stops being what is exercised, silently +# (`-q`). The `git fetch --tags` above is what normally keeps that from happening. +# +# One of those packages must be in the IMAGE rather than installed +# from a cell, because it ships a Jupyter frontend extension a running kernel cannot +# pick up: ipympl, the `%matplotlib widget` backend hhdc_viewer's rotatable 3-D view +# needs. ipywidgets -- the sliders and dropdowns both use -- is NOT named here: it +# already arrives with the `viz` extra above, because ipyleaflet requires +# `ipywidgets<9,>=7.6.0`. Naming it again would resolve only the requested set, +# ignoring that installed cap; the day ipywidgets 9 ships, pip would upgrade past +# ipyleaflet's ceiling, print a dependency-conflict warning and exit 0 -- which +# `set -euo pipefail` cannot catch -- leaving a green build with a broken map. +# moczarr itself is pure Python, so the notebooks' own `%pip install` covers it +# either way. +python -m pip install --no-cache-dir ipympl diff --git a/.gitignore b/.gitignore index 8ab0cdba3..9d47c36c1 100644 --- a/.gitignore +++ b/.gitignore @@ -70,6 +70,10 @@ out.zarr/ demo/outputs/ demo/wheels/ demo/atl03_v007_full.parquet +# hhdc_viewer.ipynb's export cells write these into notebooks/ by construction +# every time anyone runs it to the end (an .npz of o22 voxel chips and an .npz +# of the co-registered pair, tens of MB); generated artifacts, not sources. +notebooks/*.np[yz] # data/ holds builder scripts + small AOI geojsons only — the heavy neighbors # (the local ATL03 catalog clone, generated shardmaps, per-shard count tables) diff --git a/README.md b/README.md index 549fbb2a8..c826a5d29 100644 --- a/README.md +++ b/README.md @@ -127,27 +127,24 @@ The store path and output grid parameters are defined in the YAML config (`outpu ### Step 4: Visualize Results -The output Zarr is a public DGGS dataset. The included notebook rasterizes HEALPix cells to a polar stereographic grid for fast rendering with `imshow`. +The output Zarr is a public DGGS dataset. The included notebooks read it back anonymously — a paired 3-D view of one shard, and the cell-level GEDI × ATL03 digest join. ```bash -uv run jupyter notebook notebooks/rasterized_zarr.ipynb +uv run jupyter notebook notebooks/hhdc_viewer.ipynb ``` -Adjust `GRID_SPACING` in the notebook to control output resolution. - ## Example Notebooks -The notebooks under `notebooks/` run on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks) — no install, no credentials. They install `zagg[analysis]` via the `.binder/` conda config and read only synthetic in-notebook data or the **anonymous, public** [source.coop](https://source.coop/englacial/zagg/benchmarks) benchmark store. +The notebooks under `notebooks/` run on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks) — no install, no credentials. Both are **reader-only**: they import `mortie` + `moczarr[zagg]` and call no zagg public API — the only zagg code on the path is the t-digest algebra `moczarr` imports rather than vendors, which is exactly what the `[zagg]` extra is for — reading the anonymous, public [source.coop](https://source.coop/englacial/zagg) demo stores (ICESat-2 ATL03 + GEDI L1B over California and the NEON AOP sites). Each carries its own `%pip install` line, so they run outside Binder unchanged. | Notebook | What it shows | Binder | |----------|---------------|--------| -| `custom_aggregations.ipynb` | Config-driven aggregation API on synthetic data | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/custom_aggregations.ipynb) | -| `rasterized_zarr.ipynb` | Rasterize the published HEALPix store to an 8 km polar-stereo grid | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/rasterized_zarr.ipynb) | -| `jupyterhub_example.ipynb` | Drive the API from a science hub; read & visualize a published result | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/jupyterhub_example.ipynb) | -| `cryocloud_example.ipynb` | End-to-end ISMIP6 read + **AWS Lambda fan-out** on CryoCloud | **not Binder-runnable** (needs live AWS + Earthdata credentials) | -| `cost_reporting.ipynb` | Max / estimated / actual invoke cost + progress-bar dispatch wrapper | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/cost_reporting.ipynb) | +| `hhdc_viewer.ipynb` | Polygon → MOC → shard → **rotatable paired 3-D view** (ATL03 + GEDI), binned or exact → numpy tensors | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/hhdc_viewer.ipynb) | +| `waveform_viewer.ipynb` | The cell-level join: one GEDI footprint against the 2×2 ATL03 cells beneath it, both from stored digests | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/waveform_viewer.ipynb) | + +They share `notebooks/viewers.py`, which holds the drawing so the notebooks stay about the read path. The two are split because one needs `%matplotlib widget` for its rotatable 3-D view and the other `%matplotlib inline`; the backends collide in a single kernel. -`cryocloud_example.ipynb` is the only Lambda demo; it dispatches to a deployed AWS Lambda and reads private-account S3 via the CryoCloud IRSA role, so it cannot run on Binder. +**Archived notebooks.** The earlier examples (`custom_aggregations`, `rasterized_zarr`, `jupyterhub_example`, `cryocloud_example`, `cost_reporting`, `aoi_mask`, `shardmap_viewer`, `tdigest_reader_example`, `sentinel2_fusion`) were written against older APIs and had drifted out of date — stale documentation being worse than none. They are removed from `main` and preserved at [`c56221b4`](https://github.com/englacial/zagg/tree/c56221b4/notebooks), the last commit on `main` that carried them — a commit permalink rather than a branch link, because GitHub keeps a blob URL at an explicit sha forever and there is no ref anyone has to remember not to delete. ## Project Structure diff --git a/docs/aoi_mask.md b/docs/aoi_mask.md index 28bdd5d2c..e8deec2bb 100644 --- a/docs/aoi_mask.md +++ b/docs/aoi_mask.md @@ -198,7 +198,7 @@ strict = ds.where(ds["aoi_mask"]) equivalent `(lats, lons)` ring. See the runnable, data-free example in -[`notebooks/aoi_mask.ipynb`](https://github.com/englacial/zagg/blob/main/notebooks/aoi_mask.ipynb), +[`aoi_mask.ipynb`](https://github.com/englacial/zagg/blob/c56221b4/notebooks/aoi_mask.ipynb) (archived — see the note below), which builds a small HEALPix grid + AOI box and shows the mask is `True` exactly for the in-AOI cells. The notebook is self-contained (no remote data) and runs anywhere `zagg` (with `mortie>=0.8.3`) is installed. Binder launch additionally diff --git a/docs/index.md b/docs/index.md index b21aa83db..dfcf45a95 100644 --- a/docs/index.md +++ b/docs/index.md @@ -59,7 +59,7 @@ See [Lambda Deployment](deployment/lambda.md) for AWS setup. ### 3. Visualize results ```bash -uv run jupyter notebook notebooks/rasterized_zarr.ipynb +uv run jupyter notebook notebooks/hhdc_viewer.ipynb ``` ## Design Philosophy diff --git a/docs/quickstart.md b/docs/quickstart.md index 1d4351830..d1fbcda9f 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -64,7 +64,7 @@ cells. The processing step consumes this file. To inspect the chunking interactively -- shard outlines, granule footprints, and a grid that appears on zoom -- use the shard-map viewer (`pip install zagg[viz]`). See the -[shard-map viewer notebook](https://github.com/englacial/zagg/blob/main/notebooks/shardmap_viewer.ipynb), +[shard-map viewer notebook](https://github.com/englacial/zagg/blob/c56221b4/notebooks/shardmap_viewer.ipynb) (archived), which runs on a synthetic example (no network needed) and includes manual in-browser verification instructions. @@ -202,7 +202,7 @@ excludes is reported as **not applicable** — `touch_skipped_paths` in the run record — never as a failure. See `src/zagg/configs/atl06.yaml` for a complete example and the -[custom aggregations notebook](https://github.com/englacial/zagg/blob/main/notebooks/custom_aggregations.ipynb) +[custom aggregations notebook](https://github.com/englacial/zagg/blob/c56221b4/notebooks/custom_aggregations.ipynb) (archived) for customization examples. ## Reading the output @@ -212,6 +212,6 @@ reconstructs dense, spatially faithful `(64, 64, n_bins)` tensors — with an occupancy mask and a shared `(offset, gain)` z-window per block (the [reader contract](ragged_layout.md#spatially-faithful-tensors-deinterleave-blocks-mask)) — or lossless raw value vectors from the stored digests. The -[t-digest reader notebook](https://github.com/englacial/zagg/blob/main/notebooks/tdigest_reader_example.ipynb) +[t-digest reader notebook](https://github.com/englacial/zagg/blob/c56221b4/notebooks/tdigest_reader_example.ipynb) (archived) walks through `read_tensors` and `read_raw_values` end to end on a self-contained synthetic store. diff --git a/notebooks/aoi_mask.ipynb b/notebooks/aoi_mask.ipynb deleted file mode 100644 index 9ff1daf4f..000000000 --- a/notebooks/aoi_mask.ipynb +++ /dev/null @@ -1,146 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "aoi-intro", - "metadata": {}, - "source": "# Strict-AOI cell mask\n\nzagg's optional `output.aoi_mask` (default off) packages a per-cell boolean\naligned to the output cell grid — `True` where the cell falls inside the area of\ninterest (AOI). It is **\"package, don't clip\"**: no observation is dropped, and a\nflag-off run is byte-identical to today. See the narrative doc `docs/aoi_mask.md`\nfor the config and storage layout.\n\nThis notebook is fully self-contained — it needs no remote data — and runs\nanywhere `zagg` (with `mortie>=0.8.3`) is installed. Binder launch additionally\nrequires the repo-wide `.binder/` environment, which lands separately via #105; a\nBinder badge is omitted here until that infrastructure exists, since a default\nBinder build can't resolve `zagg` + `mortie>=0.8.3`.\n\nIt builds a small HEALPix grid and an AOI box, computes the mask the same way the\nshard-map build stage does (native morton, no lat/lon decode), and shows the mask\nis `True` exactly for the in-AOI cells (issue #101).\n\nThe AOI may be given either as the `[(lats, lons), ...]` ring parts used below or\nas a native **WKB/WKT geometry** (`mortie >= 0.8.3`): pass `aoi=` WKB bytes / WKT\ntext to `ShardMap.build` / `make_shardmap`, or `--aoi-wkt` / `--aoi-wkb` on the\nCLI. The last cell shows the WKB/WKT path yields the identical mask." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "aoi-build", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "from zagg.grids import HealpixGrid\n", - "\n", - "# A small HEALPix grid. parent_order is the shard order; child_order the leaf\n", - "# (cell) order the mask resolves at.\n", - "grid = HealpixGrid(parent_order=4, child_order=8, layout=\"fullsphere\")\n", - "\n", - "# An AOI box in WGS84, as the [(lats, lons), ...] parts list `coverage` takes.\n", - "lat0, lon0, lat1, lon1 = 10.0, 10.0, 20.0, 20.0\n", - "lats = np.array([lat0, lat0, lat1, lat1, lat0])\n", - "lons = np.array([lon0, lon1, lon1, lon0, lon0])\n", - "parts = [(lats, lons)]\n", - "\n", - "# 1) Compact MOC of the AOI at child_order (built once at the shard-map stage).\n", - "aoi_moc = grid.aoi_moc(parts)\n", - "print(f\"AOI MOC: {aoi_moc.size} mixed-order cells\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "aoi-expand", - "metadata": {}, - "outputs": [], - "source": [ - "# 2) Pick a shard the AOI touches and expand its per-shard sub-MOC to a per-cell\n", - "# boolean over the shard's children() — already in cell/storage order.\n", - "from mortie import moc_to_order\n", - "\n", - "flat = np.unique(np.asarray(moc_to_order(aoi_moc, grid.child_order), dtype=np.uint64))\n", - "shard_key = int(grid.shards_of(grid.cells_of(flat[:1]))[0])\n", - "children = grid.children(shard_key)\n", - "\n", - "shard_moc = grid.aoi_shard_moc(aoi_moc, shard_key)\n", - "mask = grid.aoi_mask_for_children(shard_moc, children)\n", - "\n", - "print(f\"shard {shard_key}: {len(children)} cells, {int(mask.sum())} in-AOI\")\n", - "\n", - "# The mask equals: which of this shard's children land in the AOI cover.\n", - "expected = np.isin(np.asarray(children, dtype=np.uint64), flat)\n", - "assert np.array_equal(mask, expected)\n", - "print(\"mask matches the cell-order AOI cover ✓\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "aoi-plot", - "metadata": {}, - "outputs": [], - "source": [ - "# 3) Visualize the in-AOI cells by their HEALPix center (decode only for the\n", - "# picture — the mask itself never decodes centers).\n", - "import matplotlib.pyplot as plt\n", - "from mortie import mort2geo\n", - "\n", - "clat, clon = mort2geo(np.asarray(children))\n", - "fig, ax = plt.subplots(figsize=(5, 5))\n", - "ax.scatter(clon[~mask], clat[~mask], s=8, c=\"lightgray\", label=\"out of AOI\")\n", - "ax.scatter(clon[mask], clat[mask], s=8, c=\"C0\", label=\"in AOI\")\n", - "ax.add_patch(\n", - " plt.Rectangle(\n", - " (lon0, lat0), lon1 - lon0, lat1 - lat0, fill=False, edgecolor=\"red\", lw=1.5, label=\"AOI\"\n", - " )\n", - ")\n", - "ax.set_xlabel(\"lon\")\n", - "ax.set_ylabel(\"lat\")\n", - "ax.legend(loc=\"upper left\")\n", - "ax.set_title(f\"shard {shard_key}: strict-AOI cell mask\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "aoi-read", - "metadata": {}, - "source": [ - "Once a run with `output.aoi_mask: true` has written the store, filtering to the\n", - "strict AOI is just a selection on the `aoi_mask` array:\n", - "\n", - "```python\n", - "import xarray as xr\n", - "ds = xr.open_zarr(\"s3://bucket/atl06.zarr/13\") # child_order=13 group\n", - "strict = ds.where(ds[\"aoi_mask\"]) # NaN out the overhang cells\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "c8e10bb7", - "metadata": {}, - "outputs": [], - "source": [ - "# 4) Same AOI as WKB/WKT geometry -> identical mask. The HEALPix engine rides\n", - "# mortie's public from_wkb / from_wkt cover entry points (espg/mortie#89), so a\n", - "# WKB/WKT AOI produces the same compact MOC as the ring above.\n", - "import shapely\n", - "\n", - "# WKT of the SAME box (WKT is lon-lat ordered).\n", - "wkt = \"POLYGON((%s))\" % \", \".join(\n", - " f\"{x} {y}\" for x, y in [(lon0, lat0), (lon1, lat0), (lon1, lat1), (lon0, lat1), (lon0, lat0)]\n", - ")\n", - "wkb = shapely.from_wkt(wkt).wkb\n", - "\n", - "moc_wkt = grid.aoi_moc(wkt) # str -> WKT\n", - "moc_wkb = grid.aoi_moc(wkb) # bytes -> WKB\n", - "mask_wkt = grid.aoi_mask_for_children(grid.aoi_shard_moc(moc_wkt, shard_key), children)\n", - "mask_wkb = grid.aoi_mask_for_children(grid.aoi_shard_moc(moc_wkb, shard_key), children)\n", - "\n", - "assert np.array_equal(mask, mask_wkt)\n", - "assert np.array_equal(mask, mask_wkb)\n", - "print(\"WKT and WKB AOI masks match the ring mask ✓\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/notebooks/cost_reporting.ipynb b/notebooks/cost_reporting.ipynb deleted file mode 100644 index c6210fe7a..000000000 --- a/notebooks/cost_reporting.ipynb +++ /dev/null @@ -1,235 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Invoke cost reporting: max -> progress -> actual\n", - "\n", - "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/cost_reporting.ipynb)\n", - "\n", - "_Runs end-to-end on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/cost_reporting.ipynb): it reads only a small shard-map fixture checked into the git tree -- no cloud data, no credentials. The one step Binder cannot do (invoke AWS Lambda) is replayed with a stand-in, clearly marked below._\n", - "\n", - "A zagg Lambda fan-out spends real money, so every run reports three cost figures (issue #298):\n", - "\n", - "| figure | when | meaning |\n", - "|---|---|---|\n", - "| **max** | before any invoke | hard ceiling: `n_units x rate(arch) x memory_gb x 900 s` -- every unit is one Lambda invocation billed at the configured memory for at most the function timeout |\n", - "| **estimated** | before any invoke | prior-run history estimate (pilot-first); currently a `None` placeholder -- deferred behind the stats-sidecar and template-hash issues |\n", - "| **actual** | after the run | rollup of billed durations: `sum(duration_s) x memory_gb x rate` |\n", - "\n", - "The ceiling is exact and cheap: it needs only the shard map (how many units) and the config (which worker memory)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Max cost, before any invoke\n", - "\n", - "`zagg.notebook.max_cost_preview` resolves the ceiling from a shard map + config with the same unit accounting the dispatcher uses (cell selection, windowed-unit expansion). The fixture below is a 4-shard ATL03 shard map from the test tree." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "\n", - "from zagg.config import default_config\n", - "from zagg.notebook import format_max_cost, max_cost_preview\n", - "\n", - "\n", - "def find(rel):\n", - " \"\"\"Resolve a repo file whether we run from the repo root or notebooks/.\"\"\"\n", - " for base in (Path.cwd(), Path.cwd().parent):\n", - " if (base / rel).exists():\n", - " return str(base / rel)\n", - " raise FileNotFoundError(rel)\n", - "\n", - "\n", - "shardmap = find(\"tests/data/benchmark/shardmaps/sm_healpix_o9.json\")\n", - "config = default_config(\"atl06\")\n", - "\n", - "preview = max_cost_preview(config, shardmap)\n", - "print(format_max_cost(preview))\n", - "preview" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ceiling scales with the worker-size variant the config selects (the `worker:` block, issue #235):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for memory_mb in (2048, 4096, 8192):\n", - " config.worker = {\"memory\": memory_mb}\n", - " p = max_cost_preview(config, shardmap)\n", - " print(f\"worker {memory_mb} MB -> ceiling ${p['max_cost_usd']:.4f}\")\n", - "config.worker = None # back to the default 4 GB worker" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Estimated cost (deferred)\n", - "\n", - "The estimator is an interface stub for now: **tier 1 (pilot-first)** scales prior actuals for the *same* template hash by per-shard granule/obs counts (the dev -> single-shard test -> fleet workflow means a pilot run usually exists); **tier 2** falls back to a cross-template regression on per-shard `n_obs`. Both need the per-shard stats sidecars and template-hash identity to exist first, so until those land the stub returns `None` and `cost[\"estimated_cost_usd\"]` stays a placeholder." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.dispatch import estimate_cost_usd\n", - "\n", - "print(estimate_cost_usd()) # None until the sidecar history exists" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Progress + report\n", - "\n", - "`zagg.notebook.run` wraps `zagg.agg` with a per-unit progress bar (tqdm when importable, logging otherwise; the running cost rides the postfix) and returns a `RunView` -- a summary dict with a rich HTML repr. The notebook path is **informational only**: it displays the ceiling and never prompts (the blocking yes/no gate is CLI-only).\n", - "\n", - "Binder cannot invoke AWS Lambda, so the cell below substitutes a stand-in `agg` that replays a realistic run summary while ticking the same `on_progress` callback a real Lambda fan-out drives -- the wrapper code path (ceiling display, bar, `RunView`) is exactly the one a real run exercises." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import time\n", - "from unittest.mock import patch\n", - "\n", - "from zagg.dispatch import LAMBDA_PRICE_PER_GB_SEC\n", - "\n", - "MEMORY_GB = 4.0\n", - "DURATIONS = {-231928: 41.0, -231921: 63.0, -231920: 55.0, -231913: 48.0} # simulated billed seconds\n", - "\n", - "\n", - "def replay_agg(config, *, on_progress=None, **kwargs):\n", - " \"\"\"Stand-in for zagg.agg: replays a recorded-style Lambda run summary.\"\"\"\n", - " total, cost = len(DURATIONS), 0.0\n", - " for i, duration in enumerate(DURATIONS.values(), 1):\n", - " time.sleep(0.3) # a real shard takes seconds-to-minutes\n", - " cost += duration * MEMORY_GB * LAMBDA_PRICE_PER_GB_SEC\n", - " if on_progress is not None:\n", - " on_progress(i, total, cost)\n", - " lambda_time = sum(DURATIONS.values())\n", - " actual = lambda_time * MEMORY_GB * LAMBDA_PRICE_PER_GB_SEC\n", - " return {\n", - " \"backend\": \"lambda\",\n", - " \"store_path\": \"s3://example-bucket/atl03_serc.zarr\",\n", - " \"total_cells\": total,\n", - " \"cells_with_data\": total,\n", - " \"cells_error\": 0,\n", - " \"total_obs\": 1_284_055,\n", - " \"wall_time_s\": 71.2,\n", - " \"lambda_time_s\": lambda_time,\n", - " \"cost\": {\n", - " \"max_cost_usd\": preview[\"max_cost_usd\"],\n", - " \"estimated_cost_usd\": None,\n", - " \"actual_cost_usd\": actual,\n", - " },\n", - " \"results\": [\n", - " {\"shard_key\": k, \"status_code\": 200, \"error\": None, \"lambda_duration\": d}\n", - " for k, d in DURATIONS.items()\n", - " ],\n", - " }\n", - "\n", - "\n", - "from zagg import notebook\n", - "\n", - "with patch(\"zagg.runner.agg\", replay_agg):\n", - " report = notebook.run(config, catalog=shardmap, backend=\"lambda\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "report # RunView: cost block, counters, failures render as HTML" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cost = report[\"cost\"]\n", - "assert cost[\"max_cost_usd\"] >= cost[\"actual_cost_usd\"] # the ceiling always bounds the bill\n", - "cost" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Running it for real\n", - "\n", - "Against a deployed fleet the call is the same, minus the stand-in:\n", - "\n", - "```python\n", - "from zagg import load_config, notebook\n", - "\n", - "config = load_config(\"atl03_serc.yaml\")\n", - "report = notebook.run(\n", - " config,\n", - " catalog=\"shardmap_atl03_serc.json\",\n", - " backend=\"lambda\",\n", - " store=\"s3://your-bucket/atl03_serc.zarr\",\n", - ")\n", - "report # max / estimated / actual + per-shard failures\n", - "```\n", - "\n", - "From the command line the same run **blocks on the ceiling** with a yes/no prompt; `--yes`/`-y` skips it for scripted runs:\n", - "\n", - "```console\n", - "$ python -m zagg --config atl03_serc.yaml --catalog shardmap_atl03_serc.json --backend lambda\n", - "Max cost ceiling: ~$0.19 (4 units x 4 GB x 900s, arm64)\n", - "Proceed with Lambda fan-out? [y/N] y\n", - "```" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/cryocloud_example.ipynb b/notebooks/cryocloud_example.ipynb deleted file mode 100644 index aa40e128c..000000000 --- a/notebooks/cryocloud_example.ipynb +++ /dev/null @@ -1,648 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "608fca4d", - "metadata": {}, - "source": [ - "# zagg on CryoCloud: end-to-end example\n", - "\n", - "> **Not Binder-runnable.** This notebook dispatches to a deployed AWS Lambda\n", - "> function and reads/writes private-account S3 via the CryoCloud IRSA role, so\n", - "> it needs live AWS + NASA Earthdata credentials and **cannot run on\n", - "> [Binder](https://mybinder.org)**. It is the *only* notebook that demonstrates\n", - "> the Lambda fan-out backend. For Binder-runnable tours of the API and the\n", - "> published results, see `custom_aggregations.ipynb`, `rasterized_zarr.ipynb`,\n", - "> and `jupyterhub_example.ipynb`.\n", - "\n", - "This notebook demonstrates running the [`zagg`](https://github.com/englacial/zagg) library on CryoCloud, including:\n", - "\n", - "1. **Installing** `zagg` and its dependencies into the CryoCloud user environment\n", - "2. **Reading ISMIP6 data** from a virtualizarr / icechunk repository (single variable, DataTree, and inter-model variance with a time slider)\n", - "3. **Local processing** of ICESat-2 ATL06 data with a custom aggregation\n", - "4. **Lambda processing** of the full ATL06 dataset against an AWS Lambda function in account `429435741471`\n", - "5. **Reprojecting and plotting** the result on an Antarctic Polar Stereographic grid (EPSG:3031) using `imshow`\n", - "\n", - "AWS authentication is handled automatically by the CryoCloud IRSA role (`nasa-cryo-prod`) -- no AWS credentials need to be configured manually. NASA Earthdata authentication still uses `earthaccess` (your `~/.netrc` or `earthaccess.login()`)." - ] - }, - { - "cell_type": "markdown", - "id": "a250275a", - "metadata": {}, - "source": [ - "## 0. Install zagg and configure the Lambda target\n", - "\n", - "On CryoCloud, the bare function name `process-morton-cell` would resolve in CryoCloud's own AWS account (574251165169) and fail with `AccessDeniedException`. We point boto3 at the full cross-account ARN via the `ZAGG_LAMBDA_FUNCTION_NAME` env var that `zagg.runner` reads.\n", - "\n", - "Lambda output is written to the public `source.coop` benchmarks store in the zagg owner's account -- the function's execution role already has write permission there, and the result stays colocated with the published benchmarks." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "15118dcb", - "metadata": {}, - "outputs": [], - "source": [ - "# zagg + the analysis extra (notebook runtime). On a non-CryoCloud hub install\n", - "# from PyPI instead: %pip install --quiet \"zagg[analysis]\"\n", - "%pip install --quiet \"zagg[analysis]\"\n", - "# Cryo-only extras used below (ISMIP6 icechunk read + the hvplot time slider).\n", - "# These are deliberately NOT in the Binder image -- only this notebook needs them.\n", - "%pip install --quiet icechunk \"cubed-xarray>=0.0.9\" hvplot" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "47a35a59", - "metadata": {}, - "outputs": [], - "source": [ - "import logging\n", - "import os\n", - "\n", - "# Show INFO-level zagg messages (per-cell progress, summary lines).\n", - "# Without this, only WARNING+ surfaces and runs look silent.\n", - "logging.basicConfig(level=logging.INFO, format=\"%(message)s\")\n", - "\n", - "os.environ[\"ZAGG_LAMBDA_FUNCTION_NAME\"] = (\n", - " \"arn:aws:lambda:us-west-2:429435741471:function:process-morton-cell\"\n", - ")\n", - "os.environ.setdefault(\"AWS_DEFAULT_REGION\", \"us-west-2\")\n", - "\n", - "# Public source.coop benchmarks store (anonymous read works from anywhere).\n", - "OUTPUT_BUCKET = \"us-west-2.opendata.source.coop\"\n", - "OUTPUT_KEY = \"englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr\"\n", - "OUTPUT_URL = f\"s3://{OUTPUT_BUCKET}/{OUTPUT_KEY}\"\n", - "print(f\"Lambda output -> {OUTPUT_URL}\")\n", - "\n", - "# Confirm IRSA credentials are present (sanity check)\n", - "import boto3\n", - "\n", - "print(boto3.client(\"sts\").get_caller_identity()[\"Arn\"])" - ] - }, - { - "cell_type": "markdown", - "id": "d20312ce", - "metadata": {}, - "source": [ - "## 1. Reading ISMIP6 data from the virtualizarr / icechunk store\n", - "\n", - "The ISMIP6 Antarctic ensemble is published as an [icechunk](https://icechunk.io) repository at `s3://us-west-2.opendata.source.coop/englacial/ismip6/icechunk-ais`. The icechunk session contains *virtual* references to the original NetCDF chunks plus rechunked overlays.\n", - "\n", - "We open it once anonymously and then show three access patterns:\n", - "\n", - "- **A single variable** from one model/experiment (the most common pattern)\n", - "- **The full repository as an `xarray.DataTree`** (when you want to traverse the model/experiment hierarchy programmatically)\n", - "- **An inter-model variance** of `lithk` for `exp05`, with a time slider" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f4daa648", - "metadata": {}, - "outputs": [], - "source": [ - "import warnings\n", - "\n", - "warnings.filterwarnings(\"ignore\", message=\"Numcodecs codecs are not in the Zarr version 3\")\n", - "\n", - "import icechunk\n", - "import xarray as xr\n", - "import zarr\n", - "\n", - "SOURCE_BUCKET = \"s3://us-west-2.opendata.source.coop/englacial/ismip6\"\n", - "\n", - "storage = icechunk.s3_storage(\n", - " bucket=\"us-west-2.opendata.source.coop\",\n", - " prefix=\"englacial/ismip6/icechunk-ais\",\n", - " region=\"us-west-2\",\n", - " anonymous=True,\n", - ")\n", - "\n", - "config = icechunk.RepositoryConfig.default()\n", - "config.set_virtual_chunk_container(\n", - " icechunk.VirtualChunkContainer(\n", - " SOURCE_BUCKET + \"/\",\n", - " store=icechunk.s3_store(region=\"us-west-2\", anonymous=True),\n", - " )\n", - ")\n", - "credentials = icechunk.containers_credentials({SOURCE_BUCKET + \"/\": None})\n", - "\n", - "repo = icechunk.Repository.open(\n", - " storage=storage,\n", - " config=config,\n", - " authorize_virtual_chunk_access=credentials,\n", - ")\n", - "session = repo.readonly_session(branch=\"main\")\n", - "\n", - "root = zarr.open(session.store, mode=\"r\")\n", - "print(root.tree(level=1))" - ] - }, - { - "cell_type": "markdown", - "id": "ff137877", - "metadata": {}, - "source": [ - "### 1a. Open a single variable from one model/experiment\n", - "\n", - "Each model/experiment is a zarr group under `combined//`. We open one of them and plot the ice thickness `lithk` field. The `combined` store unions state and flux variables onto a single time axis, so the very first timestep can be all-NaN for state variables -- we use `time=1` to be safe." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "de6c4cee", - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "\n", - "ds = xr.open_zarr(\n", - " session.store,\n", - " group=\"combined/AWI_PISM1/exp05\",\n", - " consolidated=False,\n", - ")\n", - "print(f\"Variables: {sorted(ds.data_vars)}\")\n", - "print(\n", - " f\"Time range: {str(ds.time.values[0])[:10]} -> {str(ds.time.values[-1])[:10]} ({len(ds.time)} steps)\"\n", - ")\n", - "\n", - "fig, ax = plt.subplots(figsize=(7, 7))\n", - "ds[\"lithk\"].isel(time=1).plot(ax=ax, cmap=\"viridis\")\n", - "ax.set_aspect(\"equal\")\n", - "ax.set_title(f\"AWI_PISM1 / exp05 / lithk @ {str(ds.time.values[1])[:10]}\")\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "29567fcd", - "metadata": {}, - "source": [ - "### 1b. Open the entire repository as an xarray DataTree\n", - "\n", - "`xr.open_datatree` traverses every group below the root, giving you a nested tree you can iterate or index into with paths. We pass `chunked_array_type=\"cubed\"` so dask is not required (`cubed` is the default chunked backend for icechunk virtual stores)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f6abbf3d", - "metadata": {}, - "outputs": [], - "source": [ - "ds_tree = xr.open_datatree(\n", - " session.store,\n", - " engine=\"zarr\",\n", - " create_default_indexes=False,\n", - " decode_times=False,\n", - " chunked_array_type=\"cubed\",\n", - ")\n", - "\n", - "# Show the top-level groups\n", - "print(ds_tree)\n", - "\n", - "# Index into it just like a dict-of-datasets\n", - "lithk_awi = ds_tree[\"combined/AWI_PISM1/exp05\"].ds[\"lithk\"]\n", - "print(f\"\\nlithk shape: {lithk_awi.shape}, dtype: {lithk_awi.dtype}\")" - ] - }, - { - "cell_type": "markdown", - "id": "ef8d1c95", - "metadata": {}, - "source": [ - "### 1c. Inter-model variance of `lithk` for `exp05` with a time slider\n", - "\n", - "Now we use the DataTree to walk every model under `combined/`, pull the `lithk` field for `exp05`, stack them along a new `model` dimension, and compute variance across models.\n", - "\n", - "ISMIP6 specifies a common 8 km Antarctic grid (761 x 761), so most submissions share spatial coordinates. We restrict to models that match the reference grid and the shortest common time axis; mismatches are skipped with a printed warning. The result is a `(time, y, x)` array that we plot with `hvplot`, which renders a year-by-year time slider." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ffa3bd2d", - "metadata": {}, - "outputs": [], - "source": [ - "# Walk every model in combined/, pull lithk for exp05\n", - "model_arrays = []\n", - "model_names = []\n", - "for model in sorted(ds_tree[\"combined\"].children):\n", - " try:\n", - " node = xr.open_zarr(\n", - " session.store,\n", - " group=f\"combined/{model}/exp05\",\n", - " consolidated=False,\n", - " )\n", - " except (KeyError, FileNotFoundError, ValueError) as e:\n", - " print(f\"skip {model}: {e}\")\n", - " continue\n", - " if \"lithk\" not in node.data_vars:\n", - " print(f\"skip {model}: no lithk\")\n", - " continue\n", - " model_arrays.append(node[\"lithk\"])\n", - " model_names.append(model)\n", - "\n", - "# Restrict to the reference grid (first model) and the shortest common time axis\n", - "ref = model_arrays[0]\n", - "ref_shape = (ref.sizes[\"y\"], ref.sizes[\"x\"])\n", - "min_time = min(a.sizes[\"time\"] for a in model_arrays)\n", - "\n", - "aligned = []\n", - "kept_names = []\n", - "for name, arr in zip(model_names, model_arrays):\n", - " if (arr.sizes[\"y\"], arr.sizes[\"x\"]) != ref_shape:\n", - " print(f\"skip {name}: grid {arr.sizes['y']}x{arr.sizes['x']} != {ref_shape}\")\n", - " continue\n", - " aligned.append(\n", - " arr.isel(time=slice(0, min_time)).assign_coords(\n", - " x=ref.x, y=ref.y, time=ref.time.isel(time=slice(0, min_time))\n", - " )\n", - " )\n", - " kept_names.append(name)\n", - "\n", - "print(f\"\\nUsing {len(aligned)} models, time={min_time} steps, grid={ref_shape}\")\n", - "print(f\"Models: {kept_names}\")\n", - "\n", - "model_dim = xr.DataArray(kept_names, dims=\"model\")\n", - "stacked = xr.concat(aligned, dim=model_dim, compat=\"override\", coords=\"minimal\")\n", - "lithk_var = stacked.var(dim=\"model\").compute()\n", - "lithk_var.name = \"lithk_variance\"\n", - "lithk_var.attrs[\"units\"] = \"m^2\"\n", - "lithk_var" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b3d21604", - "metadata": {}, - "outputs": [], - "source": [ - "import hvplot.xarray # noqa: F401 (registers the .hvplot accessor)\n", - "\n", - "lithk_var.hvplot.image(\n", - " x=\"x\",\n", - " y=\"y\",\n", - " groupby=\"time\",\n", - " cmap=\"viridis\",\n", - " clim=(0, float(lithk_var.quantile(0.99))),\n", - " aspect=\"equal\",\n", - " frame_width=500,\n", - " widget_type=\"scrubber\",\n", - " widget_location=\"bottom\",\n", - " title=\"Inter-model variance of lithk (exp05)\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "38bd2b31", - "metadata": {}, - "source": "## 2. Local processing with a custom aggregation\n\nBefore scaling out to Lambda, validate the pipeline locally on a tiny slice. We start from the bundled `atl06` config and override the aggregation block to compute a leaner set of statistics (median, min/max, variance) instead of the default weighted-mean + quantile bundle.\n\nWe use the HTTPS driver here so this also works from non-AWS environments.\n\nFirst we build the granule shard map used by both the local and Lambda runs.\n`python -m zagg.catalog` takes the grid from the pipeline `--config` (so the\nshard map can't drift from the grid the run uses) and a spatial extent\n(`--polygon`/`--bbox`); the parent order lives in the config, not on the CLI.\nThe CMR-STAC query is **anonymous** -- building the catalog needs no Earthdata\nlogin (only the `agg()` data fetch below does)." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "c59ecb56", - "metadata": {}, - "outputs": [], - "source": [ - "# Build the shard map: ICESat-2 ATL06, cycle 22, grid (incl. parent order 6)\n", - "# from the bundled atl06 config; written to --output below.\n", - "# The CMR-STAC query is anonymous (no Earthdata login needed to *find* granules;\n", - "# only fetching their pixels in agg() needs auth). The default 'auto' geometry\n", - "# backend uses the spherely fork if installed, else falls back to mortie (HEALPix\n", - "# MOC) -- so no fork is required here. --polygon is your own area-of-interest\n", - "# GeoJSON (drives the CMR query + the coverage mask); swap in any region you have.\n", - "# Run from the repo root so the --config path resolves, or point --config at an\n", - "# installed config.\n", - "!python -m zagg.catalog \\\n", - " --config src/zagg/configs/atl06.yaml \\\n", - " --short-name ATL06 \\\n", - " --cycle 22 \\\n", - " --polygon antarctica.geojson \\\n", - " --output shardmap_atl06_cycle22.json" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9e382a37", - "metadata": {}, - "outputs": [], - "source": [ - "import yaml\n", - "\n", - "from zagg import agg\n", - "from zagg.config import get_agg_fields, load_config_from_dict\n", - "\n", - "custom_yaml = \"\"\"\n", - "data_source:\n", - " reader: h5coro\n", - " groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r]\n", - " coordinates:\n", - " latitude: \"/{group}/land_ice_segments/latitude\"\n", - " longitude: \"/{group}/land_ice_segments/longitude\"\n", - " variables:\n", - " h_li: \"/{group}/land_ice_segments/h_li\"\n", - " s_li: \"/{group}/land_ice_segments/h_li_sigma\"\n", - " quality_filter:\n", - " dataset: \"/{group}/land_ice_segments/atl06_quality_summary\"\n", - " value: 0\n", - "\n", - "aggregation:\n", - " coordinates:\n", - " cell_ids: {dtype: uint64, fill_value: 0}\n", - " morton: {dtype: int64, fill_value: 0}\n", - " variables:\n", - " count: {function: len, source: h_li, dtype: int32, fill_value: 0}\n", - " h_median: {function: median, source: h_li, dtype: float32}\n", - " h_min: {function: min, source: h_li, dtype: float32}\n", - " h_max: {function: max, source: h_li, dtype: float32}\n", - " h_variance: {function: var, source: h_li, dtype: float32}\n", - "\n", - "output:\n", - " grid:\n", - " type: healpix\n", - " indexing_scheme: nested\n", - " child_order: 12\n", - "\"\"\"\n", - "\n", - "cfg_custom = load_config_from_dict(yaml.safe_load(custom_yaml))\n", - "\n", - "for name, meta in get_agg_fields(cfg_custom).items():\n", - " func = meta.get(\"function\", meta.get(\"expression\", \"N/A\"))\n", - " print(f\"{name:12s} -> {func}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4d51f1e8", - "metadata": {}, - "outputs": [], - "source": [ - "# Run on 2 cells using HTTPS to confirm the pipeline works end-to-end\n", - "results_local = agg(\n", - " cfg_custom,\n", - " catalog=\"shardmap_atl06_cycle22.json\",\n", - " store=\"./local_test.zarr\",\n", - " backend=\"local\",\n", - " driver=\"https\",\n", - " max_cells=2,\n", - " max_workers=2,\n", - " overwrite=True,\n", - ")\n", - "\n", - "print(f\"Cells processed: {results_local['cells_with_data']}\")\n", - "print(f\"Total observations: {results_local['total_obs']:,}\")\n", - "print(f\"Wall time: {results_local['wall_time_s']:.1f}s\")" - ] - }, - { - "cell_type": "markdown", - "id": "c5e0ebdc", - "metadata": {}, - "source": [ - "## 3. Full-dataset Lambda run on CryoCloud\n", - "\n", - "With the pipeline validated, we now use the default `atl06` config (which produces the richer set of variables expected by the existing rasterized output) and dispatch every parent cell in the shard map to the Lambda function. The function runs in *your* AWS account (`429435741471`), so all compute and S3 PUT charges bill there -- not to CryoCloud.\n", - "\n", - "Output is written to the public `source.coop` benchmarks store (`OUTPUT_URL` above). The Lambda's execution role (`zagg-lambda-execution`) is already scoped to that location, so writes work without any cross-account bucket-policy work.\n", - "\n", - "The IRSA role `nasa-cryo-prod` was granted `lambda:InvokeFunction` on the function ARN via the merged 2i2c PR; the Lambda's resource-based policy was extended to trust that role via `aws lambda add-permission`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "df0e6aed", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg import default_config\n", - "\n", - "cfg_full = default_config(\"atl06\")\n", - "\n", - "results_lambda = agg(\n", - " cfg_full,\n", - " catalog=\"shardmap_atl06_cycle22.json\",\n", - " store=OUTPUT_URL, # s3://us-west-2.opendata.source.coop/englacial/zagg/benchmarks/...\n", - " backend=\"lambda\",\n", - " driver=\"s3\", # Lambda runs in us-west-2; direct S3 read is cheaper\n", - " max_workers=1700, # Reserved-concurrency cap on the function\n", - " overwrite=True,\n", - ")\n", - "\n", - "print(f\"Cells with data: {results_lambda['cells_with_data']:,}\")\n", - "print(f\"Total observations: {results_lambda['total_obs']:,}\")\n", - "print(f\"Wall time: {results_lambda['wall_time_s']:.1f}s\")\n", - "print(f\"Lambda compute: {results_lambda['lambda_time_s']:.0f}s\")\n", - "print(f\"Estimated cost: ${results_lambda['estimated_cost_usd']:.2f}\")" - ] - }, - { - "cell_type": "markdown", - "id": "65188412", - "metadata": {}, - "source": [ - "## 4. Reproject to EPSG:3031 and plot with imshow\n", - "\n", - "The result is a HEALPix-indexed zarr at order 12 (~3.4 km cells). To make it map-friendly we:\n", - "\n", - "1. Convert each filled HEALPix cell center to lon/lat with `mortie._healpix.pix2ang`\n", - "2. Project to Antarctic Polar Stereographic (EPSG:3031) using `pyproj`\n", - "3. Bin onto an 8 km regular grid (`np.bincount` for sums; weighted mean for mean elevation; error propagation for sigma)\n", - "4. Display each variable with `imshow` under a `cartopy.crs.SouthPolarStereo` axis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "40f1f450", - "metadata": {}, - "outputs": [], - "source": [ - "import cartopy.crs as ccrs\n", - "import cartopy.feature as cfeature\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import xarray as xr\n", - "from matplotlib.colors import LogNorm\n", - "from mortie._healpix import pix2ang\n", - "from pyproj import Transformer\n", - "\n", - "from zagg import open_store\n", - "\n", - "CHILD_ORDER = 12\n", - "NSIDE = 2**CHILD_ORDER\n", - "\n", - "# The source.coop benchmarks store is public, so an unsigned read works anywhere;\n", - "# open_store applies zagg's read-only retry policy (issue #186).\n", - "store = open_store(OUTPUT_URL, read_only=True, skip_signature=True)\n", - "ds = xr.open_dataset(\n", - " store,\n", - " engine=\"zarr\",\n", - " consolidated=False,\n", - " zarr_format=3,\n", - " group=str(CHILD_ORDER),\n", - ")\n", - "\n", - "has_data = ds[\"count\"].values > 0\n", - "cell_ids = ds[\"cell_ids\"].values[has_data]\n", - "count = ds[\"count\"].values[has_data].astype(np.float64)\n", - "h_mean = ds[\"h_mean\"].values[has_data].astype(np.float64)\n", - "h_sigma = ds[\"h_sigma\"].values[has_data].astype(np.float64)\n", - "h_min = ds[\"h_min\"].values[has_data].astype(np.float64)\n", - "h_max = ds[\"h_max\"].values[has_data].astype(np.float64)\n", - "\n", - "print(f\"Filled HEALPix cells: {has_data.sum():,} / {has_data.size:,}\")\n", - "print(f\"Total observations: {count.sum():.0f}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "64a2fc81", - "metadata": {}, - "outputs": [], - "source": [ - "# Reproject HEALPix cell centers to EPSG:3031\n", - "# pix2ang takes the HEALPix order (depth), not nside (nside = 2**order).\n", - "lon, lat = pix2ang(CHILD_ORDER, cell_ids.astype(np.int64))\n", - "transformer = Transformer.from_crs(\"EPSG:4326\", \"EPSG:3031\", always_xy=True)\n", - "x, y = transformer.transform(lon, lat)\n", - "\n", - "# Define an 8 km regular grid covering ~south of 59 S\n", - "GRID_SPACING = 8000\n", - "EXTENT = 3_400_000\n", - "x_min, x_max = -EXTENT, EXTENT\n", - "y_min, y_max = -EXTENT, EXTENT\n", - "nx = int((x_max - x_min) / GRID_SPACING)\n", - "ny = int((y_max - y_min) / GRID_SPACING)\n", - "n_out = nx * ny\n", - "\n", - "ix = np.floor((x - x_min) / GRID_SPACING).astype(np.int64)\n", - "iy = np.floor((y - y_min) / GRID_SPACING).astype(np.int64)\n", - "inside = (ix >= 0) & (ix < nx) & (iy >= 0) & (iy < ny)\n", - "ix, iy = ix[inside], iy[inside]\n", - "flat = (iy * nx + ix).astype(np.int64)\n", - "\n", - "count_in = count[inside]\n", - "h_mean_in = h_mean[inside]\n", - "h_sigma_in = h_sigma[inside]\n", - "h_min_in = h_min[inside]\n", - "h_max_in = h_max[inside]\n", - "\n", - "# Aggregate onto the regular grid\n", - "count_out = np.bincount(flat, weights=count_in, minlength=n_out)\n", - "filled = count_out > 0\n", - "\n", - "\n", - "def weighted_mean(values):\n", - " wsum = np.bincount(flat, weights=values * count_in, minlength=n_out)\n", - " out = np.full(n_out, np.nan)\n", - " out[filled] = wsum[filled] / count_out[filled]\n", - " return out\n", - "\n", - "\n", - "h_mean_out = weighted_mean(h_mean_in)\n", - "\n", - "h_sigma_out = np.full(n_out, np.nan)\n", - "sigma_sq_wsum = np.bincount(flat, weights=count_in**2 * h_sigma_in**2, minlength=n_out)\n", - "h_sigma_out[filled] = np.sqrt(sigma_sq_wsum[filled]) / count_out[filled]\n", - "\n", - "h_min_out = np.full(n_out, np.inf)\n", - "np.minimum.at(h_min_out, flat, h_min_in)\n", - "h_min_out[~filled] = np.nan\n", - "\n", - "h_max_out = np.full(n_out, -np.inf)\n", - "np.maximum.at(h_max_out, flat, h_max_in)\n", - "h_max_out[~filled] = np.nan\n", - "\n", - "ds_out = xr.Dataset(\n", - " {\n", - " \"count\": ([\"y\", \"x\"], count_out.reshape(ny, nx).astype(np.int64)),\n", - " \"h_mean\": ([\"y\", \"x\"], h_mean_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_sigma\": ([\"y\", \"x\"], h_sigma_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_min\": ([\"y\", \"x\"], h_min_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_max\": ([\"y\", \"x\"], h_max_out.reshape(ny, nx).astype(np.float32)),\n", - " },\n", - " coords={\n", - " \"x\": np.arange(x_min + GRID_SPACING / 2, x_max, GRID_SPACING),\n", - " \"y\": np.arange(y_min + GRID_SPACING / 2, y_max, GRID_SPACING),\n", - " },\n", - " attrs={\"crs\": \"EPSG:3031\", \"grid_spacing_m\": GRID_SPACING},\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b6626226", - "metadata": {}, - "outputs": [], - "source": [ - "# Plot the four panels with imshow under SouthPolarStereo\n", - "proj = ccrs.SouthPolarStereo()\n", - "img_extent = [x_min, x_max, y_min, y_max]\n", - "mask = ds_out[\"count\"].values > 0\n", - "\n", - "fig, axes = plt.subplots(2, 2, figsize=(16, 14), subplot_kw={\"projection\": proj})\n", - "for ax in axes.flat:\n", - " ax.coastlines(resolution=\"50m\", linewidth=0.5)\n", - " ax.add_feature(cfeature.LAND, facecolor=\"lightgray\", alpha=0.3)\n", - " ax.gridlines(draw_labels=False, alpha=0.3)\n", - " ax.set_extent([-180, 180, -90, -60], crs=ccrs.PlateCarree())\n", - "\n", - "panels = [\n", - " (axes[0, 0], \"h_mean\", \"Mean elevation (m)\", dict(cmap=\"terrain\", vmin=0, vmax=4000)),\n", - " (axes[0, 1], \"count\", \"Observation count\", dict(cmap=\"viridis\", norm=LogNorm(vmin=1))),\n", - " (axes[1, 0], \"h_sigma\", \"Combined uncertainty (m)\", dict(cmap=\"plasma\", vmax=1.0)),\n", - " (axes[1, 1], None, \"Elevation range (m)\", dict(cmap=\"hot\", vmax=100)),\n", - "]\n", - "\n", - "for ax, var, title, kwargs in panels:\n", - " if var is None:\n", - " data = (ds_out[\"h_max\"].values - ds_out[\"h_min\"].values).astype(np.float64)\n", - " else:\n", - " data = ds_out[var].values.astype(np.float64)\n", - " data[~mask] = np.nan\n", - " im = ax.imshow(\n", - " data, origin=\"lower\", extent=img_extent, transform=proj, interpolation=\"nearest\", **kwargs\n", - " )\n", - " ax.set_title(title, fontsize=13, weight=\"bold\")\n", - " plt.colorbar(im, ax=ax, shrink=0.7)\n", - "\n", - "plt.suptitle(\n", - " \"ATL06 cycle 22 - 8 km Antarctic Polar Stereographic\", fontsize=15, weight=\"bold\", y=1.02\n", - ")\n", - "plt.tight_layout()\n", - "plt.show()" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/custom_aggregations.ipynb b/notebooks/custom_aggregations.ipynb deleted file mode 100644 index 89c329c2a..000000000 --- a/notebooks/custom_aggregations.ipynb +++ /dev/null @@ -1,607 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "intro-md", - "metadata": {}, - "source": [ - "# Customizing Aggregations in zagg\n", - "\n", - "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/custom_aggregations.ipynb)\n", - "\n", - "_Runs end-to-end on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/custom_aggregations.ipynb): it uses only small synthetic arrays defined in-notebook -- no cloud data, no credentials._\n", - "\n", - "\n", - "The aggregation pipeline is driven by a single YAML config file with three\n", - "required sections, plus optional fields:\n", - "\n", - "1. **`data_source`** -- what to read: reader, HDF5 groups, coordinates, variables, quality filter\n", - "2. **`aggregation`** -- what to compute: coordinate columns plus statistical variables, each wired to a numpy function or a Python expression\n", - "3. **`output`** -- where to write: grid spec (type, indexing scheme, child order) and optional store path\n", - "\n", - "Optional top-level fields:\n", - "- **`catalog`** -- path to a granule catalog JSON\n", - "- **`bounds`** -- temporal/spatial bounds for filtering\n", - "\n", - "Everything -- function dispatch, Zarr template generation, validation --\n", - "derives from this config. There is no separate function registry or\n", - "hard-coded schema class.\n", - "\n", - "This notebook walks through the default config, then shows two customization\n", - "examples:\n", - "- Replacing weighted mean with median and dropping quantile fields\n", - "- Adding a new input variable (along-track slope) with its own statistics" - ] - }, - { - "cell_type": "markdown", - "id": "part1-header", - "metadata": {}, - "source": [ - "## Part 1: The current setup\n", - "\n", - "### Loading the default config\n", - "\n", - "`default_config()` loads the built-in `atl06.yaml` shipped with the package\n", - "and returns a `PipelineConfig` dataclass with three dict fields." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "load-config", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.config import default_config\n", - "\n", - "cfg = default_config()\n", - "print(type(cfg))\n", - "print(\"Sections: data_source, aggregation, output\")" - ] - }, - { - "cell_type": "markdown", - "id": "three-sections-md", - "metadata": {}, - "source": [ - "### The three sections\n", - "\n", - "Each section is a plain dict parsed from YAML. Let's inspect them." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "show-datasource", - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "\n", - "print(\"=== data_source ===\")\n", - "print(json.dumps(cfg.data_source, indent=2))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "show-aggregation", - "metadata": {}, - "outputs": [], - "source": [ - "print(\"=== aggregation ===\")\n", - "print(json.dumps(cfg.aggregation, indent=2))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "show-output", - "metadata": {}, - "outputs": [], - "source": [ - "print(\"=== output ===\")\n", - "print(json.dumps(cfg.output, indent=2))" - ] - }, - { - "cell_type": "markdown", - "id": "yaml-md", - "metadata": {}, - "source": [ - "### The YAML source\n", - "\n", - "The config above comes from `src/zagg/configs/atl06.yaml`. Here is its\n", - "content verbatim:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "show-yaml", - "metadata": {}, - "outputs": [], - "source": [ - "from importlib import resources\n", - "\n", - "import zagg.configs\n", - "\n", - "yaml_text = resources.files(zagg.configs).joinpath(\"atl06.yaml\").read_text()\n", - "print(yaml_text)" - ] - }, - { - "cell_type": "markdown", - "id": "resolve-md", - "metadata": {}, - "source": [ - "### Function resolution\n", - "\n", - "`resolve_function()` maps a string name to a callable. The rules are:\n", - "- `\"len\"` or `\"count\"` resolves to the builtin `len`\n", - "- A bare name like `\"average\"` resolves to `np.average`\n", - "- A dotted path like `\"np.quantile\"` also works" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "resolve-demo", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "from zagg.config import resolve_function\n", - "\n", - "for name in [\"len\", \"min\", \"max\", \"average\", \"var\", \"quantile\", \"np.median\"]:\n", - " func = resolve_function(name)\n", - " if hasattr(np, func.__name__):\n", - " qualified = f\"np.{func.__name__}\"\n", - " else:\n", - " qualified = f\"builtins.{func.__name__}\"\n", - " print(f\" {name:20s} -> {qualified}\")" - ] - }, - { - "cell_type": "markdown", - "id": "expr-md", - "metadata": {}, - "source": [ - "### Expression evaluation\n", - "\n", - "Some aggregation variables use `expression:` instead of `function:`. These\n", - "are evaluated via `evaluate_expression()` in a restricted namespace\n", - "containing numpy and the cell's column arrays.\n", - "\n", - "For example, `h_sigma` in the default config uses:\n", - "```\n", - "expression: \"1.0 / np.sqrt(np.sum(1.0 / s_li**2))\"\n", - "```\n", - "\n", - "Let's try it with some synthetic data:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "expr-demo", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "from zagg.config import evaluate_expression\n", - "\n", - "columns = {\n", - " \"h_li\": np.array([120.5, 118.3, 122.1, 119.7, 121.0], dtype=np.float32),\n", - " \"s_li\": np.array([0.05, 0.10, 0.03, 0.08, 0.06], dtype=np.float32),\n", - "}\n", - "\n", - "# The h_sigma expression: inverse-variance combined uncertainty\n", - "expr = \"1.0 / np.sqrt(np.sum(1.0 / s_li**2))\"\n", - "result = evaluate_expression(expr, columns)\n", - "print(f\"h_sigma = {result:.6f}\")\n", - "\n", - "# A simpler example: RMS of h_li\n", - "expr2 = \"np.sqrt(np.mean(h_li**2))\"\n", - "result2 = evaluate_expression(expr2, columns)\n", - "print(f\"h_rms = {result2:.4f}\")" - ] - }, - { - "cell_type": "markdown", - "id": "dispatch-md", - "metadata": {}, - "source": [ - "### Dispatch via `calculate_cell_statistics()`\n", - "\n", - "`calculate_cell_statistics()` ties it all together. Given a single grid\n", - "cell's observations as a dict of column name -> numpy array, it iterates\n", - "the aggregation variables\n", - "from the config, resolves each function (or evaluates the expression), and\n", - "returns a dict of results." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "dispatch-demo", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.processing import calculate_cell_statistics\n", - "\n", - "# calculate_cell_statistics takes the per-cell eval namespace: a dict of\n", - "# column name -> equal-length numpy array (NOT a DataFrame).\n", - "# Five synthetic ICESat-2 observations: elevation (h_li) and uncertainty (s_li).\n", - "cell = {\n", - " \"h_li\": np.array([120.5, 118.3, 122.1, 119.7, 121.0], dtype=np.float32),\n", - " \"s_li\": np.array([0.05, 0.10, 0.03, 0.08, 0.06], dtype=np.float32),\n", - "}\n", - "\n", - "stats = calculate_cell_statistics(cell)\n", - "\n", - "for name, value in stats.items():\n", - " print(f\" {name:12s} = {value:.4f}\")" - ] - }, - { - "cell_type": "markdown", - "id": "dispatch-explain-md", - "metadata": {}, - "source": [ - "Notice that `h_mean` is *not* the arithmetic mean -- it's the\n", - "inverse-variance weighted mean (via `np.average` with\n", - "`weights: \"1.0 / s_li**2\"`). The observation at 122.1 has `s_li=0.03`\n", - "(the most precise), so the weighted mean is pulled toward it.\n", - "\n", - "`h_sigma` is computed by expression, not by function dispatch -- it\n", - "evaluates `1.0 / np.sqrt(np.sum(1.0 / s_li**2))` directly on the\n", - "column arrays." - ] - }, - { - "cell_type": "markdown", - "id": "part2-header", - "metadata": {}, - "source": [ - "## Part 2: Modifying aggregations via YAML\n", - "\n", - "Suppose you want a simpler output: median elevation instead of weighted\n", - "mean, and no quantile columns. You write a new YAML config (or edit a\n", - "copy of `atl06.yaml`) with only the variables you want:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "modified-config", - "metadata": {}, - "outputs": [], - "source": [ - "import yaml\n", - "\n", - "from zagg.config import get_agg_fields, load_config_from_dict, validate_config\n", - "\n", - "simple_yaml = \"\"\"\n", - "data_source:\n", - " reader: h5coro\n", - " groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r]\n", - " coordinates:\n", - " latitude: \"/{group}/land_ice_segments/latitude\"\n", - " longitude: \"/{group}/land_ice_segments/longitude\"\n", - " variables:\n", - " h_li: \"/{group}/land_ice_segments/h_li\"\n", - " s_li: \"/{group}/land_ice_segments/h_li_sigma\"\n", - " quality_filter:\n", - " dataset: \"/{group}/land_ice_segments/atl06_quality_summary\"\n", - " value: 0\n", - "\n", - "aggregation:\n", - " coordinates:\n", - " cell_ids: {dtype: uint64, fill_value: 0}\n", - " morton: {dtype: int64, fill_value: 0}\n", - " variables:\n", - " count:\n", - " function: len\n", - " source: h_li\n", - " dtype: int32\n", - " fill_value: 0\n", - " h_median:\n", - " function: median\n", - " source: h_li\n", - " dtype: float32\n", - " h_min:\n", - " function: min\n", - " source: h_li\n", - " dtype: float32\n", - " h_max:\n", - " function: max\n", - " source: h_li\n", - " dtype: float32\n", - " h_variance:\n", - " function: var\n", - " source: h_li\n", - " dtype: float32\n", - "\n", - "output:\n", - " grid:\n", - " type: healpix\n", - " indexing_scheme: nested\n", - " parent_order: 6\n", - " child_order: 12\n", - "\"\"\"\n", - "\n", - "cfg_simple = load_config_from_dict(yaml.safe_load(simple_yaml))\n", - "validate_config(cfg_simple)\n", - "print(\"Validated successfully.\\n\")\n", - "\n", - "print(\"Aggregation variables:\")\n", - "for name, meta in get_agg_fields(cfg_simple).items():\n", - " func = meta.get(\"function\", meta.get(\"expression\"))\n", - " print(f\" {name:12s} {func}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "run-modified", - "metadata": {}, - "outputs": [], - "source": [ - "stats_simple = calculate_cell_statistics(cell, config=cfg_simple)\n", - "\n", - "print(\"Simplified statistics (5 variables instead of 9):\")\n", - "for name, value in stats_simple.items():\n", - " print(f\" {name:12s} = {value:.4f}\")" - ] - }, - { - "cell_type": "markdown", - "id": "median-explain-md", - "metadata": {}, - "source": [ - "The median (120.5) differs from the inverse-variance weighted mean because\n", - "it isn't influenced by uncertainty -- it's just the middle value.\n", - "\n", - "### Zarr template adapts automatically\n", - "\n", - "The Zarr template is also driven by the config, so it adapts to the\n", - "modified variable list:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "zarr-template", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.schema import xdggs_spec\n", - "\n", - "# Default config -> 9 data vars + 2 coords = 11 arrays\n", - "spec_default = xdggs_spec(parent_order=6, child_order=12)\n", - "print(f\"Default template: {len(spec_default.members)} arrays\")\n", - "\n", - "# Simplified config -> 5 data vars + 2 coords = 7 arrays\n", - "spec_simple = xdggs_spec(parent_order=6, child_order=12, config=cfg_simple)\n", - "print(f\"Simplified template: {len(spec_simple.members)} arrays\")\n", - "print(f\" {sorted(spec_simple.members.keys())}\")" - ] - }, - { - "cell_type": "markdown", - "id": "part3-header", - "metadata": {}, - "source": [ - "## Part 3: Adding a new input variable\n", - "\n", - "The examples above all aggregate `h_li` (elevation). Suppose you also\n", - "want statistics on along-track surface slope (`dh_fit_dx`), which lives\n", - "in the same ATL06 HDF5 group.\n", - "\n", - "This touches two config sections:\n", - "- `data_source.variables` — add the new HDF5 path\n", - "- `aggregation.variables` — define statistics that reference it\n", - "\n", - "Here's the full YAML with slope added:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "extend-datasource", - "metadata": {}, - "outputs": [], - "source": [ - "extended_yaml = \"\"\"\n", - "data_source:\n", - " reader: h5coro\n", - " groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r]\n", - " coordinates:\n", - " latitude: \"/{group}/land_ice_segments/latitude\"\n", - " longitude: \"/{group}/land_ice_segments/longitude\"\n", - " variables:\n", - " h_li: \"/{group}/land_ice_segments/h_li\"\n", - " s_li: \"/{group}/land_ice_segments/h_li_sigma\"\n", - " dh_fit_dx: \"/{group}/land_ice_segments/dh_fit_dx\" # <-- new\n", - " quality_filter:\n", - " dataset: \"/{group}/land_ice_segments/atl06_quality_summary\"\n", - " value: 0\n", - "\n", - "aggregation:\n", - " coordinates:\n", - " cell_ids: {dtype: uint64, fill_value: 0}\n", - " morton: {dtype: int64, fill_value: 0}\n", - " variables:\n", - " # --- elevation stats (same as default) ---\n", - " count: {function: len, source: h_li, dtype: int32, fill_value: 0}\n", - " h_min: {function: min, source: h_li, dtype: float32}\n", - " h_max: {function: max, source: h_li, dtype: float32}\n", - " h_mean: {function: average, source: h_li, params: {weights: \"1.0 / s_li**2\"}, dtype: float32}\n", - " h_sigma: {expression: \"1.0 / np.sqrt(np.sum(1.0 / s_li**2))\", dtype: float32}\n", - " h_variance: {function: var, source: h_li, dtype: float32}\n", - " h_q25: {function: quantile, source: h_li, params: {q: 0.25}, dtype: float32}\n", - " h_q50: {function: quantile, source: h_li, params: {q: 0.50}, dtype: float32}\n", - " h_q75: {function: quantile, source: h_li, params: {q: 0.75}, dtype: float32}\n", - " # --- slope stats (new) ---\n", - " slope_mean: {function: mean, source: dh_fit_dx, dtype: float32}\n", - " slope_min: {function: min, source: dh_fit_dx, dtype: float32}\n", - " slope_max: {function: max, source: dh_fit_dx, dtype: float32}\n", - " slope_variance: {function: var, source: dh_fit_dx, dtype: float32}\n", - "\n", - "output:\n", - " grid:\n", - " type: healpix\n", - " indexing_scheme: nested\n", - " parent_order: 6\n", - " child_order: 12\n", - "\"\"\"\n", - "\n", - "cfg_extended = load_config_from_dict(yaml.safe_load(extended_yaml))\n", - "validate_config(cfg_extended)\n", - "print(\"Extended config validates successfully.\\n\")\n", - "\n", - "print(\"All aggregation variables:\")\n", - "for name in get_agg_fields(cfg_extended):\n", - " print(f\" {name}\")" - ] - }, - { - "cell_type": "markdown", - "id": "validate-fail-md", - "metadata": {}, - "source": [ - "What if we forget to add `dh_fit_dx` to `data_source.variables`?\n", - "`validate_config()` catches it:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "validate-fail", - "metadata": {}, - "outputs": [], - "source": [ - "# Same config but without dh_fit_dx in data_source.variables\n", - "broken_yaml = \"\"\"\n", - "data_source:\n", - " reader: h5coro\n", - " groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r]\n", - " coordinates:\n", - " latitude: \"/{group}/land_ice_segments/latitude\"\n", - " longitude: \"/{group}/land_ice_segments/longitude\"\n", - " variables:\n", - " h_li: \"/{group}/land_ice_segments/h_li\"\n", - " s_li: \"/{group}/land_ice_segments/h_li_sigma\"\n", - " # dh_fit_dx is missing!\n", - "\n", - "aggregation:\n", - " coordinates:\n", - " cell_ids: {dtype: uint64, fill_value: 0}\n", - " morton: {dtype: int64, fill_value: 0}\n", - " variables:\n", - " count: {function: len, source: h_li, dtype: int32, fill_value: 0}\n", - " slope_mean: {function: mean, source: dh_fit_dx, dtype: float32}\n", - "\n", - "output:\n", - " grid:\n", - " type: healpix\n", - " indexing_scheme: nested\n", - " parent_order: 6\n", - " child_order: 12\n", - "\"\"\"\n", - "\n", - "try:\n", - " validate_config(load_config_from_dict(yaml.safe_load(broken_yaml)))\n", - "except ValueError as e:\n", - " print(f\"Validation error: {e}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "run-extended", - "metadata": {}, - "outputs": [], - "source": [ - "# Run calculate_cell_statistics with the extended config on synthetic data\n", - "cell_with_slope = {\n", - " \"h_li\": np.array([120.5, 118.3, 122.1, 119.7, 121.0], dtype=np.float32),\n", - " \"s_li\": np.array([0.05, 0.10, 0.03, 0.08, 0.06], dtype=np.float32),\n", - " \"dh_fit_dx\": np.array([0.002, -0.001, 0.005, 0.003, -0.002], dtype=np.float32),\n", - "}\n", - "\n", - "stats_ext = calculate_cell_statistics(cell_with_slope, config=cfg_extended)\n", - "\n", - "print(\"Elevation statistics:\")\n", - "for name, value in stats_ext.items():\n", - " if name.startswith(\"h_\") or name == \"count\":\n", - " print(f\" {name:12s} = {value:.4f}\")\n", - "\n", - "print(\"\\nSlope statistics:\")\n", - "for name, value in stats_ext.items():\n", - " if name.startswith(\"slope_\"):\n", - " print(f\" {name:16s} = {value:.6f}\")" - ] - }, - { - "cell_type": "markdown", - "id": "extended-zarr-md", - "metadata": {}, - "source": [ - "The Zarr template also picks up the new variables:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "extended-zarr", - "metadata": {}, - "outputs": [], - "source": [ - "spec_ext = xdggs_spec(parent_order=6, child_order=12, config=cfg_extended)\n", - "print(f\"Extended template arrays ({len(spec_ext.members)}):\")\n", - "for name in sorted(spec_ext.members.keys()):\n", - " member = spec_ext.members[name]\n", - " print(f\" {name:16s} dtype={member.data_type} fill={member.fill_value}\")" - ] - }, - { - "cell_type": "markdown", - "id": "summary-md", - "metadata": {}, - "source": "### Summary\n\nAll customization flows through the YAML config (or its dict equivalent):\n\n| What to change | Config section | Key fields |\n|----------------|----------------|------------|\n| New HDF5 input variable | `data_source.variables` | `column_name: \"/{group}/path\"` |\n| New function-based statistic | `aggregation.variables` | `function:`, `source:`, `params:` |\n| New expression-based statistic | `aggregation.variables` | `expression:`, `dtype:` |\n| Output grid settings | `output.grid` | `type:`, `indexing_scheme:`, `parent_order:`, `child_order:` |\n| Output store path | `output.store` | local path or `s3://bucket/prefix` |\n| Catalog path | `catalog` | top-level, optional |\n\nRules to remember:\n- `function:` and `expression:` are **mutually exclusive** per variable\n- Bare function names (e.g. `\"min\"`) resolve to `np.min`; dotted paths\n like `\"np.median\"` also work\n- `params:` values can be numeric literals, bare column names, or\n expressions containing column names (e.g. `\"1.0 / s_li**2\"`)\n- `validate_config()` checks that all `source:` and param references\n point to columns defined in `data_source.variables`\n- For a healpix grid, `output.grid` requires both `parent_order:` and\n `child_order:`\n- The Zarr template (`xdggs_spec()`) and `calculate_cell_statistics()`\n both derive their structure from the same config" - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0547fc27-59ae-4b46-9da1-2c9eacfb7d64", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/notebooks/export.py b/notebooks/export.py new file mode 100644 index 000000000..a24c85fe5 --- /dev/null +++ b/notebooks/export.py @@ -0,0 +1,242 @@ +"""Voxel exports for `hhdc_viewer.ipynb`: isotropic chips and co-registered pairs. + +Both leave the `read_tensors` path and build from the stored digests, because each +needs something it cannot give -- an arbitrary output order for the first, and a +z axis shared between two sensors for the second. +""" + +from __future__ import annotations + +import json +import os +import time +import zipfile +from io import BytesIO + +import moczarr as mz +import numpy as np +from moczarr.hhdc import ( + block_rank, + chunk_z_range, + rank_to_rowcol, + rasterize_cell, + rowcol_to_rank, +) +from mortie import generate_morton_children +from viewers import BLOCK_ORDER, SIDE, UNITS, human_bytes + +__all__ = ["fit_window", "registered_pair", "voxel_chips"] + + +def _wq(zs, ws, q): + """Weighted quantile of pre-sorted `zs` with weights `ws`. + + Positions are the MIDPOINT of each centroid's weight interval, + ``(cumsum - w/2) / total``. Interpolating on the raw cumsum instead biases + every quantile low by half a centroid -- the median of two points at 0 and + 100 comes back 0. + """ + cum = np.cumsum(ws) + if cum[-1] <= 0: + return float(zs[0]) + return float(np.interp(q * cum[-1], cum - 0.5 * ws, zs)) + + +def fit_window(z, wt, n_bins, dz, first=0.02, step=0.05, max_drop=0.5): + """Fit `z` into a FIXED `n_bins * dz` window by trimming tails. + + Returns ``(z0, lo_drop, hi_drop, kept)`` -- window floor, the weight fraction + trimmed off each end, and the fraction that survives. + + The window is fixed because the grid is isotropic: `dz` is the cell edge, so + the cube really is a cube. When the data does not fit, something has to go, + and the only honest choices are which end and how much. + + The first cut takes `first` off whichever tail lies further from the WEIGHTED + median -- the mass sits near the median, so the far tail is the cheaper side + to lose. Cuts then alternate in `step` increments so neither end is + preferentially eaten. If `max_drop` is reached without fitting, the window is + centred on the median and `kept` reports what actually lands inside it, which + is worse but never silent. + """ + order = np.argsort(z, kind="stable") + zs, ws = np.asarray(z)[order], np.asarray(wt)[order] + span = n_bins * dz + total = float(ws.sum()) + z0 = float(np.floor(zs[0])) + # Test the window we will ACTUALLY use. `z0` is floored, so it sits up to + # 1 m below `zs[0]`; checking `zs[-1] - zs[0] <= span` passes data that then + # falls off the top of the window and is dropped while this reports kept=1. + if total <= 0 or zs[-1] < z0 + span: + return z0, 0.0, 0.0, 1.0 + + med = _wq(zs, ws, 0.5) + drop_high = (zs[-1] - med) >= (med - zs[0]) # start on the far tail + lo_d = hi_d = 0.0 + cut = first + while lo_d + hi_d + cut <= max_drop: + if drop_high: + hi_d += cut + else: + lo_d += cut + drop_high, cut = not drop_high, step + lo, hi = _wq(zs, ws, lo_d), _wq(zs, ws, 1.0 - hi_d) + if hi - lo <= span: + # `kept` is what LANDS in the window, not the nominal quantile drop: + # the window is wider than the trimmed range and recaptures some of it. + z0 = float(np.floor(lo)) + return z0, lo_d, hi_d, float(ws[(zs >= z0) & (zs < z0 + span)].sum() / total) + + z0 = float(np.floor(med - span / 2)) # give up trimming; centre on the mass + inside = float(ws[(zs >= z0) & (zs < z0 + span)].sum() / total) + return z0, None, None, inside # None, not NaN: `meta.json` has to parse + + +def voxel_chips(handles, block, sensor="atl03", order=22, side=128, path=None): + """One o12 block -> isotropic `side`**3 count chips, empty chips skipped. + + Shifting a centroid's block-local rank right by ``2 * (29 - order)`` truncates + its order-29 point word to `order` -- nested ranks are hierarchical -- which is + how the cube gets finer than the stored cells. The z bin equals the cell edge + and `n_bins` equals `side`, so every chip is a true cube; each keeps its own + `z0` and trim record in `meta.json`. + """ + store, field = handles[sensor] + n_bins = side + dz = SIDE / 2 ** (order - BLOCK_ORDER) # isotropic: z bin == cell edge + depth = order - (side.bit_length() - 1) - BLOCK_ORDER + tiles = generate_morton_children(int(block), order - (side.bit_length() - 1)) + + t0 = time.perf_counter() + got = list(mz.read_ragged(store, field, locations=True, subtree=mz.morton_decimal(int(block)))) + v = np.concatenate([np.asarray(r[1]) for r in got]) + z, wt = v[:, 0], v[:, 1] + rank = block_rank(np.concatenate([np.asarray(r[2], np.uint64) for r in got]), BLOCK_ORDER)[0] + row, col = rank_to_rowcol(rank >> np.uint64(2 * (29 - order)), order - BLOCK_ORDER) + read_s = time.perf_counter() - t0 + + # Stream each chip out and drop it: 64 x 8 MiB held at once is 512 MiB, and + # mybinder caps the container at 2 GB. An .npz is a zip of .npy members. + path = path or f"{sensor}_o{order}_chips_{mz.morton_decimal(int(block))}.npz" + t1 = time.perf_counter() + meta, kept, filled, dense, trimmed = {}, 0, 0, 0, 0 + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED, compresslevel=1) as zf: + for i in range(2**depth): + for j in range(2**depth): + m = np.flatnonzero((row // side == i) & (col // side == j)) + if not len(m): + continue + z0, lo_d, hi_d, frac = fit_window(z[m], wt[m], n_bins, dz) + # np.floor, not a bare cast: `astype` truncates TOWARD ZERO, so a + # centroid in [z0 - dz, z0) gives -0.x -> 0, passes `iz >= 0`, and + # lands in bin 0 at the wrong height. `fit_window` can put z0 above + # the minimum, so this is reachable whenever a chip is trimmed. + iz = np.floor((z[m] - z0) / dz).astype(np.int64) + keep = (iz >= 0) & (iz < n_bins) + trimmed += int(wt[m][~keep].sum()) + mk, iz = m[keep], iz[keep] + acc = np.zeros((side, side, n_bins), dtype=np.float32) + np.add.at(acc, (row[mk] % side, col[mk] % side, iz), wt[mk]) + chip = np.rint(acc).astype(np.uint32) # merged centroids weigh fractionally + name = mz.morton_decimal(int(tiles[rowcol_to_rank(i, j, depth=depth)])) + buf = BytesIO() + np.save(buf, chip) + zf.writestr(f"{name}.npy", buf.getvalue()) + meta[name] = { + "z0": z0, + "dz": dz, + "written": int(chip.sum()), + "trim_low": lo_d, + "trim_high": hi_d, + "kept": frac, + } + kept += int(chip.sum()) + filled += int(np.count_nonzero(chip)) + dense += chip.nbytes + zf.writestr("meta.json", json.dumps(meta)) # z0 + trim per chip + + n, on_disk = 2 ** (2 * depth), os.path.getsize(path) + cut = [k for k, mm in meta.items() if mm["kept"] < 1.0] + print(f"{sensor} o{order} — {dz:.3f} m isotropic voxels, {side}^3 = {dz * side:.0f} m cubes") + print(f" read {len(z):,} centroids, {wt.sum():,.0f} {UNITS[sensor]}, {read_s:.1f}s") + print( + f" wrote {len(meta)} of {n} chips, {kept:,} of {wt.sum():,.0f} {UNITS[sensor]}" + + (f" ({trimmed:,} trimmed to fit)" if trimmed else " (nothing trimmed)") + + f", {time.perf_counter() - t1:.1f}s" + ) + if cut: + worst = min(meta[k]["kept"] for k in cut) + print( + f" {len(cut)} chip(s) needed a trim; worst kept {100 * worst:.1f}% of its weight" + ) + # Dense size is the tensors in MEMORY; on-disk is deflate on runs of zeros. + print( + f" {human_bytes(dense)} in memory -> {human_bytes(on_disk)} on disk " + f"({dense / on_disk:.0f}x — {100 * filled / (dense // 4):.3f}% of voxels occupied)" + ) + return path, meta + + +def registered_pair(handles, block, order=19, n_bins=128, resolution=0.5, path=None): + """Both sensors as cubes of ONE shape: one xy lattice, one z axis. + + Each GEDI o18 cell is REPLICATED into its four o19 children rather than ATL03 + being merged up to o18 -- so a GEDI cube sums to four times its stored weight. + `chunk_z_range` is handed both sensors' digests together; derived per sensor it + puts them bins apart, and nothing downstream notices. + """ + t0 = time.perf_counter() + side = 2 ** (order - BLOCK_ORDER) + got = { + n: list(mz.read_ragged(store, field, subtree=mz.morton_decimal(int(block)))) + for n, (store, field) in handles.items() + } + read_s = time.perf_counter() - t0 + + # What each sensor would get ALONE, against the window taken over both. + window = dict( + n_bins=n_bins, resolution=resolution, bottom=0.05, top=0.95, fit="degrade_resolution" + ) + solo = { + n: chunk_z_range([np.asarray(v) for _w, v in rows], **window) for n, rows in got.items() + } + z0, n_bins, dz = chunk_z_range( + [np.asarray(v) for rows in got.values() for _w, v in rows], **window + ) + + t1 = time.perf_counter() + cubes, lines = {}, [] + for name, (_store, field) in handles.items(): + cell_order = int(field.split("/", 1)[0]) + k = 2 ** (order - cell_order) # children of one cell on the output grid + words = np.array([w for w, _v in got[name]], dtype=np.uint64) + r, c = rank_to_rowcol(block_rank(words, BLOCK_ORDER)[0], cell_order - BLOCK_ORDER) + cube = np.zeros((side, side, n_bins), dtype=np.float32) + for i, (_w, v) in enumerate(got[name]): + cube[r[i] * k : (r[i] + 1) * k, c[i] * k : (c[i] + 1) * k] = rasterize_cell( + np.asarray(v), z0, dz, n_bins + ) + cubes[name] = cube + lines.append( + f" {name}: {len(words):,} o{cell_order} cells -> {k}x{k} -> " + f"{int((cube.sum(2) > 0).sum()):,}/{side * side:,} columns, " + f"{cube.sum():,.0f} {UNITS[name]}" + (f" ({k**2}x replicated)" if k > 1 else "") + ) + + path = path or f"registered_o{order}_{mz.morton_decimal(int(block))}.npz" + np.savez_compressed(path, **cubes, z0=z0, dz=dz, order=order) + dense, on_disk = sum(c.nbytes for c in cubes.values()), os.path.getsize(path) + fit = "as asked" if abs(dz - resolution) < 1e-9 else f"DEGRADED from {resolution:g} m" + print( + f"registered o{order} — {next(iter(cubes.values())).shape} float32 each, {read_s:.1f}s read" + ) + print( + f" grid shared z = {z0:.1f} m + bin * {dz:g} m ({fit}) — alone they would be " + + " vs ".join(f"{n} {s0:.1f}/{g:g} m" for n, (s0, _b, g) in solo.items()) + ) + print("\n".join(lines)) + print( + f" wrote {human_bytes(dense)} in memory -> {human_bytes(on_disk)} on disk " + f"in {path}, {time.perf_counter() - t1:.1f}s" + ) + return path, cubes diff --git a/notebooks/hhdc_viewer.ipynb b/notebooks/hhdc_viewer.ipynb new file mode 100644 index 000000000..6025564ac --- /dev/null +++ b/notebooks/hhdc_viewer.ipynb @@ -0,0 +1,242 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c008e8d60d5c9d12", + "metadata": {}, + "source": [ + "# polygon \u2192 MOC \u2192 shard \u2192 3-D \u2192 numpy\n", + "\n", + "The whole zagg read stack in two libraries, one `%pip install`, and zero\n", + "credentials. A geojson polygon becomes a morton MOC; the MOC checks itself\n", + "against the store's own coverage; the covered shards open with timings; one\n", + "shard renders in 3-D (ATL03 + GEDI together); the current view exports to\n", + "voxel cubes on any grid you name and saves to disk.\n", + "\n", + "Everything below is reader-side and calls no zagg public API \u2014 `mortie` for\n", + "the geometry, `moczarr` for the store, plus the t-digest algebra that\n", + "`moczarr[zagg]` imports from zagg rather than vendoring (moczarr issue #19).\n", + "That algebra is the only zagg code on the path. It all runs anonymously\n", + "against public S3, binder-ready.\n", + "\n", + "Its sibling is [`waveform_viewer.ipynb`](waveform_viewer.ipynb), which takes\n", + "the same polygon and stores down to the cell-level join: one GEDI o18\n", + "footprint against the 2\u00d72 ATL03 o19 cells beneath it, both rebuilt from their\n", + "stored t-digests. The two are separate notebooks because this one needs\n", + "`%matplotlib widget` and that one `%matplotlib inline`; the backends collide\n", + "in a single kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b853ebb386902b20", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install -q mortie \"moczarr[zagg]>=0.7\" matplotlib ipympl ipywidgets\n", + "%matplotlib widget\n", + "\n", + "\n", + "import moczarr as mz\n", + "import numpy as np\n", + "\n", + "# The drawing lives in viewers.py and the exports in export.py, both beside this\n", + "# notebook, so the cells here stay about the READ path.\n", + "from export import registered_pair, voxel_chips\n", + "from mortie import moc\n", + "from viewers import densest_shard, view3d, viewer_stats\n", + "\n", + "# One store per product, each appendable. Coverage answers which ground the\n", + "# store holds; the store name never does.\n", + "STORES = {\n", + " \"atl03\": (\n", + " \"s3://us-west-2.opendata.source.coop/englacial/zagg/demo/atl03_tdigest_o9.zarr\",\n", + " \"19/h_tdigest_signal\",\n", + " ),\n", + " \"gedi\": (\n", + " \"s3://us-west-2.opendata.source.coop/englacial/zagg/demo/gedi_flux_o9.zarr\",\n", + " \"18/rx_flux\",\n", + " ),\n", + "}\n", + "S3 = {\"region\": \"us-west-2\", \"anonymous\": True}" + ] + }, + { + "cell_type": "markdown", + "id": "e36de6ce8afeb4dd", + "metadata": {}, + "source": [ + "## One polygon in, covered shards out\n", + "\n", + "The polygon below sits on SERC, the mid-Atlantic forest site the HHDC\n", + "diffusion papers are built on. Replace it with any area within California or a\n", + "NEON AOP site \u2014 it need not be a box \u2014 and the cell tests whether each store\n", + "actually covers what you asked for. An AOI in Alaska will pass\n", + "the ATL03 check and fail the GEDI one \u2014 GEDI flies on the ISS, so it sees no\n", + "higher than |lat| 51.6 and the Alaska NEON sites are outside its reach." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d395a2afc0f117b7", + "metadata": {}, + "outputs": [], + "source": [ + "aoi = {\n", + " \"features\": [\n", + " {\n", + " \"geometry\": {\n", + " \"coordinates\": [\n", + " [\n", + " [-76.5750, 38.9000],\n", + " [-76.5450, 38.9000],\n", + " [-76.5450, 38.8780],\n", + " [-76.5600, 38.8720],\n", + " [-76.5750, 38.8800],\n", + " [-76.5750, 38.9000],\n", + " ]\n", + " ]\n", + " }\n", + " }\n", + " ]\n", + "}\n", + "\n", + "q = moc(aoi)\n", + "shards = None\n", + "for name, (root, _field) in STORES.items():\n", + " assert mz.coverage_moc(root, **S3).contains(q), f\"{name} does not contain the polygon\"\n", + " ids = set(mz.candidate_shards(root, aoi=q, **S3))\n", + " shards = ids if shards is None else shards & ids\n", + "shards = sorted(shards)\n", + "\n", + "# Most GEDI granules, the rule demo/06_paired used. Named once, used below.\n", + "SHARD = densest_shard(STORES[\"gedi\"][0], shards, **S3)\n", + "print(f\"{len(shards)} shards cover the polygon: {shards}\\nworking {SHARD}\")" + ] + }, + { + "cell_type": "markdown", + "id": "28c4719bbd152619", + "metadata": {}, + "source": [ + "## Open one shard \u2014 and price a full sweep of one field\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1516849303d73cba", + "metadata": {}, + "outputs": [], + "source": [ + "# Open the leaf for each store. `handles` is all the exports need; the sweep\n", + "# below is separate, and only the viewer's dropdown labels depend on it.\n", + "handles = {n: (mz.open_leaf(root, SHARD, **S3), field) for n, (root, field) in STORES.items()}\n", + "stats = viewer_stats(handles)\n" + ] + }, + { + "cell_type": "markdown", + "id": "c15ec274445cfcbf", + "metadata": {}, + "source": [ + "## The 3-D view \u2014 both sensors, exact centroids, time-aware" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4760d9b9fdcfae44", + "metadata": {}, + "outputs": [], + "source": [ + "view = view3d(handles, SHARD, stats) # blocks ordered by coincident cells\n" + ] + }, + { + "cell_type": "markdown", + "id": "7fb27b941602401d91542211134fc71a", + "metadata": {}, + "source": [ + "## Export \u2014 voxel cubes, on a grid you choose\n", + "\n", + "Exporting tensors is computationally cheap but not simple: we have to decide what\n", + "happens when a chip's data range exceeds the range bins available at the requested\n", + "resolution, and how to align range bins between sensors when exporting jointly.\n", + "\n", + "Both are the user's call, i.e. client-side functions. Because the original locations\n", + "are stored, a tensor can be built on a 12 m, 1.5 m or 0.38 m spatial grid \u2014 down to\n", + "1.21 cm at order 29. In the examples below, range bins are isotropic with the\n", + "spatial grid by default; the joint export takes whatever z resolution you ask\n", + "for, but coarsens it when the two sensors' shared relief will not fit \u2014 and\n", + "prints that it did.\n", + "\n", + "**1 \u2014 ATL03 alone, finer than the store.** Default o22: 1.554 m voxels, z binned to\n", + "match, the block emitted as 8\u00d78 isotropic 128\u00b3 chips, empty ones skipped. Where a\n", + "chip's relief will not fit, `fit_window` trims the tail furthest from the weighted\n", + "median first and records the cost \u2014 nothing is clipped silently.\n", + "\n", + "**2 \u2014 ATL03 and GEDI co-registered**, ready to stack: one xy lattice (ATL03's o19,\n", + "each GEDI o18 cell replicated into its four children) and one z axis for both.\n", + "`read_tensors` derives its window *per sensor*, so two tensors of the same block can\n", + "differ in origin *and* in bin height \u2014 at which point they share no axis at all. The\n", + "cell prints what each sensor would have got alone beside the shared window it uses.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "acae54e37e7d407bbb7b55eff062a284", + "metadata": {}, + "outputs": [], + "source": [ + "chips, manifest = voxel_chips(handles, view.block) # the block on screen, at o22\n", + "# chips24, _ = voxel_chips(handles, view.block, order=24) # 0.389 m voxels, 1,024 chips\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a63283cbaf04dbcab1f6479b197f3a8", + "metadata": {}, + "outputs": [], + "source": [ + "pair, cubes = registered_pair(handles, view.block)\n", + "stacked = np.stack([cubes[\"atl03\"], cubes[\"gedi\"]], axis=0) # registered, so they stack\n", + "print(f\"stacked {stacked.shape} \u2014 ready for a 2-channel model\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "8dd0d8092fe74a7c96281538738b07e2", + "metadata": {}, + "source": [ + "Two libraries, one polygon \u2014 coverage, shards, timings, the paired 3-D view,\n", + "and co-registered voxel cubes on disk." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/jupyterhub_example.ipynb b/notebooks/jupyterhub_example.ipynb deleted file mode 100644 index 486a34a39..000000000 --- a/notebooks/jupyterhub_example.ipynb +++ /dev/null @@ -1,397 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "85f7f7db", - "metadata": {}, - "source": "# Running zagg from a JupyterHub\n\n[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/jupyterhub_example.ipynb)\n\nThis notebook shows how to drive the `zagg` Python API from a Jupyter\nnotebook (e.g. a science JupyterHub such as CryoCloud or Pangeo, or your own\nlaptop). It covers **local processing** -- running the aggregation on the hub\nitself with threads, no AWS Lambda required.\n\n> **Lambda fan-out** for full-scale runs lives in the\n> [cryocloud_example](cryocloud_example.ipynb) notebook, which dispatches to a\n> deployed AWS Lambda and therefore cannot run on Binder.\n\n### What runs where\n\n| Section | Needs credentials? | Runs on Binder? |\n|---------|--------------------|-----------------|\n| Build a catalog (CMR-STAC metadata query) | none -- the CMR-STAC search is anonymous | **yes** |\n| Local `agg()` -- fetch granule pixels from NSIDC | **NASA Earthdata login** | no |\n| Read & visualize a published result | none (anonymous source.coop) | **yes** |\n\nBuilding the catalog is a pure **metadata** query against CMR-STAC and needs\n**no login** -- it runs on Binder (the HEALPix `mortie` coverage backend is the\ndefault, so the non-PyPI [spherely fork](https://github.com/espg/spherely) is\n*not* required either). Only *fetching the granule pixels* -- the `agg()` step\nthat reads ICESat-2 HDF5 from NSIDC -- needs a NASA Earthdata login, so those\ncells are written to **skip cleanly** when credentials are absent. The final\nsection reads a public, anonymous result and produces the plots end-to-end." - }, - { - "cell_type": "markdown", - "id": "auth-header", - "metadata": {}, - "source": [ - "## 1. Authentication\n", - "\n", - "zagg needs two sets of credentials:\n", - "\n", - "- **NASA Earthdata** -- for reading source data (ICESat-2 HDF5 on NSIDC S3)\n", - "- **AWS** (optional) -- for writing output to S3 or invoking Lambda\n", - "\n", - "Both use standard credential discovery -- if you already have `~/.netrc`\n", - "and `~/.aws/credentials` configured, skip this section entirely.\n", - "\n", - "### Default credential locations\n", - "\n", - "| Service | File | Format |\n", - "|---------|------|--------|\n", - "| Earthdata | `~/.netrc` | `machine urs.earthdata.nasa.gov login USERNAME password PASSWORD` |\n", - "| AWS | `~/.aws/credentials` | `[default]` profile with `aws_access_key_id` / `aws_secret_access_key` |\n", - "\n", - "### Operator-managed hub (CryoCloud, Pangeo)\n", - "\n", - "Credentials are pre-configured as environment variables or mounted files.\n", - "Nothing to do -- `earthaccess` and `boto3` discover them automatically.\n", - "\n", - "### Override via environment variables\n", - "\n", - "Only needed if the standard files above are not present:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "auth-setup", - "metadata": { - "execution": { - "iopub.execute_input": "2026-04-07T23:52:04.491117Z", - "iopub.status.busy": "2026-04-07T23:52:04.490578Z", - "iopub.status.idle": "2026-04-07T23:52:04.497026Z", - "shell.execute_reply": "2026-04-07T23:52:04.495779Z" - } - }, - "outputs": [], - "source": [ - "# Only needed if ~/.netrc is not configured:\n", - "# os.environ[\"EARTHDATA_USERNAME\"] = \"your_username\"\n", - "# os.environ[\"EARTHDATA_PASSWORD\"] = \"your_password\"\n", - "\n", - "# Only needed if ~/.aws/credentials is not configured:\n", - "# os.environ[\"AWS_ACCESS_KEY_ID\"] = \"AKIA...\"\n", - "# os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"...\"\n", - "# os.environ[\"AWS_DEFAULT_REGION\"] = \"us-west-2\"" - ] - }, - { - "cell_type": "markdown", - "id": "config-header", - "metadata": {}, - "source": [ - "## 2. Load a pipeline config\n", - "\n", - "The config defines what data to read, how to aggregate it, and where to\n", - "write output. Use the built-in ATL06 config or load a custom YAML." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "load-config", - "metadata": { - "execution": { - "iopub.execute_input": "2026-04-07T23:52:04.500720Z", - "iopub.status.busy": "2026-04-07T23:52:04.500405Z", - "iopub.status.idle": "2026-04-07T23:52:05.134933Z", - "shell.execute_reply": "2026-04-07T23:52:05.133872Z" - } - }, - "outputs": [], - "source": [ - "from zagg import default_config, get_child_order\n", - "\n", - "config = default_config(\"atl06\")\n", - "\n", - "print(f\"Reader: {config.data_source['reader']}\")\n", - "print(f\"Groups: {config.data_source['groups']}\")\n", - "print(f\"Variables: {list(config.data_source['variables'].keys())}\")\n", - "print(f\"Child order: {get_child_order(config)}\")" - ] - }, - { - "cell_type": "markdown", - "id": "46615328", - "metadata": {}, - "source": "## 3. Build a catalog *(anonymous -- runs on Binder)*\n\nThe catalog (a *ShardMap*) maps grid cells to granule hrefs. It is built from a\n**CMR-STAC metadata query**, which is anonymous: no NASA Earthdata login is\nneeded to *find* the granules (only *fetching* their pixels in section 4 needs\nauth). The build is driven by the **same config** as the run so the shard grid\ncan't drift from the output grid, and the default HEALPix `mortie` coverage\nbackend means the non-PyPI [spherely fork](https://github.com/espg/spherely) is\nnot required here.\n\nThe cell below builds a small Antarctic catalog over a short date window so it\nfinishes quickly on Binder. Equivalently from the CLI:\n\n```bash\n# Bundled atl06 config; small bbox + a few days so the build is quick:\npython -m zagg.catalog --config atl06.yaml --short-name ATL06 --version 006 \\\n --start-date 2025-04-01 --end-date 2025-04-03 \\\n --bbox=-70,-78,-65,-75 --backend mortie\n# -> writes shardmap_ATL06__.json\n```\n\n(For a full cycle, swap the bbox/dates for `--cycle 22 --polygon antarctica.geojson`.)" - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b964db81", - "metadata": {}, - "outputs": [], - "source": [ - "# Build a small ATL06 catalog anonymously (CMR-STAC metadata only -- no login).\n", - "# This is the same work `python -m zagg.catalog ...` does, via the Python API.\n", - "from zagg.catalog.shardmap import ShardMap\n", - "from zagg.catalog.sources import CMRSource, Query\n", - "from zagg.grids import from_config\n", - "\n", - "catalog_path = \"shardmap_ATL06_2025-04-01_2025-04-03.json\"\n", - "grid = from_config(config)\n", - "\n", - "# Small bbox + a few days so the build is quick on Binder.\n", - "query = Query(\n", - " short_name=\"ATL06\",\n", - " version=\"006\",\n", - " start_date=\"2025-04-01\",\n", - " end_date=\"2025-04-03\",\n", - " region=(-70.0, -78.0, -65.0, -75.0), # lon_min, lat_min, lon_max, lat_max\n", - ")\n", - "\n", - "try:\n", - " cat = CMRSource().fetch(query)\n", - " print(f\"Fetched {len(cat)} granules ({query.collection})\")\n", - " # backend=\"mortie\": HEALPix coverage, so no spherely fork needed.\n", - " sm = ShardMap.build(cat, grid, backend=\"mortie\")\n", - " sm.to_json(catalog_path)\n", - " HAVE_CATALOG = True\n", - " print(\n", - " f\"ShardMap: {len(sm.shard_keys)} shards, \"\n", - " f\"{sm.metadata['total_pairs']} pairs -> {catalog_path}\"\n", - " )\n", - "except Exception as exc:\n", - " # CMR-STAC is anonymous but needs outbound network; degrade cleanly on an\n", - " # offline / egress-restricted Binder so the rest of the notebook still runs.\n", - " HAVE_CATALOG = False\n", - " print(\n", - " f\"Catalog build skipped ({type(exc).__name__}: {exc}).\\n\"\n", - " \"If this environment can't reach cmr.earthdata.nasa.gov, jump to the \"\n", - " \"public-result section, which reads source.coop directly.\"\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "d0624f49", - "metadata": {}, - "source": "## 4. Local processing *(requires NASA Earthdata login)*\n\nThe catalog above is just metadata; **fetching the granule pixels** is what\nneeds a NASA Earthdata login. Process cells locally with a thread pool -- no\nLambda needed. Good for testing, small regions, or hubs without AWS set up.\n`driver=\"https\"` reads ICESat-2 HDF5 over HTTPS (works from anywhere, still\nneeds Earthdata auth); `driver=\"s3\"` is faster but only from `us-west-2`.\n\nThese cells run only when `HAVE_CATALOG` is true (i.e. the catalog built above)\n**and** an Earthdata login is configured; without credentials they skip cleanly." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "af6ee4c3", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg import agg\n", - "\n", - "if HAVE_CATALOG:\n", - " # Dry run -- preview the plan without processing.\n", - " preview = agg(\n", - " config,\n", - " catalog=catalog_path,\n", - " store=\"./test_output.zarr\",\n", - " driver=\"https\",\n", - " dry_run=True,\n", - " )\n", - " print(preview)\n", - "else:\n", - " print(\"Skipped: no catalog (needs an Earthdata login).\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9ce98e8f", - "metadata": {}, - "outputs": [], - "source": [ - "import logging\n", - "\n", - "if HAVE_CATALOG:\n", - " logging.basicConfig(level=logging.INFO, format=\"%(message)s\")\n", - " # Process 2 cells locally, write to a local Zarr. driver=\"https\" works\n", - " # outside us-west-2.\n", - " results = agg(\n", - " config,\n", - " catalog=catalog_path,\n", - " store=\"./test_output.zarr\",\n", - " driver=\"https\",\n", - " max_cells=2,\n", - " max_workers=2,\n", - " overwrite=True,\n", - " )\n", - " print(f\"Cells with data: {results['cells_with_data']}\")\n", - " print(f\"Total observations: {results['total_obs']:,}\")\n", - " print(f\"Wall time: {results['wall_time_s']:.1f}s\")\n", - "else:\n", - " print(\"Skipped: no catalog (needs an Earthdata login).\")" - ] - }, - { - "cell_type": "markdown", - "id": "e513a9a9", - "metadata": {}, - "source": [ - "## 5. Inspect the local output\n", - "\n", - "The output is a Zarr v3 store following the DGGS convention. (Runs only if the\n", - "local processing above produced `./test_output.zarr`.)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f041ed8e", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import zarr\n", - "\n", - "if HAVE_CATALOG:\n", - " store = zarr.open_group(\"./test_output.zarr\", mode=\"r\")\n", - " child_order = get_child_order(config)\n", - " group = store[str(child_order)]\n", - "\n", - " print(f\"Arrays: {list(group.array_keys())}\")\n", - " for nm in [\"count\", \"h_mean\"]:\n", - " arr = group[nm]\n", - " print(f\" {nm}: shape={arr.shape}, dtype={arr.dtype}\")\n", - "\n", - " count = group[\"count\"][:]\n", - " has_data = count > 0\n", - " print(f\"\\nCells with data: {has_data.sum():,} / {len(count):,}\")\n", - " print(f\"Total observations: {count[has_data].sum():,}\")\n", - "else:\n", - " print(\"Skipped: no local output to inspect.\")" - ] - }, - { - "cell_type": "markdown", - "id": "bbfe8ae4", - "metadata": {}, - "source": "## 6. Read & visualize a published result *(no credentials -- runs on Binder)*\n\nThe `agg()` fetch above needs an Earthdata login. This section instead reads a\n**public, anonymous** ATL06 cycle-22 aggregation published on\n[source.coop](https://source.coop/englacial/zagg/benchmarks)\n(`englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr`), so it runs end-to-end on Binder.\n\nWe open the HEALPix order-12 store read-only over an unsigned S3 request, then\nscatter the filled cells on an Antarctic Polar Stereographic map." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4785a44f", - "metadata": {}, - "outputs": [], - "source": [ - "import xarray as xr\n", - "\n", - "from zagg import open_store\n", - "\n", - "# Public source.coop store -- unsigned (anonymous) read works from anywhere;\n", - "# open_store applies zagg's read-only retry policy (issue #186).\n", - "PUBLIC_BUCKET = \"us-west-2.opendata.source.coop\"\n", - "PUBLIC_PREFIX = \"englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr\"\n", - "CHILD_ORDER = 12\n", - "\n", - "pub = xr.open_dataset(\n", - " open_store(f\"s3://{PUBLIC_BUCKET}/{PUBLIC_PREFIX}\", read_only=True, skip_signature=True),\n", - " engine=\"zarr\",\n", - " consolidated=False,\n", - " zarr_format=3,\n", - " group=str(CHILD_ORDER),\n", - ")\n", - "\n", - "has_data = pub[\"count\"].values > 0\n", - "cell_ids = pub[\"cell_ids\"].values[has_data].astype(np.int64)\n", - "h_mean = pub[\"h_mean\"].values[has_data]\n", - "count = pub[\"count\"].values[has_data].astype(np.float64)\n", - "print(f\"Filled HEALPix cells: {has_data.sum():,} / {has_data.size:,}\")\n", - "print(f\"Total observations: {count.sum():.0f}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f76c552d", - "metadata": {}, - "outputs": [], - "source": [ - "import cartopy.crs as ccrs\n", - "import cartopy.feature as cfeature\n", - "import matplotlib.pyplot as plt\n", - "from mortie._healpix import pix2ang\n", - "\n", - "# HEALPix cell centers -> lon/lat (NESTED), let cartopy project to polar stereo.\n", - "# pix2ang takes the HEALPix order (depth), not nside.\n", - "lon, lat = pix2ang(CHILD_ORDER, cell_ids)\n", - "\n", - "proj = ccrs.SouthPolarStereo()\n", - "fig, axes = plt.subplots(1, 2, figsize=(14, 6), subplot_kw={\"projection\": proj})\n", - "for ax in axes:\n", - " ax.coastlines(resolution=\"50m\", linewidth=0.5)\n", - " ax.add_feature(cfeature.LAND, facecolor=\"lightgray\", alpha=0.3)\n", - " ax.gridlines(draw_labels=False, alpha=0.3)\n", - " ax.set_extent([-180, 180, -90, -60], crs=ccrs.PlateCarree())\n", - "\n", - "im = axes[0].scatter(\n", - " lon, lat, c=h_mean, s=0.05, cmap=\"terrain\", vmin=0, vmax=4000, transform=ccrs.PlateCarree()\n", - ")\n", - "axes[0].set_title(f\"Mean Elevation ({has_data.sum():,} cells)\")\n", - "plt.colorbar(im, ax=axes[0], label=\"m\", shrink=0.7)\n", - "\n", - "im = axes[1].scatter(\n", - " lon,\n", - " lat,\n", - " c=count,\n", - " s=0.05,\n", - " cmap=\"viridis\",\n", - " norm=plt.matplotlib.colors.LogNorm(vmin=1),\n", - " transform=ccrs.PlateCarree(),\n", - ")\n", - "axes[1].set_title(\"Observation Count\")\n", - "plt.colorbar(im, ax=axes[1], label=\"count (log)\", shrink=0.7)\n", - "\n", - "plt.suptitle(\"ATL06 cycle 22 (published source.coop result)\", fontsize=14, weight=\"bold\")\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "4eb8c43f", - "metadata": {}, - "source": [ - "## 7. Custom configs\n", - "\n", - "You can embed defaults (catalog path, store path) directly in the config object\n", - "so you don't repeat them in every `agg()` call, or load a fully custom YAML with\n", - "different variables and aggregation functions.\n", - "\n", - "See the [custom_aggregations](custom_aggregations.ipynb) notebook for examples\n", - "of modifying statistics and adding new input variables." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "custom-config", - "metadata": { - "execution": { - "iopub.execute_input": "2026-04-07T23:57:36.413553Z", - "iopub.status.busy": "2026-04-07T23:57:36.413386Z", - "iopub.status.idle": "2026-04-07T23:57:36.416784Z", - "shell.execute_reply": "2026-04-07T23:57:36.415991Z" - } - }, - "outputs": [], - "source": [ - "# Load from a custom YAML file:\n", - "# config = load_config(\"my_custom_config.yaml\")\n", - "\n", - "# Or set catalog/store in the config itself:\n", - "config.catalog = catalog_path\n", - "config.output[\"store\"] = \"./output.zarr\"\n", - "\n", - "# Then agg() uses them as defaults (CLI/kwargs still override):\n", - "# results = agg(config, max_cells=5)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/rasterized_zarr.ipynb b/notebooks/rasterized_zarr.ipynb deleted file mode 100644 index 6859b85f6..000000000 --- a/notebooks/rasterized_zarr.ipynb +++ /dev/null @@ -1,516 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# HEALPix to Polar Stereographic Raster (8 km)\n", - "\n", - "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/rasterized_zarr.ipynb)\n", - "\n", - "_Runs on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/rasterized_zarr.ipynb). It reads the public, anonymously accessible ATL06 cycle-22 aggregation published on [source.coop](https://source.coop/englacial/zagg/benchmarks) (`englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr`) -- no credentials required._\n", - "\n", - "\n", - "For each cell in an 8 km EPSG:3031 grid, find all HEALPix order-12 cells\n", - "whose centers fall within it, then aggregate.\n", - "\n", - "Source resolution ~0.86 km, target 8 km, so ~87 source cells per output pixel." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import warnings\n", - "from datetime import timedelta\n", - "\n", - "import cartopy.crs as ccrs\n", - "import cartopy.feature as cfeature\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import xarray as xr\n", - "from mortie._healpix import pix2ang\n", - "from pyproj import Transformer\n", - "from zarr import config\n", - "\n", - "from zagg import open_store\n", - "\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 1. Load zarr and filter to cells with data" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "S3_BUCKET = \"us-west-2.opendata.source.coop\"\n", - "S3_PREFIX = \"englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr\"\n", - "CHILD_ORDER = 12\n", - "NSIDE = 2**CHILD_ORDER # 4096\n", - "\n", - "config.set({\"async.concurrency\": 128})\n", - "# Anonymous read via open_store (issue #186); this notebook's bespoke client\n", - "# tuning rides through open_store's kwargs to obstore and overrides the\n", - "# read-only retry default.\n", - "store = open_store(\n", - " f\"s3://{S3_BUCKET}/{S3_PREFIX}\",\n", - " read_only=True,\n", - " skip_signature=True,\n", - " client_options={\"connect_timeout\": timedelta(seconds=60)},\n", - " retry_config={\n", - " \"max_retries\": 2,\n", - " \"retry_timeout\": timedelta(seconds=5),\n", - " \"backoff\": {\n", - " \"init_backoff\": timedelta(seconds=5),\n", - " \"max_backoff\": timedelta(seconds=30),\n", - " \"base\": 1.2,\n", - " },\n", - " },\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "# Load raw zarr -- no xdggs decode needed\n", - "ds = xr.open_dataset(\n", - " store,\n", - " engine=\"zarr\",\n", - " consolidated=False,\n", - " zarr_format=3,\n", - " group=str(CHILD_ORDER),\n", - ")\n", - "\n", - "# Filter to cells with actual observations\n", - "has_data = ds[\"count\"].values > 0\n", - "cell_ids = ds[\"cell_ids\"].values[has_data]\n", - "count = ds[\"count\"].values[has_data].astype(np.float64)\n", - "h_mean = ds[\"h_mean\"].values[has_data].astype(np.float64)\n", - "h_sigma = ds[\"h_sigma\"].values[has_data].astype(np.float64)\n", - "h_min = ds[\"h_min\"].values[has_data].astype(np.float64)\n", - "h_max = ds[\"h_max\"].values[has_data].astype(np.float64)\n", - "h_variance = ds[\"h_variance\"].values[has_data].astype(np.float64)\n", - "h_q25 = ds[\"h_q25\"].values[has_data].astype(np.float64)\n", - "h_q50 = ds[\"h_q50\"].values[has_data].astype(np.float64)\n", - "h_q75 = ds[\"h_q75\"].values[has_data].astype(np.float64)\n", - "\n", - "print(f\"Total cells in zarr: {len(has_data):,}\")\n", - "print(f\"Cells with count > 0: {has_data.sum():,}\")\n", - "print(f\"Cells with count = 0 (discarded): {(~has_data).sum():,}\")\n", - "print(f\"Total observations: {count.sum():.0f}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 2. HEALPix cell centers to polar stereographic coordinates\n", - "\n", - "Use mortie's Rust-backed `pix2ang` (returns lon/lat in degrees, NESTED scheme)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "# pix2ang takes the HEALPix order (depth), not nside (nside = 2**order).\n", - "lon, lat = pix2ang(CHILD_ORDER, cell_ids.astype(np.int64))\n", - "\n", - "print(f\"Lat range: {lat.min():.2f} to {lat.max():.2f}\")\n", - "print(f\"Lon range: {lon.min():.2f} to {lon.max():.2f}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "transformer = Transformer.from_crs(\"EPSG:4326\", \"EPSG:3031\", always_xy=True)\n", - "x, y = transformer.transform(lon, lat)\n", - "\n", - "print(f\"X range: {x.min() / 1e3:.0f} to {x.max() / 1e3:.0f} km\")\n", - "print(f\"Y range: {y.min() / 1e3:.0f} to {y.max() / 1e3:.0f} km\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 3. Define output grid and assign source cells to pixels\n", - "\n", - "Change `GRID_SPACING` below to adjust output resolution (e.g., 1000 for 1 km, 8000 for 8 km)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "GRID_SPACING = 8000 # meters — adjust for coarser/finer grids\n", - "\n", - "# Grid extent: 3400 km covers to ~59S\n", - "EXTENT = 3_400_000\n", - "x_min, x_max = -EXTENT, EXTENT\n", - "y_min, y_max = -EXTENT, EXTENT\n", - "\n", - "# Pixel indices for each source cell\n", - "ix = np.floor((x - x_min) / GRID_SPACING).astype(np.int64)\n", - "iy = np.floor((y - y_min) / GRID_SPACING).astype(np.int64)\n", - "\n", - "nx = int((x_max - x_min) / GRID_SPACING)\n", - "ny = int((y_max - y_min) / GRID_SPACING)\n", - "n_out = nx * ny\n", - "\n", - "# Keep only source cells whose center falls within the grid\n", - "inside = (ix >= 0) & (ix < nx) & (iy >= 0) & (iy < ny)\n", - "\n", - "print(f\"Output grid: {nx} x {ny} = {n_out:,} pixels at {GRID_SPACING / 1e3:.0f} km\")\n", - "print(f\"Source cells inside grid: {inside.sum():,} / {len(inside):,}\")\n", - "print(f\"Source cells outside grid: {(~inside).sum():,}\")\n", - "print(\n", - " f\"Observations captured: {count[inside].sum():.0f} / {count.sum():.0f} \"\n", - " f\"({100 * count[inside].sum() / count.sum():.2f}%)\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Apply the inside mask and compute flat index\n", - "ix = ix[inside]\n", - "iy = iy[inside]\n", - "flat = (iy * nx + ix).astype(np.int64)\n", - "\n", - "count_in = count[inside]\n", - "h_mean_in = h_mean[inside]\n", - "h_sigma_in = h_sigma[inside]\n", - "h_min_in = h_min[inside]\n", - "h_max_in = h_max[inside]\n", - "h_variance_in = h_variance[inside]\n", - "h_q25_in = h_q25[inside]\n", - "h_q50_in = h_q50[inside]\n", - "h_q75_in = h_q75[inside]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 4. Aggregate\n", - "\n", - "- **count**: sum of observation counts\n", - "- **h_mean**: weighted mean by nobs: `sum(n_i * h_i) / sum(n_i)`\n", - "- **h_sigma**: propagated uncertainty: `sqrt(sum(n_i^2 * sigma_i^2)) / sum(n_i)`\n", - "- **h_min**: minimum across source cells\n", - "- **h_max**: maximum across source cells\n", - "- **h_variance, h_q25, h_q50, h_q75**: weighted mean by nobs" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "# --- count: sum ---\n", - "count_out = np.bincount(flat, weights=count_in, minlength=n_out)\n", - "filled = count_out > 0\n", - "\n", - "\n", - "# --- weighted means (weight = nobs) using bincount ---\n", - "def weighted_mean(values, weights=count_in):\n", - " wsum = np.bincount(flat, weights=values * weights, minlength=n_out)\n", - " result = np.full(n_out, np.nan)\n", - " result[filled] = wsum[filled] / count_out[filled]\n", - " return result\n", - "\n", - "\n", - "h_mean_out = weighted_mean(h_mean_in)\n", - "h_variance_out = weighted_mean(h_variance_in)\n", - "h_q25_out = weighted_mean(h_q25_in)\n", - "h_q50_out = weighted_mean(h_q50_in)\n", - "h_q75_out = weighted_mean(h_q75_in)\n", - "\n", - "# --- sigma: error propagation through weighted mean ---\n", - "# sigma_combined = sqrt(sum(n_i^2 * sigma_i^2)) / sum(n_i)\n", - "sigma_sq_wsum = np.bincount(flat, weights=count_in**2 * h_sigma_in**2, minlength=n_out)\n", - "h_sigma_out = np.full(n_out, np.nan)\n", - "h_sigma_out[filled] = np.sqrt(sigma_sq_wsum[filled]) / count_out[filled]\n", - "\n", - "# --- min / max ---\n", - "h_min_out = np.full(n_out, np.inf)\n", - "np.minimum.at(h_min_out, flat, h_min_in)\n", - "h_min_out[~filled] = np.nan\n", - "\n", - "h_max_out = np.full(n_out, -np.inf)\n", - "np.maximum.at(h_max_out, flat, h_max_in)\n", - "h_max_out[~filled] = np.nan\n", - "\n", - "print(f\"Filled pixels: {filled.sum():,} / {n_out:,} ({100 * filled.mean():.1f}%)\")\n", - "print(f\"Mean source cells per filled pixel: {len(flat) / filled.sum():.1f}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build output dataset\n", - "x_centers = np.arange(x_min + GRID_SPACING / 2, x_max, GRID_SPACING)\n", - "y_centers = np.arange(y_min + GRID_SPACING / 2, y_max, GRID_SPACING)\n", - "\n", - "ds_out = xr.Dataset(\n", - " {\n", - " \"count\": ([\"y\", \"x\"], count_out.reshape(ny, nx).astype(np.int64)),\n", - " \"h_mean\": ([\"y\", \"x\"], h_mean_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_sigma\": ([\"y\", \"x\"], h_sigma_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_min\": ([\"y\", \"x\"], h_min_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_max\": ([\"y\", \"x\"], h_max_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_variance\": ([\"y\", \"x\"], h_variance_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_q25\": ([\"y\", \"x\"], h_q25_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_q50\": ([\"y\", \"x\"], h_q50_out.reshape(ny, nx).astype(np.float32)),\n", - " \"h_q75\": ([\"y\", \"x\"], h_q75_out.reshape(ny, nx).astype(np.float32)),\n", - " },\n", - " coords={\"x\": x_centers, \"y\": y_centers},\n", - " attrs={\n", - " \"crs\": \"EPSG:3031\",\n", - " \"grid_spacing_m\": GRID_SPACING,\n", - " \"source\": f\"s3://{S3_BUCKET}/{S3_PREFIX}\",\n", - " \"source_grid\": f\"HEALPix order {CHILD_ORDER} (nside={NSIDE})\",\n", - " \"method\": \"binning: healpix cell center in output pixel, nobs-weighted mean\",\n", - " },\n", - ")\n", - "ds_out[\"x\"].attrs = {\"units\": \"m\", \"standard_name\": \"projection_x_coordinate\"}\n", - "ds_out[\"y\"].attrs = {\"units\": \"m\", \"standard_name\": \"projection_y_coordinate\"}\n", - "print(ds_out)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 5. Visualize" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "proj = ccrs.SouthPolarStereo()\n", - "img_extent = [x_min, x_max, y_min, y_max]\n", - "\n", - "fig, axes = plt.subplots(2, 2, figsize=(18, 16), subplot_kw={\"projection\": proj})\n", - "for ax in axes.flat:\n", - " ax.coastlines(resolution=\"50m\", linewidth=0.5)\n", - " ax.add_feature(cfeature.LAND, facecolor=\"lightgray\", alpha=0.3)\n", - " ax.gridlines(draw_labels=False, alpha=0.3)\n", - " ax.set_extent([-180, 180, -90, -60], crs=ccrs.PlateCarree())\n", - "\n", - "mask = ds_out[\"count\"].values > 0\n", - "\n", - "# 1. Mean elevation\n", - "ax = axes[0, 0]\n", - "d = ds_out[\"h_mean\"].values.copy()\n", - "d[~mask] = np.nan\n", - "im = ax.imshow(\n", - " d,\n", - " origin=\"lower\",\n", - " extent=img_extent,\n", - " transform=proj,\n", - " cmap=\"terrain\",\n", - " vmin=0,\n", - " vmax=4000,\n", - " interpolation=\"nearest\",\n", - ")\n", - "ax.set_title(f\"Mean Elevation ({mask.sum():,} pixels)\", fontsize=14, weight=\"bold\")\n", - "plt.colorbar(im, ax=ax, label=\"Elevation (m)\", shrink=0.7)\n", - "\n", - "# 2. Observation count\n", - "ax = axes[0, 1]\n", - "d = ds_out[\"count\"].values.astype(np.float64)\n", - "d[~mask] = np.nan\n", - "im = ax.imshow(\n", - " d,\n", - " origin=\"lower\",\n", - " extent=img_extent,\n", - " transform=proj,\n", - " cmap=\"viridis\",\n", - " norm=plt.matplotlib.colors.LogNorm(vmin=1),\n", - " interpolation=\"nearest\",\n", - ")\n", - "ax.set_title(\"Observation Count\", fontsize=14, weight=\"bold\")\n", - "plt.colorbar(im, ax=ax, label=\"Count (log scale)\", shrink=0.7)\n", - "\n", - "# 3. Uncertainty\n", - "ax = axes[1, 0]\n", - "d = ds_out[\"h_sigma\"].values.copy()\n", - "d[~mask] = np.nan\n", - "im = ax.imshow(\n", - " d,\n", - " origin=\"lower\",\n", - " extent=img_extent,\n", - " transform=proj,\n", - " cmap=\"plasma\",\n", - " vmax=1.0,\n", - " interpolation=\"nearest\",\n", - ")\n", - "ax.set_title(\"Uncertainty (h_sigma)\", fontsize=14, weight=\"bold\")\n", - "plt.colorbar(im, ax=ax, label=\"Uncertainty (m)\", shrink=0.7)\n", - "\n", - "# 4. Elevation range\n", - "ax = axes[1, 1]\n", - "d = ds_out[\"h_max\"].values - ds_out[\"h_min\"].values\n", - "d[~mask] = np.nan\n", - "im = ax.imshow(\n", - " d,\n", - " origin=\"lower\",\n", - " extent=img_extent,\n", - " transform=proj,\n", - " cmap=\"hot\",\n", - " vmax=100,\n", - " interpolation=\"nearest\",\n", - ")\n", - "ax.set_title(\"Elevation Range (max - min)\", fontsize=14, weight=\"bold\")\n", - "plt.colorbar(im, ax=ax, label=\"Range (m)\", shrink=0.7)\n", - "\n", - "plt.suptitle(\"ATL06 - Cycle 22 (8 km Polar Stereo)\", fontsize=16, weight=\"bold\", y=1.02)\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 6. Regional Zoom - West Antarctica" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "proj = ccrs.SouthPolarStereo()\n", - "fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={\"projection\": proj})\n", - "ax.coastlines(resolution=\"10m\", linewidth=0.5)\n", - "ax.add_feature(cfeature.LAND, facecolor=\"lightgray\", alpha=0.3)\n", - "ax.gridlines(draw_labels=True, alpha=0.3)\n", - "ax.set_extent([-140, -70, -85, -70], crs=ccrs.PlateCarree())\n", - "\n", - "d = ds_out[\"h_mean\"].values.copy()\n", - "d[ds_out[\"count\"].values == 0] = np.nan\n", - "im = ax.imshow(\n", - " d,\n", - " origin=\"lower\",\n", - " extent=[x_min, x_max, y_min, y_max],\n", - " transform=proj,\n", - " cmap=\"terrain\",\n", - " vmin=0, # vmax=.005,\n", - " interpolation=\"nearest\",\n", - ")\n", - "plt.colorbar(im, ax=ax, label=\"Elevation (m)\", shrink=0.7)\n", - "ax.set_title(\"West Antarctica - Mean Elevation (8 km)\", fontsize=14, weight=\"bold\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 7. Save to Zarr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "output_path = \"atl06_polar_stereo_8km.zarr\"\n", - "ds_out.to_zarr(output_path, mode=\"w\")\n", - "print(f\"Saved to {output_path}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 8. Summary" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "filled = ds_out[\"count\"].values > 0\n", - "print(f\"Grid: {nx} x {ny} at {GRID_SPACING / 1e3:.0f} km (EPSG:3031)\")\n", - "print(f\"Filled pixels: {filled.sum():,} / {n_out:,} ({100 * filled.mean():.1f}%)\")\n", - "print(f\"Total observations: {ds_out['count'].values.sum():,}\")\n", - "print(\"\\nElevation (h_mean):\")\n", - "print(f\" Min: {np.nanmin(ds_out['h_mean'].values):.2f} m\")\n", - "print(f\" Max: {np.nanmax(ds_out['h_mean'].values):.2f} m\")\n", - "print(f\" Mean: {np.nanmean(ds_out['h_mean'].values):.2f} m\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/notebooks/sentinel2_fusion.ipynb b/notebooks/sentinel2_fusion.ipynb deleted file mode 100644 index 0fdb5cd32..000000000 --- a/notebooks/sentinel2_fusion.ipynb +++ /dev/null @@ -1,374 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7fb27b941602401d91542211134fc71a", - "metadata": {}, - "source": [ - "# Sentinel-2 ingest and ICESat-2 fusion on a shared HEALPix grid\n", - "\n", - "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/sentinel2_fusion.ipynb)\n", - "\n", - "_Runs end-to-end on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/sentinel2_fusion.ipynb):\n", - "every byte it reads is anonymous -- Sentinel-2 metadata from the\n", - "[Earth Search](https://element84.com/earth-search/) STAC API, imagery from the\n", - "public `e84-earth-search-sentinel-data` COG bucket, and ICESat-2 aggregates\n", - "from a public [source.coop](https://source.coop/englacial/zagg/benchmarks)\n", - "store. No credentials, no AWS account._\n", - "\n", - "This example is the [issue #218](https://github.com/englacial/zagg/issues/218)\n", - "template in action. It ingests Sentinel-2 L2A surface reflectance onto the\n", - "**same HEALPix order-19 cells (~12.4 m)** that zagg's ICESat-2 products\n", - "aggregate to, using **pull-NN sampling**: every cell takes the source pixel\n", - "nearest its center. That one design choice buys three things at once:\n", - "\n", - "* **dense by construction** -- every covered cell gets a value at *any* cell\n", - " order (push-hashing pixel centers instead leaves 2--4% holes, because\n", - " HEALPix cells are equal-area but not equal-shape);\n", - "* **exact DNs** -- the most-central measurement is selected, never averaged,\n", - " so uint16 digital numbers survive bit-for-bit and categorical bands (the\n", - " scene classification, `scl`) stay categorical;\n", - "* **trivial fusion** -- Sentinel-2 and ICESat-2 land in the same cell-ID\n", - " space, so coincident extraction is an integer set intersection. No\n", - " reprojection, no interpolation, no geometry at read time.\n", - "\n", - "The output layout is one **`(time, cells)`** array per band -- time chunks of\n", - "1, so each timestep x chunk is one storage object -- with a real `time`\n", - "coordinate. Timesteps are Sentinel-2 *datatakes* (`s2:datatake_id`), not item\n", - "datetimes: adjacent MGRS tiles of one datatake are a single timestep, and\n", - "cells in the 9.8 km tile overlap take the value from the **nearest tile\n", - "center**.\n" - ] - }, - { - "cell_type": "markdown", - "id": "acae54e37e7d407bbb7b55eff062a284", - "metadata": {}, - "source": [ - "## 1. The fusion target: ATL06 cycle 22, published on source.coop\n", - "\n", - "The join target is a public, anonymous ATL06 (land-ice height) aggregation\n", - "for ICESat-2 **cycle 22** (Jan--Apr 2024) at HEALPix **order 12** (~1.6 km),\n", - "nested ids, fullsphere layout. Our Sentinel-2 ingest below lands at order\n", - "**19** (~12.4 m) in the same nested-id hierarchy, so the cross-resolution\n", - "join is *integer arithmetic*: an order-19 cell's order-12 parent is\n", - "``nested_19 // 4**7`` -- one shift, no geometry.\n", - "\n", - "We work over the Larsen B embayment on the Antarctic Peninsula, in the same\n", - "austral summer the ICESat-2 cycle spans.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9a63283cbaf04dbcab1f6479b197f3a8", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "# Larsen B embayment, Antarctic Peninsula (lon_min, lat_min, lon_max, lat_max).\n", - "BBOX = (-62.45, -65.55, -62.15, -65.35)\n", - "ATL06_STORE = (\n", - " \"s3://us-west-2.opendata.source.coop/englacial/zagg/benchmarks/atl06_cycle22_fullsphere.zarr\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "8dd0d8092fe74a7c96281538738b07e2", - "metadata": {}, - "source": [ - "## 2. Catalog: query Earth Search\n", - "\n", - "`STACSource` speaks to any STAC API root. For Sentinel-2 we query the\n", - "Collection-1 reprocessing (`sentinel-2-c1-l2a`) together with\n", - "`sentinel-2-pre-c1-l2a` (identical schema; fills the archive where the\n", - "reprocessing hasn't reached yet), keep only the assets we'll ingest, and set\n", - "`time_key` so acquisition grouping happens by datatake.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "72eea5119410473aa328ad9291626812", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.catalog.sources import STACQuery, STACSource\n", - "\n", - "src = STACSource(\n", - " \"https://earth-search.aws.element84.com/v1\",\n", - " assets=[\"red\", \"nir\", \"scl\"],\n", - " time_key=\"s2:datatake_id\",\n", - ")\n", - "query = STACQuery(\n", - " collections=[\"sentinel-2-c1-l2a\", \"sentinel-2-pre-c1-l2a\"],\n", - " start_date=\"2023-11-01\",\n", - " end_date=\"2024-03-31\",\n", - " region=BBOX,\n", - " max_cloud_cover=40,\n", - ")\n", - "cat = src.fetch(query)\n", - "for rec in cat.granule_records()[:8]:\n", - " print(rec[\"id\"], rec[\"datetime\"][:19], \"->\", rec[\"time_key\"])" - ] - }, - { - "cell_type": "markdown", - "id": "8edb47106e1a46a883d545849b8ab81b", - "metadata": {}, - "source": [ - "## 3. Shardmap and ingest\n", - "\n", - "The shipped `sentinel2_l2a` config declares the whole pipeline: bands with\n", - "dtypes and CF `scale_factor`/`add_offset` attrs (reflectance = DN x 1e-4 -\n", - "0.1 -- recorded, never applied), nodata, and the parent-11/child-19 HEALPix\n", - "grid. `ShardMap.build` intersects the MGRS tile footprints with the grid\n", - "exactly as it does ICESat-2 swaths; `agg` then runs one worker per shard:\n", - "open each band's COG header, fetch only the ~3 km window of 1024-px DEFLATE\n", - "tiles the shard touches (async-tiff -- no GDAL anywhere), gather the nearest\n", - "pixel per cell, and write `(time, cells)` slabs.\n", - "\n", - "To keep Binder light we ingest the first two datatakes over the AOI.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "10185d26023b46108eb7d9f57d49d2b3", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.catalog.shardmap import ShardMap\n", - "from zagg.config import default_config, validate_config\n", - "from zagg.grids import from_config\n", - "\n", - "config = default_config(\"sentinel2_l2a\")\n", - "config.data_source[\"bands\"] = {\n", - " k: v for k, v in config.data_source[\"bands\"].items() if k in (\"red\", \"nir\", \"scl\")\n", - "}\n", - "config.output[\"store\"] = \"sentinel2_larsenb.zarr\"\n", - "validate_config(config)\n", - "\n", - "grid = from_config(config)\n", - "sm = ShardMap.build(cat, grid) # coverage from the catalog bbox\n", - "\n", - "# Keep the first two datatakes (chronological) so the Binder run stays small.\n", - "keep = sorted({e[\"time_key\"] for g in sm.granules for e in g})[:2]\n", - "sm.granules = [[e for e in g if e[\"time_key\"] in keep] for g in sm.granules]\n", - "sm.to_json(\"shardmap_s2_larsenb.json\")\n", - "print(f\"{len(sm.shard_keys)} shards, datatakes: {keep}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8763a12b2bbd4a93a75aff182afb95dc", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.runner import agg\n", - "\n", - "summary = agg(\n", - " config,\n", - " catalog=\"shardmap_s2_larsenb.json\",\n", - " backend=\"local\",\n", - " max_workers=4,\n", - " overwrite=True,\n", - ")\n", - "summary" - ] - }, - { - "cell_type": "markdown", - "id": "7623eae2785240b9bd12b16a66d81610", - "metadata": {}, - "source": [ - "## 4. Read the `(time, cells)` product\n", - "\n", - "The store is plain Zarr v3: one 2-D array per band plus `time` (microseconds\n", - "since epoch) and `cell_ids` (nested order-19 ids). The arrays logically span\n", - "the full sphere; we slice the shard ranges we wrote. Reflectance decode is\n", - "the CF attrs applied by hand -- the stored values are exact source DNs.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7cdc8c89c7104fffa095e18ddfef8986", - "metadata": {}, - "outputs": [], - "source": [ - "import xarray as xr\n", - "\n", - "from zagg import open_store\n", - "\n", - "s2 = xr.open_dataset(\n", - " open_store(config.output[\"store\"], read_only=True),\n", - " engine=\"zarr\",\n", - " consolidated=False,\n", - " zarr_format=3,\n", - " group=grid.group_path,\n", - ")\n", - "\n", - "shard_ranges = [\n", - " (int(grid.block_index(int(s))[0]) * grid.cells_per_shard, int(s)) for s in sm.shard_keys\n", - "]\n", - "\n", - "\n", - "def gather(name, t=None):\n", - " \"\"\"Concatenate the written shard ranges of one array (one timestep).\"\"\"\n", - " parts = []\n", - " for start, _s in shard_ranges:\n", - " sl = slice(start, start + grid.cells_per_shard)\n", - " parts.append((s2[name][t, sl] if t is not None else s2[name][sl]).values)\n", - " return np.concatenate(parts)\n", - "\n", - "\n", - "times = s2[\"time\"].values.astype(\"datetime64[us]\")\n", - "s2_ids = gather(\"cell_ids\")\n", - "red0 = gather(\"red\", 0)\n", - "nir0 = gather(\"nir\", 0)\n", - "scl0 = gather(\"scl\", 0)\n", - "print(\"timesteps:\", list(times))\n", - "print(f\"cells written: {s2_ids.size:,}; valid at t0: {(red0 > 0).sum():,}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b118ea5561624da68c537baed56e602f", - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "from mortie import mort2geo\n", - "\n", - "cells_morton = np.concatenate([grid.children(s) for _start, s in shard_ranges])\n", - "clats, clons = mort2geo(cells_morton)\n", - "ok = red0 > 0 # DN 0 is fill/nodata\n", - "refl = red0[ok].astype(float) * 1e-4 - 0.1\n", - "\n", - "fig, ax = plt.subplots(figsize=(8, 6))\n", - "pc = ax.scatter(clons[ok], clats[ok], c=refl, s=1.5, cmap=\"RdYlGn_r\", vmin=0, vmax=0.35)\n", - "fig.colorbar(pc, label=\"red surface reflectance\")\n", - "ax.set_title(f\"Sentinel-2 red on HEALPix order 19 -- {times[0]}\")\n", - "ax.set_xlabel(\"lon\")\n", - "ax.set_ylabel(\"lat\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "938c804e27f84196a10c8828c723f798", - "metadata": {}, - "source": [ - "## 5. Fusion: the cross-resolution join is a bit shift\n", - "\n", - "Both stores carry **nested HEALPix ids**. The order-19 Sentinel-2 cells roll\n", - "up to their order-12 ATL06 parents by integer division (``4**(19-12)``\n", - "children per parent), and the fullsphere layout means an order-12 array\n", - "index *is* the nested id -- so \"read the ICESat-2 cells under my imagery\"\n", - "is a fancy-indexed Zarr read, nothing more.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "504fb2a444614c0babb325280ed9130a", - "metadata": {}, - "outputs": [], - "source": [ - "from mortie import clip2order, mort2healpix\n", - "\n", - "# Order-12 parents of every ingested S2 cell (morton -> nested id).\n", - "parents_morton = np.unique(clip2order(12, cells_morton))\n", - "parents_nested, _order = mort2healpix(parents_morton)\n", - "parents_nested = np.sort(parents_nested.astype(np.int64))\n", - "\n", - "atl06 = xr.open_dataset(\n", - " open_store(ATL06_STORE, read_only=True, skip_signature=True),\n", - " engine=\"zarr\",\n", - " consolidated=False,\n", - " zarr_format=3,\n", - " group=\"12\",\n", - ")\n", - "h_mean = atl06[\"h_mean\"][parents_nested].values\n", - "is2_count = atl06[\"count\"][parents_nested].values\n", - "measured = is2_count > 0\n", - "print(f\"order-12 parents under the imagery: {parents_nested.size:,}\")\n", - "print(f\"...with ICESat-2 cycle-22 heights: {measured.sum():,}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "59bbdb311c014d738909a11f9e486628", - "metadata": {}, - "outputs": [], - "source": [ - "# Roll the imagery up to the same parents (analysis-time mean of the exact\n", - "# per-cell DNs) and compare against the coincident ATL06 heights.\n", - "parent_of_cell, _ = mort2healpix(clip2order(12, cells_morton))\n", - "pos = np.searchsorted(parents_nested, parent_of_cell.astype(np.int64))\n", - "\n", - "valid = red0 > 0\n", - "refl_all = red0.astype(float) * 1e-4 - 0.1\n", - "sums = np.bincount(pos[valid], weights=refl_all[valid], minlength=parents_nested.size)\n", - "ns = np.bincount(pos[valid], minlength=parents_nested.size)\n", - "mean_refl = np.divide(sums, ns, out=np.full_like(sums, np.nan), where=ns > 0)\n", - "\n", - "both = measured & (ns > 0)\n", - "print(f\"parents with BOTH imagery and altimetry: {both.sum():,}\")\n", - "\n", - "fig, ax = plt.subplots(figsize=(7, 5))\n", - "ax.scatter(mean_refl[both], h_mean[both], s=8, alpha=0.6)\n", - "ax.set_xlabel(f\"Sentinel-2 red reflectance (order-12 mean) -- {times[0]}\")\n", - "ax.set_ylabel(\"ATL06 h_mean (m), cycle 22\")\n", - "ax.set_title(\"Coincident cells, joined by integer id -- zero resampling\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "b43b363d81ae4b689946ece5c682cd59", - "metadata": {}, - "source": [ - "## 6. Where this goes\n", - "\n", - "* **Training loaders**: the band arrays are chunked `(1, 65536)` -- one\n", - " storage object per (timestep, shard). A PyTorch `DataLoader` can range-read\n", - " chunk `k` from every aligned array (`red`, `nir`, `scl`, the ICESat-2\n", - " aggregates) and get co-registered tensors with no geospatial code in the\n", - " loop: `np.asarray` of a Zarr block, straight to `torch.from_numpy`.\n", - " The scale/offset decode is one fused multiply-add on the GPU.\n", - "* **More timesteps**: appends *will* land as the standard Zarr\n", - " resize-then-write-slab pattern (a runner-owned resize; chunk keys are\n", - " `t/chunk`, so nothing already written is rewritten) -- tracked on\n", - " [issue #218](https://github.com/englacial/zagg/issues/218). Today a\n", - " changed time range is a full re-run with `overwrite=True`.\n", - "* **Same-order joins**: against an order-19 product (e.g. the ATL03\n", - " photon t-digests), the join needs no roll-up at all -- it is\n", - " `np.intersect1d` on `cell_ids`.\n", - "* **Full-scale runs**: the same config fans out one shard per AWS Lambda --\n", - " see the [deployment docs](../docs/deployment/lambda.md). The order-19 vs\n", - " order-20 storage trade and the pull-NN design measurements live on\n", - " [issue #218](https://github.com/englacial/zagg/issues/218).\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/shardmap_viewer.ipynb b/notebooks/shardmap_viewer.ipynb deleted file mode 100644 index 939f500b3..000000000 --- a/notebooks/shardmap_viewer.ipynb +++ /dev/null @@ -1,324 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "cell-0", - "metadata": {}, - "source": "# Shard-map viewer\n\n[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/shardmap_viewer.ipynb)\n\n_Runs end-to-end on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/shardmap_viewer.ipynb): it queries real ICESat-2 ATL06 granule **metadata** anonymously from NASA CMR-STAC (no Earthdata Login -- the CMR-STAC search is anonymous) and renders with `ipyleaflet`. The Binder image already provides `zagg[analysis,catalog,viz]` (incl. `stac-geoparquet` and `ipyleaflet`) via the repo's `.binder/` environment._\n\nThis notebook demonstrates the `zagg.viz` shard-map viewer (issue\n[#38](https://github.com/englacial/zagg/issues/38)) with **real ICESat-2\nATL06 granules** fetched anonymously from NASA CMR-STAC.\n\nThe viewer has two layers:\n\n- **Headless render core** (`zagg.viz.render_shardmap` and friends) — pure\n Python, no browser/widget stack. Turns a `ShardMap` (plus an optional\n `Catalog`) into WGS84 GeoJSON `FeatureCollection` dicts: shard outlines\n and granule footprints.\n- **ipyleaflet wrapper** (`zagg.viz.show_shardmap`) — builds an interactive\n context basemap + shard layer + toggleable granule-footprint layer.\n\n**Polar-aware projection.** `show_shardmap` picks the display CRS from the\nmap's extent: a polar AOI renders in NASA polar-stereographic (EPSG:3031\nAntarctic / EPSG:3413 Arctic) with a **GIBS** basemap — undistorted at the\npole — while mid-latitude AOIs stay on Web Mercator + OpenStreetMap.\n\n## Install\n\nOn Binder this is already set up by `.binder/postBuild`. Locally:\n\n```bash\npip install \"zagg[analysis,catalog,viz]\" # core + stac-geoparquet + ipyleaflet widget\n```\n\nThe `viz` extra pulls in `ipyleaflet` (the interactive map); `catalog` pulls\nin `stac-geoparquet` (the CMR-STAC catalog). **Requires an internet\nconnection** for the anonymous CMR-STAC query in section 1 — no credentials\nare needed.\n\n## Areas of interest\n\nThe notebook queries two AOIs:\n\n1. **Antarctic Peninsula** (`[-65, -70, -55, -64]` WGS84 lon/lat) — January\n 2020. Dense ICESat-2 coverage near the pole; O(10–40) granules in two\n weeks. Renders in EPSG:3031 (Antarctic Polar Stereographic).\n2. **Jakobshavn Glacier, West Greenland** (`[-52, 68, -45, 72]`) — June\n 2020. Second example showing the same pipeline on an Arctic AOI\n (EPSG:3413).\n\nBoth inputs (ShardMap JSON and STAC-geoparquet Catalog) are supported from\nday one — see section 3 for the round-trip to disk and reload." - }, - { - "cell_type": "markdown", - "id": "cell-1", - "metadata": {}, - "source": [ - "## 1. Antarctic Peninsula — fetch real ATL06 granules from NASA CMR-STAC\n", - "\n", - "`CMRSource` speaks directly to NASA's CMR-STAC endpoint (`requests`).\n", - "No Earthdata Login or credentials are needed for anonymous granule-metadata\n", - "queries. The returned `Catalog` wraps a stac-geoparquet Arrow table (one row\n", - "per granule, both S3 and HTTPS asset hrefs preserved)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-2", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.catalog.sources import CMRSource, Query\n", - "\n", - "# Antarctic Peninsula AOI — two weeks in January 2020.\n", - "AOI_BBOX = (-65.0, -70.0, -55.0, -64.0) # lon_min, lat_min, lon_max, lat_max\n", - "START_DATE = \"2020-01-01\"\n", - "END_DATE = \"2020-01-15\"\n", - "\n", - "query = Query(\n", - " short_name=\"ATL06\",\n", - " version=\"007\",\n", - " start_date=START_DATE,\n", - " end_date=END_DATE,\n", - " region=AOI_BBOX,\n", - " provider=\"NSIDC_CPRD\",\n", - ")\n", - "\n", - "catalog = CMRSource().fetch(query)\n", - "print(f\"Fetched {len(catalog)} granules for {query.collection} over {AOI_BBOX}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-3", - "metadata": {}, - "outputs": [], - "source": [ - "# Inspect the catalog schema — each row is a STAC Item with WKB geometry.\n", - "print(\"Schema columns:\", catalog.table.schema.names[:10], \"...\")\n", - "print(\"Total rows:\", catalog.table.num_rows)\n", - "\n", - "# Decode a few granule records to confirm footprints are present.\n", - "records = catalog.granule_records()\n", - "print(f\"\\n{len(records)} records with valid footprint geometry.\")\n", - "for rec in records[:3]:\n", - " print(f\" {rec['id']} | https: {(rec['https'] or 'None')[:70]}...\")" - ] - }, - { - "cell_type": "markdown", - "id": "cell-4", - "metadata": {}, - "source": [ - "## 2. Build a ShardMap on a HEALPix grid\n", - "\n", - "We use a HEALPix grid (`parent_order=6, child_order=12`) — the same\n", - "configuration as `src/zagg/configs/atl06.yaml`. The `mortie` backend is\n", - "used for HEALPix intersection and requires no extra install. If `spherely`\n", - "is installed, pass `backend='auto'` to prefer exact S2 intersections.\n", - "\n", - "`ShardMap.build` maps each grid shard (a parent-order HEALPix cell) to the\n", - "set of granules whose footprint intersects it. The manifest is self-contained\n", - "— it stores `{\"id\", \"s3\", \"https\"}` per granule, so the aggregation runner\n", - "never needs the Catalog again at run time." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-5", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.catalog.shardmap import ShardMap\n", - "from zagg.grids import HealpixGrid\n", - "\n", - "# HEALPix grid matching atl06.yaml (parent_order=6, child_order=12, fullsphere).\n", - "grid = HealpixGrid(parent_order=6, child_order=12, layout=\"fullsphere\")\n", - "\n", - "# Build: catalog footprints are intersected with the HEALPix shard cells\n", - "# that cover the AOI. Use backend='auto' to prefer spherely when available.\n", - "shardmap = ShardMap.build(catalog, grid, backend=\"mortie\")\n", - "\n", - "print(\n", - " f\"ShardMap: {len(shardmap.shard_keys)} shards, \"\n", - " f\"{shardmap.metadata['total_pairs']} granule-shard pairs\"\n", - ")\n", - "print(\n", - " f\"Build time: {shardmap.metadata['build_wall_s']:.2f}s backend: {shardmap.metadata['backend']}\"\n", - ")\n", - "print(f\"Grid signature: {shardmap.grid_signature}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-6", - "metadata": {}, - "outputs": [], - "source": [ - "# Inspect a few shard assignments.\n", - "for key, shard_granules in zip(shardmap.shard_keys[:4], shardmap.granules[:4]):\n", - " ids = [g[\"id\"] for g in shard_granules]\n", - " print(f\" shard {key:6d}: {len(ids)} granule(s) — {ids[:2]}{'...' if len(ids) > 2 else ''}\")" - ] - }, - { - "cell_type": "markdown", - "id": "cell-7", - "metadata": {}, - "source": [ - "## 3. Persist to disk and reload (round-trip)\n", - "\n", - "Both the ShardMap JSON and the STAC-geoparquet Catalog are supported as\n", - "file-path inputs to `show_shardmap`. Here we round-trip both to disk to\n", - "demonstrate the saved-file path you would use in practice (e.g. after\n", - "running `python -m zagg.catalog --config atl06.yaml ...`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-8", - "metadata": {}, - "outputs": [], - "source": [ - "import tempfile\n", - "from pathlib import Path\n", - "\n", - "tmp = Path(tempfile.mkdtemp())\n", - "sm_path = tmp / \"shardmap_ATL06_peninsula_jan2020.json\"\n", - "cat_path = tmp / \"catalog_ATL06_peninsula_jan2020.parquet\"\n", - "\n", - "shardmap.to_json(str(sm_path))\n", - "catalog.to_geoparquet(str(cat_path))\n", - "\n", - "print(f\"ShardMap -> {sm_path} ({sm_path.stat().st_size / 1024:.1f} KB)\")\n", - "print(f\"Catalog -> {cat_path} ({cat_path.stat().st_size / 1024:.1f} KB)\")\n", - "\n", - "# Reload to verify round-trip.\n", - "from zagg.catalog.sources import Catalog\n", - "\n", - "sm_rt = ShardMap.from_json(str(sm_path))\n", - "cat_rt = Catalog.from_geoparquet(str(cat_path))\n", - "print(f\"\\nRound-trip OK: {len(sm_rt.shard_keys)} shards, {len(cat_rt)} granules\")" - ] - }, - { - "cell_type": "markdown", - "id": "cell-9", - "metadata": {}, - "source": [ - "## 4. Headless render core: GeoJSON FeatureCollections\n", - "\n", - "`render_shardmap` assembles every layer into one dict of GeoJSON\n", - "`FeatureCollection`s. Pass a `catalog` to include granule footprints. Each\n", - "value is plain, JSON-serializable GeoJSON — no widgets required." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-10", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.viz import render_shardmap\n", - "\n", - "layers = render_shardmap(shardmap, catalog)\n", - "{name: (fc and len(fc[\"features\"])) for name, fc in layers.items()}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-11", - "metadata": {}, - "outputs": [], - "source": [ - "# One feature per shard, with the shard key and its granule count.\n", - "shards_fc = layers[\"shards\"]\n", - "feat = shards_fc[\"features\"][0]\n", - "print(\"geometry type:\", feat[\"geometry\"][\"type\"])\n", - "print(\"properties:\", feat[\"properties\"])\n", - "\n", - "# All populated shards in the Antarctic Peninsula window.\n", - "[f[\"properties\"] for f in shards_fc[\"features\"]][:5]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-12", - "metadata": {}, - "outputs": [], - "source": [ - "# One feature per granule footprint, decoded from the real Catalog.\n", - "print(f\"{len(layers['granules']['features'])} granule footprint(s)\")\n", - "[f[\"properties\"][\"id\"] for f in layers[\"granules\"][\"features\"]][:5]" - ] - }, - { - "cell_type": "markdown", - "id": "cell-15", - "metadata": {}, - "source": "## 5. Interactive map — Antarctic Peninsula (ipyleaflet, polar-stereographic)\n\n`show_shardmap` builds an `ipyleaflet.Map` from the saved ShardMap and\nCatalog paths. Both ShardMap JSON and STAC-geoparquet Catalog are accepted\nas file paths or in-memory objects.\n\n**Projection is auto-selected from the map's extent.** This Antarctic AOI is\nentirely south of −60°, so the viewer renders in **EPSG:3031** (Antarctic\nPolar Stereographic) with a NASA **GIBS** polar basemap instead of the\ndistorted Web Mercator default. An Arctic AOI (poleward of +60°) gets\n**EPSG:3413**; mid-latitude AOIs stay on Web Mercator + OpenStreetMap. Pass\n`crs=\"EPSG:3857\"` to force Mercator, or `crs=\"EPSG:3413\"`/`\"EPSG:3031\"` to\nforce a specific pole. Vector layers stay WGS84 GeoJSON — proj4leaflet\nreprojects them client-side.\n\n**Run in JupyterLab** (Binder already has the deps via `.binder/postBuild`;\nlocally `pip install \"zagg[analysis,catalog,viz]\"`) to see the live map.\nUnder headless `nbconvert` the Map object is constructed (no error) but\ntiles won't display.\n\n### Verification checklist\n\n1. **Basemap** — the GIBS polar basemap renders, pans, zooms, and the\n continent is undistorted at the pole (no Web-Mercator stretching).\n2. **Shard outlines** — blue polygons over the Peninsula; click one for\n its `shard_key` and `n_granules` properties.\n3. **Granule footprints toggle** — layer-control (top-right); the ICESat-2\n track footprints appear/disappear." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-16", - "metadata": {}, - "outputs": [], - "source": [ - "from zagg.viz import show_shardmap\n", - "\n", - "# File-path interface — same as after `python -m zagg.catalog ...`.\n", - "# CRS auto-selected from the AOI extent: this Antarctic map renders in EPSG:3031.\n", - "m = show_shardmap(str(sm_path), catalog=str(cat_path), zoom=5)\n", - "print(\"display CRS:\", m.crs[\"name\"])\n", - "print(\"layers:\", [getattr(layer, \"name\", type(layer).__name__) for layer in m.layers])\n", - "m" - ] - }, - { - "cell_type": "markdown", - "id": "cell-17", - "metadata": {}, - "source": [ - "## 6. Second AOI — Jakobshavn Glacier, West Greenland\n", - "\n", - "A second example using an Arctic AOI to show the same pipeline working\n", - "independently on a different region. Jakobshavn Glacier is one of\n", - "Greenland's fastest-moving glaciers and a standard ICESat-2 study region." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-18", - "metadata": {}, - "outputs": [], - "source": [ - "# render_shardmap / show_shardmap re-imported here for a standalone re-run of\n", - "# this section; CMRSource / Query / ShardMap / grid still come from sections 1-2.\n", - "from zagg.viz import render_shardmap, show_shardmap\n", - "\n", - "# Second AOI: Jakobshavn Glacier, West Greenland — June 2020.\n", - "AOI2_BBOX = (-52.0, 68.0, -45.0, 72.0)\n", - "START2, END2 = \"2020-06-01\", \"2020-06-15\"\n", - "\n", - "query2 = Query(\n", - " short_name=\"ATL06\",\n", - " version=\"007\",\n", - " start_date=START2,\n", - " end_date=END2,\n", - " region=AOI2_BBOX,\n", - " provider=\"NSIDC_CPRD\",\n", - ")\n", - "\n", - "catalog2 = CMRSource().fetch(query2)\n", - "shardmap2 = ShardMap.build(catalog2, grid, backend=\"mortie\")\n", - "\n", - "print(f\"Greenland AOI: {len(catalog2)} granules, {len(shardmap2.shard_keys)} shards\")\n", - "\n", - "# Headless render — confirm layers are populated.\n", - "layers2 = render_shardmap(shardmap2, catalog2)\n", - "print(\"Layer feature counts:\", {name: (fc and len(fc[\"features\"])) for name, fc in layers2.items()})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-19", - "metadata": {}, - "outputs": [], - "source": [ - "# Save and display interactive map for the Greenland AOI.\n", - "tmp2 = Path(tempfile.mkdtemp())\n", - "sm2_path = tmp2 / \"shardmap_ATL06_greenland_jun2020.json\"\n", - "cat2_path = tmp2 / \"catalog_ATL06_greenland_jun2020.parquet\"\n", - "\n", - "shardmap2.to_json(str(sm2_path))\n", - "catalog2.to_geoparquet(str(cat2_path))\n", - "\n", - "# Arctic AOI -> auto-selects EPSG:3413 (NSIDC Sea Ice Polar Stereographic North).\n", - "m2 = show_shardmap(str(sm2_path), catalog=str(cat2_path), zoom=5)\n", - "print(\"display CRS:\", m2.crs[\"name\"])\n", - "m2" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/notebooks/tdigest_reader_example.ipynb b/notebooks/tdigest_reader_example.ipynb deleted file mode 100644 index 386496ca9..000000000 --- a/notebooks/tdigest_reader_example.ipynb +++ /dev/null @@ -1,563 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "ac52de6c", - "metadata": {}, - "source": [ - "# Reading t-digest products: tensors and raw values\n", - "\n", - "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/tdigest_reader_example.ipynb)\n", - "\n", - "_Runs end-to-end on [Binder](https://mybinder.org/v2/gh/englacial/zagg/main?urlpath=lab/tree/notebooks/tdigest_reader_example.ipynb): it builds a small synthetic t-digest store in memory -- no cloud data, no credentials. The Binder image already provides `zagg[analysis]` (incl. matplotlib) via the repo's `.binder/` environment._\n", - "\n", - "zagg writes a gridded Zarr product where each coverage chunk is a `64 x 64`\n", - "block of child cells, and each populated cell carries a **t-digest** sketch of\n", - "its observed values (e.g. ICESat-2 photon elevations). The sketch is stored in a\n", - "ragged **variable-length-bytes** field (issue #209): one vlen array on the cell\n", - "grid whose populated cells each hold the raw little-endian bytes of their\n", - "`(k_centroids, 2)` `(mean, weight)` digest. The layout is self-describing -- see\n", - "the [ragged layout doc](https://github.com/englacial/zagg/blob/main/docs/ragged_layout.md).\n", - "\n", - "A downstream client usually wants something denser and fixed-size. The\n", - "`zagg.readers` package (issue #79) reconstructs that from the stored digests:\n", - "\n", - "- **`read_tensors`** -- a generator yielding the reader contract\n", - " `(tensor, mask, (offset, gain), morton_index)` per block (one read chunk by\n", - " default), where `tensor` has shape `(64, 64, n_bins)`: each cell's digest\n", - " rasterized into `n_bins` evenly-spaced z-bins, placed at the **spatially\n", - " faithful** deinterleave of its nested rank (issue #336; mortie spec §8), with\n", - " a uint8 occupancy `mask` and the block's shared z-window `(z_lo, resolution)`.\n", - "- **`read_raw_values`** -- the lossless companion: when a cell's digest was built\n", - " with no merges, its centroid means *are* the original samples, so the raw value\n", - " vector is recovered exactly (and it raises when that is not possible).\n", - "- **`read_cell`** -- random access to one cell's digest (2 ranged GETs on a\n", - " sharded product, never the whole shard).\n", - "\n", - "This notebook is **self-contained**: it builds a small synthetic t-digest store\n", - "in memory, so it runs anywhere (including Binder) with no external data or\n", - "credentials. The store is built with the same public API zagg uses to write the\n", - "real product (`zagg.stats.tdigest.build_tdigest` +\n", - "`zagg.processing.write_ragged_to_zarr`), so the read path exercised here is\n", - "identical to the one a real product uses." - ] - }, - { - "cell_type": "markdown", - "id": "caf66144", - "metadata": {}, - "source": "## 0. Requirements\n\nThe reader API (`read_tensors` / `read_raw_values`) needs only zagg's core\ndependencies (numpy + zarr). The one plot below uses `matplotlib`, which ships\nin zagg's `analysis` extra.\n\n- **On Binder** nothing to do -- the repo's `.binder/` environment already\n installs `zagg[analysis]` (matplotlib included), so every cell runs as-is.\n- **Standalone** (your own laptop/hub), install once with\n `pip install 'zagg[analysis]'` before running the notebook." - }, - { - "cell_type": "markdown", - "id": "521699fb", - "metadata": {}, - "source": [ - "## 1. Build a synthetic t-digest store\n", - "\n", - "We mimic two coverage chunks at real order-6 morton ids, each with a handful of\n", - "populated cells. For every populated cell we draw some samples and sketch them\n", - "with `build_tdigest`, then write the chunk's per-cell digests into the product's\n", - "ragged vlen array with `write_ragged_to_zarr` -- exactly the layout the readers\n", - "consume. Each chunk's coverage-cell morton id is recovered by the readers from\n", - "the sibling `morton` coordinate array (not from a subgroup name).\n", - "\n", - "A real product would instead point the readers at an on-disk or S3-backed Zarr\n", - "store of the same shape; nothing else about the read code changes." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "abe4f114", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:16.958772Z", - "iopub.status.busy": "2026-06-25T19:05:16.958527Z", - "iopub.status.idle": "2026-06-25T19:05:17.625190Z", - "shell.execute_reply": "2026-06-25T19:05:17.623535Z" - } - }, - "outputs": [], - "source": [ - "import numpy as np\n", - "import zarr\n", - "from zarr.storage import MemoryStore\n", - "\n", - "from zagg.config import PipelineConfig\n", - "from zagg.grids import HealpixGrid\n", - "from zagg.grids.morton import morton_word\n", - "from zagg.processing import write_ragged_to_zarr\n", - "from zagg.stats.tdigest import build_tdigest\n", - "\n", - "# One ragged t-digest field on a HEALPix grid: order-6 shards (parent), order-12\n", - "# cells (child) -> each 4096-cell coverage chunk is a 64x64 child block.\n", - "config = PipelineConfig(\n", - " data_source={\"groups\": [\"g\"]},\n", - " aggregation={\n", - " \"coordinates\": {\"morton\": {\"dtype\": \"uint64\", \"fill_value\": 0}},\n", - " \"variables\": {\n", - " \"h_tdigest\": {\n", - " \"function\": \"zagg.stats.tdigest.build_tdigest\",\n", - " \"source\": \"h\",\n", - " \"kind\": \"ragged\",\n", - " \"inner_shape\": [2], # each centroid is a (mean, weight) pair\n", - " \"dtype\": \"float32\",\n", - " \"fill_value\": 0,\n", - " }\n", - " },\n", - " },\n", - " output={\"grid\": {\"type\": \"healpix\", \"parent_order\": 6, \"child_order\": 12}},\n", - ")\n", - "grid = HealpixGrid(6, 12, layout=\"fullsphere\", config=config)\n", - "FIELD = f\"{grid.group_path}/h_tdigest\" # array path the readers consume, e.g. \"12/h_tdigest\"\n", - "rng = np.random.default_rng(0)\n", - "\n", - "\n", - "def write_chunk(store, morton_key, cell_to_values, *, delta=512):\n", - " \"\"\"Sketch each cell's samples and write the chunk into the ragged vlen array.\n", - "\n", - " Returns ``(morton_word, cell_base)`` -- the chunk's packed coverage-cell id\n", - " and the global offset of its first cell (so single cells can be addressed by\n", - " ``read_cell`` below).\n", - " \"\"\"\n", - " word = morton_word(morton_key)\n", - " block = grid.block_index(word)\n", - " base = block[0] * grid.cells_per_chunk\n", - " # The dense per-cell morton coordinate the readers use for chunk identity.\n", - " morton_arr = zarr.open_array(store, path=f\"{grid.group_path}/morton\", mode=\"r+\")\n", - " morton_arr[base : base + grid.cells_per_chunk] = grid.children(word)\n", - " cell_ids = sorted(cell_to_values)\n", - " digests = [build_tdigest(np.asarray(cell_to_values[c]), delta=delta) for c in cell_ids]\n", - " write_ragged_to_zarr({\"h_tdigest\": (digests, cell_ids)}, store, grid=grid, chunk_idx=block)\n", - " return word, base\n", - "\n", - "\n", - "store = MemoryStore()\n", - "grid.emit_template(store)\n", - "\n", - "# Chunk A: elevations clustered around ~20 m. cell_id is the cell's position in\n", - "# the chunk's row-major (64x64) children block (so 0 -> (0,0), 4095 -> (63,63)).\n", - "word_a, base_a = write_chunk(\n", - " store,\n", - " \"1121121\",\n", - " {\n", - " 0: rng.normal(20.0, 2.0, 3_000),\n", - " 5: rng.normal(22.0, 1.5, 2_000),\n", - " 4095: rng.normal(19.0, 2.5, 1_500),\n", - " },\n", - ")\n", - "\n", - "# Chunk B: a different elevation band, ~50 m.\n", - "word_b, base_b = write_chunk(\n", - " store,\n", - " \"2431123\",\n", - " {\n", - " 7: rng.normal(50.0, 3.0, 2_500),\n", - " 63: rng.normal(52.0, 2.0, 2_000),\n", - " },\n", - ")\n", - "\n", - "print(\"chunk morton ids written:\", sorted((word_a, word_b)))" - ] - }, - { - "cell_type": "markdown", - "id": "0c6d98a4", - "metadata": {}, - "source": [ - "The field is now ONE `variable_length_bytes` array on the cell grid: each\n", - "populated cell holds the raw little-endian bytes of its `(k_centroids, 2)`\n", - "`(mean, weight)` digest, and empty cells keep the `b\"\"` fill. The element\n", - "interpretation (dtype + shape) rides in the array's `ragged` attrs, so a reader\n", - "decodes exactly what the writer declared. We can peek at one cell's digest with\n", - "`read_cell`, which random-accesses a single cell." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ee99c084", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:17.627747Z", - "iopub.status.busy": "2026-06-25T19:05:17.627472Z", - "iopub.status.idle": "2026-06-25T19:05:17.637744Z", - "shell.execute_reply": "2026-06-25T19:05:17.636536Z" - } - }, - "outputs": [], - "source": [ - "from zagg.readers import read_cell\n", - "\n", - "# read_cell addresses ONE cell by its global position (chunk base + cell id).\n", - "digest = read_cell(store, FIELD, base_a + 5) # cell 5 of chunk A\n", - "print(f\"cell 5 of chunk A: digest shape {digest.shape} (k_centroids, [mean, weight]):\")\n", - "print(digest[:5])" - ] - }, - { - "cell_type": "markdown", - "id": "f1047ad2", - "metadata": {}, - "source": [ - "## 2. `read_tensors`: dense fixed-size tensors\n", - "\n", - "`read_tensors` yields `(tensor, mask, (offset, gain), morton_index)` per block\n", - "(one read chunk by default; `block_order=` assembles whole chunks into larger\n", - "blocks). For each block it:\n", - "\n", - "1. trims each cell's tails to the `bottom`/`top` quantiles (default 5%/95%),\n", - "2. anchors a fixed `n_bins * resolution` z-window at the floor of the trimmed\n", - " range -- the emitted `(offset, gain)` pair, shared by every cell in the block,\n", - "3. rasterizes every cell's digest into `n_bins` per-bin counts over that window,\n", - " placing cell rank `r` at `rank_to_rowcol(r, depth)` -- the bit deinterleave\n", - " pinned to mortie spec §8 (row = y, col = x, gridlook's texture convention),\n", - "4. emits the `(64, 64, n_bins)` tensor, the `(64, 64)` uint8 occupancy `mask`\n", - " (0 unobserved / 1 observed-no-digest / 2 observed-with-data; a flat store\n", - " like this one degrades to `{0, 2}`), and the block's morton id (recovered\n", - " from the store).\n", - "\n", - "Defaults: `n_bins=128`, `resolution=0.5` m (a 64 m window), `dtype=\"uint32\"`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b7c17ee2", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:17.640277Z", - "iopub.status.busy": "2026-06-25T19:05:17.640090Z", - "iopub.status.idle": "2026-06-25T19:05:17.660031Z", - "shell.execute_reply": "2026-06-25T19:05:17.658853Z" - } - }, - "outputs": [], - "source": [ - "from zagg.readers import read_tensors\n", - "\n", - "tensors = {morton: tensor for tensor, _mask, _scale, morton in read_tensors(store, FIELD)}\n", - "\n", - "for morton, tensor in sorted(tensors.items()):\n", - " nonzero_cells = int((tensor.sum(axis=2) > 0).sum())\n", - " print(\n", - " f\"chunk morton={morton}: tensor {tensor.shape} dtype={tensor.dtype}, \"\n", - " f\"{nonzero_cells} populated cells, total counts={int(tensor.sum())}\"\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "79a77921", - "metadata": {}, - "source": [ - "The morton id is the chunk's coverage cell, recovered from the sibling\n", - "`morton` coordinate array (`{group}/morton`, alongside the field under the same\n", - "group — the per-cell morton word coarsened to the chunk's order). Populated\n", - "cells land at the deinterleaved `(row, col)` of their nested rank (never a\n", - "row-major reshape -- nested order is a Z-order curve); everything else is zero." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fa454469", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:17.662386Z", - "iopub.status.busy": "2026-06-25T19:05:17.662201Z", - "iopub.status.idle": "2026-06-25T19:05:17.666754Z", - "shell.execute_reply": "2026-06-25T19:05:17.665324Z" - } - }, - "outputs": [], - "source": [ - "from zagg.readers import rank_to_rowcol\n", - "\n", - "tensor = tensors[word_a]\n", - "# A cell's position is the bit deinterleave of its nested rank within the\n", - "# 64x64 (depth-6) chunk: rank 5 -> (0, 3); rank 4095 -> (63, 63).\n", - "for rank, note in [(5, \"populated\"), (4095, \"populated\"), (10, \"unpopulated\")]:\n", - " row, col = rank_to_rowcol(rank, 6)\n", - " print(\n", - " f\"cell rank {rank:4d} -> (row {row:2d}, col {col:2d}) count sum:\",\n", - " int(tensor[row, col].sum()),\n", - " f\"({note})\",\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "534f924c", - "metadata": {}, - "source": [ - "### The reconstructed histogram\n", - "\n", - "Each cell's length-`n_bins` vector is a histogram reconstructed from the digest\n", - "CDF. Plotting one cell's vector against the z-bin edges shows the recovered\n", - "distribution -- here a roughly Gaussian elevation profile centred near 20 m." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2029227c", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:17.668769Z", - "iopub.status.busy": "2026-06-25T19:05:17.668595Z", - "iopub.status.idle": "2026-06-25T19:05:18.272212Z", - "shell.execute_reply": "2026-06-25T19:05:18.270382Z" - } - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "\n", - "from zagg.readers import chunk_z_range\n", - "\n", - "n_bins, resolution = 128, 0.5\n", - "tensor = tensors[word_a]\n", - "vec = tensor[0, 3] # cell rank 5 of chunk A -- deinterleaved position (0, 3)\n", - "\n", - "# Reconstruct the window the reader used: floor of the chunk's trimmed minimum.\n", - "# We recompute the floor here only to label the x-axis; the reader sets it internally.\n", - "digests = [read_cell(store, FIELD, base_a + c) for c in (0, 5, 4095)] # chunk A's cells\n", - "z_lo, n_bins_c, res_c = chunk_z_range(\n", - " digests, n_bins=n_bins, resolution=resolution, bottom=0.05, top=0.95, fit=\"raise\"\n", - ")\n", - "edges = z_lo + res_c * np.arange(n_bins_c + 1)\n", - "centers = 0.5 * (edges[:-1] + edges[1:])\n", - "\n", - "fig, ax = plt.subplots(figsize=(7, 3.5))\n", - "ax.bar(centers, vec, width=res_c, align=\"center\", color=\"steelblue\", edgecolor=\"none\")\n", - "ax.set_xlabel(\"z (m)\")\n", - "ax.set_ylabel(\"count\")\n", - "ax.set_title(f\"cell 5 of chunk A -- reconstructed histogram (z_lo={z_lo:.0f} m)\")\n", - "ax.set_xlim(z_lo, z_lo + 30) # zoom to the populated region\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "id": "9a15e5bd", - "metadata": {}, - "source": [ - "### Output dtype flag\n", - "\n", - "The client consumes `uint16`; `read_tensors` also offers `uint32` (the default,\n", - "safe for dense cells with many counts per bin) and `float32` (keeps fractional\n", - "reconstructed counts). The flag only changes the output cast." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "22da8a78", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:18.275317Z", - "iopub.status.busy": "2026-06-25T19:05:18.274890Z", - "iopub.status.idle": "2026-06-25T19:05:18.298392Z", - "shell.execute_reply": "2026-06-25T19:05:18.296964Z" - } - }, - "outputs": [], - "source": [ - "for dtype in (\"uint16\", \"uint32\", \"float32\"):\n", - " t, *_rest = next(read_tensors(store, FIELD, dtype=dtype))\n", - " print(f\"dtype={dtype:8s} -> tensor.dtype={t.dtype}, peak bin count={t.max()}\")" - ] - }, - { - "cell_type": "markdown", - "id": "3ea2fbfa", - "metadata": {}, - "source": [ - "## 3. Fit policy: when the trimmed range exceeds the window\n", - "\n", - "The window is fixed at `n_bins * resolution` (64 m by default). If a chunk's\n", - "trimmed `bottom`->`top` range is wider than that, the `fit` flag decides:\n", - "\n", - "- `\"raise\"` (default) -- raise, so the client never silently loses tails;\n", - "- `\"degrade_resolution\"` -- double `resolution` in powers of two until it fits\n", - " (keeps `n_bins`);\n", - "- `\"collapse_bins\"` -- shrink `n_bins` to the smallest power of two that fits\n", - " (keeps `resolution`).\n", - "\n", - "We build a deliberately wide chunk (elevations spanning ~400 m) to trigger it." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fb8f0800", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:18.300379Z", - "iopub.status.busy": "2026-06-25T19:05:18.300173Z", - "iopub.status.idle": "2026-06-25T19:05:18.325005Z", - "shell.execute_reply": "2026-06-25T19:05:18.323897Z" - } - }, - "outputs": [], - "source": [ - "wide = MemoryStore()\n", - "grid.emit_template(wide)\n", - "write_chunk(wide, \"1121121\", {0: rng.uniform(0.0, 400.0, 5_000)})\n", - "\n", - "# Default fit=\"raise\": the 400 m span overflows the 64 m window.\n", - "try:\n", - " next(read_tensors(wide, FIELD, bottom=0.0, top=1.0))\n", - "except ValueError as exc:\n", - " print(\"fit='raise' ->\", exc)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "00ce8977", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:18.327004Z", - "iopub.status.busy": "2026-06-25T19:05:18.326770Z", - "iopub.status.idle": "2026-06-25T19:05:18.357425Z", - "shell.execute_reply": "2026-06-25T19:05:18.356093Z" - } - }, - "outputs": [], - "source": [ - "# degrade_resolution: coarsen the bins until the full span fits in 128 bins.\n", - "t, *_rest = next(read_tensors(wide, FIELD, bottom=0.0, top=1.0, fit=\"degrade_resolution\"))\n", - "print(\"fit='degrade_resolution' -> tensor\", t.shape, \"dtype\", t.dtype)\n", - "\n", - "# collapse_bins on a *narrow* chunk: shrink n_bins to the smallest pow2 that fits.\n", - "narrow = MemoryStore()\n", - "grid.emit_template(narrow)\n", - "write_chunk(narrow, \"1121121\", {0: rng.uniform(100.0, 110.0, 4_000)})\n", - "t2, *_rest = next(read_tensors(narrow, FIELD, bottom=0.0, top=1.0, fit=\"collapse_bins\"))\n", - "print(\"fit='collapse_bins' -> tensor\", t2.shape, \"(n_bins collapsed to fit ~10 m span)\")" - ] - }, - { - "cell_type": "markdown", - "id": "608a084d", - "metadata": {}, - "source": [ - "## 4. `read_raw_values`: lossless recovery\n", - "\n", - "When a cell's digest was built with **no merges** -- every centroid weight is 1,\n", - "which happens whenever a cell has at most `delta` samples -- the centroid means\n", - "are exactly the original observations. `read_raw_values` yields\n", - "`(morton_index, (row, col), values)` for each such cell, with `values` the\n", - "recovered sample vector (sorted ascending, as the digest stores centroids by\n", - "mean) and `(row, col)` the same deinterleaved position `read_tensors` places\n", - "the cell at.\n", - "\n", - "If any cell carries a merged centroid (weight > 1), exact recovery is impossible\n", - "and the reader raises -- the same `raise`-by-default contract as `read_tensors`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8e43b0b0", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:18.359441Z", - "iopub.status.busy": "2026-06-25T19:05:18.359267Z", - "iopub.status.idle": "2026-06-25T19:05:18.375601Z", - "shell.execute_reply": "2026-06-25T19:05:18.374485Z" - } - }, - "outputs": [], - "source": [ - "from zagg.readers import read_raw_values\n", - "\n", - "lossless = MemoryStore()\n", - "grid.emit_template(lossless)\n", - "samples = np.array([3.0, 1.0, 2.0, 5.0, 4.0]) # 5 values, well under delta -> no merges\n", - "write_chunk(lossless, \"1121121\", {7: samples}, delta=512)\n", - "\n", - "for morton, (row, col), values in read_raw_values(lossless, FIELD):\n", - " print(f\"chunk {morton}, cell at (row {row}, col {col}): recovered {values}\")\n", - " print(\"matches sorted input:\", np.allclose(values, np.sort(samples)))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b02758c0", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-25T19:05:18.377464Z", - "iopub.status.busy": "2026-06-25T19:05:18.377262Z", - "iopub.status.idle": "2026-06-25T19:05:18.401378Z", - "shell.execute_reply": "2026-06-25T19:05:18.400150Z" - } - }, - "outputs": [], - "source": [ - "# A cell with many samples at a small delta forces merges -> not recoverable.\n", - "merged = MemoryStore()\n", - "grid.emit_template(merged)\n", - "write_chunk(merged, \"1121121\", {0: rng.standard_normal(5_000)}, delta=64)\n", - "try:\n", - " list(read_raw_values(merged, FIELD))\n", - "except ValueError as exc:\n", - " print(\"merged digest ->\", exc)" - ] - }, - { - "cell_type": "markdown", - "id": "ed2c8dcb", - "metadata": {}, - "source": [ - "## Summary\n", - "\n", - "- `read_tensors(store, field, ...)` -> `(tensor, mask, (offset, gain), morton)`\n", - " per block: dense, spatially faithful `(64, 64, n_bins)` tensors (deinterleaved\n", - " per mortie spec §8) with a uint8 occupancy mask and the block's shared\n", - " z-window, plus `bottom`/`top` tail-trim, a `fit` policy for wide blocks, a\n", - " `dtype` flag (`uint16`/`uint32`/`float32`), and `block_order=` to assemble\n", - " whole chunks into larger blocks (with hive leaves, the mask decodes the\n", - " `coverage.moc` occupancy sidecar).\n", - "- `read_raw_values(store, field)` -> exact samples per cell when the digest is\n", - " merge-free, raising otherwise.\n", - "- `read_cell(store, field, cell)` -> one cell's digest by global position.\n", - "\n", - "All read the ragged (variable-length-bytes) t-digest field zagg writes (issue\n", - "#209) and recover the chunk morton id from the sibling `morton` coordinate.\n", - "Point them at a real on-disk or S3 Zarr product and the same calls apply -- only\n", - "the `store` argument changes.\n", - "\n", - "> **Note:** this notebook uses a synthetic in-memory store so it stays\n", - "> Binder-runnable with no credentials. A companion check against a non-synthetic\n", - "> (real-product) t-digest store is tracked as a follow-up on issue #79." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/viewers.py b/notebooks/viewers.py new file mode 100644 index 000000000..3c3d0a53c --- /dev/null +++ b/notebooks/viewers.py @@ -0,0 +1,1050 @@ +"""Shared viewers for the reader-only demo notebooks. + +`hhdc_viewer.ipynb` and `waveform_viewer.ipynb` both open the same paired +stores and differ only in how they draw what they find, so the drawing lives +here and the notebooks stay about the READ path -- polygon, coverage, shards, +leaves, digests. + +Everything here is reader-side: `mortie` for the geometry, `moczarr` for the +store. The one zagg import (`cdf_from_tdigest`) is lazy and confined to the +waveform view -- moczarr imports the t-digest algebra from zagg rather than +vendoring it (moczarr issue #19), which is what the `moczarr[zagg]` extra +carries and what `read_tensors` has always needed. +""" + +from __future__ import annotations + +import time + +import matplotlib.pyplot as plt +import moczarr as mz +import numpy as np +from IPython.display import display +from ipywidgets import ( + Checkbox, + Dropdown, + FloatSlider, + FloatText, + HBox, + IntSlider, + Layout, + VBox, + interactive_output, +) +from matplotlib.colors import LogNorm +from moczarr.hhdc import block_rank, rank_to_rowcol +from mortie import generate_morton_children, toc2time + +#: The o12 block the views frame everything in. Cells across a block edge are +#: ``2**(cell_order - BLOCK_ORDER)`` -- 128 for ATL03's o19, 64 for GEDI's o18. +BLOCK_ORDER = 12 + +#: o12 block side in metres (equal-area, square-equivalent). +SIDE = float(np.sqrt(4 * np.pi / (12 * 4**BLOCK_ORDER)) * 6_371_000) + +#: Points drawn per sensor per block. A leaf holds millions; the 3-D view +#: subsamples deterministically so rotation stays smooth. +CAP = 20_000 + +#: What one unit of a sensor's stored WEIGHT is. ATL03 counts photons. GEDI's +#: flux is background-subtracted counts scaled by a named receiver gain, and +#: this store ships `gain: {name: "unit", value: 1.0}` -- a placeholder -- so +#: "pe" here is PROPORTIONAL to photoelectrons, not calibrated to them (the +#: published chain is ~0.4 pe/count; see `waveform_viewer.ipynb`). The two are +#: not commensurate and run orders of magnitude apart, so every count is +#: labelled and none are ever summed together. +UNITS = {"atl03": "ph", "gedi": "pe"} + +__all__ = [ + "BLOCK_ORDER", + "CAP", + "SIDE", + "UNITS", + "View", + "block_of", + "grid_xy", + "densest_shard", + "human_bytes", + "joint_cells", + "paired_blocks", + "viewer_stats", + "load", + "view3d", + "waveform_view", +] + + +def _si(n) -> str: + """1266765 -> '1.27M'. Counts below a thousand stay exact.""" + n = float(n) + for cut, suffix in ((1e9, "G"), (1e6, "M"), (1e3, "k")): + if n >= cut: + return f"{n / cut:.3g}{suffix}" + return f"{n:,.0f}" + + +def paired_blocks(handles, block_order: int = BLOCK_ORDER, n_bins: int = 256, res: float = 1.0): + """Read both sensors' block tensors and rank the blocks by coincident cells. + + Returns ``(gains, pairs)``: the per-sensor ``{block word: (offset, gain)}`` z + window, and the sorted ``(word, joint mask, A2, G2)`` list the waveform view + picks cells out of. `A2` is ATL03's column sums folded 2x2 onto GEDI's o18 + grid -- one GEDI footprint covers exactly four ATL03 cells -- so `joint` marks + the o18 cells both sensors populate. + + Holds NO tensor. One ATL03 block at these settings is (128, 128, 256) uint32 = + 16 MiB and a shard has 64 per sensor, so keeping the `read_tensors` tuples was + 1.25 GiB against mybinder's 2 GB cap. Each tensor is reduced to its column sums + (64 KiB) and dropped as it arrives, and the z gain is the only other thing the + view ever needed from it -- so peak is one block, not sixty-four. + + Totals printed here come from the TENSORS, so they count only what fell inside + each block's z window and run a few per cent under what the digests hold. + """ + t0 = time.perf_counter() + gains, cols = {}, {} + for name, (store, field) in handles.items(): + gains[name], cols[name] = {}, {} + for tensor, _mask, window, w in mz.read_tensors( + store, + field, + n_bins=n_bins, + resolution=res, + block_order=block_order, + fit="degrade_resolution", + ): + gains[name][w] = window + cols[name][w] = tensor.sum(axis=2) # the tensor is dropped here + print( + f"paired tensors in {time.perf_counter() - t0:.1f}s — " + + ", ".join(f"{len(v)} {k} blocks" for k, v in cols.items()) + ) + pairs = [] + for w, g2 in cols["gedi"].items(): + if w not in cols["atl03"]: + continue + gside = g2.shape[0] + a2 = cols["atl03"][w].reshape(gside, 2, gside, 2).sum(axis=(1, 3)) + joint = (a2 > 0) & (g2 > 0) + if joint.any(): + pairs.append((w, joint, a2, g2)) + pairs.sort(key=lambda p: -int(p[1].sum())) + for w, j, a2, g2 in pairs[:8]: + print( + f" {mz.morton_decimal(w)} {int(j.sum()):4,} joint o18 cells " + f"{int(a2[j].sum()):9,} atl03 ph* {int(g2[j].sum()):9,} gedi pe*" + ) + print(" * tensor totals: only what fell inside each block's z window") + return gains, pairs + + +def viewer_stats(handles, block_order: int = BLOCK_ORDER) -> dict: + """Sweep each field once and return what the block dropdown needs. + + One pass per sensor over ONE field's stored digests. Nothing is kept -- the + viewer reads per block on demand -- but the per-block cell and weight counts, + and the count of cells both sensors populate, fall out of the same pass for + free. Also prints what that pass cost, which is the price of reading one + array of the leaf end to end (this ATL03 leaf holds nine). + + Reading the dense `morton`/`count` sidecars instead is slower, not faster: + they span the shard's whole cell space (1,048,576 entries at o19) where this + touches only populated cells. Measured 124 s against 74 s. + """ + tally, words = {}, {} + for name, (store, field) in handles.items(): + t0 = time.perf_counter() + _, element = mz.open_ragged(store, field) + per, seen, centroids, nbytes, obs = {}, [], 0, 0, 0.0 + for word, value in mz.read_ragged(store, field): + v = np.asarray(value) + seen.append(word) + weight = float(v[:, 1].sum()) # column 1 is the centroid weight + block = int(block_of(word, block_order)) + had_cells, had_obs = per.get(block, (0, 0.0)) + per[block] = (had_cells + 1, had_obs + weight) + centroids, nbytes, obs = centroids + len(v), nbytes + v.nbytes, obs + weight + tally[name] = per + words[name] = np.asarray(seen, dtype=np.uint64) + print( + f"{name:6s} swept {field} in {time.perf_counter() - t0:5.1f}s — " + f"{len(seen):,} cells over {len(per)} blocks, {centroids:,} centroids, " + f"{obs:,.0f} {UNITS.get(name, 'obs')}, {nbytes / 2**20:.1f} MiB decoded" + ) + print(f" element {element}") + return {"tally": tally, "joint": joint_cells(words, block_order)} + + +def densest_shard(root, shards, **s3) -> str: + """The shard with the most granules, by the store's own leaf telemetry. + + `read_stats` fetches one leaf's stats sidecar by arithmetic on the manifest + -- no LIST, a few KB each -- so ranking a handful of candidates is cheap. + This is the rule `demo/06_paired` used (`max` on the GEDI granule count), + reproduced here off the published store instead of a local shardmap, so + both demos land on the same shard without anyone hardcoding an id. + """ + return max(shards, key=lambda s: (mz.read_stats(root, s, **s3) or {}).get("n_granules", 0)) + + +def joint_cells(words, block_order: int = BLOCK_ORDER): + """{block word: cells BOTH sensors populate}, from the sweep's words alone. + + The sensors are keyed at different orders -- ATL03 o19, GEDI o18 -- so each + side is folded to the COARSER of the two before intersecting: one GEDI + footprint covers 2x2 ATL03 cells, and it is joint when any of the four holds + a digest. Exact, and free: the open-and-sweep pass already visited every one + of these words, so this reads nothing and needs no tensor. + + (Checked against the tensor-derived mask `waveform_viewer` builds -- both + give 27 joint cells on block 4331422233444 -- and against the occupancy + channel `read_tensors` returns, which agrees to the cell.) + """ + if len(words) < 2: + return {} + per, coarse = {}, None + for name, w in words.items(): + w = np.asarray(w, dtype=np.uint64) + if not len(w): + return {} # a sensor with nothing here has nothing joint anywhere + order = int(np.asarray(block_rank(w[:1], block_order)[1]).ravel()[0]) + per[name] = w + coarse = order if coarse is None else min(coarse, order) + both = None + for w in per.values(): + folded = np.unique(block_of(w, coarse)) + both = folded if both is None else np.intersect1d(both, folded, assume_unique=True) + if not len(both): + return {} + blocks, counts = np.unique(block_of(both, block_order), return_counts=True) + return {int(b): int(c) for b, c in zip(blocks, counts)} + + +def human_bytes(n) -> str: + """Byte count -> '320.0 MiB' / '1.25 GiB'. Exports print both sides of a write.""" + return f"{n / 2**30:.2f} GiB" if n >= 2**30 else f"{n / 2**20:.1f} MiB" + + +def block_of(words, block_order: int = BLOCK_ORDER): + """Cell word(s) -> their order-`block_order` ancestor word, by truncation. + + A packed area word is ``[prefix | body | 6-bit suffix]`` and the suffix + carries the order, so an ancestor is the body truncated to the ancestor's + depth with the suffix rewritten -- no decode, and it vectorizes over an + array. Only valid for AREA words at or below `block_order`. + """ + sh = np.uint64(6 + 2 * (27 - block_order)) + return ((np.asarray(words, dtype=np.uint64) >> sh) << sh) | np.uint64(block_order) + + +# --------------------------------------------------------------------------- # +# geometry +# --------------------------------------------------------------------------- # +def grid_xy(words, block_order: int = BLOCK_ORDER): + """Word -> (x, y) metres inside its o12 block. Three library calls, no bits. + + ``block_rank`` recovers each word's block-local nested rank and its own + order -- normalizing POINT words to their order-29 area twins on the way, + which is the step whose absence used to make the level-28/29 digits decode + out of range. ``rank_to_rowcol`` is the bit deinterleave (mortie spec + section 8): it returns ``(row, col) = (y, x)`` with ``[0, 0]`` at the + subtree's south corner, and this returns them as ``(x, y) = (col, row)`` + so the picture matches the tensors this notebook exports. + + Those axes are NOT east and north. zagg's ``readers/_layout.py`` pins the + semantics: "``tensor[0, 0]`` is the subtree's south corner; rows advance + toward the north-WEST edge, columns toward the north-EAST edge." A HEALPix + face is a diamond, so its two local axes run along the face's edges, not + along the compass. ``view3d`` labels them for what they are. + + Note the tensors ``read_tensors`` (and so ``hhdc_viewer``'s export cells) + hand back are indexed ``(row, col, bin)``: the picture's x is the tensor's + SECOND axis. :func:`_binned_pts` applies the same swap, so the binned and + exact views agree with each other. + + A leaf mixes orders -- o29 located words beside coarser merged-centroid + words, and GEDI's cell words are o18 -- so the ranks are grouped by depth + and handed over one vectorized call per depth. + """ + rank, order = block_rank(words, block_order) + depth = order - block_order + row = np.zeros(len(rank)) + col = np.zeros(len(rank)) + for d in np.unique(depth): + m = np.flatnonzero(depth == d) + r, c = rank_to_rowcol(rank[m], int(d)) + row[m] = np.asarray(r, dtype=float) + col[m] = np.asarray(c, dtype=float) + side = 2.0 ** depth.astype(float) # cells along a block edge + return (col + 0.5) / side * SIDE, (row + 0.5) / side * SIDE + + +def load(store, field, block_order: int = BLOCK_ORDER) -> dict: + """One sensor's centroids: z, weight, xy, o12 block, acquisition days. + + xy comes from the located sibling's word where the field declares one + (ATL03 -- point-exact) and from the cell word where it does not (GEDI + flux, whose shots are unlocated by design, so a centroid renders at its + cell's centre). Both land in the block's own lattice frame. + """ + arr, _ = mz.open_ragged(store, field) + attrs = dict(arr.attrs) + locname = (attrs.get("ragged") or {}).get("locations") + tname = attrs.get("times") + zs, wts, cells, locw, seq = [], [], [], [], [] + for row in mz.read_ragged(store, field, locations=bool(locname)): + seq.append(row[0]) + v = np.asarray(row[1]) + zs.append(v[:, 0]) + wts.append(v[:, 1]) + cells.append(np.full(len(v), row[0], dtype=np.uint64)) + if locname: + locw.append(np.asarray(row[2], dtype=np.uint64)) + z, wt = np.concatenate(zs), np.concatenate(wts) + cells = np.concatenate(cells) + blocks = block_of(cells, block_order) + x, y = grid_xy(np.concatenate(locw) if locname else cells, block_order) + t = None + if tname: + tmap = { + int(r[0]): np.asarray(r[1], dtype=np.uint64).ravel() + for r in mz.read_ragged(store, field.rsplit("/", 1)[0] + "/" + tname) + } + tw = np.concatenate([tmap[int(c)] for c in seq]) + ns2018 = float( + (np.datetime64("2018-01-01") - np.datetime64("1850-01-01")) // np.timedelta64(1, "ns") + ) + t = (np.asarray(toc2time(tw)[0], dtype="float64") - ns2018) / 86.4e12 + return { + "z": z, + "wt": wt, + "x": x, + "y": y, + "blocks": blocks, + "t": t, + "xy_note": "exact xy" if locname else "cell-center xy", + } + + +# --------------------------------------------------------------------------- # +# the 3-D view +# --------------------------------------------------------------------------- # +def _binned_pts(tensor, offset, gain, cap: int = CAP): + """Occupied voxels of one block tensor -> (x, y, z, weight), subsampled. + + ``x``/``y`` come back as FRACTIONS of the block edge (the tensor's own + ``side`` is its xy shape), so the caller scales them by :data:`SIDE` the + same way the exact path does, and ``x`` is the tensor's COL axis and ``y`` + its ROW axis, the same round as :func:`grid_xy`. + + ``np.nonzero`` gives cell INDICES, so the ``+ 0.5`` that moves a point off + the cell corner onto its centre is the same one :func:`grid_xy` applies -- + without it, toggling **binned** shifts the whole cloud by half a cell, and + by a DIFFERENT half-cell per sensor (6.2 m for ATL03's o19 cells, 12.4 m + for GEDI's o18). + """ + side = tensor.shape[0] + rows, cols, zs = np.nonzero(tensor) # the tensor is indexed (row, col, bin) + counts = tensor[rows, cols, zs].astype("float64") + if len(rows) > cap: + keep = np.random.default_rng(0).choice(len(rows), cap, replace=False) + rows, cols, zs, counts = rows[keep], cols[keep], zs[keep], counts[keep] + return (cols + 0.5) / side, (rows + 0.5) / side, offset + zs * gain, counts + + +def _coverage(recs): + """Coarse sensor -> ``(fine sensor, its cells over a fine cell, its cells)``. + + One GEDI o18 footprint tiles exactly 2x2 of ATL03's o19 cells, so "does + this GEDI cell sit over a stored ATL03 SIGNAL digest" is a block-reduce of the + finer occupancy channel ANDed with the coarser one: two boolean arrays of + at most 128x128, and no extra I/O -- `read_tensors` hands the occupancy + back beside every tensor it already read. + + Occupancy is `mask == 2` ("observed, digest stored"), never + `mask.astype(bool)`: where the leaf carries an exact occupancy sidecar the + mask is 3-state and `1` means observed-but-EMPTY, which is precisely the + cell this is asking about. Counting it as populated would inflate both the + cell counts and this overlap. + """ + occ = {n: r["occ"] for n, r in recs.items() if r["occ"] is not None} + if len(occ) < 2: # a block one sensor missed entirely -- nothing to compare + return {} + fine = max(occ, key=lambda n: occ[n].shape[0]) + out = {} + for name, o in occ.items(): + side = o.shape[0] + k, rem = divmod(occ[fine].shape[0], side) + if name == fine or rem: # not a whole-cell nesting; say nothing + continue + over = occ[fine].reshape(side, k, side, k).any(axis=(1, 3)) + out[name] = (fine, int((over & o).sum()), int(o.sum())) + return out + + +class View: + """Holds what is on screen so `export` knows which block you mean.""" + + shard = None + block = None + + +def view3d( + handles, + shard, + stats=None, + block_order: int = BLOCK_ORDER, + cap: int = CAP, + n_bins: int = 256, + resolution: float = 1.0, +) -> View: + """Interactive paired 3-D view. Drag to rotate -- the two panes stay linked. + + ``handles`` is ``{sensor: (leaf_store, field)}``. Two ways to see the same + block, switched by the **binned** box: + + * **binned** (default) -- the fixed ``read_tensors`` voxels, so xy AND z + are on the store's own grid and the two sensors share a z window. + * **exact** -- the digests read straight through, z the stored float32 + centroid elevation and xy decoded from the located sibling's word where + the field carries one (ATL03, point-exact) or the cell word where it + does not (GEDI flux, footprint scale). + + Nothing is read until it is drawn. The block list is arithmetic, each + block's tensor is fetched the first time you select it and cached, and the + exact centroids are swept lazily on the first unbinned draw -- so a reader + who looks at one block pays for one block, and one who never unticks the + box never pays for the sweep at all. + + `stats` is what :func:`viewer_stats` returns; it labels and orders the block + dropdown, and costs nothing here because that pass has already been paid for. + Those are STORED totals over the whole digest. The counts each read prints, + and the ones in the pane titles, are what landed inside the tensor's finite + z window and so run lower -- the tails are trimmed at `bottom`/`top` and + anything past ``n_bins * gain`` metres of relief has nowhere to go. + + Returns the :class:`View` the widgets write to, so a later cell can export + whatever is on screen. + """ + view = View() + view.shard = shard + names = list(handles) + tally = (stats or {}).get("tally") + joint = (stats or {}).get("joint") + + # One block per read, cached. Sweeping the shard up front cost ~1.3 GiB + # resident (Binder caps at 2 GB) and ~148 s before the first frame. + voxels: dict = {n: {} for n in names} + + def _voxels(name, block): + """This sensor's drawable voxels for one block, read once and cached.""" + hit = voxels[name].get(block) + if hit is None: + store, field = handles[name] + t0 = time.perf_counter() + got = next( + iter( + mz.read_tensors( + store, + field, + n_bins=n_bins, + resolution=resolution, + block_order=block_order, + fit="degrade_resolution", + subtree=mz.morton_decimal(block), + ) + ), + None, + ) + if got is None: # the block holds nothing in this store + hit = {"pts": None, "cells": 0, "gain": resolution, "occ": None, "obs": 0.0} + else: + occ = got[1] == 2 # stored digest; see `_coverage` on why not `.astype(bool)` + hit = { + "pts": _binned_pts(got[0], *got[2], cap), + "cells": int(occ.sum()), + "gain": got[2][1], + "occ": occ, + "obs": float(got[0].sum(dtype="float64")), + "bins": got[0].shape[2], + } + voxels[name][block] = hit + # Report what the read actually MOVED, not a resident-memory delta: + # cells, the sensor's own weight unit, and the z window they fell in. + print( + f" {name} {mz.morton_decimal(block)}: " + + ( + "empty" + if hit["pts"] is None + else f"{hit['cells']:,} cells, {_si(hit['obs'])} " + f"{UNITS.get(name, 'obs')} in a {hit['bins']} x {hit['gain']:g} m z window" + ) + + f", {time.perf_counter() - t0:.1f}s", + flush=True, + ) + return hit + + exact: dict = {} # filled on the first unbinned draw + + def _exact(): + if not exact: + # One whole-shard sweep per sensor, millions of centroids -- say so, + # or the first untick reads as a frozen notebook. + print("sweeping exact centroids (once, whole shard, both sensors)...", flush=True) + t0 = time.perf_counter() + for name, (store, field) in handles.items(): + exact[name] = load(store, field, block_order) + print( + " " + + ", ".join(f"{len(exact[n]['z']):,} {n}" for n in handles) + + f" in {time.perf_counter() - t0:.0f}s", + flush=True, + ) + return exact + + # Blocks the sweep actually found, so they cannot disagree with the open + # leaf; without a tally, arithmetic off the shard word, no I/O. + within = {int(w) for w in generate_morton_children(int(mz.morton_word(shard)), block_order)} + if tally: + found = {w for per in tally.values() for w in per} + if not found & within: + raise ValueError( + f"the stats hold no block beneath shard {shard} — `handles`/`stats` and " + f"`shard` disagree. Open and view the SAME shard (one name, used twice)." + ) + blocks = found & within + else: + blocks = within + if not blocks: + raise ValueError(f"shard {shard}: no order-{block_order} blocks beneath it") + + # Rank by coincident cells -- no scaling needed, unlike the weights. + if joint: + order = sorted(blocks, key=lambda w: (-joint.get(w, 0), -w)) + else: + order = sorted(blocks, reverse=True) # descending, so ...444 heads the list + + def _label(w): + """Block label: decimal id, coincident cells, each sensor's stored size.""" + parts = [] + if joint is not None: + parts.append(f"{joint.get(w, 0):,} joint") + parts += [ + f"{n} {c:,} cells / {_si(o)} {UNITS.get(n, 'obs')}" + for n in names + for c, o in [(tally or {}).get(n, {}).get(w, (0, 0.0))] + ] + return mz.morton_decimal(w) + (" — " + " · ".join(parts) if tally or joint else "") + + dd = Dropdown( + options=[(_label(w), w) for w in order], + value=order[0], + description="block", + layout=Layout(width="700px" if (tally or joint) else "260px"), + ) + zmode = Dropdown( + options=[("independent z", "auto"), *[(f"pin z to {n}", n) for n in names]], + value="auto", + description="z extent", + ) + elev_cb = Checkbox(value=False, description="color by elevation (shared)") + time_cb = Checkbox(value=False, description="color by time (shared)") + bin_cb = Checkbox(value=True, description="binned (fixed tensors)") + # Flat by default: weight is already the colour, and letting it drive alpha + # too put 76% of a block's voxels under 0.10 and hid the single-count ones. + # `shade by weight` puts that cue back deliberately, spanning the slider's + # value down to a twentieth of it. + # `continuous_update=False` here, unlike the waveform's `nth`: this slider + # reads nothing (block voxels are cached), so the only cost is the redraw -- + # ~585 ms for 2 panes of 20k 3-D points. Firing that per drag step queues + # frames for a value nobody asked to stop on; on release it is one redraw. + alpha_sl = FloatSlider( + min=0.05, + max=1.0, + step=0.05, + value=0.55, + description="opacity", + continuous_update=False, + readout_format=".2f", + layout=Layout(width="330px"), + ) + shade_cb = Checkbox(value=False, description="shade by weight") + + said: set = set() # blocks whose coarse-over-fine line has already printed + + def _panes(block, binned): + if binned: + recs = {n: _voxels(n, block) for n in names} + cov = _coverage(recs) + if block not in said: + said.add(block) + for n, (fine, over, total) in cov.items(): + print( + f" {n} over {fine}: {over:,} of {total:,} cells " + f"({100 * over / max(total, 1):.0f}%)", + flush=True, + ) + out = [] + for n in names: + rec = recs[n] + if rec["pts"] is None: # nothing stored here for this sensor + empty = np.empty(0) + out.append( + { + "x": empty, + "y": empty, + "z": empty, + "wt": empty, + "t": None, + "label": n, + "xy_note": "no data in this block", + } + ) + continue + x, y, z, wt = rec["pts"] + note = f"binned ({rec['gain']:g} m z, {rec['cells']:,} cells" + if n in cov: + fine, over, total = cov[n] + note += f", {100 * over / max(total, 1):.0f}% over {fine}" + out.append( + { + "x": x * SIDE, + "y": y * SIDE, + "z": z, + "wt": wt, + "t": None, + "label": n, + "xy_note": note + ")", + } + ) + return out + data = _exact() + out = [] + rng = np.random.default_rng(0) + for n in names: + d = data[n] + m = np.flatnonzero(d["blocks"] == np.uint64(block)) + if len(m) > cap: + m = m[rng.choice(len(m), cap, replace=False)] + out.append( + { + **{k: d[k][m] for k in ("z", "wt", "x", "y")}, + "t": None if d["t"] is None else d["t"][m], + "label": n, + "xy_note": d["xy_note"], + } + ) + return out + + live: dict = {} # the one figure this view keeps open + + def draw(block, by_elev, by_time, binned, zmode, opacity, shade): + view.block = block + panes = _panes(block, binned) + # A block can be empty in one store and not the other (and every pane is + # empty if it is empty in both), so every range below is taken over the + # NON-empty panes only -- `.min()` and `np.percentile` both raise on an + # empty array, which would take the whole view down on a sparse block. + filled = [p for p in panes if len(p["z"])] + pinned = panes[names.index(zmode)] if zmode != "auto" else None + zlim = ( + (pinned["z"].min(), pinned["z"].max()) + if pinned is not None and len(pinned["z"]) + else None + ) + shared_elev = ( + plt.Normalize(min(p["z"].min() for p in filled), max(p["z"].max() for p in filled)) + if by_elev and not by_time and filled + else None + ) + tv = [p["t"] for p in panes if p["t"] is not None and len(p["t"])] + shared_time = ( + plt.Normalize(min(t.min() for t in tv), max(t.max() for t in tv)) + if (by_time and tv) + else None + ) + ticks = [0.0, SIDE / 2, SIDE] + labels = [f"{v:.0f} m" for v in ticks] + + # `draw` reruns on EVERY widget change, and `clear_output(wait=True)` + # clears the Output widget's display, not matplotlib's figure registry. + # An ipympl figure is a live canvas -- it holds its websocket comm and + # the `_sync` handler connected below for the kernel's lifetime -- so + # without this the 21st interaction warns and every orphan keeps + # handling mouse events. Close the PREVIOUS figure, not this one: the + # widget backend needs the current canvas alive to render it. + if (prev := live.pop("fig", None)) is not None: + plt.close(prev) + fig = plt.figure(figsize=(11, 5.2)) + live["fig"] = fig + axes = [] + for k, (p, cmap) in enumerate(zip(panes, ("viridis", "plasma"))): + ax = fig.add_subplot(1, 2, k + 1, projection="3d") + axes.append(ax) + note = f" — {p['xy_note']}" + if not len(p["z"]): # empty here; label the pane and leave the axes bare + ax.set_title(p["label"] + note, fontsize=9) + ax.set_zlabel("elevation (m)") + continue + # Weight is ALREADY the colour, log-normed. Letting it drive opacity + # linearly too erased the sparse returns: with a p98 in the hundreds + # a single-count voxel landed on the 0.08 floor and vanished, and + # single-count voxels are the interesting part of a lidar cloud -- + # canopy top, understory, anything off the ground return. Ramp on + # log(weight) instead, over a narrow band with a high floor, so + # density still reads as depth without hiding anything. + if shade: + # Ramp DOWN from the slider, not up to opaque: the slider stays the + # ceiling for the densest voxels and weight pushes the sparse ones + # toward invisible, which is the whole point of ticking the box. + # Ticking it never makes anything more opaque than the flat setting. + ref = max(np.percentile(p["wt"], 98), 1.0) + lo = 0.05 * opacity + alpha = np.clip( + lo + (opacity - lo) * np.log1p(p["wt"]) / np.log1p(ref), lo, opacity + ) + else: + alpha = opacity + if by_time and p["t"] is not None: + pts = ax.scatter( + p["x"], + p["y"], + p["z"], + c=p["t"], + s=1.5, + cmap="turbo", + norm=shared_time, + alpha=alpha, + ) + fig.colorbar(pts, shrink=0.55, pad=0.10, label="days since 2018-01-01") + elif by_time: + ax.scatter(p["x"], p["y"], p["z"], color="#9498a0", s=1.5, alpha=opacity * 0.4) + note += " — no temporal channel" + (" in binned tensors" if binned else "") + elif by_elev: + pts = ax.scatter( + p["x"], + p["y"], + p["z"], + c=p["z"], + s=1.5, + cmap="viridis", + norm=shared_elev, + alpha=alpha, + ) + fig.colorbar(pts, shrink=0.55, pad=0.10, label="elevation (m)") + else: + pts = ax.scatter( + p["x"], + p["y"], + p["z"], + c=p["wt"], + s=1.5, + cmap=cmap, + norm=LogNorm(), + alpha=alpha, + ) + fig.colorbar(pts, shrink=0.55, pad=0.10, label="weight") + if zlim is not None: + ax.set_zlim(*zlim) + ax.set_title(p["label"] + note, fontsize=9) + ax.set_zlabel("elevation (m)") + ax.set_xticks(ticks, labels, fontsize=7) + ax.set_yticks(ticks, labels, fontsize=7) + # Face-local axes from the block's SOUTH corner: col runs toward the + # north-EAST edge, row toward the north-WEST (zagg readers/_layout.py). + # A HEALPix face is a diamond -- these are not compass east/north. + ax.set_xlabel("→ NE edge", fontsize=8, labelpad=-2) + ax.set_ylabel("→ NW edge", fontsize=8, labelpad=-2) + + # Linked rotation: only while DRAGGING, and only when the angles moved -- + # a bare hover must not trigger a redraw. + def _sync(event): + if event.button is None or event.inaxes not in axes: + return + src = event.inaxes + for other in axes: + if other is not src and (other.elev != src.elev or other.azim != src.azim): + other.view_init(elev=src.elev, azim=src.azim) + fig.canvas.draw_idle() + + fig.canvas.mpl_connect("motion_notify_event", _sync) + fig.suptitle( + f"shard {shard} — block {mz.morton_decimal(block)}" + + ("" if binned else " — exact centroids"), + fontsize=11, + ) + plt.show() + + out = interactive_output( + draw, + { + "block": dd, + "by_elev": elev_cb, + "by_time": time_cb, + "binned": bin_cb, + "zmode": zmode, + "opacity": alpha_sl, + "shade": shade_cb, + }, + ) + # Three rows: an HBox does not wrap, and an overrun row squeezes the slider + # into an undraggable stub. + display( + VBox( + [ + HBox([dd, zmode]), + HBox([elev_cb, time_cb, bin_cb]), + HBox([alpha_sl, shade_cb]), + out, + ] + ) + ) + return view + + +# --------------------------------------------------------------------------- # +# the waveform view +# --------------------------------------------------------------------------- # +def _mixture(digest, z, sigma): + """Digest -> density on `z`: each centroid a Gaussian of width `sigma`.""" + mu, wt = digest[:, 0], digest[:, 1] + pdf = (wt[None, :] * np.exp(-0.5 * ((z[:, None] - mu[None, :]) / sigma) ** 2)).sum(axis=1) + return pdf / max(wt.sum(), 1e-9) / (sigma * np.sqrt(2 * np.pi)) + + +def _block_digests(store, field, block, cell_order, block_order=BLOCK_ORDER): + """Every stored digest in one block, keyed by its ``(row, col)`` in that block. + + ONE `read_ragged(subtree=...)` for the whole block instead of a `read_cell` + per cell. Same bytes, one round trip: the per-cell path cost five round trips + for every slider step (one GEDI cell plus its four ATL03 quarters), and a + single `read_cell` measured a median of 8.9 s on a poor link -- 45 s a step, + which is why the view read as frozen. Reading the block up front costs about + what one step used to, and every step after it is a dict lookup. + + Keying on ``(row, col)`` also retires the chunk arithmetic the per-cell path + needed: `read_cell` had to resolve an o19 cell to its o13 read chunk first, + with a `chunk_side` that merely happened to equal GEDI's block side on this + store pair. `read_ragged` hands back the cell word, and `block_rank` / + `rank_to_rowcol` place it directly. + """ + got = list(mz.read_ragged(store, field, subtree=mz.morton_decimal(int(block)))) + if not got: + return {} + words = np.array([w for w, _v in got], dtype=np.uint64) + rows, cols = rank_to_rowcol(block_rank(words, block_order)[0], cell_order - block_order) + return {(int(rows[i]), int(cols[i])): np.asarray(v) for i, (_w, v) in enumerate(got)} + + +def waveform_view(stores, fields, gains, pairs, shard): + """Interactive coincident-waveform view. + + ``stores``/``fields`` are ``{sensor: ...}``; ``gains`` and ``pairs`` are what + :func:`paired_blocks` returns. + + Digests are read a BLOCK at a time and cached, so moving the slider does no + I/O at all -- see :func:`_block_digests`. + """ + from zagg.stats.tdigest import cdf_from_tdigest # moczarr[zagg]; see module docstring + + orders = {n: int(f.split("/", 1)[0]) for n, f in fields.items()} + + live: dict = {} # the one figure this view keeps open + + # Cached per BLOCK, not per cell: only changing the block reads anything, so + # `nth` and `binw` are pure rendering. Two blocks kept -- one is a few MB. + cache: dict = {} + + def _block(w): + if w not in cache: + if len(cache) > 1: + cache.clear() + cache[w] = { + n: _block_digests(stores[n], fields[n], w, orders[n]) for n in ("gedi", "atl03") + } + return cache[w] + + def _digests(pair, r, c, w): + """One GEDI cell and the 2x2 ATL03 cells under it, from the block cache.""" + cells = _block(w) + gdigest = cells["gedi"].get((r, c), np.empty((0, 2))) + quarters = [cells["atl03"].get((2 * r + dr, 2 * c + dc)) for dr in (0, 1) for dc in (0, 1)] + kept = [q for q in quarters if q is not None and len(q)] + if not kept: + return gdigest, np.empty((0, 2)) + # SORT BY MEAN. Each quarter is sorted on its own, so concatenating gives + # four ascending runs, not one -- and `cdf_from_tdigest` reads the centroid + # means as `np.interp`'s x-coordinates, which numpy documents as nonsense + # unless they increase. Measured on a representative 2x2 the CDF drifted by + # 32% of the cell's weight while still looking monotonic, so it fails + # silently. The mixture is a weighted sum and does not care, which is why + # only the CDF was wrong. + merged = np.concatenate(kept) + return gdigest, merged[np.argsort(merged[:, 0], kind="stable")] + + def paired_waveform(pair=0, nth=0, binw=1.0): + w, joint, a2, g2 = pairs[pair] # a2/g2: the notebook's A2/G2, lowercased for N806 + _aoff, ag = gains["atl03"][w] + _goff, gg = gains["gedi"][w] + # Rank the joint cells by the weaker side. a2 is ATL03 PHOTON counts and + # g2 GEDI PHOTOELECTRONS -- incommensurate units that run two to three + # orders of magnitude apart on this store pair, so a raw `np.minimum` + # would collapse to a2 on essentially every cell and "the weaker member" + # would really mean "the ATL03 count". Scale each side by its own + # maximum over the JOINT cells first, so the min picks out whichever + # sensor is relatively weaker. (What makes a pick *coincident* at all is + # the `joint` mask itself -- both sensors populated -- not this min.) + a = a2 / max(int(a2[joint].max()), 1) + g = g2 / max(int(g2[joint].max()), 1) + rank = np.minimum(a, g) * joint + order = np.argsort(rank.ravel())[::-1] + r, c = np.unravel_index(int(order[min(nth, int(joint.sum()) - 1)]), rank.shape) + + gdigest, adigest = _digests(pair, int(r), int(c), w) + + # The ATL03 side can legitimately come back empty: `joint` is built from + # the TENSORS (a finite z window, possibly degraded) while this reads the + # RAW digests, and the two are not guaranteed to agree cell for cell. + # Draw the GEDI side alone and say so, rather than letting an empty + # `.min()` turn "no photons here" into a traceback inside the widget. + zmu = [gdigest[:, 0]] + ([adigest[:, 0]] if len(adigest) else []) + lo = min(a.min() for a in zmu) - 5 + hi = max(a.max() for a in zmu) + 5 + z = np.linspace(lo, hi, 700) + amu, awt = adigest[:, 0], adigest[:, 1] + + # Same figure-registry bound as `view3d.draw`: this is an + # `interactive_output` callback too, so close the previous figure + # rather than accumulating one per widget change. + if (prev := live.pop("fig", None)) is not None: + plt.close(prev) + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4.6), sharey=True) + live["fig"] = fig + ax1.plot( + _mixture(gdigest, z, gg), + z, + color="#7b3294", + lw=2, + label=f"GEDI flux ({gdigest[:, 1].sum():.0f} pe)", + ) + if len(adigest): + ax1.plot( + _mixture(adigest, z, ag), + z, + color="#008837", + lw=2, + label=f"ATL03 signal ({awt.sum():.0f} ph)", + ) + ax1.set_xlabel("normalized density") + ax1.set_ylabel("elevation (m)") + + # Top axis: the RAW ATL03 photons -- binned bars, or one dot per centroid + # when the width is 0 (cells under the store's delta budget are loss-free, + # so centroids ~ photons there; a merged centroid's weight sets its size). + ax1t = ax1.twiny() + if binw and binw > 0: + edges = np.arange(lo, hi + binw, binw) + counts, _ = np.histogram(amu, bins=edges, weights=awt) + ax1t.barh( + edges[:-1] + binw / 2, + counts, + height=binw * 0.9, + color="#008837", + alpha=0.25, + zorder=0, + ) + ax1t.set_xlabel(f"ATL03 photons / {binw:g} m bin", fontsize=9) + else: + ax1t.scatter( + np.zeros(len(amu)), + amu, + s=np.clip(awt * 8, 8, 40), + color="#008837", + alpha=0.45, + zorder=0, + ) + ax1t.set_xlim(-0.05, 1.0) + ax1t.set_xlabel("ATL03 photons (unbinned)", fontsize=9) + ax1.set_zorder(ax1t.get_zorder() + 1) + ax1.patch.set_visible(False) + ax1.set_title( + f"cell ({r},{c}) @o18 — GEDI {len(gdigest)} centroids " + + ( + f"vs ATL03 {len(adigest)} (2×2 @o19)" + if len(adigest) + else "— no ATL03 digest under this cell" + ), + fontsize=9, + ) + ax1.legend(fontsize=8) + + # cdf_from_tdigest returns CUMULATIVE WEIGHT (pe for GEDI, photons for + # ATL03) -- normalize each by its own total so both share the axis honestly. + ax2.plot( + cdf_from_tdigest(gdigest, z) / max(gdigest[:, 1].sum(), 1e-9), z, color="#7b3294", lw=2 + ) + if len(adigest): + ax2.plot( + cdf_from_tdigest(adigest, z) / max(adigest[:, 1].sum(), 1e-9), + z, + color="#008837", + lw=2, + ) + ax2.set_xlim(0, 1) + ax2.set_xlabel("CDF (probability)") + ax2.set_title("cumulative", fontsize=10) + for ax in (ax1, ax2): + ax.spines[["top", "right"]].set_visible(False) + ax.grid(alpha=0.25, lw=0.5) + fig.suptitle(f"shard {shard} — block {mz.morton_decimal(w)}", fontsize=11) + plt.tight_layout() + plt.show() + + dd = Dropdown( + options=[ + (f"{mz.morton_decimal(w)} ({int(j.sum()):,} joint cells)", i) + for i, (w, j, _, _) in enumerate(pairs) + ], + value=0, + description="block", + layout=Layout(width="360px"), + ) + # Opens on the 7th coincident cell (0-indexed): on the SERC block this + # notebook ships with, that one shows the canopy/ground structure clearly. + # Any other AOI will rank differently -- it is a nice default, not a rule. + # + # `continuous_update` left at its default (True), unlike the opacity slider: + # since `_block_digests` caches the whole block, moving this reads NOTHING -- + # a dict lookup and two 2-D line plots -- so there is nothing to suppress. + # (It used to cost five S3 round trips a step, ~45 s on a poor link, which is + # what made the control look dead.) + # + # `max` is re-set per block below. A fixed 0..40 was wrong: 30 of the 64 blocks + # on this shard hold fewer than 41 coincident cells, and `paired_waveform` + # clamps with `min(nth, joint.sum() - 1)`, so every position past the count + # selected the SAME cell and the view stopped changing. + nth = IntSlider(min=0, max=40, value=6, description="nth joint", layout=Layout(width="330px")) + binw = FloatText( + value=1.0, + step=0.5, + description="histogram bin size, ATL03 z (m)", + style={"description_width": "initial"}, + layout=Layout(width="330px"), + ) + + # Two rows, every widget sized -- see the HBox note in `view3d`. + def _fit_nth(_change=None): + """Bound `nth` to the cells this block actually has.""" + nth.max = max(int(pairs[dd.value][1].sum()) - 1, 0) + + dd.observe(_fit_nth, names="value") + _fit_nth() + display( + VBox( + [ + HBox([dd, nth]), + binw, + interactive_output(paired_waveform, {"pair": dd, "nth": nth, "binw": binw}), + ] + ) + ) diff --git a/notebooks/waveform_viewer.ipynb b/notebooks/waveform_viewer.ipynb new file mode 100644 index 000000000..dda0dc6e5 --- /dev/null +++ b/notebooks/waveform_viewer.ipynb @@ -0,0 +1,205 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7fb27b941602401d91542211134fc71a", + "metadata": {}, + "source": [ + "# polygon \u2192 MOC \u2192 shard \u2192 coincident waveforms\n", + "\n", + "The waveform half of the minimal read stack. Same two libraries, same polygon,\n", + "same public store, zero credentials \u2014 but instead of the 3-D scatter, this\n", + "notebook does the **cell-level join**: one GEDI o18 footprint against the 2\u00d72\n", + "ATL03 o19 cells beneath it, both reconstructed from their stored t-digests as\n", + "densities on a shared elevation axis.\n", + "\n", + "Its sibling is [`hhdc_viewer.ipynb`](hhdc_viewer.ipynb), which takes the same\n", + "polygon and the same stores to a rotatable paired 3-D view and numpy tensors.\n", + "The two are separate notebooks on purpose: that one needs `%matplotlib widget`\n", + "for the 3-D view, this one needs `%matplotlib inline`, and the two backends\n", + "collide in a single kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "acae54e37e7d407bbb7b55eff062a284", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install -q mortie \"moczarr[zagg]>=0.7\" matplotlib ipywidgets\n", + "%matplotlib inline\n", + "\n", + "\n", + "import moczarr as mz\n", + "from mortie import moc\n", + "\n", + "# The drawing lives in viewers.py beside this notebook, so the cells below stay\n", + "# about the READ path. It also holds the one zagg import: moczarr imports the\n", + "# t-digest algebra rather than vendoring it (moczarr issue #19), which is what\n", + "# the `moczarr[zagg]` extra carries.\n", + "from viewers import densest_shard, paired_blocks, waveform_view\n", + "\n", + "STORES = {\n", + " \"atl03\": (\n", + " \"s3://us-west-2.opendata.source.coop/englacial/zagg/demo/atl03_tdigest_o9.zarr\",\n", + " \"19/h_tdigest_signal\",\n", + " ),\n", + " \"gedi\": (\n", + " \"s3://us-west-2.opendata.source.coop/englacial/zagg/demo/gedi_flux_o9.zarr\",\n", + " \"18/rx_flux\",\n", + " ),\n", + "}\n", + "S3 = {\"region\": \"us-west-2\", \"anonymous\": True}" + ] + }, + { + "cell_type": "markdown", + "id": "9a63283cbaf04dbcab1f6479b197f3a8", + "metadata": {}, + "source": [ + "## One polygon in, covered shards out\n", + "\n", + "The polygon below sits on SERC, the mid-Atlantic forest site the HHDC\n", + "diffusion papers are built on. Replace it with any area within California or a\n", + "NEON AOP site \u2014 it need not be a box \u2014 and the cell tests whether each store\n", + "actually covers what you asked for. An AOI in Alaska will pass\n", + "the ATL03 check and fail the GEDI one \u2014 GEDI flies on the ISS, so it sees no\n", + "higher than |lat| 51.6 and the Alaska NEON sites are outside its reach." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8dd0d8092fe74a7c96281538738b07e2", + "metadata": {}, + "outputs": [], + "source": [ + "aoi = {\n", + " \"features\": [\n", + " {\n", + " \"geometry\": {\n", + " \"coordinates\": [\n", + " [\n", + " [-76.5750, 38.9000],\n", + " [-76.5450, 38.9000],\n", + " [-76.5450, 38.8780],\n", + " [-76.5600, 38.8720],\n", + " [-76.5750, 38.8800],\n", + " [-76.5750, 38.9000],\n", + " ]\n", + " ]\n", + " }\n", + " }\n", + " ]\n", + "}\n", + "\n", + "q = moc(aoi)\n", + "shards = None\n", + "for name, (root, _field) in STORES.items():\n", + " assert mz.coverage_moc(root, **S3).contains(q), f\"{name} does not contain the polygon\"\n", + " ids = set(mz.candidate_shards(root, aoi=q, **S3))\n", + " shards = ids if shards is None else shards & ids\n", + "shards = sorted(shards)\n", + "\n", + "# Most GEDI granules -- same rule as hhdc_viewer and demo/06_paired.\n", + "SHARD = densest_shard(STORES[\"gedi\"][0], shards, **S3)\n", + "print(f\"{len(shards)} shards cover the polygon: {shards}\\nworking {SHARD}\")" + ] + }, + { + "cell_type": "markdown", + "id": "72eea5119410473aa328ad9291626812", + "metadata": {}, + "source": [ + "## Paired tensors \u2014 one shard, both sensors\n", + "\n", + "`read_tensors` yields one `(tensor, mask, (offset, gain), block)` per populated\n", + "o12 block. Reading both sensors on the same z grid is what makes the two\n", + "comparable; the blocks they share are the candidates for a join." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8edb47106e1a46a883d545849b8ab81b", + "metadata": {}, + "outputs": [], + "source": [ + "handles = {n: (mz.open_leaf(root, SHARD, **S3), field) for n, (root, field) in STORES.items()}\n", + "gains, pairs = paired_blocks(handles)\n" + ] + }, + { + "cell_type": "markdown", + "id": "10185d26023b46108eb7d9f57d49d2b3", + "metadata": {}, + "source": [ + "## Coincident waveforms\n", + "\n", + "A GEDI footprint cell (o18) covers exactly four ATL03 cells (o19). This reads all\n", + "four and merges their centroids (observations) into a single digest (waveform /\n", + "pdf), then compares that against GEDI's \u2014 so both sides describe the same patch of\n", + "ground at the same scale.\n", + "\n", + "A centroid is one `(elevation, weight)` row. These stores are written with a\n", + "centroid budget far larger than any single cell fills, so nothing is merged away\n", + "at this scale: a centroid is one photon for ATL03 and one above-threshold\n", + "waveform sample for GEDI, and combining the four cells is a plain union of their\n", + "rows \u2014 nothing re-binned, re-compressed, or averaged.\n", + "\n", + "Both sensors are read **straight from their stored digests** \u2014 no tensor in the\n", + "loop. The tensors above only decide which cells are worth comparing.\n", + "\n", + "The slider orders those cells by joint signal strength, strongest first.\n", + "*Coincident* means the joint mask: both sensors present. *Strongest* needs a\n", + "little care. ATL03 counts photons; for GEDI the weight is an **estimate** of\n", + "detected photoelectrons \u2014 background-subtracted waveform counts scaled by a\n", + "named, versioned receiver gain (the published chain gives about 0.4 pe/count:\n", + "Sun et al. 2020, NASA NTRS\n", + "[20200001052](https://ntrs.nasa.gov/citations/20200001052); this store ships a\n", + "unit-gain placeholder, so its weights are *proportional* to photoelectrons rather\n", + "than calibrated to them). The two are not commensurate \u2014 GEDI's run two to three\n", + "orders of magnitude larger \u2014 so taking the smaller of the raw numbers would\n", + "always return ATL03's and rank by it alone. Each sensor is first divided by its\n", + "own maximum across the block's coincident cells, putting both on 0-1; the score\n", + "is the smaller of the two. A cell ranks high only when it is well populated for\n", + "**both** instruments." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8763a12b2bbd4a93a75aff182afb95dc", + "metadata": {}, + "outputs": [], + "source": [ + "stores = {n: store for n, (store, _f) in handles.items()}\n", + "fields = {n: field for n, (_s, field) in handles.items()}\n", + "waveform_view(stores, fields, gains, pairs, SHARD)\n" + ] + }, + { + "cell_type": "markdown", + "id": "7623eae2785240b9bd12b16a66d81610", + "metadata": {}, + "source": [ + "Two libraries, one polygon \u2014 coverage, shards, paired tensors, and the\n", + "cell-level waveform join, anonymously against public S3." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}