Refactor FilterBank

docs/citation.py

@@ -1,9 +1,7 @@
 import requests
 
-# Replace with your paper's DOI
 DOI = "10.3389/fnins.2022.999029"
 
-# API endpoint for Semantic Scholar
 API_URL = f"https://api.semanticscholar.org/v1/paper/{DOI}"
 
 

spikify/filtering/filterbank.py

@@ -61,7 +61,7 @@ def __init__(
         f_max: float,
         order: int,
         filter_type: Literal["butterworth", "gammatone", "sos"] = "butterworth",
-        **kwargs
+        **kwargs,
     ):
         """Constructor method."""
         super().__init__()
@@ -76,11 +76,6 @@ def __init__(
         )
         self.freq_poles[-1, 1] = self.fs / 2 * 0.99999
 
-        # Validate inputs
-        if self.filter_type not in ["butterworth", "gammatone", "sos"]:
-            raise ValueError("filter_type must be 'butterworth', 'gammatone', or 'sos'")
-
-        # Build filter coefficients
         self._build_filters(**kwargs)
 
     def _build_filters(self, **kwargs):
@@ -92,28 +87,36 @@ def _build_filters(self, **kwargs):
 
         :param kwargs: Additional filter parameters for specific filter types.
         :type kwargs: dict
+        :raises ValueError: If filter type is not supported.
 
         """
-        self.filters = []
+        self.filter_coeffs = []
         self.channel_frequencies = []
 
-        if self.filter_type == "butterworth":
-            for low_freq, high_freq in self.freq_poles:
-                num, den = butter(N=self.order, Wn=[low_freq, high_freq], btype="band", fs=self.fs)
-                self.filters.append((num, den))
-                self.channel_frequencies.append((low_freq, high_freq))
+        match self.filter_type:
+
+            case "butterworth":
+                for low_freq, high_freq in self.freq_poles:
+                    num, den = butter(N=self.order, Wn=[low_freq, high_freq], btype="band", fs=self.fs, **kwargs)
+                    self.filter_coeffs.append((num, den))
+                    self.channel_frequencies.append((low_freq, high_freq))
+
+            case "gammatone":
+                for freq in self.freq_centers:
+                    num, den = gammatone(order=self.order, freq=freq, ftype="fir", fs=self.fs, **kwargs)
+                    self.filter_coeffs.append((num, den))
+                    self.channel_frequencies.append(freq)
 
-        elif self.filter_type == "gammatone":
-            for freq in self.freq_centers:
-                num, den = gammatone(order=self.order, freq=freq, ftype="fir", fs=self.fs)
-                self.filters.append((num, den))
-                self.channel_frequencies.append(freq)
+            case "sos":
+                for low_freq, high_freq in self.freq_poles:
+                    sos = butter(
+                        N=self.order, Wn=[low_freq, high_freq], btype="band", output="sos", fs=self.fs, **kwargs
+                    )
+                    self.filter_coeffs.append(sos)
+                    self.channel_frequencies.append((low_freq, high_freq))
 
-        elif self.filter_type == "sos":
-            for low_freq, high_freq in self.freq_poles:
-                sos = butter(N=self.order, Wn=[low_freq, high_freq], btype="band", output="sos", fs=self.fs)
-                self.filters.append(sos)
-                self.channel_frequencies.append((low_freq, high_freq))
+            case _:
+                raise ValueError(f"Filter {self.filter_type} is not supported")
 
     def decompose(self, signal: np.ndarray) -> np.ndarray:
         """
@@ -127,30 +130,28 @@ def decompose(self, signal: np.ndarray) -> np.ndarray:
         :type signal: numpy.ndarray
         :return: Array of filtered signals with shape (timestamps, channels, features).
         :rtype: numpy.ndarray
