From 2877f27b4370d2c3f22a5a281c750f00202d6618 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Tue, 28 Nov 2023 15:27:10 +0000
Subject: [PATCH 01/26] matrixplot functionality

---
 scanpy/plotting/_baseplot_class.py |  8 ++++++--
 scanpy/plotting/_matrixplot.py     | 18 +++++++++++++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 65ebfa64b5..64c1aa844e 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -77,6 +77,7 @@ def __init__(
         adata: AnnData,
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
+        groupby_cols: str | Sequence[str] = [],
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -109,10 +110,14 @@ def __init__(
 
         self._update_var_groups()
 
+        self.groupby = [groupby] if isinstance(groupby, str) else groupby
+        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
+        self.groupby += groupby_cols
+
         self.categories, self.obs_tidy = _prepare_dataframe(
             adata,
             self.var_names,
-            groupby,
+            self.groupby,
             use_raw,
             log,
             num_categories,
@@ -139,7 +144,6 @@ def __init__(
                 return
 
         self.adata = adata
-        self.groupby = [groupby] if isinstance(groupby, str) else groupby
         self.log = log
         self.kwds = kwds
 
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index fe1b7844da..93d0ed070e 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
+import pandas as pd
 from matplotlib import pyplot as plt
 from matplotlib import rcParams
 
@@ -99,6 +100,7 @@ def __init__(
         adata: AnnData,
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
+        groupby_cols: str | Sequence[str] = [],
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -124,6 +126,7 @@ def __init__(
             adata,
             var_names,
             groupby,
+            groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -142,7 +145,6 @@ def __init__(
             norm=norm,
             **kwds,
         )
-
         if values_df is None:
             # compute mean value
             values_df = self.obs_tidy.groupby(level=0).mean()
@@ -158,6 +160,18 @@ def __init__(
             else:
                 logg.warning("Unknown type for standard_scale, ignored")
 
+            if len(groupby_cols) > 0:
+                label = values_df.index.name
+                dirty_df = values_df.reset_index()
+                dirty_df.index = pd.MultiIndex.from_tuples(
+                    dirty_df[label].str.split('_').tolist(), names=self.groupby)
+                dirty_df = dirty_df.drop(label, axis=1).unstack(level=self.groupby_cols)
+
+                temp_df = dirty_df.reset_index(drop=True)
+                temp_df.index = dirty_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                temp_df.columns = dirty_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                values_df = temp_df
+
         self.values_df = values_df
 
         self.cmap = self.DEFAULT_COLORMAP
@@ -290,6 +304,7 @@ def matrixplot(
     adata: AnnData,
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
+    groupby_cols: str | Sequence[str] = [],
     use_raw: bool | None = None,
     log: bool = False,
     num_categories: int = 7,
@@ -383,6 +398,7 @@ def matrixplot(
         adata,
         var_names,
         groupby=groupby,
+        groupby_cols=groupby_cols,
         use_raw=use_raw,
         log=log,
         num_categories=num_categories,

From dc676d4ef3f647a87ecbccd5fc5bd734b5362ed7 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Tue, 28 Nov 2023 15:58:28 +0000
Subject: [PATCH 02/26] added convert_tidy

---
 scanpy/plotting/_baseplot_class.py | 13 ++++++++++++-
 scanpy/plotting/_matrixplot.py     | 22 +++++++++++-----------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 64c1aa844e..5ccddf43b4 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -1,4 +1,4 @@
-"""BasePlot for dotplot, matrixplot and stacked_violin
+"""BasePlot for dotplot, matrixplot and stacked_violinA
 """
 from __future__ import annotations
 
@@ -9,6 +9,7 @@
 from warnings import warn
 
 import numpy as np
+import pandas as pd
 from matplotlib import gridspec
 from matplotlib import pyplot as plt
 
@@ -835,6 +836,16 @@ def savefig(self, filename: str, bbox_inches: str | None = "tight", **kwargs):
         self.make_figure()
         plt.savefig(filename, bbox_inches=bbox_inches, **kwargs)
 
+    def _convert_tidy_to_stacked(self, values_df):
+        """\
+        Utility function used to convert obs_tidy into the correct format when using a groupby_col.
+        """
+        label = values_df.index.name
+        stacked_df = values_df.reset_index()
+        stacked_df.index = pd.MultiIndex.from_tuples(
+            stacked_df[label].str.split('_').tolist(), names=self.groupby)
+        return stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
+
     def _reorder_categories_after_dendrogram(self, dendrogram):
         """\
         Function used by plotting functions that need to reorder the the groupby
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 93d0ed070e..8cace0162d 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-import pandas as pd
 from matplotlib import pyplot as plt
 from matplotlib import rcParams
 
@@ -161,16 +160,17 @@ def __init__(
                 logg.warning("Unknown type for standard_scale, ignored")
 
             if len(groupby_cols) > 0:
-                label = values_df.index.name
-                dirty_df = values_df.reset_index()
-                dirty_df.index = pd.MultiIndex.from_tuples(
-                    dirty_df[label].str.split('_').tolist(), names=self.groupby)
-                dirty_df = dirty_df.drop(label, axis=1).unstack(level=self.groupby_cols)
-
-                temp_df = dirty_df.reset_index(drop=True)
-                temp_df.index = dirty_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                temp_df.columns = dirty_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                values_df = temp_df
+                stacked_df = self._convert_tidy_to_stacked(values_df)
+                # label = values_df.index.name
+                # stacked_df = values_df.reset_index()
+                # stacked_df.index = pd.MultiIndex.from_tuples(
+                #     stacked_df[label].str.split('_').tolist(), names=self.groupby)
+                # stacked_df = stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
+
+                # recreate the original formatting of values_df
+                values_df = stacked_df.reset_index(drop=True)
+                values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
 
         self.values_df = values_df
 

From a334e9b96d2b7a7fc63008b97163c21a0a374cff Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Tue, 28 Nov 2023 16:36:05 +0000
Subject: [PATCH 03/26] Distinguish MultiIndex

---
 scanpy/plotting/_matrixplot.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 8cace0162d..9e977f8ed4 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
+import pandas as pd
 from matplotlib import pyplot as plt
 from matplotlib import rcParams
 
@@ -169,7 +170,10 @@ def __init__(
 
                 # recreate the original formatting of values_df
                 values_df = stacked_df.reset_index(drop=True)
-                values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                if isinstance(stacked_df.index, pd.MultiIndex):
+                    values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                else:
+                    values_df.index = stacked_df.index.to_series().apply(lambda x: ''.join(map(str, x))).values
                 values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
 
         self.values_df = values_df

From 0ce9b8e6886d14e26cb578a9a91abec8075411fb Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Tue, 28 Nov 2023 16:47:44 +0000
Subject: [PATCH 04/26] functionality for dotplot

---
 scanpy/plotting/_dotplot.py    | 17 +++++++++++++++++
 scanpy/plotting/_matrixplot.py |  5 -----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 8d21cf384b..eb7041eb55 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -116,6 +116,7 @@ def __init__(
         adata: AnnData,
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
+        groupby_cols: str | Sequence[str] = [],
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -144,6 +145,7 @@ def __init__(
             adata,
             var_names,
             groupby,
+            groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -178,6 +180,13 @@ def __init__(
             dot_size_df = (
                 obs_bool.groupby(level=0).sum() / obs_bool.groupby(level=0).count()
             )
+            if len(groupby_cols) > 0:
+                dot_size_df_stacked = self._convert_tidy_to_stacked(dot_size_df)
+
+                values_df = dot_size_df_stacked.reset_index(drop=True)
+                values_df.index = dot_size_df_stacked.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                values_df.columns = dot_size_df_stacked.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                dot_size_df = values_df
 
         if dot_color_df is None:
             # 2. compute mean expression value value
@@ -198,6 +207,13 @@ def __init__(
                 pass
             else:
                 logg.warning("Unknown type for standard_scale, ignored")
+            if len(groupby_cols) > 0:
+                dot_color_df_stacked = self._convert_tidy_to_stacked(dot_color_df)
+
+                values_df = dot_color_df_stacked.reset_index(drop=True)
+                values_df.index = dot_color_df_stacked.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                values_df.columns = dot_color_df_stacked.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                dot_color_df = values_df
         else:
             # check that both matrices have the same shape
             if dot_color_df.shape != dot_size_df.shape:
@@ -708,6 +724,7 @@ def _dotplot(
         # rescale size to match smallest_dot and largest_dot
         size = size * (largest_dot - smallest_dot) + smallest_dot
         normalize = check_colornorm(vmin, vmax, vcenter, norm)
+        normalize(mean_flat[~np.isnan(mean_flat)])  # circumvent unexpected behavior with nan in matplotlib
 
         if color_on == "square":
             if edge_color is None:
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 9e977f8ed4..afad94fd35 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -162,11 +162,6 @@ def __init__(
 
             if len(groupby_cols) > 0:
                 stacked_df = self._convert_tidy_to_stacked(values_df)
-                # label = values_df.index.name
-                # stacked_df = values_df.reset_index()
-                # stacked_df.index = pd.MultiIndex.from_tuples(
-                #     stacked_df[label].str.split('_').tolist(), names=self.groupby)
-                # stacked_df = stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
 
                 # recreate the original formatting of values_df
                 values_df = stacked_df.reset_index(drop=True)

From 23a00bfe0734b9f20bdc6065b09d10db3edff0e8 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Tue, 28 Nov 2023 17:12:30 +0000
Subject: [PATCH 05/26] remove copy-pasta

---
 scanpy/plotting/_baseplot_class.py | 14 ++++++++++++--
 scanpy/plotting/_dotplot.py        | 15 ++-------------
 scanpy/plotting/_matrixplot.py     | 10 +---------
 3 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 5ccddf43b4..81209b1ec6 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -1,4 +1,4 @@
-"""BasePlot for dotplot, matrixplot and stacked_violinA
+"""BasePlot for dotplot, matrixplot and stacked_violin
 """
 from __future__ import annotations
 
@@ -844,7 +844,17 @@ def _convert_tidy_to_stacked(self, values_df):
         stacked_df = values_df.reset_index()
         stacked_df.index = pd.MultiIndex.from_tuples(
             stacked_df[label].str.split('_').tolist(), names=self.groupby)
-        return stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
+        stacked_df = stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
+
+        # recreate the original formatting of values_df
+        values_df = stacked_df.reset_index(drop=True)
+        if isinstance(stacked_df.index, pd.MultiIndex):
+            values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+        else:
+            values_df.index = stacked_df.index.to_series().apply(lambda x: ''.join(map(str, x))).values
+        values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+
+        return values_df
 
     def _reorder_categories_after_dendrogram(self, dendrogram):
         """\
diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index eb7041eb55..1c2664651d 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -181,13 +181,7 @@ def __init__(
                 obs_bool.groupby(level=0).sum() / obs_bool.groupby(level=0).count()
             )
             if len(groupby_cols) > 0:
-                dot_size_df_stacked = self._convert_tidy_to_stacked(dot_size_df)
-
-                values_df = dot_size_df_stacked.reset_index(drop=True)
-                values_df.index = dot_size_df_stacked.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                values_df.columns = dot_size_df_stacked.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                dot_size_df = values_df
-
+                dot_size_df = self._convert_tidy_to_stacked(dot_size_df)
         if dot_color_df is None:
             # 2. compute mean expression value value
             if mean_only_expressed:
@@ -208,12 +202,7 @@ def __init__(
             else:
                 logg.warning("Unknown type for standard_scale, ignored")
             if len(groupby_cols) > 0:
-                dot_color_df_stacked = self._convert_tidy_to_stacked(dot_color_df)
-
-                values_df = dot_color_df_stacked.reset_index(drop=True)
-                values_df.index = dot_color_df_stacked.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                values_df.columns = dot_color_df_stacked.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                dot_color_df = values_df
+                dot_color_df = self._convert_tidy_to_stacked(dot_color_df)
         else:
             # check that both matrices have the same shape
             if dot_color_df.shape != dot_size_df.shape:
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index afad94fd35..9683b41c16 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -161,15 +161,7 @@ def __init__(
                 logg.warning("Unknown type for standard_scale, ignored")
 
             if len(groupby_cols) > 0:
-                stacked_df = self._convert_tidy_to_stacked(values_df)
-
-                # recreate the original formatting of values_df
-                values_df = stacked_df.reset_index(drop=True)
-                if isinstance(stacked_df.index, pd.MultiIndex):
-                    values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
-                else:
-                    values_df.index = stacked_df.index.to_series().apply(lambda x: ''.join(map(str, x))).values
-                values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+                values_df = self._convert_tidy_to_stacked(values_df)
 
         self.values_df = values_df
 

From 0c7b83ecd71eb31a82fdc2691efbca7e7f046a00 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Tue, 28 Nov 2023 17:59:36 +0000
Subject: [PATCH 06/26] functionality for stacked_violin

---
 scanpy/plotting/_matrixplot.py     |  1 -
 scanpy/plotting/_stacked_violin.py | 53 ++++++++++++++++++++----------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 9683b41c16..131b8ab2ae 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-import pandas as pd
 from matplotlib import pyplot as plt
 from matplotlib import rcParams
 
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index e058358400..3f4a75862e 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -135,6 +135,7 @@ def __init__(
         adata: AnnData,
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
+        groupby_cols: str | Sequence[str] = [],
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -159,6 +160,7 @@ def __init__(
             adata,
             var_names,
             groupby,
+            groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -331,6 +333,8 @@ def _mainplot(self, ax):
         _color_df = _matrix.groupby(level=0).median()
         if self.are_axes_swapped:
             _color_df = _color_df.T
+        if len(self.groupby_cols) > 0:
+            _color_df = self._convert_tidy_to_stacked(_color_df)
 
         cmap = plt.get_cmap(self.kwds.get("cmap", self.cmap))
         if "cmap" in self.kwds:
@@ -341,6 +345,7 @@ def _mainplot(self, ax):
             self.vboundnorm.vcenter,
             self.vboundnorm.norm,
         )
+        normalize(_color_df.values[~np.isnan(_color_df.values)])  # circumvent unexpected behavior with nan in matplotlib
         colormap_array = cmap(normalize(_color_df.values))
         x_spacer_size = self.plot_x_padding
         y_spacer_size = self.plot_y_padding
@@ -408,25 +413,37 @@ def _make_rows_of_violinplots(
         # This format is convenient to aggregate per gene or per category
         # while making the violin plots.
 
-        df = (
-            pd.DataFrame(_matrix.stack(dropna=False))
-            .reset_index()
-            .rename(
-                columns={
-                    "level_1": "genes",
-                    _matrix.index.name: "categories",
-                    0: "values",
-                }
+        if len(self.groupby_cols) < 1:
+            df = (
+                pd.DataFrame(_matrix.stack(dropna=False))
+                .reset_index()
+                .rename(
+                    columns={
+                        "level_1": "genes",
+                        _matrix.index.name: "categories",
+                        0: "values",
+                    }
+                )
             )
-        )
-        df["genes"] = (
-            df["genes"].astype("category").cat.reorder_categories(_matrix.columns)
-        )
-        df["categories"] = (
-            df["categories"]
-            .astype("category")
-            .cat.reorder_categories(_matrix.index.categories)
-        )
+            df["genes"] = (
+                df["genes"].astype("category").cat.reorder_categories(_matrix.columns)
+            )
+            df["categories"] = (
+                df["categories"]
+                .astype("category")
+                .cat.reorder_categories(_matrix.index.categories)
+            )
+        else:
+            # partially taken from self._convert_tidy_to_stacked
+            label = _matrix.index.name
+            stacked_df = _matrix.reset_index()
+            stacked_df.index = pd.MultiIndex.from_tuples(
+                stacked_df[label].str.split('_').tolist(), names=self.groupby)
+            stacked_df = stacked_df.drop(label, axis=1).reset_index().melt(id_vars=self.groupby)
+            stacked_df['genes'] = stacked_df[self.groupby_cols + ['variable']].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
+            stacked_df = stacked_df.drop(self.groupby_cols + ['variable'], axis=1)
+            stacked_df.columns = ['categories', 'values', 'genes']
+            df = stacked_df
 
         # the ax need to be subdivided
         # define a layout of nrows = len(categories) rows

From 78a39730b372d80e10138b038bbdac6608b9e0c2 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Wed, 29 Nov 2023 10:44:36 +0000
Subject: [PATCH 07/26] bug fixes to categories

---
 scanpy/plotting/_baseplot_class.py | 26 ++++++++++++++++++++------
 scanpy/plotting/_stacked_violin.py |  7 +++----
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 81209b1ec6..038dc3b431 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -125,6 +125,18 @@ def __init__(
             layer=layer,
             gene_symbols=gene_symbols,
         )
+        # reset categories if using groupby_cols
+        if len(self.groupby_cols) > 0:
+            self.categories, _ = _prepare_dataframe(
+                adata,
+                self.var_names,
+                self.groupby[:-len(self.groupby_cols)],
+                use_raw,
+                log,
+                num_categories,
+                layer=layer,
+                gene_symbols=gene_symbols,
+            )
         if len(self.categories) > self.MAX_NUM_CATEGORIES:
             warn(
                 f"Over {self.MAX_NUM_CATEGORIES} categories found. "
@@ -853,7 +865,6 @@ def _convert_tidy_to_stacked(self, values_df):
         else:
             values_df.index = stacked_df.index.to_series().apply(lambda x: ''.join(map(str, x))).values
         values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
-
         return values_df
 
     def _reorder_categories_after_dendrogram(self, dendrogram):
@@ -882,13 +893,16 @@ def _format_first_three_categories(_categories):
                 _categories = _categories[:3] + ["etc."]
             return ", ".join(_categories)
 
-        key = _get_dendrogram_key(self.adata, dendrogram, self.groupby)
-
+        if len(self.groupby_cols) > 0:
+            dendro_groupby = self.groupby[:-len(self.groupby_cols)]
+        else:
+            dendro_groupby = self.groupby
+        key = _get_dendrogram_key(self.adata, dendrogram, dendro_groupby)
         dendro_info = self.adata.uns[key]
-        if self.groupby != dendro_info["groupby"]:
+        if dendro_groupby != dendro_info["groupby"]:
             raise ValueError(
                 "Incompatible observations. The precomputed dendrogram contains "
-                f"information for the observation: '{self.groupby}' while the plot is "
+                f"information for the observation: '{dendro_groupby}' while the plot is "
                 f"made for the observation: '{dendro_info['groupby']}. "
                 "Please run `sc.tl.dendrogram` using the right observation.'"
             )
@@ -901,7 +915,7 @@ def _format_first_three_categories(_categories):
             raise ValueError(
                 "Incompatible observations. Dendrogram data has "
                 f"{len(categories_idx_ordered)} categories but current groupby "
-                f"observation {self.groupby!r} contains {len(self.categories)} categories. "
+                f"observation {dendro_groupby!r} contains {len(self.categories)} categories. "
                 "Most likely the underlying groupby observation changed after the "
                 "initial computation of `sc.tl.dendrogram`. "
                 "Please run `sc.tl.dendrogram` again.'"
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 3f4a75862e..6f5cfb51c3 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -440,9 +440,9 @@ def _make_rows_of_violinplots(
             stacked_df.index = pd.MultiIndex.from_tuples(
                 stacked_df[label].str.split('_').tolist(), names=self.groupby)
             stacked_df = stacked_df.drop(label, axis=1).reset_index().melt(id_vars=self.groupby)
-            stacked_df['genes'] = stacked_df[self.groupby_cols + ['variable']].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
-            stacked_df = stacked_df.drop(self.groupby_cols + ['variable'], axis=1)
-            stacked_df.columns = ['categories', 'values', 'genes']
+            stacked_df['genes'] = stacked_df[['variable'] + self.groupby_cols].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
+            stacked_df['categories'] = stacked_df[self.groupby[:-len(self.groupby_cols)]].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
+            stacked_df = stacked_df.drop(self.groupby + ['variable'], axis=1).rename(columns={'value':'values'})
             df = stacked_df
 
         # the ax need to be subdivided
@@ -481,7 +481,6 @@ def _make_rows_of_violinplots(
                 # (in _color_df the values are not renamed as those
                 # values will be used to label the ticks)
                 _df = df[df.genes == _matrix.columns[idx]]
-
             row_ax = sns.violinplot(
                 x=x,
                 y="values",

From ea6d8ee153ad9582876c1b94ff10938a79e109a4 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Wed, 29 Nov 2023 11:19:49 +0000
Subject: [PATCH 08/26] Adapt figure width to conditions

---
 scanpy/plotting/_baseplot_class.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 038dc3b431..158776e5de 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -648,7 +648,7 @@ def make_figure(self):
         if self.height is None:
             mainplot_height = len(self.categories) * category_height
             mainplot_width = (
-                len(self.var_names) * category_width + self.group_extra_size
+                len(self.var_names) * category_width  * (1+len(self.groupby_cols)) + self.group_extra_size
             )
             if self.are_axes_swapped:
                 mainplot_height, mainplot_width = mainplot_width, mainplot_height

From 90856417088b2b01e8a9f8ef5425a498cb80aac0 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Wed, 29 Nov 2023 11:24:16 +0000
Subject: [PATCH 09/26] switch to changing obs_tidy instead of categories

---
 scanpy/plotting/_baseplot_class.py | 12 +++++++-----
 scanpy/plotting/_dotplot.py        |  1 -
 scanpy/plotting/_stacked_violin.py |  8 ++++----
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 158776e5de..ece905ea86 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -113,7 +113,7 @@ def __init__(
 
         self.groupby = [groupby] if isinstance(groupby, str) else groupby
         self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
-        self.groupby += groupby_cols
+#        self.groupby += groupby_cols
 
         self.categories, self.obs_tidy = _prepare_dataframe(
             adata,
@@ -125,12 +125,12 @@ def __init__(
             layer=layer,
             gene_symbols=gene_symbols,
         )
-        # reset categories if using groupby_cols
+        # reset obs_tidy if using groupby_cols
         if len(self.groupby_cols) > 0:
-            self.categories, _ = _prepare_dataframe(
+            _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
-                self.groupby[:-len(self.groupby_cols)],
+                self.groupby + self.groupby_cols,
                 use_raw,
                 log,
                 num_categories,
@@ -369,6 +369,8 @@ def add_totals(
         _sort = True if sort is not None else False
         _ascending = True if sort == "ascending" else False
         counts_df = self.obs_tidy.index.value_counts(sort=_sort, ascending=_ascending)
+        if len(self.groupby_cols) > 0:  # could remove the previous line and only use this but this is slower
+            counts_df = self.adata.obs[self.groupby].value_counts(sort=_sort, ascending=_ascending)
 
         if _sort:
             self.categories_order = counts_df.index
@@ -855,7 +857,7 @@ def _convert_tidy_to_stacked(self, values_df):
         label = values_df.index.name
         stacked_df = values_df.reset_index()
         stacked_df.index = pd.MultiIndex.from_tuples(
-            stacked_df[label].str.split('_').tolist(), names=self.groupby)
+            stacked_df[label].str.split('_').tolist(), names=self.groupby + self.groupby_cols)
         stacked_df = stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
 
         # recreate the original formatting of values_df
diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 1c2664651d..1bd19de592 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -759,7 +759,6 @@ def _dotplot(
                 edgecolor=edge_color,
                 norm=normalize,
             )
-
             dot_ax.scatter(x, y, **kwds)
 
         y_ticks = np.arange(dot_color.shape[0]) + 0.5
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 6f5cfb51c3..658e725d92 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -438,11 +438,11 @@ def _make_rows_of_violinplots(
             label = _matrix.index.name
             stacked_df = _matrix.reset_index()
             stacked_df.index = pd.MultiIndex.from_tuples(
-                stacked_df[label].str.split('_').tolist(), names=self.groupby)
-            stacked_df = stacked_df.drop(label, axis=1).reset_index().melt(id_vars=self.groupby)
+                stacked_df[label].str.split('_').tolist(), names=self.groupby + self.groupby_cols)
+            stacked_df = stacked_df.drop(label, axis=1).reset_index().melt(id_vars=self.groupby + self.groupby_cols)
             stacked_df['genes'] = stacked_df[['variable'] + self.groupby_cols].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
-            stacked_df['categories'] = stacked_df[self.groupby[:-len(self.groupby_cols)]].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
-            stacked_df = stacked_df.drop(self.groupby + ['variable'], axis=1).rename(columns={'value':'values'})
+            stacked_df['categories'] = stacked_df[self.groupby].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
+            stacked_df = stacked_df.drop(self.groupby + self.groupby_cols + ['variable'], axis=1).rename(columns={'value':'values'})
             df = stacked_df
 
         # the ax need to be subdivided

From 10da0937ee61027ed9f3c586ba599764b6cc1d5d Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Wed, 29 Nov 2023 11:32:17 +0000
Subject: [PATCH 10/26] revert some changes

---
 scanpy/plotting/_baseplot_class.py | 22 ++++++++--------------
 1 file changed, 8 insertions(+), 14 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index ece905ea86..a237eed4ff 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -111,14 +111,10 @@ def __init__(
 
         self._update_var_groups()
 
-        self.groupby = [groupby] if isinstance(groupby, str) else groupby
-        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
-#        self.groupby += groupby_cols
-
         self.categories, self.obs_tidy = _prepare_dataframe(
             adata,
             self.var_names,
-            self.groupby,
+            groupby,
             use_raw,
             log,
             num_categories,
@@ -130,7 +126,7 @@ def __init__(
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
-                self.groupby + self.groupby_cols,
+                groupby + groupby_cols,
                 use_raw,
                 log,
                 num_categories,
@@ -157,6 +153,8 @@ def __init__(
                 return
 
         self.adata = adata
+        self.groupby = [groupby] if isinstance(groupby, str) else groupby
+        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
         self.log = log
         self.kwds = kwds
 
@@ -895,16 +893,12 @@ def _format_first_three_categories(_categories):
                 _categories = _categories[:3] + ["etc."]
             return ", ".join(_categories)
 
-        if len(self.groupby_cols) > 0:
-            dendro_groupby = self.groupby[:-len(self.groupby_cols)]
-        else:
-            dendro_groupby = self.groupby
-        key = _get_dendrogram_key(self.adata, dendrogram, dendro_groupby)
+        key = _get_dendrogram_key(self.adata, dendrogram, self.groupby)
         dendro_info = self.adata.uns[key]
-        if dendro_groupby != dendro_info["groupby"]:
+        if self.groupby != dendro_info["groupby"]:
             raise ValueError(
                 "Incompatible observations. The precomputed dendrogram contains "
-                f"information for the observation: '{dendro_groupby}' while the plot is "
+                f"information for the observation: '{self.groupby}' while the plot is "
                 f"made for the observation: '{dendro_info['groupby']}. "
                 "Please run `sc.tl.dendrogram` using the right observation.'"
             )
@@ -917,7 +911,7 @@ def _format_first_three_categories(_categories):
             raise ValueError(
                 "Incompatible observations. Dendrogram data has "
                 f"{len(categories_idx_ordered)} categories but current groupby "
-                f"observation {dendro_groupby!r} contains {len(self.categories)} categories. "
+                f"observation {self.groupby!r} contains {len(self.categories)} categories. "
                 "Most likely the underlying groupby observation changed after the "
                 "initial computation of `sc.tl.dendrogram`. "
                 "Please run `sc.tl.dendrogram` again.'"

From 170e5358305301efeba4ba8646ebf6132ba2e325 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Wed, 29 Nov 2023 11:34:51 +0000
Subject: [PATCH 11/26] bug fix

---
 scanpy/plotting/_baseplot_class.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index a237eed4ff..d33f92d476 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -122,7 +122,7 @@ def __init__(
             gene_symbols=gene_symbols,
         )
         # reset obs_tidy if using groupby_cols
-        if len(self.groupby_cols) > 0:
+        if len(groupby_cols) > 0:
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
@@ -288,7 +288,7 @@ def add_dendrogram(
         # to correctly plot the dendrogram the categories need to be ordered
         # according to the dendrogram ordering.
         self._reorder_categories_after_dendrogram(dendrogram_key)
-
+        print(self.categories)
         dendro_ticks = np.arange(len(self.categories)) + 0.5
 
         self.group_extra_size = size

From a64f09e2eeb7334a024202cf1ee8e86cacc52633 Mon Sep 17 00:00:00 2001
From: yugeji <jiyuge@gmail.com>
Date: Wed, 29 Nov 2023 11:50:28 +0000
Subject: [PATCH 12/26] bug fix

---
 scanpy/plotting/_baseplot_class.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index d33f92d476..937726cdb8 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -102,6 +102,8 @@ def __init__(
         self.var_group_positions = var_group_positions
         self.var_group_rotation = var_group_rotation
         self.width, self.height = figsize if figsize is not None else (None, None)
+        self.groupby = [groupby] if isinstance(groupby, str) else groupby
+        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
 
         self.has_var_groups = (
             True
@@ -114,7 +116,7 @@ def __init__(
         self.categories, self.obs_tidy = _prepare_dataframe(
             adata,
             self.var_names,
-            groupby,
+            self.groupby,
             use_raw,
             log,
             num_categories,
@@ -122,11 +124,11 @@ def __init__(
             gene_symbols=gene_symbols,
         )
         # reset obs_tidy if using groupby_cols
-        if len(groupby_cols) > 0:
+        if len(self.groupby_cols) > 0:
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
-                groupby + groupby_cols,
+                self.groupby + self.groupby_cols,
                 use_raw,
                 log,
                 num_categories,
@@ -153,8 +155,6 @@ def __init__(
                 return
 
         self.adata = adata
-        self.groupby = [groupby] if isinstance(groupby, str) else groupby
-        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
         self.log = log
         self.kwds = kwds
 
@@ -288,7 +288,7 @@ def add_dendrogram(
         # to correctly plot the dendrogram the categories need to be ordered
         # according to the dendrogram ordering.
         self._reorder_categories_after_dendrogram(dendrogram_key)
-        print(self.categories)
+
         dendro_ticks = np.arange(len(self.categories)) + 0.5
 
         self.group_extra_size = size

From adebf2b43211b41cf22bb909291e71e585c18673 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Thu, 30 Nov 2023 13:19:06 +0000
Subject: [PATCH 13/26] Add groupby_cols to common_plot_args

---
 scanpy/plotting/_docs.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/scanpy/plotting/_docs.py b/scanpy/plotting/_docs.py
index c1078090b0..36798737b4 100644
--- a/scanpy/plotting/_docs.py
+++ b/scanpy/plotting/_docs.py
@@ -202,6 +202,8 @@
     then the `var_group_labels` and `var_group_positions` are set.
 groupby
     The key of the observation grouping to consider.
+groupby_cols
+    The key of the observation grouping to consider for grouping columns.
 use_raw
     Use `raw` attribute of `adata` if present.
 log

From fc03fdc3baa92a90757e70509f24fd6906368ff0 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Thu, 30 Nov 2023 13:48:01 +0000
Subject: [PATCH 14/26] Show groupby_cols in docstrings

---
 scanpy/plotting/_dotplot.py        | 4 +++-
 scanpy/plotting/_stacked_violin.py | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 1bd19de592..719abacdad 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -813,6 +813,7 @@ def dotplot(
     adata: AnnData,
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
+    groupby_cols: str | Sequence[str] = [],
     use_raw: bool | None = None,
     log: bool = False,
     num_categories: int = 7,
@@ -955,7 +956,8 @@ def dotplot(
     dp = DotPlot(
         adata,
         var_names,
-        groupby,
+        groupby=groupby,
+        groupby_cols=groupby_cols,
         use_raw=use_raw,
         log=log,
         num_categories=num_categories,
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 658e725d92..f7b9a0a996 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -577,6 +577,7 @@ def stacked_violin(
     adata: AnnData,
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
+    groupby_cols: str | Sequence[str] = [],
     log: bool = False,
     use_raw: bool | None = None,
     num_categories: int = 7,
@@ -713,6 +714,7 @@ def stacked_violin(
         adata,
         var_names,
         groupby=groupby,
+        groupby_cols=groupby_cols,
         use_raw=use_raw,
         log=log,
         num_categories=num_categories,

From 408bbeb7da3482fb2be631ba157013eeb76a6f15 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Thu, 30 Nov 2023 14:58:23 +0000
Subject: [PATCH 15/26] Groupby_cols category counter for fig formatting

---
 scanpy/plotting/_baseplot_class.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 937726cdb8..a877fa5eb8 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -104,7 +104,6 @@ def __init__(
         self.width, self.height = figsize if figsize is not None else (None, None)
         self.groupby = [groupby] if isinstance(groupby, str) else groupby
         self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
-
         self.has_var_groups = (
             True
             if var_group_positions is not None and len(var_group_positions) > 0
@@ -123,8 +122,11 @@ def __init__(
             layer=layer,
             gene_symbols=gene_symbols,
         )
+        
         # reset obs_tidy if using groupby_cols
         if len(self.groupby_cols) > 0:
+            # TODO : Check if we rather need the product of categories ?
+            self.categories_cols = adata.obs.loc[:,self.groupby_cols].nunique().sum()
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
@@ -135,6 +137,8 @@ def __init__(
                 layer=layer,
                 gene_symbols=gene_symbols,
             )
+        else:
+            self.categories_cols = 0
         if len(self.categories) > self.MAX_NUM_CATEGORIES:
             warn(
                 f"Over {self.MAX_NUM_CATEGORIES} categories found. "
@@ -648,7 +652,7 @@ def make_figure(self):
         if self.height is None:
             mainplot_height = len(self.categories) * category_height
             mainplot_width = (
-                len(self.var_names) * category_width  * (1+len(self.groupby_cols)) + self.group_extra_size
+                len(self.var_names) * category_width  * (1+self.categories_cols) + self.group_extra_size
             )
             if self.are_axes_swapped:
                 mainplot_height, mainplot_width = mainplot_width, mainplot_height

From 9d6def727aeaf18b21d546d148607b0cdc291d8d Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 14:19:33 +0000
Subject: [PATCH 16/26] Add detail to docstrings

---
 scanpy/plotting/_dotplot.py        | 1 +
 scanpy/plotting/_matrixplot.py     | 1 +
 scanpy/plotting/_stacked_violin.py | 1 +
 3 files changed, 3 insertions(+)

diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 719abacdad..205530e351 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -850,6 +850,7 @@ def dotplot(
     Makes a *dot plot* of the expression values of `var_names`.
 
     For each var_name and each `groupby` category a dot is plotted.
+    Columns can optionally be grouped by specifying `groupby_cols`.
     Each dot represents two values: mean expression within each category
     (visualized by color) and fraction of cells expressing the `var_name` in the
     category (visualized by the size of the dot). If `groupby` is not given,
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 131b8ab2ae..d137704328 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -323,6 +323,7 @@ def matrixplot(
 ) -> MatrixPlot | dict | None:
     """\
     Creates a heatmap of the mean expression values per group of each var_names.
+    Columns can optionally be grouped by specifying `groupby_cols`.
 
     This function provides a convenient interface to the :class:`~scanpy.pl.MatrixPlot`
     class. If you need more flexibility, you should use :class:`~scanpy.pl.MatrixPlot`
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index f7b9a0a996..370c018b8c 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -616,6 +616,7 @@ def stacked_violin(
     Makes a compact image composed of individual violin plots
     (from :func:`~seaborn.violinplot`) stacked on top of each other.
     Useful to visualize gene expression per cluster.
+    Columns can optionally be grouped by specifying `groupby_cols`.
 
     Wraps :func:`seaborn.violinplot` for :class:`~anndata.AnnData`.
 

From f5789778b9d4722ca243a400edb0b31ba5d31577 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 15:01:49 +0000
Subject: [PATCH 17/26] Correct max if nans present

---
 scanpy/plotting/_dotplot.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 205530e351..cc0cf98a85 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -692,7 +692,7 @@ def _dotplot(
         if "cmap" in kwds:
             del kwds["cmap"]
         if dot_max is None:
-            dot_max = np.ceil(max(frac) * 10) / 10
+            dot_max = np.ceil(np.nanmax(frac) * 10) / 10
         else:
             if dot_max < 0 or dot_max > 1:
                 raise ValueError("`dot_max` value has to be between 0 and 1")

From 32b0687e976013db52df9f6074edd6da31c7d299 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 16:11:35 +0000
Subject: [PATCH 18/26] Fix row formatting for stacked_violin

---
 scanpy/plotting/_baseplot_class.py |  2 --
 scanpy/plotting/_stacked_violin.py | 17 +++++++++++------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index a877fa5eb8..d2f87d0b1e 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -122,7 +122,6 @@ def __init__(
             layer=layer,
             gene_symbols=gene_symbols,
         )
-        
         # reset obs_tidy if using groupby_cols
         if len(self.groupby_cols) > 0:
             # TODO : Check if we rather need the product of categories ?
@@ -586,7 +585,6 @@ def _plot_legend(self, legend_ax, return_ax_dict, normalize):
     def _mainplot(self, ax):
         y_labels = self.categories
         x_labels = self.var_names
-
         if self.var_names_idx_order is not None:
             x_labels = [x_labels[x] for x in self.var_names_idx_order]
 
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 370c018b8c..f19313d5e3 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Literal
+from warnings import warn
 
 import numpy as np
 import pandas as pd
@@ -327,14 +328,14 @@ def _mainplot(self, ax):
             _matrix.index = _matrix.index.reorder_categories(
                 self.categories_order, ordered=True
             )
-
         # get mean values for color and transform to color values
         # using colormap
         _color_df = _matrix.groupby(level=0).median()
-        if self.are_axes_swapped:
-            _color_df = _color_df.T
         if len(self.groupby_cols) > 0:
             _color_df = self._convert_tidy_to_stacked(_color_df)
+            self.stacked_violin_col_order = _color_df.columns
+        if self.are_axes_swapped:
+            _color_df = _color_df.T
 
         cmap = plt.get_cmap(self.kwds.get("cmap", self.cmap))
         if "cmap" in self.kwds:
@@ -404,8 +405,12 @@ def _make_rows_of_violinplots(
         # All columns should have a unique name, yet, frequently
         # gene names are repeated in self.var_names,  otherwise the
         # violin plot will not distinguish those genes
-        _matrix.columns = [f"{x}_{idx}" for idx, x in enumerate(_matrix.columns)]
-
+        # _matrix.columns = [f"{x}_{idx}" for idx, x in enumerate(_matrix.columns)] # added warning as this row was commented out
+        if _matrix.columns.value_counts().max()>1:
+            warn(
+                f"Duplicate gene or condition labels present."
+                "Results might be unexpected."
+            )
         # transform the  dataframe into a dataframe having three columns:
         # the categories name (from groupby),
         # the gene name
@@ -470,7 +475,6 @@ def _make_rows_of_violinplots(
                 palette_colors = colormap_array[idx, :]
             else:
                 palette_colors = None
-
             if not self.are_axes_swapped:
                 x = "genes"
                 _df = df[df.categories == row_label]
@@ -485,6 +489,7 @@ def _make_rows_of_violinplots(
                 x=x,
                 y="values",
                 data=_df,
+                order=self.stacked_violin_col_order,
                 orient="vertical",
                 ax=row_ax,
                 palette=palette_colors,

From 70ebe65716244b191e607035ed8cd151a30953c1 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 16:34:28 +0000
Subject: [PATCH 19/26] Fix stacked_violin swap_axes issue

---
 scanpy/plotting/_stacked_violin.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index f19313d5e3..471a9b7061 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -336,6 +336,7 @@ def _mainplot(self, ax):
             self.stacked_violin_col_order = _color_df.columns
         if self.are_axes_swapped:
             _color_df = _color_df.T
+            self.stacked_violin_col_order = _color_df.columns
 
         cmap = plt.get_cmap(self.kwds.get("cmap", self.cmap))
         if "cmap" in self.kwds:
@@ -484,7 +485,7 @@ def _make_rows_of_violinplots(
                 # we need to use this instead of the 'row_label'
                 # (in _color_df the values are not renamed as those
                 # values will be used to label the ticks)
-                _df = df[df.genes == _matrix.columns[idx]]
+                _df = df[df.genes == row_label]
             row_ax = sns.violinplot(
                 x=x,
                 y="values",

From 824d5070e4fd61f3eb70ee0305510115dc88a4e5 Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 16:50:55 +0000
Subject: [PATCH 20/26] Fix when groupby_cols missing

---
 scanpy/plotting/_stacked_violin.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 471a9b7061..90eb7d824e 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -333,7 +333,7 @@ def _mainplot(self, ax):
         _color_df = _matrix.groupby(level=0).median()
         if len(self.groupby_cols) > 0:
             _color_df = self._convert_tidy_to_stacked(_color_df)
-            self.stacked_violin_col_order = _color_df.columns
+        self.stacked_violin_col_order = _color_df.columns
         if self.are_axes_swapped:
             _color_df = _color_df.T
             self.stacked_violin_col_order = _color_df.columns

From 4d6800a7d130b7eff9bfae3a065119e7a1e7acca Mon Sep 17 00:00:00 2001
From: LMD <lisa.dratva@gmail.com>
Date: Fri, 1 Dec 2023 17:00:54 +0000
Subject: [PATCH 21/26] ValueError for overlapping arguments

---
 scanpy/plotting/_baseplot_class.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index d2f87d0b1e..15d962b8b7 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -124,6 +124,10 @@ def __init__(
         )
         # reset obs_tidy if using groupby_cols
         if len(self.groupby_cols) > 0:
+            if len(set(self.groupby).intersection(set(self.groupby_cols)))>0:
+                raise ValueError(
+                    f"`groupby` and `groupby_cols` have overlapping elements: {set(self.groupby).intersection(set(self.groupby_cols))}."
+                )
             # TODO : Check if we rather need the product of categories ?
             self.categories_cols = adata.obs.loc[:,self.groupby_cols].nunique().sum()
             _, self.obs_tidy = _prepare_dataframe(

From 91590f049c85922a37dbd579ae30c8d446c58a38 Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Mon, 18 Mar 2024 15:00:49 +0100
Subject: [PATCH 22/26] Fixup merge

---
 scanpy/plotting/_dotplot.py        |  9 ++++----
 scanpy/plotting/_matrixplot.py     |  6 ++---
 scanpy/plotting/_stacked_violin.py | 37 +++++++++++++++++++++---------
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 23630b78e7..a30c5fce80 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -141,7 +141,7 @@ def __init__(
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
         *,
-        groupby_cols: str | Sequence[str] = [],
+        groupby_cols: str | Sequence[str] = (),
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -170,7 +170,7 @@ def __init__(
             adata,
             var_names,
             groupby,
-            groupby_cols,
+            groupby_cols=groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -763,7 +763,8 @@ def _dotplot(
         # rescale size to match smallest_dot and largest_dot
         size = size * (largest_dot - smallest_dot) + smallest_dot
         normalize = check_colornorm(vmin, vmax, vcenter, norm)
-        normalize(mean_flat[~np.isnan(mean_flat)])  # circumvent unexpected behavior with nan in matplotlib
+        # circumvent unexpected behavior with nan in matplotlib
+        normalize(mean_flat[~np.isnan(mean_flat)])
 
         if color_on == "square":
             if edge_color is None:
@@ -876,7 +877,7 @@ def dotplot(
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
     *,
-    groupby_cols: str | Sequence[str] = [],
+    groupby_cols: str | Sequence[str] = (),
     use_raw: bool | None = None,
     log: bool = False,
     num_categories: int = 7,
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 2a9d13d139..05ded8a895 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -122,7 +122,7 @@ def __init__(
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
         *,
-        groupby_cols: str | Sequence[str] = [],
+        groupby_cols: str | Sequence[str] = (),
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -148,7 +148,7 @@ def __init__(
             adata,
             var_names,
             groupby,
-            groupby_cols,
+            groupby_cols=groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -343,7 +343,7 @@ def matrixplot(
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
     *,
-    groupby_cols: str | Sequence[str] = [],
+    groupby_cols: str | Sequence[str] = (),
     use_raw: bool | None = None,
     log: bool = False,
     num_categories: int = 7,
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 42f982723f..d950e02c39 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -177,7 +177,7 @@ def __init__(
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
         *,
-        groupby_cols: str | Sequence[str] = [],
+        groupby_cols: str | Sequence[str] = (),
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
@@ -202,7 +202,7 @@ def __init__(
             adata,
             var_names,
             groupby,
-            groupby_cols,
+            groupby_cols=groupby_cols,
             use_raw=use_raw,
             log=log,
             num_categories=num_categories,
@@ -413,7 +413,9 @@ def _mainplot(self, ax):
             self.vboundnorm.vcenter,
             self.vboundnorm.norm,
         )
-        normalize(_color_df.values[~np.isnan(_color_df.values)])  # circumvent unexpected behavior with nan in matplotlib
+        normalize(
+            _color_df.values[~np.isnan(_color_df.values)]
+        )  # circumvent unexpected behavior with nan in matplotlib
         colormap_array = cmap(normalize(_color_df.values))
         x_spacer_size = self.plot_x_padding
         y_spacer_size = self.plot_y_padding
@@ -471,9 +473,9 @@ def _make_rows_of_violinplots(
         # gene names are repeated in self.var_names,  otherwise the
         # violin plot will not distinguish those genes
         # _matrix.columns = [f"{x}_{idx}" for idx, x in enumerate(_matrix.columns)] # added warning as this row was commented out
-        if _matrix.columns.value_counts().max()>1:
+        if _matrix.columns.value_counts().max() > 1:
             warn(
-                f"Duplicate gene or condition labels present."
+                "Duplicate gene or condition labels present."
                 "Results might be unexpected."
             )
         # transform the  dataframe into a dataframe having three columns:
@@ -498,6 +500,7 @@ def _make_rows_of_violinplots(
                         0: "values",
                     }
                 )
+            )
             df["genes"] = (
                 df["genes"].astype("category").cat.reorder_categories(_matrix.columns)
             )
@@ -511,11 +514,23 @@ def _make_rows_of_violinplots(
             label = _matrix.index.name
             stacked_df = _matrix.reset_index()
             stacked_df.index = pd.MultiIndex.from_tuples(
-                stacked_df[label].str.split('_').tolist(), names=self.groupby + self.groupby_cols)
-            stacked_df = stacked_df.drop(label, axis=1).reset_index().melt(id_vars=self.groupby + self.groupby_cols)
-            stacked_df['genes'] = stacked_df[['variable'] + self.groupby_cols].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
-            stacked_df['categories'] = stacked_df[self.groupby].apply(lambda row: '_'.join(row.values.astype(str)), axis=1)
-            stacked_df = stacked_df.drop(self.groupby + self.groupby_cols + ['variable'], axis=1).rename(columns={'value':'values'})
+                stacked_df[label].str.split("_").tolist(),
+                names=self.groupby + self.groupby_cols,
+            )
+            stacked_df = (
+                stacked_df.drop(label, axis=1)
+                .reset_index()
+                .melt(id_vars=self.groupby + self.groupby_cols)
+            )
+            stacked_df["genes"] = stacked_df[["variable"] + self.groupby_cols].apply(
+                lambda row: "_".join(row.values.astype(str)), axis=1
+            )
+            stacked_df["categories"] = stacked_df[self.groupby].apply(
+                lambda row: "_".join(row.values.astype(str)), axis=1
+            )
+            stacked_df = stacked_df.drop(
+                self.groupby + self.groupby_cols + ["variable"], axis=1
+            ).rename(columns={"value": "values"})
             df = stacked_df
 
         # the ax need to be subdivided
@@ -671,7 +686,7 @@ def stacked_violin(
     var_names: _VarNames | Mapping[str, _VarNames],
     groupby: str | Sequence[str],
     *,
-    groupby_cols: str | Sequence[str] = [],
+    groupby_cols: str | Sequence[str] = (),
     log: bool = False,
     use_raw: bool | None = None,
     num_categories: int = 7,

From 3de4f910902cf5d95712a88f3b9a467aad57e1f0 Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Mon, 18 Mar 2024 15:04:56 +0100
Subject: [PATCH 23/26] Smaller diff

---
 scanpy/plotting/_baseplot_class.py | 44 ++++++++++++++++++++++--------
 scanpy/plotting/_dotplot.py        |  2 ++
 scanpy/plotting/_matrixplot.py     |  1 +
 scanpy/plotting/_stacked_violin.py |  1 +
 4 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 5e3aed68e1..9dbfdd70dd 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -123,7 +123,9 @@ def __init__(
         self.var_group_rotation = var_group_rotation
         self.width, self.height = figsize if figsize is not None else (None, None)
         self.groupby = [groupby] if isinstance(groupby, str) else groupby
-        self.groupby_cols = [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
+        self.groupby_cols = (
+            [groupby_cols] if isinstance(groupby_cols, str) else groupby_cols
+        )
         self.has_var_groups = (
             True
             if var_group_positions is not None and len(var_group_positions) > 0
@@ -144,12 +146,12 @@ def __init__(
         )
         # reset obs_tidy if using groupby_cols
         if len(self.groupby_cols) > 0:
-            if len(set(self.groupby).intersection(set(self.groupby_cols)))>0:
+            if overlap := (set(self.groupby) & set(self.groupby_cols)):
                 raise ValueError(
-                    f"`groupby` and `groupby_cols` have overlapping elements: {set(self.groupby).intersection(set(self.groupby_cols))}."
+                    f"`groupby` and `groupby_cols` have overlapping elements: {overlap}."
                 )
             # TODO : Check if we rather need the product of categories ?
-            self.categories_cols = adata.obs.loc[:,self.groupby_cols].nunique().sum()
+            self.categories_cols = adata.obs.loc[:, self.groupby_cols].nunique().sum()
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
@@ -394,8 +396,11 @@ def add_totals(
         _sort = True if sort is not None else False
         _ascending = True if sort == "ascending" else False
         counts_df = self.obs_tidy.index.value_counts(sort=_sort, ascending=_ascending)
-        if len(self.groupby_cols) > 0:  # could remove the previous line and only use this but this is slower
-            counts_df = self.adata.obs[self.groupby].value_counts(sort=_sort, ascending=_ascending)
+        # could remove the previous line and only use this but this is slower
+        if len(self.groupby_cols) > 0:
+            counts_df = self.adata.obs[self.groupby].value_counts(
+                sort=_sort, ascending=_ascending
+            )
 
         if _sort:
             self.categories_order = counts_df.index
@@ -613,6 +618,7 @@ def _plot_legend(self, legend_ax, return_ax_dict, normalize):
     def _mainplot(self, ax):
         y_labels = self.categories
         x_labels = self.var_names
+
         if self.var_names_idx_order is not None:
             x_labels = [x_labels[x] for x in self.var_names_idx_order]
 
@@ -678,7 +684,8 @@ def make_figure(self):
         if self.height is None:
             mainplot_height = len(self.categories) * category_height
             mainplot_width = (
-                len(self.var_names) * category_width  * (1+self.categories_cols) + self.group_extra_size
+                len(self.var_names) * category_width * (1 + self.categories_cols)
+                + self.group_extra_size
             )
             if self.are_axes_swapped:
                 mainplot_height, mainplot_width = mainplot_width, mainplot_height
@@ -880,23 +887,35 @@ def savefig(self, filename: str, bbox_inches: str | None = "tight", **kwargs):
         self.make_figure()
         plt.savefig(filename, bbox_inches=bbox_inches, **kwargs)
 
-    def _convert_tidy_to_stacked(self, values_df):
+    def _convert_tidy_to_stacked(self, values_df: pd.DataFrame) -> pd.DataFrame:
         """\
         Utility function used to convert obs_tidy into the correct format when using a groupby_col.
         """
         label = values_df.index.name
         stacked_df = values_df.reset_index()
         stacked_df.index = pd.MultiIndex.from_tuples(
-            stacked_df[label].str.split('_').tolist(), names=self.groupby + self.groupby_cols)
+            stacked_df[label].str.split("_").tolist(),
+            names=self.groupby + self.groupby_cols,
+        )
         stacked_df = stacked_df.drop(label, axis=1).unstack(level=self.groupby_cols)
 
         # recreate the original formatting of values_df
         values_df = stacked_df.reset_index(drop=True)
         if isinstance(stacked_df.index, pd.MultiIndex):
-            values_df.index = stacked_df.index.to_series().apply(lambda x: '_'.join(map(str, x))).values
+            values_df.index = (
+                stacked_df.index.to_series()
+                .apply(lambda x: "_".join(map(str, x)))
+                .values
+            )
         else:
-            values_df.index = stacked_df.index.to_series().apply(lambda x: ''.join(map(str, x))).values
-        values_df.columns = stacked_df.columns.to_series().apply(lambda x: '_'.join(map(str, x))).values
+            values_df.index = (
+                stacked_df.index.to_series()
+                .apply(lambda x: "".join(map(str, x)))
+                .values
+            )
+        values_df.columns = (
+            stacked_df.columns.to_series().apply(lambda x: "_".join(map(str, x))).values
+        )
         return values_df
 
     def _reorder_categories_after_dendrogram(self, dendrogram) -> None:
@@ -926,6 +945,7 @@ def _format_first_three_categories(_categories):
             return ", ".join(_categories)
 
         key = _get_dendrogram_key(self.adata, dendrogram, self.groupby)
+
         dendro_info = self.adata.uns[key]
         if self.groupby != dendro_info["groupby"]:
             raise ValueError(
diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index a30c5fce80..68aec78b93 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -208,6 +208,7 @@ def __init__(
             )
             if len(groupby_cols) > 0:
                 dot_size_df = self._convert_tidy_to_stacked(dot_size_df)
+
         if dot_color_df is None:
             # 2. compute mean expression value value
             if mean_only_expressed:
@@ -806,6 +807,7 @@ def _dotplot(
                 linewidth=edge_lw,
                 edgecolor=edge_color,
             )
+
             dot_ax.scatter(x, y, **kwds)
 
         y_ticks = np.arange(dot_color.shape[0]) + 0.5
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 05ded8a895..628271065d 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -167,6 +167,7 @@ def __init__(
             norm=norm,
             **kwds,
         )
+
         if values_df is None:
             # compute mean value
             values_df = (
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index d950e02c39..d8a15151e6 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -568,6 +568,7 @@ def _make_rows_of_violinplots(
                 # (in _color_df the values are not renamed as those
                 # values will be used to label the ticks)
                 _df = df[df.genes == row_label]
+
             row_ax = sns.violinplot(
                 x=x,
                 y="values",

From 5d57811b1e0bbc29518464367ccb17076ed1e91b Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Mon, 18 Mar 2024 15:07:46 +0100
Subject: [PATCH 24/26] Missed some

---
 scanpy/plotting/_baseplot_class.py | 2 +-
 scanpy/plotting/_stacked_violin.py | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 9dbfdd70dd..8608c6a315 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -98,7 +98,7 @@ def __init__(
         var_names: _VarNames | Mapping[str, _VarNames],
         groupby: str | Sequence[str],
         *,
-        groupby_cols: str | Sequence[str] = [],
+        groupby_cols: str | Sequence[str] = (),
         use_raw: bool | None = None,
         log: bool = False,
         num_categories: int = 7,
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index d8a15151e6..0f0c4da037 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -413,9 +413,8 @@ def _mainplot(self, ax):
             self.vboundnorm.vcenter,
             self.vboundnorm.norm,
         )
-        normalize(
-            _color_df.values[~np.isnan(_color_df.values)]
-        )  # circumvent unexpected behavior with nan in matplotlib
+        # circumvent unexpected behavior with nan in matplotlib
+        normalize(_color_df.values[~np.isnan(_color_df.values)])
         colormap_array = cmap(normalize(_color_df.values))
         x_spacer_size = self.plot_x_padding
         y_spacer_size = self.plot_y_padding

From e6766df3c49a8b80824f8f1e572603cf4871ec93 Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Thu, 21 Mar 2024 14:59:19 +0100
Subject: [PATCH 25/26] Cleanup

---
 scanpy/plotting/_anndata.py         |  4 ++--
 scanpy/plotting/_baseplot_class.py  |  4 ++--
 scanpy/plotting/_dotplot.py         |  2 +-
 scanpy/plotting/_matrixplot.py      |  2 +-
 scanpy/plotting/_stacked_violin.py  | 32 ++++++++++++++++-------------
 scanpy/plotting/_tools/__init__.py  |  5 +++--
 scanpy/testing/_helpers/__init__.py |  2 ++
 7 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/scanpy/plotting/_anndata.py b/scanpy/plotting/_anndata.py
index 970e56a27c..ba48b03d62 100755
--- a/scanpy/plotting/_anndata.py
+++ b/scanpy/plotting/_anndata.py
@@ -2060,7 +2060,7 @@ def _prepare_dataframe(
         # and does not need to be given
         groupby = groupby.copy()  # copy to not modify user passed parameter
         groupby.remove(groupby_index)
-    keys = list(groupby) + list(np.unique(var_names))
+    keys = [*groupby, *np.unique(var_names)]
     obs_tidy = get.obs_df(
         adata, keys=keys, layer=layer, use_raw=use_raw, gene_symbols=gene_symbols
     )
@@ -2097,7 +2097,7 @@ def _prepare_dataframe(
         categorical = categorical.cat.reorder_categories(
             sorted(categorical.cat.categories, key=lambda x: order[x])
         )
-    obs_tidy = obs_tidy[var_names].set_index(categorical)
+    obs_tidy = obs_tidy[np.unique(var_names)].set_index(categorical)
     categories = obs_tidy.index.categories
 
     if log:
diff --git a/scanpy/plotting/_baseplot_class.py b/scanpy/plotting/_baseplot_class.py
index 8608c6a315..7ad37b0d52 100644
--- a/scanpy/plotting/_baseplot_class.py
+++ b/scanpy/plotting/_baseplot_class.py
@@ -155,7 +155,7 @@ def __init__(
             _, self.obs_tidy = _prepare_dataframe(
                 adata,
                 self.var_names,
-                self.groupby + self.groupby_cols,
+                [*self.groupby, *self.groupby_cols],
                 use_raw,
                 log,
                 num_categories,
@@ -615,7 +615,7 @@ def _plot_legend(self, legend_ax, return_ax_dict, normalize):
         self._plot_colorbar(color_legend_ax, normalize)
         return_ax_dict["color_legend_ax"] = color_legend_ax
 
-    def _mainplot(self, ax):
+    def _mainplot(self, ax: Axes):
         y_labels = self.categories
         x_labels = self.var_names
 
diff --git a/scanpy/plotting/_dotplot.py b/scanpy/plotting/_dotplot.py
index 68aec78b93..d6464aa92f 100644
--- a/scanpy/plotting/_dotplot.py
+++ b/scanpy/plotting/_dotplot.py
@@ -574,7 +574,7 @@ def _plot_legend(self, legend_ax, return_ax_dict, normalize):
             self._plot_colorbar(color_legend_ax, normalize)
             return_ax_dict["color_legend_ax"] = color_legend_ax
 
-    def _mainplot(self, ax):
+    def _mainplot(self, ax: Axes):
         # work on a copy of the dataframes. This is to avoid changes
         # on the original data frames after repetitive calls to the
         # DotPlot object, for example once with swap_axes and other without
diff --git a/scanpy/plotting/_matrixplot.py b/scanpy/plotting/_matrixplot.py
index 628271065d..c970cb1692 100644
--- a/scanpy/plotting/_matrixplot.py
+++ b/scanpy/plotting/_matrixplot.py
@@ -257,7 +257,7 @@ def style(
 
         return self
 
-    def _mainplot(self, ax):
+    def _mainplot(self, ax: Axes):
         # work on a copy of the dataframes. This is to avoid changes
         # on the original data frames after repetitive calls to the
         # MatrixPlot object, for example once with swap_axes and other without
diff --git a/scanpy/plotting/_stacked_violin.py b/scanpy/plotting/_stacked_violin.py
index 0f0c4da037..3829170b8c 100644
--- a/scanpy/plotting/_stacked_violin.py
+++ b/scanpy/plotting/_stacked_violin.py
@@ -25,12 +25,10 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import (
-        Mapping,  # Special
-        Sequence,  # ABCs
-    )
+    from collections.abc import Mapping, Sequence
 
     from anndata import AnnData
+    from matplotlib.axes import Axes
 
 
 @_doc_params(common_plot_args=doc_common_plot_args)
@@ -374,7 +372,7 @@ def style(
 
         return self
 
-    def _mainplot(self, ax):
+    def _mainplot(self, ax: Axes):
         # to make the stacked violin plots, the
         # `ax` is subdivided horizontally and in each horizontal sub ax
         # a seaborn violin plot is added.
@@ -453,7 +451,13 @@ def _mainplot(self, ax):
         return normalize
 
     def _make_rows_of_violinplots(
-        self, ax, _matrix, colormap_array, _color_df, x_spacer_size, y_spacer_size
+        self,
+        ax: Axes,
+        _matrix: pd.DataFrame,
+        colormap_array: np.ndarray,
+        _color_df: pd.DataFrame,
+        x_spacer_size: float | None,
+        y_spacer_size: float | None,
     ):
         import seaborn as sns  # Slow import, only import if called
 
@@ -477,16 +481,16 @@ def _make_rows_of_violinplots(
                 "Duplicate gene or condition labels present."
                 "Results might be unexpected."
             )
-        # transform the  dataframe into a dataframe having three columns:
+        # transform the dataframe into a dataframe having three columns:
         # the categories name (from groupby),
         # the gene name
         # the expression value
         # This format is convenient to aggregate per gene or per category
         # while making the violin plots.
         if Version(pd.__version__) >= Version("2.1"):
-            stack_kwargs = {"future_stack": True}
+            stack_kwargs = dict(future_stack=True)
         else:
-            stack_kwargs = {"dropna": False}
+            stack_kwargs = dict(dropna=False)
 
         if len(self.groupby_cols) < 1:
             df = (
@@ -510,25 +514,25 @@ def _make_rows_of_violinplots(
             )
         else:
             # partially taken from self._convert_tidy_to_stacked
-            label = _matrix.index.name
+            label: str = _matrix.index.name
             stacked_df = _matrix.reset_index()
             stacked_df.index = pd.MultiIndex.from_tuples(
                 stacked_df[label].str.split("_").tolist(),
-                names=self.groupby + self.groupby_cols,
+                names=[*self.groupby, *self.groupby_cols],
             )
             stacked_df = (
                 stacked_df.drop(label, axis=1)
                 .reset_index()
-                .melt(id_vars=self.groupby + self.groupby_cols)
+                .melt(id_vars=[*self.groupby, *self.groupby_cols])
             )
-            stacked_df["genes"] = stacked_df[["variable"] + self.groupby_cols].apply(
+            stacked_df["genes"] = stacked_df[["variable", *self.groupby_cols]].apply(
                 lambda row: "_".join(row.values.astype(str)), axis=1
             )
             stacked_df["categories"] = stacked_df[self.groupby].apply(
                 lambda row: "_".join(row.values.astype(str)), axis=1
             )
             stacked_df = stacked_df.drop(
-                self.groupby + self.groupby_cols + ["variable"], axis=1
+                [*self.groupby, *self.groupby_cols, "variable"], axis=1
             ).rename(columns={"value": "values"})
             df = stacked_df
 
diff --git a/scanpy/plotting/_tools/__init__.py b/scanpy/plotting/_tools/__init__.py
index a55b5c4262..b2bc932e7d 100644
--- a/scanpy/plotting/_tools/__init__.py
+++ b/scanpy/plotting/_tools/__init__.py
@@ -1002,8 +1002,9 @@ def rank_genes_groups_stacked_violin(
     >>> adata = sc.datasets.pbmc68k_reduced()
     >>> sc.tl.rank_genes_groups(adata, 'bulk_labels')
 
-    >>> sc.pl.rank_genes_groups_stacked_violin(adata, n_genes=4,
-    ... min_logfoldchange=4, figsize=(8,6))
+    >>> sc.pl.rank_genes_groups_stacked_violin(
+    ...     adata, n_genes=4, min_logfoldchange=4, figsize=(8,6)
+    ... )
 
     """
 
diff --git a/scanpy/testing/_helpers/__init__.py b/scanpy/testing/_helpers/__init__.py
index b85faec3f3..6d44e39301 100644
--- a/scanpy/testing/_helpers/__init__.py
+++ b/scanpy/testing/_helpers/__init__.py
@@ -12,6 +12,8 @@
 
 import scanpy as sc
 
+from . import data  # noqa: F401
+
 # TODO: Report more context on the fields being compared on error
 # TODO: Allow specifying paths to ignore on comparison
 

From 1c4740ee3192e20a9dc8219cfbfb5b2d6bf3f06e Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Thu, 21 Mar 2024 15:31:48 +0100
Subject: [PATCH 26/26] undo

---
 scanpy/plotting/_anndata.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scanpy/plotting/_anndata.py b/scanpy/plotting/_anndata.py
index ba48b03d62..6daa595e44 100755
--- a/scanpy/plotting/_anndata.py
+++ b/scanpy/plotting/_anndata.py
@@ -2097,7 +2097,7 @@ def _prepare_dataframe(
         categorical = categorical.cat.reorder_categories(
             sorted(categorical.cat.categories, key=lambda x: order[x])
         )
-    obs_tidy = obs_tidy[np.unique(var_names)].set_index(categorical)
+    obs_tidy = obs_tidy[var_names].set_index(categorical)
     categories = obs_tidy.index.categories
 
     if log: