ezmsg-learn 1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/learn/process/linear_regressor.py

@@ -1,17 +1,16 @@
 from dataclasses import field
 
-import numpy as np
-from sklearn.linear_model._base import LinearModel
 import ezmsg.core as ez
-from ezmsg.sigproc.base import (
-    processor_state,
+import numpy as np
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
+    processor_state,
 )
 from ezmsg.util.messages.axisarray import AxisArray, replace
-from ezmsg.sigproc.sampler import SampleMessage
+from sklearn.linear_model._base import LinearModel
 
-from ..util import get_regressor, StaticLinearRegressor, RegressorType
+from ..util import RegressorType, StaticLinearRegressor, get_regressor
 
 
 class LinearRegressorSettings(ez.Settings):
@@ -27,9 +26,7 @@ class LinearRegressorState:
 
 
 class LinearRegressorTransformer(
-    BaseAdaptiveTransformer[
-        LinearRegressorSettings, AxisArray, AxisArray, LinearRegressorState
-    ]
+    BaseAdaptiveTransformer[LinearRegressorSettings, AxisArray, AxisArray, LinearRegressorState]
 ):
     """
     Linear regressor.
@@ -47,9 +44,7 @@ class LinearRegressorTransformer(
             with open(self.settings.settings_path, "rb") as f:
                 self.state.model = pickle.load(f)
         else:
-            regressor_klass = get_regressor(
-                RegressorType.STATIC, self.settings.model_type
-            )
+            regressor_klass = get_regressor(RegressorType.STATIC, self.settings.model_type)
             self.state.model = regressor_klass(**self.settings.model_kwargs)
 
     def _reset_state(self, message: AxisArray) -> None:
@@ -57,18 +52,18 @@ class LinearRegressorTransformer(
         #  .model and .template are initialized in __init__
         pass
 
-    def partial_fit(self, message: SampleMessage) -> None:
-        if np.any(np.isnan(message.sample.data)):
+    def partial_fit(self, message: AxisArray) -> None:
+        if np.any(np.isnan(message.data)):
             return
 
-        X = message.sample.data
-        y = message.trigger.value.data
+        X = message.data
+        y = message.attrs["trigger"].value.data
         # TODO: Resample should provide identical durations.
         self.state.model = self.state.model.fit(X[: y.shape[0]], y[: X.shape[0]])
         self.state.template = replace(
-            message.trigger.value,
+            message.attrs["trigger"].value,
             data=np.array([[]]),
-            key=message.trigger.value.key + "_pred",
+            key=message.attrs["trigger"].value.key + "_pred",
         )
 
     def _process(self, message: AxisArray) -> AxisArray:

ezmsg/learn/process/mlp_old.py

@@ -1,17 +1,16 @@
 import typing
 
+import ezmsg.core as ez
 import numpy as np
 import torch
 import torch.nn
-import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.util.messages.util import replace
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..model.mlp_old import MLP
 
@@ -24,10 +23,12 @@ class MLPSettings(ez.Settings):
     """Norm layer that will be stacked on top of the linear layer. If None this layer won’t be used."""
 
     activation_layer: typing.Callable[..., torch.nn.Module] | None = torch.nn.ReLU
-    """Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the linear layer. If None this layer won’t be used."""
+    """Activation function which will be stacked on top of the normalization layer (if not None),
+    otherwise on top of the linear layer. If None this layer won’t be used."""
 
     inplace: bool | None = None
-    """Parameter for the activation layer, which can optionally do the operation in-place. Default is None, which uses the respective default values of the activation_layer and Dropout layer."""
+    """Parameter for the activation layer, which can optionally do the operation in-place.
+    Default is None, which uses the respective default values of the activation_layer and Dropout layer."""
 
     bias: bool = True
     """Whether to use bias in the linear layer."""
@@ -58,9 +59,7 @@ class MLPState:
     device: object | None = None
 
 
-class MLPProcessor(
-    BaseAdaptiveTransformer[MLPSettings, AxisArray, AxisArray, MLPState]
-):
+class MLPProcessor(BaseAdaptiveTransformer[MLPSettings, AxisArray, AxisArray, MLPState]):
     def _hash_message(self, message: AxisArray) -> int:
         hash_items = (message.key,)
         if "ch" in message.dims:
@@ -85,39 +84,29 @@ class MLPProcessor(
                 checkpoint = torch.load(self.settings.checkpoint_path)
                 self._state.model.load_state_dict(checkpoint["model_state_dict"])
             except Exception as e:
-                raise RuntimeError(
-                    f"Failed to load checkpoint from {self.settings.checkpoint_path}: {str(e)}"
-                )
+                raise RuntimeError(f"Failed to load checkpoint from {self.settings.checkpoint_path}: {str(e)}")
 
         # Set the model to evaluation mode by default
         self._state.model.eval()
 
         # Create the optimizer
-        self._state.optimizer = torch.optim.Adam(
-            self._state.model.parameters(), lr=self.settings.learning_rate
-        )
+        self._state.optimizer = torch.optim.Adam(self._state.model.parameters(), lr=self.settings.learning_rate)
 
         # Update the optimizer from checkpoint if it exists
         if self.settings.checkpoint_path is not None:
             try:
                 checkpoint = torch.load(self.settings.checkpoint_path)
                 if "optimizer_state_dict" in checkpoint:
-                    self._state.optimizer.load_state_dict(
-                        checkpoint["optimizer_state_dict"]
-                    )
+                    self._state.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
             except Exception as e:
-                raise RuntimeError(
-                    f"Failed to load optimizer from {self.settings.checkpoint_path}: {str(e)}"
-                )
+                raise RuntimeError(f"Failed to load optimizer from {self.settings.checkpoint_path}: {str(e)}")
 
         # TODO: Should the model be moved to a device before the next line?
         self._state.device = next(self.state.model.parameters()).device
 
         # Optionally create the learning rate scheduler
         self._state.scheduler = (
-            torch.optim.lr_scheduler.ExponentialLR(
-                self._state.optimizer, gamma=self.settings.scheduler_gamma
-            )
+            torch.optim.lr_scheduler.ExponentialLR(self._state.optimizer, gamma=self.settings.scheduler_gamma)
             if self.settings.scheduler_gamma > 0.0
             else None
         )
@@ -144,14 +133,14 @@ class MLPProcessor(
         dtype = torch.float32 if self.settings.single_precision else torch.float64
         return torch.tensor(data, dtype=dtype, device=self._state.device)
 
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         self._state.model.train()
 
         # TODO: loss_fn should be determined by setting
         loss_fn = torch.nn.functional.mse_loss
 
-        X = self._to_tensor(message.sample.data)
-        y_targ = self._to_tensor(message.trigger.value)
+        X = self._to_tensor(message.data)
+        y_targ = self._to_tensor(message.attrs["trigger"].value)
 
         with torch.set_grad_enabled(True):
             self._state.model.train()
@@ -171,9 +160,7 @@ class MLPProcessor(
         if not isinstance(data, torch.Tensor):
             data = torch.tensor(
                 data,
-                dtype=torch.float32
-                if self.settings.single_precision
-                else torch.float64,
+                dtype=torch.float32 if self.settings.single_precision else torch.float64,
             )
 
         with torch.no_grad():

ezmsg/learn/process/refit_kalman.py

@@ -3,12 +3,11 @@ from pathlib import Path
 
 import ezmsg.core as ez
 import numpy as np
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
 )
-from ezmsg.sigproc.sampler import SampleMessage
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -284,27 +283,25 @@ class RefitKalmanFilterProcessor(
             key=f"{message.key}_filtered" if hasattr(message, "key") else "filtered",
         )
 
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         """
         Perform refitting using externally provided data.
 
