bugfix: None colormap

src/palantir/plot.py

@@ -1155,10 +1155,13 @@ def plot_stats(
     }
     scatter_kwargs.update(kwargs)
 
-    try:
-        cmap = matplotlib.colormaps[cmap] if isinstance(cmap, str) else cmap
-    except KeyError:
+    if cmap is None:
         cmap = matplotlib.colormaps["viridis"]
+    else:
+        try:
+            cmap = matplotlib.colormaps[cmap] if isinstance(cmap, str) else cmap
+        except KeyError:
+            cmap = matplotlib.colormaps["viridis"]
     cmap = copy(cmap)
     cmap.set_bad(na_color)
 

src/palantir/utils.py

@@ -454,10 +454,10 @@ def run_diffusion_maps(
         and returned.
     """
 
-    if isinstance(data, pd.DataFrame):
-        data_df = data
-    else:
+    if isinstance(data, sc.AnnData):
         data_df = pd.DataFrame(data.obsm[pca_key], index=data.obs_names)
+    else:
+        data_df = data
 
     if not isinstance(data_df, pd.DataFrame) and not issparse(data_df):
         raise ValueError("'data_df' should be a pd.DataFrame or sc.AnnData")

tests/plot.py

@@ -550,7 +550,7 @@ def test_plot_stats_optional_parameters(mock_anndata):
 def test_plot_stats_masking(mock_anndata):
     # Create a condition here that you want to mask
     mask_condition = mock_anndata.obs["palantir_pseudotime"] > 0.5
-    mock_anndata.obsm["branch_masks"] = mask_condition
+    mock_anndata.obsm["branch_masks"] = pd.DataFrame({"mock_branch": mask_condition})
     fig, ax = plot_stats(
         mock_anndata,
         x="palantir_pseudotime",