diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index 9a1ae02431..405a227183 100644 --- a/src/spikeinterface/postprocessing/amplitude_scalings.py +++ b/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 diff --git a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py b/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py index 0868f5238e..b9bfc09889 100644 --- a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py +++ b/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):