Merge pull request #4167 from samuelgarcia/generation_ajust_params

Fix bug in waveform generation + adapt parameters

Author: samuelgarcia

examples/tutorials/comparison/generate_erroneous_sorting.py

@@ -44,7 +44,8 @@ def generate_erroneous_sorting():
     units_err = {}
 
     # sorting_true have 10 units
-    np.random.seed(0)
+    # np.random.seed(0)
+    rng = np.random.default_rng(seed=0)
 
     # unit 1 2 are perfect
     for u in [1, 2]:
@@ -54,29 +55,29 @@ def generate_erroneous_sorting():
     # unit 3 4 (medium) 10 (low) have medium to low agreement
     for u, score in [(3, 0.8), (4, 0.75), (10, 0.3)]:
         st = sorting_true.get_unit_spike_train(u)
-        st = np.sort(np.random.choice(st, size=int(st.size * score), replace=False))
+        st = np.sort(rng.choice(st, size=int(st.size * score), replace=False))
         units_err[u] = st
 
     # unit 5 6 are over merge
     st5 = sorting_true.get_unit_spike_train(5)
     st6 = sorting_true.get_unit_spike_train(6)
     st = np.unique(np.concatenate([st5, st6]))
-    st = np.sort(np.random.choice(st, size=int(st.size * 0.7), replace=False))
+    st = np.sort(rng.choice(st, size=int(st.size * 0.7), replace=False))
     units_err[56] = st
 
     # unit 7 is over split in 2 part
     st7 = sorting_true.get_unit_spike_train(7)
     st70 = st7[::2]
     units_err[70] = st70
     st71 = st7[1::2]
-    st71 = np.sort(np.random.choice(st71, size=int(st71.size * 0.9), replace=False))
+    st71 = np.sort(rng.choice(st71, size=int(st71.size * 0.9), replace=False))
     units_err[71] = st71
 
     # unit 8 is redundant 3 times
     st8 = sorting_true.get_unit_spike_train(8)
-    st80 = np.sort(np.random.choice(st8, size=int(st8.size * 0.65), replace=False))
-    st81 = np.sort(np.random.choice(st8, size=int(st8.size * 0.6), replace=False))
-    st82 = np.sort(np.random.choice(st8, size=int(st8.size * 0.55), replace=False))
+    st80 = np.sort(rng.choice(st8, size=int(st8.size * 0.65), replace=False))
+    st81 = np.sort(rng.choice(st8, size=int(st8.size * 0.6), replace=False))
+    st82 = np.sort(rng.choice(st8, size=int(st8.size * 0.55), replace=False))
     units_err[80] = st80
     units_err[81] = st81
     units_err[82] = st82

src/spikeinterface/core/generate.py

