add to_xarray method to component and context

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

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

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)