Merge pull request #2420 from DradeAW/faster_gaussian_filter

Faster Gaussian filter implementation

Author: alejoe91

src/spikeinterface/core/__init__.py

@@ -88,6 +88,7 @@
 from .core_tools import (
     read_python,
     write_python,
+    normal_pdf,
 )
 from .job_tools import ensure_n_jobs, ensure_chunk_size, ChunkRecordingExecutor, split_job_kwargs, fix_job_kwargs
 from .recording_tools import (

src/spikeinterface/core/core_tools.py

@@ -453,3 +453,26 @@ def is_editable_mode() -> bool:
     import spikeinterface
 
     return (Path(spikeinterface.__file__).parents[2] / "README.md").exists()
+
+
+def normal_pdf(x, mu: float = 0.0, sigma: float = 1.0):
+    """
+    Manual implementation of the Normal distribution pdf (probability density function).
+    It is about 8 to 10 times faster than scipy.stats.norm.pdf().
+
+    Parameters
+    ----------
+    x: scalar or array
+        The x-axis
+    mu: float, default: 0.0
+        The mean of the Normal distribution.
+    sigma: float, default: 1.0
+        The standard deviation of the Normal distribution.
+
+    Returns
+    -------
+    normal_pdf: scalar or array (same type as 'x')
+        The pdf of the Normal distribution for the given x-axis.
+    """
+
+    return 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((x - mu) ** 2) / (2 * sigma**2))

src/spikeinterface/core/tests/test_core_tools.py

@@ -1,4 +1,5 @@
 import platform
+import math
 from multiprocessing.shared_memory import SharedMemory
 from pathlib import Path
 import importlib
@@ -10,6 +11,7 @@
     make_paths_relative,
     make_paths_absolute,
     check_paths_relative,
+    normal_pdf,
 )
 from spikeinterface.core.binaryrecordingextractor import BinaryRecordingExtractor
 from spikeinterface.core.generate import NoiseGeneratorRecording
@@ -87,5 +89,21 @@ def test_path_utils_functions():
         assert check_paths_relative(d, r"\\host\share")
 
 
+def test_normal_pdf() -> None:
+    mu = 4.160771
+    sigma = 2.9334
+    dx = 0.001
+
+    xaxis = np.arange(-15, 25, dx)
+    gauss = normal_pdf(xaxis, mu=mu, sigma=sigma)
+
+    assert math.isclose(1.0, dx * np.sum(gauss))  # Checking that sum of pdf is 1
+    assert math.isclose(mu, dx * np.sum(xaxis * gauss))  # Checking that mean is mu
+    assert math.isclose(sigma**2, dx * np.sum(xaxis**2 * gauss) - mu**2)  # Checking that variance is sigma^2
+
+    print(normal_pdf(-0.9355, mu=mu, sigma=sigma))
+    assert math.isclose(normal_pdf(-0.9355, mu=mu, sigma=sigma), 0.03006929091)
+
+
 if __name__ == "__main__":
     test_path_utils_functions()

src/spikeinterface/preprocessing/filter_gaussian.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from spikeinterface.core import BaseRecording, BaseRecordingSegment, get_chunk_with_margin
+from spikeinterface.core import BaseRecording, BaseRecordingSegment, get_chunk_with_margin, normal_pdf
 from spikeinterface.core.core_tools import define_function_from_class
 from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 
@@ -95,16 +95,14 @@ def _create_gaussian(self, N: int, cutoff_f: float):
         sf = self.parent_recording_segment.sampling_frequency
         faxis = np.fft.fftfreq(N, d=1 / sf)
 
-        from scipy.stats import norm
-
         if cutoff_f > sf / 8:  # The Fourier transform of a Gaussian with a very low sigma isn't a Gaussian.
             sigma = sf / (2 * np.pi * cutoff_f)
-            limit = int(round(6 * sigma)) + 1
+            limit = int(round(5 * sigma)) + 1
             xaxis = np.arange(-limit, limit + 1) / sigma
-            gaussian = norm.pdf(xaxis) / sigma
+            gaussian = normal_pdf(xaxis) / sigma
             gaussian = np.abs(np.fft.fft(gaussian, n=N))
         else:
-            gaussian = norm.pdf(faxis / cutoff_f) * np.sqrt(2 * np.pi)
+            gaussian = normal_pdf(faxis / cutoff_f) * np.sqrt(2 * np.pi)
 
         if cutoff_f not in self.cached_gaussian:
             self.cached_gaussian[cutoff_f] = dict()