From ab5befb8cac969174eddaddb3e124d80e2c9f8c7 Mon Sep 17 00:00:00 2001
From: Leonardo Uieda <leo@uieda.com>
Date: Fri, 23 Aug 2024 11:43:31 -0300
Subject: [PATCH] Add nearest neighbor interpolation to fill gaps

The `interpolate_missing` function now takes a `method` argument that
can used to set interpolation to nearest-neighbors instead of cubic.
This can be useful for gaps at the corners of images, where cubic
wouldn't fill them in because they're outside of the data convex hull.
---
 doc/conf.py                |  1 +
 xlandsat/_interpolation.py | 20 +++++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/doc/conf.py b/doc/conf.py
index 3d5b479..087bb93 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -40,6 +40,7 @@
     "skimage": ("https://scikit-image.org/docs/stable/", None),
     "pooch": ("https://www.fatiando.org/pooch/latest/", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
 }
 
 # Autosummary pages will be generated by sphinx-autogen instead of sphinx-build
diff --git a/xlandsat/_interpolation.py b/xlandsat/_interpolation.py
index 5c3275d..f2ab0c6 100644
--- a/xlandsat/_interpolation.py
+++ b/xlandsat/_interpolation.py
@@ -8,9 +8,9 @@
 import scipy as sp
 
 
-def interpolate_missing(scene, pixel_radius=20):
+def interpolate_missing(scene, pixel_radius=20, method="cubic"):
     """
-    Fill missing values (NaNs) in a scene by cubic interpolation
+    Fill missing values (NaNs) in a scene by interpolation
 
     Each missing value is filled by interpolating the pixels within a
     neighboring region (controlled by ``pixel_radius``) using a piecewise cubic
@@ -28,12 +28,26 @@ def interpolate_missing(scene, pixel_radius=20):
         that will be used for interpolation. Smaller values make for faster
         interpolation but may lead to bad results if many missing values are
         grouped together.
+    method : str
+        The interpolation method used. Can be one of: ``"cubic"`` for cubic
+        interpolation with
+        :class:`scipy.interpolate.CloughTocher2DInterpolator`, ``"neighbors"``
+        for nearest neighbor interpolation with
+        :class:`scipy.interpolate.NearestNDInterpolator`.
 
     Returns
     -------
     filled_scene : :class:`xarray.Dataset`
         The scene with missing values filled in.
     """
+    valid_methods = {
+        "cubic": sp.interpolate.CloughTocher2DInterpolator,
+        "neighbors": sp.interpolate.NearestNDInterpolator,
+    }
+    if method not in valid_methods:
+        raise ValueError(
+            f"Invalid interpolation method '{method}'. Must be one of: {tuple(valid_methods.keys())}"
+        )
     filled_scene = scene.copy(deep=True)
     rows, columns = np.meshgrid(
         np.arange(scene.northing.size),
@@ -47,7 +61,7 @@ def interpolate_missing(scene, pixel_radius=20):
             imin, imax = _search_range(i, pixel_radius, scene.northing.size)
             jmin, jmax = _search_range(j, pixel_radius, scene.easting.size)
             valid = ~np.isnan(values[imin:imax, jmin:jmax])
-            interpolator = sp.interpolate.CloughTocher2DInterpolator(
+            interpolator = valid_methods[method](
                 (
                     rows[imin:imax, jmin:jmax][valid],
                     columns[imin:imax, jmin:jmax][valid],