ENH: Allow elementwise coloring in background_gradient with axis=None #15204

Author: soxofaan

doc/source/whatsnew/v0.24.0.txt

@@ -14,6 +14,8 @@ Other Enhancements
 ^^^^^^^^^^^^^^^^^^
 - :func:`to_datetime` now supports the ``%Z`` and ``%z`` directive when passed into ``format`` (:issue:`13486`)
 - :func:`to_csv` now supports ``compression`` keyword when a file handle is passed. (:issue:`21227`)
+- Allow elementwise coloring in ``style.background_gradient`` with ``axis=None`` (:issue:`15204`)
+-
 -
 
 .. _whatsnew_0240.api_breaking:

pandas/io/formats/style.py

@@ -901,16 +901,21 @@ def background_gradient(self, cmap='PuBu', low=0, high=0, axis=0,
     def _background_gradient(s, cmap='PuBu', low=0, high=0):
         """Color background in a range according to the data."""
         with _mpl(Styler.background_gradient) as (plt, colors):
-            rng = s.max() - s.min()
+            smin = s.values.min()
+            smax = s.values.max()
+            rng = smax - smin
             # extend lower / upper bounds, compresses color range
-            norm = colors.Normalize(s.min() - (rng * low),
-                                    s.max() + (rng * high))
-            # matplotlib modifies inplace?
+            norm = colors.Normalize(smin - (rng * low), smax + (rng * high))
+            # matplotlib colors.Normalize modifies inplace?
             # https://github.com/matplotlib/matplotlib/issues/5427
-            normed = norm(s.values)
-            c = [colors.rgb2hex(x) for x in plt.cm.get_cmap(cmap)(normed)]
-            return ['background-color: {color}'.format(color=color)
-                    for color in c]
+            rgbs = plt.cm.get_cmap(cmap)(norm(s.values))
+
+            css = lambda x: 'background-color: {c}'.format(c=colors.rgb2hex(x))
+            if s.ndim == 1:
+                return [css(x) for x in rgbs]
+            else:
+                return pd.DataFrame([[css(x) for x in row] for row in rgbs],
+                                    index=s.index, columns=s.columns)
 
     def set_properties(self, subset=None, **kwargs):
         """

pandas/tests/io/formats/test_style.py

@@ -1033,6 +1033,34 @@ def test_background_gradient(self):
             subset=pd.IndexSlice[1, 'A'])._compute().ctx
         assert result[(1, 0)] == ['background-color: #fff7fb']
 
+    @td.skip_if_no_mpl
+    def test_background_gradient_axis(self):
+        df = pd.DataFrame([[1, 2], [2, 4]], columns=['A', 'B'])
+
+        low = ['background-color: #f7fbff']
+        high = ['background-color: #08306b']
+        mid = ['background-color: #abd0e6']
+        result = df.style.background_gradient(cmap='Blues',
+                                              axis=0)._compute().ctx
+        assert result[(0, 0)] == low
+        assert result[(0, 1)] == low
+        assert result[(1, 0)] == high
+        assert result[(1, 1)] == high
+
+        result = df.style.background_gradient(cmap='Blues',
+                                              axis=1)._compute().ctx
+        assert result[(0, 0)] == low
+        assert result[(0, 1)] == high
+        assert result[(1, 0)] == low
+        assert result[(1, 1)] == high
+
+        result = df.style.background_gradient(cmap='Blues',
+                                              axis=None)._compute().ctx
+        assert result[(0, 0)] == low
+        assert result[(0, 1)] == mid
+        assert result[(1, 0)] == mid
+        assert result[(1, 1)] == high
+
 
 def test_block_names():
     # catch accidental removal of a block