-        :raises ValueError: If signal is not 1D or 2D.
 
         """
-        if len(signal.shape) == 1:
+        # Ensure 2D processing (T, F)
+        if signal.ndim == 1:
             signal = signal.reshape(-1, 1)
-        elif len(signal.shape) != 2:
-            raise ValueError("Signal must be 1D or 2D array")
 
-        n_timestamps, n_features = signal.shape
-        n_channels = len(self.filters)
+        T, F = signal.shape
+        n_channels = len(self.filter_coeffs)
 
         # Initialize output
-        freq_components = np.zeros((n_timestamps, n_channels, n_features))
+        freq_components = np.zeros((T, n_channels, F))
 
         for ch in range(n_channels):
-            filter_coeffs = self.filters[ch]
+            filter_coeffs = self.filter_coeffs[ch]
 
             if self.filter_type == "sos":
                 # Use sosfilt for second-order sections
-                for feat in range(n_features):
+                for feat in range(F):
                     freq_components[:, ch, feat] = sosfilt(filter_coeffs, signal[:, feat])
             else:
                 # Use lfilter for b,a coefficients
-                for feat in range(n_features):
+                for feat in range(F):
                     num, den = filter_coeffs
                     freq_components[:, ch, feat] = lfilter(num, den, signal[:, feat])
 

tests/filtering/test_filterbank.py

@@ -77,18 +77,6 @@ def test_sos_decomposition(self):
         freq_components = filterbank.decompose(self.signal)
         self.assertEqual(freq_components.shape, (self.signal_length, self.channels, 1))
 
-    def test_invalid_filter_type(self):
-        """Test that an invalid filter type raises a ValueError."""
-        with self.assertRaises(ValueError):
-            FilterBank(
-                fs=self.fs,
-                channels=self.channels,
-                f_min=self.f_min,
-                f_max=self.f_max,
-                filter_type="invalid_type",
-                order=self.order,
-            )
-
     def test_signal_multiple_shape_decomposition(self):
         """Test that decomposing a too short signal raises a ValueError."""
         filterbank = FilterBank(
@@ -99,7 +87,7 @@ def test_signal_multiple_shape_decomposition(self):
             filter_type="butterworth",
             order=self.order,
         )
-        signal = np.random.randn(10, 5, 3)  # Invalid shape
+        signal = np.random.randn(10, 5, 3)
         with self.assertRaises(ValueError):
             filterbank.decompose(signal)
 
@@ -113,11 +101,23 @@ def test_signal_with_multiple_features(self):
             filter_type="butterworth",
             order=self.order,
         )
-        multi_feature_signal = np.random.randn(self.signal_length, 3)  # 3 features
+        multi_feature_signal = np.random.randn(self.signal_length, 3)
         freq_components = filterbank.decompose(multi_feature_signal)
         self.assertEqual(freq_components.shape, (self.signal_length, self.channels, 3))
 
-    def test_center_frequencies(self):
+    def test_unsupported_filter(self):
+
+        with self.assertRaises(ValueError):
+            FilterBank(
+                fs=self.fs,
+                channels=self.channels,
+                f_min=self.f_min,
+                f_max=self.f_max,
+                filter_type="gammachirp",
+                order=self.order,
+            )
+
+    def test_center_frequencies_butterworth(self):
         """Test that center frequencies are computed correctly."""
         filterbank = FilterBank(
             fs=self.fs,
@@ -127,6 +127,49 @@ def test_center_frequencies(self):
             filter_type="butterworth",
             order=self.order,
         )
-        expected_octave = (self.channels - 0.5) * np.log10(2) / np.log10(self.f_max / self.f_min)
-        expected_freq_centers = np.array([self.f_min * (2 ** (ch / expected_octave)) for ch in range(self.channels)])
-        np.testing.assert_array_almost_equal(filterbank.freq_centers, expected_freq_centers)
+
+        octave = (self.channels - 0.5) * np.log10(2) / np.log10(self.f_max / self.f_min)
+        freq_centers = np.array([self.f_min * (2 ** (ch / octave)) for ch in range(self.channels)])
+        freq_poles = np.array(
+            [(freq * (2 ** (-1 / (2 * octave))), (freq * (2 ** (1 / (2 * octave))))) for freq in freq_centers]
+        )
+        freq_poles[-1, 1] = self.fs / 2 * 0.99999
+
+        expected_freq_centers = np.array([np.mean(freqs) for freqs in freq_poles])
+        np.testing.assert_array_almost_equal(filterbank.center_frequencies, expected_freq_centers)
+
+    def test_center_frequencies_gammatone(self):
+        """Test that center frequencies are computed correctly."""
+        filterbank = FilterBank(
+            fs=self.fs,
+            channels=self.channels,
+            f_min=self.f_min,
+            f_max=self.f_max,
+            filter_type="gammatone",
+            order=self.order,
+        )
+
+        octave = (self.channels - 0.5) * np.log10(2) / np.log10(self.f_max / self.f_min)
+        expected_freq_centers = np.array([self.f_min * (2 ** (ch / octave)) for ch in range(self.channels)])
+        np.testing.assert_array_almost_equal(filterbank.center_frequencies, expected_freq_centers)
+
+    def test_center_frequencies_sos(self):
+        """Test that center frequencies are computed correctly."""
+        filterbank = FilterBank(
+            fs=self.fs,
+            channels=self.channels,
+            f_min=self.f_min,
+            f_max=self.f_max,
+            filter_type="sos",
+            order=self.order,
+        )
+
+        octave = (self.channels - 0.5) * np.log10(2) / np.log10(self.f_max / self.f_min)
+        freq_centers = np.array([self.f_min * (2 ** (ch / octave)) for ch in range(self.channels)])
+        freq_poles = np.array(
+            [(freq * (2 ** (-1 / (2 * octave))), (freq * (2 ** (1 / (2 * octave))))) for freq in freq_centers]
+        )
+        freq_poles[-1, 1] = self.fs / 2 * 0.99999
+
+        expected_freq_centers = np.array([np.mean(freqs) for freqs in freq_poles])
+        np.testing.assert_array_almost_equal(filterbank.center_frequencies, expected_freq_centers)