[WIP] Enhance distance and gpf metrics

pycrostates/cluster/_base.py

@@ -14,7 +14,7 @@
 from scipy.signal import convolve2d
 
 from ..segmentation import EpochsSegmentation, RawSegmentation
-from ..utils import _corr_vectors
+from ..utils import _correlation
 from ..utils._checks import (
     _check_picks_uniqueness,
     _check_reject_by_annotation,
@@ -820,6 +820,7 @@ def _predict_raw(
             "min_segment_length": min_segment_length,
             "reject_edges": reject_edges,
             "reject_by_annotation": reject_by_annotation,
+            "ignore_polarity": self._ignore_polarity,
         }
 
         # retrieve data for picks
@@ -839,22 +840,32 @@ def _predict_raw(
 
                 data_ = data[:, onset:end]
                 segment = _BaseCluster._segment(
-                    data_, cluster_centers_, factor, tol, half_window_size
+                    data_,
+                    cluster_centers_,
+                    self._ignore_polarity,
+                    factor,
+                    tol,
+                    half_window_size,
                 )
                 if reject_edges:
                     segment = _BaseCluster._reject_edge_segments(segment)
                 segmentation[onset:end] = segment
 
         else:
             segmentation = _BaseCluster._segment(
-                data, cluster_centers_, factor, tol, half_window_size
+                data,
+                cluster_centers_,
+                self._ignore_polarity,
+                factor,
+                tol,
+                half_window_size,
             )
             if reject_edges:
                 segmentation = _BaseCluster._reject_edge_segments(segmentation)
 
         if 0 < min_segment_length:
             segmentation = _BaseCluster._reject_short_segments(
-                segmentation, data, min_segment_length
+                segmentation, data, min_segment_length, self._ignore_polarity
             )
 
         # Provide properties to copy the arrays
@@ -883,6 +894,7 @@ def _predict_epochs(
             "half_window_size": half_window_size,
             "min_segment_length": min_segment_length,
             "reject_edges": reject_edges,
+            "ignore_polarity": self._ignore_polarity,
         }
 
         # retrieve data for picks
@@ -893,12 +905,17 @@ def _predict_epochs(
         segments = []
         for epoch_data in data:
             segment = _BaseCluster._segment(
-                epoch_data, cluster_centers_, factor, tol, half_window_size
+                epoch_data,
+                cluster_centers_,
+                self._ignore_polarity,
+                factor,
+                tol,
+                half_window_size,
             )
 
             if 0 < min_segment_length:
                 segment = _BaseCluster._reject_short_segments(
-                    segment, epoch_data, min_segment_length
+                    segment, epoch_data, min_segment_length, self._ignore_polarity
                 )
             if reject_edges:
                 segment = _BaseCluster._reject_edge_segments(segment)
@@ -919,33 +936,29 @@ def _predict_epochs(
     def _segment(
         data: ScalarFloatArray,
         states: ScalarFloatArray,
+        ignore_polarity: bool,
         factor: int,
         tol: int | float,
         half_window_size: int,
     ) -> ScalarIntArray:
         """Create segmentation. Must operate on a copy of states."""
-        data -= np.mean(data, axis=0)
-        std = np.std(data, axis=0)
-        std[std == 0] = 1  # std == 0 -> null map
-        data /= std
-
-        states -= np.mean(states, axis=1)[:, np.newaxis]
-        states /= np.std(states, axis=1)[:, np.newaxis]
-
-        labels = np.argmax(np.abs(np.dot(states, data)), axis=0)
+        corr = np.zeros((states.shape[0], data.shape[1]))
+        for k in range(0, states.shape[0]):
+            corr[k] = _correlation(data, states[k], ignore_polarity=ignore_polarity)
+        labels = np.argmax(corr, axis=0)
 
         if factor != 0:
             labels = _BaseCluster._smooth_segmentation(
-                data, states, labels, factor, tol, half_window_size
+                data, states, ignore_polarity, labels, factor, tol, half_window_size
             )
-
         return labels
 
     @staticmethod
     def _smooth_segmentation(
-        data: ScalarFloatArray,
-        states: ScalarFloatArray,
-        labels: ScalarIntArray,
+        data: NDArray[float],
+        states: NDArray[float],
+        ignore_polarity: bool,
+        labels: NDArray[int],
         factor: int,
         tol: int | float,
         half_window_size: int,
@@ -963,6 +976,8 @@ def _smooth_segmentation(
                vol. 42, no. 7, pp. 658-665, July 1995,
                https://doi.org/10.1109/10.391164.
         """
+        if not ignore_polarity:
+            raise NotImplementedError("Smoothing with polarity is not implemented yet.")
         Ne, Nt = data.shape
         Nu = states.shape[0]
         Vvar = np.sum(data * data, axis=0)
@@ -1000,7 +1015,8 @@ def _reject_short_segments(
         segmentation: ScalarIntArray,
         data: ScalarFloatArray,
         min_segment_length: int,
-    ) -> ScalarIntArray:
+        ignore_polarity: bool,
+    ) -> NDArray[int]:
         """Reject segments that are too short.
 
         Reject segments that are too short by replacing the labels with the adjacent
@@ -1027,16 +1043,16 @@ def _reject_short_segments(
 
                 while len(new_segment) != 0:
                     # compute correlation left/right side
-                    left_corr = np.abs(
-                        _corr_vectors(
-                            data[:, left - 1].T,
-                            data[:, left].T,
-                        )
-                    )
-                    right_corr = np.abs(
-                        _corr_vectors(data[:, right].T, data[:, right + 1].T)
-                    )
-
+                    left_corr = _correlation(
+                        data[:, left - 1],
+                        data[:, left],
+                        ignore_polarity=ignore_polarity,
+                    )[0]
+                    right_corr = _correlation(
+                        data[:, right],
+                        data[:, right + 1],
+                        ignore_polarity=ignore_polarity,
+                    )[0]
                     if np.abs(right_corr - left_corr) <= 1e-8:
                         # equal corr, try to do both sides
                         if len(new_segment) == 1:
@@ -1196,3 +1212,9 @@ def _check_n_clusters(n_clusters: int) -> int:
                 f"Provided: '{n_clusters}'."
             )
         return n_clusters
+
+    @staticmethod
+    def _check_ignore_polarity(ignore_polarity: bool) -> bool:
+        """Check that ignore_polarity is a boolean."""
+        _check_type(ignore_polarity, (bool,), item_name="ignore_polarity")
+        return ignore_polarity

pycrostates/cluster/aahc.py

@@ -177,8 +177,6 @@ def save(self, fname: str | Path):
             GEV_=self._GEV_,
         )
 
-    # --------------------------------------------------------------------
-    @staticmethod
     def _aahc(
         data: ScalarFloatArray,
         n_clusters: int,
@@ -193,6 +191,7 @@ def _aahc(
         gev = _compute_gev(data, maps, segmentation, ch_type)
         return gev, maps, segmentation
 
+    # --------------------------------------------------------------------
     # pylint: disable=too-many-locals
     @staticmethod
     def _compute_maps(
@@ -276,12 +275,6 @@ def fitted(self, fitted):
         if not fitted:
             self._GEV_ = None
 
-    @staticmethod
-    def _check_ignore_polarity(ignore_polarity: bool) -> bool:
-        """Check that ignore_polarity is a boolean."""
-        _check_type(ignore_polarity, (bool,), item_name="ignore_polarity")
-        return ignore_polarity
-
     @staticmethod
     def _check_normalize_input(normalize_input: bool) -> bool:
         """Check that normalize_input is a boolean."""

pycrostates/cluster/kmeans.py

@@ -11,9 +11,13 @@
 from mne.parallel import parallel_func
 from numpy.random import Generator, RandomState
 
-from pycrostates.utils.utils import _compute_gev
-
-from ..utils._checks import _check_n_jobs, _check_random_state, _check_type
+from .._typing import Picks
+from ..utils import _gev
+from ..utils._checks import (
+    _check_n_jobs,
+    _check_random_state,
+    _check_type,
+)
 from ..utils._docs import copy_doc, fill_doc
 from ..utils._logs import logger
 from ._base import _BaseCluster
@@ -74,6 +78,9 @@ def __init__(
         # fit variables
         self._GEV_ = None
 
+        # ignore polarity
+        self._ignore_polarity = True
+
     def _repr_html_(self, caption=None):
         from ..html_templates import repr_templates_env
 
@@ -117,6 +124,7 @@ def __eq__(self, other: Any) -> bool:
                 # '_random_state',
                 # TODO: think about comparison and I/O for random states
                 "_GEV_",
+                "_ignore_polarity",
             )
             for attribute in attributes:
                 try:
@@ -194,7 +202,8 @@ def fit(
                 gev, maps, segmentation, converged = ModKMeans._kmeans(
                     data,
                     self._n_clusters,
-                    self.get_channel_types(unique=True)[0],
+                    self._ignore_polarity,
+                    self.get_channel_types()[0],
                     self._max_iter,
                     init,
                     self._tol,
@@ -216,7 +225,8 @@ def fit(
                 p_fun(
                     data,
                     self._n_clusters,
-                    self.get_channel_types(unique=True)[0],
+                    self._ignore_polarity,
+                    self.get_channel_types()[0],
                     self._max_iter,
                     init,
                     self._tol,
@@ -251,7 +261,6 @@ def fit(
         self._cluster_centers_ = best_maps
         self._labels_ = best_segmentation
         self._fitted = True
-        self._ignore_polarity = True
 
     @copy_doc(_BaseCluster.save)
     def save(self, fname: str | Path):
@@ -271,32 +280,33 @@ def save(self, fname: str | Path):
             n_init=self._n_init,
             max_iter=self._max_iter,
             tol=self._tol,
+            ignore_polarity=self._ignore_polarity,
             GEV_=self._GEV_,
         )
 
-    # --------------------------------------------------------------------
-    @staticmethod
     def _kmeans(
         data: ScalarFloatArray,
         n_clusters: int,
+        ignore_polarity: bool,
         ch_type: str,
         max_iter: int,
         random_state: RandomState | Generator,
         tol: int | float,
     ) -> tuple[float, ScalarFloatArray, ScalarIntArray, bool]:
         """Run the k-means algorithm."""
-        maps, converged = ModKMeans._compute_maps(
-            data, n_clusters, max_iter, random_state, tol
+        maps, segmentation, converged = ModKMeans._compute_maps(
+            data, n_clusters, ignore_polarity, max_iter, random_state, tol
         )
-        activation = maps.dot(data)
-        segmentation = np.argmax(np.abs(activation), axis=0)
-        gev = _compute_gev(data, maps, segmentation, ch_type)
+        # Compute GEV
+        gev = _gev(data, maps, segmentation, ch_type=ch_type)
         return gev, maps, segmentation, converged
 
+    # --------------------------------------------------------------------
     @staticmethod
     def _compute_maps(
         data: ScalarFloatArray,
         n_clusters: int,
+        ignore_polarity: bool,
         max_iter: int,
         random_state: RandomState | Generator,
         tol: int | float,
@@ -329,7 +339,10 @@ def _compute_maps(
         for _ in range(max_iter):
             # Assign each sample to the best matching microstate
             activation = maps.dot(data)
-            segmentation = np.argmax(np.abs(activation), axis=0)
+            if ignore_polarity:
+                segmentation = np.argmax(np.abs(activation), axis=0)
+            else:
+                segmentation = np.argmax(activation, axis=0)
 
             # Recompute the topographic maps of the microstates, based on the
             # samples that were assigned to each state.
@@ -358,7 +371,7 @@ def _compute_maps(
         else:
             converged = False
 
-        return maps, converged
+        return maps, segmentation, converged
 
     # --------------------------------------------------------------------
     @property

pycrostates/cluster/tests/test_base.py

@@ -40,7 +40,7 @@ def test_reject_short_segments():
             [3, 3, 3, 3, 3, 3, 3, 6, 4, 5, 2, 2, 2],
         ]
     )
-    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3)
+    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3, True)
     # solo 1 should turn to 2; initial 0 should not change
     assert [0, 0, 1, 1, 1, 3, 3, 3, 2, 2, 2, 2, 2] == segmentation
 
@@ -53,7 +53,7 @@ def test_reject_short_segments():
             [3, 3, 3, 3, 3, 3, 3, 6, 4, 4, 6, 2, 2, 2],
         ]
     )
-    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3)
+    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3, True)
     assert [0, 0, 1, 1, 1, 3, 3, 3, 3, 2, 2, 2, 2, 2] == segmentation
 
     # singleton, same correlation
@@ -65,7 +65,32 @@ def test_reject_short_segments():
             [3, 3, 3, 3, 3, 3, 3, 6, 4, 6, 2, 2, 2],
         ]
     )
-    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3)
+    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3, True)
+    assert [0, 0, 1, 1, 1, 3, 3, 3, 3, 2, 2, 2, 2] == segmentation
+
+    # ignore polarity
+    segmentation = [0, 0, 1, 1, 1, 3, 3, 3, 1, 2, 2, 2, 2]
+    data = np.array(
+        [
+            [1, 1, 1, 1, 1, 1, 1, 1, -1, -1, 1, 1, 1],
+            [2, 2, 2, 2, 2, 2, 3, 0.5, -0.5, -0.5, 1, 1, 1],
+            [3, 3, 3, 3, 3, 3, 3, 4, -4, -4, 2, 2, 2],
+        ]
+    )
+    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3, False)
+    # solo 1 should turn to 2; initial 0 should not change
+    assert [0, 0, 1, 1, 1, 3, 3, 3, 2, 2, 2, 2, 2] == segmentation
+
+    segmentation = [0, 0, 1, 1, 1, 3, 3, 3, 1, 2, 2, 2, 2]
+    data = np.array(
+        [
+            [1, 1, 1, 1, 1, 1, 1, -1, -1, 1, 1, 1, 1],
+            [2, 2, 2, 2, 2, 2, 3, -0.5, -0.5, 0.5, 1, 1, 1],
+            [3, 3, 3, 3, 3, 3, 3, -4, -4, 4, 2, 2, 2],
+        ]
+    )
+    segmentation = _BaseCluster._reject_short_segments(segmentation, data, 3, False)
+    # solo 1 should turn to 2; initial 0 should not change
     assert [0, 0, 1, 1, 1, 3, 3, 3, 3, 2, 2, 2, 2] == segmentation