Update _plot_discrete

Refactor adding `"nan"` as a category.
theislab · Mar 15, 2024 · b4c2dc7 · b4c2dc7
1 parent a84eeca
commit b4c2dc7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/cellrank/estimators/terminal_states/_term_states_estimator.py b/src/cellrank/estimators/terminal_states/_term_states_estimator.py
@@ -441,8 +441,8 @@ def _plot_discrete(
  with RandomKeys(self.adata, n=1 if same_plot else len(states), where="obs") as keys:
  if same_plot:
  outline = _data.cat.categories.to_list()
- _data = _data.astype(str).astype("category").cat.reorder_categories(["nan"] + states)
- states = ["nan"] + states
+ _data = _data.cat.add_categories(["nan"]).fillna("nan")
+ states.append("nan")
  color_mapper["nan"] = "#dedede"
  self.adata.obs[keys[0]] = _data
  self.adata.uns[f"{keys[0]}_colors"] = [color_mapper[name] for name in states]