WIP: Allow Dredge AP to use online motion resolution method

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

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
 
@@ -322,66 +324,67 @@ 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(
+    full_xcorr_kw = dict(
+        rigid=rigid,
+        bin_um=bin_um,
+        max_disp_um=max_disp_um,
+        progress_bar=False,
+        device=device,
+        **xcorr_kw,
+    )
+
+    if resolution_mode == "simultaneous":
+        displacement, extra = compute_displacement_simultaneous(
             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,
+            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:
-        Ds, Cs, max_disp_um = precomputed_D_C_maxdisp
-
-    # turn Cs into weights
-    Us, wextra = weight_correlation_matrix(
-        Ds,
-        Cs,
-        windows,
-        raster,
-        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,
-    )
-    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
+        raise ValueError(f"No resolution mode called {resolution_mode}")
 
     if extra_outputs:
         extra["windows"] = windows
         extra["window_centers"] = window_centers
         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,18 +586,76 @@ def dredge_online_lfp(
         zero_threshold=1e-5,
     )
 
-    B = len(windows)
-
     if extra_outputs:
         extra = dict(window_centers=window_centers, windows=windows)
 
+    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)
+
+    if extra_outputs:
+        return motion, extra
+    else:
+        return motion
+
+
+def compute_displacement_online(
+    lfp_recording,
+    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,
+):
+
+    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
+
+    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,
+        windows,
+        spatial_bin_edges_um,
+        win_scale_um,
+        **full_xcorr_kw,
+    )
+    mincorr = weights_kw["mincorr"]
     Ss0, mincorr0 = threshold_correlation_matrix(
         Cs0,
         mincorr=mincorr,
@@ -615,30 +676,35 @@ def dredge_online_lfp(
 
     # -- 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)
-        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
         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,
@@ -672,12 +738,77 @@ def dredge_online_lfp(
         t0, t1 = t1, t2
         traces0 = traces1
 
-    motion = Motion([P_online.T], [lfp_recording.get_times(0)], window_centers, direction=direction)
+    return P_online, extra
 
-    if extra_outputs:
-        return motion, extra
+
+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:
-        return motion
+        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)
@@ -1289,7 +1420,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,
@@ -1335,7 +1465,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,
@@ -1370,7 +1500,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,
     )
@@ -1389,7 +1519,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,