ezmsg-sigproc 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/sigproc/sampler.py

@@ -1,147 +1,204 @@
 from dataclasses import dataclass, replace, field
-import logging
 import time
 
 import ezmsg.core as ez
 import numpy as np
 
-from .messages import TSMessage
+from ezmsg.util.messages.axisarray import AxisArray
 
-from typing import Optional, Any, Tuple, List, Dict
+from typing import Optional, Any, Tuple, List, Dict, AsyncGenerator
 
-@dataclass( frozen = True )
+
+@dataclass(frozen=True)
 class SampleTriggerMessage:
-    timestamp: float = field( default_factory = time.time )
-    period: Optional[ Tuple[ float, float ] ] = None
+    timestamp: float = field(default_factory=time.time)
+    period: Optional[Tuple[float, float]] = None
     value: Any = None
 
+
 @dataclass
 class SampleMessage:
     trigger: SampleTriggerMessage
-    sample: TSMessage
+    sample: AxisArray
 
-class SamplerSettings( ez.Settings ):
+
+class SamplerSettings(ez.Settings):
     buffer_dur: float
-    period: Optional[ Tuple[ float, float ] ] = None # Optional default period if unspecified in SampleTriggerMessage
-    value: Any = None # Optional default value if unspecified in SampleTriggerMessage
+    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 
+    # 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 ):
-    triggers: Dict[ SampleTriggerMessage, int ] = field( default_factory = dict )
-    last_msg: Optional[ TSMessage ] = None
-    buffer: Optional[ np.ndarray ] = None
 
-class Sampler( ez.Unit ):
+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
+
+
+class Sampler(ez.Unit):
     SETTINGS: SamplerSettings
     STATE: SamplerState
 
-    INPUT_TRIGGER = ez.InputStream( SampleTriggerMessage )
-    INPUT_SIGNAL = ez.InputStream( TSMessage )
-    OUTPUT_SAMPLE = ez.OutputStream( SampleMessage )
+    INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
+    INPUT_SETTINGS = ez.InputStream(SamplerSettings)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SAMPLE = ez.OutputStream(SampleMessage)
+
+    def initialize(self) -> None:
+        self.STATE.cur_settings = self.SETTINGS
+
+    @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:
+    @ez.subscriber(INPUT_TRIGGER)
+    async def on_trigger(self, msg: SampleTriggerMessage) -> None:
         if self.STATE.last_msg is not None:
-            fs = self.STATE.last_msg.fs
+            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.SETTINGS.period
-            value = msg.value if msg.value is not None else self.SETTINGS.value
+            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' )
+                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' )
+            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.SETTINGS.buffer_dur * fs )
-            req_buf_len = int( ( period[1] - period[0] ) * fs )
+            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.SETTINGS.buffer_dur=}' )
+                ez.logger.warning(
+                    f"Sampling failed: {period=} >= {self.STATE.cur_settings.buffer_dur=}"
+                )
                 return
 
             offset: int = 0
-            if self.SETTINGS.estimate_alignment:
+            if self.STATE.cur_settings.estimate_alignment:
                 # Do what we can with the wall clock to determine sample alignment
-                wall_delta = msg.timestamp - self.STATE.last_msg.timestamp
-                offset = int( wall_delta * fs )
+                wall_delta = msg.timestamp - last_msg_timestamp
+                offset = int(wall_delta * fs)
 
             # Check that current buffer accumulation allows for offset - period start
-            if -min( offset + start_offset, 0 ) >= self.STATE.buffer.shape[0]:
-                ez.logger.warning( 'Sampling failed: insufficient buffer accumulation for requested sample period' )
+            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
 
-            trigger = replace( msg, period = period, value = value )
-            self.STATE.triggers[ trigger ] = offset 
+            self.STATE.triggers[replace(msg, period=period, value=value)] = offset
 
-        else: ez.logger.warning( 'Sampling failed: no signal to sample yet' )
+        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: TSMessage ) -> None:
+    @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.swapaxes( msg.data, msg.time_dim, 0 )
-        last_msg_data = np.swapaxes( last_msg.data, last_msg.time_dim, 0 )
-        
-        if ( # Check if signal properties have changed in a breaking way
-            msg.fs != last_msg.fs or \
-            msg.time_dim != last_msg.time_dim or \
-            msg_data.shape[1:] != last_msg_data.shape[1:]
-        ):
+        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' )
+            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 )
+        self.STATE.buffer = (
+            msg_data
+            if self.STATE.buffer is None
+            else np.concatenate((self.STATE.buffer, msg_data), axis=0)
+        )
 
-        pub_samples: List[ SampleMessage ] = []
-        remaining_triggers: Dict[ SampleTriggerMessage, int ] = dict()
+        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
+
+        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 = offset - msg.n_time
-            start = offset + int( trigger.period[0] * msg.fs )
-            stop = offset + int( trigger.period[1] * msg.fs )
+            offset -= msg_data.shape[0]
+            start = offset + int(trigger.period[0] * fs)
+            stop = offset + int(trigger.period[1] * fs)
+
+            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)
+
+                sample_offset = buffer_offset[start]
+                sample_axis = replace(axis, offset=sample_offset)
+                sample_axes = {**msg.axes, **{axis_name: sample_axis}}
 
