Merge pull request #644 from Sichao25/pl

Debug `.pl.nneighbors()`
aristoteleo · Feb 20, 2024 · f54f524 · f54f524
2 parents 806dbe7 + 92a3550
commit f54f524
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/dynamo/plot/connectivity.py b/dynamo/plot/connectivity.py
@@ -221,6 +221,7 @@ def connectivity_base(
     available_themes = [
         "blue",
         "red",
+        "glasbey_dark",
         "green",
         "inferno",
         "fire",
@@ -288,8 +289,8 @@ def connectivity_base(
             color_key,
             color_key_cmap,
             None,
-            figsize[0] * dpi,
-            figsize[1] * dpi,
+            figsize[0],
+            figsize[1],
             True,
             sort=sort,
         )
@@ -525,22 +526,25 @@ def nneighbors(
                 i += 1
 
                 # if highligts is a list of lists - each list is relate to each color element
-                if is_list_of_lists(highlights):
-                    _highlights = highlights[color.index(cur_c)]
-                    _highlights = _highlights if all([i in _color for i in _highlights]) else None
+                if highlights is not None:
+                    if is_list_of_lists(highlights):
+                        _highlights = highlights[color.index(cur_c)]
+                        _highlights = _highlights if all([i in _color for i in _highlights]) else None
+                    else:
+                        _highlights = highlights if all([i in _color for i in highlights]) else None
                 else:
-                    _highlights = highlights if all([i in _color for i in highlights]) else None
+                    _highlights = None
 
-                connectivity_base(
+                ax = connectivity_base(
                     x_,
                     y_,
                     edge_df,
+                    _highlights,
                     edge_bundling,
                     edge_cmap,
                     show_points,
                     labels,
                     values,
-                    _highlights,
                     theme,
                     cmap,
                     color_key,
@@ -556,7 +560,8 @@ def nneighbors(
                 ax.set_ylabel(cur_b + "_2")
                 ax.set_title(cur_c)
 
-    return save_show_ret("nneighbors", save_show_or_return, save_kwargs, g)
+    return save_show_ret("nneighbors", save_show_or_return, save_kwargs, plt.gcf())
+
 
 
 def pgraph():

diff --git a/dynamo/plot/utils.py b/dynamo/plot/utils.py
@@ -930,8 +930,8 @@ def _datashade_points(
                 )
                 reorder_data = data.copy(deep=True)
                 (reorder_data.iloc[: sum(background_ids), :], reorder_data.iloc[sum(background_ids) :, :],) = (
-                    data.iloc[background_ids, :],
-                    data.iloc[highlight_ids, :],
+                    data.loc[background_ids, :],
+                    data.loc[highlight_ids, :],
                 )
                 aggregation = canvas.points(reorder_data, "x", "y", agg=ds.count_cat("label"))