diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index 13c3cebc..d74c2224 100644 --- a/flopy4/mf6/component.py +++ b/flopy4/mf6/component.py @@ -195,3 +195,6 @@ def to_dict(self, blocks: bool = False, strict: bool = False) -> dict[str, Any]: for field_name in spec.keys() if spec[field_name].block or not strict } + + def to_xarray(self): + return self.data.dataset # type: ignore diff --git a/flopy4/mf6/context.py b/flopy4/mf6/context.py index f79b9c6e..f3d38228 100644 --- a/flopy4/mf6/context.py +++ b/flopy4/mf6/context.py @@ -34,3 +34,6 @@ def load(self, format=MF6): def write(self, format=MF6): with cd(self.workspace): super().write(format=format) + + def to_xarray(self): + return self.data # type: ignore diff --git a/test/test_mf6_component.py b/test/test_mf6_component.py index 0465d52c..e1a753b6 100644 --- a/test/test_mf6_component.py +++ b/test/test_mf6_component.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd import pytest +import xarray as xr from flopy.discretization import StructuredGrid from xarray import DataTree @@ -555,3 +556,24 @@ def test_tdis_from_timestamps(): np.testing.assert_array_equal(tdis.perlen, [4.0, 10.0]) np.testing.assert_array_equal(tdis.nstp, [5, 5]) np.testing.assert_array_equal(tdis.tsmult, [1.2, 1.2]) + + +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 ds.attrs["start_date_time"] == pd.Timestamp("2020-01-01") + + +def test_to_xarray_on_context(function_tmpdir): + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) + ims = Ims(models=["gwf"]) + 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 dt.attrs["filename"] == "mfsim.nam" + assert dt.attrs["workspace"] == Path(function_tmpdir)