Make as much processing as possible use Array API

Author: cboulay

src/ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -1,7 +1,18 @@
+"""Adaptive decomposition transformers (PCA, NMF).
+
+.. note::
+    This module supports the Array API standard via
+    ``array_api_compat.get_namespace()``.  Reshaping and output allocation
+    use Array API operations; a NumPy boundary is applied before sklearn
+    ``partial_fit``/``transform`` calls.
+"""
+
+import math
 import typing
 
 import ezmsg.core as ez
 import numpy as np
+from array_api_compat import get_namespace, is_numpy_array
 from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
@@ -128,6 +139,8 @@ def _process(self, message: AxisArray) -> AxisArray:
         if in_dat.shape[ax_idx] == 0:
             return self._state.template
 
+        xp = get_namespace(in_dat)
+
         # Re-order axes
         sorted_dims_exp = [iter_axis] + off_targ_axes + targ_axes
         if message.dims != sorted_dims_exp:
@@ -137,16 +150,20 @@ def _process(self, message: AxisArray) -> AxisArray:
             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))
+        d2 = math.prod(in_dat.shape[len(off_targ_axes) + 1 :])
+        in_dat = xp.reshape(in_dat, (-1, d2))
 
         replace_kwargs = {
             "axes": {**self._state.template.axes, iter_axis: message.axes[iter_axis]},
         }
 
-        # Transform data
+        # Transform data — sklearn needs numpy
         if hasattr(self._state.estimator, "components_"):
-            decomp_dat = self._state.estimator.transform(in_dat).reshape((-1,) + self._state.template.data.shape[1:])
+            in_np = np.asarray(in_dat) if not is_numpy_array(in_dat) else in_dat
+            decomp_dat = self._state.estimator.transform(in_np)
+            # Convert back to source namespace
+            decomp_dat = xp.asarray(decomp_dat) if not is_numpy_array(in_dat) else decomp_dat
+            decomp_dat = xp.reshape(decomp_dat, (-1,) + self._state.template.data.shape[1:])
             replace_kwargs["data"] = decomp_dat
 
         return replace(self._state.template, **replace_kwargs)
@@ -165,18 +182,21 @@ def partial_fit(self, message: AxisArray) -> None:
         if in_dat.shape[ax_idx] == 0:
             return
 
+        xp = get_namespace(in_dat)
+
         # 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))
+        d2 = math.prod(in_dat.shape[len(off_targ_axes) + 1 :])
+        in_dat = xp.reshape(in_dat, (-1, d2))
 
-        # Fit the estimator
-        self._state.estimator.partial_fit(in_dat)
+        # Fit the estimator — sklearn needs numpy
+        in_np = np.asarray(in_dat) if not is_numpy_array(in_dat) else in_dat
+        self._state.estimator.partial_fit(in_np)
 
 
 class IncrementalPCASettings(AdaptiveDecompSettings):

src/ezmsg/learn/model/cca.py

@@ -1,5 +1,29 @@
+"""Incremental Canonical Correlation Analysis (CCA).
+
+.. note::
+    This module supports the Array API standard via
+    ``array_api_compat.get_namespace()``.  All linear algebra uses Array API
+    operations; ``scipy.linalg.sqrtm`` is replaced by an eigendecomposition-
+    based inverse square root (:func:`_inv_sqrtm_spd`).
+"""
+
 import numpy as np
-from scipy import linalg
+from array_api_compat import get_namespace
+from ezmsg.sigproc.util.array import array_device, xp_create
+
+
+def _inv_sqrtm_spd(xp, A):
+    """Inverse matrix square root for symmetric positive-definite matrices.
+
+    Computes ``inv(sqrtm(A)) = Q @ diag(1/sqrt(lambda)) @ Q^T`` using the
+    eigendecomposition.  This is more numerically stable than computing
+    ``inv(sqrtm(...))`` separately and uses only Array API operations.
+    """
+    eigenvalues, eigenvectors = xp.linalg.eigh(A)
+    eigenvalues = xp.clip(eigenvalues, 1e-12, None)  # avoid div-by-zero
+    inv_sqrt_eig = 1.0 / xp.sqrt(eigenvalues)
+    # Q @ diag(v) == Q * v (broadcasting), then @ Q^T
+    return (eigenvectors * inv_sqrt_eig) @ xp.linalg.matrix_transpose(eigenvectors)
 
 
 class IncrementalCCA:
