When using level-specific colormaps, sort cluster names to assign colors

complextissue · Sep 30, 2024 · beab0ff · beab0ff
1 parent 6353065
commit beab0ff
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/docs/source/example.ipynb b/docs/source/example.ipynb
diff --git a/pyclustree/_clustree.py b/pyclustree/_clustree.py
@@ -125,6 +125,7 @@ def clustree(
     ]
 
     unique_clusters = [np.unique(df_cluster_assignments[key]).tolist() for key in cluster_keys]
+    unique_clusters_sorted = [sorted(unique_clusters_level) for unique_clusters_level in unique_clusters]
 
     if order_clusters:
         unique_clusters = order_unique_clusters(unique_clusters, transition_matrices)
@@ -211,7 +212,13 @@ def clustree(
                 plt.cm.get_cmap(node_colormap[i]) if isinstance(node_colormap[i], str) else node_colormap[i]
             )
             norm_level = plt.Normalize(vmin=0, vmax=len(unique_clusters[i]) - 1)
-            node_colors.extend([node_colormap_level(norm_level(j)) for j in range(len(unique_clusters[i]))])
+
+            node_colors.extend(
+                [
+                    node_colormap_level(norm_level(unique_clusters_sorted[i].index(unique_cluster)))
+                    for unique_cluster in unique_clusters[i]
+                ]
+            )
 
     # If a node color gene is provided, use the expression of the gene to color the nodes
     if node_color_gene is not None:

diff --git a/pyclustree/_utils.py b/pyclustree/_utils.py
@@ -40,7 +40,10 @@ def get_centered_positions(
     x_start = -(n_nodes - 1) / 2 * x_spacing  # Center the nodes on the x-axis
     pos = {}
     for i, node in enumerate(nodes):
-        pos[node] = (x_start + i * x_spacing, -level * y_spacing)  # Calculate and assign position
+        pos[node] = (
+            x_start + i * x_spacing,
+            -level * y_spacing,
+        )  # Calculate and assign position
     return pos