Repository: theislab/scanpy (python)

Issue:
Plotting error when NaN in the categorical column and there are repetitive colours
- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of scanpy.
- [x] (optional) I have confirmed this bug exists on the master branch of scanpy.

---

Related issues: #1975.

### Minimal code sample

The reproducible example below features a few observations but the real-world scenario for this issue is very likely to arise when  there are a lot of categories (larger than the number of colours in the palette) and there is also at least one `NaN` value.

```python
import numpy as np
import pandas as pd
import scanpy as sc

# simulate a categorical column with a missing value
v = pd.Series(np.arange(10).astype(str))
v[0] = np.nan
v = v.astype("category")

ad = sc.AnnData(obs=pd.DataFrame(v, columns=["value"]))
ad.obsm["X_umap"] = np.random.normal(size=(ad.n_obs, 2))

sc.pl.umap(ad, color="value")
# colours are unique

# simulate scenario with non-unique colours
ad.uns['value_colors'][1] = ad.uns['value_colors'][0]

sc.pl.umap(ad, color="value")
# => AttributeError
```

```pytb
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-89-b7bdec198ea2> in <module>
----> 1 sc.pl.umap(ad, color="value")

site-packages/scanpy/plotting/_tools/scatterplots.py in umap(adata, **kwargs)
    657     tl.umap
    658     """
--> 659     return embedding(adata, 'umap', **kwargs)
    660 
    661 

site-packages/scanpy/plotting/_tools/scatterplots.py in embedding(adata, basis, color, gene_symbols, use_raw, sort_order, edges, edges_width, edges_color, neighbors_key, arrows, arrows_kwds, groups, components, layer, projection, scale_factor, color_map, cmap, palette, na_color, na_in_legend, size, frameon, legend_fontsize, legend_fontweight, legend_loc, legend_fontoutline, vmax, vmin, vcenter, norm, add_outline, outline_width, outline_color, ncols, hspace, wspace, title, show, save, ax, return_fig, **kwargs)
    250             groups=groups,
    251         )
--> 252         color_vector, categorical = _color_vector(
    253             adata,
    254             value_to_plot,

site-packages/scanpy/plotting/_tools/scatterplots.py in _color_vector(adata, values_key, values, palette, na_color)
   1279            # if hasattr(color_vector, "add_categories"):
   1280            #     color_vector = color_vector.add_categories([to_hex(na_color)])
-> 1281             color_vector = color_vector.add_categories([to_hex(na_color)])
   1282             color_vector = color_vector.fillna(to_hex(na_color))
   1283         return color_vector, True

AttributeError: 'Index' object has no attribute 'add_categories'
```

#### Versions

<details>

anndata             0.7.4, 0.7.9

numpy               1.22.2

pandas              1.4.0

scanpy              1.5.1, 1.8.2

</details>


### Workaround

Modify `_color_vector()` to only call `.add_categories()` when possible:

https://github.com/theislab/scanpy/blob/b69015e9e7007193c9ac461d5c6fbf845b3d6962/scanpy/plotting/_tools/scatterplots.py#L1275-L1278


```py
   1275  # Set color to 'missing color' for all missing values 
   1276  if color_vector.isna().any(): 
   1277      if hasattr(color_vector, "add_categories"):
   1278          color_vector = color_vector.add_categories([to_hex(na_color)]) 
   1279      color_vector = color_vector.fillna(to_hex(na_color))       
```


Pull request description:
Cast color vector back to categorical
Fixes #2133

Would be faster if we do the mapping in a may that maintains a categorical array.

Patch:
diff --git a/scanpy/plotting/_tools/scatterplots.py b/scanpy/plotting/_tools/scatterplots.py
index bc6a1a2a..61c7d48f 100644
--- a/scanpy/plotting/_tools/scatterplots.py
+++ b/scanpy/plotting/_tools/scatterplots.py
@@ -1270,7 +1270,7 @@ def _color_vector(
         }
         # If color_map does not have unique values, this can be slow as the
         # result is not categorical
-        color_vector = values.map(color_map)
+        color_vector = pd.Categorical(values.map(color_map))
 
         # Set color to 'missing color' for all missing values
         if color_vector.isna().any():


Test patch:
diff --git a/scanpy/tests/test_plotting.py b/scanpy/tests/test_plotting.py
index b0acd839..1b74f02f 100644
--- a/scanpy/tests/test_plotting.py
+++ b/scanpy/tests/test_plotting.py
@@ -1454,6 +1454,22 @@ def test_color_cycler(caplog):
     assert caplog.text == ""
 
 
+def test_repeated_colors_w_missing_value():
+    # https://github.com/theislab/scanpy/issues/2133
+    v = pd.Series(np.arange(10).astype(str))
+    v[0] = np.nan
+    v = v.astype("category")
+
+    ad = sc.AnnData(obs=pd.DataFrame(v, columns=["value"]))
+    ad.obsm["X_umap"] = np.random.normal(size=(ad.n_obs, 2))
+
+    sc.pl.umap(ad, color="value")
+
+    ad.uns['value_colors'][1] = ad.uns['value_colors'][0]
+
+    sc.pl.umap(ad, color="value")
+
+
 @pytest.mark.parametrize(
     "plot",
     (