From 123b5d674c913e0f19a7f959e42dff82ee2f2550 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Mon, 19 Jun 2023 23:28:58 +0800 Subject: [PATCH 01/15] support multigrid topologies mesh --- hydromt/models/model_mesh.py | 253 ++++++++++++++++++++++++++++++----- tests/test_model.py | 9 +- 2 files changed, 221 insertions(+), 41 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 606369455..833c1d78a 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -10,6 +10,7 @@ import pandas as pd import xarray as xr import xugrid as xu +from pyproj import CRS from shapely.geometry import box from .. import workflows @@ -32,18 +33,19 @@ def __init__(self, *args, **kwargs) -> None: self._mesh = None ## general setup methods - def setup_mesh_from_rasterdataset( + def setup_mesh2d_from_rasterdataset( self, raster_fn: Union[str, Path, xr.DataArray, xr.Dataset], + grid_name: Optional[str] = "mesh2d", variables: Optional[list] = None, fill_method: Optional[str] = None, resampling_method: Optional[str] = "mean", all_touched: Optional[bool] = True, rename: Optional[Dict] = dict(), ) -> List[str]: - """HYDROMT CORE METHOD: Add data variable(s) from ``raster_fn`` to mesh object. + """HYDROMT CORE METHOD: Add data variable(s) from ``raster_fn`` to 2D ``grid_name`` in mesh object. - Raster data is interpolated to the mesh grid using the ``resampling_method``. + Raster data is interpolated to the mesh ``grid_name`` using the ``resampling_method``. If raster is a dataset, all variables will be added unless ``variables`` list is specified. @@ -55,6 +57,8 @@ def setup_mesh_from_rasterdataset( ---------- raster_fn: str, Path, xr.DataArray, xr.Dataset Data catalog key, path to raster file or raster xarray data object. + grid_name: str, optional + Name of the mesh grid to add the data to. By default 'mesh2d'. variables: list, optional List of variables to add to mesh from raster_fn. By default all. fill_method : str, optional @@ -75,11 +79,14 @@ def setup_mesh_from_rasterdataset( ------- list List of variables added to mesh. - """ + """ # noqa: E501 self.logger.info(f"Preparing mesh data from raster source {raster_fn}") + # Check if grid name in self.mesh + if grid_name not in self.mesh_names: + raise ValueError(f"Grid name {grid_name} not in mesh ({self.mesh_names}).") # Read raster data and select variables ds = self.data_catalog.get_rasterdataset( - raster_fn, geom=self.region, buffer=2, variables=variables + raster_fn, bbox=self.bounds[grid_name], buffer=2, variables=variables ) if isinstance(ds, xr.DataArray): ds = ds.to_dataset() @@ -90,23 +97,26 @@ def setup_mesh_from_rasterdataset( # Convert mesh grid as geodataframe for sampling # Reprojection happens to gdf inside of zonal_stats method ds_sample = ds.raster.zonal_stats( - gdf=self.mesh_gdf, stats=resampling_method, all_touched=all_touched + gdf=self.mesh_gdf[grid_name], + stats=resampling_method, + all_touched=all_touched, ) # Rename variables rm_dict = {f"{var}_{resampling_method}": var for var in ds.data_vars} ds_sample = ds_sample.rename(rm_dict).rename(rename) # Convert to UgridDataset - uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh.ugrid.grid) + uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh_grids[grid_name]) - self.set_mesh(uds_sample) + self.set_mesh(uds_sample, grid_name=grid_name, overwrite_grid=False) return list(ds_sample.data_vars.keys()) - def setup_mesh_from_raster_reclass( + def setup_mesh2d_from_raster_reclass( self, raster_fn: Union[str, Path, xr.DataArray], reclass_table_fn: Union[str, Path, pd.DataFrame], reclass_variables: list, + grid_name: Optional[str] = "mesh2d", variable: Optional[str] = None, fill_nodata: Optional[str] = None, resampling_method: Optional[Union[str, list]] = "mean", @@ -114,7 +124,7 @@ def setup_mesh_from_raster_reclass( rename: Optional[Dict] = dict(), **kwargs, ) -> List[str]: - """HYDROMT CORE METHOD: Add data variable(s) to mesh object by reclassifying the data in ``raster_fn`` based on ``reclass_table_fn``. + """HYDROMT CORE METHOD: Add data variable(s) to 2D ``grid_name`` in mesh object by reclassifying the data in ``raster_fn`` based on ``reclass_table_fn``. The reclassified raster data are subsequently interpolated to the mesh using `resampling_method`. @@ -135,6 +145,8 @@ def setup_mesh_from_raster_reclass( reclass_variables : list List of reclass_variables from the reclass_table_fn table to add to the mesh. The index column should match values in raster_fn. + grid_name : str, optional + Name of the mesh grid to add the data to. By default 'mesh2d'. variable : str, optional Name of the raster dataset variable to use. This is only required when reading datasets with multiple variables. By default, None. @@ -173,9 +185,16 @@ def setup_mesh_from_raster_reclass( f"Preparing mesh data by reclassifying the data in {raster_fn} " f"based on {reclass_table_fn}." ) + # Check if grid name in self.mesh + if grid_name not in self.mesh_names: + raise ValueError(f"Grid name {grid_name} not in mesh ({self.mesh_names}).") # Read raster data and mapping table da = self.data_catalog.get_rasterdataset( - raster_fn, geom=self.region, buffer=2, variables=variable, **kwargs + raster_fn, + bbox=self.bounds[grid_name], + buffer=2, + variables=variable, + **kwargs, ) if not isinstance(da, xr.DataArray): raise ValueError( @@ -195,7 +214,7 @@ def setup_mesh_from_raster_reclass( # Convert mesh grid as geodataframe for sampling # Reprojection happens to gdf inside of zonal_stats method ds_sample = ds_vars.raster.zonal_stats( - gdf=self.mesh_gdf, + gdf=self.mesh_gdf[grid_name], stats=np.unique(np.atleast_1d(resampling_method)), all_touched=all_touched, ) @@ -209,9 +228,9 @@ def setup_mesh_from_raster_reclass( ds_sample = ds_sample.rename(rm_dict).rename(rename) ds_sample = ds_sample[reclass_variables] # Convert to UgridDataset - uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh.ugrid.grid) + uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh_grids[grid_name]) - self.set_mesh(uds_sample) + self.set_mesh(uds_sample, grid_name=grid_name, overwrite_grid=False) return list(ds_sample.data_vars.keys()) @@ -227,6 +246,8 @@ def set_mesh( self, data: Union[xu.UgridDataArray, xu.UgridDataset], name: Optional[str] = None, + grid_name: Optional[str] = None, + overwrite_grid: Optional[bool] = False, ) -> None: """Add data to mesh. @@ -235,11 +256,17 @@ def set_mesh( Parameters ---------- data: xugrid.UgridDataArray or xugrid.UgridDataset - new layer to add to mesh + new layer to add to mesh, TODO support one grid only or multiple grids? name: str, optional Name of new object layer, this is used to overwrite the name of a UgridDataArray. + grid_name: str, optional + Name of the mesh grid to add data to. If None, inferred from data. + Can be used for renaming the grid. + overwrite_grid: bool, optional + If True, overwrite the grid with the same name as the grid in self.mesh. """ + # Checks on data if not isinstance(data, (xu.UgridDataArray, xu.UgridDataset)): raise ValueError( "New mesh data in set_mesh should be of type xu.UgridDataArray" @@ -253,13 +280,107 @@ def set_mesh( f"Cannot set mesh from {str(type(data).__name__)} without a name." ) data = data.to_dataset() + + # Checks on grid topology + # TODO: check if we support setting multiple grids at once. For now just one + if len(data.ugrid.grids) > 1: + raise ValueError( + "set_mesh methods only supports adding data to one grid at a time." + ) + if grid_name is None: + grid_name = data.ugrid.grid.name + elif grid_name != data.ugrid.grid.name: + data = data.ugrid.rename(name=grid_name) + + # Adding to mesh if self._mesh is None: # NOTE: mesh is initialized with None self._mesh = data else: - for dvar in data.data_vars: - if dvar in self._mesh: - self.logger.warning(f"Replacing mesh parameter: {dvar}") - self._mesh[dvar] = data[dvar] + # Check on crs + if not data.ugrid.grid.crs == self.crs: + raise ValueError("Data and self.mesh should have the same CRS.") + # Check on new grid topology + if grid_name in self.mesh_names: + # check if the two grids are the same + if ( + not self.mesh_grids[grid_name] + .to_dataset() + .equals(data.ugrid.grid.to_dataset()) + ): + if not overwrite_grid: + raise ValueError( + f"Grid {grid_name} already exists in mesh" + " and has a different topology. " + "Use overwrite_grid=True to overwrite the grid" + " topology and related data." + ) + else: + # Remove grid and all corresponding data variables from mesh + self.logger.warning( + f"Overwriting grid {grid_name} and the corresponding" + " data variables in mesh." + ) + grids = [ + self.mesh_datasets[g].ugrid.to_dataset() + for g in self.mesh_names + if g != grid_name + ] + # Re-define _mesh + grids = xr.merge(grids) + self._mesh = xu.UgridDataset(grids) + # Check again mesh_names, could have changed if overwrite_grid=True + if grid_name in self.mesh_names: + for dvar in data.data_vars: + if dvar in self._mesh: + self.logger.warning(f"Replacing mesh parameter: {dvar}") + self._mesh[dvar] = data[dvar] + else: + # We are potentially adding a new grid without any data variables + self._mesh = xu.UgridDataset( + xr.merge([self.mesh.ugrid.to_dataset(), data.ugrid.to_dataset()]) + ) + + def get_mesh( + self, grid_name: str, include_data: bool = False + ) -> Union[xu.UgridDataArray, xu.UgridDataset]: + """ + Return a specific grid topology from mesh based on grid_name. + + If include_data is True, the data variables for that specific + grid are also included. + + Parameters + ---------- + grid_name : str + Name of the grid to return. + include_data : bool, optional + If True, also include data variables, by default False. + + Returns + ------- + uds: Union[xu.UgridDataArray, xu.UgridDataset] + Grid topology with or without data variables. + """ + if self.mesh is None: + raise ValueError("Mesh is not set, please use set_mesh first.") + if grid_name not in self.mesh_names: + raise ValueError(f"Grid {grid_name} not found in mesh.") + if include_data: + uds = xu.UgridDataset(grids=self.mesh_grids[grid_name]) + # Look for data_vars that are defined on grid_name + for var in self.mesh.data_vars: + if hasattr(self.mesh[var], "ugrid"): + if self.mesh[var].ugrid.grid.name == grid_name: + uds[var] = self.mesh[var] + # additionnal topology properties + elif var.startswith(grid_name): + uds[var] = self.mesh[var] + # else is global property (not grid specific) + + return uds + + else: + return self.mesh_grids[grid_name] def read_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: """Read model mesh data at / and add to mesh property. @@ -309,12 +430,55 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: ds_out = ds_out.rio.write_crs(self.mesh.ugrid.grid.crs) ds_out.to_netcdf(_fn, **kwargs) + # Other mesh properties + @property + def mesh_grids(self) -> Dict: + """Dictionnary of grid names and Ugrid topologies in mesh.""" + grids = dict() + if self.mesh is not None: + for grid in self.mesh.ugrid.grids: + grids[grid.name] = grid + + return grids + + @property + def mesh_datasets(self) -> Dict: + """Dictionnary of grid names and corresponding UgridDataset topology and data variables in mesh.""" # noqa: E501 + datasets = dict() + if self.mesh is not None: + for grid in self.mesh.ugrid.grids: + datasets[grid.name] = self.get_mesh( + grid_name=grid.name, include_data=True + ) + + return datasets + + @property + def mesh_names(self) -> List[str]: + """List of grid names in mesh.""" + if self.mesh is not None: + return list(self.mesh_grids.keys()) + else: + return [] + + @property + def mesh_gdf(self) -> Dict: + """Returns dict of geometry of grids in mesh as a gpd.GeoDataFrame.""" + mesh_gdf = dict() + if self._mesh is not None: + for k, v in self.mesh_datasets.items(): + # works better on a DataArray + # name = [n for n in self.mesh.data_vars][0] + mesh_gdf[k] = v.ugrid.to_geodataframe() + + return mesh_gdf + class MeshModel(MeshMixin, Model): """Model class Mesh Model for mesh models in HydroMT.""" - _CLI_ARGS = {"region": "setup_mesh", "res": "setup_mesh"} + _CLI_ARGS = {"region": "setup_mesh2d", "res": "setup_mesh2d"} _NAME = "mesh_model" def __init__( @@ -325,7 +489,7 @@ def __init__( data_libs: List[str] = None, logger=logger, ): - """Initialize a MeshModel for distributed models with an unstructured grid.""" + """Initialize a MeshModel for models with an unstructured grid.""" super().__init__( root=root, mode=mode, @@ -335,11 +499,12 @@ def __init__( ) ## general setup methods - def setup_mesh( + def setup_mesh2d( self, region: dict, res: Optional[float] = None, crs: int = None, + grid_name: str = "mesh2d", ) -> xu.UgridDataset: """HYDROMT CORE METHOD: Create an 2D unstructured mesh or reads an existing 2D mesh according UGRID conventions. @@ -352,7 +517,7 @@ def setup_mesh( Adds/Updates model layers: - * **mesh** mesh topology: add mesh topology to mesh object + * **grid_name** mesh topology: add grid_name 2D topology to mesh object Parameters ---------- @@ -370,6 +535,8 @@ def setup_mesh( crs : EPSG code, int, optional Optional EPSG code of the model. If None using the one from region, and else 4326. + grid_name : str, optional + Name of the 2D grid in mesh, by default "mesh2d". Returns ------- @@ -470,7 +637,7 @@ def setup_mesh( self.logger.info(f"Reprojecting mesh to crs {crs}") mesh2d.ugrid.grid.to_crs(self.crs) - self.set_mesh(mesh2d) + self.set_mesh(mesh2d, grid_name=grid_name) # This setup method returns region so that it can be wrapped for models # which require more information @@ -521,24 +688,36 @@ def write( def bounds(self) -> Tuple: """Returns model mesh bounds.""" if self._mesh is not None: - return self._mesh.ugrid.grid.bounds + return self._mesh.ugrid.bounds + + @property + def crs(self) -> CRS: + """Returns model mesh crs.""" + if self._mesh is not None: + grid_crs = self._mesh.ugrid.crs + # Check if all the same + crs = None + for k, v in grid_crs.items(): + if crs is None: + crs = v + if v == crs: + continue + else: + raise ValueError( + f"Mesh crs is not uniform, please check {grid_crs}" + ) + return crs + else: + return None @property def region(self) -> gpd.GeoDataFrame: - """Returns geometry of region of the model area of interest.""" + """Returns geometry of region of the model area of interest based on mesh total bounds.""" # noqa: E501 region = gpd.GeoDataFrame() if "region" in self.geoms: region = self.geoms["region"] elif self.mesh is not None: - crs = self.mesh.ugrid.grid.crs - if crs is None and hasattr(crs, "to_epsg"): - crs = crs.to_epsg() # not all CRS have an EPSG code - region = gpd.GeoDataFrame(geometry=[box(*self.bounds)], crs=crs) + region = gpd.GeoDataFrame( + geometry=[box(*self.mesh.ugrid.total_bounds)], crs=self.crs + ) return region - - @property - def mesh_gdf(self) -> gpd.GeoDataFrame: - """Returns geometry of mesh as a gpd.GeoDataFrame.""" - if self._mesh is not None: - name = [n for n in self.mesh.data_vars][0] # works better on a DataArray - return self._mesh[name].ugrid.to_geodataframe() diff --git a/tests/test_model.py b/tests/test_model.py index ac41f2806..9064716e8 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -480,18 +480,19 @@ def test_meshmodel_setup(griduda, world): mod = MeshModel(data_libs=["artifact_data", dc_param_fn]) mod.setup_config(**{"header": {"setting": "value"}}) region = {"geom": world[world.name == "Italy"]} - mod.setup_mesh(region, res=10000, crs=3857) + mod.setup_mesh2d(region, res=10000, crs=3857, grid_name="mesh2d") mod.region region = {"mesh": griduda.ugrid.to_dataset()} mod1 = MeshModel(data_libs=["artifact_data", dc_param_fn]) - mod1.setup_mesh(region) - mod1.setup_mesh_from_rasterdataset("vito") + mod1.setup_mesh2d(region, grid_name="mesh2d") + mod1.setup_mesh2d_from_rasterdataset("vito", grid_name="mesh2d") assert "vito" in mod1.mesh.data_vars - mod1.setup_mesh_from_raster_reclass( + mod1.setup_mesh2d_from_raster_reclass( raster_fn="vito", reclass_table_fn="vito_mapping", reclass_variables=["roughness_manning"], resampling_method="mean", + grid_name="mesh2d", ) assert "roughness_manning" in mod1.mesh.data_vars From fe1ad094a19174795493284b561d71f7fd88c3cd Mon Sep 17 00:00:00 2001 From: hboisgon Date: Mon, 19 Jun 2023 23:30:05 +0800 Subject: [PATCH 02/15] fix relative path to naturalearth --- tests/conftest.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 465456240..b7f2cf9af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,9 @@ import os os.environ["USE_PYGEOS"] = "0" + +from os.path import abspath, dirname, join + import geopandas as gpd import numpy as np import pandas as pd @@ -21,6 +24,8 @@ ) from hydromt.data_catalog import DataCatalog +DATADIR = join(dirname(abspath(__file__)), "data") + @pytest.fixture() def rioda(): @@ -83,7 +88,7 @@ def geodf(df): @pytest.fixture() def world(): - world = gpd.read_file("tests/data/naturalearth_lowres.geojson") + world = gpd.read_file(join(DATADIR, "naturalearth_lowres.geojson")) return world From a7ccf7fe5cada857c4ab30e011e74ab3802fbc5c Mon Sep 17 00:00:00 2001 From: hboisgon Date: Tue, 11 Jul 2023 19:46:58 +0800 Subject: [PATCH 03/15] resend after pre-commit --- hydromt/models/model_mesh.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 833c1d78a..586677203 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -3,7 +3,7 @@ import os from os.path import dirname, isdir, isfile, join from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Union import geopandas as gpd import numpy as np @@ -522,7 +522,7 @@ def setup_mesh2d( Parameters ---------- region : dict - Dictionary describing region of interest, e.g.: + Dictionary describing region of interest, e.g.: TODO support bounds in region for type mesh * {'bbox': [xmin, ymin, xmax, ymax]} @@ -632,6 +632,18 @@ def setup_mesh2d( mesh2d.ugrid.grid.set_crs(crs) mesh2d = mesh2d.drop_vars(GEO_MAP_COORD, errors="ignore") + # TODO if bounds clip + # Check if intersects with region + # xmin, ymin, xmax, ymax = self.bounds + # subset = mesh2d.ugrid.sel(y=slice(ymin, ymax), x=slice(xmin, xmax)) + # err = "RasterDataset: No data within model region." + # 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) + # self._mesh = subset + # Reproject to user crs option if needed if mesh2d.ugrid.grid.crs != crs and crs is not None: self.logger.info(f"Reprojecting mesh to crs {crs}") @@ -685,7 +697,7 @@ def write( # MeshModel properties @property - def bounds(self) -> Tuple: + def bounds(self) -> Dict: """Returns model mesh bounds.""" if self._mesh is not None: return self._mesh.ugrid.bounds From aaac6420cdfcc851aaee914e949480c7a0cd1245 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Thu, 20 Jul 2023 17:06:45 +0800 Subject: [PATCH 04/15] some changes and adapt after test with dflowfm --- hydromt/models/model_mesh.py | 147 ++++++---------------- hydromt/workflows/__init__.py | 5 +- hydromt/workflows/grid.py | 4 +- hydromt/workflows/mesh.py | 230 ++++++++++++++++++++++++++++++++++ 4 files changed, 271 insertions(+), 115 deletions(-) create mode 100644 hydromt/workflows/mesh.py diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 586677203..e94ddfc46 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -1,7 +1,7 @@ """Implementations for model mesh workloads.""" import logging import os -from os.path import dirname, isdir, isfile, join +from os.path import dirname, isdir, join from pathlib import Path from typing import Dict, List, Optional, Union @@ -252,6 +252,7 @@ def set_mesh( """Add data to mesh. All layers of mesh have identical spatial coordinates in Ugrid conventions. + Also updates self.region if grid_name is new or overwrite_grid is True. Parameters ---------- @@ -266,6 +267,12 @@ def set_mesh( overwrite_grid: bool, optional If True, overwrite the grid with the same name as the grid in self.mesh. """ + # Check if new grid_name + if grid_name not in self.mesh_names: + new_grid = True + else: + new_grid = False + # Checks on data if not isinstance(data, (xu.UgridDataArray, xu.UgridDataset)): raise ValueError( @@ -290,7 +297,7 @@ def set_mesh( if grid_name is None: grid_name = data.ugrid.grid.name elif grid_name != data.ugrid.grid.name: - data = data.ugrid.rename(name=grid_name) + data = workflows.rename_mesh(data, name=grid_name) # Adding to mesh if self._mesh is None: # NOTE: mesh is initialized with None @@ -340,6 +347,12 @@ def set_mesh( xr.merge([self.mesh.ugrid.to_dataset(), data.ugrid.to_dataset()]) ) + # update related geoms if necessary: region + if overwrite_grid or new_grid: + # add / updates region + self._geoms.pop("region", None) + self.region + def get_mesh( self, grid_name: str, include_data: bool = False ) -> Union[xu.UgridDataArray, xu.UgridDataset]: @@ -425,9 +438,9 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: os.makedirs(dirname(_fn)) self.logger.debug(f"Writing file {fn}") ds_out = self.mesh.ugrid.to_dataset() - if self.mesh.ugrid.grid.crs is not None: + if self.crs is not None: # save crs to spatial_ref coordinate - ds_out = ds_out.rio.write_crs(self.mesh.ugrid.grid.crs) + ds_out = ds_out.rio.write_crs(self.crs) ds_out.to_netcdf(_fn, **kwargs) # Other mesh properties @@ -510,7 +523,8 @@ def setup_mesh2d( Grids are read according to UGRID conventions. An 2D unstructured mesh will be created as 2D rectangular grid from a geometry (geom_fn) or bbox. - If an existing 2D mesh is given, then no new mesh will be generated + If an existing 2D mesh is given, then no new mesh will be generated but an extent + 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. @@ -522,19 +536,22 @@ def setup_mesh2d( Parameters ---------- region : dict - Dictionary describing region of interest, e.g.: TODO support bounds in region for type mesh + Dictionary describing region of interest, bounds can be provided for type 'mesh'. + CRS for 'bbox' and 'bounds' should be 4326; e.g.: * {'bbox': [xmin, ymin, xmax, ymax]} * {'geom': 'path/to/polygon_geometry'} * {'mesh': 'path/to/2dmesh_file'} + + * {'mesh': 'path/to/2dmesh_file', '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'. crs : EPSG code, int, optional - Optional EPSG code of the model. If None using the one from region, - and else 4326. + Optional EPSG code of the model or "utm" to let hydromt find the closest projected CRS. + If None using the one from region, and else 4326. grid_name : str, optional Name of the 2D grid in mesh, by default "mesh2d". @@ -546,112 +563,17 @@ def setup_mesh2d( """ # noqa: E501 self.logger.info("Preparing 2D mesh.") - if "mesh" not in region: - if not isinstance(res, (int, float)): - raise ValueError("res argument required") - kind, region = workflows.parse_region(region, logger=self.logger) - if kind == "bbox": - bbox = region["bbox"] - geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=4326) - elif kind == "geom": - geom = region["geom"] - if geom.crs is None: - raise ValueError('Model region "geom" has no CRS') - else: - raise ValueError( - f"Region for mesh must of kind [bbox, geom, mesh], kind {kind} " - "not understood." - ) - if crs is not None: - geom = geom.to_crs(crs) - # Generate grid based on res for region bbox - xmin, ymin, xmax, ymax = geom.total_bounds - # note we flood the number of faces within bounds - ncol = int((xmax - xmin) // res) - nrow = int((ymax - ymin) // res) - dx, dy = res, -res - faces = [] - for i in range(nrow): - top = ymax + i * dy - bottom = ymax + (i + 1) * dy - for j in range(ncol): - left = xmin + j * dx - right = xmin + (j + 1) * dx - faces.append(box(left, bottom, right, top)) - 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" - ).index_right.notna() - ].reset_index() - # Create mesh from grid - grid.index.name = "mesh2d_nFaces" - mesh2d = xu.UgridDataset.from_geodataframe(grid) - mesh2d.ugrid.grid.set_crs(grid.crs) - - else: - mesh2d_fn = region["mesh"] - if isinstance(mesh2d_fn, (str, Path)) and isfile(mesh2d_fn): - self.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." - ) - - # Continues with a 2D grid - mesh2d = xu.UgridDataset(ds) - # Check crs and reproject to model crs - if crs is None: - crs = 4326 - if ds.rio.crs is not None: # parse crs - mesh2d.ugrid.grid.set_crs(ds.raster.crs) - else: - # Assume model crs - self.logger.warning( - f"Mesh data from {mesh2d_fn} 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") - - # TODO if bounds clip - # Check if intersects with region - # xmin, ymin, xmax, ymax = self.bounds - # subset = mesh2d.ugrid.sel(y=slice(ymin, ymax), x=slice(xmin, xmax)) - # err = "RasterDataset: No data within model region." - # 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) - # self._mesh = subset - - # Reproject to user crs option if needed - if mesh2d.ugrid.grid.crs != crs and crs is not None: - self.logger.info(f"Reprojecting mesh to crs {crs}") - mesh2d.ugrid.grid.to_crs(self.crs) - + # Create mesh2d + mesh2d = workflows.create_mesh2d( + region=region, + res=res, + crs=crs, + logger=self.logger, + ) + # Add mesh2d to self.mesh self.set_mesh(mesh2d, grid_name=grid_name) - # This setup method returns region so that it can be wrapped for models + # This setup method returns mesh2d so that it can be wrapped for models # which require more information return mesh2d @@ -732,4 +654,5 @@ def region(self) -> gpd.GeoDataFrame: region = gpd.GeoDataFrame( geometry=[box(*self.mesh.ugrid.total_bounds)], crs=self.crs ) + self.set_geoms(region, name="region") return region diff --git a/hydromt/workflows/__init__.py b/hydromt/workflows/__init__.py index 8365ebb6a..1f97c719a 100644 --- a/hydromt/workflows/__init__.py +++ b/hydromt/workflows/__init__.py @@ -1,7 +1,10 @@ # -*- coding: utf-8 -*- """HydroMT workflows.""" - +from .. import _compat from .basin_mask import * from .forcing import * from .grid import * from .rivers import * + +if _compat.HAS_XUGRID: + from .mesh import * diff --git a/hydromt/workflows/grid.py b/hydromt/workflows/grid.py index 61c03be5f..5f1b90c62 100644 --- a/hydromt/workflows/grid.py +++ b/hydromt/workflows/grid.py @@ -8,7 +8,7 @@ import xarray as xr from shapely.geometry import Polygon -import hydromt +from ..raster import full logger = logging.getLogger(__name__) @@ -52,7 +52,7 @@ def grid_from_constant( da: xr.DataArray Grid with constant value. """ - da = hydromt.raster.full( + da = full( coords=grid_like.raster.coords, nodata=constant, dtype=dtype, diff --git a/hydromt/workflows/mesh.py b/hydromt/workflows/mesh.py new file mode 100644 index 000000000..17d46211f --- /dev/null +++ b/hydromt/workflows/mesh.py @@ -0,0 +1,230 @@ +"""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 shapely.geometry import box +from xugrid.ugrid import conventions + +from .. import gis_utils +from ..raster import GEO_MAP_COORD +from .basin_mask import parse_region + +logger = logging.getLogger(__name__) + +__all__ = [ + "create_mesh2d", + "rename_mesh", +] + + +def create_mesh2d( + region: Dict, + res: Optional[float] = None, + crs: Optional[int] = None, + align: bool = True, + logger=logger, +) -> xu.UgridDataset: + """ + Create an 2D unstructured mesh or reads an existing 2D mesh. + + Grids are read according to UGRID conventions. An 2D unstructured mesh + will be created as 2D rectangular grid from a geometry (geom_fn) or bbox. + If an existing 2D mesh is given, then no new mesh will be generated but an extent + 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: + + * **grid_name** mesh topology: add grid_name 2D topology to mesh object + + Parameters + ---------- + region : dict + Dictionary describing region of interest, bounds can be provided + for type 'mesh'. CRS for 'bbox' and 'bounds' should be 4326; e.g.: + + * {'bbox': [xmin, ymin, xmax, ymax]} + + * {'geom': 'path/to/polygon_geometry'} + + * {'mesh': 'path/to/2dmesh_file'} + + * {'mesh': 'path/to/2dmesh_file', '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'. + crs : EPSG code, int, optional + Optional EPSG code of the model or "utm" to let hydromt find the closest + projected CRS. If None using the one from region, and else 4326. + align : bool, optional + Align the mesh to the resolution, by default True. + + Returns + ------- + mesh2d : xu.UgridDataset + Generated mesh2d. + """ + if "mesh" not in region: + 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) + elif kind == "geom": + geom = region["geom"] + if geom.crs is None: + raise ValueError('Model region "geom" has no CRS') + else: + raise ValueError( + "Region for mesh must be of kind [bbox, geom, mesh], " + f"kind {kind} not understood." + ) + # Parse crs and reproject geom if needed + if crs is not None: + crs = gis_utils.parse_crs(crs, bbox=geom.total_bounds) + geom = geom.to_crs(crs) + # Generate grid based on res for region bbox + xmin, ymin, xmax, ymax = geom.total_bounds + if align: + xmin = round(xmin / res) * res + ymin = round(ymin / res) * res + xmax = round(xmax / res) * res + ymax = round(ymax / res) * res + # note we flood the number of faces within bounds + ncol = round((xmax - xmin) / res) # int((xmax - xmin) // res) + nrow = round((ymax - ymin) / res) # int((ymax - ymin) // res) + dx, dy = res, -res + faces = [] + for i in range(nrow): + top = ymax + i * dy + bottom = ymax + (i + 1) * dy + for j in range(ncol): + left = xmin + j * dx + right = xmin + (j + 1) * dx + faces.append(box(left, bottom, right, top)) + 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" + ).index_right.notna() + ].reset_index() + # Create mesh from grid + grid.index.name = "mesh2d_nFaces" + mesh2d = xu.UgridDataset.from_geodataframe(grid) + 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." + ) + + # Continues with a 2D grid + mesh2d = xu.UgridDataset(ds) + # Check crs and reproject to model crs + if crs is not None: + crs = gis_utils.parse_crs(crs, bbox=mesh2d.ugrid.grid.bounds) + else: + crs = 4326 + if ds.raster.crs is not None: # parse crs + mesh2d.ugrid.grid.set_crs(ds.raster.crs) + else: + # Assume model crs + logger.warning( + f"Mesh data from {mesh2d_fn} 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") + + # Reproject to user crs option if needed + if mesh2d.ugrid.grid.crs != crs and crs is not None: + 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 + + +def rename_mesh(mesh: Union[xu.UgridDataArray, xu.UgridDataset], name: str): + """ + Rename all grid variables in mesh according to UGRID conventions. + + Note: adapted from xugrid.ugrid.grid.rename to also work on + UgridDataset and UgridDataArray + """ + # Get the old and the new names. Their keys are the same. + old_attrs = mesh.ugrid.grid._attrs + new_attrs = conventions.default_topology_attrs( + name, mesh.ugrid.grid.topology_dimension + ) + + # The attrs will have some roles joined together, e.g. node_coordinates + # will contain x and y as "mesh2d_node_x mesh2d_node_y". + name_dict = {mesh.ugrid.grid.name: name} + skip = ("cf_role", "long_name", "topology_dimension") + for key, value in old_attrs.items(): + if key in new_attrs and key not in skip: + split_new = new_attrs[key].split() + split_old = value.split() + if len(split_new) != len(split_old): + raise ValueError( + f"Number of entries does not match on {key}: " + f"{split_new} versus {split_old}" + ) + for name_key, name_value in zip(split_old, split_new): + name_dict[name_key] = name_value + + new = mesh.copy() + new.ugrid.grid.name = name + new.ugrid.grid._attrs = new_attrs + new.ugrid.grid._indexes = { + k: name_dict[v] for k, v in new.ugrid.grid._indexes.items() + } + + to_rename = tuple(new.data_vars) + tuple(new.coords) + tuple(new.dims) + new = new.rename({k: v for k, v in name_dict.items() if k in to_rename}) + + return new From fa4561d15dfdc75c8f1722feb8a9e9a467c708de Mon Sep 17 00:00:00 2001 From: hboisgon Date: Thu, 20 Jul 2023 17:07:24 +0800 Subject: [PATCH 05/15] add more tests for setup_mesh2d --- tests/test_model.py | 71 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/tests/test_model.py b/tests/test_model.py index 9064716e8..179e9c3ab 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -473,6 +473,77 @@ def test_meshmodel(mesh_model, tmpdir): assert equal, errors +@pytest.mark.skipif(not hasattr(hydromt, "MeshModel"), reason="Xugrid not installed.") +def test_setup_mesh(tmpdir, griduda): + MeshModel = MODELS.load("mesh_model") + # Initialize model + model = MeshModel( + root=join(tmpdir, "mesh_model"), + data_libs=["artifact_data"], + mode="w", + ) + # wrong region kind + with pytest.raises(ValueError, match="Region for mesh must be of kind "): + model.setup_mesh2d( + region={"basin": [12.5, 45.5]}, + res=0.05, + ) + # bbox + bbox = [12.05, 45.30, 12.85, 45.65] + with pytest.raises( + ValueError, match="res argument required for kind 'bbox', 'geom'" + ): + model.setup_mesh2d({"bbox": bbox}) + model.setup_mesh2d( + region={"bbox": bbox}, + res=0.05, + crs=4326, + grid_name="mesh2d", + ) + assert "mesh2d" in model.mesh_names + assert model.crs.to_epsg() == 4326 + assert np.all(np.round(model.region.total_bounds, 3) == bbox) + assert model.mesh.ugrid.grid.n_node == 136 + model._mesh = None # remove old mesh + + # geom + region = model._geoms.pop("region") + model.setup_mesh2d( + region={"geom": region}, + res=10000, + crs="utm", + grid_name="mesh2d", + ) + assert model.crs.to_epsg() == 32633 + assert model.mesh.ugrid.grid.n_node == 35 + model._mesh = None # remove old mesh + + # mesh + # create mesh file + mesh_fn = str(tmpdir.join("mesh2d.nc")) + gridda = griduda.ugrid.to_dataset() + gridda = gridda.rio.write_crs(griduda.ugrid.grid.crs) + gridda.to_netcdf(mesh_fn) + + model.setup_mesh2d( + region={"mesh": mesh_fn}, + grid_name="mesh2d", + ) + assert np.all(griduda.ugrid.total_bounds == model.region.total_bounds) + assert model.mesh.ugrid.grid.n_node == 169 + model._mesh = None # remove old mesh + + # mesh with bounds + bounds = [12.095, 46.495, 12.10, 46.50] + model.setup_mesh2d( + {"mesh": mesh_fn, "bounds": bounds}, + grid_name="mesh1", + ) + assert "mesh1" in model.mesh_names + assert model.mesh.ugrid.grid.n_node == 49 + assert np.all(np.round(model.region.total_bounds, 3) == bounds) + + @pytest.mark.skipif(not hasattr(hydromt, "MeshModel"), reason="Xugrid not installed.") def test_meshmodel_setup(griduda, world): MeshModel = MODELS.load("mesh_model") From 85d31de204e0583c691587d2874664cac1dde34b Mon Sep 17 00:00:00 2001 From: hboisgon Date: Thu, 20 Jul 2023 17:24:43 +0800 Subject: [PATCH 06/15] update docs --- docs/api.rst | 19 +++++++++++++++++-- docs/changelog.rst | 2 +- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 74d1a6fb6..9824d0bf8 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -322,8 +322,13 @@ Components and attributes :toctree: _generated MeshModel.mesh + MeshModel.mesh_names + MeshModel.mesh_grids + MeshModel.mesh_datasets + MeshModel.mesh_gdf MeshModel.crs MeshModel.region + MeshModel.bounds General methods --------------- @@ -332,6 +337,7 @@ General methods :toctree: _generated MeshModel.set_mesh + MeshModel.get_mesh MeshModel.read_mesh MeshModel.write_mesh @@ -344,10 +350,11 @@ Setup methods MeshModel.setup_config MeshModel.setup_region + MeshModel.setup_mesh2d + MeshModel.setup_mesh2d_from_rasterdataset + MeshModel.setup_mesh2d_from_raster_reclass MeshModel.setup_maps_from_rasterdataset MeshModel.setup_maps_from_raster_reclass - MeshModel.setup_mesh_from_rasterdataset - MeshModel.setup_mesh_from_raster_reclass ========= @@ -365,6 +372,14 @@ Grid workflows.grid.grid_from_raster_reclass workflows.grid.grid_from_geodataframe +Mesh +==== + +.. autosummary:: + :toctree: _generated + + workflows.mesh.create_mesh2d + Basin mask ========== diff --git a/docs/changelog.rst b/docs/changelog.rst index d8a07cbfb..edde2e9f3 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -15,7 +15,7 @@ Added Changed ------- -- +- Updated ``MeshModel`` and related methods to support multigrids instead of one single 2D grid. PR #412 Fixed ----- From f5b1f2a6ee5411c7a1cb2438a7077f16d2897657 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Mon, 24 Jul 2023 18:06:12 +0800 Subject: [PATCH 07/15] additionnal bugfixes --- hydromt/models/model_mesh.py | 24 ++++++++++++++++++++---- hydromt/workflows/mesh.py | 1 + 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index e94ddfc46..e50f18a30 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -306,6 +306,8 @@ def set_mesh( # Check on crs if not data.ugrid.grid.crs == self.crs: raise ValueError("Data and self.mesh should have the same CRS.") + # Save crs as it will be lost when converting to xarray + crs = self.crs # Check on new grid topology if grid_name in self.mesh_names: # check if the two grids are the same @@ -346,6 +348,9 @@ def set_mesh( self._mesh = xu.UgridDataset( xr.merge([self.mesh.ugrid.to_dataset(), data.ugrid.to_dataset()]) ) + # Restore crs + for grid in self._mesh.ugrid.grids: + grid.set_crs(crs) # update related geoms if necessary: region if overwrite_grid or new_grid: @@ -379,7 +384,9 @@ def get_mesh( if grid_name not in self.mesh_names: raise ValueError(f"Grid {grid_name} not found in mesh.") if include_data: - uds = xu.UgridDataset(grids=self.mesh_grids[grid_name]) + grid = self.mesh_grids[grid_name] + uds = xu.UgridDataset(grid.to_dataset()) + uds.ugrid.grid.set_crs(self.crs) # Look for data_vars that are defined on grid_name for var in self.mesh.data_vars: if hasattr(self.mesh[var], "ugrid"): @@ -458,7 +465,7 @@ def mesh_grids(self) -> Dict: def mesh_datasets(self) -> Dict: """Dictionnary of grid names and corresponding UgridDataset topology and data variables in mesh.""" # noqa: E501 datasets = dict() - if self.mesh is not None: + if self._mesh is not None: for grid in self.mesh.ugrid.grids: datasets[grid.name] = self.get_mesh( grid_name=grid.name, include_data=True @@ -480,8 +487,17 @@ def mesh_gdf(self) -> Dict: mesh_gdf = dict() if self._mesh is not None: for k, v in self.mesh_datasets.items(): - # works better on a DataArray - # name = [n for n in self.mesh.data_vars][0] + # works better on a Dataarray / Dataset + # (need to add dummy variable if empty) + if len(v.data_vars) == 0: + grid = v.ugrid.grid + if grid.topology_dimension == 1: + dim = grid.edge_dimension + elif grid.topology_dimension == 2: + dim = grid.face_dimension + v[f"{grid.name}_node_id"] = xr.DataArray( + v[dim].values.astype(str), dims=dim + ) mesh_gdf[k] = v.ugrid.to_geodataframe() return mesh_gdf diff --git a/hydromt/workflows/mesh.py b/hydromt/workflows/mesh.py index 17d46211f..01e87d025 100644 --- a/hydromt/workflows/mesh.py +++ b/hydromt/workflows/mesh.py @@ -122,6 +122,7 @@ def create_mesh2d( # Create mesh from grid grid.index.name = "mesh2d_nFaces" mesh2d = xu.UgridDataset.from_geodataframe(grid) + mesh2d = mesh2d.ugrid.assign_face_coords() mesh2d.ugrid.grid.set_crs(grid.crs) else: From 350c5972c5821d47e46a8c7c8cf7167f2a0a5959 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Tue, 25 Jul 2023 17:49:49 +0800 Subject: [PATCH 08/15] 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 From 92db0dc78b2878766ac23ef9ad0308459a4414e0 Mon Sep 17 00:00:00 2001 From: Xiaohan Li Date: Fri, 4 Aug 2023 17:44:04 +0200 Subject: [PATCH 09/15] review 1d --- hydromt/models/model_mesh.py | 57 ++++++++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 13 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 4663ff5eb..3d635b88b 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -307,11 +307,18 @@ def set_mesh( # Adding to mesh if self._mesh is None: # NOTE: mesh is initialized with None + # Check on crs + if not data.ugrid.grid.crs: + raise ValueError( + "Data should have CRS." + ) # FIXME new data coming in should have crs self._mesh = data else: # Check on crs if not data.ugrid.grid.crs == self.crs: - raise ValueError("Data and self.mesh should have the same CRS.") + raise ValueError( + "Data and self.mesh should have the same CRS." + ) # FIXME since data and self.mesh have the same crs, can it be a global property? # Save crs as it will be lost when converting to xarray crs = self.crs # Check on new grid topology @@ -391,8 +398,10 @@ def get_mesh( raise ValueError(f"Grid {grid_name} not found in mesh.") if include_data: grid = self.mesh_grids[grid_name] - uds = xu.UgridDataset(grid.to_dataset()) - uds.ugrid.grid.set_crs(self.crs) + uds = xu.UgridDataset( + grid.to_dataset(optional_attributes=True) + ) # FIXME: it seems optional_attributes is needed when creating a new object for all variables about grid to be written, not needed when reading + uds.ugrid.grid.set_crs(grid.crs) # Look for data_vars that are defined on grid_name for var in self.mesh.data_vars: if hasattr(self.mesh[var], "ugrid"): @@ -408,25 +417,41 @@ def get_mesh( else: return self.mesh_grids[grid_name] - def read_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: + def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None: """Read model mesh data at / and add to mesh property. - key-word arguments are passed to :py:func:`xr.open_dataset` + key-word arguments are passed to :py:func:`xu.open_dataset` Parameters ---------- fn : str, optional filename relative to model root, by default 'mesh/mesh.nc' + crs : CRS, optional + Coordinate Reference System (CRS) object representing the spatial reference system of the mesh file. **kwargs : dict - Additional keyword arguments to be passed to the `_read_nc` method. + Additional keyword arguments to be passed to the `open_dataset` method. """ + # FIXME: general question, why there is no force overwrite argument for model update? self._assert_read_mode - for ds in self._read_nc(fn, **kwargs).values(): - uds = xu.UgridDataset(ds) - if ds.rio.crs is not None: # parse crs - uds.ugrid.grid.set_crs(ds.raster.crs) - uds = uds.drop_vars(GEO_MAP_COORD, errors="ignore") - self.set_mesh(uds) + uds = xu.open_dataset( + fn, **kwargs + ) # FIXME: switching self._read_nc to xu.open_dataset. is the idea to support single file? or also multiple files? e.g. xu.open_mfdataset + if not any( + uds.ugrid.crs.values() + ): # FIXME: crs cannot be write/read correctly by xugrid https://github.com/Deltares/xugrid/issues/138 + if not crs: + raise ValueError( + "no crs is found in the file nor passed to the reader." + ) + else: + uds.ugrid.set_crs(crs) + self.logger.info( + "no crs is found in the file, assigning from user input." + ) + else: + self.logger.info("crs is read from file") + self._mesh = uds + # self.set_mesh(uds) # FIXME: infinit loop because self.mesh property is called in self.set_mesh, in which read_mesh is called. def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: """Write model grid data to a netCDF file at /. @@ -504,7 +529,10 @@ def mesh_gdf(self) -> Dict: v[f"{grid.name}_node_id"] = xr.DataArray( v[dim].values.astype(str), dims=dim ) - mesh_gdf[k] = v.ugrid.to_geodataframe() + gdf = ( + v.ugrid.to_geodataframe() + ) # FIXME xugrid crs not passed on correctly: https://github.com/Deltares/xugrid/issues/138 + mesh_gdf[k] = gdf.set_crs(v.ugrid.crs[k]) return mesh_gdf @@ -513,6 +541,8 @@ class MeshModel(MeshMixin, Model): """Model class Mesh Model for mesh models in HydroMT.""" + # FIXME: general remark: good to make it clear that the class supports only mesh created using UGrid and CF convention, and add a link in the doc to the convention + _CLI_ARGS = {"region": "setup_mesh2d", "res": "setup_mesh2d"} _NAME = "mesh_model" @@ -532,6 +562,7 @@ def __init__( data_libs=data_libs, logger=logger, ) + # FIXME: crs as global property? ## general setup methods def setup_mesh2d( From 4b0287a22deac36dd86c3042578ed9e86d361263 Mon Sep 17 00:00:00 2001 From: Xiaohan Li Date: Mon, 7 Aug 2023 13:36:05 +0200 Subject: [PATCH 10/15] review 1d discussed --- hydromt/models/model_mesh.py | 45 ++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 23 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 3d635b88b..56189f742 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -309,16 +309,12 @@ def set_mesh( if self._mesh is None: # NOTE: mesh is initialized with None # Check on crs if not data.ugrid.grid.crs: - raise ValueError( - "Data should have CRS." - ) # FIXME new data coming in should have crs + raise ValueError("Data should have CRS.") self._mesh = data else: # Check on crs if not data.ugrid.grid.crs == self.crs: - raise ValueError( - "Data and self.mesh should have the same CRS." - ) # FIXME since data and self.mesh have the same crs, can it be a global property? + raise ValueError("Data and self.mesh should have the same CRS.") # Save crs as it will be lost when converting to xarray crs = self.crs # Check on new grid topology @@ -400,7 +396,7 @@ def get_mesh( grid = self.mesh_grids[grid_name] uds = xu.UgridDataset( grid.to_dataset(optional_attributes=True) - ) # FIXME: it seems optional_attributes is needed when creating a new object for all variables about grid to be written, not needed when reading + ) # FIXME: would be nice if always write all attributes https://github.com/Deltares/xugrid/issues/140 uds.ugrid.grid.set_crs(grid.crs) # Look for data_vars that are defined on grid_name for var in self.mesh.data_vars: @@ -420,7 +416,7 @@ def get_mesh( def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None: """Read model mesh data at / and add to mesh property. - key-word arguments are passed to :py:func:`xu.open_dataset` + key-word arguments are passed to :py:func:`xr.open_dataset` # FIXME make doc consistent Parameters ---------- @@ -429,16 +425,17 @@ def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None crs : CRS, optional Coordinate Reference System (CRS) object representing the spatial reference system of the mesh file. **kwargs : dict - Additional keyword arguments to be passed to the `open_dataset` method. + Additional keyword arguments to be passed to the `_read_nc` method. # FIXME make doc consistent """ - # FIXME: general question, why there is no force overwrite argument for model update? + # FIXME: check how read_mesh behaves when multiple files are read self._assert_read_mode - uds = xu.open_dataset( - fn, **kwargs - ) # FIXME: switching self._read_nc to xu.open_dataset. is the idea to support single file? or also multiple files? e.g. xu.open_mfdataset - if not any( - uds.ugrid.crs.values() - ): # FIXME: crs cannot be write/read correctly by xugrid https://github.com/Deltares/xugrid/issues/138 + for ds in self._read_nc(fn, **kwargs).values(): + uds = xu.UgridDataset(ds) + if ds.rio.crs is not None: # parse crs + uds.ugrid.grid.set_crs(ds.raster.crs) + uds = uds.drop_vars(GEO_MAP_COORD, errors="ignore") + # FIXME how to apply crs if crs is missing? + if not any(uds.ugrid.crs.values()): if not crs: raise ValueError( "no crs is found in the file nor passed to the reader." @@ -451,7 +448,8 @@ def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None else: self.logger.info("crs is read from file") self._mesh = uds - # self.set_mesh(uds) # FIXME: infinit loop because self.mesh property is called in self.set_mesh, in which read_mesh is called. + # FIXME check how self.set_mesh(uds) behaves in the latest fix + # self.set_mesh(uds) def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: """Write model grid data to a netCDF file at /. @@ -507,6 +505,7 @@ def mesh_datasets(self) -> Dict: @property def mesh_names(self) -> List[str]: """List of grid names in mesh.""" + # FIXME would be nice if xugrid supports it: https://github.com/Deltares/xugrid/issues/141 if self.mesh is not None: return list(self.mesh_grids.keys()) else: @@ -529,9 +528,8 @@ def mesh_gdf(self) -> Dict: v[f"{grid.name}_node_id"] = xr.DataArray( v[dim].values.astype(str), dims=dim ) - gdf = ( - v.ugrid.to_geodataframe() - ) # FIXME xugrid crs not passed on correctly: https://github.com/Deltares/xugrid/issues/138 + gdf = v.ugrid.to_geodataframe() + # FIXME xugrid crs not passed on correctly: https://github.com/Deltares/xugrid/issues/138 mesh_gdf[k] = gdf.set_crs(v.ugrid.crs[k]) return mesh_gdf @@ -539,9 +537,11 @@ def mesh_gdf(self) -> Dict: class MeshModel(MeshMixin, Model): - """Model class Mesh Model for mesh models in HydroMT.""" + """Model class Mesh Model for mesh models in HydroMT. - # FIXME: general remark: good to make it clear that the class supports only mesh created using UGrid and CF convention, and add a link in the doc to the convention + Uses xugrid for working with unstructured grids, for data and topology stored according to UGRID conventions. + See also: xugrid + """ _CLI_ARGS = {"region": "setup_mesh2d", "res": "setup_mesh2d"} _NAME = "mesh_model" @@ -562,7 +562,6 @@ def __init__( data_libs=data_libs, logger=logger, ) - # FIXME: crs as global property? ## general setup methods def setup_mesh2d( From 764136d90c0bcd22a59ad53ec083932c3b312b84 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Tue, 8 Aug 2023 10:19:57 +0800 Subject: [PATCH 11/15] add review comments --- hydromt/models/model_mesh.py | 58 +++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 56189f742..18cbf2b8d 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -292,7 +292,7 @@ def set_mesh( raise ValueError( f"Cannot set mesh from {str(type(data).__name__)} without a name." ) - data = data.to_dataset() + data = data.to_dataset(optional_attributes=True) # Checks on grid topology # TODO: check if we support setting multiple grids at once. For now just one @@ -339,7 +339,9 @@ def set_mesh( " data variables in mesh." ) grids = [ - self.mesh_datasets[g].ugrid.to_dataset() + self.mesh_datasets[g].ugrid.to_dataset( + optional_attributes=True + ) for g in self.mesh_names if g != grid_name ] @@ -355,7 +357,12 @@ def set_mesh( else: # We are potentially adding a new grid without any data variables self._mesh = xu.UgridDataset( - xr.merge([self.mesh.ugrid.to_dataset(), data.ugrid.to_dataset()]) + xr.merge( + [ + self.mesh.ugrid.to_dataset(optional_attributes=True), + data.ugrid.to_dataset(optional_attributes=True), + ] + ) ) # Restore crs for grid in self._mesh.ugrid.grids: @@ -416,16 +423,19 @@ def get_mesh( def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None: """Read model mesh data at / and add to mesh property. - key-word arguments are passed to :py:func:`xr.open_dataset` # FIXME make doc consistent + key-word arguments are passed to :py:func:`xr.open_dataset` + # FIXME make doc consistent Parameters ---------- fn : str, optional filename relative to model root, by default 'mesh/mesh.nc' crs : CRS, optional - Coordinate Reference System (CRS) object representing the spatial reference system of the mesh file. + Coordinate Reference System (CRS) object representing the spatial reference + system of the mesh file. **kwargs : dict - Additional keyword arguments to be passed to the `_read_nc` method. # FIXME make doc consistent + Additional keyword arguments to be passed to the `_read_nc` method. + # FIXME make doc consistent """ # FIXME: check how read_mesh behaves when multiple files are read self._assert_read_mode @@ -473,7 +483,7 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: if not isdir(dirname(_fn)): os.makedirs(dirname(_fn)) self.logger.debug(f"Writing file {fn}") - ds_out = self.mesh.ugrid.to_dataset() + ds_out = self.mesh.ugrid.to_dataset(optional_attributes=True) if self.crs is not None: # save crs to spatial_ref coordinate ds_out = ds_out.rio.write_crs(self.crs) @@ -481,9 +491,10 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: # Other mesh properties @property - def mesh_grids(self) -> Dict: + def mesh_grids(self) -> Dict[str, Union[xu.Ugrid1d, xu.Ugrid2d]]: """Dictionnary of grid names and Ugrid topologies in mesh.""" grids = dict() + # FIXME to be replaced by uds.ugrid.topology. See https://github.com/Deltares/xugrid/issues/141 if self.mesh is not None: for grid in self.mesh.ugrid.grids: grids[grid.name] = grid @@ -491,7 +502,7 @@ def mesh_grids(self) -> Dict: return grids @property - def mesh_datasets(self) -> Dict: + def mesh_datasets(self) -> Dict[str, xu.UgridDataset]: """Dictionnary of grid names and corresponding UgridDataset topology and data variables in mesh.""" # noqa: E501 datasets = dict() if self._mesh is not None: @@ -505,7 +516,7 @@ def mesh_datasets(self) -> Dict: @property def mesh_names(self) -> List[str]: """List of grid names in mesh.""" - # FIXME would be nice if xugrid supports it: https://github.com/Deltares/xugrid/issues/141 + # FIXME to be replaced by uds.ugrid.name/uds.ugrid.name: https://github.com/Deltares/xugrid/issues/141 if self.mesh is not None: return list(self.mesh_grids.keys()) else: @@ -516,21 +527,18 @@ def mesh_gdf(self) -> Dict: """Returns dict of geometry of grids in mesh as a gpd.GeoDataFrame.""" mesh_gdf = dict() if self._mesh is not None: - for k, v in self.mesh_datasets.items(): + for k, grid in self.mesh_grids.items(): # works better on a Dataarray / Dataset # (need to add dummy variable if empty) - if len(v.data_vars) == 0: - grid = v.ugrid.grid - if grid.topology_dimension == 1: - dim = grid.edge_dimension - elif grid.topology_dimension == 2: - dim = grid.face_dimension - v[f"{grid.name}_node_id"] = xr.DataArray( - v[dim].values.astype(str), dims=dim - ) - gdf = v.ugrid.to_geodataframe() - # FIXME xugrid crs not passed on correctly: https://github.com/Deltares/xugrid/issues/138 - mesh_gdf[k] = gdf.set_crs(v.ugrid.crs[k]) + if grid.topology_dimension == 1: + dim = grid.edge_dimension + elif grid.topology_dimension == 2: + dim = grid.face_dimension + gdf = gpd.GeoDataFrame( + index=grid.to_dataset()[dim].values.astype(str), + geometry=grid.to_shapely(dim), + ) + mesh_gdf[k] = gdf.set_crs(grid.crs) return mesh_gdf @@ -539,7 +547,9 @@ class MeshModel(MeshMixin, Model): """Model class Mesh Model for mesh models in HydroMT. - Uses xugrid for working with unstructured grids, for data and topology stored according to UGRID conventions. + Uses xugrid for working with unstructured grids, for data and topology stored + according to UGRID conventions. + See also: xugrid """ From 11453ac42fe7525c1647b595bf6658cce239ef38 Mon Sep 17 00:00:00 2001 From: hboisgon Date: Tue, 8 Aug 2023 11:02:18 +0800 Subject: [PATCH 12/15] update read and write_mesh --- hydromt/models/model_mesh.py | 63 ++++++++++++++++++------------------ tests/test_model.py | 1 - 2 files changed, 31 insertions(+), 33 deletions(-) diff --git a/hydromt/models/model_mesh.py b/hydromt/models/model_mesh.py index 18cbf2b8d..ef528cef7 100644 --- a/hydromt/models/model_mesh.py +++ b/hydromt/models/model_mesh.py @@ -263,7 +263,7 @@ def set_mesh( Parameters ---------- data: xugrid.UgridDataArray or xugrid.UgridDataset - new layer to add to mesh, TODO support one grid only or multiple grids? + new layer to add to mesh, should contain only one grid topology. name: str, optional Name of new object layer, this is used to overwrite the name of a UgridDataArray. @@ -401,9 +401,7 @@ def get_mesh( raise ValueError(f"Grid {grid_name} not found in mesh.") if include_data: grid = self.mesh_grids[grid_name] - uds = xu.UgridDataset( - grid.to_dataset(optional_attributes=True) - ) # FIXME: would be nice if always write all attributes https://github.com/Deltares/xugrid/issues/140 + uds = xu.UgridDataset(grid.to_dataset(optional_attributes=True)) uds.ugrid.grid.set_crs(grid.crs) # Look for data_vars that are defined on grid_name for var in self.mesh.data_vars: @@ -420,32 +418,31 @@ def get_mesh( else: return self.mesh_grids[grid_name] - def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None: + def read_mesh( + self, fn: str = "mesh/mesh.nc", crs: Union[CRS, int] = None, **kwargs + ) -> None: """Read model mesh data at / and add to mesh property. key-word arguments are passed to :py:func:`xr.open_dataset` - # FIXME make doc consistent Parameters ---------- fn : str, optional filename relative to model root, by default 'mesh/mesh.nc' - crs : CRS, optional - Coordinate Reference System (CRS) object representing the spatial reference - system of the mesh file. + crs : CRS or int, optional + Coordinate Reference System (CRS) object or EPSG code representing the + spatial reference system of the mesh file. Only used if the CRS is not + found when reading the mesh file. **kwargs : dict Additional keyword arguments to be passed to the `_read_nc` method. - # FIXME make doc consistent """ - # FIXME: check how read_mesh behaves when multiple files are read self._assert_read_mode - for ds in self._read_nc(fn, **kwargs).values(): - uds = xu.UgridDataset(ds) - if ds.rio.crs is not None: # parse crs - uds.ugrid.grid.set_crs(ds.raster.crs) - uds = uds.drop_vars(GEO_MAP_COORD, errors="ignore") - # FIXME how to apply crs if crs is missing? - if not any(uds.ugrid.crs.values()): + ds = xr.merge(self._read_nc(fn, **kwargs).values()) + uds = xu.UgridDataset(ds) + if ds.rio.crs is not None: # parse crs + uds.ugrid.set_crs(ds.raster.crs) + uds = uds.drop_vars(GEO_MAP_COORD, errors="ignore") + else: if not crs: raise ValueError( "no crs is found in the file nor passed to the reader." @@ -455,24 +452,28 @@ def read_mesh(self, fn: str = "mesh/mesh.nc", crs: CRS = None, **kwargs) -> None self.logger.info( "no crs is found in the file, assigning from user input." ) - else: - self.logger.info("crs is read from file") self._mesh = uds - # FIXME check how self.set_mesh(uds) behaves in the latest fix - # self.set_mesh(uds) - def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: + def write_mesh( + self, + fn: str = "mesh/mesh.nc", + write_optional_ugrid_attributes: bool = True, + **kwargs, + ) -> None: """Write model grid data to a netCDF file at /. - Keyword arguments are passed to :py:meth:`xarray.Dataset.ugrid.to_netcdf`. + Keyword arguments are passed to :py:meth:`xarray.Dataset.to_netcdf`. Parameters ---------- fn : str, optional Filename relative to the model root directory, by default 'grid/grid.nc'. + write_optional_ugrid_attributes : bool, optional + If True, write optional ugrid attributes to the netCDF file, by default + True. **kwargs : dict Additional keyword arguments to be passed to the - `xarray.Dataset.ugrid.to_netcdf` method. + `xarray.Dataset.to_netcdf` method. """ if self._mesh is None: self.logger.debug("No mesh data found, skip writing.") @@ -483,7 +484,9 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: if not isdir(dirname(_fn)): os.makedirs(dirname(_fn)) self.logger.debug(f"Writing file {fn}") - ds_out = self.mesh.ugrid.to_dataset(optional_attributes=True) + ds_out = self.mesh.ugrid.to_dataset( + optional_attributes=write_optional_ugrid_attributes, + ) if self.crs is not None: # save crs to spatial_ref coordinate ds_out = ds_out.rio.write_crs(self.crs) @@ -494,7 +497,6 @@ def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None: def mesh_grids(self) -> Dict[str, Union[xu.Ugrid1d, xu.Ugrid2d]]: """Dictionnary of grid names and Ugrid topologies in mesh.""" grids = dict() - # FIXME to be replaced by uds.ugrid.topology. See https://github.com/Deltares/xugrid/issues/141 if self.mesh is not None: for grid in self.mesh.ugrid.grids: grids[grid.name] = grid @@ -516,20 +518,17 @@ def mesh_datasets(self) -> Dict[str, xu.UgridDataset]: @property def mesh_names(self) -> List[str]: """List of grid names in mesh.""" - # FIXME to be replaced by uds.ugrid.name/uds.ugrid.name: https://github.com/Deltares/xugrid/issues/141 if self.mesh is not None: - return list(self.mesh_grids.keys()) + return [grid.name for grid in self.mesh.ugrid.grids] else: return [] @property - def mesh_gdf(self) -> Dict: + def mesh_gdf(self) -> Dict[str, gpd.GeoDataFrame]: """Returns dict of geometry of grids in mesh as a gpd.GeoDataFrame.""" mesh_gdf = dict() if self._mesh is not None: for k, grid in self.mesh_grids.items(): - # works better on a Dataarray / Dataset - # (need to add dummy variable if empty) if grid.topology_dimension == 1: dim = grid.edge_dimension elif grid.topology_dimension == 2: diff --git a/tests/test_model.py b/tests/test_model.py index 179e9c3ab..71c696b02 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -198,7 +198,6 @@ def test_setup_region(model, demda, tmpdir): demda.raster.to_raster(grid_fn) model.setup_region({"grid": grid_fn}) assert np.all(demda.raster.bounds == model.region.total_bounds) - # # TODO model once we have registered the Model class entrypoint # basin model._geoms.pop("region") # remove old region model.setup_region({"basin": [12.2, 45.833333333333329]}) From 0a6988acda30444573a642b07171b24e19d67e4b Mon Sep 17 00:00:00 2001 From: hboisgon Date: Mon, 21 Aug 2023 13:14:24 +0800 Subject: [PATCH 13/15] import xugrid at the top --- hydromt/workflows/basin_mask.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/hydromt/workflows/basin_mask.py b/hydromt/workflows/basin_mask.py index 69aa8078b..7d3f8ccd1 100644 --- a/hydromt/workflows/basin_mask.py +++ b/hydromt/workflows/basin_mask.py @@ -22,6 +22,9 @@ from ..io import open_raster from ..models import MODELS +if _compat.HAS_XUGRID: + import xugrid as xu + logger = logging.getLogger(__name__) __all__ = ["get_basin_geometry", "parse_region"] @@ -143,8 +146,6 @@ def parse_region(region, logger=logger): 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): From 059160ac188d480ce8fa5d77229643aa7ff941e6 Mon Sep 17 00:00:00 2001 From: Sam Vente Date: Mon, 21 Aug 2023 09:03:40 +0200 Subject: [PATCH 14/15] double test timeout limimt --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ba73bd39f..b5314084f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ include = ["hydromt"] exclude = ["docs", "examples", "envs", "tests", "binder", ".github"] [tool.pytest.ini_options] -addopts = "--ff --timeout=60" +addopts = "--ff --timeout=120" testpaths = ["tests"] filterwarnings = [ From 73a3c7450ee0dd79008fcf1efacd97701f042651 Mon Sep 17 00:00:00 2001 From: Sam Vente Date: Mon, 21 Aug 2023 09:22:20 +0200 Subject: [PATCH 15/15] remove failfast from ci so we can still test all versions --- .github/workflows/tests.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f91b1aa6c..e5b776975 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -23,6 +23,7 @@ jobs: run: shell: bash -l {0} strategy: + fail-fast: false matrix: os: [ubuntu-latest] python-version: ['3.9','3.10','3.11']