Add sdr.design_bandstop_fir()

src/sdr/_filter/_design_fir.py

@@ -60,6 +60,16 @@ def _ideal_bandpass(order: int, center_freq: float, bandwidth: float) -> npt.NDA
     return h_ideal
 
 
+def _ideal_bandstop(order: int, center_freq: float, bandwidth: float) -> npt.NDArray[np.float_]:
+    """
+    Returns the ideal bandstop filter impulse response.
+    """
+    h_lp = _ideal_lowpass(order, center_freq - bandwidth / 2)
+    h_hp = _ideal_highpass(order, center_freq + bandwidth / 2)
+    h_ideal = h_lp + h_hp
+    return h_ideal
+
+
 def _window(
     order: int,
     window: Literal["hamming", "hann", "blackman", "blackman-harris", "chebyshev", "kaiser"]
@@ -391,3 +401,112 @@ def design_bandpass_fir(
     h = _normalize_passband(h, center_freq)
 
     return h
+
+
+@export
+def design_bandstop_fir(
+    order: int,
+    center_freq: float,
+    bandwidth: float,
+    window: None
+    | Literal["hamming", "hann", "blackman", "blackman-harris", "chebyshev", "kaiser"]
+    | npt.ArrayLike = "hamming",
+    atten: float = 60,
+) -> npt.NDArray[np.float_]:
+    r"""
+    Designs a bandstop FIR filter impulse response $h[n]$ using the window method.
+
+    Arguments:
+        order: The filter order $N$. Must be even.
+        center_freq: The center frequency $f_{center}$, normalized to the Nyquist frequency $f_s / 2$.
+        bandwidth: The two-sided bandwidth about $f_{center}$, normalized to the Nyquist frequency $f_s / 2$.
+        window: The time-domain window to use.
+
+            - `None`: No windowing. Equivalently, a length-$N + 1$ vector of ones.
+            - `"hamming"`: Hamming window, see :func:`scipy.signal.windows.hamming`.
+            - `"hann"`: Hann window, see :func:`scipy.signal.windows.hann`.
+            - `"blackman"`: Blackman window, see :func:`scipy.signal.windows.blackman`.
+            - `"blackman-harris"`: Blackman-Harris window, see :func:`scipy.signal.windows.blackmanharris`.
+            - `"chebyshev"`: Chebyshev window, see :func:`scipy.signal.windows.chebwin`.
+            - `"kaiser"`: Kaiser window, see :func:`scipy.signal.windows.kaiser`.
+            - `npt.ArrayLike`: A custom window. Must be a length-$N + 1$ vector.
+
+        atten: The sidelobe attenuation in dB. Only used if `window` is `"chebyshev"` or `"kaiser"`.
+
+    Returns:
+        The filter impulse response $h[n]$ with length $N + 1$. The center of the larger passband has 0 dB gain.
+
+    References:
+        - https://www.mathworks.com/help/dsp/ref/designbandstopfir.html
+
+    Examples:
+        Design a length-101 bandstop FIR filter with center frequency $f_{center} = 0.4 \cdot f_s / 2$
+        and bandwidth $0.75 \cdot f_s / 2$, using a Hamming window.
+
+        .. ipython:: python
+
+            h_hamming = sdr.design_bandstop_fir(100, 0.4, 0.75, window="hamming")
+
+            @savefig sdr_design_bandstop_fir_1.png
+            plt.figure(figsize=(8, 4)); \
+            sdr.plot.impulse_response(h_hamming);
+
+            @savefig sdr_design_bandstop_fir_2.png
+            plt.figure(figsize=(8, 4)); \
+            sdr.plot.magnitude_response(h_hamming);
+
+        Compare filter designs using different windows.
+
+        .. ipython:: python
+
+            h_hann = sdr.design_bandstop_fir(100, 0.4, 0.75, window="hann"); \
+            h_blackman = sdr.design_bandstop_fir(100, 0.4, 0.75, window="blackman"); \
+            h_blackman_harris = sdr.design_bandstop_fir(100, 0.4, 0.75, window="blackman-harris"); \
+            h_chebyshev = sdr.design_bandstop_fir(100, 0.4, 0.75, window="chebyshev"); \
+            h_kaiser = sdr.design_bandstop_fir(100, 0.4, 0.75, window="kaiser")
+
+            @savefig sdr_design_bandstop_fir_3.png
+            plt.figure(figsize=(8, 4)); \
+            sdr.plot.magnitude_response(h_hamming, label="Hamming"); \
+            sdr.plot.magnitude_response(h_hann, label="Hann"); \
+            sdr.plot.magnitude_response(h_blackman, label="Blackman"); \
+            sdr.plot.magnitude_response(h_blackman_harris, label="Blackman-Harris"); \
+            sdr.plot.magnitude_response(h_chebyshev, label="Chebyshev"); \
+            sdr.plot.magnitude_response(h_kaiser, label="Kaiser"); \
+            plt.legend(); \
+            plt.ylim(-100, 10);
+
+    Group:
+        dsp-fir-filtering
+    """
+    if not isinstance(order, int):
+        raise TypeError(f"Argument 'order' must be an integer, not {type(order).__name__}.")
+    if not order % 2 == 0:
+        raise ValueError(f"Argument 'order' must be even, not {order}.")
+
+    if not isinstance(center_freq, (int, float)):
+        raise TypeError(f"Argument 'center_freq' must be a number, not {type(center_freq).__name__}.")
+    if not 0 <= center_freq <= 1:
+        raise ValueError(f"Argument 'center_freq' must be between 0 and 1, not {center_freq}.")
+
+    if not isinstance(bandwidth, (int, float)):
+        raise TypeError(f"Argument 'bandwidth' must be a number, not {type(bandwidth).__name__}.")
+    if not 0 <= center_freq + bandwidth / 2 <= 1:
+        raise ValueError(f"Argument 'bandwidth' must be between 0 and 0.5 - 'center_freq', not {bandwidth}.")
+    if not 0 <= center_freq - bandwidth / 2 <= 1:
+        raise ValueError(f"Argument 'bandwidth' must be between 0 and 'center_freq', not {bandwidth}.")
+
+    if not isinstance(atten, (int, float)):
+        raise TypeError(f"Argument 'atten' must be a number, not {type(atten).__name__}.")
+    if not 0 <= atten:
+        raise ValueError(f"Argument 'atten' must be non-negative, not {atten}.")
+
+    h_ideal = _ideal_bandstop(order, center_freq, bandwidth)
+    h_window = _window(order, window, atten=atten)
+    h = h_ideal * h_window
+    if center_freq > 0.5:
+        h = _normalize_passband(h, 0)
+    else:
+        h = _normalize_passband(h, 1)
+
+    return h