add xarray support

[andyrich]

"Xarray makes working with labelled multi-dimensional arrays in Python simple, efficient, and fun!" website

Xarray is very well suited for spatial, timeseries, data import/export and other analyses used in the flopy framework. The package seems to be well supported.

I imagine the get_xarray could be part of the utils in flopy but it could also be part of the BinaryHeadFile class to access it directly from that object.

Here's some examples:

import xarray as xr
import flopy


def get_xarray(hdobj, ml):
    times = hdobj.get_times()
    head = [hdobj.get_data(totim=t) for t in times]  # shape: (ntimes, nlay, nrow, ncol)
    head = np.array(head)

    xl, yl = ml.modelgrid.xcellcenters, ml.modelgrid.ycellcenters

    datetimes = pd.to_datetime(ml.start_datetime) + pd.to_timedelta(times, unit='D')
    # build xarray Dataset
    ds = xr.Dataset(
        data_vars=dict(
            head=(["time", "layer", "y", "x"], head)
        ),
        coords=dict(
            time=("time", datetimes),
            layer=("layer", list(range(ml.nlay))),
            x=(["y", "x"], xl),
            y=(["y", "x"], yl),
        ),
    )

    return ds

create a Dataset from your model and headfileobject:

heads = get_xarray(hdobj, ml)

heads

plotting is easy and (somewhat) intuitive

plot heads for layer a date

h= heads['head'].sel(time='10/1/2024', layer = 0)
h.plot.contourf()

plot average heads for a period and layer

# subset to date range (assuming `time` is datetime64 or convertible)
ds_sub = heads.sel(time=slice("2000-01-01", "2000-12-31"), layer = 0)

# average across time
avg_lay1 = ds_sub.mean(dim="time")

# plot average head values for a time period for layer 1
avg_lay1['head'].plot(vmin = -50, vmax = 2000)
plt.title("Average Head, Layer 1, 2000")
plt.show()

plot a timeseries of average head values of layer 1

ts = heads['head'].sel(layer=0).mean(dim=("y", "x"))

ts.plot()
plt.title("Average head in layer 1")
plt.ylabel("Head (m)")
plt.show()

plot a grid of layer heads

g_simple = heads.isel(time = 10)['head'].plot(x="x", y="y", col="layer", col_wrap=3,vmin = -100,vmax = 1000)

[wpbonelli]

Thanks @andyrich! Yeah, xarray is great. We are designing flopy 4.x around it, not only for utils but as the core data model. Maybe in the meantime we can add some dataset-returning utilities to 3.x like the one you provided.