Add rhs argument to enable "poisson interpolation" in laplace_interpo…

imod/prepare/spatial.py

@@ -112,7 +112,15 @@ def fill(da, invalid=None, by=None):
 
 
 def laplace_interpolate(
-    source, ibound=None, close=0.01, mxiter=5, iter1=50, relax=0.98
+    source,
+    ibound=None,
+    close=0.01,
+    mxiter=5,
+    iter1=250,
+    relax=0.98,
+    rhs=None,
+    robin_value=None,
+    robin_coefficient=None,
 ):
     """
     Fills gaps in `source` by interpolating from existing values using Laplace
@@ -121,44 +129,64 @@ def laplace_interpolate(
     Parameters
     ----------
     source : xr.DataArray of floats with dims (y, x)
-        Data values to interpolate.
+        Data values to interpolate. (MODFLOW CHD values.)
     ibound : xr.DataArray of bool with dims (y, x)
         Precomputed array which marks where to interpolate.
     close : float
         Closure criteration of iterative solver. Should be one to two orders
         of magnitude smaller than desired accuracy.
     mxiter : int
-        Outer iterations of iterative solver.
+        Outer iterations of iterative solver. Only required to avoid numerical
+        breakdown.
     iter1 : int
-        Inner iterations of iterative solver. Should not exceed 50.
+        Inner iterations of iterative solver.
     relax : float
         Iterative solver relaxation parameter. Should be between 0 and 1.
+    rhs : xr.DataArray of floats with dims (y, x).
+        Right-hand-side of the linear system of equations: Laplace equation if
+        zero, Poisson equation if nonzero. (MODFLOW WEL rate value.)
+    robin_value : xr.DataArray of floats with dims (y, x).
+        Robin boundary value. (MODFLOW GHB head values.)
+    robin_coefficient: xr.DataArray of floats with dims (y, x).
+        Robin boundary coefficient. (MODFLOW GHB conductance values.)
 
     Returns
     -------
     interpolated : xr.DataArray with dims (y, x)
         source, with interpolated values where ibound equals 1
     """
-    solver = pcg.PreconditionedConjugateGradientSolver(
-        close, close * 1.0e6, mxiter, iter1, relax
-    )
 
-    if not source.dims == ("y", "x"):
+    def _check(da: xr.DataArray, name: str) -> None:
+        if da.dims != ("y", "x"):
+            raise ValueError(f'{name} dims must be ("y", "x")')
+        if da.shape != source.shape:
+            raise ValueError(f"source and {name} must have the same shape")
+
+    if source.dims != ("y", "x"):
         raise ValueError('source dims must be ("y", "x")')
+    if rhs is not None:
+        _check(rhs, "rhs")
+
+    if (robin_value is None) ^ (robin_coefficient is None):
+        raise ValueError("Provide both robin_value and robin_coefficient, or neither.")
+    if robin_value is not None:
+        _check(robin_value, "robin_value")
+        _check(robin_coefficient, "robin_coefficient")
+        if (robin_value.isnull() != robin_coefficient.isnull()).any():
+            raise ValueError(
+                "robin_value and robin_coefficient must be consistent in terms "
+                "of nodata (NaN) values."
+            )
 
     # expand dims to make 3d
     source3d = source.expand_dims("layer")
     if ibound is None:
-        iboundv = xr.full_like(source3d, 1, dtype=np.int32).values
+        iboundv = xr.full_like(source3d, 1, dtype=np.int32).to_numpy()
+        iboundv[source3d.notnull().values] = -1
     else:
-        if not ibound.dims == ("y", "x"):
-            raise ValueError('ibound dims must be ("y", "x")')
-        if not ibound.shape == source.shape:
-            raise ValueError("ibound and source must have the same shape")
-        iboundv = ibound.expand_dims("layer").astype(np.int32).values
-
-    has_data = source3d.notnull().values
-    iboundv[has_data] = -1
+        _check(ibound, "ibound")
+        iboundv = ibound.expand_dims("layer").astype(np.int32).to_numpy()
+
     hnew = source3d.fillna(0.0).values.astype(
         np.float64
     )  # Set start interpolated estimate to 0.0
@@ -171,8 +199,28 @@ def laplace_interpolate(
     cc = np.ones(shape)
     cr = np.ones(shape)
     cv = np.ones(shape)
-    rhs = np.zeros(shape)
-    hcof = np.zeros(shape)
+    if rhs is None:
+        rhsv = np.zeros(shape)
+    else:
+        rhsv = rhs.expand_dims("layer").astype(np.float64).to_numpy().copy()
+
+    if robin_value is None:
+        hcof = np.zeros(shape)
+    else:
+        hcof = (
+            -robin_coefficient.fillna(0.0)
+            .expand_dims("layer")
+            .astype(np.float64)
+            .to_numpy()
+        )
+        rhsv -= (
+            (robin_coefficient * robin_value)
+            .fillna(0.0)
+            .expand_dims("layer")
+            .astype(np.float64)
+            .to_numpy()
+        )
+
     # Solver work arrays
     res = np.zeros(nodes)
     cd = np.zeros(nodes)
@@ -181,6 +229,10 @@ def laplace_interpolate(
     p = np.zeros(nodes)
 
     # Picard iteration
+    solver = pcg.PreconditionedConjugateGradientSolver(
+        close, close * 1.0e6, mxiter, iter1, relax
+    )
+
     converged = False
     outer_iteration = 0
     while not converged and outer_iteration < mxiter:
@@ -191,7 +243,7 @@ def laplace_interpolate(
             cr=cr,
             cv=cv,
             ibound=iboundv,
-            rhs=rhs,
+            rhs=rhsv,
             hcof=hcof,
             res=res,
             cd=cd,