ezmsg-sigproc 1.1.1__tar.gz → 1.2.0__tar.gz

{ezmsg-sigproc-1.1.1/ezmsg_sigproc.egg-info → ezmsg-sigproc-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ezmsg-sigproc
-Version: 1.1.1
+Version: 1.2.0
 Summary: Timeseries signal processing implementations in ezmsg
 Home-page: https://github.com/iscoe/ezmsg
 Author: Griffin Milsap

ezmsg-sigproc-1.2.0/ezmsg/sigproc/__init__.py

@@ -0,0 +1 @@
+from .__version__ import __version__

ezmsg-sigproc-1.2.0/ezmsg/sigproc/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.2.0"

{ezmsg-sigproc-1.1.1 → ezmsg-sigproc-1.2.0}/ezmsg/sigproc/butterworthfilter.py

@@ -1,49 +1,40 @@
-from dataclasses import dataclass, field
-import logging
-
 import ezmsg.core as ez
 import scipy.signal
 import numpy as np
 
-from .filter import Filter, FilterState, FilterSettings
+from .filter import Filter, FilterState, FilterSettingsBase
 
 from typing import Optional, Tuple, Union
 
-logger = logging.getLogger('ezmsg')
-
 
-@dataclass( frozen = True )
-class ButterworthFilterDesign:
+class ButterworthFilterSettings(FilterSettingsBase):
     order: int = 0
     cuton: Optional[float] = None  # Hz
     cutoff: Optional[float] = None  # Hz
 
-    def filter_specs( self ) -> Optional[ Tuple[ str, Union[ float, Tuple[ float, float ] ] ] ]:
+    def filter_specs(self) -> Optional[Tuple[str, Union[float, Tuple[float, float]]]]:
         if self.cuton is None and self.cutoff is None:
             return None
         elif self.cuton is None and self.cutoff is not None:
-            return 'lowpass', self.cutoff
+            return "lowpass", self.cutoff
         elif self.cuton is not None and self.cutoff is None:
-            return 'highpass', self.cuton
+            return "highpass", self.cuton
         elif self.cuton is not None and self.cutoff is not None:
-            if self.cuton <= self.cutoff: 
-                return 'bandpass', ( self.cuton, self.cutoff ) 
-            else: return 'bandstop', ( self.cutoff, self.cuton )
-
-
-class ButterworthFilterSettings(ButterworthFilterDesign, FilterSettings):
-    ...
+            if self.cuton <= self.cutoff:
+                return "bandpass", (self.cuton, self.cutoff)
+            else:
+                return "bandstop", (self.cutoff, self.cuton)
 
 
 class ButterworthFilterState(FilterState):
-    design: ButterworthFilterDesign
+    design: ButterworthFilterSettings
 
 
 class ButterworthFilter(Filter):
     SETTINGS: ButterworthFilterSettings
     STATE: ButterworthFilterState
 
-    INPUT_FILTER = ez.InputStream(ButterworthFilterDesign)
+    INPUT_FILTER = ez.InputStream(ButterworthFilterSettings)
 
     def initialize(self) -> None:
         self.STATE.design = self.SETTINGS
@@ -55,18 +46,17 @@ class ButterworthFilter(Filter):
         if self.STATE.design.order > 0 and specs is not None:
             btype, cut = specs
             return scipy.signal.butter(
-                self.STATE.design.order, 
-                Wn=cut, 
-                btype=btype, 
-                fs=self.STATE.fs, 
-                output="ba"
+                self.STATE.design.order,
+                Wn=cut,
+                btype=btype,
+                fs=self.STATE.fs,
+                output="ba",
             )
         else:
             return None
 
-
     @ez.subscriber(INPUT_FILTER)
-    async def redesign(self, message: ButterworthFilterDesign) -> None:
+    async def redesign(self, message: ButterworthFilterSettings) -> None:
         if self.STATE.design.order != message.order:
             self.STATE.zi = None
         self.STATE.design = message

{ezmsg-sigproc-1.1.1 → ezmsg-sigproc-1.2.0}/ezmsg/sigproc/decimate.py

@@ -2,23 +2,22 @@ import ezmsg.core as ez
 
 import scipy.signal
 
+from ezmsg.util.messages.axisarray import AxisArray
+
 from .downsample import Downsample, DownsampleSettings
 from .filter import Filter, FilterCoefficients, FilterSettings
-from .messages import TSMessage as TimeSeriesMessage
 
 
 class Decimate(ez.Collection):
-
     SETTINGS: DownsampleSettings
 
-    INPUT_SIGNAL = ez.InputStream(TimeSeriesMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(TimeSeriesMessage)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     FILTER = Filter()
     DOWNSAMPLE = Downsample()
 
     def configure(self) -> None:
-
         self.DOWNSAMPLE.apply_settings(self.SETTINGS)
 
         if self.SETTINGS.factor < 1:
@@ -27,11 +26,9 @@ class Decimate(ez.Collection):
             filt = FilterCoefficients()
         else:
             # See scipy.signal.decimate for IIR Filter Condition
-            system = scipy.signal.dlti(
-                *scipy.signal.cheby1(8, 0.05, 0.8 / self.SETTINGS.factor)
-            )
-
-            filt = FilterCoefficients(b=system.num, a=system.den)
+            b, a = scipy.signal.cheby1(8, 0.05, 0.8 / self.SETTINGS.factor)
+            system = scipy.signal.dlti(b, a)
+            filt = FilterCoefficients(b=system.num, a=system.den)  # type: ignore
 
         self.FILTER.apply_settings(FilterSettings(filt=filt))
 

ezmsg-sigproc-1.2.0/ezmsg/sigproc/downsample.py

@@ -0,0 +1,63 @@
+from dataclasses import replace
+
+from ezmsg.util.messages.axisarray import AxisArray
+
+import ezmsg.core as ez
+import numpy as np
+
+from typing import (
+    AsyncGenerator,
+    Optional,
+)
+
+
+class DownsampleSettings(ez.Settings):
+    axis: Optional[str] = None
+    factor: int = 1
+
+
+class DownsampleState(ez.State):
+    cur_settings: DownsampleSettings
+    s_idx: int = 0
+
+
+class Downsample(ez.Unit):
+    SETTINGS: DownsampleSettings
+    STATE: DownsampleState
+
+    INPUT_SETTINGS = ez.InputStream(DownsampleSettings)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    def initialize(self) -> None:
+        self.STATE.cur_settings = self.SETTINGS
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: DownsampleSettings) -> None:
+        self.STATE.cur_settings = msg
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
+        if self.STATE.cur_settings.factor < 1:
+            raise ValueError("Downsample factor must be at least 1 (no downsampling)")
+
+        axis_name = self.STATE.cur_settings.axis
+        if axis_name is None:
+            axis_name = msg.dims[0]
+        axis = msg.get_axis(axis_name)
+        axis_idx = msg.get_axis_idx(axis_name)
+
+        samples = np.arange(msg.data.shape[axis_idx]) + self.STATE.s_idx
+        samples = samples % self.STATE.cur_settings.factor
+        self.STATE.s_idx = samples[-1] + 1
+
+        pub_samples = np.where(samples == 0)[0]
+        if len(pub_samples) != 0:
+            new_axes = {ax_name: msg.get_axis(ax_name) for ax_name in msg.dims}
+            new_offset = axis.offset + (axis.gain * pub_samples[0].item())
+            new_gain = axis.gain * self.STATE.cur_settings.factor
+            new_axes[axis_name] = replace(axis, gain=new_gain, offset=new_offset)
+            down_data = np.take(msg.data, pub_samples, axis_idx)
+            out_msg = replace(msg, data=down_data, dims=msg.dims, axes=new_axes)
+            yield self.OUTPUT_SIGNAL, out_msg

ezmsg-sigproc-1.2.0/ezmsg/sigproc/ewmfilter.py

@@ -0,0 +1,127 @@
+import asyncio
+from dataclasses import replace
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+
+import numpy as np
+
+from .window import Window, WindowSettings
+
+from typing import AsyncGenerator, Optional
+
+
+class EWMSettings(ez.Settings):
+    axis: Optional[str] = None
+    zero_offset: bool = True  # If true, we assume zero DC offset
+
+
+class EWMState(ez.State):
+    buffer_queue: "asyncio.Queue[AxisArray]"
+    signal_queue: "asyncio.Queue[AxisArray]"
+
+
+class EWM(ez.Unit):
+    """
+    Exponentially Weighted Moving Average Standardization
+
+    References https://stackoverflow.com/a/42926270
+    """
+
+    SETTINGS: EWMSettings
+    STATE: EWMState
+
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    INPUT_BUFFER = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    def initialize(self) -> None:
+        self.STATE.signal_queue = asyncio.Queue()
+        self.STATE.buffer_queue = asyncio.Queue()
+
+    @ez.subscriber(INPUT_SIGNAL)
+    async def on_signal(self, message: AxisArray) -> None:
+        self.STATE.signal_queue.put_nowait(message)
+
+    @ez.subscriber(INPUT_BUFFER)
+    async def on_buffer(self, message: AxisArray) -> None:
+        self.STATE.buffer_queue.put_nowait(message)
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def sync_output(self) -> AsyncGenerator:
+        while True:
+            signal = await self.STATE.signal_queue.get()
+            buffer = await self.STATE.buffer_queue.get()  # includes signal
+
+            axis_name = self.SETTINGS.axis
+            if axis_name is None:
+                axis_name = signal.dims[0]
+
+            axis_idx = signal.get_axis_idx(axis_name)
+
+            buffer_len = buffer.shape[axis_idx]
+            block_len = signal.shape[axis_idx]
+            window = buffer_len - block_len
+
+            alpha = 2 / (window + 1.0)
+            alpha_rev = 1 - alpha
+
+            pows = alpha_rev ** (np.arange(buffer_len + 1))
+            scale_arr = 1 / pows[:-1]
+            pw0 = alpha * alpha_rev ** (buffer_len - 1)
+
+            buffer_data = buffer.data
+            buffer_data = np.moveaxis(buffer_data, axis_idx, 0)
+
+            def ewma(data: np.ndarray) -> np.ndarray:
+                mult = scale_arr[:, np.newaxis] * data * pw0
+                out = scale_arr[::-1, np.newaxis] * mult.cumsum(axis=0)
+
+                if not self.SETTINGS.zero_offset:
+                    out = (data[0, :, np.newaxis] * pows[1:]).T + out
+
+                return out
+
+            mean = ewma(buffer_data)
+            std = ewma((buffer_data - mean) ** 2.0)
+
+            standardized = (buffer_data - mean) / np.sqrt(std).clip(1e-4)
+            standardized = standardized[-signal.shape[axis_idx] :, ...]
+            standardized = np.moveaxis(standardized, axis_idx, 0)
+
+            yield self.OUTPUT_SIGNAL, replace(signal, data=standardized)
+
+
+class EWMFilterSettings(ez.Settings):
+    history_dur: float  # previous data to accumulate for standardization
+    axis: Optional[str] = None
+    zero_offset: bool = True  # If true, we assume zero DC offset for input data
+
+
+class EWMFilter(ez.Collection):
+    SETTINGS: EWMFilterSettings
+
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    WINDOW = Window()
+    EWM = EWM()
+
+    def configure(self) -> None:
+        self.EWM.apply_settings(EWMSettings(axis=self.SETTINGS.axis, zero_offset=True))
+
+        self.WINDOW.apply_settings(
+            WindowSettings(
+                axis=self.SETTINGS.axis,
+                window_dur=self.SETTINGS.history_dur,
+                window_shift=None,  # 1:1 mode
+            )
+        )
+
+    def network(self) -> ez.NetworkDefinition:
+        return (
+            (self.INPUT_SIGNAL, self.WINDOW.INPUT_SIGNAL),
+            (self.WINDOW.OUTPUT_SIGNAL, self.EWM.INPUT_BUFFER),
+            (self.INPUT_SIGNAL, self.EWM.INPUT_SIGNAL),
+            (self.EWM.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL),
+        )

{ezmsg-sigproc-1.1.1 → ezmsg-sigproc-1.2.0}/ezmsg/sigproc/filter.py

@@ -4,26 +4,30 @@ import ezmsg.core as ez
 import scipy.signal
 import numpy as np
 import asyncio
-import logging
 
-from .messages import TSMessage as TimeSeriesMessage
+from ezmsg.util.messages.axisarray import AxisArray
 
 from typing import AsyncGenerator, Optional, Tuple
 
 
 @dataclass
 class FilterCoefficients:
-    b: np.ndarray = field(default_factory = lambda: np.array([1.0, 0.0]))
-    a: np.ndarray = field(default_factory = lambda: np.array([1.0, 0.0]))
+    b: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0]))
+    a: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0]))
 
 
-class FilterSettings(ez.Settings):
+class FilterSettingsBase(ez.Settings):
+    axis: Optional[str] = None
+    fs: Optional[float] = None
+
+
+class FilterSettings(FilterSettingsBase):
     # If you'd like to statically design a filter, define it in settings
     filt: Optional[FilterCoefficients] = None
