diff --git a/benchmarks/region_timing.py b/benchmarks/region_timing.py new file mode 100644 index 00000000..151a4e5c --- /dev/null +++ b/benchmarks/region_timing.py @@ -0,0 +1,408 @@ +"""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 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: + + * ``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 != -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.) + +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 +import and construct cleanly and stay ruff-clean regardless. + +Run (in a credentialed session):: + + uv run python benchmarks/region_timing.py --windows 1y --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 + +import numpy as np + +from zagg.config import default_config, get_data_vars +from zagg.runner import agg + +# 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 = { + "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. +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 template x handoff) measurement.""" + + region: str + window: str + template: 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 _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. + + Works for both scalar (1-D) and vector (N-D) variables; the trailing + dimension(s) are preserved so callers can check shapes. + """ + 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' 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] + 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, + date_range: tuple[str, str], + *, + 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_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}/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 + 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(), + 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_waveform_counts", + ) + # Shape check: waveform_counts must carry the 128-element trailing dim. + wf = arrays["pandas"]["waveform_counts"] + if 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}{'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.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}" + ) + 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: both). Choices: " + ", ".join(WINDOWS), + ) + 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) + + 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], + version=args.version, + max_cells=args.max_cells, + 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) + if args.out: + with open(args.out, "a") as f: + f.write(table + "\n") + + +if __name__ == "__main__": + main() 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.yaml b/src/zagg/configs/atl03.yaml new file mode 100644 index 00000000..984697aa --- /dev/null +++ b/src/zagg/configs/atl03.yaml @@ -0,0 +1,95 @@ +# 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 drops only TEP (transmit-echo path) photons and keeps +# everything else: signal_conf_ph != -2 (column 0, the land surface type), 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 (n_photons x 5 surface types) and TEP is flagged -2 +# uniformly across every surface type, so the land column (0) is operationally +# equivalent to any other column for the TEP drop. The structured filter list +# 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 + 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" + filters: + - dataset: "/{group}/heights/signal_conf_ph" + column: 0 # land surface type; TEP is uniform across columns + 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} + link: null + segments: + path: "/{group}/geolocation" + coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} + 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: + 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/configs/atl03_waveform_counts.yaml b/src/zagg/configs/atl03_waveform_counts.yaml new file mode 100644 index 00000000..881b2e05 --- /dev/null +++ b/src/zagg/configs/atl03_waveform_counts.yaml @@ -0,0 +1,83 @@ +# 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 h_ph itself. Each bin +# spans 2 m, giving a 256 m window [median(h_ph) - 128 m, median(h_ph) + 128 m). +# Out-of-range photons are dropped (not counted). The companion scalar +# bin_start records median(h_ph) so the absolute elevation of bin 0 can be +# recovered downstream. +# +# Note (issue #30 thread): an earlier draft of this template tried to centre +# on np.median(dem_h), but dem_h lives in /{group}/geophys_corr/dem_h at the +# segment level (one value per ~100 photons), not the heights group; centring +# 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). 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 + 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" + filters: + - dataset: "/{group}/heights/signal_conf_ph" + column: 0 # land surface type; TEP is uniform across columns + 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} + link: null + segments: + path: "/{group}/geolocation" + coordinates: {latitude: reference_photon_lat, longitude: reference_photon_lon} + 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: + waveform_counts: + kind: vector + trailing_shape: 128 + expression: "np.histogram(h_ph - np.median(h_ph), 128, (-128.0, 128.0))[0].astype('uint32')" + source: h_ph + dtype: uint32 + fill_value: 0 + bin_start: + function: median + 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 4c22c5d9..f4ff44a8 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,261 @@ 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 + + # ``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`` 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( + 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 +1077,62 @@ 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). + """ + 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) + + +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"] @@ -1247,7 +1550,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 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/src/zagg/runner.py b/src/zagg/runner.py index 28028f3d..d2144234 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, ) # write_dataframe_to_zarr no-ops on an empty carrier (DataFrame or Arrow # table), so no carrier-specific emptiness check is needed here. @@ -298,7 +306,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"]) @@ -344,7 +353,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 } diff --git a/tests/test_config.py b/tests/test_config.py index a5cdf842..7d137e58 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -99,6 +99,71 @@ 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_tep(self, atl03_config): + # The ATL03 template carries one structured TEP filter: keep photons where + # signal_conf_ph[:, 0] (land surface type) != -2. TEP is uniform across + # surface types per the ATL03 v3 data dictionary, so column 0 is + # operationally equivalent to any other column for the TEP drop. + filters = atl03_config.data_source["filters"] + assert len(filters) == 1 + f = filters[0] + assert f["value"] == -2 + assert f["op"] == "ne" # keep signal_conf_ph != -2 (drop only TEP) + assert f["column"] == 0 + assert f["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 + + 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 # --------------------------------------------------------------------------- @@ -972,6 +1037,118 @@ def test_vector_list_signature(self): assert sig["inner_shape"] == () +# --------------------------------------------------------------------------- +# 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_only(self, cfg): + # Option A: the histogram is centered on np.median(h_ph), so dem_h is + # not needed (and the segment-level ``geophys_corr/dem_h`` path was the + # wrong group anyway -- see #30 thread). + ds_vars = cfg.data_source["variables"] + assert "h_ph" in ds_vars + assert "dem_h" not in ds_vars + assert ds_vars["h_ph"].endswith("h_ph") + + 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 ±128 m of their own median; all should be counted. + from zagg.processing import calculate_cell_statistics + + np.random.seed(0) + h_ph = np.random.uniform(-100.0, 100.0, 50).astype("float32") + result = calculate_cell_statistics( + {"h_ph": h_ph, "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 + + # Two photons clustered near 0, one far outlier at 500 m. The cell median + # is ~5 m, so the outlier sits beyond ±128 m and falls outside the hist. + h_ph = np.array([0.0, 10.0, 500.0], dtype="float32") + result = calculate_cell_statistics( + {"h_ph": h_ph, "leaf_id": np.arange(3)}, config=cfg + ) + wc = result["waveform_counts"] + assert int(wc.sum()) == 2, "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([]), "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): + # Both templates carry the same TEP filter expressed in the structured + # ``filters:`` list form (op: ne, value: -2, column: 0 -- land surface + # type; TEP is uniform across columns per the v3 data dictionary). + filters = cfg.data_source["filters"] + assert len(filters) == 1 + f = filters[0] + assert f["op"] == "ne" + assert f["value"] == -2 + assert f["column"] == 0 + assert f["dataset"].endswith("signal_conf_ph") + + def test_rectilinear_grid(self, cfg): + grid = cfg.output["grid"] + 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 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"] + 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" + + # --------------------------------------------------------------------------- # Ragged output kind (issue #48, phase 1) # --------------------------------------------------------------------------- diff --git a/tests/test_processing.py b/tests/test_processing.py index d153c061..657fb001 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( { @@ -1623,3 +1661,381 @@ 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) + + 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() + + 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] 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 = [] 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"