diff --git a/docs/examples/circle.py b/docs/examples/circle.py index 5934216a..e381fc86 100644 --- a/docs/examples/circle.py +++ b/docs/examples/circle.py @@ -118,9 +118,9 @@ def plot_head_ugrid(head, cbc, workspace): ) # Build the xugrid Ugrid2d mesh from the Disv package. -# `disv.to_grid()` returns a `VertexGrid`; `.ugrid` converts it to an -# `xu.Ugrid2d` object suitable for xugrid operations. -grid = disv.to_grid().ugrid +# `disv.to_grid()` returns a `VertexGrid`; `.to_xarray()` returns an +# `xu.UgridDataset`; `.grids[0]` extracts the `xu.Ugrid2d` mesh object. +grid = disv.to_grid().to_xarray().grids[0] # `dims` captures array shapes needed by packages that pre-allocate xarray storage. dims = {"nper": nper, "nlay": nlay, "ncpl": ncpl, "nvert": len(vertices), "nodes": nlay * ncpl} diff --git a/docs/examples/frenchman-flat.py b/docs/examples/frenchman-flat.py index 012d5e49..7879f29e 100644 --- a/docs/examples/frenchman-flat.py +++ b/docs/examples/frenchman-flat.py @@ -79,7 +79,7 @@ def plot_head_ugrid(head, cbc, grid, workspace): import xarray as xr import xugrid as xu - ugrid = grid.ugrid + ugrid = grid.to_xarray(netcdf_format=NetCDFFormat.LAYERED_MESH).grids[0] facedim = ugrid.face_dimension # Select first timestep and first layer; flatten (y, x) -> face dimension @@ -113,7 +113,7 @@ def plot_head_ugrid(head, cbc, grid, workspace): ds = ds.ugrid.assign_face_coords() fig, ax = plt.subplots(figsize=(10, 8)) - ds.plot.quiver(x="mesh2d_face_x", y="mesh2d_face_y", u="u", v="v", color="black", scale=100) + ds.plot.quiver(x="mesh_face_x", y="mesh_face_y", u="u", v="v", color="black", scale=100) xu.plot.line(ugrid, ax=ax, color="black", linewidth=0.2) ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") ax.set_title("Frenchman Flat flow vectors overlaid on Mesh") diff --git a/flopy4/mf6/netcdf.py b/flopy4/mf6/netcdf.py index b29ffe2f..ca78f6be 100644 --- a/flopy4/mf6/netcdf.py +++ b/flopy4/mf6/netcdf.py @@ -3,6 +3,7 @@ import numpy as np import xarray as xr +import xugrid as xu from pydantic import ( BaseModel, ConfigDict, @@ -102,7 +103,7 @@ def from_dict(cls, meta: dict, context: dict | None): pass @abc.abstractmethod - def to_xarray(self) -> xr.Dataset: + def to_xarray(self) -> "xr.Dataset | xu.UgridDataset": """create xarray dataset.""" pass @@ -228,11 +229,11 @@ def from_model( nc_model.time = time return nc_model - def to_xarray(self) -> xr.Dataset: + def to_xarray(self) -> "xr.Dataset | xu.UgridDataset": import datetime - dss = [] meta = self.model_dump(by_alias=True) + grid_result = None if self._grid is not None and self._time is not None: # type: ignore conventions = "CF-1.11" # type: ignore @@ -243,14 +244,20 @@ def to_xarray(self) -> xr.Dataset: if meta["attrs"]["mesh"] is not None else NetCDFFormat.STRUCTURED ) - dss.append(self._grid.to_xarray(netcdf_format=_fmt, modeltime=self._time)) + grid_result = self._grid.to_xarray(netcdf_format=_fmt, modeltime=self._time) meta["attrs"]["Conventions"] = conventions + pkg_dss = [] for p in self.packages: p._context["grid"] = self.grid - dss.append(p.to_xarray()) + pkg_dss.append(p.to_xarray()) - ds = xr.merge(dss) + if isinstance(grid_result, xu.UgridDataset): + merged = xr.merge([grid_result.obj, *pkg_dss]) + ds = xu.UgridDataset(merged, grids=grid_result.grids) + else: + all_dss = ([grid_result] if grid_result is not None else []) + pkg_dss + ds = xr.merge(all_dss) if all_dss else xr.Dataset() dt = datetime.datetime.now() timestamp = dt.strftime("%m/%d/%Y %H:%M:%S") @@ -263,7 +270,31 @@ def to_xarray(self) -> xr.Dataset: return ds def to_netcdf(self, path: str | PathLike) -> None: - self.to_xarray().to_netcdf(path) + result = self.to_xarray() + if isinstance(result, xu.UgridDataset): + grid = result.grids[0] + topo = grid.assign_face_coords(grid.to_dataset()) + # Rename xugrid's generated dimension names to MF6's UGRID convention + # so input files are consistent with MODFLOW 6 output files. + _dim_rename = { + k: v + for k, v in { + "mesh_nFaces": "nmesh_face", + "mesh_nNodes": "nmesh_node", + "mesh_nMax_face_nodes": "max_nmesh_face_nodes", + }.items() + if k in topo.dims + } + if _dim_rename: + topo = topo.rename(_dim_rename) + for attr in ("face_dimension", "node_dimension", "max_face_nodes_dimension"): + if topo["mesh"].attrs.get(attr) in _dim_rename: + topo["mesh"].attrs[attr] = _dim_rename[topo["mesh"].attrs[attr]] + merged = topo.merge(result.obj) + merged.attrs = result.obj.attrs + merged.to_netcdf(path) + else: + result.to_netcdf(path) @property def meta(self): diff --git a/flopy4/mf6/utils/codegen/make.py b/flopy4/mf6/utils/codegen/make.py index c0b725ed..e895ccc5 100644 --- a/flopy4/mf6/utils/codegen/make.py +++ b/flopy4/mf6/utils/codegen/make.py @@ -191,8 +191,7 @@ def _expand_record_field( spec_call="", generatable=False, skip_reason=( - f"positional sub-fields not yet supported: " - f"{', '.join(unexpandable_optional)}" + f"positional sub-fields not yet supported: {', '.join(unexpandable_optional)}" ), ) ) diff --git a/flopy4/mf6/utils/grid.py b/flopy4/mf6/utils/grid.py index 91f24604..5ea8eef5 100644 --- a/flopy4/mf6/utils/grid.py +++ b/flopy4/mf6/utils/grid.py @@ -491,134 +491,87 @@ def to_xarray( modeltime=None, netcdf_format: NetCDFFormat = NetCDFFormat.STRUCTURED, configuration=None, - ): - """ - modeltime : FloPy ModelTime object - netcdf_format: NetCDFFormat.LAYERED_MESH or NetCDFFormat.STRUCTURED - configuration: configuration dictionary - """ - if modeltime is None: - raise ValueError("modeltime required for dataset timeseries") - + ) -> "xr.Dataset | xu.UgridDataset": self.legacy = True try: + if netcdf_format == NetCDFFormat.LAYERED_MESH: + return self._layered_mesh_dataset(modeltime, configuration) + if modeltime is None: + raise ValueError("modeltime required for structured dataset timeseries") ds = xr.Dataset() ds.attrs["modflow_grid"] = "STRUCTURED" - - if netcdf_format == NetCDFFormat.LAYERED_MESH: - ds = self._layered_mesh_dataset(ds, modeltime, configuration) - else: - ds = self._structured_dataset(ds, modeltime, configuration) - - return ds + return self._structured_dataset(ds, modeltime, configuration) finally: self.legacy = False - def _layered_mesh_dataset(self, ds, modeltime=None, configuration=None): - lenunits = {0: "unknown", 1: "ft", 2: "m", 3: "cm"} - - # All topology arrays computed once; self.legacy is already True here. + def _layered_mesh_dataset(self, modeltime=None, configuration=None) -> xu.UgridDataset: topo = self._topology() + ugrid2d = xu.Ugrid2d( + topo["node_x"], + topo["node_y"], + FILL_INT64, + topo["face_nodes"], + name="mesh", + projected=True, + start_index=1, + ) - # create dataset coordinate vars - # Use cumulative per-period time (one value per stress period) - var_d = { - "time": (["time"], np.cumsum(modeltime.perlen)), - } - ds = ds.assign(var_d) - ds["time"].attrs["calendar"] = "standard" - _tunits = modeltime.time_units if modeltime.time_units not in (None, "unknown") else "days" - ds["time"].attrs["units"] = f"{_tunits} since {modeltime.start_datetime}" - ds["time"].attrs["axis"] = "T" - ds["time"].attrs["standard_name"] = "time" - ds["time"].attrs["long_name"] = "time" - ds["time"].encoding["_FillValue"] = None - ds = ds.assign_coords({"layer": ("layer", np.arange(1, self.nlay + 1))}) - ds["layer"].attrs["long_name"] = "model layer" - ds["layer"].attrs["units"] = "1" - ds["layer"].attrs["positive"] = "down" - ds["layer"].attrs["axis"] = "Z" - ds["layer"].encoding["_FillValue"] = None - - # mesh container variable - ds = ds.assign({"mesh": ([], np.int64(1))}) - ds["mesh"].attrs["cf_role"] = "mesh_topology" - ds["mesh"].attrs["long_name"] = "2D mesh topology" - ds["mesh"].attrs["topology_dimension"] = np.int64(2) - ds["mesh"].attrs["face_dimension"] = "nmesh_face" - ds["mesh"].attrs["node_coordinates"] = "mesh_node_x mesh_node_y" - ds["mesh"].attrs["face_coordinates"] = "mesh_face_x mesh_face_y" - ds["mesh"].attrs["face_node_connectivity"] = "mesh_face_nodes" - - # mesh node x and y - var_d = { - "mesh_node_x": (["nmesh_node"], topo["node_x"]), - "mesh_node_y": (["nmesh_node"], topo["node_y"]), - } - ds = ds.assign(var_d) - ds["mesh_node_x"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_node_x"].attrs["standard_name"] = "projection_x_coordinate" - ds["mesh_node_x"].attrs["long_name"] = "Easting" - ds["mesh_node_x"].encoding["_FillValue"] = None - ds["mesh_node_y"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_node_y"].attrs["standard_name"] = "projection_y_coordinate" - ds["mesh_node_y"].attrs["long_name"] = "Northing" - ds["mesh_node_y"].encoding["_FillValue"] = None - - # mesh face x, y and bounds - var_d = { - "mesh_face_x": (["nmesh_face"], topo["face_x"]), - "mesh_face_xbnds": (["nmesh_face", "max_nmesh_face_nodes"], topo["x_bnds"]), - "mesh_face_y": (["nmesh_face"], topo["face_y"]), - "mesh_face_ybnds": (["nmesh_face", "max_nmesh_face_nodes"], topo["y_bnds"]), - } - ds = ds.assign(var_d) - ds["mesh_face_x"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_face_x"].attrs["standard_name"] = "projection_x_coordinate" - ds["mesh_face_x"].attrs["long_name"] = "Easting" - ds["mesh_face_x"].attrs["bounds"] = "mesh_face_xbnds" - ds["mesh_face_x"].encoding["_FillValue"] = None - ds["mesh_face_y"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_face_y"].attrs["standard_name"] = "projection_y_coordinate" - ds["mesh_face_y"].attrs["long_name"] = "Northing" - ds["mesh_face_y"].attrs["bounds"] = "mesh_face_ybnds" - ds["mesh_face_y"].encoding["_FillValue"] = None - - # mesh face nodes - var_d = { - "mesh_face_nodes": (["nmesh_face", "max_nmesh_face_nodes"], topo["face_nodes"]), - } - ds = ds.assign(var_d) - ds["mesh_face_nodes"].attrs["cf_role"] = "face_node_connectivity" - ds["mesh_face_nodes"].attrs["long_name"] = "Vertices bounding cell (counterclockwise)" - ds["mesh_face_nodes"].attrs["start_index"] = np.int64(1) - ds["mesh_face_nodes"].encoding["_FillValue"] = FILL_INT64 + ds = xr.Dataset() + ds.attrs["modflow_grid"] = "STRUCTURED" + + ds = ds.assign( + { + "mesh_face_xbnds": xr.DataArray( + topo["x_bnds"], dims=["nmesh_face", "max_nmesh_face_nodes"] + ), + "mesh_face_ybnds": xr.DataArray( + topo["y_bnds"], dims=["nmesh_face", "max_nmesh_face_nodes"] + ), + } + ) + ds["mesh_face_xbnds"].attrs["long_name"] = "x bounds of mesh faces" + ds["mesh_face_ybnds"].attrs["long_name"] = "y bounds of mesh faces" + ds["mesh_face_xbnds"].encoding["_FillValue"] = None + ds["mesh_face_ybnds"].encoding["_FillValue"] = None + + if modeltime is not None: + _tunits = ( + modeltime.time_units if modeltime.time_units not in (None, "unknown") else "days" + ) + ds = ds.assign({"time": (["time"], np.cumsum(modeltime.perlen))}) + ds["time"].attrs["calendar"] = "standard" + ds["time"].attrs["units"] = f"{_tunits} since {modeltime.start_datetime}" + ds["time"].attrs["axis"] = "T" + ds["time"].attrs["standard_name"] = "time" + ds["time"].attrs["long_name"] = "time" + ds["time"].encoding["_FillValue"] = None + ds = ds.assign_coords({"layer": ("layer", np.arange(1, self.nlay + 1))}) + ds["layer"].attrs["long_name"] = "model layer" + ds["layer"].attrs["units"] = "1" + ds["layer"].attrs["positive"] = "down" + ds["layer"].attrs["axis"] = "Z" + ds["layer"].encoding["_FillValue"] = None wkt_configured = ( configuration is not None and "wkt" in configuration and configuration["wkt"] is not None ) - if wkt_configured or self.crs is not None: - ds["mesh_node_x"].attrs["grid_mapping"] = "projection" - ds["mesh_node_y"].attrs["grid_mapping"] = "projection" - ds["mesh_face_x"].attrs["grid_mapping"] = "projection" - ds["mesh_face_y"].attrs["grid_mapping"] = "projection" - ds = ds.assign({"projection": ([], np.int64(1))}) from pyproj import CRS as ProjCRS from pyproj.enums import WktVersion _crs = ProjCRS.from_wkt(configuration["wkt"]) if wkt_configured else self.crs _wkt1 = _crs.to_wkt(WktVersion.WKT1_GDAL) _wkt2 = _crs.to_wkt(WktVersion.WKT2_2019) + ds = ds.assign({"projection": ([], np.int64(1))}) ds["projection"].attrs["wkt"] = _wkt1 ds["projection"].attrs["crs_wkt"] = _wkt2 _gmn = _crs.to_cf().get("grid_mapping_name") if _gmn: ds["projection"].attrs["grid_mapping_name"] = _gmn - return ds + return xu.UgridDataset(ds, grids=ugrid2d) def _structured_dataset(self, ds, modeltime=None, configuration=None): # Produces a conventional CF structured dataset (x/y dimension coords, @@ -785,32 +738,6 @@ def _topology(self) -> dict: "y_bnds": y_bnds, } - @property - def ugrid(self) -> xu.Ugrid2d: - """ - Build a :class:`xugrid.Ugrid2d` mesh from this structured grid. - - Returns - ------- - xu.Ugrid2d - A 2-D unstructured grid object whose faces correspond to the - structured grid cells in row-major order. - """ - self.legacy = True - try: - topo = self._topology() - return xu.Ugrid2d( - topo["node_x"], - topo["node_y"], - FILL_INT64, - topo["face_nodes"], - projected=True, - crs=self.crs, - start_index=1, - ) - finally: - self.legacy = False - class VertexGrid(LegacyVertexGrid): """ @@ -1132,129 +1059,78 @@ def to_xarray( modeltime=None, netcdf_format: NetCDFFormat = NetCDFFormat.LAYERED_MESH, configuration=None, - ): - """ - modeltime : FloPy ModelTime object - netcdf_format: must be NetCDFFormat.LAYERED_MESH; VertexGrid only - supports UGRID layered-mesh output - configuration: configuration dictionary - """ + ) -> xu.UgridDataset: if netcdf_format != NetCDFFormat.LAYERED_MESH: raise ValueError("VertexGrid only supports NetCDFFormat.LAYERED_MESH output") - if modeltime is None: - raise ValueError("modeltime required for dataset timeseries") - - lenunits = {0: "unknown", 1: "ft", 2: "m", 3: "cm"} - self.legacy = True try: - # All topology arrays computed once from cell2d/verts. topo = self._topology() + ugrid2d = xu.Ugrid2d( + topo["node_x"], + topo["node_y"], + FILL_INT64, + topo["face_nodes"], + name="mesh", + projected=True, + start_index=1, + ) ds = xr.Dataset() ds.attrs["modflow_grid"] = "VERTEX" - # create dataset coordinate vars - # Use cumulative per-period time (one value per stress period) - var_d = { - "time": (["time"], np.cumsum(modeltime.perlen)), - } - ds = ds.assign(var_d) - ds["time"].attrs["calendar"] = "standard" - _tu = modeltime.time_units - _tunits = _tu if _tu not in (None, "unknown") else "days" - ds["time"].attrs["units"] = f"{_tunits} since {modeltime.start_datetime}" - ds["time"].attrs["axis"] = "T" - ds["time"].attrs["standard_name"] = "time" - ds["time"].attrs["long_name"] = "time" - ds["time"].encoding["_FillValue"] = None - ds = ds.assign_coords({"layer": ("layer", np.arange(1, self.nlay + 1))}) - ds["layer"].attrs["long_name"] = "model layer" - ds["layer"].attrs["units"] = "1" - ds["layer"].attrs["positive"] = "down" - ds["layer"].attrs["axis"] = "Z" - ds["layer"].encoding["_FillValue"] = None - - # mesh container variable - ds = ds.assign({"mesh": ([], np.int64(1))}) - ds["mesh"].attrs["cf_role"] = "mesh_topology" - ds["mesh"].attrs["long_name"] = "2D mesh topology" - ds["mesh"].attrs["topology_dimension"] = np.int64(2) - ds["mesh"].attrs["face_dimension"] = "nmesh_face" - ds["mesh"].attrs["node_coordinates"] = "mesh_node_x mesh_node_y" - ds["mesh"].attrs["face_coordinates"] = "mesh_face_x mesh_face_y" - ds["mesh"].attrs["face_node_connectivity"] = "mesh_face_nodes" - - # mesh node x and y - var_d = { - "mesh_node_x": (["nmesh_node"], topo["node_x"]), - "mesh_node_y": (["nmesh_node"], topo["node_y"]), - } - ds = ds.assign(var_d) - ds["mesh_node_x"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_node_x"].attrs["standard_name"] = "projection_x_coordinate" - ds["mesh_node_x"].attrs["long_name"] = "Easting" - ds["mesh_node_x"].encoding["_FillValue"] = None - ds["mesh_node_y"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_node_y"].attrs["standard_name"] = "projection_y_coordinate" - ds["mesh_node_y"].attrs["long_name"] = "Northing" - ds["mesh_node_y"].encoding["_FillValue"] = None - - # mesh face x, y and bounds - var_d = { - "mesh_face_x": (["nmesh_face"], topo["face_x"]), - "mesh_face_xbnds": (["nmesh_face", "max_nmesh_face_nodes"], topo["x_bnds"]), - "mesh_face_y": (["nmesh_face"], topo["face_y"]), - "mesh_face_ybnds": (["nmesh_face", "max_nmesh_face_nodes"], topo["y_bnds"]), - } - ds = ds.assign(var_d) - ds["mesh_face_x"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_face_x"].attrs["standard_name"] = "projection_x_coordinate" - ds["mesh_face_x"].attrs["long_name"] = "Easting" - ds["mesh_face_x"].attrs["bounds"] = "mesh_face_xbnds" - ds["mesh_face_x"].encoding["_FillValue"] = None - ds["mesh_face_y"].attrs["units"] = lenunits.get(self.lenuni, "unknown") - ds["mesh_face_y"].attrs["standard_name"] = "projection_y_coordinate" - ds["mesh_face_y"].attrs["long_name"] = "Northing" - ds["mesh_face_y"].attrs["bounds"] = "mesh_face_ybnds" - ds["mesh_face_y"].encoding["_FillValue"] = None - - # mesh face nodes - var_d = { - "mesh_face_nodes": (["nmesh_face", "max_nmesh_face_nodes"], topo["face_nodes"]), - } - ds = ds.assign(var_d) - ds["mesh_face_nodes"].attrs["cf_role"] = "face_node_connectivity" - ds["mesh_face_nodes"].attrs["long_name"] = "Vertices bounding cell (counterclockwise)" - ds["mesh_face_nodes"].attrs["start_index"] = np.int64(1) - ds["mesh_face_nodes"].encoding["_FillValue"] = FILL_INT64 + ds = ds.assign( + { + "mesh_face_xbnds": xr.DataArray( + topo["x_bnds"], dims=["nmesh_face", "max_nmesh_face_nodes"] + ), + "mesh_face_ybnds": xr.DataArray( + topo["y_bnds"], dims=["nmesh_face", "max_nmesh_face_nodes"] + ), + } + ) + ds["mesh_face_xbnds"].attrs["long_name"] = "x bounds of mesh faces" + ds["mesh_face_ybnds"].attrs["long_name"] = "y bounds of mesh faces" + ds["mesh_face_xbnds"].encoding["_FillValue"] = FILL_INT64 + ds["mesh_face_ybnds"].encoding["_FillValue"] = FILL_INT64 + + if modeltime is not None: + _tu = modeltime.time_units + _tunits = _tu if _tu not in (None, "unknown") else "days" + ds = ds.assign({"time": (["time"], np.cumsum(modeltime.perlen))}) + ds["time"].attrs["calendar"] = "standard" + ds["time"].attrs["units"] = f"{_tunits} since {modeltime.start_datetime}" + ds["time"].attrs["axis"] = "T" + ds["time"].attrs["standard_name"] = "time" + ds["time"].attrs["long_name"] = "time" + ds["time"].encoding["_FillValue"] = None + ds = ds.assign_coords({"layer": ("layer", np.arange(1, self.nlay + 1))}) + ds["layer"].attrs["long_name"] = "model layer" + ds["layer"].attrs["units"] = "1" + ds["layer"].attrs["positive"] = "down" + ds["layer"].attrs["axis"] = "Z" + ds["layer"].encoding["_FillValue"] = None wkt_configured = ( configuration is not None and "wkt" in configuration and configuration["wkt"] is not None ) - if wkt_configured or self.crs is not None: - ds["mesh_node_x"].attrs["grid_mapping"] = "projection" - ds["mesh_node_y"].attrs["grid_mapping"] = "projection" - ds["mesh_face_x"].attrs["grid_mapping"] = "projection" - ds["mesh_face_y"].attrs["grid_mapping"] = "projection" - ds = ds.assign({"projection": ([], np.int64(1))}) from pyproj import CRS as ProjCRS from pyproj.enums import WktVersion _crs = ProjCRS.from_wkt(configuration["wkt"]) if wkt_configured else self.crs _wkt1 = _crs.to_wkt(WktVersion.WKT1_GDAL) _wkt2 = _crs.to_wkt(WktVersion.WKT2_2019) + ds = ds.assign({"projection": ([], np.int64(1))}) ds["projection"].attrs["wkt"] = _wkt1 ds["projection"].attrs["crs_wkt"] = _wkt2 _gmn = _crs.to_cf().get("grid_mapping_name") if _gmn: ds["projection"].attrs["grid_mapping_name"] = _gmn - return ds + return xu.UgridDataset(ds, grids=ugrid2d) finally: self.legacy = False @@ -1307,32 +1183,6 @@ def _topology(self) -> dict: "y_bnds": y_bnds, } - @property - def ugrid(self) -> xu.Ugrid2d: - """ - Build a :class:`xugrid.Ugrid2d` mesh from this vertex grid. - - Returns - ------- - xu.Ugrid2d - A 2-D unstructured grid object whose faces correspond to the - vertex grid cells. - """ - self.legacy = True - try: - topo = self._topology() - return xu.Ugrid2d( - topo["node_x"], - topo["node_y"], - FILL_INT64, - topo["face_nodes"], - projected=True, - crs=self.crs, - start_index=1, - ) - finally: - self.legacy = False - def get_coords(grid: LegacyStructuredGrid) -> dict[str, Any]: # unpack tuples diff --git a/test/test_mf6_component.py b/test/test_mf6_component.py index cfa5f547..b7a4be60 100644 --- a/test/test_mf6_component.py +++ b/test/test_mf6_component.py @@ -6,6 +6,7 @@ import pandas as pd import pytest import xarray as xr +import xugrid from xarray import DataTree from flopy4.mf6.component import COMPONENTS @@ -1114,7 +1115,9 @@ def test_ugrid_from_dis_factory(): botm_near_x15 = grid.botm.sel(x=15.0, method="nearest") assert botm_near_x15.shape == (2, 5) - ugrid = grid.ugrid + uds = grid.to_xarray(netcdf_format=NetCDFFormat.LAYERED_MESH) + assert isinstance(uds, xugrid.UgridDataset) + ugrid = uds.grids[0] assert ugrid.n_face == nrow * ncol assert ugrid.n_node == (nrow + 1) * (ncol + 1) @@ -1217,30 +1220,26 @@ def test_ugrid_from_disv_factory(): np.testing.assert_allclose(grid.dataset.coords["z"].values[1], np.full((ncpl), -15.0)) np.testing.assert_allclose(grid.dataset.coords["z"].values[2], np.full((ncpl), -25.0)) - ugrid = grid.ugrid + uds = grid.to_xarray() + assert isinstance(uds, xugrid.UgridDataset) + ugrid = uds.grids[0] assert ugrid.n_face == ncpl assert ugrid.n_node == nvert udata = xugrid.UgridDataArray(dis.top, grid=ugrid) udataset = xugrid.UgridDataset(dis.data.dataset, grids=ugrid) - - # udata.to_netcdf("./udata.nc") - # drop global attributes, or filter? - uds = udataset.drop_attrs() - # drop objects that need serialization - uds = uds.drop_vars(["cell2ddata"]) - # uds.to_netcdf("./udataset.nc") + udataset = udataset.drop_attrs() + udataset = udataset.drop_vars(["cell2ddata"]) def test_ugrid_from_dis_uniform(): - """StructuredGrid.uniform().ugrid should have the correct face/node counts.""" - import xugrid - + """StructuredGrid.uniform().to_xarray() should return an UgridDataset.""" nrow, ncol = 4, 6 grid = StructuredGrid.uniform(nlay=2, nrow=nrow, ncol=ncol, delr=50.0, delc=50.0) - ugrid = grid.ugrid - assert isinstance(ugrid, xugrid.Ugrid2d) + uds = grid.to_xarray(netcdf_format=NetCDFFormat.LAYERED_MESH) + assert isinstance(uds, xugrid.UgridDataset) + ugrid = uds.grids[0] assert ugrid.n_face == nrow * ncol assert ugrid.n_node == (nrow + 1) * (ncol + 1) diff --git a/test/test_netcdf_cf_conventions.py b/test/test_netcdf_cf_conventions.py index da7e7357..be8e7223 100644 --- a/test/test_netcdf_cf_conventions.py +++ b/test/test_netcdf_cf_conventions.py @@ -13,6 +13,7 @@ import numpy as np import pytest import xarray as xr +import xugrid as xu from flopy.discretization.modeltime import ModelTime from flopy4.mf6.enums import NetCDFFormat @@ -127,10 +128,13 @@ def test_structured_dis_no_latlon(structured_grid, modeltime): def test_mesh_dis_topology_variable(structured_grid, modeltime): - """Mesh topology variable must be present with cf_role = mesh_topology.""" - ds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH) - assert "mesh" in ds, "mesh topology variable missing" - assert ds["mesh"].attrs.get("cf_role") == "mesh_topology" + """UGRID topology must be present as an xu.Ugrid2d grid on the returned UgridDataset.""" + uds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH) + assert isinstance(uds, xu.UgridDataset), "expected UgridDataset for layered-mesh format" + assert len(uds.grids) == 1, "expected exactly one UGRID grid" + assert isinstance(uds.grids[0], xu.Ugrid2d) + assert uds.grids[0].name == "mesh", "mesh topology variable missing" + assert uds.grids[0].to_dataset()["mesh"].attrs.get("cf_role") == "mesh_topology" def test_mesh_dis_projection_crs_attrs(structured_grid, modeltime): @@ -164,10 +168,13 @@ def test_mesh_dis_data_var_location_attr(structured_grid): def test_mesh_disv_topology_variable(vertex_grid, modeltime): - """Mesh topology variable must be present with cf_role = mesh_topology.""" - ds = vertex_grid.to_xarray(modeltime=modeltime) - assert "mesh" in ds - assert ds["mesh"].attrs.get("cf_role") == "mesh_topology" + """UGRID topology must be present as an xu.Ugrid2d grid on the returned UgridDataset.""" + uds = vertex_grid.to_xarray(modeltime=modeltime) + assert isinstance(uds, xu.UgridDataset), "expected UgridDataset for vertex grid" + assert len(uds.grids) == 1, "expected exactly one UGRID grid" + assert isinstance(uds.grids[0], xu.Ugrid2d) + assert uds.grids[0].name == "mesh" + assert uds.grids[0].to_dataset()["mesh"].attrs.get("cf_role") == "mesh_topology" def test_mesh_disv_projection_crs_attrs(vertex_grid, modeltime): @@ -229,27 +236,37 @@ def test_mesh_disv_layer_coord(vertex_grid, modeltime): assert ds["layer"].attrs.get("positive") == "down" -def test_mesh_dis_face_nodes_fill_in_encoding(structured_grid, modeltime): - """mesh_face_nodes _FillValue must be in encoding, not attrs (DIS path). +def test_mesh_dis_face_bounds(structured_grid, modeltime): + """mesh_face_xbnds/ybnds must be present with shape (nfaces, 4) for a structured grid.""" + uds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH) + nfaces = structured_grid.nrow * structured_grid.ncol + for var in ("mesh_face_xbnds", "mesh_face_ybnds"): + assert var in uds, f"{var} missing from layered-mesh DIS dataset" + assert uds[var].shape == (nfaces, 4), f"{var} shape mismatch: {uds[var].shape}" - xarray only writes a proper NetCDF _FillValue declaration when the fill - value is in encoding. If it lands in attrs it becomes a plain variable - attribute and NetCDF4/UGRID readers will not recognise padding slots. - """ + +def test_mesh_disv_face_bounds(vertex_grid, modeltime): + """mesh_face_xbnds/ybnds must be present with shape (ncpl, max_nodes) for a vertex grid.""" + uds = vertex_grid.to_xarray(modeltime=modeltime) + for var in ("mesh_face_xbnds", "mesh_face_ybnds"): + assert var in uds, f"{var} missing from layered-mesh DISV dataset" + assert uds[var].shape[0] == vertex_grid.ncpl, f"{var} face dimension mismatch" + + +def test_mesh_dis_face_nodes_fill_in_encoding(structured_grid, modeltime): + """The UGRID fill value passed to xu.Ugrid2d must match FILL_INT64 (DIS path).""" from flopy4.mf6.constants import FILL_INT64 - ds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH) - assert "_FillValue" not in ds["mesh_face_nodes"].attrs, "_FillValue must not be in attrs" - assert ds["mesh_face_nodes"].encoding.get("_FillValue") == FILL_INT64 + uds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH) + assert uds.grids[0].fill_value == FILL_INT64 def test_mesh_disv_face_nodes_fill_in_encoding(vertex_grid, modeltime): - """mesh_face_nodes _FillValue must be in encoding, not attrs (DISV path).""" + """The UGRID fill value passed to xu.Ugrid2d must match FILL_INT64 (DISV path).""" from flopy4.mf6.constants import FILL_INT64 - ds = vertex_grid.to_xarray(modeltime=modeltime) - assert "_FillValue" not in ds["mesh_face_nodes"].attrs, "_FillValue must not be in attrs" - assert ds["mesh_face_nodes"].encoding.get("_FillValue") == FILL_INT64 + uds = vertex_grid.to_xarray(modeltime=modeltime) + assert uds.grids[0].fill_value == FILL_INT64 def test_structured_merged_sel_layer(structured_grid, modeltime): @@ -277,22 +294,13 @@ def test_structured_merged_sel_layer(structured_grid, modeltime): assert sliced["npf_k"].dims == ("y", "x") -def _raw_var_attrs(path, varname: str) -> dict: - """Read raw NetCDF variable attributes bypassing xarray's CF decoder.""" - nc4 = pytest.importorskip("netCDF4") - ds = nc4.Dataset(path) - try: - return {a: ds.variables[varname].getncattr(a) for a in ds.variables[varname].ncattrs()} - finally: - ds.close() +def test_structured_data_var_no_redundant_coordinates(structured_grid, modeltime): + """Structured data vars must not carry non-dimension coords for y/x. - -def test_structured_data_var_no_redundant_coordinates(structured_grid, modeltime, tmp_path): - """Structured data vars must not carry a coordinates attr pointing to dimension coords. - - y and x are dimension coordinates — CF tools find them via the dimension names. - A redundant coordinates attr would rely on xarray encoding internals to survive - to_netcdf and adds no information for any CF-aware consumer. + y and x are dimension coordinates — CF tools find them via the dimension + names. A redundant coordinates attr would cause xarray to emit a + coordinates attribute on the variable during to_netcdf, which adds no + information for any CF-aware consumer and can confuse some tools. """ grid_ds = structured_grid.to_xarray(modeltime=modeltime, netcdf_format=NetCDFFormat.STRUCTURED) context = { @@ -306,65 +314,10 @@ def test_structured_data_var_no_redundant_coordinates(structured_grid, modeltime param = NetCDFParam.from_dict({"name": "k", "attrs": {}}, context=context) param._context["grid"] = structured_grid merged = xr.merge([grid_ds, param.to_xarray()]) - path = tmp_path / "structured_roundtrip.nc" - merged.to_netcdf(path) - attrs = _raw_var_attrs(path, "npf_k") - coords_attr = attrs.get("coordinates", "") - assert coords_attr == "" or all( - c not in coords_attr for c in ["y", "x"] - ), f"structured data var should not carry redundant coordinates attr: {coords_attr!r}" - - -def test_mesh_data_var_coordinates_survives_roundtrip(structured_grid, modeltime, tmp_path): - """coordinates = 'mesh_face_x mesh_face_y' on mesh face vars must survive round-trip. - - Verified via raw netCDF4 because xarray strips the coordinates attr when reopening. - """ - grid_ds = structured_grid.to_xarray( - modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH + extra_coords = set(merged["npf_k"].coords) - set(merged["npf_k"].dims) + assert not any(c in extra_coords for c in ["y", "x"]), ( + f"structured data var carries redundant coordinates: {extra_coords}" ) - context = { - "mesh": "layered", - "modelname": "gwfmodel", - "gridtype": "structured", - "package_name": "npf", - "package_type": "gwf-npf", - "dims": [2, 2, 3, 3], - } - param = NetCDFParam.from_dict({"name": "k", "attrs": {"layer": 1}}, context=context) - param._context["grid"] = structured_grid - merged = xr.merge([grid_ds, param.to_xarray()]) - path = tmp_path / "mesh_roundtrip.nc" - merged.to_netcdf(path) - attrs = _raw_var_attrs(path, "npf_k_l1") - coords_attr = attrs.get("coordinates", "") - assert ( - "mesh_face_x" in coords_attr and "mesh_face_y" in coords_attr - ), f"coordinates attr missing or wrong on raw disk: {coords_attr!r}" - - -def test_mesh_data_var_cf_attrs_survive_roundtrip(structured_grid, modeltime, tmp_path): - """mesh, location, and grid_mapping attrs on face data vars must survive round-trip.""" - grid_ds = structured_grid.to_xarray( - modeltime=modeltime, netcdf_format=NetCDFFormat.LAYERED_MESH - ) - context = { - "mesh": "layered", - "modelname": "gwfmodel", - "gridtype": "structured", - "package_name": "npf", - "package_type": "gwf-npf", - "dims": [2, 2, 3, 3], - } - param = NetCDFParam.from_dict({"name": "k", "attrs": {"layer": 1}}, context=context) - param._context["grid"] = structured_grid - merged = xr.merge([grid_ds, param.to_xarray()]) - path = tmp_path / "mesh_cf_attrs_roundtrip.nc" - merged.to_netcdf(path) - attrs = _raw_var_attrs(path, "npf_k_l1") - assert attrs.get("mesh") == "mesh", f"mesh attr missing after round-trip: {attrs}" - assert attrs.get("location") == "face", f"location attr missing after round-trip: {attrs}" - assert attrs.get("grid_mapping") == "projection", f"grid_mapping missing: {attrs}" def test_structured_dis_gdal_geotransform(structured_grid, modeltime, tmp_path):