-            if stop < 0: # We should be able to dispatch a sample
-                sample_data = self.STATE.buffer[ start : stop, ... ]
-                sample_data = np.swapaxes( sample_data, msg.time_dim, 0 )
-                
-                pub_samples.append( SampleMessage( 
-                    trigger = trigger,
-                    sample = replace( msg, data = sample_data )
-                ) )
+                pub_samples.append(
+                    SampleMessage(
+                        trigger=trigger,
+                        sample=replace(msg, data=sample_data, axes=sample_axes),
+                    )
+                )
 
-            else: remaining_triggers[ trigger ] = offset
+            else:
+                remaining_triggers[trigger] = offset
 
         for sample in pub_samples:
-            yield self.OUTPUT_SAMPLE, sample 
-            
+            yield self.OUTPUT_SAMPLE, sample
+
         self.STATE.triggers = remaining_triggers
 
-        buf_len = int( self.SETTINGS.buffer_dur * msg.fs )
-        self.STATE.buffer = self.STATE.buffer[ -buf_len:, ... ]
+        buf_len = int(self.STATE.cur_settings.buffer_dur * fs)
+        self.STATE.buffer = self.STATE.buffer[-buf_len:, ...]
         self.STATE.last_msg = msg
 
 
@@ -153,99 +210,78 @@ from ezmsg.sigproc.synth import Oscillator, OscillatorSettings
 
 from typing import AsyncGenerator
 
-class SampleFormatter( ez.Unit ):
-
-    INPUT = ez.InputStream( SampleMessage )
-    OUTPUT = ez.OutputStream( str )
-
-    @ez.subscriber( INPUT )
-    @ez.publisher( OUTPUT )
-    async def format( self, msg: SampleMessage ) -> AsyncGenerator:
-        str_msg = f'Trigger: {msg.trigger.value}, '
-        str_msg += f'{msg.sample.n_time} samples @ {msg.sample.fs} Hz, '
-        
-        time_axis = np.arange(msg.sample.n_time) / msg.sample.fs
-        time_axis = time_axis + msg.trigger.period[0]
-        str_msg += f'[{time_axis[0]},{time_axis[-1]})'
-        
-        yield self.OUTPUT, str_msg
 
-class TriggerGeneratorSettings( ez.Settings ):
-    period: Tuple[ float, float ] # sec
-    prewait: float = 0.5 # sec
-    publish_period: float = 5.0 # sec
+class TriggerGeneratorSettings(ez.Settings):
+    period: Tuple[float, float]  # sec
+    prewait: float = 0.5  # sec
+    publish_period: float = 5.0  # sec
 
-class TriggerGenerator( ez.Unit ):
 
+class TriggerGenerator(ez.Unit):
     SETTINGS: TriggerGeneratorSettings
 
-    OUTPUT_TRIGGER = ez.OutputStream( SampleTriggerMessage )
+    OUTPUT_TRIGGER = ez.OutputStream(SampleTriggerMessage)
 
-    @ez.publisher( OUTPUT_TRIGGER )
-    async def generate( self ) -> AsyncGenerator:
-        await asyncio.sleep( self.SETTINGS.prewait )
+    @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
+                period=self.SETTINGS.period, value=output
             )
 
-            await asyncio.sleep( self.SETTINGS.publish_period )
+            await asyncio.sleep(self.SETTINGS.publish_period)
             output += 1
 
-class SamplerTestSystemSettings( ez.Settings ):
+
+class SamplerTestSystemSettings(ez.Settings):
     sampler_settings: SamplerSettings
     trigger_settings: TriggerGeneratorSettings
 
-class SamplerTestSystem( ez.System ):
 
+class SamplerTestSystem(ez.System):
     SETTINGS: SamplerTestSystemSettings
 
     OSC = Oscillator()
     SAMPLER = Sampler()
     TRIGGER = TriggerGenerator()
-    FORMATTER = SampleFormatter()
     DEBUG = DebugLog()
 
-    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
-        ) )
-
-    def network( self ) -> ez.NetworkDefinition:
-        return (
-            ( self.OSC.OUTPUT_SIGNAL, self.SAMPLER.INPUT_SIGNAL ),
-            ( self.TRIGGER.OUTPUT_TRIGGER, self.SAMPLER.INPUT_TRIGGER ),
+    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
+            )
+        )
 
-            ( self.TRIGGER.OUTPUT_TRIGGER, self.DEBUG.INPUT ),
-            ( self.SAMPLER.OUTPUT_SAMPLE, self.FORMATTER.INPUT ),
-            ( self.FORMATTER.OUTPUT, self.DEBUG.INPUT )
+    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),
         )
 
-if __name__ == '__main__':
 
-    settings = SamplerTestSystemSettings( 
-        sampler_settings = SamplerSettings(
-            buffer_dur = 5.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
         ),
-        trigger_settings = TriggerGeneratorSettings(
-            period = ( 1.0, 2.0 ),
-            prewait = 0.5,
-            publish_period = 5.0
-        )
     )
 
