diff --git a/pyproject.toml b/pyproject.toml index 91362bc..35ac40d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,11 @@ include-package-data = true where = ["src"] exclude = ["tests*"] +[tool.setuptools.package-data] +# Ship the NextGen per-VPU bounding-box index so AOI->VPU resolution needs no +# network round-trip. (Other runtime data is fetched on demand via pooch/S3.) +"fimbox.nextgen" = ["data/*.json"] + [tool.ruff] line-length = 88 target-version = "py310" diff --git a/src/fimbox/__init__.py b/src/fimbox/__init__.py index 1d011ba..724700e 100644 --- a/src/fimbox/__init__.py +++ b/src/fimbox/__init__.py @@ -284,3 +284,29 @@ # On-the-fly access to the bundled reference/calibration lookup tables, fetched # anonymously from the public SDML S3 bucket via pooch (imported at module top). __all__ += ["fetch_data"] + +# NextGen-in-a-Box bridge: resolve an AOI to the NextGen v2.2 hydrofabric +# catchments and pull ngen/t-route discharge from the public CIROH community +# datastream bucket. Heavy/optional deps (s3fs, xarray, geopandas) are imported +# lazily inside the functions, so guard the import to keep a partial install +# from breaking `import fimbox`. +try: + from .nextgen import ( + AOIHydrofabric, + NextGenAOI, + NextGenDatastream, + NextGenHydrofabric, + NextGenResult, + getNextGenAOI, + ) + + __all__ += [ + "getNextGenAOI", + "NextGenAOI", + "NextGenResult", + "NextGenHydrofabric", + "AOIHydrofabric", + "NextGenDatastream", + ] +except Exception: + pass diff --git a/src/fimbox/nextgen/README.md b/src/fimbox/nextgen/README.md new file mode 100644 index 0000000..6037ada --- /dev/null +++ b/src/fimbox/nextgen/README.md @@ -0,0 +1,135 @@ +### NextGen-in-a-Box Integration + +This part of fimbox connects your **area of interest** (a map outline of the place you care about) to NextGen streamflow data, so that flood maps can be produced for that area. + +In simple terms, you hand it a shapefile or geopackage of an area, and it figures out which rivers and drainage areas fall inside that shape, then goes and fetches how much water is flowing in each of them. + +**What you give it** + +- An **area of interest** — a boundary file (a `.shp` shapefile or a `.gpkg` GeoPackage) outlining the region you want. For example, a city, a watershed, or a study site. + +**What you get back**, saved in a folder named after your area: + +- A **catchments** map file — the drainage areas that fall inside your boundary. +- A **discharge** table — how much water is flowing in each of those areas (in cubic metres per second). +- The **stream IDs** — the identifiers that link each drainage area to its river reach. + +Everything is downloaded from a [free public s3 bucket](https://ciroh-community-ngen-datastream.s3.amazonaws.com/index.html). + +### How to run it + +The simplest way, from a terminal: + +```bash +python -m fimbox.nextgen my_area.shp --out-dir out +``` + +Replace `my_area.shp` with the path to your boundary file. Results land in `out/my_area/`. + +Or from Python, in one line: + +```python +import fimbox + +result = fimbox.getNextGenAOI("my_area.shp", out_dir="out") + +print(result.catchments_path) # the catchments map file +print(result.discharge_csvs) # the water-flow table(s) +print(result.feature_ids) # the stream IDs +``` + +By default it grabs the **most recent** available forecast. + +**Output folder structure** + +``` +out/my_area/ + feature_id.csv the list of stream IDs + hydrofabric/ + aoi_catchments.gpkg the catchments (drainage areas) map + aoi_flowpaths.gpkg the rivers/streams inside your area + network_crosswalk.csv a lookup table linking IDs together + discharge-inputs/ + nextgen_..._maximum.csv the water-flow numbers for each stream +``` + +### A few useful options + +You can add these to the command above: + +- `--no-discharge` — just find the catchments and IDs, skip downloading the water-flow data (faster). +- `--date 20260721 --cycle 09` — use a specific day/time instead of the latest. +- `--forecast medium_range` — use a longer-range forecast (default is `short_range`). +- `--sortby maximum` — for each stream, keep the **peak** flow over the forecast (the default, best for flood extent). Other choices: `minimum`, `mean`, or `none` to keep every hour. + + +--- + +### Developer reference + +`fimbox.nextgen` maps an AOI to the NOAA/OWP **NextGen v2.2 hydrofabric** and the ngen/t-route **discharge** published on the public [CIROH community NextGen DataStream](https://ciroh-community-ngen-datastream.s3.amazonaws.com) S3 bucket. + +**How the crosswalk works.** Each NextGen divide (`divide_id` = `cat-`) drains to exactly one flowpath (`id` = `wb-`). The integer `` is the `feature_id` used in the ngen/t-route outputs, so catchment → flowpath → discharge is a direct join. The `network` layer also carries `hf_id` (the NWM COMID), available as an optional crosswalk (`nwm_crosswalk=True`, slower because it scans that table over S3). + +```mermaid +flowchart LR + A[AOI shp/gpkg] --> B[Resolve VPU
cached bbox index] + B --> C[Read intersecting
divides from S3 gpkg] + C --> D[cat-* -> wb-* -> feature_id
+ optional NWM hf_id] + D --> E[Locate ngen run
model/forecast/date/cycle/VPU] + E --> F[Read t-route flow
parquet or tar.gz/netCDF] + F --> G[FIM-ready CSVs
feature_id, discharge_cms] + D --> H[aoi_catchments.gpkg
network_crosswalk.csv] +``` + +Newer runs store discharge as `ngen-run/outputs/troute/troute_output_*.parquet`; older runs as an `ngen-run.tar.gz` containing `outputs/troute/*.nc`. Both are handled transparently. + +**Module contents** + +| File | What it contains | +|---|---| +| `hydrofabric.py` | `NextGenHydrofabric` / `AOIHydrofabric`: AOI → VPU resolution, intersecting `divides`, and the `cat`/`wb`/`feature_id`(+`hf_id`) crosswalk. `build_vpu_index()` regenerates the cached VPU bbox index. | +| `datastream.py` | `NextGenDatastream` / `DischargeRun`: locate an ngen run on the bucket (default `cfe_nom` / `short_range` / latest), read t-route `flow` for the feature ids, and write FIM-ready CSVs. | +| `pipeline.py` | `NextGenAOI` and `getNextGenAOI`: AOI → catchments + discharge in one call, written into the standard fimbox AOI layout. Also the module CLI. | +| `_common.py` | Bucket constants, anonymous S3 handle, AOI-layout helpers, the cached VPU bbox index, and `wb`→`feature_id` parsing. | +| `data/vpu_bbox.json` | Per-VPU bounding boxes (EPSG:5070) so AOI→VPU resolution needs no network round-trip. | + +**Full Python API** + +```python +import fimbox + +res = fimbox.getNextGenAOI( + "my_aoi.shp", # AOI shapefile / GeoPackage / GeoJSON + out_dir="out", # AOI folder created at out/ + # aoi_layer=None, # layer name for a multi-layer gpkg + # predicate="intersects", # or "within" (catchments fully inside the AOI) + # nwm_crosswalk=False, # also resolve NWM hf_id/COMID (slower S3 read) + # model="cfe_nom", # ngen model output set + # forecast="short_range", # short_range | medium_range | analysis_assim_extend + # date=None, cycle=None, # YYYYMMDD + cycle hour; default latest available + # sortby="maximum", # horizon aggregation: maximum|minimum|mean|None + # at_time=None, # single timestamp instead of aggregating + # fetch_discharge=True, # False -> resolve catchments/ids only +) + +res.catchments_path # /hydrofabric/aoi_catchments.gpkg +res.feature_ids # NextGen feature_ids == streamflow network ids +res.network_ids # wb-* flowpath ids +res.discharge # long DataFrame [feature_id, time, flow] (flow in cms) +res.discharge_csvs # FIM-ready CSVs in /discharge-inputs/ +res.run # e.g. "cfe_nom/short_range/ngen.20260721/09/VPU_03N" + +# Step-by-step (class-based) +from fimbox.nextgen import NextGenHydrofabric, NextGenDatastream + +hf = NextGenHydrofabric("my_aoi.shp").resolve() # AOIHydrofabric +ds = NextGenDatastream(hf.vpus[0]) # latest cfe_nom short_range run +flow = ds.read_discharge(hf.feature_ids) # [feature_id, time, flow] +csvs = ds.to_fim_inputs("out/my_aoi", hf.feature_ids, sortby="maximum") +``` + +**Notes** + +- The bucket is read anonymously; no AWS credentials are required. +- `feature_id.csv` and `discharge-inputs/` match the `fimbox.streamflow` layout, so NextGen discharge is a drop-in alternative to the NWM/GEOGLOWS sources for FIM generation. diff --git a/src/fimbox/nextgen/__init__.py b/src/fimbox/nextgen/__init__.py new file mode 100644 index 0000000..a2bd820 --- /dev/null +++ b/src/fimbox/nextgen/__init__.py @@ -0,0 +1,47 @@ +""" +Author: Manjila Singh +Date Updated: July 2026 + +NextGen-in-a-Box integration for fimbox. + +Map an AOI to the NOAA/OWP NextGen v2.2 hydrofabric and the +ngen/t-route discharge published on the public CIROH community NextGen +DataStream bucket (https://ciroh-community-ngen-datastream.s3.amazonaws.com), +so the resulting streamflow can drive FIM generation. + +Resolution +---------- +NextGenHydrofabric AOI (shp/gpkg) -> intersecting catchments + network-id crosswalk +NextGenDatastream NextGen network ids -> ngen/t-route discharge (S3) + +Orchestration +------------- +NextGenAOI / getNextGenAOI AOI -> catchments + discharge in one call, written + into the standard fimbox AOI layout. + +Example +------- + import fimbox + + res = fimbox.getNextGenAOI("my_aoi.shp", out_dir="out") + res.catchments_path # /hydrofabric/aoi_catchments.gpkg + res.feature_ids # NextGen feature_ids (== streamflow network ids) + res.discharge_csvs # FIM-ready /discharge-inputs/*.csv +""" + +from __future__ import annotations + +from .datastream import DischargeRun, NextGenDatastream +from .hydrofabric import AOIHydrofabric, NextGenHydrofabric, build_vpu_index +from .pipeline import NextGenAOI, NextGenResult, getNextGenAOI + +__all__ = [ + "NextGenHydrofabric", + "AOIHydrofabric", + "NextGenDatastream", + "DischargeRun", + "NextGenAOI", + "NextGenResult", + "getNextGenAOI", + "build_vpu_index", +] diff --git a/src/fimbox/nextgen/__main__.py b/src/fimbox/nextgen/__main__.py new file mode 100644 index 0000000..f5b0c18 --- /dev/null +++ b/src/fimbox/nextgen/__main__.py @@ -0,0 +1,17 @@ +""" +Author: Manjila Singh +Date Created: July 2026 + +Package entry point so ``python -m fimbox.nextgen`` runs the pipeline CLI. + +Kept separate from ``pipeline.py`` (and not imported by ``__init__``) so that +``python -m fimbox.nextgen`` does not trip the runpy "found in sys.modules" +double-import warning. +""" + +from __future__ import annotations + +from .pipeline import _main + +if __name__ == "__main__": + _main() diff --git a/src/fimbox/nextgen/_common.py b/src/fimbox/nextgen/_common.py new file mode 100644 index 0000000..8deeb79 --- /dev/null +++ b/src/fimbox/nextgen/_common.py @@ -0,0 +1,113 @@ +""" +Shared helpers for the ``fimbox.nextgen`` subpackage: constants for the public +CIROH community NextGen DataStream bucket, an anonymous s3fs handle, AOI-layout +resolution, and the cached per-VPU bounding-box index used to resolve an area +of interest to its hydrofabric VPU(s) without opening every GeoPackage. + +The bucket is read anonymously (``--no-sign-request`` equivalent) so it works +for any caller regardless of whether AWS credentials are configured. +""" + +from __future__ import annotations + +import importlib +import json +import logging +from functools import lru_cache +from pathlib import Path +from typing import Optional, Union + +from ..logging_utils import aoi_root, attach_case_log + +log = logging.getLogger(__name__) + +PathLike = Union[str, Path] + +BUCKET = "ciroh-community-ngen-datastream" +S3_BUCKET_URL = f"https://{BUCKET}.s3.amazonaws.com" + +HF_VERSION = "v2.2" +HF_GEOPACKAGES_PREFIX = f"{BUCKET}/resources/{HF_VERSION}_hydrofabric/geopackages" + +# outputs//_hydrofabric/ngen.///VPU_/ +OUTPUTS_PREFIX = f"{BUCKET}/outputs" + +HF_EPSG = 5070 # CONUS Albers — all spatial ops happen in this CRS +DISCHARGE_COL = "discharge_cms" # column name the FIM Inundator expects + +HYDROFABRIC_DIR = "hydrofabric" +DISCHARGE_INPUTS_DIR = "discharge-inputs" + + +def require(module: str): + """Import an optional dependency, raising a clear install hint if missing.""" + try: + return importlib.import_module(module) + except ImportError as exc: # pragma: no cover - trivial + raise ImportError( + f"'{module}' is required for fimbox.nextgen. " + f"Install it with: pip install {module}" + ) from exc + + +@lru_cache(maxsize=1) +def s3() -> "object": + """Anonymous s3fs filesystem for the public bucket (cached).""" + s3fs = require("s3fs") + return s3fs.S3FileSystem(anon=True) + + +def resolve_aoi(aoi_dir: PathLike) -> Path: + """Return the AOI root for any directory the caller passes (the root itself + or its ``watershed-data`` subfolder).""" + return aoi_root(Path(aoi_dir)) + + +def hydrofabric_dir(aoi_dir: PathLike) -> Path: + """``/hydrofabric`` — subset catchments / flowpaths / crosswalk land + here. Created on demand.""" + d = resolve_aoi(aoi_dir) / HYDROFABRIC_DIR + d.mkdir(parents=True, exist_ok=True) + return d + + +def discharge_inputs_dir(aoi_dir: PathLike) -> Path: + """``/discharge-inputs`` — the FIM-ready discharge CSVs the generator + iterates (shared with the streamflow subpackage).""" + d = resolve_aoi(aoi_dir) / DISCHARGE_INPUTS_DIR + d.mkdir(parents=True, exist_ok=True) + return d + + +def attach_log(aoi_dir: PathLike) -> None: + """Route nextgen log records into the AOI's combined processing.log.""" + attach_case_log(aoi_dir) + + +_DATA_DIR = Path(__file__).parent / "data" +_VPU_INDEX_PATH = _DATA_DIR / "vpu_bbox.json" + + +@lru_cache(maxsize=1) +def vpu_bbox_index() -> dict[str, list[float]]: + """Load the cached ``{VPU_id: [minx, miny, maxx, maxy]}`` index (EPSG:5070). + + Shipped as package data so AOI->VPU resolution needs no network round-trip. + """ + payload = json.loads(_VPU_INDEX_PATH.read_text()) + return {k: list(map(float, v)) for k, v in payload["bboxes"].items()} + + +def vpu_gpkg_uri(vpu: str) -> str: + """``/vsis3`` URI for a VPU's NextGen GeoPackage (read anonymously).""" + return f"/vsis3/{HF_GEOPACKAGES_PREFIX}/{vpu}/nextgen_{vpu}.gpkg" + + +def feature_id_from_wb(wb_id: str) -> Optional[int]: + """``"wb-2855078"`` -> ``2855078``. Returns None for non-``wb`` ids.""" + if not isinstance(wb_id, str) or "-" not in wb_id: + return None + prefix, _, num = wb_id.partition("-") + if prefix != "wb" or not num.isdigit(): + return None + return int(num) diff --git a/src/fimbox/nextgen/data/vpu_bbox.json b/src/fimbox/nextgen/data/vpu_bbox.json new file mode 100644 index 0000000..8a70251 --- /dev/null +++ b/src/fimbox/nextgen/data/vpu_bbox.json @@ -0,0 +1,28 @@ +{ + "_comment": "NextGen v2.2 hydrofabric per-VPU bounding boxes in EPSG:5070 [minx, miny, maxx, maxy]. Derived from the divides layer total_bounds of each nextgen_VPU_.gpkg in s3://ciroh-community-ngen-datastream/resources/v2.2_hydrofabric/geopackages/. Used to resolve an AOI to its VPU(s) without opening every gpkg. Regenerate with fimbox.nextgen.hydrofabric.build_vpu_index().", + "crs": "EPSG:5070", + "hydrofabric_version": "v2.2", + "bboxes": { + "VPU_01": [1824824.9997, 2215544.9958, 2258234.9955, 3086925.0003], + "VPU_02": [1350645.0003, 1687635.0, 1991054.9979, 2720864.9997], + "VPU_03N": [1115894.9997, 1023914.9979, 1833435.0045, 1740465.0], + "VPU_03S": [1061444.9997, 340785.0, 1598235.0003, 1331925.0003], + "VPU_03W": [578895.0003, 788805.0, 1161915.0003, 1402994.9997], + "VPU_04": [213255.0, 1990604.9997, 1744665.0003, 2824454.9997], + "VPU_05": [588854.9997, 1413285.0003, 1488224.9997, 2292855.0003], + "VPU_06": [645014.9997, 1272255.0003, 1294065.0, 1672544.9997], + "VPU_07": [-106125.0003, 1565505.0, 837735.0003, 2754015.0003], + "VPU_08": [150975.0, 673394.9976, 689775.0003, 1661985.0], + "VPU_09": [-646274.9997, 2498744.9997, 455654.9997, 3083384.9997], + "VPU_10L": [-1074915.0, 1556235.0, 505784.9997, 2300715.0], + "VPU_10U": [-1389615.0003, 2117324.9997, 31545.0, 3039464.9997], + "VPU_11": [-907575.0003, 908984.9997, 518355.0, 1866134.9997], + "VPU_12": [-729254.9997, 312795.0, 271244.9997, 1319204.9997], + "VPU_13": [-1206315.0, 209715.0003, -115845.0003, 1759995.0], + "VPU_14": [-1425105.0, 1456994.9997, -811784.9997, 2356695.0], + "VPU_15": [-1772655.0003, 975944.9997, -1081784.9997, 1971974.9997], + "VPU_16": [-2063565.0, 1528364.9997, -1192725.0, 2368545.0003], + "VPU_17": [-2294145.0, 2191574.9997, -1102155.0003, 3506235.0003], + "VPU_18": [-2355824.9997, 1202325.0003, -1696155.0003, 2538315.0] + } +} diff --git a/src/fimbox/nextgen/datastream.py b/src/fimbox/nextgen/datastream.py new file mode 100644 index 0000000..24b55bc --- /dev/null +++ b/src/fimbox/nextgen/datastream.py @@ -0,0 +1,312 @@ +""" +Retrieve NextGen-in-a-Box discharge from the public CIROH community NextGen +DataStream bucket for a set of hydrofabric ``feature_id``s. + +Output layout on the bucket:: + + outputs//v2.2_hydrofabric/ngen.///VPU_/ + ngen-run/outputs/troute/troute_output_.parquet (newer runs) + ngen-run.tar.gz -> ngen-run/outputs/troute/*.nc (older runs) + +t-route output is the routed streamflow that drives FIM. It is keyed by +``feature_id`` (the integer of the ``wb-`` flowpath id) with an hourly +``flow`` series in cubic metres per second. This module locates the run +(defaulting to model ``cfe_nom``, ``short_range``, and the latest available +date/cycle that has the requested VPU), reads the ``flow`` series for the +requested feature ids, and emits FIM-ready CSVs (``feature_id, discharge_cms``) +into ``/discharge-inputs/``. +""" + +from __future__ import annotations + +import io +import logging +import tarfile +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Union + +import pandas as pd + +from . import _common as C + +log = logging.getLogger(__name__) + +PathLike = Union[str, Path] + +DEFAULT_MODEL = "cfe_nom" +DEFAULT_FORECAST = "short_range" +_HF_SEG = f"{C.HF_VERSION}_hydrofabric" + + +@dataclass +class DischargeRun: + """A located ngen datastream run.""" + + model: str + forecast: str + date: str # YYYYMMDD + cycle: str # e.g. "00" + vpu: str + + @property + def prefix(self) -> str: + return ( + f"{C.OUTPUTS_PREFIX}/{self.model}/{_HF_SEG}/ngen.{self.date}/" + f"{self.forecast}/{self.cycle}/{self.vpu}" + ) + + def __str__(self) -> str: + return f"{self.model}/{self.forecast}/ngen.{self.date}/{self.cycle}/{self.vpu}" + + +class NextGenDatastream: + """Locate and read NextGen datastream discharge for one VPU. + + Parameters + ---------- + vpu : str + Hydrofabric VPU id, e.g. ``"VPU_03N"``. + model : str + ngen model output set. Default ``"cfe_nom"``. + forecast : str + Forecast product: ``"short_range"`` (default), ``"medium_range"``, + ``"analysis_assim_extend"``, ... + date : str, optional + Run date ``YYYYMMDD`` (or ``YYYY-MM-DD``). Default: latest available. + cycle : str, optional + Cycle hour, e.g. ``"06"``. Default: latest cycle that has ``vpu``. + """ + + def __init__( + self, + vpu: str, + *, + model: str = DEFAULT_MODEL, + forecast: str = DEFAULT_FORECAST, + date: Optional[str] = None, + cycle: Optional[str] = None, + ): + self.vpu = vpu + self.model = model + self.forecast = forecast + self._fs = C.s3() + self.run = self._locate(date, cycle) + + def _model_root(self) -> str: + return f"{C.OUTPUTS_PREFIX}/{self.model}/{_HF_SEG}" + + def _ls_names(self, prefix: str) -> list[str]: + try: + return [p.split("/")[-1] for p in self._fs.ls(prefix)] + except FileNotFoundError: + return [] + + def _available_dates(self) -> list[str]: + dates = [ + n.split("ngen.")[-1] + for n in self._ls_names(self._model_root()) + if n.startswith("ngen.") + ] + return sorted(dates, reverse=True) + + def _locate(self, date: Optional[str], cycle: Optional[str]) -> DischargeRun: + date = date.replace("-", "") if date else None + + def cycle_has_vpu(d: str, c: str) -> bool: + base = f"{self._model_root()}/ngen.{d}/{self.forecast}/{c}/{self.vpu}" + return self._fs.exists(base) + + if date and cycle: + if not cycle_has_vpu(date, cycle): + raise FileNotFoundError( + f"No {self.vpu} output for {self.model}/{self.forecast} " + f"ngen.{date} cycle {cycle}." + ) + return DischargeRun(self.model, self.forecast, date, cycle, self.vpu) + + candidate_dates = [date] if date else self._available_dates() + if not candidate_dates: + raise FileNotFoundError( + f"No dated runs under {self._model_root()} (model={self.model})." + ) + for d in candidate_dates: + cycles = self._ls_names(f"{self._model_root()}/ngen.{d}/{self.forecast}") + wanted = [cycle] if cycle else sorted(cycles, reverse=True) + for c in wanted: + if c in cycles and cycle_has_vpu(d, c): + log.info( + "Located discharge run: %s/%s ngen.%s cycle %s %s", + self.model, + self.forecast, + d, + c, + self.vpu, + ) + return DischargeRun(self.model, self.forecast, d, c, self.vpu) + raise FileNotFoundError( + f"No {self.forecast} run with {self.vpu} found for model " + f"{self.model} (searched {len(candidate_dates)} date(s))." + ) + + def _troute_dir(self) -> str: + return f"{self.run.prefix}/ngen-run/outputs/troute" + + def read_discharge(self, feature_ids: Optional[list[int]] = None) -> pd.DataFrame: + """Return a long DataFrame ``[feature_id, time, flow]`` (flow in cms), + filtered to ``feature_ids`` when given. + + Handles both the extracted parquet/netCDF layout (newer runs) and the + ``ngen-run.tar.gz`` archive layout (older runs). + """ + want = set(int(f) for f in feature_ids) if feature_ids is not None else None + troute_dir = self._troute_dir() + + if self._fs.exists(troute_dir): + files = self._fs.ls(troute_dir) + df = self._read_extracted(files) + elif self._fs.exists(f"{self.run.prefix}/ngen-run.tar.gz"): + df = self._read_tarball(f"{self.run.prefix}/ngen-run.tar.gz") + else: + raise FileNotFoundError( + f"No troute output (extracted or tarball) under {self.run.prefix}" + ) + + if want is not None: + df = df[df["feature_id"].isin(want)].copy() + df = df.sort_values(["feature_id", "time"]).reset_index(drop=True) + log.info( + "Discharge: %d reaches x %d timesteps from %s", + df["feature_id"].nunique(), + df["time"].nunique(), + self.run, + ) + return df + + def _read_extracted(self, files: list[str]) -> pd.DataFrame: + parquet = [f for f in files if f.endswith(".parquet")] + netcdf = [f for f in files if f.endswith(".nc")] + if parquet: + frames = [pd.read_parquet(self._fs.open(f)) for f in parquet] + return self._normalize(pd.concat(frames, ignore_index=True)) + if netcdf: + frames = [self._read_netcdf(self._fs.open(f)) for f in netcdf] + return self._normalize(pd.concat(frames, ignore_index=True)) + raise FileNotFoundError("troute dir has no .parquet or .nc output") + + def _read_tarball(self, key: str) -> pd.DataFrame: + log.info("Downloading archive %s ...", key.split("/")[-1]) + raw = self._fs.cat_file(key) + frames = [] + with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf: + members = [ + m + for m in tf.getmembers() + if "outputs/troute" in m.name and m.name.endswith((".nc", ".parquet")) + ] + for m in members: + fh = tf.extractfile(m) + buf = io.BytesIO(fh.read()) + if m.name.endswith(".parquet"): + frames.append(pd.read_parquet(buf)) + else: + frames.append(self._read_netcdf(buf)) + if not frames: + raise FileNotFoundError(f"{key} has no troute output members") + return self._normalize(pd.concat(frames, ignore_index=True)) + + @staticmethod + def _read_netcdf(fileobj) -> pd.DataFrame: + xr = C.require("xarray") + # netCDF4 can't read a file-like object; materialise to a temp file + import tempfile + + data = fileobj.read() if hasattr(fileobj, "read") else fileobj + with tempfile.NamedTemporaryFile(suffix=".nc") as tmp: + tmp.write(data) + tmp.flush() + with xr.open_dataset(tmp.name) as ds: + keep = [v for v in ("flow", "velocity", "depth") if v in ds] + df = ds[keep].to_dataframe().reset_index() + return df + + @staticmethod + def _normalize(df: pd.DataFrame) -> pd.DataFrame: + """Coerce whatever troute layout we read into ``[feature_id, time, flow]``.""" + cols = {c.lower(): c for c in df.columns} + fid = cols.get("feature_id") or cols.get("featureid") + flow = cols.get("flow") or cols.get("streamflow") or cols.get("q") + time = cols.get("time") or cols.get("current_time") + if fid is None or flow is None: + raise ValueError(f"Unrecognized troute columns: {list(df.columns)}") + # keep flowpaths (wb), drop nexus rows when a 'type' column is present + if "type" in cols: + df = df[df[cols["type"]].astype(str).str.lower() == "wb"] + out = df.rename(columns={fid: "feature_id", flow: "flow"}) + if time is not None: + out = out.rename(columns={time: "time"}) + else: + out["time"] = pd.NaT + out["feature_id"] = out["feature_id"].astype("int64") + out["time"] = pd.to_datetime(out["time"]) + return out[["feature_id", "time", "flow"]] + + def to_fim_inputs( + self, + aoi_dir: PathLike, + feature_ids: list[int], + *, + sortby: Optional[str] = "maximum", + at_time: Optional[str] = None, + ) -> list[Path]: + """Write FIM-ready discharge CSV(s) to ``/discharge-inputs/``. + + ``sortby`` in {``"maximum"``, ``"minimum"``, ``"mean"``} collapses the + forecast horizon into one CSV of that statistic per reach (default + ``"maximum"`` — the peak, most relevant for flood extent). ``at_time`` + selects a single timestamp instead. ``sortby=None`` and no ``at_time`` + writes one CSV per timestep. + """ + C.attach_log(aoi_dir) + df = self.read_discharge(feature_ids) + if df.empty: + log.warning("No discharge matched the requested feature_ids.") + return [] + out_dir = C.discharge_inputs_dir(aoi_dir) + tag = ( + f"nextgen_{self.model}_{self.run.forecast}_{self.run.date}{self.run.cycle}" + ) + written: list[Path] = [] + + if at_time is not None: + ts = pd.to_datetime(at_time) + sub = df[df["time"] == ts] + if sub.empty: + raise ValueError( + f"No timestep at {at_time}; available: " + f"{df['time'].min()}..{df['time'].max()}" + ) + written.append( + self._write_csv(sub, out_dir / f"{tag}_{C.DISCHARGE_COL}.csv") + ) + elif sortby: + agg = {"maximum": "max", "minimum": "min", "mean": "mean"} + if sortby not in agg: + raise ValueError("sortby must be maximum, minimum, mean, or None") + g = df.groupby("feature_id")["flow"].agg(agg[sortby]).reset_index() + written.append(self._write_csv(g, out_dir / f"{tag}_{sortby}.csv")) + else: + for ts, sub in df.groupby("time"): + stamp = pd.to_datetime(ts).strftime("%Y%m%dT%H%M") + written.append(self._write_csv(sub, out_dir / f"{tag}_{stamp}.csv")) + return written + + @staticmethod + def _write_csv(df: pd.DataFrame, path: Path) -> Path: + out = df.rename(columns={"flow": C.DISCHARGE_COL})[ + ["feature_id", C.DISCHARGE_COL] + ].copy() + out["feature_id"] = out["feature_id"].astype("int64") + out.to_csv(path, index=False) + log.info("FIM-ready discharge (%d reaches) --> %s", len(out), path.name) + return path diff --git a/src/fimbox/nextgen/hydrofabric.py b/src/fimbox/nextgen/hydrofabric.py new file mode 100644 index 0000000..30cac10 --- /dev/null +++ b/src/fimbox/nextgen/hydrofabric.py @@ -0,0 +1,358 @@ +""" +Resolve an area of interest (AOI) to the NextGen v2.2 hydrofabric. + +Given an AOI (shapefile / GeoPackage / any vector fiona can read), this module +finds which hydrofabric VPU(s) the AOI falls in, reads the catchment +(``divides``) polygons that intersect the AOI directly from the public +per-VPU GeoPackage on the CIROH community NextGen DataStream bucket, and +crosswalks them to the NextGen network ids used by the ngen/t-route outputs. + +Each NextGen divide (``divide_id`` = ``cat-``) drains to exactly one +flowpath (``id`` = ``wb-``); the integer ```` is the ``feature_id`` used +in the t-route discharge outputs. The ``network`` layer additionally carries +``hf_id`` — the NOAA reference-fabric / NWM COMID — kept as a crosswalk to the +NWM feature ids the rest of fimbox uses. + +Outputs (written under ``/hydrofabric/`` when ``save=True``): + aoi_catchments.gpkg -- intersecting divides (catchment polygons) + aoi_flowpaths.gpkg -- the flowpaths draining those catchments + network_crosswalk.csv -- divide_id, wb_id, feature_id, hf_id (NWM), vpu + feature_id.csv -- unique feature_ids (drop-in for the streamflow step) +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Union + +import geopandas as gpd +import pandas as pd +from shapely.geometry import box + +from . import _common as C + +log = logging.getLogger(__name__) + +PathLike = Union[str, Path] + +# GDAL/pyogrio env needed to read the public gpkg anonymously over /vsis3. +_VSIS3_ENV = { + "AWS_NO_SIGN_REQUEST": "YES", + "AWS_S3_ENDPOINT": "s3.amazonaws.com", + "CPL_VSIL_CURL_ALLOWED_EXTENSIONS": ".gpkg", + "GDAL_DISABLE_READDIR_ON_OPEN": "EMPTY_DIR", + "VSI_CACHE": "TRUE", + "GDAL_HTTP_MAX_RETRY": "3", + "GDAL_HTTP_RETRY_DELAY": "1", +} + + +@contextmanager +def _vsis3_env(): + """Temporarily set the GDAL env vars required for anonymous /vsis3 reads.""" + old = {k: os.environ.get(k) for k in _VSIS3_ENV} + os.environ.update(_VSIS3_ENV) + try: + yield + finally: + for k, v in old.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +@dataclass +class AOIHydrofabric: + """Result of an AOI -> hydrofabric resolution. + + Attributes + ---------- + catchments : GeoDataFrame + Intersecting ``divides`` (catchment polygons), EPSG:5070. + flowpaths : GeoDataFrame + The flowpaths (``wb-*``) draining those catchments, EPSG:5070. + crosswalk : DataFrame + One row per catchment: divide_id, wb_id, feature_id, hf_id, vpu. + vpus : list[str] + VPU(s) the AOI intersects (usually one). + """ + + catchments: gpd.GeoDataFrame + flowpaths: gpd.GeoDataFrame + crosswalk: pd.DataFrame + vpus: list[str] = field(default_factory=list) + + @property + def feature_ids(self) -> list[int]: + """Unique t-route ``feature_id``s (integers of the ``wb-*`` ids).""" + return sorted( + {int(f) for f in self.crosswalk["feature_id"].dropna().astype("int64")} + ) + + @property + def catchment_ids(self) -> list[str]: + """``cat-*`` divide ids.""" + return self.crosswalk["divide_id"].dropna().astype(str).tolist() + + @property + def network_ids(self) -> list[str]: + """``wb-*`` NextGen flowpath (network) ids.""" + return self.crosswalk["wb_id"].dropna().astype(str).tolist() + + @property + def nwm_feature_ids(self) -> list[int]: + """NWM/NHD reference COMIDs (``hf_id``) crosswalked to the catchments.""" + return sorted( + {int(f) for f in self.crosswalk["hf_id"].dropna().astype("int64")} + ) + + +class NextGenHydrofabric: + """AOI -> NextGen v2.2 catchments + network-id crosswalk. + + Parameters + ---------- + aoi : str or Path or GeoDataFrame + Area of interest: a path to a shapefile / GeoPackage / GeoJSON, or an + already-loaded GeoDataFrame. + aoi_layer : str, optional + Layer name when ``aoi`` is a multi-layer GeoPackage. + predicate : str, optional + Spatial test used to select catchments. ``"intersects"`` (default) + keeps every catchment that touches the AOI; ``"within"`` keeps only + catchments fully inside the AOI. + """ + + def __init__( + self, + aoi: Union[PathLike, gpd.GeoDataFrame], + *, + aoi_layer: Optional[str] = None, + predicate: str = "intersects", + nwm_crosswalk: bool = False, + ): + if predicate not in ("intersects", "within"): + raise ValueError("predicate must be 'intersects' or 'within'") + self.predicate = predicate + # opt-in: reading the non-spatial `network` table scans the whole layer + self.nwm_crosswalk = nwm_crosswalk + self.aoi = self._load_aoi(aoi, aoi_layer).to_crs(epsg=C.HF_EPSG) + self._aoi_geom = self.aoi.geometry.union_all() + + @staticmethod + def _load_aoi( + aoi: Union[PathLike, gpd.GeoDataFrame], layer: Optional[str] + ) -> gpd.GeoDataFrame: + if isinstance(aoi, gpd.GeoDataFrame): + gdf = aoi.copy() + else: + gdf = gpd.read_file(aoi, layer=layer) if layer else gpd.read_file(aoi) + if gdf.empty: + raise ValueError("AOI is empty.") + if gdf.crs is None: + raise ValueError("AOI has no CRS; set one before resolving.") + return gdf + + def candidate_vpus(self) -> list[str]: + """VPU(s) whose bounding box overlaps the AOI (bbox pre-filter).""" + aoi_box = box(*self.aoi.total_bounds) + hits = [ + vpu + for vpu, bbox in C.vpu_bbox_index().items() + if box(*bbox).intersects(aoi_box) + ] + if not hits: + raise ValueError( + "AOI does not overlap any CONUS hydrofabric VPU bounding box. " + "Check the AOI CRS / extent (must be within CONUS)." + ) + return hits + + def _read_divides(self, vpu: str) -> gpd.GeoDataFrame: + """Read ``divides`` in the AOI bbox from a VPU gpkg over /vsis3, then + apply the precise spatial predicate against the AOI geometry.""" + uri = C.vpu_gpkg_uri(vpu) + bounds = tuple(self.aoi.total_bounds) + with _vsis3_env(): + divides = gpd.read_file(uri, layer="divides", bbox=bounds) + if divides.empty: + return divides + if divides.crs is None: + divides.set_crs(epsg=C.HF_EPSG, inplace=True) + else: + divides = divides.to_crs(epsg=C.HF_EPSG) + if self.predicate == "within": + mask = divides.within(self._aoi_geom) + else: + mask = divides.intersects(self._aoi_geom) + hit = divides[mask].copy() + hit["vpuid"] = hit.get("vpuid", vpu) + hit["vpu"] = vpu + return hit + + def _read_flowpaths(self, vpu: str, wb_ids: list[str]) -> gpd.GeoDataFrame: + """Read the flowpaths draining the selected catchments from a VPU gpkg. + + Uses the AOI bbox spatial filter (fast, index-backed) then keeps only + the ``wb`` ids we selected — an attribute ``WHERE id IN (...)`` over + /vsis3 would scan the whole layer instead.""" + if not wb_ids: + return gpd.GeoDataFrame(geometry=[], crs=f"EPSG:{C.HF_EPSG}") + uri = C.vpu_gpkg_uri(vpu) + bounds = tuple(self.aoi.total_bounds) + with _vsis3_env(): + fp = gpd.read_file(uri, layer="flowpaths", bbox=bounds) + if fp.empty: + return fp + if fp.crs is not None: + fp = fp.to_crs(epsg=C.HF_EPSG) + fp = fp[fp["id"].astype(str).isin(set(wb_ids))].copy() + return fp + + def _read_network_crosswalk(self, vpu: str, wb_ids: list[str]) -> pd.DataFrame: + """Read the NWM ``hf_id`` crosswalk for the selected ``wb`` ids from the + ``network`` layer (non-spatial). One hf_id per flowpath is kept.""" + if not wb_ids: + return pd.DataFrame(columns=["wb_id", "hf_id"]) + uri = C.vpu_gpkg_uri(vpu) + where = "id IN ({})".format(",".join(f"'{i}'" for i in wb_ids)) + try: + import pyogrio + + with _vsis3_env(): + net = pyogrio.read_dataframe( + uri, + layer="network", + where=where, + read_geometry=False, + columns=["id", "hf_id"], + ) + except Exception as exc: + log.warning("network crosswalk read failed for %s: %s", vpu, exc) + return pd.DataFrame(columns=["wb_id", "hf_id"]) + net = net.rename(columns={"id": "wb_id"}) + net = net.dropna(subset=["wb_id"]).drop_duplicates(subset=["wb_id"]) + return net[["wb_id", "hf_id"]] + + def resolve(self) -> AOIHydrofabric: + """Run the full AOI -> catchments + crosswalk resolution.""" + cats, fps, xwalks, vpus = [], [], [], [] + for vpu in self.candidate_vpus(): + hit = self._read_divides(vpu) + if hit.empty: + log.info("VPU %s: no catchments intersect the AOI", vpu) + continue + vpus.append(vpu) + wb_ids = hit["id"].dropna().astype(str).tolist() + log.info("VPU %s: %d catchments intersect the AOI", vpu, len(hit)) + + xwalk = pd.DataFrame( + { + "divide_id": hit["divide_id"].astype(str).values, + "wb_id": hit["id"].astype(str).values, + "vpu": vpu, + } + ) + xwalk["feature_id"] = xwalk["wb_id"].map(C.feature_id_from_wb) + if self.nwm_crosswalk: + nwm = self._read_network_crosswalk(vpu, wb_ids) + xwalk = xwalk.merge(nwm, on="wb_id", how="left") + else: + xwalk["hf_id"] = pd.NA + + cats.append(hit) + fps.append(self._read_flowpaths(vpu, wb_ids)) + xwalks.append(xwalk) + + if not cats: + raise ValueError( + "No hydrofabric catchments intersect the AOI in any candidate VPU." + ) + + catchments = gpd.GeoDataFrame( + pd.concat(cats, ignore_index=True), crs=f"EPSG:{C.HF_EPSG}" + ) + flowpaths = ( + gpd.GeoDataFrame( + pd.concat([f for f in fps if not f.empty], ignore_index=True), + crs=f"EPSG:{C.HF_EPSG}", + ) + if any(not f.empty for f in fps) + else gpd.GeoDataFrame(geometry=[], crs=f"EPSG:{C.HF_EPSG}") + ) + crosswalk = pd.concat(xwalks, ignore_index=True).drop_duplicates( + subset=["divide_id"] + ) + log.info( + "AOI hydrofabric: %d catchments across VPU(s) %s (%d unique feature_ids)", + len(catchments), + ", ".join(vpus), + crosswalk["feature_id"].nunique(), + ) + return AOIHydrofabric(catchments, flowpaths, crosswalk, vpus) + + def resolve_and_save(self, aoi_dir: PathLike) -> AOIHydrofabric: + """Resolve and write the catchment/flowpath/crosswalk outputs under + ``/hydrofabric/`` plus a ``feature_id.csv`` at the AOI root.""" + C.attach_log(aoi_dir) + result = self.resolve() + out = C.hydrofabric_dir(aoi_dir) + + cat_path = out / "aoi_catchments.gpkg" + result.catchments.to_file(cat_path, driver="GPKG") + log.info("catchments (%d) --> %s", len(result.catchments), cat_path.name) + + if not result.flowpaths.empty: + fp_path = out / "aoi_flowpaths.gpkg" + result.flowpaths.to_file(fp_path, driver="GPKG") + log.info("flowpaths (%d) --> %s", len(result.flowpaths), fp_path.name) + + xwalk_path = out / "network_crosswalk.csv" + result.crosswalk.to_csv(xwalk_path, index=False) + log.info("network crosswalk --> %s", xwalk_path.name) + + fid_path = C.resolve_aoi(aoi_dir) / "feature_id.csv" + pd.DataFrame({"feature_id": result.feature_ids}).to_csv(fid_path, index=False) + log.info("feature ids (%d) --> %s", len(result.feature_ids), fid_path.name) + return result + + +def build_vpu_index(out_path: Optional[PathLike] = None) -> dict[str, list[float]]: + """Regenerate the cached per-VPU bounding-box index from the public bucket. + + Reads each ``nextgen_VPU_.gpkg`` ``divides`` layer's ``total_bounds`` + (metadata only) over /vsis3 and writes them to ``data/vpu_bbox.json``. Run + this only if the hydrofabric VPU set or extents change upstream. + """ + import json + + import pyogrio + + fs = C.s3() + vpus = sorted( + p.split("/")[-1] + for p in fs.ls(C.HF_GEOPACKAGES_PREFIX) + if "VPU_" in p.split("/")[-1] + ) + bboxes: dict[str, list[float]] = {} + with _vsis3_env(): + for vpu in vpus: + info = pyogrio.read_info(C.vpu_gpkg_uri(vpu), layer="divides") + bboxes[vpu] = [float(x) for x in info["total_bounds"]] + log.info("VPU %s bounds %s", vpu, bboxes[vpu]) + payload = { + "_comment": "EPSG:5070 [minx, miny, maxx, maxy] per VPU; regenerate via " + "fimbox.nextgen.hydrofabric.build_vpu_index().", + "crs": "EPSG:5070", + "hydrofabric_version": C.HF_VERSION, + "bboxes": bboxes, + } + dest = Path(out_path) if out_path else C._VPU_INDEX_PATH + dest.write_text(json.dumps(payload, indent=2)) + C.vpu_bbox_index.cache_clear() + return bboxes diff --git a/src/fimbox/nextgen/pipeline.py b/src/fimbox/nextgen/pipeline.py new file mode 100644 index 0000000..45c78a4 --- /dev/null +++ b/src/fimbox/nextgen/pipeline.py @@ -0,0 +1,321 @@ +""" +NextGen-in-a-Box -> fimbox bridge. + +One entry point takes an area of interest (shapefile / GeoPackage) and returns +everything downstream FIM generation needs from the NextGen ecosystem: + + 1. the hydrofabric catchments (``divides``) that intersect the AOI, + 2. the NextGen network ids (``wb-*`` flowpaths / integer ``feature_id``s) and + their NWM ``hf_id`` crosswalk, and + 3. the ngen/t-route discharge for those reaches, pulled from the public + CIROH community NextGen DataStream S3 bucket and written as FIM-ready + ``feature_id, discharge_cms`` CSVs. + +Everything lands in the standard fimbox AOI layout:: + + / + feature_id.csv -- unique NextGen feature_ids + hydrofabric/ + aoi_catchments.gpkg -- intersecting catchment polygons + aoi_flowpaths.gpkg -- their flowpaths + network_crosswalk.csv -- divide_id, wb_id, feature_id, hf_id, vpu + discharge-inputs/ + nextgen____.csv +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Union + +import geopandas as gpd +import pandas as pd + +from . import _common as C +from .datastream import DEFAULT_FORECAST, DEFAULT_MODEL, NextGenDatastream +from .hydrofabric import AOIHydrofabric, NextGenHydrofabric + +log = logging.getLogger(__name__) + +PathLike = Union[str, Path] + + +@dataclass +class NextGenResult: + """Everything the AOI resolution produced. + + Attributes + ---------- + aoi_dir : Path + AOI root the outputs were written under. + hydrofabric : AOIHydrofabric + Catchments, flowpaths, and the network crosswalk. + catchments_path : Path + Saved catchment GeoPackage (the hydrofabric "shapefile"). + crosswalk_path : Path + Saved ``network_crosswalk.csv``. + discharge : DataFrame + Long ``[feature_id, time, flow]`` discharge series (cms). + discharge_csvs : list[Path] + FIM-ready ``feature_id, discharge_cms`` CSVs. + run : str + The datastream run used (``model/forecast/ngen.date/cycle/VPU``). + """ + + aoi_dir: Path + hydrofabric: AOIHydrofabric + catchments_path: Path + crosswalk_path: Path + discharge: pd.DataFrame + discharge_csvs: list[Path] + run: str + + @property + def feature_ids(self) -> list[int]: + return self.hydrofabric.feature_ids + + @property + def network_ids(self) -> list[str]: + return self.hydrofabric.network_ids + + @property + def catchments(self) -> gpd.GeoDataFrame: + return self.hydrofabric.catchments + + +def getNextGenAOI( + aoi: Union[PathLike, gpd.GeoDataFrame], + out_dir: Optional[PathLike] = None, + *, + aoi_layer: Optional[str] = None, + predicate: str = "intersects", + nwm_crosswalk: bool = False, + model: str = DEFAULT_MODEL, + forecast: str = DEFAULT_FORECAST, + date: Optional[str] = None, + cycle: Optional[str] = None, + sortby: Optional[str] = "maximum", + at_time: Optional[str] = None, + fetch_discharge: bool = True, +) -> NextGenResult: + """Resolve an AOI to NextGen catchments + discharge in one call. + + Parameters + ---------- + aoi : path or GeoDataFrame + Area of interest (shapefile / GeoPackage / GeoJSON, or a GeoDataFrame). + out_dir : path, optional + Root output directory; an AOI folder ``/`` is created + underneath. Defaults to ``./out``. + aoi_layer : str, optional + Layer name when the AOI is a multi-layer GeoPackage. + predicate : str + ``"intersects"`` (default) or ``"within"`` for catchment selection. + nwm_crosswalk : bool + Also resolve each reach's NWM/NHD ``hf_id`` (COMID) from the hydrofabric + ``network`` table. Off by default because that read scans the whole + table over S3. + model, forecast, date, cycle + Discharge selection on the datastream bucket. Defaults: ``cfe_nom`` / + ``short_range`` / latest date / latest cycle that has the AOI's VPU. + sortby : str or None + Horizon aggregation for the FIM-ready CSV: ``"maximum"`` (default), + ``"minimum"``, ``"mean"``, or ``None`` for one CSV per timestep. + at_time : str, optional + Single timestamp to slice instead of aggregating. + fetch_discharge : bool + Set False to resolve catchments/ids only (no S3 discharge download). + + Returns + ------- + NextGenResult + """ + return NextGenAOI( + aoi, + out_dir, + aoi_layer=aoi_layer, + predicate=predicate, + nwm_crosswalk=nwm_crosswalk, + ).run( + model=model, + forecast=forecast, + date=date, + cycle=cycle, + sortby=sortby, + at_time=at_time, + fetch_discharge=fetch_discharge, + ) + + +class NextGenAOI: + """AOI -> NextGen hydrofabric catchments + datastream discharge.""" + + def __init__( + self, + aoi: Union[PathLike, gpd.GeoDataFrame], + out_dir: Optional[PathLike] = None, + *, + aoi_layer: Optional[str] = None, + predicate: str = "intersects", + nwm_crosswalk: bool = False, + ): + self.aoi = aoi + self.aoi_layer = aoi_layer + self.predicate = predicate + self.nwm_crosswalk = nwm_crosswalk + + stem = Path(aoi).stem if not isinstance(aoi, gpd.GeoDataFrame) else "aoi" + root = Path(out_dir) if out_dir else (Path.cwd() / "out") + self.aoi_dir = (root / stem).resolve() + self.aoi_dir.mkdir(parents=True, exist_ok=True) + C.attach_log(self.aoi_dir) + + def run( + self, + *, + model: str = DEFAULT_MODEL, + forecast: str = DEFAULT_FORECAST, + date: Optional[str] = None, + cycle: Optional[str] = None, + sortby: Optional[str] = "maximum", + at_time: Optional[str] = None, + fetch_discharge: bool = True, + ) -> NextGenResult: + log.info("=== NextGen AOI: %s ===", self.aoi_dir.name) + + log.info("--- Hydrofabric (AOI -> catchments) ---") + hf = NextGenHydrofabric( + self.aoi, + aoi_layer=self.aoi_layer, + predicate=self.predicate, + nwm_crosswalk=self.nwm_crosswalk, + ) + result = hf.resolve_and_save(self.aoi_dir) + cat_path = C.hydrofabric_dir(self.aoi_dir) / "aoi_catchments.gpkg" + xwalk_path = C.hydrofabric_dir(self.aoi_dir) / "network_crosswalk.csv" + + discharge = pd.DataFrame(columns=["feature_id", "time", "flow"]) + csvs: list[Path] = [] + run_label = "(discharge not fetched)" + + if fetch_discharge: + if len(result.vpus) > 1: + log.warning( + "AOI spans %d VPUs (%s); fetching discharge from each.", + len(result.vpus), + ", ".join(result.vpus), + ) + parts, labels = [], [] + for vpu in result.vpus: + log.info("--- Discharge (%s) ---", vpu) + fids = [ + int(f) + for f in result.crosswalk.loc[ + result.crosswalk["vpu"] == vpu, "feature_id" + ].dropna() + ] + ds = NextGenDatastream( + vpu, model=model, forecast=forecast, date=date, cycle=cycle + ) + parts.append(ds.read_discharge(fids)) + csvs += ds.to_fim_inputs( + self.aoi_dir, fids, sortby=sortby, at_time=at_time + ) + labels.append(str(ds.run)) + if parts: + discharge = pd.concat(parts, ignore_index=True) + run_label = "; ".join(labels) + + log.info("=== DONE: %s ===", self.aoi_dir.name) + return NextGenResult( + aoi_dir=self.aoi_dir, + hydrofabric=result, + catchments_path=cat_path, + crosswalk_path=xwalk_path, + discharge=discharge, + discharge_csvs=csvs, + run=run_label, + ) + + +def _main() -> None: + import argparse + + from ..logging_utils import configure_cli_logging + + configure_cli_logging() + + p = argparse.ArgumentParser( + description="Resolve an AOI to NextGen hydrofabric catchments and pull " + "ngen/t-route discharge from the CIROH community datastream bucket." + ) + p.add_argument("aoi", help="AOI shapefile / GeoPackage / GeoJSON path") + p.add_argument("--aoi-layer", default=None, help="AOI layer (multi-layer gpkg)") + p.add_argument("--out-dir", default="out", help="Root output dir (default ./out)") + p.add_argument( + "--predicate", + default="intersects", + choices=["intersects", "within"], + help="Catchment selection test (default intersects)", + ) + p.add_argument( + "--nwm-crosswalk", + action="store_true", + help="Also resolve NWM/NHD hf_id (COMID) crosswalk (slower S3 read)", + ) + p.add_argument( + "--model", default=DEFAULT_MODEL, help="ngen model (default cfe_nom)" + ) + p.add_argument( + "--forecast", + default=DEFAULT_FORECAST, + help="Forecast product (short_range, medium_range, analysis_assim_extend)", + ) + p.add_argument("--date", default=None, help="Run date YYYYMMDD (default latest)") + p.add_argument("--cycle", default=None, help="Cycle hour, e.g. 06 (default latest)") + p.add_argument( + "--sortby", + default="maximum", + help="Horizon aggregation: maximum|minimum|mean|none (default maximum)", + ) + p.add_argument("--at-time", default=None, help="Single timestamp to slice") + p.add_argument( + "--no-discharge", + action="store_true", + help="Resolve catchments/ids only; skip discharge download", + ) + args = p.parse_args() + + sortby = None if str(args.sortby).lower() == "none" else args.sortby + res = getNextGenAOI( + args.aoi, + args.out_dir, + aoi_layer=args.aoi_layer, + predicate=args.predicate, + nwm_crosswalk=args.nwm_crosswalk, + model=args.model, + forecast=args.forecast, + date=args.date, + cycle=args.cycle, + sortby=sortby, + at_time=args.at_time, + fetch_discharge=not args.no_discharge, + ) + log.info( + "AOI %s: %d catchments, %d feature_ids, VPU(s) %s", + res.aoi_dir.name, + len(res.catchments), + len(res.feature_ids), + ", ".join(res.hydrofabric.vpus), + ) + log.info("Discharge run(s): %s", res.run) + log.info("Catchments: %s", res.catchments_path) + log.info("Crosswalk: %s", res.crosswalk_path) + for c in res.discharge_csvs: + log.info("Discharge: %s", c) + + +if __name__ == "__main__": + _main()