Improve getting time coordinates from groundwater model objects

nlmod/dims/time.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import pandas as pd
+from xarray import IndexVariable
 
 logger = logging.getLogger(__name__)
 
@@ -145,7 +146,7 @@ def estimate_nstp(
     ----------
     forcing : array-like
         Array with a forcing value for each stress period. Forcing can be
-        for example a pumping rate of a rainfall intensity.
+        for example a pumping rate or a rainfall intensity.
     perlen : float or array of floats (nper)
         An array of the stress period lengths.
     tsmult : float or array of floats (nper)
@@ -216,3 +217,36 @@ def estimate_nstp(
 
     else:
         return nstp_ceiled
+
+
+def ds_time_from_model(gwf):
+    """Get time index variable from model (gwf or gwt).
+
+    Parameters
+    ----------
+    gwf : flopy MFModel object
+        groundwater flow or groundwater transport model
+
+    Returns
+    -------
+    IndexVariable
+        time coordinate for xarray data-array or dataset
+    """
+
+    start_datetime = gwf.simulation.get_package("TDIS").start_date_time.data
+
+    if start_datetime is not None:
+        dt = pd.to_timedelta(
+            np.cumsum(gwf.dimensions.simulation_time.get_perioddata()["perlen"]), "D"
+        )
+        times = pd.Timestamp(start_datetime) + dt
+
+    else:
+        times = np.cumsum(gwf.dimensions.simulation_time.get_perioddata()["perlen"])
+
+    time = IndexVariable(["time"], times)
+    time.attrs["time_units"] = gwf.dimensions.simulation_time.get_time_units()
+    if start_datetime is not None:
+        time.attrs["start"] = str(start_datetime)
+
+    return time

nlmod/gwf/output.py

@@ -10,14 +10,15 @@
 
 from ..dims.grid import modelgrid_from_ds
 from ..dims.resample import get_affine, get_xy_mid_structured
+from ..dims.time import ds_time_from_model
 
 logger = logging.getLogger(__name__)
 
 
 def get_heads_da(ds=None, gwf=None, fname_hds=None):
     """Reads heads file given either a dataset or a groundwater flow object.
 
-    Note: Calling this function with ds is currently prevered over calling it
+    Note: Calling this function with ds is currently preferred over calling it
     with gwf, because the layer and time coordinates can not be fully
     reconstructed from gwf.
 
@@ -86,11 +87,12 @@ def get_heads_da(ds=None, gwf=None, fname_hds=None):
     else:
         assert 0, "Gridtype not supported"
 
-    if ds is not None:
+    # set layer and time coordinates
+    if gwf is not None:
+        head_ar.coords["layer"] = np.arange(gwf.modelgrid.nlay)
+        head_ar.coords["time"] = ds_time_from_model(gwf)
+    else:
         head_ar.coords["layer"] = ds.layer
-
-        # TODO: temporarily only add time for when ds is passed because unable to
-        # exactly recreate ds.time from gwf.
         head_ar.coords["time"] = ds.time
 
     if ds is not None and "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
@@ -114,6 +116,96 @@ def get_heads_da(ds=None, gwf=None, fname_hds=None):
     return head_ar
 
 
+def get_budget_da(text, ds=None, gwf=None, fname_cbc=None, kstpkper=None):
+    """Reads budget file given either a dataset or a groundwater flow object.
+
+    Parameters
+    ----------
+    text : str
+        record to get from budget file
+    ds : xarray.Dataset, optional
+        xarray dataset with model data. One of ds or gwf must be provided.
+    gwf : flopy ModflowGwf, optional
+        Flopy groundwaterflow object. One of ds or gwf must be provided.
+    fname_cbc : path, optional
+        specify the budget file to load, if not provided budget file will
+        be obtained from ds or gwf.
+
+    Returns
+    -------
+    q_ar : xarray.DataArray
+        budget data array.
+    """
+    cbfobj = _get_cbf(ds=ds, gwf=gwf, fname_cbc=fname_cbc)
+
+    q = cbfobj.get_data(text=text, kstpkper=kstpkper, full3D=True)
+    q = np.stack(q)
+
+    if gwf is not None:
+        gridtype = gwf.modelgrid.grid_type
+    else:
+        gridtype = ds.gridtype
+
+    if gridtype == "vertex":
+        q_ar = xr.DataArray(
+            data=q,
+            dims=("time", "layer", "icell2d"),
+            coords={},
+            attrs={"units": "m3/d"},
+        )
+
+    elif gridtype == "structured":
+        if gwf is not None:
+            delr = np.unique(gwf.modelgrid.delr).item()
+            delc = np.unique(gwf.modelgrid.delc).item()
+            extent = gwf.modelgrid.extent
+            x, y = get_xy_mid_structured(extent, delr, delc)
+
+        else:
+            x = ds.x
+            y = ds.y
+
+        q_ar = xr.DataArray(
+            data=q,
+            dims=("time", "layer", "y", "x"),
+            coords={
+                "x": x,
+                "y": y,
+            },
+            attrs={"units": "m3/d"},
+        )
+    else:
+        assert 0, "Gridtype not supported"
+
+    # set layer and time coordinates
+    if gwf is not None:
+        q_ar.coords["layer"] = np.arange(gwf.modelgrid.nlay)
+        q_ar.coords["time"] = ds_time_from_model(gwf)
+    else:
+        q_ar.coords["layer"] = ds.layer
+        q_ar.coords["time"] = ds.time
+
+    if ds is not None and "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
+        affine = get_affine(ds)
+        q_ar.rio.write_transform(affine, inplace=True)
+
+    elif gwf is not None and gwf.modelgrid.angrot != 0.0:
+        attrs = dict(
+            delr=np.unique(gwf.modelgrid.delr).item(),
+            delc=np.unique(gwf.modelgrid.delc).item(),
+            xorigin=gwf.modelgrid.xoffset,
+            yorigin=gwf.modelgrid.yoffset,
+            angrot=gwf.modelgrid.angrot,
+            extent=gwf.modelgrid.extent,
+        )
+        affine = get_affine(attrs)
+        q_ar.rio.write_transform(affine, inplace=True)
+
+    q_ar.rio.write_crs("EPSG:28992", inplace=True)
+
+    return q_ar
+
+
 def _get_hds(ds=None, gwf=None, fname_hds=None):
     msg = "Load the heads using either the ds, gwf or fname_hds"
     assert ((ds is not None) + (gwf is not None) + (fname_hds is not None)) >= 1, msg
@@ -129,7 +221,7 @@ def _get_hds(ds=None, gwf=None, fname_hds=None):
     return headobj
 
 
-def _get_cbc(ds=None, gwf=None, fname_cbc=None):
+def _get_cbf(ds=None, gwf=None, fname_cbc=None):
     msg = "Load the budgets using either the ds or the gwf"
     assert ((ds is not None) + (gwf is not None)) == 1, msg
 

nlmod/gwt/output.py

@@ -3,11 +3,11 @@
 
 import flopy
 import numpy as np
-import pandas as pd
 import xarray as xr
 
 from ..dims.layers import calculate_thickness
 from ..dims.resample import get_affine, get_xy_mid_structured
+from ..dims.time import ds_time_from_model
 
 logger = logging.getLogger(__name__)
 
@@ -97,20 +97,13 @@ def get_concentration_da(ds=None, gwt=None, fname_conc=None):
     else:
         assert 0, "Gridtype not supported"
 
-    if ds is not None:
+    # set layer and time coordinates
+    if gwt is not None:
+        conc_ar.coords["layer"] = np.arange(gwt.modelgrid.nlay)
+        conc_ar.coords["time"] = ds_time_from_model(gwt)
+    else:
         conc_ar.coords["layer"] = ds.layer
-
-        # TODO: temporarily only add time for when ds is passed because unable to
-        # exactly recreate ds.time from gwt.
-        times = np.array(
-            [
-                pd.Timestamp(ds.time.start)
-                + pd.Timedelta(t, unit=ds.time.time_units[0])
-                for t in concobj.get_times()
-            ],
-            dtype=np.datetime64,
-        )
-        conc_ar.coords["time"] = times
+        conc_ar.coords["time"] = ds.time
 
     if ds is not None and "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
         affine = get_affine(ds)

nlmod/sim/sim.py

@@ -182,7 +182,7 @@ def tdis(ds, sim, pname="tdis"):
         pname=pname,
         time_units=ds.time.time_units,
         nper=len(ds.time),
-        # start_date_time=ds.time.start, # disable until fix in modpath
+        start_date_time=pd.Timestamp(ds.time.start).isoformat(),
         perioddata=tdis_perioddata,
     )