Resolve the operating dimension from chunk_dim across the processors

pyproject.toml

@@ -14,8 +14,8 @@ dependencies = [
     # left to the transitive requirement: uv only enables pre-releases for a
     # package named with a pre-release marker in *this* file.
     "ezmsg>=3.10.0b2",
-    "ezmsg-baseproc>=1.12.0",
-    "ezmsg-sigproc>=3.0.0",
+    "ezmsg-baseproc>=1.13.0",
+    "ezmsg-sigproc>=3.8.1",  # Window/Resample must resolve chunk_dim from axis=None
     "numpy",
     "scipy",
     "array-api-compat",

src/ezmsg/learn/collection/sample_adapt_regressor.py

@@ -17,6 +17,8 @@
     BaseTransformerUnit,
     SampleTriggerMessage,
     processor_state,
+    suppress_axis_deprecation,
+    warn_axis_deprecated,
 )
 from ezmsg.sigproc.resample import ResampleSettings, ResampleUnit
 from ezmsg.sigproc.window import Window, WindowSettings
@@ -171,8 +173,14 @@ class SampleAdaptRegressorSettings(ez.Settings):
     generic ``ch0..chN``."""
 
     # Resampling settings
-    resample_axis: str = "time"
-    """Axis to resample along."""
+    resample_axis: str | None = None
+    """.. deprecated:: 1.6
+        Scheduled for removal in 2.0. Resampling buffers along the dimension
+        messages accumulate along, which now comes from
+        :attr:`~ezmsg.util.messages.axisarray.AxisArray.chunk_dim`."""
+
+    def __post_init__(self) -> None:
+        warn_axis_deprecated(self, "resample_axis", package="ezmsg-learn", removal="2.0")
 
     resample_buffer_duration: float = 2.0
     """Duration of the buffer for resampling in seconds."""
@@ -274,7 +282,8 @@ def configure(self) -> None:
             if use_window:
                 self.WINDOW.apply_settings(
                     WindowSettings(
-                        axis="time",
+                        # No `axis`: Window follows the stream's chunk_dim, which
+                        # is what "time" was standing in for.
                         newaxis="win",
                         window_dur=self.SETTINGS.decode_window_dur,
                         window_shift=self.SETTINGS.decode_window_shift,
@@ -292,14 +301,17 @@ def configure(self) -> None:
                     )
                 )
             if use_sample_path:
-                self.RESAMPLE.apply_settings(
-                    ResampleSettings(
-                        axis=self.SETTINGS.resample_axis,
-                        max_chunk_delay=float("inf"),
-                        fill_value="extrapolate",
-                        buffer_duration=self.SETTINGS.resample_buffer_duration,
+                # Forwarding our own already-warned setting; warning again would
+                # name a sigproc class for something set on ours.
+                with suppress_axis_deprecation():
+                    self.RESAMPLE.apply_settings(
+                        ResampleSettings(
+                            axis=self.SETTINGS.resample_axis,
+                            max_chunk_delay=float("inf"),
+                            fill_value="extrapolate",
+                            buffer_duration=self.SETTINGS.resample_buffer_duration,
+                        )
                     )
-                )
                 self.SEQSEQSAMPLER.apply_settings(
                     SeqSeqSamplerSettings(
                         max_buffer_dur=self.SETTINGS.sampler_max_buffer_dur,

src/ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -17,6 +17,8 @@
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
+    resolve_chunk_dim,
+    warn_axis_deprecated,
 )
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
@@ -30,7 +32,25 @@
 
 
 class AdaptiveDecompSettings(ez.Settings):
-    axis: str = "!time"
+    axis: str | None = None
+    """Which dimension to decompose.
+
+    ``None`` (default) decomposes every dimension except the one messages
+    accumulate along, iterating over that one. Naming a dimension (e.g.
+    ``"ch"``) decomposes it and iterates over the chunk dimension instead.
+
+    .. deprecated:: 1.6
+        The ``"!time"`` spelling -- "iterate over time" -- is scheduled for
+        removal in 2.0. It hardcodes what
+        :attr:`~ezmsg.util.messages.axisarray.AxisArray.chunk_dim` now
+        answers; leave this unset for the same behaviour."""
+
+    def __post_init__(self) -> None:
+        # Only the "!" spelling is going away. Naming a target axis is a real
+        # choice about what to decompose, and stays.
+        if self.axis is not None and self.axis.startswith("!"):
+            warn_axis_deprecated(self, package="ezmsg-learn", removal="2.0")
+
     n_components: int = 2
 
 
@@ -76,26 +96,33 @@ def _create_estimator(self) -> EstimatorType:
         return estimator_klass(**estimator_settings)
 
     def _calculate_axis_groups(self, message: AxisArray):
-        if self.settings.axis.startswith("!"):
+        axis = self.settings.axis
+        if axis is None:
+            # Iterate over the dimension messages accumulate along and collapse
+            # every other one -- what "!time" spelled, with the dimension read
+            # off the stream instead of assumed.
+            iter_axis = resolve_chunk_dim(message, self.STREAMING_DIMS)
+            it_ax_ix = message.get_axis_idx(iter_axis)
+            targ_axes = message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :]
+            off_targ_axes = []
+        elif axis.startswith("!"):
             # Iterate over the !axis and collapse all other axes
-            iter_axis = self.settings.axis[1:]
+            iter_axis = axis[1:]
             it_ax_ix = message.get_axis_idx(iter_axis)
             targ_axes = message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :]
             off_targ_axes = []
         else:
             # Do PCA on the parameterized axis
-            targ_axes = [self.settings.axis]
-            # Iterate over streaming axis
-            iter_axis = "win" if "win" in message.dims else "time"
-            if iter_axis == self.settings.axis:
-                raise ValueError(
-                    f"Iterating axis ({iter_axis}) cannot be the same as the target axis ({self.settings.axis})"
-                )
+            targ_axes = [axis]
+            # Iterate over the dimension messages accumulate along. This was a
+            # hand-rolled `"win" if "win" in dims else "time"` guess, which is
+            # exactly what chunk_dim exists to answer.
+            iter_axis = resolve_chunk_dim(message, self.STREAMING_DIMS)
+            if iter_axis == axis:
+                raise ValueError(f"Iterating axis ({iter_axis}) cannot be the same as the target axis ({axis})")
             it_ax_ix = message.get_axis_idx(iter_axis)
             # Remaining axes are to be treated independently
-            off_targ_axes = [
-                _ for _ in (message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :]) if _ != self.settings.axis
-            ]
+            off_targ_axes = [_ for _ in (message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :]) if _ != axis]
         self._state.axis_groups = iter_axis, targ_axes, off_targ_axes
 
     def _reset_state(self, message: AxisArray) -> None:

src/ezmsg/learn/dim_reduce/incremental_decomp.py

@@ -6,6 +6,7 @@
     BaseStatefulProcessor,
     BaseTransformerUnit,
     CompositeProcessor,
+    warn_axis_deprecated,
 )
 from ezmsg.sigproc.window import WindowTransformer
 from ezmsg.util.messages.axisarray import AxisArray, replace
@@ -19,7 +20,18 @@
 
 
 class IncrementalDecompSettings(ez.Settings):
-    axis: str = "!time"
+    axis: str | None = None
+    """Which dimension to decompose. See
+    :obj:`~ezmsg.learn.dim_reduce.adaptive_decomp.AdaptiveDecompSettings.axis`.
+
+    .. deprecated:: 1.6
+        The ``"!time"`` spelling is scheduled for removal in 2.0; leave this
+        unset for the same behaviour."""
+
+    def __post_init__(self) -> None:
+        if self.axis is not None and self.axis.startswith("!"):
+            warn_axis_deprecated(self, package="ezmsg-learn", removal="2.0")
+
     n_components: int = 2
     update_interval: float = 0.0
     method: str = "pca"
@@ -72,8 +84,12 @@ def _initialize_processors(
 
         # Create windowing processor if update_interval is specified
         if settings.update_interval > 0:
-            # TODO: This `iter_axis` is likely incorrect.
-            iter_axis = settings.axis[1:] if settings.axis.startswith("!") else "time"
+            # Only the "!axis" spelling names the iteration dimension outright.
+            # Otherwise leave it to Window, which resolves chunk_dim from the
+            # message -- this is the "likely incorrect" hardcoded "time" that
+            # used to be here, and there is no message to resolve from at this
+            # point anyway.
+            iter_axis = settings.axis[1:] if (settings.axis or "").startswith("!") else None
             windowing = WindowTransformer(
                 axis=iter_axis,
                 window_dur=settings.update_interval,
@@ -98,7 +114,14 @@ def _partial_fit_windowed(self, train_msg: AxisArray) -> None:
             axis_idx = train_msg.get_axis_idx("win")
             win_axis = train_msg.axes["win"]
             offsets = win_axis.value(np.asarray(range(train_msg.data.shape[axis_idx])))
-            for ix, _msg in enumerate(train_msg.iter_over_axis("win")):
+            # Slicing "win" away leaves each sub-message no longer a chunk along
+            # it. Newer ezmsg clears the declaration for us, but say what these
+            # slices *are* chunks along rather than leaving them undeclared:
+            # successive windows advance along the within-window axis, which is
+            # what the offset fix-up below re-anchors. Clearing it first keeps
+            # this working on ezmsg versions that do not.
+            unbundled = replace(train_msg, chunk_dim=None)
+            for ix, _msg in enumerate(unbundled.iter_over_axis("win")):
                 _msg = replace(
                     _msg,
                     axes={
@@ -108,6 +131,7 @@ def _partial_fit_windowed(self, train_msg: AxisArray) -> None:
                             offset=_msg.axes["time"].offset + offsets[ix],
                         ),
                     },
+                    chunk_dim="time",
                 )
                 self._procs["decomp"].partial_fit(_msg)
 

src/ezmsg/learn/process/flatten.py

@@ -26,6 +26,7 @@
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
+    resolve_chunk_dim,
 )
 from ezmsg.sigproc.flatten import (
     FlattenSettings as SigprocFlattenSettings,
@@ -135,7 +136,7 @@ class FlattenTransformer(BaseStatefulTransformer[FlattenSettings, AxisArray, Axi
     """
 
     def _reset_state(self, message: AxisArray) -> None:
