diff --git a/src/ezmsg/sigproc/align.py b/src/ezmsg/sigproc/align.py new file mode 100644 index 00000000..3bd018fa --- /dev/null +++ b/src/ezmsg/sigproc/align.py @@ -0,0 +1,241 @@ +"""Time-align two AxisArray streams, outputting paired aligned chunks.""" + +from __future__ import annotations + +import math +import typing + +import ezmsg.core as ez +import numpy as np +from ezmsg.baseproc.protocols import processor_state +from ezmsg.baseproc.stateful import BaseStatefulTransformer +from ezmsg.util.messages.axisarray import AxisArray + +from .util.axisarray_buffer import HybridAxisArrayBuffer + + +class AlignAlongAxisSettings(ez.Settings): + axis: str = "time" + """Axis used for alignment (typically the time axis).""" + + buffer_dur: float = 10.0 + """Buffer duration in seconds for each input stream.""" + + +@processor_state +class AlignAlongAxisState: + gain: float | None = None + align_axis: str | None = None + aligned: bool = False + buf_a: HybridAxisArrayBuffer | None = None + buf_b: HybridAxisArrayBuffer | None = None + # Per-input non-alignment shape for reset detection. + a_shape_sig: tuple[int, ...] | None = None + b_shape_sig: tuple[int, ...] | None = None + + +_AlignPair = tuple[AxisArray, AxisArray] + + +class AlignAlongAxisProcessor( + BaseStatefulTransformer[ + AlignAlongAxisSettings, + AxisArray, + _AlignPair | None, + AlignAlongAxisState, + ] +): + """Processor that time-aligns two AxisArray streams. + + Input A flows through ``__call__`` / ``_process`` with automatic + hash-based reset. Input B flows through :meth:`push_b`. + + Returns ``(aligned_a, aligned_b)`` when alignment succeeds, else ``None``. + """ + + # -- Helpers ------------------------------------------------------------- + + def _extract_gain(self, message: AxisArray) -> float | None: + align_name = self.settings.axis or message.dims[0] + ax = message.axes.get(align_name) + if ax is not None and hasattr(ax, "gain"): + return ax.gain + if ax is not None and hasattr(ax, "data") and len(ax.data) > 1: + return float(ax.data[-1] - ax.data[0]) / (len(ax.data) - 1) + return None + + @staticmethod + def _non_align_shape(message: AxisArray, align_axis: str) -> tuple[int, ...]: + align_idx = message.dims.index(align_axis) + return tuple(s for i, s in enumerate(message.data.shape) if i != align_idx) + + # -- Reset helpers ------------------------------------------------------- + + def _full_reset(self, align_axis: str) -> None: + """ + Reset state. Called either on input A (through default __call__ path) + or on Input B. + Args: + align_axis: + + Returns: + + """ + self._state.align_axis = align_axis + self._state.buf_a = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) + self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) + self._state.gain = None + self._state.aligned = False + self._state.a_shape_sig = None + self._state.b_shape_sig = None + + def _reset_a_state(self) -> None: + self._state.buf_a = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=self._state.align_axis) + + def _reset_b_state(self) -> None: + self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=self._state.align_axis) + + # -- BaseStatefulTransformer interface ------------------------------------ + + def _hash_message(self, message: AxisArray) -> int: + return hash(self._extract_gain(message)) + + def _reset_state(self, message: AxisArray) -> None: + align_axis = self.settings.axis or message.dims[0] + self._full_reset(align_axis) + + def _process(self, message: AxisArray) -> _AlignPair | None: + """Process input A: detect shape changes, buffer, try align.""" + shape_sig = self._non_align_shape(message, self._state.align_axis) + if self._state.a_shape_sig is not None and shape_sig != self._state.a_shape_sig: + self._reset_a_state() + self._state.aligned = False + self._state.a_shape_sig = shape_sig + + self._state.buf_a.write(message) + if self._state.gain is None: + self._state.gain = self._state.buf_a.axis_gain + return self._try_align() + + # -- Input B entry point ------------------------------------------------ + + def push_b(self, message: AxisArray) -> _AlignPair | None: + """Process input B: check gain, detect shape changes, buffer, try align.""" + align_axis = self.settings.axis or message.dims[0] + + # Gain compatibility check. + b_gain = self._extract_gain(message) + if self._state.gain is not None and not math.isclose(b_gain, self._state.gain): + self._full_reset(align_axis) + self._hash = self._hash_message(message) + + # Lazy-create buf_b if B arrives before A. + if self._state.buf_b is None: + if self._state.align_axis is None: + self._state.align_axis = align_axis + self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) + + shape_sig = self._non_align_shape(message, align_axis) + if self._state.b_shape_sig is not None and shape_sig != self._state.b_shape_sig: + self._reset_b_state() + self._state.aligned = False + self._state.b_shape_sig = shape_sig + + self._state.buf_b.write(message) + if self._state.gain is None: + self._state.gain = self._state.buf_b.axis_gain + return self._try_align() + + # -- Core alignment logic ----------------------------------------------- + + def _try_align(self) -> _AlignPair | None: + """Align and read from both buffers, returning the pair ``(a, b)``.""" + if self._state.buf_a is None or self._state.buf_b is None: + return None + if self._state.buf_a.is_empty() or self._state.buf_b.is_empty(): + return None + + gain = self._state.gain + + # --- Initial alignment (runs once) --- + if not self._state.aligned: + first_a = self._state.buf_a.axis_first_value + final_a = self._state.buf_a.axis_final_value + first_b = self._state.buf_b.axis_first_value + final_b = self._state.buf_b.axis_final_value + + overlap_start = max(first_a, first_b) + overlap_end = min(final_a, final_b) + + if overlap_end < overlap_start - gain / 2: + if final_a < first_b: + self._state.buf_a.seek(self._state.buf_a.available()) + elif final_b < first_a: + self._state.buf_b.seek(self._state.buf_b.available()) + return None + + if first_a < overlap_start - gain / 2: + self._state.buf_a.seek(int(round((overlap_start - first_a) / gain))) + if first_b < overlap_start - gain / 2: + self._state.buf_b.seek(int(round((overlap_start - first_b) / gain))) + + # --- Read aligned samples --- + n_read = min(self._state.buf_a.available(), self._state.buf_b.available()) + if n_read <= 0: + return None + + aa_a = self._state.buf_a.read(n_read) + aa_b = self._state.buf_b.read(n_read) + if aa_a is None or aa_b is None: + return None + + if not self._state.aligned: + axis_a = aa_a.axes.get(self._state.align_axis) + axis_b = aa_b.axes.get(self._state.align_axis) + if axis_a is not None and axis_b is not None: + off_a = axis_a.value(0) if hasattr(axis_a, "value") else None + off_b = axis_b.value(0) if hasattr(axis_b, "value") else None + if off_a is not None and off_b is not None: + if not np.isclose(off_a, off_b, atol=abs(gain) * 1e-6): + raise RuntimeError( + f"Offset mismatch after alignment: " f"off_a={off_a}, off_b={off_b}, gain={gain}" + ) + self._state.aligned = True + + return aa_a, aa_b + + +class AlignAlongAxis(ez.Unit): + """Time-align two AxisArray streams and output paired aligned chunks. + + Each subscriber can publish to *both* output streams; when alignment + succeeds, a paired (A, B) result is yielded to the respective outputs. + """ + + SETTINGS = AlignAlongAxisSettings + + INPUT_SIGNAL_A = ez.InputStream(AxisArray) + INPUT_SIGNAL_B = ez.InputStream(AxisArray) + OUTPUT_SIGNAL_A = ez.OutputStream(AxisArray) + OUTPUT_SIGNAL_B = ez.OutputStream(AxisArray) + + async def initialize(self) -> None: + self.processor = AlignAlongAxisProcessor(settings=self.SETTINGS) + + @ez.subscriber(INPUT_SIGNAL_A, zero_copy=True) + @ez.publisher(OUTPUT_SIGNAL_A) + @ez.publisher(OUTPUT_SIGNAL_B) + async def on_a(self, msg: AxisArray) -> typing.AsyncGenerator: + pair = await self.processor.__acall__(msg) + if pair is not None: + yield self.OUTPUT_SIGNAL_A, pair[0] + yield self.OUTPUT_SIGNAL_B, pair[1] + + @ez.subscriber(INPUT_SIGNAL_B, zero_copy=True) + @ez.publisher(OUTPUT_SIGNAL_A) + @ez.publisher(OUTPUT_SIGNAL_B) + async def on_b(self, msg: AxisArray) -> typing.AsyncGenerator: + pair = self.processor.push_b(msg) + if pair is not None: + yield self.OUTPUT_SIGNAL_A, pair[0] + yield self.OUTPUT_SIGNAL_B, pair[1] diff --git a/src/ezmsg/sigproc/concat.py b/src/ezmsg/sigproc/concat.py new file mode 100644 index 00000000..7b5381d6 --- /dev/null +++ b/src/ezmsg/sigproc/concat.py @@ -0,0 +1,338 @@ +"""Concatenate two AxisArray streams along an existing or new axis.""" + +from __future__ import annotations + +import asyncio +import typing +from dataclasses import dataclass, field + +import ezmsg.core as ez +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, AxisBase, CoordinateAxis +from ezmsg.util.messages.util import replace + +# --------------------------------------------------------------------------- +# Shared helpers (also used by merge.py) +# --------------------------------------------------------------------------- + + +def _build_merged_coordinate_axis( + axis_a: CoordinateAxis, + axis_b: CoordinateAxis, + relabel: bool, + label_a: str, + label_b: str, +) -> CoordinateAxis: + """Build a merged CoordinateAxis from two per-input axes. + + Handles both simple (string/numeric) and structured (numpy struct) dtypes. + When *relabel* is True and the dtype is structured, only the ``"label"`` + field is modified (or created if absent). + """ + data_a = axis_a.data + data_b = axis_b.data + + if data_a.dtype.names is not None or data_b.dtype.names is not None: + return _merge_struct_axes(data_a, data_b, relabel, label_a, label_b, axis_a) + + # Simple (non-struct) path — current behaviour. + if relabel: + labels_a = np.array([str(lbl) + label_a for lbl in data_a]) + labels_b = np.array([str(lbl) + label_b for lbl in data_b]) + else: + labels_a = data_a + labels_b = data_b + return CoordinateAxis( + data=np.concatenate([labels_a, labels_b]), + dims=axis_a.dims, + unit=axis_a.unit, + ) + + +def _merge_struct_axes( + data_a: np.ndarray, + data_b: np.ndarray, + relabel: bool, + label_a: str, + label_b: str, + ref_axis: CoordinateAxis, +) -> CoordinateAxis: + """Merge two structured-dtype coordinate arrays, preserving all fields.""" + names_a = set(data_a.dtype.names or ()) + names_b = set(data_b.dtype.names or ()) + + # Build the union dtype. Shared fields must have compatible sub-dtypes. + union_fields: list[tuple[str, np.dtype]] = [] + seen: set[str] = set() + + for src_names, src_dtype in [(data_a.dtype.names or (), data_a.dtype), (data_b.dtype.names or (), data_b.dtype)]: + for name in src_names: + if name in seen: + continue + seen.add(name) + dt_a = data_a.dtype[name] if name in names_a else None + dt_b = data_b.dtype[name] if name in names_b else None + if dt_a is not None and dt_b is not None: + resolved = _resolve_field_dtype(name, dt_a, dt_b) + else: + resolved = dt_a if dt_a is not None else dt_b + union_fields.append((name, resolved)) + + # If relabel and "label" is not already a field, add it. + has_label = "label" in seen + if relabel and not has_label: + max_len = max( + max((len(str(i)) for i in range(len(data_a))), default=1), + max((len(str(i)) for i in range(len(data_b))), default=1), + ) + suffix_len = max(len(label_a), len(label_b)) + union_fields.append(("label", np.dtype(f"U{max_len + suffix_len}"))) + has_label = True + + union_dtype = np.dtype(union_fields) + merged = np.zeros(len(data_a) + len(data_b), dtype=union_dtype) + + # Copy values from A. + for name in data_a.dtype.names or (): + merged[name][: len(data_a)] = data_a[name] + # Copy values from B. + for name in data_b.dtype.names or (): + merged[name][len(data_a) :] = data_b[name] + + # Relabel only the "label" field. + if relabel and has_label: + for i in range(len(data_a)): + src = str(data_a[i]["label"]) if "label" in names_a else str(i) + merged[i]["label"] = src + label_a + for j in range(len(data_b)): + src = str(data_b[j]["label"]) if "label" in names_b else str(j) + merged[len(data_a) + j]["label"] = src + label_b + + return CoordinateAxis(data=merged, dims=ref_axis.dims, unit=ref_axis.unit) + + +def _resolve_field_dtype(name: str, dt_a: np.dtype, dt_b: np.dtype) -> np.dtype: + """Resolve a shared struct field's dtype. String fields use the wider width.""" + if dt_a == dt_b: + return dt_a + if dt_a.kind == "U" and dt_b.kind == "U": + return np.dtype(f"U{max(dt_a.itemsize // 4, dt_b.itemsize // 4)}") + raise ValueError(f"Incompatible dtypes for shared struct field {name!r}: {dt_a} vs {dt_b}") + + +def _validate_shared_axes( + a: AxisArray, + b: AxisArray, + concat_dim: str, + align_dim: str | None, + assert_flag: bool, +) -> None: + """Raise ValueError if shared CoordinateAxis .data arrays differ.""" + if not assert_flag: + return + skip = {concat_dim, align_dim} + for name in a.axes: + if name in skip or name not in b.axes: + continue + ax_a, ax_b = a.axes[name], b.axes[name] + if hasattr(ax_a, "data") and hasattr(ax_b, "data"): + if not np.array_equal(ax_a.data, ax_b.data): + raise ValueError(f"Shared axis {name!r} has different .data between inputs A and B") + if hasattr(ax_a, "gain") and hasattr(ax_b, "gain"): + if ax_a.gain != ax_b.gain: + raise ValueError(f"Shared axis {name!r} has different gain: {ax_a.gain} vs {ax_b.gain}") + + +def _build_cached_axes( + a: AxisArray, + concat_dim: str, + align_dim: str | None, + merged_concat_axis: CoordinateAxis | None, +) -> dict[str, AxisBase]: + """Build the output axes dict (everything except the alignment axis).""" + axes: dict[str, AxisBase] = {} + for name, ax in a.axes.items(): + if name == align_dim: + continue + if name == concat_dim and merged_concat_axis is not None: + axes[name] = merged_concat_axis + else: + axes[name] = ax + if concat_dim not in axes and merged_concat_axis is not None: + axes[concat_dim] = merged_concat_axis + return axes + + +# --------------------------------------------------------------------------- +# ConcatProcessor / Concat unit +# --------------------------------------------------------------------------- + + +class ConcatSettings(ez.Settings): + axis: str = "ch" + """Axis along which to concatenate the two signals.""" + + relabel_axis: bool = True + """Whether to relabel coordinate axis labels to ensure uniqueness.""" + + label_a: str = "_a" + """Suffix appended to signal A labels when relabel_axis is True.""" + + label_b: str = "_b" + """Suffix appended to signal B labels when relabel_axis is True.""" + + assert_identical_shared_axes: bool = False + """If True, raise ValueError when shared CoordinateAxis .data arrays differ.""" + + new_key: str | None = None + """Output AxisArray key. If None, uses the key from signal A.""" + + +@dataclass +class ConcatState: + queue_a: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue) + queue_b: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue) + merged_concat_axis: CoordinateAxis | None = None + cached_axes: dict[str, AxisBase] | None = None + # Fingerprints for cache invalidation. + a_fingerprint: tuple | None = None + b_fingerprint: tuple | None = None + + +class ConcatProcessor: + """Concatenate paired AxisArray messages from two input queues. + + Uses FIFO queue pairing (like :class:`~ezmsg.sigproc.math.add.AddProcessor`). + No time-alignment or buffering — inputs are assumed pre-synchronized. + """ + + def __init__(self, settings: ConcatSettings): + self.settings = settings + self._state = ConcatState() + + @property + def state(self) -> ConcatState: + return self._state + + @state.setter + def state(self, state: ConcatState | bytes | None) -> None: + if state is not None: + self._state = state + + def push_a(self, msg: AxisArray) -> None: + self._state.queue_a.put_nowait(msg) + + def push_b(self, msg: AxisArray) -> None: + self._state.queue_b.put_nowait(msg) + + async def __acall__(self) -> AxisArray: + a = await self._state.queue_a.get() + b = await self._state.queue_b.get() + return self._concat(a, b) + + def _concat(self, a: AxisArray, b: AxisArray) -> AxisArray: + """Concatenate *a* and *b* along the configured axis.""" + concat_dim = self.settings.axis + fp_a = self._fingerprint(a) + fp_b = self._fingerprint(b) + if fp_a != self._state.a_fingerprint or fp_b != self._state.b_fingerprint: + self._rebuild_cache(a, b) + self._state.a_fingerprint = fp_a + self._state.b_fingerprint = fp_b + + new_axis = concat_dim not in a.dims + + # expand_dims for new-axis concatenation. + if new_axis: + a = replace(a, data=np.expand_dims(a.data, axis=-1), dims=[*a.dims, concat_dim]) + b = replace(b, data=np.expand_dims(b.data, axis=-1), dims=[*b.dims, concat_dim]) + + concat_idx = a.dims.index(concat_dim) + data = np.concatenate([a.data, b.data], axis=concat_idx) + + # Build axes: use cached axes + live alignment axis from a. + axes = dict(self._state.cached_axes) if self._state.cached_axes is not None else dict(a.axes) + # Re-insert any axis that changes per-message (e.g. time offset). + for name, ax in a.axes.items(): + if name not in axes: + axes[name] = ax + + key = self.settings.new_key if self.settings.new_key is not None else a.key + return AxisArray(data, dims=list(a.dims), axes=axes, key=key) + + def _fingerprint(self, msg: AxisArray) -> tuple: + concat_dim = self.settings.axis + ax = msg.axes.get(concat_dim) + ax_hash = hash(ax.data.tobytes()) if ax is not None and hasattr(ax, "data") else None + return (tuple(msg.dims), msg.data.shape, ax_hash) + + def _rebuild_cache(self, a: AxisArray, b: AxisArray) -> None: + concat_dim = self.settings.axis + + # Validate shared axes. + _validate_shared_axes( + a, + b, + concat_dim, + align_dim=None, + assert_flag=self.settings.assert_identical_shared_axes, + ) + + # New-axis validation: all other dims must match. + if concat_dim not in a.dims or concat_dim not in b.dims: + for i, (d, sa, sb) in enumerate(zip(a.dims, a.data.shape, b.data.shape)): + if sa != sb: + raise ValueError( + f"Cannot concatenate along new axis {concat_dim!r}: " + f"dimension {d!r} has size {sa} in A but {sb} in B" + ) + + # Build merged concat axis. + ax_a = a.axes.get(concat_dim) + ax_b = b.axes.get(concat_dim) + if ax_a is not None and ax_b is not None and hasattr(ax_a, "data") and hasattr(ax_b, "data"): + self._state.merged_concat_axis = _build_merged_coordinate_axis( + ax_a, + ax_b, + relabel=self.settings.relabel_axis, + label_a=self.settings.label_a, + label_b=self.settings.label_b, + ) + else: + self._state.merged_concat_axis = None + + self._state.cached_axes = _build_cached_axes( + a, + concat_dim, + align_dim=None, + merged_concat_axis=self._state.merged_concat_axis, + ) + + +class Concat(ez.Unit): + """Concatenate two AxisArray streams along an axis. + + Pairs messages by arrival order (FIFO). No time-alignment. + """ + + SETTINGS = ConcatSettings + + INPUT_SIGNAL_A = ez.InputStream(AxisArray) + INPUT_SIGNAL_B = ez.InputStream(AxisArray) + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + async def initialize(self) -> None: + self.processor = ConcatProcessor(self.SETTINGS) + + @ez.subscriber(INPUT_SIGNAL_A) + async def on_a(self, msg: AxisArray) -> None: + self.processor.push_a(msg) + + @ez.subscriber(INPUT_SIGNAL_B) + async def on_b(self, msg: AxisArray) -> None: + self.processor.push_b(msg) + + @ez.publisher(OUTPUT_SIGNAL) + async def output(self) -> typing.AsyncGenerator: + while True: + yield self.OUTPUT_SIGNAL, await self.processor.__acall__() diff --git a/src/ezmsg/sigproc/merge.py b/src/ezmsg/sigproc/merge.py index 70afe935..27c86a9f 100644 --- a/src/ezmsg/sigproc/merge.py +++ b/src/ezmsg/sigproc/merge.py @@ -1,20 +1,17 @@ -"""Time-aligned merge of two AxisArray streams along a non-time axis.""" +"""Time-aligned merge of two AxisArray streams along a non-time axis. -from __future__ import annotations +``Merge`` is an :class:`ez.Collection` that composes +:class:`~ezmsg.sigproc.align.AlignAlongAxis` (time-alignment) with +:class:`~ezmsg.sigproc.concat.Concat` (axis-aware concatenation). +""" -import math -import typing +from __future__ import annotations import ezmsg.core as ez -import numpy as np -from array_api_compat import get_namespace -from ezmsg.baseproc.protocols import processor_state -from ezmsg.baseproc.stateful import BaseStatefulTransformer -from ezmsg.baseproc.units import BaseProcessorUnit -from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis -from ezmsg.util.messages.util import replace +from ezmsg.util.messages.axisarray import AxisArray -from .util.axisarray_buffer import HybridAxisArrayBuffer +from .align import AlignAlongAxis, AlignAlongAxisProcessor, AlignAlongAxisSettings +from .concat import Concat, ConcatProcessor, ConcatSettings class MergeSettings(ez.Settings): @@ -36,303 +33,73 @@ class MergeSettings(ez.Settings): label_b: str = "_b" """Suffix appended to signal B labels when relabel_axis is True.""" + assert_identical_shared_axes: bool = False + """If True, raise ValueError when shared CoordinateAxis .data arrays differ.""" + new_key: str | None = None """Output AxisArray key. If None, uses the key from signal A.""" -@processor_state -class MergeState: - # Common state - gain: float | None = None - align_axis: str | None = None - aligned: bool = False - merged_concat_axis: CoordinateAxis | None = None - - # A state - buf_a: HybridAxisArrayBuffer | None = None - concat_axis_a: CoordinateAxis | None = None - a_concat_dim: int | None = None - a_other_dims: tuple[int, ...] | None = None - - # B state - buf_b: HybridAxisArrayBuffer | None = None - concat_axis_b: CoordinateAxis | None = None - b_concat_dim: int | None = None - b_other_dims: tuple[int, ...] | None = None - - -class MergeProcessor(BaseStatefulTransformer[MergeSettings, AxisArray, AxisArray | None, MergeState]): - """Processor that time-aligns two AxisArray streams and concatenates them. +class MergeProcessor: + """Convenience processor that composes alignment + concatenation. - Input A flows through the standard ``__call__`` / ``_process`` path, - getting automatic ``_hash_message`` / ``_reset_state`` handling from - :class:`BaseStatefulTransformer`. Input B flows through :meth:`push_b`, - which independently tracks its own structure. - - Invalidation rules: - - - Gain mismatch (either input vs stored common gain) → full reset. - - Concat-axis dimensionality change → per-input buffer reset + - alignment and merged-axis cache invalidation. - - Non-align/non-concat axis shape change → per-input buffer reset + - alignment invalidation. + Preserves the same call interface as the previous monolithic processor + so that existing code using ``proc(msg_a)`` / ``proc.push_b(msg_b)`` + continues to work unchanged. """ - # -- Structural extraction helpers --------------------------------------- - - def _extract_gain(self, message: AxisArray) -> float | None: - """Extract the align-axis gain from a message.""" - align_name = self.settings.align_axis or message.dims[0] - ax = message.axes.get(align_name) - if ax is not None and hasattr(ax, "gain"): - return ax.gain - if ax is not None and hasattr(ax, "data") and len(ax.data) > 1: - return float(ax.data[-1] - ax.data[0]) / (len(ax.data) - 1) - return None - - # -- Reset helpers ------------------------------------------------------- - - def _full_reset(self, align_axis: str) -> None: - """Reset all state — both inputs and common merge state.""" - self._state.align_axis = align_axis - self._state.buf_a = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) - self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) - self._state.gain = None - self._state.aligned = False - self._state.concat_axis_a = None - self._state.concat_axis_b = None - self._state.merged_concat_axis = None - self._state.a_concat_dim = None - self._state.a_other_dims = None - self._state.b_concat_dim = None - self._state.b_other_dims = None - - def _reset_a_state(self) -> None: - """Reset input-A buffer and concat-axis cache.""" - self._state.buf_a = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=self._state.align_axis) - self._state.concat_axis_a = None - - def _reset_b_state(self) -> None: - """Reset input-B buffer and concat-axis cache.""" - self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=self._state.align_axis) - self._state.concat_axis_b = None - - # -- BaseStatefulTransformer interface ------------------------------------ - - def _hash_message(self, message: AxisArray) -> int: - """Hash the align-axis gain only. - - Gain changes trigger a full reset via ``_reset_state``. Concat-axis - and non-merge dimension changes are handled as partial resets inside - ``_process`` and ``push_b``. - """ - return hash(self._extract_gain(message)) - - def _reset_state(self, message: AxisArray) -> None: - """Full reset — called by the base class when gain changes.""" - align_axis = self.settings.align_axis or message.dims[0] - self._full_reset(align_axis) - - def _process(self, message: AxisArray) -> AxisArray | None: - """Process input A: detect structural changes, buffer, try merge.""" - # Detect per-input structural changes. - align_idx = message.dims.index(self._state.align_axis) - concat_idx = message.dims.index(self.settings.axis) if self.settings.axis in message.dims else None - concat_dim = message.data.shape[concat_idx] if concat_idx is not None else None - other_dims = tuple(s for i, s in enumerate(message.data.shape) if i != align_idx and i != concat_idx) - - if self._state.a_concat_dim is not None and concat_dim != self._state.a_concat_dim: - self._reset_a_state() - self._state.aligned = False - self._state.merged_concat_axis = None - elif self._state.a_other_dims is not None and other_dims != self._state.a_other_dims: - self._reset_a_state() - self._state.aligned = False - - self._state.a_concat_dim = concat_dim - self._state.a_other_dims = other_dims - - self._state.buf_a.write(message) - if self._state.gain is None: - self._state.gain = self._state.buf_a.axis_gain - self._update_concat_axis(message, "a") - return self._try_merge() - - # -- Input B entry point ------------------------------------------------ - - def push_b(self, message: AxisArray) -> AxisArray | None: - """Process input B: check gain, detect structural changes, buffer, try merge.""" - align_axis = self.settings.align_axis or message.dims[0] - - # Gain compatibility check. - b_gain = self._extract_gain(message) - if self._state.gain is not None and not math.isclose(b_gain, self._state.gain): - self._full_reset(align_axis) - # Set the base-class hash so the next compatible A goes straight - # to _process instead of triggering another full reset. - self._hash = self._hash_message(message) - - # Lazy-create buf_b if B arrives before A. - if self._state.buf_b is None: - if self._state.align_axis is None: - self._state.align_axis = align_axis - self._state.buf_b = HybridAxisArrayBuffer(duration=self.settings.buffer_dur, axis=align_axis) - - # Detect per-input structural changes. - align_idx = message.dims.index(align_axis) - concat_idx = message.dims.index(self.settings.axis) if self.settings.axis in message.dims else None - concat_dim = message.data.shape[concat_idx] if concat_idx is not None else None - other_dims = tuple(s for i, s in enumerate(message.data.shape) if i != align_idx and i != concat_idx) - - if self._state.b_concat_dim is not None and concat_dim != self._state.b_concat_dim: - self._reset_b_state() - self._state.aligned = False - self._state.merged_concat_axis = None - elif self._state.b_other_dims is not None and other_dims != self._state.b_other_dims: - self._reset_b_state() - self._state.aligned = False - - self._state.b_concat_dim = concat_dim - self._state.b_other_dims = other_dims - - self._state.buf_b.write(message) - if self._state.gain is None: - self._state.gain = self._state.buf_b.axis_gain - self._update_concat_axis(message, "b") - return self._try_merge() - - # -- Concat-axis caching ------------------------------------------------ - - def _update_concat_axis(self, message: AxisArray, which: str) -> None: - """Track each input's concat-axis labels; invalidate cache on change.""" - concat_dim = self.settings.axis - if concat_dim not in message.axes: - return - ax = message.axes[concat_dim] - if not hasattr(ax, "data"): - return - - if which == "a": - if self._state.concat_axis_a is None or not np.array_equal(self._state.concat_axis_a.data, ax.data): - self._state.concat_axis_a = ax - self._state.merged_concat_axis = None - else: - if self._state.concat_axis_b is None or not np.array_equal(self._state.concat_axis_b.data, ax.data): - self._state.concat_axis_b = ax - self._state.merged_concat_axis = None - - def _build_merged_concat_axis(self) -> CoordinateAxis | None: - """Build the merged CoordinateAxis from the two cached per-input axes.""" - if self._state.concat_axis_a is None or self._state.concat_axis_b is None: - return None - if self.settings.relabel_axis: - labels_a = np.array([str(lbl) + self.settings.label_a for lbl in self._state.concat_axis_a.data]) - labels_b = np.array([str(lbl) + self.settings.label_b for lbl in self._state.concat_axis_b.data]) - else: - labels_a = self._state.concat_axis_a.data - labels_b = self._state.concat_axis_b.data - return CoordinateAxis( - data=np.concatenate([labels_a, labels_b]), - dims=self._state.concat_axis_a.dims, - unit=self._state.concat_axis_a.unit, + def __init__(self, settings: MergeSettings): + self.settings = settings + self._align = AlignAlongAxisProcessor( + settings=AlignAlongAxisSettings( + axis=settings.align_axis or "time", + buffer_dur=settings.buffer_dur, + ) + ) + self._concat = ConcatProcessor( + settings=ConcatSettings( + axis=settings.axis, + relabel_axis=settings.relabel_axis, + label_a=settings.label_a, + label_b=settings.label_b, + assert_identical_shared_axes=settings.assert_identical_shared_axes, + new_key=settings.new_key, + ) ) - # -- Core merge logic --------------------------------------------------- - - def _try_merge(self) -> AxisArray | None: - """Align and read from both buffers, returning the merged result. - - Initial alignment is performed once. After the first successful - merge the two streams are assumed to share a common clock and - never drop samples, so we simply read - ``min(available_a, available_b)`` on every subsequent call. - """ - if self._state.buf_a is None or self._state.buf_b is None: - return None - if self._state.buf_a.is_empty() or self._state.buf_b.is_empty(): - return None - - gain = self._state.gain - - # --- Initial alignment (runs only until the first successful merge) --- - if not self._state.aligned: - first_a = self._state.buf_a.axis_first_value - final_a = self._state.buf_a.axis_final_value - first_b = self._state.buf_b.axis_first_value - final_b = self._state.buf_b.axis_final_value - - overlap_start = max(first_a, first_b) - overlap_end = min(final_a, final_b) - - if overlap_end < overlap_start - gain / 2: - if final_a < first_b: - self._state.buf_a.seek(self._state.buf_a.available()) - elif final_b < first_a: - self._state.buf_b.seek(self._state.buf_b.available()) - return None - - if first_a < overlap_start - gain / 2: - self._state.buf_a.seek(int(round((overlap_start - first_a) / gain))) - if first_b < overlap_start - gain / 2: - self._state.buf_b.seek(int(round((overlap_start - first_b) / gain))) - - # --- Read aligned samples --- - n_read = min(self._state.buf_a.available(), self._state.buf_b.available()) - if n_read <= 0: - return None - - aa_a = self._state.buf_a.read(n_read) - aa_b = self._state.buf_b.read(n_read) - if aa_a is None or aa_b is None: - return None - - if not self._state.aligned: - axis_a = aa_a.axes.get(self._state.align_axis) - axis_b = aa_b.axes.get(self._state.align_axis) - if axis_a is not None and axis_b is not None: - off_a = axis_a.value(0) if hasattr(axis_a, "value") else None - off_b = axis_b.value(0) if hasattr(axis_b, "value") else None - if off_a is not None and off_b is not None: - if not np.isclose(off_a, off_b, atol=abs(gain) * 1e-6): - raise RuntimeError( - f"Offset mismatch after alignment: " f"off_a={off_a}, off_b={off_b}, gain={gain}" - ) - self._state.aligned = True - - return self._concat(aa_a, aa_b) - - def _concat(self, a: AxisArray, b: AxisArray) -> AxisArray: - """Concatenate *a* and *b* along the configured merge axis.""" - merge_dim = self.settings.axis + @property + def align_state(self): + """Expose alignment state for introspection / tests.""" + return self._align.state - # If the merge dim doesn't exist in an input, add it as a trailing axis. - if merge_dim not in a.dims: - xp = get_namespace(a.data) - a = replace(a, data=xp.expand_dims(a.data, axis=-1), dims=[*a.dims, merge_dim]) - if merge_dim not in b.dims: - xp = get_namespace(b.data) - b = replace(b, data=xp.expand_dims(b.data, axis=-1), dims=[*b.dims, merge_dim]) + @property + def concat_state(self): + """Expose concatenation state for introspection / tests.""" + return self._concat.state - # Use the cached merged axis (rebuilt lazily when labels change). - if self._state.merged_concat_axis is None: - self._state.merged_concat_axis = self._build_merged_concat_axis() + def __call__(self, msg_a: AxisArray) -> AxisArray | None: + pair = self._align(msg_a) + if pair is not None: + return self._concat._concat(*pair) + return None - key = self.settings.new_key if self.settings.new_key is not None else a.key - result = AxisArray.concatenate(a, b, dim=merge_dim, axis=self._state.merged_concat_axis) - if key != result.key: - result = replace(result, key=key) - return result + async def __acall__(self, msg_a: AxisArray) -> AxisArray | None: + pair = await self._align.__acall__(msg_a) + if pair is not None: + return self._concat._concat(*pair) + return None + def push_b(self, msg_b: AxisArray) -> AxisArray | None: + pair = self._align.push_b(msg_b) + if pair is not None: + return self._concat._concat(*pair) + return None -class Merge(BaseProcessorUnit[MergeSettings]): - """Merge two AxisArray streams by time-aligning and concatenating along a non-time axis. - Input A routes through the processor's ``__acall__`` (triggering - hash-based reset when the stream structure changes). Input B - routes through ``push_b`` which independently tracks its own structure. +class Merge(ez.Collection): + """Merge two AxisArray streams by time-aligning and concatenating. - Inherits ``INPUT_SETTINGS`` and ``on_settings`` → ``create_processor`` - from :class:`BaseProcessorUnit`. + Composes :class:`AlignAlongAxis` → :class:`Concat`. """ SETTINGS = MergeSettings @@ -341,19 +108,32 @@ class Merge(BaseProcessorUnit[MergeSettings]): INPUT_SIGNAL_B = ez.InputStream(AxisArray) OUTPUT_SIGNAL = ez.OutputStream(AxisArray) - def create_processor(self) -> None: - self.processor = MergeProcessor(settings=self.SETTINGS) + ALIGN = AlignAlongAxis() + CONCAT = Concat() - @ez.subscriber(INPUT_SIGNAL_A, zero_copy=True) - @ez.publisher(OUTPUT_SIGNAL) - async def on_a(self, msg: AxisArray) -> typing.AsyncGenerator: - result = await self.processor.__acall__(msg) - if result is not None: - yield self.OUTPUT_SIGNAL, result + def configure(self) -> None: + self.ALIGN.apply_settings( + AlignAlongAxisSettings( + axis=self.SETTINGS.align_axis or "time", + buffer_dur=self.SETTINGS.buffer_dur, + ) + ) + self.CONCAT.apply_settings( + ConcatSettings( + axis=self.SETTINGS.axis, + relabel_axis=self.SETTINGS.relabel_axis, + label_a=self.SETTINGS.label_a, + label_b=self.SETTINGS.label_b, + assert_identical_shared_axes=self.SETTINGS.assert_identical_shared_axes, + new_key=self.SETTINGS.new_key, + ) + ) - @ez.subscriber(INPUT_SIGNAL_B, zero_copy=True) - @ez.publisher(OUTPUT_SIGNAL) - async def on_b(self, msg: AxisArray) -> typing.AsyncGenerator: - result = self.processor.push_b(msg) - if result is not None: - yield self.OUTPUT_SIGNAL, result + def network(self) -> ez.NetworkDefinition: + return ( + (self.INPUT_SIGNAL_A, self.ALIGN.INPUT_SIGNAL_A), + (self.INPUT_SIGNAL_B, self.ALIGN.INPUT_SIGNAL_B), + (self.ALIGN.OUTPUT_SIGNAL_A, self.CONCAT.INPUT_SIGNAL_A), + (self.ALIGN.OUTPUT_SIGNAL_B, self.CONCAT.INPUT_SIGNAL_B), + (self.CONCAT.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), + ) diff --git a/src/ezmsg/sigproc/sampler.py b/src/ezmsg/sigproc/sampler.py index 9cd304db..17fad08d 100644 --- a/src/ezmsg/sigproc/sampler.py +++ b/src/ezmsg/sigproc/sampler.py @@ -190,14 +190,13 @@ class Sampler(BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, Sampler SETTINGS = SamplerSettings INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage) - OUTPUT_SIGNAL = ez.OutputStream(AxisArray) @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) + @ez.publisher(BaseTransformerUnit.OUTPUT_SIGNAL) @profile_subpub(trace_oldest=False) async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator: try: diff --git a/tests/unit/test_align.py b/tests/unit/test_align.py new file mode 100644 index 00000000..a1f6f3ca --- /dev/null +++ b/tests/unit/test_align.py @@ -0,0 +1,194 @@ +"""Unit tests for ezmsg.sigproc.align module.""" + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis +from frozendict import frozendict + +from ezmsg.sigproc.align import AlignAlongAxisProcessor, AlignAlongAxisSettings + + +def _make_msg( + data: np.ndarray, + fs: float = 100.0, + offset: float = 0.0, + ch_labels: list[str] | None = None, + key: str = "", +) -> AxisArray: + """Helper to build a (time, ch) AxisArray.""" + n_ch = data.shape[1] if data.ndim > 1 else 1 + if data.ndim == 1: + data = data[:, None] + if ch_labels is None: + ch_labels = [f"Ch{i}" for i in range(n_ch)] + ch_axis = CoordinateAxis(data=np.array(ch_labels), dims=["ch"], unit="label") + time_axis = AxisArray.TimeAxis(fs=fs, offset=offset) + return AxisArray( + data, + dims=["time", "ch"], + axes=frozendict({"time": time_axis, "ch": ch_axis}), + key=key, + ) + + +class TestAlignedPairs: + """Both streams perfectly aligned — verify paired output.""" + + def test_basic(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + fs = 100.0 + n = 10 + + data_a = np.arange(n * 2, dtype=float).reshape(n, 2) + data_b = np.arange(n * 3, dtype=float).reshape(n, 3) + 100 + + msg_a = _make_msg(data_a, fs=fs, offset=0.0, ch_labels=["A0", "A1"]) + msg_b = _make_msg(data_b, fs=fs, offset=0.0, ch_labels=["B0", "B1", "B2"]) + + pair_from_a = proc(msg_a) + assert pair_from_a is None # Only one side present. + + pair = proc.push_b(msg_b) + assert pair is not None + aa_a, aa_b = pair + assert aa_a.data.shape == (n, 2) + assert aa_b.data.shape == (n, 3) + np.testing.assert_array_equal(aa_a.data, data_a) + np.testing.assert_array_equal(aa_b.data, data_b) + + def test_multiple_chunks(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + fs = 100.0 + chunk = 5 + + for i in range(4): + offset = i * chunk / fs + data_a = np.ones((chunk, 2)) * i + data_b = np.ones((chunk, 3)) * (i + 10) + msg_a = _make_msg(data_a, fs=fs, offset=offset, ch_labels=["A0", "A1"]) + msg_b = _make_msg(data_b, fs=fs, offset=offset, ch_labels=["B0", "B1", "B2"]) + proc(msg_a) + pair = proc.push_b(msg_b) + assert pair is not None + assert pair[0].data.shape == (chunk, 2) + assert pair[1].data.shape == (chunk, 3) + + +class TestStaggeredArrival: + def test_a_arrives_first(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + fs = 100.0 + chunk = 10 + + for i in range(3): + offset = i * chunk / fs + msg_a = _make_msg(np.ones((chunk, 2)) * i, fs=fs, offset=offset) + assert proc(msg_a) is None + + b_offset = 1 * chunk / fs + msg_b = _make_msg(np.ones((chunk, 3)) * 99, fs=fs, offset=b_offset) + pair = proc.push_b(msg_b) + assert pair is not None + aa_a, aa_b = pair + assert aa_a.data.shape[0] == chunk + assert aa_b.data.shape[0] == chunk + np.testing.assert_allclose(aa_a.data, 1.0) + np.testing.assert_allclose(aa_b.data, 99.0) + + +class TestFloatingPointOffset: + def test_tiny_epsilon(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + fs = 1000.0 + n = 20 + offset = 1.0 + eps = 1e-14 + + msg_a = _make_msg(np.ones((n, 2)), fs=fs, offset=offset) + msg_b = _make_msg(np.ones((n, 2)) * 2, fs=fs, offset=offset + eps) + proc(msg_a) + pair = proc.push_b(msg_b) + assert pair is not None + assert pair[0].data.shape == (n, 2) + + +class TestGainMismatch: + def test_b_different_gain_resets(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + n = 10 + + msg_a = _make_msg(np.ones((n, 2)), fs=100.0) + msg_b_bad = _make_msg(np.ones((n, 3)), fs=200.0) + proc(msg_a) + assert proc.push_b(msg_b_bad) is None + + msg_a2 = _make_msg(np.ones((n, 2)) * 2, fs=200.0, offset=0.0) + pair = proc(msg_a2) + assert pair is not None + assert pair[0].data.shape == (n, 2) + assert pair[1].data.shape == (n, 3) + + def test_a_gain_change(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + n = 5 + + msg_a1 = _make_msg(np.ones((n, 2)), fs=100.0) + msg_b1 = _make_msg(np.ones((n, 3)), fs=100.0) + proc(msg_a1) + assert proc.push_b(msg_b1) is not None + + # A switches to 200 Hz — triggers reset. + msg_a2 = _make_msg(np.ones((n, 2)) * 2, fs=200.0, offset=0.5) + assert proc(msg_a2) is None + + msg_b2 = _make_msg(np.ones((n, 3)) * 3, fs=200.0, offset=0.5) + pair = proc.push_b(msg_b2) + assert pair is not None + + def test_b_gain_change(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + n = 5 + + msg_a1 = _make_msg(np.ones((n, 2)), fs=100.0) + msg_b1 = _make_msg(np.ones((n, 3)), fs=100.0) + proc(msg_a1) + assert proc.push_b(msg_b1) is not None + + msg_b2 = _make_msg(np.ones((n, 3)) * 2, fs=200.0, offset=0.5) + assert proc.push_b(msg_b2) is None + + msg_a2 = _make_msg(np.ones((n, 2)) * 3, fs=200.0, offset=0.5) + pair = proc(msg_a2) + assert pair is not None + + +class TestShapeChange: + def test_a_shape_change_resets_a(self): + settings = AlignAlongAxisSettings(axis="time") + proc = AlignAlongAxisProcessor(settings) + fs = 100.0 + n = 5 + + msg_a1 = _make_msg(np.ones((n, 2)), fs=fs, offset=0.0) + msg_b1 = _make_msg(np.ones((n, 3)), fs=fs, offset=0.0) + proc(msg_a1) + pair = proc.push_b(msg_b1) + assert pair is not None + + # Push B at next offset. + offset_1 = n / fs + msg_b2 = _make_msg(np.ones((n, 3)) * 2, fs=fs, offset=offset_1) + assert proc.push_b(msg_b2) is None + + # A changes to 4 channels at same offset — shape change, re-alignment. + msg_a2 = _make_msg(np.ones((n, 4)) * 3, fs=fs, offset=offset_1, ch_labels=["A0", "A1", "A2", "A3"]) + pair = proc(msg_a2) + assert pair is not None + assert pair[0].data.shape == (n, 4) + assert pair[1].data.shape == (n, 3) diff --git a/tests/unit/test_concat.py b/tests/unit/test_concat.py new file mode 100644 index 00000000..8a5d5175 --- /dev/null +++ b/tests/unit/test_concat.py @@ -0,0 +1,339 @@ +"""Unit tests for ezmsg.sigproc.concat module.""" + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis +from frozendict import frozendict + +from ezmsg.sigproc.concat import ( + ConcatProcessor, + ConcatSettings, + _build_merged_coordinate_axis, + _validate_shared_axes, +) + + +def _make_msg( + data: np.ndarray, + fs: float = 100.0, + offset: float = 0.0, + ch_labels: list[str] | None = None, + ch_axis: CoordinateAxis | None = None, + key: str = "", +) -> AxisArray: + """Helper to build a (time, ch) AxisArray.""" + n_ch = data.shape[1] if data.ndim > 1 else 1 + if data.ndim == 1: + data = data[:, None] + if ch_axis is None: + if ch_labels is None: + ch_labels = [f"Ch{i}" for i in range(n_ch)] + ch_axis = CoordinateAxis(data=np.array(ch_labels), dims=["ch"], unit="label") + time_axis = AxisArray.TimeAxis(fs=fs, offset=offset) + return AxisArray( + data, + dims=["time", "ch"], + axes=frozendict({"time": time_axis, "ch": ch_axis}), + key=key, + ) + + +# --------------------------------------------------------------------------- +# Shared helper tests +# --------------------------------------------------------------------------- + + +class TestBuildMergedCoordinateAxis: + def test_simple_labels_relabel(self): + ax_a = CoordinateAxis(data=np.array(["X", "Y"]), dims=["ch"]) + ax_b = CoordinateAxis(data=np.array(["X", "Y", "Z"]), dims=["ch"]) + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=True, label_a="_a", label_b="_b") + assert list(merged.data) == ["X_a", "Y_a", "X_b", "Y_b", "Z_b"] + + def test_simple_labels_no_relabel(self): + ax_a = CoordinateAxis(data=np.array(["A0", "A1"]), dims=["ch"]) + ax_b = CoordinateAxis(data=np.array(["B0", "B1"]), dims=["ch"]) + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=False, label_a="", label_b="") + assert list(merged.data) == ["A0", "A1", "B0", "B1"] + + def test_struct_preserves_fields(self): + dt = np.dtype([("label", "U8"), ("x", "f4"), ("y", "f4")]) + data_a = np.array([("ch0", 1.0, 2.0), ("ch1", 3.0, 4.0)], dtype=dt) + data_b = np.array([("ch2", 5.0, 6.0)], dtype=dt) + ax_a = CoordinateAxis(data=data_a, dims=["ch"]) + ax_b = CoordinateAxis(data=data_b, dims=["ch"]) + + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=True, label_a="_L", label_b="_R") + assert len(merged.data) == 3 + assert merged.data[0]["label"] == "ch0_L" + assert merged.data[2]["label"] == "ch2_R" + # Non-label fields preserved. + np.testing.assert_allclose(merged.data[0]["x"], 1.0) + np.testing.assert_allclose(merged.data[2]["y"], 6.0) + + def test_struct_no_relabel(self): + dt = np.dtype([("label", "U8"), ("x", "f4")]) + data_a = np.array([("ch0", 1.0)], dtype=dt) + data_b = np.array([("ch1", 2.0)], dtype=dt) + ax_a = CoordinateAxis(data=data_a, dims=["ch"]) + ax_b = CoordinateAxis(data=data_b, dims=["ch"]) + + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=False, label_a="", label_b="") + assert merged.data[0]["label"] == "ch0" + assert merged.data[1]["label"] == "ch1" + np.testing.assert_allclose(merged.data[0]["x"], 1.0) + + def test_struct_union_fields(self): + """Inputs with different struct fields produce a union dtype.""" + dt_a = np.dtype([("label", "U8"), ("x", "f4")]) + dt_b = np.dtype([("label", "U8"), ("bank", "i4")]) + data_a = np.array([("ch0", 1.0)], dtype=dt_a) + data_b = np.array([("ch1", 42)], dtype=dt_b) + ax_a = CoordinateAxis(data=data_a, dims=["ch"]) + ax_b = CoordinateAxis(data=data_b, dims=["ch"]) + + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=False, label_a="", label_b="") + assert "label" in merged.data.dtype.names + assert "x" in merged.data.dtype.names + assert "bank" in merged.data.dtype.names + np.testing.assert_allclose(merged.data[0]["x"], 1.0) + assert merged.data[1]["bank"] == 42 + # Missing field in B gets default (0.0 for float). + np.testing.assert_allclose(merged.data[1]["x"], 0.0) + + def test_struct_incompatible_dtypes_raises(self): + dt_a = np.dtype([("label", "U8"), ("x", "f4")]) + dt_b = np.dtype([("label", "U8"), ("x", "i4")]) + data_a = np.array([("ch0", 1.0)], dtype=dt_a) + data_b = np.array([("ch1", 2)], dtype=dt_b) + ax_a = CoordinateAxis(data=data_a, dims=["ch"]) + ax_b = CoordinateAxis(data=data_b, dims=["ch"]) + + with pytest.raises(ValueError, match="Incompatible dtypes"): + _build_merged_coordinate_axis(ax_a, ax_b, relabel=False, label_a="", label_b="") + + def test_struct_label_created_when_absent(self): + """If relabel=True and no 'label' field exists, one is added.""" + dt = np.dtype([("x", "f4"), ("y", "f4")]) + data_a = np.array([(1.0, 2.0)], dtype=dt) + data_b = np.array([(3.0, 4.0)], dtype=dt) + ax_a = CoordinateAxis(data=data_a, dims=["ch"]) + ax_b = CoordinateAxis(data=data_b, dims=["ch"]) + + merged = _build_merged_coordinate_axis(ax_a, ax_b, relabel=True, label_a="_L", label_b="_R") + assert "label" in merged.data.dtype.names + assert merged.data[0]["label"] == "0_L" + assert merged.data[1]["label"] == "0_R" + + +class TestValidateSharedAxes: + def test_identical_axes_pass(self): + ch_ax = CoordinateAxis(data=np.array(["C0", "C1"]), dims=["ch"]) + a = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ch_ax}) + b = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ch_ax}) + # Should not raise. + _validate_shared_axes(a, b, concat_dim="feature", align_dim="time", assert_flag=True) + + def test_different_axes_raise(self): + ax_a = CoordinateAxis(data=np.array(["C0", "C1"]), dims=["ch"]) + ax_b = CoordinateAxis(data=np.array(["X0", "X1"]), dims=["ch"]) + a = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ax_a}) + b = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ax_b}) + with pytest.raises(ValueError, match="Shared axis 'ch'"): + _validate_shared_axes(a, b, concat_dim="feature", align_dim="time", assert_flag=True) + + def test_flag_false_skips(self): + ax_a = CoordinateAxis(data=np.array(["C0", "C1"]), dims=["ch"]) + ax_b = CoordinateAxis(data=np.array(["X0", "X1"]), dims=["ch"]) + a = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ax_a}) + b = AxisArray(np.ones((5, 2)), dims=["time", "ch"], axes={"ch": ax_b}) + # Should not raise when flag is False. + _validate_shared_axes(a, b, concat_dim="feature", align_dim="time", assert_flag=False) + + +# --------------------------------------------------------------------------- +# ConcatProcessor tests +# --------------------------------------------------------------------------- + + +class TestBasicConcat: + def test_concat_along_ch(self): + settings = ConcatSettings(axis="ch") + proc = ConcatProcessor(settings) + n, fs = 10, 100.0 + + data_a = np.arange(n * 2, dtype=float).reshape(n, 2) + data_b = np.arange(n * 3, dtype=float).reshape(n, 3) + 100 + msg_a = _make_msg(data_a, fs=fs, ch_labels=["A0", "A1"]) + msg_b = _make_msg(data_b, fs=fs, ch_labels=["B0", "B1", "B2"]) + + proc.push_a(msg_a) + proc.push_b(msg_b) + + import asyncio + + result = asyncio.get_event_loop().run_until_complete(proc.__acall__()) + assert result.data.shape == (n, 5) + np.testing.assert_array_equal(result.data[:, :2], data_a) + np.testing.assert_array_equal(result.data[:, 2:], data_b) + + def test_concat_direct(self): + """Test _concat directly (synchronous path used by MergeProcessor).""" + settings = ConcatSettings(axis="ch", relabel_axis=False) + proc = ConcatProcessor(settings) + n = 5 + + data_a = np.ones((n, 2)) + data_b = np.ones((n, 3)) * 2 + msg_a = _make_msg(data_a, ch_labels=["A0", "A1"]) + msg_b = _make_msg(data_b, ch_labels=["B0", "B1", "B2"]) + + result = proc._concat(msg_a, msg_b) + assert result.data.shape == (n, 5) + labels = list(result.axes["ch"].data) + assert labels == ["A0", "A1", "B0", "B1", "B2"] + + +class TestConcatRelabel: + def test_default_suffix(self): + settings = ConcatSettings(axis="ch", label_a="_left", label_b="_right") + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["X", "Y"]) + msg_b = _make_msg(np.ones((n, 3)), ch_labels=["X", "Y", "Z"]) + + result = proc._concat(msg_a, msg_b) + labels = list(result.axes["ch"].data) + assert labels == ["X_left", "Y_left", "X_right", "Y_right", "Z_right"] + + def test_no_relabel(self): + settings = ConcatSettings(axis="ch", relabel_axis=False) + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["A0", "A1"]) + msg_b = _make_msg(np.ones((n, 3)), ch_labels=["B0", "B1", "B2"]) + + result = proc._concat(msg_a, msg_b) + labels = list(result.axes["ch"].data) + assert labels == ["A0", "A1", "B0", "B1", "B2"] + + +class TestNewAxisConcat: + def test_new_feature_axis(self): + settings = ConcatSettings(axis="feature", relabel_axis=False) + proc = ConcatProcessor(settings) + n = 5 + + data_a = np.ones((n, 3)) + data_b = np.ones((n, 3)) * 2 + msg_a = _make_msg(data_a, ch_labels=["C0", "C1", "C2"]) + msg_b = _make_msg(data_b, ch_labels=["C0", "C1", "C2"]) + + result = proc._concat(msg_a, msg_b) + assert result.data.shape == (n, 3, 2) + np.testing.assert_array_equal(result.data[:, :, 0], 1.0) + np.testing.assert_array_equal(result.data[:, :, 1], 2.0) + assert "feature" in result.dims + + def test_new_axis_dim_mismatch_raises(self): + settings = ConcatSettings(axis="feature", relabel_axis=False) + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 3)), ch_labels=["C0", "C1", "C2"]) + msg_b = _make_msg(np.ones((n, 4)), ch_labels=["C0", "C1", "C2", "C3"]) + + with pytest.raises(ValueError, match="Cannot concatenate along new axis"): + proc._concat(msg_a, msg_b) + + +class TestAssertIdenticalSharedAxes: + def test_identical_passes(self): + settings = ConcatSettings(axis="feature", assert_identical_shared_axes=True) + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["C0", "C1"]) + msg_b = _make_msg(np.ones((n, 2)) * 2, ch_labels=["C0", "C1"]) + + result = proc._concat(msg_a, msg_b) + assert result.data.shape == (n, 2, 2) + + def test_different_raises(self): + settings = ConcatSettings(axis="feature", assert_identical_shared_axes=True) + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["C0", "C1"]) + msg_b = _make_msg(np.ones((n, 2)) * 2, ch_labels=["X0", "X1"]) + + with pytest.raises(ValueError, match="Shared axis 'ch'"): + proc._concat(msg_a, msg_b) + + def test_flag_false_allows_different(self): + settings = ConcatSettings(axis="feature", assert_identical_shared_axes=False) + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["C0", "C1"]) + msg_b = _make_msg(np.ones((n, 2)) * 2, ch_labels=["X0", "X1"]) + + result = proc._concat(msg_a, msg_b) + assert result.data.shape == (n, 2, 2) + + +class TestStructAwareConcat: + def test_struct_axis_concat(self): + settings = ConcatSettings(axis="ch", label_a="_L", label_b="_R") + proc = ConcatProcessor(settings) + n = 5 + + dt = np.dtype([("label", "U8"), ("x", "f4"), ("y", "f4")]) + ch_a = np.array([("ch0", 1.0, 2.0), ("ch1", 3.0, 4.0)], dtype=dt) + ch_b = np.array([("ch2", 5.0, 6.0)], dtype=dt) + ax_a = CoordinateAxis(data=ch_a, dims=["ch"]) + ax_b = CoordinateAxis(data=ch_b, dims=["ch"]) + + msg_a = _make_msg(np.ones((n, 2)), ch_axis=ax_a) + msg_b = _make_msg(np.ones((n, 1)) * 2, ch_axis=ax_b) + + result = proc._concat(msg_a, msg_b) + assert result.data.shape == (n, 3) + ch_out = result.axes["ch"].data + assert ch_out[0]["label"] == "ch0_L" + assert ch_out[2]["label"] == "ch2_R" + np.testing.assert_allclose(ch_out[0]["x"], 1.0) + np.testing.assert_allclose(ch_out[2]["y"], 6.0) + + +class TestCachedAxes: + def test_cache_reused(self): + settings = ConcatSettings(axis="ch") + proc = ConcatProcessor(settings) + n = 5 + + msg_a = _make_msg(np.ones((n, 2)), ch_labels=["A0", "A1"]) + msg_b = _make_msg(np.ones((n, 2)), ch_labels=["B0", "B1"]) + + proc._concat(msg_a, msg_b) + first_cache = proc.state.cached_axes + + proc._concat(msg_a, msg_b) + assert proc.state.cached_axes is first_cache # Same dict object. + + def test_cache_invalidated_on_shape_change(self): + settings = ConcatSettings(axis="ch", relabel_axis=False) + proc = ConcatProcessor(settings) + n = 5 + + msg_a1 = _make_msg(np.ones((n, 2)), ch_labels=["A0", "A1"]) + msg_b1 = _make_msg(np.ones((n, 3)), ch_labels=["B0", "B1", "B2"]) + proc._concat(msg_a1, msg_b1) + first_cache = proc.state.cached_axes + + # Change A to 4 channels. + msg_a2 = _make_msg(np.ones((n, 4)), ch_labels=["A0", "A1", "A2", "A3"]) + proc._concat(msg_a2, msg_b1) + assert proc.state.cached_axes is not first_cache diff --git a/tests/unit/test_merge.py b/tests/unit/test_merge.py index 9da1887b..8af871a3 100644 --- a/tests/unit/test_merge.py +++ b/tests/unit/test_merge.py @@ -319,7 +319,7 @@ def test_cached_axis_reused(self): proc.push_b(msg_b) # The cached axis should have been built once and reused. - assert proc.state.merged_concat_axis is not None + assert proc.concat_state.merged_concat_axis is not None class TestNoRelabel: