diff --git a/docs/examples/image/quickstart.png b/docs/examples/image/quickstart.png deleted file mode 100644 index d775d78e..00000000 Binary files a/docs/examples/image/quickstart.png and /dev/null differ diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index d001872c..f179e291 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -44,13 +44,13 @@ sim.run(verbose=True) assert chd.data["head"][0, 0] == 1.0 -assert chd.data.head.sel(per=0)[99] == 0.0 +assert chd.data.head.sel(kper=0)[99] == 0.0 assert np.allclose(chd.data.head[:, 1:99], np.full(98, 3e30)) assert gwf.dis.data.botm.sel(lay=0, col=0, row=0) == 0.0 assert oc.data["save_head"][0] == "all" -assert oc.data.save_head.sel(per=0) == "all" +assert oc.data.save_head.sel(kper=0) == "all" # get head and budget results budget = gwf.output.budget.squeeze() @@ -67,4 +67,4 @@ head.plot.imshow(ax=ax) head.plot.contour(ax=ax, levels=[0.2, 0.4, 0.6, 0.8], linewidths=3.0) budget.plot.quiver(x="x", y="y", u="npf-qx", v="npf-qy", ax=ax, color="white") -fig.savefig(workspace / ".." / "image" / "quickstart.png") +fig.savefig(workspace / ".." / "quickstart.png") diff --git a/flopy4/mf6/indexes.py b/flopy4/mf6/indexes.py index 0ef5f3bc..51ca5a06 100644 --- a/flopy4/mf6/indexes.py +++ b/flopy4/mf6/indexes.py @@ -54,4 +54,4 @@ def grid_index(dataset: xr.Dataset) -> MetaIndex: def time_index(dataset: xr.Dataset) -> PandasIndex: - return alias(dataset, "nper", "per") + return alias(dataset, "nper", "kper") diff --git a/flopy4/mf6/package.py b/flopy4/mf6/package.py index 8471c58d..fa73b103 100644 --- a/flopy4/mf6/package.py +++ b/flopy4/mf6/package.py @@ -1,8 +1,11 @@ from abc import ABC +import numpy as np +import pandas as pd from xattree import xattree from flopy4.mf6.component import Component +from flopy4.mf6.constants import FILL_DNODATA @xattree @@ -11,3 +14,183 @@ def default_filename(self) -> str: name = self.parent.name if self.parent else self.name # type: ignore cls_name = self.__class__.__name__.lower() return f"{name}.{cls_name}" + + @property + def stress_period_data(self) -> pd.DataFrame: + """ + Get combined stress period data for all period data fields. + + Returns a DataFrame with columns: 'kper' (stress period), spatial + coordinates, and all period data field values (e.g., 'head', 'elev', 'cond'). + + Spatial coordinates are automatically determined based on grid type: + - Structured grids: 'layer', 'row', 'col' columns + - Unstructured grids: 'node' column + + Returns + ------- + pd.DataFrame + DataFrame with stress period data for all fields. + + Examples + -------- + >>> # Structured grid - uses layer/row/col + >>> chd = Chd(parent=gwf, head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}}) + >>> df = chd.stress_period_data + >>> print(df) + kper layer row col head + 0 0 0 0 0 1.0 + 1 0 0 9 9 0.0 + + >>> # Multi-field package + >>> drn = Drn(parent=gwf, elev={0: {(0, 7, 5): 10.0}}, cond={0: {(0, 7, 5): 1.0}}) + >>> df = drn.stress_period_data + >>> print(df) + kper layer row col elev cond + 0 0 0 7 5 10.0 1.0 + + Notes + ----- + This property is read-only. Setting stress period data via the + initializer or attribute assignment will be supported after PR #266. + + The coordinate format depends on grid information from the parent model: + - If structured grid dimensions (nlay, nrow, ncol) are available from the + parent, the DataFrame will use layer/row/col columns + - Otherwise, it will use node indices + """ + from attrs import fields + + # Find all period block fields + period_fields = [] + for f in fields(self.__class__): # type: ignore + if f.metadata.get("block") == "period" and f.metadata.get("xattree", {}).get("dims"): + period_fields.append(f.name) + + if not period_fields: + raise TypeError("No period block fields found in package") + + # Determine spatial coordinate format based on available grid info + # If parent has structured grid dims, use layer/row/col + # Otherwise use node indices + # TODO generalize this, maybe a `grid_type` property somewhere + # like flopy3 has + has_structured_grid = False + nlay = nrow = ncol = None + + if hasattr(self, "parent") and self.parent is not None: + # Try to get grid dimensions from parent model + if ( + hasattr(self.parent, "nlay") + and hasattr(self.parent, "nrow") + and hasattr(self.parent, "ncol") + ): + nlay = self.parent.nlay + nrow = self.parent.nrow + ncol = self.parent.ncol + has_structured_grid = True + + # Build combined DataFrame + all_records = [] + coord_columns = None + + for field_name in period_fields: + data = getattr(self, field_name) + if data is None: + continue + + # Convert field data to records + for kper in range(data.shape[0]): + per_data = data[kper] + + # Handle sparse arrays + try: + import sparse + + if isinstance(per_data, sparse.COO): + # Convert to dense for processing + per_data = per_data.todense() + except ImportError: + pass + + # Find non-empty cells + # Handle different dtypes for the mask + if np.issubdtype(per_data.dtype, np.str_) or np.issubdtype( + per_data.dtype, np.bytes_ + ): + # For string fields, check for non-empty and non-fill strings + mask = (per_data != "") & (per_data != str(FILL_DNODATA)) + else: + # For numeric fields, use standard FILL_DNODATA + mask = per_data != FILL_DNODATA + + indices = np.where(mask) + + if len(indices) == 0 or indices[0].size == 0: + continue + + values = per_data[mask] + + for i in range(len(values)): + # Extract scalar value from xarray if needed + val = values[i] + if hasattr(val, "item"): + val = val.item() + + if len(indices) == 1: # 1D array (nodes) + node = int(indices[0][i]) + + # Convert to layer/row/col if structured grid info available + if has_structured_grid and nlay and nrow and ncol: + layer = node // (nrow * ncol) + row = (node % (nrow * ncol)) // ncol + col = node % ncol + record = { + "kper": kper, + "layer": int(layer), + "row": int(row), + "col": int(col), + field_name: val, + } + if coord_columns is None: + coord_columns = ["kper", "layer", "row", "col"] + else: + # Use node index if no structured grid info + record = {"kper": kper, "node": node, field_name: val} + if coord_columns is None: + coord_columns = ["kper", "node"] + elif len(indices) == 3: # 3D array (layer, row, col) + layer, row, col = ( + indices[0][i], + indices[1][i], + indices[2][i], + ) + record = { + "kper": kper, + "layer": int(layer), + "row": int(row), + "col": int(col), + field_name: val, + } + if coord_columns is None: + coord_columns = ["kper", "layer", "row", "col"] + else: + continue + all_records.append(record) + + if not all_records: + # Return empty DataFrame with appropriate columns + cols = coord_columns or ["kper", "layer", "row", "col"] + cols.extend(period_fields) + return pd.DataFrame(columns=cols) + + # Create DataFrame from records + df = pd.DataFrame(all_records) + + # For multi-field packages, merge fields with same coordinates + # Single-field packages can skip the groupby for better performance + if len(period_fields) > 1 and coord_columns: + # Fill NaN for fields that don't have data at certain coordinates + df = df.groupby(coord_columns, as_index=False).first() + + return df diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index 5e362ca4..33d12de8 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -20,7 +20,7 @@ class PeriodData: nstp: int tsmult: float - nper: int = dim(block="dimensions", coord="per", default=1, scope=ROOT) + nper: int = dim(block="dimensions", coord="kper", default=1, scope=ROOT) time_units: Optional[str] = field(block="options", default=None) start_date_time: Optional[datetime] = field(block="options", default=None) perlen: NDArray[np.float64] = array( diff --git a/test/test_dataframe_api.py b/test/test_dataframe_api.py new file mode 100644 index 00000000..51853e28 --- /dev/null +++ b/test/test_dataframe_api.py @@ -0,0 +1,252 @@ +"""Tests for stress_period_data property.""" + +import pandas as pd + +from flopy4.mf6.gwf import Chd, Drn, Gwf, Wel + + +def test_chd_stress_period_data(): + """Test stress_period_data property for CHD package.""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + chd = Chd(dims=dims, head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}}) + df = chd.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + assert "kper" in df.columns + assert "node" in df.columns + assert "head" in df.columns + + # Check first record (node 0 = cell (0,0,0)) + assert df.iloc[0]["kper"] == 0 + assert df.iloc[0]["node"] == 0 + assert df.iloc[0]["head"] == 1.0 + + # Check second record (node 99 = cell (0,9,9)) + assert df.iloc[1]["kper"] == 0 + assert df.iloc[1]["node"] == 99 + assert df.iloc[1]["head"] == 0.0 + + +def test_wel_stress_period_data(): + """Test stress_period_data property for WEL package.""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + wel = Wel( + dims=dims, + q={0: {(0, 5, 5): -100.0, (0, 8, 8): 50.0}}, + ) + df = wel.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + assert "kper" in df.columns + assert "q" in df.columns + + # Check records + assert df.iloc[0]["q"] == -100.0 + assert df.iloc[1]["q"] == 50.0 + + +def test_drn_stress_period_data_multifield(): + """Test stress_period_data property for DRN package (multi-field).""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + drn = Drn( + dims=dims, + elev={0: {(0, 7, 5): 10.0}}, + cond={0: {(0, 7, 5): 1.0}}, + ) + + df = drn.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 1 + + # Should have both elev and cond columns + assert "node" in df.columns + assert "elev" in df.columns + assert "cond" in df.columns + assert df.iloc[0]["elev"] == 10.0 + assert df.iloc[0]["cond"] == 1.0 + + +def test_multi_period_stress_period_data(): + """Test stress_period_data property with multiple stress periods.""" + dims = {"nper": 3, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + chd = Chd( + dims=dims, + head={ + 0: {(0, 0, 0): 1.0}, + 1: {(0, 0, 0): 0.9}, + 2: {(0, 0, 0): 0.8}, + }, + ) + + df = chd.stress_period_data + + assert len(df) == 3 + assert df[df["kper"] == 0].iloc[0]["head"] == 1.0 + assert df[df["kper"] == 1].iloc[0]["head"] == 0.9 + assert df[df["kper"] == 2].iloc[0]["head"] == 0.8 + + +def test_stress_period_data_multiple_cells(): + """Test stress_period_data property with multiple cells per period.""" + dims = {"nper": 2, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + chd = Chd( + dims=dims, + head={ + 0: {(0, 0, 0): 1.0, (0, 5, 5): 0.5, (0, 9, 9): 0.0}, + 1: {(0, 0, 0): 0.9, (0, 5, 5): 0.45, (0, 9, 9): 0.0}, + }, + ) + + df = chd.stress_period_data + + # Should have 6 records (3 cells x 2 periods) + assert len(df) == 6 + + # Verify period 0 + per0 = df[df["kper"] == 0] + assert len(per0) == 3 + assert set(per0["head"].values) == {1.0, 0.5, 0.0} + + # Verify period 1 + per1 = df[df["kper"] == 1] + assert len(per1) == 3 + assert set(per1["head"].values) == {0.9, 0.45, 0.0} + + +def test_empty_stress_period_data(): + """Test stress_period_data property with no data.""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + chd = Chd(dims=dims) # No head data + df = chd.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 0 + # Should have coordinate and field columns even if empty + assert "kper" in df.columns + assert "head" in df.columns + + +def test_stress_period_data_different_cells_per_field(): + """Test stress_period_data when different fields have data at different cells.""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + # Node 75 = (0, 7, 5), Node 33 = (0, 3, 3) + drn = Drn( + dims=dims, + elev={0: {(0, 7, 5): 10.0, (0, 3, 3): 12.0}}, # 2 cells + cond={0: {(0, 7, 5): 1.0}}, # 1 cell (overlapping) + ) + + df = drn.stress_period_data + + # Should have 2 rows (one for each unique cell location) + assert len(df) == 2 + + # Node 75 (cell 0,7,5) should have both elev and cond + row1 = df[df["node"] == 75] + assert len(row1) == 1 + assert row1.iloc[0]["elev"] == 10.0 + assert row1.iloc[0]["cond"] == 1.0 + + # Node 33 (cell 0,3,3) should have elev but NaN for cond + row2 = df[df["node"] == 33] + assert len(row2) == 1 + assert row2.iloc[0]["elev"] == 12.0 + assert pd.isna(row2.iloc[0]["cond"]) + + +def test_stress_period_data_with_aux_and_boundname(): + """Test stress_period_data includes aux and boundname fields.""" + dims = {"nper": 1, "nlay": 1, "nrow": 10, "ncol": 10, "nodes": 100} + + chd = Chd( + dims=dims, + head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}}, + aux={0: {(0, 0, 0): 100.0, (0, 9, 9): 200.0}}, + boundname={0: {(0, 0, 0): "INLET", (0, 9, 9): "OUTLET"}}, + ) + + df = chd.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + + # Check that all fields are present + assert "kper" in df.columns + assert "node" in df.columns + assert "head" in df.columns + assert "aux" in df.columns + assert "boundname" in df.columns + + # Check first record + assert df.iloc[0]["kper"] == 0 + assert df.iloc[0]["node"] == 0 + assert df.iloc[0]["head"] == 1.0 + assert df.iloc[0]["aux"] == 100.0 + assert df.iloc[0]["boundname"] == "INLET" + + # Check second record + assert df.iloc[1]["kper"] == 0 + assert df.iloc[1]["node"] == 99 + assert df.iloc[1]["head"] == 0.0 + assert df.iloc[1]["aux"] == 200.0 + assert df.iloc[1]["boundname"] == "OUTLET" + + +def test_stress_period_data_with_structured_grid_parent(): + """Test stress_period_data uses layer/row/col when parent has grid info.""" + from flopy4.mf6 import Simulation + from flopy4.mf6.gwf import Dis + + # Create a real model with DIS (structured grid) package + sim = Simulation() + + # Create DIS package with structured grid dimensions + dis = Dis( + nlay=1, + nrow=10, + ncol=10, + ) + + # Create Gwf model with the DIS package + gwf = Gwf(parent=sim, dis=dis) + + chd = Chd( + parent=gwf, + head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}}, + ) + + df = chd.stress_period_data + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + + # Should have layer/row/col columns, not node + assert "kper" in df.columns + assert "layer" in df.columns + assert "row" in df.columns + assert "col" in df.columns + assert "node" not in df.columns + + # Check first record (node 0 = layer 0, row 0, col 0) + assert df.iloc[0]["kper"] == 0 + assert df.iloc[0]["layer"] == 0 + assert df.iloc[0]["row"] == 0 + assert df.iloc[0]["col"] == 0 + assert df.iloc[0]["head"] == 1.0 + + # Check second record (node 99 = layer 0, row 9, col 9) + assert df.iloc[1]["kper"] == 0 + assert df.iloc[1]["layer"] == 0 + assert df.iloc[1]["row"] == 9 + assert df.iloc[1]["col"] == 9 + assert df.iloc[1]["head"] == 0.0 diff --git a/test/test_mf6_component.py b/test/test_mf6_component.py index e4417339..ca06e72a 100644 --- a/test/test_mf6_component.py +++ b/test/test_mf6_component.py @@ -612,8 +612,8 @@ def test_to_xarray_on_component(): tdis = Tdis.from_timestamps(["2020-01-01", "2020-01-05", "2020-01-15"], nstp=5, tsmult=1.2) ds = tdis.to_xarray() assert isinstance(ds, xr.Dataset) - assert isinstance(ds.per, xr.DataArray) - assert np.array_equal(ds.per, [0, 1]) + assert isinstance(ds.kper, xr.DataArray) + assert np.array_equal(ds.kper, [0, 1]) assert ds.attrs["start_date_time"] == pd.Timestamp("2020-01-01") @@ -623,7 +623,7 @@ def test_to_xarray_on_context(function_tmpdir): sim = Simulation(tdis=time, solutions={"ims": ims}, workspace=function_tmpdir) dt = sim.to_xarray() assert isinstance(dt, xr.DataTree) - assert isinstance(dt.per, xr.DataArray) - assert np.array_equal(dt.per, [0]) + assert isinstance(dt.kper, xr.DataArray) + assert np.array_equal(dt.kper, [0]) assert dt.attrs["filename"] == "mfsim.nam" assert dt.attrs["workspace"] == Path(function_tmpdir)