diff --git a/regridding/_fill/_fill.py b/regridding/_fill/_fill.py index 578f4b7..89cb226 100644 --- a/regridding/_fill/_fill.py +++ b/regridding/_fill/_fill.py @@ -34,6 +34,7 @@ def fill( kwargs Additional method-specific keyword arguments. For the Gauss-Seidel method, the valid keyword arguments are: + - ``guess=np.median(a[where], axis)``, the first guess at the fill value. - ``num_iterations=100``, the number of red-black Gauss-Seidel iterations to perform. Examples diff --git a/regridding/_fill/_gauss_seidel.py b/regridding/_fill/_gauss_seidel.py index b4278f7..40c3644 100644 --- a/regridding/_fill/_gauss_seidel.py +++ b/regridding/_fill/_gauss_seidel.py @@ -12,14 +12,18 @@ def fill_gauss_seidel( a: np.ndarray, where: np.ndarray, axis: None | int | Sequence[int], + guess: None | float | np.ndarray = None, num_iterations: int = 100, ) -> np.ndarray: a = a.copy() - a, where = np.broadcast_arrays(a, where, subok=True) + if guess is None: + guess = np.median(a[where], axis=axis) - a[where] = 0 + a, where, guess = np.broadcast_arrays(a, where, guess, subok=True) + + a[where] = guess[where] axis = regridding._util._normalize_axis(axis=axis, ndim=a.ndim) axis_numba = ~np.arange(len(axis))[::-1]