Issue #1729 fix convert_pointwaterhead_freshwaterhead for ugrid

docs/api/changelog.rst

@@ -41,6 +41,9 @@ Fixed
   unaligned.
 - Fixed bug where :class:`imod.mf6.Lake` package did not pass ``budgetfile``,
   ``budgetcsvfile``, ``stagefile`` options to the written MODFLOW 6 package.
+- Fixed bug where :func:`imod.evaluate.convert_pointwaterhead_freshwaterhead`
+  produced incorrect results when point water heads were below elevation levels
+  for unstructured grids.
 
 Changed
 ~~~~~~~

imod/evaluate/head.py

@@ -4,10 +4,15 @@
 import pandas as pd
 import xarray as xr
 
+from imod.typing import GridDataArray
+
 
 def convert_pointwaterhead_freshwaterhead(
-    pointwaterhead, density, elevation, density_fresh=1000.0
-):
+    pointwaterhead: GridDataArray,
+    density: GridDataArray,
+    elevation: GridDataArray,
+    density_fresh: float = 1000.0,
+) -> GridDataArray:
     r"""Function to convert point water head (as outputted by seawat)
     into freshwater head, using Eq.3 from Guo, W., & Langevin, C. D. (2002):
 
@@ -21,19 +26,19 @@ def convert_pointwaterhead_freshwaterhead(
 
     Parameters
     ----------
-    pointwaterhead : float or xr.DataArray of floats
+    pointwaterhead : float or xr.DataArray/xu.UgridDataArray of floats
         the point water head as outputted by SEAWAT, in m.
-    density : float or xr.DataArray of floats
+    density : float or xr.DataArray/xu.UgridDataArray of floats
         the water density at the same locations as `pointwaterhead`.
-    elevation : float or xr.DataArray of floats
+    elevation : float or xr.DataArray/xu.UgridDataArray of floats
         elevation at the same locations as `pointwaterhead`, in m.
     density_fresh : float, optional
         the density of freshwater (1000 kg/m3), or a different value if
         different units are used, or a different density reference is required.
 
     Returns
     -------
-    freshwaterhead : float or xr.DataArray of floats
+    freshwaterhead : float or xr.xr.DataArray/xu.UgridDataArray of floats
     """
 
     freshwaterhead = (
@@ -43,8 +48,8 @@ def convert_pointwaterhead_freshwaterhead(
 
     # edge case: point water head below z
     # return freshwater head of top underlying cell where elevation < pointwaterhead
-    # only for xr.DataArrays
-    if isinstance(pointwaterhead, xr.DataArray) and "layer" in pointwaterhead.dims:
+    # only for grids with layers
+    if isinstance(pointwaterhead, GridDataArray) and "layer" in pointwaterhead.dims:
         freshwaterhead = freshwaterhead.where(pointwaterhead > elevation)
         freshwaterhead = freshwaterhead.bfill(dim="layer")
 

imod/tests/test_evaluate/test_head.py

@@ -2,8 +2,11 @@
 import pandas as pd
 import pytest
 import xarray as xr
+import xugrid as xu
+from pytest_cases import parametrize_with_cases
 
 import imod
+from imod.typing.grid import full_like
 
 
 def test_convert_pointwaterhead_freshwaterhead_scalar():
@@ -20,18 +23,51 @@ def test_convert_pointwaterhead_freshwaterhead_scalar():
     )
 
 
-def test_convert_pointwaterhead_freshwaterhead_da():
-    data = np.ones((3, 2, 1))
-    coords = {"layer": [1, 2, 3], "y": [0.5, 1.5], "x": [0.5]}
+def structured_grid():
+    data = np.ones((3, 2, 2))
+    coords = {"layer": [1, 2, 3], "y": [0.5, 1.5], "x": [0.5, 1.5]}
     dims = ("layer", "y", "x")
 
-    da = xr.DataArray(data, coords, dims)
-    pwh = xr.full_like(da, 4.0)
-    dens = xr.full_like(da, 1025.0)
-    z = xr.full_like(da, 1.0)
-    fwh = xr.full_like(da, 4.075)
-    fwh2 = imod.evaluate.convert_pointwaterhead_freshwaterhead(pwh, dens, z)
-    assert fwh.equals(fwh2.round(5))
+    return xr.DataArray(data, coords, dims)
+
+
+class GridCases:
+    def case_structured(self):
+        return structured_grid()
+
+    def case_unstructured(self):
+        structured = structured_grid()
+        return xu.UgridDataArray.from_structured(structured)
+
+
+@parametrize_with_cases("da", cases=GridCases)
+def test_convert_pointwaterhead_freshwaterhead_da(da):
+    pwh = full_like(da, 4.0)
+    dens = full_like(da, 1025.0)
+    z = full_like(da, 1.0)
+    fwh_expected = full_like(da, 4.075)
+    fwh_actual = imod.evaluate.convert_pointwaterhead_freshwaterhead(pwh, dens, z)
+    assert type(fwh_actual) is type(da)
+    assert fwh_expected.equals(fwh_actual.round(5))
+
+
+@parametrize_with_cases("da", cases=GridCases)
+def test_convert_pointwaterhead_freshwaterhead_backfill(da):
+    # edge case: point water head below z
+    # return freshwater head of top underlying cell where elevation < pointwaterhead
+    # only for grids with layers
+    #
+    # Test this by setting point water head in top layer to 0.0. This should
+    # trigger a backfill with the value from layer 2, which is 4.075. So in the
+    # end, test results should be 4.075 everywhere.
+    pwh = full_like(da, 4.0)
+    pwh[0, ...] = 0.0
+    dens = full_like(da, 1025.0)
+    z = full_like(da, 1.0)
+    fwh_expected = full_like(da, 4.075)
+    fwh_actual = imod.evaluate.convert_pointwaterhead_freshwaterhead(pwh, dens, z)
+    assert type(fwh_actual) is type(da)
+    assert fwh_expected.equals(fwh_actual.round(5))
 
 
 @pytest.mark.parametrize("chunk", [False, True])