Repository: scverse/scanpy

Issue:
when plotting, _prepare_dataframe can't look up gene_symbols in raw
When giving a plotting function the `gene_symbols` argument to specify that it should look in a column of `var` for `var_names` rather than look for them in the index, the underlying `_prepare_dataframe` function tries to find the `var_names` in `adata.var` rather than `adata.raw.var`, even when looking for the data itself in raw. For example, this code:

<!-- Put a minimal reproducible example that reproduces the bug in the code block below: -->
```python
adata.var['varnames'] = list(adata.var.index)
adata.raw.var['varnames'] = list(adata.raw.var.index)
'ENSGALG00000048305' in adata.raw.var['varnames'] # returns true
sc.pl.heatmap(
    adata,
    var_names=marker_genes_table.iloc[:, :5].values.flatten(),
    groupby='cluster_anno',
    show_gene_labels=True,
    swap_axes=True,
    gene_symbols='varnames',
    save=True,
    use_raw=True,
)
```
produces this error:
<!-- Put your Error output in this code block (if applicable, else delete the block): -->
```pytb
ERROR: Gene symbol 'ENSGALG00000048305' not found in given gene_symbols column: 'varnames'

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-40-80ce653c9d2e> in <module>
----> 1 sc.pl.heatmap(
      2     adata,
      3     var_names=marker_genes_table.iloc[:, :5].values.flatten(),
      4     groupby='cluster_anno',
      5     show_gene_labels=True,

~/anaconda3/envs/scanpy/lib/python3.8/site-packages/scanpy/plotting/_anndata.py in heatmap(adata, var_names, groupby, use_raw, log, num_categories, dendrogram, gene_symbols, var_group_positions, var_group_labels, var_group_rotation, layer, standard_scale, swap_axes, show_gene_labels, show, save, figsize, **kwds)
   1413     )
   1414 
-> 1415     categories, obs_tidy = _prepare_dataframe(
   1416         adata,
   1417         var_names,

TypeError: cannot unpack non-iterable NoneType object
```

#### Versions:
<!-- Output of scanpy.logging.print_versions() -->
```
scanpy==1.4.6 anndata==0.7.1 umap==0.4.3 numpy==1.18.4 scipy==1.4.1 pandas==1.0.3 scikit-learn==0.23.0 statsmodels==0.11.1 python-igraph==0.8.2
```



Discussion:

Resolving patch:
diff --git a/scanpy/plotting/_anndata.py b/scanpy/plotting/_anndata.py
--- a/scanpy/plotting/_anndata.py
+++ b/scanpy/plotting/_anndata.py
@@ -2890,15 +2890,21 @@ def _prepare_dataframe(
         # translate gene_symbols to var_names
         # slow method but gives a meaningful error if no gene symbol is found:
         translated_var_names = []
+        # if we're using raw to plot, we should also do gene symbol translations
+        # using raw
+        if use_raw:
+            adata_or_raw = adata.raw
+        else:
+            adata_or_raw = adata
         for symbol in var_names:
-            if symbol not in adata.var[gene_symbols].values:
+            if symbol not in adata_or_raw.var[gene_symbols].values:
                 logg.error(
                     f"Gene symbol {symbol!r} not found in given "
                     f"gene_symbols column: {gene_symbols!r}"
                 )
                 return
             translated_var_names.append(
-                adata.var[adata.var[gene_symbols] == symbol].index[0]
+                adata_or_raw.var[adata_or_raw.var[gene_symbols] == symbol].index[0]
             )
         symbols = var_names
         var_names = translated_var_names


Test patch: