ezmsg-baseproc 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/baseproc/__init__.py

@@ -15,6 +15,14 @@ from .clock import (
     ClockState,
 )
 
+# Clock-driven producers
+from .clockdriven import (
+    BaseClockDrivenProducer,
+    ClockDrivenSettings,
+    ClockDrivenSettingsType,
+    ClockDrivenState,
+)
+
 # Composite processor classes
 from .composite import (
     CompositeProcessor,
@@ -74,15 +82,18 @@ from .stateful import (
 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 @@ __all__ = [
     "ConsumerType",
     "TransformerType",
     "AdaptiveTransformerType",
+    "ClockDrivenProducerType",
     # Decorators
     "processor_state",
     # Base processor classes
@@ -131,6 +143,11 @@ __all__ = [
     "BaseStatefulTransformer",
     "BaseAdaptiveTransformer",
     "BaseAsyncTransformer",
+    # Clock-driven producers
+    "BaseClockDrivenProducer",
+    "ClockDrivenSettings",
+    "ClockDrivenSettingsType",
+    "ClockDrivenState",
     # Composite classes
     "CompositeStateful",
     "CompositeProcessor",
@@ -141,12 +158,14 @@ __all__ = [
     "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",

ezmsg/baseproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.1.0'
-__version_tuple__ = version_tuple = (1, 1, 0)
+__version__ = version = '1.2.0'
+__version_tuple__ = version_tuple = (1, 2, 0)
 
 __commit_id__ = commit_id = None

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

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

ezmsg/baseproc/units.py

@@ -7,8 +7,9 @@ from abc import ABC, abstractmethod
 
 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 @@ TransformerType = typing.TypeVar(
     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 @@ class BaseAdaptiveTransformerUnit(
         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

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}")

{ezmsg_baseproc-1.1.0.dist-info → ezmsg_baseproc-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-baseproc
-Version: 1.1.0
+Version: 1.2.0
 Summary: Base processor classes and protocols for ezmsg signal processing pipelines
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

{ezmsg_baseproc-1.1.0.dist-info → ezmsg_baseproc-1.2.0.dist-info}/RECORD

@@ -1,18 +1,19 @@
-ezmsg/baseproc/__init__.py,sha256=OLpOqiGBjQSrLPjAtWEHdCS8Sm_2tJ55I7bpCtXntYY,4243
-ezmsg/baseproc/__version__.py,sha256=ePNVzJOkxR8FY5bezqKQ_fgBRbzH1G7QTaRDHvGQRAY,704
+ezmsg/baseproc/__init__.py,sha256=cSkyXkfPvIiAxe6maHiZvJpFbdRw3xTEitBClXeehec,4755
+ezmsg/baseproc/__version__.py,sha256=-uLONazCO1SzFfcY-K6A1keL--LIVfTYccGX6ciADac,704
 ezmsg/baseproc/clock.py,sha256=jRJGxaWWBek503OcGRYsW4Qc7Lt4m-tVl1wn_v-qCIk,3254
+ezmsg/baseproc/clockdriven.py,sha256=ckPZSVHZfYfjFRHDCERUWjDwyQgOu-aRTNsJ3EDFBaI,6161
 ezmsg/baseproc/composite.py,sha256=Lin4K_rmS2Tnxt-m8daP-PUyeeqL4id5JkVh-AUNrQw,14901
-ezmsg/baseproc/counter.py,sha256=3JR4AQ_wwZhdjKQdPlLhjhNVLHE5fj9ggwUPtJhWDS4,4601
+ezmsg/baseproc/counter.py,sha256=kcBPiVxMPULp4ojnVESNw7mn_4v0xSODfASHrL83GtM,2168
 ezmsg/baseproc/processor.py,sha256=Ir9FtNuVG4yc-frwNxoYrlld99ff1mXwwGWaHxEJ6tY,8056
 ezmsg/baseproc/protocols.py,sha256=O3Qp0ymE9Ovlmh8t22v-lMmFzuWK0D93REAYMnJV3xA,5106
 ezmsg/baseproc/stateful.py,sha256=-jjAZIyJA5eiTECi1fSfazfqgv__RtyqPp1ZvLFFIDI,11424
-ezmsg/baseproc/units.py,sha256=TRhjDKw0lqUUst7BHYJKP3AhGpRd6mvcdKxULfeWjA0,10283
+ezmsg/baseproc/units.py,sha256=byFijVLEZFO145HE74sZk1_qpCu6nFjB8-vSYz9Grds,12077
 ezmsg/baseproc/util/__init__.py,sha256=hvMUJOBuqioER50GZ5-GZiQbQ9NtQYEze13ZlR2jbMA,37
 ezmsg/baseproc/util/asio.py,sha256=0sF5oDc58DSLlcEgoUpNiqjjcbqnZhjSpQrXn6IdosM,4960
 ezmsg/baseproc/util/message.py,sha256=l_b1b6bXX8N6VF9RbUELzsHs73cKkDURBdIr0lt3CY0,909
 ezmsg/baseproc/util/profile.py,sha256=MOQDsFsW6ddXT0uAOgytW3aK_AZW5ieA16Pz2hWuE2o,6189
-ezmsg/baseproc/util/typeresolution.py,sha256=WCHHYIrMMZ1CfwJWVlJPQgFyY2gnGRNFJVQynAsee7Y,3113
-ezmsg_baseproc-1.1.0.dist-info/METADATA,sha256=zv5si0HBsR2t8lzzmwzeuLiik4SeanZvmV-JNWqOnKg,3415
-ezmsg_baseproc-1.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ezmsg_baseproc-1.1.0.dist-info/licenses/LICENSE,sha256=BDD8rfac1Ur7mp0_3izEdr6fHgSA3Or6U1Kb0ZAWsow,1066
-ezmsg_baseproc-1.1.0.dist-info/RECORD,,
+ezmsg/baseproc/util/typeresolution.py,sha256=5on4QcrYd1rxsRoDEqivNjuWT5BkU-Wg7XdTNaOircI,3485
+ezmsg_baseproc-1.2.0.dist-info/METADATA,sha256=H2jTn5VSw0pEwQWLdd3Hu0G1OCbbKM-On4AmKyyyfm4,3415
+ezmsg_baseproc-1.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ezmsg_baseproc-1.2.0.dist-info/licenses/LICENSE,sha256=BDD8rfac1Ur7mp0_3izEdr6fHgSA3Or6U1Kb0ZAWsow,1066
+ezmsg_baseproc-1.2.0.dist-info/RECORD,,