Nearby unit vis and high level functions

Author: cwindolf

src/dartsort/cluster/merge.py

@@ -18,6 +18,7 @@ def merge_templates(
     motion_est=None,
     max_shift_samples=20,
     superres_linkage=np.max,
+    sym_function=np.minimum,
     merge_distance_threshold=0.25,
     temporal_upsampling_factor=8,
     amplitude_scaling_variance=0.0,
@@ -68,6 +69,50 @@ def merge_templates(
             save_npz_name=template_npz_filename,
         )
 
+    units, dists, shifts, template_snrs = calculate_merge_distances(
+        template_data,
+        superres_linkage=superres_linkage,
+        sym_function=sym_function,
+        max_shift_samples=max_shift_samples,
+        temporal_upsampling_factor=temporal_upsampling_factor,
+        amplitude_scaling_variance=amplitude_scaling_variance,
+        amplitude_scaling_boundary=amplitude_scaling_boundary,
+        svd_compression_rank=svd_compression_rank,
+        min_channel_amplitude=min_channel_amplitude,
+        conv_batch_size=conv_batch_size,
+        units_batch_size=units_batch_size,
+        device=device,
+        n_jobs=n_jobs,
+        show_progress=show_progress,
+    )
+
+    # now run hierarchical clustering
+    return recluster(
+        sorting,
+        units,
+        dists,
+        shifts,
+        template_snrs,
+        merge_distance_threshold=merge_distance_threshold,
+    )
+
+
+def calculate_merge_distances(
+    template_data,
+    superres_linkage=np.max,
+    sym_function=np.minimum,
+    max_shift_samples=20,
+    temporal_upsampling_factor=8,
+    amplitude_scaling_variance=0.0,
+    amplitude_scaling_boundary=0.5,
+    svd_compression_rank=10,
+    min_channel_amplitude=0.0,
+    conv_batch_size=128,
+    units_batch_size=8,
+    device=None,
+    n_jobs=0,
+    show_progress=True,
+):
     # allocate distance + shift matrices. shifts[i,j] is trough[j]-trough[i].
     n_templates = template_data.templates.shape[0]
     sup_dists = np.full((n_templates, n_templates), np.inf)
@@ -116,15 +161,9 @@ def merge_templates(
             template_data.templates.ptp(1).max(1) / template_data.spike_counts
         )
 
-    # now run hierarchical clustering
-    return recluster(
-        sorting,
-        units,
-        dists,
-        shifts,
-        template_snrs,
-        merge_distance_threshold=merge_distance_threshold,
-    )
+    dists = sym_function(dists, dists.T)
+
+    return units, dists, shifts, template_snrs
 
 
 def recluster(
@@ -134,9 +173,7 @@ def recluster(
     shifts,
     template_snrs,
     merge_distance_threshold=0.25,
-    sym_function=np.minimum,
 ):
-    dists = sym_function(dists, dists.T)
 
     # upper triangle not including diagonal, aka condensed distance matrix in scipy
     pdist = dists[np.triu_indices(dists.shape[0], k=1)]

src/dartsort/config.py

@@ -242,3 +242,4 @@ class SplitMergeConfig:
 default_split_merge_config = SplitMergeConfig()
 coarse_template_config = TemplateConfig(superres_templates=False)
 default_matching_config = MatchingConfig()
+default_motion_estimation_config = MotionEstimationConfig()

src/dartsort/main.py

@@ -1,3 +1,4 @@
+from dataclasses import asdict
 from pathlib import Path
 
 from dartsort.cluster.initial import ensemble_chunks
@@ -6,6 +7,7 @@
 from dartsort.config import (default_clustering_config,
                              default_featurization_config,
                              default_matching_config,
+                             default_motion_estimation_config,
                              default_split_merge_config,
                              default_subtraction_config,
                              default_template_config)
@@ -14,6 +16,7 @@
 from dartsort.templates import TemplateData
 from dartsort.util.data_util import check_recording
 from dartsort.util.peel_util import run_peeler
+from dartsort.util.registration_util import estimate_motion
 
 
 def dartsort_from_config(
@@ -27,6 +30,7 @@ def dartsort(
     recording,
     output_directory,
     featurization_config=default_featurization_config,
+    motion_estimation_config=default_motion_estimation_config,
     subtraction_config=default_subtraction_config,
     matching_config=default_subtraction_config,
     template_config=default_template_config,
@@ -50,8 +54,9 @@ def dartsort(
         device=device,
     )
     if motion_est is None:
-        # TODO
-        motion_est = estimate_motion()
+        motion_est = estimate_motion(
+            recording, sorting, **asdict(motion_estimation_config)
+        )
     sorting = cluster(
         sub_h5,
         recording,

src/dartsort/util/analysis.py

@@ -9,7 +9,7 @@
 """
 from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import Optional
+from typing import Callable, Optional
 
 import h5py
 import numpy as np
@@ -19,7 +19,7 @@
 from sklearn.decomposition import PCA
 from spikeinterface.comparison import GroundTruthComparison
 
-from ..cluster import relocate
+from ..cluster import merge, relocate
 from ..templates import TemplateData
 from ..transform import WaveformPipeline
 from .data_util import DARTsortSorting
@@ -56,6 +56,11 @@ class DARTsortAnalysis:
     tpca_features_dataset = "collisioncleaned_tpca_features"
     template_indices_dataset = "collisioncleaned_tpca_features"
 
+    # configuration for analysis computations not included in above objects
+    device: Optional[str, torch.device] = None
+    merge_distance_templates_kind: str = "coarse"
+    merge_superres_linkage: Callable[[np.ndarray], float] = np.max
+
     # helper constructors
 
     @classmethod
@@ -110,7 +115,9 @@ def from_peeling_paths(
     def __post_init__(self):
         if self.featurization_pipeline is not None:
             assert not self.featurization_pipeline.needs_fit()
-        assert np.isin(self.template_data.unit_ids, np.unique(self.sorting.labels)).all()
+        assert np.isin(
+            self.template_data.unit_ids, np.unique(self.sorting.labels)
+        ).all()
 
         assert self.hdf5_path.exists()
         self.coarse_template_data = self.template_data.coarsen()
@@ -127,6 +134,7 @@ def __post_init__(self):
                 self.motion_est is not None
                 and self.template_data.registered_geom is not None
             )
+        assert self.coarse_template_data.unit_ids == self.unit_ids
 
         # cached hdf5 pointer
         self._h5 = None
@@ -145,6 +153,7 @@ def clear_cache(self):
         self._sklearn_tpca = None
         self._unit_ids = None
         self._spike_counts = None
+        self._merge_dist = None
         self._feats = {}
 
     def __getstate__(self):
@@ -209,6 +218,12 @@ def sklearn_tpca(self):
             self._sklearn_tpca = tpca_feature[0].to_sklearn()
         return self._sklearn_tpca
 
+    @property
+    def merge_dist(self):
+        if self._merge_dist is None:
+            self._merge_dist = self._calc_merge_dist()
+        return self._merge_dist
+
     # spike train helpers
 
     @property
@@ -236,6 +251,16 @@ def in_template(self, template_index):
     def unit_template_indices(self, unit_id):
         return np.flatnonzero(self.template_data.unit_ids == self.unit_id)
 
+    @property
+    def show_geom(self):
+        show_geom = self.template_data.registered_geom
+        if show_geom is None:
+            show_geom = self.recording.get_channel_locations()
+        return show_geom
+
+    def show_channel_index(self, radius_um=50):
+        return make_channel_index(self.show_geom, radius_um)
+
     # spike feature loading methods
 
     def named_feature(self, name, which=slice(None)):
@@ -330,7 +355,12 @@ def unit_raw_waveforms(
         if not self.shifting:
             return which, waveforms
 
-        waveforms, max_chan, show_geom, show_channel_index = self.unit_shift_or_relocate_channels(
+        (
+            waveforms,
+            max_chan,
+            show_geom,
+            show_channel_index,
+        ) = self.unit_shift_or_relocate_channels(
             unit_id,
             which,
             waveforms,
@@ -367,7 +397,12 @@ def unit_tpca_waveforms(
         t = waveforms.shape[1]
         waveforms = waveforms.reshape(n, c, t).transpose(0, 2, 1)
 
-        waveforms, max_chan, show_geom, show_channel_index = self.unit_shift_or_relocate_channels(
+        (
+            waveforms,
+            max_chan,
+            show_geom,
+            show_channel_index,
+        ) = self.unit_shift_or_relocate_channels(
             unit_id,
             which,
             waveforms,
@@ -378,9 +413,21 @@ def unit_tpca_waveforms(
         return which, waveforms, max_chan, show_geom, show_channel_index
 
     def unit_pca_features(
-        self, unit_id, relocated=True, rank=2, pca_radius_um=75, random_seed=0, max_count=500
+        self,
+        unit_id,
+        relocated=True,
+        rank=2,
+        pca_radius_um=75,
+        random_seed=0,
+        max_count=500,
     ):
-        which, waveforms, max_chan, show_geom, show_channel_index = self.unit_tpca_waveforms(
+        (
+            which,
+            waveforms,
+            max_chan,
+            show_geom,
+            show_channel_index,
+        ) = self.unit_tpca_waveforms(
             unit_id,
             relocated=relocated,
             show_radius_um=pca_radius_um,
@@ -439,7 +486,9 @@ def unit_shift_or_relocate_channels(
         show_channel_index = make_channel_index(show_geom, show_radius_um)
         show_chans = show_channel_index[max_chan]
         show_chans = show_chans[show_chans < len(show_geom)]
-        show_channel_index = np.broadcast_to(show_chans[None], (len(show_geom), show_chans.size))
+        show_channel_index = np.broadcast_to(
+            show_chans[None], (len(show_geom), show_chans.size)
+        )
 
         if not self.shifting:
             return waveforms, max_chan, show_geom, show_channel_index
@@ -478,6 +527,34 @@ def unit_shift_or_relocate_channels(
 
         return waveforms, max_chan, show_geom, show_channel_index
 
+    def nearby_coarse_templates(self, unit_id, n_neighbors=5):
+        unit_ix = np.searchsorted(self.unit_ids, unit_id)
+        unit_dists = self.merge_dist[unit_ix]
+        distance_order = np.argsort(unit_dists)
+        assert distance_order[0] == unit_ix
+        neighbor_ixs = distance_order[:n_neighbors]
+        neighbor_ids = self.unit_ids[:n_neighbors]
+        neighbor_dists = self.merge_dist[neighbor_ixs[:, None], neighbor_ixs[None, :]]
+        neighbor_coarse_templates = self.coarse_template_data.templates[neighbor_ixs]
+        return neighbor_ids, neighbor_dists, neighbor_coarse_templates
+
+    # computation
+
+    def _calc_merge_dist(self):
+        """Compute the merge distance matrix"""
+        merge_td = self.template_data
+        if self.merge_distance_templates_kind == "coarse":
+            merge_td = self.coarse_template_data
+
+        units, dists, shifts, template_snrs = merge.calculate_merge_distances(
+            merge_td,
+            superres_linkage=self.merge_superres_linkage,
+            device=self.device,
+            n_jobs=1,
+        )
+        assert np.array_equal(units, self.unit_ids)
+        self._merge_dist = dists
+
 
 @dataclass
 class DARTsortGroundTruthComparison:

src/dartsort/util/registration_util.py

@@ -1,14 +1,65 @@
+from typing import Optional
+
+import numpy as np
+
 try:
     from dredge import dredge_ap
+
     have_dredge = True
 except ImportError:
     have_dredge = False
     pass
 
 
-def estimate_motion(recording, sorting, motion_estimation_config=None, localizations_dataset_name="point_source_localizations"):
-    if not motion_estimation_config.do_motion_estimation:
+def estimate_motion(
+    recording,
+    sorting,
+    do_motion_estimation=True,
+    probe_boundary_padding_um=100.0,
+    spatial_bin_length_um: float = 1.0,
+    temporal_bin_length_s: float = 1.0,
+    window_step_um: float = 400.0,
+    window_scale_um: float = 450.0,
+    window_margin_um: Optional[float] = None,
+    max_dt_s: float = 0.1,
+    max_disp_um: Optional[float] = None,
+    localizations_dataset_name="point_source_localizations",
+    amplitudes_dataset_name="denoised_ptp_amplitudes",
+):
+    if not do_motion_estimation:
         return None
 
     if not have_dredge:
         raise ValueError("Please install DREDge to use motion estimation.")
+
+    x = getattr(sorting, localizations_dataset_name)[:, 0]
+    z = getattr(sorting, localizations_dataset_name)[:, 1]
+    geom = recording.get_channel_locations()
+    xmin = geom[:, 0].min() - probe_boundary_padding_um
+    xmax = geom[:, 0].max() + probe_boundary_padding_um
+    zmin = geom[:, 1].min() - probe_boundary_padding_um
+    zmax = geom[:, 1].max() + probe_boundary_padding_um
+    xvalid = x == np.clip(x, xmin, xmax)
+    zvalid = z == np.clip(z, zmin, zmax)
+    valid = np.flatnonzero(xvalid & zvalid)
+
+    # features for registration
+    z = z[valid]
+    t_s = sorting.times_seconds[valid]
+    a = getattr(sorting, amplitudes_dataset_name)[valid]
+
+    # run registration
+    motion_est, info = dredge_ap.register(
+        a,
+        z,
+        t_s,
+        window_step_um=window_step_um,
+        bin_um=spatial_bin_length_um,
+        bin_s=temporal_bin_length_s,
+        window_scale_um=window_scale_um,
+        window_margin_um=window_margin_um,
+        max_disp_um=max_disp_um,
+        max_dt_s=max_dt_s,
+    )
+
+    return motion_est

src/dartsort/vis/analysis_plots.py

@@ -0,0 +1,2 @@
+def distance_matrix_dendro():
+    pass