diff --git a/src/scloop/plotting/__init__.py b/src/scloop/plotting/__init__.py index e4c7ed9..6f7c565 100644 --- a/src/scloop/plotting/__init__.py +++ b/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", diff --git a/src/scloop/plotting/_homology.py b/src/scloop/plotting/_homology.py index 7f4a345..0ab35b0 100644 --- a/src/scloop/plotting/_homology.py +++ b/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 {}), )