diff --git a/pyproject.toml b/pyproject.toml index d07c19d..b34f522 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/ezmsg/learn/collection/sample_adapt_regressor.py b/src/ezmsg/learn/collection/sample_adapt_regressor.py index 3db795b..405ef82 100644 --- a/src/ezmsg/learn/collection/sample_adapt_regressor.py +++ b/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, diff --git a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py b/src/ezmsg/learn/dim_reduce/adaptive_decomp.py index 74edfb8..f416c29 100644 --- a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py +++ b/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: diff --git a/src/ezmsg/learn/dim_reduce/incremental_decomp.py b/src/ezmsg/learn/dim_reduce/incremental_decomp.py index 25906bf..3b87c49 100644 --- a/src/ezmsg/learn/dim_reduce/incremental_decomp.py +++ b/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) diff --git a/src/ezmsg/learn/process/flatten.py b/src/ezmsg/learn/process/flatten.py index 128d0ed..6e247b2 100644 --- a/src/ezmsg/learn/process/flatten.py +++ b/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 diff --git a/src/ezmsg/learn/process/sgd.py b/src/ezmsg/learn/process/sgd.py index 3d5ab67..bf4474b 100644 --- a/src/ezmsg/learn/process/sgd.py +++ b/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, diff --git a/src/ezmsg/learn/process/slda.py b/src/ezmsg/learn/process/slda.py index fb346f6..3840bc4 100644 --- a/src/ezmsg/learn/process/slda.py +++ b/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,8 +126,8 @@ 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, @@ -120,7 +135,7 @@ def _process(self, message: AxisArray) -> ClassifierMessage: 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: diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index b51d9f9..593e6d1 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -53,6 +53,7 @@ BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit, processor_state, + resolve_feature_dim, ) from ezmsg.baseproc.protocols import SettingsType, StateType from ezmsg.sigproc.affinetransform import ( @@ -149,7 +150,7 @@ class SelfSupervisedRegressionTransformer( # -- message hash / state management ------------------------------------ def _reset_state(self, message: AxisArray) -> None: - axis = self.settings.axis or message.dims[-1] + axis = self.settings.axis or resolve_feature_dim(message) axis_idx = message.get_axis_idx(axis) n_channels = message.data.shape[axis_idx] @@ -358,7 +359,7 @@ def partial_fit(self, message: AxisArray) -> None: # type: ignore[override] self._reset_state(message) self._hash = msg_hash - axis = self.settings.axis or message.dims[-1] + axis = self.settings.axis or resolve_feature_dim(message) axis_idx = message.get_axis_idx(axis) data = message.data @@ -510,7 +511,7 @@ def _on_weights_updated(self) -> None: # -- transform ----------------------------------------------------------- def _process(self, message: AxisArray) -> AxisArray: - axis = self.settings.axis or message.dims[-1] + axis = self.settings.axis or resolve_feature_dim(message) if message.data.shape[message.get_axis_idx(axis)] == 0: # No channels (e.g. a fully sliced-out hub): nothing to rereference. # Pass the 0-channel message through unchanged -- building an affine diff --git a/tests/dim_reduce/test_incremental_decomp.py b/tests/dim_reduce/test_incremental_decomp.py index 6c9f794..ef0ef9c 100644 --- a/tests/dim_reduce/test_incremental_decomp.py +++ b/tests/dim_reduce/test_incremental_decomp.py @@ -135,7 +135,10 @@ def test_initialization_pca(self, pca_test_data, update_interval): if update_interval > 0: assert "windowing" in transformer._procs win = transformer._procs["windowing"] - assert win.settings.axis == "time" + # Unset on purpose: Window resolves the stream's chunk_dim, which is + # what the old hardcoded "time" was standing in for (and got wrong + # downstream of another windowing stage). + assert win.settings.axis is None assert win.settings.window_dur == update_interval assert win.settings.window_shift == update_interval assert win.settings.zero_pad_until == "none"