Merge pull request #3 from dschaub95/main

Add spatial stability and more clustering algos

Author: dschaub95

poetry.lock

@@ -21,7 +21,8 @@ dependencies = [
     "jupyter",
     "leidenalg",
     "igraph",
-    "harmonypy (>=0.0.10,<0.0.11)"
+    "harmonypy (>=0.0.10,<0.0.11)",
+    "louvain (>=0.8.2,<0.9.0)"
 ]
 
 

pyproject.toml

@@ -11,6 +11,8 @@
     homogeneity_score,
     completeness_score,
 )
+from sklearn.cluster import KMeans
+from scipy.sparse import issparse
 
 from scale.config import Config
 
@@ -21,11 +23,13 @@ def calc_clusterings(
     n_jobs=20,
     ensure_unique=True,
     emb_prefix="X_gnn",
+    method="leiden",
     **kwargs,
 ):
     resolutions = np.arange(
         cfg.resolution_set.start, cfg.resolution_set.stop, cfg.resolution_set.step
     ).round(4)
+    n_repeats = 1 if cfg.stability_spatial else cfg.n_repeats
 
     all_clusterings = pd.DataFrame(index=adata.obs_names)
 
@@ -47,11 +51,12 @@ def calc_clusterings(
         sparam_str = "dist" if "dist" in emb_key else "knn"
         sparam = emb_key.split(f"{sparam_str}_")[-1].split("_lam")[0]
 
-        for i in tqdm(range(cfg.n_repeats), desc="Calculating clusterings"):
-            parallel_leiden(
+        for i in tqdm(range(n_repeats), desc="Calculating clusterings"):
+            parallel_clustering(
                 ad_tmp,
                 resolutions,
                 key_added=f"leiden_rep_{i}_{sparam_str}_{sparam}",
+                method=method,
                 n_jobs=n_jobs,
                 verbose=kwargs.get("verbose", False),
                 random_state=i,
@@ -124,6 +129,71 @@ def loop(r, adata):
     return adata
 
 
+def parallel_clustering(
+    adata,
+    resolutions,
+    method="leiden",
+    key_added="scale",
+    n_jobs=10,
+    verbose=True,
+    random_state=0,
+    **kwargs,
+):
+    def to_key(r):
+        return key_added + "_res_" + str(r)
+
+    def loop(adata, r=None, **kwargs):
+        if method == "leiden":
+            flavor = kwargs.pop("flavor", "igraph")
+            n_iterations = kwargs.pop("n_iterations", 2)
+            key = to_key(r)
+            sc.tl.leiden(
+                adata,
+                resolution=r,
+                key_added=key,
+                random_state=random_state,
+                flavor=flavor,
+                n_iterations=n_iterations,
+                **kwargs,
+            )
+            if verbose:
+                print(f"Resolution = {r} Done!")
+            return adata.obs[key]
+        elif method == "louvain":
+            key = to_key(r)
+            sc.tl.louvain(
+                adata,
+                resolution=r,
+                key_added=key,
+                random_state=random_state,
+                flavor="vtraag",
+            )
+            if verbose:
+                print(f"Resolution = {r} Done!")
+            return adata.obs[key]
+        elif method == "kmeans":
+            k = int(r)
+            X = adata.obsm["X_pca"] if "X_pca" in adata.obsm else adata.X
+            X = X.A if issparse(X) else X
+            km = KMeans(n_clusters=k, random_state=random_state)
+            labels = km.fit_predict(X)
+            key = to_key(k)
+            adata.obs[key] = pd.Categorical(labels.astype(str))
+            if verbose:
+                print(f"K value = {k} Done!")
+            return adata.obs[key]
+        else:
+            raise ValueError(f"Invalid method: {method}")
+
+    clusterings = Parallel(n_jobs=n_jobs)(
+        delayed(loop)(adata, r, **kwargs) for r in resolutions
+    )
+    for clustering in clusterings:
+        adata.obs[clustering.name] = clustering
+
+    return adata
+
+
 def calc_cluster_metrics(
     labels_true,
     labels_pred,

scale/clustering.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Literal
 
+
 class BaseConfig(dict):
     """Dict that also supports attribute access (recursively)."""
 
@@ -66,14 +67,18 @@ class Config(BaseConfig):
     n_heads: int = 5
     max_epoch: int = 500
     lr: float = 0.01
-    n_sample: int = None # number of maximum edges in case of distance graph (randomyl selected)
+    n_sample: int = (
+        None  # number of maximum edges in case of distance graph (randomyl selected)
+    )
     sample_key: str = None
     preprocess: bool = False
     device: str | None = None
     distance_set: dict | list = field(
         default_factory=lambda: {"start": 15, "stop": 60, "step": 5}
     )
-    knn_set: dict | list = field(default_factory=lambda: {"start": 5, "stop": 40, "step": 5})
+    knn_set: dict | list = field(
+        default_factory=lambda: {"start": 5, "stop": 40, "step": 5}
+    )
     lambda_set: list = field(
         default_factory=lambda: [
             1e-6,
@@ -100,6 +105,8 @@ class Config(BaseConfig):
     spatial_graph_method: Literal["distance", "knn"] = "distance"
     repeated_negative_sampling: bool = False
     y_aggregated: bool = False
+    stability_spatial: bool = False
+    stability_delta: float = 1.0
 
 
 def load_config(**kwargs):

scale/config.py

@@ -107,12 +107,14 @@ def run_scale(
         calc_clusterings(
             ad_tmp,
             cfg=cfg,
+            method=kwargs.get("method", "leiden"),
             flavor=kwargs.get("flavor", "igraph"),
             n_iterations=kwargs.get("n_iterations", 2),
         )
 
         calc_stability(
             ad_tmp,
+            cfg=cfg,
             verbose=kwargs.get("verbose", True),
             n_repeat=kwargs.get("n_repeat", 4),
             min_dist=kwargs.get("min_dist", 15),
@@ -122,7 +124,7 @@ def run_scale(
             min_res=kwargs.get("min_res", None),
             max_res=kwargs.get("max_res", None),
         )
-        results = calc_entropy(
+        calc_entropy(
             ad_tmp,
             n_levels=kwargs.get("n_levels", 2),
             top_n=kwargs.get("top_n", 0.15),