Add spike_ampltiudes to units table

Author: bjhardcastle

src/npc_ephys/spikeinterface.py

@@ -447,8 +447,9 @@ def decoder_label(self, probe: str) -> npt.NDArray[np.str_]:
         )
 
     @functools.cache
-    def spike_amplitudes(self, probe: str) -> npt.NDArray[np.floating]:
-        return np.load(
+    def spike_amplitudes(self, probe: str) -> tuple[npt.NDArray[np.floating], ...]:
+        """array of amplitudes for each unit in order of unique('unit_indexes')"""
+        spike_amplitudes = np.load(
             io.BytesIO(
                 self.get_correct_path(
                     self.postprocessed(probe),
@@ -457,7 +458,14 @@ def spike_amplitudes(self, probe: str) -> npt.NDArray[np.floating]:
                 ).read_bytes()
             )
         )
-
+        unit_indexes = self.unit_indexes(probe)
+        spike_amplitudes_by_unit: list[npt.NDArray[np.floating]] = []
+        for index in sorted(np.unique(unit_indexes)):
+            spike_amplitudes_by_unit.append(
+                spike_amplitudes[np.where(unit_indexes == index)[0]]
+            )
+        return tuple(spike_amplitudes_by_unit)
+    
     @functools.cache
     def unit_locations(self, probe: str) -> npt.NDArray[np.floating]:
         return np.load(

src/npc_ephys/units.py

@@ -224,6 +224,8 @@ def _device_helper(
         [len(unit) for unit in units_x_spike_times],
     ), "Mismatch between rows in spike_times and metrics.csv"
     df_device_metrics["spike_times"] = units_x_spike_times
+    df_device_metrics["spike_amplitudes"] = spike_interface_data.spike_amplitudes(electrode_group_name)
+    assert all(len(df_device_metrics["spike_amplitudes"][i]) == len(spike_times) for i, spike_times in enumerate(units_x_spike_times)), "Mismatch between spike_times and spike_amplitudes"
 
     return df_device_metrics