diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index dacf8e9c9..346ec5846 100644 --- a/docs/api/changelog.rst +++ b/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 ~~~~~~~ diff --git a/imod/evaluate/budget.py b/imod/evaluate/budget.py index aed2065be..47c083192 100644 --- a/imod/evaluate/budget.py +++ b/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) diff --git a/imod/tests/test_evaluate/test_budget.py b/imod/tests/test_evaluate/test_budget.py index 17f7e6406..aa12e8cc7 100644 --- a/imod/tests/test_evaluate/test_budget.py +++ b/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()