perf(postprocessing): compute only the amplitude scaling coefficients

src/spikeinterface/postprocessing/amplitude_scalings.py

@@ -46,7 +46,7 @@ class ComputeAmplitudeScalings(BaseSpikeVectorExtension):
     handle_collisions: bool, default: True
         Whether to handle collisions between spikes. If True, the amplitude scaling of colliding spikes
         (defined as spikes within `delta_collision_ms` ms and with overlapping sparsity) is computed by fitting a
-        multi-linear regression model (with `sklearn.LinearRegression`). If False, each spike is fitted independently.
+        non-negative multi-linear regression model. If False, each spike is fitted independently.
     delta_collision_ms: float, default: 2
         The maximum time difference in ms before and after a spike to gather colliding spikes.
     """
@@ -227,8 +227,6 @@ def get_dtype(self):
         return self._dtype
 
     def compute(self, traces, peaks):
-        from scipy.stats import linregress
-
         gains = self._gains
         offsets = self._offsets
         all_templates = self._all_templates
@@ -292,21 +290,7 @@ def compute(self, traces, peaks):
                 local_waveform = local_waveform.astype("float32") * gains[sparse_indices] + offsets[sparse_indices]
             assert template.shape == local_waveform.shape
 
-            # here we use linregress, which is equivalent to using sklearn LinearRegression with fit_intercept=True
-            # y = local_waveform.flatten()
-            # X = template.flatten()[:, np.newaxis]
-            # reg = LinearRegression(positive=True, fit_intercept=True).fit(X, y)
-            # scalings[spike_index] = reg.coef_[0]
-
-            # closed form: W = (X' * X)^-1 X' y
-            # y = local_waveform.flatten()[:, None]
-            # X = np.ones((len(y), 2))
-            # X[:, 0] = template.flatten()
-            # W = np.linalg.inv(X.T @ X) @ X.T @ y
-            # scalings[spike_index] = W[0, 0]
-
-            linregress_res = linregress(template.flatten(), local_waveform.flatten())
-            scalings[spike_index] = linregress_res[0]
+            scalings[spike_index] = _ordinary_scaling_slope(template, local_waveform)
 
         # deal with collisions
         if len(collisions) > 0:
@@ -360,6 +344,43 @@ def _are_units_spatially_overlapping(sparsity_mask, i, j):
         return False
 
 
+def _ordinary_scaling_slope(template, local_waveform):
+    """
+    Fit the scaling factor of a single, non-colliding spike against its unit template.
+
+    Equivalent to the slope from ``scipy.stats.linregress(template, local_waveform)``,
+    without its unused statistics (intercept, r-value, p-value, standard errors).
+    The centered covariance/variance are always accumulated in float64: SciPy versions
+    before its array-API rewrite compute ``linregress`` via ``np.cov``, which promotes
+    to float64 regardless of input dtype, while newer SciPy preserves the input dtype.
+    Matching float64 here avoids losing precision relative to either supported version.
+
+    Parameters
+    ----------
+    template : np.ndarray
+        The unit template, cut out to the same window as `local_waveform`.
+    local_waveform : np.ndarray
+        The observed waveform to fit against the template.
+
+    Returns
+    -------
+    float
+        The fitted scaling factor.
+    """
+    template = template.astype(np.float64, copy=False).reshape(-1)
+    local_waveform = local_waveform.astype(np.float64, copy=False).reshape(-1)
+    template_centered = template - np.mean(template)
+    waveform_centered = local_waveform - np.mean(local_waveform)
+    num_samples = template.size
+    template_variance = np.vecdot(template_centered, template_centered) / num_samples
+    if template_variance == 0:
+        from scipy.stats import linregress
+
+        return linregress(template, local_waveform).slope
+    covariance = np.vecdot(template_centered, waveform_centered) / num_samples
+    return covariance / template_variance
+
+
 def find_collisions(spikes, spikes_within_margin, delta_collision_samples, sparsity_mask):
     """
     Finds the collisions between spikes.
