Make universal scale...

Author: cwindolf

src/dartsort/peel/matching.py

@@ -37,6 +37,7 @@ def __init__(
         featurization_pipeline,
         motion_est=None,
         pairwise_conv_db=None,
+        low_rank_templates=None,
         svd_compression_rank=10,
         coarse_objective=True,
         temporal_upsampling_factor=8,
@@ -78,6 +79,7 @@ def __init__(
         # main properties
         self.template_data = template_data
         self.pairwise_conv_db = pairwise_conv_db
+        self.low_rank_templates = low_rank_templates
         self.coarse_objective = coarse_objective
         self.temporal_upsampling_factor = temporal_upsampling_factor
         self.upsampling_peak_window_radius = upsampling_peak_window_radius
@@ -244,14 +246,15 @@ def build_template_data(
         if not have_groups:
             # the logic of coarsening is not needed.
             self.coarse_objective = False
-        low_rank_templates = template_util.svd_compress_templates(
-            template_data,
-            min_channel_amplitude=min_channel_amplitude,
-            rank=svd_compression_rank,
-        )
-        temporal_components = low_rank_templates.temporal_components.astype(dtype)
-        singular_values = low_rank_templates.singular_values.astype(dtype)
-        spatial_components = low_rank_templates.spatial_components.astype(dtype)
+        if self.low_rank_templates is None:
+            self.low_rank_templates = template_util.svd_compress_templates(
+                template_data,
+                min_channel_amplitude=min_channel_amplitude,
+                rank=svd_compression_rank,
+            )
+        temporal_components = self.low_rank_templates.temporal_components.astype(dtype)
+        singular_values = self.low_rank_templates.singular_values.astype(dtype)
+        spatial_components = self.low_rank_templates.spatial_components.astype(dtype)
         self.register_buffer("temporal_components", torch.tensor(temporal_components))
         self.register_buffer("singular_values", torch.tensor(singular_values))
         self.register_buffer("spatial_components", torch.tensor(spatial_components))
@@ -292,7 +295,7 @@ def build_template_data(
             self.obj_n_templates = spatial_components.shape[0]
         else:
             objective_temp_data = template_data
-            objective_low_rank_temps = low_rank_templates
+            objective_low_rank_temps = self.low_rank_templates
             self.objective_template_depths_um = self.registered_template_depths_um
             self.register_buffer(
                 "objective_temporal_components", self.temporal_components
@@ -302,6 +305,7 @@ def build_template_data(
                 "objective_spatial_components", self.spatial_components
             )
             self.obj_n_templates = self.n_templates
+
         self.handle_template_groups(
             objective_temp_data.unit_ids, self.template_data.unit_ids
         )
@@ -322,13 +326,14 @@ def build_template_data(
         chunk_centers_s = self.recording._recording_segments[0].sample_index_to_time(
             chunk_centers_samples
         )
+
         if self.pairwise_conv_db is None:
             self.pairwise_conv_db = CompressedPairwiseConv.from_template_data(
                 save_folder / "pconv.h5",
                 template_data=objective_temp_data,
                 low_rank_templates=objective_low_rank_temps,
                 template_data_b=template_data,
-                low_rank_templates_b=low_rank_templates,
+                low_rank_templates_b=self.low_rank_templates,
                 compressed_upsampled_temporal=compressed_upsampled_temporal,
                 chunk_time_centers_s=chunk_centers_s,
                 motion_est=self.motion_est,

src/dartsort/peel/universal.py

@@ -1,9 +1,11 @@
 import torch
+import numpy as np
 
 from ..util import universal_util, waveform_util
 from ..transform import WaveformPipeline
 from .matching import ObjectiveUpdateTemplateMatchingPeeler
 from ..templates.pairwise import SeparablePairwiseConv
+from ..templates.template_util import LowRankTemplates
 
 
 class UniversalTemplatesMatchingPeeler(ObjectiveUpdateTemplateMatchingPeeler):
@@ -30,7 +32,6 @@ def __init__(
         threshold=50.0,
         trough_offset_samples=42,
         spike_length_samples=121,
-        svd_compression_rank=10,
         amplitude_scaling_variance=100.0,
         amplitude_scaling_boundary=500.0,
         detection_threshold=6.0,
@@ -73,17 +74,34 @@ def __init__(
                 kmeanspp_initial="random",
             )
         )
+
+        Nf = len(footprints)
+        Ns = len(shapes)
+        shapes_ixd = torch.asarray(shapes)[None]
+        shapes_ixd = shapes_ixd.broadcast_to((Nf, Ns, *shapes.shape[1:]))
+        shapes_ixd = shapes_ixd.reshape(Nf * Ns, *shapes.shape[1:], 1)
+        footprints_ixd = torch.asarray(footprints)[:, None]
+        footprints_ixd = footprints_ixd.broadcast_to((Nf, Ns, *footprints.shape[1:]))
+        footprints_ixd = footprints_ixd.reshape(Nf * Ns, 1, *footprints.shape[1:])
+        low_rank_templates = LowRankTemplates(
+            temporal_components=shapes_ixd.numpy(),
+            singular_values=shapes_ixd.new_ones(Nf * Ns, 1).numpy(),
+            spatial_components=footprints_ixd.numpy(),
+            spike_counts_by_channel=np.broadcast_to(
+                np.atleast_2d([100]), (Nf * Ns, footprints.shape[1])
+            ),
+        )
         pairwise_conv_db = SeparablePairwiseConv(footprints, shapes)
         super().__init__(
             recording,
             template_data,
             channel_index,
             featurization_pipeline,
             pairwise_conv_db=pairwise_conv_db,
+            low_rank_templates=low_rank_templates,
             threshold=threshold,
             amplitude_scaling_variance=amplitude_scaling_variance,
             amplitude_scaling_boundary=amplitude_scaling_boundary,
-            svd_compression_rank=svd_compression_rank,
             # usual gizmos
             trough_offset_samples=trough_offset_samples,
             chunk_length_samples=chunk_length_samples,
@@ -94,10 +112,11 @@ def __init__(
             fit_sampling=fit_sampling,
             dtype=dtype,
             # matching params which don't need setting
-            min_channel_amplitude=1.0,
+            svd_compression_rank=1,
+            min_channel_amplitude=0.0,
             motion_est=None,
             coarse_approx_error_threshold=0.0,
-            conv_ignore_threshold=5.0,
+            conv_ignore_threshold=0.0,
             coarse_objective=True,
             temporal_upsampling_factor=1,
             refractory_radius_frames=10,
@@ -127,6 +146,7 @@ def from_config(
         )
         return cls(
             recording,
+            threshold=subtraction_config.universal_threshold,
             channel_index=channel_index,
             featurization_pipeline=featurization_pipeline,
             trough_offset_samples=trough_offset_samples,

src/dartsort/templates/pairwise.py

@@ -4,6 +4,7 @@
 import h5py
 import numpy as np
 import torch
+from torch._functorch.vmap import _as_tuple
 import torch.nn.functional as F
 
 from .pairwise_util import compressed_convolve_to_h5
@@ -282,9 +283,9 @@ def __init__(self, spatial_footprints, temporal_shapes):
         Let Nf = len(spatial_footprints), Ns = len(temporal_shapes). Then
         indexing is footprint-major, so that
 
-            template[i] = spatial_footprints[i // Nf] * temporal_shapes[i - Nf * (i // Nf)]
+            template[i] = spatial_footprints[i // Ns] * temporal_shapes[i - Ns * (i // Ns)]
 
-        Let f(i) = i // Nf and s(i) = i - Nf * (i // Nf). Then the channel-summed
+        Let f(i) = i // Ns and s(i) = i - Ns * (i // Ns). Then the channel-summed
         convolution of templates i and j is given by
 
             conv(t; i, j) = (
@@ -299,6 +300,7 @@ def __init__(self, spatial_footprints, temporal_shapes):
         self.register_buffer("spatial_footprints", torch.asarray(spatial_footprints))
         self.register_buffer("temporal_shapes", torch.asarray(temporal_shapes))
         self.Nf = len(spatial_footprints)
+        self.Ns, self.nt = temporal_shapes.shape
 
         # convolve all pairs of temporal shapes
         # i is data, j is filter
@@ -311,7 +313,8 @@ def __init__(self, spatial_footprints, temporal_shapes):
         # spatial component
         sdot = self.spatial_footprints @ self.spatial_footprints.T
         self.register_buffer("sdot", sdot)
-        self.tia = torch.arange(len(temporal_shapes))
+        self.overlap = (self.sdot > 0).cpu()
+        self.tia = torch.arange(self.Ns * self.Nf)
 
     def query(
         self,
@@ -329,24 +332,35 @@ def query(
         if device is not None and device != self.spatial_footprints.device:
             self.to(device)
         assert shifts_a is shifts_b is None
-        assert upsampling_indices_b is None
-        del return_zero_convs  # choose not to implement this
+        assert upsampling_indices_b is None or (upsampling_indices_b == 0).all()
+        assert not return_zero_convs  # choose not to implement this
         assert grid  # only this case here. can probably do the same above.
         if template_indices_a is None:
             template_indices_a = self.tia.to(template_indices_b)
 
-        f_i = template_indices_b // self.Nf
-        f_j = template_indices_a // self.Nf
-        s_i = template_indices_b - self.Nf * f_i
-        s_j = template_indices_a - self.Nf * f_j
+        f_i = template_indices_a // self.Ns
+        f_j = template_indices_b // self.Ns
+        keep_i, keep_j = self.overlap[f_i[:, None], f_j[None, :]].nonzero(as_tuple=True)
 
-        pconvs = (
-            self.sdot[f_i[:, None], f_j[None, :]]
-            * self.tconv[s_i[:, None], s_j[None, :]]
-        )
+        template_indices_a = template_indices_a[keep_i]
+        template_indices_b = template_indices_b[keep_j]
+        if scalings_b is not None:
+            scalings_b = scalings_b[keep_j]
+        if times_b is not None:
+            times_b = times_b[keep_j]
 
+        f_i = template_indices_a // self.Ns
+        f_j = template_indices_b // self.Ns
+        s_i = template_indices_a - self.Ns * f_i
+        s_j = template_indices_b - self.Ns * f_j
+
+        sdot = self.sdot[f_i, f_j]
+        assert sdot.ndim == 1
         if scalings_b is not None:
-            pconvs.mul_(scalings_b[:, None])
+            sdot = sdot * scalings_b
+        tconv = self.tconv[s_i, s_j]
+        assert tconv.shape == (len(sdot), 2 * self.nt - 1)
+        pconvs = sdot.unsqueeze(1) * tconv
 
         if times_b is not None:
             return template_indices_a, template_indices_b, times_b, pconvs

src/dartsort/util/internal_config.py

@@ -161,6 +161,7 @@ class SubtractionConfig:
     n_singlechan_templates: int = 10
     singlechan_alignment_padding_ms: float = 0.7
     use_universal_templates: bool = False
+    universal_threshold: float = 50.0
 
     # how will waveforms be denoised before subtraction?
     # users can also save waveforms/features during subtraction

src/dartsort/util/universal_util.py

@@ -133,7 +133,7 @@ def get_singlechan_waveforms(
 
 
 def spatial_footprint_bank(
-    geom, n_sigmas=5, min_template_size=10.0, max_distance=32.0, dx=32.0
+    geom, n_sigmas=5, min_template_size=10.0, max_distance=32.0, dx=32.0, eps=0.025
 ):
     # this is just a single shank version, since I don't plan to use
     # this in production. but it is copied from KS' code and can be
@@ -168,6 +168,8 @@ def spatial_footprint_bank(
     sigmas = min_template_size * (1 + np.arange(n_sigmas))
     spatial_profiles = np.exp(-distsq[:, None, :] / sigmas[:, None] ** 2)
     spatial_profiles = spatial_profiles.reshape(n_centers * n_sigmas, len(geom))
+    spatial_profiles /= np.linalg.norm(spatial_profiles, axis=1, keepdims=True)
+    spatial_profiles[spatial_profiles < eps] = 0.0
 
     return spatial_profiles
 
@@ -189,6 +191,7 @@ def singlechan_to_library(
         max_distance=max_distance,
         dx=dx,
     )
+    footprints = footprints.astype(str(singlechan_templates.dtype).split(".")[1])
     if torch.is_tensor(singlechan_templates):
         singlechan_templates = singlechan_templates.numpy(force=True)
     nf, nc = footprints.shape
@@ -202,8 +205,8 @@ def singlechan_to_library(
 
     template_data = TemplateData(
         templates,
-        unit_ids=np.arange(nsct * nt),
-        spike_counts=np.ones(nsct * nt, dtype=int),
+        unit_ids=np.arange(nsct * nf),
+        spike_counts=np.ones(nsct * nf, dtype=int),
     )
     return footprints, template_data