diff --git a/.github/ISSUE_TEMPLATE/capability.md b/.github/ISSUE_TEMPLATE/capability.md new file mode 100644 index 0000000..41378f6 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/capability.md @@ -0,0 +1,38 @@ +--- +name: JIRA Story/Capability +about: The structure for outlining work being done on a JIRA story +labels: JIRA Story +--- + +# Capability + + + +## Task + + + +## Plan/Outline + + + + +### TODOS + + +- [ ] + +### Additional components / Context + + diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 69d40bb..309ec95 100755 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -18,7 +18,6 @@ ## Screenshots - ## Notes - diff --git a/.github/workflows/cicd.yaml b/.github/workflows/cicd.yaml new file mode 100644 index 0000000..0250fc4 --- /dev/null +++ b/.github/workflows/cicd.yaml @@ -0,0 +1,56 @@ +name: CI/CD + +on: + push: + branches: + - main + pull_request: + +jobs: + ruff: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: Set up Python + uses: astral-sh/setup-uv@v5 + with: + python-version: "3.13" + + - name: Install dependencies + run: uv sync + + - name: Run ruff check + run: | + uv run ruff check --config pyproject.toml --output-format=github + + - name: Run ruff format check + run: | + uv run ruff format --config pyproject.toml + + pytests: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.12', '3.13'] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: uv sync --all-extras + + - name: Run tests + run: uv run pytest tests diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..a970a7f --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,30 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + exclude: LICENSE|\.csv$ + - id: end-of-file-fixer + exclude: LICENSE|\.csv$ + - id: check-yaml + - id: debug-statements + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.0 + hooks: + - id: ruff-check + types_or: [ python, pyi ] + args: [ --fix ] + - id: ruff-format + types_or: [ python, pyi ] + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.17.0 + hooks: + - id: mypy + additional_dependencies: ["types-PyYAML==6.0.12.20250516"] + + - repo: https://github.com/kynan/nbstripout + rev: 0.8.1 + hooks: + - id: nbstripout diff --git a/LICENSE b/LICENSE index e8ab96d..6a4ae71 100644 --- a/LICENSE +++ b/LICENSE @@ -201,11 +201,11 @@ See the License for the specific language governing permissions and limitations under the License. -- - - - - - - - - - - - - - +- - - - - - - - - - - - - - - -“Software code created by U.S. Government employees is not subject to copyright in the +"Software code created by U.S. Government employees is not subject to copyright in the United States (17 U.S.C. §105). The United States/Department of Commerce reserves all rights to seek and obtain copyright protection in countries other than the United States for Software authored in its entirety by the Department of Commerce. To this end, the Department of Commerce hereby grants to Recipient a royalty-free, nonexclusive license -to use, copy, and create derivative works of the Software outside of the United States.” +to use, copy, and create derivative works of the Software outside of the United States." diff --git a/NOTICE.txt b/NOTICE.txt new file mode 100644 index 0000000..546a10d --- /dev/null +++ b/NOTICE.txt @@ -0,0 +1,7 @@ +-- + +Inspiration and assistance in the creation, and storage of network graphs comes from the following repos. Credit to the authors: +- https://github.com/DeepGroundwater/ddr/blob/ab4c3962c2c119e6a9182a77f2a9faceec19f2e0/engine/adjacency.py +- https://github.com/CIROH-UA/NGIAB_data_preprocess/blob/main/modules/data_processing/graph_utils.py + +-- diff --git a/README.md b/README.md index 2159d62..9c7fa93 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,9 @@ 3. Update the README, replacing the contents below as prescribed. 4. Add any libraries, assets, or hard dependencies whose source code will be included in the project's repository to the _Exceptions_ section in the [TERMS](TERMS.md). - - If no exceptions are needed, remove that section from TERMS. + +- If no exceptions are needed, remove that section from TERMS. + 5. If working with an existing code base, answer the questions on the [open source checklist](opensource-checklist.md) 6. Delete these instructions and everything up to the _Project Title_ from the README. 7. Write some great software and tell people about it. @@ -20,28 +22,26 @@ To install all of the template files, run the following script from the root of bash -c "$(curl -s https://raw.githubusercontent.com/NOAA-OWP/owp-open-source-project-template/open_source_template.sh)" ``` ----- +--- # Project Title -**Description**: Put a meaningful, short, plain-language description of what +**Description**: Put a meaningful, short, plain-language description of what this project is trying to accomplish and why it matters. Describe the problem(s) this project solves. Describe how this software can improve the lives of its audience. Other things to include: - - **Technology stack**: Indicate the technological nature of the software, including primary programming language(s) and whether the software is intended as standalone or as a module in a framework or other ecosystem. - - **Status**: Alpha, Beta, 1.1, etc. It's OK to write a sentence, too. The goal is to let interested people know where this project is at. This is also a good place to link to the [CHANGELOG](CHANGELOG.md). - - **Links to production or demo instances** - - Describe what sets this apart from related-projects. Linking to another doc or page is OK if this can't be expressed in a sentence or two. - +- **Technology stack**: Indicate the technological nature of the software, including primary programming language(s) and whether the software is intended as standalone or as a module in a framework or other ecosystem. +- **Status**: Alpha, Beta, 1.1, etc. It's OK to write a sentence, too. The goal is to let interested people know where this project is at. This is also a good place to link to the [CHANGELOG](CHANGELOG.md). +- **Links to production or demo instances** +- Describe what sets this apart from related-projects. Linking to another doc or page is OK if this can't be expressed in a sentence or two. **Screenshot**: If the software has visual components, place a screenshot after the description; e.g., ![](https://raw.githubusercontent.com/NOAA-OWP/owp-open-source-project-template/master/doc/Screenshot.png) - ## Dependencies Describe any dependencies that must be installed for this software to work. @@ -88,8 +88,7 @@ important pieces, etc. General instructions on _how_ to contribute should be stated with a link to [CONTRIBUTING](CONTRIBUTING.md). - ----- +--- ## Open source licensing info @@ -99,8 +98,7 @@ as is, but remove this sentence): 1. [TERMS](TERMS.md) 2. [LICENSE](LICENSE) - ----- +--- ## Credits and references diff --git a/builds/build_reference.py b/builds/build_reference.py new file mode 100644 index 0000000..d019baf --- /dev/null +++ b/builds/build_reference.py @@ -0,0 +1,86 @@ +"""An end-to-end build file that will take the NHD ReferenceConfig and turn it into a reference fabric""" + +import argparse +import logging + +from pydantic import ValidationError + +from reference_builds.configs import BaseDataset, ReferenceConfig +from reference_builds.local_runner import LocalRunner +from reference_builds.pipeline import ( + build_geoglows_graphs, + build_geoglows_reference, + build_nhd_graphs, + build_nhd_reference, + build_usgs_hf_graphs, + build_usgs_hf_reference, + download_geoglows_data, + download_nhd_data, + download_usgs_hf_data, + write_reference, +) + +logger = logging.getLogger(__name__) + + +def main() -> int: + """Main entry point for the hydrofabric-build pipeline CLI. + + Returns + ------- + int + Exit code: 0 for success, 1 for failure. + """ + parser = argparse.ArgumentParser(description="A local runner for hydrofabric data processing") + parser.add_argument("--config", required=False, help="Config file") + args = parser.parse_args() + + try: + config = ReferenceConfig.from_yaml(args.config) + except ValidationError as e: + print("Configuration validation failed:") + for error in e.errors(): + print(f" {error['loc']}: {error['msg']}") + return 1 + except FileNotFoundError: + logger.error(f"Config file not found: {args.config}") + return 1 + except TypeError as e: + logger.error("Config file not specified.") + raise TypeError("Config file not specified.") from e + + with LocalRunner(config) as runner: + if config.base_dataset == BaseDataset.NHD: + runner.run_task(task_id="download", python_callable=download_nhd_data, op_kwargs={}) + runner.run_task(task_id="build_nhd_graphs", python_callable=build_nhd_graphs, op_kwargs={}) + runner.run_task(task_id="build_reference", python_callable=build_nhd_reference, op_kwargs={}) + runner.run_task(task_id="write_reference", python_callable=write_reference, op_kwargs={}) + elif config.base_dataset == BaseDataset.GEOGLOWS: + runner.run_task(task_id="download", python_callable=download_geoglows_data, op_kwargs={}) + runner.run_task( + task_id="build_geoglows_graphs", python_callable=build_geoglows_graphs, op_kwargs={} + ) + runner.run_task(task_id="build_reference", python_callable=build_geoglows_reference, op_kwargs={}) + runner.run_task(task_id="write_reference", python_callable=write_reference, op_kwargs={}) + elif config.base_dataset == BaseDataset.USGS_HF: + runner.run_task(task_id="download", python_callable=download_usgs_hf_data, op_kwargs={}) + runner.run_task( + task_id="build_usgs_hf_graphs", python_callable=build_usgs_hf_graphs, op_kwargs={} + ) + runner.run_task(task_id="build_reference", python_callable=build_usgs_hf_reference, op_kwargs={}) + runner.run_task(task_id="write_reference", python_callable=write_reference, op_kwargs={}) + else: + raise NotImplementedError("Base Dataset not implemented") + + print("Pipeline completed") + print("=" * 60) + for task_id, info in runner.results.items(): + status = "✓" if info["status"] == "success" else "✗" + print(f" {status} {task_id}: {info['status']}") + print("=" * 60) + + return 0 + + +if __name__ == "__main__": + exit(main()) diff --git a/config/example_ak.yaml b/config/example_ak.yaml new file mode 100644 index 0000000..f8754f1 --- /dev/null +++ b/config/example_ak.yaml @@ -0,0 +1,5 @@ +domain: ak +base_dataset: usgs-reference-hf +input_file_regex: AK/reference_19.gpkg +vpu_id: "19" +crs: EPSG:3338 \ No newline at end of file diff --git a/config/example_hi.yaml b/config/example_hi.yaml new file mode 100644 index 0000000..047662e --- /dev/null +++ b/config/example_hi.yaml @@ -0,0 +1,18 @@ +domain: hi +input_file_regex: HI/NHDPLUS_H_20*_HU4_GPKG +vpu_id: "20" +crs: EPSG:32604 +permitted_fcodes: + - 'Canal/Ditch' + - 'Stream/River: Hydrographic Category = Perennial' + - 'Stream/River: Hydrographic Category = Intermittent' + - 'Artificial Path' + - 'Pipeline: Pipeline Type = Siphon' + - 'Pipeline: Pipeline Type = Aqueduct' + - 'Pipeline' + - 'Connector' + - 'Canal Ditch: Canal Ditch Type = Stormwater' + # These FCODES are currently disabled in HI due to incorrect NULL `toid` and `dn_hydroseq` values. + # - 'Pipeline: Pipeline Type = Aqueduct; Relationship to Surface = At or Near' + # - 'Pipeline: Pipeline Type = Aqueduct; Relationship to Surface = Underground' + # - 'Pipeline: Pipeline Type = Penstock; Relationship to Surface = At or Near' diff --git a/config/example_prvi.yaml b/config/example_prvi.yaml new file mode 100644 index 0000000..15c9b88 --- /dev/null +++ b/config/example_prvi.yaml @@ -0,0 +1,17 @@ +domain: prvi +input_file_regex: PRVI/NHDPLUS_H_21*_HU4_GPKG +vpu_id: "21" +crs: EPSG:6566 +permitted_fcodes: + - 'Canal/Ditch' + - 'Stream/River: Hydrographic Category = Perennial' + - 'Stream/River: Hydrographic Category = Intermittent' + - 'Artificial Path' + - 'Pipeline: Pipeline Type = Aqueduct; Relationship to Surface = At or Near' + - 'Pipeline: Pipeline Type = Aqueduct; Relationship to Surface = Underground' + - 'Pipeline: Pipeline Type = Siphon' + - 'Pipeline: Pipeline Type = Aqueduct' + - 'Pipeline' + - 'Connector' + - 'Canal Ditch: Canal Ditch Type = Stormwater' + - 'Pipeline: Pipeline Type = Penstock; Relationship to Surface = At or Near' diff --git a/docs/img/hierarchy.png b/docs/img/hierarchy.png new file mode 100644 index 0000000..629074e Binary files /dev/null and b/docs/img/hierarchy.png differ diff --git a/licenses.md b/licenses.md new file mode 100644 index 0000000..5f1f570 --- /dev/null +++ b/licenses.md @@ -0,0 +1,60 @@ +# Licenses + +## CC BY 4.0 License + +The GEOGloWS Hydrologic Model Version 2.0 (GEOGloWS) is licensed under the Creative Commons Attribution 4.0 +International License. To view a copy of this license, visit [http://creativecommons.org/licenses/by/4.0/](http://creativecommons.org/licenses/by/4.0/). + +You are free to: + +Share — copy and redistribute the material in any medium or format for any purpose, even commercially. + +Adapt — remix, transform, and build upon the material for any purpose, even commercially. The licensor cannot revoke these freedoms as long as you follow the license terms. + +Under the following terms: + +Attribution - You must give appropriate credit , provide a link to the license, and indicate if changes were made . You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use. + +No additional restrictions - You may not apply legal terms or technological measures that legally restrict others from doing anything the license permits. + +Notices: + +You do not have to comply with the license for elements of the material in the public domain or where your use is permitted by an applicable exception or limitation . + +No warranties are given. The license may not give you all of the permissions necessary for your intended use. For example, other rights such as publicity, privacy, or moral rights may limit how you use the material. + + +## Licenses of Open Datasets +GEOGloWS uses several licensed open source datasets. The following table summarizes the licenses for these data. + +| Dataset | Purpose | License | +|---------------------------------------------------|--------------------------------------------------------|-----------------------------------------------------------------| +| [TDX-Hydro](https://nga.mil) | Geospatial stream center line and catchment boundaries | [CC BY-SA 4.0](https://creativecommons.org/licenses/by-sa/4.0/) | +| [ECMWF IFS](https://doi.org/10.21957/fv6k37c49h) | Gridded runoff forecast data | [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) | +| [ERA5](https://doi.org/10.24381/cds.143582cf) | Gridded runoff retrospective data | [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) | + +### Licensing statements required by ECMWF Terms of Use + +The following statements are required by ECMWF in compliance with their [terms of use](https://doi.org/10.21957/open-data). + +The following wording shall be attached to the use of this ECMWF data product: +1. Copyright statement: Copyright "© [2023] European Centre for Medium-Range Weather Forecasts (ECMWF)". +2. Source www.ecmwf.int +3. Licence Statement: This data is published under a Creative Commons Attribution 4.0 International (CC BY 4.0). https://creativecommons.org/licenses/by/4.0/ +4. Disclaimer: ECMWF does not accept any liability whatsoever for any error or omission in the data, their availability, or for any loss or damage arising from their use. +5. Where applicable, an indication if the material has been modified and an indication of previous modifications. GEOGloWS makes no modifications to datasets provided by ECMWF. + +## Acknowledgement of Open Source Code +Several open source code projects are used to build GEOGloWS. They do not require license statements or redistribution but we acknowledge their role. +1. [GDAL](https://gdal.org/) - Geospatial Data Abstraction Library +2. [PROJ](https://proj.org/) - Cartographic Projections Library +3. [Geopackage](https://www.geopackage.org/) - GeoPackage is an OGC open source, standards-based, platform-independent, portable, self-describing, compact format for transferring geospatial information. +4. [NetCDF](https://www.unidata.ucar.edu/software/netcdf/) - Network Common Data Form +5. [Zarr](https://zarr.readthedocs.io/en/stable/) - Zarr is a format for the storage of chunked, compressed, N-dimensional arrays +6. [NCO](http://nco.sourceforge.net/) - NetCDF Operators (NCO) are a suite of tools for manipulation of NetCDF files +7. [Xarray](http://xarray.pydata.org/en/stable/) - N-D labeled arrays and datasets in Python +8. [Dask](https://dask.org/) - Dask is a flexible library for parallel computing in Python +9. [Numpy](https://numpy.org/) - The fundamental package for scientific computing with Python +10. [Pandas](https://pandas.pydata.org/) - Powerful data structures for data analysis, time series, and statistics +11. [GeoPandas](https://geopandas.org/) - GeoPandas is an open source project to make working with geospatial data in python easier +12. [RAPID](https://rapid-hub.org/) - RAPID is an implementation of the musking river routing method in fortran optimized for large river networks. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..85f6b1a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,144 @@ +[build-system] +build-backend = "hatchling.build" +requires = ["hatchling", "hatch-vcs"] + +[tool.hatch] +version.source = "vcs" +build.hooks.vcs.version-file = "src/reference_builds/_version.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/reference_builds"] + +[project] +name = "reference-builds" +dynamic = ["version"] +description = "Building Reference Datasets + Processing Ancillary Data" +readme = "README.md" +requires-python = ">=3.12" +license = { file = "LICENSE" } +authors = [ + { name = "Brodie Alexander", email = "Brodie.A.Alexander@rtx.com" }, + { name = "Tadd Bindas", email = "tadd.bindas@ertcorp.com" }, + { name = "Quercus Hamlin", email = "qhamlin@asrcfederal.com" }, + { name = "Dylan Lee", email = "dylan.lee@entarian.com" }, + { name = "Farshid Rahmani", email = "Farshid.Rahmani@rtx.com" }, +] +maintainers = [ + { name = "Brodie Alexander", email = "Brodie.A.Alexander@rtx.com" }, + { name = "Quercus Hamlin", email = "qhamlin@asrcfederal.com" }, + { name = "Dylan Lee", email = "dylan.lee@entarian.com" }, +] + +dependencies = [ + "python-dotenv==1.1.0", + "geopandas==1.1.1", + "boto3==1.40.45", + "pyarrow==20.0.0", + "pyiceberg[s3fs,glue,sql-sqlite]==0.9.1", + "pyprojroot==0.3.0", + "tqdm==4.67.1", + "folium==0.20.0", + "matplotlib==3.10.5", + "mapclassify==2.10.0", + "numpy==2.2.6", + "netCDF4==1.7.2", + "pandas==2.3.0", + "polars==1.34.0", + "shapely==2.1.1", + "xarray==2025.7.1", + "rioxarray==0.19.0", + "rasterio==1.4.3", + "rustworkx==0.17.1", + "exactextract==0.2.2", + "rasterstats>=0.20.0", +] + +[dependency-groups] +dev = [ + "pre-commit==3.8.0", + "ruff==0.11.13", + "mypy==1.15.0", + "nbstripout==0.8.1", + "types-PyYaml", + "types-requests==2.32.4.20250611", +] +examples = ["ipykernel==6.29.5", "jupyterlab==4.4.3"] +tests = ["pytest==8.4.1", "pytest-cov==6.1.1", "astropy==7.1.0"] + +[tool.uv] +default-groups = ["dev", "examples", "tests"] + +[tool.ruff] +line-length = 110 +exclude = [".csv", "LICENSE", ".tf", ".tfvars", "_version.py"] +lint.select = [ + "F", # Errors detected by Pyflakes + "E", # Error detected by Pycodestyle + "W", # Warning detected by Pycodestyle + "I", # isort + "D", # pydocstyle + "B", # flake8-bugbear + "Q", # flake8-quotes + "TID", # flake8-tidy-imports + "C4", # flake8-comprehensions + "BLE", # flake8-blind-except + "UP", # pyupgrade + "RUF100", # Report unused noqa directives +] +lint.ignore = [ + # line too long -> we accept long comment lines; black gets rid of long code lines + "E501", + # Do not assign a lambda expression, use a def -> lambda expression assignments are convenient + "E731", + # allow I, O, l as variable names -> I is the identity matrix + "E741", + # Missing docstring in public package + "D104", + # Missing docstring in public module + "D100", + # Missing docstring in __init__ + "D107", + # Errors from function calls in argument defaults. These are fine when the result is immutable. + "B008", + # __magic__ methods are are often self-explanatory, allow missing docstrings + "D105", + # first line should end with a period [Bug: doesn't work with single-line docstrings] + "D400", + # First line should be in imperative mood; try rephrasing + "D401", + ## Disable one in each pair of mutually incompatible rules + # We don't want a blank line before a class docstring + "D203", + # We want docstrings to start immediately after the opening triple quote + "D213", + # Bare except okay for passing + "E722", +] + +[tool.ruff.lint.pydocstyle] +convention = "numpy" + +[tool.ruff.lint.per-file-ignores] +"docs/*" = ["I"] +"tests/*" = ["D"] +"*/__init__.py" = ["F401"] + +[tool.mypy] +python_version = "3.11" +warn_return_any = false +disallow_any_unimported = false +warn_unused_configs = true +strict_optional = true +ignore_missing_imports = true +check_untyped_defs = true +disallow_untyped_defs = true +no_implicit_optional = true +show_error_codes = true +warn_unused_ignores = true + +[tool.pytest.ini_options] +filterwarnings = [ + "ignore::DeprecationWarning:pyogrio", + "ignore:The 'shapely.geos' module is deprecated:DeprecationWarning", + "ignore:The behavior of DataFrame concatenation with empty or all-NA entries is deprecated:FutureWarning", +] diff --git a/src/reference_builds/__init__.py b/src/reference_builds/__init__.py new file mode 100644 index 0000000..26d23ba --- /dev/null +++ b/src/reference_builds/__init__.py @@ -0,0 +1,3 @@ +from ._version import __version__ + +__all__ = ["__version__"] diff --git a/src/reference_builds/configs/__init__.py b/src/reference_builds/configs/__init__.py new file mode 100644 index 0000000..1116bc1 --- /dev/null +++ b/src/reference_builds/configs/__init__.py @@ -0,0 +1,3 @@ +from .reference_config import BaseDataset, ReferenceConfig + +__all__ = ["BaseDataset", "ReferenceConfig"] diff --git a/src/reference_builds/configs/reference_config.py b/src/reference_builds/configs/reference_config.py new file mode 100644 index 0000000..a377fc4 --- /dev/null +++ b/src/reference_builds/configs/reference_config.py @@ -0,0 +1,104 @@ +"""A file to host all Hydrofabric Schemas""" + +from enum import Enum +from pathlib import Path +from typing import Self + +import yaml +from pydantic import BaseModel, Field +from pyprojroot import here + +from reference_builds import __version__ + + +class BaseDataset(str, Enum): + """Enum for the base dataset used in reference builds""" + + NHD = "nhd" + GEOGLOWS = "geoglows" + USGS_HF = "usgs-reference-hf" + + +class ReferenceConfig(BaseModel): + """Configs for building the ReferenceConfig reference""" + + output_dir: Path = Field( + default=here() / "data/", + description="The directory for output files to be saved from Hydrofabric builds", + ) + + base_dataset: BaseDataset = Field( + default=BaseDataset.NHD, + description="The base dataset to use for the reference build", + ) + + domain: str = Field(description="The domain used for the building your reference") + + input_file_regex: str = Field( + description="regex to find input files to be converted into a reference product", + ) + geoglows_catchment_regex: str | None = Field( + default=None, + description="regex to file catchment files from geoglows", + ) + + crs: str = Field( + default="EPSG:4326", + description="Coordinate Reference System for the domain reference builds. Defaults to https://epsg.io/4326", + ) + + vpu_id: str = Field(description="The VPUID for the domain") + + write_gpkg: bool = Field( + default=True, description="Writes a geopackage in addition to parquet files for output" + ) + + permitted_fcodes: list[str] = Field( + default_factory=lambda: [ + "Stream/River: Hydrographic Category = Intermittent", + "Artificial Path", + "Connector", + "Stream/River: Hydrographic Category = Perennial", + "Canal/Ditch", + "Canal Ditch: Canal Ditch Type = Stormwater", + ], + description="The permitted fcode descriptions for the reference", + ) + + output_reference_gpkg_path: Path = Field( + default_factory=lambda data: data["output_dir"] / f"{data['domain']}_{__version__}_reference.gpkg", + description="Save directory for the domain's reference (in .gpkg form)", + ) + + output_reference_divides_path: Path = Field( + default_factory=lambda data: ( + data["output_dir"] / f"{data['domain']}_{__version__}_reference_divides.parquet" + ), + description="Save directory for the domain's reference divides", + ) + + output_reference_flowpaths_path: Path = Field( + default_factory=lambda data: ( + data["output_dir"] / f"{data['domain']}_{__version__}_reference_flowpaths.parquet" + ), + description="Save directory for the domain's reference flowpaths", + ) + + @classmethod + def from_yaml(cls, path: str | Path) -> Self: + """An internal method to read a config from a YAML file + + Parameters + ---------- + path : str | Path + The path to the provided YAML file + + Returns + ------- + HFConfig + A configuration object validated + """ + with open(path) as f: + data = yaml.safe_load(f) + + return cls(**data) diff --git a/src/reference_builds/local_runner.py b/src/reference_builds/local_runner.py new file mode 100644 index 0000000..c864189 --- /dev/null +++ b/src/reference_builds/local_runner.py @@ -0,0 +1,128 @@ +"""A file to hold the LocalRunner class""" + +from collections.abc import Callable +from datetime import datetime +from typing import Any, Self + +from reference_builds.configs import ReferenceConfig +from reference_builds.logs import setup_logging +from reference_builds.task_instance import TaskInstance + + +class LocalRunner: + """Execute pipeline tasks locally with Airflow-like interface. + + Parameters + ---------- + config : HFConfig + Pipeline configuration containing build settings and parameters. + run_id : str or None, default=None + Unique identifier for this pipeline run. If None, generated from + current timestamp in format 'YYYYMMDD_HHMMSS'. + + Attributes + ---------- + config : HFConfig + The pipeline configuration. + run_id : str + Unique identifier for this run. + ti : TaskInstance + TaskInstance for XCom operations. + results : dict[str, dict[str, Any]] + Execution results for each task, keyed by task_id. + """ + + def __init__( + self, + config: ReferenceConfig, + run_id: str | None = None, + ) -> None: + """Initialize the LocalRunner. + + Parameters + ---------- + config : HFConfig + Pipeline configuration. + run_id : str or None, default=None + Optional run identifier. Auto-generated if not provided. + """ + self.config: ReferenceConfig = config + self.run_id: str = run_id or datetime.now().strftime("%Y%m%d_%H%M%S") + self.ti: TaskInstance = TaskInstance() + self.results: dict[str, dict[str, Any]] = {} + self.logger = setup_logging() + + def cleanup(self) -> None: + """Clean up resources""" + self.logger.info("runner: Closing processes") + + def __enter__(self: Self) -> Self: + """Context manager entry.""" + return self + + def __exit__(self: Self, *args: str, **kwargs: str) -> None: + """Context manager exit - ensures cleanup.""" + self.cleanup() + + def run_task( + self, + task_id: str, + python_callable: Callable[..., Any], + op_kwargs: dict[str, Any] | None = None, + ) -> Any: + """Execute a single task. + + Parameters + ---------- + task_id : str + Unique identifier for this task. Used in XCom keys and result tracking. + python_callable : Callable[..., Any] + The function to execute. Must accept **kwargs to receive context. + op_kwargs : dict[str, Any] or None, default=None + Additional keyword arguments to pass to the callable. + + Returns + ------- + Any + The return value from the callable. + """ + self.logger.info(f"Running task: {task_id}") + + context: dict[str, Any] = { + "ti": self.ti, + "task_id": task_id, + "run_id": self.run_id, + "ds": datetime.now().strftime("%Y-%m-%d"), + "execution_date": datetime.now(), + "config": self.config, + } + + kwargs = {**(op_kwargs or {}), **context} + + result = python_callable(**kwargs) + + for k, v in result.items(): + self.ti.xcom_push(f"{task_id}.{k}", v) + self.results[task_id] = {"status": "success", "result": result} + + self.logger.info(f"✓ Task {task_id} completed") + return result + + def get_result(self, task_id: str) -> dict[str, Any]: + """Retrieve execution results for a specific task. + + Parameters + ---------- + task_id : str + The identifier of the task to get results for. + + Returns + ------- + dict[str, Any] or None + Dictionary containing 'status' and either 'result' (on success) + or 'error' (on failure). Returns None if task_id not found. + """ + result = self.results.get(task_id) + if result is None: + raise ValueError("Cannot find result from task") + return result diff --git a/src/reference_builds/logs.py b/src/reference_builds/logs.py new file mode 100644 index 0000000..c7fefe7 --- /dev/null +++ b/src/reference_builds/logs.py @@ -0,0 +1,31 @@ +import logging +import logging.handlers +import os + +from dotenv import load_dotenv +from pyprojroot import here + + +def setup_logging() -> logging.Logger: + """Configures the reference builds logging""" + load_dotenv(here() / ".env") + + logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") + + logger = logging.getLogger(__name__) + logging.getLogger("rasterio").setLevel(logging.WARNING) # turning off rasterio INFO logging + + log_file_path = here() / "logs/" + log_file_path.mkdir(exist_ok=True) + max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760)) + backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5)) + + file_handler = logging.handlers.RotatingFileHandler( + log_file_path / "reference_builds.log", maxBytes=max_bytes, backupCount=backup_count + ) + + formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") + file_handler.setFormatter(formatter) + file_handler.setLevel(logging.DEBUG) + logging.getLogger().addHandler(file_handler) + return logger diff --git a/src/reference_builds/pipeline/__init__.py b/src/reference_builds/pipeline/__init__.py new file mode 100644 index 0000000..c4f544e --- /dev/null +++ b/src/reference_builds/pipeline/__init__.py @@ -0,0 +1,17 @@ +from .build_reference import build_geoglows_reference, build_nhd_reference, build_usgs_hf_reference +from .download import download_geoglows_data, download_nhd_data, download_usgs_hf_data +from .processing import build_geoglows_graphs, build_nhd_graphs, build_usgs_hf_graphs +from .write import write_reference + +__all__ = [ + "build_usgs_hf_graphs", + "build_geoglows_graphs", + "build_nhd_graphs", + "download_geoglows_data", + "build_nhd_reference", + "build_geoglows_reference", + "build_usgs_hf_reference", + "download_nhd_data", + "download_usgs_hf_data", + "write_reference", +] diff --git a/src/reference_builds/pipeline/build_reference.py b/src/reference_builds/pipeline/build_reference.py new file mode 100644 index 0000000..11bb786 --- /dev/null +++ b/src/reference_builds/pipeline/build_reference.py @@ -0,0 +1,1066 @@ +"""Contains all code for building a reference fabric from the reference graph object""" + +import logging +from typing import Any, cast + +import geopandas as gpd +import pandas as pd +import polars as pl +import rustworkx as rx + +from reference_builds.configs import ReferenceConfig +from reference_builds.task_instance import TaskInstance +from reference_builds.utils.geometries import _orient_flowpath_downstream + +logger = logging.getLogger(__name__) + + +def _trace_attributes( + graph: rx.PyDiGraph, + node_indices: dict[str, int], + flowpaths: gpd.GeoDataFrame, + divides: gpd.GeoDataFrame, + vpu_id: str, +) -> gpd.GeoDataFrame: + """Trace flowpath attributes for the entire graph. + + Parameters + ---------- + graph : rx.PyDiGraph + The rustworkx directed graph (may contain multiple disconnected subgraphs) + node_indices : dict[str, int] + Mapping from NHDPlusID (as string) to node index + flowpaths : gpd.GeoDataFrame + The flowpaths GeoDataFrame with LengthKM + divides : gpd.GeoDataFrame + The divides GeoDataFrame with AreaSqKm + + Returns + ------- + pl.DataFrame + Traced attributes: totdasqkm, mainstemlp, pathlength, dnhydroseq, hydroseq, stream_order + """ + flowpaths_lookup = flowpaths.set_index("NHDPlusID")["LengthKM"].to_dict() + divides_lookup = divides.set_index("NHDPlusID")["AreaSqKm"].to_dict() + fp_geom_lookup = flowpaths.set_index("NHDPlusID")["geometry"].to_dict() + fp_fcode_lookup = flowpaths.set_index("NHDPlusID")["fcode_description"].to_dict() + for node_idx in graph.node_indices(): + flowpath_id = str(graph[node_idx]) + nhd_id = int(flowpath_id) + + graph[node_idx] = { + "flowpath_id": flowpath_id, + "areasqkm": divides_lookup.get(nhd_id, 0.0), + "lengthkm": flowpaths_lookup.get(nhd_id, 0.0), + "totdasqkm": 0.0, + "mainstemlp": None, + "pathlength": 0.0, + "dnhydroseq": None, + "hydroseq": None, + "streamorder": None, + "fcode_description": fp_fcode_lookup[nhd_id], + "geometry": fp_geom_lookup[nhd_id], + } + + # Find all outlets (nodes with no downstream connections) + outlets = [idx for idx in graph.node_indices() if graph.out_degree(idx) == 0] + logger.info(f"build_nhd_reference task: Found {len(outlets)} outlets (disconnected subgraphs)") + + # Get topological order for entire graph + try: + topo_order = rx.topological_sort(graph) + except rx.DAGHasCycle as e: + raise AssertionError("Graph contains cycles") from e + + # PASS 1: Calculate pathlength and hydroseq (reverse topo order - upstream from outlets) + current_hydroseq = 1 + + # Initialize outlets + for outlet_idx in outlets: + graph[outlet_idx]["pathlength"] = 0.0 + graph[outlet_idx]["dnhydroseq"] = 0 + + # Traverse in reverse topo order + for node_idx in reversed(topo_order): + # Assign hydroseq + graph[node_idx]["hydroseq"] = current_hydroseq + current_hydroseq += 1 + + # Calculate pathlength based on downstream node + out_edges = graph.out_edges(node_idx) + if out_edges: + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + if downstream_nodes: + downstream_idx = max(downstream_nodes, key=lambda idx: graph[idx]["pathlength"]) + graph[node_idx]["pathlength"] = ( + graph[downstream_idx]["pathlength"] + graph[downstream_idx]["lengthkm"] + ) + + # Trace mainstems for each outlet's basin + current_mainstem_id = 1 + processed: set[int] = set() + + for outlet_idx in outlets: + # Trace main mainstem (longest path from outlet to headwater) + current_idx = outlet_idx + mainstem_nodes = [] + + while current_idx not in processed: + mainstem_nodes.append(current_idx) + graph[current_idx]["mainstemlp"] = current_mainstem_id + processed.add(current_idx) + + in_edges = list(graph.in_edges(current_idx)) + if not in_edges: + break + + upstream_candidates = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + if not upstream_candidates: + break + + current_idx = max( + upstream_candidates, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + current_mainstem_id += 1 + + # Assign tributary mainstems for remaining nodes + for node_idx in graph.node_indices(): + if node_idx not in processed: + tributary_id = current_mainstem_id + current_mainstem_id += 1 + + trib_current = node_idx + while trib_current not in processed: + graph[trib_current]["mainstemlp"] = tributary_id + processed.add(trib_current) + + in_edges = list(graph.in_edges(trib_current)) + upstream_in_basin = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + + if not upstream_in_basin: + break + + trib_current = max( + upstream_in_basin, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + # Assign dnhydroseq and flowpath_toid based on graph edges + for node_idx in graph.node_indices(): + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + if downstream_nodes: + downstream_idx = downstream_nodes[0] + graph[node_idx]["dnhydroseq"] = graph[downstream_idx]["hydroseq"] + graph[node_idx]["flowpath_toid"] = graph[downstream_idx]["flowpath_id"] + else: + graph[node_idx]["dnhydroseq"] = 0 + graph[node_idx]["flowpath_toid"] = "0" + + # PASS 2: Calculate totdasqkm and stream_order (forward topo order - downstream from headwaters) + for node_idx in topo_order: + in_edges = list(graph.in_edges(node_idx)) + + # Accumulate upstream drainage area + upstream_total = sum(graph[src_idx]["totdasqkm"] for src_idx, _, _ in in_edges) + graph[node_idx]["totdasqkm"] = upstream_total + graph[node_idx]["areasqkm"] + + # Calculate Strahler stream order + if not in_edges: + graph[node_idx]["streamorder"] = 1 + else: + upstream_orders = [graph[src_idx]["streamorder"] for src_idx, _, _ in in_edges] + max_order = max(upstream_orders) + count_max = upstream_orders.count(max_order) + + if count_max >= 2: + graph[node_idx]["streamorder"] = max_order + 1 + else: + graph[node_idx]["streamorder"] = max_order + + # PASS 3: Orient all flowpath geometries so they flow upstream -> downstream + for node_idx in graph.node_indices(): + in_edges = graph.in_edges(node_idx) + upstream_nodes = [src_idx for src_idx, _, _ in in_edges] + + # Check if this is an outlet (flowpath_toid == "0" means no downstream) + is_outlet = graph[node_idx]["flowpath_toid"] == "0" + + if is_outlet and upstream_nodes: + # Outlet: use upstream geometry + upstream_idx = upstream_nodes[0] + us_geom = graph[upstream_idx]["geometry"] + graph[node_idx]["geometry"] = _orient_flowpath_downstream( + graph[node_idx]["geometry"], ds_geom=None, us_geom=us_geom + ) + else: + # Normal case: use downstream geometry + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + ds_geom = None + if downstream_nodes: + downstream_idx = downstream_nodes[0] + ds_geom = graph[downstream_idx]["geometry"] + + graph[node_idx]["geometry"] = _orient_flowpath_downstream( + graph[node_idx]["geometry"], ds_geom=ds_geom, us_geom=None + ) + + # Extract results + flowpath_ids = [] + flowpath_toids = [] + vpu_ids = [] + das = [] + lengthkms = [] + total_das = [] + mainstems = [] + pathlengths = [] + dnhydroseqs = [] + hydroseqs = [] + streamorders = [] + fcodes = [] + geometries = [] + + for node_idx in graph.node_indices(): + node_data = graph[node_idx] + flowpath_ids.append(node_data["flowpath_id"]) + flowpath_toids.append(node_data["flowpath_toid"]) + vpu_ids.append(vpu_id) + das.append(node_data["areasqkm"]) + lengthkms.append(node_data["lengthkm"]) + total_das.append(node_data["totdasqkm"]) + mainstems.append(node_data["mainstemlp"]) + pathlengths.append(node_data["pathlength"]) + dnhydroseqs.append(node_data["dnhydroseq"]) + hydroseqs.append(node_data["hydroseq"]) + streamorders.append(node_data["streamorder"]) + fcodes.append(node_data["fcode_description"]) + geometries.append(node_data["geometry"]) + + return gpd.GeoDataFrame( + { + "flowpath_id": flowpath_ids, + "flowpath_toid": flowpath_toids, + "VPUID": vpu_ids, + "lengthkm": lengthkms, + "areasqkm": das, + "totdasqkm": total_das, + "mainstemlp": mainstems, + "pathlength": pathlengths, + "dnhydroseq": dnhydroseqs, + "hydroseq": hydroseqs, + "streamorder": streamorders, + "fcode_description": fcodes, + }, + geometry=geometries, + crs="EPSG:4269", + ) + + +def _trace_geoglows_attributes( + graph: rx.PyDiGraph, + node_indices: dict[str, int], + flowpaths: gpd.GeoDataFrame, + catchments: gpd.GeoDataFrame, + vpu_id: str, +) -> gpd.GeoDataFrame: + """Trace flowpath attributes for the entire GeoGLOWS graph. + + Parameters + ---------- + graph : rx.PyDiGraph + The rustworkx directed graph (may contain multiple disconnected subgraphs) + node_indices : dict[str, int] + Mapping from LINKNO (as string) to node index + flowpaths : gpd.GeoDataFrame + The GeoGLOWS flowpaths GeoDataFrame with LengthKM, strmOrder + catchments : gpd.GeoDataFrame + The GeoGLOWS catchments GeoDataFrame with linkno, areasqkm, and geometry + vpu_id : str + The VPUID for the domain + + Returns + ------- + gpd.GeoDataFrame + Traced attributes: totdasqkm, mainstemlp, pathlength, dnhydroseq, hydroseq, stream_order + """ + length_lookup = flowpaths.set_index("LINKNO")["LengthKM"].to_dict() + order_lookup = flowpaths.set_index("LINKNO")["strmOrder"].to_dict() + fp_geom_lookup = flowpaths.set_index("LINKNO")["geometry"].to_dict() + + # Build catchment lookups (area and geometry keyed by linkno) + catchment_area_lookup = catchments.set_index("linkno")["areasqkm"].to_dict() + catchment_geom_lookup = catchments.set_index("linkno")["geometry"].to_dict() + + for node_idx in graph.node_indices(): + flowpath_id = str(graph[node_idx]) + link_id = int(flowpath_id) + + # Get length in km (already calculated from geometry) + length_km = length_lookup.get(link_id, 0.0) + + # Get local catchment area (from catchments, keyed by linkno) + local_areasqkm = catchment_area_lookup.get(link_id, 0.0) + + # Get geometries + catchment_geometry = catchment_geom_lookup.get(link_id) + flowpath_geometry = fp_geom_lookup.get(link_id) + + graph[node_idx] = { + "flowpath_id": flowpath_id, + "areasqkm": local_areasqkm, + "lengthkm": length_km, + "totdasqkm": 0.0, # Will be accumulated in PASS 2 + "mainstemlp": None, + "pathlength": 0.0, + "dnhydroseq": None, + "hydroseq": None, + "streamorder": order_lookup.get(link_id, 1), + "flowpath_geometry": flowpath_geometry, + "catchment_geometry": catchment_geometry, + } + + # Find all outlets (nodes with no downstream connections) + outlets = [idx for idx in graph.node_indices() if graph.out_degree(idx) == 0] + logger.info(f"build_geoglows_reference task: Found {len(outlets)} outlets (disconnected subgraphs)") + + # Get topological order for entire graph + try: + topo_order = rx.topological_sort(graph) + except rx.DAGHasCycle as e: + raise AssertionError("Graph contains cycles") from e + + # PASS 1: Calculate pathlength and hydroseq (reverse topo order - upstream from outlets) + current_hydroseq = 1 + + # Initialize outlets + for outlet_idx in outlets: + graph[outlet_idx]["pathlength"] = 0.0 + graph[outlet_idx]["dnhydroseq"] = 0 + + # Traverse in reverse topo order + for node_idx in reversed(topo_order): + # Assign hydroseq + graph[node_idx]["hydroseq"] = current_hydroseq + current_hydroseq += 1 + + # Calculate pathlength based on downstream node + out_edges = graph.out_edges(node_idx) + if out_edges: + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + if downstream_nodes: + downstream_idx = max(downstream_nodes, key=lambda idx: graph[idx]["pathlength"]) + graph[node_idx]["pathlength"] = ( + graph[downstream_idx]["pathlength"] + graph[downstream_idx]["lengthkm"] + ) + + # PASS 2: Calculate totdasqkm and stream_order (forward topo order - downstream from headwaters) + for node_idx in topo_order: + in_edges = list(graph.in_edges(node_idx)) + + # Accumulate upstream drainage area + upstream_total = sum(graph[src_idx]["totdasqkm"] for src_idx, _, _ in in_edges) + graph[node_idx]["totdasqkm"] = upstream_total + graph[node_idx]["areasqkm"] + + # Trace mainstems for each outlet's basin + current_mainstem_id = 1 + processed: set[int] = set() + + for outlet_idx in outlets: + # Trace main mainstem (longest path from outlet to headwater) + current_idx = outlet_idx + + while current_idx not in processed: + graph[current_idx]["mainstemlp"] = current_mainstem_id + processed.add(current_idx) + + in_edges = list(graph.in_edges(current_idx)) + if not in_edges: + break + + upstream_candidates = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + if not upstream_candidates: + break + + current_idx = max( + upstream_candidates, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + current_mainstem_id += 1 + + # Assign tributary mainstems for remaining nodes + for node_idx in graph.node_indices(): + if node_idx not in processed: + tributary_id = current_mainstem_id + current_mainstem_id += 1 + + trib_current = node_idx + while trib_current not in processed: + graph[trib_current]["mainstemlp"] = tributary_id + processed.add(trib_current) + + in_edges = list(graph.in_edges(trib_current)) + upstream_in_basin = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + + if not upstream_in_basin: + break + + trib_current = max( + upstream_in_basin, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + # Assign dnhydroseq and flowpath_toid based on graph edges + for node_idx in graph.node_indices(): + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + if downstream_nodes: + downstream_idx = downstream_nodes[0] + graph[node_idx]["dnhydroseq"] = graph[downstream_idx]["hydroseq"] + graph[node_idx]["flowpath_toid"] = graph[downstream_idx]["flowpath_id"] + else: + graph[node_idx]["dnhydroseq"] = 0 + graph[node_idx]["flowpath_toid"] = "0" + + # PASS 3: Orient all flowpath geometries so they flow upstream -> downstream + for node_idx in graph.node_indices(): + in_edges = graph.in_edges(node_idx) + upstream_nodes = [src_idx for src_idx, _, _ in in_edges] + + # Check if this is an outlet (flowpath_toid == "0" means no downstream) + is_outlet = graph[node_idx]["flowpath_toid"] == "0" + + if is_outlet and upstream_nodes: + # Outlet: use upstream geometry + upstream_idx = upstream_nodes[0] + us_geom = graph[upstream_idx]["flowpath_geometry"] + graph[node_idx]["flowpath_geometry"] = _orient_flowpath_downstream( + graph[node_idx]["flowpath_geometry"], ds_geom=None, us_geom=us_geom + ) + else: + # Normal case: use downstream geometry + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + ds_geom = None + if downstream_nodes: + downstream_idx = downstream_nodes[0] + ds_geom = graph[downstream_idx]["flowpath_geometry"] + + graph[node_idx]["flowpath_geometry"] = _orient_flowpath_downstream( + graph[node_idx]["flowpath_geometry"], ds_geom=ds_geom, us_geom=None + ) + + # Extract results for flowpaths + flowpath_ids = [] + flowpath_toids = [] + vpu_ids = [] + das = [] + lengthkms = [] + total_das = [] + mainstems = [] + pathlengths = [] + dnhydroseqs = [] + hydroseqs = [] + streamorders = [] + flowpath_geometries = [] + + for node_idx in graph.node_indices(): + node_data = graph[node_idx] + flowpath_ids.append(node_data["flowpath_id"]) + flowpath_toids.append(node_data["flowpath_toid"]) + vpu_ids.append(vpu_id) + das.append(node_data["areasqkm"]) + lengthkms.append(node_data["lengthkm"]) + total_das.append(node_data["totdasqkm"]) + mainstems.append(node_data["mainstemlp"]) + pathlengths.append(node_data["pathlength"]) + dnhydroseqs.append(node_data["dnhydroseq"]) + hydroseqs.append(node_data["hydroseq"]) + streamorders.append(node_data["streamorder"]) + flowpath_geometries.append(node_data["flowpath_geometry"]) + + return gpd.GeoDataFrame( + { + "flowpath_id": flowpath_ids, + "flowpath_toid": flowpath_toids, + "VPUID": vpu_ids, + "lengthkm": lengthkms, + "areasqkm": das, + "totdasqkm": total_das, + "mainstemlp": mainstems, + "pathlength": pathlengths, + "dnhydroseq": dnhydroseqs, + "hydroseq": hydroseqs, + "streamorder": streamorders, + }, + geometry=flowpath_geometries, + crs="EPSG:3857", + ) + + +def _trace_usgs_hf_attributes( + graph: rx.PyDiGraph, + node_indices: dict[str, int], + flowpaths: gpd.GeoDataFrame, + catchments: gpd.GeoDataFrame, + vpu_id: str, +) -> gpd.GeoDataFrame: + """Trace flowpath attributes for the entire USGS Reference Hydrofabric graph. + + Parameters + ---------- + graph : rx.PyDiGraph + The rustworkx directed graph (may contain multiple disconnected subgraphs) + node_indices : dict[str, int] + Mapping from hydroseq (as string) to node index + flowpaths : gpd.GeoDataFrame + The USGS Reference Hydrofabric flowpaths GeoDataFrame with LengthKM, streamorde + catchments : gpd.GeoDataFrame + The USGS Reference Hydrofabric catchments GeoDataFrame with hydroseq, areasqkm, and geometry + vpu_id : str + The VPUID for the domain + + Returns + ------- + gpd.GeoDataFrame + Traced attributes: totdasqkm, mainstemlp, pathlength, dnhydroseq, hydroseq, stream_order + """ + length_lookup = flowpaths.set_index("hydroseq")["LengthKM"].to_dict() + order_lookup = flowpaths.set_index("hydroseq")["streamorde"].to_dict() + fp_geom_lookup = flowpaths.set_index("hydroseq")["geometry"].to_dict() + comid_lookup = flowpaths.set_index("hydroseq")["comid"].to_dict() + + # Build catchment lookups (area and geometry keyed by linkno) + catchment_area_lookup = catchments.set_index("COMID")["areasqkm"].to_dict() + catchment_geom_lookup = catchments.set_index("COMID")["geometry"].to_dict() + + for node_idx in graph.node_indices(): + flowpath_id = str(graph[node_idx]) + link_id = int(flowpath_id) + + comid = comid_lookup.get(link_id) + + # Get length in km (already calculated from geometry) + length_km = length_lookup.get(link_id, 0.0) + + # Get local catchment area (from catchments, keyed by comid) + local_areasqkm = catchment_area_lookup.get(comid, 0.0) + + # Get geometries + catchment_geometry = catchment_geom_lookup.get(comid) + flowpath_geometry = fp_geom_lookup.get(link_id) + + graph[node_idx] = { + "flowpath_id": flowpath_id, + "areasqkm": local_areasqkm, + "lengthkm": length_km, + "totdasqkm": 0.0, # Will be accumulated in PASS 2 + "mainstemlp": None, + "pathlength": 0.0, + "dnhydroseq": None, + "hydroseq": None, + "streamorder": order_lookup.get(link_id, 1), + "flowpath_geometry": flowpath_geometry, + "catchment_geometry": catchment_geometry, + } + + # Find all outlets (nodes with no downstream connections) + outlets = [idx for idx in graph.node_indices() if graph.out_degree(idx) == 0] + logger.info(f"build_usgs_hf_reference task: Found {len(outlets)} outlets (disconnected subgraphs)") + + # Get topological order for entire graph + try: + topo_order = rx.topological_sort(graph) + except rx.DAGHasCycle as e: + raise AssertionError("Graph contains cycles") from e + + # PASS 1: Calculate pathlength and hydroseq (reverse topo order - upstream from outlets) + current_hydroseq = 1 + + # Initialize outlets + for outlet_idx in outlets: + graph[outlet_idx]["pathlength"] = 0.0 + graph[outlet_idx]["dnhydroseq"] = 0 + + # Traverse in reverse topo order + for node_idx in reversed(topo_order): + # Assign hydroseq + graph[node_idx]["hydroseq"] = current_hydroseq + current_hydroseq += 1 + + # Calculate pathlength based on downstream node + out_edges = graph.out_edges(node_idx) + if out_edges: + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + if downstream_nodes: + downstream_idx = max(downstream_nodes, key=lambda idx: graph[idx]["pathlength"]) + graph[node_idx]["pathlength"] = ( + graph[downstream_idx]["pathlength"] + graph[downstream_idx]["lengthkm"] + ) + + # PASS 2: Calculate totdasqkm and stream_order (forward topo order - downstream from headwaters) + for node_idx in topo_order: + in_edges = list(graph.in_edges(node_idx)) + + # Accumulate upstream drainage area + upstream_total = sum(graph[src_idx]["totdasqkm"] for src_idx, _, _ in in_edges) + graph[node_idx]["totdasqkm"] = upstream_total + graph[node_idx]["areasqkm"] + + # Trace mainstems for each outlet's basin + current_mainstem_id = 1 + processed: set[int] = set() + + for outlet_idx in outlets: + # Trace main mainstem (longest path from outlet to headwater) + current_idx = outlet_idx + + while current_idx not in processed: + graph[current_idx]["mainstemlp"] = current_mainstem_id + processed.add(current_idx) + + in_edges = list(graph.in_edges(current_idx)) + if not in_edges: + break + + upstream_candidates = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + if not upstream_candidates: + break + + current_idx = max( + upstream_candidates, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + current_mainstem_id += 1 + + # Assign tributary mainstems for remaining nodes + for node_idx in graph.node_indices(): + if node_idx not in processed: + tributary_id = current_mainstem_id + current_mainstem_id += 1 + + trib_current = node_idx + while trib_current not in processed: + graph[trib_current]["mainstemlp"] = tributary_id + processed.add(trib_current) + + in_edges = list(graph.in_edges(trib_current)) + upstream_in_basin = [src_idx for src_idx, _, _ in in_edges if src_idx not in processed] + + if not upstream_in_basin: + break + + trib_current = max( + upstream_in_basin, + key=lambda idx: (graph[idx]["pathlength"], graph[idx]["totdasqkm"]), + ) + + # Assign dnhydroseq and flowpath_toid based on graph edges + for node_idx in graph.node_indices(): + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + if downstream_nodes: + downstream_idx = downstream_nodes[0] + graph[node_idx]["dnhydroseq"] = graph[downstream_idx]["hydroseq"] + graph[node_idx]["flowpath_toid"] = graph[downstream_idx]["flowpath_id"] + else: + graph[node_idx]["dnhydroseq"] = 0 + graph[node_idx]["flowpath_toid"] = "0" + + # PASS 3: Orient all flowpath geometries so they flow upstream -> downstream + for node_idx in graph.node_indices(): + in_edges = graph.in_edges(node_idx) + upstream_nodes = [src_idx for src_idx, _, _ in in_edges] + + # Check if this is an outlet (flowpath_toid == "0" means no downstream) + is_outlet = graph[node_idx]["flowpath_toid"] == "0" + + if is_outlet and upstream_nodes: + # Outlet: use upstream geometry + upstream_idx = upstream_nodes[0] + us_geom = graph[upstream_idx]["flowpath_geometry"] + if not graph[node_idx]["flowpath_geometry"]: + print(f"NONETYPE: {graph[node_idx]}") + graph[node_idx]["flowpath_geometry"] = _orient_flowpath_downstream( + graph[node_idx]["flowpath_geometry"], ds_geom=None, us_geom=us_geom + ) + else: + # Normal case: use downstream geometry + out_edges = graph.out_edges(node_idx) + downstream_nodes = [tgt_idx for _, tgt_idx, _ in out_edges] + + ds_geom = None + if downstream_nodes: + downstream_idx = downstream_nodes[0] + ds_geom = graph[downstream_idx]["flowpath_geometry"] + + graph[node_idx]["flowpath_geometry"] = _orient_flowpath_downstream( + graph[node_idx]["flowpath_geometry"], ds_geom=ds_geom, us_geom=None + ) + + # Extract results for flowpaths + flowpath_ids = [] + flowpath_toids = [] + vpu_ids = [] + das = [] + lengthkms = [] + total_das = [] + mainstems = [] + pathlengths = [] + dnhydroseqs = [] + hydroseqs = [] + streamorders = [] + flowpath_geometries = [] + + for node_idx in graph.node_indices(): + node_data = graph[node_idx] + flowpath_ids.append(node_data["flowpath_id"]) + flowpath_toids.append(node_data["flowpath_toid"]) + vpu_ids.append(vpu_id) + das.append(node_data["areasqkm"]) + lengthkms.append(node_data["lengthkm"]) + total_das.append(node_data["totdasqkm"]) + mainstems.append(node_data["mainstemlp"]) + pathlengths.append(node_data["pathlength"]) + dnhydroseqs.append(node_data["dnhydroseq"]) + hydroseqs.append(node_data["hydroseq"]) + streamorders.append(node_data["streamorder"]) + flowpath_geometries.append(node_data["flowpath_geometry"]) + + return gpd.GeoDataFrame( + { + "flowpath_id": flowpath_ids, + "flowpath_toid": flowpath_toids, + "VPUID": vpu_ids, + "lengthkm": lengthkms, + "areasqkm": das, + "totdasqkm": total_das, + "mainstemlp": mainstems, + "pathlength": pathlengths, + "dnhydroseq": dnhydroseqs, + "hydroseq": hydroseqs, + "streamorder": streamorders, + }, + geometry=flowpath_geometries, + crs="EPSG:4326", + ) + + +def _create_reference_divides( + divides_df: gpd.GeoDataFrame, reference_flowpaths: gpd.GeoDataFrame, vpu_id: str +) -> gpd.GeoDataFrame: + """A function to create the reference divides table + + Parameters + ---------- + divides_df : gpd.GeoDataFrame + the NHDCatchments table + reference_flowpaths : gpd.GeoDataFrame + The reference flowpaths + vpu_id : str + the VPUID we're working in + + Returns + ------- + gpd.GeoDataFrame + the outputted reference_divides + """ + reference_divides = divides_df.rename( + columns={"NHDPlusID": "divide_id", "VPUID": "vpuid", "AreaSqKm": "areasqkm"} + ) + reference_divides["divide_id"] = reference_divides["divide_id"].astype(int).astype(str) + reference_divides["vpuid"] = vpu_id + mask = reference_divides["divide_id"].isin(reference_flowpaths["flowpath_id"]) + reference_divides["has_flowpath"] = mask + reference_divides["flowpath_id"] = pd.NA + reference_divides.loc[mask, "flowpath_id"] = reference_divides.loc[mask, "divide_id"] + return reference_divides + + +def _create_geoglows_reference_divides( + catchments_df: gpd.GeoDataFrame, reference_flowpaths: gpd.GeoDataFrame, vpu_id: str +) -> gpd.GeoDataFrame: + """A function to create the reference divides table from GeoGLOWS catchments + + Parameters + ---------- + catchments_df : gpd.GeoDataFrame + The GeoGLOWS catchments table with linkno, areasqkm, and geometry + reference_flowpaths : gpd.GeoDataFrame + The reference flowpaths + vpu_id : str + The VPUID we're working in + + Returns + ------- + gpd.GeoDataFrame + The outputted reference_divides with catchment geometries + """ + reference_divides = catchments_df.copy() + reference_divides = reference_divides.rename(columns={"linkno": "divide_id"}) + reference_divides["divide_id"] = reference_divides["divide_id"].astype(int).astype(str) + reference_divides["vpuid"] = vpu_id + + # Filter to only include catchments that have a corresponding flowpath + mask = reference_divides["divide_id"].isin(reference_flowpaths["flowpath_id"]) + reference_divides["has_flowpath"] = mask + reference_divides["flowpath_id"] = pd.NA + reference_divides.loc[mask, "flowpath_id"] = reference_divides.loc[mask, "divide_id"] + + return reference_divides + + +def _create_usgs_hf_reference_divides( + catchments_df: gpd.GeoDataFrame, reference_flowpaths: gpd.GeoDataFrame, vpu_id: str +) -> gpd.GeoDataFrame: + """A function to create the reference divides table from USGS Reference Hydrofabric catchments + + Parameters + ---------- + catchments_df : gpd.GeoDataFrame + The USGS Reference Hydrofabric catchments table with hydroseq, areasqkm, and geometry + reference_flowpaths : gpd.GeoDataFrame + The reference flowpaths + vpu_id : str + The VPUID we're working in + + Returns + ------- + gpd.GeoDataFrame + The outputted reference_divides with catchment geometries + """ + reference_divides = catchments_df.copy() + reference_divides = reference_divides.rename(columns={"COMID": "divide_id"}) + reference_divides["divide_id"] = reference_divides["divide_id"].astype(int).astype(str) + reference_divides["vpuid"] = vpu_id + + # Filter to only include catchments that have a corresponding flowpath + mask = reference_divides["divide_id"].isin(reference_flowpaths["flowpath_id"]) + reference_divides["has_flowpath"] = mask + reference_divides["flowpath_id"] = pd.NA + reference_divides.loc[mask, "flowpath_id"] = reference_divides.loc[mask, "divide_id"] + + return reference_divides + + +def build_nhd_reference(**context: dict[str, Any]) -> dict[str, Any]: + """Opens local / downloads for the reference-build process + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The reference flowpath and divides references in memory + """ + ti = cast(TaskInstance, context["ti"]) + cfg = cast(ReferenceConfig, context["config"]) + graph: rx.PyDiGraph = ti.xcom_pull(task_id="build_nhd_graphs", key="graph") + node_indices: dict[str, int] = ti.xcom_pull(task_id="build_nhd_graphs", key="node_indices") + _flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="nhd_flowpaths") + _divides: pl.DataFrame = ti.xcom_pull(task_id="download", key="nhd_divides") + cycles_iter = rx.simple_cycles(graph) + cycles: list[list[str]] = [] + cycle_ids: set[str] = set() + for cycle in cycles_iter: + _ids: list[Any] = [graph.get_node_data(node_idx) for node_idx in cycle] + cycles.append(_ids) + cycle_ids.update(_ids) + if cycle_ids: + raise NotImplementedError("Cycle Detected. Please create method for removing") + + _flowpaths_df = gpd.GeoDataFrame( + _flowpaths.select( + [ + pl.col("NHDPlusID"), + pl.col("VPUID"), + pl.col("LengthKM"), + pl.col("fcode_description"), + ] + ).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_flowpaths["geometry"]), + crs="EPSG:4269", + ) + + _divides_df = gpd.GeoDataFrame( + _divides.select( + [ + pl.col("NHDPlusID"), + pl.col("VPUID"), + pl.col("AreaSqKm"), + ] + ).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_divides["geometry"]), + crs="EPSG:4269", + ) + reference_flowpaths = _trace_attributes(graph, node_indices, _flowpaths_df, _divides_df, cfg.vpu_id) + reference_divides = _create_reference_divides(_divides_df, reference_flowpaths, cfg.vpu_id) + + return {"reference_flowpaths": reference_flowpaths, "reference_divides": reference_divides} + + +def build_geoglows_reference(**context: dict[str, Any]) -> dict[str, Any]: + """Builds reference fabric from GeoGLOWS data + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The reference flowpaths and divides in memory + """ + ti = cast(TaskInstance, context["ti"]) + cfg = cast(ReferenceConfig, context["config"]) + graph: rx.PyDiGraph = ti.xcom_pull(task_id="build_geoglows_graphs", key="graph") + node_indices: dict[str, int] = ti.xcom_pull(task_id="build_geoglows_graphs", key="node_indices") + _flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="geoglows_flowpaths") + _catchments: pl.DataFrame = ti.xcom_pull(task_id="download", key="geoglows_divides") + + # Check for cycles + cycles_iter = rx.simple_cycles(graph) + cycles: list[list[str]] = [] + cycle_ids: set[str] = set() + for cycle in cycles_iter: + _ids: list[Any] = [graph.get_node_data(node_idx) for node_idx in cycle] + cycles.append(_ids) + cycle_ids.update(_ids) + if cycle_ids: + raise NotImplementedError("Cycle Detected. Please create method for removing") + + _flowpaths_df = gpd.GeoDataFrame( + _flowpaths.select( + [ + pl.col("LINKNO"), + pl.col("DSLINKNO"), + pl.col("strmOrder"), + ] + ).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_flowpaths["geometry"]), + crs="EPSG:3857", + ) + + _flowpaths_df_projected = _flowpaths_df.to_crs(cfg.crs) + _flowpaths_df["LengthKM"] = _flowpaths_df_projected.geometry.length / 1000 + + _catchments_df = gpd.GeoDataFrame( + _catchments.select([pl.col("linkno")]).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_catchments["geometry"]), + crs="EPSG:4326", + ) + _catchments_df_projected = _catchments_df.to_crs(cfg.crs) + _catchments_df["areasqkm"] = _catchments_df_projected.geometry.area / 1e6 + + # Log any flowpaths without matching catchments + flowpath_linkno_set = set(_flowpaths_df["LINKNO"].tolist()) + catchment_linkno_set = set(_catchments_df["linkno"].tolist()) + missing_catchments = flowpath_linkno_set - catchment_linkno_set + if missing_catchments: + logger.warning( + f"build_geoglows_reference task: {len(missing_catchments)} flowpaths have no matching catchment" + ) + + reference_flowpaths = _trace_geoglows_attributes( + graph, node_indices, _flowpaths_df, _catchments_df, cfg.vpu_id + ) + reference_divides = _create_geoglows_reference_divides(_catchments_df, reference_flowpaths, cfg.vpu_id) + + return {"reference_flowpaths": reference_flowpaths, "reference_divides": reference_divides} + + +def build_usgs_hf_reference(**context: dict[str, Any]) -> dict[str, Any]: + """Builds reference fabric from USGS Reference Hydrofabric data + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The reference flowpaths and divides in memory + """ + ti = cast(TaskInstance, context["ti"]) + cfg = cast(ReferenceConfig, context["config"]) + graph: rx.PyDiGraph = ti.xcom_pull(task_id="build_usgs_hf_graphs", key="graph") + node_indices: dict[str, int] = ti.xcom_pull(task_id="build_usgs_hf_graphs", key="node_indices") + _flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="usgs_flowpaths") + _catchments: pl.DataFrame = ti.xcom_pull(task_id="download", key="usgs_divides") + + # Check for cycles + cycles_iter = rx.simple_cycles(graph) + cycles: list[list[str]] = [] + cycle_ids: set[str] = set() + for cycle in cycles_iter: + _ids: list[Any] = [graph.get_node_data(node_idx) for node_idx in cycle] + cycles.append(_ids) + cycle_ids.update(_ids) + if cycle_ids: + raise NotImplementedError("Cycle Detected. Please create method for removing") + + _flowpaths_df = gpd.GeoDataFrame( + _flowpaths.select( + [pl.col("hydroseq"), pl.col("dnhydroseq"), pl.col("streamorde"), pl.col("comid")] + ).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_flowpaths["geometry"]), + crs="EPSG:4326", + ) + + _flowpaths_df_projected = _flowpaths_df.to_crs(cfg.crs) + _flowpaths_df["LengthKM"] = _flowpaths_df_projected.geometry.length / 1000 + + _catchments_df = gpd.GeoDataFrame( + _catchments.select([pl.col("COMID")]).to_pandas(), + geometry=gpd.GeoSeries.from_wkb(_catchments["geometry"]), + crs="EPSG:4326", + ) + _catchments_df_projected = _catchments_df.to_crs(cfg.crs) + _catchments_df["areasqkm"] = _catchments_df_projected.geometry.area / 1e6 + + # Log any flowpaths without matching catchments + flowpath_linkno_set = set(_flowpaths_df["comid"].tolist()) + catchment_linkno_set = set(_catchments_df["COMID"].tolist()) + missing_catchments = flowpath_linkno_set - catchment_linkno_set + if missing_catchments: + logger.warning( + f"build_usgs_hf_reference task: {len(missing_catchments)} flowpaths have no matching catchment" + ) + + reference_flowpaths = _trace_usgs_hf_attributes( + graph, node_indices, _flowpaths_df, _catchments_df, cfg.vpu_id + ) + reference_divides = _create_usgs_hf_reference_divides(_catchments_df, reference_flowpaths, cfg.vpu_id) + + return {"reference_flowpaths": reference_flowpaths, "reference_divides": reference_divides} diff --git a/src/reference_builds/pipeline/download.py b/src/reference_builds/pipeline/download.py new file mode 100644 index 0000000..db95018 --- /dev/null +++ b/src/reference_builds/pipeline/download.py @@ -0,0 +1,264 @@ +"""Contains all code for downloading hydrofabric data""" + +import logging +from pathlib import Path +from typing import Any, cast + +import geopandas as gpd +import pandas as pd +import polars as pl +from shapely.geometry import MultiLineString + +from reference_builds.configs import ReferenceConfig +from reference_builds.utils import _validate_and_fix_geometries +from reference_builds.utils.geometries import _fix_divide_exclaves + +logger = logging.getLogger(__name__) + + +def _load_and_concat_layers(gpkg_files: list[Path], layer_name: str | None) -> gpd.GeoDataFrame: + """Load a specific layer from all gpkg files and concatenate.""" + gdfs = [] + for gpkg_path in gpkg_files: + if layer_name is None: + gdf = gpd.read_file(gpkg_path, driver="GPKG") + else: + gdf = gpd.read_file(gpkg_path, layer=layer_name) + gdfs.append(gdf) + return pd.concat(gdfs, ignore_index=True) + + +def _load_and_concat_parquet(parquet_files: list[Path]) -> gpd.GeoDataFrame: + """Load a specific layer from all parquet files and concatenate.""" + gdfs = [] + for parquet_path in parquet_files: + gdf = gpd.read_parquet(parquet_path) + gdfs.append(gdf) + return pd.concat(gdfs, ignore_index=True) + + +def _merge_flowpaths_without_catchments( + flowpaths: gpd.GeoDataFrame, + catchments: gpd.GeoDataFrame, + connectivity: pd.DataFrame, +) -> tuple[gpd.GeoDataFrame, pd.DataFrame]: + """Merge flowpaths without catchments to its downstream neighbor. + + FIXME: This currently breaks connectivity info as well as path length calculations among other things. Needs significant rework. + + Parameters + ---------- + flowpaths : gpd.GeoDataFrame + The flowpath geodataframe, must contain 'DnHydroSeq' and 'HydroSeq' columns + for upstream/downstream connectivity and 'NHDPlusID' to match with its catchment. + catchments : gpd.GeoDataFrame + The catchment geodataframe, must contain 'NHDPlusID' column + + Returns + ------- + tuple[gpd.GeoDataFrame, pd.DataFrame] + The updated flowpath geodataframe and connectivity dataframe after merging flowpaths without catchments + """ + _flowpaths = flowpaths.copy() + _rename_mapping = {key: f"{key}_VAA" for key in connectivity.columns if key != "NHDPlusID"} + _connectivity = connectivity.copy().rename(columns=_rename_mapping) + + # merge flowpaths with connectivity info + n_connect = _flowpaths["NHDPlusID"].isin(_connectivity["NHDPlusID"]).sum() + n_flowpaths = len(_flowpaths.index) + if n_connect != n_flowpaths: + n_missing = n_flowpaths - n_connect + logger.warning( + f"{n_missing}/{n_flowpaths} flowpaths are missing connectivity info. These flowpaths will be dropped from the reference build." + ) + + _flowpaths = _flowpaths.merge(_connectivity, on="NHDPlusID", how="inner") + + # identify flowpaths without catchments + _has_catchment = _flowpaths["NHDPlusID"].isin(catchments["NHDPlusID"]) + _flowpaths_without_catchments = _flowpaths[~_has_catchment] + + logger.info(f"Merging {len(_flowpaths_without_catchments)} flowpaths without catchments") + + for _, row in _flowpaths_without_catchments.iterrows(): + hydroseq = row["HydroSeq_VAA"] + dnhydroseq = row["DnHydroSeq_VAA"] + downstream_flowpath = _flowpaths[_flowpaths["HydroSeq_VAA"] == dnhydroseq] + + if len(downstream_flowpath) == 0 or row["TerminalFl_VAA"] == 1: + continue + elif len(downstream_flowpath) > 1: + logger.warning(f"Multiple downstream flowpaths found for {row['NHDPlusID']}") + continue + else: + # merge geometry + _gdf = gpd.GeoDataFrame([row, downstream_flowpath.iloc[0]], geometry="geometry") + merged_geom = _gdf.geometry.union_all() + # update geometry of downstream flowpath + dn_idx = downstream_flowpath.index[0] + _flowpaths.at[dn_idx, "geometry"] = merged_geom + # update connectivity + upstream_flowpaths = _flowpaths[_flowpaths["DnHydroSeq_VAA"] == hydroseq] + if len(upstream_flowpaths) >= 1: + _flowpaths.loc[upstream_flowpaths.index, "DnHydroSeq_VAA"] = dnhydroseq + _flowpaths.loc[dn_idx, "FromNode_VAA"] = upstream_flowpaths.iloc[0]["ToNode_VAA"] + # TODO: check if other attributes need to be updated?? + + # drop flowpaths without catchments after merging + _flowpaths = _flowpaths[_has_catchment] + + # derive updated connectivity from merged flowpaths + connectivity = _flowpaths[list(_rename_mapping.values())].rename( + columns={v: k for k, v in _rename_mapping.items()} + ) + connectivity["NHDPlusID"] = _flowpaths["NHDPlusID"] + flowpaths = _flowpaths.drop(columns=list(_rename_mapping.values())) + + return flowpaths, connectivity + + +def download_geoglows_data(**context: dict[str, Any]) -> dict[str, pl.DataFrame]: + """Opens local / downloads for the reference-build process + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, pl.DataFrame] + The reference flowpath and divides references in memory + """ + cfg = cast(ReferenceConfig, context["config"]) + + # find the gpkg files from the ScienceBase NHD folders + gpkg_files = list(cfg.output_dir.glob(cfg.input_file_regex)) + + assert cfg.geoglows_catchment_regex is not None, "Need to specify where the catchment parquet files are" + parquet_files = list(cfg.output_dir.glob(cfg.geoglows_catchment_regex)) + # load layers + __flowpaths = _load_and_concat_layers(gpkg_files, layer_name=None) + __catchments = _load_and_concat_parquet(parquet_files) + # filter/validate layers + flowpaths = _validate_and_fix_geometries(__flowpaths, geom_type="flowpaths") + catchments = _validate_and_fix_geometries(__catchments, geom_type="divides") + + return { + "geoglows_flowpaths": pl.from_pandas(flowpaths.to_wkb()), + "geoglows_divides": pl.from_pandas(catchments.to_wkb()), + } + + +def download_nhd_data(**context: dict[str, Any]) -> dict[str, pl.DataFrame]: + """Opens local / downloads for the reference-build process + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, pl.DataFrame] + The reference flowpath and divides references in memory + """ + cfg = cast(ReferenceConfig, context["config"]) + + # find the gpkg files from the ScienceBase NHD folders + matching_folders = list(cfg.output_dir.glob(cfg.input_file_regex)) + gpkg_files: list[Path] = [] + for folder in matching_folders: + if folder.is_dir(): + gpkg_files.extend(folder.glob("*.gpkg")) + + # load layers + layers = [ + "NHDFlowline", + "NHDPlusCatchment", + "NHDPlusFlowlineVAA", + ] + data = {layer: _load_and_concat_layers(gpkg_files, layer) for layer in layers} + + # filter/validate layers + _flowpaths = _validate_and_fix_geometries(data["NHDFlowline"], geom_type="flowpaths") + catchments = _validate_and_fix_geometries(data["NHDPlusCatchment"], geom_type="divides") + + _flowpaths_with_catchments = _flowpaths[_flowpaths["NHDPlusID"].isin(catchments["NHDPlusID"])] + flowpaths = _flowpaths_with_catchments[ + _flowpaths_with_catchments["fcode_description"].isin(cfg.permitted_fcodes) + ] + + return { + "nhd_flowpaths": pl.from_pandas(flowpaths.to_wkb()), + "nhd_divides": pl.from_pandas(catchments.to_wkb()), + "nhd_connectivity": pl.from_pandas(data["NHDPlusFlowlineVAA"]), + } + + +def download_usgs_hf_data(**context: dict[str, Any]) -> dict[str, pl.DataFrame]: + """Opens local / downloads for the reference-build process + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, pl.DataFrame] + The reference flowpath and divides references in memory + """ + cfg = cast(ReferenceConfig, context["config"]) + + # find the gpkg files + gpkg_files = list(cfg.output_dir.glob(cfg.input_file_regex)) + + # load layers + flowpaths = _load_and_concat_layers(gpkg_files, layer_name="reference_flowline").to_crs("EPSG:4326") + + flowpaths["geometry"] = flowpaths["geometry"].apply( + lambda x: x if x.geom_type == "MultiLineString" else MultiLineString([x]) + ) + + valid_hydroseq = set(flowpaths["hydroseq"].unique()) + flowpaths.loc[~flowpaths["dnhydroseq"].isin(valid_hydroseq), "dnhydroseq"] = 0 + + catchments = _load_and_concat_layers(gpkg_files, layer_name="reference_catchments").to_crs("EPSG:4326") + + flowpaths["dnhydroseq"] = flowpaths["dnhydroseq"].fillna(0) + + hydroseq_lookup = flowpaths.set_index("comid")["hydroseq"].to_dict() + + flowpaths["comid"] = flowpaths["comid"].map(hydroseq_lookup) + catchments["COMID"] = catchments["COMID"].map(hydroseq_lookup) + + # filter/validate layers + _flowpaths = _validate_and_fix_geometries(flowpaths, geom_type="flowpaths") + _flowpaths = _flowpaths[_flowpaths["comid"].isin(catchments["COMID"])] + + catchments = _validate_and_fix_geometries(catchments, geom_type="divides") + catchments = _fix_divide_exclaves(catchments.to_crs("EPSG:3338")).to_crs("EPSG:4326") + + return { + "usgs_flowpaths": pl.from_pandas(_flowpaths.to_wkb()), + "usgs_divides": pl.from_pandas(catchments.to_wkb()), + } diff --git a/src/reference_builds/pipeline/processing.py b/src/reference_builds/pipeline/processing.py new file mode 100644 index 0000000..55a3073 --- /dev/null +++ b/src/reference_builds/pipeline/processing.py @@ -0,0 +1,125 @@ +"""Contains all pipeline code for processing reference data""" + +import logging +from typing import Any, cast + +import polars as pl +import rustworkx as rx + +from reference_builds.task_instance import TaskInstance +from reference_builds.utils.geoglows_graph import _build_geoglows_graph +from reference_builds.utils.nhd_graph import _build_graph +from reference_builds.utils.usgs_graph import _build_usgs_hf_graph + +logger = logging.getLogger(__name__) + + +def _build_rustworkx_object( + upstream_network: dict[str, list[str]] | dict[int, list[int]], +) -> tuple[rx.PyDiGraph, dict[str, int] | dict[int, int]]: + """Build a RustWorkX directed graph from upstream network dictionary. + + Parameters + ---------- + upstream_network : dict[str, list[str]] | dict[int, list[int]] + Dictionary mapping downstream flowpath IDs to lists of upstream flowpath IDs + + Returns + ------- + tuple[rx.PyDiGraph, dict[str, int] | dict[int, int]] + The flowpaths object in graph form and node indices for each object in the graph + """ + graph = rx.PyDiGraph(check_cycle=True) + node_indices: dict[Any, int] = {} + if None in upstream_network: + upstream_network.pop(None) # type: ignore + for to_edge in sorted(upstream_network.keys()): + from_edges = upstream_network[to_edge] # type: ignore + if to_edge not in node_indices: + node_indices[to_edge] = graph.add_node(to_edge) + for from_edge in from_edges: + if from_edge not in node_indices: + node_indices[from_edge] = graph.add_node(from_edge) + for to_edge, from_edges in upstream_network.items(): + for from_edge in from_edges: + graph.add_edge(node_indices[from_edge], node_indices[to_edge], None) + return graph, node_indices + + +def build_nhd_graphs(**context: dict[str, Any]) -> dict[str, Any]: + """Builds and processes graphs from NHD data + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The rustworkx graph object and node_indices for the NHD + """ + ti = cast(TaskInstance, context["ti"]) + flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="nhd_flowpaths") + connectivity: pl.DataFrame = ti.xcom_pull(task_id="download", key="nhd_connectivity") + upstream_network = _build_graph(connectivity, flowpaths) + graph, node_indices = _build_rustworkx_object(upstream_network) + return {"graph": graph, "node_indices": node_indices} + + +def build_geoglows_graphs(**context: dict[str, Any]) -> dict[str, Any]: + """Builds and processes graphs from NHD data + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The rustworkx graph object and node_indices for the NHD + """ + ti = cast(TaskInstance, context["ti"]) + flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="geoglows_flowpaths") + upstream_network = _build_geoglows_graph(flowpaths) + graph, node_indices = _build_rustworkx_object(upstream_network) + return {"graph": graph, "node_indices": node_indices} + + +def build_usgs_hf_graphs(**context: dict[str, Any]) -> dict[str, Any]: + """Builds and processes graphs from NHD data + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The rustworkx graph object and node_indices for the NHD + """ + ti = cast(TaskInstance, context["ti"]) + flowpaths: pl.DataFrame = ti.xcom_pull(task_id="download", key="usgs_flowpaths") + upstream_network = _build_usgs_hf_graph(flowpaths) + graph, node_indices = _build_rustworkx_object(upstream_network) + return {"graph": graph, "node_indices": node_indices} diff --git a/src/reference_builds/pipeline/write.py b/src/reference_builds/pipeline/write.py new file mode 100644 index 0000000..7152e31 --- /dev/null +++ b/src/reference_builds/pipeline/write.py @@ -0,0 +1,50 @@ +"""Contains all code for downloading hydrofabric data""" + +import logging +from typing import Any, cast + +from reference_builds.configs import ReferenceConfig +from reference_builds.task_instance import TaskInstance + +logger = logging.getLogger(__name__) + + +def write_reference(**context: dict[str, Any]) -> dict[str, Any]: + """Opens local / downloads for the reference-build process + + Parameters + ---------- + **context : dict + Airflow-compatible context containing: + - ti : TaskInstance for XCom operations + - config : HFConfig with pipeline configuration + - task_id : str identifier for this task + - run_id : str identifier for this pipeline run + - ds : str execution date + - execution_date : datetime object + + Returns + ------- + dict[str, Any] + The reference flowpath and divides references in memory + """ + cfg = cast(ReferenceConfig, context["config"]) + ti = cast(TaskInstance, context["ti"]) + cfg.output_reference_flowpaths_path.unlink(missing_ok=True) # deletes files that exist with the same name + cfg.output_reference_divides_path.unlink(missing_ok=True) # deletes files that exist with the same name + + final_flowpaths = ti.xcom_pull(task_id="build_reference", key="reference_flowpaths") + final_divides = ti.xcom_pull(task_id="build_reference", key="reference_divides") + + final_flowpaths = final_flowpaths.to_crs(cfg.crs) + final_divides = final_divides.to_crs(cfg.crs) + + if cfg.write_gpkg: + cfg.output_reference_gpkg_path.unlink(missing_ok=True) + final_flowpaths.to_file(cfg.output_reference_gpkg_path, layer="reference_flowpaths", driver="GPKG") + final_divides.to_file(cfg.output_reference_gpkg_path, layer="reference_divides", driver="GPKG") + logger.info(f"write_reference task: wrote geopackage layers to {cfg.output_reference_gpkg_path}") + + final_flowpaths.to_parquet(cfg.output_reference_flowpaths_path) + final_divides.to_parquet(cfg.output_reference_divides_path) + return {} diff --git a/src/reference_builds/task_instance.py b/src/reference_builds/task_instance.py new file mode 100644 index 0000000..9d3f795 --- /dev/null +++ b/src/reference_builds/task_instance.py @@ -0,0 +1,42 @@ +"""Mocks a local instance for building runners""" + +from typing import Any + + +class TaskInstance: + """A Mock TaskInstance for local runners similar to Apache Airflow.""" + + def __init__(self) -> None: + """Initialize the TaskInstance with empty XCom storage.""" + self.xcom_storage: dict[str, Any] = {} + + def xcom_push(self, key: str, value: Any) -> None: + """ + Store a value in XCom for retrieval by downstream tasks. + + Parameters + ---------- + key : str + Unique identifier for the stored value. Convention is to use '{task_id}.{key_name}' format for namespacing. + value : Any + The data to store. Can be any Python object. + """ + self.xcom_storage[key] = value + + def xcom_pull(self, task_id: str, key: str = "return_value") -> Any: + """ + Retrieve a value from XCom that was pushed by an upstream task. + + Parameters + ---------- + task_id : str + The task_id of the task that pushed the value. + key : str, default='return_value' + The key used when the value was pushed. Default 'return_value' is used for values returned from task functions. + + Returns + ------- + Any + The stored value, or None if the key doesn't exist. + """ + return self.xcom_storage.get(f"{task_id}.{key}") diff --git a/src/reference_builds/utils/__init__.py b/src/reference_builds/utils/__init__.py new file mode 100644 index 0000000..7a6e1e2 --- /dev/null +++ b/src/reference_builds/utils/__init__.py @@ -0,0 +1,4 @@ +from .geometries import _validate_and_fix_geometries +from .v22_graph import build_v22_graph + +__all__ = ["_validate_and_fix_geometries", "v22_graph"] diff --git a/src/reference_builds/utils/geoglows_graph.py b/src/reference_builds/utils/geoglows_graph.py new file mode 100644 index 0000000..66f7089 --- /dev/null +++ b/src/reference_builds/utils/geoglows_graph.py @@ -0,0 +1,52 @@ +"""Contains all code for processing nhd data""" + +import logging + +import polars as pl + +logger = logging.getLogger(__name__) + + +def _build_geoglows_graph(flowpaths: pl.DataFrame) -> dict[str, list[str]]: + """Build a graph of upstream flowpath connections from GeoGLOWS data. + + Parameters + ---------- + flowpaths : pl.DataFrame + The GeoGLOWS flowpaths with LINKNO and DSLINKNO columns + + Returns + ------- + dict[str, list[str]] + The upstream dictionary containing upstream and downstream connections + Key is the downstream flowpath ID, values are the upstream flowpath IDs + """ + # Filter out terminal links (DSLINKNO == -1) for building upstream connections + connectivity = flowpaths.select( + [ + pl.col("LINKNO").cast(pl.Int64), + pl.col("DSLINKNO").cast(pl.Int64), + ] + ).filter(pl.col("DSLINKNO") != -1) + + # Build upstream network: group by downstream link to get all upstream links + upstream_network_df = connectivity.group_by(pl.col("DSLINKNO").cast(pl.Utf8).alias("downstream_id")).agg( + pl.col("LINKNO").cast(pl.Utf8).alias("upstream_list") + ) + + upstream_dict: dict[str, list[str]] = dict( + zip( + upstream_network_df["downstream_id"].to_list(), + upstream_network_df["upstream_list"].to_list(), + strict=False, + ) + ) + + # Ensure all flowpath IDs are in the dictionary (even those with no upstream) + all_flowpath_ids = flowpaths.select(pl.col("LINKNO").cast(pl.Utf8))["LINKNO"].to_list() + + for fp_id in all_flowpath_ids: + if fp_id not in upstream_dict: + upstream_dict[fp_id] = [] + + return upstream_dict diff --git a/src/reference_builds/utils/geometries.py b/src/reference_builds/utils/geometries.py new file mode 100644 index 0000000..58d9ce9 --- /dev/null +++ b/src/reference_builds/utils/geometries.py @@ -0,0 +1,185 @@ +"""A file for all geometry related internal functions""" + +import logging + +import geopandas as gpd +import pandas as pd +from shapely import Geometry, wkb +from shapely.geometry import LineString, MultiLineString, MultiPolygon, Point + +logger = logging.getLogger(__name__) + + +def _ensure_geometry(geom): # type: ignore[no-untyped-def] + """Convert bytes to Shapely geometry if needed.""" + if isinstance(geom, bytes): + return wkb.loads(geom) + return geom + + +def _get_endpoints(geom): # type: ignore[no-untyped-def] + """Get start and end points of a line or multiline geometry.""" + geom = _ensure_geometry(geom) + if geom.geom_type == "MultiLineString": + start_coord = list(geom.geoms)[0].coords[0] + end_coord = list(geom.geoms)[-1].coords[-1] + else: + start_coord = geom.coords[0] + end_coord = geom.coords[-1] + return start_coord, end_coord + + +def _reverse_line(geom): # type: ignore[no-untyped-def] + """Reverse a LineString or MultiLineString.""" + geom = _ensure_geometry(geom) + if geom.geom_type == "MultiLineString": + # Reverse each component and reverse the order of components + reversed_parts = [LineString(part.coords[::-1]) for part in reversed(geom.geoms)] + return MultiLineString(reversed_parts) + else: + return LineString(geom.coords[::-1]) + + +def _orient_flowpath_downstream(geom, ds_geom=None, us_geom=None): # type: ignore[no-untyped-def] + """ + Orient a flowpath so coords go from upstream to downstream. + + Parameters + ---------- + geom : geometry + The flowpath geometry + ds_geom : geometry, optional + The downstream flowpath geometry for direction detection + us_geom : geometry, optional + The upstream flowpath geometry (used for outlets when ds_geom is None) + + Returns + ------- + geometry + The oriented geometry (start = upstream, end = downstream) + """ + geom = _ensure_geometry(geom) + start_coord, end_coord = _get_endpoints(geom) + start_pt = Point(start_coord) + end_pt = Point(end_coord) + + # Primary: use downstream geometry + if ds_geom is not None: + ds_geom = _ensure_geometry(ds_geom) + dist_start = start_pt.distance(ds_geom) + dist_end = end_pt.distance(ds_geom) + + # If start is closer to downstream, the line is reversed + if dist_start < dist_end: + return _reverse_line(geom) + return geom + + # Fallback for outlets: use upstream geometry + if us_geom is not None: + us_geom = _ensure_geometry(us_geom) + dist_start = start_pt.distance(us_geom) + dist_end = end_pt.distance(us_geom) + + # Upstream end should be CLOSER to upstream geometry + # So if end is closer to upstream, line is reversed + if dist_end < dist_start: + return _reverse_line(geom) + return geom + + # No reference at all - return as-is + return geom + + +def _drop_exclaves(geom: Geometry) -> Geometry: + """Find and destroy non-contiguous parts of MultiPolygons""" + if geom.geom_type != "MultiPolygon": + return geom + main_part = geom.geoms[0] + for part in geom.geoms: + if part.area > main_part.area: + main_part = part + main_part_buffered = main_part.buffer(1.0) + for part in geom.geoms: + if part.intersects(main_part_buffered): + main_part = main_part.union(part) + return main_part + + +def _find_exclaves(geom: Geometry) -> pd.Series: + """Find and exclude non-contiguous parts of MultiPolygons, appending them to a list to be resolved later""" + exclaves = [] + if geom.geom_type != "MultiPolygon": + return pd.Series(data={"geometry": geom, "exclaves": exclaves}, index=["geometry", "exclaves"]) + + main_part = geom.geoms[0] + for part in geom.geoms: + if part.area > main_part.area: + main_part = part + included_parts = [main_part] + main_part_buffered = main_part.buffer(1.0) + for part in geom.geoms: + if part.intersects(main_part_buffered): + included_parts.append(part) + main_part = MultiPolygon(included_parts) + else: + exclaves.append(part) + return pd.Series(data={"geometry": main_part, "exclaves": exclaves}, index=["geometry", "exclaves"]) + + +def _fix_divide_exclaves(gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame: + """Remove exclaves from catchment geometries, merging them in with the neighbor it intersects the most with""" + exclaves = gdf["geometry"].apply(_find_exclaves) + + fixed_count = 0 + + gdf["geometry"] = exclaves["geometry"] + for idx, e in exclaves["exclaves"].items(): + for exclave in e: + buffer = exclave.buffer(1.0) + intersection_areas = gdf.intersection(buffer).area + intersection_areas[idx] = 0.0 + best_idx = intersection_areas.argmax() + gdf.loc[best_idx, "geometry"] = gdf["geometry"][best_idx].union(exclave) + fixed_count += 1 + + logger.info(f"fix_divide_exclaves: fixed/merged {fixed_count} polygons") + return gdf + + +def _validate_and_fix_geometries(gdf: gpd.GeoDataFrame, geom_type: str) -> gpd.GeoDataFrame: + """Validate and fix invalid geometries in a GeoDataFrame. + + Parameters + ---------- + gdf : gpd.GeoDataFrame + GeoDataFrame to validate + geom_type : str + Description for logging (e.g., "flowpaths", "divides") + + Returns + ------- + gpd.GeoDataFrame + GeoDataFrame with fixed geometries + + Raises + ------ + ValueError + If geometries cannot be fixed or invalid geometries remain + """ + invalid_mask = ~gdf.geometry.is_valid + invalid_count = invalid_mask.sum() + + if invalid_count == 0: + return gdf # No invalid geometries + + geometries = gdf[invalid_mask].geometry + gdf.loc[invalid_mask, "geometry"] = geometries.make_valid(method="structure") + + if len(gdf[~gdf.geometry.is_valid]) > 0: + raise ValueError(f"Could not fix invalid geometries in {geom_type}") + + still_invalid = (~gdf.geometry.is_valid).sum() + if still_invalid > 0: + raise ValueError(f"Invalid Geometries remain: {gdf[~gdf.geometry.is_valid]}") + + return gdf diff --git a/src/reference_builds/utils/nhd_graph.py b/src/reference_builds/utils/nhd_graph.py new file mode 100644 index 0000000..67f18b4 --- /dev/null +++ b/src/reference_builds/utils/nhd_graph.py @@ -0,0 +1,72 @@ +"""Contains all code for processing nhd data""" + +import logging + +import polars as pl + +logger = logging.getLogger(__name__) + + +def _build_graph(connectivity: pl.DataFrame, flowpaths: pl.DataFrame) -> dict[str, list[str]]: + """Build a graph of upstream flowpath connections. + + Parameters + ---------- + connectivity : pl.DataFrame + The connectivity/flow table with FromNode and ToNode columns + flowpaths : pl.DataFrame + The reference flowpaths to filter to + + Returns + ------- + dict[str, list[str]] + The upstream dictionary containing upstream and downstream connections + Key is the downstream flowpath ID, values are the upstream flowpath IDs + """ + valid_ids = flowpaths.select(pl.col("NHDPlusID").cast(pl.Int64))["NHDPlusID"] + + filtered_connectivity = connectivity.select( + [ + pl.col("NHDPlusID").cast(pl.Int64), + pl.col("FromNode").cast(pl.Int64), + pl.col("ToNode").cast(pl.Int64), + ] + ).filter(pl.col("NHDPlusID").is_in(valid_ids)) + + tonode_lookup = filtered_connectivity.select( + [ + pl.col("ToNode"), + pl.col("NHDPlusID").cast(pl.Utf8).alias("upstream_id"), + ] + ) + + fromnode_lookup = filtered_connectivity.select( + [ + pl.col("FromNode"), + pl.col("NHDPlusID").cast(pl.Utf8).alias("downstream_id"), + ] + ) + + merged = tonode_lookup.join(fromnode_lookup, left_on="ToNode", right_on="FromNode", how="inner").select( + ["upstream_id", "downstream_id"] + ) + + upstream_network_df = merged.group_by("downstream_id").agg(pl.col("upstream_id").alias("upstream_list")) + + upstream_dict: dict[str, list[str]] = dict( + zip( + upstream_network_df["downstream_id"].to_list(), + upstream_network_df["upstream_list"].to_list(), + strict=False, + ) + ) + + all_flowpath_ids = flowpaths.select(pl.col("NHDPlusID").cast(pl.Int64).cast(pl.Utf8))[ + "NHDPlusID" + ].to_list() + + for fp_id in all_flowpath_ids: + if fp_id not in upstream_dict: + upstream_dict[fp_id] = [] + + return upstream_dict diff --git a/src/reference_builds/utils/usgs_graph.py b/src/reference_builds/utils/usgs_graph.py new file mode 100644 index 0000000..92f531d --- /dev/null +++ b/src/reference_builds/utils/usgs_graph.py @@ -0,0 +1,52 @@ +"""Contains all code for processing USGS Reference Hydrofabric data""" + +import logging + +import polars as pl + +logger = logging.getLogger(__name__) + + +def _build_usgs_hf_graph(flowpaths: pl.DataFrame) -> dict[str, list[str]]: + """Build a graph of upstream flowpath connections from USGS Reference Hydrofabric data. + + Parameters + ---------- + flowpaths : pl.DataFrame + The USGS Reference Hydrofabric flowpaths with hydroseq and dnhydroseq columns + + Returns + ------- + dict[str, list[str]] + The upstream dictionary containing upstream and downstream connections + Key is the downstream flowpath ID, values are the upstream flowpath IDs + """ + # Filter out terminal links (dnhydroseq == 0) for building upstream connections + connectivity = flowpaths.select( + [ + pl.col("hydroseq").cast(pl.Int64), + pl.col("dnhydroseq").cast(pl.Int64), + ] + ).filter(pl.col("dnhydroseq") != 0) + + # Build upstream network: group by downstream link to get all upstream links + upstream_network_df = connectivity.group_by( + pl.col("dnhydroseq").cast(pl.Utf8).alias("downstream_id") + ).agg(pl.col("hydroseq").cast(pl.Utf8).alias("upstream_list")) + + upstream_dict: dict[str, list[str]] = dict( + zip( + upstream_network_df["downstream_id"].to_list(), + upstream_network_df["upstream_list"].to_list(), + strict=False, + ) + ) + + # Ensure all flowpath IDs are in the dictionary (even those with no upstream) + all_flowpath_ids = flowpaths.select(pl.col("hydroseq").cast(pl.Utf8))["hydroseq"].to_list() + + for fp_id in all_flowpath_ids: + if fp_id not in upstream_dict: + upstream_dict[fp_id] = [] + + return upstream_dict diff --git a/src/reference_builds/utils/v22_graph.py b/src/reference_builds/utils/v22_graph.py new file mode 100644 index 0000000..d77f3fb --- /dev/null +++ b/src/reference_builds/utils/v22_graph.py @@ -0,0 +1,109 @@ +""" +Builds a graph object from the v2.2 Hydrofabric + +Kudos to Nels Fraizer for assistance in making the graph building code: +https://github.com/DeepGroundwater/ddr/blob/ab4c3962c2c119e6a9182a77f2a9faceec19f2e0/engine/adjacency.py +""" + +import sqlite3 +from pathlib import Path +from typing import Any + +import polars as pl +import rustworkx as rx +from tqdm import tqdm + + +def _find_outlets_by_hydroseq(reference_flowpaths: pl.DataFrame) -> list[str]: + """Find outlets for the river using hydroseq. + + Parameters + ---------- + reference_flowpaths : pl.DataFrame + The flowpath reference + + Returns + ------- + list[str] + All outlets from the reference + """ + df_pl = reference_flowpaths.select(pl.col(["flowpath_id", "hydroseq", "dnhydroseq", "totdasqkm"])) + + df_with_str_id = df_pl.with_columns( + pl.col("flowpath_id").cast(pl.Float64).cast(pl.Int64).cast(pl.Utf8).alias("flowpath_id_str") + ) + + hydroseq_set: set[Any] = set(df_pl["hydroseq"].to_list()) + + outlets_df = df_with_str_id.filter( + (pl.col("dnhydroseq") == 0) | ~pl.col("dnhydroseq").is_in(hydroseq_set) + ).sort("flowpath_id_str") # dnhydroseq is 0, or doesn't exist in hydroseq + + # outlets_sorted = outlets_df.sort("totdasqkm", descending=True) # Commenting out until production + outlets: list[str] = outlets_df["flowpath_id_str"].to_list() + + return outlets + + +def create_matrix(fp: pl.LazyFrame, network: pl.LazyFrame) -> tuple[rx.PyDiGraph, dict[str, int]]: + """ + Create a directed graph from flowpaths and network dataframes. + + Parameters + ---------- + fp : pl.LazyFrame + Flowpaths dataframe with 'toid' column indicating downstream nexus IDs. + network : pl.LazyFrame + Network dataframe with 'toid' column indicating downstream flowpath IDs. + + Returns + ------- + tuple[rx.PyDiGraph, dict[str, int]] + tuple[0]: A rustworkx directed graph + tuple[1]: Mapping of flowpath IDs to graph node indices + """ + fp = fp.with_row_index(name="idx").collect() + network = network.collect().unique(subset=["id"]) + _values = zip(fp["idx"], fp["toid"], strict=False) + fp = dict(zip(fp["id"], _values, strict=True)) + network = dict(zip(network["id"], network["toid"], strict=True)) + + graph = rx.PyDiGraph(check_cycle=False, node_count_hint=len(fp), edge_count_hint=len(fp)) + gidx = graph.add_nodes_from(fp.keys()) + + # Create mapping from flowpath ID to graph node index + node_index = {graph.get_node_data(idx): idx for idx in gidx} + + for idx in tqdm(gidx, desc="Building network graph"): + id = graph.get_node_data(idx) + nex = fp[id][1] # the downstream nexus id + ds_wb = network.get(nex) + if ds_wb is not None: + graph.add_edge(idx, node_index[ds_wb], nex) + + return graph, node_index + + +def build_v22_graph(file_path: Path) -> tuple[rx.PyDiGraph, dict[str, int]]: + """Builds a graph from the v2.2 hydrofabric + + Parameters + ---------- + file_path : Path + The path to the v2.2 geopackage + + Returns + ------- + tuple[rx.PyDiGraph, dict[str, int]] + tuple[0]: A rustworkx directed graph + tuple[1]: Mapping of flowpath IDs to graph node indices + """ + # Read hydrofabric geopackage using sqlite + query = "SELECT id,toid FROM flowpaths" + conn = sqlite3.connect(file_path) + fp = pl.read_database(query=query, connection=conn) + fp = fp.extend(pl.DataFrame({"id": ["wb-0"], "toid": [None]})).lazy() + query = "SELECT id,toid FROM network" + network = pl.read_database(query=query, connection=conn).lazy() + network = network.filter(pl.col("id").str.starts_with("wb-").not_()) + return create_matrix(fp, network) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7d8ea0a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,304 @@ +"""Conftests for Pytest Suite""" + +import geopandas as gpd +import pytest +import rustworkx as rx +from shapely.geometry import LineString, MultiLineString, MultiPolygon, Polygon + + +@pytest.fixture +def sample_graph() -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a simple directed graph for testing. + + Graph structure: + 1 -> 3 + 2 -> 3 + 3 -> 4 + 4 -> 5 (outlet) + """ + graph = rx.PyDiGraph() + + # Add nodes (using NHDPlusID as string) + node_data = [ + "85000100000001", + "85000100000002", + "85000100000003", + "85000100000004", + "85000100000005", + ] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + # Add edges (upstream -> downstream) + graph.add_edge(node_indices["85000100000001"], node_indices["85000100000003"], None) + graph.add_edge(node_indices["85000100000002"], node_indices["85000100000003"], None) + graph.add_edge(node_indices["85000100000003"], node_indices["85000100000004"], None) + graph.add_edge(node_indices["85000100000004"], node_indices["85000100000005"], None) + + return graph, node_indices + + +@pytest.fixture +def sample_flowpaths() -> gpd.GeoDataFrame: + """Create sample flowpaths GeoDataFrame.""" + data = { + "NHDPlusID": [ + 85000100000001, + 85000100000002, + 85000100000003, + 85000100000004, + 85000100000005, + ], + "VPUID": ["2101", "2101", "2101", "2101", "2101"], + "LengthKM": [0.333, 0.5, 1.201, 0.182, 0.387], + "fcode_description": [ + "Stream/River: Hydrographic Category = Intermittent", + "Artificial Path", + "Stream/River: Hydrographic Category = Intermittent", + "Artificial Path", + "Stream/River: Hydrographic Category = Intermittent", + ], + } + + # Create simple linestring geometries + geometries = [MultiLineString([LineString([(0, i), (1, i)])]) for i in range(5)] + + return gpd.GeoDataFrame(data, geometry=geometries, crs="EPSG:4269") + + +@pytest.fixture +def sample_divides() -> gpd.GeoDataFrame: + """Create sample divides GeoDataFrame.""" + data = { + "NHDPlusID": [ + 85000100000001, + 85000100000002, + 85000100000003, + 85000100000004, + 85000100000005, + ], + "VPUID": ["2101", "2101", "2101", "2101", "2101"], + "AreaSqKm": [0.1602, 0.6949, 0.0248, 0.1413, 0.9395], + } + + # Create simple polygon geometries + geometries = [MultiPolygon([Polygon([(i, 0), (i + 1, 0), (i + 1, 1), (i, 1)])]) for i in range(5)] + + return gpd.GeoDataFrame(data, geometry=geometries, crs="EPSG:4269") + + +@pytest.fixture +def disconnected_graph() -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a graph with multiple disconnected subgraphs (multiple outlets). + + Subgraph 1: + 1 -> 2 (outlet) + + Subgraph 2: + 3 -> 4 -> 5 (outlet) + """ + graph = rx.PyDiGraph() + + node_data = [ + "85000100000001", + "85000100000002", + "85000100000003", + "85000100000004", + "85000100000005", + ] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + # Subgraph 1 + graph.add_edge(node_indices["85000100000001"], node_indices["85000100000002"], None) + + # Subgraph 2 + graph.add_edge(node_indices["85000100000003"], node_indices["85000100000004"], None) + graph.add_edge(node_indices["85000100000004"], node_indices["85000100000005"], None) + + return graph, node_indices + + +@pytest.fixture +def sample_geoglows_graph() -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a sample GeoGLOWS graph for testing. + + Network topology: + 1 ──┐ + ├──> 3 ──> 4 ──> 5 (outlet) + 2 ──┘ + + Where: + - 1, 2 are headwaters (order 1) + - 3 is confluence (order 2) + - 4, 5 continue downstream (order 2) + """ + graph = rx.PyDiGraph() + + node_data = [ + "810000001", + "810000002", + "810000003", + "810000004", + "810000005", + ] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + # Add edges (upstream -> downstream) + graph.add_edge(node_indices["810000001"], node_indices["810000003"], None) + graph.add_edge(node_indices["810000002"], node_indices["810000003"], None) + graph.add_edge(node_indices["810000003"], node_indices["810000004"], None) + graph.add_edge(node_indices["810000004"], node_indices["810000005"], None) + + return graph, node_indices + + +@pytest.fixture +def sample_geoglows_flowpaths() -> gpd.GeoDataFrame: + """Create sample GeoGLOWS flowpaths GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "LINKNO": [810000001, 810000002, 810000003, 810000004, 810000005], + "DSLINKNO": [810000003, 810000003, 810000004, 810000005, -1], + "strmOrder": [1, 1, 2, 2, 2], + "LengthKM": [0.5, 0.6, 0.8, 0.7, 0.4], + }, + geometry=[ + LineString([(0, 0), (1, 1)]), + LineString([(2, 0), (1, 1)]), + LineString([(1, 1), (1, 2)]), + LineString([(1, 2), (1, 3)]), + LineString([(1, 3), (1, 4)]), + ], + crs="EPSG:3857", + ) + + +@pytest.fixture +def sample_geoglows_catchments() -> gpd.GeoDataFrame: + """Create sample GeoGLOWS catchments GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "linkno": [810000001, 810000002, 810000003, 810000004, 810000005], + "areasqkm": [1.0, 1.2, 0.8, 0.9, 0.6], + }, + geometry=[ + MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])]), + MultiPolygon([Polygon([(1, 0), (2, 0), (2, 1), (1, 1)])]), + MultiPolygon([Polygon([(0, 1), (2, 1), (2, 2), (0, 2)])]), + MultiPolygon([Polygon([(0, 2), (2, 2), (2, 3), (0, 3)])]), + MultiPolygon([Polygon([(0, 3), (2, 3), (2, 4), (0, 4)])]), + ], + crs="EPSG:3857", + ) + + +@pytest.fixture +def disconnected_geoglows_graph() -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a disconnected GeoGLOWS graph with two separate networks. + + Network 1: 1 -> 3 (outlet) + Network 2: 2 -> 4 -> 5 (outlet) + """ + graph = rx.PyDiGraph() + + node_data = [ + "810000001", + "810000002", + "810000003", + "810000004", + "810000005", + ] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + # Network 1 + graph.add_edge(node_indices["810000001"], node_indices["810000003"], None) + # Network 2 + graph.add_edge(node_indices["810000002"], node_indices["810000004"], None) + graph.add_edge(node_indices["810000004"], node_indices["810000005"], None) + + return graph, node_indices + + +@pytest.fixture +def reversed_flowpaths() -> gpd.GeoDataFrame: + """Create sample flowpaths with reversed digitization (downstream to upstream).""" + data = { + "NHDPlusID": [ + 85000100000001, + 85000100000002, + 85000100000003, + 85000100000004, + 85000100000005, + ], + "VPUID": ["2101", "2101", "2101", "2101", "2101"], + "LengthKM": [0.333, 0.5, 1.201, 0.182, 0.387], + "fcode_description": [ + "Stream/River", + "Artificial Path", + "Stream/River", + "Artificial Path", + "Stream/River", + ], + } + + # Create geometries that connect but are digitized downstream-to-upstream + # Network: 1,2 -> 3 -> 4 -> 5 (outlet) + geometries = [ + MultiLineString([LineString([(1, 3), (0, 4)])]), # 1: connects to 3 at (1,3) + MultiLineString([LineString([(1, 3), (2, 4)])]), # 2: connects to 3 at (1,3) + MultiLineString([LineString([(1, 2), (1, 3)])]), # 3: connects to 4 at (1,2) + MultiLineString([LineString([(1, 1), (1, 2)])]), # 4: connects to 5 at (1,1) + MultiLineString([LineString([(1, 0), (1, 1)])]), # 5: outlet at (1,0) + ] + + return gpd.GeoDataFrame(data, geometry=geometries, crs="EPSG:4269") + + +@pytest.fixture +def correctly_oriented_flowpaths() -> gpd.GeoDataFrame: + """Create sample flowpaths with correct digitization (upstream to downstream).""" + data = { + "NHDPlusID": [ + 85000100000001, + 85000100000002, + 85000100000003, + 85000100000004, + 85000100000005, + ], + "VPUID": ["2101", "2101", "2101", "2101", "2101"], + "LengthKM": [0.333, 0.5, 1.201, 0.182, 0.387], + "fcode_description": [ + "Stream/River", + "Artificial Path", + "Stream/River", + "Artificial Path", + "Stream/River", + ], + } + + # Create geometries digitized upstream-to-downstream + # Network: 1,2 -> 3 -> 4 -> 5 (outlet) + geometries = [ + MultiLineString([LineString([(0, 4), (1, 3)])]), # 1: ends at confluence (1,3) + MultiLineString([LineString([(2, 4), (1, 3)])]), # 2: ends at confluence (1,3) + MultiLineString([LineString([(1, 3), (1, 2)])]), # 3: from confluence to (1,2) + MultiLineString([LineString([(1, 2), (1, 1)])]), # 4: continues downstream + MultiLineString([LineString([(1, 1), (1, 0)])]), # 5: outlet ends at (1,0) + ] + + return gpd.GeoDataFrame(data, geometry=geometries, crs="EPSG:4269") diff --git a/tests/test_builds.py b/tests/test_builds.py new file mode 100644 index 0000000..4ed45ef --- /dev/null +++ b/tests/test_builds.py @@ -0,0 +1,946 @@ +"""Tests for build_nhd_reference module""" + +import geopandas as gpd +import numpy as np +import pytest +import rustworkx as rx +from shapely.geometry import LineString, MultiLineString, MultiPolygon, Point, Polygon + +from reference_builds.pipeline.build_reference import ( + _create_geoglows_reference_divides, + _create_reference_divides, + _trace_attributes, + _trace_geoglows_attributes, +) + + +class TestTraceGeoglowsAttributes: + """Tests for _trace_geoglows_attributes function.""" + + def test_output_columns( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that output has expected columns.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + expected_columns = { + "flowpath_id", + "flowpath_toid", + "VPUID", + "lengthkm", + "areasqkm", + "totdasqkm", + "mainstemlp", + "pathlength", + "dnhydroseq", + "hydroseq", + "streamorder", + "geometry", + } + + assert set(result.columns) == expected_columns + + def test_all_flowpaths_traced( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that all flowpaths in graph are traced.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + assert len(result) == graph.num_nodes() + + def test_hydroseq_unique( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that hydroseq values are unique.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + assert result["hydroseq"].nunique() == len(result) + + def test_totdasqkm_accumulates_downstream( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that total drainage area accumulates downstream.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + outlet_row = result[result["flowpath_id"] == "810000005"] + max_da = result["totdasqkm"].max() + + assert outlet_row["totdasqkm"].iloc[0] == max_da + + def test_totdasqkm_equals_sum_of_all_areas_at_outlet( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that outlet totdasqkm equals sum of all upstream areas.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + outlet_row = result[result["flowpath_id"] == "810000005"] + total_area = sample_geoglows_catchments["areasqkm"].sum() + + assert np.isclose(outlet_row["totdasqkm"].iloc[0], total_area) + + def test_pathlength_increases_upstream( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that pathlength increases going upstream.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + headwater_row = result[result["flowpath_id"] == "810000001"] + outlet_row = result[result["flowpath_id"] == "810000005"] + + assert headwater_row["pathlength"].iloc[0] > outlet_row["pathlength"].iloc[0] + + def test_pathlength_calculation_correct( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that pathlength is calculated correctly as sum of downstream lengths.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + # Headwater 1's pathlength should be sum of downstream lengths (5->4->3) + # 0.4 + 0.7 + 0.8 = 1.9 + headwater_row = result[result["flowpath_id"] == "810000001"] + expected_pathlength = 0.4 + 0.7 + 0.8 + + assert np.isclose(headwater_row["pathlength"].iloc[0], expected_pathlength) + + def test_mainstemlp_assigned_to_all_nodes( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that all nodes have a mainstem level path assigned.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + assert result["mainstemlp"].notna().all() + + def test_multiple_outlets_handled( + self, + disconnected_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that disconnected subgraphs with multiple outlets are handled.""" + graph, node_indices = disconnected_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + outlets = result[result["dnhydroseq"] == 0] + assert len(outlets) == 2 + + def test_vpuid_assigned_correctly( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that VPUID is assigned correctly to all rows.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + assert (result["VPUID"] == "701").all() + + def test_lengthkm_preserved_from_input( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that lengthkm values are preserved from input.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + for _, row in sample_geoglows_flowpaths.iterrows(): + linkno = str(row["LINKNO"]) + expected_length = row["LengthKM"] + actual_length = result[result["flowpath_id"] == linkno]["lengthkm"].iloc[0] + assert np.isclose(actual_length, expected_length) + + def test_areasqkm_from_catchments( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that areasqkm values come from catchments.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + for _, row in sample_geoglows_catchments.iterrows(): + linkno = str(row["linkno"]) + expected_area = row["areasqkm"] + actual_area = result[result["flowpath_id"] == linkno]["areasqkm"].iloc[0] + assert np.isclose(actual_area, expected_area) + + +class TestCreateGeoglowsReferenceDivides: + """Tests for _create_geoglows_reference_divides function.""" + + @pytest.fixture + def sample_geoglows_reference_flowpaths( + self, sample_geoglows_flowpaths: gpd.GeoDataFrame + ) -> gpd.GeoDataFrame: + """Create sample reference flowpaths (subset of catchments).""" + return gpd.GeoDataFrame( + { + "flowpath_id": ["810000001", "810000003", "810000005"], + "VPUID": ["701", "701", "701"], + }, + geometry=sample_geoglows_flowpaths.geometry.iloc[:3].values, + crs="EPSG:3857", + ) + + def test_returns_geodataframe( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that _create_geoglows_reference_divides returns a GeoDataFrame.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert isinstance(result, gpd.GeoDataFrame) + + def test_all_catchments_included( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that all catchments are included in output.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert len(result) == len(sample_geoglows_catchments) + + def test_has_flowpath_flag_correct( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that has_flowpath flag is set correctly.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert result["has_flowpath"].sum() == 3 + + def test_flowpath_id_assigned_when_has_flowpath( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that flowpath_id is assigned when has_flowpath is True.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + with_flowpath = result[result["has_flowpath"]] + assert with_flowpath["flowpath_id"].notna().all() + assert (with_flowpath["flowpath_id"] == with_flowpath["divide_id"]).all() + + def test_flowpath_id_na_when_no_flowpath( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that flowpath_id is NA when has_flowpath is False.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + without_flowpath = result[~result["has_flowpath"]] + assert without_flowpath["flowpath_id"].isna().all() + + def test_vpuid_assigned_correctly( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that vpuid is assigned correctly to all rows.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert (result["vpuid"] == "701").all() + + def test_divide_id_is_string( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that divide_id is converted to string.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert result["divide_id"].dtype == object + + def test_areasqkm_preserved( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that areasqkm values are preserved.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert "areasqkm" in result.columns + original_areas = set(sample_geoglows_catchments["areasqkm"].tolist()) + result_areas = set(result["areasqkm"].tolist()) + assert original_areas == result_areas + + def test_geometry_preserved( + self, + sample_geoglows_catchments: gpd.GeoDataFrame, + sample_geoglows_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that catchment geometries are preserved.""" + result = _create_geoglows_reference_divides( + sample_geoglows_catchments, sample_geoglows_reference_flowpaths, "701" + ) + + assert result.geometry is not None + assert len(result.geometry) == len(sample_geoglows_catchments) + + +# ============================================================================= +# Tests for GeoGLOWS Edge Cases +# ============================================================================= + + +class TestGeoglowsEdgeCases: + """Tests for GeoGLOWS edge cases.""" + + @pytest.fixture + def single_node_geoglows_graph(self) -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a graph with a single node (headwater that is also outlet).""" + graph = rx.PyDiGraph() + idx = graph.add_node("810000001") + return graph, {"810000001": idx} + + @pytest.fixture + def single_geoglows_flowpath(self) -> gpd.GeoDataFrame: + """Create single GeoGLOWS flowpath GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "LINKNO": [810000001], + "DSLINKNO": [-1], + "strmOrder": [1], + "LengthKM": [0.5], + }, + geometry=[LineString([(0, 0), (1, 1)])], + crs="EPSG:3857", + ) + + @pytest.fixture + def single_geoglows_catchment(self) -> gpd.GeoDataFrame: + """Create single GeoGLOWS catchment GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "linkno": [810000001], + "areasqkm": [1.0], + }, + geometry=[MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])])], + crs="EPSG:3857", + ) + + def test_single_node_graph( + self, + single_node_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + single_geoglows_flowpath: gpd.GeoDataFrame, + single_geoglows_catchment: gpd.GeoDataFrame, + ) -> None: + """Test that single node graph is handled correctly.""" + graph, node_indices = single_node_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, single_geoglows_flowpath, single_geoglows_catchment, "701" + ) + + assert len(result) == 1 + assert result["streamorder"].iloc[0] == 1 + assert result["dnhydroseq"].iloc[0] == 0 + assert result["pathlength"].iloc[0] == 0.0 + assert result["totdasqkm"].iloc[0] == 1.0 + + +class TestGeoglowsGraphWithCycles: + """Tests for handling GeoGLOWS graphs with cycles.""" + + @pytest.fixture + def cyclic_geoglows_graph(self) -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a GeoGLOWS graph with a cycle.""" + graph = rx.PyDiGraph() + + node_data = ["810000001", "810000002", "810000003"] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + graph.add_edge(node_indices["810000001"], node_indices["810000002"], None) + graph.add_edge(node_indices["810000002"], node_indices["810000003"], None) + graph.add_edge(node_indices["810000003"], node_indices["810000001"], None) + + return graph, node_indices + + def test_trace_geoglows_attributes_raises_on_cycle( + self, + cyclic_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that _trace_geoglows_attributes raises AssertionError on cyclic graph.""" + graph, node_indices = cyclic_geoglows_graph + + with pytest.raises(AssertionError, match="Graph contains cycles"): + _trace_geoglows_attributes( + graph, node_indices, sample_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + +class TestTraceAttributes: + """Tests for _trace_attributes function.""" + + def test_output_columns( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that output has expected columns.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + expected_columns = { + "flowpath_id", + "flowpath_toid", + "VPUID", + "lengthkm", + "areasqkm", + "totdasqkm", + "mainstemlp", + "pathlength", + "dnhydroseq", + "hydroseq", + "streamorder", + "fcode_description", + "geometry", + } + + assert set(result.columns) == expected_columns + + def test_all_flowpaths_traced( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that all flowpaths in graph are traced.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + assert len(result) == graph.num_nodes() + + def test_hydroseq_unique( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that hydroseq values are unique.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + assert result["hydroseq"].nunique() == len(result) + + def test_outlet_has_zero_dnhydroseq( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that outlet node has dnhydroseq of 0.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Outlet is 85000100000005 + outlet_row = result[result["flowpath_id"] == "85000100000005"] + assert outlet_row["dnhydroseq"].iloc[0] == 0 + + def test_outlet_has_zero_pathlength( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that outlet node has pathlength of 0.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + outlet_row = result[result["flowpath_id"] == "85000100000005"] + assert outlet_row["pathlength"].iloc[0] == 0.0 + + def test_headwaters_have_stream_order_1( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that headwater nodes have stream order 1.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Headwaters are 85000100000001 and 85000100000002 + headwater_rows = result[result["flowpath_id"].isin(["85000100000001", "85000100000002"])] + assert (headwater_rows["streamorder"] == 1).all() + + def test_strahler_order_increases_at_confluence( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that stream order increases when two streams of same order meet.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Node 3 is confluence of two order-1 streams, should be order 2 + confluence_row = result[result["flowpath_id"] == "85000100000003"] + assert confluence_row["streamorder"].iloc[0] == 2 + + def test_totdasqkm_accumulates_downstream( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that total drainage area accumulates downstream.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Outlet should have largest totdasqkm + outlet_row = result[result["flowpath_id"] == "85000100000005"] + max_da = result["totdasqkm"].max() + + assert outlet_row["totdasqkm"].iloc[0] == max_da + + def test_totdasqkm_equals_sum_of_all_areas_at_outlet( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that outlet totdasqkm equals sum of all upstream areas.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + outlet_row = result[result["flowpath_id"] == "85000100000005"] + total_area = sample_divides["AreaSqKm"].sum() + + assert np.isclose(outlet_row["totdasqkm"].iloc[0], total_area) + + def test_pathlength_increases_upstream( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that pathlength increases going upstream.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Headwaters should have larger pathlength than outlet + headwater_row = result[result["flowpath_id"] == "85000100000001"] + outlet_row = result[result["flowpath_id"] == "85000100000005"] + + assert headwater_row["pathlength"].iloc[0] > outlet_row["pathlength"].iloc[0] + + def test_mainstemlp_assigned_to_all_nodes( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that all nodes have a mainstem level path assigned.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + assert result["mainstemlp"].notna().all() + + def test_multiple_outlets_handled( + self, + disconnected_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that disconnected subgraphs with multiple outlets are handled.""" + graph, node_indices = disconnected_graph + result = _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + # Should have two outlets with dnhydroseq = 0 + outlets = result[result["dnhydroseq"] == 0] + assert len(outlets) == 2 + + +class TestCreateReferenceDivides: + """Tests for _create_reference_divides function.""" + + @pytest.fixture + def sample_reference_flowpaths(self, sample_flowpaths: gpd.GeoDataFrame) -> gpd.GeoDataFrame: + """Create sample reference flowpaths (subset of divides).""" + # Only include some flowpaths (simulating filtered network) + return gpd.GeoDataFrame( + { + "flowpath_id": ["85000100000001", "85000100000003", "85000100000005"], + "VPUID": ["21", "21", "21"], + }, + geometry=sample_flowpaths.geometry.iloc[:3].values, + crs="EPSG:4269", + ) + + def test_returns_geodataframe( + self, + sample_divides: gpd.GeoDataFrame, + sample_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that _create_reference_divides returns a GeoDataFrame.""" + result = _create_reference_divides(sample_divides, sample_reference_flowpaths, "21") + + assert isinstance(result, gpd.GeoDataFrame) + + def test_all_divides_included( + self, + sample_divides: gpd.GeoDataFrame, + sample_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that all divides are included in output.""" + result = _create_reference_divides(sample_divides, sample_reference_flowpaths, "21") + + assert len(result) == len(sample_divides) + + def test_has_flowpath_flag_correct( + self, + sample_divides: gpd.GeoDataFrame, + sample_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that has_flowpath flag is set correctly.""" + result = _create_reference_divides(sample_divides, sample_reference_flowpaths, "21") + + # 3 divides should have flowpaths + assert result["has_flowpath"].sum() == 3 + + def test_flowpath_id_assigned_when_has_flowpath( + self, + sample_divides: gpd.GeoDataFrame, + sample_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that flowpath_id is assigned when has_flowpath is True.""" + result = _create_reference_divides(sample_divides, sample_reference_flowpaths, "21") + + with_flowpath = result[result["has_flowpath"]] + assert with_flowpath["flowpath_id"].notna().all() + assert (with_flowpath["flowpath_id"] == with_flowpath["divide_id"]).all() + + def test_flowpath_id_na_when_no_flowpath( + self, + sample_divides: gpd.GeoDataFrame, + sample_reference_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that flowpath_id is NA when has_flowpath is False.""" + result = _create_reference_divides(sample_divides, sample_reference_flowpaths, "21") + + without_flowpath = result[~result["has_flowpath"]] + assert without_flowpath["flowpath_id"].isna().all() + + +class TestGraphWithCycles: + """Tests for handling graphs with cycles.""" + + @pytest.fixture + def cyclic_graph(self) -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a graph with a cycle. + + 1 -> 2 -> 3 -> 1 (cycle) + """ + graph = rx.PyDiGraph() + + node_data = ["85000100000001", "85000100000002", "85000100000003"] + + node_indices = {} + for fp_id in node_data: + idx = graph.add_node(fp_id) + node_indices[fp_id] = idx + + graph.add_edge(node_indices["85000100000001"], node_indices["85000100000002"], None) + graph.add_edge(node_indices["85000100000002"], node_indices["85000100000003"], None) + graph.add_edge(node_indices["85000100000003"], node_indices["85000100000001"], None) + + return graph, node_indices + + def test_trace_attributes_raises_on_cycle( + self, + cyclic_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that _trace_attributes raises AssertionError on cyclic graph.""" + graph, node_indices = cyclic_graph + + with pytest.raises(AssertionError, match="Graph contains cycles"): + _trace_attributes(graph, node_indices, sample_flowpaths, sample_divides, "21") + + +class TestEdgeCases: + """Tests for edge cases.""" + + @pytest.fixture + def single_node_graph(self) -> tuple[rx.PyDiGraph, dict[str, int]]: + """Create a graph with a single node (headwater that is also outlet).""" + graph = rx.PyDiGraph() + idx = graph.add_node("85000100000001") + return graph, {"85000100000001": idx} + + @pytest.fixture + def single_flowpath(self) -> gpd.GeoDataFrame: + """Create single flowpath GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "NHDPlusID": [85000100000001], + "VPUID": ["2101"], + "LengthKM": [0.5], + "fcode_description": ["Stream/River"], + }, + geometry=[MultiLineString([LineString([(0, 0), (1, 1)])])], + crs="EPSG:4269", + ) + + @pytest.fixture + def single_divide(self) -> gpd.GeoDataFrame: + """Create single divide GeoDataFrame.""" + return gpd.GeoDataFrame( + { + "NHDPlusID": [85000100000001], + "VPUID": ["2101"], + "AreaSqKm": [1.0], + }, + geometry=[MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])])], + crs="EPSG:4269", + ) + + def test_single_node_graph( + self, + single_node_graph: tuple[rx.PyDiGraph, dict[str, int]], + single_flowpath: gpd.GeoDataFrame, + single_divide: gpd.GeoDataFrame, + ) -> None: + """Test that single node graph is handled correctly.""" + graph, node_indices = single_node_graph + result = _trace_attributes(graph, node_indices, single_flowpath, single_divide, "21") + + assert len(result) == 1 + assert result["streamorder"].iloc[0] == 1 + assert result["dnhydroseq"].iloc[0] == 0 + assert result["pathlength"].iloc[0] == 0.0 + assert result["totdasqkm"].iloc[0] == 1.0 + + def test_missing_divide_defaults_to_zero_area( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + sample_flowpaths: gpd.GeoDataFrame, + ) -> None: + """Test that missing divides default to zero area.""" + graph, node_indices = sample_graph + + # Create divides missing one entry + partial_divides = gpd.GeoDataFrame( + { + "NHDPlusID": [85000100000001, 85000100000002], + "VPUID": ["2101", "2101"], + "AreaSqKm": [0.5, 0.5], + }, + geometry=[ + MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])]), + MultiPolygon([Polygon([(1, 0), (2, 0), (2, 1), (1, 1)])]), + ], + crs="EPSG:4269", + ) + + result = _trace_attributes(graph, node_indices, sample_flowpaths, partial_divides, "21") + + # Nodes without divides should have areasqkm = 0 + missing_divide_rows = result[~result["flowpath_id"].isin(["85000100000001", "85000100000002"])] + assert (missing_divide_rows["areasqkm"] == 0.0).all() + + +class TestFlowpathOrientation: + """Tests for flowpath geometry orientation in _trace_attributes.""" + + def test_reversed_flowpaths_are_corrected( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + reversed_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that reversed flowpaths are oriented correctly.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, reversed_flowpaths, sample_divides, "21") + + # Check that flowpath 4 ends at the confluence with 5 + fp4 = result[result["flowpath_id"] == "85000100000004"].iloc[0] + fp5 = result[result["flowpath_id"] == "85000100000005"].iloc[0] + + # Get endpoints + from reference_builds.utils.geometries import _get_endpoints + + _, fp4_end = _get_endpoints(fp4.geometry) + fp5_start, _ = _get_endpoints(fp5.geometry) + + # fp4 should end where fp5 starts (or very close) + assert Point(fp4_end).distance(Point(fp5_start)) < 0.001 + + def test_correctly_oriented_flowpaths_unchanged( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + correctly_oriented_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that correctly oriented flowpaths remain unchanged.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, correctly_oriented_flowpaths, sample_divides, "21") + + from reference_builds.utils.geometries import _get_endpoints + + # Check original and result have same orientation + for _, row in correctly_oriented_flowpaths.iterrows(): + fp_id = str(int(row["NHDPlusID"])) + original_start, original_end = _get_endpoints(row.geometry) + + result_row = result[result["flowpath_id"] == fp_id].iloc[0] + result_start, result_end = _get_endpoints(result_row.geometry) + + assert original_start == result_start + assert original_end == result_end + + def test_outlet_oriented_correctly( + self, + sample_graph: tuple[rx.PyDiGraph, dict[str, int]], + reversed_flowpaths: gpd.GeoDataFrame, + sample_divides: gpd.GeoDataFrame, + ) -> None: + """Test that outlet flowpath is oriented correctly using upstream.""" + graph, node_indices = sample_graph + result = _trace_attributes(graph, node_indices, reversed_flowpaths, sample_divides, "21") + + from reference_builds.utils.geometries import _get_endpoints + + # Outlet is 85000100000005 + outlet = result[result["flowpath_id"] == "85000100000005"].iloc[0] + upstream = result[result["flowpath_id"] == "85000100000004"].iloc[0] + + _, upstream_end = _get_endpoints(upstream.geometry) + outlet_start, _ = _get_endpoints(outlet.geometry) + + # Outlet should start where upstream ends + assert Point(outlet_start).distance(Point(upstream_end)) < 0.001 + + +class TestGeoglowsFlowpathOrientation: + """Tests for flowpath orientation in _trace_geoglows_attributes.""" + + @pytest.fixture + def reversed_geoglows_flowpaths(self) -> gpd.GeoDataFrame: + """Create GeoGLOWS flowpaths with reversed digitization.""" + return gpd.GeoDataFrame( + { + "LINKNO": [810000001, 810000002, 810000003, 810000004, 810000005], + "DSLINKNO": [810000003, 810000003, 810000004, 810000005, -1], + "strmOrder": [1, 1, 2, 2, 2], + "LengthKM": [0.5, 0.6, 0.8, 0.7, 0.4], + }, + geometry=[ + LineString([(1, 1), (0, 0)]), # Reversed: should be (0,0)->(1,1) + LineString([(1, 1), (2, 0)]), # Reversed: should be (2,0)->(1,1) + LineString([(1, 2), (1, 1)]), # Reversed: should be (1,1)->(1,2) + LineString([(1, 3), (1, 2)]), # Reversed: should be (1,2)->(1,3) + LineString([(1, 4), (1, 3)]), # Reversed: should be (1,3)->(1,4) + ], + crs="EPSG:3857", + ) + + def test_reversed_geoglows_flowpaths_are_corrected( + self, + sample_geoglows_graph: tuple[rx.PyDiGraph, dict[str, int]], + reversed_geoglows_flowpaths: gpd.GeoDataFrame, + sample_geoglows_catchments: gpd.GeoDataFrame, + ) -> None: + """Test that reversed GeoGLOWS flowpaths are oriented correctly.""" + graph, node_indices = sample_geoglows_graph + result = _trace_geoglows_attributes( + graph, node_indices, reversed_geoglows_flowpaths, sample_geoglows_catchments, "701" + ) + + from reference_builds.utils.geometries import _get_endpoints + + # Check that flowpath 4 ends where flowpath 5 starts + fp4 = result[result["flowpath_id"] == "810000004"].iloc[0] + fp5 = result[result["flowpath_id"] == "810000005"].iloc[0] + + _, fp4_end = _get_endpoints(fp4.geometry) + fp5_start, _ = _get_endpoints(fp5.geometry) + + assert Point(fp4_end).distance(Point(fp5_start)) < 0.001 diff --git a/tests/test_geometry_utils.py b/tests/test_geometry_utils.py new file mode 100644 index 0000000..0165399 --- /dev/null +++ b/tests/test_geometry_utils.py @@ -0,0 +1,266 @@ +"""Tests for geometry utility functions""" + +from shapely import wkb +from shapely.geometry import LineString, MultiLineString + +from reference_builds.utils.geometries import ( + _ensure_geometry, + _get_endpoints, + _orient_flowpath_downstream, + _reverse_line, +) + + +class TestEnsureGeometry: + """Tests for _ensure_geometry function.""" + + def test_returns_geometry_unchanged(self) -> None: + """Test that Shapely geometry is returned unchanged.""" + line = LineString([(0, 0), (1, 1)]) + result = _ensure_geometry(line) + assert result == line + + def test_converts_wkb_bytes_to_geometry(self) -> None: + """Test that WKB bytes are converted to Shapely geometry.""" + line = LineString([(0, 0), (1, 1)]) + wkb_bytes = wkb.dumps(line) + result = _ensure_geometry(wkb_bytes) + assert result.equals(line) + + def test_handles_multilinestring(self) -> None: + """Test that MultiLineString is handled correctly.""" + multi = MultiLineString([[(0, 0), (1, 1)], [(2, 2), (3, 3)]]) + result = _ensure_geometry(multi) + assert result == multi + + def test_handles_multilinestring_wkb(self) -> None: + """Test that MultiLineString WKB is converted correctly.""" + multi = MultiLineString([[(0, 0), (1, 1)], [(2, 2), (3, 3)]]) + wkb_bytes = wkb.dumps(multi) + result = _ensure_geometry(wkb_bytes) + assert result.equals(multi) + + +class TestGetEndpoints: + """Tests for _get_endpoints function.""" + + def test_linestring_endpoints(self) -> None: + """Test getting endpoints from LineString.""" + line = LineString([(0, 0), (1, 1), (2, 2)]) + start, end = _get_endpoints(line) + assert start == (0, 0) + assert end == (2, 2) + + def test_multilinestring_endpoints(self) -> None: + """Test getting endpoints from MultiLineString.""" + multi = MultiLineString([[(0, 0), (1, 1)], [(2, 2), (3, 3)]]) + start, end = _get_endpoints(multi) + assert start == (0, 0) + assert end == (3, 3) + + def test_single_segment_linestring(self) -> None: + """Test LineString with only two points.""" + line = LineString([(5, 5), (10, 10)]) + start, end = _get_endpoints(line) + assert start == (5, 5) + assert end == (10, 10) + + def test_handles_wkb_input(self) -> None: + """Test that WKB input is handled.""" + line = LineString([(0, 0), (1, 1)]) + wkb_bytes = wkb.dumps(line) + start, end = _get_endpoints(wkb_bytes) + assert start == (0, 0) + assert end == (1, 1) + + +class TestReverseLine: + """Tests for _reverse_line function.""" + + def test_reverse_linestring(self) -> None: + """Test reversing a LineString.""" + line = LineString([(0, 0), (1, 1), (2, 2)]) + result = _reverse_line(line) + assert list(result.coords) == [(2, 2), (1, 1), (0, 0)] + + def test_reverse_multilinestring(self) -> None: + """Test reversing a MultiLineString.""" + multi = MultiLineString([[(0, 0), (1, 1)], [(2, 2), (3, 3)]]) + result = _reverse_line(multi) + + # Should reverse order of parts and coords within each part + parts = list(result.geoms) + assert len(parts) == 2 + assert list(parts[0].coords) == [(3, 3), (2, 2)] + assert list(parts[1].coords) == [(1, 1), (0, 0)] + + def test_reverse_preserves_geometry_type(self) -> None: + """Test that reversed geometry has same type.""" + line = LineString([(0, 0), (1, 1)]) + result = _reverse_line(line) + assert result.geom_type == "LineString" + + multi = MultiLineString([[(0, 0), (1, 1)]]) + result = _reverse_line(multi) + assert result.geom_type == "MultiLineString" + + def test_handles_wkb_input(self) -> None: + """Test that WKB input is handled.""" + line = LineString([(0, 0), (1, 1)]) + wkb_bytes = wkb.dumps(line) + result = _reverse_line(wkb_bytes) + assert list(result.coords) == [(1, 1), (0, 0)] + + +class TestOrientFlowpathDownstream: + """Tests for _orient_flowpath_downstream function.""" + + def test_already_correct_orientation_with_downstream(self) -> None: + """Test that correctly oriented line is unchanged.""" + # Line goes from (0,0) to (1,1), downstream is at (1,1) to (1,2) + line = LineString([(0, 0), (1, 1)]) + ds_geom = LineString([(1, 1), (1, 2)]) + + result = _orient_flowpath_downstream(line, ds_geom=ds_geom) + + # End point (1,1) is closer to downstream, so no change + assert list(result.coords) == [(0, 0), (1, 1)] + + def test_reversed_orientation_with_downstream(self) -> None: + """Test that reversed line is corrected.""" + # Line goes from (1,1) to (0,0), but downstream is at (1,1) + line = LineString([(1, 1), (0, 0)]) + ds_geom = LineString([(1, 1), (1, 2)]) + + result = _orient_flowpath_downstream(line, ds_geom=ds_geom) + + # Start point (1,1) is closer to downstream, so should reverse + assert list(result.coords) == [(0, 0), (1, 1)] + + def test_outlet_with_upstream_geometry_correct(self) -> None: + """Test outlet orientation using upstream geometry (already correct).""" + # Outlet line goes from (1,1) to (1,2), upstream ends at (1,1) + line = LineString([(1, 1), (1, 2)]) + us_geom = LineString([(0, 0), (1, 1)]) + + result = _orient_flowpath_downstream(line, ds_geom=None, us_geom=us_geom) + + # Start (1,1) is closer to upstream, which is correct + assert list(result.coords) == [(1, 1), (1, 2)] + + def test_outlet_with_upstream_geometry_reversed(self) -> None: + """Test outlet orientation using upstream geometry (needs reversal).""" + # Outlet line goes from (1,2) to (1,1), but upstream ends at (1,1) + line = LineString([(1, 2), (1, 1)]) + us_geom = LineString([(0, 0), (1, 1)]) + + result = _orient_flowpath_downstream(line, ds_geom=None, us_geom=us_geom) + + # End (1,1) is closer to upstream, so should reverse + assert list(result.coords) == [(1, 1), (1, 2)] + + def test_no_reference_returns_unchanged(self) -> None: + """Test that line is unchanged when no reference geometry provided.""" + line = LineString([(0, 0), (1, 1)]) + result = _orient_flowpath_downstream(line, ds_geom=None, us_geom=None) + assert list(result.coords) == [(0, 0), (1, 1)] + + def test_multilinestring_orientation(self) -> None: + """Test orientation of MultiLineString.""" + # Multi goes from (0,0) to (2,2), downstream starts at (2,2) + multi = MultiLineString([[(0, 0), (1, 1)], [(1, 1), (2, 2)]]) + ds_geom = LineString([(2, 2), (3, 3)]) + + result = _orient_flowpath_downstream(multi, ds_geom=ds_geom) + + # End (2,2) is closer to downstream, so no change + start, end = _get_endpoints(result) + assert start == (0, 0) + assert end == (2, 2) + + def test_multilinestring_reversed(self) -> None: + """Test reversal of MultiLineString.""" + # Multi goes from (2,2) to (0,0), but downstream is at (2,2) + multi = MultiLineString([[(2, 2), (1, 1)], [(1, 1), (0, 0)]]) + ds_geom = LineString([(2, 2), (3, 3)]) + + result = _orient_flowpath_downstream(multi, ds_geom=ds_geom) + + # Start (2,2) is closer to downstream, so should reverse + start, end = _get_endpoints(result) + assert start == (0, 0) + assert end == (2, 2) + + def test_handles_wkb_input(self) -> None: + """Test that WKB input is handled for both geometries.""" + line = LineString([(0, 0), (1, 1)]) + ds_geom = LineString([(1, 1), (1, 2)]) + + line_wkb = wkb.dumps(line) + ds_wkb = wkb.dumps(ds_geom) + + result = _orient_flowpath_downstream(line_wkb, ds_geom=ds_wkb) + assert list(result.coords) == [(0, 0), (1, 1)] + + def test_downstream_takes_priority_over_upstream(self) -> None: + """Test that downstream geometry is used when both are provided.""" + line = LineString([(1, 1), (0, 0)]) + ds_geom = LineString([(1, 1), (1, 2)]) # Would cause reversal + us_geom = LineString([(0, 0), (0, -1)]) # Would not cause reversal + + result = _orient_flowpath_downstream(line, ds_geom=ds_geom, us_geom=us_geom) + + # Should use downstream, so should reverse + assert list(result.coords) == [(0, 0), (1, 1)] + + +class TestOrientFlowpathDownstreamEdgeCases: + """Edge case tests for _orient_flowpath_downstream.""" + + def test_identical_endpoints(self) -> None: + """Test handling when both endpoints are equidistant from downstream.""" + # This is a degenerate case - line perpendicular to downstream + line = LineString([(0, 0), (2, 0)]) + ds_geom = LineString([(1, 0), (1, 1)]) + + # Should return unchanged (or reversed, but consistently) + result = _orient_flowpath_downstream(line, ds_geom=ds_geom) + assert result.geom_type == "LineString" + + def test_touching_downstream(self) -> None: + """Test when line endpoint touches downstream geometry.""" + line = LineString([(0, 0), (1, 1)]) + ds_geom = LineString([(1, 1), (2, 2)]) # Starts exactly at line end + + result = _orient_flowpath_downstream(line, ds_geom=ds_geom) + + # End touches downstream (distance 0), so should not reverse + assert list(result.coords) == [(0, 0), (1, 1)] + + def test_touching_upstream(self) -> None: + """Test outlet when line start touches upstream geometry.""" + line = LineString([(1, 1), (2, 2)]) + us_geom = LineString([(0, 0), (1, 1)]) # Ends exactly at line start + + result = _orient_flowpath_downstream(line, ds_geom=None, us_geom=us_geom) + + # Start touches upstream (distance 0), so should not reverse + assert list(result.coords) == [(1, 1), (2, 2)] + + def test_long_multilinestring(self) -> None: + """Test with MultiLineString with many segments.""" + multi = MultiLineString( + [ + [(0, 0), (1, 0)], + [(1, 0), (2, 0)], + [(2, 0), (3, 0)], + [(3, 0), (4, 0)], + ] + ) + ds_geom = LineString([(4, 0), (5, 0)]) + + result = _orient_flowpath_downstream(multi, ds_geom=ds_geom) + + start, end = _get_endpoints(result) + assert start == (0, 0) + assert end == (4, 0)