-    system = SamplerTestSystem( settings )
+    system = SamplerTestSystem(settings)
 
-    ez.run_system( system )
+    ez.run_system(system)

ezmsg/sigproc/spectral.py

@@ -0,0 +1,132 @@
+import enum
+
+from dataclasses import replace
+
+import numpy as np
+import ezmsg.core as ez
+
+from ezmsg.util.messages.axisarray import AxisArray
+
+from typing import Optional, AsyncGenerator
+
+
+class OptionsEnum(enum.Enum):
+    @classmethod
+    def options(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+class WindowFunction(OptionsEnum):
+    NONE = "None (Rectangular)"
+    HAMMING = "Hamming"
+    HANNING = "Hanning"
+    BARTLETT = "Bartlett"
+    BLACKMAN = "Blackman"
+
+
+WINDOWS = {
+    WindowFunction.NONE: np.ones,
+    WindowFunction.HAMMING: np.hamming,
+    WindowFunction.HANNING: np.hanning,
+    WindowFunction.BARTLETT: np.bartlett,
+    WindowFunction.BLACKMAN: np.blackman,
+}
+
+
+class SpectralTransform(OptionsEnum):
+    RAW_COMPLEX = "Complex FFT Output"
+    REAL = "Real Component of FFT"
+    IMAG = "Imaginary Component of FFT"
+    REL_POWER = "Relative Power"
+    REL_DB = "Log Power (Relative dB)"
+
+
+class SpectralOutput(OptionsEnum):
+    FULL = "Full Spectrum"
+    POSITIVE = "Positive Frequencies"
+    NEGATIVE = "Negative Frequencies"
+
+
+class SpectrumSettings(ez.Settings):
+    axis: Optional[str] = None
+    # n: Optional[int] = None # n parameter for fft
+    out_axis: Optional[str] = "freq"  # If none; don't change dim name
+    window: WindowFunction = WindowFunction.HAMMING
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    output: SpectralOutput = SpectralOutput.POSITIVE
+
+
+class SpectrumState(ez.State):
+    cur_settings: SpectrumSettings
+
+
+class Spectrum(ez.Unit):
+    SETTINGS: SpectrumSettings
+    STATE: SpectrumState
+
+    INPUT_SETTINGS = ez.InputStream(SpectrumSettings)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    def initialize(self) -> None:
+        self.STATE.cur_settings = self.SETTINGS
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: SpectrumSettings):
+        self.STATE.cur_settings = msg
+
+    @ez.subscriber(INPUT_SIGNAL)
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def on_data(self, message: AxisArray) -> AsyncGenerator:
+        axis_name = self.STATE.cur_settings.axis
+        if axis_name is None:
+            axis_name = message.dims[0]
+        axis_idx = message.get_axis_idx(axis_name)
+        axis = message.get_axis(axis_name)
+
+        spectrum = np.moveaxis(message.data, axis_idx, -1)
+
+        n_time = message.data.shape[axis_idx]
+        window = WINDOWS[self.STATE.cur_settings.window](n_time)
+
+        spectrum = np.fft.fft(spectrum * window) / n_time
+        spectrum = np.fft.fftshift(spectrum, axes=-1)
+        freqs = np.fft.fftshift(np.fft.fftfreq(n_time, d=axis.gain), axes=-1)
+
+        if self.STATE.cur_settings.transform != SpectralTransform.RAW_COMPLEX:
+            if self.STATE.cur_settings.transform == SpectralTransform.REAL:
+                spectrum = spectrum.real
+            elif self.STATE.cur_settings.transform == SpectralTransform.IMAG:
+                spectrum = spectrum.imag
+            else:
+                scale = np.sum(window**2.0) * axis.gain
+                spectrum = (2.0 * (np.abs(spectrum) ** 2.0)) / scale
+
+                if self.STATE.cur_settings.transform == SpectralTransform.REL_DB:
+                    spectrum = 10 * np.log10(spectrum)
+
+        axis_offset = freqs[0]
+        if self.STATE.cur_settings.output == SpectralOutput.POSITIVE:
+            axis_offset = freqs[n_time // 2]
+            spectrum = spectrum[..., n_time // 2 :]
+        elif self.STATE.cur_settings.output == SpectralOutput.NEGATIVE:
+            spectrum = spectrum[..., : n_time // 2]
+
+        spectrum = np.moveaxis(spectrum, axis_idx, -1)
+
+        out_axis = self.SETTINGS.out_axis
+        if out_axis is None:
+            out_axis = axis_name
+
+        freq_axis = AxisArray.Axis(
+            unit="Hz", gain=1.0 / (axis.gain * n_time), offset=axis_offset
+        )
+        new_axes = {**message.axes, **{out_axis: freq_axis}}
+
+        new_dims = [d for d in message.dims]
+        if self.SETTINGS.out_axis is not None:
+            new_dims[axis_idx] = self.SETTINGS.out_axis
+
+        out_msg = replace(message, data=spectrum, dims=new_dims, axes=new_axes)
+
+        yield self.OUTPUT_SIGNAL, out_msg