ezmsg-sigproc 1.5.0__py3-none-any.whl → 1.7.0__py3-none-any.whl

ezmsg/sigproc/scaler.py

@@ -1,9 +1,12 @@
+import functools
 import typing
 
 import numpy as np
 import numpy.typing as npt
+import scipy.signal
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -27,9 +30,139 @@ def _alpha_from_tau(tau: float, dt: float) -> float:
     return 1 - np.exp(-dt / tau)
 
 
+def ewma_step(
+    sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None
+):
+    """
+    Do an exponentially weighted moving average step.
+
+    Args:
+        sample: The new sample.
+        zi: The output of the previous step.
+        alpha: Fading factor.
+        beta: Persisting factor. If None, it is calculated as 1-alpha.
+
+    Returns:
+        alpha * sample + beta * zi
+
+    """
+    # Potential micro-optimization:
+    #  Current: scalar-arr multiplication, scalar-arr multiplication, arr-arr addition
+    #  Alternative: arr-arr subtraction, arr-arr multiplication, arr-arr addition
+    # return zi + alpha * (new_sample - zi)
+    beta = beta or (1 - alpha)
+    return alpha * sample + beta * zi
+
+
+class EWMA:
+    def __init__(self, alpha: float):
+        self.beta = 1 - alpha
+        self._filt_func = functools.partial(
+            scipy.signal.lfilter, [alpha], [1.0, alpha - 1.0], axis=0
+        )
+        self.prev = None
+
+    def compute(self, arr: npt.NDArray) -> npt.NDArray:
+        if self.prev is None:
+            self.prev = self.beta * arr[:1]
+        expected, self.prev = self._filt_func(arr, zi=self.prev)
+        return expected
+
+
+class EWMA_Deprecated:
+    """
+    Grabbed these methods from https://stackoverflow.com/a/70998068 and other answers in that topic,
+    but they ended up being slower than the scipy.signal.lfilter method.
+    Additionally, `compute` and `compute2` suffer from potential errors as the vector length increases
+    and beta**n approaches zero.
+    """
+
+    def __init__(self, alpha: float, max_len: int):
+        self.alpha = alpha
+        self.beta = 1 - alpha
+        self.prev: npt.NDArray | None = None
+        self.weights = np.empty((max_len + 1,), float)
+        self._precalc_weights(max_len)
+        self._step_func = functools.partial(ewma_step, alpha=self.alpha, beta=self.beta)
+
+    def _precalc_weights(self, n: int):
+        #   (1-α)^0, (1-α)^1, (1-α)^2, ..., (1-α)^n
+        np.power(self.beta, np.arange(n + 1), out=self.weights)
+
+    def compute(self, arr: npt.NDArray, out: npt.NDArray | None = None) -> npt.NDArray:
+        if out is None:
+            out = np.empty(arr.shape, arr.dtype)
+
+        n = arr.shape[0]
+        weights = self.weights[:n]
+        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
+
+        #   α*P0, α*P1, α*P2, ..., α*Pn
+        np.multiply(self.alpha, arr, out)
+
+        #   α*P0/(1-α)^0, α*P1/(1-α)^1, α*P2/(1-α)^2, ..., α*Pn/(1-α)^n
+        np.divide(out, weights, out)
+
+        #   α*P0/(1-α)^0, α*P0/(1-α)^0 + α*P1/(1-α)^1, ...
+        np.cumsum(out, axis=0, out=out)
+
+        #   (α*P0/(1-α)^0)*(1-α)^0, (α*P0/(1-α)^0 + α*P1/(1-α)^1)*(1-α)^1, ...
+        np.multiply(out, weights, out)
+
+        # Add the previous output
+        if self.prev is None:
+            self.prev = arr[:1]
+
+        out += self.prev * np.expand_dims(
+            self.weights[1 : n + 1], list(range(1, arr.ndim))
+        )
+
+        self.prev = out[-1:]
+
+        return out
+
+    def compute2(self, arr: npt.NDArray) -> npt.NDArray:
+        """
+        Compute the Exponentially Weighted Moving Average (EWMA) of the input array.
+
+        Args:
+            arr: The input array to be smoothed.
+
+        Returns:
+            The smoothed array.
+        """
+        n = arr.shape[0]
+        if n > len(self.weights):
+            self._precalc_weights(n)
+        weights = self.weights[:n][::-1]
+        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
+
+        result = np.cumsum(self.alpha * weights * arr, axis=0)
+        result = result / weights
+
+        # Handle the first call when prev is unset
+        if self.prev is None:
+            self.prev = arr[:1]
+
+        result += self.prev * np.expand_dims(
+            self.weights[1 : n + 1], list(range(1, arr.ndim))
+        )
+
+        # Store the result back into prev
+        self.prev = result[-1]
+
+        return result
+
+    def compute_sample(self, new_sample: npt.NDArray) -> npt.NDArray:
+        if self.prev is None:
+            self.prev = new_sample
+        self.prev = self._step_func(new_sample, self.prev)
+        return self.prev
+
+
 @consumer
 def scaler(
-    time_constant: float = 1.0, axis: typing.Optional[str] = None
+    time_constant: float = 1.0, axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Apply the adaptive standard scaler from https://riverml.xyz/latest/api/preprocessing/AdaptiveStandardScaler/
@@ -77,7 +210,7 @@ def scaler(
 
 @consumer
 def scaler_np(
-    time_constant: float = 1.0, axis: typing.Optional[str] = None
+    time_constant: float = 1.0, axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Create a generator function that applies an adaptive standard scaler.
@@ -95,10 +228,8 @@ def scaler_np(
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    alpha: float = 0.0
-    means: typing.Optional[npt.NDArray] = None
-    vars_means: typing.Optional[npt.NDArray] = None
-    vars_sq_means: typing.Optional[npt.NDArray] = None
+    samps_ewma: EWMA | None = None
+    vars_sq_ewma: EWMA | None = None
 
     # Reset if input changes
     check_input = {
@@ -107,45 +238,32 @@ def scaler_np(
         "key": None,  # Key change implies buffered means/vars are invalid.
     }
 
-    def _ew_update(arr, prev, _alpha):
-        if np.all(prev == 0):
-            return arr
-        # return _alpha * arr + (1 - _alpha) * prev
-        # Micro-optimization: sub, mult, add (below) is faster than sub, mult, mult, add (above)
-        return prev + _alpha * (arr - prev)
-
     while True:
         msg_in: AxisArray = yield msg_out
 
         axis = axis or msg_in.dims[0]
         axis_idx = msg_in.get_axis_idx(axis)
 
-        if msg_in.axes[axis].gain != check_input["gain"]:
-            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
-            check_input["gain"] = msg_in.axes[axis].gain
-
         data: npt.NDArray = np.moveaxis(msg_in.data, axis_idx, 0)
         b_reset = data.shape[1:] != check_input["shape"]
-        b_reset |= msg_in.key != check_input["key"]
+        b_reset = b_reset or msg_in.axes[axis].gain != check_input["gain"]
+        b_reset = b_reset or msg_in.key != check_input["key"]
         if b_reset:
             check_input["shape"] = data.shape[1:]
+            check_input["gain"] = msg_in.axes[axis].gain
             check_input["key"] = msg_in.key
-            vars_sq_means = np.zeros_like(data[0], dtype=float)
-            vars_means = np.zeros_like(data[0], dtype=float)
-            means = np.zeros_like(data[0], dtype=float)
-
-        result = np.zeros_like(data)
-        for sample_ix in range(data.shape[0]):
-            sample = data[sample_ix]
-            # Update step
-            vars_means = _ew_update(sample, vars_means, alpha)
-            vars_sq_means = _ew_update(sample**2, vars_sq_means, alpha)
-            means = _ew_update(sample, means, alpha)
-            # Get step
-            varis = vars_sq_means - vars_means**2
-            y = (sample - means) / (varis**0.5)
-            result[sample_ix] = y
+            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
+            samps_ewma = EWMA(alpha=alpha)
+            vars_sq_ewma = EWMA(alpha=alpha)
+
+        # Update step
+        means = samps_ewma.compute(data)
+        vars_sq_means = vars_sq_ewma.compute(data**2)
 
+        # Get step
+        varis = vars_sq_means - means**2
+        with np.errstate(divide="ignore", invalid="ignore"):
+            result = (data - means) / (varis**0.5)
         result[np.isnan(result)] = 0.0
         result = np.moveaxis(result, 0, axis_idx)
         msg_out = replace(msg_in, data=result)
@@ -158,7 +276,7 @@ class AdaptiveStandardScalerSettings(ez.Settings):
     """
 
     time_constant: float = 1.0
-    axis: typing.Optional[str] = None
+    axis: str | None = None
 
 
 class AdaptiveStandardScaler(GenAxisArray):

ezmsg/sigproc/signalinjector.py

@@ -1,21 +1,22 @@
 import typing
 
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 import numpy as np
 import numpy.typing as npt
 
 
 class SignalInjectorSettings(ez.Settings):
     time_dim: str = "time"  # Input signal needs a time dimension with units in sec.
-    frequency: typing.Optional[float] = None  # Hz
+    frequency: float | None = None  # Hz
     amplitude: float = 1.0
-    mixing_seed: typing.Optional[int] = None
+    mixing_seed: int | None = None
 
 
 class SignalInjectorState(ez.State):
-    cur_shape: typing.Optional[typing.Tuple[int, ...]] = None
-    cur_frequency: typing.Optional[float] = None
+    cur_shape: tuple[int, ...] | None = None
+    cur_frequency: float | None = None
     cur_amplitude: float
     mixing: npt.NDArray
 
@@ -29,7 +30,7 @@ class SignalInjector(ez.Unit):
     SETTINGS = SignalInjectorSettings
     STATE = SignalInjectorState
 
-    INPUT_FREQUENCY = ez.InputStream(typing.Optional[float])
+    INPUT_FREQUENCY = ez.InputStream(float | None)
     INPUT_AMPLITUDE = ez.InputStream(float)
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -40,7 +41,7 @@ class SignalInjector(ez.Unit):
         self.STATE.mixing = np.array([])
 
     @ez.subscriber(INPUT_FREQUENCY)
-    async def on_frequency(self, msg: typing.Optional[float]) -> None:
+    async def on_frequency(self, msg: float | None) -> None:
         self.STATE.cur_frequency = msg
 
     @ez.subscriber(INPUT_AMPLITUDE)

ezmsg/sigproc/slicer.py

@@ -21,8 +21,8 @@ Slicer:Select a subset of data along a particular axis.
 
 def parse_slice(
     s: str,
-    axinfo: typing.Optional[AxisArray.CoordinateAxis] = None,
-) -> typing.Tuple[typing.Union[slice, int], ...]:
+    axinfo: AxisArray.CoordinateAxis | None = None,
+) -> tuple[slice | int, ...]:
     """
     Parses a string representation of a slice and returns a tuple of slice objects.
 
@@ -63,7 +63,7 @@ def parse_slice(
 
 @consumer
 def slicer(
-    selection: str = "", axis: typing.Optional[str] = None
+    selection: str = "", axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Slice along a particular axis.
@@ -80,8 +80,8 @@ def slicer(
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    _slice: typing.Optional[typing.Union[slice, npt.NDArray]] = None
-    new_axis: typing.Optional[AxisBase] = None
+    _slice: slice | npt.NDArray | None = None
+    new_axis: AxisBase | None = None
     b_change_dims: bool = False  # If number of dimensions changes when slicing
 
     # Reset if input changes
@@ -154,7 +154,7 @@ def slicer(
 
 class SlicerSettings(ez.Settings):
     selection: str = ""
-    axis: typing.Optional[str] = None
+    axis: str | None = None
 
 
 class Slicer(GenAxisArray):

ezmsg/sigproc/spectrogram.py

@@ -12,12 +12,12 @@ from .base import GenAxisArray
 
 @consumer
 def spectrogram(
-    window_dur: typing.Optional[float] = None,
-    window_shift: typing.Optional[float] = None,
+    window_dur: float | None = None,
+    window_shift: float | None = None,
     window: WindowFunction = WindowFunction.HANNING,
     transform: SpectralTransform = SpectralTransform.REL_DB,
     output: SpectralOutput = SpectralOutput.POSITIVE,
-) -> typing.Generator[typing.Optional[AxisArray], AxisArray, None]:
+) -> typing.Generator[AxisArray | None, AxisArray, None]:
     """
     Calculate a spectrogram on streaming data.
 
@@ -50,7 +50,7 @@ def spectrogram(
     )
 
     # State variables
-    msg_out: typing.Optional[AxisArray] = None
+    msg_out: AxisArray | None = None
 
     while True:
         msg_in: AxisArray = yield msg_out
@@ -63,8 +63,8 @@ class SpectrogramSettings(ez.Settings):
     See :obj:`spectrogram` for a description of the parameters.
     """
 
-    window_dur: typing.Optional[float] = None  # window duration in seconds
-    window_shift: typing.Optional[float] = None
+    window_dur: float | None = None  # window duration in seconds
+    window_shift: float | None = None
     """"window step in seconds. If None, window_shift == window_dur"""
 
     # See SpectrumSettings for details of following settings:

ezmsg/sigproc/spectrum.py

@@ -68,14 +68,14 @@ class SpectralOutput(OptionsEnum):
 
 @consumer
 def spectrum(
-    axis: typing.Optional[str] = None,
-    out_axis: typing.Optional[str] = "freq",
+    axis: str | None = None,
+    out_axis: str | None = "freq",
     window: WindowFunction = WindowFunction.HANNING,
     transform: SpectralTransform = SpectralTransform.REL_DB,
     output: SpectralOutput = SpectralOutput.POSITIVE,
-    norm: typing.Optional[str] = "forward",
+    norm: str | None = "forward",
     do_fftshift: bool = True,
-    nfft: typing.Optional[int] = None,
+    nfft: int | None = None,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Calculate a spectrum on a data slice.
@@ -105,10 +105,10 @@ def spectrum(
     apply_window = window != WindowFunction.NONE
     do_fftshift &= output == SpectralOutput.FULL
     f_sl = slice(None)
-    freq_axis: typing.Optional[AxisArray.LinearAxis] = None
-    fftfun: typing.Optional[typing.Callable] = None
-    f_transform: typing.Optional[typing.Callable] = None
-    new_dims: typing.Optional[typing.List[str]] = None
+    freq_axis: AxisArray.LinearAxis | None = None
+    fftfun: typing.Callable | None = None
+    f_transform: typing.Callable | None = None
+    new_dims: list[str] | None = None
 
     # Reset if input changes substantially
     check_input = {
@@ -238,9 +238,9 @@ class SpectrumSettings(ez.Settings):
     See :obj:`spectrum` for a description of the parameters.
     """
 
-    axis: typing.Optional[str] = None
-    # n: typing.Optional[int] = None # n parameter for fft
-    out_axis: typing.Optional[str] = "freq"  # If none; don't change dim name
+    axis: str | None = None
+    # n: int | None = None # n parameter for fft
+    out_axis: str | None = "freq"  # If none; don't change dim name
     window: WindowFunction = WindowFunction.HAMMING
     transform: SpectralTransform = SpectralTransform.REL_DB
     output: SpectralOutput = SpectralOutput.POSITIVE

ezmsg/sigproc/synth.py

@@ -1,18 +1,19 @@
 import asyncio
 from dataclasses import field
 import time
-from typing import Optional, Generator, AsyncGenerator, Union
+import typing
 
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from .butterworthfilter import ButterworthFilter, ButterworthFilterSettings
 from .base import GenAxisArray
 
 
-def clock(dispatch_rate: Optional[float]) -> Generator[ez.Flag, None, None]:
+def clock(dispatch_rate: float | None) -> typing.Generator[ez.Flag, None, None]:
     """
     Construct a generator that yields events at a specified rate.
 
@@ -32,7 +33,7 @@ def clock(dispatch_rate: Optional[float]) -> Generator[ez.Flag, None, None]:
         yield ez.Flag()
 
 
-async def aclock(dispatch_rate: Optional[float]) -> AsyncGenerator[ez.Flag, None]:
+async def aclock(dispatch_rate: float | None) -> typing.AsyncGenerator[ez.Flag, None]:
     """
     ``asyncio`` version of :obj:`clock`.
 
@@ -53,12 +54,12 @@ class ClockSettings(ez.Settings):
     """Settings for :obj:`Clock`. See :obj:`clock` for parameter description."""
 
     # Message dispatch rate (Hz), or None (fast as possible)
-    dispatch_rate: Optional[float]
+    dispatch_rate: float | None
 
 
 class ClockState(ez.State):
     cur_settings: ClockSettings
-    gen: AsyncGenerator
+    gen: typing.AsyncGenerator
 
 
 class Clock(ez.Unit):
@@ -83,7 +84,7 @@ class Clock(ez.Unit):
         self.construct_generator()
 
     @ez.publisher(OUTPUT_CLOCK)
-    async def generate(self) -> AsyncGenerator:
+    async def generate(self) -> typing.AsyncGenerator:
         while True:
             out = await self.STATE.gen.__anext__()
             if out:
@@ -93,11 +94,11 @@ class Clock(ez.Unit):
 # COUNTER - Generate incrementing integer. fs and dispatch_rate parameters combine to give many options. #
 async def acounter(
     n_time: int,
-    fs: Optional[float],
+    fs: float | None,
     n_ch: int = 1,
-    dispatch_rate: Optional[Union[float, str]] = None,
-    mod: Optional[int] = None,
-) -> AsyncGenerator[AxisArray, None]:
+    dispatch_rate: float | str | None = None,
+    mod: int | None = None,
+) -> typing.AsyncGenerator[AxisArray, None]:
     """
     Construct an asynchronous generator to generate AxisArray objects at a specified rate
     and with the specified sampling rate.
@@ -206,14 +207,14 @@ class CounterSettings(ez.Settings):
     # Message dispatch rate (Hz), 'realtime', 'ext_clock', or None (fast as possible)
     #  Note: if dispatch_rate is a float then time offsets will be synthetic and the
     #  system will run faster or slower than wall clock time.
-    dispatch_rate: Optional[Union[float, str]] = None
+    dispatch_rate: float | str | None = None
 
     # If set to an integer, counter will rollover
-    mod: Optional[int] = None
+    mod: int | None = None
 
 
 class CounterState(ez.State):
-    gen: AsyncGenerator[AxisArray, Optional[ez.Flag]]
+    gen: typing.AsyncGenerator[AxisArray, ez.Flag | None]
     cur_settings: CounterSettings
     new_generator: asyncio.Event
 
@@ -262,7 +263,7 @@ class Counter(ez.Unit):
             yield self.OUTPUT_SIGNAL, out
 
     @ez.publisher(OUTPUT_SIGNAL)
-    async def run_generator(self) -> AsyncGenerator:
+    async def run_generator(self) -> typing.AsyncGenerator:
         while True:
             await self.STATE.new_generator.wait()
             self.STATE.new_generator.clear()
@@ -277,11 +278,11 @@ class Counter(ez.Unit):
 
 @consumer
 def sin(
-    axis: Optional[str] = "time",
+    axis: str | None = "time",
     freq: float = 1.0,
     amp: float = 1.0,
     phase: float = 0.0,
-) -> Generator[AxisArray, AxisArray, None]:
+) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Construct a generator of sinusoidal waveforms in AxisArray objects.
 
@@ -319,7 +320,7 @@ class SinGeneratorSettings(ez.Settings):
     See :obj:`sin` for parameter descriptions.
     """
 
-    time_axis: Optional[str] = "time"
+    time_axis: str | None = "time"
     freq: float = 1.0  # Oscillation frequency in Hz
     amp: float = 1.0  # Amplitude
     phase: float = 0.0  # Phase offset (in radians)
@@ -353,7 +354,7 @@ class OscillatorSettings(ez.Settings):
     n_ch: int = 1
     """Number of channels to output per block"""
 
-    dispatch_rate: Optional[Union[float, str]] = None
+    dispatch_rate: float | str | None = None
     """(Hz) | 'realtime' | 'ext_clock'"""
 
     freq: float = 1.0
@@ -435,7 +436,7 @@ class RandomGenerator(ez.Unit):
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def generate(self, msg: AxisArray) -> AsyncGenerator:
+    async def generate(self, msg: AxisArray) -> typing.AsyncGenerator:
         random_data = np.random.normal(
             size=msg.shape, loc=self.SETTINGS.loc, scale=self.SETTINGS.scale
         )
@@ -451,7 +452,7 @@ class NoiseSettings(ez.Settings):
     n_time: int  # Number of samples to output per block
     fs: float  # Sampling rate of signal output in Hz
     n_ch: int = 1  # Number of channels to output
-    dispatch_rate: Optional[Union[float, str]] = None
+    dispatch_rate: float | str | None = None
     """(Hz), 'realtime', or 'ext_clock'"""
     loc: float = 0.0  # DC offset
     scale: float = 1.0  # Scale (in standard deviations)
@@ -553,12 +554,12 @@ class Add(ez.Unit):
         self.STATE.queue_b.put_nowait(msg)
 
     @ez.publisher(OUTPUT_SIGNAL)
-    async def output(self) -> AsyncGenerator:
+    async def output(self) -> typing.AsyncGenerator:
         while True:
             a = await self.STATE.queue_a.get()
             b = await self.STATE.queue_b.get()
 
-            yield (self.OUTPUT_SIGNAL, replace(a, data=a.data + b.data))
+            yield self.OUTPUT_SIGNAL, replace(a, data=a.data + b.data)
 
 
 class EEGSynthSettings(ez.Settings):

ezmsg/sigproc/wavelets.py

@@ -4,7 +4,8 @@ import numpy as np
 import numpy.typing as npt
 import pywt
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -13,44 +14,61 @@ from .filterbank import filterbank, FilterbankMode, MinPhaseMode
 
 @consumer
 def cwt(
-    scales: typing.Union[list, tuple, npt.NDArray],
-    wavelet: typing.Union[str, pywt.ContinuousWavelet, pywt.Wavelet],
+    frequencies: list | tuple | npt.NDArray | None,
+    wavelet: str | pywt.ContinuousWavelet | pywt.Wavelet,
     min_phase: MinPhaseMode = MinPhaseMode.NONE,
     axis: str = "time",
+    scales: list | tuple | npt.NDArray | None = None,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Perform a continuous wavelet transform.
     The function is equivalent to the :obj:`pywt.cwt` function, but is designed to work with streaming data.
 
     Args:
-        scales: The wavelet scales to use. Note: Scales will be sorted from largest to smallest.
+        frequencies: The wavelet frequencies to use in Hz. If `None` provided then the scales will be used.
+          Note: frequencies will be sorted from smallest to largest.
         wavelet: Wavelet object or name of wavelet to use.
         min_phase: See filterbank MinPhaseMode for details.
         axis: The target axis for operation. Note that this will be moved to the -1th dimension
           because fft and matrix multiplication is much faster on the last axis.
           This axis must be in the msg.axes and it must be of type AxisArray.LinearAxis.
+        scales: The scales to use. If None, the scales will be calculated from the frequencies.
+          Note: Scales will be sorted from largest to smallest.
+          Note: Use of scales is deprecated in favor of frequencies. Convert scales to frequencies using
+            `pywt.scale2frequency(wavelet, scales, precision=10) * fs` where fs is the sampling frequency.
 
     Returns:
         A primed Generator object that expects an :obj:`AxisArray` via `.send(axis_array)` of continuous data
         and yields an :obj:`AxisArray` with a continuous wavelet transform in its data.
     """
-    msg_out: typing.Optional[AxisArray] = None
+    precision = 10
+    msg_out: AxisArray | None = None
 
     # Check parameters
-    scales = np.sort(scales)[::-1]
-    assert np.all(scales > 0), "Scales must be positive."
-    assert scales.ndim == 1, "Scales must be a 1D list, tuple, or array."
+    if frequencies is None and scales is None:
+        raise ValueError("Either frequencies or scales must be provided.")
+    if frequencies is not None and scales is not None:
+        raise ValueError("Only one of frequencies or scales can be provided.")
+    if scales is not None:
+        scales = np.sort(scales)[::-1]
+        assert np.all(scales > 0), "scales must be positive."
+        assert scales.ndim == 1, "scales must be a 1D list, tuple, or array."
+
     if not isinstance(wavelet, (pywt.ContinuousWavelet, pywt.Wavelet)):
         wavelet = pywt.DiscreteContinuousWavelet(wavelet)
-    precision = 10
+
+    if frequencies is not None:
+        frequencies = np.sort(frequencies)
+        assert np.all(frequencies > 0), "frequencies must be positive."
+        assert frequencies.ndim == 1, "frequencies must be a 1D list, tuple, or array."
 
     # State variables
-    neg_rt_scales = -np.sqrt(scales)[:, None]
+    neg_rt_scales: npt.NDArray | None = None
     int_psi, wave_xvec = pywt.integrate_wavelet(wavelet, precision=precision)
     int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
-    template: typing.Optional[AxisArray] = None
-    fbgen: typing.Optional[typing.Generator[AxisArray, AxisArray, None]] = None
-    last_conv_samp: typing.Optional[npt.NDArray] = None
+    template: AxisArray | None = None
+    fbgen: typing.Generator[AxisArray, AxisArray, None] | None = None
+    last_conv_samp: npt.NDArray | None = None
 
     # Reset if input changed
     check_input = {
@@ -76,6 +94,12 @@ def cwt(
             check_input["shape"] = in_shape
             check_input["key"] = msg_in.key
 
+            if frequencies is not None:
+                scales = pywt.frequency2scale(
+                    wavelet, frequencies * msg_in.axes[axis].gain, precision=precision
+                )
+            neg_rt_scales = -np.sqrt(scales)[:, None]
+
             # convert int_psi, wave_xvec to the same precision as the data
             dt_data = msg_in.data.dtype  # _check_dtype(msg_in.data)
             dt_cplx = np.result_type(dt_data, np.complex64)
@@ -148,8 +172,8 @@ class CWTSettings(ez.Settings):
     See :obj:`cwt` for argument details.
     """
 
-    scales: typing.Union[list, tuple, npt.NDArray]
-    wavelet: typing.Union[str, pywt.ContinuousWavelet, pywt.Wavelet]
+    scales: list | tuple | npt.NDArray
+    wavelet: str | pywt.ContinuousWavelet | pywt.Wavelet
     min_phase: MinPhaseMode = MinPhaseMode.NONE
     axis: str = "time"