transform: be lazier with the featurization in fit()

Author: cwindolf

src/dartsort/transform/amplitudes.py

@@ -10,6 +10,7 @@ class AmplitudeFeatures(BaseWaveformFeaturizer):
     """Extract spike amplitudes."""
 
     is_multi = True
+    featurize_in_fit = True
 
     def __init__(
         self,
@@ -122,6 +123,7 @@ def transform(self, waveforms, **unused):
 
 class AmplitudeVector(BaseWaveformFeaturizer):
     default_name = "amplitude_vectors"
+    featurize_in_fit = True
 
     def __init__(
         self,
@@ -153,6 +155,7 @@ def transform(self, waveforms, **unused):
 
 class MaxAmplitude(BaseWaveformFeaturizer):
     default_name = "amplitudes"
+    featurize_in_fit = True
     shape = ()
 
     def __init__(

src/dartsort/transform/mixture_classifier.py

@@ -12,9 +12,10 @@
 """
 
 import sys
+from collections.abc import Sequence
 from dataclasses import replace
 from math import fabs
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch

src/dartsort/transform/pipeline.py

@@ -1,8 +1,8 @@
 """A class which manages pipelines of denoisers and featurizers"""
 
+from collections.abc import Sequence
 from copy import deepcopy
 from pathlib import Path
-from typing import Sequence
 
 import torch
 from spikeinterface.core import BaseRecording
@@ -341,7 +341,7 @@ def fit(
             if transformer.is_featurizer and transformer.is_denoiser:
                 waveforms, new_features = transformer(**features)
                 features.update(waveforms=waveforms, **new_features)
-            elif transformer.is_featurizer:
+            elif transformer.is_featurizer and transformer.featurize_in_fit:
                 assert isinstance(transformer, BaseWaveformFeaturizer)
                 features.update(transformer.transform(**features))
             elif transformer.is_denoiser:
@@ -372,7 +372,7 @@ def precompute(self):
     def transform_to_disk(
         self,
         hdf5_filename: str | Path,
-        waveforms_dataset_name: str | None = "waveforms",
+        waveforms_dataset_name: str = "waveforms",
         other_dset_names: Sequence[str] | None = None,
         start_index: int | None = None,
         up_to_index: int | None = None,
@@ -409,7 +409,6 @@ def transform_to_disk(
         with File(hdf5_filename, mode="r+", libver="latest", locking=False) as h5:
             if all(ds.name in h5 for ds in datasets):
                 return
-            wfs = h5[waveforms_dataset_name] if waveforms_dataset_name else None
             other_dsets = {od: h5[od] for od in other_dset_names}
             outs = {
                 ds.name: h5.create_dataset(
@@ -418,7 +417,9 @@ def transform_to_disk(
                 for ds in datasets
             }
             for sli, chk in yield_chunks(
-                wfs, desc_prefix="Transform to disk", show_progress=False
+                h5[waveforms_dataset_name],
+                desc_prefix="Transform to disk",
+                show_progress=False,
             ):
                 chk_fp = {k: v[sli].to(device=dev) for k, v in fixed_properties.items()}
                 other_fp = {

src/dartsort/transform/reduction.py

@@ -48,18 +48,15 @@ def __init__(
             pfx = ""
         else:
             pfx = name_prefix + "_"
-        if not self.online:
-            names = [f"{pfx}{_my_name}"]
-        else:
-            names = []
+        names = [] if self.online else [f"{pfx}{_my_name}"]
         super().__init__(geom=geom, channel_index=channel_index, name=names)
         self.name_prefix = name_prefix
         self.reduction = reduction
         self.n_units = n_units
         self.feature_dim = feature_dim
         self.output_channels = output_channels
-        self.shape = [(self.feature_dim, output_channels)]
-        self.dtype = [dtype]
+        self.shape = [] if self.online else [(self.feature_dim, output_channels)]
+        self.dtype = [] if self.online else [dtype]
         self.with_raw_std_dev = with_raw_std_dev
         self._initialize((self.feature_dim, output_channels))
 

src/dartsort/transform/transform_base.py

@@ -1,5 +1,6 @@
+from collections.abc import Iterable
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterable, Self
+from typing import TYPE_CHECKING, Any, Self
 
 import torch
 from spikeinterface.core import BaseRecording
@@ -25,6 +26,7 @@ class BaseWaveformModule(BModule):
     needs_residual = False
     fits_from_disk = False
     needs_more_features = False
+    featurize_in_fit = False
 
     def __init__(
         self,
@@ -231,7 +233,7 @@ def spike_datasets(self, force_save: bool = False) -> Iterable[SpikeDataset]:
                     shape_per_spike=s,
                     dtype=str(d).split(".")[1],
                 )
-                for n, s, d in zip(self.name, self.shape, self.dtype)
+                for n, s, d in zip(self.name, self.shape, self.dtype, strict=True)
             ]
             return datasets
         else:
@@ -321,6 +323,7 @@ def forward(self, waveforms, **spike_data):
 
 class Waveform(BaseWaveformFeaturizer):
     default_name = "waveforms"
+    featurize_in_fit = True
 
     def __init__(
         self,

tests/test_clustering.py

@@ -86,6 +86,7 @@ def test_clustering(simulations, sim_name, featkw, cluskw):
     if cluskw["cluster_strategy"] == "density_peaks_uhdversion":
         if not featkw["use_amplitude"]:
             return
+        featkw = featkw | {"need_xyza": True}
     sim = simulations[sim_name]
     recording = sim["recording"]
     sorting = sim["sorting"]