Skip to content

Commit ddc1cbe

Browse files
authored
add to_xarray method to component and context (#257)
return xr.Dataset for Component and xr.DataTree for Context. I figure since pandas does to_xarray and you get either xr.Dataset or xr.DataArray depending if the object called on was a series or a dataframe, we might as well follow. but no strong feelings
1 parent 606a219 commit ddc1cbe

File tree

3 files changed

+28
-0
lines changed

3 files changed

+28
-0
lines changed

flopy4/mf6/component.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,3 +195,6 @@ def to_dict(self, blocks: bool = False, strict: bool = False) -> dict[str, Any]:
195195
for field_name in spec.keys()
196196
if spec[field_name].block or not strict
197197
}
198+
199+
def to_xarray(self):
200+
return self.data.dataset # type: ignore

flopy4/mf6/context.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,6 @@ def load(self, format=MF6):
3434
def write(self, format=MF6):
3535
with cd(self.workspace):
3636
super().write(format=format)
37+
38+
def to_xarray(self):
39+
return self.data # type: ignore

test/test_mf6_component.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import numpy as np
66
import pandas as pd
77
import pytest
8+
import xarray as xr
89
from flopy.discretization import StructuredGrid
910
from xarray import DataTree
1011

@@ -555,3 +556,24 @@ def test_tdis_from_timestamps():
555556
np.testing.assert_array_equal(tdis.perlen, [4.0, 10.0])
556557
np.testing.assert_array_equal(tdis.nstp, [5, 5])
557558
np.testing.assert_array_equal(tdis.tsmult, [1.2, 1.2])
559+
560+
561+
def test_to_xarray_on_component():
562+
tdis = Tdis.from_timestamps(["2020-01-01", "2020-01-05", "2020-01-15"], nstp=5, tsmult=1.2)
563+
ds = tdis.to_xarray()
564+
assert isinstance(ds, xr.Dataset)
565+
assert isinstance(ds.per, xr.DataArray)
566+
assert np.array_equal(ds.per, [0, 1])
567+
assert ds.attrs["start_date_time"] == pd.Timestamp("2020-01-01")
568+
569+
570+
def test_to_xarray_on_context(function_tmpdir):
571+
time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0])
572+
ims = Ims(models=["gwf"])
573+
sim = Simulation(tdis=time, solutions={"ims": ims}, workspace=function_tmpdir)
574+
dt = sim.to_xarray()
575+
assert isinstance(dt, xr.DataTree)
576+
assert isinstance(dt.per, xr.DataArray)
577+
assert np.array_equal(dt.per, [0])
578+
assert dt.attrs["filename"] == "mfsim.nam"
579+
assert dt.attrs["workspace"] == Path(function_tmpdir)

0 commit comments

Comments
 (0)