compute theil-sen estimator manually (exact and faster)

Author: ecobost

src/spikeinterface/metrics/template/metrics.py

@@ -283,22 +283,27 @@ def sort_template_and_locations(template, channel_locations, depth_direction="y"
     sort_indices = np.argsort(channel_locations[:, depth_dim])
     return template[:, sort_indices], channel_locations[sort_indices, :]
 
-
-def fit_line_robust(x, y):
+def fit_line_robust(x, y, eps=1e-12):
     """
     Fit line using robust Theil-Sen estimator (median of pairwise slopes).
     """
-    import sklearn.linear_model as lm
-
-    # Center data to improve numerical stability
-    X = (x - x.mean()).reshape(-1, 1)
-    y = y - y.mean()
-
-    theil = lm.TheilSenRegressor()
-    theil.fit(X, y)
-    slope = theil.coef_[0]
-    score = theil.score(X, y)  # R^2 score
-    return slope, score
+    import itertools
+
+    # Calculate slope and bias using Theil-Sen estimator
+    slopes = []
+    for (xs, ys) in zip(itertools.combinations(x, 2), itertools.combinations(y, 2)):
+        if np.abs(xs[0] - xs[1]) > eps:
+            slopes.append((ys[1] - ys[0]) / (xs[1] - xs[0]))
+    if len(slopes) == 0:  # all x are identical
+        return np.nan, -np.inf
+    slope = np.median(slopes)
+    bias = np.median(y - slope * x)
+
+    # Calculate R2 score
+    y_pred = slope * x + bias
+    r2_score = 1 - ((y - y_pred)**2).sum() / (((y - y.mean())**2).sum() + eps)
+
+    return slope, r2_score
 
 
 def get_velocity_fits(template, channel_locations, sampling_frequency, **kwargs):