diffusion hyperplane

src/scloop/plotting/__init__.py

@@ -4,6 +4,7 @@
 from ._homology import (
     bar_lifetimes,
     hist_lifetimes,
+    loop_embedding,
     loops,
     persistence_diagram,
 )
@@ -12,6 +13,7 @@
 __all__ = [
     "bar_lifetimes",
     "hist_lifetimes",
+    "loop_embedding",
     "loops",
     "loop_edge_embedding",
     "loop_edge_overlay",

src/scloop/plotting/_homology.py

@@ -8,6 +8,7 @@
 from anndata import AnnData
 from matplotlib.axes import Axes
 from pydantic import ConfigDict, validate_call
+from scipy.linalg import svd
 
 from ..data.analysis_containers import BootstrapAnalysis
 from ..data.constants import DEFAULT_DPI, DEFAULT_FIGSIZE, SCLOOP_UNS_KEY
@@ -19,9 +20,12 @@
     "hist_lifetimes",
     "bar_lifetimes",
     "persistence_diagram",
+    "loop_embedding",
     "loops",
 ]
 
+DEFAULT_GLASBEY_BLOCK_SIZE = 5
+
 
 # ugly function, fix that
 def _get_track_loop(
@@ -59,6 +63,52 @@ def _get_track_loop(
     return tracked_pairs
 
 
+def loop_embedding(
+    adata: AnnData,
+    basis: str,
+    track_ids: int | list[int],
+    ndims: int = 2,
+    key_added: str = "loops",
+    key_homology: str = SCLOOP_UNS_KEY,
+) -> None:
+    hdata: HomologyData = adata.uns[key_homology]
+    emb = np.asarray(adata.obsm[basis])
+    track_list = [track_ids] if isinstance(track_ids, int) else list(track_ids)
+
+    planes = []
+    for tid in track_list:
+        X = np.concatenate(
+            hdata._get_loop_embedding(
+                selector=tid, embedding_alt=emb, include_bootstrap=True
+            ),
+            axis=0,
+        )
+        X = X - X.mean(axis=0, keepdims=True)
+        U, _, _ = svd(X.T, full_matrices=False)
+        planes.append(U[:, :ndims])
+
+    if len(planes) == 1:
+        plane = planes[0]
+        score = 1.0
+    else:
+        """
+        Steps to figure out a consensus plane:
+           1. Given per-loop diffusion planes Us, the projection matricies (to the target diffusion plane) are U^t. We want to figure out a consensus plane P.
+           2. The goal is max sum_k(tr(P^t Uk Uk^t P)) such that the projection of P to all diffusion planes are maximized.
+               - max sum_k(tr(P^t Uk Uk^t P))
+               - max tr(P^t sum_k(Uk Uk^t) P)
+           3. Let P = [ v1 v2 ], where v1 and v2 are top eigenvecs of sum_k(Uk Uk^t)
+           4. tr(P^t sum_k(Uk Uk^t) P) = tr([ [ l1 0 ] [ 0 l2 ] ]) = l1 + l2
+        """
+        M = sum(U @ U.T for U in planes)
+        evals, evecs = np.linalg.eigh(M)  # ascending
+        plane = evecs[:, ::-1][:, :ndims]  # top ndims, descending
+        score = float(evals[::-1][:ndims].sum() / (ndims * len(planes)))
+
+    adata.obsm[f"X_{key_added}"] = (emb - emb.mean(axis=0, keepdims=True)) @ plane
+    adata.uns[f"{key_added}_loop_embedding"] = {"plane": plane, "score": score}
+
+
 @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
 def hist_lifetimes(
     adata: AnnData,
@@ -208,7 +258,7 @@ def bar_lifetimes(
     track_ids = track_ids or []
     n_tracks = len(track_ids)
     if n_tracks > 0:
-        block_size = 5
+        block_size = DEFAULT_GLASBEY_BLOCK_SIZE
         cmap = glasbey.create_block_palette(block_sizes=[block_size] * n_tracks)
         cmap = [cmap[i : i + block_size] for i in range(0, len(cmap), block_size)]
         for i, src_tid in enumerate(track_ids):
@@ -321,7 +371,7 @@ def persistence_diagram(
     track_ids = track_ids or []
     n_tracks = len(track_ids)
     if n_tracks > 0:
-        block_size = 5
+        block_size = DEFAULT_GLASBEY_BLOCK_SIZE
         cmap = glasbey.create_block_palette(block_sizes=[block_size] * n_tracks)
         cmap = [cmap[i : i + block_size] for i in range(0, len(cmap), block_size)]
         for i, src_tid in enumerate(track_ids):
@@ -431,7 +481,7 @@ def loops(
 
     n_selectors = len(selectors)
     if n_selectors > 0:
-        block_size = 5
+        block_size = DEFAULT_GLASBEY_BLOCK_SIZE
         cmap = glasbey.create_block_palette(block_sizes=[block_size] * n_selectors)
         cmap = [cmap[i : i + block_size] for i in range(0, len(cmap), block_size)]
     else:
@@ -522,7 +572,7 @@ def _loops_for_selector(
             ax.plot(
                 loop[:, components[0]],
                 loop[:, components[1]],
-                color=cmap[i][j % block_size],
+                color=cmap[i][(block_size - 1) - j % block_size],
                 **(kwargs_scatter or {}),
             )