ezmsg-sigproc 2.13.0__py3-none-any.whl → 2.13.1__py3-none-any.whl

ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.13.0'
-__version_tuple__ = version_tuple = (2, 13, 0)
+__version__ = version = '2.13.1'
+__version_tuple__ = version_tuple = (2, 13, 1)
 
 __commit_id__ = commit_id = None

ezmsg/sigproc/affinetransform.py

@@ -215,58 +215,11 @@ class AffineTransformTransformer(
         if weights is not None:
             weights = np.ascontiguousarray(weights)
 
-        self._state.weights = weights
-
-        # Note: If weights were scipy.sparse BSR then maybe we could use automate this next part.
-        #  However, that would break compatibility with Array API.
-
-        # --- Block-diagonal cluster detection ---
-        # Clusters are a list of (input_indices, output_indices) tuples.
-        n_in, n_out = weights.shape
-        if self.settings.channel_clusters is not None:
-            # Validate input index bounds
-            all_in = np.concatenate([np.asarray(group) for group in self.settings.channel_clusters])
-            if np.any((all_in < 0) | (all_in >= n_in)):
-                raise ValueError(
-                    "channel_clusters contains out-of-range input indices " f"(valid range: 0..{n_in - 1})"
-                )
-
-            # Derive output indices from non-zero weights for each input cluster
-            clusters = []
-            for group in self.settings.channel_clusters:
-                in_idx = np.asarray(group)
-                out_idx = np.where(np.any(weights[in_idx, :] != 0, axis=0))[0]
-                clusters.append((in_idx, out_idx))
-
-            max_cross = _max_cross_cluster_weight(weights, clusters)
-            if max_cross > 0:
-                ez.logger.warning(
-                    f"Non-zero cross-cluster weights detected (max abs: {max_cross:.2e}). "
-                    "These will be ignored in block-diagonal multiplication."
-                )
-        else:
-            clusters = _find_block_diagonal_clusters(weights)
-            if clusters is not None:
-                ez.logger.info(
-                    f"Auto-detected {len(clusters)} block-diagonal clusters "
-                    f"(sizes: {[(len(i), len(o)) for i, o in clusters]})"
-                )
-
-        # Merge small clusters to avoid excessive loop overhead
-        if clusters is not None:
-            clusters = _merge_small_clusters(clusters, self.settings.min_cluster_size)
-
-        if clusters is not None and len(clusters) > 1:
-            self._state.n_out = n_out
-            self._state.clusters = [
-                (in_idx, out_idx, np.ascontiguousarray(weights[np.ix_(in_idx, out_idx)]))
-                for in_idx, out_idx in clusters
-            ]
-            self._state.weights = None
-        else:
-            self._state.clusters = None
+        # Cluster detection + weight storage (delegated)
+        self.set_weights(weights, recalc_clusters=True)
 
         # --- Axis label handling (for non-square transforms, non-cluster path) ---
+        n_in, n_out = weights.shape
         axis = self.settings.axis or message.dims[-1]
         if axis in message.axes and hasattr(message.axes[axis], "data") and n_in != n_out:
             in_labels = message.axes[axis].data
@@ -302,6 +255,80 @@ class AffineTransformTransformer(
                 for in_idx, out_idx, sub_w in self._state.clusters
             ]
 
+    def set_weights(self, weights, *, recalc_clusters=False) -> None:
+        """Replace weight values, optionally recalculating cluster decomposition.
+
+        *weights* must be in **canonical orientation** (``right_multiply``
+        already applied by the caller or by ``_reset_state``).  The array may
+        live in any Array-API namespace (NumPy, CuPy, etc.).
+
+        Args:
+            weights: Weight matrix in canonical orientation.
+            recalc_clusters: When True, re-run block-diagonal cluster detection
+                and store the new decomposition.  When False (default), reuse
+                the existing cluster structure and only update weight values.
+        """
+        if recalc_clusters:
+            # Note: If weights were scipy.sparse BSR then maybe we could automate this next part.
+            #  However, that would break compatibility with Array API.
+
+            # --- Block-diagonal cluster detection ---
+            # Clusters are a list of (input_indices, output_indices) tuples.
+            w_np = np.ascontiguousarray(weights)
+            n_in, n_out = w_np.shape
+            if self.settings.channel_clusters is not None:
+                # Validate input index bounds
+                all_in = np.concatenate([np.asarray(group) for group in self.settings.channel_clusters])
+                if np.any((all_in < 0) | (all_in >= n_in)):
+                    raise ValueError(
+                        "channel_clusters contains out-of-range input indices " f"(valid range: 0..{n_in - 1})"
+                    )
+
+                # Derive output indices from non-zero weights for each input cluster
+                clusters = []
+                for group in self.settings.channel_clusters:
+                    in_idx = np.asarray(group)
+                    out_idx = np.where(np.any(w_np[in_idx, :] != 0, axis=0))[0]
+                    clusters.append((in_idx, out_idx))
+
+                max_cross = _max_cross_cluster_weight(w_np, clusters)
+                if max_cross > 0:
+                    ez.logger.warning(
+                        f"Non-zero cross-cluster weights detected (max abs: {max_cross:.2e}). "
+                        "These will be ignored in block-diagonal multiplication."
+                    )
+            else:
+                clusters = _find_block_diagonal_clusters(w_np)
+                if clusters is not None:
+                    ez.logger.info(
+                        f"Auto-detected {len(clusters)} block-diagonal clusters "
+                        f"(sizes: {[(len(i), len(o)) for i, o in clusters]})"
+                    )
+
+            # Merge small clusters to avoid excessive loop overhead
+            if clusters is not None:
+                clusters = _merge_small_clusters(clusters, self.settings.min_cluster_size)
+
+            if clusters is not None and len(clusters) > 1:
+                self._state.n_out = n_out
+                self._state.clusters = [
+                    (in_idx, out_idx, np.ascontiguousarray(w_np[np.ix_(in_idx, out_idx)]))
+                    for in_idx, out_idx in clusters
+                ]
+                self._state.weights = None
+            else:
+                self._state.weights = weights
+                self._state.clusters = None
+        else:
+            xp = get_namespace(weights)
+            if self._state.clusters is not None:
+                self._state.clusters = [
+                    (in_idx, out_idx, xp.take(xp.take(weights, in_idx, axis=0), out_idx, axis=1))
+                    for in_idx, out_idx, _ in self._state.clusters
+                ]
+            else:
+                self._state.weights = weights
+
     def _block_diagonal_matmul(self, xp, data, axis_idx):
         """Perform matmul using block-diagonal decomposition.
 

{ezmsg_sigproc-2.13.0.dist-info → ezmsg_sigproc-2.13.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.13.0
+Version: 2.13.1
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

{ezmsg_sigproc-2.13.0.dist-info → ezmsg_sigproc-2.13.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 ezmsg/sigproc/__init__.py,sha256=8K4IcOA3-pfzadoM6s2Sfg5460KlJUocGgyTJTJl96U,52
-ezmsg/sigproc/__version__.py,sha256=_4LOjlEcfZzfuqIlglDZmVBPO4LyQ8P97qO716YoUL8,706
+ezmsg/sigproc/__version__.py,sha256=LrOgVsOxSmuj7RbXPLo3yIvC77lH9VlW4tk7Ihs46rY,706
 ezmsg/sigproc/activation.py,sha256=83vnTa3ZcC4Q3VSWcGfaqhCEqYRNySUOyVpMHZXfz-c,2755
 ezmsg/sigproc/adaptive_lattice_notch.py,sha256=ThUR48mbSHuThkimtD0j4IXNMrOVcpZgGhE7PCYfXhU,8818
-ezmsg/sigproc/affinetransform.py,sha256=ZugiQg89Ly1I9SDgf0ZzgU2XdwVDmPrU7-orO9yrt7w,20210
+ezmsg/sigproc/affinetransform.py,sha256=mjA21DRVYm0kS2tK7dNR_mU5XAxDbJGuuhnuzz0gtw4,21679
 ezmsg/sigproc/aggregate.py,sha256=7Hdz1m-S6Cl9h0oRQHeS_UTGBemhOB4XdFyX6cGcdHo,9362
 ezmsg/sigproc/bandpower.py,sha256=dAhH56sUrXNhcRFymTTwjdM_KcU5OxFzrR_sxIPAxyw,2264
 ezmsg/sigproc/base.py,sha256=SJvKEb8gw6mUMwlV5sH0iPG0bXrgS8tvkPwhI-j89MQ,3672
@@ -62,7 +62,7 @@ ezmsg/sigproc/util/message.py,sha256=ppN3IYtIAwrxWG9JOvgWFn1wDdIumkEzYFfqpH9VQkY
 ezmsg/sigproc/util/profile.py,sha256=eVOo9pXgusrnH1yfRdd2RsM7Dbe2UpyC0LJ9MfGpB08,416
 ezmsg/sigproc/util/sparse.py,sha256=NjbJitCtO0B6CENTlyd9c-lHEJwoCan-T3DIgPyeShw,4834
 ezmsg/sigproc/util/typeresolution.py,sha256=fMFzLi63dqCIclGFLcMdM870OYxJnkeWw6aWKNMk718,362
-ezmsg_sigproc-2.13.0.dist-info/METADATA,sha256=RXENX541lABAic8oUDuT8vQwx9nlWY9JETyXYKxdeTQ,1909
-ezmsg_sigproc-2.13.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ezmsg_sigproc-2.13.0.dist-info/licenses/LICENSE,sha256=seu0tKhhAMPCUgc1XpXGGaCxY1YaYvFJwqFuQZAl2go,1100
-ezmsg_sigproc-2.13.0.dist-info/RECORD,,
+ezmsg_sigproc-2.13.1.dist-info/METADATA,sha256=mjyCiq1zy3JbgayLzq3BQLB-i1LOouUeGHe0tn3v7MY,1909
+ezmsg_sigproc-2.13.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ezmsg_sigproc-2.13.1.dist-info/licenses/LICENSE,sha256=seu0tKhhAMPCUgc1XpXGGaCxY1YaYvFJwqFuQZAl2go,1100
+ezmsg_sigproc-2.13.1.dist-info/RECORD,,