generate.py - add extra outputs.

Author: JoeZiminski

src/spikeinterface/core/generate.py

@@ -95,6 +95,7 @@ def generate_sorting(
     add_spikes_on_borders=False,
     num_spikes_per_border=3,
     border_size_samples=20,
+    extra_outputs=False,
     seed=None,
 ):
     """
@@ -135,10 +136,14 @@ def generate_sorting(
     num_segments = len(durations)
     unit_ids = np.arange(num_units)
 
+    extra_outputs_dict = {
+        "firing_rates": [],
+    }
+
     spikes = []
     for segment_index in range(num_segments):
         num_samples = int(sampling_frequency * durations[segment_index])
-        samples, labels = synthesize_poisson_spike_vector(
+        samples, labels, firing_rates_array = synthesize_poisson_spike_vector(
             num_units=num_units,
             sampling_frequency=sampling_frequency,
             duration=durations[segment_index],
@@ -172,12 +177,17 @@ def generate_sorting(
             )
             spikes.append(spikes_on_borders)
 
+        extra_outputs_dict["firing_rates"].append(firing_rates_array)
+
     spikes = np.concatenate(spikes)
     spikes = spikes[np.lexsort((spikes["sample_index"], spikes["segment_index"]))]
 
     sorting = NumpySorting(spikes, sampling_frequency, unit_ids)
 
-    return sorting
+    if extra_outputs:
+        return sorting, extra_outputs_dict
+    else:
+        return sorting
 
 
 def add_synchrony_to_sorting(sorting, sync_event_ratio=0.3, seed=None):
@@ -776,7 +786,7 @@ def synthesize_poisson_spike_vector(
     unit_indices = unit_indices[sort_indices]
     spike_frames = spike_frames[sort_indices]
 
-    return spike_frames, unit_indices
+    return spike_frames, unit_indices, firing_rates
 
 
 def synthesize_random_firings(