From e92c67f4d1936a1255168bccc057ff8135fee21f Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 11:40:29 +0100 Subject: [PATCH 01/20] Remove pysheds --- dev-requirements.txt | 35 +---- doc-requirements.txt | 35 +---- pyproject.toml | 3 +- requirements.txt | 33 +---- swmmanywhere/filepaths.py | 2 +- swmmanywhere/geospatial_utilities.py | 208 ++++++--------------------- swmmanywhere/preprocessing.py | 2 +- swmmanywhere/swmmanywhere.py | 5 +- tests/test_geospatial_utilities.py | 4 +- 9 files changed, 55 insertions(+), 272 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index a71a4920..4e934777 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with Python 3.11 # by the following command: # -# pip-compile --extra=dev --output-file=dev-requirements.txt +# pip-compile --extra=dev --output-file=dev-requirements.txt pyproject.toml # aenum==3.1.11 # via @@ -11,7 +11,6 @@ aenum==3.1.11 affine==2.4.0 # via # pyflwdir - # pysheds # rasterio annotated-types==0.7.0 # via pydantic @@ -86,8 +85,6 @@ fsspec==2024.6.1 # via fastparquet geographiclib==2.0 # via geopy -geojson==3.1.0 - # via pysheds geopandas==1.0.1 # via # osmnx @@ -102,8 +99,6 @@ identify==2.6.0 # via pre-commit idna==3.7 # via requests -imageio==2.34.2 - # via scikit-image iniconfig==2.0.0 # via pytest joblib==1.4.2 @@ -116,8 +111,6 @@ jsonschema-specifications==2023.12.1 # via jsonschema julian==0.14 # via pyswmm -lazy-loader==0.4 - # via scikit-image llvmlite==0.43.0 # via numba loguru==0.7.2 @@ -138,20 +131,16 @@ networkx==3.3 # via # netcomp # osmnx - # scikit-image # swmmanywhere (pyproject.toml) nodeenv==1.9.1 # via pre-commit numba==0.60.0 - # via - # pyflwdir - # pysheds + # via pyflwdir numpy==1.26.4 # via # cftime # fastparquet # geopandas - # imageio # netcdf4 # netcomp # numba @@ -160,15 +149,12 @@ numpy==1.26.4 # pyarrow # pyflwdir # pyogrio - # pysheds # rasterio # rioxarray - # scikit-image # scipy # shapely # snuggs # swmmanywhere (pyproject.toml) - # tifffile # xarray osmnx==1.9.3 # via swmmanywhere (pyproject.toml) @@ -177,26 +163,19 @@ packaging==24.1 # build # fastparquet # geopandas - # lazy-loader # planetary-computer # pyogrio # pyswmm # pytest # rioxarray - # scikit-image # xarray pandas==2.2.2 # via # fastparquet # geopandas # osmnx - # pysheds # swmmanywhere (pyproject.toml) # xarray -pillow==10.4.0 - # via - # imageio - # scikit-image pip-tools==7.4.1 # via swmmanywhere (pyproject.toml) planetary-computer==1.0.0 @@ -224,14 +203,11 @@ pyparsing==3.1.2 pyproj==3.6.1 # via # geopandas - # pysheds # rioxarray pyproject-hooks==1.1.0 # via # build # pip-tools -pysheds==0.3.5 - # via swmmanywhere (pyproject.toml) pystac[validation]==1.10.1 # via # planetary-computer @@ -273,7 +249,6 @@ pyyaml==6.0.1 # swmmanywhere (pyproject.toml) rasterio==1.3.10 # via - # pysheds # rioxarray # swmmanywhere (pyproject.toml) referencing==0.35.1 @@ -296,14 +271,10 @@ rpds-py==0.19.1 # referencing ruff==0.5.5 # via swmmanywhere (pyproject.toml) -scikit-image==0.24.0 - # via pysheds scipy==1.14.0 # via # netcomp # pyflwdir - # pysheds - # scikit-image # swmmanywhere (pyproject.toml) shapely==2.0.5 # via @@ -320,8 +291,6 @@ snuggs==1.4.7 # via rasterio swmm-toolkit==0.15.5 # via pyswmm -tifffile==2024.7.2 - # via scikit-image toolz==0.12.1 # via cytoolz tqdm==4.66.4 diff --git a/doc-requirements.txt b/doc-requirements.txt index bcd8c604..949da9a8 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with Python 3.11 # by the following command: # -# pip-compile --extra=doc --output-file=doc-requirements.txt +# pip-compile --extra=doc --output-file=doc-requirements.txt pyproject.toml # aenum==3.1.11 # via @@ -11,7 +11,6 @@ aenum==3.1.11 affine==2.4.0 # via # pyflwdir - # pysheds # rasterio annotated-types==0.7.0 # via pydantic @@ -98,8 +97,6 @@ fsspec==2024.6.1 # via fastparquet geographiclib==2.0 # via geopy -geojson==3.1.0 - # via pysheds geopandas==1.0.1 # via # osmnx @@ -116,8 +113,6 @@ griffe==0.47.0 # via mkdocstrings-python idna==3.7 # via requests -imageio==2.34.2 - # via scikit-image ipykernel==6.29.5 # via mkdocs-jupyter ipython==8.26.0 @@ -156,8 +151,6 @@ jupyterlab-pygments==0.3.0 # via nbconvert jupytext==1.16.3 # via mkdocs-jupyter -lazy-loader==0.4 - # via scikit-image llvmlite==0.43.0 # via numba loguru==0.7.2 @@ -248,18 +241,14 @@ networkx==3.3 # via # netcomp # osmnx - # scikit-image # swmmanywhere (pyproject.toml) numba==0.60.0 - # via - # pyflwdir - # pysheds + # via pyflwdir numpy==1.26.4 # via # cftime # fastparquet # geopandas - # imageio # netcdf4 # netcomp # numba @@ -268,15 +257,12 @@ numpy==1.26.4 # pyarrow # pyflwdir # pyogrio - # pysheds # rasterio # rioxarray - # scikit-image # scipy # shapely # snuggs # swmmanywhere (pyproject.toml) - # tifffile # xarray osmnx==1.9.3 # via swmmanywhere (pyproject.toml) @@ -286,14 +272,12 @@ packaging==24.1 # geopandas # ipykernel # jupytext - # lazy-loader # mkdocs # nbconvert # planetary-computer # pyogrio # pyswmm # rioxarray - # scikit-image # xarray paginate==0.5.6 # via mkdocs-material @@ -302,7 +286,6 @@ pandas==2.2.2 # fastparquet # geopandas # osmnx - # pysheds # swmmanywhere (pyproject.toml) # xarray pandocfilters==1.5.1 @@ -311,10 +294,6 @@ parso==0.8.4 # via jedi pathspec==0.12.1 # via mkdocs -pillow==10.4.0 - # via - # imageio - # scikit-image planetary-computer==1.0.0 # via swmmanywhere (pyproject.toml) platformdirs==4.2.2 @@ -355,10 +334,7 @@ pyparsing==3.1.2 pyproj==3.6.1 # via # geopandas - # pysheds # rioxarray -pysheds==0.3.5 - # via swmmanywhere (pyproject.toml) pystac[validation]==1.10.1 # via # planetary-computer @@ -402,7 +378,6 @@ pyzmq==26.0.3 # jupyter-client rasterio==1.3.10 # via - # pysheds # rioxarray # swmmanywhere (pyproject.toml) referencing==0.35.1 @@ -426,14 +401,10 @@ rpds-py==0.19.1 # via # jsonschema # referencing -scikit-image==0.24.0 - # via pysheds scipy==1.14.0 # via # netcomp # pyflwdir - # pysheds - # scikit-image # swmmanywhere (pyproject.toml) shapely==2.0.5 # via @@ -456,8 +427,6 @@ stack-data==0.6.3 # via ipython swmm-toolkit==0.15.5 # via pyswmm -tifffile==2024.7.2 - # via scikit-image tinycss2==1.3.0 # via nbconvert toolz==0.12.1 diff --git a/pyproject.toml b/pyproject.toml index 19d5dd05..db61d414 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,14 +36,13 @@ dependencies = [ "netcdf4", "netcomp@ git+https://github.com/barneydobson/NetComp.git", "networkx>=3", - "numpy<2.0.0", + "numpy", "osmnx", "pandas", "planetary_computer", "pyarrow", "pydantic", "pyflwdir", - "pysheds==0.3.5", "pystac_client", "pyswmm", "PyYAML", diff --git a/requirements.txt b/requirements.txt index e65bc1f8..5306ed80 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,6 @@ aenum==3.1.11 affine==2.4.0 # via # pyflwdir - # pysheds # rasterio annotated-types==0.7.0 # via pydantic @@ -70,8 +69,6 @@ fsspec==2024.6.1 # via fastparquet geographiclib==2.0 # via geopy -geojson==3.1.0 - # via pysheds geopandas==1.0.1 # via # osmnx @@ -84,8 +81,6 @@ gitpython==3.1.43 # via swmmanywhere (pyproject.toml) idna==3.7 # via requests -imageio==2.34.2 - # via scikit-image joblib==1.4.2 # via swmmanywhere (pyproject.toml) jsonschema==4.23.0 @@ -96,8 +91,6 @@ jsonschema-specifications==2023.12.1 # via jsonschema julian==0.14 # via pyswmm -lazy-loader==0.4 - # via scikit-image llvmlite==0.43.0 # via numba loguru==0.7.2 @@ -112,18 +105,14 @@ networkx==3.3 # via # netcomp # osmnx - # scikit-image # swmmanywhere (pyproject.toml) numba==0.60.0 - # via - # pyflwdir - # pysheds + # via pyflwdir numpy==1.26.4 # via # cftime # fastparquet # geopandas - # imageio # netcdf4 # netcomp # numba @@ -132,15 +121,12 @@ numpy==1.26.4 # pyarrow # pyflwdir # pyogrio - # pysheds # rasterio # rioxarray - # scikit-image # scipy # shapely # snuggs # swmmanywhere (pyproject.toml) - # tifffile # xarray osmnx==1.9.3 # via swmmanywhere (pyproject.toml) @@ -148,25 +134,18 @@ packaging==24.1 # via # fastparquet # geopandas - # lazy-loader # planetary-computer # pyogrio # pyswmm # rioxarray - # scikit-image # xarray pandas==2.2.2 # via # fastparquet # geopandas # osmnx - # pysheds # swmmanywhere (pyproject.toml) # xarray -pillow==10.4.0 - # via - # imageio - # scikit-image planetary-computer==1.0.0 # via swmmanywhere (pyproject.toml) pyarrow==16.1.0 @@ -186,10 +165,7 @@ pyparsing==3.1.2 pyproj==3.6.1 # via # geopandas - # pysheds # rioxarray -pysheds==0.3.5 - # via swmmanywhere (pyproject.toml) pystac[validation]==1.10.1 # via # planetary-computer @@ -217,7 +193,6 @@ pyyaml==6.0.1 # via swmmanywhere (pyproject.toml) rasterio==1.3.10 # via - # pysheds # rioxarray # swmmanywhere (pyproject.toml) referencing==0.35.1 @@ -238,14 +213,10 @@ rpds-py==0.19.1 # via # jsonschema # referencing -scikit-image==0.24.0 - # via pysheds scipy==1.14.0 # via # netcomp # pyflwdir - # pysheds - # scikit-image # swmmanywhere (pyproject.toml) shapely==2.0.5 # via @@ -262,8 +233,6 @@ snuggs==1.4.7 # via rasterio swmm-toolkit==0.15.5 # via pyswmm -tifffile==2024.7.2 - # via scikit-image toolz==0.12.1 # via cytoolz tqdm==4.66.4 diff --git a/swmmanywhere/filepaths.py b/swmmanywhere/filepaths.py index dabeedde..d3099996 100644 --- a/swmmanywhere/filepaths.py +++ b/swmmanywhere/filepaths.py @@ -289,7 +289,7 @@ def filepaths_from_yaml(f: Path): """Get file paths from a yaml file.""" address_dict = yaml_load(f.read_text()) address_dict["base_dir"] = Path(address_dict["base_dir"]) - overrides = address_dict.pop("overrides") + overrides = address_dict.pop("overrides", {}) addresses = FilePaths(**address_dict, **overrides) return addresses diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index e94071b0..34cb5474 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -8,9 +8,7 @@ import itertools import json import math -import operator import os -from copy import deepcopy from functools import lru_cache from pathlib import Path from typing import List, Optional @@ -27,7 +25,6 @@ from scipy.interpolate import RegularGridInterpolator from scipy.spatial import KDTree from shapely import geometry as sgeom -from shapely import ops as sops from shapely.strtree import STRtree from tqdm.auto import tqdm @@ -35,8 +32,6 @@ os.environ["NUMBA_NUM_THREADS"] = "1" import pyflwdir # noqa: E402 -import pysheds # noqa: E402 -from pysheds import grid as pgrid # noqa: E402 TransformerFromCRS = lru_cache(pyproj.transformer.Transformer.from_crs) @@ -368,117 +363,6 @@ def burn_shape_in_raster( dest.write(data, 1) -def condition_dem( - grid: pysheds.sgrid.sGrid, dem: pysheds.sview.Raster -) -> pysheds.sview.Raster: - """Condition a DEM with pysheds. - - Args: - grid (pysheds.sgrid.sGrid): The grid object. - dem (pysheds.sview.Raster): The input DEM. - - Returns: - pysheds.sview.Raster: The conditioned DEM. - """ - # Fill pits, depressions, and resolve flats in the DEM - pit_filled_dem = grid.fill_pits(dem) - flooded_dem = grid.fill_depressions(pit_filled_dem) - inflated_dem = grid.resolve_flats(flooded_dem) - - return inflated_dem - - -def compute_flow_directions( - grid: pysheds.sgrid.sGrid, inflated_dem: pysheds.sview.Raster -) -> tuple[pysheds.sview.Raster, tuple]: - """Compute flow directions. - - Args: - grid (pysheds.sgrid.sGrid): The grid object. - inflated_dem (pysheds.sview.Raster): The input DEM. - - Returns: - pysheds.sview.Raster: Flow directions. - tuple: Direction mapping. - """ - dirmap = (64, 128, 1, 2, 4, 8, 16, 32) - flow_dir = grid.flowdir(inflated_dem, dirmap=dirmap) - return flow_dir, dirmap - - -def calculate_flow_accumulation( - grid: pysheds.sgrid.sGrid, flow_dir: pysheds.sview.Raster, dirmap: tuple -) -> pysheds.sview.Raster: - """Calculate flow accumulation. - - Args: - grid (pysheds.sgrid.sGrid): The grid object. - flow_dir (pysheds.sview.Raster): Flow directions. - dirmap (tuple): Direction mapping. - - Returns: - pysheds.sview.Raster: Flow accumulations. - """ - flow_acc = grid.accumulation(flow_dir, dirmap=dirmap) - return flow_acc - - -def delineate_catchment( - grid: pysheds.sgrid.sGrid, - flow_acc: pysheds.sview.Raster, - flow_dir: pysheds.sview.Raster, - dirmap: tuple, - G: nx.Graph, -) -> gpd.GeoDataFrame: - """Delineate catchments. - - Args: - grid (pysheds.sgrid.Grid): The grid object. - flow_acc (pysheds.sview.Raster): Flow accumulations. - flow_dir (pysheds.sview.Raster): Flow directions. - dirmap (tuple): Direction mapping. - G (nx.Graph): The input graph with nodes containing 'x' and 'y'. - - Returns: - gpd.GeoDataFrame: A GeoDataFrame containing polygons with columns: - 'geometry', 'area', and 'id'. Sorted by area in descending order. - """ - polys = [] - # Iterate over the nodes in the graph - for id, data in tqdm(G.nodes(data=True), total=len(G.nodes), disable=not verbose()): - # Snap the node to the nearest grid cell - x, y = data["x"], data["y"] - grid_ = deepcopy(grid) - x_snap, y_snap = grid_.snap_to_mask(flow_acc >= 0, (x, y)) - - # Delineate the catchment - catch = grid_.catchment( - x=x_snap, - y=y_snap, - fdir=flow_dir, - dirmap=dirmap, - xytype="coordinate", - algorithm="recursive", - ) - # n.b. recursive algorithm is not recommended, but crashes with a seg - # fault occasionally otherwise. - - grid_.clip_to(catch) - - # Polygonize the catchment - shapes = grid_.polygonize() - catchment_polygon = sops.unary_union( - [sgeom.shape(shape) for shape, value in shapes] - ) - - # Add the catchment to the list - polys.append( - {"id": id, "geometry": catchment_polygon, "area": catchment_polygon.area} - ) - polys.sort(key=operator.itemgetter("area"), reverse=True) - return gpd.GeoDataFrame(polys, crs=grid.crs) - - def remove_intersections(polys: gpd.GeoDataFrame) -> gpd.GeoDataFrame: """Remove intersections from a GeoDataFrame of polygons. @@ -573,14 +457,14 @@ def attach_unconnected_subareas( def calculate_slope( - polys_gdf: gpd.GeoDataFrame, grid: pysheds.sgrid.sGrid, cell_slopes: np.ndarray + polys_gdf: gpd.GeoDataFrame, grid: dict, cell_slopes: np.ndarray ) -> gpd.GeoDataFrame: """Calculate the average slope of each polygon. Args: polys_gdf (gpd.GeoDataFrame): A GeoDataFrame containing polygons with columns: 'geometry', 'area', and 'id'. - grid (pysheds.sgrid.sGrid): The grid object. + grid (dict): Information of the raster (affine, shape, crs, bbox) cell_slopes (np.ndarray): The slopes of each cell in the grid. Returns: @@ -590,7 +474,7 @@ def calculate_slope( polys_gdf["slope"] = None for idx, row in polys_gdf.iterrows(): mask = features.geometry_mask( - [row.geometry], grid.shape, grid.affine, invert=True + [row.geometry], grid["shape"], grid["affine"], invert=True ) average_slope = cell_slopes[mask].mean().mean() polys_gdf.loc[idx, "slope"] = max(float(average_slope), 0) @@ -633,7 +517,7 @@ def vectorize( def delineate_catchment_pyflwdir( - grid: pysheds.sgrid.sGrid, flow_dir: pysheds.sview.Raster, G: nx.Graph + grid: dict, flow_dir: np.array, G: nx.Graph ) -> gpd.GeoDataFrame: """Derive subcatchments from the nodes on a graph and a DEM. @@ -641,8 +525,8 @@ def delineate_catchment_pyflwdir( faster than delineate_catchment. Args: - grid (pysheds.sgrid.Grid): The grid object. - flow_dir (pysheds.sview.Raster): Flow directions. + grid (dict): Information of the raster (affine, shape, crs, bbox). + flow_dir (np.array): Flow directions. G (nx.Graph): The input graph with nodes containing 'x' and 'y'. Returns: @@ -653,9 +537,9 @@ def delineate_catchment_pyflwdir( flow_dir, ftype="d8", check_ftype=False, - transform=grid.affine, + transform=grid["affine"], ) - bbox = sgeom.box(*grid.bbox) + bbox = sgeom.box(*grid["bbox"]) u, x, y = zip( *[ (u, float(p["x"]), float(p["y"])) @@ -691,21 +575,21 @@ def derive_subbasins_streamorder( gpd.GeoDataFrame: A GeoDataFrame containing polygons. """ # Load and process the DEM - grid, flow_dir, _, _ = load_and_process_dem(fid) + grid, flow_dir, _ = load_and_process_dem(fid) flw = pyflwdir.from_array( flow_dir, ftype="d8", check_ftype=False, - transform=grid.affine, + transform=grid["affine"], ) xy = [ (x_, y_) for x_, y_ in zip(x, y) - if (x_ > grid.bbox[0]) - and (x_ < grid.bbox[2]) - and (y_ > grid.bbox[1]) - and (y_ < grid.bbox[3]) + if (x_ > grid["bbox"][0]) + and (x_ < grid["bbox"][2]) + and (y_ > grid["bbox"][1]) + and (y_ < grid["bbox"][3]) ] idxs, _ = flw.snap(xy=list(zip(*xy))) @@ -726,7 +610,7 @@ def derive_subbasins_streamorder( subbasins = subbasins_ gdf_bas = vectorize( - subbasins.astype(np.int32), 0, flw.transform, grid.crs, name="basin" + subbasins.astype(np.int32), 0, flw.transform, grid["crs"], name="basin" ) return gdf_bas @@ -734,65 +618,57 @@ def derive_subbasins_streamorder( def load_and_process_dem( fid: Path, -) -> tuple[pysheds.sgrid.sGrid, pysheds.sview.Raster, tuple, pysheds.sview.Raster]: +) -> tuple[dict, np.array, np.array]: """Load and condition a DEM. Args: fid (Path): Filepath to the DEM. Returns: - tuple: A tuple containing the grid, flow directions, direction mapping, - and cell slopes. + tuple: A tuple containing the grid, flow directions, and cell slopes. """ - # Initialise pysheds grids - grid = pgrid.Grid.from_raster(str(fid)) - dem = grid.read_raster(str(fid)) - - # Condition the DEM - inflated_dem = condition_dem(grid, dem) + with rst.open(fid, "r") as src: + elevtn = src.read(1).astype(float) + nodata = float(src.nodata) + transform = src.transform + crs = src.crs + + flw = pyflwdir.from_dem( + data=elevtn, + nodata=nodata, + transform=transform, + latlon=crs.is_geographic, + ) + flow_dir = flw.to_array(ftype="d8").astype(int) - # Compute flow directions - flow_dir, dirmap = compute_flow_directions(grid, inflated_dem) + cell_slopes = pyflwdir.dem.slope( + elevtn, + nodata=nodata, + transform=transform, + latlon=crs.is_geographic, + ) - # Calculate slopes - cell_slopes = grid.cell_slopes(dem, flow_dir) + grid = {"affine": transform, "shape": elevtn.shape, "crs": crs, "bbox": src.bounds} - return grid, flow_dir, dirmap, cell_slopes + return grid, flow_dir, cell_slopes -def derive_subcatchments(G: nx.Graph, fid: Path, method="pyflwdir") -> gpd.GeoDataFrame: +def derive_subcatchments(G: nx.Graph, fid: Path) -> gpd.GeoDataFrame: """Derive subcatchments from the nodes on a graph and a DEM. Args: G (nx.Graph): The input graph with nodes containing 'x' and 'y'. fid (Path): Filepath to the DEM. - method (str, optional): The method to use for delineating catchments. - Defaults to 'pyflwdir'. Can also be `pysheds` to use the old - method. Returns: gpd.GeoDataFrame: A GeoDataFrame containing polygons with columns: 'geometry', 'area', 'id', 'width', and 'slope'. """ - if method not in ["pyflwdir", "pysheds"]: - raise ValueError("Invalid method. Must be 'pyflwdir' or 'pysheds'.") - # Load and process the DEM - grid, flow_dir, dirmap, cell_slopes = load_and_process_dem(fid) - - if method == "pysheds": - # Calculate flow accumulations - flow_acc = calculate_flow_accumulation(grid, flow_dir, dirmap) - - # Delineate catchments - polys = delineate_catchment(grid, flow_acc, flow_dir, dirmap, G) - - # Remove intersections - result_polygons = remove_intersections(polys) + grid, flow_dir, cell_slopes = load_and_process_dem(fid) - elif method == "pyflwdir": - # Delineate catchments - result_polygons = delineate_catchment_pyflwdir(grid, flow_dir, G) + # Delineate catchments + result_polygons = delineate_catchment_pyflwdir(grid, flow_dir, G) # Convert to GeoDataFrame polys_gdf = result_polygons.dropna(subset=["geometry"]) diff --git a/swmmanywhere/preprocessing.py b/swmmanywhere/preprocessing.py index e38b79df..47867c1b 100644 --- a/swmmanywhere/preprocessing.py +++ b/swmmanywhere/preprocessing.py @@ -32,7 +32,7 @@ def write_df(df: pd.DataFrame | gpd.GeoDataFrame, fid: Path): """ if fid.suffix in (".geoparquet", ".parquet"): df.to_parquet(fid) - elif fid.suffix == ".json": + elif fid.suffix in (".geojson", ".json"): if isinstance(df, gpd.GeoDataFrame): df.to_file(fid, driver="GeoJSON") else: diff --git a/swmmanywhere/swmmanywhere.py b/swmmanywhere/swmmanywhere.py index 2c712d29..49e89c58 100644 --- a/swmmanywhere/swmmanywhere.py +++ b/swmmanywhere/swmmanywhere.py @@ -78,6 +78,7 @@ def swmmanywhere(config: dict) -> tuple[Path, dict | None]: config["bbox"], config.get("bbox_number", None), config.get("model_number", None), + config.get("extension", "parquet"), **config.get("address_overrides", {}), ) @@ -323,7 +324,7 @@ def check_and_register_custom_graphfcns(config: dict): spec.loader.exec_module(custom_graphfcn_module) # Validate the import - validate_graphfcn_list(config["graphfcn_list"]) + validate_graphfcn_list(config.get("graphfcn_list", [])) return config @@ -354,7 +355,7 @@ def check_and_register_custom_metrics(config: dict): spec.loader.exec_module(custom_metric_module) # Validate metric list - validate_metric_list(config["metric_list"]) + validate_metric_list(config.get("metric_list", [])) return config diff --git a/tests/test_geospatial_utilities.py b/tests/test_geospatial_utilities.py index d70daf81..9ef6df6b 100644 --- a/tests/test_geospatial_utilities.py +++ b/tests/test_geospatial_utilities.py @@ -231,8 +231,8 @@ def test_burn_shape_in_raster(): def test_derive_subcatchments(street_network): """Test the derive_subcatchments function.""" elev_fid = Path(__file__).parent / "test_data" / "elevation.tif" - for method in ["pysheds", "pyflwdir"]: - polys = go.derive_subcatchments(street_network, elev_fid, method=method) + for method in ["pyflwdir"]: + polys = go.derive_subcatchments(street_network, elev_fid) assert "slope" in polys.columns assert "area" in polys.columns assert "geometry" in polys.columns From d04c4988364400439f1eb907c80e9cd842b16acf Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 13:42:39 +0100 Subject: [PATCH 02/20] add whitebox --- dev-requirements.txt | 3 ++ doc-requirements.txt | 3 ++ pyproject.toml | 1 + requirements.txt | 3 ++ swmmanywhere/geospatial_utilities.py | 81 ++++++++++++++++++++++++---- tests/test_geospatial_utilities.py | 17 +++--- tests/test_graph_utilities.py | 6 +-- 7 files changed, 96 insertions(+), 18 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index 4e934777..e0374099 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -50,6 +50,7 @@ click==8.1.7 # pip-tools # planetary-computer # rasterio + # whitebox click-plugins==1.1.1 # via # fiona @@ -312,6 +313,8 @@ virtualenv==20.26.3 # via pre-commit wheel==0.43.0 # via pip-tools +whitebox==2.3.5 + # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/doc-requirements.txt b/doc-requirements.txt index 949da9a8..4e443788 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -56,6 +56,7 @@ click==8.1.7 # mkdocstrings # planetary-computer # rasterio + # whitebox click-plugins==1.1.1 # via # fiona @@ -471,6 +472,8 @@ webencodings==0.5.1 # via # bleach # tinycss2 +whitebox==2.3.5 + # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/pyproject.toml b/pyproject.toml index db61d414..b661d683 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ dependencies = [ "SciPy", "shapely", "tqdm", + "whitebox", "xarray", ] [project.optional-dependencies] diff --git a/requirements.txt b/requirements.txt index 5306ed80..eeee7e6b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,6 +44,7 @@ click==8.1.7 # fiona # planetary-computer # rasterio + # whitebox click-plugins==1.1.1 # via # fiona @@ -249,6 +250,8 @@ tzdata==2024.1 # via pandas urllib3==2.2.2 # via requests +whitebox==2.3.5 + # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 34cb5474..0b5197ca 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -9,6 +9,8 @@ import json import math import os +import shutil +import tempfile from functools import lru_cache from pathlib import Path from typing import List, Optional @@ -27,6 +29,7 @@ from shapely import geometry as sgeom from shapely.strtree import STRtree from tqdm.auto import tqdm +from whitebox import WhiteboxTools from swmmanywhere.logging import logger, verbose @@ -616,13 +619,64 @@ def derive_subbasins_streamorder( return gdf_bas +def flwdir_whitebox(fid: Path) -> np.array: + """Calculate flow direction using WhiteboxTools. + + Args: + fid (Path): Filepath to the DEM. + + Returns: + np.array: Flow directions. + """ + # Initialize WhiteboxTools + wbt = WhiteboxTools() + wbt.verbose = False + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + # Set the working directory + wbt.work_dir = temp_dir + + # Copy raster to working directory + dem = str(temp_path / "dem.tif") + shutil.copy(fid, dem) + + # Condition + breached_dem = str(temp_path / "breached_dem.tif") + wbt.breach_depressions(dem, breached_dem) + + # Calculate flow direction using the D inf algorithm + fdir = str(temp_path / "flow_direction.tif") + wbt.d8_pointer(breached_dem, fdir) + + with rst.open(fdir) as src: + flow_dir = src.read(1) + + return flow_dir + + +def whitebox_to_pyflwdir_mapping(x: int) -> int: + """Map WhiteboxTools flow directions to pyflwdir flow directions. + + Args: + x (int): WhiteboxTools flow direction. + + Returns: + int: pyflwdir flow direction. + """ + mapping = {1: 128, 2: 1, 4: 2, 8: 4, 16: 8, 32: 16, 64: 32, 128: 64} + return mapping.get(x, x) + + def load_and_process_dem( fid: Path, + method: str = "whitebox", ) -> tuple[dict, np.array, np.array]: """Load and condition a DEM. Args: fid (Path): Filepath to the DEM. + method (str, optional): The method to use for conditioning. Defaults to + "whitebox". Returns: tuple: A tuple containing the grid, flow directions, and cell slopes. @@ -633,13 +687,19 @@ def load_and_process_dem( transform = src.transform crs = src.crs - flw = pyflwdir.from_dem( - data=elevtn, - nodata=nodata, - transform=transform, - latlon=crs.is_geographic, - ) - flow_dir = flw.to_array(ftype="d8").astype(int) + if method == "whitebox": + flow_dir = flwdir_whitebox(fid) + flow_dir = np.vectorize(whitebox_to_pyflwdir_mapping)(flow_dir) + elif method == "pyflwdir": + flw = pyflwdir.from_dem( + data=elevtn, + nodata=nodata, + transform=transform, + latlon=crs.is_geographic, + ) + flow_dir = flw.to_array(ftype="d8").astype(int) + else: + raise ValueError("Method must be 'whitebox' or 'pyflwdir'.") cell_slopes = pyflwdir.dem.slope( elevtn, @@ -653,19 +713,22 @@ def load_and_process_dem( return grid, flow_dir, cell_slopes -def derive_subcatchments(G: nx.Graph, fid: Path) -> gpd.GeoDataFrame: +def derive_subcatchments( + G: nx.Graph, fid: Path, method: str = "whitebox" +) -> gpd.GeoDataFrame: """Derive subcatchments from the nodes on a graph and a DEM. Args: G (nx.Graph): The input graph with nodes containing 'x' and 'y'. fid (Path): Filepath to the DEM. + method (str, optional): The method to use for conditioning. Returns: gpd.GeoDataFrame: A GeoDataFrame containing polygons with columns: 'geometry', 'area', 'id', 'width', and 'slope'. """ # Load and process the DEM - grid, flow_dir, cell_slopes = load_and_process_dem(fid) + grid, flow_dir, cell_slopes = load_and_process_dem(fid, method) # Delineate catchments result_polygons = delineate_catchment_pyflwdir(grid, flow_dir, G) diff --git a/tests/test_geospatial_utilities.py b/tests/test_geospatial_utilities.py index 9ef6df6b..7d2283c2 100644 --- a/tests/test_geospatial_utilities.py +++ b/tests/test_geospatial_utilities.py @@ -231,8 +231,12 @@ def test_burn_shape_in_raster(): def test_derive_subcatchments(street_network): """Test the derive_subcatchments function.""" elev_fid = Path(__file__).parent / "test_data" / "elevation.tif" - for method in ["pyflwdir"]: - polys = go.derive_subcatchments(street_network, elev_fid) + methods = { + "pyflwdir": {"area": 2498, "slope": 0.1187, "width": 28.202}, + "whitebox": {"area": 2998, "slope": 0.1102, "width": 30.894}, + } + for method, results in methods.items(): + polys = go.derive_subcatchments(street_network, elev_fid, method) assert "slope" in polys.columns assert "area" in polys.columns assert "geometry" in polys.columns @@ -241,13 +245,14 @@ def test_derive_subcatchments(street_network): assert polys.dropna().shape == polys.shape assert polys.crs == street_network.graph["crs"] - # Pyflwdir and pysheds catchment derivation aren't absolutely identical - assert almost_equal(polys.set_index("id").loc[2623975694, "area"], 1499, tol=1) assert almost_equal( - polys.set_index("id").loc[2623975694, "slope"], 0.06145, tol=0.001 + polys.set_index("id").loc[2623975694, "area"], results["area"], tol=1 + ) + assert almost_equal( + polys.set_index("id").loc[2623975694, "slope"], results["slope"], tol=0.001 ) assert almost_equal( - polys.set_index("id").loc[2623975694, "width"], 21.845, tol=0.001 + polys.set_index("id").loc[2623975694, "width"], results["width"], tol=0.001 ) diff --git a/tests/test_graph_utilities.py b/tests/test_graph_utilities.py index fc744973..dd14868a 100644 --- a/tests/test_graph_utilities.py +++ b/tests/test_graph_utilities.py @@ -687,7 +687,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 9 + assert len(G_.edges) == 7 # Test default clipping streamorder subcatchment_derivation = parameters.SubcatchmentDerivation() @@ -695,7 +695,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 4 + assert len(G_.edges) == 2 # Test clipping subcatchment_derivation = parameters.SubcatchmentDerivation( @@ -706,7 +706,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 30 + assert len(G_.edges) == 31 # Test clipping with different params subcatchment_derivation = parameters.SubcatchmentDerivation( From 518f81385c23a7b6f7c880bac5a40e331933af68 Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 15:25:48 +0100 Subject: [PATCH 03/20] Attempt to fix tempfile issue --- swmmanywhere/geospatial_utilities.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 0b5197ca..98fdc30b 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -648,6 +648,12 @@ def flwdir_whitebox(fid: Path) -> np.array: fdir = str(temp_path / "flow_direction.tif") wbt.d8_pointer(breached_dem, fdir) + # Force the OS to flush filesystem buffers + os.fsync(Path(fdir).open("r")) + + if not Path(fdir).exists(): + raise ValueError("Flow direction raster not created.") + with rst.open(fdir) as src: flow_dir = src.read(1) From ac7cbc7b82e540b7e2e6eda5c9acc711243d907e Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 15:31:40 +0100 Subject: [PATCH 04/20] try fix tempfile issue --- swmmanywhere/geospatial_utilities.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 98fdc30b..207ce0a6 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -13,6 +13,7 @@ import tempfile from functools import lru_cache from pathlib import Path +from time import sleep from typing import List, Optional import geopandas as gpd @@ -648,8 +649,8 @@ def flwdir_whitebox(fid: Path) -> np.array: fdir = str(temp_path / "flow_direction.tif") wbt.d8_pointer(breached_dem, fdir) - # Force the OS to flush filesystem buffers - os.fsync(Path(fdir).open("r")) + # Give time for filesystem to flush + sleep(5) if not Path(fdir).exists(): raise ValueError("Flow direction raster not created.") From 8157a059b9c4f5a0276b5672204b7301ec7373de Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 15:57:21 +0100 Subject: [PATCH 05/20] provide default instead of fixing --- swmmanywhere/geospatial_utilities.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 207ce0a6..46f2da40 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -13,7 +13,6 @@ import tempfile from functools import lru_cache from pathlib import Path -from time import sleep from typing import List, Optional import geopandas as gpd @@ -620,18 +619,18 @@ def derive_subbasins_streamorder( return gdf_bas -def flwdir_whitebox(fid: Path) -> np.array: +def flwdir_whitebox(fid: Path) -> np.array | None: """Calculate flow direction using WhiteboxTools. Args: fid (Path): Filepath to the DEM. Returns: - np.array: Flow directions. + np.array: Flow directions. Seems to occasionally fail (randomly) """ # Initialize WhiteboxTools wbt = WhiteboxTools() - wbt.verbose = False + wbt.verbose = verbose() with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) # Set the working directory @@ -649,15 +648,13 @@ def flwdir_whitebox(fid: Path) -> np.array: fdir = str(temp_path / "flow_direction.tif") wbt.d8_pointer(breached_dem, fdir) - # Give time for filesystem to flush - sleep(5) - if not Path(fdir).exists(): - raise ValueError("Flow direction raster not created.") + logger.warning("Flow direction raster not created.") + return None with rst.open(fdir) as src: flow_dir = src.read(1) - + flow_dir = np.vectorize(whitebox_to_pyflwdir_mapping)(flow_dir) return flow_dir @@ -694,10 +691,11 @@ def load_and_process_dem( transform = src.transform crs = src.crs + flow_dir = None if method == "whitebox": flow_dir = flwdir_whitebox(fid) - flow_dir = np.vectorize(whitebox_to_pyflwdir_mapping)(flow_dir) - elif method == "pyflwdir": + + if (method == "pyflwdir") or (flow_dir is None): flw = pyflwdir.from_dem( data=elevtn, nodata=nodata, From f63c20714ad2dac310a0ff75c558380533a71e17 Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 16:09:16 +0100 Subject: [PATCH 06/20] Fall back on pyflwdir by default --- swmmanywhere/geospatial_utilities.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 46f2da40..8118fc45 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -691,6 +691,9 @@ def load_and_process_dem( transform = src.transform crs = src.crs + if method not in ["whitebox", "pyflwdir"]: + raise ValueError("Method must be 'whitebox' or 'pyflwdir'.") + flow_dir = None if method == "whitebox": flow_dir = flwdir_whitebox(fid) @@ -703,8 +706,6 @@ def load_and_process_dem( latlon=crs.is_geographic, ) flow_dir = flw.to_array(ftype="d8").astype(int) - else: - raise ValueError("Method must be 'whitebox' or 'pyflwdir'.") cell_slopes = pyflwdir.dem.slope( elevtn, From debe9444b48fa77218ebe4702ad226820f73a986 Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 16:10:12 +0100 Subject: [PATCH 07/20] one more go at debug --- swmmanywhere/geospatial_utilities.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 8118fc45..61cf7fed 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -649,8 +649,9 @@ def flwdir_whitebox(fid: Path) -> np.array | None: wbt.d8_pointer(breached_dem, fdir) if not Path(fdir).exists(): - logger.warning("Flow direction raster not created.") - return None + # logger.warning("Flow direction raster not created.") + # return None + raise ValueError("Flow direction raster not created.") with rst.open(fdir) as src: flow_dir = src.read(1) From e045a2c069d2dcf5b28f0ca9128af6cba79f5c8a Mon Sep 17 00:00:00 2001 From: Dobson Date: Fri, 23 Aug 2024 16:51:42 +0100 Subject: [PATCH 08/20] Fall back on pyflwdir by default --- swmmanywhere/geospatial_utilities.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 61cf7fed..239c2994 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -649,9 +649,9 @@ def flwdir_whitebox(fid: Path) -> np.array | None: wbt.d8_pointer(breached_dem, fdir) if not Path(fdir).exists(): - # logger.warning("Flow direction raster not created.") - # return None - raise ValueError("Flow direction raster not created.") + logger.warning("Flow direction raster not created.") + return None + # raise ValueError("Flow direction raster not created.") with rst.open(fdir) as src: flow_dir = src.read(1) @@ -692,7 +692,7 @@ def load_and_process_dem( transform = src.transform crs = src.crs - if method not in ["whitebox", "pyflwdir"]: + if method not in ("whitebox", "pyflwdir"): raise ValueError("Method must be 'whitebox' or 'pyflwdir'.") flow_dir = None From 7c32b7e675e3c0089df7ce0d0f38c014dccc4cb4 Mon Sep 17 00:00:00 2001 From: Dobson Date: Tue, 27 Aug 2024 14:23:36 +0100 Subject: [PATCH 09/20] Fix bug --- swmmanywhere/geospatial_utilities.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 239c2994..43e10f5a 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -631,7 +631,7 @@ def flwdir_whitebox(fid: Path) -> np.array | None: # Initialize WhiteboxTools wbt = WhiteboxTools() wbt.verbose = verbose() - with tempfile.TemporaryDirectory() as temp_dir: + with tempfile.TemporaryDirectory(dir=fid.parent) as temp_dir: temp_path = Path(temp_dir) # Set the working directory wbt.work_dir = temp_dir @@ -649,9 +649,9 @@ def flwdir_whitebox(fid: Path) -> np.array | None: wbt.d8_pointer(breached_dem, fdir) if not Path(fdir).exists(): - logger.warning("Flow direction raster not created.") - return None - # raise ValueError("Flow direction raster not created.") + # logger.warning("Flow direction raster not created.") + # return None + raise ValueError("Flow direction raster not created.") with rst.open(fdir) as src: flow_dir = src.read(1) From e23ae8b81c43b75e14a60c0232999b0e49207f06 Mon Sep 17 00:00:00 2001 From: Dobson Date: Tue, 27 Aug 2024 14:37:50 +0100 Subject: [PATCH 10/20] fix the fix --- swmmanywhere/geospatial_utilities.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 43e10f5a..17f451fd 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -631,7 +631,7 @@ def flwdir_whitebox(fid: Path) -> np.array | None: # Initialize WhiteboxTools wbt = WhiteboxTools() wbt.verbose = verbose() - with tempfile.TemporaryDirectory(dir=fid.parent) as temp_dir: + with tempfile.TemporaryDirectory(dir=str(fid.parent)) as temp_dir: temp_path = Path(temp_dir) # Set the working directory wbt.work_dir = temp_dir From 3130f30784b3d5ec102ed559f35b26a88a2bad6a Mon Sep 17 00:00:00 2001 From: Dobson Date: Tue, 27 Aug 2024 14:47:06 +0100 Subject: [PATCH 11/20] parameterise test --- tests/test_geospatial_utilities.py | 41 +++++++++++++++--------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/tests/test_geospatial_utilities.py b/tests/test_geospatial_utilities.py index 7d2283c2..806d114d 100644 --- a/tests/test_geospatial_utilities.py +++ b/tests/test_geospatial_utilities.py @@ -228,32 +228,33 @@ def test_burn_shape_in_raster(): new_raster_fid.unlink(missing_ok=True) -def test_derive_subcatchments(street_network): +@pytest.mark.parametrize("method", ["pyflwdir", "whitebox"]) +def test_derive_subcatchments(street_network, method): """Test the derive_subcatchments function.""" elev_fid = Path(__file__).parent / "test_data" / "elevation.tif" methods = { "pyflwdir": {"area": 2498, "slope": 0.1187, "width": 28.202}, "whitebox": {"area": 2998, "slope": 0.1102, "width": 30.894}, } - for method, results in methods.items(): - polys = go.derive_subcatchments(street_network, elev_fid, method) - assert "slope" in polys.columns - assert "area" in polys.columns - assert "geometry" in polys.columns - assert "id" in polys.columns - assert polys.shape[0] > 0 - assert polys.dropna().shape == polys.shape - assert polys.crs == street_network.graph["crs"] - - assert almost_equal( - polys.set_index("id").loc[2623975694, "area"], results["area"], tol=1 - ) - assert almost_equal( - polys.set_index("id").loc[2623975694, "slope"], results["slope"], tol=0.001 - ) - assert almost_equal( - polys.set_index("id").loc[2623975694, "width"], results["width"], tol=0.001 - ) + results = methods[method] + polys = go.derive_subcatchments(street_network, elev_fid, method) + assert "slope" in polys.columns + assert "area" in polys.columns + assert "geometry" in polys.columns + assert "id" in polys.columns + assert polys.shape[0] > 0 + assert polys.dropna().shape == polys.shape + assert polys.crs == street_network.graph["crs"] + + assert almost_equal( + polys.set_index("id").loc[2623975694, "area"], results["area"], tol=1 + ) + assert almost_equal( + polys.set_index("id").loc[2623975694, "slope"], results["slope"], tol=0.001 + ) + assert almost_equal( + polys.set_index("id").loc[2623975694, "width"], results["width"], tol=0.001 + ) def test_derive_rc(street_network): From 7aaf2dcbd8bc952f56cb73893b96283af0647432 Mon Sep 17 00:00:00 2001 From: Dobson Date: Tue, 27 Aug 2024 14:59:39 +0100 Subject: [PATCH 12/20] Chagne grid to struct --- swmmanywhere/geospatial_utilities.py | 48 +++++++++++++++++++--------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 17f451fd..a96e640d 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -460,14 +460,14 @@ def attach_unconnected_subareas( def calculate_slope( - polys_gdf: gpd.GeoDataFrame, grid: dict, cell_slopes: np.ndarray + polys_gdf: gpd.GeoDataFrame, grid: Grid, cell_slopes: np.ndarray ) -> gpd.GeoDataFrame: """Calculate the average slope of each polygon. Args: polys_gdf (gpd.GeoDataFrame): A GeoDataFrame containing polygons with columns: 'geometry', 'area', and 'id'. - grid (dict): Information of the raster (affine, shape, crs, bbox) + grid (Grid): Information of the raster (affine, shape, crs, bbox) cell_slopes (np.ndarray): The slopes of each cell in the grid. Returns: @@ -477,7 +477,7 @@ def calculate_slope( polys_gdf["slope"] = None for idx, row in polys_gdf.iterrows(): mask = features.geometry_mask( - [row.geometry], grid["shape"], grid["affine"], invert=True + [row.geometry], grid.shape, grid.affine, invert=True ) average_slope = cell_slopes[mask].mean().mean() polys_gdf.loc[idx, "slope"] = max(float(average_slope), 0) @@ -520,7 +520,7 @@ def vectorize( def delineate_catchment_pyflwdir( - grid: dict, flow_dir: np.array, G: nx.Graph + grid: Grid, flow_dir: np.array, G: nx.Graph ) -> gpd.GeoDataFrame: """Derive subcatchments from the nodes on a graph and a DEM. @@ -528,7 +528,7 @@ def delineate_catchment_pyflwdir( faster than delineate_catchment. Args: - grid (dict): Information of the raster (affine, shape, crs, bbox). + grid (Grid): Information of the raster (affine, shape, crs, bbox). flow_dir (np.array): Flow directions. G (nx.Graph): The input graph with nodes containing 'x' and 'y'. @@ -540,9 +540,9 @@ def delineate_catchment_pyflwdir( flow_dir, ftype="d8", check_ftype=False, - transform=grid["affine"], + transform=grid.affine, ) - bbox = sgeom.box(*grid["bbox"]) + bbox = sgeom.box(*grid.bbox) u, x, y = zip( *[ (u, float(p["x"]), float(p["y"])) @@ -584,15 +584,15 @@ def derive_subbasins_streamorder( flow_dir, ftype="d8", check_ftype=False, - transform=grid["affine"], + transform=grid.affine, ) xy = [ (x_, y_) for x_, y_ in zip(x, y) - if (x_ > grid["bbox"][0]) - and (x_ < grid["bbox"][2]) - and (y_ > grid["bbox"][1]) - and (y_ < grid["bbox"][3]) + if (x_ > grid.bbox[0]) + and (x_ < grid.bbox[2]) + and (y_ > grid.bbox[1]) + and (y_ < grid.bbox[3]) ] idxs, _ = flw.snap(xy=list(zip(*xy))) @@ -613,7 +613,7 @@ def derive_subbasins_streamorder( subbasins = subbasins_ gdf_bas = vectorize( - subbasins.astype(np.int32), 0, flw.transform, grid["crs"], name="basin" + subbasins.astype(np.int32), 0, flw.transform, grid.crs, name="basin" ) return gdf_bas @@ -672,10 +672,28 @@ def whitebox_to_pyflwdir_mapping(x: int) -> int: return mapping.get(x, x) +class Grid: + """A class to represent a grid.""" + + def __init__(self, affine: rst.Affine, shape: tuple, crs: int, bbox: tuple): + """Initialize the Grid class. + + Args: + affine (rst.Affine): The affine transformation. + shape (tuple): The shape of the grid. + crs (int): The CRS of the grid. + bbox (tuple): The bounding box of the grid. + """ + self.affine = affine + self.shape = shape + self.crs = crs + self.bbox = bbox + + def load_and_process_dem( fid: Path, method: str = "whitebox", -) -> tuple[dict, np.array, np.array]: +) -> tuple[Grid, np.array, np.array]: """Load and condition a DEM. Args: @@ -715,7 +733,7 @@ def load_and_process_dem( latlon=crs.is_geographic, ) - grid = {"affine": transform, "shape": elevtn.shape, "crs": crs, "bbox": src.bounds} + grid = Grid(transform, elevtn.shape, crs, src.bounds) return grid, flow_dir, cell_slopes From e8567fb36158d848fde6ed5d40c979be192ccc41 Mon Sep 17 00:00:00 2001 From: Dobson Date: Tue, 27 Aug 2024 16:03:22 +0100 Subject: [PATCH 13/20] Fix mapping --- swmmanywhere/geospatial_utilities.py | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index a96e640d..f721da2b 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -619,14 +619,14 @@ def derive_subbasins_streamorder( return gdf_bas -def flwdir_whitebox(fid: Path) -> np.array | None: +def flwdir_whitebox(fid: Path) -> np.array: """Calculate flow direction using WhiteboxTools. Args: fid (Path): Filepath to the DEM. Returns: - np.array: Flow directions. Seems to occasionally fail (randomly) + np.array: Flow directions. """ # Initialize WhiteboxTools wbt = WhiteboxTools() @@ -649,27 +649,16 @@ def flwdir_whitebox(fid: Path) -> np.array | None: wbt.d8_pointer(breached_dem, fdir) if not Path(fdir).exists(): - # logger.warning("Flow direction raster not created.") - # return None raise ValueError("Flow direction raster not created.") with rst.open(fdir) as src: flow_dir = src.read(1) - flow_dir = np.vectorize(whitebox_to_pyflwdir_mapping)(flow_dir) - return flow_dir - - -def whitebox_to_pyflwdir_mapping(x: int) -> int: - """Map WhiteboxTools flow directions to pyflwdir flow directions. - - Args: - x (int): WhiteboxTools flow direction. - Returns: - int: pyflwdir flow direction. - """ + # Adjust mapping from WhiteboxTools to pyflwdir mapping = {1: 128, 2: 1, 4: 2, 8: 4, 16: 8, 32: 16, 64: 32, 128: 64} - return mapping.get(x, x) + get_flow_dir = np.vectorize(mapping.get, excluded=["default"]) + flow_dir = get_flow_dir(flow_dir, np.nan) + return flow_dir class Grid: @@ -713,11 +702,9 @@ def load_and_process_dem( if method not in ("whitebox", "pyflwdir"): raise ValueError("Method must be 'whitebox' or 'pyflwdir'.") - flow_dir = None if method == "whitebox": flow_dir = flwdir_whitebox(fid) - - if (method == "pyflwdir") or (flow_dir is None): + elif method == "pyflwdir": flw = pyflwdir.from_dem( data=elevtn, nodata=nodata, From 2ce36bdaacc1b611d9fc14037086c8fd64b587fe Mon Sep 17 00:00:00 2001 From: barneydobson Date: Wed, 28 Aug 2024 09:06:58 +0100 Subject: [PATCH 14/20] Add verbose whitebox test --- swmmanywhere/logging.py | 5 +++++ tests/test_geospatial_utilities.py | 6 +++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/swmmanywhere/logging.py b/swmmanywhere/logging.py index 29af226b..cae70f24 100644 --- a/swmmanywhere/logging.py +++ b/swmmanywhere/logging.py @@ -25,6 +25,11 @@ def verbose() -> bool: return os.getenv("SWMMANYWHERE_VERBOSE", "false").lower() == "true" +def set_verbose(verbose: bool): + """Set the verbosity.""" + os.environ["SWMMANYWHERE_VERBOSE"] = str(verbose).lower() + + def dynamic_filter(record): """A dynamic filter.""" return verbose() diff --git a/tests/test_geospatial_utilities.py b/tests/test_geospatial_utilities.py index 806d114d..c987bdf0 100644 --- a/tests/test_geospatial_utilities.py +++ b/tests/test_geospatial_utilities.py @@ -15,6 +15,7 @@ from swmmanywhere import geospatial_utilities as go from swmmanywhere import graph_utilities as ge +from swmmanywhere.logging import set_verbose from swmmanywhere.misc.debug_derive_rc import derive_rc_alt @@ -229,8 +230,11 @@ def test_burn_shape_in_raster(): @pytest.mark.parametrize("method", ["pyflwdir", "whitebox"]) -def test_derive_subcatchments(street_network, method): +@pytest.mark.parametrize("verbose", [True, False]) +def test_derive_subcatchments(street_network, method, verbose): """Test the derive_subcatchments function.""" + set_verbose(verbose) + elev_fid = Path(__file__).parent / "test_data" / "elevation.tif" methods = { "pyflwdir": {"area": 2498, "slope": 0.1187, "width": 28.202}, From 33ce12bcd5f4dd5ab2d9edbe229342949a79093d Mon Sep 17 00:00:00 2001 From: barneydobson Date: Wed, 28 Aug 2024 09:58:05 +0100 Subject: [PATCH 15/20] Set max wbt procs --- swmmanywhere/geospatial_utilities.py | 1 + 1 file changed, 1 insertion(+) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index f721da2b..74cb8ac2 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -630,6 +630,7 @@ def flwdir_whitebox(fid: Path) -> np.array: """ # Initialize WhiteboxTools wbt = WhiteboxTools() + wbt.set_max_procs(1) wbt.verbose = verbose() with tempfile.TemporaryDirectory(dir=str(fid.parent)) as temp_dir: temp_path = Path(temp_dir) From 0f7ab3cd73804abe99285d01d29cb2561936ed59 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Thu, 29 Aug 2024 09:28:23 +0100 Subject: [PATCH 16/20] Use Taher's custom whitebox wrapper --- swmmanywhere/geospatial_utilities.py | 30 ++--- swmmanywhere/whitebox_utils.py | 162 +++++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 13 deletions(-) create mode 100644 swmmanywhere/whitebox_utils.py diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 74cb8ac2..cec8c314 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -29,9 +29,9 @@ from shapely import geometry as sgeom from shapely.strtree import STRtree from tqdm.auto import tqdm -from whitebox import WhiteboxTools from swmmanywhere.logging import logger, verbose +from swmmanywhere.whitebox_utils import whitebox_tools os.environ["NUMBA_NUM_THREADS"] = "1" import pyflwdir # noqa: E402 @@ -629,26 +629,30 @@ def flwdir_whitebox(fid: Path) -> np.array: np.array: Flow directions. """ # Initialize WhiteboxTools - wbt = WhiteboxTools() - wbt.set_max_procs(1) - wbt.verbose = verbose() with tempfile.TemporaryDirectory(dir=str(fid.parent)) as temp_dir: temp_path = Path(temp_dir) - # Set the working directory - wbt.work_dir = temp_dir # Copy raster to working directory - dem = str(temp_path / "dem.tif") + dem = temp_path / "dem.tif" shutil.copy(fid, dem) # Condition - breached_dem = str(temp_path / "breached_dem.tif") - wbt.breach_depressions(dem, breached_dem) - - # Calculate flow direction using the D inf algorithm - fdir = str(temp_path / "flow_direction.tif") - wbt.d8_pointer(breached_dem, fdir) + whitebox_tools( + "BreachDepressionsLeastCost", + ["-i=dem.tif", "-o=dem_corr.tif"], + work_dir=temp_path, + verbose=verbose(), + wbt_root=temp_path / "WBT", + ) + whitebox_tools( + "D8Pointer", + ["-i=dem_corr.tif", "-o=fdir.tif"], + work_dir=temp_path, + verbose=verbose(), + wbt_root=temp_path / "WBT", + ) + fdir = temp_path / "fdir.tif" if not Path(fdir).exists(): raise ValueError("Flow direction raster not created.") diff --git a/swmmanywhere/whitebox_utils.py b/swmmanywhere/whitebox_utils.py new file mode 100644 index 00000000..3795eab5 --- /dev/null +++ b/swmmanywhere/whitebox_utils.py @@ -0,0 +1,162 @@ +"""Whitebox utilities wrapper. + +@author: cheginit +""" +from __future__ import annotations + +import platform +import shutil +import stat +import subprocess +import tempfile +import zipfile +from pathlib import Path + +import requests + + +def _download_wbt(wbt_root: str | Path = "WBT") -> None: + """Download the WhiteboxTools executable for the current platform.""" + base_url = "https://www.whiteboxgeo.com/WBT_{}/WhiteboxTools_{}.zip" + + system = platform.system() + if system not in ("Windows", "Darwin", "Linux"): + raise ValueError(f"Unsupported operating system: {system}") + + if system == "Windows": + platform_suffix = "win_amd64" + elif system == "Darwin": + if platform.machine() == "arm64": + platform_suffix = "darwin_m_series" + else: + platform_suffix = "darwin_amd64" + elif system == "Linux": + if "musl" in platform.libc_ver()[0].lower(): + platform_suffix = "linux_musl" + else: + platform_suffix = "linux_amd64" + + url = base_url.format(system, platform_suffix) + wbt_root = Path(wbt_root) + + exe_name = ( + "whitebox_tools.exe" if platform.system() == "Windows" else "whitebox_tools" + ) + if (wbt_root / exe_name).exists(): + shutil.rmtree(wbt_root) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + zip_path = temp_path / "whitebox_tools.zip" + + try: + response = requests.get(url) + response.raise_for_status() + zip_path.write_bytes(response.content) + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Failed to download WhiteboxTools: {e!s}") from e + + try: + with zipfile.ZipFile(zip_path, "r") as zip_ref: + for file_info in zip_ref.infolist(): + if "/WBT/" in file_info.filename: + extracted_path = Path( + *Path(file_info.filename).parts[ + Path(file_info.filename).parts.index("WBT") + 1 : + ] + ) + if extracted_path.parts: + source = zip_ref.extract(file_info, path=temp_dir) + dest = wbt_root / extracted_path + dest.parent.mkdir(parents=True, exist_ok=True) + shutil.move(source, dest) + except zipfile.BadZipFile: + raise RuntimeError("Downloaded file is not a valid zip file.") + + # Set executable permissions for non-Windows platforms + if system != "Windows": + executable_names = ["whitebox_tools", "whitebox_runner"] + for exec_name in executable_names: + exec_path = wbt_root / exec_name + if exec_path.exists(): + exec_path.chmod( + exec_path.stat().st_mode + | stat.S_IXUSR + | stat.S_IXGRP + | stat.S_IXOTH + ) + + +def whitebox_tools( + tool_name: str, + args: list[str] | tuple[str, ...], + wbt_root: str | Path = "WBT", + work_dir: str | Path = "", + verbose: bool = False, + compress_rasters: bool = False, + refresh_download: bool = False, +) -> None: + """Run a WhiteboxTools (not Whitebox Runner) tool with specified arguments. + + Parameters + ---------- + tool_name : str + Name of the WhiteboxTools to run. + args : list or tuple + List of arguments for the tool. + wbt_root : str or Path, optional + Path to the root directory containing the Whitebox executables + (default is "WBT"). + work_dir : str or Path, optional + Working directory for the tool (default is current directory). + verbose : bool, optional + Whether to print verbose output (default is False). + compress_rasters : bool, optional + Whether to compress output rasters (default is False). + refresh_download : bool, optional + Whether to refresh the download if WhiteboxTools is found (default is False). + + Raises: + ------ + subprocess.CalledProcessError + If the tool execution fails. + Exception + For any other unexpected errors. + + Notes: + ----- + This function will run the specified WhiteboxTools tool and handle its output. + If verbose is True, all output will be printed. + """ + wbt_root = Path(wbt_root) + work_dir = Path(work_dir) if work_dir else Path.cwd() + + exe_name = ( + "whitebox_tools.exe" if platform.system() == "Windows" else "whitebox_tools" + ) + exe_path = wbt_root / exe_name + + if not exe_path.exists() or refresh_download: + _download_wbt(wbt_root) + + command = [ + str(exe_path), + f"--run={tool_name}", + f"--wd={work_dir}", + "-v" if verbose else "-v=false", + f"--compress_rasters={'true' if compress_rasters else 'false'}", + ] + command.extend(args) + + if verbose: + print(" ".join(map(str, command))) + + try: + process = subprocess.run(command, check=True, text=True, capture_output=True) + if verbose: + print(process.stdout) + except subprocess.CalledProcessError as e: + print(f"Error output: {e.stdout}") + raise Exception(f"Error running tool: {e}") + except Exception as e: + raise Exception(f"Unexpected error: {e}") From 26201fb72e87b87e5be7a2979c77bdc5838a7399 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Thu, 29 Aug 2024 09:35:21 +0100 Subject: [PATCH 17/20] Update reqs and fix tests --- dev-requirements.txt | 3 --- doc-requirements.txt | 3 --- pyproject.toml | 1 - requirements.txt | 3 --- tests/test_graph_utilities.py | 8 ++++---- 5 files changed, 4 insertions(+), 14 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index e0374099..4e934777 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -50,7 +50,6 @@ click==8.1.7 # pip-tools # planetary-computer # rasterio - # whitebox click-plugins==1.1.1 # via # fiona @@ -313,8 +312,6 @@ virtualenv==20.26.3 # via pre-commit wheel==0.43.0 # via pip-tools -whitebox==2.3.5 - # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/doc-requirements.txt b/doc-requirements.txt index 4e443788..949da9a8 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -56,7 +56,6 @@ click==8.1.7 # mkdocstrings # planetary-computer # rasterio - # whitebox click-plugins==1.1.1 # via # fiona @@ -472,8 +471,6 @@ webencodings==0.5.1 # via # bleach # tinycss2 -whitebox==2.3.5 - # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/pyproject.toml b/pyproject.toml index b661d683..db61d414 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,6 @@ dependencies = [ "SciPy", "shapely", "tqdm", - "whitebox", "xarray", ] [project.optional-dependencies] diff --git a/requirements.txt b/requirements.txt index eeee7e6b..5306ed80 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,6 @@ click==8.1.7 # fiona # planetary-computer # rasterio - # whitebox click-plugins==1.1.1 # via # fiona @@ -250,8 +249,6 @@ tzdata==2024.1 # via pandas urllib3==2.2.2 # via requests -whitebox==2.3.5 - # via swmmanywhere (pyproject.toml) win32-setctime==1.1.0 # via loguru xarray==2024.6.0 diff --git a/tests/test_graph_utilities.py b/tests/test_graph_utilities.py index 447bb694..90c9c5fb 100644 --- a/tests/test_graph_utilities.py +++ b/tests/test_graph_utilities.py @@ -688,7 +688,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 7 + assert len(G_.edges) == 10 # Test default clipping streamorder subcatchment_derivation = parameters.SubcatchmentDerivation() @@ -696,7 +696,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 2 + assert len(G_.edges) == 10 # Test clipping subcatchment_derivation = parameters.SubcatchmentDerivation( @@ -707,7 +707,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 31 + assert len(G_.edges) == 30 # Test clipping with different params subcatchment_derivation = parameters.SubcatchmentDerivation( @@ -718,7 +718,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 38 + assert len(G_.edges) == 39 # Test no cuts subcatchment_derivation = parameters.SubcatchmentDerivation( From dcc74b37295bd570f20f8246e91c0129a15f064b Mon Sep 17 00:00:00 2001 From: barneydobson Date: Thu, 29 Aug 2024 12:53:42 +0100 Subject: [PATCH 18/20] Revert to BreachDepressions and tests --- swmmanywhere/geospatial_utilities.py | 2 +- tests/test_graph_utilities.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index cec8c314..0775b87c 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -638,7 +638,7 @@ def flwdir_whitebox(fid: Path) -> np.array: # Condition whitebox_tools( - "BreachDepressionsLeastCost", + "BreachDepressions", ["-i=dem.tif", "-o=dem_corr.tif"], work_dir=temp_path, verbose=verbose(), diff --git a/tests/test_graph_utilities.py b/tests/test_graph_utilities.py index 90c9c5fb..447bb694 100644 --- a/tests/test_graph_utilities.py +++ b/tests/test_graph_utilities.py @@ -688,7 +688,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 10 + assert len(G_.edges) == 7 # Test default clipping streamorder subcatchment_derivation = parameters.SubcatchmentDerivation() @@ -696,7 +696,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 10 + assert len(G_.edges) == 2 # Test clipping subcatchment_derivation = parameters.SubcatchmentDerivation( @@ -707,7 +707,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 30 + assert len(G_.edges) == 31 # Test clipping with different params subcatchment_derivation = parameters.SubcatchmentDerivation( @@ -718,7 +718,7 @@ def test_clip_to_catchments(street_network): G_ = gu.clip_to_catchments( G, addresses=addresses, subcatchment_derivation=subcatchment_derivation ) - assert len(G_.edges) == 39 + assert len(G_.edges) == 38 # Test no cuts subcatchment_derivation = parameters.SubcatchmentDerivation( From 89c7f9716f90638e08eb0bc487bd4f4ef97961c8 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Thu, 29 Aug 2024 13:11:47 +0100 Subject: [PATCH 19/20] Finish review comments --- swmmanywhere/geospatial_utilities.py | 2 +- tests/test_geospatial_utilities.py | 21 +++++++++------------ 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 0775b87c..69c34f5e 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -662,7 +662,7 @@ def flwdir_whitebox(fid: Path) -> np.array: # Adjust mapping from WhiteboxTools to pyflwdir mapping = {1: 128, 2: 1, 4: 2, 8: 4, 16: 8, 32: 16, 64: 32, 128: 64} get_flow_dir = np.vectorize(mapping.get, excluded=["default"]) - flow_dir = get_flow_dir(flow_dir, np.nan) + flow_dir = get_flow_dir(flow_dir, 0) return flow_dir diff --git a/tests/test_geospatial_utilities.py b/tests/test_geospatial_utilities.py index c987bdf0..1ffdbe33 100644 --- a/tests/test_geospatial_utilities.py +++ b/tests/test_geospatial_utilities.py @@ -229,18 +229,17 @@ def test_burn_shape_in_raster(): new_raster_fid.unlink(missing_ok=True) -@pytest.mark.parametrize("method", ["pyflwdir", "whitebox"]) +@pytest.mark.parametrize( + "method,area,slope,width", + [("pyflwdir", 2498, 0.1187, 28.202), ("whitebox", 2998, 0.1102, 30.894)], +) @pytest.mark.parametrize("verbose", [True, False]) -def test_derive_subcatchments(street_network, method, verbose): +def test_derive_subcatchments(street_network, method, area, slope, width, verbose): """Test the derive_subcatchments function.""" set_verbose(verbose) elev_fid = Path(__file__).parent / "test_data" / "elevation.tif" - methods = { - "pyflwdir": {"area": 2498, "slope": 0.1187, "width": 28.202}, - "whitebox": {"area": 2998, "slope": 0.1102, "width": 30.894}, - } - results = methods[method] + polys = go.derive_subcatchments(street_network, elev_fid, method) assert "slope" in polys.columns assert "area" in polys.columns @@ -250,14 +249,12 @@ def test_derive_subcatchments(street_network, method, verbose): assert polys.dropna().shape == polys.shape assert polys.crs == street_network.graph["crs"] + assert almost_equal(polys.set_index("id").loc[2623975694, "area"], area, tol=1) assert almost_equal( - polys.set_index("id").loc[2623975694, "area"], results["area"], tol=1 - ) - assert almost_equal( - polys.set_index("id").loc[2623975694, "slope"], results["slope"], tol=0.001 + polys.set_index("id").loc[2623975694, "slope"], slope, tol=0.001 ) assert almost_equal( - polys.set_index("id").loc[2623975694, "width"], results["width"], tol=0.001 + polys.set_index("id").loc[2623975694, "width"], width, tol=0.001 ) From e1b71714c8f5eab8a3626f3775290b247e9f7ac3 Mon Sep 17 00:00:00 2001 From: Dobson Date: Mon, 2 Sep 2024 08:45:38 +0100 Subject: [PATCH 20/20] Use cheginits pywbt --- dev-requirements.txt | 2 + doc-requirements.txt | 2 + pyproject.toml | 1 + requirements.txt | 2 + swmmanywhere/geospatial_utilities.py | 17 ++- swmmanywhere/whitebox_utils.py | 162 --------------------------- 6 files changed, 14 insertions(+), 172 deletions(-) delete mode 100644 swmmanywhere/whitebox_utils.py diff --git a/dev-requirements.txt b/dev-requirements.txt index 4e934777..6fa8e17b 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -243,6 +243,8 @@ pytz==2024.1 # multiurl # pandas # planetary-computer +pywbt==0.1.1 + # via swmmanywhere (pyproject.toml) pyyaml==6.0.1 # via # pre-commit diff --git a/doc-requirements.txt b/doc-requirements.txt index e17a0c16..8f847875 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -360,6 +360,8 @@ pytz==2024.1 # multiurl # pandas # planetary-computer +pywbt==0.1.1 + # via swmmanywhere (pyproject.toml) pywin32==306 # via jupyter-core pyyaml==6.0.1 diff --git a/pyproject.toml b/pyproject.toml index db61d414..98a9d170 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "pyflwdir", "pystac_client", "pyswmm", + "pywbt", "PyYAML", "rasterio", "rioxarray", diff --git a/requirements.txt b/requirements.txt index 5306ed80..d2d4fce5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -189,6 +189,8 @@ pytz==2024.1 # multiurl # pandas # planetary-computer +pywbt==0.1.1 + # via swmmanywhere (pyproject.toml) pyyaml==6.0.1 # via swmmanywhere (pyproject.toml) rasterio==1.3.10 diff --git a/swmmanywhere/geospatial_utilities.py b/swmmanywhere/geospatial_utilities.py index 69c34f5e..c990d970 100644 --- a/swmmanywhere/geospatial_utilities.py +++ b/swmmanywhere/geospatial_utilities.py @@ -23,6 +23,7 @@ import rasterio as rst import rioxarray import shapely +from pywbt import whitebox_tools from rasterio import features from scipy.interpolate import RegularGridInterpolator from scipy.spatial import KDTree @@ -31,7 +32,6 @@ from tqdm.auto import tqdm from swmmanywhere.logging import logger, verbose -from swmmanywhere.whitebox_utils import whitebox_tools os.environ["NUMBA_NUM_THREADS"] = "1" import pyflwdir # noqa: E402 @@ -637,19 +637,16 @@ def flwdir_whitebox(fid: Path) -> np.array: shutil.copy(fid, dem) # Condition + wbt_args = { + "BreachDepressions": ["-i=dem.tif", "--fillpits", "-o=dem_corr.tif"], + "D8Pointer": ["-i=dem_corr.tif", "-o=fdir.tif"], + } whitebox_tools( - "BreachDepressions", - ["-i=dem.tif", "-o=dem_corr.tif"], - work_dir=temp_path, - verbose=verbose(), - wbt_root=temp_path / "WBT", - ) - whitebox_tools( - "D8Pointer", - ["-i=dem_corr.tif", "-o=fdir.tif"], + wbt_args, work_dir=temp_path, verbose=verbose(), wbt_root=temp_path / "WBT", + max_procs=1, ) fdir = temp_path / "fdir.tif" diff --git a/swmmanywhere/whitebox_utils.py b/swmmanywhere/whitebox_utils.py deleted file mode 100644 index 3795eab5..00000000 --- a/swmmanywhere/whitebox_utils.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Whitebox utilities wrapper. - -@author: cheginit -""" -from __future__ import annotations - -import platform -import shutil -import stat -import subprocess -import tempfile -import zipfile -from pathlib import Path - -import requests - - -def _download_wbt(wbt_root: str | Path = "WBT") -> None: - """Download the WhiteboxTools executable for the current platform.""" - base_url = "https://www.whiteboxgeo.com/WBT_{}/WhiteboxTools_{}.zip" - - system = platform.system() - if system not in ("Windows", "Darwin", "Linux"): - raise ValueError(f"Unsupported operating system: {system}") - - if system == "Windows": - platform_suffix = "win_amd64" - elif system == "Darwin": - if platform.machine() == "arm64": - platform_suffix = "darwin_m_series" - else: - platform_suffix = "darwin_amd64" - elif system == "Linux": - if "musl" in platform.libc_ver()[0].lower(): - platform_suffix = "linux_musl" - else: - platform_suffix = "linux_amd64" - - url = base_url.format(system, platform_suffix) - wbt_root = Path(wbt_root) - - exe_name = ( - "whitebox_tools.exe" if platform.system() == "Windows" else "whitebox_tools" - ) - if (wbt_root / exe_name).exists(): - shutil.rmtree(wbt_root) - - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - zip_path = temp_path / "whitebox_tools.zip" - - try: - response = requests.get(url) - response.raise_for_status() - zip_path.write_bytes(response.content) - except requests.exceptions.RequestException as e: - raise RuntimeError(f"Failed to download WhiteboxTools: {e!s}") from e - - try: - with zipfile.ZipFile(zip_path, "r") as zip_ref: - for file_info in zip_ref.infolist(): - if "/WBT/" in file_info.filename: - extracted_path = Path( - *Path(file_info.filename).parts[ - Path(file_info.filename).parts.index("WBT") + 1 : - ] - ) - if extracted_path.parts: - source = zip_ref.extract(file_info, path=temp_dir) - dest = wbt_root / extracted_path - dest.parent.mkdir(parents=True, exist_ok=True) - shutil.move(source, dest) - except zipfile.BadZipFile: - raise RuntimeError("Downloaded file is not a valid zip file.") - - # Set executable permissions for non-Windows platforms - if system != "Windows": - executable_names = ["whitebox_tools", "whitebox_runner"] - for exec_name in executable_names: - exec_path = wbt_root / exec_name - if exec_path.exists(): - exec_path.chmod( - exec_path.stat().st_mode - | stat.S_IXUSR - | stat.S_IXGRP - | stat.S_IXOTH - ) - - -def whitebox_tools( - tool_name: str, - args: list[str] | tuple[str, ...], - wbt_root: str | Path = "WBT", - work_dir: str | Path = "", - verbose: bool = False, - compress_rasters: bool = False, - refresh_download: bool = False, -) -> None: - """Run a WhiteboxTools (not Whitebox Runner) tool with specified arguments. - - Parameters - ---------- - tool_name : str - Name of the WhiteboxTools to run. - args : list or tuple - List of arguments for the tool. - wbt_root : str or Path, optional - Path to the root directory containing the Whitebox executables - (default is "WBT"). - work_dir : str or Path, optional - Working directory for the tool (default is current directory). - verbose : bool, optional - Whether to print verbose output (default is False). - compress_rasters : bool, optional - Whether to compress output rasters (default is False). - refresh_download : bool, optional - Whether to refresh the download if WhiteboxTools is found (default is False). - - Raises: - ------ - subprocess.CalledProcessError - If the tool execution fails. - Exception - For any other unexpected errors. - - Notes: - ----- - This function will run the specified WhiteboxTools tool and handle its output. - If verbose is True, all output will be printed. - """ - wbt_root = Path(wbt_root) - work_dir = Path(work_dir) if work_dir else Path.cwd() - - exe_name = ( - "whitebox_tools.exe" if platform.system() == "Windows" else "whitebox_tools" - ) - exe_path = wbt_root / exe_name - - if not exe_path.exists() or refresh_download: - _download_wbt(wbt_root) - - command = [ - str(exe_path), - f"--run={tool_name}", - f"--wd={work_dir}", - "-v" if verbose else "-v=false", - f"--compress_rasters={'true' if compress_rasters else 'false'}", - ] - command.extend(args) - - if verbose: - print(" ".join(map(str, command))) - - try: - process = subprocess.run(command, check=True, text=True, capture_output=True) - if verbose: - print(process.stdout) - except subprocess.CalledProcessError as e: - print(f"Error output: {e.stdout}") - raise Exception(f"Error running tool: {e}") - except Exception as e: - raise Exception(f"Unexpected error: {e}")