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

ezmsg/sigproc/sampler.py

@@ -1,293 +1,244 @@
-import asyncio  # Dev/test apparatus
-from collections import deque
-from dataclasses import dataclass, field
-import time
+import asyncio
+import copy
+import traceback
 import typing
+from collections import deque
 
-import numpy as np
-import numpy.typing as npt
 import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import (
+    BaseConsumerUnit,
+    BaseProducerUnit,
+    BaseStatefulProducer,
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
 )
 from ezmsg.util.messages.util import replace
-from ezmsg.util.generator import consumer
-
 
-@dataclass(unsafe_hash=True)
-class SampleTriggerMessage:
-    timestamp: float = field(default_factory=time.time)
-    """Time of the trigger, in seconds. The Clock depends on the input but defaults to time.time"""
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+from .util.buffer import UpdateStrategy
+from .util.message import SampleMessage, SampleTriggerMessage
+from .util.profile import profile_subpub
 
-    period: tuple[float, float] | None = None
-    """The period around the timestamp, in seconds"""
 
-    value: typing.Any = None
-    """A value or 'label' associated with the trigger."""
+class SamplerSettings(ez.Settings):
+    """
+    Settings for :obj:`Sampler`.
+    See :obj:`sampler` for a description of the fields.
+    """
 
+    buffer_dur: float
+    """
+     The duration of the buffer in seconds. The buffer must be long enough to store the oldest
+        sample to be included in a window. e.g., a trigger lagged by 0.5 seconds with a period of (-1.0, +1.5) will
+        need a buffer of 0.5 + (1.5 - -1.0) = 3.0 seconds. It is best to at least double your estimate if memory allows.
+    """
 
-@dataclass
-class SampleMessage:
-    trigger: SampleTriggerMessage
-    """The time, window, and value (if any) associated with the trigger."""
+    axis: str | None = None
+    """
+    The axis along which to sample the data.
+        None (default) will choose the first axis in the first input.
+        Note: (for now) the axis must exist in the msg .axes and be of type AxisArray.LinearAxis
+    """
 
-    sample: AxisArray
-    """The data sampled around the trigger."""
+    period: tuple[float, float] | None = None
+    """Optional default period (in seconds) if unspecified in SampleTriggerMessage."""
 
+    value: typing.Any = None
+    """Optional default value if unspecified in SampleTriggerMessage"""
 
-@consumer
-def sampler(
-    buffer_dur: float,
-    axis: str | None = None,
-    period: tuple[float, float] | None = None,
-    value: typing.Any = None,
-    estimate_alignment: bool = True,
-) -> typing.Generator[list[SampleMessage], AxisArray | SampleTriggerMessage, None]:
+    estimate_alignment: bool = True
+    """
+    If true, use message timestamp fields and reported sampling rate to estimate
+     sample-accurate alignment for samples.
+    If false, sampling will be limited to incoming message rate -- "Block timing"
+    NOTE: For faster-than-realtime playback --  Incoming timestamps must reflect
+    "realtime" operation for estimate_alignment to operate correctly.
     """
-    Sample data into a buffer, accept triggers, and return slices of sampled
-    data around the trigger time.
-
-    Args:
-        buffer_dur: The duration of the buffer in seconds. The buffer must be long enough to store the oldest
-            sample to be included in a window. e.g., a trigger lagged by 0.5 seconds with a period of (-1.0, +1.5) will
-            need a buffer of 0.5 + (1.5 - -1.0) = 3.0 seconds. It is best to at least double your estimate if memory allows.
-        axis: The axis along which to sample the data.
-            None (default) will choose the first axis in the first input.
-            Note: (for now) the axis must exist in the msg .axes and be of type AxisArray.LinearAxis
-        period: The period in seconds during which to sample the data.
-            Defaults to None. Only used if not None and the trigger message does not define its own period.
-        value: The value to sample. Defaults to None.
-        estimate_alignment: Whether to estimate the sample alignment. Defaults to True.
-            If True, the trigger timestamp field is used to slice the buffer.
-            If False, the trigger timestamp is ignored and the next signal's .offset is used.
-            NOTE: For faster-than-realtime playback -- Signals and triggers must share the same (fast) clock for
-            estimate_alignment to operate correctly.
 
