vis/eval: various fixes

Author: cwindolf

src/dartsort/__init__.py

@@ -76,7 +76,12 @@
     get_logger,
     set_log_level,
 )
-from .util.motion import MotionInfo, get_motion_info, try_load_motion_info
+from .util.motion import (
+    MotionInfo,
+    detect_for_motion,
+    get_motion_info,
+    try_load_motion_info,
+)
 from .util.noise_util import EmbeddedNoise
 from .util.preprocess_util import preprocess
 from .util.py_util import databag, ensure_path

src/dartsort/clustering/agglomerate.py

@@ -1156,7 +1156,7 @@ def _calc_coentropy(
         u = cands[i]
         q = resps[i]
         log_q = np.log(q)
-        np.nan_to_num(log_q, copy=False, neginf=0.0)
+        np.nan_to_num(log_q, copy=False, neginf=0.0, posinf=np.inf)
         dh = q * log_q
 
         ui0 = u[0]

src/dartsort/clustering/density.py

@@ -188,7 +188,7 @@ def get_smoothed_density_ratio(
     else:
         dens_ = dens[0]
         dens_ /= dens[1]
-        np.nan_to_num(dens_, out=dens_)  # type: ignore
+        np.nan_to_num(dens_, out=dens_, posinf=np.inf, neginf=-np.inf)  # type: ignore
         return dens_
 
 
@@ -407,7 +407,7 @@ def kdt_density(
             desc=f"KDTdens[{n_jobs}]",
         ):
             density[i0:i1] = dens
-    np.nan_to_num(density, copy=False)
+    np.nan_to_num(density, posinf=np.inf, neginf=-np.inf, copy=False)
     return density
 
 

src/dartsort/clustering/merge.py

@@ -358,21 +358,47 @@ def get_deconv_resid_decrease_iter(
     )
 
 
-def combine_templates(template_data_a, template_data_b):
+def combine_templates(template_data_a: TemplateData, template_data_b: TemplateData):
     rgeom = template_data_a.registered_geom
     if rgeom is not None:
+        assert template_data_b.registered_geom is not None
         if not np.array_equal(rgeom, template_data_b.registered_geom):
             raise ValueError(
                 f"Template data had different registered geoms: "
                 f"{template_data_a.registered_geom=} {template_data_b.registered_geom=}"
             )
 
+    ta = template_data_a.templates
+    tb = template_data_b.templates
+    assert ta.shape[2] == tb.shape[2]
+    if ta.shape[1] > tb.shape[1]:
+        i0 = (
+            template_data_a.trough_offset_samples
+            - template_data_b.trough_offset_samples
+        )
+        assert i0 >= 0
+        i1 = i0 + template_data_b.spike_length_samples
+        ta = ta[:, i0:i1]
+        trough_offset_samples = template_data_b.trough_offset_samples
+        sampling_frequency = template_data_b.sampling_frequency
+    elif tb.shape[1] > ta.shape[1]:
+        i0 = (
+            template_data_b.trough_offset_samples
+            - template_data_a.trough_offset_samples
+        )
+        assert i0 >= 0
+        i1 = i0 + template_data_a.spike_length_samples
+        tb = tb[:, i0:i1]
+        trough_offset_samples = template_data_a.trough_offset_samples
+        sampling_frequency = template_data_a.sampling_frequency
+    else:
+        trough_offset_samples = template_data_a.trough_offset_samples
+        sampling_frequency = template_data_a.sampling_frequency
+
     ids_a = template_data_a.unit_ids
     ids_b = template_data_b.unit_ids + ids_a.max() + 1
     unit_ids = np.concatenate((ids_a, ids_b))
-    templates = np.concatenate(
-        (template_data_a.templates, template_data_b.templates), axis=0
-    )
+    templates = np.concatenate((ta, tb), axis=0)
     spike_counts = np.concatenate(
         (template_data_a.spike_counts, template_data_b.spike_counts)
     )
@@ -389,8 +415,8 @@ def combine_templates(template_data_a, template_data_b):
         spike_counts=spike_counts,
         registered_geom=rgeom,
         spike_counts_by_channel=spike_counts_by_channel,
-        trough_offset_samples=template_data_a.trough_offset_samples,
-        sampling_frequency=template_data_a.sampling_frequency,
+        trough_offset_samples=trough_offset_samples,
+        sampling_frequency=sampling_frequency,
     )
 
     cross_mask = np.logical_and(

src/dartsort/evaluate/analysis.py

@@ -136,7 +136,7 @@ def from_sorting(
                 tpca = get_tpca(
                     sorting, featurization_pipeline_pt=featurization_pipeline_pt
                 )
-            except ValueError:
+            except (ValueError, FileNotFoundError):
                 tpca = None
         else:
             tpca = None

src/dartsort/evaluate/comparison.py

@@ -259,7 +259,9 @@ def _unit_info_dataframe(self, force_distances=False, perf_only=False):
         except ValueError:
             pass
         if self.has_templates and (force_distances or self.compute_distances):
-            dist = np.nan_to_num(self.template_distances, nan=np.inf).min(axis=1)
+            dist = np.nan_to_num(
+                self.template_distances, nan=np.inf, posinf=np.inf, neginf=-np.inf
+            ).min(axis=1)
             df["min_temp_dist"] = dist
         rec = []
         for uid in df.index:
@@ -334,7 +336,9 @@ def unit_matched_misalignment_rms(self):
                 if udt.size:
                     match_dt_rms[j] = np.sqrt(np.square(udt).mean())
             except ValueError as e:
-                warnings.warn(f"ValueError in misalignment. SI matching bug. {e=}", stacklevel=2)
+                warnings.warn(
+                    f"ValueError in misalignment. SI matching bug. {e=}", stacklevel=2
+                )
         return match_dt_rms
 
     def matched_misalignment(self, gt_unit_id):
@@ -477,7 +481,8 @@ def get_spikes_by_category(self, gt_unit, tested_unit=None):
                 warnings.warn(
                     f"Strange match sizes for {gt_unit=} {tested_unit=}: "
                     f"{matched_gt_indices.shape=} {matched_tested_indices.shape=} "
-                    f"{matched_tested_mask.sum()=}", stacklevel=2
+                    f"{matched_tested_mask.sum()=}",
+                    stacklevel=2,
                 )
         else:
             matched_tested_indices = np.zeros(shape=(0,), dtype=np.int64)
@@ -741,7 +746,8 @@ def __init__(
         self.other_analyses = other_analyses
         self.other_names = [oa.name for oa in other_analyses]
         self.other_templates = [
-            (oa.name or f"Test{c}") for oa, c in zip(other_analyses, self.default_ids, strict=True)
+            (oa.name or f"Test{c}")
+            for oa, c in zip(other_analyses, self.default_ids, strict=True)
         ]
         self.other_sortings = [oa.sorting for oa in other_analyses]
 

src/dartsort/evaluate/hybridkit.py

@@ -1,5 +1,6 @@
 import json
 from copy import deepcopy
+from dataclasses import replace
 from pathlib import Path
 from typing import Any, Literal
 
@@ -14,6 +15,7 @@
 from ..util.job_util import ensure_computation_config
 from ..util.motion import MotionInfo, resample_motion
 from ..util.py_util import databag
+from .analysis import DARTsortAnalysis
 from .sim_template_tools import TemplateLibrarySimulator, get_template_simulator
 from .simlib import InjectSpikesPreprocessor
 
@@ -28,6 +30,28 @@ class HybridDataset:
     motion: MotionInfo
     metadata: dict[str, Any]
 
+    def gt_analysis(
+        self, trough_offset: int = 42, spike_length: int = 121
+    ) -> DARTsortAnalysis:
+        if spike_length != self.gt_templates.spike_length_samples:
+            i0 = self.gt_templates.trough_offset_samples - trough_offset
+            assert i0 >= 0
+            template_data = replace(
+                self.gt_templates,
+                templates=self.gt_templates.templates[:, i0 : i0 + spike_length],
+                trough_offset_samples=trough_offset,
+            )
+        else:
+            assert trough_offset == self.gt_templates.spike_length_samples
+            template_data = self.gt_templates
+        return DARTsortAnalysis.from_sorting(
+            recording=self.recording,
+            sorting=self.gt_sorting,
+            template_data=template_data,
+            motion=self.motion,
+            name="GT",
+        )
+
 
 def load_hybrid_recording(folder: str | Path) -> HybridDataset | None:
     folder = ensure_path(folder)

src/dartsort/templates/templates.py

@@ -1,7 +1,7 @@
 import gc
 from dataclasses import replace
 from pathlib import Path
-from typing import ClassVar
+from typing import ClassVar, Self
 
 import numpy as np
 import torch
@@ -83,19 +83,19 @@ def __post_init__(self):
                 assert v.shape[0] == ntemp
 
     @property
-    def spike_length_samples(self):
+    def spike_length_samples(self) -> int:
         return self.templates.shape[1]
 
-    def snrs_by_channel(self):
+    def snrs_by_channel(self) -> np.ndarray:
         amp_vecs = np.nan_to_num(np.ptp(self.templates, axis=1), nan=-np.inf)
         if self.spike_counts_by_channel is not None:
             amp_vecs *= np.sqrt(self.spike_counts_by_channel)
         return amp_vecs
 
-    def main_channels(self):
+    def main_channels(self) -> np.ndarray:
         return self.snrs_by_channel().argmax(axis=1)
 
-    def template_locations(self, mode="channel", radius=100.0):
+    def template_locations(self, mode="channel", radius=100.0) -> np.ndarray:
         assert mode in ("localization", "channel")
 
         if mode == "channel":
@@ -112,11 +112,11 @@ def template_locations(self, mode="channel", radius=100.0):
         rdepths = np.c_[rdepths["x"], rdepths["z_abs"]]
         return rdepths
 
-    def registered_depths_um(self, mode="channel", radius=100.0):
+    def registered_depths_um(self, mode="channel", radius=100.0) -> np.ndarray:
         return self.template_locations(mode=mode, radius=radius)[:, 1]
 
     @classmethod
-    def from_npz(cls, npz_path):
+    def from_npz(cls, npz_path) -> Self:
         with np.load(npz_path, allow_pickle=True) as data:
             data = dict(**data)
             data["whiten_strategy"] = str(data["whiten_strategy"])
@@ -182,7 +182,7 @@ def to_npz(self, npz_path):
                 to_save[f"__prop_{k}"] = p
         np.savez(npz_path, **to_save)  # type: ignore
 
-    def __getitem__(self, subset):
+    def __getitem__(self, subset) -> Self:
         if not np.array_equal(self.unit_ids, np.arange(len(self.unit_ids))):
             subset_ixs = np.searchsorted(self.unit_ids, subset, side="right") - 1
             matched = self.unit_ids[subset_ixs] == subset
@@ -212,7 +212,7 @@ def __getitem__(self, subset):
             whiten_strategy=self.whiten_strategy,
         )
 
-    def coarsen(self):
+    def coarsen(self) -> Self:
         """Weighted average all templates that share a unit id."""
         # update templates
         unit_ids_unique, flat_ids = np.unique(self.unit_ids, return_inverse=True)
@@ -231,10 +231,10 @@ def coarsen(self):
             tsvd=self.tsvd,
         )
 
-    def unit_mask(self, unit_id):
+    def unit_mask(self, unit_id) -> np.ndarray:
         return np.isin(self.unit_ids, unit_id)
 
-    def unit_templates(self, unit_id):
+    def unit_templates(self, unit_id) -> np.ndarray:
         return self.templates[self.unit_mask(unit_id)]
 
     def __init_subclass__(cls):
@@ -258,7 +258,7 @@ def from_config(
         featurization_basis=None,
         computation_cfg: ComputationConfig | None = None,
         show_progress: bool = True,
-    ) -> "TemplateData":
+    ) -> Self:
         # load if saved already and not overwriting
         if save_folder is not None:
             save_folder = Path(save_folder)
@@ -326,7 +326,7 @@ def _from_config(
         whitener: Whitener | None = None,
         tsvd=None,
         computation_cfg: ComputationConfig | None = None,
-    ) -> "TemplateData":
+    ) -> Self:
         raise NotImplementedError
 
 
@@ -353,4 +353,4 @@ def _try_reload_svd(
         logger.dartsortdebug(f"Reloading TSVD from {tnpz}")
     else:
         logger.dartsortdebug(f"No TSVD to reload in {tnpz}")
-    return tsvd
+    return tsvd  # ty: ignore[invalid-return-type]