Fix AttributeError in facebudget when front is omitted (fixes #1097)

docs/api/changelog.rst

@@ -74,6 +74,10 @@ Fixed
   settings were specified as dask array, which could happen when loading a model
   lazily with :meth:`imod.mf6.Modflow6Simulation.from_file` and not
   computing the data before writing.
+- Fixed bug where :func:`imod.evaluate.facebudget` raised an error when the
+  ``front`` budget was left out, even though you only need to provide one of
+  ``front``, ``lower`` or ``right``. Leaving out ``front`` now works as
+  described in the documentation.
 
 Changed
 ~~~~~~~

imod/evaluate/budget.py

@@ -269,6 +269,11 @@ def facebudget(budgetzone, front=None, lower=None, right=None, netflow=True):
             if da_shape != shape:
                 raise ValueError(f"Shape of {daname} does not match budgetzone")
 
+    # At least one of front, lower, right is guaranteed to be provided (checked
+    # above), but any of them may be omitted. Use the first one that is present
+    # as the reference for the time axis and the array shape.
+    reference = next(da for da in (front, lower, right) if da is not None)
+
     # Determine control surface
     face = _outer_edge(budgetzone)
     # Convert indices array to dask array, otherwise garbage collection gets
@@ -289,7 +294,7 @@ def facebudget(budgetzone, front=None, lower=None, right=None, netflow=True):
         results_right = []
 
         if "time" in dims:
-            for itime in range(front.coords["time"].size):
+            for itime in range(reference.coords["time"].size):
                 if front is not None:
                     F = front.isel(time=itime)
                 if lower is not None:
@@ -317,9 +322,9 @@ def facebudget(budgetzone, front=None, lower=None, right=None, netflow=True):
             dask_front, dask_lower, dask_right = delayed_collect(indices, F, L, R)
     else:
         chunks = (1, *da_shape)
-        dask_front = dask.array.full(front.shape, np.nan, chunks=chunks)
-        dask_lower = dask.array.full(lower.shape, np.nan, chunks=chunks)
-        dask_right = dask.array.full(right.shape, np.nan, chunks=chunks)
+        dask_front = dask.array.full(reference.shape, np.nan, chunks=chunks)
+        dask_lower = dask.array.full(reference.shape, np.nan, chunks=chunks)
+        dask_right = dask.array.full(reference.shape, np.nan, chunks=chunks)
 
     if netflow:
         return xr.DataArray(dask_front + dask_lower + dask_right, coords, dims)

imod/tests/test_evaluate/test_budget.py

@@ -216,3 +216,62 @@ def test_flowlower_up_big_budgetzonenr():
     assert rf.sum() == 0.0
     assert rl.sum() == 2.0
     assert rr.sum() == 0.0
+
+
+def test_omit_front():
+    # facebudget only requires one of front/lower/right; omitting `front`
+    # previously raised AttributeError because the time loop and the
+    # no-zones branch dereferenced `front` unconditionally.
+    data = np.zeros((3, 2, 1))
+    coords = {"layer": [1, 2, 3], "y": [0.5, 1.5], "x": [0.5]}
+    dims = ("layer", "y", "x")
+
+    da = xr.DataArray(data, coords, dims)
+    lower = xr.full_like(da, 0.0)
+    budgetzone = xr.full_like(da, np.nan)
+
+    lower[:2] = 1.0
+    budgetzone[:2] = 1
+    # front and right omitted (default None)
+    assert imod.evaluate.facebudget(budgetzone, lower=lower).sum() == 2.0
+    rf, rl, rr = imod.evaluate.facebudget(budgetzone, lower=lower, netflow=False)
+    assert rf.sum() == 0.0
+    assert rl.sum() == 2.0
+    assert rr.sum() == 0.0
+
+
+def test_omit_front_transient():
+    # Same as above, but with a time dimension (exercises the time loop, which
+    # took its length from front.coords["time"]) and the no-zones branch.
+    data = np.zeros((4, 3, 2, 1))
+    coords = {
+        "time": pd.date_range("2000-01-01", "2000-01-04"),
+        "layer": [1, 2, 3],
+        "y": [0.5, 1.5],
+        "x": [0.5],
+    }
+    dims = ("time", "layer", "y", "x")
+
+    lower = xr.DataArray(data, coords, dims)
+    # A partial zone creates a control surface, so indices.size > 0 and the
+    # time loop runs -- this loop took its length from front.coords["time"].
+    budgetzone = xr.DataArray(
+        data=np.array([[[1]], [[1]], [[np.nan]]]),
+        coords={"layer": [1, 2, 3], "y": [0.5], "x": [0.5]},
+        dims=("layer", "y", "x"),
+    )
+    lower = lower.isel(y=[0])
+    netflow = imod.evaluate.facebudget(budgetzone, lower=lower)
+    assert netflow.shape == (4, 3, 1, 1)
+
+    # Without zones defined: the no-zones branch sized its output from
+    # front.shape, which is None when front is omitted.
+    lower = xr.DataArray(data, coords, dims)
+    budgetzone = xr.DataArray(
+        data=np.zeros((3, 2, 1)),
+        coords={"layer": [1, 2, 3], "y": [0.5, 1.5], "x": [0.5]},
+        dims=("layer", "y", "x"),
+    )
+    netflow = imod.evaluate.facebudget(budgetzone, lower=lower)
+    assert netflow.shape == (4, 3, 2, 1)
+    assert np.isnan(netflow.values).all()