compute_synchrony_metrics update

src/spikeinterface/core/generate.py

@@ -616,6 +616,8 @@ def synthesize_poisson_spike_vector(
     refractory_period_ms=4.0,
     firing_rates=3.0,
     seed=0,
+    insertions=None,
+    insertions_replace=True,
 ):
     """
     Generate random spike frames for neuronal units using a Poisson process.
@@ -640,6 +642,11 @@ def synthesize_poisson_spike_vector(
         each element being the firing rate for one unit
     seed : int, default: 0
         Seed for random number generator
+    insertions: array-like
+        List of manually inserted spikes, each spike formatted as [sample index, unit index]. Intended for
+        use in testing.
+    insertions_replace: bool
+        If True, randomly replace generated spikes. If False, add to generated spikes.
 
     Returns
     -------
@@ -709,6 +716,12 @@ def synthesize_poisson_spike_vector(
     spike_frames[:num_correct_frames] = spike_frames[mask]  # Avoids a malloc
     unit_indices = unit_indices[mask]
 
+    # Add insertions
+    if insertions is not None:
+        spike_frames, unit_indices = _add_insertions(
+            [spike_frames, unit_indices], insertions=insertions, insertions_replace=insertions_replace
+        )
+
     # Sort globaly
     spike_frames = spike_frames[:num_correct_frames]
     sort_indices = np.argsort(spike_frames, kind="stable")  # I profiled the different kinds, this is the fastest.
@@ -727,6 +740,8 @@ def synthesize_random_firings(
     firing_rates=3.0,
     add_shift_shuffle=False,
     seed=None,
+    insertions=None,
+    insertions_replace=False,
 ):
     """ "
     Generate some spiketrain with random firing for one segment.
