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

ezmsg/sigproc/sampler.py

@@ -1,287 +1,278 @@
-from dataclasses import dataclass, replace, field
-import time
+import asyncio
+import copy
+import traceback
+import typing
+from collections import deque
 
 import ezmsg.core as ez
 import numpy as np
-
-from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import Optional, Any, Tuple, List, Dict, AsyncGenerator
-
-
-@dataclass(unsafe_hash = True)
-class SampleTriggerMessage:
-    timestamp: float = field(default_factory=time.time)
-    period: Optional[Tuple[float, float]] = None
-    value: Any = None
-
-
-@dataclass
-class SampleMessage:
-    trigger: SampleTriggerMessage
-    sample: AxisArray
+from ezmsg.baseproc import (
+    BaseConsumerUnit,
+    BaseProducerUnit,
+    BaseStatefulProducer,
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+)
+from ezmsg.util.messages.util import replace
+
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+from .util.buffer import UpdateStrategy
+from .util.message import SampleMessage, SampleTriggerMessage
+from .util.profile import profile_subpub
 
 
 class SamplerSettings(ez.Settings):
-    buffer_dur: float
-    axis: Optional[str] = None
-    period: Optional[
-        Tuple[float, float]
-    ] = None  # Optional default period if unspecified in SampleTriggerMessage
-    value: Any = None  # Optional default value if unspecified in SampleTriggerMessage
-
-    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.
-
-
-class SamplerState(ez.State):
-    cur_settings: SamplerSettings
-    triggers: Dict[SampleTriggerMessage, int] = field(default_factory=dict)
-    last_msg: Optional[AxisArray] = None
-    buffer: Optional[np.ndarray] = None
-
+    """
+    Settings for :obj:`Sampler`.
+    See :obj:`sampler` for a description of the fields.
+    """
 
-class Sampler(ez.Unit):
-    SETTINGS: SamplerSettings
-    STATE: SamplerState
+    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.
+    """
 
-    INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
-    INPUT_SETTINGS = ez.InputStream(SamplerSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SAMPLE = ez.OutputStream(SampleMessage)
+    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
+    """
 
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
+    period: tuple[float, float] | None = None
+    """Optional default period (in seconds) if unspecified in SampleTriggerMessage."""
 
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: SamplerSettings) -> None:
-        self.STATE.cur_settings = msg
-
-    @ez.subscriber(INPUT_TRIGGER)
-    async def on_trigger(self, msg: SampleTriggerMessage) -> None:
-        if self.STATE.last_msg is not None:
-            axis_name = self.STATE.cur_settings.axis
-            if axis_name is None:
-                axis_name = self.STATE.last_msg.dims[0]
-            axis = self.STATE.last_msg.get_axis(axis_name)
-            axis_idx = self.STATE.last_msg.get_axis_idx(axis_name)
-
-            fs = 1.0 / axis.gain
-            last_msg_timestamp = axis.offset + (
-                self.STATE.last_msg.shape[axis_idx] / fs
-            )
-
-            period = (
-                msg.period if msg.period is not None else self.STATE.cur_settings.period
-            )
-            value = (
-                msg.value if msg.value is not None else self.STATE.cur_settings.value
-            )
-
-            if period is None:
-                ez.logger.warning(f"Sampling failed: period not specified")
-                return
-
-            # Check that period is valid
-            start_offset = int(period[0] * fs)
-            stop_offset = int(period[1] * fs)
-            if (stop_offset - start_offset) <= 0:
-                ez.logger.warning(f"Sampling failed: invalid period requested")
-                return
-
-            # Check that period is compatible with buffer duration
-            max_buf_len = int(self.STATE.cur_settings.buffer_dur * fs)
-            req_buf_len = int((period[1] - period[0]) * fs)
-            if req_buf_len >= max_buf_len:
-                ez.logger.warning(
-                    f"Sampling failed: {period=} >= {self.STATE.cur_settings.buffer_dur=}"
-                )
-                return
-
-            offset: int = 0
-            if self.STATE.cur_settings.estimate_alignment:
-                # Do what we can with the wall clock to determine sample alignment
-                wall_delta = msg.timestamp - last_msg_timestamp
-                offset = int(wall_delta * fs)
-
-            # Check that current buffer accumulation allows for offset - period start
-            if (
-                self.STATE.buffer is None
-                or -min(offset + start_offset, 0) >= self.STATE.buffer.shape[0]
-            ):
-                ez.logger.warning(
-                    "Sampling failed: insufficient buffer accumulation for requested sample period"
-                )
-                return
-
-            self.STATE.triggers[replace(msg, period=period, value=value)] = offset
+    value: typing.Any = None
+    """Optional default value if unspecified in SampleTriggerMessage"""
 