@@ -296,7 +296,7 @@ def generate_sorting_to_inject(
             injected_spike_train = injected_spike_train[~violations]
 
             if len(injected_spike_train) > n_injection:
-                injected_spike_train = np.sort(np.random.choice(injected_spike_train, n_injection, replace=False))
+                injected_spike_train = np.sort(rng.choice(injected_spike_train, n_injection, replace=False))
 
             injected_spike_trains[segment_index][unit_id] = injected_spike_train
 
@@ -1519,7 +1519,7 @@ def exp_growth(start_amp, end_amp, duration_ms, tau_ms, sampling_frequency, flip
     return y[:-1]
 
 
-def get_ellipse(positions, center, b=1, c=1, x_angle=0, y_angle=0, z_angle=0):
+def get_ellipse(positions, center, x_factor=1, y_factor=1, x_angle=0, y_angle=0, z_angle=0):
     """
     Compute the distances to a particular ellipsoid in order to take into account
     spatial inhomogeneities while generating the template. In a carthesian, centered
@@ -1537,7 +1537,7 @@ def get_ellipse(positions, center, b=1, c=1, x_angle=0, y_angle=0, z_angle=0):
                             z - z0
 
     In this new space, we can compute the radius of the ellipsoidal shape given the same formula
-        R = X**2 + (Y/b)**2 + (Z/c)**2
+        R = (X/x_factor)**2 + (Y/y_factor)**2 + (Z/1)**2
 
     and thus obtain putative amplitudes given the ellipsoidal projections. Note that in case of a=b=1 and
     no rotation, the distance is the same as the euclidean distance
@@ -1555,7 +1555,7 @@ def get_ellipse(positions, center, b=1, c=1, x_angle=0, y_angle=0, z_angle=0):
     Rx = np.zeros((3, 3))
     Rx[0, 0] = 1
     Rx[1, 1] = np.cos(-x_angle)
-    Rx[1, 0] = -np.sin(-x_angle)
+    Rx[1, 2] = -np.sin(-x_angle)
     Rx[2, 1] = np.sin(-x_angle)
     Rx[2, 2] = np.cos(-x_angle)
 
@@ -1573,10 +1573,12 @@ def get_ellipse(positions, center, b=1, c=1, x_angle=0, y_angle=0, z_angle=0):
     Rz[1, 0] = np.sin(-z_angle)
     Rz[1, 1] = np.cos(-z_angle)
 
-    inv_matrix = np.dot(Rx, Ry, Rz)
-    P = np.dot(inv_matrix, p)
+    rot_matrix = Rx @ Ry @ Rz
+    P = rot_matrix @ p
 
-    return np.sqrt(P[0] ** 2 + (P[1] / b) ** 2 + (P[2] / c) ** 2)
+    distances = np.sqrt((P[0] / x_factor) ** 2 + (P[1] / y_factor) ** 2 + (P[2] / 1) ** 2)
+
+    return distances
 
 
 def generate_single_fake_waveform(
@@ -1632,7 +1634,10 @@ def generate_single_fake_waveform(
     smooth_kernel = np.exp(-(bins**2) / (2 * smooth_size**2))
     smooth_kernel /= np.sum(smooth_kernel)
     # smooth_kernel = smooth_kernel[4:]
+    old_max = np.max(np.abs(wf))
     wf = np.convolve(wf, smooth_kernel, mode="same")
+    new_max = np.max(np.abs(wf))
+    wf *= old_max / new_max
 
     # ensure the the peak to be extatly at nbefore (smooth can modify this)
     ind = np.argmin(wf)
@@ -1653,13 +1658,10 @@ def generate_single_fake_waveform(
     recovery_ms=(1.0, 1.5),
     positive_amplitude=(0.1, 0.25),
     smooth_ms=(0.03, 0.07),
-    spatial_decay=(20, 40),
+    spatial_decay=(10.0, 45.0),
     propagation_speed=(250.0, 350.0),  # um  / ms
-    b=(0.1, 1),
-    c=(0.1, 1),
-    x_angle=(0, np.pi),
-    y_angle=(0, np.pi),
-    z_angle=(0, np.pi),
+    ellipse_shrink=(0.4, 1),
+    ellipse_angle=(0, np.pi * 2),
 )
 
 
@@ -1813,21 +1815,21 @@ def generate_templates(
             distances = get_ellipse(
                 channel_locations,
                 units_locations[u],
-                1,
-                1,
-                0,
-                0,
-                0,
+                x_factor=1,
+                y_factor=1,
+                x_angle=0,
+                y_angle=0,
+                z_angle=0,
             )
         elif mode == "ellipsoid":
             distances = get_ellipse(
                 channel_locations,
                 units_locations[u],
-                params["b"][u],
-                params["c"][u],
-                params["x_angle"][u],
-                params["y_angle"][u],
-                params["z_angle"][u],
+                x_factor=1,
+                y_factor=params["ellipse_shrink"][u],
+                x_angle=0,
+                y_angle=0,
+                z_angle=params["ellipse_angle"][u],
             )
 
         channel_factors = alpha * np.exp(-distances / spatial_decay)
@@ -2166,7 +2168,7 @@ def _generate_multimodal(rng, size, num_modes, lim0, lim1):
         sigma = mode_step / 5.0
         prob += np.exp(-((bins - center) ** 2) / (2 * sigma**2))
     prob /= np.sum(prob)
-    choices = np.random.choice(np.arange(bins.size), size, p=prob)
+    choices = rng.choice(np.arange(bins.size), size, p=prob)
     values = bins[choices] + rng.uniform(low=-bin_step / 2, high=bin_step / 2, size=size)
     return values
 

src/spikeinterface/curation/curation_tools.py

@@ -46,7 +46,7 @@ def _find_duplicated_spikes_numpy(
 
 def _find_duplicated_spikes_random(spike_train: np.ndarray, censored_period: int, seed: int) -> np.ndarray:
     # random seed
-    rng = np.random.RandomState(seed=seed)
+    rng = np.random.default_rng(seed=seed)
 
     indices_of_duplicates = []
     while not np.all(np.diff(np.delete(spike_train, indices_of_duplicates)) > censored_period):

src/spikeinterface/curation/tests/test_sortingview_curation.py

@@ -212,7 +212,8 @@ def test_label_inheritance_str():
     num_timepoints = int(sampling_frequency * duration)
     num_spikes = 1000
     times = np.int_(np.sort(np.random.uniform(0, num_timepoints, num_spikes)))
-    labels = np.random.choice(["a", "b", "c", "d", "e", "f", "g"], size=num_spikes)
+    rng = np.random.default_rng(seed=None)
+    labels = rng.choice(["a", "b", "c", "d", "e", "f", "g"], size=num_spikes)
 
     sorting = se.NumpySorting.from_samples_and_labels(times, labels, sampling_frequency)
     # print(f"Sorting: {sorting.get_unit_ids()}")

src/spikeinterface/generation/drifting_generator.py

@@ -18,6 +18,7 @@
     generate_sorting,
     generate_templates,
     _ensure_unit_params,
+    _ensure_seed,
 )
 from .drift_tools import DriftingTemplates, make_linear_displacement, InjectDriftingTemplatesRecording
 from .noise_tools import generate_noise
@@ -136,7 +137,7 @@ def make_one_displacement_vector(
 
         min_bump_interval, max_bump_interval = bump_interval_s
 
-        rg = np.random.RandomState(seed=seed)
+        rg = np.random.default_rng(seed=seed)
         diff = rg.uniform(min_bump_interval, max_bump_interval, size=int(duration / min_bump_interval))
         bumps_times = np.cumsum(diff) + t_start_drift
         bumps_times = bumps_times[bumps_times < t_end_drift]
@@ -152,8 +153,8 @@ def make_one_displacement_vector(
                 displacement_vector[ind0:ind1] = -0.5
 
     elif drift_mode == "random_walk":
-        rg = np.random.RandomState(seed=seed)
-        steps = rg.random_integers(low=0, high=1, size=num_samples)
+        rg = np.random.default_rng(seed=seed)
+        steps = rg.integers(low=0, high=1, size=num_samples, endpoint=True)
         steps = steps.astype("float64")
         # 0 -> -1 and 1 -> 1
         steps = steps * 2 - 1
@@ -340,12 +341,14 @@ def generate_drifting_recording(
         ms_after=3.0,
         mode="ellipsoid",
         unit_params=dict(
-            alpha=(150.0, 500.0),
+            alpha=(100.0, 500.0),
             spatial_decay=(10, 45),
+            ellipse_shrink=(0.4, 1),
+            ellipse_angle=(0, np.pi * 2),
         ),
     ),
     generate_sorting_kwargs=dict(firing_rates=(2.0, 8.0), refractory_period_ms=4.0),
-    generate_noise_kwargs=dict(noise_levels=(12.0, 15.0), spatial_decay=25.0),
+    generate_noise_kwargs=dict(noise_levels=(6.0, 8.0), spatial_decay=25.0),
     extra_outputs=False,
     seed=None,
 ):
@@ -400,6 +403,9 @@ def generate_drifting_recording(
 
         This can be helpfull for motion benchmark.
     """
+
+    seed = _ensure_seed(seed)
+
     # probe
     if generate_probe_kwargs is None:
         generate_probe_kwargs = _toy_probes[probe_name]

src/spikeinterface/generation/splitting_tools.py

@@ -31,7 +31,7 @@ def split_sorting_by_times(
     """
 
     sorting = sorting_analyzer.sorting
-    rng = np.random.RandomState(seed)
+    rng = np.random.default_rng(seed)
     fs = sorting_analyzer.sampling_frequency
 
     nb_splits = int(splitting_probability * len(sorting.unit_ids))
@@ -60,7 +60,7 @@ def split_sorting_by_times(
         ind_mask = spike_indices[0][unit_id]
         m = np.median(spikes[0][ind_mask]["sample_index"])
         time_mask = spikes[0][ind_mask]["sample_index"] > m
-        mask = time_mask & (rng.rand(len(ind_mask)) <= partial_split_prob).astype(bool)
+        mask = time_mask & (rng.uniform(size=len(ind_mask)) <= partial_split_prob).astype(bool)
         new_index = int(unit_id) * np.ones(len(mask), dtype=bool)
         new_index[mask] = max_index + 1
         new_spikes["unit_index"][ind_mask] = new_index
@@ -102,7 +102,7 @@ def split_sorting_by_amplitudes(
     if sorting_analyzer.get_extension("spike_amplitudes") is None:
         sorting_analyzer.compute("spike_amplitudes")
 
-    rng = np.random.RandomState(seed)
+    rng = np.random.default_rng(seed)
     fs = sorting_analyzer.sampling_frequency
     from spikeinterface.core.template_tools import get_template_extremum_channel
 
@@ -135,7 +135,7 @@ def split_sorting_by_amplitudes(
         ind_mask = spike_indices[0][unit_id]
         thresh = np.median(amplitudes[ind_mask])
         amplitude_mask = amplitudes[ind_mask] > thresh
-        mask = amplitude_mask & (rng.rand(len(ind_mask)) <= partial_split_prob).astype(bool)
+        mask = amplitude_mask & (rng.uniform(size=len(ind_mask)) <= partial_split_prob).astype(bool)
         new_index = int(unit_id) * np.ones(len(mask))
         new_index[mask] = max_index + 1
         new_spikes["unit_index"][ind_mask] = new_index

src/spikeinterface/preprocessing/tests/test_astype.py

@@ -11,8 +11,8 @@
 
 
 def test_astype():
-    rng = np.random.RandomState(0)
-    traces = (rng.randn(10000, 4) * 100).astype("float32")
+    rng = np.random.default_rng(0)
+    traces = (rng.normal(size=(10000, 4)) * 100).astype("float32")
     rec_float32 = NumpyRecording(traces, sampling_frequency=30000)
     traces_int16 = traces.astype("int16")
     np.testing.assert_array_equal(traces_int16, astype(rec_float32, "int16", round=False).get_traces())

src/spikeinterface/preprocessing/tests/test_detect_bad_channels.py

@@ -168,9 +168,8 @@ def test_detect_bad_channels_ibl(num_channels):
     recording.set_channel_offsets(0)
 
     # Generate random channels to be dead / noisy
-    is_bad = np.random.choice(
-        np.arange(num_channels - 3), size=np.random.randint(5, int(num_channels * 0.25)), replace=False
-    )
+    rng = np.random.default_rng(seed=None)
+    is_bad = rng.choice(np.arange(num_channels - 3), size=np.random.randint(5, int(num_channels * 0.25)), replace=False)
     is_noisy, is_dead = np.array_split(is_bad, 2)
     not_noisy = np.delete(np.arange(num_channels), is_noisy)
 
@@ -230,8 +229,9 @@ def test_detect_bad_channels_ibl(num_channels):
     assert np.array_equal(recording.ids_to_indices(bad_channel_ids), np.where(bad_channel_labels_ibl != 0)[0])
 
     # Test on randomly sorted channels
+    rng = np.random.default_rng(seed=None)
     recording_scrambled = recording.channel_slice(
-        np.random.choice(recording.channel_ids, len(recording.channel_ids), replace=False)
+        rng.choice(recording.channel_ids, len(recording.channel_ids), replace=False)
     )
     bad_channel_ids_scrambled, bad_channel_label_scrambled = detect_bad_channels(
         recording_scrambled,

src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py

@@ -92,7 +92,7 @@ def test_highpass_spatial_filter_synthetic_data(num_channels, ntr_pad, ntr_tap,
     options = dict(lagc=lagc, ntr_pad=ntr_pad, ntr_tap=ntr_tap, butter_kwargs=butter_kwargs)
 
     durations = [2, 2]
-    rng = np.random.RandomState(seed=100)
+    rng = np.random.default_rng(seed=100)
     si_recording = generate_recording(num_channels=num_channels, durations=durations)
 
     _, si_highpass_spatial_filter = run_si_highpass_filter(si_recording, get_traces=False, **options)

src/spikeinterface/preprocessing/tests/test_interpolate_bad_channels.py

@@ -75,7 +75,8 @@ def test_compare_real_data_with_ibl():
     )
 
     num_channels = si_recording.get_num_channels()
-    bad_channel_indexes = np.random.choice(num_channels, 10, replace=False)
+    rng = np.random.default_rng(seed=None)
+    bad_channel_indexes = rng.choice(num_channels, 10, replace=False)
     bad_channel_ids = si_recording.channel_ids[bad_channel_indexes]
     si_recording = spre.scale(si_recording, dtype="float32")
 
@@ -123,12 +124,13 @@ def test_compare_input_argument_ranges_against_ibl(shanks, p, sigma_um, num_chan
     recording = generate_recording(num_channels=num_channels, durations=[1])
 
     # distribute default probe locations across 4 shanks if set
-    x = np.random.choice(shanks, num_channels)
+    rng = np.random.default_rng(seed=None)
+    x = rng.choice(shanks, num_channels)
     for idx, __ in enumerate(recording._properties["contact_vector"]):
         recording._properties["contact_vector"][idx][1] = x[idx]
 
     # generate random bad channel locations
-    bad_channel_indexes = np.random.choice(num_channels, np.random.randint(1, int(num_channels / 5)), replace=False)
+    bad_channel_indexes = rng.choice(num_channels, rng.randint(1, int(num_channels / 5)), replace=False)
     bad_channel_ids = recording.channel_ids[bad_channel_indexes]
 
     # Run SI and IBL interpolation and check against eachother