diff --git a/autotest/test_export.py b/autotest/test_export.py index 1812418ef..f3f6eb4ff 100644 --- a/autotest/test_export.py +++ b/autotest/test_export.py @@ -977,6 +977,77 @@ def test_export_huge_shapefile(function_tmpdir): gdf.to_file(function_tmpdir / "huge.shp") +@requires_pkg("geopandas") +def test_to_geodataframe_incomplete_stress_period_data(): + sim = flopy.mf6.MFSimulation() + tdis = flopy.mf6.ModflowTdis( + sim, + nper=2, + perioddata=[(1, 1, 1), (1, 1, 1)], + ) + ims = flopy.mf6.ModflowIms(sim) + + gwf = flopy.mf6.ModflowGwf(sim, modelname="dev_gdf") + + dis = flopy.mf6.ModflowGwfdis( + gwf, nlay=1, nrow=10, ncol=11, delc=100, delr=100, top=100, botm=0, idomain=1 + ) + + npf = flopy.mf6.ModflowGwfnpf( + gwf, + k=10, + ) + + ic = flopy.mf6.ModflowGwfic(gwf, strt=99) + + chd_rec = [(0, i, 0, 95) for i in range(10)] + chd = flopy.mf6.ModflowGwfchd(gwf, stress_period_data={0: chd_rec}) + + ghb_rec = [(0, i, 10, 85.0, 10.0) for i in range(10)] + ghb = flopy.mf6.ModflowGwfghb(gwf, stress_period_data={0: ghb_rec, 1: ghb_rec}) + + rch_rec = np.full((10, 11), 0.0005) + rch_rec[:, 0] = 0 + rch_rec[:, -1] = 0 + rch = flopy.mf6.ModflowGwfrcha(gwf, recharge={0: rch_rec}) + + wel_rec = [ + (0, 4, 5, -1500.0), + ] + wel = flopy.mf6.ModflowGwfwel(gwf, stress_period_data={0: wel_rec}) + + recharge = rch.recharge.array[0].ravel() + gdf = rch.to_geodataframe(kper=1) + + np.testing.assert_allclose( + recharge, + gdf["rcha_recharge_1"].values, + err_msg="GeoDataFrame does not match recharge values from package", + ) + + wel_data = wel.stress_period_data.to_array(kper=0, mask=True)["q"].ravel() + gdf = wel.to_geodataframe(kper=1) + np.testing.assert_allclose( + wel_data, + gdf["wel_q_0_1"].values, + err_msg="GeoDataFrame does not match pumping values from wel package", + ) + + gdf = gwf.to_geodataframe(kper=1) + + np.testing.assert_allclose( + recharge, + gdf["rcha_recharge_1"].values, + err_msg="GeoDataFrame from gwf does not match recharge values from package", + ) + + np.testing.assert_allclose( + wel_data, + gdf["wel_q_0_1"].values, + err_msg="GeoDataFrame from gwf does not match pumping values from wel package", + ) + + @requires_pkg("netCDF4", "pyproj") def test_polygon_from_ij(function_tmpdir): """test creation of a polygon from an i, j location using get_vertices().""" diff --git a/flopy/mf6/data/mfdataarray.py b/flopy/mf6/data/mfdataarray.py index 0f192263e..a7b596fe1 100644 --- a/flopy/mf6/data/mfdataarray.py +++ b/flopy/mf6/data/mfdataarray.py @@ -823,6 +823,7 @@ def _get_data(self, layer=None, apply_mult=False, **kwargs): and kwargs["array"] and isinstance(self, MFTransientArray) and data is not [] # noqa: F632 + and data is not None ): data = np.expand_dims(data, 0) return data @@ -1889,7 +1890,10 @@ def _get_array(self, num_sp, apply_mult, **kwargs): if sp in self._data_storage: self.get_data_prep(sp) data = super().get_data(apply_mult=apply_mult, **kwargs) - data = np.expand_dims(data, 0) + if data is not None: + data = np.expand_dims(data, 0) + else: + data = output else: # if there is no previous data provide array of # zeros, otherwise provide last array of data found @@ -2057,6 +2061,13 @@ def to_geodataframe(self, gdf=None, kper=0, full_grid=True, shorten_attr=False, name = f"{self.path[1]}_{self.name}" data = self.get_data(key=kper, apply_mult=True) + if data is None: + per_with_data = np.array([i for i, v in self.empty_keys.items() if not v]) + per_with_data = per_with_data[per_with_data < kper] + if len(per_with_data) == 0: + return gdf + data = self.get_data(key=per_with_data[-1], apply_mult=True) + if data.size == ncpl: name = f"{name}_{kper}" gdf[name] = data.ravel() diff --git a/flopy/mf6/data/mfdatalist.py b/flopy/mf6/data/mfdatalist.py index 86ce4251e..c821b34f9 100644 --- a/flopy/mf6/data/mfdatalist.py +++ b/flopy/mf6/data/mfdatalist.py @@ -1697,9 +1697,17 @@ def to_geodataframe(self, gdf=None, kper=0, full_grid=True, shorten_attr=False, if gdf is None: gdf = modelgrid.to_geodataframe() + if self.data is None: + return gdf + data = self.to_array(kper=kper, mask=True) if data is None: - return gdf + # get data from the last stress period where data was specified + per_with_data = np.array([i for i, v in self.data.items() if v is not None]) + per_with_data = per_with_data[per_with_data < kper] + if len(per_with_data) == 0: + return gdf + data = self.to_array(kper=per_with_data[-1], mask=True) col_names = [] for name, array3d in data.items(): diff --git a/flopy/mf6/data/mfdataplist.py b/flopy/mf6/data/mfdataplist.py index 669b9af66..3b7472fde 100644 --- a/flopy/mf6/data/mfdataplist.py +++ b/flopy/mf6/data/mfdataplist.py @@ -1968,7 +1968,17 @@ def to_geodataframe( if gdf is None: gdf = modelgrid.to_geodataframe() + if self.data is None: + return gdf + data = self.to_array(kper=kper, mask=True) + if data is None: + # get data from the last stress period where data was specified + per_with_data = np.array([i for i, v in self.data.items() if v is not None]) + per_with_data = per_with_data[per_with_data < kper] + if len(per_with_data) == 0: + return gdf + data = self.to_array(kper=per_with_data[-1], mask=True) col_names = [] for name, array3d in data.items(): diff --git a/flopy/utils/rasters.py b/flopy/utils/rasters.py index caae365f1..4c31c484e 100644 --- a/flopy/utils/rasters.py +++ b/flopy/utils/rasters.py @@ -1,5 +1,6 @@ import warnings from os import PathLike +from pathlib import Path from typing import Union import numpy as np @@ -828,21 +829,21 @@ def get_array(self, band, masked=True): return array - def write(self, name): + def write(self, name: Union[str, PathLike]): """ Method to write raster data to a .tif file Parameters ---------- - name : str + name : PathLike output raster .tif file name """ rasterio = import_optional_dependency("rasterio") - if not name.endswith(".tif"): - name += ".tif" + if not str(name).endswith(".tif"): + name = Path(f"{name!s}.tif") with rasterio.open(name, "w", **self._meta) as foo: for band, arr in self.__arr_dict.items():