ezmsg-learn 1.1.0__py3-none-any.whl

ezmsg/learn/model/mlp.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn
+
+
+class MLP(torch.nn.Module):
+    """
+    A simple Multi-Layer Perceptron (MLP) model. Adapted from Ezmsg MLP.
+
+    Attributes:
+        feature_extractor (torch.nn.Sequential): The sequential feature extractor part of the MLP.
+        heads (torch.nn.ModuleDict): A dictionary of output linear layers for each output head.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int | list[int],
+        num_layers: int | None = None,
+        output_heads: int | dict[str, int] = 2,
+        norm_layer: str | None = None,
+        activation_layer: str | None = "ReLU",
+        inplace: bool | None = None,
+        bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        """
+        Initialize the MLP model.
+        Args:
+            input_size (int): The size of the input features.
+            hidden_size (int | list[int]): The sizes of the hidden layers. If a list, num_layers must be None or the
+                length of the list. If a single integer, num_layers must be specified and determines the number of
+                hidden layers.
+            num_layers (int, optional): The number of hidden layers. Length of hidden_size if None. Default is None.
+            output_heads (int | dict[str, int], optional): Number of output features or classes if single head output
+                or a dictionary mapping head names to output sizes if multi-head output. Default is 2 (single head).
+            norm_layer (str, optional): A normalization layer to be applied after each linear layer. Default is None.
+                Common choices are "BatchNorm1d" or "LayerNorm".
+            activation_layer (str, optional): An activation function to be applied after each normalization
+                layer. Default is "ReLU".
+            inplace (bool, optional): Whether the activation function is performed in-place. Default is None.
+            bias (bool, optional): Whether to use bias in the linear layers. Default is True.
+            dropout (float, optional): The dropout rate to be applied after each linear layer. Default is 0.0.
+        """
+        super().__init__()
+        if isinstance(hidden_size, int):
+            if num_layers is None:
+                raise ValueError("If hidden_size is an integer, num_layers must be specified.")
+            hidden_size = [hidden_size] * num_layers
+        if len(hidden_size) == 0:
+            raise ValueError("hidden_size must have at least one element")
+        if any(not isinstance(x, int) for x in hidden_size):
+            raise ValueError("hidden_size must contain only integers")
+        if num_layers is not None and len(hidden_size) != num_layers:
+            raise ValueError("Length of hidden_size must match num_layers if num_layers is specified.")
+
+        params = {} if inplace is None else {"inplace": inplace}
+
+        layers = []
+        in_dim = input_size
+
+        def _get_layer_class(layer_name: str):
+            if layer_name is not None and "torch.nn" in layer_name:
+                return getattr(torch.nn, layer_name.rsplit(".", 1)[1])
+            return None
+
+        norm_layer_class = _get_layer_class(norm_layer)
+        activation_layer_class = _get_layer_class(activation_layer)
+        for hidden_dim in hidden_size[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
+            if norm_layer_class is not None:
+                layers.append(norm_layer_class(hidden_dim))
+            if activation_layer_class is not None:
+                layers.append(activation_layer_class(**params))
+            layers.append(torch.nn.Dropout(dropout, **params))
+            in_dim = hidden_dim
+
+        layers.append(torch.nn.Linear(in_dim, hidden_size[-1], bias=bias))
+
+        self.feature_extractor = torch.nn.Sequential(*layers)
+
+        if isinstance(output_heads, int):
+            output_heads = {"output": output_heads}
+        self.heads = torch.nn.ModuleDict(
+            {name: torch.nn.Linear(hidden_size[-1], output_size) for name, output_size in output_heads.items()}
+        )
+
+    @classmethod
+    def infer_config_from_state_dict(cls, state_dict: dict) -> dict[str, int | float]:
+        """
+        Infer the configuration from the state dict.
+
+        Args:
+            state_dict: The state dict of the model.
+
+        Returns:
+            dict[str, int | float]: A dictionary containing the inferred configuration.
+        """
+        input_size = state_dict["feature_extractor.0.weight"].shape[1]
+        hidden_size = [
+            param.shape[0]
+            for key, param in state_dict.items()
+            if key.startswith("feature_extractor.") and key.endswith(".weight")
+        ]
+        output_heads = {
+            key.split(".")[1]: param.shape[0]
+            for key, param in state_dict.items()
+            if key.startswith("heads.") and key.endswith(".bias")
+        }
+
+        return {
+            "input_size": input_size,
+            "hidden_size": hidden_size,
+            "output_heads": output_heads,
+        }
+
+    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        Forward pass through the MLP.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch, seq_len, input_size).
+
+        Returns:
+            dict[str, torch.Tensor]: A dictionary mapping head names to output tensors.
+        """
+        x = self.feature_extractor(x)
+        return {name: head(x) for name, head in self.heads.items()}

