From 777a0440ca56844680c68c05bddf1c7ac2f06361 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 21 Aug 2026 09:41:09 +0100 Subject: [PATCH 1/3] wip: refactor so compute P is isolated --- .../sortingcomponents/motion/dredge.py | 68 ++++++++++++------- 1 file changed, 44 insertions(+), 24 deletions(-) diff --git a/src/spikeinterface/sortingcomponents/motion/dredge.py b/src/spikeinterface/sortingcomponents/motion/dredge.py index bc3fa48246..2c6f5435f1 100644 --- a/src/spikeinterface/sortingcomponents/motion/dredge.py +++ b/src/spikeinterface/sortingcomponents/motion/dredge.py @@ -588,6 +588,33 @@ def dredge_online_lfp( if extra_outputs: extra = dict(window_centers=window_centers, windows=windows) + P_online = compute_P_online + + motion = Motion([P_online.T], [lfp_recording.get_times(0)], window_centers, direction=direction) + + if extra_outputs: + return motion, extra + else: + return motion + + +def compute_P_online( + lfp_recording, + B, + T_total, + T_chunk, + windows, + contact_depths, + # parameters + win_scale_um, + full_xcorr_kw, + mincorr, + mincorr_percentile, + threshold_kw, + thomas_kw, + chunk_len_s, +): + # -- allocate output and initialize first chunk P_online = np.empty((B, T_total), dtype=np.float32) # below, t0 is start of prev chunk, t1 start of cur chunk, t2 end of cur @@ -601,15 +628,15 @@ def dredge_online_lfp( mincorr_percentile=mincorr_percentile, **threshold_kw, ) - if extra_outputs: - extra["D"] = [Ds0] - extra["C"] = [Cs0] - extra["S"] = [Ss0] - extra["D01"] = [] - extra["C01"] = [] - extra["S01"] = [] - extra["mincorrs"] = [mincorr0] - extra["max_disp_um"] = max_disp_um + # if extra_outputs: + # extra["D"] = [Ds0] + # extra["C"] = [Cs0] + # extra["S"] = [Ss0] + # extra["D01"] = [] + # extra["C01"] = [] + # extra["S01"] = [] + # extra["mincorrs"] = [mincorr0] + # extra["max_disp_um"] = max_disp_um P_online[:, t0:t1], _ = thomas_solve(Ds0, Ss0, **thomas_kw) @@ -647,14 +674,14 @@ def dredge_online_lfp( ) Ss10, _ = threshold_correlation_matrix(Cs10, mincorr=mincorr1, t_offset_bins=T_chunk, **threshold_kw) - if extra_outputs: - extra["mincorrs"].append(mincorr1) - extra["D"].append(Ds1) - extra["C"].append(Cs1) - extra["S"].append(Ss1) - extra["D01"].append(Ds10) - extra["C01"].append(Cs10) - extra["S01"].append(Ss10) + # if extra_outputs: + # extra["mincorrs"].append(mincorr1) + # extra["D"].append(Ds1) + # extra["C"].append(Cs1) + # extra["S"].append(Ss1) + # extra["D01"].append(Ds10) + # extra["C01"].append(Cs10) + # extra["S01"].append(Ss10) # solve online problem P_online[:, t1:t2], _ = thomas_solve( @@ -672,13 +699,6 @@ def dredge_online_lfp( t0, t1 = t1, t2 traces0 = traces1 - motion = Motion([P_online.T], [lfp_recording.get_times(0)], window_centers, direction=direction) - - if extra_outputs: - return motion, extra - else: - return motion - dredge_online_lfp.__doc__ = dredge_online_lfp.__doc__.format(DredgeLfpRegistration.params_doc) From 03b74decf1f6c22424421afbb9c70fc169b2afb3 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 28 Aug 2026 22:42:41 +0100 Subject: [PATCH 2/3] refactor computation of displacements to unify arguments --- .../sortingcomponents/motion/dredge.py | 246 +++++++++++------- 1 file changed, 159 insertions(+), 87 deletions(-) diff --git a/src/spikeinterface/sortingcomponents/motion/dredge.py b/src/spikeinterface/sortingcomponents/motion/dredge.py index 2c6f5435f1..c77723f2ce 100644 --- a/src/spikeinterface/sortingcomponents/motion/dredge.py +++ b/src/spikeinterface/sortingcomponents/motion/dredge.py @@ -322,59 +322,30 @@ def dredge_ap( # if extra_outputs and count_masked_correlation: # extra["counts"] = counts - # cross-correlate to get D and C - if precomputed_D_C_maxdisp is None: - Ds, Cs, max_disp_um = xcorr_windows( - raster, - windows, - spatial_bin_edges_um, - win_scale_um, - rigid=rigid, - bin_um=bin_um, - max_disp_um=max_disp_um, - progress_bar=progress_bar, - device=device, - # TODO charlie : put back the count for the mask - # masks=(counts > 0) if count_masked_correlation else None, - **xcorr_kw, - ) - else: - Ds, Cs, max_disp_um = precomputed_D_C_maxdisp + full_xcorr_kw = dict( + rigid=rigid, + bin_um=np.median(np.diff(contact_depths)), + max_disp_um=max_disp_um, + progress_bar=False, + device=device, + **xcorr_kw, + ) - # turn Cs into weights - Us, wextra = weight_correlation_matrix( - Ds, - Cs, - windows, + displacement, extra = compute_displacement_simultaneous( raster, + windows, spatial_bin_edges_um, - time_bin_edges_s, - # raster_kw, #@charlie this is removed - post_transform=post_transform, # @charlie this isnew - lambda_t=thomas_kw.get("lambda_t", DEFAULT_LAMBDA_T), - eps=thomas_kw.get("eps", DEFAULT_EPS), - progress_bar=progress_bar, - in_place=not extra_outputs, - **weights_kw, + win_scale_um, + bin_s, + mincorr_percentile, + extra, + extra_outputs, + thomas_kw, + weights_kw, + full_xcorr_kw, + precomputed_D_C_maxdisp, + post_transform, ) - extra.update({k: wextra[k] for k in wextra if k not in ("S", "U")}) - if extra_outputs: - extra.update({k: wextra[k] for k in wextra if k in ("S", "U")}) - del wextra - if extra_outputs: - extra["D"] = Ds - extra["C"] = Cs - del Cs - - # @charlie : is this needed ? - gc.collect() - - # solve for P - # now we can do our tridiag solve - displacement, textra = thomas_solve(Ds, Us, progress_bar=progress_bar, **thomas_kw) - if extra_outputs: - extra.update(textra) - del textra if extra_outputs: extra["windows"] = windows @@ -382,6 +353,7 @@ def dredge_ap( extra["max_disp_um"] = max_disp_um time_bin_centers = 0.5 * (time_bin_edges_s[1:] + time_bin_edges_s[:-1]) + motion = Motion([displacement.T], [time_bin_centers], window_centers, direction=direction) if extra_outputs: @@ -583,12 +555,31 @@ def dredge_online_lfp( zero_threshold=1e-5, ) - B = len(windows) - if extra_outputs: extra = dict(window_centers=window_centers, windows=windows) - P_online = compute_P_online + weights_kw = dict( + mincorr=mincorr, + time_horizon_s=time_horizon_s, + ) + + bin_s = 1 / lfp_recording.sampling_frequency + P_online = compute_displacement_online( + lfp_recording, + windows, + T_total, + T_chunk, + contact_depths, + win_scale_um, + bin_s, + mincorr_percentile, + extra, + extra_outputs, + thomas_kw, + weights_kw, + full_xcorr_kw, + threshold_kw, + ) motion = Motion([P_online.T], [lfp_recording.get_times(0)], window_centers, direction=direction) @@ -598,56 +589,67 @@ def dredge_online_lfp( return motion -def compute_P_online( +def compute_displacement_online( lfp_recording, - B, + windows, T_total, T_chunk, - windows, - contact_depths, - # parameters + spatial_bin_edges_um, win_scale_um, - full_xcorr_kw, - mincorr, + bin_s, mincorr_percentile, - threshold_kw, + extra, + extra_outputs, thomas_kw, - chunk_len_s, + weights_kw, + full_xcorr_kw, + threshold_kw, ): + B = len(windows) + # -- allocate output and initialize first chunk P_online = np.empty((B, T_total), dtype=np.float32) # below, t0 is start of prev chunk, t1 start of cur chunk, t2 end of cur t0, t1 = 0, T_chunk + traces0 = lfp_recording.get_traces(start_frame=t0, end_frame=t1) - Ds0, Cs0, max_disp_um = xcorr_windows(traces0.T, windows, contact_depths, win_scale_um, **full_xcorr_kw) - full_xcorr_kw["max_disp_um"] = max_disp_um + + Ds0, Cs0, max_disp_um = xcorr_windows( + traces0.T, + windows, + spatial_bin_edges_um, + win_scale_um, + **full_xcorr_kw, + ) + mincorr = weights_kw["mincorr"] Ss0, mincorr0 = threshold_correlation_matrix( Cs0, mincorr=mincorr, mincorr_percentile=mincorr_percentile, **threshold_kw, ) - # if extra_outputs: - # extra["D"] = [Ds0] - # extra["C"] = [Cs0] - # extra["S"] = [Ss0] - # extra["D01"] = [] - # extra["C01"] = [] - # extra["S01"] = [] - # extra["mincorrs"] = [mincorr0] - # extra["max_disp_um"] = max_disp_um + if extra_outputs: + extra["D"] = [Ds0] + extra["C"] = [Cs0] + extra["S"] = [Ss0] + extra["D01"] = [] + extra["C01"] = [] + extra["S01"] = [] + extra["mincorrs"] = [mincorr0] + extra["max_disp_um"] = max_disp_um P_online[:, t0:t1], _ = thomas_solve(Ds0, Ss0, **thomas_kw) # -- loop through chunks chunk_starts = range(T_chunk, T_total, T_chunk) + progress_bar = full_xcorr_kw["progress_bar"] if progress_bar: chunk_starts = trange( T_chunk, T_total, T_chunk, - desc=f"Online chunks [{chunk_len_s}s each]", + desc=f"Online chunks [{T_chunk}s each]", ) for t1 in chunk_starts: t2 = min(T_total, t1 + T_chunk) @@ -658,14 +660,14 @@ def compute_P_online( Ds10, Cs10, _ = xcorr_windows( traces1.T, windows, - contact_depths, + spatial_bin_edges_um, win_scale_um, raster_b=traces0.T, **full_xcorr_kw, ) # cross-correlation in current chunk - Ds1, Cs1, _ = xcorr_windows(traces1.T, windows, contact_depths, win_scale_um, **full_xcorr_kw) + Ds1, Cs1, _ = xcorr_windows(traces1.T, windows, spatial_bin_edges_um, win_scale_um, **full_xcorr_kw) Ss1, mincorr1 = threshold_correlation_matrix( Cs1, mincorr_percentile=mincorr_percentile, @@ -674,14 +676,14 @@ def compute_P_online( ) Ss10, _ = threshold_correlation_matrix(Cs10, mincorr=mincorr1, t_offset_bins=T_chunk, **threshold_kw) - # if extra_outputs: - # extra["mincorrs"].append(mincorr1) - # extra["D"].append(Ds1) - # extra["C"].append(Cs1) - # extra["S"].append(Ss1) - # extra["D01"].append(Ds10) - # extra["C01"].append(Cs10) - # extra["S01"].append(Ss10) + if extra_outputs: + extra["mincorrs"].append(mincorr1) + extra["D"].append(Ds1) + extra["C"].append(Cs1) + extra["S"].append(Ss1) + extra["D01"].append(Ds10) + extra["C01"].append(Cs10) + extra["S01"].append(Ss10) # solve online problem P_online[:, t1:t2], _ = thomas_solve( @@ -699,6 +701,78 @@ def compute_P_online( t0, t1 = t1, t2 traces0 = traces1 + return P_online + + +def compute_displacement_simultaneous( + raster, + windows, + spatial_bin_edges_um, + win_scale_um, + bin_s, + mincorr_percentile, + extra, + extra_outputs, + thomas_kw, + weights_kw, + full_xcorr_kw, + precomputed_D_C_maxdisp, + post_transform, +): + + # cross-correlate to get D and C + if precomputed_D_C_maxdisp is None: + Ds, Cs, max_disp_um = xcorr_windows( + raster, + windows, + spatial_bin_edges_um, + win_scale_um, + # TODO charlie : put back the count for the mask + # masks=(counts > 0) if count_masked_correlation else None, + **full_xcorr_kw, + ) + else: + Ds, Cs, max_disp_um = precomputed_D_C_maxdisp + + # turn Cs into weights + progress_bar = full_xcorr_kw["progress_bar"] + Us, wextra = weight_correlation_matrix( + Ds, + Cs, + windows, + raster, + spatial_bin_edges_um, + bin_s, + # raster_kw, #@charlie this is removed + post_transform=post_transform, # @charlie this isnew + lambda_t=thomas_kw.get("lambda_t", DEFAULT_LAMBDA_T), + eps=thomas_kw.get("eps", DEFAULT_EPS), + progress_bar=progress_bar, + in_place=not extra_outputs, + mincorr_percentile=mincorr_percentile, + **weights_kw, + ) + extra.update({k: wextra[k] for k in wextra if k not in ("S", "U")}) + if extra_outputs: + extra.update({k: wextra[k] for k in wextra if k in ("S", "U")}) + del wextra + if extra_outputs: + extra["D"] = Ds + extra["C"] = Cs + del Cs + + # @charlie : is this needed ? + gc.collect() + + # solve for P + # now we can do our tridiag solve + displacement, textra = thomas_solve(Ds, Us, progress_bar=progress_bar, **thomas_kw) + if extra_outputs: + extra.update(textra) + del textra + + return displacement, extra + dredge_online_lfp.__doc__ = dredge_online_lfp.__doc__.format(DredgeLfpRegistration.params_doc) @@ -1309,7 +1383,6 @@ def get_weights( windows, raster, dbe, - tbe, # @charlie raster_kw is removed in favor of post_transform only is this OK ??? # raster_kw, post_transform=np.log1p, @@ -1355,7 +1428,7 @@ def weight_correlation_matrix( windows, raster, depth_bin_edges, - time_bin_edges, + bin_s, # @charlie raster_kw is remove in favor of post_transform only # raster_kw, post_transform=np.log1p, @@ -1390,7 +1463,7 @@ def weight_correlation_matrix( mincorr_percentile=mincorr_percentile, mincorr_percentile_nneighbs=mincorr_percentile_nneighbs, time_horizon_s=time_horizon_s, - bin_s=time_bin_edges[1] - time_bin_edges[0], + bin_s=bin_s, T=T, in_place=in_place, ) @@ -1409,7 +1482,6 @@ def weight_correlation_matrix( windows, raster, depth_bin_edges, - time_bin_edges, # raster_kw, post_transform=post_transform, weights_threshold_low=weights_threshold_low, From 7dd3c84741deffb7a898f40f6540f25d1b151b42 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 28 Aug 2026 23:19:55 +0100 Subject: [PATCH 3/3] WIP: allow for choice of resolution mode --- src/spikeinterface/preprocessing/motion.py | 1 + .../sortingcomponents/motion/dredge.py | 77 ++++++++++++++----- 2 files changed, 58 insertions(+), 20 deletions(-) diff --git a/src/spikeinterface/preprocessing/motion.py b/src/spikeinterface/preprocessing/motion.py index 62bdbfd7a9..273927a1d5 100644 --- a/src/spikeinterface/preprocessing/motion.py +++ b/src/spikeinterface/preprocessing/motion.py @@ -38,6 +38,7 @@ win_step_um=400.0, win_scale_um=400.0, win_margin_um=None, + resolution_mode="simulataneous", ), "interpolate_motion_kwargs": dict( border_mode="force_extrapolate", spatial_interpolation_method="kriging", sigma_um=20.0, p=2 diff --git a/src/spikeinterface/sortingcomponents/motion/dredge.py b/src/spikeinterface/sortingcomponents/motion/dredge.py index c77723f2ce..e8da72a10d 100644 --- a/src/spikeinterface/sortingcomponents/motion/dredge.py +++ b/src/spikeinterface/sortingcomponents/motion/dredge.py @@ -29,6 +29,7 @@ import numpy as np from tqdm.auto import trange +from spikeinterface.core import BaseRecording from spikeinterface.core.motion import Motion from .motion_utils import ( get_spatial_bin_edges, @@ -173,6 +174,7 @@ def dredge_ap( progress_bar=True, extra_outputs=False, precomputed_D_C_maxdisp=None, + resolution_mode="simultaneous", ): """Estimate motion from spikes @@ -324,28 +326,57 @@ def dredge_ap( full_xcorr_kw = dict( rigid=rigid, - bin_um=np.median(np.diff(contact_depths)), + bin_um=bin_um, max_disp_um=max_disp_um, progress_bar=False, device=device, **xcorr_kw, ) - displacement, extra = compute_displacement_simultaneous( - raster, - windows, - spatial_bin_edges_um, - win_scale_um, - bin_s, - mincorr_percentile, - extra, - extra_outputs, - thomas_kw, - weights_kw, - full_xcorr_kw, - precomputed_D_C_maxdisp, - post_transform, - ) + if resolution_mode == "simultaneous": + displacement, extra = compute_displacement_simultaneous( + raster, + windows, + spatial_bin_edges_um, + win_scale_um, + bin_s, + mincorr_percentile, + extra, + extra_outputs, + thomas_kw, + weights_kw, + full_xcorr_kw, + precomputed_D_C_maxdisp, + post_transform, + ) + elif resolution_mode == "online": + T_total = raster.shape[1] + T_chunk = 10 + threshold_kw = dict( + mincorr_percentile_nneighbs=mincorr_percentile_nneighbs, + in_place=True, + soft=False, + # time_horizon_s=weights_kw["time_horizon_s"], # max_dt not implemented for lfp at this point + time_horizon_s=time_horizon_s, + bin_s=10, + ) + displacement, extra = compute_displacement_online( + raster, + windows, + T_total, + T_chunk, + spatial_bin_edges_um, + win_scale_um, + mincorr_percentile, + extra, + extra_outputs, + thomas_kw, + weights_kw, + full_xcorr_kw, + threshold_kw, + ) + else: + raise ValueError(f"No resolution mode called {resolution_mode}") if extra_outputs: extra["windows"] = windows @@ -596,7 +627,6 @@ def compute_displacement_online( T_chunk, spatial_bin_edges_um, win_scale_um, - bin_s, mincorr_percentile, extra, extra_outputs, @@ -613,7 +643,10 @@ def compute_displacement_online( # below, t0 is start of prev chunk, t1 start of cur chunk, t2 end of cur t0, t1 = 0, T_chunk - traces0 = lfp_recording.get_traces(start_frame=t0, end_frame=t1) + if isinstance(lfp_recording, BaseRecording): + traces0 = lfp_recording.get_traces(start_frame=t0, end_frame=t1) + else: + traces0 = lfp_recording[:, t0:t1].T Ds0, Cs0, max_disp_um = xcorr_windows( traces0.T, @@ -653,7 +686,11 @@ def compute_displacement_online( ) for t1 in chunk_starts: t2 = min(T_total, t1 + T_chunk) - traces1 = lfp_recording.get_traces(start_frame=t1, end_frame=t2) + + if isinstance(lfp_recording, BaseRecording): + traces1 = lfp_recording.get_traces(start_frame=t1, end_frame=t2) + else: + traces1 = lfp_recording[:, t1:t2].T # cross-correlations between prev/cur chunks # these are T1, T0 shaped @@ -701,7 +738,7 @@ def compute_displacement_online( t0, t1 = t1, t2 traces0 = traces1 - return P_online + return P_online, extra def compute_displacement_simultaneous(