Clean up use of windows in plot functions

mhostetter · Jul 19, 2024 · 801caa7 · 801caa7
1 parent 43870a4
commit 801caa7
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 12 deletions.
diff --git a/src/sdr/_filter/_applications.py b/src/sdr/_filter/_applications.py
@@ -198,8 +198,7 @@ def __init__(
         h[order // 2] = 0
 
         if window is not None:
-            h_win = scipy.signal.windows.get_window(window, order + 1, fftbins=False)
-            h *= h_win
+            h *= scipy.signal.windows.get_window(window, order + 1, fftbins=False)
 
         super().__init__(h, streaming=streaming)
 

diff --git a/src/sdr/plot/_freq_domain.py b/src/sdr/plot/_freq_domain.py
@@ -20,7 +20,7 @@
 def dft(
     x: npt.ArrayLike,
     sample_rate: float | None = None,
-    window: str | None = None,
+    window: str | float | tuple | None = None,
     size: int | None = None,
     oversample: int | None = None,
     fast: bool = False,
@@ -39,7 +39,8 @@ def dft(
         x: The time-domain signal $x[n]$.
         sample_rate: The sample rate $f_s$ of the signal in samples/s. If `None`, the x-axis will
             be labeled as "Normalized frequency".
-        window: The windowing function to use. This can be a string or a vector of length `length`.
+        window: The SciPy window definition. See :func:`scipy.signal.windows.get_window` for details.
+            If `None`, no window is applied.
         size: The number of points to use for the DFT. If `None`, the length of the signal is used.
         oversample: The factor to oversample the DFT. If `None`, the DFT is not oversampled. This is only considered
             if `size` is `None`.
@@ -128,8 +129,7 @@ def dft(
             ax = plt.gca()
 
         if window is not None:
-            w = scipy.signal.windows.get_window(window, x.size)
-            x = x * w
+            x *= scipy.signal.windows.get_window(window, x.size)
 
         if size is None:
             if oversample is None:
@@ -176,7 +176,7 @@ def dft(
 def dtft(
     x: npt.ArrayLike,
     sample_rate: float | None = None,
-    window: str | None = None,
+    window: str | float | tuple | None = None,
     size: int = int(2**20),  # ~1 million points
     centered: bool = True,
     ax: plt.Axes | None = None,
@@ -191,7 +191,8 @@ def dtft(
         x: The time-domain signal $x[n]$.
         sample_rate: The sample rate $f_s$ of the signal in samples/s. If `None`, the x-axis will
             be labeled as "Normalized frequency".
-        window: The windowing function to use. This can be a string or a vector of length `length`.
+        window: The SciPy window definition. See :func:`scipy.signal.windows.get_window` for details.
+            If `None`, no window is applied.
         size: The number of points to use for the DTFT. The actual size used will be the nearest power of 2.
         centered: Indicates whether to center the DTFT about 0.
         ax: The axis to plot on. If `None`, the current axis is used.

diff --git a/src/sdr/plot/_spectral_estimation.py b/src/sdr/plot/_spectral_estimation.py
@@ -21,7 +21,7 @@
 def periodogram(
     x: npt.ArrayLike,
     sample_rate: float | None = None,
-    window: str | npt.ArrayLike = "hann",
+    window: str | float | tuple | None = "hann",
     length: int | None = None,
     overlap: int | None = None,
     fft: int | None = None,
@@ -43,7 +43,8 @@ def periodogram(
         x: The time-domain signal $x[n]$.
         sample_rate: The sample rate $f_s$ of the signal in samples/s. If `None`, the x-axis will
             be labeled as "Normalized frequency".
-        window: The windowing function to use. This can be a string or a vector of length `length`.
+        window: The SciPy window definition. See :func:`scipy.signal.windows.get_window` for details.
+            If `None`, no window is applied.
         length: The length of each segment in samples. If `None`, the length is set to 256.
         overlap: The number of samples to overlap between segments. If `None`, the overlap is set to `length // 2`.
         fft: The number of points to use in the FFT. If `None`, the FFT length is set to `length`.
@@ -128,7 +129,7 @@ def periodogram(
 def spectrogram(
     x: npt.ArrayLike,
     sample_rate: float | None = None,
-    window: str | npt.ArrayLike = "hann",
+    window: str | float | tuple | None = "hann",
     length: int | None = None,
     overlap: int | None = None,
     fft: int | None = None,
@@ -149,7 +150,8 @@ def spectrogram(
         x: The time-domain signal $x[n]$.
         sample_rate: The sample rate $f_s$ of the signal in samples/s. If `None`, the x-axis will
             be label as "Samples" and the y-axis as "Normalized frequency".
-        window: The windowing function to use. This can be a string or a vector of length `length`.
+        window: The SciPy window definition. See :func:`scipy.signal.windows.get_window` for details.
+            If `None`, no window is applied.
         length: The length of each segment in samples. If `None`, the length is set to 256.
         overlap: The number of samples to overlap between segments. If `None`, the overlap is set to `length // 2`.
         fft: The number of points to use in the FFT. If `None`, the FFT length is set to `length`.