+    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.
+    """
+
+    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.
+    """
+
+
+@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:
-            ez.logger.warning("Sampling failed: no signal to sample yet")
-
-    @ez.subscriber(INPUT_SIGNAL)
-    @ez.publisher(OUTPUT_SAMPLE)
-    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
-        axis_name = self.STATE.cur_settings.axis
-        if axis_name is None:
-            axis_name = msg.dims[0]
-        axis = msg.get_axis(axis_name)
-
-        fs = 1.0 / axis.gain
-
-        if self.STATE.last_msg is None:
-            self.STATE.last_msg = msg
-
-        # Easier to deal with timeseries on axis 0
-        last_msg = self.STATE.last_msg
-        msg_data = np.moveaxis(msg.data, msg.get_axis_idx(axis_name), 0)
-        last_msg_data = np.moveaxis(last_msg.data, last_msg.get_axis_idx(axis_name), 0)
-        last_msg_axis = last_msg.get_axis(axis_name)
-        last_msg_fs = 1.0 / last_msg_axis.gain
-
-        # Check if signal properties have changed in a breaking way
-        if fs != last_msg_fs or msg_data.shape[1:] != last_msg_data.shape[1:]:
-            # Data stream changed meaningfully -- flush buffer, stop sampling
-            if len(self.STATE.triggers) > 0:
-                ez.logger.warning("Sampling failed: Discarding all triggers")
-            ez.logger.warning("Flushing buffer: signal properties changed")
-            self.STATE.buffer = None
-            self.STATE.triggers = dict()
-
-        # Accumulate buffer ( time dim => dim 0 )
-        self.STATE.buffer = (
-            msg_data
-            if self.STATE.buffer is None
-            else np.concatenate((self.STATE.buffer, msg_data), axis=0)
+            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()
 
-        buffer_offset = np.arange(self.STATE.buffer.shape[0] + msg_data.shape[0])
-        buffer_offset -= self.STATE.buffer.shape[0] + 1
-        buffer_offset = (buffer_offset * axis.gain) + axis.offset
+    def _process(self, message: AxisArray) -> list[SampleMessage]:
+        self._state.buffer.write(message)
 
-        pub_samples: List[SampleMessage] = []
-        remaining_triggers: Dict[SampleTriggerMessage, int] = dict()
-        for trigger, offset in self.STATE.triggers.items():
-            if trigger.period is None:
-                continue
-
-            # trigger_offset points to t = 0 within buffer
-            offset -= msg_data.shape[0]
-            start = offset + int(trigger.period[0] * fs)
-            stop = offset + int(trigger.period[1] * fs)
+        # How much data in the buffer?
+        buff_t_range = (
+            self._state.buffer.axis_first_value,
+            self._state.buffer.axis_final_value,
+        )
 
-            if stop < 0:  # We should be able to dispatch a sample
-                sample_data = self.STATE.buffer[start:stop, ...]
-                sample_data = np.moveaxis(sample_data, msg.get_axis_idx(axis_name), 0)
+        # 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
 
-                sample_offset = buffer_offset[start]
-                sample_axis = replace(axis, offset=sample_offset)
-                sample_axes = {**msg.axes, **{axis_name: sample_axis}}
+            trig_range = trig.timestamp + np.array(trig.period)
 
-                pub_samples.append(
-                    SampleMessage(
-                        trigger=trigger,
-                        sample=replace(msg, data=sample_data, axes=sample_axes),
-                    )
+            # 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: Buffer span {buff_t_range} begins beyond the "
+                    f"requested sample period start: {trig_range[0]}"
                 )
+                del self._state.triggers[trig_ix]
+                continue
 
-            else:
-                remaining_triggers[trigger] = offset
-
-        for sample in pub_samples:
-            yield self.OUTPUT_SAMPLE, sample
+            if trig_range[1] > buff_t_range[1]:
+                # We don't *yet* have enough data to satisfy this trigger.
+                continue
 
-        self.STATE.triggers = remaining_triggers
+            # 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.
 
-        buf_len = int(self.STATE.cur_settings.buffer_dur * fs)
-        self.STATE.buffer = self.STATE.buffer[-buf_len:, ...]
-        self.STATE.last_msg = msg
+            # Prepare output and drop trigger
+            msgs_out.append(SampleMessage(trigger=copy.copy(trig), sample=buff_axarr))
+            del self._state.triggers[trig_ix]
 
+        msgs_out.reverse()  # in-place
+        return msgs_out
 
-## Dev/test apparatus
-import asyncio
+    def push_trigger(self, message: SampleTriggerMessage) -> list[SampleMessage]:
+        # Input is a trigger message that we will use to sample the buffer.
 
-from ezmsg.testing.debuglog import DebugLog
-from ezmsg.sigproc.synth import Oscillator, OscillatorSettings
+        if self._state.buffer is None:
+            # We've yet to see any data; drop the trigger.
+            return []
 
-from typing import AsyncGenerator
+        _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
 
+        if _period is None:
+            ez.logger.warning("Sampling failed: period not specified")
+            return []
 
-class TriggerGeneratorSettings(ez.Settings):
-    period: Tuple[float, float]  # sec
-    prewait: float = 0.5  # sec
-    publish_period: float = 5.0  # sec
+        # 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 TriggerGenerator(ez.Unit):
-    SETTINGS: TriggerGeneratorSettings
+        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
 
-    OUTPUT_TRIGGER = ez.OutputStream(SampleTriggerMessage)
+        new_trig_msg = replace(message, timestamp=trigger_ts, period=_period, value=_value)
+        self._state.triggers.append(new_trig_msg)
+        return []
 
-    @ez.publisher(OUTPUT_TRIGGER)
-    async def generate(self) -> AsyncGenerator:
-        await asyncio.sleep(self.SETTINGS.prewait)
 
-        output = 0
-        while True:
-            yield self.OUTPUT_TRIGGER, SampleTriggerMessage(
-                period=self.SETTINGS.period, value=output
-            )
+class Sampler(BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, SamplerTransformer]):
+    SETTINGS = SamplerSettings
 