ezmsg/learn/model/mlp_old.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn
+
+
+class MLP(torch.nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: list[int],
+        norm_layer: torch.nn.Module | None = None,
+        activation_layer: torch.nn.Module | None = torch.nn.ReLU,
+        inplace: bool | None = None,
+        bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        """
+        Copy-pasted from torchvision MLP
+
+        :param in_channels: Number of input channels
+        :param hidden_channels: List of the hidden channel dimensions
+        :param norm_layer: Norm layer that will be stacked on top of the linear layer. If None this layer won’t be used.
+        :param activation_layer: 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.
+        :param inplace: 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.
+        :param bias: Whether to use bias in the linear layer.
+        :param dropout: The probability for the dropout layer.
+        """
+        if len(hidden_channels) == 0:
+            raise ValueError("hidden_channels must have at least one element")
+        if any(not isinstance(x, int) for x in hidden_channels):
+            raise ValueError("hidden_channels must contain only integers")
+
+        params = {} if inplace is None else {"inplace": inplace}
+
+        layers = []
+        in_dim = in_channels
+        for hidden_dim in hidden_channels[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
+            if norm_layer is not None:
+                layers.append(norm_layer(hidden_dim))
+            if activation_layer is not None:
+                layers.append(activation_layer(**params))
+            layers.append(torch.nn.Dropout(dropout, **params))
+            in_dim = hidden_dim
+
+        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias))
+
+        super().__init__(*layers)

ezmsg/learn/model/refit_kalman.py