@@ -748,6 +763,11 @@ def synthesize_random_firings(
         Optionally add a small shuffle on half of the spikes to make the autocorrelogram less flat.
     seed: int, default: None
         seed for the generator
+    insertions: array-like
+        List of manually inserted spikes, each spike formatted as [sample index, unit index]. Intended for
+        use in testing.
+    insertions_replace: bool, default: False
+        If True, randomly replace generated spikes. If False, add to generated spikes.
 
     Returns
     -------
@@ -801,13 +821,63 @@ def synthesize_random_firings(
     times = np.concatenate(times)
     labels = np.concatenate(labels)
 
+    if insertions is not None:
+        times, labels = _add_insertions([times, labels], insertions=insertions, insertions_replace=insertions_replace)
+
     sort_inds = np.argsort(times)
     times = times[sort_inds]
     labels = labels[sort_inds]
 
     return (times, labels)
 
 
+def _add_insertions(spike_train, insertions=None, insertions_replace=False, rng=None, seed=None):
+    """
+    Add specified insertions into a spike train
+
+    Parameters
+    ----------
+    spike_train: array-like
+        Taking the form [spike_times, spike_units, segment_indices]. Segments are optional.
+    insertions: array-like
+        List of manually inserted spikes, each spike formatted as [sample index, unit index, segment_index]. Intended for
+        use in testing. Segment index is optional.
+    insertions_replace: bool, default: False
+        If True, randomly replace generated spikes. If False, add to generated spikes.
+    rng: numpy.random.Generator
+        A random number generator
+    seed: int, default: None
+        seed for the generator
+
+    Returns
+    -------
+        spike_train: np.ndarray
+        Numpy array in same form as spike_train input, including insertions.
+
+    """
+
+    if insertions is None:
+        return spike_train
+
+    new_spike_train = np.array(spike_train)
+    insertions = np.array(insertions)
+
+    assert np.shape(new_spike_train)[0] == len(
+        insertions[0]
+    ), "Number of spike train and insertions indices are not equal."
+
+    if rng is None:
+        rng = np.random.default_rng(seed=seed)
+
+    if insertions_replace:
+        replacement_indices = rng.choice(len(spike_train), len(insertions), replace=False)
+        new_spike_train[:, replacement_indices] = np.transpose(insertions)
+    else:
+        new_spike_train = np.concatenate((new_spike_train, np.transpose(insertions)), axis=1)
+
+    return new_spike_train
+
+
 def clean_refractory_period(times, refractory_period):
     """
     Remove spike that violate the refractory period in a given spike train.

src/spikeinterface/core/tests/test_generate.py

@@ -19,6 +19,8 @@
     generate_unit_locations,
     generate_ground_truth_recording,
     generate_sorting_to_inject,
+    _add_insertions,
+    synthesize_random_firings,
 )
 
 from spikeinterface.core.numpyextractors import NumpySorting
@@ -530,6 +532,88 @@ def test_generate_sorting_to_inject():
         assert num_injected_spikes[unit_id] <= num_spikes[unit_id]
 
 
+def test_synthesize_random_firings_length():
+
+    firing_rates = [2.0, 3.0]
+    duration = 2
+    num_units = 2
+
+    spike_times, spike_units = synthesize_random_firings(
+        num_units=num_units, duration=duration, firing_rates=firing_rates
+    )
+
+    assert len(spike_times) == int(np.sum(firing_rates) * duration)
+
+    units, counts = np.unique(spike_units, return_counts=True)
+
+    assert len(units) == num_units
+    assert np.sum(counts) == int(np.sum(firing_rates) * duration)
+
+
+def test_synthesize_random_firings_length_with_insertion():
+
+    firing_rates = [2.0, 3.0]
+    duration = 2
+    num_units = 2
+
+    insertions = [[15000, 12], [1, 2]]
+
+    spike_times, _ = synthesize_random_firings(
+        num_units=num_units,
+        duration=duration,
+        firing_rates=firing_rates,
+        insertions=insertions,
+        insertions_replace=False,
+    )
+
+    assert len(spike_times) == int(np.sum(firing_rates) * duration) + 2
+
+    spike_times, _ = synthesize_random_firings(
+        num_units=num_units,
+        duration=duration,
+        firing_rates=firing_rates,
+        insertions=insertions,
+        insertions_replace=True,
+    )
+
+    assert len(spike_times) == int(np.sum(firing_rates) * duration)
+
+
+def test_add_insertions_replacement():
+
+    train_length = 10
+
+    spike_train = [np.zeros(train_length), np.zeros(train_length), np.zeros(train_length)]
+
+    insertion_1 = (15, 12, 0)
+    insertion_2 = (1, 2, 3)
+
+    insertions = [insertion_1, insertion_2]
+
+    spike_train_replace = _add_insertions(spike_train, insertions=insertions, insertions_replace=True)
+
+    assert np.any(np.all(np.transpose(spike_train_replace) == np.array(insertion_1), axis=1))
+    assert np.shape(spike_train_replace)[1] == train_length
+
+
+def test_add_insertions_no_replacement():
+
+    train_length = 10
+
+    spike_train = [np.zeros(train_length), np.zeros(train_length), np.zeros(train_length)]
+
+    insertion_1 = [15, 12, 0]
+    insertion_2 = [1, 2, 3]
+
+    insertions = [insertion_1, insertion_2]
+
+    spike_train_replace = _add_insertions(spike_train, insertions=insertions, insertions_replace=False)
+
+    assert np.any(np.all(np.transpose(spike_train_replace) == np.array(insertion_1), axis=1))
+
+    assert np.shape(spike_train_replace)[1] == train_length + len(insertions)
+
+
 if __name__ == "__main__":
     strategy = "tile_pregenerated"
     # strategy = "on_the_fly"

src/spikeinterface/qualitymetrics/__init__.py

@@ -6,3 +6,4 @@
     get_default_qm_params,
 )
 from .pca_metrics import get_quality_pca_metric_list
+from .misc_metrics import get_synchrony_counts

src/spikeinterface/qualitymetrics/misc_metrics.py

@@ -496,7 +496,51 @@ def compute_sliding_rp_violations(
 )
 
 
-def compute_synchrony_metrics(sorting_analyzer, synchrony_sizes=(2, 4, 8), unit_ids=None, **kwargs):
+def get_synchrony_counts(spikes, synchrony_sizes, all_unit_ids):
+    """Compute synchrony counts, the number of simultaneous spikes with sizes `synchrony_sizes`
+
+    Parameters
+    ----------
+    spikes : np.array
+        Structured numpy array with fields ("sample_index", "unit_index", "segment_index").
+    synchrony_sizes : numpy array
+        The synchrony sizes to compute. Should be pre-sorted.
+    unit_ids : list or None, default: None
+        List of unit ids to compute the synchrony metrics. Expecting all units.
+
+    Returns
+    -------
+    synchrony_counts : dict
+        The synchrony counts for the synchrony sizes.
+
+    References
+    ----------
+    Based on concepts described in [Gruen]_
+    This code was adapted from `Elephant - Electrophysiology Analysis Toolkit <https://github.com/NeuralEnsemble/elephant/blob/master/elephant/spike_train_synchrony.py#L245>`_
+    """
+
+    synchrony_counts = np.zeros((np.size(synchrony_sizes), len(all_unit_ids)), dtype=np.int64)
+
+    # compute the occurrence of each sample_index. Count >2 means there's synchrony
+    _, unique_spike_index, counts = np.unique(spikes["sample_index"], return_index=True, return_counts=True)
+
+    sync_indices = unique_spike_index[counts >= 2]
+    sync_counts = counts[counts >= 2]
+
+    for i, sync_index in enumerate(sync_indices):
+
+        num_of_syncs = sync_counts[i]
+        units_with_sync = [spikes[sync_index + a][1] for a in range(0, num_of_syncs)]
+
+        # Counts inclusively. E.g. if there are 3 simultaneous spikes, these are also added
+        # to the 2 simultaneous spike bins.
+        how_many_bins_to_add_to = np.size(synchrony_sizes[synchrony_sizes <= num_of_syncs])
+        synchrony_counts[:how_many_bins_to_add_to, units_with_sync] += 1
+
+    return synchrony_counts
+
+
+def compute_synchrony_metrics(sorting_analyzer, synchrony_sizes=(2, 4, 8), unit_ids=None):
     """Compute synchrony metrics. Synchrony metrics represent the rate of occurrences of
     "synchrony_size" spikes at the exact same sample index.
 
@@ -521,49 +565,39 @@ def compute_synchrony_metrics(sorting_analyzer, synchrony_sizes=(2, 4, 8), unit_
     This code was adapted from `Elephant - Electrophysiology Analysis Toolkit <https://github.com/NeuralEnsemble/elephant/blob/master/elephant/spike_train_synchrony.py#L245>`_
     """
     assert min(synchrony_sizes) > 1, "Synchrony sizes must be greater than 1"
-    spike_counts = sorting_analyzer.sorting.count_num_spikes_per_unit(outputs="dict")
-    sorting = sorting_analyzer.sorting
-    spikes = sorting.to_spike_vector(concatenated=False)
+    # Sort the synchrony times so we can slice numpy arrays, instead of using dicts
+    synchrony_sizes_np = np.array(synchrony_sizes, dtype=np.int64)
+    synchrony_sizes_np.sort()
 
-    if unit_ids is None:
-        unit_ids = sorting_analyzer.unit_ids
+    res = namedtuple("synchrony", [f"sync_spike_{size}" for size in synchrony_sizes_np])
 
-    # Pre-allocate synchrony counts
-    synchrony_counts = {}
-    for synchrony_size in synchrony_sizes:
-        synchrony_counts[synchrony_size] = np.zeros(len(sorting_analyzer.unit_ids), dtype=np.int64)
+    sorting = sorting_analyzer.sorting
 
-    all_unit_ids = list(sorting.unit_ids)
-    for segment_index in range(sorting.get_num_segments()):
-        spikes_in_segment = spikes[segment_index]
+    spike_counts = sorting.count_num_spikes_per_unit(outputs="dict")
 
-        # we compute just by counting the occurrence of each sample_index
-        unique_spike_index, complexity = np.unique(spikes_in_segment["sample_index"], return_counts=True)
+    spikes = sorting.to_spike_vector()
+    all_unit_ids = sorting.unit_ids
+    synchrony_counts = get_synchrony_counts(spikes, synchrony_sizes_np, all_unit_ids)
+
+    synchrony_metrics_dict = {}
+    for sync_idx, synchrony_size in enumerate(synchrony_sizes_np):
+        synch_id_metrics_dict = {}
+        for i, unit_id in enumerate(all_unit_ids):
+            if spike_counts[unit_id] != 0:
+                synch_id_metrics_dict[unit_id] = synchrony_counts[sync_idx][i] / spike_counts[unit_id]
+            else:
+                synch_id_metrics_dict[unit_id] = 0
+        synchrony_metrics_dict[f"sync_spike_{synchrony_size}"] = synch_id_metrics_dict
 
-        # add counts for this segment
-        for unit_id in unit_ids:
-            unit_index = all_unit_ids.index(unit_id)
-            spikes_per_unit = spikes_in_segment[spikes_in_segment["unit_index"] == unit_index]
-            # some segments/units might have no spikes
-            if len(spikes_per_unit) == 0:
-                continue
-            spike_complexity = complexity[np.isin(unique_spike_index, spikes_per_unit["sample_index"])]
-            for synchrony_size in synchrony_sizes:
-                synchrony_counts[synchrony_size][unit_index] += np.count_nonzero(spike_complexity >= synchrony_size)
-
-    # add counts for this segment
-    synchrony_metrics_dict = {
-        f"sync_spike_{synchrony_size}": {
-            unit_id: synchrony_counts[synchrony_size][all_unit_ids.index(unit_id)] / spike_counts[unit_id]
-            for unit_id in unit_ids
-        }
-        for synchrony_size in synchrony_sizes
-    }
-
-    # Convert dict to named tuple
-    synchrony_metrics_tuple = namedtuple("synchrony_metrics", synchrony_metrics_dict.keys())
-    synchrony_metrics = synchrony_metrics_tuple(**synchrony_metrics_dict)
-    return synchrony_metrics
+    if np.all(unit_ids == None) or (len(unit_ids) == len(all_unit_ids)):
+        return res(**synchrony_metrics_dict)
+    else:
+        reduced_synchrony_metrics_dict = {}
+        for key in synchrony_metrics_dict:
+            reduced_synchrony_metrics_dict[key] = {
+                unit_id: synchrony_metrics_dict[key][unit_id] for unit_id in unit_ids
+            }
+        return res(**reduced_synchrony_metrics_dict)
 
 
 _default_params["synchrony"] = dict(synchrony_sizes=(2, 4, 8))