diff --git a/src/ezmsg/sigproc/aggregate.py b/src/ezmsg/sigproc/aggregate.py index c3687a27..8ade3882 100644 --- a/src/ezmsg/sigproc/aggregate.py +++ b/src/ezmsg/sigproc/aggregate.py @@ -125,7 +125,7 @@ def _reset_state(self, message: AxisArray) -> None: ax_idx = message.get_axis_idx(axis) if hasattr(target_axis, "data"): - self._state.ax_vec = target_axis.data + self._state.ax_vec = np.array(target_axis.data) else: self._state.ax_vec = target_axis.value(np.arange(message.data.shape[ax_idx])) @@ -136,9 +136,9 @@ def _reset_state(self, message: AxisArray) -> None: slices.append(slice(int(inds[0]), int(inds[-1]) + 1)) if hasattr(target_axis, "data"): if self._state.ax_vec.dtype.type is np.str_: - sl_dat = f"{self._state.ax_vec[start]} - {self._state.ax_vec[stop]}" + sl_dat = f"{self._state.ax_vec[inds[0]]} - {self._state.ax_vec[inds[-1]]}" else: - ax_dat.append(np.mean(self._state.ax_vec[inds])) + sl_dat = np.mean(self._state.ax_vec[inds]) else: sl_dat = target_axis.value(np.mean(inds)) ax_dat.append(sl_dat) diff --git a/src/ezmsg/sigproc/concat.py b/src/ezmsg/sigproc/concat.py index 10ef6749..1d6125d9 100644 --- a/src/ezmsg/sigproc/concat.py +++ b/src/ezmsg/sigproc/concat.py @@ -5,6 +5,7 @@ import asyncio import logging import typing +from copy import deepcopy from dataclasses import dataclass, field import ezmsg.core as ez @@ -332,7 +333,12 @@ def _build_cached_axes( align_dim: str | None, merged_concat_axis: CoordinateAxis | None, ) -> dict[str, AxisBase]: - """Build the output axes dict (everything except the alignment axis).""" + """Build an owned output-axis cache (everything except the alignment axis). + + Input axes may be views into an ezmsg transport buffer whose lifetime ends + after the current subscriber callback. Axes kept across calls must therefore + be copied into processor-owned memory. + """ axes: dict[str, AxisBase] = {} for name, ax in a.axes.items(): if name == align_dim: @@ -340,7 +346,7 @@ def _build_cached_axes( if name == concat_dim and merged_concat_axis is not None: axes[name] = merged_concat_axis else: - axes[name] = ax + axes[name] = deepcopy(ax) if concat_dim not in axes and merged_concat_axis is not None: axes[concat_dim] = merged_concat_axis return axes diff --git a/src/ezmsg/sigproc/filterbank.py b/src/ezmsg/sigproc/filterbank.py index de638f61..6162cde7 100644 --- a/src/ezmsg/sigproc/filterbank.py +++ b/src/ezmsg/sigproc/filterbank.py @@ -3,6 +3,7 @@ import functools import math import typing +from copy import deepcopy import ezmsg.core as ez import numpy as np @@ -137,12 +138,12 @@ def _reset_state(self, message: AxisArray) -> None: tail_shape = in_shape + (len(kernels), self._state.overlap) self._state.tail = np.zeros(tail_shape, dtype="complex" if b_complex else "float") - # Prepare output template -- kernels axis immediately before the target axis + # Prepare output template -- kernels axis immediately before the target axis. dummy_shape = in_shape + (len(kernels), 0) self._state.template = AxisArray( data=np.zeros(dummy_shape, dtype="complex" if b_complex else "float"), dims=message.dims[:targ_ax_ix] + message.dims[targ_ax_ix + 1 :] + [self.settings.new_axis, axis], - axes=message.axes.copy(), + axes={k: deepcopy(v) for k, v in message.axes.items()}, key=message.key, ) diff --git a/src/ezmsg/sigproc/flatten.py b/src/ezmsg/sigproc/flatten.py index ad52fc1e..4fff796c 100644 --- a/src/ezmsg/sigproc/flatten.py +++ b/src/ezmsg/sigproc/flatten.py @@ -114,8 +114,9 @@ class FlattenState: # None means input dims already match (preserve, *flatten, *rest). perm: tuple[int, ...] | None = None - # Final shape after permute_dims + reshape: (n_preserve, n_flat, *rest_shape). - target_shape: tuple[int, ...] = () + # Shape after the preserve axis: (n_flat, *rest_shape). The preserve-axis + # length is supplied from each live message because chunk sizes may vary. + target_inner_shape: tuple[int, ...] = () # Precomputed merged-axis CoordinateAxis (structured array). output_axis_obj: CoordinateAxis | None = None @@ -240,7 +241,9 @@ def _expand(arr: np.ndarray, axis_idx: int) -> np.ndarray: class FlattenTransformer(BaseStatefulTransformer[FlattenSettings, AxisArray, AxisArray, FlattenState]): def _hash_message(self, message: AxisArray) -> int: - return hash((tuple(message.dims), tuple(message.data.shape))) + preserve_axis = self.settings.preserve_axis or message.dims[0] + non_preserve_shape = tuple(size for dim, size in zip(message.dims, message.data.shape) if dim != preserve_axis) + return hash((tuple(message.dims), non_preserve_shape)) def _reset_state(self, message: AxisArray) -> None: preserve_axis = self.settings.preserve_axis or message.dims[0] @@ -271,11 +274,10 @@ def _reset_state(self, message: AxisArray) -> None: perm = None permuted_shape = tuple(message.data.shape) - n_preserve = permuted_shape[0] flatten_sizes = permuted_shape[1 : 1 + len(flatten_axes)] n_flat = int(math.prod(flatten_sizes)) if flatten_sizes else 1 rest_shape = permuted_shape[1 + len(flatten_axes) :] - target_shape = (n_preserve, n_flat, *rest_shape) + target_inner_shape = (n_flat, *rest_shape) output_axis_obj = _build_merged_axis( message, @@ -292,7 +294,7 @@ def _reset_state(self, message: AxisArray) -> None: st.flatten_axes = flatten_axes st.rest_axes = rest_axes st.perm = perm - st.target_shape = target_shape + st.target_inner_shape = target_inner_shape st.output_axis_obj = output_axis_obj st.output_dims = (sample_axis, output_axis, *rest_axes) @@ -304,7 +306,7 @@ def _process(self, message: AxisArray) -> AxisArray: data = xp.permute_dims(message.data, st.perm) else: data = message.data - data = xp.reshape(data, st.target_shape) + data = xp.reshape(data, (data.shape[0], *st.target_inner_shape)) # Carry the live preserve axis through (its gain/offset/data may # advance per message). Rename to sample_axis on the output if diff --git a/src/ezmsg/sigproc/util/axisarray_buffer.py b/src/ezmsg/sigproc/util/axisarray_buffer.py index d331e3a9..7ff60a21 100644 --- a/src/ezmsg/sigproc/util/axisarray_buffer.py +++ b/src/ezmsg/sigproc/util/axisarray_buffer.py @@ -3,6 +3,7 @@ import math import typing import warnings +from copy import deepcopy import numpy as np from array_api_compat import get_namespace @@ -289,11 +290,11 @@ def axis_final_value(self) -> float | None: def _initialize(self, first_msg: AxisArray) -> None: # Create a template message that has everything except the data are length 0 - # and the target axis is missing. + # and the target axis is missing. self._template_msg = replace( first_msg, data=first_msg.data[:0], - axes={k: v for k, v in first_msg.axes.items() if k != self._axis}, + axes={k: deepcopy(v) for k, v in first_msg.axes.items() if k != self._axis}, ) in_axis = first_msg.axes[self._axis] diff --git a/src/ezmsg/sigproc/wavelets.py b/src/ezmsg/sigproc/wavelets.py index b9c18fab..71b87385 100644 --- a/src/ezmsg/sigproc/wavelets.py +++ b/src/ezmsg/sigproc/wavelets.py @@ -1,6 +1,7 @@ """Continuous wavelet transform (CWT) for streaming time-frequency analysis.""" import typing +from copy import deepcopy import ezmsg.core as ez import numpy as np @@ -110,7 +111,7 @@ def _reset_state(self, message: AxisArray) -> None: axis=self.settings.axis, ) - # Create output template + # Create output template. ax_idx = message.get_axis_idx(self.settings.axis) in_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] freqs = pywt.scale2frequency(wavelet, scales, precision) / message.axes[self.settings.axis].gain @@ -119,7 +120,7 @@ def _reset_state(self, message: AxisArray) -> None: np.zeros(dummy_shape, dtype=dt_cplx if wavelet.complex_cwt else dt_data), dims=message.dims[:ax_idx] + message.dims[ax_idx + 1 :] + ["freq", self.settings.axis], axes={ - **message.axes, + **{k: deepcopy(v) for k, v in message.axes.items()}, "freq": AxisArray.CoordinateAxis(unit="Hz", data=freqs, dims=["freq"]), }, key=message.key, diff --git a/tests/unit/buffer/test_axisarray_buffer.py b/tests/unit/buffer/test_axisarray_buffer.py index a22e5a33..57021656 100644 --- a/tests/unit/buffer/test_axisarray_buffer.py +++ b/tests/unit/buffer/test_axisarray_buffer.py @@ -431,6 +431,39 @@ def test_deferred_initialization_coordinate(coordinate_axis_message): assert buf._axis_buffer.capacity == 100 +def test_template_axes_are_cached_and_owned(): + """The output template reuses owned non-time axes across messages.""" + shared_axes = { + "ch": CoordinateAxis(data=np.array([1, 2]), dims=["ch"]), + "feature": CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]), + } + messages = [ + AxisArray( + data=np.full((n_samples, 2, 2), fill, dtype=np.float64), + dims=["time", "ch", "feature"], + axes={ + "time": LinearAxis(gain=0.01, offset=offset, unit="s"), + **shared_axes, + }, + ) + for n_samples, fill, offset in ((3, 1.0, 0.00), (9, 2.0, 0.03), (5, 3.0, 0.12)) + ] + for axis_name in shared_axes: + assert all(message.axes[axis_name] is messages[0].axes[axis_name] for message in messages) + + buffer = HybridAxisArrayBuffer(duration=1.0, axis="time", update_strategy="immediate") + outputs = [] + for message in messages: + buffer.write(message) + output = buffer.read() + assert output is not None + outputs.append(output) + + for axis_name in shared_axes: + assert all(output.axes[axis_name] is outputs[0].axes[axis_name] for output in outputs) + assert all(output.axes[axis_name] is not message.axes[axis_name] for output, message in zip(outputs, messages)) + + def test_add_and_get_linear(linear_axis_message): buf = HybridAxisArrayBuffer(duration=1.0, update_strategy="immediate") msg1 = linear_axis_message(samples=10, fs=100.0, offset=0.0) diff --git a/tests/unit/test_aggregate.py b/tests/unit/test_aggregate.py index e1332bf2..33c754e1 100644 --- a/tests/unit/test_aggregate.py +++ b/tests/unit/test_aggregate.py @@ -532,3 +532,63 @@ def run(): import mlx.core as mx assert isinstance(result.data, mx.array) + + +def test_ranged_aggregate_coordinate_axis_vector_owned(): + """The cached coordinate vector and output axis must not alias input axis memory.""" + freq_data = np.arange(8, dtype=float) + shared_freq_axis = AxisArray.CoordinateAxis(data=freq_data, dims=["freq"]) + messages = [ + AxisArray( + data=np.arange(4 * 2 * 8, dtype=float).reshape(4, 2, 8), + dims=["time", "ch", "freq"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0, offset=offset), + "freq": shared_freq_axis, + }, + key="test_owned_ax_vec", + ) + for offset in (0.0, 0.04) + ] + + xformer = RangedAggregateTransformer( + RangedAggregateSettings( + axis="freq", + bands=[(1.0, 3.0), (4.0, 6.0)], + operation=AggregationFunction.TRAPEZOID, + ) + ) + out_first = xformer(messages[0]) + assert not np.shares_memory(xformer._state.ax_vec, freq_data) + assert np.array_equal(out_first.axes["freq"].data, np.array([2.0, 5.0])) + + # Simulate transport-buffer reuse: the first message's axis memory is overwritten. + freq_data[:] = -1.0 + out_second = xformer(messages[1]) + + assert out_second.axes["freq"] is out_first.axes["freq"] + assert np.array_equal(out_second.data, out_first.data) + + +def test_ranged_aggregate_coordinate_axis_string_labels(): + """String coordinate axes get 'first - last' labels from the matched entries.""" + label_axis = AxisArray.CoordinateAxis(data=np.array(["a", "b", "c", "d"]), dims=["ch"]) + msg = AxisArray( + data=np.arange(3 * 4, dtype=float).reshape(3, 4), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0, offset=0.0), + "ch": label_axis, + }, + key="test_str_labels", + ) + xformer = RangedAggregateTransformer( + RangedAggregateSettings( + axis="ch", + bands=[("a", "b"), ("c", "d")], + operation=AggregationFunction.MEAN, + ) + ) + out = xformer(msg) + assert list(out.axes["ch"].data) == ["a - b", "c - d"] + assert np.array_equal(out.data, np.stack([msg.data[:, :2].mean(axis=1), msg.data[:, 2:].mean(axis=1)], axis=1)) diff --git a/tests/unit/test_concat.py b/tests/unit/test_concat.py index 18569e5d..bfa7b18b 100644 --- a/tests/unit/test_concat.py +++ b/tests/unit/test_concat.py @@ -517,6 +517,44 @@ def test_cache_invalidated_on_shape_change(self): proc._concat(msg_a2, msg_b1) assert proc.state.cached_axes is not first_cache + def test_cached_axes_are_reused_and_owned(self): + """Concat reuses owned non-time output axes across messages.""" + shared_feature_axis = CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]) + shared_ch_a = CoordinateAxis(data=np.array(["A"]), dims=["ch"]) + shared_ch_b = CoordinateAxis(data=np.array(["B"]), dims=["ch"]) + + def _messages(fill: float, ch_axis: CoordinateAxis) -> list[AxisArray]: + return [ + AxisArray( + data=np.full((3, 1, 2), fill + index, dtype=np.float64), + dims=["time", "ch", "feature"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0, offset=index * 0.03), + "ch": ch_axis, + "feature": shared_feature_axis, + }, + ) + for index in range(3) + ] + + messages_a = _messages(1.0, shared_ch_a) + messages_b = _messages(2.0, shared_ch_b) + for messages in (messages_a, messages_b): + for axis_name in ("ch", "feature"): + assert all(message.axes[axis_name] is messages[0].axes[axis_name] for message in messages) + + proc = ConcatProcessor(ConcatSettings(axis="ch", relabel_axis=False, align_axis="time")) + outputs = [proc._concat(message_a, message_b) for message_a, message_b in zip(messages_a, messages_b)] + + for axis_name in ("ch", "feature"): + assert all(output.axes[axis_name] is outputs[0].axes[axis_name] for output in outputs) + assert all( + output.axes[axis_name] is not message.axes[axis_name] for output, message in zip(outputs, messages_a) + ) + assert all( + output.axes[axis_name] is not message.axes[axis_name] for output, message in zip(outputs, messages_b) + ) + # --------------------------------------------------------------------------- # Attrs merge + promotion diff --git a/tests/unit/test_filterbank.py b/tests/unit/test_filterbank.py index 09b0de7c..771cc67b 100644 --- a/tests/unit/test_filterbank.py +++ b/tests/unit/test_filterbank.py @@ -104,3 +104,30 @@ def test_filterbank(mode: str, kernel_type: str): axes[1, ch_ix].imshow(tmp[ch_ix], aspect="auto", origin="lower") axes[2, ch_ix].imshow(np.abs(tmp[ch_ix]), aspect="auto", origin="lower") plt.show() + + +def test_template_axes_cached_and_owned(): + """Filterbank reuses processor-owned non-target axes across messages.""" + fs = 100.0 + n_time = 20 + shared_ch_axis = AxisArray.CoordinateAxis(data=np.array(["Ch0", "Ch1"]), dims=["ch"]) + messages = [ + AxisArray( + data=np.arange(2 * n_time, dtype=float).reshape(2, n_time) + i, + dims=["ch", "time"], + axes={ + "ch": shared_ch_axis, + "time": AxisArray.TimeAxis(fs=fs, offset=i * n_time / fs), + }, + key="test_filterbank_axes", + ) + for i in range(3) + ] + assert all(message.axes["ch"] is shared_ch_axis for message in messages) + + kernels = [np.array([0.5, 0.5]), np.array([1.0, -1.0])] + proc = FilterbankTransformer(settings=FilterbankSettings(kernels=kernels, mode=FilterbankMode.CONV, axis="time")) + outputs = [proc(message) for message in messages] + + assert all(output.axes["ch"] is outputs[0].axes["ch"] for output in outputs) + assert all(output.axes["ch"] is not message.axes["ch"] for output, message in zip(outputs, messages)) diff --git a/tests/unit/test_flatten.py b/tests/unit/test_flatten.py index 9ee22390..55a03895 100644 --- a/tests/unit/test_flatten.py +++ b/tests/unit/test_flatten.py @@ -126,6 +126,33 @@ def test_preserve_axis_kept_through(self): time_ax = out.axes["time"] assert time_ax.gain == pytest.approx(1.0 / 200.0) + def test_output_axis_cached_across_preserve_axis_sizes_and_owned(self): + """Variable time chunks reuse one processor-owned merged axis.""" + shared_axes = { + "ch": CoordinateAxis(data=np.array(["c1", "c2", "c3"]), dims=["ch"]), + "feature": CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]), + } + messages = [ + AxisArray( + data=np.arange(n_time * 3 * 2, dtype=float).reshape(n_time, 3, 2), + dims=["time", "ch", "feature"], + axes={ + "time": AxisArray.TimeAxis(fs=50.0, offset=offset), + **shared_axes, + }, + ) + for n_time, offset in ((3, 0.00), (9, 0.06), (5, 0.24)) + ] + for axis_name in shared_axes: + assert all(message.axes[axis_name] is messages[0].axes[axis_name] for message in messages) + + proc = FlattenTransformer(FlattenSettings()) + outputs = [proc(message) for message in messages] + + assert [output.data.shape for output in outputs] == [(3, 6), (9, 6), (5, 6)] + assert all(output.axes["ch"] is outputs[0].axes["ch"] for output in outputs) + assert all(output.axes["ch"] is not message.axes["ch"] for output, message in zip(outputs, messages)) + class TestStructFieldPassthrough: def test_struct_axis_propagates_all_fields(self): diff --git a/tests/unit/test_wavelets.py b/tests/unit/test_wavelets.py index 2b6f1e5c..8d080218 100644 --- a/tests/unit/test_wavelets.py +++ b/tests/unit/test_wavelets.py @@ -188,3 +188,29 @@ def test_cwt_complex_wavelet_min_phase_informative_error(): ) with pytest.raises(ValueError, match="[Cc]omplex.*min[_ ]phase|min_phase.*complex"): proc(msg) + + +def test_template_axes_cached_and_owned(): + """CWT reuses processor-owned non-target axes across messages.""" + fs = 200.0 + n_time = 32 + shared_ch_axis = AxisArray.CoordinateAxis(data=np.array(["Ch0", "Ch1"]), dims=["ch"]) + messages = [ + AxisArray( + data=np.arange(2 * n_time, dtype=float).reshape(2, n_time) + i, + dims=["ch", "time"], + axes={ + "ch": shared_ch_axis, + "time": AxisArray.TimeAxis(fs=fs, offset=i * n_time / fs), + }, + key="test_cwt_axes", + ) + for i in range(3) + ] + assert all(message.axes["ch"] is shared_ch_axis for message in messages) + + proc = CWTTransformer(CWTSettings(frequencies=np.array([10.0, 20.0, 40.0]), wavelet="morl", axis="time")) + outputs = [proc(message) for message in messages] + + assert all(output.axes["ch"] is outputs[0].axes["ch"] for output in outputs) + assert all(output.axes["ch"] is not message.axes["ch"] for output, message in zip(outputs, messages))