-    fs: Optional[float] = None
 
 
 class FilterState(ez.State):
+    axis: Optional[str] = None
     zi: Optional[np.ndarray] = None
     filt_designed: bool = False
     filt: Optional[FilterCoefficients] = None
@@ -33,21 +37,25 @@ class FilterState(ez.State):
 
 
 class Filter(ez.Unit):
-    SETTINGS: FilterSettings
+    SETTINGS: FilterSettingsBase
     STATE: FilterState
 
     INPUT_FILTER = ez.InputStream(FilterCoefficients)
-    INPUT_SIGNAL = ez.InputStream(TimeSeriesMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(TimeSeriesMessage)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     def design_filter(self) -> Optional[Tuple[np.ndarray, np.ndarray]]:
         raise NotImplementedError("Must implement 'design_filter' in Unit subclass!")
 
     # Set up filter with static initialization if specified
     def initialize(self) -> None:
-        if self.SETTINGS.filt is not None:
-            self.STATE.filt = self.SETTINGS.filt
-            self.STATE.filt_set.set()
+        if self.SETTINGS.axis is not None:
+            self.STATE.axis = self.SETTINGS.axis
+
+        if isinstance(self.SETTINGS, FilterSettings):
+            if self.SETTINGS.filt is not None:
+                self.STATE.filt = self.SETTINGS.filt
+                self.STATE.filt_set.set()
         else:
             self.STATE.filt_set.clear()
 
@@ -64,7 +72,9 @@ class Filter(ez.Unit):
     def update_filter(self):
         try:
             coefs = self.design_filter()
-            self.STATE.filt = FilterCoefficients() if coefs is None else FilterCoefficients( *coefs )
+            self.STATE.filt = (
+                FilterCoefficients() if coefs is None else FilterCoefficients(*coefs)
+            )
             self.STATE.filt_set.set()
             self.STATE.filt_designed = True
         except NotImplementedError as e:
@@ -74,30 +84,36 @@ class Filter(ez.Unit):
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def apply_filter(self, message: TimeSeriesMessage) -> AsyncGenerator:
-        if self.STATE.fs != message.fs and self.STATE.filt_designed is True:
-            self.STATE.fs = message.fs
+    async def apply_filter(self, msg: AxisArray) -> AsyncGenerator:
+        axis_name = msg.dims[0] if self.STATE.axis is None else self.STATE.axis
+        axis_idx = msg.get_axis_idx(axis_name)
+        axis = msg.get_axis(axis_name)
+        fs = 1.0 / axis.gain
+
+        if self.STATE.fs != fs and self.STATE.filt_designed is True:
+            self.STATE.fs = fs
             self.update_filter()
 
         # Ensure filter is defined
+        # TODO: Maybe have me be a passthrough filter until coefficients are received
         if self.STATE.filt is None:
             self.STATE.filt_set.clear()
             ez.logger.info("Awaiting filter coefficients...")
             await self.STATE.filt_set.wait()
             ez.logger.info("Filter coefficients received.")
 
-        arr_in: np.ndarray
+        assert self.STATE.filt is not None
+
+        arr_in = msg.data
 
         # If the array is one dimensional, add a temporary second dimension so that the math works out
         one_dimensional = False
-        if message.data.ndim == 1:
-            arr_in = np.expand_dims(message.data, axis=1)
+        if arr_in.ndim == 1:
+            arr_in = np.expand_dims(arr_in, axis=1)
             one_dimensional = True
-        else:
-            arr_in = message.data
 
         # We will perform filter with time dimension as last axis
-        arr_in = np.moveaxis(arr_in, message.time_dim, -1)
+        arr_in = np.moveaxis(arr_in, axis_idx, -1)
         samp_shape = arr_in[..., 0].shape
 
         # Re-calculate/reset zi if necessary
@@ -115,10 +131,10 @@ class Filter(ez.Unit):
             self.STATE.filt.b, self.STATE.filt.a, arr_in, zi=self.STATE.zi
         )
 
-        arr_out = np.moveaxis(arr_out, -1, message.time_dim)
+        arr_out = np.moveaxis(arr_out, -1, axis_idx)
 
         # Remove temporary first dimension if necessary
         if one_dimensional:
             arr_out = np.squeeze(arr_out, axis=1)
 
-        yield (self.OUTPUT_SIGNAL, replace(message, data=arr_out))
+        yield self.OUTPUT_SIGNAL, replace(msg, data=arr_out),

ezmsg-sigproc-1.2.0/ezmsg/sigproc/messages.py

@@ -0,0 +1,31 @@
+import warnings
+import time
+
+import numpy.typing as npt
+
+from ezmsg.util.messages.axisarray import AxisArray
+
+from typing import Optional
+
+# UPCOMING: TSMessage Deprecation
+# TSMessage is deprecated because it doesn't handle multiple time axes well.
+# AxisArray has an incompatible API but supports a superset of functionality.
+warnings.warn(
+    "TimeSeriesMessage/TSMessage is deprecated.  Please use ezmsg.utils.AxisArray",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+
+def TSMessage(
+    data: npt.NDArray,
+    fs: float = 1.0,
+    time_dim: int = 0,
+    timestamp: Optional[float] = None,
+) -> AxisArray:
+    dims = [f"dim_{i}" for i in range(data.ndim)]
+    dims[time_dim] = "time"
+    offset = time.time() if timestamp is None else timestamp
+    offset_adj = data.shape[time_dim] / fs  # offset corresponds to idx[0] on time_dim
+    axis = AxisArray.Axis.TimeAxis(fs, offset=offset - offset_adj)
+    return AxisArray(data, dims=dims, axes=dict(time=axis))