diff --git a/docs/source/guides/ProcessorsBase.md b/docs/source/guides/ProcessorsBase.md index 4bc94ca..9e48184 100644 --- a/docs/source/guides/ProcessorsBase.md +++ b/docs/source/guides/ProcessorsBase.md @@ -6,12 +6,13 @@ The `ezmsg.baseproc` module contains the base classes for message processors. Th ### Generic TypeVars -| Idx | Class | Description | -|-----|-----------------------|----------------------------------------------------------------------------| -| 1 | `MessageInType` (Mi) | for messages passed to a consumer, processor, or transformer | -| 2 | `MessageOutType` (Mo) | for messages returned by a producer, processor, or transformer | -| 3 | `SettingsType` | bound to ez.Settings | -| 4 | `StateType` (St) | bound to ProcessorState which is simply ez.State with a `hash: int` field. | +| Idx | Class | Description | +|-----|----------------------------|----------------------------------------------------------------------------| +| 1 | `MessageInType` (Mi) | for messages passed to a consumer, processor, or transformer | +| 2 | `MessageOutType` (Mo) | for messages returned by a producer, processor, or transformer | +| 3 | `SettingsType` | bound to ez.Settings | +| 4 | `StateType` (St) | bound to ProcessorState which is simply ez.State with a `hash: int` field. | +| 5 | `ClockDrivenSettingsType` | bound to `ClockDrivenSettings` (provides `fs` and `n_time`) | ### Protocols @@ -47,6 +48,7 @@ Note: `__call__` and `partial_fit` both have asynchronous alternatives: `__acall | 10 | `BaseAsyncTransformer` | 8 | 8 | `__acall__` wraps abstract `_aprocess`; `__call__` runs `__acall__`. | | 11 | `CompositeProcessor` | 1 | 5 | Methods iterate over sequence of processors created in `_initialize_processors`. | | 12 | `CompositeProducer` | 2 | 6 | Similar to `CompositeProcessor`, but first processor must be a producer. | +| 13 | `BaseClockDrivenProducer` | 5 | 8 | Clock-driven data generator. Implements `_produce(n_samples, time_axis)`. | NOTES: 1. Producers do not inherit from `BaseProcessor`, so concrete implementations should subclass `BaseProducer` or `BaseStatefulProducer`. @@ -60,25 +62,27 @@ do not inherit from `BaseStatefulProcessor` and `BaseStatefulProducer`. They acc ### Generic TypeVars for ezmsg Units -| Idx | Class | Description | -|-----|---------------------------|------------------------------------------------------------------------------------------------------------------| -| 5 | `ProducerType` | bound to `BaseProducer` (hence, also `BaseStatefulProducer`, `CompositeProducer`) | -| 6 | `ConsumerType` | bound to `BaseConsumer`, `BaseStatefulConsumer` | -| 7 | `TransformerType` | bound to `BaseTransformer`, `BaseStatefulTransformer`, `CompositeProcessor` (hence, also `BaseAsyncTransformer`) | -| 8 | `AdaptiveTransformerType` | bound to `BaseAdaptiveTransformer` | +| Idx | Class | Description | +|-----|----------------------------|------------------------------------------------------------------------------------------------------------------| +| 5 | `ProducerType` | bound to `BaseProducer` (hence, also `BaseStatefulProducer`, `CompositeProducer`) | +| 6 | `ConsumerType` | bound to `BaseConsumer`, `BaseStatefulConsumer` | +| 7 | `TransformerType` | bound to `BaseTransformer`, `BaseStatefulTransformer`, `CompositeProcessor` (hence, also `BaseAsyncTransformer`) | +| 8 | `AdaptiveTransformerType` | bound to `BaseAdaptiveTransformer` | +| 9 | `ClockDrivenProducerType` | bound to `BaseClockDrivenProducer` | ### Abstract implementations (Base Classes) for ezmsg Units using processors: -| Idx | Class | Parents | Expected TypeVars | -|-----|-------------------------------|---------|---------------------------| -| 1 | `BaseProcessorUnit` | - | - | -| 2 | `BaseProducerUnit` | - | `ProducerType` | -| 3 | `BaseConsumerUnit` | 1 | `ConsumerType` | -| 4 | `BaseTransformerUnit` | 1 | `TransformerType` | -| 5 | `BaseAdaptiveTransformerUnit` | 1 | `AdaptiveTransformerType` | +| Idx | Class | Parents | Expected TypeVars | +|-----|--------------------------------|---------|----------------------------| +| 1 | `BaseProcessorUnit` | - | - | +| 2 | `BaseProducerUnit` | - | `ProducerType` | +| 3 | `BaseConsumerUnit` | 1 | `ConsumerType` | +| 4 | `BaseTransformerUnit` | 1 | `TransformerType` | +| 5 | `BaseAdaptiveTransformerUnit` | 1 | `AdaptiveTransformerType` | +| 6 | `BaseClockDrivenProducerUnit` | 1 | `ClockDrivenProducerType` | -Note, it is strongly recommended to use `BaseConsumerUnit`, `BaseTransformerUnit`, or `BaseAdaptiveTransformerUnit` for implementing concrete subclasses rather than `BaseProcessorUnit`. +Note, it is strongly recommended to use `BaseConsumerUnit`, `BaseTransformerUnit`, `BaseAdaptiveTransformerUnit`, or `BaseClockDrivenProducerUnit` for implementing concrete subclasses rather than `BaseProcessorUnit`. ## Implementing a custom standalone processor @@ -125,6 +129,7 @@ flowchart TD * For stateful processors that need to respond to a change in the incoming data, implement `_hash_message`. * For adaptive transformers, implement `partial_fit`. * For chains of processors (`CompositeProcessor`/ `CompositeProducer`), need to implement `_initialize_processors`. + * For clock-driven producers (`BaseClockDrivenProducer`), implement `_reset_state(time_axis)` and `_produce(n_samples, time_axis)`. See the [clock-driven how-to guide](how-tos/processors/clockdriven.rst). * See processors in `ezmsg.sigproc` for signal processing examples, or `ezmsg.learn` for machine learning examples. 5. Override non-abstract methods if you need special behaviour. diff --git a/docs/source/guides/how-tos/processors/clockdriven.rst b/docs/source/guides/how-tos/processors/clockdriven.rst new file mode 100644 index 0000000..65bbbbe --- /dev/null +++ b/docs/source/guides/how-tos/processors/clockdriven.rst @@ -0,0 +1,224 @@ +How to implement a clock-driven producer? +######################################### + +Clock-driven producers generate data synchronized to clock ticks. They are useful +for signal generators, simulators, and other components that need to produce +timed data streams. + +The ``BaseClockDrivenProducer`` base class simplifies this pattern by handling +all the timing and sample counting logic internally. You only need to implement +the data generation. + +When to use BaseClockDrivenProducer +=================================== + +Use ``BaseClockDrivenProducer`` when you need to: + +- Generate synthetic signals (sine waves, noise, test patterns) +- Simulate sensor data at a specific sample rate +- Produce timed data streams driven by a ``Clock`` + +This base class eliminates the need for the ``Clock → Counter → Generator`` +pattern by combining the counter functionality into the generator. + +Basic Structure +=============== + +A clock-driven producer consists of three parts: + +1. **Settings** - Extends ``ClockDrivenSettings`` (which provides ``fs`` and ``n_time``) +2. **State** - Extends ``ClockDrivenState`` (which provides ``counter`` and ``fractional_samples``) +3. **Producer** - Extends ``BaseClockDrivenProducer`` and implements ``_reset_state`` and ``_produce`` + +Example: Sine Wave Generator +============================ + +Here's a complete example of a sine wave generator: + +.. code-block:: python + + import numpy as np + from ezmsg.util.messages.axisarray import AxisArray, LinearAxis + + from ezmsg.baseproc import ( + BaseClockDrivenProducer, + BaseClockDrivenProducerUnit, + ClockDrivenSettings, + ClockDrivenState, + processor_state, + ) + + + class SinGeneratorSettings(ClockDrivenSettings): + """ + Settings for the sine wave generator. + + Inherits from ClockDrivenSettings which provides: + - fs: Output sampling rate in Hz + - n_time: Samples per block (optional, derived from clock if None) + """ + + freq: float = 1.0 + """Frequency of the sine wave in Hz.""" + + amp: float = 1.0 + """Amplitude of the sine wave.""" + + phase: float = 0.0 + """Initial phase in radians.""" + + + @processor_state + class SinGeneratorState(ClockDrivenState): + """ + State for the sine wave generator. + + Inherits from ClockDrivenState which provides: + - counter: Current sample counter (total samples produced) + - fractional_samples: For accumulating sub-sample timing + """ + + ang_freq: float = 0.0 + """Pre-computed angular frequency (2 * pi * freq).""" + + + class SinGenerator( + BaseClockDrivenProducer[SinGeneratorSettings, SinGeneratorState] + ): + """ + Generates sine wave data synchronized to clock ticks. + """ + + def _reset_state(self, time_axis: LinearAxis) -> None: + """ + Initialize state. Called once before first production. + + Use this to pre-compute values that don't change between chunks. + """ + self._state.ang_freq = 2 * np.pi * self.settings.freq + + def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: + """ + Generate sine wave data for this chunk. + + Args: + n_samples: Number of samples to generate + time_axis: LinearAxis with correct offset and gain (1/fs) + + Returns: + AxisArray containing the sine wave data + """ + # Calculate time values using the internal counter + t = (np.arange(n_samples) + self._state.counter) * time_axis.gain + + # Generate sine wave + data = self.settings.amp * np.sin( + self._state.ang_freq * t + self.settings.phase + ) + + return AxisArray( + data=data, + dims=["time"], + axes={"time": time_axis}, + ) + + + class SinGeneratorUnit( + BaseClockDrivenProducerUnit[SinGeneratorSettings, SinGenerator] + ): + """ + ezmsg Unit wrapper for SinGenerator. + + Receives clock ticks on INPUT_CLOCK and outputs AxisArray on OUTPUT_SIGNAL. + """ + + SETTINGS = SinGeneratorSettings + + +Key Points +========== + +**Settings inheritance**: Your settings class should extend ``ClockDrivenSettings``, +which provides: + +- ``fs``: The output sampling rate in Hz +- ``n_time``: Optional fixed chunk size. If ``None``, chunk size is derived from + the clock's gain (``fs * clock.gain``) + +**State inheritance**: Your state class should extend ``ClockDrivenState``, +which provides: + +- ``counter``: Tracks total samples produced (use this for continuous signals) +- ``fractional_samples``: Accumulates sub-sample timing for accurate chunk sizes + +**The _produce method**: This is where you generate data. You receive: + +- ``n_samples``: How many samples to generate this chunk +- ``time_axis``: A ``LinearAxis`` with the correct ``offset`` and ``gain`` (1/fs) + +The base class automatically: + +- Computes ``n_samples`` from clock timing or settings +- Manages the sample counter (incremented after ``_produce`` returns) +- Handles fractional sample accumulation for non-integer chunk sizes +- Supports both fixed ``n_time`` and variable chunk modes + +Using Standalone (Outside ezmsg) +================================ + +Clock-driven producers can be used standalone for testing or offline processing: + +.. code-block:: python + + from ezmsg.util.messages.axisarray import AxisArray + + # Create the producer + producer = SinGenerator(SinGeneratorSettings( + fs=1000.0, # 1000 Hz sample rate + n_time=100, # 100 samples per chunk + freq=10.0, # 10 Hz sine wave + amp=1.0, + )) + + # Simulate clock ticks (LinearAxis with gain=1/dispatch_rate, offset=timestamp) + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=0.0) # 10 Hz dispatch + + # Generate data + result = producer(clock_tick) + print(f"Shape: {result.data.shape}") # (100,) + print(f"Sample rate: {1/result.axes['time'].gain} Hz") # 1000.0 Hz + + +Using with ezmsg +================ + +In an ezmsg pipeline, connect a ``Clock`` to your generator's ``INPUT_CLOCK``: + +.. code-block:: python + + import ezmsg.core as ez + from ezmsg.baseproc import Clock, ClockSettings + + + class SinPipeline(ez.Collection): + SETTINGS = SinGeneratorSettings + + CLOCK = Clock() + GENERATOR = SinGeneratorUnit() + + def configure(self) -> None: + self.CLOCK.apply_settings(ClockSettings(dispatch_rate=10.0)) + self.GENERATOR.apply_settings(self.SETTINGS) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.CLOCK.OUTPUT_SIGNAL, self.GENERATOR.INPUT_CLOCK), + ) + + +See Also +======== + +- :doc:`API Reference for clockdriven module <../../../api/generated/ezmsg.baseproc.clockdriven>` +- :doc:`stateful` - For general stateful processor patterns +- :doc:`unit` - For converting processors to ezmsg Units diff --git a/docs/source/guides/how-tos/processors/content-processors.rst b/docs/source/guides/how-tos/processors/content-processors.rst index 88f2385..d545b24 100644 --- a/docs/source/guides/how-tos/processors/content-processors.rst +++ b/docs/source/guides/how-tos/processors/content-processors.rst @@ -10,4 +10,5 @@ Processor HOW TOs adaptive composite unit + clockdriven checkpoint diff --git a/src/ezmsg/baseproc/__init__.py b/src/ezmsg/baseproc/__init__.py index bea6d73..949bf00 100644 --- a/src/ezmsg/baseproc/__init__.py +++ b/src/ezmsg/baseproc/__init__.py @@ -15,6 +15,14 @@ ClockState, ) +# Clock-driven producers +from .clockdriven import ( + BaseClockDrivenProducer, + ClockDrivenSettings, + ClockDrivenSettingsType, + ClockDrivenState, +) + # Composite processor classes from .composite import ( CompositeProcessor, @@ -74,15 +82,18 @@ from .units import ( AdaptiveTransformerType, BaseAdaptiveTransformerUnit, + BaseClockDrivenProducerUnit, BaseConsumerUnit, BaseProcessorUnit, BaseProducerUnit, BaseTransformerUnit, + ClockDrivenProducerType, ConsumerType, GenAxisArray, ProducerType, TransformerType, get_base_adaptive_transformer_type, + get_base_clockdriven_producer_type, get_base_consumer_type, get_base_producer_type, get_base_transformer_type, @@ -116,6 +127,7 @@ "ConsumerType", "TransformerType", "AdaptiveTransformerType", + "ClockDrivenProducerType", # Decorators "processor_state", # Base processor classes @@ -131,6 +143,11 @@ "BaseStatefulTransformer", "BaseAdaptiveTransformer", "BaseAsyncTransformer", + # Clock-driven producers + "BaseClockDrivenProducer", + "ClockDrivenSettings", + "ClockDrivenSettingsType", + "ClockDrivenState", # Composite classes "CompositeStateful", "CompositeProcessor", @@ -141,12 +158,14 @@ "BaseConsumerUnit", "BaseTransformerUnit", "BaseAdaptiveTransformerUnit", + "BaseClockDrivenProducerUnit", "GenAxisArray", # Type resolution helpers "get_base_producer_type", "get_base_consumer_type", "get_base_transformer_type", "get_base_adaptive_transformer_type", + "get_base_clockdriven_producer_type", "_get_base_processor_settings_type", "_get_base_processor_message_in_type", "_get_base_processor_message_out_type", diff --git a/src/ezmsg/baseproc/clockdriven.py b/src/ezmsg/baseproc/clockdriven.py new file mode 100644 index 0000000..78c7ef7 --- /dev/null +++ b/src/ezmsg/baseproc/clockdriven.py @@ -0,0 +1,179 @@ +"""Clock-driven producer base classes for generating data synchronized to clock ticks.""" + +import typing +from abc import abstractmethod + +import ezmsg.core as ez +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis + +from .protocols import StateType, processor_state +from .stateful import BaseStatefulProcessor + + +class ClockDrivenSettings(ez.Settings): + """ + Base settings for clock-driven producers. + + Subclass this to add your own settings while inheriting fs and n_time. + + Example:: + + class SinGeneratorSettings(ClockDrivenSettings): + freq: float = 1.0 + amp: float = 1.0 + """ + + fs: float + """Output sampling rate in Hz.""" + + n_time: int | None = None + """ + Samples per block. + - If specified: fixed chunk size (clock gain is ignored for determining chunk size) + - If None: derived from clock gain (fs * clock.gain), with fractional sample tracking + """ + + +# Type variable for settings that extend ClockDrivenSettings +ClockDrivenSettingsType = typing.TypeVar("ClockDrivenSettingsType", bound=ClockDrivenSettings) + + +@processor_state +class ClockDrivenState: + """ + Internal state for clock-driven producers. + + Tracks sample counting and fractional sample accumulation. + Subclasses should extend this if they need additional state. + """ + + counter: int = 0 + """Current sample counter (total samples produced).""" + + fractional_samples: float = 0.0 + """Accumulated fractional samples for variable chunk mode.""" + + +class BaseClockDrivenProducer( + BaseStatefulProcessor[ClockDrivenSettingsType, AxisArray.LinearAxis, AxisArray, StateType], + typing.Generic[ClockDrivenSettingsType, StateType], +): + """ + Base class for clock-driven data producers. + + Accepts clock ticks (LinearAxis) as input and produces AxisArray output. + Handles all the timing/counter logic internally, so subclasses only need + to implement the data generation logic. + + This eliminates the need for the Clock → Counter → Generator pattern + by combining the Counter functionality into the generator base class. + + Subclasses must implement: + - ``_reset_state(time_axis)``: Initialize any state needed for production + - ``_produce(n_samples, time_axis)``: Generate the actual output data + + Example:: + + @processor_state + class SinState(ClockDrivenState): + ang_freq: float = 0.0 + + class SinProducer(BaseClockDrivenProducer[SinSettings, SinState]): + def _reset_state(self, time_axis: AxisArray.TimeAxis) -> None: + self._state.ang_freq = 2 * np.pi * self.settings.fs + + def _produce(self, n_samples: int, time_axis: AxisArray.TimeAxis) -> AxisArray: + t = (np.arange(n_samples) + self._state.counter) * time_axis.gain + data = np.sin(self._state.ang_freq * t) + return AxisArray(data=data, dims=["time"], axes={"time": time_axis}) + """ + + def _hash_message(self, message: AxisArray.LinearAxis) -> int: + # Return constant hash - state should not reset based on clock rate changes. + # The producer maintains continuity regardless of clock rate changes. + return 0 + + def _compute_samples_and_offset(self, clock_tick: AxisArray.LinearAxis) -> tuple[int, float] | None: + """ + Compute number of samples and time offset from a clock tick. + + Returns: + Tuple of (n_samples, offset) or None if no samples to produce yet. + + Raises: + ValueError: If clock gain is 0 (AFAP mode) and n_time is not specified. + """ + if self.settings.n_time is not None: + # Fixed chunk size mode + n_samples = self.settings.n_time + if clock_tick.gain == 0.0: + # AFAP mode - synthetic offset based on counter + offset = self._state.counter / self.settings.fs + else: + # Use clock's timestamp + offset = clock_tick.offset + else: + # Variable chunk size mode - derive from clock gain + if clock_tick.gain == 0.0: + raise ValueError("Cannot use clock with gain=0 (AFAP) without specifying n_time") + + # Calculate samples including fractional accumulation + samples_float = self.settings.fs * clock_tick.gain + self._state.fractional_samples + n_samples = int(samples_float + 1e-9) + self._state.fractional_samples = samples_float - n_samples + + if n_samples == 0: + return None + + offset = clock_tick.offset + + return n_samples, offset + + @abstractmethod + def _reset_state(self, time_axis: LinearAxis) -> None: + """ + Reset/initialize state for production. + + Called once before the first call to _produce, or when state needs resetting. + Use this to pre-compute values, create templates, etc. + + Args: + time_axis: TimeAxis with the output sampling rate (fs) and initial offset. + """ + ... + + @abstractmethod + def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: + """ + Generate output data for this chunk. + + Args: + n_samples: Number of samples to generate. + time_axis: TimeAxis with correct offset and gain (1/fs) for this chunk. + + Returns: + AxisArray containing the generated data. The time axis should use + the provided time_axis or one derived from it. + """ + ... + + def _process(self, clock_tick: LinearAxis) -> AxisArray | None: + """ + Process a clock tick and produce output. + + Handles all the counter/timing logic internally, then calls _produce. + """ + result = self._compute_samples_and_offset(clock_tick) + if result is None: + return None + + n_samples, offset = result + time_axis = AxisArray.TimeAxis(fs=self.settings.fs, offset=offset) + + # Call subclass production method + output = self._produce(n_samples, time_axis) + + # Update counter + self._state.counter += n_samples + + return output diff --git a/src/ezmsg/baseproc/counter.py b/src/ezmsg/baseproc/counter.py index 40c2d8e..aa2cdd8 100644 --- a/src/ezmsg/baseproc/counter.py +++ b/src/ezmsg/baseproc/counter.py @@ -1,47 +1,32 @@ """Counter generator for sample counting and timing.""" -import ezmsg.core as ez import numpy as np -from ezmsg.util.messages.axisarray import AxisArray, replace +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis, replace +from .clockdriven import ( + BaseClockDrivenProducer, + ClockDrivenSettings, + ClockDrivenState, +) from .protocols import processor_state -from .stateful import BaseStatefulTransformer -from .units import BaseTransformerUnit +from .units import BaseClockDrivenProducerUnit -class CounterSettings(ez.Settings): +class CounterSettings(ClockDrivenSettings): """Settings for :obj:`Counter` and :obj:`CounterTransformer`.""" - fs: float - """Sampling rate in Hz.""" - - n_time: int | None = None - """ - Samples per block. - - If specified: fixed chunk size (clock gain is ignored) - - If None: derived from clock gain (fs * clock.gain), with fractional sample tracking - """ - mod: int | None = None """If set, counter values rollover at this modulus.""" @processor_state -class CounterTransformerState: +class CounterTransformerState(ClockDrivenState): """State for :obj:`CounterTransformer`.""" - counter: int = 0 - """Current counter value (next sample index).""" - - fractional_samples: float = 0.0 - """Accumulated fractional samples for variable chunk mode.""" - template: AxisArray | None = None -class CounterTransformer( - BaseStatefulTransformer[CounterSettings, AxisArray.LinearAxis, AxisArray, CounterTransformerState] -): +class CounterTransformer(BaseClockDrivenProducer[CounterSettings, CounterTransformerState]): """ Transforms clock ticks (LinearAxis) into AxisArray counter values. @@ -49,80 +34,34 @@ class CounterTransformer( fixed (n_time setting) or derived from the clock's gain (fs * gain). """ - def _reset_state(self, message: AxisArray.LinearAxis) -> None: - """Reset state - counter transformer state is simple, just reset values.""" - self._state.counter = 0 - self._state.fractional_samples = 0.0 + def _reset_state(self, time_axis: LinearAxis) -> None: + """Reset state - initialize template for counter output.""" self._state.template = AxisArray( data=np.array([], dtype=int), dims=["time"], - axes={ - "time": AxisArray.TimeAxis(fs=self.settings.fs, offset=message.offset), - }, + axes={"time": time_axis}, key="counter", ) - def _hash_message(self, message: AxisArray.LinearAxis) -> int: - # Return constant hash - counter state should never reset based on message content. - # The counter maintains continuity regardless of clock rate changes. - return 0 - - def _process(self, clock_tick: AxisArray.LinearAxis) -> AxisArray | None: - """Transform a clock tick into counter AxisArray.""" - # Determine number of samples for this block - if self.settings.n_time is not None: - # Fixed chunk size mode - n_samples = self.settings.n_time - # Use wall clock or synthetic offset based on clock gain - if clock_tick.gain == 0.0: - # AFAP mode - synthetic offset - offset = self.state.counter / self.settings.fs - else: - # Use clock's timestamp - offset = clock_tick.offset - else: - # Variable chunk size mode - derive from clock gain - if clock_tick.gain == 0.0: - # AFAP with no fixed n_time - this is an error - raise ValueError("Cannot use clock with gain=0 (AFAP) without specifying n_time") - - # Calculate samples including fractional accumulation - # Add small epsilon to avoid floating point truncation errors (e.g., 0.9999999 -> 0) - samples_float = self.settings.fs * clock_tick.gain + self.state.fractional_samples - n_samples = int(samples_float + 1e-9) - self.state.fractional_samples = samples_float - n_samples - - if n_samples == 0: - # Not enough samples accumulated yet - # TODO: Return empty array. What should offset be? - return None - - # Use clock's timestamp for offset - offset = clock_tick.offset - - # Generate counter data - block_samp = np.arange(self.state.counter, self.state.counter + n_samples) + def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: + """Generate counter values for this chunk.""" + # Generate counter data (using pre-increment counter value) + block_samp = np.arange(self._state.counter, self._state.counter + n_samples) if self.settings.mod is not None: block_samp = block_samp % self.settings.mod - # Create output AxisArray - result = replace( + return replace( self._state.template, data=block_samp, - axes={"time": replace(self._state.template.axes["time"], offset=offset)}, + axes={"time": time_axis}, ) - # Update state - self.state.counter += n_samples - - return result - -class Counter(BaseTransformerUnit[CounterSettings, AxisArray.LinearAxis, AxisArray, CounterTransformer]): +class Counter(BaseClockDrivenProducerUnit[CounterSettings, CounterTransformer]): """ Transforms clock ticks into monotonically increasing counter values as AxisArray. - Receives timing from INPUT_SIGNAL (LinearAxis from Clock) and outputs AxisArray. + Receives timing from INPUT_CLOCK (LinearAxis from Clock) and outputs AxisArray. """ SETTINGS = CounterSettings diff --git a/src/ezmsg/baseproc/units.py b/src/ezmsg/baseproc/units.py index 77b075a..560fef9 100644 --- a/src/ezmsg/baseproc/units.py +++ b/src/ezmsg/baseproc/units.py @@ -7,8 +7,9 @@ import ezmsg.core as ez from ezmsg.util.generator import GenState -from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis +from .clockdriven import BaseClockDrivenProducer from .composite import CompositeProcessor from .processor import BaseConsumer, BaseProducer, BaseTransformer from .protocols import MessageInType, MessageOutType, SettingsType @@ -25,6 +26,7 @@ bound=BaseTransformer | BaseStatefulTransformer | CompositeProcessor, ) AdaptiveTransformerType = typing.TypeVar("AdaptiveTransformerType", bound=BaseAdaptiveTransformer) +ClockDrivenProducerType = typing.TypeVar("ClockDrivenProducerType", bound=BaseClockDrivenProducer) def get_base_producer_type(cls: type) -> type: @@ -43,6 +45,10 @@ def get_base_adaptive_transformer_type(cls: type) -> type: return resolve_typevar(cls, AdaptiveTransformerType) +def get_base_clockdriven_producer_type(cls: type) -> type: + return resolve_typevar(cls, ClockDrivenProducerType) + + # --- Base classes for ezmsg Unit with specific processing capabilities --- class BaseProducerUnit(ez.Unit, ABC, typing.Generic[SettingsType, MessageOutType, ProducerType]): """ @@ -240,6 +246,47 @@ async def on_sample(self, msg: SampleMessage) -> None: await self.processor.apartial_fit(msg) +class BaseClockDrivenProducerUnit( + BaseProcessorUnit[SettingsType], + ABC, + typing.Generic[SettingsType, ClockDrivenProducerType], +): + """ + Base class for clock-driven producer units. + + These units receive clock ticks (LinearAxis) and produce AxisArray output. + This simplifies the Clock → Counter → Generator pattern by combining + the counter functionality into the generator. + + Implement a new Unit as follows:: + + class SinGeneratorUnit(BaseClockDrivenProducerUnit[ + SinGeneratorSettings, # SettingsType (must extend ClockDrivenSettings) + SinProducer, # ClockDrivenProducerType + ]): + SETTINGS = SinGeneratorSettings + + Where SinGeneratorSettings extends ClockDrivenSettings and SinProducer + extends BaseClockDrivenProducer. + """ + + INPUT_CLOCK = ez.InputStream(LinearAxis) + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + def create_processor(self) -> None: + """Create the clock-driven producer instance from settings.""" + producer_type = get_base_clockdriven_producer_type(self.__class__) + self.processor = producer_type(settings=self.SETTINGS) + + @ez.subscriber(INPUT_CLOCK, zero_copy=True) + @ez.publisher(OUTPUT_SIGNAL) + @profile_subpub(trace_oldest=False) + async def on_clock(self, clock_tick: LinearAxis) -> typing.AsyncGenerator: + result = await self.processor.__acall__(clock_tick) + if result is not None: + yield self.OUTPUT_SIGNAL, result + + # Legacy class class GenAxisArray(ez.Unit): STATE = GenState diff --git a/src/ezmsg/baseproc/util/typeresolution.py b/src/ezmsg/baseproc/util/typeresolution.py index 7ebb47f..9787edc 100644 --- a/src/ezmsg/baseproc/util/typeresolution.py +++ b/src/ezmsg/baseproc/util/typeresolution.py @@ -11,6 +11,10 @@ def resolve_typevar(cls: type, target_typevar: typing.TypeVar) -> type: and checks the original bases of each class in the MRO for the TypeVar. If the TypeVar is found, it returns the concrete type bound to it. If the TypeVar is not found, it raises a TypeError. + + If the resolved type is itself a TypeVar, this function recursively + resolves it until a concrete type is found. + Args: cls (type): The class to inspect. target_typevar (typing.TypeVar): The TypeVar to resolve. @@ -30,7 +34,11 @@ def resolve_typevar(cls: type, target_typevar: typing.TypeVar) -> type: index = params.index(target_typevar) args = typing.get_args(orig_base) try: - return args[index] + resolved = args[index] + # If the resolved type is itself a TypeVar, resolve it recursively + if isinstance(resolved, typing.TypeVar): + return resolve_typevar(cls, resolved) + return resolved except IndexError: pass raise TypeError(f"Could not resolve {target_typevar} in {cls}") diff --git a/tests/test_clockdriven.py b/tests/test_clockdriven.py new file mode 100644 index 0000000..a9da2c6 --- /dev/null +++ b/tests/test_clockdriven.py @@ -0,0 +1,373 @@ +"""Unit tests for ezmsg.baseproc.clockdriven module.""" + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray + +from ezmsg.baseproc import ( + BaseClockDrivenProducer, + ClockDrivenSettings, + ClockDrivenState, + processor_state, +) + +# --- Test implementations --- + + +class CounterProducerSettings(ClockDrivenSettings): + """Simple counter producer settings for testing.""" + + mod: int | None = None + """If set, counter values rollover at this modulus.""" + + +class CounterProducer(BaseClockDrivenProducer[CounterProducerSettings, ClockDrivenState]): + """ + Simple counter producer for testing. + + Outputs AxisArray with monotonically increasing counter values. + This is similar to CounterTransformer but uses BaseClockDrivenProducer. + """ + + def _reset_state(self, time_axis: AxisArray.LinearAxis) -> None: + """Reset state - nothing special needed for counter.""" + pass + + def _produce(self, n_samples: int, time_axis: AxisArray.LinearAxis) -> AxisArray: + """Generate counter values.""" + # Generate counter data (using pre-increment counter value) + block_samp = np.arange(self._state.counter, self._state.counter + n_samples) + if self.settings.mod is not None: + block_samp = block_samp % self.settings.mod + + return AxisArray( + data=block_samp, + dims=["time"], + axes={"time": time_axis}, + key="counter", + ) + + +@processor_state +class SinProducerState(ClockDrivenState): + """State for sine wave producer.""" + + ang_freq: float = 0.0 + amp: float = 1.0 + phase: float = 0.0 + + +class SinProducerSettings(ClockDrivenSettings): + """Sine wave producer settings.""" + + freq: float = 1.0 + """Frequency in Hz.""" + + amp: float = 1.0 + """Amplitude.""" + + phase: float = 0.0 + """Initial phase in radians.""" + + +class SinProducer(BaseClockDrivenProducer[SinProducerSettings, SinProducerState]): + """ + Sine wave producer for testing. + + Demonstrates a more complex clock-driven producer with custom state. + """ + + def _reset_state(self, time_axis: AxisArray.LinearAxis) -> None: + """Pre-compute angular frequency.""" + self._state.ang_freq = 2 * np.pi * self.settings.freq + self._state.amp = self.settings.amp + self._state.phase = self.settings.phase + + def _produce(self, n_samples: int, time_axis: AxisArray.LinearAxis) -> AxisArray: + """Generate sine wave data.""" + # Calculate time values for this chunk + t = (np.arange(n_samples) + self._state.counter) * time_axis.gain + data = self._state.amp * np.sin(self._state.ang_freq * t + self._state.phase) + + return AxisArray( + data=data, + dims=["time"], + axes={"time": time_axis}, + key="sine", + ) + + +# --- Tests --- + + +class TestBaseClockDrivenProducer: + """Tests for BaseClockDrivenProducer.""" + + def test_fixed_n_time_mode(self): + """Test producer with fixed n_time.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=100, mod=None)) + + # Create clock tick with gain = 0.1 (10 Hz dispatch rate) + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=1.0) + + result = producer(clock_tick) + + assert isinstance(result, AxisArray) + assert result.data.shape == (100,) + assert result.dims == ["time"] + # TimeAxis has gain = 1/fs + assert result.axes["time"].gain == 1 / 1000.0 + assert result.axes["time"].offset == 1.0 # Uses clock's offset + np.testing.assert_array_equal(result.data, np.arange(100)) + + def test_fixed_n_time_with_afap_clock(self): + """Test producer with fixed n_time and AFAP clock (gain=0).""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=50, mod=None)) + + # AFAP clock has gain=0 + clock_tick = AxisArray.LinearAxis(gain=0.0, offset=123.456) + + result = producer(clock_tick) + + assert isinstance(result, AxisArray) + assert result.data.shape == (50,) + # With AFAP clock, offset is synthetic (counter / fs) + assert result.axes["time"].offset == 0.0 # First block starts at 0 + + # Second call + result2 = producer(clock_tick) + assert result2.axes["time"].offset == 50 / 1000.0 # 0.05 seconds + + def test_variable_n_time_mode(self): + """Test producer with n_time derived from clock gain.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=None, mod=None)) + + # Clock at 10 Hz with fs=1000 -> 100 samples per tick + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=5.0) + + result = producer(clock_tick) + + assert isinstance(result, AxisArray) + assert result.data.shape == (100,) # 1000 * 0.1 = 100 + assert result.axes["time"].offset == 5.0 + + def test_variable_n_time_fractional_accumulation(self): + """Test fractional sample accumulation in variable mode.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=None, mod=None)) + + # Clock at 3 Hz with fs=1000 -> 333.33... samples per tick + # Over 3 ticks we should get exactly 1000 samples (333 + 333 + 334) + clock_tick = AxisArray.LinearAxis(gain=1.0 / 3.0, offset=0.0) + + # First tick: 333.33 -> 333 samples, 0.33 fractional + result1 = producer(clock_tick) + assert result1.data.shape == (333,) + + # Second tick: 333.33 + 0.33 = 666.66 -> 333 samples, 0.66 fractional + result2 = producer(clock_tick) + assert result2.data.shape == (333,) + + # Third tick: 333.33 + 0.66 = 999.99... ≈ 1000 -> 334 samples + result3 = producer(clock_tick) + assert result3.data.shape == (334,) + + # Verify total is exactly 1000 (3 ticks * 333.33... = 1000) + total_samples = sum(r.data.shape[0] for r in [result1, result2, result3]) + assert total_samples == 1000 + + # Fourth tick starts fresh cycle: 333 samples + result4 = producer(clock_tick) + assert result4.data.shape == (333,) + + def test_variable_n_time_returns_none_when_no_samples(self): + """Test that producer returns None when not enough samples accumulated.""" + producer = CounterProducer( + CounterProducerSettings(fs=10.0, n_time=None, mod=None) # Low fs + ) + + # Clock at 100 Hz with fs=10 -> 0.1 samples per tick + clock_tick = AxisArray.LinearAxis(gain=0.01, offset=0.0) + + # Need 10 ticks to accumulate 1 sample + for _ in range(9): + result = producer(clock_tick) + assert result is None + + # 10th tick should produce 1 sample + result = producer(clock_tick) + assert result is not None + assert result.data.shape == (1,) + + def test_variable_n_time_afap_raises_error(self): + """Test that variable mode with AFAP clock raises error.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=None, mod=None)) + + clock_tick = AxisArray.LinearAxis(gain=0.0, offset=0.0) + + with pytest.raises(ValueError, match="Cannot use clock with gain=0"): + producer(clock_tick) + + def test_mod_rollover(self): + """Test counter rollover with mod.""" + producer = CounterProducer(CounterProducerSettings(fs=100.0, n_time=10, mod=8)) + + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=0.0) + + result = producer(clock_tick) + np.testing.assert_array_equal(result.data, [0, 1, 2, 3, 4, 5, 6, 7, 0, 1]) + + def test_continuity_across_calls(self): + """Test counter continuity across multiple calls.""" + producer = CounterProducer(CounterProducerSettings(fs=100.0, n_time=5, mod=None)) + + clock_tick = AxisArray.LinearAxis(gain=0.05, offset=0.0) + + results = [producer(clock_tick) for _ in range(4)] + agg = AxisArray.concatenate(*results, dim="time") + np.testing.assert_array_equal(agg.data, np.arange(20)) + + +class TestSinProducer: + """Tests for the example SinProducer implementation.""" + + def test_basic_sine_generation(self): + """Test basic sine wave generation.""" + fs = 1000.0 + freq = 10.0 + producer = SinProducer(SinProducerSettings(fs=fs, n_time=100, freq=freq, amp=1.0, phase=0.0)) + + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=0.0) + + result = producer(clock_tick) + + assert isinstance(result, AxisArray) + assert result.data.shape == (100,) + + # Verify it's actually a sine wave + t = np.arange(100) / fs + expected = np.sin(2 * np.pi * freq * t) + np.testing.assert_allclose(result.data, expected, rtol=1e-10) + + def test_sine_continuity_across_chunks(self): + """Test that sine wave is continuous across multiple chunks.""" + fs = 1000.0 + freq = 10.0 + producer = SinProducer(SinProducerSettings(fs=fs, n_time=50, freq=freq, amp=1.0, phase=0.0)) + + clock_tick = AxisArray.LinearAxis(gain=0.05, offset=0.0) + + # Generate 4 chunks + results = [producer(clock_tick) for _ in range(4)] + agg = AxisArray.concatenate(*results, dim="time") + + # Should be continuous sine wave + t = np.arange(200) / fs + expected = np.sin(2 * np.pi * freq * t) + np.testing.assert_allclose(agg.data, expected, rtol=1e-10) + + def test_sine_with_amplitude_and_phase(self): + """Test sine wave with custom amplitude and phase.""" + fs = 1000.0 + freq = 5.0 + amp = 2.5 + phase = np.pi / 4 + producer = SinProducer(SinProducerSettings(fs=fs, n_time=100, freq=freq, amp=amp, phase=phase)) + + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=0.0) + + result = producer(clock_tick) + + t = np.arange(100) / fs + expected = amp * np.sin(2 * np.pi * freq * t + phase) + np.testing.assert_allclose(result.data, expected, rtol=1e-10) + + +class TestClockDrivenProducerWithExternalClock: + """Tests simulating external clock patterns.""" + + def test_external_clock_with_fixed_n_time(self): + """Test external clock mode with fixed n_time.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=100, mod=None)) + + # Simulate external clock ticks + timestamps = [1.0, 1.1, 1.2, 1.3, 1.4] + clock_ticks = [AxisArray.LinearAxis(gain=0.1, offset=ts) for ts in timestamps] + + results = [producer(tick) for tick in clock_ticks] + + # Verify offsets match clock timestamps + offsets = [r.axes["time"].offset for r in results] + np.testing.assert_array_equal(offsets, timestamps) + + # Verify data continuity + agg = AxisArray.concatenate(*results, dim="time") + np.testing.assert_array_equal(agg.data, np.arange(500)) + + def test_external_clock_variable_chunk_sizes(self): + """Test external clock mode with variable chunk sizes.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=None, mod=None)) + + # Clock with varying rates + clock_ticks = [ + AxisArray.LinearAxis(gain=0.1, offset=0.0), # 100 samples + AxisArray.LinearAxis(gain=0.05, offset=0.1), # 50 samples + AxisArray.LinearAxis(gain=0.2, offset=0.15), # 200 samples + ] + + results = [producer(tick) for tick in clock_ticks] + + assert results[0].data.shape == (100,) + assert results[1].data.shape == (50,) + assert results[2].data.shape == (200,) + + # Verify data continuity + agg = AxisArray.concatenate(*results, dim="time") + np.testing.assert_array_equal(agg.data, np.arange(350)) + + +class TestClockDrivenStateManagement: + """Tests for state management behavior.""" + + def test_state_counter_increments(self): + """Test that internal counter increments correctly.""" + producer = CounterProducer(CounterProducerSettings(fs=100.0, n_time=10, mod=None)) + + clock_tick = AxisArray.LinearAxis(gain=0.1, offset=0.0) + + assert producer._state.counter == 0 + + producer(clock_tick) + assert producer._state.counter == 10 + + producer(clock_tick) + assert producer._state.counter == 20 + + def test_fractional_samples_accumulate(self): + """Test that fractional samples accumulate correctly.""" + producer = CounterProducer(CounterProducerSettings(fs=1000.0, n_time=None, mod=None)) + + # Clock at 3 Hz -> 333.33... samples per tick + clock_tick = AxisArray.LinearAxis(gain=1.0 / 3.0, offset=0.0) + + assert producer._state.fractional_samples == 0.0 + + producer(clock_tick) + # After first tick: 333.33... - 333 ≈ 0.33... + assert 0.33 < producer._state.fractional_samples < 0.34 + + producer(clock_tick) + # After second tick: 333.33... + 0.33... - 333 ≈ 0.66... + assert 0.66 < producer._state.fractional_samples < 0.67 + + def test_hash_returns_constant(self): + """Test that _hash_message returns constant (no state reset on clock changes).""" + producer = CounterProducer(CounterProducerSettings(fs=100.0, n_time=10, mod=None)) + + tick1 = AxisArray.LinearAxis(gain=0.1, offset=0.0) + tick2 = AxisArray.LinearAxis(gain=0.2, offset=1.0) + tick3 = AxisArray.LinearAxis(gain=0.0, offset=5.0) + + # All should return same hash (0) + assert producer._hash_message(tick1) == 0 + assert producer._hash_message(tick2) == 0 + assert producer._hash_message(tick3) == 0