Add interpolation_kwargs to interpolate_templates (#4665)

Author: chrishalcrow

src/spikeinterface/generation/drift_tools.py

@@ -13,6 +13,7 @@ def interpolate_templates(
     source_locations: np.ndarray,
     dest_locations: np.ndarray,
     interpolation_method: Literal["cubic", "linear", "nearest", "thin_plate"] = "cubic",
+    interpolation_kwargs: None | dict = None,
 ):
     """
     Interpolate templates_array to new positions.
@@ -33,6 +34,9 @@ def interpolate_templates(
         shape = (num_channels, 2) or (num_motions, num_channels, 2)
     interpolation_method : str, default "cubic"
         The interpolation method.
+    interpolation_kwargs
+        Kwargs that are passed to RBFInterpolator (if interpolation_method = "thin_plate") or to
+        griddata (otherwise).
 
     Returns
     -------
@@ -41,6 +45,9 @@ def interpolate_templates(
     """
     from scipy.interpolate import griddata, RBFInterpolator
 
+    if interpolation_kwargs is None:
+        interpolation_kwargs = dict()
+
     source_locations = np.asarray(source_locations)
     dest_locations = np.asarray(dest_locations)
 
@@ -65,7 +72,15 @@ def interpolate_templates(
 
             if interpolation_method == "thin_plate":
 
-                tps_interpolator = RBFInterpolator(source_locations, template, kernel="thin_plate_spline", neighbors=12)
+                if "neighbors" not in interpolation_kwargs:
+                    # If neighbors it not passed, `RBFInterpolator` uses all channels.
+                    # This works poorly for standard ephys template interpolation.
+                    # Depending on the probe, you might want to decrease this value.
+                    interpolation_kwargs["neighbors"] = 12
+
+                tps_interpolator = RBFInterpolator(
+                    source_locations, template, kernel="thin_plate_spline", **interpolation_kwargs
+                )
                 if dest_locations_dims == 2:
                     interp_template = tps_interpolator(dest_locations)
                 elif dest_locations_dims == 3:
@@ -75,7 +90,12 @@ def interpolate_templates(
 
             else:
                 interp_template = griddata(
-                    source_locations, template, dest_locations, method=interpolation_method, fill_value=0
+                    source_locations,
+                    template,
+                    dest_locations,
+                    method=interpolation_method,
+                    fill_value=0,
+                    **interpolation_kwargs,
                 )
 
             if dest_locations_dims == 2:

src/spikeinterface/generation/tests/test_drift_tools.py

@@ -83,6 +83,14 @@ def test_interpolate_templates():
             templates.templates_array, source_locations, dest_locations, interpolation_method=interpolation_method
         )
 
+    interpolate_templates(
+        templates.templates_array,
+        source_locations,
+        dest_locations,
+        interpolation_method="thin_plate",
+        interpolation_kwargs={"neighbors": 12},
+    )
+
 
 def test_move_dense_templates():
     templates = make_some_templates()