Own axes retained in processor state; Flatten variable-chunk hash; RangedAggregate coordinate-axis fixes

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)

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,15 +333,20 @@ 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:
             continue
         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

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,
         )
 

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

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]

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,

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)

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

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

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

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):