-    Returns:
-        A generator that expects `.send` either an :obj:`AxisArray` containing streaming data messages,
-        or a :obj:`SampleTriggerMessage` containing a trigger, and yields the list of :obj:`SampleMessage` s.
+    buffer_update_strategy: UpdateStrategy = "immediate"
+    """
+    The buffer update strategy. See :obj:`ezmsg.sigproc.util.buffer.UpdateStrategy`.
+    If you expect to push data much more frequently than triggers, then "on_demand"
+    might be more efficient. For most other scenarios, "immediate" is best.
     """
-    msg_out: list[SampleMessage] = []
 
-    # State variables (most shared between trigger- and data-processing.
-    triggers: deque[SampleTriggerMessage] = deque()
-    buffer: npt.NDArray | None = None
-    n_samples: int = 0
-    offset: float = 0.0
 
-    check_inputs = {
-        "fs": None,  # Also a state variable
-        "key": None,
-        "shape": None,
-    }
+@processor_state
+class SamplerState:
+    buffer: HybridAxisArrayBuffer | None = None
+    triggers: deque[SampleTriggerMessage] | None = None
+
+
+class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, AxisArray, SamplerState]):
+    def __call__(self, message: AxisArray | SampleTriggerMessage) -> list[SampleMessage]:
+        # TODO: Currently we have a single entry point that accepts both
+        #  data and trigger messages and we choose a code path based on
+        #  the message type. However, in the future we will likely replace
+        #  SampleTriggerMessage with an agumented form of AxisArray,
+        #  leveraging its attrs field, which makes this a bit harder.
+        #  We should probably force callers of this object to explicitly
+        #  call `push_trigger` for trigger messages. This will also
+        #  simplify typing somewhat because `push_trigger` should not
+        #  return anything yet we currently have it returning an empty
+        #  list just to be compatible with __call__.
+        if isinstance(message, AxisArray):
+            return super().__call__(message)
+        else:
+            return self.push_trigger(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        # Compute hash based on message properties that require state reset
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((sample_shape, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.buffer = HybridAxisArrayBuffer(
+            duration=self.settings.buffer_dur,
+            axis=self.settings.axis or message.dims[0],
+            update_strategy=self.settings.buffer_update_strategy,
+            overflow_strategy="warn-overwrite",  # True circular buffer
+        )
+        if self._state.triggers is None:
+            self._state.triggers = deque()
+        self._state.triggers.clear()
 
-    while True:
-        msg_in = yield msg_out
-        msg_out = []
+    def _process(self, message: AxisArray) -> list[SampleMessage]:
+        self._state.buffer.write(message)
 
-        if isinstance(msg_in, SampleTriggerMessage):
-            # Input is a trigger message that we will use to sample the buffer.
+        # How much data in the buffer?
+        buff_t_range = (
+            self._state.buffer.axis_first_value,
+            self._state.buffer.axis_final_value,
+        )
 
-            if buffer is None or check_inputs["fs"] is None:
-                # We've yet to see any data; drop the trigger.
+        # Process in reverse order so that we can remove triggers safely as we iterate.
+        msgs_out: list[SampleMessage] = []
+        for trig_ix in range(len(self._state.triggers) - 1, -1, -1):
+            trig = self._state.triggers[trig_ix]
+            if trig.period is None:
+                ez.logger.warning("Sampling failed: trigger period not specified")
+                del self._state.triggers[trig_ix]
                 continue
 
-            _period = msg_in.period if msg_in.period is not None else period
-            _value = msg_in.value if msg_in.value is not None else value
+            trig_range = trig.timestamp + np.array(trig.period)
 
-            if _period is None:
-                ez.logger.warning("Sampling failed: period not specified")
-                continue
-
-            # Check that period is valid
-            if _period[0] >= _period[1]:
+            # If the previous iteration had insufficient data for the trigger timestamp + period,
+            #  and buffer-management removed data required for the trigger, then we will never be able
+            #  to accommodate this trigger. Discard it. An increase in buffer_dur is recommended.
+            if trig_range[0] < buff_t_range[0]:
                 ez.logger.warning(
-                    f"Sampling failed: invalid period requested ({_period})"
+                    f"Sampling failed: Buffer span {buff_t_range} begins beyond the "
+                    f"requested sample period start: {trig_range[0]}"
                 )
+                del self._state.triggers[trig_ix]
                 continue
 
-            # Check that period is compatible with buffer duration.
-            max_buf_len = int(np.round(buffer_dur * check_inputs["fs"]))
-            req_buf_len = int(np.round((_period[1] - _period[0]) * check_inputs["fs"]))
-            if req_buf_len >= max_buf_len:
-                ez.logger.warning(f"Sampling failed: {period=} >= {buffer_dur=}")
+            if trig_range[1] > buff_t_range[1]:
+                # We don't *yet* have enough data to satisfy this trigger.
                 continue
 
-            trigger_ts: float = msg_in.timestamp
-            if not estimate_alignment:
-                # Override the trigger timestamp with the next sample's likely timestamp.
-                trigger_ts = offset + (n_samples + 1) / check_inputs["fs"]
+            # We know we have enough data in the buffer to satisfy this trigger.
+            buff_idx = self._state.buffer.axis_searchsorted(trig_range, side="right")
+            self._state.buffer.seek(buff_idx[0])  # FFWD to starting position.
+            buff_axarr = self._state.buffer.peek(buff_idx[1] - buff_idx[0])
+            self._state.buffer.seek(-buff_idx[0])  # Rewind it back.
+            # Note: buffer will trim itself as needed based on buffer_dur.
 
-            new_trig_msg = replace(
-                msg_in, timestamp=trigger_ts, period=_period, value=_value
-            )
-            triggers.append(new_trig_msg)
-
-        elif isinstance(msg_in, AxisArray):
-            # Get properties from message
-            axis = axis or msg_in.dims[0]
-            axis_idx = msg_in.get_axis_idx(axis)
-            axis_info = msg_in.get_axis(axis)
-            fs = 1.0 / axis_info.gain
-            sample_shape = (
-                msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1 :]
-            )
+            # Prepare output and drop trigger
+            msgs_out.append(SampleMessage(trigger=copy.copy(trig), sample=buff_axarr))
+            del self._state.triggers[trig_ix]
 
-            # TODO: We could accommodate change in dim order.
-            # if axis_idx != check_inputs["axis_idx"]:
-            #     msg_in = replace(
-            #         msg_in,
-            #         data=np.moveaxis(msg_in.data, axis_idx, check_inputs["axis_idx"]),
-            #         dims=TODO...
-            #     )
-            #    axis_idx = check_inputs["axis_idx"]
-
-            # If the properties have changed in a breaking way then reset buffer and triggers.
-            b_reset = fs != check_inputs["fs"]
-            b_reset = b_reset or sample_shape != check_inputs["shape"]
-            # TODO: Skip next line if we do np.moveaxis above
-            b_reset = b_reset or axis_idx != check_inputs["axis_idx"]
-            b_reset = b_reset or msg_in.key != check_inputs["key"]
-            if b_reset:
-                check_inputs["fs"] = fs
-                check_inputs["shape"] = sample_shape
-                check_inputs["axis_idx"] = axis_idx
-                check_inputs["key"] = msg_in.key
-                n_samples = msg_in.data.shape[axis_idx]
-                buffer = None
-                if len(triggers) > 0:
-                    ez.logger.warning("Data stream changed: Discarding all triggers")
-                triggers.clear()
-
-            # Save some info for trigger processing
-            offset = axis_info.offset
-
-            # Update buffer
-            buffer = (
-                msg_in.data
-                if buffer is None
-                else np.concatenate((buffer, msg_in.data), axis=axis_idx)
-            )
+        msgs_out.reverse()  # in-place
+        return msgs_out
 
-            # Calculate timestamps associated with buffer.
-            buffer_offset = np.arange(buffer.shape[axis_idx], dtype=float)
-            buffer_offset -= buffer_offset[-msg_in.data.shape[axis_idx]]
-            buffer_offset *= axis_info.gain
-            buffer_offset += axis_info.offset
-
-            # ... for each trigger, collect the message (if possible) and append to msg_out
-            for trig in list(triggers):
-                if trig.period is None:
-                    # This trigger was malformed; drop it.
-                    triggers.remove(trig)
-
-                # If the previous iteration had insufficient data for the trigger timestamp + period,
-                #  and buffer-management removed data required for the trigger, then we will never be able
-                #  to accommodate this trigger. Discard it. An increase in buffer_dur is recommended.
-                if (trig.timestamp + trig.period[0]) < buffer_offset[0]:
-                    ez.logger.warning(
-                        f"Sampling failed: Buffer span {buffer_offset[0]} is beyond the "
-                        f"requested sample period start: {trig.timestamp + trig.period[0]}"
-                    )
-                    triggers.remove(trig)
-
-                t_start = trig.timestamp + trig.period[0]
-                if t_start >= buffer_offset[0]:
-                    start = np.searchsorted(buffer_offset, t_start)
-                    stop = start + int(np.round(fs * (trig.period[1] - trig.period[0])))
-                    if buffer.shape[axis_idx] > stop:
-                        # Trigger period fully enclosed in buffer.
-                        msg_out.append(
-                            SampleMessage(
-                                trigger=trig,
-                                sample=replace(
-                                    msg_in,
-                                    data=slice_along_axis(
-                                        buffer, slice(start, stop), axis_idx
-                                    ),
-                                    axes={
-                                        **msg_in.axes,
-                                        axis: replace(
-                                            axis_info, offset=buffer_offset[start]
-                                        ),
-                                    },
-                                ),
-                            )
-                        )
-                        triggers.remove(trig)
-
-            buf_len = int(buffer_dur * fs)
-            buffer = slice_along_axis(buffer, np.s_[-buf_len:], axis_idx)
+    def push_trigger(self, message: SampleTriggerMessage) -> list[SampleMessage]:
+        # Input is a trigger message that we will use to sample the buffer.
 
+        if self._state.buffer is None:
+            # We've yet to see any data; drop the trigger.
+            return []
 
-class SamplerSettings(ez.Settings):
-    """
-    Settings for :obj:`Sampler`.
-    See :obj:`sampler` for a description of the fields.
-    """
-
-    buffer_dur: float
-    axis: str | None = None
-    period: tuple[float, float] | None = None
-    """Optional default period if unspecified in SampleTriggerMessage"""
+        _period = message.period if message.period is not None else self.settings.period
+        _value = message.value if message.value is not None else self.settings.value
 
-    value: typing.Any = None
-    """Optional default value if unspecified in SampleTriggerMessage"""
+        if _period is None:
+            ez.logger.warning("Sampling failed: period not specified")
+            return []
 
-    estimate_alignment: bool = True
-    """
-    If true, use message timestamp fields and reported sampling rate to estimate sample-accurate alignment for samples.
-    If false, sampling will be limited to incoming message rate -- "Block timing"
-    NOTE: For faster-than-realtime playback --  Incoming timestamps must reflect
-    "realtime" operation for estimate_alignment to operate correctly.
-    """
+        # Check that period is valid
+        if _period[0] >= _period[1]:
+            ez.logger.warning(f"Sampling failed: invalid period requested ({_period})")
+            return []
 
+        # Check that period is compatible with buffer duration.
+        if (_period[1] - _period[0]) > self.settings.buffer_dur:
+            ez.logger.warning(
+                f"Sampling failed: trigger period {_period=} >= buffer capacity {self.settings.buffer_dur=}"
+            )
+            return []
 
-class SamplerState(ez.State):
-    cur_settings: SamplerSettings
-    gen: typing.Generator[AxisArray | SampleTriggerMessage, list[SampleMessage], None]
+        trigger_ts: float = message.timestamp
+        if not self.settings.estimate_alignment:
+            # Override the trigger timestamp with the next sample's likely timestamp.
+            trigger_ts = self._state.buffer.axis_final_value + self._state.buffer.axis_gain
 
+        new_trig_msg = replace(message, timestamp=trigger_ts, period=_period, value=_value)
+        self._state.triggers.append(new_trig_msg)
+        return []
 
-class Sampler(ez.Unit):
-    """An :obj:`Unit` for :obj:`sampler`."""
 
+class Sampler(BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, SamplerTransformer]):
     SETTINGS = SamplerSettings
-    STATE = SamplerState
 
     INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
-    INPUT_SETTINGS = ez.InputStream(SamplerSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SAMPLE = ez.OutputStream(SampleMessage)
-
-    def construct_generator(self):
-        self.STATE.gen = sampler(
-            buffer_dur=self.STATE.cur_settings.buffer_dur,
-            axis=self.STATE.cur_settings.axis,
-            period=self.STATE.cur_settings.period,
-            value=self.STATE.cur_settings.value,
-            estimate_alignment=self.STATE.cur_settings.estimate_alignment,
-        )
-
-    async def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
-        self.construct_generator()
-
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: SamplerSettings) -> None:
-        self.STATE.cur_settings = msg
-        self.construct_generator()
+    OUTPUT_SIGNAL = ez.OutputStream(SampleMessage)
 
     @ez.subscriber(INPUT_TRIGGER)
     async def on_trigger(self, msg: SampleTriggerMessage) -> None:
-        _ = self.STATE.gen.send(msg)
+        _ = self.processor.push_trigger(msg)
+
+    @ez.subscriber(BaseConsumerUnit.INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        try:
+            for sample in self.processor(message):
+                yield self.OUTPUT_SIGNAL, sample
+        except Exception as e:
+            ez.logger.info(f"{traceback.format_exc()} - {e}")
+
+
+def sampler(
+    buffer_dur: float,
+    axis: str | None = None,
+    period: tuple[float, float] | None = None,
+    value: typing.Any = None,
+    estimate_alignment: bool = True,
+) -> SamplerTransformer:
+    """
+    Sample data into a buffer, accept triggers, and return slices of sampled
+    data around the trigger time.
 
-    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
-    @ez.publisher(OUTPUT_SAMPLE)
-    async def on_signal(self, msg: AxisArray) -> typing.AsyncGenerator:
-        pub_samples = self.STATE.gen.send(msg)
-        for sample in pub_samples:
-            yield self.OUTPUT_SAMPLE, sample
+    Returns:
+        A generator that expects `.send` either an :obj:`AxisArray` containing streaming data messages,
+        or a :obj:`SampleTriggerMessage` containing a trigger, and yields the list of :obj:`SampleMessage` s.
+    """
+    return SamplerTransformer(
+        settings=SamplerSettings(
+            buffer_dur=buffer_dur,
+            axis=axis,
+            period=period,
+            value=value,
+            estimate_alignment=estimate_alignment,
+        )
+    )
 
 
 class TriggerGeneratorSettings(ez.Settings):
@@ -301,23 +252,27 @@ class TriggerGeneratorSettings(ez.Settings):
     """The period between triggers (sec)"""
 
 
-class TriggerGenerator(ez.Unit):
-    """
-    A unit to generate triggers every `publish_period` interval.
-    """
+@processor_state
+class TriggerGeneratorState:
+    output: int = 0
 
-    SETTINGS = TriggerGeneratorSettings
 
-    OUTPUT_TRIGGER = ez.OutputStream(SampleTriggerMessage)
+class TriggerProducer(BaseStatefulProducer[TriggerGeneratorSettings, SampleTriggerMessage, TriggerGeneratorState]):
+    def _reset_state(self) -> None:
+        self._state.output = 0
 
-    @ez.publisher(OUTPUT_TRIGGER)
-    async def generate(self) -> typing.AsyncGenerator:
-        await asyncio.sleep(self.SETTINGS.prewait)
+    async def _produce(self) -> SampleTriggerMessage:
+        await asyncio.sleep(self.settings.publish_period)
+        out_msg = SampleTriggerMessage(period=self.settings.period, value=self._state.output)
+        self._state.output += 1
+        return out_msg
 
-        output = 0
-        while True:
-            out_msg = SampleTriggerMessage(period=self.SETTINGS.period, value=output)
-            yield self.OUTPUT_TRIGGER, out_msg
 
-            await asyncio.sleep(self.SETTINGS.publish_period)
-            output += 1
+class TriggerGenerator(
+    BaseProducerUnit[
+        TriggerGeneratorSettings,
+        SampleTriggerMessage,
+        TriggerProducer,
+    ]
+):
+    SETTINGS = TriggerGeneratorSettings