@@ -502,7 +523,7 @@ def fit_collision(
     np.ndarray
         The fitted scaling factors for the colliding spikes.
     """
-    from sklearn.linear_model import LinearRegression
+    from scipy.optimize import nnls
 
     # Find the first and last spike peak index
     # from the set of colliding spikes.
@@ -554,8 +575,11 @@ def fit_collision(
 
         X[:, i] = full_template.T.flatten()
 
-    reg = LinearRegression(fit_intercept=True, positive=True).fit(X, y)
-    scalings = reg.coef_
+    # Centering reproduces fit_intercept=True; NNLS reproduces positive=True.
+    y = y.astype(X.dtype, copy=False)
+    X -= np.mean(X, axis=0)
+    y -= np.mean(y)
+    scalings = nnls(X, y)[0]
     return scalings
 
 

src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py

@@ -4,6 +4,63 @@
 from spikeinterface.postprocessing.tests.common_extension_tests import AnalyzerExtensionCommonTestSuite
 
 from spikeinterface.postprocessing import ComputeAmplitudeScalings
+from spikeinterface.postprocessing.amplitude_scalings import _ordinary_scaling_slope, fit_collision
+
+
+def test_ordinary_scaling_slope_float32_precision():
+    """
+    The closed-form slope must not lose precision for float32 template/waveform
+    inputs relative to an independent float64 `linregress` oracle. At amplitude-scale
+    magnitudes (raw ADC/uV range), accumulating in float32 instead of float64 produces
+    an error orders of magnitude above float64 rounding noise.
+    """
+    from scipy.stats import linregress
+
+    rng = np.random.default_rng(2205)
+    template = rng.normal(scale=300, size=90).astype(np.float32)
+    local_waveform = (1.4 * template + rng.normal(scale=20, size=90).astype(np.float32)).astype(np.float32)
+
+    slope = _ordinary_scaling_slope(template.copy(), local_waveform.copy())
+    oracle = linregress(template.astype(np.float64), local_waveform.astype(np.float64)).slope
+
+    assert abs(slope - oracle) < 1e-9
+
+
+def test_fit_collision_recovers_positive_coefficients():
+    """
+    `fit_collision` must recover the known, positive scaling factors of two temporally
+    overlapping spikes, and must be insensitive to a constant offset added to the traces
+    (this exercises the centered non-negative least-squares fit: `positive=True` plus
+    `fit_intercept=True` reproduced by centering before `scipy.optimize.nnls`).
+    """
+    cut_out_before, cut_out_after = 5, 10
+    nbefore = cut_out_before
+    template_length = nbefore + cut_out_after
+
+    rng = np.random.default_rng(7)
+    template_0 = rng.normal(scale=200, size=template_length).astype(np.float32)
+    template_1 = rng.normal(scale=200, size=template_length).astype(np.float32)
+    all_templates = np.stack([template_0, template_1])[:, :, np.newaxis]
+    sparsity_mask = np.ones((2, 1), dtype=bool)
+
+    true_scalings = np.array([1.3, 0.7])
+    spike_0_sample, spike_1_sample = 20, 23  # 3 samples apart: their cut-out windows overlap
+
+    traces = np.zeros((50, 1), dtype=np.float32)
+    traces[spike_0_sample - cut_out_before : spike_0_sample + cut_out_after, 0] += true_scalings[0] * template_0
+    traces[spike_1_sample - cut_out_before : spike_1_sample + cut_out_after, 0] += true_scalings[1] * template_1
+    traces += 50.0  # constant offset: must not bias the fit if centering is correct
+    traces += rng.normal(scale=0.5, size=traces.shape).astype(np.float32)  # small noise
+
+    collision = np.array(
+        [(spike_0_sample, 0), (spike_1_sample, 1)],
+        dtype=[("sample_index", "int64"), ("unit_index", "int64")],
+    )
+
+    recovered = fit_collision(collision, traces, nbefore, all_templates, sparsity_mask, cut_out_before, cut_out_after)
+
+    assert np.all(recovered >= 0)  # positive=True
+    np.testing.assert_allclose(recovered, true_scalings, atol=0.05)
 
 
 class TestAmplitudeScalingsExtension(AnalyzerExtensionCommonTestSuite):