From 350c5972c5821d47e46a8c7c8cf7167f2a0a5959 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Tue, 25 Jul 2023 17:49:49 +0800 Subject: [PATCH] grid selection in create_mesh2d + update docs --- docs/user_guide/model_overview.rst | 5 +- docs/user_guide/model_region.rst | 2 + hydromt/models/model_mesh.py | 13 +++- hydromt/workflows/basin_mask.py | 17 +++++ hydromt/workflows/mesh.py | 112 +++++++++++++++-------------- 5 files changed, 91 insertions(+), 58 deletions(-) diff --git a/docs/user_guide/model_overview.rst b/docs/user_guide/model_overview.rst index a7006deac..30cf1b189 100644 --- a/docs/user_guide/model_overview.rst +++ b/docs/user_guide/model_overview.rst @@ -37,7 +37,7 @@ While a generalized model class can readily be used, it can also be tailored to .. NOTE:: - As of version 0.6.0, the grid model (distributed grid model), lumped model (semi-distributed and lumped models), mesh model (unstructured grid models) have been implemented. Other model classes such as network models will follow in future versions. + As of version 0.6.0, the grid model (distributed grid model), lumped model (semi-distributed and lumped models), mesh model (unstructured grid(s) models) have been implemented. Other model classes such as network models will follow in future versions. The table below lists the base model components common to all model classes. All base model attributes and methods can be found the :ref:`API reference ` @@ -112,9 +112,10 @@ For each generalized model class, the respective computational unit components e | :py:func:`~LumpedModel.write_response_units` * - mesh - :ref:`MeshModel ` - - Static mesh (unstructured grid) data + - Static mesh (unstructured grid(s)) data - | :py:attr:`~MeshModel.mesh` | :py:func:`~MeshModel.set_mesh` + | :py:func:`~MeshModel.get_mesh` | :py:func:`~MeshModel.read_mesh` | :py:func:`~MeshModel.write_mesh` diff --git a/docs/user_guide/model_region.rst b/docs/user_guide/model_region.rst index 3950ab925..8d69b8497 100644 --- a/docs/user_guide/model_region.rst +++ b/docs/user_guide/model_region.rst @@ -22,6 +22,8 @@ Geospatial region Based on a raster file: ``{'grid': '/path/to/raster_file'}`` + Based on a mesh file: ``{'mesh': '/path/to/mesh_file'}`` + Hydrographic region ^^^^^^^^^^^^^^^^^^^ diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index e50f18a30..4663ff5eb 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -236,7 +236,13 @@ def setup_mesh2d_from_raster_reclass( @property def mesh(self) -> Union[xu.UgridDataArray, xu.UgridDataset]: - """Model static mesh data. Returns a xarray.Dataset.""" + """ + Model static mesh data. It returns a xugrid.UgridDataset. + + Mesh can contain several grids (1D, 2D, 3D) defined according + to UGRID conventions. To extract a specific grid, use get_mesh + method. + """ # XU grid data type Xarray dataset with xu sampling. if self._mesh is None and self._read: self.read_mesh() @@ -543,7 +549,6 @@ def setup_mesh2d( can be extracted using the `bounds` argument of region. Note Only existing meshed with only 2D grid can be read. - #FIXME: read existing 1D2D network file and extract 2D part. Adds/Updates model layers: @@ -553,6 +558,8 @@ def setup_mesh2d( ---------- region : dict Dictionary describing region of interest, bounds can be provided for type 'mesh'. + In case of 'mesh', if the file includes several grids, the specific 2D grid can + be selected using the 'grid_name' argument. CRS for 'bbox' and 'bounds' should be 4326; e.g.: * {'bbox': [xmin, ymin, xmax, ymax]} @@ -561,7 +568,7 @@ def setup_mesh2d( * {'mesh': 'path/to/2dmesh_file'} - * {'mesh': 'path/to/2dmesh_file', 'bounds': [xmin, ymin, xmax, ymax]} + * {'mesh': 'path/to/mesh_file', 'grid_name': 'mesh2d', 'bounds': [xmin, ymin, xmax, ymax]} res: float Resolution used to generate 2D mesh [unit of the CRS], required if region is not based on 'mesh'. diff --git a/hydromt/workflows/basin_mask.py b/hydromt/workflows/basin_mask.py index f087c9799..69aa8078b 100644 --- a/hydromt/workflows/basin_mask.py +++ b/hydromt/workflows/basin_mask.py @@ -14,6 +14,7 @@ import xarray as xr from shapely.geometry import box +from .. import _compat from ..data_adapter import GeoDataFrameAdapter from ..flw import basin_map, flwdir_from_da, outlet_map, stream_map @@ -48,6 +49,10 @@ def parse_region(region, logger=logger): * {'grid': /path/to/raster} + For a region based on a mesh grid of a mesh file: + + * {'mesh': /path/to/mesh} + Entire basin can be defined based on an ID, one or multiple point location (x, y), or a region of interest (bounding box or geometry) for which the basin IDs are looked up. The basins withint the area of interest can be further @@ -123,6 +128,7 @@ def parse_region(region, logger=logger): "geom": ["geom"], "bbox": ["bbox"], "grid": ["RasterDataArray"], + "mesh": ["UgridDataArray"], } kind = next(iter(kwargs)) # first key of region value0 = kwargs.pop(kind) @@ -135,6 +141,17 @@ def parse_region(region, logger=logger): kwargs = dict(grid=open_raster(value0, **kwargs)) elif isinstance(value0, (xr.Dataset, xr.DataArray)): kwargs = dict(grid=value0) + elif kind == "mesh": + if _compat.HAS_XUGRID: + import xugrid as xu + + if isinstance(value0, (str, Path)) and isfile(value0): + kwarg = dict(mesh=xu.open_dataset(value0)) + elif isinstance(value0, xu.UgridDataset): + kwarg = dict(mesh=value0) + kwargs.update(kwarg) + else: + raise ImportError("xugrid is required to read mesh files.") elif kind not in options: k_lst = '", "'.join(list(options.keys()) + list(MODELS)) raise ValueError(f'Region key "{kind}" not understood, select from "{k_lst}"') diff --git a/hydromt/workflows/mesh.py b/hydromt/workflows/mesh.py index 01e87d025..74bdcda3b 100644 --- a/hydromt/workflows/mesh.py +++ b/hydromt/workflows/mesh.py @@ -1,12 +1,10 @@ """Implementation for mesh based workflows.""" import logging -from os.path import isfile -from pathlib import Path from typing import Dict, Optional, Union import geopandas as gpd -import xarray as xr import xugrid as xu +from pyproj import CRS from shapely.geometry import box from xugrid.ugrid import conventions @@ -38,7 +36,6 @@ def create_mesh2d( can be extracted using the `bounds` argument of region. Note Only existing meshed with only 2D grid can be read. - #FIXME: read existing 1D2D network file and extract 2D part. Adds/Updates model layers: @@ -48,7 +45,9 @@ def create_mesh2d( ---------- region : dict Dictionary describing region of interest, bounds can be provided - for type 'mesh'. CRS for 'bbox' and 'bounds' should be 4326; e.g.: + for type 'mesh'. In case of 'mesh', if the file includes several + grids, the specific 2D grid can be selected using the 'grid_name' argument. + CRS for 'bbox' and 'bounds' should be 4326; e.g.: * {'bbox': [xmin, ymin, xmax, ymax]} @@ -56,7 +55,7 @@ def create_mesh2d( * {'mesh': 'path/to/2dmesh_file'} - * {'mesh': 'path/to/2dmesh_file', 'bounds': [xmin, ymin, xmax, ymax]} + * {'mesh': 'path/to/2dmesh_file', 'grid_name': 'mesh2d', 'bounds': [xmin, ymin, xmax, ymax]} res: float Resolution used to generate 2D mesh [unit of the CRS], required if region is not based on 'mesh'. @@ -70,11 +69,11 @@ def create_mesh2d( ------- mesh2d : xu.UgridDataset Generated mesh2d. - """ - if "mesh" not in region: + """ # noqa: E501 + kind, region = parse_region(region, logger=logger) + if kind != "mesh": if not isinstance(res, (int, float)): raise ValueError("res argument required for kind 'bbox', 'geom'") - kind, region = parse_region(region, logger=logger) if kind == "bbox": bbox = region["bbox"] geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=4326) @@ -113,7 +112,6 @@ def create_mesh2d( grid = gpd.GeoDataFrame(geometry=faces, crs=geom.crs) # If needed clip to geom if kind != "bbox": - # TODO: grid.intersects(geom) does not seem to work ? grid = grid.loc[ gpd.sjoin( grid, geom, how="left", predicate="intersects" @@ -126,66 +124,74 @@ def create_mesh2d( mesh2d.ugrid.grid.set_crs(grid.crs) else: - mesh2d_fn = region["mesh"] - if isinstance(mesh2d_fn, (str, Path)) and isfile(mesh2d_fn): - logger.info("An existing 2D grid is used to prepare 2D mesh.") - - ds = xr.open_dataset(mesh2d_fn, mask_and_scale=False) - elif isinstance(mesh2d_fn, xr.Dataset): - ds = mesh2d_fn - else: - raise ValueError(f"Region 'mesh' file {mesh2d_fn} not found, please check") - topologies = [ - k for k in ds.data_vars if ds[k].attrs.get("cf_role") == "mesh_topology" - ] - for topology in topologies: - topodim = ds[topology].attrs["topology_dimension"] - if topodim != 2: # chek if 2d mesh file else throw error - raise NotImplementedError( - f"{mesh2d_fn} cannot be opened. Please check if the existing" - " grid is an 2D mesh and not 1D2D mesh. " - " This option is not yet available for 1D2D meshes." + # Read existing mesh + uds = region["mesh"] + + # Select grid and/or check single grid + grids = dict() + for grid in uds.grids: + grids[grid.name] = grid + # Select specific topology if given + if "grid_name" in region: + if region["grid_name"] not in grids: + raise ValueError( + f"Mesh file does not contain grid {region['grid_name']}." ) - + grid = grids[region["grid_name"]] + elif len(grids) == 1: + grid = list(grids.values())[0] + else: + raise ValueError( + "Mesh file contains several grids. " + "Use grid_name argument of region to select a single one." + ) + # Chek if 2d mesh file else throw error + if grid.topology_dimension != 2: + raise ValueError("Grid in mesh file for create_mesh2d is not 2D.") # Continues with a 2D grid - mesh2d = xu.UgridDataset(ds) + mesh2d = xu.UgridDataset(grid.to_dataset()) + # Check crs and reproject to model crs + grid_crs = grid.crs if crs is not None: - crs = gis_utils.parse_crs(crs, bbox=mesh2d.ugrid.grid.bounds) + bbox = None + if grid_crs is not None: + if grid_crs.to_epsg() == 4326: + bbox = mesh2d.ugrid.grid.bounds + crs = gis_utils.parse_crs(crs, bbox=bbox) else: - crs = 4326 - if ds.raster.crs is not None: # parse crs - mesh2d.ugrid.grid.set_crs(ds.raster.crs) + crs = CRS.from_user_input(4326) + if grid_crs is not None: # parse crs + mesh2d.ugrid.grid.set_crs(grid_crs) else: # Assume model crs logger.warning( - f"Mesh data from {mesh2d_fn} doesn't have a CRS." + "Mesh data from mesh file doesn't have a CRS." f" Assuming crs option {crs}" ) mesh2d.ugrid.grid.set_crs(crs) mesh2d = mesh2d.drop_vars(GEO_MAP_COORD, errors="ignore") + # If bounds are provided in region, extract mesh for bounds + if "bounds" in region: + bounds = gpd.GeoDataFrame(geometry=[box(*region["bounds"])], crs=4326) + xmin, ymin, xmax, ymax = bounds.to_crs(mesh2d.ugrid.grid.crs).total_bounds + subset = mesh2d.ugrid.sel(y=slice(ymin, ymax), x=slice(xmin, xmax)) + # Check if still cells after clipping + err = ( + "MeshDataset: No data within model region." + "Check that bounds were given in the correct crs 4326" + ) + subset = subset.ugrid.assign_node_coords() + if subset.ugrid.grid.node_x.size == 0 or subset.ugrid.grid.node_y.size == 0: + raise IndexError(err) + mesh2d = subset + # Reproject to user crs option if needed - if mesh2d.ugrid.grid.crs != crs and crs is not None: + if mesh2d.ugrid.grid.crs != crs: logger.info(f"Reprojecting mesh to crs {crs}") mesh2d.ugrid.grid.to_crs(crs) - # If bounds are provided in region for mesh, extract mesh for bounds - if "bounds" in region and "mesh" in region: - bounds = gpd.GeoDataFrame(geometry=[box(*region["bounds"])], crs=4326) - xmin, ymin, xmax, ymax = bounds.to_crs(crs).total_bounds - subset = mesh2d.ugrid.sel(y=slice(ymin, ymax), x=slice(xmin, xmax)) - # Check if still cells after clipping - err = ( - "RasterDataset: No data within model region." - "Check that bounds were given in the correct crs 4326" - ) - subset = subset.ugrid.assign_node_coords() - if subset.ugrid.grid.node_x.size == 0 or subset.ugrid.grid.node_y.size == 0: - raise IndexError(err) - # reinitialise mesh2d grid (set_mesh is used in super) - mesh2d = subset - return mesh2d