-        Expects message.sample.data (neural input) and message.trigger.value as a dict with:
+        Expects message.data (neural input) and message.attrs["trigger"].value as a dict with:
             - Y_state: (n_samples, n_states) array
             - intention_velocity_indices: Optional[int]
             - target_positions: Optional[np.ndarray]
             - cursor_positions: Optional[np.ndarray]
             - hold_flags: Optional[list[bool]]
         """
-        if not hasattr(message, "sample") or not hasattr(message, "trigger"):
+        if "trigger" not in message.attrs:
             raise ValueError("Invalid message format for partial_fit.")
 
-        X = np.array(message.sample.data)
-        values = message.trigger.value
+        X = np.array(message.data)
+        values = message.attrs["trigger"].value
 
         if not isinstance(values, dict) or "Y_state" not in values:
-            raise ValueError(
-                "partial_fit expects trigger.value to include at least 'Y_state'."
-            )
+            raise ValueError("partial_fit expects trigger.value to include at least 'Y_state'.")
 
         kwargs = {
             "X_neural": X,
@@ -319,9 +316,7 @@ class RefitKalmanFilterProcessor(
             "hold_flags",
         ]:
             if key in values and values[key] is not None:
-                kwargs[key if key != "hold_flags" else "hold_indices"] = np.array(
-                    values[key]
-                )
+                kwargs[key if key != "hold_flags" else "hold_indices"] = np.array(values[key])
 
         # Call model refit
         self._state.model.refit(**kwargs)

ezmsg/learn/process/rnn.py

@@ -3,9 +3,8 @@ import typing
 import ezmsg.core as ez
 import numpy as np
 import torch
-from ezmsg.sigproc.base import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.sigproc.util.profile import profile_subpub
+from ezmsg.baseproc import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit
+from ezmsg.baseproc.util.profile import profile_subpub
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -47,9 +46,7 @@ class RNNProcessor(
     TorchProcessorMixin,
     ModelInitMixin,
 ):
-    def _infer_output_sizes(
-        self, model: torch.nn.Module, n_input: int
-    ) -> dict[str, int]:
+    def _infer_output_sizes(self, model: torch.nn.Module, n_input: int) -> dict[str, int]:
         """Simple inference to get output channel size."""
         dummy_input = torch.zeros(1, 50, n_input, device=self._state.device)
         with torch.no_grad():
@@ -78,9 +75,7 @@ class RNNProcessor(
                 preserve_state = True
             elif "time" not in axes or "win" not in axes:
                 # Default fallback
-                ez.logger.warning(
-                    "Missing 'time' or 'win' axis for auto preserve-state logic. Defaulting to reset."
-                )
+                ez.logger.warning("Missing 'time' or 'win' axis for auto preserve-state logic. Defaulting to reset.")
                 preserve_state = False
             else:
                 # Calculate stride between windows (assuming uniform spacing)
@@ -89,9 +84,7 @@ class RNNProcessor(
                 time_len = message.data.shape[message.get_axis_idx("time")]
                 gain = getattr(axes["time"], "gain", None)
                 if gain is None:
-                    ez.logger.warning(
-                        "Time axis gain not found, using default gain of 1.0."
-                    )
+                    ez.logger.warning("Time axis gain not found, using default gain of 1.0.")
                     gain = 1.0  # fallback default
                 win_len = time_len * gain
                 # Determine if we should preserve state
@@ -102,15 +95,9 @@ class RNNProcessor(
             self.reset_hidden(batch_size)
         else:
             # If preserving state, only reset if batch size isn't 1
-            hx_batch_size = (
-                self._state.hx[0].shape[1]
-                if isinstance(self._state.hx, tuple)
-                else self._state.hx.shape[1]
-            )
+            hx_batch_size = self._state.hx[0].shape[1] if isinstance(self._state.hx, tuple) else self._state.hx.shape[1]
             if hx_batch_size != 1:
-                ez.logger.debug(
-                    f"Resetting hidden state due to batch size mismatch (hx: {hx_batch_size}, new: 1)"
-                )
+                ez.logger.debug(f"Resetting hidden state due to batch size mismatch (hx: {hx_batch_size}, new: 1)")
                 self.reset_hidden(1)
         return preserve_state
 
@@ -119,9 +106,7 @@ class RNNProcessor(
         if not isinstance(x, torch.Tensor):
             x = torch.tensor(
                 x,
-                dtype=torch.float32
-                if self.settings.single_precision
-                else torch.float64,
+                dtype=torch.float32 if self.settings.single_precision else torch.float64,
                 device=self._state.device,
             )
 
@@ -143,18 +128,11 @@ class RNNProcessor(
                             y_data[key] = []
                         y_data[key].append(out.cpu().numpy())
                 # Concatenate outputs for each key
-                y_data = {
-                    key: np.concatenate(outputs, axis=0)
-                    for key, outputs in y_data.items()
-                }
+                y_data = {key: np.concatenate(outputs, axis=0) for key, outputs in y_data.items()}
             else:
                 y, self._state.hx = self._state.model(x, hx=self._state.hx)
                 y_data = {
-                    key: (
-                        out.cpu().numpy().squeeze(0)
-                        if added_batch_dim
-                        else out.cpu().numpy()
-                    )
+                    key: (out.cpu().numpy().squeeze(0) if added_batch_dim else out.cpu().numpy())
                     for key, out in y.items()
                 }
 
@@ -205,18 +183,18 @@ class RNNProcessor(
         if self._state.scheduler is not None:
             self._state.scheduler.step()
 
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         self._state.model.train()
 
-        X = self._to_tensor(message.sample.data)
+        X = self._to_tensor(message.data)
 
         # Add batch dimension if missing
         X, batched = self._ensure_batched(X)
 
         batch_size = X.shape[0]
-        preserve_state = self._maybe_reset_state(message.sample, batch_size)
+        preserve_state = self._maybe_reset_state(message, batch_size)
 
-        y_targ = message.trigger.value
+        y_targ = message.attrs["trigger"].value
         if not isinstance(y_targ, dict):
             y_targ = {"output": y_targ}
         y_targ = {key: self._to_tensor(value) for key, value in y_targ.items()}

ezmsg/learn/process/sgd.py

@@ -2,9 +2,11 @@ import typing
 
 import ezmsg.core as ez
 import numpy as np
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.sigproc.base import GenAxisArray
-from ezmsg.util.generator import consumer
+from ezmsg.baseproc import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 from sklearn.exceptions import NotFittedError
@@ -13,103 +15,6 @@ from sklearn.linear_model import SGDClassifier
 from ..util import ClassifierMessage
 
 
-@consumer
-def sgd_decoder(
-    alpha: float = 1.5e-5,
-    eta0: float = 1e-7,  # Lower than what you'd use for offline training.
-    loss: str = "squared_hinge",
-    label_weights: dict[str, float] | None = None,
-    settings_path: str | None = None,
-) -> typing.Generator[AxisArray | SampleMessage, ClassifierMessage | None, None]:
-    """
-    Passive Aggressive Classifier
-    Online Passive-Aggressive Algorithms <http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf>
-    K. Crammer, O. Dekel, J. Keshat, S. Shalev-Shwartz, Y. Singer - JMLR (2006)
-
-    Args:
-        alpha: Maximum step size (regularization)
-        eta0: The initial learning rate for the 'adaptive’ schedules.
-        loss: The loss function to be used:
-            hinge: equivalent to PA-I in the reference paper.
-            squared_hinge: equivalent to PA-II in the reference paper.
-        label_weights: An optional dictionary of label names and their relative weight.
-            e.g., {'Go': 31.0, 'Stop': 0.5}
-            If this is None then settings_path must be provided and the pre-trained model
-        settings_path: Path to the stored sklearn model pkl file.
-
-    Returns:
-        Generator that accepts `SampleMessage` for incremental training (`partial_fit`) and yields None,
-        or `AxisArray` for inference (`predict`) and yields a `ClassifierMessage`.
-    """
-    # pre-init inputs and outputs
-    msg_out = ClassifierMessage(data=np.array([]), dims=[""])
-
-    # State variables:
-
-    if settings_path is not None:
-        import pickle
-
-        with open(settings_path, "rb") as f:
-            model = pickle.load(f)
-            if label_weights is not None:
-                model.class_weight = label_weights
-            # Overwrite eta0, probably with a value lower than what was used online.
-            model.eta0 = eta0
-    else:
-        model = SGDClassifier(
-            loss=loss,
-            alpha=alpha,
-            penalty="elasticnet",
-            learning_rate="adaptive",
-            eta0=eta0,
-            early_stopping=False,
-            class_weight=label_weights,
-        )
-
-    b_first_train = True
-    # TODO: template_out
-
-    while True:
-        msg_in: AxisArray | SampleMessage = yield msg_out
-
-        msg_out = None
-        if type(msg_in) is SampleMessage:
-            # SampleMessage used for training.
-            if not np.any(np.isnan(msg_in.sample.data)):
-                train_sample = msg_in.sample.data.reshape(1, -1)
-                if b_first_train:
-                    model.partial_fit(
-                        train_sample,
-                        [msg_in.trigger.value],
-                        classes=list(label_weights.keys()),
-                    )
-                    b_first_train = False
-                else:
-                    model.partial_fit(train_sample, [msg_in.trigger.value])
-        elif msg_in.data.size:
-            # AxisArray used for inference
-            if not np.any(np.isnan(msg_in.data)):
-                try:
-                    X = msg_in.data.reshape((msg_in.data.shape[0], -1))
-                    result = model._predict_proba_lr(X)
-                except NotFittedError:
-                    result = None
-                if result is not None:
-                    out_axes = {}
-                    if msg_in.dims[0] in msg_in.axes:
-                        out_axes[msg_in.dims[0]] = replace(
-                            msg_in.axes[msg_in.dims[0]],
-                            offset=msg_in.axes[msg_in.dims[0]].offset,
-                        )
-                    msg_out = ClassifierMessage(
-                        data=result,
-                        dims=msg_in.dims[:1] + ["labels"],
-                        axes=out_axes,
-                        labels=list(model.class_weight.keys()),
-                        key=msg_in.key,
-                    )
-
-
 class SGDDecoderSettings(ez.Settings):
     alpha: float = 1e-5
     eta0: float = 3e-4
@@ -118,14 +23,94 @@ class SGDDecoderSettings(ez.Settings):
     settings_path: str | None = None
 
 
-class SGDDecoder(GenAxisArray):
-    SETTINGS = SGDDecoderSettings
-    INPUT_SAMPLE = ez.InputStream(SampleMessage)
+@processor_state
+class SGDDecoderState:
+    model: typing.Any = None
+    b_first_train: bool = True
 
-    # Method to be implemented by subclasses to construct the specific generator
-    def construct_generator(self):
-        self.STATE.gen = sgd_decoder(**self.SETTINGS.__dict__)
 
-    @ez.subscriber(INPUT_SAMPLE)
-    async def on_sample(self, msg: SampleMessage) -> None:
-        _ = self.STATE.gen.send(msg)
+class SGDDecoderTransformer(BaseAdaptiveTransformer[SGDDecoderSettings, AxisArray, ClassifierMessage, SGDDecoderState]):
+    """
+    SGD-based online classifier.
+
+    Online Passive-Aggressive Algorithms
+    <http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf>
+    K. Crammer, O. Dekel, J. Keshat, S. Shalev-Shwartz, Y. Singer - JMLR (2006)
+    """
+
+    def _refreshed_model(self):
+        if self.settings.settings_path is not None:
+            import pickle
+
+            with open(self.settings.settings_path, "rb") as f:
+                model = pickle.load(f)
+                if self.settings.label_weights is not None:
+                    model.class_weight = self.settings.label_weights
+                model.eta0 = self.settings.eta0
+        else:
+            model = SGDClassifier(
+                loss=self.settings.loss,
+                alpha=self.settings.alpha,
+                penalty="elasticnet",
+                learning_rate="adaptive",
+                eta0=self.settings.eta0,
+                early_stopping=False,
+                class_weight=self.settings.label_weights,
+            )
+        return model
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.model = self._refreshed_model()
+
+    def _process(self, message: AxisArray) -> ClassifierMessage | None:
+        if self._state.model is None or not message.data.size:
+            return None
+        if np.any(np.isnan(message.data)):
+            return None
+        try:
+            X = message.data.reshape((message.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,
+            )
+        return ClassifierMessage(
+            data=result,
+            dims=message.dims[:1] + ["labels"],
+            axes=out_axes,
+            labels=list(self._state.model.class_weight.keys()),
+            key=message.key,
+        )
+
+    def partial_fit(self, message: AxisArray) -> None:
+        if self._hash != 0:
+            self._reset_state(message)
+            self._hash = 0
+
+        if np.any(np.isnan(message.data)):
+            return
+        train_sample = message.data.reshape(1, -1)
+        if self._state.b_first_train:
+            self._state.model.partial_fit(
+                train_sample,
+                [message.attrs["trigger"].value],
+                classes=list(self.settings.label_weights.keys()),
+            )
+            self._state.b_first_train = False
+        else:
+            self._state.model.partial_fit(train_sample, [message.attrs["trigger"].value])
+
+
+class SGDDecoder(
+    BaseAdaptiveTransformerUnit[
+        SGDDecoderSettings,
+        AxisArray,
+        ClassifierMessage,
+        SGDDecoderTransformer,
+    ]
+):
+    SETTINGS = SGDDecoderSettings