wip on subtract and localize numpy

requirements.txt

@@ -1,5 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 click
+dredge
 h5py
 hdbscan
 matplotlib

src/spike_psvae/chunk_features.py

@@ -76,7 +76,7 @@ def handle_which_wfs(self, subtracted_wfs, cleaned_wfs, denoised_wfs):
                 f"which_waveforms={self.which_waveforms} not in ('subtracted', 'cleaned', denoised)"
             )
         if not self.tensor_ok and torch.is_tensor(wfs):
-            wfs = wfs.cpu().numpy()
+            wfs = wfs.cpu().detach().numpy()
         return wfs
 
 

src/spike_psvae/denoise.py

@@ -11,6 +11,7 @@
 from sklearn.decomposition import PCA
 from torch import nn
 from tqdm.auto import trange
+import dartsort
 
 try:
     from .denoise_temporal_decrease import (
@@ -20,9 +21,7 @@
     pass
 
 
-pretrained_path = (
-    Path(__file__).parent.parent.parent / "pretrained/single_chan_denoiser.pt"
-)
+pretrained_path = Path(dartsort.__file__).parent.joinpath('pretrained', 'single_chan_denoiser.pt')
 
 
 class SingleChanDenoiser(nn.Module):
@@ -775,7 +774,7 @@ def enforce_decrease_shells(
     # compute original ptps and allocate storage for decreasing ones
     is_torch = False
     if torch.is_tensor(waveforms):
-        orig_ptps = (waveforms.max(dim=1).values - waveforms.min(dim=1).values).cpu().numpy()
+        orig_ptps = (waveforms.max(dim=1).values - waveforms.min(dim=1).values).cpu().detach().numpy()
         is_torch = True
     else:
         orig_ptps = waveforms.ptp(axis=1)

src/spike_psvae/subtract.py

@@ -1730,6 +1730,7 @@ def subtract_and_localize_numpy(
     radial_parents=None,
     tpca=None,
     device=None,
+    enfdec=None,
     probe="np1",
     trough_offset=42,
     spike_length_samples=121,
@@ -1772,6 +1773,7 @@ def subtract_and_localize_numpy(
             residual,
             threshold,
             radial_parents,
+            enfdec,  # enfdec
             tpca,
             dedup_channel_index,
             extract_channel_index,
@@ -1852,4 +1854,5 @@ def subtract_and_localize_numpy(
         ],
         columns=["sample", "trace", "x", "y", "z", "alpha"],
     )
-    return df_localisation, cleaned_wfs.to("cpu")
+    np_waveforms = cleaned_wfs.to("cpu").detach().numpy()
+    return df_localisation, np_waveforms