diff --git a/MANIFEST.in b/MANIFEST.in index 365e64b..ae0bbb3 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -5,3 +5,9 @@ prune data prune docs prune workflows exclude uv.lock + +# dist/ is committed to git, so setuptools-scm's file finder would otherwise +# pack every previously released wheel and tarball into the new sdist -- each +# build roughly doubling the one before it. +prune dist +global-exclude *.whl *.tar.gz diff --git a/dist/fimbox-0.1.12-py3-none-any.whl b/dist/fimbox-0.1.12-py3-none-any.whl new file mode 100644 index 0000000..ac75a15 Binary files /dev/null and b/dist/fimbox-0.1.12-py3-none-any.whl differ diff --git a/dist/fimbox-0.1.12.tar.gz b/dist/fimbox-0.1.12.tar.gz new file mode 100644 index 0000000..a0e14b0 Binary files /dev/null and b/dist/fimbox-0.1.12.tar.gz differ diff --git a/docs/images/fimbox.png b/docs/images/fimbox.png old mode 100755 new mode 100644 index 50b874d..be12b79 Binary files a/docs/images/fimbox.png and b/docs/images/fimbox.png differ diff --git a/pyproject.toml b/pyproject.toml index 6799985..b9aea8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "fimbox" -version = "0.1.11" +version = "0.1.12" description = "A modular open-source testbed framework to standardize Flood Inundation Mapping simulations and evaluation with custom datasets and hydrologic parameters with reproducible workflows." authors = [ { name = "Surface Dynamics Modeling Lab (SDML), The University of Alabama" }, diff --git a/src/fimbox/_dask.py b/src/fimbox/_dask.py index 1f017a2..3b86453 100644 --- a/src/fimbox/_dask.py +++ b/src/fimbox/_dask.py @@ -18,9 +18,11 @@ The new defaults: -1. ``n_workers`` is the minimum of CPU count and (free RAM GB / 6), - floored at 2. On a 32 GB / 8-core box that gives 5 workers with - ~6 GB each — enough for a HUC8 branch with margin. +1. ``n_workers`` is the minimum of CPU count and (RAM GB / 8), floored + at 2. On a 64 GB / 12-core box that gives 8 workers with ~8 GB each + — enough for a HUC8 branch with margin. A caller-supplied count is + capped by the same ceiling (see ``fimbox._workers``), so + ``n_workers=64`` on a laptop lands on what the laptop can feed. 2. Dask's "pause-at-80% / terminate-at-95%" governor is disabled. Branches run with stable RSS once the rasters are loaded; the transient peaks were causing Dask to pause workers, which made the @@ -46,6 +48,8 @@ import threading from typing import Optional +from ._workers import resolve_workers, system_ram_gb + log = logging.getLogger(__name__) _lock = threading.Lock() @@ -62,26 +66,22 @@ def _system_ram_gb() -> float: """Best-effort total RAM in GB. Returns 8.0 if it can't be detected.""" - try: - import psutil # type: ignore + return system_ram_gb() - return psutil.virtual_memory().total / (1024**3) - except Exception: - pass - # POSIX fallback (Linux + macOS): sysconf - try: - return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024**3) - except (ValueError, AttributeError, OSError): - return 8.0 +def _ram_per_branch_gb() -> float: + return float( + os.environ.get("FIMBOX_DASK_RAM_PER_BRANCH", _DEFAULT_RAM_PER_BRANCH_GB) + ) -def _resolve_n_workers() -> int: - """Pick worker count that fits comfortably in RAM. - Explicit env override wins. Otherwise: min(cpu_count, RAM_GB / - RAM_PER_BRANCH_GB), floored at 2 so very small machines still get - some parallelism, capped at cpu_count so we don't oversubscribe - cores. +def _resolve_n_workers(requested: Optional[int] = None) -> int: + """Worker count that fits comfortably in RAM. + + ``FIMBOX_DASK_WORKERS`` still wins outright — it's the escape hatch for + people who know their box. Otherwise ``requested`` is honoured only as far + as CPU and RAM allow: asking for 64 workers on a 12-core laptop gets you 12 + fighting over 8 GB slices, which is slower than the 8 the machine can feed. """ env = os.environ.get("FIMBOX_DASK_WORKERS") if env: @@ -92,23 +92,11 @@ def _resolve_n_workers() -> int: except ValueError: log.warning("FIMBOX_DASK_WORKERS=%r is not an int; auto-sizing", env) - cpu = max(1, os.cpu_count() or 1) - ram_gb = _system_ram_gb() - - per_branch = float( - os.environ.get("FIMBOX_DASK_RAM_PER_BRANCH", _DEFAULT_RAM_PER_BRANCH_GB) - ) - by_ram = max(2, int(ram_gb // per_branch)) - n = min(cpu, by_ram) - log.info( - "Dask worker sizing: cpu_count=%d, system_ram=%.1f GB, " - "ram_per_branch=%.1f GB => n_workers=%d", - cpu, - ram_gb, - per_branch, - n, + return resolve_workers( + requested, + ram_per_worker_gb=_ram_per_branch_gb(), + label="Dask worker sizing", ) - return n def _resolve_memory_limit(): @@ -153,7 +141,8 @@ def get_client(n_workers: Optional[int] = None): Pass ``n_workers`` to override the env/CPU default on first call only — once the cluster exists, later calls return the same client - regardless of the argument. + regardless of the argument. The request is bounded by what the machine + can feed; see :func:`_resolve_n_workers`. """ global _client, _cluster @@ -168,7 +157,7 @@ def get_client(n_workers: Optional[int] = None): from distributed import Client, LocalCluster - n = n_workers if n_workers else _resolve_n_workers() + n = _resolve_n_workers(n_workers) mem = _resolve_memory_limit() log.info( diff --git a/src/fimbox/_workers.py b/src/fimbox/_workers.py new file mode 100644 index 0000000..b446d1f --- /dev/null +++ b/src/fimbox/_workers.py @@ -0,0 +1,113 @@ +""" +One place that answers "how many workers can this machine actually run?". + +Callers hand over whatever the user gave them — nothing, ``None``, ``0``, or a +number picked optimistically — and get back a count the box can sustain: + +* ``None`` / ``0`` / negative -> auto-size to the machine (the common case) +* ``1`` -> serial, always honoured; it's how you debug +* anything bigger -> honoured up to what CPU + RAM allow, then clamped +* never more workers than there are branches to chew through + +RAM is the binding constraint, not core count. Every branch worker is its own +process holding its own rasters and tables, so 16 workers on a 16 GB laptop buy +swap thrash instead of speed — auto-sizing is ``min(cpu, RAM / per_worker)``. +The per-worker budget differs by workload, hence the three constants below. + +Advanced overrides: +- ``FIMBOX_WORKERS`` — pin the count for every pool (still capped by task count) +- ``FIMBOX_RAM_PER_WORKER`` — GB to budget per worker when auto-sizing +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional + +log = logging.getLogger(__name__) + +# Per-worker RAM budgets, by how raster-hungry the step is. Branch processing +# peaks hardest (AGREE DEM conditioning on a HUC8 at 10m), inundation loads a +# HAND + catchment raster pair, and the SRC/calibration steps are pandas tables. +RAM_PER_WORKER_BRANCH_GB = 8.0 +RAM_PER_WORKER_FIM_GB = 4.0 +RAM_PER_WORKER_TABLE_GB = 2.0 + + +def system_ram_gb() -> float: + """Best-effort total RAM in GB. Returns 8.0 if it can't be detected.""" + try: + import psutil # type: ignore + + return psutil.virtual_memory().total / (1024**3) + except Exception: + pass + # POSIX fallback (Linux + macOS): sysconf + try: + return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024**3) + except (ValueError, AttributeError, OSError): + return 8.0 + + +def _cpu_count() -> int: + # process_cpu_count honours cgroup/affinity limits (containers, SLURM); + # cpu_count reports the whole host and would oversubscribe a cpu-limited job. + getter = getattr(os, "process_cpu_count", None) + return max(1, (getter() if getter else os.cpu_count()) or 1) + + +def max_supported_workers(ram_per_worker_gb: float = RAM_PER_WORKER_FIM_GB) -> int: + """Most workers this machine can feed without going to swap. + + Floored at 2 so a modest laptop still gets some parallelism, then capped at + the core count so we never oversubscribe CPUs. + """ + per_worker = float(os.environ.get("FIMBOX_RAM_PER_WORKER", ram_per_worker_gb)) + by_ram = max(2, int(system_ram_gb() // max(per_worker, 0.5))) + return min(_cpu_count(), by_ram) + + +def resolve_workers( + requested: Optional[int], + *, + n_tasks: Optional[int] = None, + ram_per_worker_gb: float = RAM_PER_WORKER_FIM_GB, + label: str = "workers", +) -> int: + """Turn a user-supplied worker count into one this machine can honour. + + ``1`` is passed straight through — an explicit request for serial is a + debugging tool, not a number to second-guess. Everything else is bounded by + :func:`max_supported_workers` and by ``n_tasks``, since idle processes still + cost a spawn and a full package re-import. + """ + env = os.environ.get("FIMBOX_WORKERS") + if env: + try: + pinned = int(env) + if pinned > 0: + requested = pinned + except ValueError: + log.warning("FIMBOX_WORKERS=%r is not an int; ignoring", env) + + if requested == 1: + return 1 + + ceiling = max_supported_workers(ram_per_worker_gb) + + if requested is None or requested <= 0: + n = ceiling + why = f"auto-sized (cpu={_cpu_count()}, ram={system_ram_gb():.0f} GB)" + elif requested > ceiling: + n = ceiling + why = f"requested {requested}, clamped to what this machine supports" + else: + n = requested + why = "" # honouring the given number needs no explanation + + if n_tasks is not None: + n = max(1, min(n, n_tasks)) + + log.info("%s: %d worker(s)%s", label, n, f" — {why}" if why else "") + return n diff --git a/src/fimbox/fimgeneration/pipeline.py b/src/fimbox/fimgeneration/pipeline.py index 14e8cdb..d2e6821 100644 --- a/src/fimbox/fimgeneration/pipeline.py +++ b/src/fimbox/fimgeneration/pipeline.py @@ -35,6 +35,7 @@ import pandas as pd +from .._workers import RAM_PER_WORKER_FIM_GB, resolve_workers from ..logging_utils import WATERSHED_DIR_NAME, aoi_root from .inundator import InundationResult, Inundator from .mosaic import BranchMosaic, MosaicResult @@ -107,10 +108,10 @@ class FimGenerator: # Whether to mosaic per-branch outputs into AOI-level rasters at the end. Default True. mosaic: bool = True - # Parallel workers for the per-branch loop. 1 = serial. When use_dask - # is enabled n_workers is ignored — the shared Dask LocalCluster sizes - # itself from the machine. - n_workers: int = 1 + # Parallel workers for the per-branch loop. None/0 = size to the machine, + # 1 = serial, anything above what CPU + RAM support is clamped down. When + # use_dask is enabled this is a ceiling on the shared LocalCluster instead. + n_workers: Optional[int] = None # When True, dispatch the branch loop through the process-wide Dask # LocalCluster (auto-sized to the machine). When None, Dask is used @@ -177,10 +178,13 @@ def run(self) -> FimGenerationResult: ) use_dask = False - # n_workers=None means "auto": let ProcessPoolExecutor size to the - # machine (only an explicit <=1 runs the serial loop). Guarding here - # keeps the comparison from raising on None. - serial = self.n_workers is not None and self.n_workers <= 1 + n_workers = resolve_workers( + self.n_workers, + n_tasks=len(bids), + ram_per_worker_gb=RAM_PER_WORKER_FIM_GB, + label="FimGenerator", + ) + serial = n_workers <= 1 log.info( f"FimGenerator: AOI={self.aoi_dir.name} branches={len(bids)} " @@ -189,7 +193,9 @@ def run(self) -> FimGenerationResult: ) if use_dask: - results = self._run_with_dask(bids, branch_root, forecast_df, tmp_dir) + results = self._run_with_dask( + bids, branch_root, forecast_df, tmp_dir, n_workers + ) elif serial: results = [ _run_one_branch( @@ -205,7 +211,7 @@ def run(self) -> FimGenerationResult: ] else: results = [] - with ProcessPoolExecutor(max_workers=self.n_workers) as pool: + with ProcessPoolExecutor(max_workers=n_workers) as pool: fut_to_bid = { pool.submit( _run_one_branch, @@ -282,12 +288,15 @@ def _run_with_dask( branch_root: Path, forecast_df: pd.DataFrame, tmp_dir: Path, + n_workers: int, ) -> list[InundationResult]: # Submit every branch to the shared LocalCluster. retries=0 # mirrors the process_branches setup: an OOMed branch should # surface immediately as 'failed' so siblings finish instead # of cascading into a KilledWorker storm. - client = _get_dask_client() + # n_workers only bites if this call is what creates the cluster — + # a pool the preprocessing run already warmed up is reused as-is. + client = _get_dask_client(n_workers) log.info( "FimGenerator (dask): %d branches -> %d workers (dashboard %s)", len(bids), @@ -471,7 +480,7 @@ class generateFIM: aoi_dir: PathLike feature_id_csv: Optional[PathLike] = None - n_workers: int = 4 + n_workers: Optional[int] = None # None -> size to the machine int16_mode: bool = True depth: bool = False @@ -616,7 +625,12 @@ def generate(self, discharge_csvs: Sequence[PathLike]) -> list[FimGenerationResu help="Step 2: a single discharge CSV. When omitted, every CSV under " "/discharge-inputs/ is processed.", ) - parser.add_argument("--workers", type=int, default=4) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Parallel branch workers. Omit to size to the machine; 1 for serial.", + ) parser.add_argument( "--int16", action="store_true", diff --git a/src/fimbox/preprocessing/calculate_branch/_d8.py b/src/fimbox/preprocessing/calculate_branch/_d8.py new file mode 100644 index 0000000..1bdf150 --- /dev/null +++ b/src/fimbox/preprocessing/calculate_branch/_d8.py @@ -0,0 +1,53 @@ +"""D8 flow-direction plumbing shared by the raster traversals. + +Both the label propagation in :mod:`gage_catchments` and the flow accumulation in +:mod:`flowacc_dem` need the same thing first: for every cell, where does its water +go. One table, one builder, so the two can't drift apart. +""" + +from __future__ import annotations + +import numpy as np + +# WBT D8 pointer encoding: powers-of-2 --> (row_offset, col_offset) +# 64=N 128=NE 1=E 2=SE 4=S 8=SW 16=W 32=NW +D8_OFFSETS: dict[int, tuple[int, int]] = { + 1: (0, 1), + 2: (1, 1), + 4: (1, 0), + 8: (1, -1), + 16: (0, -1), + 32: (-1, -1), + 64: (-1, 0), + 128: (-1, 1), +} + + +def downstream_index(d8: np.ndarray) -> np.ndarray: + """Flat index of each cell's downstream neighbour; itself where flow stops. + + Nodata must already be zeroed by the caller. Uses shifted slices rather than + per-cell row/col arithmetic, so a HUC8 costs a few views instead of several + whole-grid index copies. + """ + rows, cols = d8.shape + codes = d8.astype(np.int32, copy=False) + flat_base = np.arange(rows * cols, dtype=np.int32).reshape(rows, cols) + ds = flat_base.copy() + + for code, (dr, dc) in D8_OFFSETS.items(): + sel = codes == code + if not sel.any(): + continue + # Only the window whose neighbour lands on the grid can flow; cells + # pointing off the edge keep their self-loop. + r0, r1 = max(0, -dr), rows - max(0, dr) + c0, c1 = max(0, -dc), cols - max(0, dc) + if r0 >= r1 or c0 >= c1: + continue + np.copyto( + ds[r0:r1, c0:c1], + flat_base[r0 + dr : r1 + dr, c0 + dc : c1 + dc], + where=sel[r0:r1, c0:c1], + ) + return ds.ravel() diff --git a/src/fimbox/preprocessing/calculate_branch/calculate_allbranches.py b/src/fimbox/preprocessing/calculate_branch/calculate_allbranches.py index a43be42..84394d3 100644 --- a/src/fimbox/preprocessing/calculate_branch/calculate_allbranches.py +++ b/src/fimbox/preprocessing/calculate_branch/calculate_allbranches.py @@ -54,7 +54,6 @@ from .process_branches import ( AOIProcessingConfig, BranchResult, - _process_single_branch, _resolve_paths, process_branches, ) @@ -104,10 +103,14 @@ def calculate_allbranches( Steps: 1. Write branch_ids.csv with the full branch inventory (branch 0 first, then every id from branch_ids.lst). - 2. Run BranchZero for the whole AOI (branch_id="0") in the main process, - exactly as the step-by-step test does. Skip when ``run_branch_zero=False`` + 2. Publish the shared AOI-root rasters (the clipped DEM and bridge + elevation diff every other branch clips from) in the main process, so + they land before the fan-out. Skip when ``run_branch_zero=False`` (outputs already exist from a previous run). - 3. Dispatch every non-zero branch to Dask in parallel via process_branches. + 3. Dispatch every branch to Dask in parallel via process_branches, branch + zero included: its raster prep and CreateHAND are the longest work in + the run and no sibling reads their outputs, so they overlap with the + siblings instead of holding the pool idle. 4. Apply the AOI-level deny-list cleanup to remove intermediates. Returns an :class:`AllBranchesResult` summarising what got recorded and @@ -147,17 +150,24 @@ def calculate_allbranches( f"(0 + {len(all_ids) - 1} from {Path(branch_list_path).name})" ) - # Step 2: BranchZero (whole-AOI, serial) then CreateHAND on branch 0. - # _process_single_branch skips BranchZero re-run via sentinel file. + # Step 2: publish the shared AOI-root rasters, and nothing more. This is the + # only part of branch zero the siblings genuinely wait on — /dem.tif + # and bridge_elev_diff.tif, which every non-zero branch clips from (and both + # are already in place whenever dem_path is the AOI-root dem, so this usually + # costs nothing). Branch zero's raster conditioning — levee burn, AGREE, pit + # fill, D8 — used to run here too, holding the whole pool idle for a minute + # and a half on a HUC8, even though nothing outside branches/0/ reads any of + # it. It now runs inside branch zero's own pool task in step 3. cfg = _resolve_paths(cfg) - b0_result: Optional[BranchResult] = None + branch_zero_prepared = False if run_branch_zero: dem = cfg.dem_path or (aoi_dir / "dem.tif") if not Path(dem).is_file(): log.info(f"BranchZero skipped — dem not found at {dem}") else: log.info( - f"--- BranchZero (whole-AOI, branch_id={cfg.branch_zero_id!r}) ---" + f"--- Shared AOI rasters (for branch_id={cfg.branch_zero_id!r} " + f"and every sibling) ---" ) BranchZero( dem_path=cfg.dem_path, @@ -178,25 +188,24 @@ def calculate_allbranches( agree_smooth_drop=cfg.agree_smooth_drop, agree_sharp_drop=cfg.agree_sharp_drop, branch_zero_id=cfg.branch_zero_id, - ).run() - log.info(f"--- CreateHAND branch_id={cfg.branch_zero_id!r} ---") - b0_result = _process_single_branch(cfg, cfg.branch_zero_id) - log.info( - f"--- Branch {cfg.branch_zero_id} complete " - f"(status={b0_result.status}) ---" - ) + ).publish_shared_inputs() + branch_zero_prepared = True else: log.info( f"run_branch_zero=False — skipping BranchZero + branch-zero CreateHAND " f"(branch_id={cfg.branch_zero_id!r}); assuming outputs already exist." ) - # Step 3: parallel non-zero branch loop via Dask. - results = process_branches(cfg) - n_non_zero_recorded = sum(1 for r in results if r.status == "ok") + # Step 3: every branch — zero included — through the parallel loop. + all_results = process_branches(cfg, include_branch_zero=branch_zero_prepared) - # Prepend branch-zero result so branch_results covers all branches. - all_results = ([b0_result] if b0_result is not None else []) + results + b0_result: Optional[BranchResult] = next( + (r for r in all_results if r.branch_id == cfg.branch_zero_id), None + ) + if b0_result is not None: + log.info(f"--- Branch {cfg.branch_zero_id} status={b0_result.status} ---") + results = [r for r in all_results if r.branch_id != cfg.branch_zero_id] + n_non_zero_recorded = sum(1 for r in results if r.status == "ok") # AOI-level deny-list cleanup. Default behaviour deletes the AOI # intermediates listed in ``deny_unit.lst``; pass ``delete_deny_list=False`` @@ -262,7 +271,12 @@ def calculate_allbranches( parser.add_argument("--aoi-dir", required=True) parser.add_argument("--aoi-id", required=True) parser.add_argument("--branch-list", default=None) - parser.add_argument("--workers", type=int, default=1) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Parallel branch workers. Omit to size to the machine; 1 for serial.", + ) parser.add_argument( "--deny-unit-list", default=None, diff --git a/src/fimbox/preprocessing/calculate_branch/calculate_branchzero.py b/src/fimbox/preprocessing/calculate_branch/calculate_branchzero.py index aa04092..629f053 100644 --- a/src/fimbox/preprocessing/calculate_branch/calculate_branchzero.py +++ b/src/fimbox/preprocessing/calculate_branch/calculate_branchzero.py @@ -143,6 +143,43 @@ def run(self) -> dict: log.exception("BranchZero failed") raise + def publish_shared_inputs(self) -> dict: + """Clip just the AOI-root rasters — steps 1-2, the only ones anything + outside ``branches/0/`` reads. + + Split out of :meth:`run` so a caller can satisfy the siblings' one real + prerequisite and then start the fan-out. The rest of branch zero's stack + (levee burn, AGREE, pit fill, D8) writes only into its own branch + directory, so it belongs in branch zero's pool task rather than in front + of it. Typically a no-op: when ``dem_path`` already *is* + ``/dem.tif`` both clips skip. + """ + from ...logging_utils import attach_case_log + + attach_case_log(self.out_dir) + crs, res = self._resolve_crs_and_res() + + outputs: dict = {} + dem_clipped = self.out_dir / "dem.tif" + _rasterio_clip_reproject( + self.dem_path, self.boundary_gpkg, dem_clipped, crs=crs, res=res + ) + log.info("DEM clipped --> %s", dem_clipped.name) + outputs["dem"] = dem_clipped + + if self.bridge_elev_diff_path and self.bridge_elev_diff_path.exists(): + bridge_clipped = self.out_dir / "bridge_elev_diff.tif" + _rasterio_clip_reproject( + self.bridge_elev_diff_path, + self.boundary_gpkg, + bridge_clipped, + crs=crs, + res=res, + ) + log.info("Bridge elev diff clipped --> %s", bridge_clipped.name) + outputs["bridge_elev_diff"] = bridge_clipped + return outputs + def _run(self) -> dict: bid = self.branch_zero_id branch_dir = self.out_dir / "branches" / bid @@ -156,36 +193,34 @@ def _run(self) -> dict: # on the AOI-root dem.tif corrupts the shared input. Non-zero # branches write their per-branch clip directly into branch_dir. is_branch_zero = bid == "0" + bridge_clipped: Optional[Path] = None if is_branch_zero: - dem_clipped = self.out_dir / "dem.tif" + shared = self.publish_shared_inputs() + dem_clipped = shared["dem"] + bridge_clipped = shared.get("bridge_elev_diff") else: dem_clipped = branch_dir / f"dem_{bid}.tif" - - _rasterio_clip_reproject( - self.dem_path, self.boundary_gpkg, dem_clipped, crs=crs, res=res - ) - log.info("DEM clipped --> %s", dem_clipped.name) + _rasterio_clip_reproject( + self.dem_path, self.boundary_gpkg, dem_clipped, crs=crs, res=res + ) + log.info("DEM clipped --> %s", dem_clipped.name) + if self.bridge_elev_diff_path and self.bridge_elev_diff_path.exists(): + bridge_clipped = branch_dir / f"bridge_elev_diff_{bid}.tif" + _rasterio_clip_reproject( + self.bridge_elev_diff_path, + self.boundary_gpkg, + bridge_clipped, + crs=crs, + res=res, + ) + log.info("Bridge elev diff clipped --> %s", bridge_clipped.name) dem_branch = branch_dir / f"dem_{bid}.tif" if dem_clipped != dem_branch: shutil.copy2(dem_clipped, dem_branch) - # clip bridge elev diff once, then copy into branch subdirectory - bridge_clipped: Optional[Path] = None bridge_branch: Optional[Path] = None - if self.bridge_elev_diff_path and self.bridge_elev_diff_path.exists(): - if is_branch_zero: - bridge_clipped = self.out_dir / "bridge_elev_diff.tif" - else: - bridge_clipped = branch_dir / f"bridge_elev_diff_{bid}.tif" - _rasterio_clip_reproject( - self.bridge_elev_diff_path, - self.boundary_gpkg, - bridge_clipped, - crs=crs, - res=res, - ) - log.info("Bridge elev diff clipped --> %s", bridge_clipped.name) + if bridge_clipped is not None: bridge_branch = branch_dir / f"bridge_elev_diff_{bid}.tif" if bridge_clipped != bridge_branch: shutil.copy2(bridge_clipped, bridge_branch) diff --git a/src/fimbox/preprocessing/calculate_branch/flowacc_dem.py b/src/fimbox/preprocessing/calculate_branch/flowacc_dem.py index 5c24700..82bf735 100644 --- a/src/fimbox/preprocessing/calculate_branch/flowacc_dem.py +++ b/src/fimbox/preprocessing/calculate_branch/flowacc_dem.py @@ -18,27 +18,16 @@ from __future__ import annotations import logging -from collections import deque from dataclasses import dataclass from pathlib import Path from typing import Optional import numpy as np +from numba import njit -log = logging.getLogger(__name__) +from ._d8 import downstream_index -# WBT D8 pointer encoding: powers-of-2 --> (row_offset, col_offset) -# 64=N 128=NE 1=E 2=SE 4=S 8=SW 16=W 32=NW -_D8_OFFSETS: dict[int, tuple[int, int]] = { - 1: (0, 1), - 2: (1, 1), - 4: (1, 0), - 8: (1, -1), - 16: (0, -1), - 32: (-1, -1), - 64: (-1, 0), - 128: (-1, 1), -} +log = logging.getLogger(__name__) @dataclass @@ -46,8 +35,7 @@ class FlowAccDEM: """ Headwater-weighted D8 flow accumulation using a topological BFS. - No external dependencies beyond numpy and rasterio — avoids the - numba/llvmlite build requirement that pyflwdir carries. + Hand-rolled on numpy + a compiled traversal rather than pulling in pyflwdir. Parameters ---------- @@ -223,38 +211,47 @@ def _d8_flow_accum(d8: np.ndarray, hw: np.ndarray) -> np.ndarray: accum : float32 accumulated headwater count at each cell """ rows, cols = d8.shape - n = rows * cols - flat_d8 = d8.ravel() - - # build flat downstream index; self-loop marks outlet / no-flow - ds = np.arange(n, dtype=np.int32) - for code, (dr, dc) in _D8_OFFSETS.items(): - mask = flat_d8 == code - if not mask.any(): - continue - idxs = np.where(mask)[0] - r = idxs // cols - c = idxs % cols - nr = r + dr - nc = c + dc - valid = (nr >= 0) & (nr < rows) & (nc >= 0) & (nc < cols) - ds[idxs[valid]] = (nr[valid] * cols + nc[valid]).astype(np.int32) - - # in-degree: number of upstream cells draining into each cell - non_self = ds != np.arange(n, dtype=np.int32) + + accum = hw.ravel().astype(np.float32).copy() + _accumulate_downstream(downstream_index(d8), accum) + return accum.reshape(rows, cols) + + +@njit(cache=True) +def _accumulate_downstream(ds: np.ndarray, accum: np.ndarray) -> None: + """Push each cell's running total into its downstream neighbour, Kahn order. + + In-degree and the source list are counted here rather than with bincount so + nothing wider than the grid itself gets allocated — that temporary is over a + gigabyte on a HUC8 and grows with the area. + + A cell is released only once every upstream contributor has been added in, so + each is touched exactly once, and each joins the queue exactly when its + in-degree hits zero — so the preallocated queue can never overflow. + """ + n = ds.size + in_deg = np.zeros(n, dtype=np.int16) - np.add.at(in_deg, ds[non_self], 1) - - # BFS from source cells (nothing flows into them) - accum = hw.ravel().astype(np.float32) - queue: deque[int] = deque(np.where(in_deg == 0)[0].tolist()) - while queue: - i = queue.popleft() - j = int(ds[i]) + for i in range(n): + j = ds[i] + if j != i: + in_deg[j] += 1 + + queue = np.empty(n, dtype=np.int32) + tail = 0 + for i in range(n): + if in_deg[i] == 0: + queue[tail] = i + tail += 1 + + head = 0 + while head < tail: + i = queue[head] + head += 1 + j = ds[i] if j != i: accum[j] += accum[i] in_deg[j] -= 1 if in_deg[j] == 0: - queue.append(j) - - return accum.reshape(rows, cols) + queue[tail] = j + tail += 1 diff --git a/src/fimbox/preprocessing/calculate_branch/gage_catchments.py b/src/fimbox/preprocessing/calculate_branch/gage_catchments.py index 371ce0e..3a9153e 100644 --- a/src/fimbox/preprocessing/calculate_branch/gage_catchments.py +++ b/src/fimbox/preprocessing/calculate_branch/gage_catchments.py @@ -18,7 +18,6 @@ from __future__ import annotations import logging -from collections import deque from dataclasses import dataclass from pathlib import Path from typing import Optional @@ -28,22 +27,13 @@ import rasterio import rasterio.features import rasterio.sample +from numba import njit from shapely import ops as shapely_ops from shapely.geometry import Point -log = logging.getLogger(__name__) +from ._d8 import downstream_index -# WBT D8 pointer: power-of-2 code --> (row_offset, col_offset) -_D8_OFFSETS: dict[int, tuple[int, int]] = { - 1: (0, 1), - 2: (1, 1), - 4: (1, 0), - 8: (1, -1), - 16: (0, -1), - 32: (-1, -1), - 64: (-1, 0), - 128: (-1, 1), -} +log = logging.getLogger(__name__) # Stream pixel centroids @@ -211,36 +201,56 @@ def _fill_interior_holes(labels: np.ndarray) -> np.ndarray: return out +def _orthogonal_neighbours(lab: np.ndarray, flat_idx: np.ndarray) -> np.ndarray: + """The 4 edge-neighbour labels at `flat_idx`, as (4, m); 0 off the grid. + + Order is south, north, east, west — the order the majority vote breaks ties in. + """ + rows, cols = lab.shape + flat = lab.ravel() + r, c = flat_idx // cols, flat_idx % cols + out = np.zeros((4, flat_idx.size), dtype=lab.dtype) + for k, (dr, dc) in enumerate(((1, 0), (-1, 0), (0, 1), (0, -1))): + rr, cc = r + dr, c + dc + on_grid = (rr >= 0) & (rr < rows) & (cc >= 0) & (cc < cols) + out[k, on_grid] = flat[rr[on_grid] * cols + cc[on_grid]] + return out + + def _declutter_boundaries(labels: np.ndarray, max_iter: int = 50) -> np.ndarray: """ Remove artefacts along catchment boundaries. Reassign corner-only cells to the dominant 4-neighbour label until stable. + + Teeth are a handful of cells out of hundreds of millions, so the pass finds + them first and votes only there — the whole-grid neighbour stack it used to + build cost gigabytes of copying per iteration. """ lab = labels.astype(np.int32).copy() background = lab == 0 for iteration in range(max_iter): - up = np.roll(lab, -1, 0) - dn = np.roll(lab, 1, 0) - lf = np.roll(lab, -1, 1) - rt = np.roll(lab, 1, 1) - # rolled wrap-around edges are not real neighbours — treat as background. - up[-1, :] = 0 - dn[0, :] = 0 - lf[:, -1] = 0 - rt[:, 0] = 0 + # A cell keeps its label if any edge-neighbour matches it. Off-grid + # neighbours count as background, which only ever matters for background + # cells — and those are excluded from teeth anyway. + same_label_edge = np.zeros(lab.shape, dtype=bool) + eq_vertical = lab[1:, :] == lab[:-1, :] + same_label_edge[:-1, :] |= eq_vertical + same_label_edge[1:, :] |= eq_vertical + eq_horizontal = lab[:, 1:] == lab[:, :-1] + same_label_edge[:, :-1] |= eq_horizontal + same_label_edge[:, 1:] |= eq_horizontal # "teeth": labelled cells sharing no edge (only corners) with their own label. - same_label_edge = (up == lab) | (dn == lab) | (lf == lab) | (rt == lab) - teeth = (~background) & (~same_label_edge) - if not teeth.any(): + teeth = np.flatnonzero(~background.ravel() & ~same_label_edge.ravel()) + if teeth.size == 0: break # Majority vote over the 4 orthogonal neighbours, ignoring background. - neighbours = np.stack([up, dn, lf, rt]) + neighbours = _orthogonal_neighbours(lab, teeth) valid = neighbours > 0 - best_label = np.zeros_like(lab) - best_count = np.zeros_like(lab) + best_label = np.zeros(teeth.size, dtype=lab.dtype) + best_count = np.zeros(teeth.size, dtype=np.int64) for k in range(4): candidate = neighbours[k] count = ((neighbours == candidate) & valid).sum(axis=0) * (candidate > 0) @@ -248,14 +258,10 @@ def _declutter_boundaries(labels: np.ndarray, max_iter: int = 50) -> np.ndarray: best_count = np.where(take, count, best_count) best_label = np.where(take, candidate, best_label) - apply = teeth & (best_count > 0) + apply = best_count > 0 if not apply.any(): break # remaining teeth are fully isolated (no labelled edge neighbour). - new_lab = lab.copy() - new_lab[apply] = best_label[apply] - if np.array_equal(new_lab, lab): - break - lab = new_lab + lab.flat[teeth[apply]] = best_label[apply] n_changed = int((lab != labels).sum()) log.info( @@ -336,6 +342,43 @@ def _enforce_connectivity(labels: np.ndarray, max_iter: int = 8) -> np.ndarray: return lab +@njit(cache=True) +def _label_downstream_paths( + ds: np.ndarray, result: np.ndarray, state: np.ndarray +) -> None: + """Give every cell the label of the first seed on its downstream path. + + Seeds arrive already marked done in ``state``, which is what stops the walk. + Each path is walked twice — once to find the label, once to stamp it — and + every cell is marked done on the way, so the whole grid costs O(n) rather + than one traversal per cell. + """ + n = ds.size + for start in range(n): + if state[start] == 2: + continue + + j = start + while state[j] == 0: + state[j] = 1 + nxt = ds[j] + if nxt == j: # pit or off-grid: nothing downstream to inherit from + result[j] = 0 + state[j] = 2 + break + j = nxt + + # state 1 here means the walk came back on itself — a flow cycle, which + # the label can't come out of, so the whole path stays unlabelled. + label = result[j] if state[j] == 2 else 0 + + j = start + while state[j] != 2: + result[j] = label + state[j] = 2 + j = ds[j] + + def _gage_watershed( d8: np.ndarray, outlet_rc_ids: list[tuple[int, int, int]], @@ -344,87 +387,30 @@ def _gage_watershed( ) -> np.ndarray: """ Algorithm: label propagation in topological downstream-->upstream order. - Each cell inherits the label of its downstream neighbor. We process - cells in the order they are visited during a BFS seeded from the outlet - points — a cell is only enqueued once its downstream neighbor is already - labeled, guaranteeing the label is ready when we process it. + Each cell inherits the label of its downstream neighbor, so its label is the + first outlet found by walking downstream; cells whose path reaches no outlet + stay 0. Runtime O(n): every cell is visited exactly once. """ rows, cols = d8.shape n = rows * cols - flat_d8 = d8.ravel().astype(np.int32) - if nodata_d8 is not None: - flat_d8[flat_d8 == int(nodata_d8)] = 0 - - # Build downstream flat-index (self-loop = no valid downstream) - ds = np.arange(n, dtype=np.int64) - r_all = np.arange(n, dtype=np.int64) // cols - c_all = np.arange(n, dtype=np.int64) % cols - - for d8_val, (dr, dc) in _D8_OFFSETS.items(): - sel = flat_d8 == d8_val - if not sel.any(): - continue - r_ds = r_all + dr - c_ds = c_all + dc - in_bounds = sel & (r_ds >= 0) & (r_ds < rows) & (c_ds >= 0) & (c_ds < cols) - ds[in_bounds] = (r_ds * cols + c_ds)[in_bounds] - - # Build upstream adjacency: us[j] = list of cells that drain into j - idx_all = np.arange(n, dtype=np.int64) - non_self = ds != idx_all - src_cells = idx_all[non_self] - dst_cells = ds[non_self] - sort_order = np.argsort(dst_cells, kind="stable") - dst_sorted = dst_cells[sort_order] - src_sorted = src_cells[sort_order] - split_pts = np.flatnonzero(np.diff(dst_sorted)) + 1 - boundaries = np.concatenate([[0], split_pts, [len(dst_sorted)]]) - groups = np.split(src_sorted, split_pts) - us: dict[int, np.ndarray] = { - int(dst_sorted[boundaries[k]]): groups[k] for k in range(len(groups)) - } + flat_d8 = d8 if nodata_d8 is None else np.where(d8 == int(nodata_d8), 0, d8) + ds = downstream_index(flat_d8) result = np.zeros(n, dtype=np.int32) - visited = np.zeros(n, dtype=bool) - is_seed = np.zeros(n, dtype=bool) + state = np.zeros(n, dtype=np.uint8) # 0 unseen, 1 on the current path, 2 settled - # Seed pass 1: fix every outlet cell to its own id. Seeds are adjacent along - # the channel, so this must finish for ALL seeds before propagation — else - # an upstream seed gets overwritten by a downstream seed's label. + # Fix every outlet cell to its own id first. Seeds sit next to each other + # along the channel, so they must all be pinned before propagation — else an + # upstream seed gets overwritten by a downstream seed's label. for r, c, hid in outlet_rc_ids: idx = int(r) * cols + int(c) result[idx] = hid - visited[idx] = True - is_seed[idx] = True + state[idx] = 2 - # Seed pass 2: enqueue the upstream neighbours of every seed. - queue: deque[int] = deque() - for r, c, _ in outlet_rc_ids: - idx = int(r) * cols + int(c) - for upstream_cell in us.get(idx, []): - if not visited[int(upstream_cell)]: - queue.append(int(upstream_cell)) - - # Process upstream: cell i inherits label from its downstream neighbor ds[i], - # which is already labeled because it was processed before i was enqueued. - while queue: - i = int(queue.popleft()) - if visited[i] or is_seed[i]: - continue - j = int(ds[i]) - lbl = result[j] - if lbl == 0: - # downstream not labeled yet — re-enqueue and retry (confluences). - queue.append(i) - continue - visited[i] = True - result[i] = lbl - for upstream_cell in us.get(i, []): - if not visited[int(upstream_cell)]: - queue.append(int(upstream_cell)) + _label_downstream_paths(ds, result, state) log.debug( "gage_watershed: %d/%d cells labeled", diff --git a/src/fimbox/preprocessing/calculate_branch/process_branches.py b/src/fimbox/preprocessing/calculate_branch/process_branches.py index a8746a8..b9b0900 100644 --- a/src/fimbox/preprocessing/calculate_branch/process_branches.py +++ b/src/fimbox/preprocessing/calculate_branch/process_branches.py @@ -59,6 +59,7 @@ from pathlib import Path from typing import Optional, Sequence, Union +from ..._workers import RAM_PER_WORKER_BRANCH_GB, resolve_workers from ...logging_utils import aoi_root, attach_case_log from ..source_naming import detect_identifier, resolve_source, source_name from .adjust_floodplains import adjust_floodplains @@ -170,8 +171,9 @@ def __init__( keep_failed_branches: bool = False, deny_branch_zero_list: Optional[Path] = None, deny_branches_list: Optional[Path] = None, - # parallelism - n_workers: int = 1, + # parallelism — None/0 sizes to the machine, 1 is serial, and a number + # bigger than CPU/RAM can feed is clamped instead of obeyed. + n_workers: Optional[int] = None, timeout_seconds: Optional[int] = None, ): self.aoi_dir = _pick_one("aoi_dir", aoi_dir, "huc_dir", huc_dir, required=True) @@ -265,12 +267,24 @@ def _pick_one(name_a: str, val_a, name_b: str, val_b, *, required: bool): HucProcessingConfig = AOIProcessingConfig -def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: - """Run every non-zero branch in parallel. +def process_branches( + cfg: AOIProcessingConfig, *, include_branch_zero: bool = False +) -> list[BranchResult]: + """Run the branch loop in parallel. Pure branch calculation: BranchZero, adjust_floodplains, CreateHAND, USGS crosswalk, per-branch, branch-zero cleanup. + With ``include_branch_zero=True`` branch zero joins the same pool instead of + being run to completion ahead of it. Its raster prep and CreateHAND cover the + whole AOI, so together they are the longest work in the run and nothing else + reads their outputs — holding the other branches back until they finish + wastes every other core. The caller only has to have published the shared + AOI-root rasters first (``BranchZero.publish_shared_inputs``), which is what + the siblings clip from. Note that ``cfg.timeout_seconds``, if set, now also + applies to branch zero, and it needs a far longer leash than a single level + path. + Calibration is NOT invoked here — run it explicitly via ``fimbox.run_calibration`` once the branch loop and any deny-list cleanups are complete. @@ -281,7 +295,6 @@ def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: start = time.time() log.info(f"=== process_branches: {cfg.aoi_id} ===") log.info(f"Branch list: {cfg.branch_list_path}") - log.info(f"Workers: {cfg.n_workers}") # Stage 0: AOI-level USGS gage assignment if cfg.usgs_gages_gpkg and cfg.levelpaths_gpkg: @@ -301,23 +314,36 @@ def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: except Exception as exc: log.error(f"USGS gage assignment failed: {exc}", exc_info=True) - # If outputs already exist, we run the branch-zero post-steps here. - _run_branch_zero_post_steps(cfg) - branch_ids = _read_branch_list(cfg.branch_list_path, cfg.branch_zero_id) log.info(f"Branches to process (excluding branch zero): {len(branch_ids)}") if not branch_ids: log.warning("No non-zero branches found — only branch zero will exist.") + # Branch zero goes in first: it is the whole-AOI task, so starting the + # longest job before the short ones keeps the tail of the run from being + # one lonely branch on one core (longest-processing-time-first). + dispatch_ids = ( + [cfg.branch_zero_id] + branch_ids if include_branch_zero else branch_ids + ) + + # Resolved against the machine: None/0 -> as many as CPU+RAM allow, an + # over-ambitious request -> clamped, 1 -> the serial path below. Never more + # workers than branches, so a 3-branch AOI doesn't spin up 12 processes. + n_workers = resolve_workers( + cfg.n_workers, + n_tasks=len(dispatch_ids) or None, + ram_per_worker_gb=RAM_PER_WORKER_BRANCH_GB, + label="Branch processing", + ) + results: list[BranchResult] = [] - if branch_ids and cfg.n_workers is not None and cfg.n_workers <= 1: + if dispatch_ids and n_workers <= 1: # True serial path: one branch at a time, no Dask. Use n_workers=1 to # isolate concurrency effects from deterministic per-branch behaviour. - # n_workers=None means "auto" (all cores) -> takes the Dask path below. - log.info("Processing %d branches serially (n_workers=1)", len(branch_ids)) - for bid in branch_ids: + log.info("Processing %d branches serially (n_workers=1)", len(dispatch_ids)) + for bid in dispatch_ids: results.append(_process_single_branch(cfg, bid)) - elif branch_ids: + elif dispatch_ids: from distributed import as_completed as dask_as_completed from ..._dask import get_client @@ -325,10 +351,10 @@ def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: _ensure_wbt_source(cfg.wbt_path if hasattr(cfg, "wbt_path") else None) - client = get_client(n_workers=cfg.n_workers) + client = get_client(n_workers=n_workers) log.info( "Dispatching %d branches to Dask (%d workers, dashboard %s)", - len(branch_ids), + len(dispatch_ids), len(client.scheduler_info()["workers"]), client.dashboard_link, ) @@ -339,12 +365,12 @@ def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: # 'failed' result and move on to siblings. futures = client.map( _process_single_branch, - [cfg] * len(branch_ids), - branch_ids, + [cfg] * len(dispatch_ids), + dispatch_ids, pure=False, retries=0, ) - future_to_bid = {fut.key: bid for fut, bid in zip(futures, branch_ids)} + future_to_bid = {fut.key: bid for fut, bid in zip(futures, dispatch_ids)} for fut in dask_as_completed(futures): bid = future_to_bid.get(fut.key, "?") try: @@ -371,6 +397,10 @@ def process_branches(cfg: AOIProcessingConfig) -> list[BranchResult]: branches_elapsed = time.time() - start _log_branch_summary(results, branches_elapsed) + # Branch-zero crosswalk waits until the loop drains — it reads branch zero's + # CreateHAND outputs, which may have just been produced inside the pool. + _run_branch_zero_post_steps(cfg) + # Per-branch + branch-zero deny-list cleanup runs once here, after # every branch has finished, instead of inside each worker. Tests # that consume intermediates (e.g. test_branchprocessing.py) can @@ -960,7 +990,12 @@ def _log_branch_summary(results: Sequence[BranchResult], elapsed_s: float) -> No parser.add_argument("--aoi-dir", required=True) parser.add_argument("--aoi-id", required=True) parser.add_argument("--branch-list", default=None) - parser.add_argument("--workers", type=int, default=1) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Parallel branch workers. Omit to size to the machine; 1 for serial.", + ) parser.add_argument("--fema-nfhl", default=None) parser.add_argument("--usgs-gages", default=None) parser.add_argument("--ahps", default=None) diff --git a/src/fimbox/preprocessing/calibrate_ratingcurve/__init__.py b/src/fimbox/preprocessing/calibrate_ratingcurve/__init__.py index 56d7eae..2900d02 100644 --- a/src/fimbox/preprocessing/calibrate_ratingcurve/__init__.py +++ b/src/fimbox/preprocessing/calibrate_ratingcurve/__init__.py @@ -24,6 +24,8 @@ from __future__ import annotations +from typing import Optional + from ._common import CalibrationNotImplemented from .aggregate import BranchAggregator, aggregate_branches from .dem_adjust import ( @@ -47,13 +49,15 @@ # Function-style aliases so callers that prefer ``identify_src_bankfull(...)`` # over ``SrcBankfull(...).run()`` get a one-line entry point. Each just # instantiates the class and calls .run(). -def identify_src_bankfull(aoi_dir, bankfull_flows_file, *, n_workers: int = 1): +def identify_src_bankfull( + aoi_dir, bankfull_flows_file, *, n_workers: Optional[int] = None +): return SrcBankfull( aoi_dir=aoi_dir, bankfull_flows_file=bankfull_flows_file, n_workers=n_workers ).run() -def subdiv_chan_obank_src(aoi_dir, vmann_table, *, n_workers: int = 1): +def subdiv_chan_obank_src(aoi_dir, vmann_table, *, n_workers: Optional[int] = None): return SrcSubdiv( aoi_dir=aoi_dir, vmann_table=vmann_table, n_workers=n_workers ).run() @@ -63,11 +67,11 @@ def nonmonotonic_src_adjustment(aoi_dir): return SrcNonmonotonic(aoi_dir=aoi_dir).run() -def thalweg_notches_adjustment(aoi_dir, *, n_workers: int = 1): +def thalweg_notches_adjustment(aoi_dir, *, n_workers: Optional[int] = None): return ThalwegNotchesAdjustment(aoi_dir=aoi_dir, n_workers=n_workers).run() -def longitudinal_flow_adjustment(aoi_dir, *, n_workers: int = 1): +def longitudinal_flow_adjustment(aoi_dir, *, n_workers: Optional[int] = None): return LongitudinalFlowFilter(aoi_dir=aoi_dir, n_workers=n_workers).run() @@ -94,7 +98,7 @@ def src_adjust_usgs_rating_trace( nwm_recur_file, usgs_acceptable_gages=None, *, - n_workers: int = 1, + n_workers: Optional[int] = None, ): return UsgsRatingCalibrator( aoi_dir=aoi_dir, @@ -106,7 +110,7 @@ def src_adjust_usgs_rating_trace( def src_adjust_ras2fim_rating( - aoi_dir, ras_rating_curve_csv, nwm_recur_file, *, n_workers: int = 1 + aoi_dir, ras_rating_curve_csv, nwm_recur_file, *, n_workers: Optional[int] = None ): return Ras2fimCalibrator( aoi_dir=aoi_dir, @@ -116,7 +120,9 @@ def src_adjust_ras2fim_rating( ).run() -def src_adjust_spatial_obs(aoi_dir, calib_points_file=None, *, n_workers: int = 1): +def src_adjust_spatial_obs( + aoi_dir, calib_points_file=None, *, n_workers: Optional[int] = None +): return SpatialObsCalibrator( aoi_dir=aoi_dir, calib_points_file=calib_points_file, n_workers=n_workers ).run() diff --git a/src/fimbox/preprocessing/calibrate_ratingcurve/dem_adjust.py b/src/fimbox/preprocessing/calibrate_ratingcurve/dem_adjust.py index e5d6a27..e4df112 100644 --- a/src/fimbox/preprocessing/calibrate_ratingcurve/dem_adjust.py +++ b/src/fimbox/preprocessing/calibrate_ratingcurve/dem_adjust.py @@ -15,7 +15,7 @@ import logging from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import NamedTuple, Optional import numpy as np import pandas as pd @@ -56,7 +56,7 @@ class ThalwegNotchesAdjustment: # Non-zero branches only; rewrites the SRC. aoi_dir: PathLike - n_workers: int = 1 # branch-parallel workers + n_workers: Optional[int] = None # branch-parallel workers stage_interval_m: float = 0.3048 n_stages: int = 84 # ladder length extrap_rows: int = 3 # trailing rows fit for extrapolation @@ -65,10 +65,7 @@ def run(self) -> dict[str, str]: aoi_dir = resolve_aoi_dir(self.aoi_dir) aoi_id = aoi_id_of(aoi_dir) branches = list(iter_branches(aoi_dir, exclude_zero=True)) - log.info( - f"ThalwegNotchesAdjustment: {aoi_id} " - f"({len(branches)} branches, {self.n_workers} workers)" - ) + log.info(f"ThalwegNotchesAdjustment: {aoi_id} " f"({len(branches)} branches)") return _run_branches( branches, _thalweg_one_branch, @@ -193,17 +190,36 @@ def _extend_linear( _LONGITUDINAL_SMOOTH_N = 2 -def _low_pct_ignore_zeros(arr): - nz = np.asarray(arr)[np.asarray(arr) > 0] - return np.percentile(nz, 10) if nz.size else 0.0 +def _min_pct_filter(values: np.ndarray) -> np.ndarray: + """10th percentile of the positive values in a 4-wide sliding window. + + At most four samples go into each window, so at the 10th percentile only the + two smallest positives can matter — which turns the whole pass into a sort of + four shifted copies instead of a Python callback per cell. Window offsets + (-2..+1) and reflected edges match scipy's ``generic_filter(size=4)``. + """ + padded = np.pad(np.asarray(values, dtype=float), (2, 1), mode="symmetric") + n = values.size + win = np.stack([padded[k : k + n] for k in range(4)], axis=1) + + positive = win > 0 # NaN compares False, so lake reaches drop out + counts = positive.sum(axis=1) + ranked = np.sort(np.where(positive, win, np.inf), axis=1) + + # Collapse the empty/single-value windows onto the same expression so no + # arithmetic ever touches the +inf padding. + lo = np.where(counts >= 1, ranked[:, 0], 0.0) + second = np.where(counts >= 2, ranked[:, 1], lo) + # np.percentile interpolates 0.1*(n-1) of the way from the smallest positive + # toward the next one. + return lo + (second - lo) * (0.1 * (counts - 1)) def _filter_voi(voi_array): # Min (10th-pct) then gaussian, along the chain. - from scipy.ndimage import gaussian_filter1d, generic_filter + from scipy.ndimage import gaussian_filter1d - minfilter = generic_filter(voi_array, _low_pct_ignore_zeros, size=4) - return gaussian_filter1d(minfilter, sigma=2, radius=2) + return gaussian_filter1d(_min_pct_filter(voi_array), sigma=2, radius=2) @dataclass @@ -212,17 +228,14 @@ class LongitudinalFlowFilter: # Lakes keep original Q. Non-zero branches only. aoi_dir: PathLike - n_workers: int = 1 # branch-parallel workers + n_workers: Optional[int] = None # branch-parallel workers n_stages: int = 84 def run(self) -> dict[str, str]: aoi_dir = resolve_aoi_dir(self.aoi_dir) aoi_id = aoi_id_of(aoi_dir) branches = list(iter_branches(aoi_dir, exclude_zero=True)) - log.info( - f"LongitudinalFlowFilter: {aoi_id} " - f"({len(branches)} branches, {self.n_workers} workers)" - ) + log.info(f"LongitudinalFlowFilter: {aoi_id} " f"({len(branches)} branches)") return _run_branches( branches, _longitudinal_one_branch, @@ -270,9 +283,11 @@ def _longitudinal_one_branch(branch_dir: Path, bid: str, n_stages: int) -> str: return "SKIP no multi-reach chains" # Smooth geometry, write back for non-lake reaches. + smooth_keys = _LONGITUDINAL_KEYS[:_LONGITUDINAL_SMOOTH_N] + ladders = _reach_ladders(src, smooth_keys) filtered = {} - for key in _LONGITUDINAL_KEYS[:_LONGITUDINAL_SMOOTH_N]: - filtered[key] = _filter_key(src, chains, stages, key) + for key in smooth_keys: + filtered[key] = _filter_key(src, chains, stages, key, ladders) for key in _LONGITUDINAL_KEYS[:_LONGITUDINAL_SMOOTH_N]: adj_col = f"{key}_longitudinalAdjusted" src = src.merge(filtered[key], on=["HydroID", "Stage"], how="left") @@ -345,19 +360,48 @@ def _build_chains(catch) -> list[list[int]]: next_ids = catch["NextDownID"].astype(int) headwaters_rows = catch.loc[~catch["HydroID"].isin(next_ids)] headwaters = list(headwaters_rows[headwaters_rows["LakeID"] < 0]["HydroID"]) + # HydroID -> NextDownID up front: the walk used to re-scan the whole table + # twice per step, which is what made big AOIs crawl. + next_of = dict( + zip(catch["HydroID"].astype(int).tolist(), next_ids.tolist()) + ) chains: list[list[int]] = [] for hw in headwaters: chain = [hw] nxt = hw - while catch["HydroID"].isin([nxt]).any(): - nxt = int(catch.loc[catch["HydroID"] == nxt, "NextDownID"].item()) + while nxt in next_of: + nxt = next_of[nxt] chain.append(nxt) if len(chain[:-1]) > 2: chains.append(chain) return chains -def _filter_key(src, chains, stages, key) -> pd.DataFrame: +class _ReachLadder(NamedTuple): + """One reach's rating curve: its stage ladder and the columns being smoothed.""" + + stages: np.ndarray + lake_id: float + columns: dict + + +def _reach_ladders(src, keys) -> dict[int, _ReachLadder]: + """Split the SRC into one ladder per HydroID, for O(1) lookup by reach. + + Chain smoothing asks for one (reach, stage, key) value at a time; without this + it re-filters the whole SRC table for every single one. + """ + return { + int(hid): _ReachLadder( + stages=sub["Stage"].to_numpy(), + lake_id=float(sub["LakeID"].iloc[0]), + columns={k: sub[k].to_numpy() for k in keys}, + ) + for hid, sub in src.groupby("HydroID", sort=False) + } + + +def _filter_key(src, chains, stages, key, ladders) -> pd.DataFrame: # Per chain: build a HydroID x stage matrix, smooth each stage column, # return long (HydroID, Stage, adj). stage_cols = [str(s) for s in stages] @@ -365,7 +409,7 @@ def _filter_key(src, chains, stages, key) -> pd.DataFrame: for chain in chains: rows = {} for pos, hid in enumerate(chain[:-1]): - rows[hid] = [_interp(src, hid, key, s) for s in stages] + [pos] + rows[hid] = _interp_ladder(ladders.get(int(hid)), key, stages) + [pos] mat = pd.DataFrame.from_dict( rows, orient="index", columns=stage_cols + ["long_position"] ) @@ -393,11 +437,12 @@ def _filter_key(src, chains, stages, key) -> pd.DataFrame: ) -def _interp(src, hydroid, key, stage): - sub = src.loc[src["HydroID"] == hydroid] - if sub.empty or sub["LakeID"].iloc[0] > 0: - return np.nan - return round(float(np.interp(stage, sub["Stage"], sub[key])), 3) +def _interp_ladder(ladder: Optional[_ReachLadder], key, stages) -> list: + """One reach's key interpolated onto `stages`; NaN for lakes and unknown ids.""" + if ladder is None or ladder.lake_id > 0: + return [np.nan] * len(stages) + values = np.interp(stages, ladder.stages, ladder.columns[key]) + return [round(float(v), 3) for v in values] @dataclass @@ -411,7 +456,7 @@ class BathymetricAdjustment: bathy_file_aibased: Optional[PathLike] = None ai_toggle: int = 0 ai_strm_order: int = 4 - n_workers: int = 1 + n_workers: Optional[int] = None def run(self) -> dict[str, str]: aoi_dir = resolve_aoi_dir(self.aoi_dir) @@ -436,9 +481,7 @@ def run(self) -> dict[str, str]: def _apply(self, aoi_dir: Path, bathy: pd.DataFrame) -> str: # Inject the missing-geometry table into every branch SRC, in parallel. branches = list(iter_branches(aoi_dir, exclude_zero=False)) - log.info( - f"BathymetricAdjustment: {len(branches)} branches, {self.n_workers} workers" - ) + log.info(f"BathymetricAdjustment: {len(branches)} branches") out = _run_branches( branches, _bathy_one_branch, diff --git a/src/fimbox/preprocessing/calibrate_ratingcurve/pipeline.py b/src/fimbox/preprocessing/calibrate_ratingcurve/pipeline.py index a9d7a73..0ee3144 100644 --- a/src/fimbox/preprocessing/calibrate_ratingcurve/pipeline.py +++ b/src/fimbox/preprocessing/calibrate_ratingcurve/pipeline.py @@ -23,8 +23,9 @@ src_bankfull_toggle=True, bankfull_flows_file="bankfull.csv", src_subdiv_toggle=True, vmann_input_file="mannings.csv", nonmonotonic_src_adjustment=True, - job_branch_limit=4, )) + # job_branch_limit is left unset above: the branch-parallel steps size + # themselves to the machine. Set it only to cap the pool (or 1 to go serial). Each ``Calibrator`` step can also be run on its own (see the individual classes in this subpackage), which is what the step-by-step tests exercise. @@ -118,8 +119,10 @@ class CalibrationConfig: scan_logs: bool = False # --- execution --- - # Worker count for the branch-parallel routines. - job_branch_limit: int = 1 + # Worker count for the branch-parallel routines. Leave it alone (or pass + # None / 0) to use everything the machine can feed; a number larger than + # that is clamped rather than obeyed. Pass 1 for a serial, debuggable run. + job_branch_limit: Optional[int] = None # When True, toggled-on routines that aren't ported yet warn and skip instead of raising CalibrationNotImplemented. skip_unimplemented: bool = False diff --git a/src/fimbox/preprocessing/calibrate_ratingcurve/src_adjust.py b/src/fimbox/preprocessing/calibrate_ratingcurve/src_adjust.py index 863da7e..d72c8d5 100644 --- a/src/fimbox/preprocessing/calibrate_ratingcurve/src_adjust.py +++ b/src/fimbox/preprocessing/calibrate_ratingcurve/src_adjust.py @@ -28,11 +28,12 @@ from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass from pathlib import Path -from typing import Callable +from typing import Callable, Optional import numpy as np import pandas as pd +from ..._workers import RAM_PER_WORKER_TABLE_GB, resolve_workers from ._common import ( BEDAREA_VAR, HRADIUS_VAR, @@ -51,22 +52,29 @@ def _run_branches( branches: list[tuple[str, Path]], worker: Callable[..., str], worker_args: tuple, - n_workers: int, + n_workers: Optional[int], label: str, ) -> dict[str, str]: """Run ``worker(branch_dir, bid, *worker_args)`` over every branch. - Serial when ``n_workers <= 1``. Otherwise a ProcessPoolExecutor, with a - fallback to serial if the pool breaks (``BrokenProcessPool`` — seen when a - native lib can't fork, common on macOS). A broken pool would otherwise - record every branch as a silent "FAIL", so we re-run serially instead of - publishing an empty calibration. Individual branch exceptions are caught - and recorded as ``FAIL ...`` without sinking the batch. + ``n_workers`` is resolved against the machine first: ``None``/``0`` auto-size + to every core the RAM can feed, an over-ambitious number is clamped, and + ``n_workers=1`` stays serial for debugging. Parallel runs use a + ProcessPoolExecutor with a fallback to serial if the pool breaks + (``BrokenProcessPool`` — seen when a native lib can't fork, common on + macOS). A broken pool would otherwise record every branch as a silent + "FAIL", so we re-run serially instead of publishing an empty calibration. + Individual branch exceptions are caught and recorded as ``FAIL ...`` without + sinking the batch. """ results: dict[str, str] = {} - # n_workers=None means "auto" (let ProcessPoolExecutor size to cpu count) -> - # takes the parallel path below; only an explicit <=1 runs serially. - if n_workers is not None and n_workers <= 1: + n_workers = resolve_workers( + n_workers, + n_tasks=len(branches), + ram_per_worker_gb=RAM_PER_WORKER_TABLE_GB, + label=label, + ) + if n_workers <= 1: for bid, bp in branches: try: results[bid] = worker(bp, bid, *worker_args) @@ -107,7 +115,7 @@ class SrcBankfull: aoi_dir: PathLike bankfull_flows_file: PathLike - n_workers: int = 1 + n_workers: Optional[int] = None include_branch_zero: bool = True def run(self) -> dict[str, str]: @@ -123,7 +131,7 @@ def run(self) -> dict[str, str]: log.warning(f"SrcBankfull: no branches found under {aoi_dir}") return {} - log.info(f"SrcBankfull: {len(branches)} branches, {self.n_workers} workers") + log.info(f"SrcBankfull: {len(branches)} branches") return _run_branches( branches, _bankfull_one_branch, @@ -249,7 +257,7 @@ class SrcSubdiv: aoi_dir: PathLike vmann_table: PathLike # CSV/Parquet keyed on feature_id, channel_n, overbank_n - n_workers: int = 1 + n_workers: Optional[int] = None include_branch_zero: bool = True default_channel_n: float = 0.06 default_overbank_n: float = 0.12 @@ -265,7 +273,7 @@ def run(self) -> dict[str, str]: branches = list( iter_branches(aoi_dir, exclude_zero=not self.include_branch_zero) ) - log.info(f"SrcSubdiv: {len(branches)} branches, {self.n_workers} workers") + log.info(f"SrcSubdiv: {len(branches)} branches") return _run_branches( branches, _subdiv_one_branch, @@ -473,7 +481,7 @@ class SrcNonmonotonic: aoi_dir: PathLike stream_order_min: int = 4 include_branch_zero: bool = True - n_workers: int = 1 + n_workers: Optional[int] = None def run(self) -> dict[str, str]: aoi_dir = resolve_aoi_dir(self.aoi_dir) @@ -482,9 +490,7 @@ def run(self) -> dict[str, str]: self._normalize_branch_zero(aoi_dir) branches = list(iter_branches(aoi_dir, exclude_zero=True)) - log.info( - f"SrcNonmonotonic: {len(branches)} non-zero branches, {self.n_workers} workers" - ) + log.info(f"SrcNonmonotonic: {len(branches)} non-zero branches") return _run_branches( branches, _nonmonotonic_one_branch, diff --git a/src/fimbox/preprocessing/calibrate_ratingcurve/src_calibrate.py b/src/fimbox/preprocessing/calibrate_ratingcurve/src_calibrate.py index 030c39c..6570d12 100644 --- a/src/fimbox/preprocessing/calibrate_ratingcurve/src_calibrate.py +++ b/src/fimbox/preprocessing/calibrate_ratingcurve/src_calibrate.py @@ -382,7 +382,7 @@ class UsgsRatingCalibrator: usgs_rating_curve_csv: PathLike nwm_recur_file: PathLike usgs_acceptable_gages: Optional[PathLike] = None # optional quality filter - n_workers: int = 1 + n_workers: Optional[int] = None debug_outputs: bool = False def run(self) -> dict[str, str]: @@ -420,9 +420,7 @@ def run(self) -> dict[str, str]: usgs_df["levpa_id"] = usgs_df["levpa_id"].astype("int64").astype(str) branches = list(iter_branches(aoi_dir, exclude_zero=False)) - log.info( - f"UsgsRatingCalibrator: {aoi_id} ({len(branches)} branches, {self.n_workers} workers)" - ) + log.info(f"UsgsRatingCalibrator: {aoi_id} ({len(branches)} branches)") return _run_branches( branches, _usgs_one_branch, @@ -515,7 +513,7 @@ class SpatialObsCalibrator: aoi_dir: PathLike calib_points_file: Optional[PathLike] = None - n_workers: int = 1 + n_workers: Optional[int] = None down_dist_thresh: float = DOWNSTREAM_THRESHOLD debug_outputs: bool = False @@ -539,9 +537,7 @@ def run(self) -> dict[str, str]: return {} branches = list(iter_branches(aoi_dir, exclude_zero=False)) - log.info( - f"SpatialObsCalibrator: {aoi_id} ({len(branches)} branches, {self.n_workers} workers)" - ) + log.info(f"SpatialObsCalibrator: {aoi_id} ({len(branches)} branches)") return _run_branches( branches, _spatial_one_branch, @@ -557,7 +553,7 @@ class Ras2fimCalibrator: aoi_dir: PathLike ras_rating_curve_csv: PathLike nwm_recur_file: PathLike - n_workers: int = 1 + n_workers: Optional[int] = None def run(self) -> None: not_yet_ported("Ras2fimCalibrator") diff --git a/src/fimbox/preprocessing/download_data/nld_data.py b/src/fimbox/preprocessing/download_data/nld_data.py index 47874e5..93643b8 100644 --- a/src/fimbox/preprocessing/download_data/nld_data.py +++ b/src/fimbox/preprocessing/download_data/nld_data.py @@ -24,6 +24,15 @@ log = logging.getLogger(__name__) +class NLDQueryError(RuntimeError): + """A NLD layer could not be downloaded. + + Kept distinct from an empty result on purpose: most AOIs genuinely have no + levees, and a failed request must never be reported as one of them — a + missing levee burn quietly changes the HAND surface. + """ + + class ESRI_REST: """ A robust utility for querying ESRI Feature Services. @@ -78,6 +87,7 @@ def _get_with_retry(self, params: dict, timeout: int = 120): requests.exceptions.ChunkedEncodingError, requests.exceptions.ConnectionError, requests.exceptions.Timeout, + requests.exceptions.HTTPError, ) as exc: last_exc = exc if attempt < self.MAX_PAGE_RETRIES - 1: @@ -92,9 +102,13 @@ def _get_with_retry(self, params: dict, timeout: int = 120): time.sleep(wait) raise last_exc - def _execute_query(self, params: dict) -> gpd.GeoDataFrame: + def _execute_query(self, params: dict, with_z: bool = False) -> gpd.GeoDataFrame: """Paginated ESRI Feature Service query with retry-on-truncation. + ``with_z=True`` asks the server for Z coordinates, which levee lines need + for the DEM burn. Geometry is parsed straight out of ESRI JSON, so a + third ordinate simply flows through into the shapely coords. + Uses ESRI native JSON (``f=json``) and parses ``features[].geometry`` directly into shapely instead of round-tripping through GeoJSON. This avoids the brittle ``gpd.read_file(resp.text)`` path that fails on @@ -115,9 +129,14 @@ def _execute_query(self, params: dict) -> gpd.GeoDataFrame: return gpd.GeoDataFrame() if self.verbose: - log.info(f"ESRI query: {total_features} features to download") + log.info( + f"ESRI query{' (with Z)' if with_z else ''}: " + f"{total_features} features to download" + ) base_params = {**params, "f": "json"} + if with_z: + base_params["returnZ"] = "true" results: list[gpd.GeoDataFrame] = [] offset = 0 page_size = self.INITIAL_PAGE_SIZE @@ -154,8 +173,20 @@ def _execute_query(self, params: dict) -> gpd.GeoDataFrame: page_size = used_size # remember the working size if not results: - return gpd.GeoDataFrame() - return pd.concat(results, ignore_index=True) + # The server counted features and then handed back none of them. + # That is a failed download, not an empty area. + raise NLDQueryError( + f"service reported {total_features} features but returned none" + ) + + gdf = gpd.GeoDataFrame(pd.concat(results, ignore_index=True)) + if len(gdf) < total_features: + log.warning( + "ESRI query returned %d of %d features — paging stopped early", + len(gdf), + total_features, + ) + return gdf def _fetch_page( self, base_params: dict, offset: int, page_size: int @@ -250,67 +281,14 @@ def _esri_geom_to_shapely(geom: Optional[dict]): return None def _execute_query_with_z(self, params: dict) -> gpd.GeoDataFrame: - """ - Like _execute_query but uses f=json + returnZ=true so Z coordinates are - preserved. Parses ESRI JSON 'paths' directly into shapely LineStrings. - GeoJSON silently strips Z, so this path is required for levee lines. - """ - total_features = self._get_metadata(params) - if total_features == 0: - return gpd.GeoDataFrame() - - base_params = {**params, "f": "json", "returnZ": "true"} - if self.verbose: - log.info(f"ESRI query (with Z): {total_features} features to download") - - results = [] - offset = 0 - limit_reached = True - - with tqdm( - total=total_features, disable=not self.verbose, desc="Downloading" - ) as pbar: - while limit_reached: - current_params = {**base_params, "resultOffset": offset} - resp = self._get_with_retry(current_params, timeout=120) - data = resp.json() - if "error" in data: - raise Exception( - f"ESRI Error {data['error']['code']}: {data['error']['message']}" - ) - - features = data.get("features", []) - if not features: - break - - rows = [] - for feat in features: - attrs = feat.get("attributes", {}) - paths = feat.get("geometry", {}).get("paths", []) - if not paths: - continue - lines = [LineString(path) for path in paths if len(path) >= 2] - if not lines: - continue - geom = lines[0] if len(lines) == 1 else MultiLineString(lines) - attrs["geometry"] = geom - rows.append(attrs) + """Levee lines, Z coordinates included. - if rows: - batch = gpd.GeoDataFrame(rows, geometry="geometry") - results.append(batch) - - offset += len(features) - pbar.update(len(features)) - limit_reached = data.get("exceededTransferLimit", False) - if not limit_reached and offset < total_features: - limit_reached = True - elif offset >= total_features: - limit_reached = False - - if not results: - return gpd.GeoDataFrame() - return pd.concat(results, ignore_index=True) + GeoJSON silently strips Z, so this asks for ESRI JSON with returnZ. It + now shares the paged, shrink-on-failure fetch used for polygons — the + levee layer is the one that most needs it, since a truncated response on + a levee-dense AOI used to surface as "no levees here". + """ + return self._execute_query(params, with_z=True) class DownloadNLD: @@ -319,11 +297,57 @@ class DownloadNLD: using the new geospatial.sec.usace.army.mil endpoint. """ - # UPDATED URLs and LAYER IDs BASE_SERVICE_URL = "https://geospatial.sec.usace.army.mil/dls/rest/services/NLD/Public/FeatureServer" + + # Levee centrelines live in "System Routes", leveed/protected areas in + # "Leveed Areas". The ids below are today's numbering and are used as a + # fallback; the real ids are looked up by name at runtime so a USACE + # renumbering shows up as a warning instead of a silently empty download. + LAYER_NAMES = {"lines": "System Routes", "polys": "Leveed Areas"} + FALLBACK_LAYER_IDS = {"lines": 15, "polys": 16} + + _layer_ids_cache: Optional[dict] = None + LINE_URL = f"{BASE_SERVICE_URL}/15/query" POLY_URL = f"{BASE_SERVICE_URL}/16/query" + @classmethod + def resolve_layer_ids(cls) -> dict: + """Map our two layers to live service ids, by name. Cached per process.""" + if cls._layer_ids_cache is not None: + return cls._layer_ids_cache + + ids = dict(cls.FALLBACK_LAYER_IDS) + try: + resp = requests.get(cls.BASE_SERVICE_URL, params={"f": "json"}, timeout=60) + resp.raise_for_status() + by_name = { + str(layer.get("name", "")).strip().lower(): layer.get("id") + for layer in resp.json().get("layers", []) + } + for key, name in cls.LAYER_NAMES.items(): + found = by_name.get(name.lower()) + if found is None: + log.warning( + "NLD service has no '%s' layer — falling back to id %s", + name, + ids[key], + ) + elif found != ids[key]: + log.info( + "NLD '%s' moved to layer %s (was %s)", name, found, ids[key] + ) + ids[key] = found + except Exception as exc: + log.warning( + "Could not read NLD service metadata (%s) — using layer ids %s", + exc, + ids, + ) + + cls._layer_ids_cache = ids + return ids + def __init__( self, boundary: Union[str, Path, gpd.GeoDataFrame, Polygon, MultiPolygon], @@ -339,6 +363,15 @@ def __init__( self.logger = log + layer_ids = self.resolve_layer_ids() + self.layer_ids = layer_ids + self.LINE_URL = f"{self.BASE_SERVICE_URL}/{layer_ids['lines']}/query" + self.POLY_URL = f"{self.BASE_SERVICE_URL}/{layer_ids['polys']}/query" + + # Layers that errored out, as opposed to came back empty. Callers read + # this to tell "no levees here" apart from "we never found out". + self.failed_layers: list[str] = [] + from ...logging_utils import default_output_dir self.output_dir = Path(out_dir) if out_dir else default_output_dir() @@ -369,7 +402,8 @@ def _prepare_boundary(self, boundary, layer_name) -> gpd.GeoDataFrame: return gdf.to_crs(epsg=self.epsg) def _query_nld(self, url: str, is_poly: bool = False) -> gpd.GeoDataFrame: - layer_type = "Polygons (Layer 16)" if is_poly else "Lines (Layer 15)" + key = "polys" if is_poly else "lines" + layer_type = f"{self.LAYER_NAMES[key]} (Layer {self.layer_ids[key]})" # Filter server-side on the AOI envelope so only nearby levees come back # instead of the whole national layer (~6200 features per layer). The @@ -397,14 +431,20 @@ def _query_nld(self, url: str, is_poly: bool = False) -> gpd.GeoDataFrame: "inSR": "4269", } + log.info(f"Downloading NLD {layer_type} within the AOI envelope") + rest = ESRI_REST(url) try: - log.info(f"Downloading NLD {layer_type} within the AOI envelope") - rest = ESRI_REST(url) if is_poly: gdf_raw = rest._execute_query(base_params) else: gdf_raw = rest._execute_query_with_z(base_params) + except Exception as exc: + # Deliberately not swallowed into an empty frame: the caller records + # this layer as failed so nothing downstream reads the gap as + # "this AOI has no levees". + raise NLDQueryError(f"NLD {layer_type} download failed: {exc}") from exc + try: if gdf_raw.empty: # No data is a normal outcome (many AOIs have no levees), not a # failure — report it plainly at INFO. @@ -430,26 +470,37 @@ def _query_nld(self, url: str, is_poly: bool = False) -> gpd.GeoDataFrame: # to_crs() preserves Z when the geometry already carries it. return selected.to_crs(epsg=self.epsg) - except Exception as e: - log.error(f"NLD {layer_type} failed: {e}", exc_info=True) - return gpd.GeoDataFrame() + except Exception as exc: + raise NLDQueryError(f"NLD {layer_type} processing failed: {exc}") from exc def run(self): log.info("--- NLD download ---") - lines = self._query_nld(self.LINE_URL, is_poly=False) - if not lines.empty: - lines.to_file(self.output_dir / self.lines_name, driver="GPKG") - log.info(f"NLD levee lines ({len(lines)} features) --> {self.lines_name}") - - polys = self._query_nld(self.POLY_URL, is_poly=True) - if not polys.empty: - polys.to_file(self.output_dir / self.polys_name, driver="GPKG") - log.info( - f"NLD leveed-area polygons ({len(polys)} features) --> {self.polys_name}" + for key, url, out_name, label in ( + ("lines", self.LINE_URL, self.lines_name, "levee lines"), + ("polys", self.POLY_URL, self.polys_name, "leveed-area polygons"), + ): + # Each layer stands on its own: a failure on one is recorded and the + # other is still attempted. + try: + gdf = self._query_nld(url, is_poly=(key == "polys")) + except NLDQueryError as exc: + self.failed_layers.append(key) + log.error("%s", exc, exc_info=True) + continue + + if not gdf.empty: + gdf.to_file(self.output_dir / out_name, driver="GPKG") + log.info(f"NLD {label} ({len(gdf)} features) --> {out_name}") + + if self.failed_layers: + log.error( + "NLD download incomplete — %s could not be retrieved. " + "Absence of these files does NOT mean the area has no levees.", + ", ".join(self.LAYER_NAMES[k] for k in self.failed_layers), ) - - log.info("NLD download complete.") + else: + log.info("NLD download complete.") # CLI diff --git a/src/fimbox/preprocessing/download_data/osm_data.py b/src/fimbox/preprocessing/download_data/osm_data.py index 570a43c..8c24fb4 100644 --- a/src/fimbox/preprocessing/download_data/osm_data.py +++ b/src/fimbox/preprocessing/download_data/osm_data.py @@ -1,43 +1,69 @@ """ Author: Supath Dhital -Date updated: April 2026 - -Download major road segments from OpenStreetMap (OSM) within a user-provided boundary. - -- Queries Overpass API for highway types: motorway, trunk, primary, secondary, tertiary -- Explicitly EXCLUDES bridges (ways with bridge=*) to avoid unrealistic flood depth calcs -- Boundary input can be: shapefile/gpkg/geojson path, GeoDataFrame/GeoSeries, shapely geometry, or bbox -- Output is saved to GeoPackage in EPSG:5070 -- User can pass out_dir, out_name (or ourfile), out_layer (or ourlayer); defaults used otherwise -- Large areas are automatically split into tiles fetched in parallel via ThreadPoolExecutor - -AND - -Download bridge features from OpenStreetMap (OSM) within a user-provided boundary. -- Uses OSMnx features_from_polygon with {"bridge": True} -- Boundary input can be: shapefile/gpkg/geojson path, GeoDataFrame/GeoSeries, shapely geometry, or bbox -- Converts non-LineString geometries to LineStrings when possible (Polygon -> exterior; Point -> skipped by default) -- Removes abandoned/proposed/demolished bridges based on bridge_type (highway-* / railway-*) -- Dissolves touching bridge segments (buffer + graph connectivity) to form continuous bridge lines -- Output is saved to GeoPackage in EPSG:5070 -- User can pass out_dir, out_name (or ourfile), out_layer (or ourlayer); defaults used otherwise +Date updated: July 2026 + +Download major road segments AND bridge features from OpenStreetMap (OSM) within +a user-provided boundary. + +Both downloaders share one Overpass client (:class:`_OverpassClient`) that: + - rotates over several public Overpass mirrors and *permanently drops* a host + for the session as soon as it refuses a connection or times out on connect, + so a dead mirror costs one failed socket instead of every retry; + - probes all mirrors once, in parallel, before the first real query, and keeps + only the reachable ones (some networks block *.overpass-api.de outright); + - treats an Overpass ``remark`` ("runtime error: Query timed out") as a + failure. Overpass returns those inside an HTTP 200 body, so accepting the + response at face value silently yields partial data. + +Coverage strategy (identical for roads and bridges): + - the AOI bbox is split into tiles of ~``_TILE_AREA_DEG_SQ`` sq-degrees; + - tiles that do not touch the AOI polygon are dropped before any request + (a watershed is never a rectangle — this is free speed); + - tiles are fetched concurrently, spread round-robin across live mirrors; + - a tile that fails every mirror is split into quadrants and retried + (``_MAX_SUBSPLIT_DEPTH`` times) — this is what makes large AOIs work; + - if tiles still fail, the download raises instead of writing a file that + looks complete but has holes in it (set ``allow_partial=True`` to override). + +Roads + - highway types: motorway, trunk, primary, secondary, tertiary + - explicitly EXCLUDES bridges (ways with bridge=*) to avoid unrealistic flood + depth calcs +Bridges + - ways carrying a bridge tag (bridge=no excluded), queried straight from + Overpass rather than through osmnx, so bridges get the same tiling, + mirror failover and quadrant-retry as roads + - keeps a fixed, curated tag schema (no list columns, no duplicate-cased + FIXME columns — the two things that used to break GPKG writes) + - removes abandoned/proposed/demolished bridges based on bridge_type + - dissolves touching bridge segments (buffer + graph connectivity) to form + continuous bridge lines + +Shared + - boundary input can be: shapefile/gpkg/geojson path, GeoDataFrame/GeoSeries, + shapely geometry, or bbox + - output is saved to GeoPackage in EPSG:5070 + - user can pass out_dir, out_name (or ourfile), out_layer (or ourlayer) """ import logging import math import random +import threading import time import warnings from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass +from concurrent.futures import TimeoutError as FutureTimeout +from dataclasses import dataclass, field +from datetime import datetime, timezone from pathlib import Path -from typing import Any, List, Optional, Sequence, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import geopandas as gpd -import osmnx as ox import pandas as pd import requests from networkx import Graph, connected_components +from requests.adapters import HTTPAdapter from shapely.geometry import LineString, MultiPolygon, Polygon, box from shapely.ops import unary_union from tqdm import tqdm @@ -47,7 +73,292 @@ log = logging.getLogger(__name__) -# shared boundary with IO helpers (used by both roads + bridges) +# Public planet-wide Overpass instances, best-first. Availability varies by +# network: some campus/VPN routes refuse *.overpass-api.de outright, so the +# non-.de mirrors are not optional extras — they are what keeps the download +# working. Only add hosts that serve the whole planet: a region-limited +# instance (overpass.osm.ch, for example, is Switzerland only) answers HTTP 200 +# with zero features for a US query, which reads as "no roads here". +OVERPASS_MIRRORS: Tuple[str, ...] = ( + "https://overpass-api.de/api/interpreter", + "https://overpass.openstreetmap.fr/api/interpreter", + "https://overpass.kumi.systems/api/interpreter", + "https://lz4.overpass-api.de/api/interpreter", + "https://z.overpass-api.de/api/interpreter", + "https://overpass.private.coffee/api/interpreter", +) + +_USER_AGENT = "fimbox/0.1 (+https://github.com/sdmlua/fimbox)" + +# A mirror whose OSM base timestamp is missing, unparseable, or older than this +# is out of sync with the planet; its answers are not trustworthy. +_MAX_MIRROR_LAG_DAYS = 21 + + +class _TileFailed(Exception): + """One tile could not be fetched from any live mirror.""" + + +class _Transient(Exception): + """Mirror answered but the answer is unusable (rate limit, timeout, junk).""" + + +class _StaleMirror(Exception): + """Mirror is up but its data is out of date or not planet-wide.""" + + +class _OverpassClient: + """Mirror-failover Overpass client. + + Mirror health is tracked at class level so the roads download and the + bridges download that follows it share one probe and one dead-host list — + bridges never re-learn that a host is down. + """ + + # Tiny query over a busy US interchange: cheap, and a mirror that serves + # only a non-US region answers it with zero ways, which unmasks it. + _PROBE = ( + "[out:json][timeout:10];" + 'way["highway"](33.7520,-84.3900,33.7620,-84.3800);out count;' + ) + + _PROBE_BUDGET_S = 6.0 # total wall clock allowed for the one-off probe + + _state_lock = threading.Lock() + _live: Optional[List[str]] = None # probed, reachable mirrors + _dead: set = set() + _rr = 0 # round-robin cursor, shared so parallel tiles spread out + + def __init__( + self, + timeout: int = 180, + max_attempts: int = 4, + sleep_base: float = 2.0, + mirrors: Optional[Sequence[str]] = None, + ): + self.timeout = timeout + self.max_attempts = max_attempts + self.sleep_base = sleep_base + self.mirrors = list(mirrors) if mirrors else list(OVERPASS_MIRRORS) + + self._session = requests.Session() + self._session.headers["User-Agent"] = _USER_AGENT + self._session.headers["Accept-Encoding"] = "gzip, deflate" + # one pooled connection per worker per mirror + adapter = HTTPAdapter(pool_connections=len(self.mirrors), pool_maxsize=16) + self._session.mount("https://", adapter) + self._session.mount("http://", adapter) + + # mirror health + + def _probe_one(self, mirror: str) -> bool: + """True if the host is usable: reachable, in sync with the planet, and + actually holding data outside its own region. A 429 still counts as + usable — rate-limited now, likely free by the time we need it.""" + try: + r = self._session.get(mirror, params={"data": self._PROBE}, timeout=(3, 6)) + if r.status_code in (429, 502, 503, 504): + return True # up, just busy + if r.status_code >= 400: + return False + data = r.json() + _assert_fresh(data) + elems = data.get("elements") or [] + ways = int((elems[0].get("tags") or {}).get("ways", 0)) if elems else 0 + if ways == 0: + log.warning( + "Overpass: %s returned no data for a known-populated area " + "(region-limited mirror?) — skipping it", + _host(mirror), + ) + return False + return True + except _StaleMirror as exc: + log.warning("Overpass: %s %s — skipping it", _host(mirror), exc) + return False + except Exception: + return False + + def live_mirrors(self) -> List[str]: + """Reachable mirrors, probed once per process and shared from then on. + + The probe runs on a stopwatch: hosts that answer inside + ``_PROBE_BUDGET_S`` decide the list, and a host still thinking when the + budget runs out is simply left out of it. Two of the public mirrors + accept the connection and then hang, and waiting on them would cost more + than the download itself. + """ + with self._state_lock: + if _OverpassClient._live is not None: + live = [m for m in _OverpassClient._live if m not in self._dead] + return live or list(self.mirrors) + + live: List[str] = [] + answered: set = set() + pool = ThreadPoolExecutor(max_workers=len(self.mirrors)) + futures = {pool.submit(self._probe_one, m): m for m in self.mirrors} + try: + for fut in as_completed(futures, timeout=self._PROBE_BUDGET_S): + mirror = futures[fut] + answered.add(mirror) + if fut.result(): + live.append(mirror) # fastest responders end up first + except FutureTimeout: + pass + finally: + pool.shutdown(wait=False, cancel_futures=True) + + with self._state_lock: + _OverpassClient._live = live + # Only hosts that actually answered are condemned; the silent ones + # stay in reserve for the fallback path below. + _OverpassClient._dead |= answered - set(live) + + if live: + log.info( + "Overpass: %d/%d mirror(s) reachable (%s)", + len(live), + len(self.mirrors), + ", ".join(_host(m) for m in live), + ) + else: + # Every probe failed. Do not give up here — the probe may have hit a + # transient outage; fall back to the full list and let real queries + # produce a real error message. + log.warning( + "Overpass: no mirror answered the health probe — trying all %d anyway", + len(self.mirrors), + ) + return list(self.mirrors) + return live + + def _mark_dead(self, mirror: str) -> None: + with self._state_lock: + if mirror not in _OverpassClient._dead: + _OverpassClient._dead.add(mirror) + log.warning("Overpass: dropping unreachable mirror %s", _host(mirror)) + if _OverpassClient._live is not None: + _OverpassClient._live = [ + m for m in _OverpassClient._live if m != mirror + ] + + def _next_start(self) -> int: + with self._state_lock: + _OverpassClient._rr += 1 + return _OverpassClient._rr + + # fetch + + def fetch(self, query: str, avoid: Optional[str] = None) -> Tuple[dict, str]: + """Run ``query`` against live mirrors in turn. + + Returns ``(json, mirror_used)``. ``avoid`` skips one mirror, which is how + an empty answer gets a second opinion from a different host. + Raises ``_TileFailed`` when every mirror fails. + """ + mirrors = [m for m in self.live_mirrors() if m != avoid] + if not mirrors: + raise _TileFailed("no Overpass mirror is reachable") + + attempts = max(self.max_attempts, len(mirrors)) + start = self._next_start() + errors: List[str] = [] + + for attempt in range(attempts): + mirror = mirrors[(start + attempt) % len(mirrors)] + if mirror in self._dead: + continue + try: + r = self._session.get( + mirror, + params={"data": query}, + timeout=(10, self.timeout + 60), + ) + if r.status_code in (429, 502, 503, 504): + raise _Transient(f"HTTP {r.status_code}") + r.raise_for_status() + data = r.json() # ValueError if the body was truncated + remark = str(data.get("remark", "")) + if remark and ("error" in remark.lower() or "timed out" in remark): + # Overpass reports query timeouts *inside* a 200 body. The + # payload is partial, so it must not be treated as data. + raise _Transient(f"overpass remark: {remark}") + if "elements" not in data: + raise _Transient("response has no 'elements'") + _assert_fresh(data) + return data, mirror + + except _StaleMirror as exc: + self._mark_dead(mirror) + errors.append(f"{_host(mirror)}: {exc}") + continue # its data is wrong, not just late + except requests.exceptions.ReadTimeout as exc: + # Host is alive, the query is too heavy for it — keep the mirror, + # let the caller shrink the tile. + errors.append(f"{_host(mirror)}: read timeout") + _ = exc + except ( + requests.exceptions.ConnectionError, + requests.exceptions.ConnectTimeout, + ) as exc: + self._mark_dead(mirror) + errors.append(f"{_host(mirror)}: unreachable ({_brief(exc)})") + continue # dead host: no backoff, move on immediately + except Exception as exc: + errors.append(f"{_host(mirror)}: {_brief(exc)}") + + if attempt < attempts - 1: + time.sleep(self.sleep_base * (attempt + 1) + random.uniform(0, 1.0)) + + raise _TileFailed("; ".join(errors) or "all mirrors failed") + + +def _assert_fresh(data: dict) -> None: + """Reject a response whose OSM base timestamp is missing, malformed or old. + + Every healthy instance stamps ``osm3s.timestamp_osm_base`` with the moment + its planet copy was last updated. A mirror that is broken or serving a stale + regional extract gets this wrong, and its empty answers look exactly like + real ones — so the timestamp is the cheapest honest signal we have. + """ + raw = str((data.get("osm3s") or {}).get("timestamp_osm_base", "")).strip() + if not raw: + raise _StaleMirror("gave no OSM base timestamp") + try: + base = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + raise _StaleMirror(f"gave a malformed OSM base timestamp ({raw!r})") from None + if base.tzinfo is None: + base = base.replace(tzinfo=timezone.utc) + lag = (datetime.now(timezone.utc) - base).days + if lag > _MAX_MIRROR_LAG_DAYS: + raise _StaleMirror(f"is {lag} days out of date (base {raw})") + + +def _host(url: str) -> str: + return url.split("//", 1)[-1].split("/", 1)[0] + + +def _brief(exc: Exception, limit: int = 120) -> str: + msg = f"{type(exc).__name__}: {exc}".replace("\n", " ") + return msg if len(msg) <= limit else msg[: limit - 3] + "..." + + +BBox = Tuple[float, float, float, float] + + +def _quad_split(bbox: BBox) -> List[BBox]: + minx, miny, maxx, maxy = bbox + mx, my = (minx + maxx) / 2.0, (miny + maxy) / 2.0 + return [ + (minx, miny, mx, my), + (mx, miny, maxx, my), + (minx, my, mx, maxy), + (mx, my, maxx, maxy), + ] + + +# shared boundary IO + tiled fetching (used by both roads and bridges) @dataclass class _OSMBoundaryIO: out_sr: int = 5070 @@ -152,35 +463,40 @@ def _write_gpkg( @dataclass -class DownloadOSMRoads(_OSMBoundaryIO): - timeout: int = 300 # per-tile Overpass timeout (seconds) - max_attempts: int = 5 # retries per tile - sleep_base: float = ( - 3.0 # base backoff (seconds); actual wait = base * attempt + jitter - ) - - # Tile area target in sq-degrees. Tiles are sized adaptively from the actual bbox: - _TILE_AREA_DEG_SQ: float = 0.25 - _MAX_WORKERS: int = 8 - - _OVERPASS_MIRRORS: List[str] = None +class _TiledOverpassDownloader(_OSMBoundaryIO): + """Tiling, concurrency and failed-tile recovery shared by roads/bridges.""" - def __post_init__(self): - self._OVERPASS_MIRRORS = [ - "https://overpass-api.de/api/interpreter", - "https://lz4.overpass-api.de/api/interpreter", - "https://z.overpass-api.de/api/interpreter", - ] - self._session = requests.Session() - self._session.headers["User-Agent"] = ( - "fimbox/0.1 (+https://github.com/sdmlua/fimbox)" - ) + timeout: int = 180 # per-tile Overpass timeout (seconds) + max_attempts: int = 4 # mirror attempts per tile + sleep_base: float = 2.0 # backoff between attempts (seconds) + allow_partial: bool = False # write what we got even if tiles failed + verify_empty_tiles: bool = True # confirm "no features" on a second mirror + mirrors: Optional[Sequence[str]] = None - # adaptive tiling + # Tile area target in sq-degrees; tiles are sized adaptively from the bbox. + _TILE_AREA_DEG_SQ: float = 0.25 + _MAX_WORKERS: int = 6 + _MAX_SUBSPLIT_DEPTH: int = 3 + + _client: Optional[_OverpassClient] = field(default=None, repr=False) + + # subclasses set these + _label: str = "features" + + @property + def client(self) -> _OverpassClient: + if self._client is None: + self._client = _OverpassClient( + timeout=self.timeout, + max_attempts=self.max_attempts, + sleep_base=self.sleep_base, + mirrors=self.mirrors, + ) + return self._client - def _make_tiles( - self, minx: float, miny: float, maxx: float, maxy: float - ) -> List[Tuple[float, float, float, float]]: + # tiling + def _make_tiles(self, minx: float, miny: float, maxx: float, maxy: float + ) -> List[BBox]: area = (maxx - minx) * (maxy - miny) n = max(1, math.ceil(area / self._TILE_AREA_DEG_SQ)) # distribute tiles to respect the bbox aspect ratio @@ -195,101 +511,202 @@ def _make_tiles( for j in range(ny) ] + def _tiles_for_geom(self, geom4326) -> List[BBox]: + """Tiles covering the AOI, minus the ones the AOI never touches.""" + tiles = self._make_tiles(*geom4326.bounds) + if len(tiles) == 1: + return tiles + kept = [t for t in tiles if geom4326.intersects(box(*t))] + if len(kept) < len(tiles): + log.info( + "OSM %s: %d/%d tile(s) touch the AOI", + self._label, + len(kept), + len(tiles), + ) + return kept or tiles + def _n_workers(self, n_tiles: int) -> int: - return min(n_tiles, self._MAX_WORKERS) + n_live = max(1, len(self.client.live_mirrors())) + # ~2 concurrent requests per live mirror keeps us fast without getting + # rate-limited off a single host. + return max(1, min(self._MAX_WORKERS, n_tiles, 2 * n_live)) + + # per-tile hook implemented by subclasses + def _build_query(self, bbox: BBox) -> str: # pragma: no cover - abstract + raise NotImplementedError + + def _parse(self, osm_json: dict) -> gpd.GeoDataFrame: # pragma: no cover + raise NotImplementedError + + def _fetch_tile(self, bbox: BBox) -> gpd.GeoDataFrame: + query = self._build_query(bbox) + data, mirror = self.client.fetch(query) + gdf = self._parse(data) + + # An empty tile is often genuine (no bridges out here), but it is also + # what a misbehaving mirror returns. When another mirror is available, + # ask it before believing the hole. + if gdf.empty and self.verify_empty_tiles: + others = [m for m in self.client.live_mirrors() if m != mirror] + if others: + try: + data2, mirror2 = self.client.fetch(query, avoid=mirror) + except _TileFailed: + return gdf + second = self._parse(data2) + if not second.empty: + log.warning( + "Overpass: %s reported tile %s empty but %s returned %d " + "feature(s) — trusting %s", + _host(mirror), + _fmt_bbox(bbox), + _host(mirror2), + len(second), + _host(mirror2), + ) + self.client._mark_dead(mirror) + return second + return gdf - # Overpass - def _overpass_query(self, bbox: Tuple[float, float, float, float]) -> dict: - # Overpass bbox order: S,W,N,E - minx, miny, maxx, maxy = bbox - query = ( - f"[out:json][timeout:{self.timeout}];" - f'(way["highway"~"^motorway$|^trunk$|^primary$|^secondary$|^tertiary$"]' - f'[!"bridge"]({miny},{minx},{maxy},{maxx}););' - f"out body;>;out skel qt;" - ) - # Must use GET with params={'data': ...}; POST with raw body returns 406 on overpass-api.de - last_exc: Exception = RuntimeError("no attempt") - for attempt in range(1, self.max_attempts + 1): - mirror = self._OVERPASS_MIRRORS[(attempt - 1) % len(self._OVERPASS_MIRRORS)] + def _run_pool( + self, tiles: List[BBox], desc: str + ) -> Tuple[List[gpd.GeoDataFrame], List[BBox]]: + parts: List[gpd.GeoDataFrame] = [] + failed: List[BBox] = [] + workers = self._n_workers(len(tiles)) + + def one(t: BBox): try: - r = self._session.get( - mirror, params={"data": query}, timeout=self.timeout + 60 - ) - if r.status_code in (429, 502, 503, 504): - raise RuntimeError(f"HTTP {r.status_code} from {mirror}") - r.raise_for_status() - return r.json() + return t, self._fetch_tile(t), None except Exception as exc: - last_exc = exc - if attempt < self.max_attempts: - time.sleep(self.sleep_base * attempt + random.uniform(0, 2.0)) - raise RuntimeError( - f"Overpass query failed after {self.max_attempts} attempts: {last_exc}" - ) from last_exc - - # JSON --> GDF + return t, None, exc + + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(one, t) for t in tiles] + with tqdm(total=len(tiles), desc=desc, unit="tile") as pbar: + for fut in as_completed(futures): + tile, gdf, exc = fut.result() + if exc is not None: + failed.append(tile) + tqdm.write(f" [warn] tile {_fmt_bbox(tile)} failed: {exc}") + elif gdf is not None and not gdf.empty: + parts.append(gdf) + pbar.update(1) + + return parts, failed + + def _fetch_all(self, tiles: List[BBox]) -> gpd.GeoDataFrame: + log.info( + "OSM %s: %d tile(s), %d worker(s)", + self._label, + len(tiles), + self._n_workers(len(tiles)), + ) + parts, failed = self._run_pool(tiles, f"OSM {self._label}") + + # A failed tile is usually "too much data for this mirror" — quartering + # it makes each request small enough to succeed. This is what lets big + # AOIs finish instead of losing whole blocks of the domain. + depth = 0 + while failed and depth < self._MAX_SUBSPLIT_DEPTH: + depth += 1 + sub = [s for t in failed for s in _quad_split(t)] + log.warning( + "OSM %s: %d tile(s) failed — retrying as %d smaller tiles (pass %d)", + self._label, + len(failed), + len(sub), + depth, + ) + more, failed = self._run_pool(sub, f"OSM {self._label} retry {depth}") + parts.extend(more) + + if failed: + msg = ( + f"OSM {self._label}: {len(failed)} tile(s) could not be downloaded " + f"after {self._MAX_SUBSPLIT_DEPTH} sub-split retries " + f"(e.g. {_fmt_bbox(failed[0])})" + ) + if not self.allow_partial: + # Refuse to write a layer with holes in it — downstream FIM + # would read the gaps as "no roads/bridges here". + raise RuntimeError(msg + "; set allow_partial=True to keep partial data") + log.warning(msg + " — keeping partial data (allow_partial=True)") + + if not parts: + return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") + return gpd.GeoDataFrame(pd.concat(parts, ignore_index=True), crs="EPSG:4326") + + # shared JSON -> lines parsing @staticmethod - def _json_to_lines_gdf(osm_json: dict) -> gpd.GeoDataFrame: + def _ways_to_rows(osm_json: dict, tag_cols: Sequence[str]) -> List[dict]: + """Ways -> row dicts with a fixed column set. + + Handles both ``out geom`` (coords inline on the way) and the + ``out body;>;out skel qt`` node-lookup form, so the parser survives a + mirror that ignores the geometry modifier. + """ elems = osm_json.get("elements", []) nodes = { - e["id"]: (e["lon"], e["lat"]) for e in elems if e.get("type") == "node" + e["id"]: (e["lon"], e["lat"]) + for e in elems + if e.get("type") == "node" and "lon" in e and "lat" in e } - rows = [] + + rows: List[dict] = [] for e in elems: if e.get("type") != "way": continue - coords = [nodes[nid] for nid in e.get("nodes", []) if nid in nodes] + geom = e.get("geometry") + if geom: + coords = [ + (p["lon"], p["lat"]) + for p in geom + if p and p.get("lon") is not None and p.get("lat") is not None + ] + else: + coords = [nodes[n] for n in e.get("nodes", []) if n in nodes] if len(coords) < 2: continue tags = e.get("tags") or {} - rows.append( - { - "osmid": str(e["id"]), - "highway": tags.get("highway", "unknown"), - "name": tags.get("name", ""), - "ref": tags.get("ref", ""), - "surface": tags.get("surface", ""), - "lanes": tags.get("lanes", ""), - "geometry": LineString(coords), - } - ) - if not rows: - return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") - return gpd.GeoDataFrame(rows, crs="EPSG:4326") + row: Dict[str, Any] = {"osmid": str(e["id"])} + for col in tag_cols: + row[col] = tags.get(col, "") + row["geometry"] = LineString(coords) + rows.append(row) + return rows - # parallel fetch - def _fetch_tile(self, tile: Tuple[float, float, float, float]) -> gpd.GeoDataFrame: - try: - return self._json_to_lines_gdf(self._overpass_query(tile)) - except Exception as exc: - tqdm.write(f" [warn] tile {tile} skipped: {exc}") - return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") - def _fetch_all( - self, tiles: List[Tuple[float, float, float, float]] - ) -> gpd.GeoDataFrame: - workers = self._n_workers(len(tiles)) - parts: List[gpd.GeoDataFrame] = [] +def _fmt_bbox(bbox: BBox) -> str: + return "({:.3f}, {:.3f}, {:.3f}, {:.3f})".format(*bbox) - if workers == 1: - for tile in tqdm(tiles, desc="OSM roads", unit="tile"): - r = self._fetch_tile(tile) - if not r.empty: - parts.append(r) - else: - with ThreadPoolExecutor(max_workers=workers) as pool: - futures = {pool.submit(self._fetch_tile, t): t for t in tiles} - with tqdm(total=len(tiles), desc="OSM road tiles", unit="tile") as pbar: - for fut in as_completed(futures): - r = fut.result() - if not r.empty: - parts.append(r) - pbar.update(1) - if not parts: +@dataclass +class DownloadOSMRoads(_TiledOverpassDownloader): + _label: str = "roads" + + _HIGHWAY_RE: str = "^motorway$|^trunk$|^primary$|^secondary$|^tertiary$" + _TAG_COLS: Tuple[str, ...] = ("highway", "name", "ref", "surface", "lanes") + + def _build_query(self, bbox: BBox) -> str: + # Overpass bbox order: S,W,N,E. `out geom` returns each way's coords + # inline — smaller payload and no recursive node download. + minx, miny, maxx, maxy = bbox + return ( + f"[out:json][timeout:{self.timeout}];" + f'way["highway"~"{self._HIGHWAY_RE}"][!"bridge"]' + f"({miny},{minx},{maxy},{maxx});" + f"out geom;" + ) + + def _parse(self, osm_json: dict) -> gpd.GeoDataFrame: + rows = self._ways_to_rows(osm_json, self._TAG_COLS) + if not rows: return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") - return gpd.GeoDataFrame(pd.concat(parts, ignore_index=True), crs="EPSG:4326") + gdf = gpd.GeoDataFrame(rows, crs="EPSG:4326") + gdf["highway"] = gdf["highway"].replace("", "unknown") + return gdf # public API def query_to_gdf( @@ -309,13 +726,7 @@ def query_to_gdf( restrict_to_boundary: bool = True, ) -> gpd.GeoDataFrame: geom4326 = self._boundary_to_geom4326(boundary, boundary_layer, boundary_crs) - minx, miny, maxx, maxy = geom4326.bounds - tiles = self._make_tiles(minx, miny, maxx, maxy) - log.info( - f"OSM roads: {len(tiles)} tile(s), {self._n_workers(len(tiles))} worker(s)" - ) - - gdf = self._fetch_all(tiles) + gdf = self._fetch_all(self._tiles_for_geom(geom4326)) if gdf.empty: log.warning("No OSM road features returned.") return gpd.GeoDataFrame(geometry=[], crs=f"EPSG:{self.out_sr}") @@ -369,24 +780,38 @@ def download( @dataclass -class DownloadOSMBridges(_OSMBoundaryIO): - requests_timeout: int = 300 - max_attempts: int = 5 - sleep_base: float = 2.0 +class DownloadOSMBridges(_TiledOverpassDownloader): + _label: str = "bridges" + + requests_timeout: int = 180 # kept for backwards compatibility dissolve_buffer: float = 0.0001 # dissolve happens in EPSG:4326 - drop_list_columns: bool = True + drop_list_columns: bool = True # retained; schema is curated already - # OSMnx hits a single Overpass endpoint by default; rotate across mirrors - # (same hosts the roads downloader uses) so one dead host doesn't fail bridges. - # NOTE: ox.settings.overpass_url is the base; OSMnx appends "/interpreter". - _OVERPASS_MIRRORS: List[str] = None + # Only these tags are carried through. osmnx returned the raw OSM schema, + # which differs per area and routinely broke the GPKG write with list-valued + # and duplicate-cased columns; a fixed schema removes that failure mode. + _TAG_COLS: Tuple[str, ...] = ("bridge", "highway", "railway", "name", "layer") def __post_init__(self): - self._OVERPASS_MIRRORS = [ - "https://overpass-api.de/api", - "https://lz4.overpass-api.de/api", - "https://z.overpass-api.de/api", - ] + # honour the legacy field name if a caller sets it explicitly + if self.requests_timeout and self.requests_timeout != self.timeout: + self.timeout = self.requests_timeout + + def _build_query(self, bbox: BBox) -> str: + # bridge=no means "explicitly not a bridge", so it is excluded. + minx, miny, maxx, maxy = bbox + return ( + f"[out:json][timeout:{self.timeout}];" + f'way["bridge"]["bridge"!="no"]' + f"({miny},{minx},{maxy},{maxx});" + f"out geom;" + ) + + def _parse(self, osm_json: dict) -> gpd.GeoDataFrame: + rows = self._ways_to_rows(osm_json, self._TAG_COLS) + if not rows: + return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") + return gpd.GeoDataFrame(rows, crs="EPSG:4326") @staticmethod def _find_touching_groups(gdf: gpd.GeoDataFrame) -> List[set]: @@ -449,14 +874,15 @@ def _make_bridge_type(gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame: if "railway" not in gdf.columns: gdf["railway"] = None - gdf["bridge_type"] = gdf.apply( - lambda row: ( - f"highway-{row['highway']}" - if pd.notna(row["highway"]) - else f"railway-{row['railway']}" - ), - axis=1, - ) + def _label(row) -> str: + hw, rw = row["highway"], row["railway"] + if pd.notna(hw) and str(hw) != "": + return f"highway-{hw}" + if pd.notna(rw) and str(rw) != "": + return f"railway-{rw}" + return "bridge-other" + + gdf["bridge_type"] = gdf.apply(_label, axis=1) return gdf @staticmethod @@ -505,68 +931,28 @@ def to_line(geom): gdf = gdf[gdf.geometry.notna()].copy() return gdf - def _pull_bridges_osmnx( - self, geom4326: Union[Polygon, MultiPolygon] - ) -> gpd.GeoDataFrame: - # osmnx reads better in 4326; we output 5070 later - ox.settings.requests_timeout = self.requests_timeout - - for attempt in range(1, self.max_attempts + 1): - # rotate Overpass mirror each attempt so one dead host doesn't fail bridges - ox.settings.overpass_url = self._OVERPASS_MIRRORS[ - (attempt - 1) % len(self._OVERPASS_MIRRORS) - ] - try: - gdf = ox.features_from_polygon(geom4326, {"bridge": True}) - if gdf is None or len(gdf) == 0: - return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") - - # OSMnx returns multiindex (element, id). Keep id as osmid and drop element. - if ( - isinstance(gdf.index, pd.MultiIndex) - and "element" in gdf.index.names - ): - gdf = gdf.droplevel("element") - - gdf = gdf.copy() - gdf["osmid"] = gdf.index.astype(str) - - gdf = gdf.reset_index(drop=True) - if gdf.crs is None: - gdf = gdf.set_crs("EPSG:4326") - else: - gdf = gdf.to_crs("EPSG:4326") - - return gdf - - except Exception as e: - wait = self.sleep_base * attempt + random.uniform(0, 1.5) - time.sleep(wait) - if attempt == self.max_attempts: - raise RuntimeError( - f"osmnx bridges query failed after {self.max_attempts} attempts: {e}" - ) from e - - return gpd.GeoDataFrame(geometry=[], crs="EPSG:4326") - def _dissolve_touching(self, gdf_lines_4326: gpd.GeoDataFrame) -> gpd.GeoDataFrame: if gdf_lines_4326.empty: return gdf_lines_4326 - buffered = gdf_lines_4326.copy() - buffered["geometry"] = buffered.geometry.buffer(self.dissolve_buffer) - - groups = self._find_touching_groups(buffered) - - warnings.filterwarnings("ignore") - dissolved_groups = [] - for grp in groups: - gg = buffered.loc[list(grp)] - if gg.empty: - continue - d = gg.dissolve() - d = d.explode(index_parts=False) - dissolved_groups.append(d) + # Buffering in degrees is deliberate: 0.0001 deg is ~11 m, the tolerance + # that joins two decks of the same crossing. The CRS warning that comes + # with it is expected, so it is silenced here and nowhere else. + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + buffered = gdf_lines_4326.copy() + buffered["geometry"] = buffered.geometry.buffer(self.dissolve_buffer) + + groups = self._find_touching_groups(buffered) + + dissolved_groups = [] + for grp in groups: + gg = buffered.loc[list(grp)] + if gg.empty: + continue + d = gg.dissolve() + d = d.explode(index_parts=False) + dissolved_groups.append(d) if not dissolved_groups: out = buffered.copy() @@ -604,10 +990,15 @@ def query_to_gdf( ) -> gpd.GeoDataFrame: geom4326 = self._boundary_to_geom4326(boundary, boundary_layer, boundary_crs) - gdf = self._pull_bridges_osmnx(geom4326) + gdf = self._fetch_all(self._tiles_for_geom(geom4326)) if gdf.empty: + log.warning("No OSM bridge features returned.") return gpd.GeoDataFrame(geometry=[], crs=f"EPSG:{self.out_sr}") + # a bridge on a tile seam comes back from both tiles + gdf = gdf.drop_duplicates(subset=["osmid"]).reset_index(drop=True) + log.info(f"OSM bridges: {len(gdf)} unique ways after dedup") + gdf = self._clean_schema(gdf, drop_list_columns=self.drop_list_columns) gdf = self._make_bridge_type(gdf) gdf = self._filter_unwanted_bridge_types(gdf) @@ -630,6 +1021,7 @@ def query_to_gdf( if gdf.empty: return gpd.GeoDataFrame(geometry=[], crs=f"EPSG:{self.out_sr}") + log.info(f"OSM bridges: {len(gdf)} bridge lines in the AOI") return gdf.to_crs(epsg=self.out_sr) def download( @@ -663,7 +1055,8 @@ def download( boundary_crs=boundary_crs, restrict_to_boundary=True, ) - self._write_gpkg(gdf, out_dir, out_name, out_layer) + out_path = self._write_gpkg(gdf, out_dir, out_name, out_layer) + log.info(f"{out_layer} --> {out_path.name}") return gdf @@ -702,6 +1095,11 @@ def download( p.add_argument( "--boundary_crs", default=None, help="CRS for bbox/shapely boundary, e.g. 4326" ) + p.add_argument( + "--allow_partial", + action="store_true", + help="Write the layer even if some tiles could not be downloaded", + ) args = p.parse_args() boundary_val: Any = args.boundary @@ -717,7 +1115,7 @@ def download( pass if args.mode == "roads": - dl = DownloadOSMRoads(out_sr=5070) + dl = DownloadOSMRoads(out_sr=5070, allow_partial=args.allow_partial) dl.download( boundary=boundary_val, out_dir=args.out_dir, @@ -727,7 +1125,7 @@ def download( boundary_crs=boundary_crs_val, ) else: - dl = DownloadOSMBridges(out_sr=5070) + dl = DownloadOSMBridges(out_sr=5070, allow_partial=args.allow_partial) dl.download( boundary=boundary_val, out_dir=args.out_dir, diff --git a/src/fimbox/preprocessing/preprocess_area.py b/src/fimbox/preprocessing/preprocess_area.py index 4bf2130..6ae1b4e 100644 --- a/src/fimbox/preprocessing/preprocess_area.py +++ b/src/fimbox/preprocessing/preprocess_area.py @@ -892,15 +892,17 @@ def run_nld(self): return self.logger.info("--- NLD ---") + nld_failed: list = [] if not lines_exist: try: - DownloadNLD( + nld = DownloadNLD( boundary=self.buffer_gdf, out_dir=str(self.case_dir), epsg=self.epsg, lines_name=_FILENAMES["levee_lines"], polys_name=_FILENAMES["levee_protected_areas"], ) + nld_failed = list(nld.failed_layers) except Exception as exc: self.logger.error(f"NLD download failed: {exc}", exc_info=True) return @@ -925,6 +927,12 @@ def run_nld(self): f"Levee burn lines ({len(burned)} features) --> " f"{_FILENAMES['levee_lines_burned']}" ) + elif "lines" in nld_failed: + self.logger.error( + "NLD: levee lines could not be downloaded — levee burn " + "skipped. This is a download failure, not an area " + "without levees; re-run before using this HAND output." + ) else: self.logger.info( "NLD: no levee lines in this area — skipping levee burn." @@ -940,6 +948,12 @@ def run_nld(self): self.logger.info( f"Levee protected areas --> {_FILENAMES['levee_protected_areas']}" ) + elif "polys" in nld_failed: + self.logger.error( + "NLD: leveed-area polygons could not be downloaded — " + "levee masking will be missing, and this is a download " + "failure rather than an area without leveed areas." + ) else: self.logger.info("NLD: no levee protected areas in this area.") diff --git a/src/fimbox/preprocessing/process_bridgedem/__init__.py b/src/fimbox/preprocessing/process_bridgedem/__init__.py index eef8063..485de6c 100644 --- a/src/fimbox/preprocessing/process_bridgedem/__init__.py +++ b/src/fimbox/preprocessing/process_bridgedem/__init__.py @@ -1,4 +1,5 @@ from .bridge_dem_diff import BridgeDEMDiff from .bridge_lidar_raster import generateBridgeRaster +from .bridge_source import resolve_bridge_gpkg -__all__ = ["generateBridgeRaster", "BridgeDEMDiff"] +__all__ = ["generateBridgeRaster", "BridgeDEMDiff", "resolve_bridge_gpkg"] diff --git a/src/fimbox/preprocessing/process_bridgedem/bridge_dem_diff.py b/src/fimbox/preprocessing/process_bridgedem/bridge_dem_diff.py index 26e2405..321947f 100644 --- a/src/fimbox/preprocessing/process_bridgedem/bridge_dem_diff.py +++ b/src/fimbox/preprocessing/process_bridgedem/bridge_dem_diff.py @@ -24,6 +24,8 @@ from shapely.geometry import Point from tqdm import tqdm +from .bridge_source import resolve_bridge_gpkg + log = logging.getLogger(__name__) @@ -160,6 +162,7 @@ def _run(self) -> Path: return out_path def _load_bridges(self) -> gpd.GeoDataFrame: + self.bridge_gpkg = resolve_bridge_gpkg(self.bridge_gpkg) gdf = gpd.read_file(self.bridge_gpkg) if "osmid" in gdf.columns: col = "osmid" diff --git a/src/fimbox/preprocessing/process_bridgedem/bridge_lidar_raster.py b/src/fimbox/preprocessing/process_bridgedem/bridge_lidar_raster.py index f8f2906..96e50d1 100644 --- a/src/fimbox/preprocessing/process_bridgedem/bridge_lidar_raster.py +++ b/src/fimbox/preprocessing/process_bridgedem/bridge_lidar_raster.py @@ -11,10 +11,12 @@ from __future__ import annotations +import io import logging import os import tempfile import threading as _threading +from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path @@ -27,23 +29,64 @@ from scipy.spatial import KDTree from tqdm import tqdm +from .bridge_source import resolve_bridge_gpkg + log = logging.getLogger(__name__) # LAS classification 13=bridge deck, 17=bridge deck _BRIDGE_CLASSES = {13, 17} _ENTWINE_INDEX_URL = "https://raw.githubusercontent.com/hobuinc/usgs-lidar/master/boundaries/boundaries.topojson" -# Per-thread session — avoids connection pool exhaustion when many dask threads run concurrently. _thread_local = _threading.local() +_SESSION: Optional[requests.Session] = None +_SESSION_LOCK = Lock() + def _session() -> requests.Session: - if not hasattr(_thread_local, "session"): - s = requests.Session() - adapter = requests.adapters.HTTPAdapter(pool_connections=4, pool_maxsize=16) - s.mount("https://", adapter) - _thread_local.session = s - return _thread_local.session + """One pooled, retrying session shared by every download in the process. + + A session per thread means a fresh DNS lookup and TLS handshake for each of + the thousands of short-lived tile threads — more wall-clock than the tile + download itself, and enough resolver traffic that S3 starts handing back + intermittent NXDOMAIN. Retries cover those transients so a hiccup costs a + backoff instead of the whole bridge. + """ + global _SESSION + with _SESSION_LOCK: + if _SESSION is None: + from urllib3.util.retry import Retry + + s = requests.Session() + adapter = requests.adapters.HTTPAdapter( + pool_connections=16, + pool_maxsize=64, + max_retries=Retry( + total=4, + connect=4, + read=2, + backoff_factor=0.4, + status_forcelist=(429, 500, 502, 503, 504), + allowed_methods=frozenset(["GET"]), + ), + ) + s.mount("https://", adapter) + s.mount("http://", adapter) + _SESSION = s + return _SESSION + + +def _transformer(src: str, dst: str): + """Thread-local pyproj Transformer — building one costs more than it should + to repeat 3k times, and PROJ objects can't be shared across threads.""" + cache = getattr(_thread_local, "transformers", None) + if cache is None: + cache = _thread_local.transformers = {} + if (src, dst) not in cache: + from pyproj import Transformer + + cache[(src, dst)] = Transformer.from_crs(src, dst, always_xy=True) + return cache[(src, dst)] # EPT manifest + hierarchy cache: fetch once per unique EPT URL, reuse for all bridges @@ -64,80 +107,155 @@ def _ept_meta(base: str) -> tuple[dict, dict]: return _EPT_CACHE[base] -def _fetch_one_tile(args) -> Optional[np.ndarray]: - """Download + decode one EPT .laz tile. Returns (N,4) [x,y,z,cls] in EPSG:3857 or None.""" +# Decoded-tile cache. Bridges are neighbours far more often than not, so the +# same octree tiles keep coming back: on a HUC8 with ~2.9k bridges, 7.8k tile +# requests resolve to under 2k distinct tiles, and the coarse depth-6 ones are +# asked for over a hundred times each. Tiles are small (~0.5 MB decoded), so +# holding a working set in memory turns most of those requests into a dict hit. +_TILE_MISS = object() +_TILE_CACHE: "OrderedDict[tuple, Optional[np.ndarray]]" = OrderedDict() +_TILE_CACHE_LOCK = Lock() +_TILE_CACHE_BYTES = 0 +_TILE_CACHE_MAX_BYTES = 1024 * 1024 * 1024 +_TILE_FETCH_LOCKS: dict = {} +_TILE_REQUESTS = 0 +_TILE_DOWNLOADS = 0 + + +def _tile_cache_get(key: tuple): + with _TILE_CACHE_LOCK: + if key in _TILE_CACHE: + _TILE_CACHE.move_to_end(key) + return _TILE_CACHE[key] + return _TILE_MISS + + +def _tile_cache_put(key: tuple, pts: Optional[np.ndarray]) -> None: + global _TILE_CACHE_BYTES + with _TILE_CACHE_LOCK: + old = _TILE_CACHE.pop(key, None) + if old is not None: + _TILE_CACHE_BYTES -= old.nbytes + _TILE_CACHE[key] = pts + _TILE_CACHE_BYTES += 0 if pts is None else pts.nbytes + while _TILE_CACHE_BYTES > _TILE_CACHE_MAX_BYTES and len(_TILE_CACHE) > 1: + _, evicted = _TILE_CACHE.popitem(last=False) + if evicted is not None: + _TILE_CACHE_BYTES -= evicted.nbytes + + +def _tile_fetch_lock(key: tuple) -> Lock: + with _TILE_CACHE_LOCK: + lk = _TILE_FETCH_LOCKS.get(key) + if lk is None: + lk = _TILE_FETCH_LOCKS[key] = Lock() + return lk + + +def _download_tile(base: str, tile_key: str) -> Optional[np.ndarray]: + """Download + decode one EPT .laz tile. + + Returns every last-return point in the tile as (N,4) [x,y,z,cls] in + EPSG:3857 — unclipped, so the result is reusable by any bridge that lands + in this tile — or None if the tile is absent or has no last returns. + """ import laspy - tile_key, base, qxmin, qymin, qxmax, qymax = args - url = f"{base}/ept-data/{tile_key}.laz" - resp = _session().get(url, timeout=60) + global _TILE_DOWNLOADS + resp = _session().get(f"{base}/ept-data/{tile_key}.laz", timeout=60) if resp.status_code == 404: return None resp.raise_for_status() + with _TILE_CACHE_LOCK: + _TILE_DOWNLOADS += 1 + + # Decoding from memory: the tile is already fully buffered, so a temp file + # only adds a write and a read. + las = laspy.read(io.BytesIO(resp.content)) + last = np.asarray(las.return_number) == np.asarray(las.number_of_returns) + if not last.any(): + return None + return np.column_stack( + [ + np.asarray(las.x)[last], + np.asarray(las.y)[last], + np.asarray(las.z)[last], + np.asarray(las.classification)[last].astype(float), + ] + ) - with tempfile.NamedTemporaryFile(suffix=".laz", delete=False) as f: - f.write(resp.content) - tmp = f.name - try: - las = laspy.read(tmp) - x = np.array(las.x) - y = np.array(las.y) - z = np.array(las.z) - cls = np.array(las.classification) - ret = np.array(las.return_number) - nr = np.array(las.number_of_returns) - last = ret == nr - inbox = (x >= qxmin) & (x <= qxmax) & (y >= qymin) & (y <= qymax) - mask = last & inbox - if not mask.any(): - return None - return np.column_stack([x[mask], y[mask], z[mask], cls[mask].astype(float)]) - finally: - os.unlink(tmp) +def _tile_points(base: str, tile_key: str) -> Optional[np.ndarray]: + """Cached last-return points for one EPT tile, downloading it at most once.""" + global _TILE_REQUESTS + key = (base, tile_key) + with _TILE_CACHE_LOCK: + _TILE_REQUESTS += 1 + + pts = _tile_cache_get(key) + if pts is not _TILE_MISS: + return pts + + # Per-tile lock so concurrent bridges wanting the same tile fetch it once. + with _tile_fetch_lock(key): + pts = _tile_cache_get(key) + if pts is not _TILE_MISS: + return pts + pts = _download_tile(base, tile_key) + _tile_cache_put(key, pts) + return pts -def _fetch_ept_points( - ept_url: str, - bounds: tuple, - out_crs: str, - tile_workers: int = 8, - min_depth: int = 6, -) -> Optional[np.ndarray]: - """ - Fetch last-return LiDAR points from EPT within `bounds` (EPSG:4326). - Returns (N,4) [x,y,z,cls] reprojected to out_crs, or None. + +def _bridge_tiles(ept_url: str, bounds: tuple, min_depth: int = 6) -> tuple: + """Resolve a bridge's 4326 `bounds` to (base, tile keys, query bbox in 3857). `min_depth` skips coarse octree tiles (depth < min_depth) that contain almost no points inside a tiny bridge bbox, saving several tile downloads. """ - from pyproj import Transformer - base = ept_url.rstrip("/").replace("/ept.json", "").rstrip("/") manifest, hierarchy = _ept_meta(base) ept_bounds = manifest["bounds"] # [xmin,ymin,zmin,xmax,ymax,zmax] in EPSG:3857 - tr = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True) + tr = _transformer("EPSG:4326", "EPSG:3857") qxmin, qymin = tr.transform(bounds[0], bounds[1]) qxmax, qymax = tr.transform(bounds[2], bounds[3]) + query = (qxmin, qymin, qxmax, qymax) + return base, tuple(_intersecting_tiles(hierarchy, ept_bounds, query, min_depth)), query - tiles = _intersecting_tiles( - hierarchy, ept_bounds, (qxmin, qymin, qxmax, qymax), min_depth - ) + +def _fetch_ept_points( + base: str, + tiles: tuple, + query: tuple, + out_crs: str, + pool: ThreadPoolExecutor, +) -> Optional[np.ndarray]: + """ + Fetch last-return LiDAR points from EPT within `query` (EPSG:3857 bbox). + Returns (N,4) [x,y,z,cls] reprojected to out_crs, or None. + """ if not tiles: return None - args = [(t, base, qxmin, qymin, qxmax, qymax) for t in tiles] + qxmin, qymin, qxmax, qymax = query all_pts = [] - with ThreadPoolExecutor(max_workers=min(tile_workers, len(tiles))) as pool: - for arr in pool.map(_fetch_one_tile, args): - if arr is not None: - all_pts.append(arr) + for arr in pool.map(lambda t: _tile_points(base, t), tiles): + if arr is None: + continue + inbox = ( + (arr[:, 0] >= qxmin) + & (arr[:, 0] <= qxmax) + & (arr[:, 1] >= qymin) + & (arr[:, 1] <= qymax) + ) + if inbox.any(): + all_pts.append(arr[inbox]) if not all_pts: return None pts_3857 = np.vstack(all_pts) - tr_out = Transformer.from_crs("EPSG:3857", out_crs, always_xy=True) + tr_out = _transformer("EPSG:3857", out_crs) ox, oy = tr_out.transform(pts_3857[:, 0], pts_3857[:, 1]) return np.column_stack([ox, oy, pts_3857[:, 2], pts_3857[:, 3]]) @@ -197,7 +315,9 @@ class generateBridgeRaster: user-supplied value; uses row index if not found. skip_ids : ID values to skip n_workers : parallel worker threads for bridge-level processing (default: all CPUs) - tile_workers : threads for per-bridge EPT tile downloads (default 8) + tile_workers : threads per bridge for EPT tile downloads (default 8). Backed by one + shared pool of n_workers * tile_workers threads, so connections and + downloaded tiles are reused across bridges. min_tile_depth: skip EPT octree tiles shallower than this depth (default 6). Coarse tiles cover huge areas; almost zero bridge points fall in them. bridge_cls_threshold: fraction of points that must be class 13/17 to use only those; @@ -205,6 +325,8 @@ class generateBridgeRaster: that don't classify bridge decks). skip_existing: if True (default), skip bridges whose output .tif already exists so re-runs only process new bridges instead of re-downloading everything. + tile_cache_mb: memory budget for the decoded-tile cache (default 1024 MB). Neighbouring + bridges share octree tiles; cached tiles are ~4x fewer downloads. """ bridge_gpkg: Union[str, Path] @@ -216,6 +338,7 @@ class generateBridgeRaster: min_tile_depth: int = 6 bridge_cls_threshold: float = 0.05 skip_existing: bool = True + tile_cache_mb: float = 1024.0 id_col: Optional[str] = None skip_ids: list = field(default_factory=lambda: ["229091666"]) @@ -280,6 +403,7 @@ def run(self) -> Path: return self._tif_dir def _load_bridges(self) -> gpd.GeoDataFrame: + self.bridge_gpkg = resolve_bridge_gpkg(self.bridge_gpkg) gdf = gpd.read_file(self.bridge_gpkg) if "osmid" in gdf.columns: col = "osmid" @@ -340,45 +464,111 @@ def _assign_lidar_urls(self, footprints, index) -> gpd.GeoDataFrame: ) return joined + def _plan(self, footprints: gpd.GeoDataFrame) -> list: + """Resolve every bridge to its EPT tiles, then order them for tile reuse. + + Doing the octree lookup up front costs nothing extra (the hierarchy is + already cached per survey) and lets bridges be visited in tile order, so + a shared tile stays hot in cache while the bridges that need it go by + instead of being re-downloaded after it has been evicted. + """ + jobs = [] + for _, row in footprints.iterrows(): + base, tiles, query = _bridge_tiles( + row["url"], row.geometry.bounds, self.min_tile_depth + ) + jobs.append( + { + "bridge_id": row["_bridge_id"], + "base": base, + "tiles": tiles, + "query": query, + } + ) + jobs.sort(key=lambda j: (j["base"], j["tiles"])) + requests_n = sum(len(j["tiles"]) for j in jobs) + unique_n = len({(j["base"], t) for j in jobs for t in j["tiles"]}) + log.info( + f"Tile plan: {requests_n} tile lookups over {unique_n} unique tiles " + f"({requests_n / max(1, unique_n):.1f}x reuse), cache {self.tile_cache_mb:.0f} MB" + ) + return jobs + def _process_parallel(self, footprints: gpd.GeoDataFrame): + global _TILE_CACHE_MAX_BYTES, _TILE_REQUESTS, _TILE_DOWNLOADS + _TILE_CACHE_MAX_BYTES = int(self.tile_cache_mb * 1024 * 1024) + # Counters are per-run; the cache itself is kept, so a second AOI in the + # same process still gets to reuse any tiles it shares with the first. + with _TILE_CACHE_LOCK: + _TILE_REQUESTS = _TILE_DOWNLOADS = 0 + tif_dir = str(self._tif_dir) - n = len(footprints) - ok = failed = skipped = 0 - - with ThreadPoolExecutor(max_workers=self.n_workers) as executor: - fmap = { - executor.submit( - _process_one_bridge, - bridge_id=row["_bridge_id"], - bounds=row.geometry.bounds, - lidar_url=row["url"], - tif_dir=tif_dir, - resolution=self.resolution, - tile_workers=self.tile_workers, - min_tile_depth=self.min_tile_depth, - bridge_cls_threshold=self.bridge_cls_threshold, - skip_existing=self.skip_existing, - ): row["_bridge_id"] - for _, row in footprints.iterrows() - } - with tqdm( - total=n, desc="Bridges", unit="bridge", dynamic_ncols=True - ) as pbar: - for future in as_completed(fmap): - bid = fmap[future] - try: - result = future.result() - if result == "skipped": - skipped += 1 - else: - ok += 1 - except Exception as exc: - log.warning(f"bridge {bid} failed: {exc}") - failed += 1 - pbar.update(1) - pbar.set_postfix(ok=ok, skip=skipped, fail=failed, refresh=False) - - log.info(f"Completed — {ok} processed, {skipped} skipped, {failed} failed") + skipped = 0 + if self.skip_existing: + # Filter here as well as in the worker: a re-run shouldn't pay for the + # tile lookup, or make the progress bar count work it never does. + done = {f.stem for f in self._tif_dir.glob("*.tif")} + n_before = len(footprints) + footprints = footprints[~footprints["_bridge_id"].isin(done)] + skipped = n_before - len(footprints) + if skipped: + log.info(f"{skipped} bridges already have rasters — skipping") + + jobs = self._plan(footprints) + n = len(jobs) + ok = failed = empty = 0 + + # One download pool for the whole run, not one per bridge: creating ~3k + # short-lived pools meant ~3k cold TLS connections to S3. + pool = ThreadPoolExecutor( + max_workers=max(1, min(32, self.n_workers * self.tile_workers)), + thread_name_prefix="ept-tile", + ) + try: + with ThreadPoolExecutor(max_workers=self.n_workers) as executor: + fmap = { + executor.submit( + _process_one_bridge, + bridge_id=job["bridge_id"], + base=job["base"], + tiles=job["tiles"], + query=job["query"], + tif_dir=tif_dir, + resolution=self.resolution, + pool=pool, + bridge_cls_threshold=self.bridge_cls_threshold, + skip_existing=self.skip_existing, + ): job["bridge_id"] + for job in jobs + } + with tqdm( + total=n, desc="Bridges", unit="bridge", dynamic_ncols=True + ) as pbar: + for future in as_completed(fmap): + bid = fmap[future] + try: + result = future.result() + if result == "skipped": + skipped += 1 + elif result == "empty": + empty += 1 + else: + ok += 1 + except Exception as exc: + log.warning(f"bridge {bid} failed: {exc}") + failed += 1 + pbar.update(1) + pbar.set_postfix( + ok=ok, skip=skipped, none=empty, fail=failed, refresh=False + ) + finally: + pool.shutdown(wait=True) + + log.info( + f"Completed — {ok} processed, {skipped} skipped, {empty} without LiDAR, " + f"{failed} failed; {_TILE_DOWNLOADS} tiles downloaded of " + f"{_TILE_REQUESTS} requested" + ) def _idw_rasterize( @@ -445,12 +635,12 @@ def _idw_rasterize( def _process_one_bridge( bridge_id: str, - bounds: tuple, - lidar_url: str, + base: str, + tiles: tuple, + query: tuple, tif_dir: str, resolution: float, - tile_workers: int = 8, - min_tile_depth: int = 6, + pool: ThreadPoolExecutor, bridge_cls_threshold: float = 0.05, skip_existing: bool = True, ): @@ -463,65 +653,53 @@ def _process_one_bridge( if skip_existing and os.path.exists(tif_path): return "skipped" - try: - pts = _fetch_ept_points( - lidar_url, - bounds, - out_crs, - tile_workers=tile_workers, - min_depth=min_tile_depth, - ) - if pts is None or len(pts) == 0: - return - - xy = pts[:, :2] - z = pts[:, 2].copy() - cls = pts[:, 3].astype(int) - - bridge_mask = np.isin(cls, list(_BRIDGE_CLASSES)) - if bridge_mask.sum() / len(cls) >= bridge_cls_threshold: - # survey has bridge-deck classifications — replace non-bridge z with nearest bridge z - if (~bridge_mask).any(): - n_bridge = bridge_mask.sum() - k = min(2, n_bridge) - tree = KDTree(xy[bridge_mask]) - _, idx = tree.query(xy[~bridge_mask], k=k) - if k == 1: - idx = idx.reshape(-1, 1) - z[~bridge_mask] = z[bridge_mask][idx].mean(axis=1) - - # else: survey doesn't classify bridge decks — use all last-return points as-is - - # Compute grid bounds from point cloud extent (not EPT query bbox) - xmin, ymin = xy[:, 0].min(), xy[:, 1].min() - xmax, ymax = xy[:, 0].max(), xy[:, 1].max() - # Ensure at least one pixel - if xmax <= xmin: - xmax = xmin + resolution - if ymax <= ymin: - ymax = ymin + resolution - - grid, transform, nodata = _idw_rasterize( - xy, z, (xmin, ymin, xmax, ymax), resolution - ) - - with rasterio.open( - tif_path, - "w", - driver="GTiff", - height=grid.shape[0], - width=grid.shape[1], - count=1, - dtype="float32", - crs=_CRS.from_string(out_crs), - transform=transform, - nodata=nodata, - compress="lzw", - ) as dst: - dst.write(grid, 1) - - except Exception as exc: - log.warning(f"bridge {bridge_id} failed: {exc}") + pts = _fetch_ept_points(base, tiles, query, out_crs, pool) + if pts is None or len(pts) == 0: + return "empty" + + xy = pts[:, :2] + z = pts[:, 2].copy() + cls = pts[:, 3].astype(int) + + bridge_mask = np.isin(cls, list(_BRIDGE_CLASSES)) + if bridge_mask.sum() / len(cls) >= bridge_cls_threshold: + # survey has bridge-deck classifications — replace non-bridge z with nearest bridge z + if (~bridge_mask).any(): + n_bridge = bridge_mask.sum() + k = min(2, n_bridge) + tree = KDTree(xy[bridge_mask]) + _, idx = tree.query(xy[~bridge_mask], k=k) + if k == 1: + idx = idx.reshape(-1, 1) + z[~bridge_mask] = z[bridge_mask][idx].mean(axis=1) + + # else: survey doesn't classify bridge decks — use all last-return points as-is + + # Compute grid bounds from point cloud extent (not EPT query bbox) + xmin, ymin = xy[:, 0].min(), xy[:, 1].min() + xmax, ymax = xy[:, 0].max(), xy[:, 1].max() + # Ensure at least one pixel + if xmax <= xmin: + xmax = xmin + resolution + if ymax <= ymin: + ymax = ymin + resolution + + grid, transform, nodata = _idw_rasterize(xy, z, (xmin, ymin, xmax, ymax), resolution) + + with rasterio.open( + tif_path, + "w", + driver="GTiff", + height=grid.shape[0], + width=grid.shape[1], + count=1, + dtype="float32", + crs=_CRS.from_string(out_crs), + transform=transform, + nodata=nodata, + compress="lzw", + ) as dst: + dst.write(grid, 1) # CLI @@ -552,6 +730,7 @@ def _process_one_bridge( p.add_argument("--tile_workers", type=int, default=8) p.add_argument("--min_tile_depth", type=int, default=6) p.add_argument("--bridge_cls_threshold", type=float, default=0.05) + p.add_argument("--tile_cache_mb", type=float, default=1024.0) p.add_argument("--id_col", default=None) p.add_argument("--skip_ids", nargs="*", default=["229091666"]) args = p.parse_args() @@ -564,6 +743,7 @@ def _process_one_bridge( tile_workers=args.tile_workers, min_tile_depth=args.min_tile_depth, bridge_cls_threshold=args.bridge_cls_threshold, + tile_cache_mb=args.tile_cache_mb, id_col=args.id_col, skip_ids=args.skip_ids, ) diff --git a/src/fimbox/preprocessing/process_bridgedem/bridge_source.py b/src/fimbox/preprocessing/process_bridgedem/bridge_source.py new file mode 100644 index 0000000..267964c --- /dev/null +++ b/src/fimbox/preprocessing/process_bridgedem/bridge_source.py @@ -0,0 +1,66 @@ +""" +Author: Supath Dhital +Date updated: July 2026 +--------------------- +Resolve the bridge GeoPackage the bridge-DEM stages read. + +A missing bridge layer is recoverable: the AOI folder that lacks it still holds +the boundary the layer is derived from, so it is downloaded from OSM instead of +failing the stage. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Optional, Union + +log = logging.getLogger(__name__) + +# Boundary written by preprocessing, best-first. The buffered one is what every +# other OSM download in the pipeline is clipped to. +_BOUNDARIES = ("wbd_buffered.gpkg", "wbd.gpkg", "wbd8_clp.gpkg") + + +def resolve_bridge_gpkg( + bridge_gpkg: Union[str, Path], + boundary: Optional[Union[str, Path]] = None, + out_layer: str = "osm_bridges", +) -> Path: + """Path to a bridge GeoPackage, downloaded from OSM if it is not there yet.""" + path = Path(bridge_gpkg) + if path.exists(): + return path + log.warning(f"Bridge GeoPackage not found: {path}") + + # The pipeline's other name for the same layer — reuse it before downloading. + alt = path.parent / "osm_bridges.gpkg" + if alt != path and alt.exists(): + log.info(f"Using the bridge layer already staged --> {alt.name}") + return alt + + if boundary is None: + boundary = next( + (p for n in _BOUNDARIES if (p := path.parent / n).exists()), None + ) + if boundary is None: + raise FileNotFoundError( + f"{path} is missing and {path.parent} holds no boundary " + f"({', '.join(_BOUNDARIES)}) to download bridges for — " + f"run preprocessing for this area first." + ) + + # Imported here so a run with the layer staged never loads the Overpass stack. + from ..download_data.osm_data import DownloadOSMBridges + + log.info(f"Downloading OSM bridges for {Path(boundary).name} --> {path.name}") + gdf = DownloadOSMBridges().download( + boundary=boundary, + out_dir=path.parent, + out_name=path.name, + out_layer=out_layer, + ) + if gdf.empty: + log.warning("OSM returned no bridges in this area.") + log.info(f"{out_layer} --> {path.name}") + return path diff --git a/tests/test_fimgeneration.py b/tests/test_fimgeneration.py index b01881c..38e05c5 100644 --- a/tests/test_fimgeneration.py +++ b/tests/test_fimgeneration.py @@ -28,7 +28,7 @@ # AOI_DIR = Path(__file__).resolve().parents[2] / "out" / "test_smallB" AOI_DIR = Path(__file__).resolve().parents[2] / "out" / "nwm_11239459and2more" -N_WORKERS = 4 +N_WORKERS = None # Optional selection filters (edit to match the CSVs you have). EVENT_DATE = "2020-05-20 12:00:00" diff --git a/tests/test_getallinputdata.py b/tests/test_getallinputdata.py index 51955b4..76f683f 100644 --- a/tests/test_getallinputdata.py +++ b/tests/test_getallinputdata.py @@ -8,7 +8,7 @@ test_boundary = PKG_ROOT / "docs" / "test_boundary" / "test_smallB.shp" OUT_DIR = REPO_ROOT / "out" -test_huc8 = "08060202" # Yazoo River basin, MS +test_huc8 = "08060202" # 08060202- Yazoo River basin, MS # If Just wated to test with NWM reach IDs test_nwm_ids = [11239459, 11239689, 11235965] @@ -20,18 +20,18 @@ # Combined preprocessing pipeline tests # Run full pipeline from a boundary shapefile -# def test_preprocess_all_from_boundary(): -# pp = fimbox.getAllInputData( -# boundary=test_boundary, -# out_dir=OUT_DIR, -# buffer_m=5000, # metres to buffer boundary for data downloads -# headwater_buffer_cells=8, # pixels to shrink buffer for headwater clip -# get_flowlines=True, # set False to use your own flowlines and corresponding catchments -# get_catchments=True, # set False to skip NWM catchments--> use -# source="nwmmedium", # "nwmhigh" -> NHDPlus HR via pynhd; "ngen" -> NextGen hydrofabric. Lakes always NWM. -# identifier="nwmmr", # filename prefix for ALL source files; flows download->processing. Default "nwm". -# ) -# pp.run() +def test_preprocess_all_from_boundary(): + pp = fimbox.getAllInputData( + huc8=test_huc8, + out_dir=OUT_DIR, + buffer_m=5000, # metres to buffer boundary for data downloads + headwater_buffer_cells=8, # pixels to shrink buffer for headwater clip + get_flowlines=True, # set False to use your own flowlines and corresponding catchments + get_catchments=True, # set False to skip NWM catchments--> use + source="nwmmedium", # "nwmhigh" -> NHDPlus HR via pynhd; "ngen" -> NextGen hydrofabric. Lakes always NWM. + identifier="nwmmr", # filename prefix for ALL source files; flows download->processing. Default "nwm". + ) + pp.run() # Bring your own flowlines / catchments / DEM (any column names, any source). diff --git a/tests/test_nwmstreamflow.py b/tests/test_nwmstreamflow.py index b233801..e236d88 100644 --- a/tests/test_nwmstreamflow.py +++ b/tests/test_nwmstreamflow.py @@ -18,16 +18,16 @@ ) # AOI_DIR = Path(__file__).resolve().parents[2] / "out" / "test_smallB" -AOI_DIR = Path(__file__).resolve().parents[2] / "out" / "nwm_11239459and2more" +AOI_DIR = Path(__file__).resolve().parents[2] / "out" / "HUC08060202" START = "2016-10-05" END = "2016-10-20" EVENT = "2020-10-10 21:00:00" -# retrospective — different extraction combinations -def test_retrospective_event_date(): - getNWMretrospective(AOI_DIR, date=EVENT) +# # retrospective — different extraction combinations +# def test_retrospective_event_date(): +# getNWMretrospective(AOI_DIR, date=EVENT) # def test_retrospective_range_continuous(): @@ -35,9 +35,9 @@ def test_retrospective_event_date(): # getNWMretrospective(AOI_DIR, start=START, end=END) -# def test_retrospective_range_sortby(): -# # start + end + sortby -> one aggregated CSV -# getNWMretrospective(AOI_DIR, start=START, end=END, sortby="maximum") +def test_retrospective_range_sortby(): + # start + end + sortby -> one aggregated CSV + getNWMretrospective(AOI_DIR, start=START, end=END, sortby="maximum") # def test_retrospective_feature_ids_list():