Merge pull request scipy#18375 from sagi-ezri/add-fwind1-filter

ENH: signal: Add `firwin_2d` filter

Author: DietBru

scipy/signal/__init__.py

@@ -87,6 +87,8 @@
                     -- defined as pass and stop bands.
    firwin2       -- Windowed FIR filter design, with arbitrary frequency
                     -- response.
+   firwin_2d        -- Windowed FIR filter design, with frequency response for 
+                    -- 2D using 1D design.
    freqs         -- Analog filter frequency response from TF coefficients.
    freqs_zpk     -- Analog filter frequency response from ZPK coefficients.
    freqz         -- Digital filter frequency response from TF coefficients.

scipy/signal/_fir_filter_design.py

@@ -14,8 +14,9 @@
 
 from . import _sigtools
 
+
 __all__ = ['kaiser_beta', 'kaiser_atten', 'kaiserord',
-           'firwin', 'firwin2', 'remez', 'firls', 'minimum_phase']
+           'firwin', 'firwin2', 'firwin_2d', 'remez', 'firls', 'minimum_phase']
 
 
 # Some notes on function parameters:
@@ -323,6 +324,7 @@ def firwin(numtaps, cutoff, *, width=None, window='hamming', pass_zero=True,
     See Also
     --------
     firwin2
+    firwin_2d
     firls
     minimum_phase
     remez
@@ -1284,3 +1286,148 @@ def minimum_phase(h: np.ndarray,
         h_minimum = h_temp.real
     n_out = (n_half + len(h) % 2) if half else len(h)
     return h_minimum[:n_out]
+
+
+def firwin_2d(hsize, window, *, fc=None, fs=2, circular=False,
+              pass_zero=True, scale=True):
+    """
+    2D FIR filter design using the window method.
+
+    This function computes the coefficients of a 2D finite impulse response
+    filter. The filter is separable with linear phase; it will be designed
+    as a product of two 1D filters with dimensions defined by `hsize`.
+    Additionally, it can create approximately circularly symmetric 2-D windows.
+    
+    Parameters
+    ----------
+    hsize : tuple or list of length 2
+        Lengths of the filter in each dimension. `hsize[0]` specifies the
+        number of coefficients in the row direction and `hsize[1]` specifies
+        the number of coefficients in the column direction.
+    window : tuple or list of length 2 or string
+        Desired window to use for each 1D filter or a single window type 
+        for creating circularly symmetric 2-D windows. Each element should be
+        a string or tuple of string and parameter values. See
+        `~scipy.signal.get_window` for a list of windows and required
+        parameters.
+    fc : float or 1-D array_like, optional
+        Cutoff frequency of the filter in the same units as `fs`. This defines
+        the frequency at which the filter's gain drops to approximately -6 dB
+        (half power) in a low-pass or high-pass filter. For multi-band filters,
+        `fc` can be an array of cutoff frequencies (i.e., band edges) in the 
+        range [0, fs/2], with each band specified in pairs. Required if 
+        `circular` is False.
+    fs : float, optional
+        The sampling frequency of the signal. Default is 2.
+    circular : bool, optional
+        Whether to create a circularly symmetric 2-D window. Default is ``False``.
+    pass_zero : {True, False, 'bandpass', 'lowpass', 'highpass', 'bandstop'}, optional
+        This parameter is directly passed to `firwin` for each scalar frequency axis.
+        Hence, if ``True``, the DC gain, i.e., the gain at frequency (0, 0), is 1.
+        If ``False``, the DC gain is 0 at frequency (0, 0) if `circular` is ``True``. 
+        If `circular` is ``False`` the frequencies (0, f1) and (f0, 0) will
+        have gain 0.
+        It can also be a string argument for the desired filter type 
+        (equivalent to ``btype`` in IIR design functions).
+    scale : bool, optional
+        This parameter is directly passed to `firwin` for each scalar frequency axis.
+        Set to ``True`` to scale the coefficients so that the frequency
+        response is exactly unity at a certain frequency on one frequency axis.
+        That frequency is either:
+        
+        - 0 (DC) if the first passband starts at 0 (i.e. pass_zero is ``True``)
+        - `fs`/2 (the Nyquist frequency) if the first passband ends at `fs`/2
+          (i.e., the filter is a single band highpass filter);
+          center of first passband otherwise
+
+    Returns
+    -------
+    filter_2d : (hsize[0], hsize[1]) ndarray
+        Coefficients of 2D FIR filter.
+
+    Raises
+    ------
+    ValueError
+        - If `hsize` and `window` are not 2-element tuples or lists.
+        - If `cutoff` is None when `circular` is True.
+        - If `cutoff` is outside the range [0, `fs`/2] and `circular` is ``False``.
+        - If any of the elements in `window` are not recognized.
+    RuntimeError
+        If `firwin` fails to converge when designing the filter.
+
+    See Also
+    --------
+    firwin: FIR filter design using the window method for 1d arrays.
+    get_window: Return a window of a given length and type.
+
+    Examples
+    --------
+    Generate a 5x5 low-pass filter with cutoff frequency 0.1:
+
+    >>> import numpy as np
+    >>> from scipy.signal import get_window
+    >>> from scipy.signal import firwin_2d
+    >>> hsize = (5, 5)
+    >>> window = (("kaiser", 5.0), ("kaiser", 5.0))
+    >>> fc = 0.1
+    >>> filter_2d = firwin_2d(hsize, window, fc=fc)
+    >>> filter_2d
+    array([[0.00025366, 0.00401662, 0.00738617, 0.00401662, 0.00025366],
+           [0.00401662, 0.06360159, 0.11695714, 0.06360159, 0.00401662],
+           [0.00738617, 0.11695714, 0.21507283, 0.11695714, 0.00738617],
+           [0.00401662, 0.06360159, 0.11695714, 0.06360159, 0.00401662],
+           [0.00025366, 0.00401662, 0.00738617, 0.00401662, 0.00025366]])
+
+    Generate a circularly symmetric 5x5 low-pass filter with Hamming window:
+
+    >>> filter_2d = firwin_2d((5, 5), 'hamming', fc=fc, circular=True)
+    >>> filter_2d
+    array([[-0.00020354, -0.00020354, -0.00020354, -0.00020354, -0.00020354],
+           [-0.00020354,  0.01506844,  0.09907658,  0.01506844, -0.00020354],
+           [-0.00020354,  0.09907658, -0.00020354,  0.09907658, -0.00020354],
+           [-0.00020354,  0.01506844,  0.09907658,  0.01506844, -0.00020354],
+           [-0.00020354, -0.00020354, -0.00020354, -0.00020354, -0.00020354]])
+
+    Generate Plots comparing the product of two 1d filters with a circular
+    symmetric filter:
+
+    >>> import matplotlib.pyplot as plt
+    >>> hsize, fc = (50, 50), 0.05
+    >>> window = (("kaiser", 5.0), ("kaiser", 5.0))
+    >>> filter0_2d = firwin_2d(hsize, window, fc=fc)
+    >>> filter1_2d = firwin_2d((50, 50), 'hamming', fc=fc, circular=True)
+    ...
+    >>> fg, (ax0, ax1) = plt.subplots(1, 2, tight_layout=True, figsize=(6.5, 3.5))
+    >>> ax0.set_title("Product of 2 Windows")
+    >>> im0 = ax0.imshow(filter0_2d, cmap='viridis', origin='lower', aspect='equal')
+    >>> fg.colorbar(im0, ax=ax0, shrink=0.7)
+    >>> ax1.set_title("Circular Window")
+    >>> im1 = ax1.imshow(filter1_2d, cmap='plasma', origin='lower', aspect='equal')
+    >>> fg.colorbar(im1, ax=ax1, shrink=0.7)
+    >>> plt.show()
+    """
+    if len(hsize) != 2:
+            raise ValueError("hsize must be a 2-element tuple or list")
+
+    if circular:
+        if fc is None:
+            raise ValueError("Cutoff frequency `fc` must be "
+                             "provided when `circular` is True")
+        
+        n_r = max(hsize[0], hsize[1]) * 8  # oversample 1d window by factor 8
+        
+        win_r = firwin(n_r, cutoff=fc, window=window, fs=fs)
+
+        f1, f2 = np.meshgrid(np.linspace(-1, 1, hsize[0]), np.linspace(-1, 1, hsize[1]))
+        r = np.sqrt(f1**2 + f2**2)
+
+        win_2d = np.interp(r, np.linspace(0, 1, n_r), win_r)
+        return win_2d
+
+    if len(window) != 2:
+        raise ValueError("window must be a 2-element tuple or list")
+
+    row_filter = firwin(hsize[0], cutoff=fc, window=window[0], fs=fs)
+    col_filter = firwin(hsize[1], cutoff=fc, window=window[1], fs=fs)
+
+    return np.outer(row_filter, col_filter)

scipy/signal/fir_filter_design.py

@@ -7,6 +7,7 @@
 __all__ = [  # noqa: F822
     'kaiser_beta', 'kaiser_atten', 'kaiserord',
     'firwin', 'firwin2', 'remez', 'firls', 'minimum_phase',
+    'firwin_2d',
 ]
 
 

scipy/signal/tests/test_fir_filter_design.py

@@ -7,12 +7,12 @@
 from pytest import raises as assert_raises
 import pytest
 
-from scipy.fft import fft
+from scipy.fft import fft, fft2
 from scipy.special import sinc
-from scipy.signal import (kaiser_beta, kaiser_atten, kaiserord,
-    firwin, firwin2, freqz, remez, firls, minimum_phase
-)
-
+from scipy.signal import kaiser_beta, kaiser_atten, kaiserord, \
+    firwin, firwin2, freqz, remez, firls, minimum_phase, \
+    convolve2d
+from scipy.signal._fir_filter_design import firwin_2d
 
 def test_kaiser_beta():
     b = kaiser_beta(58.7)
@@ -652,3 +652,108 @@ def test_hilbert(self):
              -0.014977068692269, -0.158416139047557]
         m = minimum_phase(h, 'hilbert', n_fft=2**19)
         xp_assert_close(m, k, rtol=2e-3)
+
+
+class Testfirwin_2d:
+    def test_invalid_args(self):
+        with pytest.raises(ValueError, 
+                           match="hsize must be a 2-element tuple or list"):
+            firwin_2d((50,), window=(("kaiser", 5.0), "boxcar"), fc=0.4)
+        
+        with pytest.raises(ValueError, 
+                           match="window must be a 2-element tuple or list"):
+            firwin_2d((51, 51), window=("hamming",), fc=0.5)
+        
+        with pytest.raises(ValueError, 
+                           match="window must be a 2-element tuple or list"):
+            firwin_2d((51, 51), window="invalid_window", fc=0.5)
+
+    def test_filter_design(self):
+        hsize = (51, 51)
+        window = (("kaiser", 8.0), ("kaiser", 8.0))
+        fc = 0.4
+        taps_kaiser = firwin_2d(hsize, window, fc=fc)
+        assert taps_kaiser.shape == (51, 51)
+
+        window = ("hamming", "hamming")
+        taps_hamming = firwin_2d(hsize, window, fc=fc)
+        assert taps_hamming.shape == (51, 51)
+
+    def test_impulse_response(self):
+        hsize = (31, 31)
+        window = ("hamming", "hamming")
+        fc = 0.4
+        taps = firwin_2d(hsize, window, fc=fc)
+
+        impulse = np.zeros((63, 63))
+        impulse[31, 31] = 1
+
+        response = convolve2d(impulse, taps, mode='same')
+
+        expected_response = taps
+        xp_assert_close(response[16:47, 16:47], expected_response, rtol=1e-5)
+
+    def test_frequency_response(self):
+        """Compare 1d and 2d frequency response. """
+        hsize = (31, 31)
+        windows = ("hamming", "hamming")
+        fc = 0.4
+        taps_1d = firwin(numtaps=hsize[0], cutoff=fc, window=windows[0])
+        taps_2d = firwin_2d(hsize, windows, fc=fc)
+
+        f_resp_1d = fft(taps_1d)
+        f_resp_2d = fft2(taps_2d)
+
+        xp_assert_close(f_resp_2d[0, :], f_resp_1d,
+                        err_msg='DC Gain at (0, f1) is not unity!')
+        xp_assert_close(f_resp_2d[:, 0], f_resp_1d,
+                        err_msg='DC Gain at (f0, 0) is not unity!')
+        xp_assert_close(f_resp_2d, np.outer(f_resp_1d, f_resp_1d),
+                        atol=np.finfo(f_resp_2d.dtype).resolution,
+                        err_msg='2d frequency response is not product of 1d responses')
+
+    def test_symmetry(self):
+        hsize = (51, 51)
+        window = ("hamming", "hamming")
+        fc = 0.4
+        taps = firwin_2d(hsize, window, fc=fc)
+        xp_assert_close(taps, np.flip(taps), rtol=1e-5)
+
+    def test_circular_symmetry(self):
+        hsize = (51, 51)
+        window = "hamming"
+        taps = firwin_2d(hsize, window, circular=True, fc=0.5)
+        center = hsize[0] // 2
+        for i in range(hsize[0]):
+            for j in range(hsize[1]):
+                xp_assert_close(taps[i, j], 
+                                taps[center - (i - center), center - (j - center)], 
+                                rtol=1e-5)
+
+    def test_edge_case_circular(self):
+        hsize = (3, 3)
+        window = "hamming"
+        taps_small = firwin_2d(hsize, window, circular=True, fc=0.5)
+        assert taps_small.shape == (3, 3)
+
+        hsize = (101, 101)
+        taps_large = firwin_2d(hsize, window, circular=True, fc=0.5)
+        assert taps_large.shape == (101, 101)
+
+    def test_known_result(self):
+        hsize = (5, 5)
+        window = ('kaiser', 8.0)
+        fc = 0.1
+        fs = 2
+
+        row_filter = firwin(hsize[0], cutoff=fc, window=window, fs=fs)
+        col_filter = firwin(hsize[1], cutoff=fc, window=window, fs=fs)
+        known_result = np.outer(row_filter, col_filter)
+
+        taps = firwin_2d(hsize, (window, window), fc=fc)
+        assert taps.shape == known_result.shape, (
+            f"Shape mismatch: {taps.shape} vs {known_result.shape}"
+        )
+        assert np.allclose(taps, known_result, rtol=1e-1), (
+            f"Filter shape mismatch: {taps} vs {known_result}"
+        )