@@ -0,0 +1,369 @@
+# refit_kalman.py
+
+import numpy as np
+from numpy.linalg import LinAlgError
+from scipy.linalg import solve_discrete_are
+
+
+class RefitKalmanFilter:
+    """
+    Refit Kalman filter for adaptive neural decoding.
+
+    This class implements a Kalman filter that can be refitted online during operation.
+    Unlike the standard Kalman filter, this version can adapt its observation model
+    (H and Q matrices) based on new data while maintaining the state transition model
+    (A and W matrices). This is particularly useful for brain-computer interfaces
+    where the relationship between neural activity and intended movements may change
+    over time.
+
+    The filter operates in two phases:
+    1. Initial fitting: Learns all system matrices (A, W, H, Q) from training data
+    2. Refitting: Updates only the observation model (H, Q) based on new data
+
+    Attributes:
+        A_state_transition_matrix: The state transition matrix A (n_states x n_states).
+        W_process_noise_covariance: The process noise covariance matrix W (n_states x n_states).
+        H_observation_matrix: The observation matrix H (n_observations x n_states).
+        Q_measurement_noise_covariance: The measurement noise covariance matrix Q (n_observations x n_observations).
+        K_kalman_gain: The Kalman gain matrix (n_states x n_observations).
+        P_state_covariance: The state error covariance matrix (n_states x n_states).
+        steady_state: Whether to use steady-state Kalman gain computation.
+        is_fitted: Whether the model has been fitted with data.
+
+    Example:
+        >>> # Create and fit the filter
+        >>> rkf = RefitKalmanFilter(steady_state=True)
+        >>> rkf.fit(X_train, y_train)
+        >>>
+        >>> # Refit with new data
+        >>> rkf.refit(X_new, Y_state, velocity_indices, targets, cursors, holds)
+        >>>
+        >>> # Predict with updated model
+        >>> x_updated = rkf.predict_and_update(measurement, current_state)
+    """
+
+    def __init__(
+        self,
+        A_state_transition_matrix=None,
+        W_process_noise_covariance=None,
+        H_observation_matrix=None,
+        Q_measurement_noise_covariance=None,
+        steady_state=False,
+        enforce_state_structure=False,
+        alpha_fading_memory=1.000,
+        process_noise_scale=1,
+        measurement_noise_scale=1.2,
+    ):
+        self.A_state_transition_matrix = A_state_transition_matrix
+        self.W_process_noise_covariance = W_process_noise_covariance
+        self.H_observation_matrix = H_observation_matrix
+        self.Q_measurement_noise_covariance = Q_measurement_noise_covariance
+        self.K_kalman_gain = None
+        self.P_state_covariance = None
+        self.alpha_fading_memory = alpha_fading_memory
+
+        # Noise scaling factors for smoothing control
+        self.process_noise_scale = process_noise_scale
+        self.measurement_noise_scale = measurement_noise_scale
+
+        self.steady_state = steady_state
+        self.enforce_state_structure = enforce_state_structure
+        self.is_fitted = False
+
+    def _validate_state_vector(self, Y_state):
+        """
+        Validate that the state vector has proper dimensions.
+
+        Args:
+            Y_state: State vector to validate
+
+        Raises:
+            ValueError: If state vector has invalid dimensions
+        """
+        if Y_state.ndim != 2:
+            raise ValueError(f"State vector must be 2D, got {Y_state.ndim}D")
+
+        if not hasattr(self, "H_observation_matrix") or self.H_observation_matrix is None:
+            raise ValueError("Model must be fitted before refitting")
+
+        expected_states = self.H_observation_matrix.shape[1]
+        if Y_state.shape[1] != expected_states:
+            raise ValueError(f"State vector has {Y_state.shape[1]} dimensions, expected {expected_states}")
+
+    def fit(self, X_train, y_train):
+        """
+        Fit the Refit Kalman filter to the training data.
+
+        This method learns all system matrices (A, W, H, Q) from training data
+        using least-squares estimation, then computes the steady-state solution.
+        This is the initial fitting phase that establishes the baseline model.
+
+        Args:
+            X_train: Neural activity (n_samples, n_neurons).
+            y_train: Outputs being predicted (n_samples, n_states).
+
+        Raises:
+            ValueError: If training data has invalid dimensions.
+            LinAlgError: If matrix operations fail during fitting.
+        """
+        # self._validate_state_vector(y_train)
+
+        X = np.array(y_train)
+        Z = np.array(X_train)
+        n_samples = X.shape[0]
+
+        # Calculate the transition matrix (from x_t to x_t+1) using least-squares
+        X2 = X[1:, :]  # x_{t+1}
+        X1 = X[:-1, :]  # x_t
+        A = X2.T @ X1 @ np.linalg.inv(X1.T @ X1)  # Transition matrix
+        W = (X2 - X1 @ A.T).T @ (X2 - X1 @ A.T) / (n_samples - 1)  # Covariance of transition matrix
+
+        # Calculate the measurement matrix (from x_t to z_t) using least-squares
+        H = Z.T @ X @ np.linalg.inv(X.T @ X)  # Measurement matrix
+        Q = (Z - X @ H.T).T @ (Z - X @ H.T) / Z.shape[0]  # Covariance of measurement matrix
+
+        self.A_state_transition_matrix = A
+        self.W_process_noise_covariance = W * self.process_noise_scale
+        self.H_observation_matrix = H
+        self.Q_measurement_noise_covariance = Q * self.measurement_noise_scale
+
+        self._compute_gain()
+        self.is_fitted = True
+
+    def refit(
+        self,
+        X_neural: np.ndarray,
+        Y_state: np.ndarray,
+        intention_velocity_indices: int | None = None,
+        target_positions: np.ndarray | None = None,
+        cursor_positions: np.ndarray | None = None,
+        hold_indices: np.ndarray | None = None,
+    ):
+        """
+        Refit the observation model based on new data.
+
+        This method updates only the observation model (H and Q matrices) while
+        keeping the state transition model (A and W matrices) unchanged. The refitting
+        process modifies the intended states based on target positions and hold flags
+        to better align with user intentions.
+
+        The refitting process:
+        1. Modifies intended states based on target positions and hold flags
+        2. Recalculates the observation matrix H using least-squares
+        3. Recalculates the measurement noise covariance Q
+        4. Updates the Kalman gain accordingly
+
+        Args:
+            X_neural: Neural activity data (n_samples, n_neurons).
+            Y_state: State estimates (n_samples, n_states).
+            intention_velocity_indices: Index of velocity components in state vector.
+            target_positions: Target positions for each sample (n_samples, 2).
+            cursor_positions: Current cursor positions (n_samples, 2).
+            hold_indices: Boolean flags indicating hold periods (n_samples,).
+
+        Raises:
+            ValueError: If input data has invalid dimensions or the model is not fitted.
+        """
+        self._validate_state_vector(Y_state)
+
+        # Check if velocity indices are provided
+        if intention_velocity_indices is None:
+            # Assume (x, y, vx, vy)
+            vel_idx = 2 if Y_state.shape[1] >= 4 else 0
+            print(f"[RefitKalmanFilter] No velocity index provided — defaulting to {vel_idx}")
+        else:
+            if isinstance(intention_velocity_indices, (list, tuple)):
+                if len(intention_velocity_indices) != 1:
+                    raise ValueError("Only one velocity start index should be provided.")
+                vel_idx = intention_velocity_indices[0]
+            else:
+                vel_idx = intention_velocity_indices
+
+        # Only remap velocity if target and cursor positions are provided
+        if target_positions is None or cursor_positions is None:
+            intended_states = Y_state.copy()
+        else:
+            intended_states = Y_state.copy()
+            # Calculate intended velocities for each sample
+            for i, (state, pos, target) in enumerate(zip(Y_state, cursor_positions, target_positions)):
+                is_hold = hold_indices[i] if hold_indices is not None else False
+
+                if is_hold:
+                    # During hold periods, intended velocity is zero
+                    intended_states[i, vel_idx : vel_idx + 2] = 0.0
+                    if i > 0:
+                        intended_states[i, :2] = intended_states[i - 1, :2]  # Same position as previous
+                else:
+                    # Calculate direction to target
+                    to_target = target - pos
+                    target_distance = np.linalg.norm(to_target)
+
+                    if target_distance > 1e-5:  # Avoid division by zero
+                        # Get current decoded velocity magnitude
+                        current_velocity = state[vel_idx : vel_idx + 2]
+                        current_speed = np.linalg.norm(current_velocity)
+
+                        # Calculate intended velocity: same speed, but toward target
+                        target_direction = to_target / target_distance
+                        intended_velocity = target_direction * current_speed
+
+                        # Update intended state with new velocity
+                        intended_states[i, vel_idx : vel_idx + 2] = intended_velocity
+                    # If target is very close, keep original velocity
+                    else:
+                        intended_states[i, vel_idx : vel_idx + 2] = state[vel_idx : vel_idx + 2]
+
+        intended_states = np.array(intended_states)
+        Z = np.array(X_neural)
+
+        # Recalculate observation matrix and noise covariance
+        H = (
+            Z.T @ intended_states @ np.linalg.pinv(intended_states.T @ intended_states)
+        )  # Using pinv() instead of inv() to avoid singular matrix errors
+        Q = (Z - intended_states @ H.T).T @ (Z - intended_states @ H.T) / Z.shape[0]
+
+        self.H_observation_matrix = H
+        self.Q_measurement_noise_covariance = Q
+
+        self._compute_gain()
+
+    def _compute_gain(self):
+        """
+        Compute the Kalman gain matrix.
+
+        This method computes the Kalman gain matrix based on the current system
+        parameters. In steady-state mode, it solves the discrete-time algebraic
+        Riccati equation to find the optimal steady-state gain. In non-steady-state
+        mode, it computes the gain using the current covariance matrix.
+
+        Raises:
+            LinAlgError: If the Riccati equation cannot be solved or matrix operations fail.
+        """
+        # TODO: consider removing non-steady-state for compute_gain() -
+        #  non_steady_state updates will occur during predict() and update()
+        # if self.steady_state:
+        try:
+            # Try with original matrices
+            self.P_state_covariance = solve_discrete_are(
+                self.A_state_transition_matrix.T,
+                self.H_observation_matrix.T,
+                self.W_process_noise_covariance,
+                self.Q_measurement_noise_covariance,
+            )
+            self.K_kalman_gain = (
+                self.P_state_covariance
+                @ self.H_observation_matrix.T
+                @ np.linalg.inv(
+                    self.H_observation_matrix @ self.P_state_covariance @ self.H_observation_matrix.T
+                    + self.Q_measurement_noise_covariance
+                )
+            )
+        except LinAlgError:
+            # Apply regularization and retry
+            # A_reg = self.A_state_transition_matrix * 0.999  # Slight damping
+            # W_reg = self.W_process_noise_covariance + 1e-7 * np.eye(
+            # self.W_process_noise_covariance.shape[0]
+            # )
+            Q_reg = self.Q_measurement_noise_covariance + 1e-7 * np.eye(self.Q_measurement_noise_covariance.shape[0])
+
+            try:
+                self.P_state_covariance = solve_discrete_are(
+                    self.A_state_transition_matrix.T,
+                    self.H_observation_matrix.T,
+                    self.W_process_noise_covariance,
+                    Q_reg,
+                )
+                self.K_kalman_gain = (
+                    self.P_state_covariance
+                    @ self.H_observation_matrix.T
+                    @ np.linalg.inv(
+                        self.H_observation_matrix @ self.P_state_covariance @ self.H_observation_matrix.T + Q_reg
+                    )
+                )
+                print("Warning: Used regularized matrices for DARE solution")
+            except LinAlgError:
+                # Fallback to identity or manual initialization
+                print("Warning: DARE failed, using identity covariance")
+                self.P_state_covariance = np.eye(self.A_state_transition_matrix.shape[0])
+
+        # else:
+        #     n_states = self.A_state_transition_matrix.shape[0]
+        #     self.P_state_covariance = (
+        #         np.eye(n_states) * 1000
+        #     )  # Large initial uncertainty
+
+        #     P_m = (
+        #         self.A_state_transition_matrix
+        #         @ self.P_state_covariance
+        #         @ self.A_state_transition_matrix.T
+        #         + self.W_process_noise_covariance
+        #     )
+
+        #     S = (
+        #         self.H_observation_matrix @ P_m @ self.H_observation_matrix.T
+        #         + self.Q_measurement_noise_covariance
+        #     )
+
+        #     self.K_kalman_gain = P_m @ self.H_observation_matrix.T @ np.linalg.pinv(S)
+
+        #     I_mat = np.eye(self.A_state_transition_matrix.shape[0])
+        #     self.P_state_covariance = (
+        #         I_mat - self.K_kalman_gain @ self.H_observation_matrix
+        #     ) @ P_m
+
+    def predict(self, x_current: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Predict the next state and covariance.
+
+        This method predicts the next state and covariance using the current state.
+        """
+        x_predicted = self.A_state_transition_matrix @ x_current
+        if self.steady_state is True:
+            return x_predicted, None
+        else:
+            P_predicted = self.alpha_fading_memory**2 * (
+                self.A_state_transition_matrix @ self.P_state_covariance @ self.A_state_transition_matrix.T
+                + self.W_process_noise_covariance
+            )
+            return x_predicted, P_predicted
+
+    def update(
+        self,
+        z_measurement: np.ndarray,
+        x_predicted: np.ndarray,
+        P_predicted: np.ndarray | None = None,
+    ) -> np.ndarray:
+        """Update state estimate and covariance based on measurement z."""
+
+        # Compute residual
+        innovation = z_measurement - self.H_observation_matrix @ x_predicted
+
+        if self.steady_state:
+            x_updated = x_predicted + self.K_kalman_gain @ innovation
+            return x_updated
+
+        if P_predicted is None:
+            raise ValueError("P_predicted must be provided for non-steady-state mode")
+
+        # Non-steady-state mode
+        # System uncertainty
+        S = self.H_observation_matrix @ P_predicted @ self.H_observation_matrix.T + self.Q_measurement_noise_covariance
+
+        # Kalman gain
+        K = P_predicted @ self.H_observation_matrix.T @ np.linalg.pinv(S)
+
+        # Updated state
+        x_updated = x_predicted + K @ innovation
+
+        # Covariance update
+        I_mat = np.eye(self.A_state_transition_matrix.shape[0])
+        P_updated = (I_mat - K @ self.H_observation_matrix) @ P_predicted @ (
+            I_mat - K @ self.H_observation_matrix
+        ).T + K @ self.Q_measurement_noise_covariance @ K.T
+
+        # Save updated values
+        self.P_state_covariance = P_updated
+        self.K_kalman_gain = K
+        # self.S = S  # Optional: for diagnostics
+
+        return x_updated