From fb8ef189cd764b03bb2adc65029d5849cda509e6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 19:18:57 +0000 Subject: [PATCH 01/16] phase 1 of issue #30 --- src/zagg/configs/atl03.yaml | 61 +++++++++++++++++++++++++++++++++++++ src/zagg/processing.py | 38 ++++++++++++++++++++++- tests/test_config.py | 36 ++++++++++++++++++++++ tests/test_processing.py | 21 +++++++++++++ 4 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 src/zagg/configs/atl03.yaml diff --git a/src/zagg/configs/atl03.yaml b/src/zagg/configs/atl03.yaml new file mode 100644 index 00000000..9748c80d --- /dev/null +++ b/src/zagg/configs/atl03.yaml @@ -0,0 +1,61 @@ +# ATL03 photon-height aggregation template. +# +# Scalar-only stats (non-scalar outputs are deferred to #29): per-cell density +# (photon count), min/max elevation, mean and median height, and variance. +# +# Confidence filter keeps everything except the noise flag: signal_conf_ph != 0 +# (retains flags 1-4), per the ATL03 v3 data dictionary +# (https://nsidc.org/sites/default/files/icesat2_atl03_data_dict_v003_0.pdf). +# +# signal_conf_ph is per-surface-type in ATL03; this template reads the land +# surface_type column (index 0). Grid is rectilinear (HEALPix order-19, the +# ~10 m match, waits on mortie #35). +data_source: + reader: h5coro + driver: s3 + groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r] + coordinates: + latitude: "/{group}/heights/lat_ph" + longitude: "/{group}/heights/lon_ph" + variables: + h_ph: "/{group}/heights/h_ph" + quality_filter: + dataset: "/{group}/heights/signal_conf_ph" + value: 0 + op: ne + +aggregation: + variables: + count: + function: len + source: h_ph + dtype: int32 + fill_value: 0 + h_min: + function: min + source: h_ph + dtype: float32 + h_max: + function: max + source: h_ph + dtype: float32 + h_mean: + function: mean + source: h_ph + dtype: float32 + h_median: + function: median + source: h_ph + dtype: float32 + h_variance: + function: var + source: h_ph + dtype: float32 + +output: + grid: + type: rectilinear + crs: EPSG:4326 + resolution: 0.0001 # degrees (~10 m at the equator) + bounds: [-180, -90, 180, 90] + chunk_shape: [256, 256] diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 49aec1d9..e6f89fe2 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -486,6 +486,42 @@ def _kernel_aggregate( # -- end EXPERIMENTAL kernel path --------------------------------------------- +def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: + """Build a keep-mask from a quality-flag array and a ``quality_filter`` spec. + + The comparison is controlled by ``quality_filter["op"]`` (default ``"eq"``): + + - ``"eq"`` keeps rows where the flag equals ``value`` (ATL06: keep + ``atl06_quality_summary == 0``). + - ``"ne"`` keeps rows where the flag differs from ``value`` (ATL03: keep + ``signal_conf_ph != 0``, dropping only the noise flag). + + Parameters + ---------- + q_flag : np.ndarray + Quality-flag values for the candidate rows. + quality_filter : dict + Filter spec with ``value`` and optional ``op``. + + Returns + ------- + np.ndarray + Boolean keep-mask, same length as ``q_flag``. + + Raises + ------ + ValueError + If ``op`` is not ``"eq"`` or ``"ne"``. + """ + value = quality_filter["value"] + op = quality_filter.get("op", "eq") + if op == "eq": + return q_flag == value + if op == "ne": + return q_flag != value + raise ValueError(f"Unknown quality_filter op '{op}' (expected 'eq' or 'ne')") + + def _read_group( h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False ): @@ -542,7 +578,7 @@ def _read_group( if quality_filter is not None: qf_path = quality_filter["dataset"].format(group=group) q_flag = data[qf_path][mask_sliced] - quality_mask = q_flag == quality_filter["value"] + quality_mask = _quality_mask(q_flag, quality_filter) if np.sum(quality_mask) == 0: return None else: diff --git a/tests/test_config.py b/tests/test_config.py index 3a5122e1..1c94dccc 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -88,6 +88,42 @@ def test_nonexistent_raises(self): default_config("nonexistent") +# --------------------------------------------------------------------------- +# ATL03 template +# --------------------------------------------------------------------------- + +class TestATL03Template: + @pytest.fixture + def atl03_config(self): + return default_config("atl03") + + def test_loads_and_validates(self, atl03_config): + # default_config already runs validate_config; assert it round-trips. + validate_config(atl03_config) + assert atl03_config.data_source["reader"] == "h5coro" + assert len(atl03_config.data_source["groups"]) == 6 + + def test_scalar_variables(self, atl03_config): + dvars = set(get_data_vars(atl03_config)) + assert dvars == {"count", "h_min", "h_max", "h_mean", "h_median", "h_variance"} + + def test_functions_resolve(self, atl03_config): + for meta in get_agg_fields(atl03_config).values(): + assert "expression" not in meta # scalar-only; non-scalar is #29 + resolve_function(meta["function"]) # raises on failure + + def test_confidence_filter_drops_noise(self, atl03_config): + qf = atl03_config.data_source["quality_filter"] + assert qf["value"] == 0 + assert qf["op"] == "ne" # keep signal_conf_ph != 0 (flags 1-4) + assert qf["dataset"].endswith("signal_conf_ph") + + def test_rectilinear_grid(self, atl03_config): + grid = atl03_config.output["grid"] + assert grid["type"] == "rectilinear" + assert len(grid["bounds"]) == 4 + + # --------------------------------------------------------------------------- # Function resolution # --------------------------------------------------------------------------- diff --git a/tests/test_processing.py b/tests/test_processing.py index 88fe881d..34625ae1 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -13,6 +13,7 @@ _group_columns, _kernel_able, _kernel_aggregate, + _quality_mask, calculate_cell_statistics, process_shard, write_dataframe_to_zarr, @@ -645,3 +646,23 @@ def test_group_template_substitution(self): ds = default_config().data_source path = ds["coordinates"]["latitude"].format(group="gt2r") assert path == "/gt2r/land_ice_segments/latitude" + + +class TestQualityMask: + """``_quality_mask`` builds the keep-mask for the configured comparison op.""" + + def test_eq_default_keeps_equal(self): + # ATL06 semantics: keep atl06_quality_summary == 0 (op defaults to eq). + q_flag = np.array([0, 1, 0, 2], dtype=np.int8) + mask = _quality_mask(q_flag, {"value": 0}) + np.testing.assert_array_equal(mask, [True, False, True, False]) + + def test_ne_keeps_nonmatching(self): + # ATL03 semantics: keep signal_conf_ph != 0 (drop only the noise flag). + q_flag = np.array([0, 1, 4, 0], dtype=np.int8) + mask = _quality_mask(q_flag, {"value": 0, "op": "ne"}) + np.testing.assert_array_equal(mask, [False, True, True, False]) + + def test_unknown_op_raises(self): + with pytest.raises(ValueError, match="Unknown quality_filter op"): + _quality_mask(np.array([0]), {"value": 0, "op": "gt"}) From ca52b907e49f038a112ecdfd93ab170893441f56 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 19:24:04 +0000 Subject: [PATCH 02/16] fold review: 2-D signal_conf_ph column selection (#30) --- src/zagg/configs/atl03.yaml | 7 ++++--- src/zagg/processing.py | 28 +++++++++++++++++++++++----- tests/test_config.py | 1 + tests/test_processing.py | 24 ++++++++++++++++++++++++ 4 files changed, 52 insertions(+), 8 deletions(-) diff --git a/src/zagg/configs/atl03.yaml b/src/zagg/configs/atl03.yaml index 9748c80d..78765e58 100644 --- a/src/zagg/configs/atl03.yaml +++ b/src/zagg/configs/atl03.yaml @@ -7,9 +7,9 @@ # (retains flags 1-4), per the ATL03 v3 data dictionary # (https://nsidc.org/sites/default/files/icesat2_atl03_data_dict_v003_0.pdf). # -# signal_conf_ph is per-surface-type in ATL03; this template reads the land -# surface_type column (index 0). Grid is rectilinear (HEALPix order-19, the -# ~10 m match, waits on mortie #35). +# signal_conf_ph is 2-D in ATL03 (n_photons x 5 surface types); this template +# selects the land column (index 0) via quality_filter.column. Grid is +# rectilinear (HEALPix order-19, the ~10 m match, waits on mortie #35). data_source: reader: h5coro driver: s3 @@ -23,6 +23,7 @@ data_source: dataset: "/{group}/heights/signal_conf_ph" value: 0 op: ne + column: 0 # land surface_type (signal_conf_ph is 2-D) aggregation: variables: diff --git a/src/zagg/processing.py b/src/zagg/processing.py index e6f89fe2..82c09fff 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -487,7 +487,7 @@ def _kernel_aggregate( def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: - """Build a keep-mask from a quality-flag array and a ``quality_filter`` spec. + """Build a 1-D keep-mask from a quality-flag array and a ``quality_filter`` spec. The comparison is controlled by ``quality_filter["op"]`` (default ``"eq"``): @@ -496,23 +496,41 @@ def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: - ``"ne"`` keeps rows where the flag differs from ``value`` (ATL03: keep ``signal_conf_ph != 0``, dropping only the noise flag). + When ``q_flag`` is 2-D (e.g. ATL03 ``signal_conf_ph`` is ``(n_photons, + n_surface_types)``), ``quality_filter["column"]`` selects the surface-type + column to compare (e.g. ``0`` for land); it is required in that case and + must not be set for a 1-D flag. + Parameters ---------- q_flag : np.ndarray - Quality-flag values for the candidate rows. + Quality-flag values for the candidate rows (1-D, or 2-D with a + ``column`` selector). quality_filter : dict - Filter spec with ``value`` and optional ``op``. + Filter spec with ``value``, optional ``op``, and optional ``column``. Returns ------- np.ndarray - Boolean keep-mask, same length as ``q_flag``. + Boolean keep-mask, length ``q_flag.shape[0]``. Raises ------ ValueError - If ``op`` is not ``"eq"`` or ``"ne"``. + If ``op`` is not ``"eq"``/``"ne"``, or if the flag dimensionality and + the ``column`` selector disagree. """ + column = quality_filter.get("column") + if q_flag.ndim == 2: + if column is None: + raise ValueError( + "quality_filter on a 2-D flag requires a 'column' index " + f"(got shape {q_flag.shape})" + ) + q_flag = q_flag[:, column] + elif column is not None: + raise ValueError("quality_filter 'column' is only valid for a 2-D flag") + value = quality_filter["value"] op = quality_filter.get("op", "eq") if op == "eq": diff --git a/tests/test_config.py b/tests/test_config.py index 1c94dccc..a4b41712 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -116,6 +116,7 @@ def test_confidence_filter_drops_noise(self, atl03_config): qf = atl03_config.data_source["quality_filter"] assert qf["value"] == 0 assert qf["op"] == "ne" # keep signal_conf_ph != 0 (flags 1-4) + assert qf["column"] == 0 # signal_conf_ph is 2-D; select land column assert qf["dataset"].endswith("signal_conf_ph") def test_rectilinear_grid(self, atl03_config): diff --git a/tests/test_processing.py b/tests/test_processing.py index 34625ae1..c6fcb828 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -666,3 +666,27 @@ def test_ne_keeps_nonmatching(self): def test_unknown_op_raises(self): with pytest.raises(ValueError, match="Unknown quality_filter op"): _quality_mask(np.array([0]), {"value": 0, "op": "gt"}) + + def test_2d_column_selects_surface_type(self): + # ATL03 signal_conf_ph is (n_photons, n_surface_types); column picks one. + q_flag = np.array([[0, 2], [3, 0], [0, 4]], dtype=np.int8) + mask = _quality_mask(q_flag, {"value": 0, "op": "ne", "column": 0}) + np.testing.assert_array_equal(mask, [False, True, False]) + mask_col1 = _quality_mask(q_flag, {"value": 0, "op": "ne", "column": 1}) + np.testing.assert_array_equal(mask_col1, [True, False, True]) + + def test_2d_without_column_raises(self): + q_flag = np.zeros((3, 5), dtype=np.int8) + with pytest.raises(ValueError, match="requires a 'column'"): + _quality_mask(q_flag, {"value": 0, "op": "ne"}) + + def test_column_on_1d_raises(self): + with pytest.raises(ValueError, match="only valid for a 2-D flag"): + _quality_mask(np.array([0, 1]), {"value": 0, "column": 0}) + + def test_atl03_template_filter_runs(self): + # The shipped template's quality_filter must apply cleanly to a 2-D flag. + qf = default_config("atl03").data_source["quality_filter"] + q_flag = np.array([[0, 1], [1, 0], [4, 0]], dtype=np.int8) + mask = _quality_mask(q_flag, qf) + np.testing.assert_array_equal(mask, [False, True, True]) From 3893d2d83028c6b2dd973fec8d4545af8ebbdba1 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 19:25:02 +0000 Subject: [PATCH 03/16] fold review: validate quality_filter op (#30) --- src/zagg/config.py | 7 +++++++ tests/test_config.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/zagg/config.py b/src/zagg/config.py index 9acd02a5..7efaf440 100644 --- a/src/zagg/config.py +++ b/src/zagg/config.py @@ -169,6 +169,13 @@ def validate_config(config: PipelineConfig) -> None: if "start_date" not in temporal or "end_date" not in temporal: raise ValueError("bounds.temporal requires start_date and end_date") + # Validate quality_filter op (default eq; eq/ne supported by _quality_mask) + qf = config.data_source.get("quality_filter") + if qf is not None and qf.get("op", "eq") not in ("eq", "ne"): + raise ValueError( + f"data_source.quality_filter.op must be 'eq' or 'ne' (got {qf['op']!r})" + ) + ds_vars = set(config.data_source.get("variables", {}).keys()) agg_vars = config.aggregation.get("variables", {}) diff --git a/tests/test_config.py b/tests/test_config.py index a4b41712..c52054d8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -202,6 +202,20 @@ def test_missing_source(self): with pytest.raises(ValueError, match="source.*nonexistent"): validate_config(cfg) + def test_bad_quality_filter_op(self): + cfg = PipelineConfig( + data_source={ + "variables": {"h_li": "/path"}, + "quality_filter": {"dataset": "/q", "value": 0, "op": "gt"}, + }, + aggregation={"variables": { + "c": {"function": "len", "source": "h_li", "dtype": "int32"}, + }}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="quality_filter.op must be"): + validate_config(cfg) + def test_missing_weights_column(self): cfg = PipelineConfig( data_source={"variables": {"h_li": "/path"}}, From 9956c0ab2b0e5204c24ffa651a8af1ed8d466107 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 19:31:25 +0000 Subject: [PATCH 04/16] phase 2 of issue #30: region timing benchmark --- benchmarks/region_timing.py | 316 ++++++++++++++++++++++++++++++++++++ tests/test_runner.py | 38 +++++ 2 files changed, 354 insertions(+) create mode 100644 benchmarks/region_timing.py diff --git a/benchmarks/region_timing.py b/benchmarks/region_timing.py new file mode 100644 index 00000000..b26dd054 --- /dev/null +++ b/benchmarks/region_timing.py @@ -0,0 +1,316 @@ +"""Real-data region timing benchmark for the ATL03 aggregation handoff (issue #30). + +This is the real-data (phase-3) counterpart to the synthetic ``handoff_bench.py``: +it drives the actual catalog -> shard-map -> ``runner.agg()`` path against +ICESat-2 ATL03 granules over three 10 km x 10 km AOIs, swept across a set of +time windows, and records wall-time, peak RSS, observation/cell counts, and +output bytes for the ``pandas`` and ``arrow`` handoff carriers -- asserting their +scalar outputs are byte-for-byte identical (#30's parity criterion on real data). + +Three regions span the density regimes the synthetic harness can't reproduce: + + * ``neon_maryland`` -- NEON site, land/vegetation (moderate photon return) + * ``russell_glacier`` -- Kangerlussuaq / Russell Glacier, on the ice (high return) + * ``bahamas`` -- shallow-water bathymetry (sparse, attenuated subsurface) + +The optional ``--hard`` flag adds an ``antarctica_88s`` AOI at 88 deg S, where RGT +convergence maximizes per-cell overlap (the worst case for the old mask loop). It +is deferred / behind the flag by default (per @espg) because it is far more +expensive than the other three. + +Confidence filter and grid come straight from the shipped ``atl03`` template: +``signal_conf_ph != 0`` (keep flags 1-4, drop only noise), rectilinear grid. The +``output.grid.bounds`` are overridden per region to the AOI's 10 km box so the +grid covers just the patch. (HEALPix order-19 -- the ~10 m match -- waits on +mortie #35, so this first pass is rectilinear only.) + +**This script is NOT run in CI**: it needs ``earthaccess``/NSIDC-S3 credentials +(CMR-STAC query + byte-range HDF5 reads) and is slow. It lives under +``benchmarks/`` (not ``tests/``) and is meant for a credentialed session. It must +import and construct cleanly and stay ruff-clean regardless. + +Run (in a credentialed session):: + + uv run python benchmarks/region_timing.py --windows 1d,5d --max-cells 4 + uv run python benchmarks/region_timing.py --hard --out results.txt +""" + +import argparse +import copy +import resource +import sys +import tempfile +import time +from dataclasses import dataclass +from datetime import date, timedelta + +import numpy as np + +from zagg.config import default_config, get_data_vars +from zagg.runner import agg + +# Time windows (per @espg). Each is the span ending at ``--end-date`` that drives +# the catalog query; longer windows pull more granules -> more overlap. +WINDOWS = { + "1d": 1, + "5d": 5, + "15d": 15, + "30d": 30, + "91d": 91, + "1y": 365, +} + +# Carriers compared head-to-head. Both must produce identical scalar outputs. +HANDOFFS = ("pandas", "arrow") + + +@dataclass(frozen=True) +class Region: + """A 10 km x 10 km AOI, as a ``(lon_min, lat_min, lon_max, lat_max)`` bbox.""" + + name: str + lon_min: float + lat_min: float + lon_max: float + lat_max: float + + @property + def bbox(self) -> tuple[float, float, float, float]: + return (self.lon_min, self.lat_min, self.lon_max, self.lat_max) + + +def _box(name: str, lon: float, lat: float) -> Region: + """A ~10 km x 10 km box centered on (lon, lat). + + 0.09 deg of latitude is ~10 km; the longitude span is scaled by + ``1/cos(lat)`` so the east-west extent stays ~10 km at high latitude. + """ + half_lat = 0.045 + half_lon = 0.045 / max(np.cos(np.radians(lat)), 1e-3) + return Region(name, lon - half_lon, lat - half_lat, lon + half_lon, lat + half_lat) + + +# Region centers. Maryland NEON (NEON SCBI/SERC area), Russell Glacier on the ice +# east of Kangerlussuaq, and a Bahamas bank for bathymetry. +REGIONS = [ + _box("neon_maryland", lon=-76.56, lat=38.89), + _box("russell_glacier", lon=-50.0, lat=67.09), + _box("bahamas", lon=-76.0, lat=24.0), +] + +# Deferred / behind --hard: 88 deg S, where RGT convergence maximizes overlap. +HARD_REGION = _box("antarctica_88s", lon=0.0, lat=-88.0) + + +@dataclass +class Record: + """One (region x window x handoff) measurement.""" + + region: str + window: str + handoff: str + wall_s: float + peak_rss_mb: float + total_obs: int + cells_with_data: int + output_bytes: int + + +def _peak_rss_mb() -> float: + """Process peak resident set size in MB (ru_maxrss is KB on Linux, bytes on macOS).""" + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return rss / 1e3 if sys.platform != "darwin" else rss / 1e6 + + +def _region_config(region: Region): + """An ``atl03`` config whose rectilinear grid is clipped to the region bbox.""" + cfg = default_config("atl03") + cfg = copy.deepcopy(cfg) + cfg.output["grid"]["bounds"] = list(region.bbox) + return cfg + + +def _store_scalars(store_path: str, cfg) -> dict[str, np.ndarray]: + """Read each aggregated data variable out of a written store as a dense array. + + Data is written under the grid's ``group_path`` (see + ``write_dataframe_to_zarr``), so resolve it from the config's grid. + """ + import zarr + + from zagg.grids import from_config + + group_path = from_config(cfg).group_path + grp = zarr.open_group(store_path, mode="r") + return {name: np.asarray(grp[f"{group_path}/{name}"][:]) for name in get_data_vars(cfg)} + + +def _assert_parity(a: dict[str, np.ndarray], b: dict[str, np.ndarray], ctx: str) -> None: + """Assert two stores' scalar outputs are byte-for-byte identical (NaN-aware).""" + assert a.keys() == b.keys(), f"{ctx}: variable mismatch {a.keys()} vs {b.keys()}" + for name in a: + x, y = a[name], b[name] + if not np.array_equal(x, y, equal_nan=np.issubdtype(x.dtype, np.floating)): + raise AssertionError(f"{ctx}: '{name}' differs between pandas and arrow") + + +def _output_bytes(store_path: str) -> int: + """Total bytes written under a local store directory.""" + import os + + total = 0 + for root, _dirs, files in os.walk(store_path): + for f in files: + total += os.path.getsize(os.path.join(root, f)) + return total + + +def run_one( + region: Region, + window: str, + days: int, + *, + end_date: date, + version: str, + max_cells: int | None, + max_workers: int | None, +) -> list[Record]: + """Build the catalog for (region, window), run both carriers, assert parity.""" + from zagg.catalog import make_shardmap + from zagg.catalog.sources import Query + from zagg.grids import from_config + + cfg = _region_config(region) + start = (end_date - timedelta(days=days)).isoformat() + query = Query("ATL03", version, start, end_date.isoformat(), region=region.bbox) + + with tempfile.TemporaryDirectory() as tmp: + catalog_path = f"{tmp}/shardmap.json" + grid = from_config(cfg) + make_shardmap(query, grid).to_json(catalog_path) + + records: list[Record] = [] + scalars: dict[str, dict] = {} + for handoff in HANDOFFS: + store_path = f"{tmp}/out_{handoff}.zarr" + t0 = time.perf_counter() + summary = agg( + cfg, + catalog=catalog_path, + store=store_path, + handoff=handoff, + max_cells=max_cells, + max_workers=max_workers, + overwrite=True, + ) + wall = time.perf_counter() - t0 + scalars[handoff] = _store_scalars(store_path, cfg) + records.append( + Record( + region=region.name, + window=window, + handoff=handoff, + wall_s=wall, + peak_rss_mb=_peak_rss_mb(), + total_obs=int(summary.get("total_obs", 0)), + cells_with_data=int(summary.get("cells_with_data", 0)), + output_bytes=_output_bytes(store_path), + ) + ) + _assert_parity( + scalars["pandas"], + scalars["arrow"], + ctx=f"{region.name}/{window}", + ) + return records + + +def format_table(records: list[Record]) -> str: + """Render the collected records as a fixed-width text table.""" + header = ( + f"{'region':<18}{'window':>8}{'handoff':>10}{'wall_s':>10}" + f"{'rss_MB':>10}{'obs':>12}{'cells':>10}{'out_MB':>10}" + ) + lines = [header, "-" * len(header)] + for r in records: + lines.append( + f"{r.region:<18}{r.window:>8}{r.handoff:>10}{r.wall_s:>10.3f}" + f"{r.peak_rss_mb:>10.1f}{r.total_obs:>12,}{r.cells_with_data:>10,}" + f"{r.output_bytes / 1e6:>10.2f}" + ) + return "\n".join(lines) + + +def parse_args(argv=None): + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument( + "--windows", + default=",".join(WINDOWS), + help="Comma-separated time windows (default: all). Choices: " + ", ".join(WINDOWS), + ) + ap.add_argument( + "--end-date", + default=date.today().isoformat(), + help="End of each time window, YYYY-MM-DD (default: today).", + ) + ap.add_argument( + "--hard", + action="store_true", + help="Also run the deferred 88 deg S max-overlap AOI (slow).", + ) + ap.add_argument( + "--max-cells", + type=int, + default=None, + help="Cap cells processed per run (for a quick smoke).", + ) + ap.add_argument("--max-workers", type=int, default=None, help="Local worker cap.") + ap.add_argument("--version", default="007", help="ATL03 product version (default: 007).") + ap.add_argument( + "--out", + default=None, + help="Also append the results table to this file.", + ) + return ap.parse_args(argv) + + +def main(argv=None): + args = parse_args(argv) + + windows = [] + for w in args.windows.split(","): + w = w.strip() + if w not in WINDOWS: + raise SystemExit(f"unknown window {w!r}; choices: {', '.join(WINDOWS)}") + windows.append(w) + + end_date = date.fromisoformat(args.end_date) + regions = list(REGIONS) + ([HARD_REGION] if args.hard else []) + + records: list[Record] = [] + for region in regions: + for window in windows: + records.extend( + run_one( + region, + window, + WINDOWS[window], + end_date=end_date, + version=args.version, + max_cells=args.max_cells, + max_workers=args.max_workers, + ) + ) + + table = format_table(records) + print(table) + if args.out: + with open(args.out, "a") as f: + f.write(table + "\n") + + +if __name__ == "__main__": + main() diff --git a/tests/test_runner.py b/tests/test_runner.py index ed46d543..2866b5bc 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -281,3 +281,41 @@ def test_non_healpix_event_omits_child_order(self): assert event["shard_key"] == 12345 assert "child_order" not in event assert "parent_morton" not in event + + +class TestHandoffPassthrough: + """`agg(handoff=...)` threads the carrier choice down to process_shard.""" + + def test_process_and_write_forwards_handoff(self, monkeypatch, atl06_config): + from zagg import runner + + captured = {} + + def fake_process_shard(grid, shard_key, urls, **kwargs): + import pandas as pd + captured["handoff"] = kwargs.get("handoff") + return pd.DataFrame(), {"shard_key": shard_key, "error": None} + + monkeypatch.setattr(runner, "process_shard", fake_process_shard) + runner._process_and_write( + 0, (0,), [_rec(1)], grid=None, s3_creds={}, zarr_store=None, + config=atl06_config, driver="s3", handoff="arrow", + ) + assert captured["handoff"] == "arrow" + + def test_default_handoff_is_pandas(self, monkeypatch, atl06_config): + from zagg import runner + + captured = {} + + def fake_process_shard(grid, shard_key, urls, **kwargs): + import pandas as pd + captured["handoff"] = kwargs.get("handoff") + return pd.DataFrame(), {"shard_key": shard_key, "error": None} + + monkeypatch.setattr(runner, "process_shard", fake_process_shard) + runner._process_and_write( + 0, (0,), [_rec(1)], grid=None, s3_creds={}, zarr_store=None, + config=atl06_config, driver="s3", + ) + assert captured["handoff"] == "pandas" From 4427b7fee6cebdfa5c9c11d6657c1ef91127c2c7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 19:32:54 +0000 Subject: [PATCH 05/16] phase 2: thread handoff through agg/runner --- src/zagg/runner.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/zagg/runner.py b/src/zagg/runner.py index 7d56cbc2..73b79c11 100644 --- a/src/zagg/runner.py +++ b/src/zagg/runner.py @@ -69,6 +69,7 @@ def agg( region: str = "us-west-2", output_credentials: dict | None = None, output_endpoint_url: str | None = None, + handoff: str = "pandas", ) -> dict: """Run the aggregation pipeline. @@ -112,6 +113,10 @@ def agg( output_endpoint_url : str, optional Custom S3-compatible endpoint for the output store (e.g. R2, MinIO). Overrides ``output.endpoint_url`` in the config. + handoff : str + Per-cell aggregation carrier: ``"pandas"`` (default) or ``"arrow"``. + Both produce byte-for-byte identical scalar outputs (#30); ``"arrow"`` + is opt-in for benchmarking. Only honored by the ``"local"`` backend. Returns ------- @@ -158,6 +163,7 @@ def agg( dry_run=dry_run, region=region, driver=resolved_driver, output_credentials=output_credentials, output_endpoint_url=resolved_endpoint, + handoff=handoff, ) elif backend == "lambda": if max_workers is None: @@ -276,7 +282,8 @@ def _check_signature(grid, catalog_data: dict) -> None: def _process_and_write(shard_key, chunk_idx, records, grid, - s3_creds, zarr_store, config, driver=None): + s3_creds, zarr_store, config, driver=None, + handoff="pandas"): """Process a single shard and write results to store.""" df_out, metadata = process_shard( grid, @@ -285,6 +292,7 @@ def _process_and_write(shard_key, chunk_idx, records, grid, s3_credentials=s3_creds, config=config, driver=driver, + handoff=handoff, ) if not df_out.empty: write_dataframe_to_zarr( @@ -297,7 +305,8 @@ def _process_and_write(shard_key, chunk_idx, records, grid, def _run_local(config, catalog_data, store_path, child_order, *, max_cells, morton_cell, max_workers, overwrite, dry_run, region, - driver="s3", output_credentials=None, output_endpoint_url=None): + driver="s3", output_credentials=None, output_endpoint_url=None, + handoff="pandas"): """Run processing locally with ThreadPoolExecutor.""" all_shards = list(catalog_data["shard_keys"]) @@ -343,7 +352,7 @@ def _run_local(config, catalog_data, store_path, child_order, *, shard_key, grid.block_index(int(shard_key)), records, grid, s3_creds, zarr_store, config, - driver=driver, + driver=driver, handoff=handoff, ): shard_key for shard_key, records in cells } From cae77fa2555df45ae02d12018add46b0795c2869 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 17 Jun 2026 01:13:08 +0000 Subject: [PATCH 06/16] address review: signal_conf TEP filter across surface types (#30) --- src/zagg/configs/atl03.yaml | 15 +++++++------ src/zagg/processing.py | 43 +++++++++++++++++++++---------------- tests/test_config.py | 9 ++++---- tests/test_processing.py | 24 +++++++++++++++------ 4 files changed, 56 insertions(+), 35 deletions(-) diff --git a/src/zagg/configs/atl03.yaml b/src/zagg/configs/atl03.yaml index 78765e58..83187917 100644 --- a/src/zagg/configs/atl03.yaml +++ b/src/zagg/configs/atl03.yaml @@ -3,13 +3,15 @@ # Scalar-only stats (non-scalar outputs are deferred to #29): per-cell density # (photon count), min/max elevation, mean and median height, and variance. # -# Confidence filter keeps everything except the noise flag: signal_conf_ph != 0 -# (retains flags 1-4), per the ATL03 v3 data dictionary +# Confidence filter drops only TEP (transmit-echo path) photons and keeps +# everything else: signal_conf_ph != -2, per the ATL03 v3 data dictionary # (https://nsidc.org/sites/default/files/icesat2_atl03_data_dict_v003_0.pdf). # -# signal_conf_ph is 2-D in ATL03 (n_photons x 5 surface types); this template -# selects the land column (index 0) via quality_filter.column. Grid is -# rectilinear (HEALPix order-19, the ~10 m match, waits on mortie #35). +# signal_conf_ph is 2-D in ATL03 (n_photons x 5 surface types). TEP photons are +# flagged -2 across every surface type, so with no column selector the != -2 +# filter is reduced across all columns (a photon is dropped only when it is TEP +# for all surface types). Grid is rectilinear (HEALPix order-19, the ~10 m +# match, waits on mortie #35). data_source: reader: h5coro driver: s3 @@ -21,9 +23,8 @@ data_source: h_ph: "/{group}/heights/h_ph" quality_filter: dataset: "/{group}/heights/signal_conf_ph" - value: 0 + value: -2 # TEP flag; reduced across all surface types op: ne - column: 0 # land surface_type (signal_conf_ph is 2-D) aggregation: variables: diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 82c09fff..9b4d1f67 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -497,15 +497,19 @@ def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: ``signal_conf_ph != 0``, dropping only the noise flag). When ``q_flag`` is 2-D (e.g. ATL03 ``signal_conf_ph`` is ``(n_photons, - n_surface_types)``), ``quality_filter["column"]`` selects the surface-type - column to compare (e.g. ``0`` for land); it is required in that case and - must not be set for a 1-D flag. + n_surface_types)``), the comparison is reduced **across all surface-type + columns** rather than keying on one column: a photon's flag matches a given + surface type if any column matches. This is what the ATL03 template needs to + drop TEP photons (``signal_conf_ph == -2``, set across every surface type) + while keeping everything else. Optionally ``quality_filter["column"]`` + restricts the comparison to a single surface-type column (e.g. ``0`` for + land); it must not be set for a 1-D flag. Parameters ---------- q_flag : np.ndarray - Quality-flag values for the candidate rows (1-D, or 2-D with a - ``column`` selector). + Quality-flag values for the candidate rows (1-D, or 2-D reduced across + surface-type columns). quality_filter : dict Filter spec with ``value``, optional ``op``, and optional ``column``. @@ -517,27 +521,30 @@ def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: Raises ------ ValueError - If ``op`` is not ``"eq"``/``"ne"``, or if the flag dimensionality and - the ``column`` selector disagree. + If ``op`` is not ``"eq"``/``"ne"``, or if ``column`` is set on a 1-D flag. """ column = quality_filter.get("column") - if q_flag.ndim == 2: - if column is None: - raise ValueError( - "quality_filter on a 2-D flag requires a 'column' index " - f"(got shape {q_flag.shape})" - ) + if q_flag.ndim == 2 and column is not None: q_flag = q_flag[:, column] - elif column is not None: + elif q_flag.ndim != 2 and column is not None: raise ValueError("quality_filter 'column' is only valid for a 2-D flag") value = quality_filter["value"] op = quality_filter.get("op", "eq") if op == "eq": - return q_flag == value - if op == "ne": - return q_flag != value - raise ValueError(f"Unknown quality_filter op '{op}' (expected 'eq' or 'ne')") + match = q_flag == value + elif op == "ne": + match = q_flag != value + else: + raise ValueError(f"Unknown quality_filter op '{op}' (expected 'eq' or 'ne')") + + # 2-D flag with no column selector: reduce across surface-type columns. ``ne`` + # keeps a photon if any surface type is non-matching (drops only all-column + # matches, e.g. TEP photons flagged -2 across every surface type); ``eq`` + # keeps a photon if any surface type matches. + if match.ndim == 2: + return match.any(axis=1) + return match def _read_group( diff --git a/tests/test_config.py b/tests/test_config.py index c52054d8..59a55c59 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -112,11 +112,12 @@ def test_functions_resolve(self, atl03_config): assert "expression" not in meta # scalar-only; non-scalar is #29 resolve_function(meta["function"]) # raises on failure - def test_confidence_filter_drops_noise(self, atl03_config): + def test_confidence_filter_drops_tep(self, atl03_config): qf = atl03_config.data_source["quality_filter"] - assert qf["value"] == 0 - assert qf["op"] == "ne" # keep signal_conf_ph != 0 (flags 1-4) - assert qf["column"] == 0 # signal_conf_ph is 2-D; select land column + assert qf["value"] == -2 + assert qf["op"] == "ne" # keep signal_conf_ph != -2 (drop only TEP) + # no column selector: TEP filter is reduced across all surface types + assert "column" not in qf assert qf["dataset"].endswith("signal_conf_ph") def test_rectilinear_grid(self, atl03_config): diff --git a/tests/test_processing.py b/tests/test_processing.py index c6fcb828..1ba05ff3 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -675,18 +675,30 @@ def test_2d_column_selects_surface_type(self): mask_col1 = _quality_mask(q_flag, {"value": 0, "op": "ne", "column": 1}) np.testing.assert_array_equal(mask_col1, [True, False, True]) - def test_2d_without_column_raises(self): - q_flag = np.zeros((3, 5), dtype=np.int8) - with pytest.raises(ValueError, match="requires a 'column'"): - _quality_mask(q_flag, {"value": 0, "op": "ne"}) + def test_2d_ne_reduces_across_surface_types(self): + # ATL03 TEP filter: signal_conf_ph == -2 across every surface type marks a + # TEP photon. With no column selector, != -2 keeps a photon if ANY surface + # type is non-TEP, dropping only all-column (true TEP) photons. + q_flag = np.array( + [[-2, -2, -2], [4, 4, 4], [-2, 3, -2]], # TEP / signal / mixed + dtype=np.int8, + ) + mask = _quality_mask(q_flag, {"value": -2, "op": "ne"}) + np.testing.assert_array_equal(mask, [False, True, True]) + + def test_2d_eq_reduces_across_surface_types(self): + # eq with no column keeps a photon if ANY surface type matches value. + q_flag = np.array([[0, 1], [2, 3], [1, 0]], dtype=np.int8) + mask = _quality_mask(q_flag, {"value": 0, "op": "eq"}) + np.testing.assert_array_equal(mask, [True, False, True]) def test_column_on_1d_raises(self): with pytest.raises(ValueError, match="only valid for a 2-D flag"): _quality_mask(np.array([0, 1]), {"value": 0, "column": 0}) def test_atl03_template_filter_runs(self): - # The shipped template's quality_filter must apply cleanly to a 2-D flag. + # The shipped template drops only TEP (-2) across all surface types. qf = default_config("atl03").data_source["quality_filter"] - q_flag = np.array([[0, 1], [1, 0], [4, 0]], dtype=np.int8) + q_flag = np.array([[-2, -2], [1, 4], [-2, 0]], dtype=np.int8) mask = _quality_mask(q_flag, qf) np.testing.assert_array_equal(mask, [False, True, True]) From 6b6c74a4bf3f6ca12459ffce0dee3b46804a6bdb Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 17 Jun 2026 01:13:19 +0000 Subject: [PATCH 07/16] address review: benchmark windows 1y + all (#30) --- benchmarks/region_timing.py | 55 +++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/benchmarks/region_timing.py b/benchmarks/region_timing.py index b26dd054..74c1caa2 100644 --- a/benchmarks/region_timing.py +++ b/benchmarks/region_timing.py @@ -2,10 +2,11 @@ This is the real-data (phase-3) counterpart to the synthetic ``handoff_bench.py``: it drives the actual catalog -> shard-map -> ``runner.agg()`` path against -ICESat-2 ATL03 granules over three 10 km x 10 km AOIs, swept across a set of -time windows, and records wall-time, peak RSS, observation/cell counts, and -output bytes for the ``pandas`` and ``arrow`` handoff carriers -- asserting their -scalar outputs are byte-for-byte identical (#30's parity criterion on real data). +ICESat-2 ATL03 granules over three 10 km x 10 km AOIs, swept across two date +ranges (calendar 2025 and the full 2018 -> Jan 2026 mission), and records +wall-time, peak RSS, observation/cell counts, and output bytes for the +``pandas`` and ``arrow`` handoff carriers -- asserting their scalar outputs are +byte-for-byte identical (#30's parity criterion on real data). Three regions span the density regimes the synthetic harness can't reproduce: @@ -19,10 +20,10 @@ expensive than the other three. Confidence filter and grid come straight from the shipped ``atl03`` template: -``signal_conf_ph != 0`` (keep flags 1-4, drop only noise), rectilinear grid. The -``output.grid.bounds`` are overridden per region to the AOI's 10 km box so the -grid covers just the patch. (HEALPix order-19 -- the ~10 m match -- waits on -mortie #35, so this first pass is rectilinear only.) +``signal_conf_ph != -2`` (drop only TEP photons, across all surface types), +rectilinear grid. The ``output.grid.bounds`` are overridden per region to the +AOI's 10 km box so the grid covers just the patch. (HEALPix order-19 -- the +~10 m match -- waits on mortie #35, so this first pass is rectilinear only.) **This script is NOT run in CI**: it needs ``earthaccess``/NSIDC-S3 credentials (CMR-STAC query + byte-range HDF5 reads) and is slow. It lives under @@ -31,7 +32,7 @@ Run (in a credentialed session):: - uv run python benchmarks/region_timing.py --windows 1d,5d --max-cells 4 + uv run python benchmarks/region_timing.py --windows 1y --max-cells 4 uv run python benchmarks/region_timing.py --hard --out results.txt """ @@ -42,22 +43,24 @@ import tempfile import time from dataclasses import dataclass -from datetime import date, timedelta import numpy as np from zagg.config import default_config, get_data_vars from zagg.runner import agg -# Time windows (per @espg). Each is the span ending at ``--end-date`` that drives -# the catalog query; longer windows pull more granules -> more overlap. +# Time windows (per @espg). Each is an explicit ``(start, end)`` date range +# (inclusive start, exclusive end) driving the catalog query. The AOI only has +# ~10 granules for a whole year, so the earlier {1d,5d,15d,30d,91d} sweep was too +# sparse to be meaningful; we benchmark two ranges instead: +# +# * ``1y`` -- calendar 2025 (2025-01-01 .. 2026-01-01). +# * ``all`` -- the full mission, 2018-01-01 .. 2026-01-01. The upper bound is +# CLIPPED at 2026-01-01 so future runs stay stable and don't drift as new +# granules land (otherwise benchmark numbers would creep every run). WINDOWS = { - "1d": 1, - "5d": 5, - "15d": 15, - "30d": 30, - "91d": 91, - "1y": 365, + "1y": ("2025-01-01", "2026-01-01"), + "all": ("2018-01-01", "2026-01-01"), } # Carriers compared head-to-head. Both must produce identical scalar outputs. @@ -168,9 +171,8 @@ def _output_bytes(store_path: str) -> int: def run_one( region: Region, window: str, - days: int, + date_range: tuple[str, str], *, - end_date: date, version: str, max_cells: int | None, max_workers: int | None, @@ -181,8 +183,8 @@ def run_one( from zagg.grids import from_config cfg = _region_config(region) - start = (end_date - timedelta(days=days)).isoformat() - query = Query("ATL03", version, start, end_date.isoformat(), region=region.bbox) + start, end = date_range + query = Query("ATL03", version, start, end, region=region.bbox) with tempfile.TemporaryDirectory() as tmp: catalog_path = f"{tmp}/shardmap.json" @@ -249,12 +251,7 @@ def parse_args(argv=None): ap.add_argument( "--windows", default=",".join(WINDOWS), - help="Comma-separated time windows (default: all). Choices: " + ", ".join(WINDOWS), - ) - ap.add_argument( - "--end-date", - default=date.today().isoformat(), - help="End of each time window, YYYY-MM-DD (default: today).", + help="Comma-separated time windows (default: both). Choices: " + ", ".join(WINDOWS), ) ap.add_argument( "--hard", @@ -287,7 +284,6 @@ def main(argv=None): raise SystemExit(f"unknown window {w!r}; choices: {', '.join(WINDOWS)}") windows.append(w) - end_date = date.fromisoformat(args.end_date) regions = list(REGIONS) + ([HARD_REGION] if args.hard else []) records: list[Record] = [] @@ -298,7 +294,6 @@ def main(argv=None): region, window, WINDOWS[window], - end_date=end_date, version=args.version, max_cells=args.max_cells, max_workers=args.max_workers, From 771870e0de0217d2487d54a71a5b45542dbc8808 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 17 Jun 2026 01:15:16 +0000 Subject: [PATCH 08/16] fold review: refresh _quality_mask docstring TEP example (#30) --- src/zagg/processing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 9b4d1f67..a9445489 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -494,7 +494,7 @@ def _quality_mask(q_flag: np.ndarray, quality_filter: dict) -> np.ndarray: - ``"eq"`` keeps rows where the flag equals ``value`` (ATL06: keep ``atl06_quality_summary == 0``). - ``"ne"`` keeps rows where the flag differs from ``value`` (ATL03: keep - ``signal_conf_ph != 0``, dropping only the noise flag). + ``signal_conf_ph != -2``, dropping only TEP photons). When ``q_flag`` is 2-D (e.g. ATL03 ``signal_conf_ph`` is ``(n_photons, n_surface_types)``), the comparison is reduced **across all surface-type From 871f3dab14683d84508ef5f9bf3be1d4a876de67 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 00:46:01 +0000 Subject: [PATCH 09/16] fix trailing newline in test_processing.py (ruff W292) Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01MznHvCXwruEszMEffRTMyw --- tests/test_processing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_processing.py b/tests/test_processing.py index ccea4e22..4db60871 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1061,4 +1061,4 @@ class _OneChunkGrid: with pytest.raises(ValueError, match="one whole chunk"): write_dataframe_to_zarr( table, store, grid=_OneChunkGrid(), chunk_idx=(0,) - ) \ No newline at end of file + ) From 7976b8e3dc9df7790551ce3b538647db44846f52 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 00:51:16 +0000 Subject: [PATCH 10/16] phase 3 of issue #30 Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01MznHvCXwruEszMEffRTMyw --- benchmarks/region_timing.py | 121 ++++++++++++++++++-- src/zagg/configs/atl03_waveform_counts.yaml | 49 ++++++++ tests/test_config.py | 97 +++++++++++++++- tests/test_processing.py | 17 +-- 4 files changed, 258 insertions(+), 26 deletions(-) create mode 100644 src/zagg/configs/atl03_waveform_counts.yaml diff --git a/benchmarks/region_timing.py b/benchmarks/region_timing.py index 74c1caa2..52402a82 100644 --- a/benchmarks/region_timing.py +++ b/benchmarks/region_timing.py @@ -25,6 +25,12 @@ AOI's 10 km box so the grid covers just the patch. (HEALPix order-19 -- the ~10 m match -- waits on mortie #35, so this first pass is rectilinear only.) +Phase 3 adds a second template sweep -- ``atl03_waveform_counts`` -- alongside +the scalar ``atl03`` run: wall-times for both templates are printed side-by-side +so the overhead of the 128-bin vector histogram over plain scalars is visible. +The waveform store is also checked to confirm the ``waveform_counts`` array has +the expected trailing shape of 128. + **This script is NOT run in CI**: it needs ``earthaccess``/NSIDC-S3 credentials (CMR-STAC query + byte-range HDF5 reads) and is slow. It lives under ``benchmarks/`` (not ``tests/``) and is meant for a credentialed session. It must @@ -107,10 +113,11 @@ def _box(name: str, lon: float, lat: float) -> Region: @dataclass class Record: - """One (region x window x handoff) measurement.""" + """One (region x window x template x handoff) measurement.""" region: str window: str + template: str handoff: str wall_s: float peak_rss_mb: float @@ -133,11 +140,19 @@ def _region_config(region: Region): return cfg -def _store_scalars(store_path: str, cfg) -> dict[str, np.ndarray]: +def _region_waveform_config(region: Region): + """An ``atl03_waveform_counts`` config clipped to the region bbox.""" + cfg = default_config("atl03_waveform_counts") + cfg = copy.deepcopy(cfg) + cfg.output["grid"]["bounds"] = list(region.bbox) + return cfg + + +def _store_arrays(store_path: str, cfg) -> dict[str, np.ndarray]: """Read each aggregated data variable out of a written store as a dense array. - Data is written under the grid's ``group_path`` (see - ``write_dataframe_to_zarr``), so resolve it from the config's grid. + Works for both scalar (1-D) and vector (N-D) variables; the trailing + dimension(s) are preserved so callers can check shapes. """ import zarr @@ -149,7 +164,7 @@ def _store_scalars(store_path: str, cfg) -> dict[str, np.ndarray]: def _assert_parity(a: dict[str, np.ndarray], b: dict[str, np.ndarray], ctx: str) -> None: - """Assert two stores' scalar outputs are byte-for-byte identical (NaN-aware).""" + """Assert two stores' outputs are byte-for-byte identical (NaN-aware for floats).""" assert a.keys() == b.keys(), f"{ctx}: variable mismatch {a.keys()} vs {b.keys()}" for name in a: x, y = a[name], b[name] @@ -192,7 +207,7 @@ def run_one( make_shardmap(query, grid).to_json(catalog_path) records: list[Record] = [] - scalars: dict[str, dict] = {} + arrays: dict[str, dict] = {} for handoff in HANDOFFS: store_path = f"{tmp}/out_{handoff}.zarr" t0 = time.perf_counter() @@ -206,11 +221,77 @@ def run_one( overwrite=True, ) wall = time.perf_counter() - t0 - scalars[handoff] = _store_scalars(store_path, cfg) + arrays[handoff] = _store_arrays(store_path, cfg) + records.append( + Record( + region=region.name, + window=window, + template="atl03", + handoff=handoff, + wall_s=wall, + peak_rss_mb=_peak_rss_mb(), + total_obs=int(summary.get("total_obs", 0)), + cells_with_data=int(summary.get("cells_with_data", 0)), + output_bytes=_output_bytes(store_path), + ) + ) + _assert_parity( + arrays["pandas"], + arrays["arrow"], + ctx=f"{region.name}/{window}/atl03", + ) + return records + + +def run_one_waveform( + region: Region, + window: str, + date_range: tuple[str, str], + *, + version: str, + max_cells: int | None, + max_workers: int | None, +) -> list[Record]: + """Run the ``atl03_waveform_counts`` template for (region, window). + + Runs both carriers and asserts parity (both produce identical vector output). + Also confirms the ``waveform_counts`` array carries the expected 128-element + trailing dimension so the shape contract is exercised on real data. + """ + from zagg.catalog import make_shardmap + from zagg.catalog.sources import Query + from zagg.grids import from_config + + cfg = _region_waveform_config(region) + start, end = date_range + query = Query("ATL03", version, start, end, region=region.bbox) + + with tempfile.TemporaryDirectory() as tmp: + catalog_path = f"{tmp}/shardmap.json" + grid = from_config(cfg) + make_shardmap(query, grid).to_json(catalog_path) + + records: list[Record] = [] + arrays: dict[str, dict] = {} + for handoff in HANDOFFS: + store_path = f"{tmp}/wf_{handoff}.zarr" + t0 = time.perf_counter() + summary = agg( + cfg, + catalog=catalog_path, + store=store_path, + handoff=handoff, + max_cells=max_cells, + max_workers=max_workers, + overwrite=True, + ) + wall = time.perf_counter() - t0 + arrays[handoff] = _store_arrays(store_path, cfg) records.append( Record( region=region.name, window=window, + template="atl03_waveform", handoff=handoff, wall_s=wall, peak_rss_mb=_peak_rss_mb(), @@ -220,23 +301,29 @@ def run_one( ) ) _assert_parity( - scalars["pandas"], - scalars["arrow"], - ctx=f"{region.name}/{window}", + arrays["pandas"], + arrays["arrow"], + ctx=f"{region.name}/{window}/atl03_waveform_counts", ) + # Shape check: waveform_counts must carry the 128-element trailing dim. + wf = arrays["pandas"]["waveform_counts"] + if wf.ndim != 2 or wf.shape[1] != 128: + raise AssertionError( + f"{region.name}/{window}: waveform_counts shape {wf.shape} expected (..., 128)" + ) return records def format_table(records: list[Record]) -> str: """Render the collected records as a fixed-width text table.""" header = ( - f"{'region':<18}{'window':>8}{'handoff':>10}{'wall_s':>10}" + f"{'region':<18}{'window':>8}{'template':>16}{'handoff':>10}{'wall_s':>10}" f"{'rss_MB':>10}{'obs':>12}{'cells':>10}{'out_MB':>10}" ) lines = [header, "-" * len(header)] for r in records: lines.append( - f"{r.region:<18}{r.window:>8}{r.handoff:>10}{r.wall_s:>10.3f}" + f"{r.region:<18}{r.window:>8}{r.template:>16}{r.handoff:>10}{r.wall_s:>10.3f}" f"{r.peak_rss_mb:>10.1f}{r.total_obs:>12,}{r.cells_with_data:>10,}" f"{r.output_bytes / 1e6:>10.2f}" ) @@ -299,6 +386,16 @@ def main(argv=None): max_workers=args.max_workers, ) ) + records.extend( + run_one_waveform( + region, + window, + WINDOWS[window], + version=args.version, + max_cells=args.max_cells, + max_workers=args.max_workers, + ) + ) table = format_table(records) print(table) diff --git a/src/zagg/configs/atl03_waveform_counts.yaml b/src/zagg/configs/atl03_waveform_counts.yaml new file mode 100644 index 00000000..00753a3f --- /dev/null +++ b/src/zagg/configs/atl03_waveform_counts.yaml @@ -0,0 +1,49 @@ +# ATL03 per-cell waveform-counts template. +# +# Aggregates photon heights (h_ph) into a 128-bin fixed-width histogram at 2 m +# vertical spacing, centred on the per-cell median of the ICESat-2 reference +# DEM height (dem_h). Each bin spans 2 m, giving a 256 m window +# [median(dem_h) - 128 m, median(dem_h) + 128 m). Out-of-range photons are +# dropped (not counted). The companion scalar bin_start records median(dem_h) +# so the absolute elevation of bin 0 can be recovered downstream. +# +# Both h_ph and dem_h live in the ATL03 heights group (NSIDC ATL03 v3 data +# dictionary). Confidence filter is the same as atl03.yaml: drop only TEP +# photons (signal_conf_ph == -2 across all surface types). Grid is rectilinear, +# matching atl03.yaml. +data_source: + reader: h5coro + driver: s3 + groups: [gt1l, gt1r, gt2l, gt2r, gt3l, gt3r] + coordinates: + latitude: "/{group}/heights/lat_ph" + longitude: "/{group}/heights/lon_ph" + variables: + h_ph: "/{group}/heights/h_ph" + dem_h: "/{group}/heights/dem_h" + quality_filter: + dataset: "/{group}/heights/signal_conf_ph" + value: -2 # TEP flag; reduced across all surface types + op: ne + +aggregation: + variables: + waveform_counts: + kind: vector + trailing_shape: 128 + expression: "np.histogram(h_ph - np.median(dem_h), 128, (-128.0, 128.0))[0].astype('uint32')" + source: h_ph + dtype: uint32 + fill_value: 0 + bin_start: + function: median + source: dem_h + dtype: float32 + +output: + grid: + type: rectilinear + crs: EPSG:4326 + resolution: 0.0001 # degrees (~10 m at the equator) + bounds: [-180, -90, 180, 90] + chunk_shape: [256, 256] diff --git a/tests/test_config.py b/tests/test_config.py index 718c6ea0..b6b2a704 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -99,6 +99,7 @@ def test_nonexistent_raises(self): # ATL03 template # --------------------------------------------------------------------------- + class TestATL03Template: @pytest.fixture def atl03_config(self): @@ -221,9 +222,11 @@ def test_bad_quality_filter_op(self): "variables": {"h_li": "/path"}, "quality_filter": {"dataset": "/q", "value": 0, "op": "gt"}, }, - aggregation={"variables": { - "c": {"function": "len", "source": "h_li", "dtype": "int32"}, - }}, + aggregation={ + "variables": { + "c": {"function": "len", "source": "h_li", "dtype": "int32"}, + } + }, output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, ) with pytest.raises(ValueError, match="quality_filter.op must be"): @@ -699,3 +702,91 @@ def test_vector_list_signature(self): sig = get_output_signature({"kind": "vector", "trailing_shape": [16, 2]}) assert sig["kind"] == "vector" assert sig["trailing_shape"] == (16, 2) + + +# --------------------------------------------------------------------------- +# atl03_waveform_counts template (issue #30, phase 3) +# --------------------------------------------------------------------------- + + +class TestATL03WaveformCountsTemplate: + @pytest.fixture + def cfg(self): + return default_config("atl03_waveform_counts") + + def test_loads_and_validates(self, cfg): + # default_config already calls validate_config; just confirm round-trip. + validate_config(cfg) + assert cfg.data_source["reader"] == "h5coro" + assert len(cfg.data_source["groups"]) == 6 + + def test_variables_include_h_ph_and_dem_h(self, cfg): + ds_vars = cfg.data_source["variables"] + assert "h_ph" in ds_vars + assert "dem_h" in ds_vars + assert ds_vars["h_ph"].endswith("h_ph") + assert ds_vars["dem_h"].endswith("dem_h") + + def test_waveform_counts_field_is_vector(self, cfg): + fields = get_agg_fields(cfg) + meta = fields["waveform_counts"] + sig = get_output_signature(meta) + assert sig["kind"] == "vector" + assert sig["trailing_shape"] == (128,) + assert sig["dtype"] == "uint32" + + def test_bin_start_field_is_scalar(self, cfg): + fields = get_agg_fields(cfg) + meta = fields["bin_start"] + sig = get_output_signature(meta) + assert sig["kind"] == "scalar" + assert sig["trailing_shape"] == () + + def test_waveform_counts_expression_with_synthetic_data(self, cfg): + # Photons all within range; all should be counted. + from zagg.processing import calculate_cell_statistics + + np.random.seed(0) + dem_h = np.full(50, 0.0, dtype="float32") + h_ph = np.random.uniform(-100.0, 100.0, 50).astype("float32") + result = calculate_cell_statistics( + {"h_ph": h_ph, "dem_h": dem_h, "leaf_id": np.arange(50)}, config=cfg + ) + wc = result["waveform_counts"] + assert wc.shape == (128,) + assert wc.dtype == np.dtype("uint32") + assert int(wc.sum()) == 50, "all in-range photons must be counted" + + def test_out_of_range_photons_dropped(self, cfg): + from zagg.processing import calculate_cell_statistics + + # One photon within range, one beyond ±128 m of dem_h median. + dem_h = np.array([0.0, 0.0], dtype="float32") + h_ph = np.array([10.0, 500.0], dtype="float32") + result = calculate_cell_statistics( + {"h_ph": h_ph, "dem_h": dem_h, "leaf_id": np.arange(2)}, config=cfg + ) + wc = result["waveform_counts"] + assert int(wc.sum()) == 1, "out-of-range photon must not appear in any bin" + + def test_empty_cell_returns_zero_filled_vector(self, cfg): + from zagg.processing import calculate_cell_statistics + + result = calculate_cell_statistics( + {"h_ph": np.array([]), "dem_h": np.array([]), "leaf_id": np.array([])}, config=cfg + ) + wc = result["waveform_counts"] + assert wc.shape == (128,) + assert np.all(wc == 0), "empty cell sentinel must be all-zero for uint32/fill_value:0" + + def test_confidence_filter_same_as_atl03(self, cfg): + qf = cfg.data_source["quality_filter"] + assert qf["value"] == -2 + assert qf["op"] == "ne" + assert "column" not in qf + assert qf["dataset"].endswith("signal_conf_ph") + + def test_rectilinear_grid(self, cfg): + grid = cfg.output["grid"] + assert grid["type"] == "rectilinear" + assert len(grid["bounds"]) == 4 diff --git a/tests/test_processing.py b/tests/test_processing.py index 4db60871..82ae27f3 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -650,7 +650,7 @@ def fake_read_group(*args, **kwargs): monkeypatch.setattr("zagg.processing._read_group", fake_read_group) monkeypatch.setattr("zagg.processing.h5coro.H5Coro", lambda *a, **k: object()) # Avoid resolving a real h5coro driver (s3driver import / creds plumbing). - monkeypatch.setattr("zagg.processing._make_url_rewriter", lambda driver: (lambda u: u)) + monkeypatch.setattr("zagg.processing._make_url_rewriter", lambda driver: lambda u: u) def test_kernel_branch_matches_default_path(self, monkeypatch): """process_shard(handoff="arrow-kernel") agrees with the default path on the @@ -783,7 +783,7 @@ def _vector_cfg(): "params": {"minlength": 3}, }, } - } + }, ) def test_has_vector_fields(self): @@ -952,6 +952,7 @@ def test_atl03_template_filter_runs(self): mask = _quality_mask(q_flag, qf) np.testing.assert_array_equal(mask, [False, True, True]) + class TestVectorRoundTrip: """Issue #29 phase 6: a vector field written to a real Zarr template reads back through the trailing-dim block, and NaN-padded empty cells are skipped @@ -975,9 +976,7 @@ def _vector_cfg(): } from zagg.config import PipelineConfig - return PipelineConfig( - data_source=cfg.data_source, aggregation=agg, output=cfg.output - ) + return PipelineConfig(data_source=cfg.data_source, aggregation=agg, output=cfg.output) def test_vector_leaf_to_zarr_to_read(self): pytest.importorskip("pyarrow") @@ -1054,11 +1053,7 @@ class _OneChunkGrid: dtype="float32", fill_value=np.float32("nan"), ) - edges = pa.FixedSizeListArray.from_arrays( - pa.array(np.arange(8.0, dtype="float32")), 4 - ) + edges = pa.FixedSizeListArray.from_arrays(pa.array(np.arange(8.0, dtype="float32")), 4) table = pa.table({"edges": edges}) with pytest.raises(ValueError, match="one whole chunk"): - write_dataframe_to_zarr( - table, store, grid=_OneChunkGrid(), chunk_idx=(0,) - ) + write_dataframe_to_zarr(table, store, grid=_OneChunkGrid(), chunk_idx=(0,)) From 18ad7bd2253913b207ccbbce13e6c4029a62950c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 00:59:21 +0000 Subject: [PATCH 11/16] fix total_obs KeyError and waveform shape check (phase 3 self-review) Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01MznHvCXwruEszMEffRTMyw --- benchmarks/region_timing.py | 2 +- src/zagg/processing.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/region_timing.py b/benchmarks/region_timing.py index 52402a82..151a4e5c 100644 --- a/benchmarks/region_timing.py +++ b/benchmarks/region_timing.py @@ -307,7 +307,7 @@ def run_one_waveform( ) # Shape check: waveform_counts must carry the 128-element trailing dim. wf = arrays["pandas"]["waveform_counts"] - if wf.ndim != 2 or wf.shape[1] != 128: + if wf.shape[-1] != 128: raise AssertionError( f"{region.name}/{window}: waveform_counts shape {wf.shape} expected (..., 128)" ) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 1735a1b2..09f92109 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -1036,7 +1036,7 @@ def process_shard( logger.info(f"Completed shard {shard_key} in {duration:.1f}s") metadata["cells_with_data"] = cells_with_data - metadata["total_obs"] = int(stats_arrays["count"].sum()) + metadata["total_obs"] = n_obs_total metadata["duration_s"] = duration return df_out, metadata From 9b111ee7cde77b86ab6a0b8187fccad35c7a9f3d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 16:52:41 +0000 Subject: [PATCH 12/16] fold review: validate flat quality_filter, add atl03 read-path test, yaml caveat --- src/zagg/config.py | 31 +++++++++++++---- src/zagg/configs/atl03_waveform_counts.yaml | 8 +++++ tests/test_processing.py | 38 +++++++++++++++++++++ 3 files changed, 70 insertions(+), 7 deletions(-) diff --git a/src/zagg/config.py b/src/zagg/config.py index 17dfd73c..5bd841c2 100644 --- a/src/zagg/config.py +++ b/src/zagg/config.py @@ -455,14 +455,31 @@ def _normalize_filter(f: dict) -> dict: def _validate_filters(data_source: dict) -> None: - """Validate the ``filters`` list (and that a flat ``quality_filter`` is sane). - - Raises ``ValueError`` on: unknown op, missing ``dataset``, ``in``/``not_in`` - without a list ``values``, scalar ops without ``value``, non-int ``column``, - a non-base-level ``expression`` filter, or wrong ``value`` type. ``column`` is - required for the N-D flag case but cannot be checked against array rank here - (no data); rank checks happen at read time. + """Validate the ``filters`` list and the flat ``quality_filter`` sugar. + + For the structured ``filters:`` list, raises ``ValueError`` on: unknown op, + missing ``dataset``, ``in``/``not_in`` without a list ``values``, scalar ops + without ``value``, non-int ``column``, a non-base-level ``expression`` filter, + or wrong ``value`` type. ``column`` is required for the N-D flag case but + cannot be checked against array rank here (no data); rank checks happen at + read time. + + For the flat ``quality_filter`` sugar, only ``dataset`` and ``value`` are + honored (``filters_from_data_source`` synthesizes ``op: eq, column: null``). + Reject any extra keys at load time so a user-typoed ``op: gt`` or stray + ``column:`` is not silently dropped on the floor — the structured ``filters:`` + list is the right form when those knobs are wanted. """ + qf = data_source.get("quality_filter") + if qf is not None: + allowed = {"dataset", "value"} + unknown = set(qf) - allowed + if unknown: + raise ValueError( + f"data_source.quality_filter only honors {sorted(allowed)} " + f"(got extra keys {sorted(unknown)}); use the structured " + "'filters:' list to set 'op', 'column', 'keep', etc." + ) filters = data_source.get("filters") if filters is None: return diff --git a/src/zagg/configs/atl03_waveform_counts.yaml b/src/zagg/configs/atl03_waveform_counts.yaml index 588c48d6..3d5b4a40 100644 --- a/src/zagg/configs/atl03_waveform_counts.yaml +++ b/src/zagg/configs/atl03_waveform_counts.yaml @@ -13,6 +13,14 @@ # on the cell's own median(h_ph) is more natural for a photon-height density # template and avoids the cross-group / resolution-mismatch read. # +# Caveat: median(h_ph) is a *per-cell* anchor, not a stable absolute reference. +# In a bimodal cell (e.g. ground + canopy returns) the median can land between +# the two clusters, splitting both into wings of the histogram; that is a real +# tradeoff of centring on the photon distribution itself. bin_start records the +# centring value so the bin-0 absolute elevation = (bin_start - 128 m) is still +# recoverable, but comparing waveforms across cells in the same absolute frame +# requires shifting each cell's bins by its own bin_start. +# # Confidence filter is the same as atl03.yaml: drop only TEP photons by the # structured ``filters:`` list (signal_conf_ph[:, 0] != -2; TEP is uniform # across surface types). Grid is rectilinear, matching atl03.yaml. diff --git a/tests/test_processing.py b/tests/test_processing.py index d153c061..4b333355 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1349,6 +1349,44 @@ def test_in_op_integer_column(self): df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) assert df["h"].tolist() == [0.0, 2.0, 4.0] + def test_atl03_shipped_template_2d_signal_conf(self): + # Drives the shipped atl03.yaml structured filter through the read path + # against a realistic (n_photons, 5) signal_conf_ph: one TEP photon + # (-2 across every surface type) plus four non-TEP photons of varying + # confidence. Only the TEP row is dropped (column 0, op: ne, value: -2). + from zagg.config import default_config + + atl03_filters = default_config("atl03").data_source["filters"] + conf = np.array( + [ + [4, 4, 4, 4, 4], # all-high-confidence (kept) + [-2, -2, -2, -2, -2], # TEP across all surface types (dropped) + [0, 0, 0, 0, 0], # noise across all (kept; -2 is the TEP flag) + [3, 2, 1, 0, 4], # mixed confidence (kept) + [-2, 4, 4, 4, 4], # column 0 is TEP but others aren't -- with + # column: 0 this row is dropped, even though + # the photon is a valid land-ice return on + # surface type 3. See PR #47 review thread: + # this is the operational tradeoff of moving + # from .any(axis=1) to a single-column key. + ] + ) + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.arange(5.0, dtype=np.float32), + "/conf": conf, + } + ) + # Rewrite the shipped template's filter dataset path to match the fake h5. + f = dict(atl03_filters[0]) + f["dataset"] = "/conf" + ds = self._data_source(filters=[f]) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + # Rows 1 (true TEP) and 4 (column-0 TEP only) are dropped. + assert df["h"].tolist() == [0.0, 2.0, 3.0] + def test_expression_filter_base_level(self): h5 = _FakeH5( { From 02764fc64cfa68036202a66e02bb261c8601ebc1 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 17:23:03 +0000 Subject: [PATCH 13/16] phase 5 of issue #30 --- src/zagg/processing.py | 287 ++++++++++++++++++++++++++++++++++++++- src/zagg/read_plan.py | 10 +- tests/test_processing.py | 229 +++++++++++++++++++++++++++++++ tests/test_read_plan.py | 13 +- 4 files changed, 526 insertions(+), 13 deletions(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 018802e6..d2de78fe 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -25,6 +25,7 @@ get_data_vars, get_output_signature, ) +from zagg.read_plan import execute_read_plan, plan_read from zagg.schema import ProcessingMetadata logger = logging.getLogger(__name__) @@ -814,6 +815,247 @@ def _predicate_mask(arr: np.ndarray, f: dict) -> np.ndarray: return mask +def _level_coord_paths(level: dict, group: str) -> tuple[str, str]: + """Resolve ``(latitude, longitude)`` HDF5 paths for a coarse-level spatial index. + + The level's ``coordinates`` field is a ``{latitude, longitude}`` dict of names + relative to the level's ``path`` template (matching the schema in #43's issue + body). Both halves are required for the ``read_plan`` to compute an AOI box. + """ + coords = level.get("coordinates") + if not isinstance(coords, dict) or "latitude" not in coords or "longitude" not in coords: + raise ValueError( + "read_plan.spatial_index level requires " + "'coordinates: {latitude: , longitude: }'" + ) + base = level["path"].format(group=group).rstrip("/") + lat_name = coords["latitude"] + lon_name = coords["longitude"] + # Allow either a relative name (joined to the level path) or an absolute path + # template (already group-substituted on .format above? no -- coords names + # don't carry templates; keep them simple). Absolute paths win as-is. + lat_path = lat_name if lat_name.startswith("/") else f"{base}/{lat_name}" + lon_path = lon_name if lon_name.startswith("/") else f"{base}/{lon_name}" + return lat_path, lon_path + + +def _planned_read_group( + h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False +): + """Planned (AOI-bounded) read of one HDF5 group via the coarse spatial index. + + Issue #43 Phase C: when ``data_source.read_plan.spatial_index`` names a coarse + level whose ``link`` points at the base level, we read the coarse coordinates + + link arrays once (small), call :func:`zagg.read_plan.plan_read` to compute + which base-rate slices the AOI bbox actually touches, and read base-rate + coords + variables + filter datasets only over those slices via + :func:`zagg.read_plan.execute_read_plan`. This avoids the + ``lat_ph`` + ``lon_ph`` full-coord read (up to ~245 MB per ATL03 beam) that + drives Lambda OOMs (issue #43 motivation). + + Falls back transparently to :func:`_read_group` when: + - the empty-AOI short-circuit fires (no parents match) → return ``None``; + - ``plan_read`` flags ``full_read=True`` (selectivity above threshold); + - the cell ``signal_conf_ph``-style 2-D structured filter would be re-read + via the planned slices either way (the helper handles that uniformly). + + Returns the same ``pandas.DataFrame`` / ``pyarrow.Table`` / ``None`` contract + as :func:`_read_group`. Output rows are in plan-slice / spatial-mask / + filter order — which matches the full-read path's row ordering because the + plan's runs are emitted in increasing parent index. + """ + coordinates = data_source["coordinates"] + variables = data_source["variables"] + levels = data_source["levels"] + base_level_key = data_source["base_level"] + rp = data_source["read_plan"] + spatial_index_level = rp["spatial_index"] + pad = int(rp.get("pad", 1)) + full_read_threshold = float(rp.get("full_read_threshold", 0.9)) + + si_lvl = levels[spatial_index_level] + link = si_lvl.get("link") + if not isinstance(link, dict): + raise ValueError( + f"read_plan.spatial_index level {spatial_index_level!r} requires a 'link'" + ) + if link["to"] != base_level_key: + raise ValueError( + f"read_plan.spatial_index level {spatial_index_level!r} must link " + f"directly to base level {base_level_key!r} (got link.to={link['to']!r})" + ) + index_base = int(link.get("index_base", 0)) + + # Read coarse-level coordinates + link arrays in one go (small — geolocation + # rate is ~30x lighter than photon rate on ATL03). + si_lat_path, si_lon_path = _level_coord_paths(si_lvl, group) + ibeg_path = link["index_beg"].format(group=group) + cnt_path = link["count"].format(group=group) + coarse_data = h5obj.readDatasets([si_lat_path, si_lon_path, ibeg_path, cnt_path]) + coarse_lats = coarse_data[si_lat_path] + coarse_lons = coarse_data[si_lon_path] + ibeg_arr = coarse_data[ibeg_path] + cnt_arr = coarse_data[cnt_path] + + if len(coarse_lats) == 0: + return None + + # Under the contiguity assumption (#43): ``sum(count) == n_base``, so the + # last parent's tail gives the total length without an extra header read. + n_base = int(ibeg_arr[-1]) - index_base + int(cnt_arr[-1]) + if n_base <= 0: + return None + + # Compute the shard's WGS84 bbox from the grid (every grid's shard_footprint + # returns a shapely Polygon). + poly = grid.shard_footprint(shard_key) + min_lon, min_lat, max_lon, max_lat = poly.bounds + bbox = (float(min_lon), float(min_lat), float(max_lon), float(max_lat)) + + plan = plan_read( + np.asarray(coarse_lats), + np.asarray(coarse_lons), + np.asarray(ibeg_arr), + np.asarray(cnt_arr), + n_base, + bbox, + index_base=index_base, + pad=pad, + full_read_threshold=full_read_threshold, + ) + + if not plan.parent_runs: + return None # empty AOI -- no parent intersects, skip the group entirely + + if plan.full_read: + # Selectivity above threshold: many small reads would still sum to most + # of the file. Defer to the full-coord-read path; semantics identical. + return _read_group_full(h5obj, group, data_source, shard_key, grid, arrow=arrow) + + # h5coro-compatible reader callback for execute_read_plan. + def _read_fn(path, hyperslice=None): + if hyperslice is None: + return h5obj.readDatasets([path])[path] + return h5obj.readDatasets([{"dataset": path, "hyperslice": hyperslice}])[path] + + # ---- Read base coords + variables + filter datasets over the planned slices. + filters = filters_from_data_source(data_source) + base_structured = [ + f + for f in filters + if "expression" not in f and (f.get("level") is None or f.get("level") == base_level_key) + ] + coarse_structured = [ + f + for f in filters + if "expression" not in f + and f.get("level") is not None + and f.get("level") != base_level_key + ] + expressions = [f for f in filters if "expression" in f] + + lat_path = coordinates["latitude"].format(group=group) + lon_path = coordinates["longitude"].format(group=group) + lats = execute_read_plan(plan, _read_fn, lat_path, np.float64) + lons = execute_read_plan(plan, _read_fn, lon_path, np.float64) + + if len(lats) == 0: + return None + + # Apply spatial / shard mask over the concatenated planned reads. + leaf_ids = grid.assign(lats, lons) + mask_spatial = grid.shards_of(leaf_ids) == shard_key + if np.sum(mask_spatial) == 0: + return None + + # Read the variables and base-level filter datasets via the same plan. Read + # each distinct path once (the variable and filter dataset paths can coincide). + var_paths = {col: tmpl.format(group=group) for col, tmpl in variables.items()} + filter_paths = {id(f): f["dataset"].format(group=group) for f in base_structured} + paths_seen: set[str] = set() + arrays_by_path: dict[str, np.ndarray] = {} + for path in list(var_paths.values()) + list(filter_paths.values()): + if path in paths_seen: + continue + paths_seen.add(path) + # dtype hint isn't load-bearing -- execute_read_plan dtype-casts via + # np.asarray, which is a no-op when the source dtype already matches. + arrays_by_path[path] = execute_read_plan(plan, _read_fn, path, None) + + # Base-level structured filters: ANDed keep-masks over the concatenated reads. + keep_mask: np.ndarray | None = None + for f in base_structured: + flag = arrays_by_path[filter_paths[id(f)]][mask_spatial] + fmask = _predicate_mask(flag, f) + keep_mask = fmask if keep_mask is None else (keep_mask & fmask) + + # Cross-level (Phase B) filters: read coarse flags fully, expand to base + # rate (length n_base), then subset to the planned indices. + if coarse_structured: + # Build the global base-index array once: which original-base positions + # are present in the concatenated planned read. + global_idx = np.concatenate( + [np.arange(s, e, dtype=np.int64) for s, e in plan.base_slices] + ) + cross_full: np.ndarray | None = None + for f in coarse_structured: + level_key = f["level"] + cf_lvl = levels[level_key] + cf_link = cf_lvl["link"] + cf_index_base = int(cf_link.get("index_base", 0)) + cf_flag_path = f["dataset"].format(group=group) + cf_ibeg_path = cf_link["index_beg"].format(group=group) + cf_cnt_path = cf_link["count"].format(group=group) + cf_data = h5obj.readDatasets([cf_flag_path, cf_ibeg_path, cf_cnt_path]) + cf_flag = cf_data[cf_flag_path] + cf_ibeg = cf_data[cf_ibeg_path] + cf_cnt = cf_data[cf_cnt_path] + coarse_fmask = _predicate_mask(cf_flag, f) + expanded = _expand_mask_to_base(coarse_fmask, cf_ibeg, cf_cnt, cf_index_base, n_base) + cross_full = expanded if cross_full is None else (cross_full & expanded) + # Subset the full-length mask to the concatenated planned indices, then + # to the spatial keep window so it lines up with keep_mask above. + cross_planned = cross_full[global_idx][mask_spatial] + keep_mask = cross_planned if keep_mask is None else (keep_mask & cross_planned) + + if keep_mask is not None and np.sum(keep_mask) == 0: + return None + + # Build the data dict (variables sliced to mask_spatial, then to keep_mask). + leaf_after_spatial = leaf_ids[mask_spatial] + data_dict: dict[str, np.ndarray] = {} + for col_name, path in var_paths.items(): + values = arrays_by_path[path][mask_spatial] + if keep_mask is not None: + values = values[keep_mask] + data_dict[col_name] = values + data_dict["leaf_id"] = leaf_after_spatial[keep_mask] if keep_mask is not None else leaf_after_spatial + + # Base-level expression filters (aggregation-time escape hatch, no pushdown). + for f in expressions: + cols = {c: data_dict[c] for c in variables if c in data_dict} + try: + emask = evaluate_filter_expression(f["expression"], cols) + except NameError as e: + raise NameError( + f"expression filter {f['expression']!r} references an undefined name: {e}" + ) from e + if emask.shape != data_dict["leaf_id"].shape: + raise ValueError( + f"expression filter {f['expression']!r} must yield a per-row " + f"boolean mask (got shape {emask.shape})" + ) + if np.sum(emask) == 0: + return None + data_dict = {k: v[emask] for k, v in data_dict.items()} + + if arrow: + import pyarrow as pa + + return pa.table(data_dict) + return pd.DataFrame(data_dict) + + def _read_group(h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False): """Read and spatially filter one HDF5 group. @@ -821,15 +1063,50 @@ def _read_group(h5obj, group: str, data_source: dict, shard_key: int, grid, arro ``pyarrow.Table`` carrying the identical columns. Returns ``None`` when the group has no observations in this shard. - Supports two modes (issue #43, Phase B): + Supports three modes (issues #43 Phase A/B/C): *Flat* (no ``levels``/``base_level`` in ``data_source``): unchanged from Phase A — all structured filters are applied directly to base-rate data. - *Hierarchical* (``levels`` + ``base_level`` present): structured filters whose - normalized ``level`` key names a non-base level are applied at coarse rate, then - expanded to base-rate via the level's ``link`` arrays (``_expand_mask_to_base``). - Base-level structured filters and expression filters are unchanged. + *Hierarchical filtering* (``levels`` + ``base_level`` present): structured + filters whose normalized ``level`` key names a non-base level are applied at + coarse rate, then expanded to base-rate via the level's ``link`` arrays + (``_expand_mask_to_base``). Base-level structured filters and expression + filters are unchanged. + + *Hierarchical (planned) read* (``read_plan.spatial_index`` set, in addition + to ``levels``/``base_level``): the AOI bbox is computed from the grid's + shard footprint, the coarse-level spatial-index coordinates are read fully + (cheap), and base-rate coords + variables + filter datasets are read only + over the planned hyperslices via :func:`zagg.read_plan.execute_read_plan`. + Empty-AOI groups short-circuit to ``None``. Selectivity above the configured + threshold falls back to the full-read path; the planned and full paths + produce row-for-row identical output (#43 Phase C parity). + """ + if ( + isinstance(data_source.get("read_plan"), dict) + and data_source["read_plan"].get("spatial_index") + and data_source.get("levels") + and data_source.get("base_level") + ): + return _planned_read_group(h5obj, group, data_source, shard_key, grid, arrow=arrow) + return _read_group_full(h5obj, group, data_source, shard_key, grid, arrow=arrow) + + +def _read_group_full( + h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False +): + """Full-coord-read variant of :func:`_read_group` (the pre-#49-Phase-C path). + + Reads the base-rate coordinate arrays in full, computes the spatial mask, + then hyperslices variables + base-level filter datasets to the matched + ``[min_idx, max_idx]`` range. Cross-level structured filters are read fully + at coarse rate and expanded to base-rate via ``_expand_mask_to_base``. + Expression filters apply over already-read variable columns. + + Kept as the explicit fallback for: groups whose ``data_source`` declares no + ``read_plan.spatial_index``; ``plan_read``'s selectivity fallback + (``full_read=True``); and the legacy flat (no-levels) form. """ coordinates = data_source["coordinates"] variables = data_source["variables"] diff --git a/src/zagg/read_plan.py b/src/zagg/read_plan.py index 6756dd59..122a34d3 100644 --- a/src/zagg/read_plan.py +++ b/src/zagg/read_plan.py @@ -28,8 +28,10 @@ class ReadPlan: base_slices : list of (int, int) Corresponding base-level ``[start, end)`` half-open slices. chunk_lists : list of list of (int, int) - h5coro-style ``[(start, end)]`` hyperslice lists (inclusive end, 0-based) - for each run. + h5coro-style ``[(start, end)]`` hyperslice lists (half-open + ``[start, end)``, 0-based — matches h5coro's ``h5dataset.py`` "must + provide as list of ranges [x,y)" contract). Mirrors ``base_slices`` + one-to-one. coarse_flag_ranges : list of (int, int) Reserved for future use (coarse-level flag read ranges). full_read : bool @@ -159,7 +161,7 @@ def plan_read( if base_end <= base_start: continue base_slices.append((base_start, base_end)) - chunk_lists.append([(base_start, base_end - 1)]) # h5coro inclusive end + chunk_lists.append([(base_start, base_end)]) # h5coro half-open [start, end) total_base += base_end - base_start # -- Selectivity fallback -- @@ -167,7 +169,7 @@ def plan_read( return ReadPlan( parent_runs=[(0, n_coarse - 1)], base_slices=[(0, n_base)], - chunk_lists=[[(0, n_base - 1)]], + chunk_lists=[[(0, n_base)]], full_read=True, ) diff --git a/tests/test_processing.py b/tests/test_processing.py index 4b333355..acb22b6f 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1661,3 +1661,232 @@ def test_flat_form_unchanged(self): } df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) assert df["h"].tolist() == [1.0, 3.0] + + +# --------------------------------------------------------------------------- +# Planned-read path (issue #43, Phase C — read_plan wiring into _read_group) +# --------------------------------------------------------------------------- + + +class _BboxGrid: + """Permissive grid stub: ``shard_footprint`` returns the bbox polygon, + every photon read is in shard 0. Keeps tests focused on what the planned + read returns (the IO-bounded slice + filters), not on a spatial-mask + re-filter we'd need to model separately. + """ + + def __init__(self, bbox, shard_key=0): + from shapely.geometry import box as _box + + self.bbox = tuple(float(v) for v in bbox) + self._poly = _box(*self.bbox) + self._shard_key = shard_key + + def shard_footprint(self, shard_key): + return self._poly + + def assign(self, lats, lons): + return np.arange(len(lats)) + + def shards_of(self, leaf_ids): + return np.full(len(leaf_ids), self._shard_key, dtype=int) + + +class _LatBboxGrid(_BboxGrid): + """Strict variant: ``shards_of`` keeps a photon only when its lat (carried + via ``assign``'s returned leaf id) falls inside the bbox lat range. Used + for the planned-vs-full parity test, where the spatial mask must agree + between paths.""" + + def assign(self, lats, lons): + # Stash lat as the leaf id; `shards_of` decodes it. Works because the + # test fixture has distinct lats. Real grids use cell ids. + return np.asarray(lats, dtype=np.float64) + + def shards_of(self, leaf_ids): + min_lon, min_lat, max_lon, max_lat = self.bbox + in_shard = (leaf_ids >= min_lat) & (leaf_ids <= max_lat) + out = np.full(len(leaf_ids), -1, dtype=int) + out[in_shard] = self._shard_key + return out + + +def _planned_read_data_source(*, with_base_filter=False, with_coarse_filter=False): + """Multi-level data source for the planned-read tests. + + Six segments × 2 photons/segment = 12 photons. The segment-level + rep-point coordinates live at /seg/lat,/seg/lon; the base-level photon + coords at /heights/lat_ph,/heights/lon_ph; the link arrays at + /seg/ph_index_beg + /seg/segment_ph_cnt (0-based contiguous). + """ + ds = { + "coordinates": { + "latitude": "/heights/lat_ph", + "longitude": "/heights/lon_ph", + }, + "variables": {"h": "/heights/h"}, + "base_level": "photons", + "levels": { + "photons": { + "path": "/heights", + "coordinates": {"latitude": "lat_ph", "longitude": "lon_ph"}, + "variables": {"h": "h"}, + "link": None, + }, + "segments": { + "path": "/seg", + "coordinates": {"latitude": "lat", "longitude": "lon"}, + "variables": {}, + "link": { + "to": "photons", + "index_beg": "/seg/ph_index_beg", + "count": "/seg/segment_ph_cnt", + "index_base": 0, + }, + }, + }, + "read_plan": {"spatial_index": "segments", "pad": 0}, + } + filters = [] + if with_base_filter: + filters.append({"dataset": "/heights/qs", "op": "eq", "value": 0}) + if with_coarse_filter: + filters.append( + {"dataset": "/seg/podppd", "op": "eq", "value": 0, "level": "segments"} + ) + if filters: + ds["filters"] = filters + return ds + + +def _planned_read_h5(*, qs=None, podppd=None): + """Six-segment / 12-photon HDF5 stub with optional base/coarse flag arrays. + + Segments live at lats 0,100,200,300,400,500 (lon 0); photons at lats + 0,50,100,150,200,250,...,550 (lon 0). The wide segment spacing keeps the + ``plan_read`` linestring-crossing check from sweeping unrelated segments + into the matched range -- a narrow bbox between two rep-points stays + bounded by the immediate neighbours. + """ + seg_lats = np.array([0.0, 100.0, 200.0, 300.0, 400.0, 500.0]) + seg_lons = np.zeros(6) + ibeg = np.arange(0, 12, 2, dtype=np.int64) + cnt = np.full(6, 2, dtype=np.int64) + ph_lats = np.array( + [0.0, 50.0, 100.0, 150.0, 200.0, 250.0, 300.0, 350.0, 400.0, 450.0, 500.0, 550.0] + ) + ph_lons = np.zeros(12) + h = np.arange(12.0, dtype=np.float32) * 10.0 + arrays = { + "/seg/lat": seg_lats, + "/seg/lon": seg_lons, + "/seg/ph_index_beg": ibeg, + "/seg/segment_ph_cnt": cnt, + "/heights/lat_ph": ph_lats, + "/heights/lon_ph": ph_lons, + "/heights/h": h, + } + if qs is not None: + arrays["/heights/qs"] = np.asarray(qs) + if podppd is not None: + arrays["/seg/podppd"] = np.asarray(podppd) + return _FakeH5(arrays) + + +class TestPlannedReadGroup: + """Phase C: ``_read_group`` dispatches to ``_planned_read_group`` when + ``data_source.read_plan.spatial_index`` is set, bounding the base-rate IO + via the coarse-level rep-point coords + link arrays. + + The shared fixture lays out 6 segments at lats ``[0, 100, 200, 300, 400, + 500]`` covering 12 photons (2 each). The wide spacing keeps ``plan_read``'s + linestring-crossing sweep bounded: a bbox between two adjacent rep-points + pulls in exactly its two neighbours.""" + + def test_planned_path_bounds_io_to_matched_segments(self): + # Bbox (-0.1, 175, 0.1, 225) directly contains segment 2 (lat=200); + # segment 1's (lat 100 -> 200) linestring crosses the lower edge so + # plan_read sweeps segment 1 in too. Two adjacent segments -> one + # contiguous run -> photons 2..5 in the base array. + ds = _planned_read_data_source() + h5 = _planned_read_h5() + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + assert df["h"].tolist() == [20.0, 30.0, 40.0, 50.0] + + def test_empty_aoi_returns_none(self): + # Bbox far from any segment rep-point or linestring -> no parents + # match -> short-circuit return None before any base-rate read. + ds = _planned_read_data_source() + h5 = _planned_read_h5() + grid = _BboxGrid((10000.0, 10000.0, 10001.0, 10001.0)) + assert _read_group(h5, "gt1l", ds, 0, grid) is None + + def test_full_read_fallback_on_high_selectivity(self): + # full_read_threshold lowered so any plan covering >=10% of n_base + # (>=2/12 photons) triggers the fallback. Same bbox as the basic test + # selects 4/12 = 33% -> falls through to _read_group_full and reads + # everything; the permissive grid keeps all 12. + ds = _planned_read_data_source() + ds["read_plan"]["full_read_threshold"] = 0.1 + h5 = _planned_read_h5() + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + assert df["h"].tolist() == [float(i * 10) for i in range(12)] + + def test_parity_with_full_read(self): + # Both paths produce the same row set when the spatial mask is keyed + # on lat: the planned read narrows IO to photons 2..5 (via plan_read); + # _LatBboxGrid.shards_of further restricts to photons with lat in + # bbox range (photon 4, lat=200). qs drops nothing in-shard. + qs = np.array([0] * 12, dtype=np.int8) + h5 = _planned_read_h5(qs=qs) + grid = _LatBboxGrid((-0.1, 175.0, 0.1, 225.0)) + + ds_planned = _planned_read_data_source(with_base_filter=True) + ds_full = { + "coordinates": { + "latitude": "/heights/lat_ph", + "longitude": "/heights/lon_ph", + }, + "variables": {"h": "/heights/h"}, + "filters": [{"dataset": "/heights/qs", "op": "eq", "value": 0}], + } + + df_planned = _read_group(h5, "gt1l", ds_planned, 0, grid) + df_full = _read_group(h5, "gt1l", ds_full, 0, grid) + # Photon 4 (lat=200, h=40) is the only one in the bbox lat range. + assert df_planned["h"].tolist() == [40.0] + assert df_full["h"].tolist() == [40.0] + + def test_coarse_filter_via_planned_path(self): + # Cross-level (Phase B) filter ANDs with the planned path: drop + # segment 1 via podppd; segment 2 (also pulled in by the linestring + # sweep) survives. Photons 2,3 dropped; 4,5 kept. + ds = _planned_read_data_source(with_coarse_filter=True) + podppd = np.array([0, 1, 0, 0, 0, 0], dtype=np.int8) + h5 = _planned_read_h5(podppd=podppd) + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + assert df["h"].tolist() == [40.0, 50.0] + + def test_pad_extends_selection(self): + # bbox (490..510) covers segment 5 (last, lat=500) directly; segment + # 4's linestring (400 -> 500) crosses the lower edge. With pad=0: + # segments 4,5 -> photons 8..11. With pad=1: segments 3,4,5,6(clamped + # back to 5) -> photons 6..11. + ds = _planned_read_data_source() + ds["read_plan"]["pad"] = 1 + h5 = _planned_read_h5() + grid = _BboxGrid((-0.1, 490.0, 0.1, 510.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + assert df["h"].tolist() == [60.0, 70.0, 80.0, 90.0, 100.0, 110.0] + + def test_invalid_link_target_raises(self): + # The spatial_index level's link must point at the base level. + ds = _planned_read_data_source() + ds["levels"]["segments"]["link"]["to"] = "not_a_level" + h5 = _planned_read_h5() + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + with pytest.raises(ValueError, match="must link directly to base level"): + _read_group(h5, "gt1l", ds, 0, grid) diff --git a/tests/test_read_plan.py b/tests/test_read_plan.py index 83e5358d..cda76326 100644 --- a/tests/test_read_plan.py +++ b/tests/test_read_plan.py @@ -136,7 +136,10 @@ def test_index_base_1_atl03_style(self): class TestExecuteReadPlan: def _make_plan(self, slices, full_read=False): runs = [(s, e - 1) for s, e in slices] - chunks = [[(s, e - 1)] for s, e in slices] + # chunk_lists use h5coro's half-open ``[start, end)`` convention -- + # mirrors base_slices exactly. The runs are inclusive index pairs in + # the coarse array (different domain than chunks). + chunks = [[slc] for slc in slices] return ReadPlan( parent_runs=runs, base_slices=list(slices), @@ -161,9 +164,11 @@ def test_single_slice_reads_correct_range(self): plan = self._make_plan([(10, 20)]) def read_fn(path, hyperslice=None): + # h5coro hyperslice is half-open ``[lo, hi)`` per h5dataset.py + # ("must provide as list of ranges [x,y)") -- Python slicing matches. assert hyperslice is not None lo, hi = hyperslice[0] - return data[lo : hi + 1] # h5coro inclusive end + return data[lo:hi] out = execute_read_plan(plan, read_fn, "/h", np.float32) np.testing.assert_array_equal(out, data[10:20]) @@ -174,7 +179,7 @@ def test_multiple_slices_concatenated(self): def read_fn(path, hyperslice=None): lo, hi = hyperslice[0] - return data[lo : hi + 1] + return data[lo:hi] out = execute_read_plan(plan, read_fn, "/h", np.float32) expected = np.concatenate([data[5:10], data[20:25]]) @@ -185,7 +190,7 @@ def test_full_read_plan_calls_without_hyperslice(self): plan = ReadPlan( parent_runs=[(0, 4)], base_slices=[(0, 50)], - chunk_lists=[[(0, 49)]], + chunk_lists=[[(0, 50)]], full_read=True, ) called_with_none = [] From 5486494cc43f2dcfaf265c6aa18859619b1f09b5 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 17:25:52 +0000 Subject: [PATCH 14/16] phase 6 of issue #30 --- src/zagg/configs/atl03.yaml | 25 +++++++++++++++ src/zagg/configs/atl03_waveform_counts.yaml | 23 +++++++++++++- tests/test_config.py | 35 +++++++++++++++++++++ 3 files changed, 82 insertions(+), 1 deletion(-) diff --git a/src/zagg/configs/atl03.yaml b/src/zagg/configs/atl03.yaml index 7286fb9d..2d9e3d28 100644 --- a/src/zagg/configs/atl03.yaml +++ b/src/zagg/configs/atl03.yaml @@ -14,6 +14,12 @@ # replaces the older flat ``quality_filter`` form to access the ``column`` and # ``op: ne`` knobs (issue #43, Phase A). Grid is rectilinear (HEALPix order-19, # the ~10 m match, waits on mortie #35). +# +# Multi-level form (issue #43 Phase B + C) is declared so the read path can +# bound base-rate IO via the coarse ``segments`` rep-point coords + the +# ``ph_index_beg``/``segment_ph_cnt`` link. Without this, ``lat_ph``/``lon_ph`` +# would be full-read (~245 MB per beam) before any spatial filter, which is the +# dominant driver of the Lambda OOMs on ATL03 region runs (#43 motivation). data_source: reader: h5coro driver: s3 @@ -28,6 +34,25 @@ data_source: column: 0 # land surface type; TEP is uniform across columns op: ne value: -2 + base_level: photons + levels: + photons: + path: "/{group}/heights" + coordinates: {latitude: lat_ph, longitude: lon_ph} + variables: {h_ph: h_ph, signal_conf_ph: signal_conf_ph} + link: null + segments: + path: "/{group}/geolocation" + coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} + variables: {ph_index_beg: ph_index_beg, segment_ph_cnt: segment_ph_cnt} + link: + to: photons + index_beg: "/{group}/geolocation/ph_index_beg" + count: "/{group}/geolocation/segment_ph_cnt" + index_base: 1 # ATL03 ph_index_beg is 1-based per the v3 dict + read_plan: + spatial_index: segments + pad: 1 # one segment of padding on each side per #43 aggregation: variables: diff --git a/src/zagg/configs/atl03_waveform_counts.yaml b/src/zagg/configs/atl03_waveform_counts.yaml index 3d5b4a40..fe63f471 100644 --- a/src/zagg/configs/atl03_waveform_counts.yaml +++ b/src/zagg/configs/atl03_waveform_counts.yaml @@ -23,7 +23,9 @@ # # Confidence filter is the same as atl03.yaml: drop only TEP photons by the # structured ``filters:`` list (signal_conf_ph[:, 0] != -2; TEP is uniform -# across surface types). Grid is rectilinear, matching atl03.yaml. +# across surface types). Multi-level form + ``read_plan`` also identical to +# atl03.yaml so the planned-IO benefits (#43 Phase C) apply to the waveform +# template equally. Grid is rectilinear, matching atl03.yaml. data_source: reader: h5coro driver: s3 @@ -38,6 +40,25 @@ data_source: column: 0 # land surface type; TEP is uniform across columns op: ne value: -2 + base_level: photons + levels: + photons: + path: "/{group}/heights" + coordinates: {latitude: lat_ph, longitude: lon_ph} + variables: {h_ph: h_ph, signal_conf_ph: signal_conf_ph} + link: null + segments: + path: "/{group}/geolocation" + coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} + variables: {ph_index_beg: ph_index_beg, segment_ph_cnt: segment_ph_cnt} + link: + to: photons + index_beg: "/{group}/geolocation/ph_index_beg" + count: "/{group}/geolocation/segment_ph_cnt" + index_base: 1 # ATL03 ph_index_beg is 1-based per the v3 dict + read_plan: + spatial_index: segments + pad: 1 # one segment of padding on each side per #43 aggregation: variables: diff --git a/tests/test_config.py b/tests/test_config.py index 22d0f45f..7ffc3dd9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -142,6 +142,27 @@ def test_rectilinear_grid(self, atl03_config): assert grid["type"] == "rectilinear" assert len(grid["bounds"]) == 4 + def test_multi_level_form_for_planned_reads(self, atl03_config): + # Phase 6: the template declares the ``photons`` (base) + ``segments`` + # (coarse) levels and the link arrays so the #43 Phase C read_plan can + # bound base-rate IO. Without this the ATL03 region runs OOM on Lambda + # (245 MB-per-beam coord-read floor; see #43). + ds = atl03_config.data_source + assert ds["base_level"] == "photons" + levels = ds["levels"] + assert set(levels) == {"photons", "segments"} + assert levels["photons"]["link"] is None + seg_link = levels["segments"]["link"] + assert seg_link["to"] == "photons" + assert seg_link["index_beg"].endswith("ph_index_beg") + assert seg_link["count"].endswith("segment_ph_cnt") + assert seg_link["index_base"] == 1 # ATL03 ph_index_beg is 1-based + + def test_read_plan_targets_segments_level(self, atl03_config): + rp = atl03_config.data_source["read_plan"] + assert rp["spatial_index"] == "segments" + assert rp["pad"] == 1 + # --------------------------------------------------------------------------- # Function resolution @@ -1109,6 +1130,20 @@ def test_rectilinear_grid(self, cfg): assert grid["type"] == "rectilinear" assert len(grid["bounds"]) == 4 + def test_multi_level_form_matches_atl03(self, cfg): + # Phase 6: waveform template carries the same multi-level form + + # read_plan as atl03.yaml so the planned-IO benefits apply equally. + ds = cfg.data_source + assert ds["base_level"] == "photons" + assert set(ds["levels"]) == {"photons", "segments"} + rp = ds["read_plan"] + assert rp["spatial_index"] == "segments" + assert rp["pad"] == 1 + # Cross-check the link is identical to atl03's so the two templates + # share the same plan_read inputs on the same granule. + atl03_link = default_config("atl03").data_source["levels"]["segments"]["link"] + assert ds["levels"]["segments"]["link"] == atl03_link + # --------------------------------------------------------------------------- # Ragged output kind (issue #48, phase 1) From f79ad6b094cbce2043b0ed1d6c18e77f00d32a40 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 17:31:42 +0000 Subject: [PATCH 15/16] fold phase 5 review: antimeridian guard, n_base contiguity, dispatch hardening, +5 tests --- src/zagg/processing.py | 48 +++++++++++++++----- tests/test_processing.py | 97 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 11 deletions(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index d2de78fe..f4ff44a8 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -900,16 +900,30 @@ def _planned_read_group( if len(coarse_lats) == 0: return None - # Under the contiguity assumption (#43): ``sum(count) == n_base``, so the - # last parent's tail gives the total length without an extra header read. - n_base = int(ibeg_arr[-1]) - index_base + int(cnt_arr[-1]) + # ``n_base`` under #43's contiguity assumption ("ranges do not overlap and + # together tile the full base array" -- :func:`_expand_mask_to_base`). + # ``int(cnt_arr.sum())`` makes the assumption explicit and is identical to + # ``ibeg_arr[-1] - index_base + cnt_arr[-1]`` when contiguity holds. If a + # future granule format drops trailing photons or gaps between parents, + # either form under- or over-estimates -- track via a follow-up to #43. + n_base = int(np.asarray(cnt_arr).sum()) if n_base <= 0: return None - # Compute the shard's WGS84 bbox from the grid (every grid's shard_footprint - # returns a shapely Polygon). + # Compute the shard's WGS84 bbox from the grid (every grid's + # ``shard_footprint`` returns a shapely ``Polygon`` or ``MultiPolygon``). + # An antimeridian-crossing HEALPix shard's footprint can come back as a + # split ``MultiPolygon`` (see ``zagg.viz.shardmap._split_antimeridian``), + # in which case ``.bounds`` spans ~360 deg in lon and would neutralize + # the IO bound (the AOI would intersect every segment). Same for + # globe-spanning polar caps. Detect the wide-bbox case up front and fall + # back to ``_read_group_full`` so we don't pretend to optimize. poly = grid.shard_footprint(shard_key) min_lon, min_lat, max_lon, max_lat = poly.bounds + if (max_lon - min_lon) >= 180.0: + # Hand off to the full-read path; the planned-IO benefit is gone for + # this shard and trying to plan would waste the coarse-coord read. + return _read_group_full(h5obj, group, data_source, shard_key, grid, arrow=arrow) bbox = (float(min_lon), float(min_lat), float(max_lon), float(max_lat)) plan = plan_read( @@ -1083,12 +1097,24 @@ def _read_group(h5obj, group: str, data_source: dict, shard_key: int, grid, arro threshold falls back to the full-read path; the planned and full paths produce row-for-row identical output (#43 Phase C parity). """ - if ( - isinstance(data_source.get("read_plan"), dict) - and data_source["read_plan"].get("spatial_index") - and data_source.get("levels") - and data_source.get("base_level") - ): + rp = data_source.get("read_plan") + levels = data_source.get("levels") + base_level = data_source.get("base_level") + # Truthy-checking ``levels``/``base_level`` would route an empty ``{}`` (a + # config typo, easy to do) back to the full-read path silently. Reject + # incomplete configurations explicitly instead -- the planned path is + # gated only when ``spatial_index`` is set, and *then* requires a real + # multi-level structure to operate on. + if isinstance(rp, dict) and rp.get("spatial_index"): + if not isinstance(levels, dict) or not levels: + raise ValueError( + "data_source.read_plan.spatial_index requires a non-empty " + "'levels' mapping" + ) + if not base_level: + raise ValueError( + "data_source.read_plan.spatial_index requires 'base_level'" + ) return _planned_read_group(h5obj, group, data_source, shard_key, grid, arrow=arrow) return _read_group_full(h5obj, group, data_source, shard_key, grid, arrow=arrow) diff --git a/tests/test_processing.py b/tests/test_processing.py index acb22b6f..52255769 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1890,3 +1890,100 @@ def test_invalid_link_target_raises(self): grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) with pytest.raises(ValueError, match="must link directly to base level"): _read_group(h5, "gt1l", ds, 0, grid) + + def test_multi_slice_plan_global_idx_alignment(self): + # Force a plan with two disjoint base-slices (one ATL03 track that + # crosses the AOI lat band twice). Fixture: 10 segments × 1 photon + # each, lats wave from 0 -> 100 -> 0 -> 100 -> 0 so plan_read's + # linestring + containment check picks up segments {2,3,4} and + # {6,7,8} but not {0,1,5,9}. The plan therefore has two runs and + # ``global_idx = [2,3,4,6,7,8]``. A cross-level podppd filter that + # drops segment 3 must align correctly through that global_idx + # (otherwise photon 3's drop hits the wrong row). + seg_lats = np.array([0.0, 0.0, 50.0, 100.0, 100.0, 0.0, 0.0, 100.0, 100.0, 0.0]) + seg_lons = np.zeros(10) + ibeg = np.arange(10, dtype=np.int64) + cnt = np.ones(10, dtype=np.int64) + ph_lats = seg_lats.copy() + ph_lons = np.zeros(10) + h = np.arange(10.0, dtype=np.float32) * 10.0 + podppd = np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0], dtype=np.int8) + h5 = _FakeH5( + { + "/seg/lat": seg_lats, + "/seg/lon": seg_lons, + "/seg/ph_index_beg": ibeg, + "/seg/segment_ph_cnt": cnt, + "/seg/podppd": podppd, + "/heights/lat_ph": ph_lats, + "/heights/lon_ph": ph_lons, + "/heights/h": h, + } + ) + ds = _planned_read_data_source(with_coarse_filter=True) + grid = _BboxGrid((-0.1, 95.0, 0.1, 105.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + # plan covers segments {2,3,4} and {6,7,8} -> base_slices [(2,5),(6,9)] + # -> global_idx [2,3,4,6,7,8] -> base photons h = [20,30,40,60,70,80] + # cross-level drops segment 3 -> drop photon 3 (h=30) only. + assert df["h"].tolist() == [20.0, 40.0, 60.0, 70.0, 80.0] + + def test_full_read_fallback_carries_filters(self): + # The selectivity-fallback path must produce the same row set as the + # planned path would, including base-level structured filters: drop + # photons 5,6 via qs=1 below. Bbox + low threshold -> fallback to + # _read_group_full, which still applies the qs filter. + ds = _planned_read_data_source(with_base_filter=True) + ds["read_plan"]["full_read_threshold"] = 0.1 + qs = np.array([0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0], dtype=np.int8) + h5 = _planned_read_h5(qs=qs) + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + # Full-coord path keeps all 12 photons (permissive grid); qs drops 5,6. + assert df["h"].tolist() == [0.0, 10.0, 20.0, 30.0, 40.0, 70.0, 80.0, 90.0, 100.0, 110.0] + + def test_antimeridian_grid_falls_back_to_full_read(self): + # A grid whose shard_footprint spans ~360 deg in lon (HEALPix + # antimeridian / polar cap) gives plan_read a useless bbox. The + # planned path detects this and falls back so the full-read path is + # used instead -- otherwise the AOI would intersect every segment. + ds = _planned_read_data_source() + h5 = _planned_read_h5() + grid = _BboxGrid((-180.0, -10.0, 180.0, 10.0)) # 360 deg lon span + df = _read_group(h5, "gt1l", ds, 0, grid) + # Full-coord path with permissive grid keeps all 12 photons. + assert df["h"].tolist() == [float(i * 10) for i in range(12)] + + def test_dispatch_rejects_empty_levels(self): + # An incomplete config -- ``read_plan.spatial_index`` set but + # ``levels`` empty -- raises rather than silently routing to the + # full-read path (which would pretend nothing was wrong). + ds = _planned_read_data_source() + ds["levels"] = {} + h5 = _planned_read_h5() + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + with pytest.raises(ValueError, match="non-empty 'levels' mapping"): + _read_group(h5, "gt1l", ds, 0, grid) + + def test_parity_with_full_read_includes_leaf_id(self): + # Strengthen the parity check: row ORDER (via leaf_id) must agree + # between paths, not just the value column. + qs = np.array([0] * 12, dtype=np.int8) + h5 = _planned_read_h5(qs=qs) + grid = _LatBboxGrid((-0.1, 175.0, 0.1, 225.0)) + + ds_planned = _planned_read_data_source(with_base_filter=True) + ds_full = { + "coordinates": { + "latitude": "/heights/lat_ph", + "longitude": "/heights/lon_ph", + }, + "variables": {"h": "/heights/h"}, + "filters": [{"dataset": "/heights/qs", "op": "eq", "value": 0}], + } + + df_planned = _read_group(h5, "gt1l", ds_planned, 0, grid) + df_full = _read_group(h5, "gt1l", ds_full, 0, grid) + # Same row -- leaf_id == lat under _LatBboxGrid.assign. + assert df_planned["leaf_id"].tolist() == df_full["leaf_id"].tolist() + assert df_planned["h"].tolist() == df_full["h"].tolist() From 9741eaee8c50dd15939538f31b7a048947eb012c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 17:34:59 +0000 Subject: [PATCH 16/16] fold phase 6 review: drop misleading level vars, integration test for shipped template --- src/zagg/configs/atl03.yaml | 8 +++- src/zagg/configs/atl03_waveform_counts.yaml | 4 +- tests/test_config.py | 10 ++-- tests/test_processing.py | 52 +++++++++++++++++++++ 4 files changed, 67 insertions(+), 7 deletions(-) diff --git a/src/zagg/configs/atl03.yaml b/src/zagg/configs/atl03.yaml index 2d9e3d28..984697aa 100644 --- a/src/zagg/configs/atl03.yaml +++ b/src/zagg/configs/atl03.yaml @@ -35,16 +35,20 @@ data_source: op: ne value: -2 base_level: photons + # ``coordinates`` uses bare names joined to the level's ``path``; link + # ``index_beg`` / ``count`` use absolute path templates so they remain + # readable when the link arrays live outside the level's path (the read + # path resolves both forms -- see ``_level_coord_paths`` in processing.py). + # ``variables`` at the level scope is documentation only; the read path + # consults the top-level ``data_source.variables`` + filter datasets. levels: photons: path: "/{group}/heights" coordinates: {latitude: lat_ph, longitude: lon_ph} - variables: {h_ph: h_ph, signal_conf_ph: signal_conf_ph} link: null segments: path: "/{group}/geolocation" coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} - variables: {ph_index_beg: ph_index_beg, segment_ph_cnt: segment_ph_cnt} link: to: photons index_beg: "/{group}/geolocation/ph_index_beg" diff --git a/src/zagg/configs/atl03_waveform_counts.yaml b/src/zagg/configs/atl03_waveform_counts.yaml index fe63f471..881b2e05 100644 --- a/src/zagg/configs/atl03_waveform_counts.yaml +++ b/src/zagg/configs/atl03_waveform_counts.yaml @@ -41,16 +41,16 @@ data_source: op: ne value: -2 base_level: photons + # See atl03.yaml for the rationale behind the bare-name coordinates + + # absolute link path templates + documentation-only level ``variables``. levels: photons: path: "/{group}/heights" coordinates: {latitude: lat_ph, longitude: lon_ph} - variables: {h_ph: h_ph, signal_conf_ph: signal_conf_ph} link: null segments: path: "/{group}/geolocation" coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} - variables: {ph_index_beg: ph_index_beg, segment_ph_cnt: segment_ph_cnt} link: to: photons index_beg: "/{group}/geolocation/ph_index_beg" diff --git a/tests/test_config.py b/tests/test_config.py index 7ffc3dd9..7d137e58 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1139,10 +1139,14 @@ def test_multi_level_form_matches_atl03(self, cfg): rp = ds["read_plan"] assert rp["spatial_index"] == "segments" assert rp["pad"] == 1 - # Cross-check the link is identical to atl03's so the two templates - # share the same plan_read inputs on the same granule. + # Cross-check only the fields that drive ``plan_read`` parity: the + # link's source/target arrays + index_base. Other level fields + # (documentation-only ``variables``, coord names, formatting) can + # legitimately diverge across templates without affecting the plan. atl03_link = default_config("atl03").data_source["levels"]["segments"]["link"] - assert ds["levels"]["segments"]["link"] == atl03_link + wf_link = ds["levels"]["segments"]["link"] + for key in ("to", "index_beg", "count", "index_base"): + assert wf_link[key] == atl03_link[key], f"link.{key} diverges from atl03.yaml" # --------------------------------------------------------------------------- diff --git a/tests/test_processing.py b/tests/test_processing.py index 52255769..657fb001 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -1987,3 +1987,55 @@ def test_parity_with_full_read_includes_leaf_id(self): # Same row -- leaf_id == lat under _LatBboxGrid.assign. assert df_planned["leaf_id"].tolist() == df_full["leaf_id"].tolist() assert df_planned["h"].tolist() == df_full["h"].tolist() + + def test_shipped_atl03_template_through_planned_read(self): + # End-to-end coverage of the shipped ``atl03`` template against an + # ATL03-shaped ``_FakeH5`` stub: real ``{group}`` path templates, + # ``index_base: 1`` arithmetic, ``pad: 1``, and the 2-D + # ``signal_conf_ph`` TEP filter all run through ``_planned_read_group`` + # together. This is the integration test the phase 6 review flagged + # was missing -- without it a future YAML edit dropping ``index_base`` + # (defaulting to 0) lands green on the synthetic ``_planned_read_h5`` + # fixture even though it'd be wrong on real ATL03. + from zagg.config import default_config + + cfg = default_config("atl03") + ds = cfg.data_source + # 4 segments × 2 photons = 8 photons, 1-based ph_index_beg. + seg_lats = np.array([0.0, 100.0, 200.0, 300.0]) + seg_lons = np.zeros(4) + ibeg = np.array([1, 3, 5, 7], dtype=np.int64) # 1-based per ATL03 v3 dict + cnt = np.array([2, 2, 2, 2], dtype=np.int64) + ph_lats = np.array([0.0, 50.0, 100.0, 150.0, 200.0, 250.0, 300.0, 350.0]) + ph_lons = np.zeros(8) + h_ph = np.arange(8.0, dtype=np.float32) * 10.0 + # 2-D signal_conf_ph (n_photons × 5 surface types). One TEP photon + # (uniform -2) at index 4 -- filter must drop it. + signal_conf = np.full((8, 5), 4, dtype=np.int8) + signal_conf[4, :] = -2 + h5 = _FakeH5( + { + "/gt1l/heights/lat_ph": ph_lats, + "/gt1l/heights/lon_ph": ph_lons, + "/gt1l/heights/h_ph": h_ph, + "/gt1l/heights/signal_conf_ph": signal_conf, + "/gt1l/geolocation/reference_photon_lat": seg_lats, + "/gt1l/geolocation/reference_photon_lon": seg_lons, + "/gt1l/geolocation/ph_index_beg": ibeg, + "/gt1l/geolocation/segment_ph_cnt": cnt, + } + ) + # Bbox around segment 2 (lat=200); pad=1 (the shipped default) widens + # the matched parent run by one on each side. The permissive grid + # accepts every photon read so the planned path's output is the + # post-filter photon set. + grid = _BboxGrid((-0.1, 175.0, 0.1, 225.0)) + df = _read_group(h5, "gt1l", ds, 0, grid) + assert df is not None + # Without pad the bbox (lat 175..225) matches seg 2 (lat=200) directly + # and seg 1 via the lstr 1->2 sweep (lat 100->200 crosses y=175); seg + # 2's own lstr 2->3 (200->300) pulls in seg 2 again. With pad=1 the + # run [1, 2] widens to [0, 3]. Plan covers photons 0..7. The 2-D + # signal_conf_ph filter drops photon 4 (uniform TEP -2 across all 5 + # surface types). Survivors: 0, 1, 2, 3, 5, 6, 7. + assert df["h_ph"].tolist() == [0.0, 10.0, 20.0, 30.0, 50.0, 60.0, 70.0]