ezmsg-learn 1.0__py3-none-any.whl

ezmsg/learn/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from ezmsg-learn!"

ezmsg/learn/__version__.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '1.0'
+__version_tuple__ = version_tuple = (1, 0)
+
+__commit_id__ = commit_id = None

ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -0,0 +1,284 @@
+import typing
+
+from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.sigproc.base import (
+    processor_state,
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+)
+from ezmsg.util.messages.axisarray import AxisArray, replace
+
+
+class AdaptiveDecompSettings(ez.Settings):
+    axis: str = "!time"
+    n_components: int = 2
+
+
+@processor_state
+class AdaptiveDecompState:
+    template: AxisArray | None = None
+    axis_groups: tuple[str, list[str], list[str]] | None = None
+    estimator: typing.Any = None
+
+
+EstimatorType = typing.TypeVar(
+    "EstimatorType", bound=typing.Union[IncrementalPCA, MiniBatchNMF]
+)
+
+
+class AdaptiveDecompTransformer(
+    BaseAdaptiveTransformer[
+        AdaptiveDecompSettings, AxisArray, AxisArray, AdaptiveDecompState
+    ],
+    typing.Generic[EstimatorType],
+):
+    """
+    Base class for adaptive decomposition transformers. See IncrementalPCATransformer and MiniBatchNMFTransformer
+    for concrete implementations.
+
+    Note that for these classes, adaptation is not automatic. The user must call partial_fit on the transformer.
+    For automated adaptation, see IncrementalDecompTransformer.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._state.estimator = self._create_estimator()
+
+    @classmethod
+    def get_message_type(cls, dir: str) -> typing.Type[AxisArray]:
+        # Override because we don't reuse the generic types.
+        return AxisArray
+
+    @classmethod
+    def get_estimator_type(cls) -> typing.Type[EstimatorType]:
+        return typing.get_args(cls.__orig_bases__[0])[0]
+
+    def _create_estimator(self) -> EstimatorType:
+        estimator_klass = self.get_estimator_type()
+        estimator_settings = self.settings.__dict__.copy()
+        estimator_settings.pop("axis")
+        return estimator_klass(**estimator_settings)
+
+    def _calculate_axis_groups(self, message: AxisArray):
+        if self.settings.axis.startswith("!"):
+            # Iterate over the !axis and collapse all other axes
+            iter_axis = self.settings.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})"
+                )
+            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
+            ]
+        self._state.axis_groups = iter_axis, targ_axes, off_targ_axes
+
+    def _hash_message(self, message: AxisArray) -> int:
+        iter_axis = (
+            self.settings.axis[1:]
+            if self.settings.axis.startswith("!")
+            else ("win" if "win" in message.dims else "time")
+        )
+        ax_idx = message.get_axis_idx(iter_axis)
+        sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+        return hash((sample_shape, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        """Reset state"""
+        self._calculate_axis_groups(message)
+        iter_axis, targ_axes, off_targ_axes = self._state.axis_groups
+
+        # Template
+        out_dims = [iter_axis] + off_targ_axes
+        out_axes = {
+            iter_axis: message.axes[iter_axis],
+            **{k: message.axes[k] for k in off_targ_axes},
+        }
+        if len(targ_axes) == 1:
+            targ_ax_name = targ_axes[0]
+        else:
+            targ_ax_name = "components"
+        out_dims += [targ_ax_name]
+        out_axes[targ_ax_name] = AxisArray.CoordinateAxis(
+            data=np.arange(self.settings.n_components).astype(str),
+            dims=[targ_ax_name],
+            unit="component",
+        )
+        out_shape = [message.data.shape[message.get_axis_idx(_)] for _ in off_targ_axes]
+        out_shape = (0,) + tuple(out_shape) + (self.settings.n_components,)
+        self._state.template = replace(
+            message,
+            data=np.zeros(out_shape, dtype=float),
+            dims=out_dims,
+            axes=out_axes,
+        )
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        iter_axis, targ_axes, off_targ_axes = self._state.axis_groups
+        ax_idx = message.get_axis_idx(iter_axis)
+        in_dat = message.data
+
+        if in_dat.shape[ax_idx] == 0:
+            return self._state.template
+
+        # Re-order axes
+        sorted_dims_exp = [iter_axis] + off_targ_axes + targ_axes
+        if message.dims != sorted_dims_exp:
+            # TODO: Implement axes transposition if needed
+            # re_order = [ax_idx] + off_targ_inds + targ_inds
+            # np.transpose(in_dat, re_order)
+            pass
+
+        # fold [iter_axis] + off_targ_axes together and fold targ_axes together
+        d2 = np.prod(in_dat.shape[len(off_targ_axes) + 1 :])
+        in_dat = in_dat.reshape((-1, d2))
+
+        replace_kwargs = {
+            "axes": {**self._state.template.axes, iter_axis: message.axes[iter_axis]},
+        }
+
+        # Transform data
+        if hasattr(self._state.estimator, "components_"):
+            decomp_dat = self._state.estimator.transform(in_dat).reshape(
+                (-1,) + self._state.template.data.shape[1:]
+            )
+            replace_kwargs["data"] = decomp_dat
+
+        return replace(self._state.template, **replace_kwargs)
+
+    def partial_fit(self, message: AxisArray) -> None:
+        # Check if we need to reset state
+        msg_hash = self._hash_message(message)
+        if self._hash != msg_hash:
+            self._reset_state(message)
+            self._hash = msg_hash
+
+        iter_axis, targ_axes, off_targ_axes = self._state.axis_groups
+        ax_idx = message.get_axis_idx(iter_axis)
+        in_dat = message.data
+
+        if in_dat.shape[ax_idx] == 0:
+            return
+
+        # Re-order axes if needed
+        sorted_dims_exp = [iter_axis] + off_targ_axes + targ_axes
+        if message.dims != sorted_dims_exp:
+            # TODO: Implement axes transposition if needed
+            pass
+
+        # fold [iter_axis] + off_targ_axes together and fold targ_axes together
+        d2 = np.prod(in_dat.shape[len(off_targ_axes) + 1 :])
+        in_dat = in_dat.reshape((-1, d2))
+
+        # Fit the estimator
+        self._state.estimator.partial_fit(in_dat)
+
+
+class IncrementalPCASettings(AdaptiveDecompSettings):
+    # Additional settings specific to PCA
+    whiten: bool = False
+    batch_size: typing.Optional[int] = None
+
+
+class IncrementalPCATransformer(AdaptiveDecompTransformer[IncrementalPCA]):
+    pass
+
+
+class MiniBatchNMFSettings(AdaptiveDecompSettings):
+    # Additional settings specific to NMF
+    init: typing.Optional[str] = "random"
+    """
+    'random', 'nndsvd', 'nndsvda', 'nndsvdar', 'custom', or None
+    """
+
+    batch_size: int = 1024
+    """
+    batch_size is used only when doing a full fit (i.e., a reset),
+    or as the exponent to forget_factor, where a very small batch_size
+    will cause the model to update more slowly.
+    It is better to set batch_size to a larger number than the expected
+    chunk size and instead use forget_factor to control the learning rate.
+    """
+
+    beta_loss: typing.Union[str, float] = "frobenius"
+    """
+    'frobenius', 'kullback-leibler', 'itakura-saito'
+    Note that values different from 'frobenius'
+        (or 2) and 'kullback-leibler' (or 1) lead to significantly slower
+        fits. Note that for `beta_loss <= 0` (or 'itakura-saito'), the input
+        matrix `X` cannot contain zeros.
+    """
+
+    tol: float = 1e-4
+
+    max_no_improvement: typing.Optional[int] = None
+
+    max_iter: int = 200
+
+    alpha_W: float = 0.0
+
+    alpha_H: typing.Union[float, str] = "same"
+
+    l1_ratio: float = 0.0
+
+    forget_factor: float = 0.7
+
+
+class MiniBatchNMFTransformer(AdaptiveDecompTransformer[MiniBatchNMF]):
+    pass
+
+
+SettingsType = typing.TypeVar(
+    "SettingsType", bound=typing.Union[IncrementalPCASettings, MiniBatchNMFSettings]
+)
+TransformerType = typing.TypeVar(
+    "TransformerType",
+    bound=typing.Union[IncrementalPCATransformer, MiniBatchNMFTransformer],
+)
+
+
+class BaseAdaptiveDecompUnit(
+    BaseAdaptiveTransformerUnit[
+        SettingsType,
+        AxisArray,
+        AxisArray,
+        TransformerType,
+    ],
+    typing.Generic[SettingsType, TransformerType],
+):
+    INPUT_SAMPLE = ez.InputStream(AxisArray)
+
+    @ez.subscriber(INPUT_SAMPLE)
+    async def on_sample(self, msg: AxisArray) -> None:
+        await self.processor.apartial_fit(msg)
+
+
+class IncrementalPCAUnit(
+    BaseAdaptiveDecompUnit[
+        IncrementalPCASettings,
+        IncrementalPCATransformer,
+    ]
+):
+    SETTINGS = IncrementalPCASettings
+
+
+class MiniBatchNMFUnit(
+    BaseAdaptiveDecompUnit[
+        MiniBatchNMFSettings,
+        MiniBatchNMFTransformer,
+    ]
+):
+    SETTINGS = MiniBatchNMFSettings

ezmsg/learn/dim_reduce/incremental_decomp.py

@@ -0,0 +1,181 @@
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.sigproc.base import (
+    CompositeProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+)
+from ezmsg.sigproc.window import WindowTransformer
+
+from .adaptive_decomp import (
+    IncrementalPCASettings,
+    IncrementalPCATransformer,
+    MiniBatchNMFSettings,
+    MiniBatchNMFTransformer,
+)
+
+
+class IncrementalDecompSettings(ez.Settings):
+    axis: str = "!time"
+    n_components: int = 2
+    update_interval: float = 0.0
+    method: str = "pca"
+    batch_size: typing.Optional[int] = None
+    # PCA specific settings
+    whiten: bool = False
+    # NMF specific settings
+    init: str = "random"
+    beta_loss: str = "frobenius"
+    tol: float = 1e-3
+    alpha_W: float = 0.0
+    alpha_H: typing.Union[float, str] = "same"
+    l1_ratio: float = 0.0
+    forget_factor: float = 0.7
+
+
+class IncrementalDecompTransformer(
+    CompositeProcessor[IncrementalDecompSettings, AxisArray, AxisArray]
+):
+    """
+    Automates usage of IncrementalPCATransformer and MiniBatchNMFTransformer by using a WindowTransformer
+    to extract training samples then calls partial_fit on the decomposition transformer.
+    """
+
+    @staticmethod
+    def _initialize_processors(
+        settings: IncrementalDecompSettings,
+    ) -> dict[str, BaseStatefulProcessor]:
+        # Create the appropriate decomposition transformer
+        if settings.method == "pca":
+            decomp_settings = IncrementalPCASettings(
+                axis=settings.axis,
+                n_components=settings.n_components,
+                batch_size=settings.batch_size,
+                whiten=settings.whiten,
+            )
+            decomp = IncrementalPCATransformer(settings=decomp_settings)
+        else:  # nmf
+            decomp_settings = MiniBatchNMFSettings(
+                axis=settings.axis,
+                n_components=settings.n_components,
+                batch_size=settings.batch_size if settings.batch_size else 1024,
+                init=settings.init,
+                beta_loss=settings.beta_loss,
+                tol=settings.tol,
+                alpha_W=settings.alpha_W,
+                alpha_H=settings.alpha_H,
+                l1_ratio=settings.l1_ratio,
+                forget_factor=settings.forget_factor,
+            )
+            decomp = MiniBatchNMFTransformer(settings=decomp_settings)
+
+        # 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"
+            windowing = WindowTransformer(
+                axis=iter_axis,
+                window_dur=settings.update_interval,
+                window_shift=settings.update_interval,
+                zero_pad_until="none",
+            )
+
+            return {
+                "decomp": decomp,
+                "windowing": windowing,
+            }
+
+        return {"decomp": decomp}
+
+    def _partial_fit_windowed(self, train_msg: AxisArray) -> None:
+        """
+        Helper function to do the partial_fit on the windowed message.
+        """
+        if np.prod(train_msg.data.shape) > 0:
+            # Windowing created a new "win" axis, but we don't actually want to use that
+            #  in the message we send to the decomp processor.
+            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")):
+                _msg = replace(
+                    _msg,
+                    axes={
+                        **_msg.axes,
+                        "time": replace(
+                            _msg.axes["time"],
+                            offset=_msg.axes["time"].offset + offsets[ix],
+                        ),
+                    },
+                )
+                self._procs["decomp"].partial_fit(_msg)
+
+    def stateful_op(
+        self,
+        state: dict[str, tuple[typing.Any, int]] | None,
+        message: AxisArray,
+    ) -> tuple[dict[str, tuple[typing.Any, int]], AxisArray]:
+        state = state or {}
+
+        estim = self._procs["decomp"]._state.estimator
+        if not hasattr(estim, "components_") or estim.components_ is None:
+            # If the estimator has not been trained once, train it with the first message
+            self._procs["decomp"].partial_fit(message)
+        elif "windowing" in self._procs:
+            state["windowing"], train_msg = self._procs["windowing"].stateful_op(
+                state.get("windowing", None), message
+            )
+            self._partial_fit_windowed(train_msg)
+
+        # Process the incoming message
+        state["decomp"], result = self._procs["decomp"].stateful_op(
+            state.get("decomp", None), message
+        )
+
+        return state, result
+
+    async def _aprocess(self, message: AxisArray) -> AxisArray:
+        """
+        Asynchronously process the incoming message.
+        This is nearly identical to the _process method, but the processors
+        are called asynchronously.
+        """
+        estim = self._procs["decomp"]._state.estimator
+        if not hasattr(estim, "components_") or estim.components_ is None:
+            # If the estimator has not been trained once, train it with the first message
+            self._procs["decomp"].partial_fit(message)
+        elif "windowing" in self._procs:
+            # If windowing is enabled, extract training samples and perform partial_fit
+            train_msg = await self._procs["windowing"].__acall__(message)
+            self._partial_fit_windowed(train_msg)  # Non async
+
+        # Process the incoming message
+        decomp_result = await self._procs["decomp"].__acall__(message)
+
+        return decomp_result
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        estim = self._procs["decomp"]._state.estimator
+        if not hasattr(estim, "components_") or estim.components_ is None:
+            # If the estimator has not been trained once, train it with the first message
+            self._procs["decomp"].partial_fit(message)
+        elif "windowing" in self._procs:
+            # If windowing is enabled, extract training samples and perform partial_fit
+            train_msg = self._procs["windowing"](message)
+            self._partial_fit_windowed(train_msg)
+
+        # Process the incoming message
+        decomp_result = self._procs["decomp"](message)
+
+        return decomp_result
+
+
+class IncrementalDecompUnit(
+    BaseTransformerUnit[
+        IncrementalDecompSettings, AxisArray, AxisArray, IncrementalDecompTransformer
+    ]
+):
+    SETTINGS = IncrementalDecompSettings

ezmsg/learn/linear_model/__init__.py

@@ -0,0 +1 @@
+# Use of this module is deprecated. Please use `ezmsg.learn.process` instead.

ezmsg/learn/linear_model/adaptive_linear_regressor.py

@@ -0,0 +1,6 @@
+from ..process.adaptive_linear_regressor import (
+    AdaptiveLinearRegressorSettings as AdaptiveLinearRegressorSettings,
+    AdaptiveLinearRegressorState as AdaptiveLinearRegressorState,
+    AdaptiveLinearRegressorTransformer as AdaptiveLinearRegressorTransformer,
+    AdaptiveLinearRegressorUnit as AdaptiveLinearRegressorUnit,
+)

ezmsg/learn/linear_model/cca.py

@@ -0,0 +1 @@
+from ..model.cca import IncrementalCCA as IncrementalCCA

ezmsg/learn/linear_model/linear_regressor.py

@@ -0,0 +1,5 @@
+from ..process.linear_regressor import (
+    LinearRegressorSettings as LinearRegressorSettings,
+    LinearRegressorState as LinearRegressorState,
+    LinearRegressorTransformer as LinearRegressorTransformer,
+)

ezmsg/learn/linear_model/sgd.py

@@ -0,0 +1,5 @@
+from ..process.sgd import (
+    sgd_decoder as sgd_decoder,
+    SGDDecoderSettings as SGDDecoderSettings,
+    SGDDecoder as SGDDecoder,
+)

ezmsg/learn/linear_model/slda.py

@@ -0,0 +1,6 @@
+from ..process.slda import (
+    SLDASettings as SLDASettings,
+    SLDAState as SLDAState,
+    SLDATransformer as SLDATransformer,
+    SLDA as SLDA,
+)

ezmsg/learn/model/cca.py

@@ -0,0 +1,122 @@
+import numpy as np
+from scipy import linalg
+
+
+class IncrementalCCA:
+    def __init__(
+        self,
+        n_components=2,
+        base_smoothing=0.95,
+        min_smoothing=0.5,
+        max_smoothing=0.99,
+        adaptation_rate=0.1,
+    ):
+        """
+        Parameters:
+        -----------
+        n_components : int
+            Number of canonical components to compute
+        base_smoothing : float
+            Base smoothing factor (will be adapted)
+        min_smoothing : float
+            Minimum allowed smoothing factor
+        max_smoothing : float
+            Maximum allowed smoothing factor
+        adaptation_rate : float
+            How quickly to adjust smoothing factor (between 0 and 1)
+        """
+        self.n_components = n_components
+        self.base_smoothing = base_smoothing
+        self.current_smoothing = base_smoothing
+        self.min_smoothing = min_smoothing
+        self.max_smoothing = max_smoothing
+        self.adaptation_rate = adaptation_rate
+        self.initialized = False
+
+    def initialize(self, d1, d2):
+        """Initialize the necessary matrices"""
+        self.d1 = d1
+        self.d2 = d2
+
+        # Initialize correlation matrices
+        self.C11 = np.zeros((d1, d1))
+        self.C22 = np.zeros((d2, d2))
+        self.C12 = np.zeros((d1, d2))
+
+        self.initialized = True
+
+    def _compute_change_magnitude(self, C11_new, C22_new, C12_new):
+        """Compute magnitude of change in correlation structure"""
+        # Frobenius norm of differences
+        diff11 = np.linalg.norm(C11_new - self.C11)
+        diff22 = np.linalg.norm(C22_new - self.C22)
+        diff12 = np.linalg.norm(C12_new - self.C12)
+
+        # Normalize by matrix sizes
+        diff11 /= self.d1 * self.d1
+        diff22 /= self.d2 * self.d2
+        diff12 /= self.d1 * self.d2
+
+        return (diff11 + diff22 + diff12) / 3
+
+    def _adapt_smoothing(self, change_magnitude):
+        """Adapt smoothing factor based on detected changes"""
+        # If change is large, decrease smoothing factor
+        target_smoothing = self.base_smoothing * (1.0 - change_magnitude)
+        target_smoothing = np.clip(
+            target_smoothing, self.min_smoothing, self.max_smoothing
+        )
+
+        # Smooth the adaptation itself
+        self.current_smoothing = (
+            1 - self.adaptation_rate
+        ) * self.current_smoothing + self.adaptation_rate * target_smoothing
+
+    def partial_fit(self, X1, X2, update_projections=True):
+        """Update the model with new samples using adaptive smoothing
+        Assumes X1 and X2 are already centered and scaled"""
+        if not self.initialized:
+            self.initialize(X1.shape[1], X2.shape[1])
+
+        # Compute new correlation matrices from current batch
+        C11_new = X1.T @ X1 / X1.shape[0]
+        C22_new = X2.T @ X2 / X2.shape[0]
+        C12_new = X1.T @ X2 / X1.shape[0]
+
+        # Detect changes and adapt smoothing factor
+        if self.C11.any():  # Skip first update
+            change_magnitude = self._compute_change_magnitude(C11_new, C22_new, C12_new)
+            self._adapt_smoothing(change_magnitude)
+
+        # Update with current smoothing factor
+        alpha = self.current_smoothing
+        self.C11 = alpha * self.C11 + (1 - alpha) * C11_new
+        self.C22 = alpha * self.C22 + (1 - alpha) * C22_new
+        self.C12 = alpha * self.C12 + (1 - alpha) * C12_new
+
+        if update_projections:
+            self._update_projections()
+
+    def _update_projections(self):
+        """Update canonical vectors and correlations"""
+        eps = 1e-8
+        C11_reg = self.C11 + eps * np.eye(self.d1)
+        C22_reg = self.C22 + eps * np.eye(self.d2)
+
+        K = (
+            linalg.inv(linalg.sqrtm(C11_reg))
+            @ self.C12
+            @ linalg.inv(linalg.sqrtm(C22_reg))
+        )
+        U, self.correlations_, V = linalg.svd(K)
+
+        self.x_weights_ = linalg.inv(linalg.sqrtm(C11_reg)) @ U[:, : self.n_components]
+        self.y_weights_ = (
+            linalg.inv(linalg.sqrtm(C22_reg)) @ V.T[:, : self.n_components]
+        )
+
+    def transform(self, X1, X2):
+        """Project data onto canonical components"""
+        X1_proj = X1 @ self.x_weights_
+        X2_proj = X2 @ self.y_weights_
+        return X1_proj, X2_proj