@@ -33,58 +57,74 @@ def __init__(
         self.adaptation_rate = adaptation_rate
         self.initialized = False
 
-    def initialize(self, d1, d2):
-        """Initialize the necessary matrices"""
+    def initialize(self, d1, d2, *, ref_array=None):
+        """Initialize the necessary matrices.
+
+        Args:
+            d1: Dimensionality of the first dataset.
+            d2: Dimensionality of the second dataset.
+            ref_array: Optional reference array to derive array namespace
+                and device from.  If ``None``, defaults to NumPy.
+        """
         self.d1 = d1
         self.d2 = d2
 
+        if ref_array is not None:
+            xp = get_namespace(ref_array)
+            dev = array_device(ref_array)
+        else:
+            xp, dev = np, None
+
         # Initialize correlation matrices
-        self.C11 = np.zeros((d1, d1))
-        self.C22 = np.zeros((d2, d2))
-        self.C12 = np.zeros((d1, d2))
+        self.C11 = xp_create(xp.zeros, (d1, d1), dtype=xp.float64, device=dev)
+        self.C22 = xp_create(xp.zeros, (d2, d2), dtype=xp.float64, device=dev)
+        self.C12 = xp_create(xp.zeros, (d1, d2), dtype=xp.float64, device=dev)
 
         self.initialized = True
 
     def _compute_change_magnitude(self, C11_new, C22_new, C12_new):
-        """Compute magnitude of change in correlation structure"""
+        """Compute magnitude of change in correlation structure."""
+        xp = get_namespace(self.C11)
+
         # 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)
+        diff11 = xp.linalg.matrix_norm(C11_new - self.C11)
+        diff22 = xp.linalg.matrix_norm(C22_new - self.C22)
+        diff12 = xp.linalg.matrix_norm(C12_new - self.C12)
 
         # Normalize by matrix sizes
-        diff11 /= self.d1 * self.d1
-        diff22 /= self.d2 * self.d2
-        diff12 /= self.d1 * self.d2
+        diff11 = diff11 / (self.d1 * self.d1)
+        diff22 = diff22 / (self.d2 * self.d2)
+        diff12 = diff12 / (self.d1 * self.d2)
 
-        return (diff11 + diff22 + diff12) / 3
+        return float((diff11 + diff22 + diff12) / 3)
 
     def _adapt_smoothing(self, change_magnitude):
-        """Adapt smoothing factor based on detected changes"""
+        """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
-        )
+        target_smoothing = max(self.min_smoothing, min(target_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"""
+        """Update the model with new samples using adaptive smoothing.
+        Assumes X1 and X2 are already centered and scaled."""
+        xp = get_namespace(X1, X2)
+        _mT = xp.linalg.matrix_transpose
+
         if not self.initialized:
-            self.initialize(X1.shape[1], X2.shape[1])
+            self.initialize(X1.shape[1], X2.shape[1], ref_array=X1)
 
         # 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]
+        C11_new = _mT(X1) @ X1 / X1.shape[0]
+        C22_new = _mT(X2) @ X2 / X2.shape[0]
+        C12_new = _mT(X1) @ X2 / X1.shape[0]
 
         # Detect changes and adapt smoothing factor
-        if self.C11.any():  # Skip first update
+        if bool(xp.any(self.C11 != 0)):  # Skip first update
             change_magnitude = self._compute_change_magnitude(C11_new, C22_new, C12_new)
             self._adapt_smoothing(change_magnitude)
 
@@ -98,25 +138,26 @@ def partial_fit(self, X1, X2, update_projections=True):
             self._update_projections()
 
     def _update_projections(self):
-        """Update canonical vectors and correlations"""
+        """Update canonical vectors and correlations."""
+        xp = get_namespace(self.C11)
+        dev = array_device(self.C11)
+        _mT = xp.linalg.matrix_transpose
+
         eps = 1e-8
-        C11_reg = self.C11 + eps * np.eye(self.d1)
-        C22_reg = self.C22 + eps * np.eye(self.d2)
+        C11_reg = self.C11 + eps * xp_create(xp.eye, self.d1, dtype=self.C11.dtype, device=dev)
+        C22_reg = self.C22 + eps * xp_create(xp.eye, self.d2, dtype=self.C22.dtype, device=dev)
+
+        inv_sqrt_C11 = _inv_sqrtm_spd(xp, C11_reg)
+        inv_sqrt_C22 = _inv_sqrtm_spd(xp, C22_reg)
 
-        K = (
-            linalg.inv(linalg.sqrtm(C11_reg))
-            @ self.C12
-            @ linalg.inv(linalg.sqrtm(C22_reg))
-        )
-        U, self.correlations_, V = linalg.svd(K)
+        K = inv_sqrt_C11 @ self.C12 @ inv_sqrt_C22
+        U, self.correlations_, Vh = xp.linalg.svd(K, full_matrices=False)
 
-        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]
-        )
+        self.x_weights_ = inv_sqrt_C11 @ U[:, : self.n_components]
+        self.y_weights_ = inv_sqrt_C22 @ _mT(Vh)[:, : self.n_components]
 
     def transform(self, X1, X2):
-        """Project data onto canonical components"""
+        """Project data onto canonical components."""
         X1_proj = X1 @ self.x_weights_
         X2_proj = X2 @ self.y_weights_
         return X1_proj, X2_proj