-        preserve_axis = self.settings.preserve_axis or message.dims[0]
+        preserve_axis = self.settings.preserve_axis or resolve_chunk_dim(message, self.STREAMING_DIMS)
         sample_axis = self.settings.sample_axis or preserve_axis
         feature_axis = self.settings.feature_axis
 

src/ezmsg/learn/process/sgd.py

@@ -6,6 +6,7 @@
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
+    resolve_chunk_dim,
 )
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
@@ -35,6 +36,11 @@ class SGDDecoderState:
 
 
 class SGDDecoderTransformer(BaseAdaptiveTransformer[SGDDecoderSettings, AxisArray, ClassifierMessage, SGDDecoderState]):
+    STREAMING_DIMS = ("win", "time")
+    """This decoder is fed windows, so a producer that declares no ``chunk_dim``
+    is accumulating along ``win`` rather than ``time``. The base default would
+    guess ``time`` and flatten the windows into the feature vector."""
+
     """
     SGD-based online classifier.
 
@@ -90,19 +96,22 @@ def _process(self, message: AxisArray) -> ClassifierMessage | None:
         if np.any(np.isnan(message.data)):
             return None
         try:
-            X = message.data.reshape((message.data.shape[0], -1))
+            chunk = resolve_chunk_dim(message, self.STREAMING_DIMS)
+            chunk_idx = message.get_axis_idx(chunk)
+            data = message.data if chunk_idx == 0 else np.moveaxis(message.data, chunk_idx, 0)
+            X = data.reshape((data.shape[0], -1))
             result = self._state.model._predict_proba_lr(X)
         except NotFittedError:
             return None
         out_axes = {}
-        if message.dims[0] in message.axes:
-            out_axes[message.dims[0]] = replace(
-                message.axes[message.dims[0]],
-                offset=message.axes[message.dims[0]].offset,
+        if chunk in message.axes:
+            out_axes[chunk] = replace(
+                message.axes[chunk],
+                offset=message.axes[chunk].offset,
             )
         return ClassifierMessage(
             data=result,
-            dims=message.dims[:1] + ["labels"],
+            dims=[chunk, "labels"],
             axes=out_axes,
             labels=list(self._state.model.class_weight.keys()),
             key=message.key,

src/ezmsg/learn/process/slda.py

@@ -16,6 +16,8 @@
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
+    resolve_configured_chunk_dim,
+    warn_axis_deprecated,
 )
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
@@ -31,17 +33,30 @@
 
 class SLDASettings(ez.Settings):
     settings_path: str
-    axis: str = "time"
+
+    axis: str | None = None
+    """.. deprecated:: 1.6
+        Scheduled for removal in 2.0. The samples this classifies accumulate
+        along one dimension, and the cached output template is keyed to it;
+        that dimension now comes from
+        :attr:`~ezmsg.util.messages.axisarray.AxisArray.chunk_dim`."""
+
+    def __post_init__(self) -> None:
+        warn_axis_deprecated(self, package="ezmsg-learn", removal="2.0")
 
 
 @processor_state
 class SLDAState:
+    axis: str = ""
+    """The resolved chunk dimension, fixed at reset so every later use agrees."""
+
     lda: LDA
     out_template: typing.Optional[ClassifierMessage] = None
 
 
 class SLDATransformer(BaseStatefulTransformer[SLDASettings, AxisArray, ClassifierMessage, SLDAState]):
     def _reset_state(self, message: AxisArray) -> None:
+        self.state.axis = resolve_configured_chunk_dim(self, message, self.settings.axis, legacy_default="time")
         if self.settings.settings_path[-4:] == ".mat":
             # Expects a very specific format from a specific project. Not for general use.
             import scipy.io as sio
@@ -77,9 +92,9 @@ def _reset_state(self, message: AxisArray) -> None:
         zero_shape = (0, len(out_labels))
         self.state.out_template = ClassifierMessage(
             data=np.zeros(zero_shape, dtype=message.data.dtype),
-            dims=[self.settings.axis, "classes"],
+            dims=[self.state.axis, "classes"],
             axes={
-                self.settings.axis: message.axes[self.settings.axis],
+                self.state.axis: message.axes[self.state.axis],
                 "classes": with_fingerprint(AxisArray.CoordinateAxis(data=np.array(out_labels), dims=["classes"])),
             },
             labels=out_labels,
@@ -88,7 +103,7 @@ def _reset_state(self, message: AxisArray) -> None:
 
     def _process(self, message: AxisArray) -> ClassifierMessage:
         xp = get_namespace(message.data)
-        samp_ax_idx = message.dims.index(self.settings.axis)
+        samp_ax_idx = message.dims.index(self.state.axis)
 
         # Move sample axis to front
         perm = (samp_ax_idx,) + tuple(i for i in range(message.data.ndim) if i != samp_ax_idx)
@@ -111,16 +126,16 @@ def _process(self, message: AxisArray) -> ClassifierMessage:
                 X_np = X_np.reshape(X_np.shape[0], -1)
                 pred_probas = self.state.lda.predict_proba(X_np)
 
-            update_ax = self.state.out_template.axes[self.settings.axis]
-            update_ax.offset = message.axes[self.settings.axis].offset
+            update_ax = self.state.out_template.axes[self.state.axis]
+            update_ax.offset = message.axes[self.state.axis].offset
 
             return replace(
                 self.state.out_template,
                 data=pred_probas,
                 axes={
                     **self.state.out_template.axes,
                     # `replace` will copy the minimal set of fields
-                    self.settings.axis: replace(update_ax, offset=update_ax.offset),
+                    self.state.axis: replace(update_ax, offset=update_ax.offset),
                 },
             )
         else: