Source code for ezmsg.sigproc.singlebandpow

"""
Time-domain single-band power estimation.

Two methods are provided:

1. **RMS Band Power** — Bandpass filter, square, window into bins, take the mean, optionally take the square root.
2. **Square-Law + LPF Band Power** — Bandpass filter, square, lowpass filter (smoothing), downsample.
"""

from dataclasses import field

import ezmsg.core as ez
from ezmsg.baseproc import (
    BaseProcessor,
    BaseStatefulProcessor,
    BaseTransformerUnit,
    CompositeProcessor,
)
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.util.messages.modify import ModifyAxisSettings, ModifyAxisTransformer

from .aggregate import AggregateSettings, AggregateTransformer, AggregationFunction
from .butterworthfilter import ButterworthFilterSettings, ButterworthFilterTransformer
from .downsample import DownsampleSettings, DownsampleTransformer
from .materialize import MaterializeMode, materialize_array
from .math.pow import PowSettings, PowTransformer
from .window import WindowTransformer




class RMSBandPowerSettings(ez.Settings):
    """Settings for :obj:`RMSBandPowerTransformer`."""

    bandpass: ButterworthFilterSettings = field(
        default_factory=lambda: ButterworthFilterSettings(order=4, coef_type="sos")
    )
    """Butterworth bandpass filter settings. Set ``cuton`` and ``cutoff`` to define the band."""

    bin_duration: float = 0.05
    """Duration of each non-overlapping bin in seconds."""

    apply_sqrt: bool = True
    """If True, output is RMS (root-mean-square). If False, output is mean-square power."""

    materialize: MaterializeMode = MaterializeMode.ASYNC
    """How to evaluate the output on a lazy backend (MLX); see
    :obj:`~ezmsg.sigproc.materialize.MaterializeMode`. No-op elsewhere.

    Defaults to :obj:`~ezmsg.sigproc.materialize.MaterializeMode.ASYNC`, which
    keeps a lazy graph from accumulating without blocking the caller for values
    it has no use for. Set ``OFF`` if a downstream node already materializes
    every cycle, or ``SYNC`` to time this stage's work as its own."""





class RMSBandPowerTransformer(CompositeProcessor[RMSBandPowerSettings, AxisArray, AxisArray]):
    """
    RMS band power estimation.

    Pipeline: bandpass -> square -> window(bins) -> mean(time) -> rename bin->time -> [sqrt]
    """

    @staticmethod
    def _initialize_processors(
        settings: RMSBandPowerSettings,
    ) -> dict[str, BaseProcessor | BaseStatefulProcessor]:
        procs: dict[str, BaseProcessor | BaseStatefulProcessor] = {
            "bandpass": ButterworthFilterTransformer(settings.bandpass),
            "square": PowTransformer(PowSettings(exponent=2.0)),
            "window": WindowTransformer(
                axis="time",
                newaxis="bin",
                window_dur=settings.bin_duration,
                window_shift=settings.bin_duration,
                zero_pad_until="none",
            ),
            "aggregate": AggregateTransformer(AggregateSettings(axis="time", operation=AggregationFunction.MEAN)),
            "rename": ModifyAxisTransformer(settings=ModifyAxisSettings(name_map={"bin": "time"})),
        }
        if settings.apply_sqrt:
            procs["sqrt"] = PowTransformer(PowSettings(exponent=0.5))
        return procs

    def _post_process(self, result: AxisArray | None) -> AxisArray | None:
        if result is not None:
            materialize_array(result.data, self.settings.materialize)
        return result





class RMSBandPower(BaseTransformerUnit[RMSBandPowerSettings, AxisArray, AxisArray, RMSBandPowerTransformer]):
    SETTINGS = RMSBandPowerSettings





class SquareLawBandPowerSettings(ez.Settings):
    """Settings for :obj:`SquareLawBandPowerTransformer`."""

    bandpass: ButterworthFilterSettings = field(
        default_factory=lambda: ButterworthFilterSettings(order=4, coef_type="sos")
    )
    """Butterworth bandpass filter settings. Set ``cuton`` and ``cutoff`` to define the band."""

    lowpass: ButterworthFilterSettings = field(
        default_factory=lambda: ButterworthFilterSettings(order=4, coef_type="sos")
    )
    """Butterworth lowpass filter settings for smoothing the squared signal."""

    downsample: DownsampleSettings = field(default_factory=DownsampleSettings)
    """Downsample settings for rate reduction after lowpass smoothing."""





class SquareLawBandPowerTransformer(CompositeProcessor[SquareLawBandPowerSettings, AxisArray, AxisArray]):
    """
    Square-law + LPF band power estimation.

    Pipeline: bandpass -> square -> lowpass -> downsample
    """

    @staticmethod
    def _initialize_processors(
        settings: SquareLawBandPowerSettings,
    ) -> dict[str, BaseProcessor | BaseStatefulProcessor]:
        return {
            "bandpass": ButterworthFilterTransformer(settings.bandpass),
            "square": PowTransformer(PowSettings(exponent=2.0)),
            "lowpass": ButterworthFilterTransformer(settings.lowpass),
            "downsample": DownsampleTransformer(settings.downsample),
        }





class SquareLawBandPower(
    BaseTransformerUnit[SquareLawBandPowerSettings, AxisArray, AxisArray, SquareLawBandPowerTransformer]
):
    SETTINGS = SquareLawBandPowerSettings