Fix bug in waveform generation + adapt parameters

src/spikeinterface/core/generate.py

@@ -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_shink=(0.2, 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_shink"][u],
+                x_angle=0,
+                y_angle=0,
+                z_angle=params["ellipse_angle"][u],
             )
 
         channel_factors = alpha * np.exp(-distances / spatial_decay)

src/spikeinterface/generation/drifting_generator.py

@@ -340,12 +340,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_shink=(0.2, 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,
 ):