-            await asyncio.sleep(self.SETTINGS.publish_period)
-            output += 1
+    INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
+    OUTPUT_SIGNAL = ez.OutputStream(SampleMessage)
 
+    @ez.subscriber(INPUT_TRIGGER)
+    async def on_trigger(self, msg: SampleTriggerMessage) -> None:
+        _ = 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.
+
+    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 SamplerTestSystemSettings(ez.Settings):
-    sampler_settings: SamplerSettings
-    trigger_settings: TriggerGeneratorSettings
 
+class TriggerGeneratorSettings(ez.Settings):
+    period: tuple[float, float]
+    """The period around the trigger event."""
 
-class SamplerTestSystem(ez.Collection):
-    SETTINGS: SamplerTestSystemSettings
+    prewait: float = 0.5
+    """The time before the first trigger (sec)"""
 
-    OSC = Oscillator()
-    SAMPLER = Sampler()
-    TRIGGER = TriggerGenerator()
-    DEBUG = DebugLog()
+    publish_period: float = 5.0
+    """The period between triggers (sec)"""
 
-    def configure(self) -> None:
-        self.SAMPLER.apply_settings(self.SETTINGS.sampler_settings)
-        self.TRIGGER.apply_settings(self.SETTINGS.trigger_settings)
 
-        self.OSC.apply_settings(
-            OscillatorSettings(
-                n_time=2,  # Number of samples to output per block
-                fs=10,  # Sampling rate of signal output in Hz
-                dispatch_rate="realtime",
-                freq=2.0,  # Oscillation frequency in Hz
-                amp=1.0,  # Amplitude
-                phase=0.0,  # Phase offset (in radians)
-                sync=True,  # Adjust `freq` to sync with sampling rate
-            )
-        )
+@processor_state
+class TriggerGeneratorState:
+    output: int = 0
 
-    def network(self) -> ez.NetworkDefinition:
-        return (
-            (self.OSC.OUTPUT_SIGNAL, self.SAMPLER.INPUT_SIGNAL),
-            (self.TRIGGER.OUTPUT_TRIGGER, self.SAMPLER.INPUT_TRIGGER),
-            (self.TRIGGER.OUTPUT_TRIGGER, self.DEBUG.INPUT),
-            (self.SAMPLER.OUTPUT_SAMPLE, self.DEBUG.INPUT),
-        )
 
+class TriggerProducer(BaseStatefulProducer[TriggerGeneratorSettings, SampleTriggerMessage, TriggerGeneratorState]):
+    def _reset_state(self) -> None:
+        self._state.output = 0
 
-if __name__ == "__main__":
-    settings = SamplerTestSystemSettings(
-        sampler_settings=SamplerSettings(buffer_dur=5.0),
-        trigger_settings=TriggerGeneratorSettings(
-            period=(1.0, 2.0), prewait=0.5, publish_period=5.0
-        ),
-    )
+    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
 
-    system = SamplerTestSystem(settings)
 
-    ez.run(SYSTEM = system)
+class TriggerGenerator(
+    BaseProducerUnit[
+        TriggerGeneratorSettings,
+        SampleTriggerMessage,
+        TriggerProducer,
+    ]
+):
+    SETTINGS = TriggerGeneratorSettings