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

Author: soxofaan

doc/source/whatsnew/v0.24.0.txt

@@ -16,7 +16,9 @@ Other Enhancements
 - :func:`Series.mode` and :func:`DataFrame.mode` now support the ``dropna`` parameter which can be used to specify whether NaN/NaT values should be considered (:issue:`17534`)
 - :func:`to_csv` now supports ``compression`` keyword when a file handle is passed. (:issue:`21227`)
 - :meth:`Index.droplevel` is now implemented also for flat indexes, for compatibility with MultiIndex (:issue:`21115`)
-
+- Allow elementwise coloring in ``style.background_gradient`` with ``axis=None`` (:issue:`15204`)
+-
+-
 
 .. _whatsnew_0240.api_breaking:
 

pandas/io/formats/style.py

@@ -913,21 +913,22 @@ def background_gradient(self, cmap='PuBu', low=0, high=0, axis=0,
     def _background_gradient(s, cmap='PuBu', low=0, high=0,
                              text_color_threshold=0.408):
         """Color background in a range according to the data."""
+        if (not isinstance(text_color_threshold, (float, int)) or
+                not 0 <= text_color_threshold <= 1):
+            msg = "`text_color_threshold` must be a value from 0 to 1."
+            raise ValueError(msg)
+
         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)]
-            if (not isinstance(text_color_threshold, (float, int)) or
-                    not 0 <= text_color_threshold <= 1):
-                msg = "`text_color_threshold` must be a value from 0 to 1."
-                raise ValueError(msg)
+            rgbas = plt.cm.get_cmap(cmap)(norm(s.values))
 
-            def relative_luminance(color):
+            def relative_luminance(rgba):
                 """
                 Calculate relative luminance of a color.
 
@@ -936,25 +937,33 @@ def relative_luminance(color):
 
                 Parameters
                 ----------
-                color : matplotlib color
-                    Hex code, rgb-tuple, or HTML color name.
+                color : rgb or rgba tuple
 
                 Returns
                 -------
                 float
                     The relative luminance as a value from 0 to 1
                 """
-                rgb = colors.colorConverter.to_rgba_array(color)[:, :3]
-                rgb = np.where(rgb <= .03928, rgb / 12.92,
-                               ((rgb + .055) / 1.055) ** 2.4)
-                lum = rgb.dot([.2126, .7152, .0722])
-                return lum.item()
-
-            text_colors = ['#f1f1f1' if relative_luminance(x) <
-                           text_color_threshold else '#000000' for x in c]
-
-            return ['background-color: {color};color: {tc}'.format(
-                    color=color, tc=tc) for color, tc in zip(c, text_colors)]
+                r, g, b = (
+                    x / 12.92 if x <= 0.03928 else ((x + 0.055) / 1.055 ** 2.4)
+                    for x in rgba[:3]
+                )
+                return 0.2126 * r + 0.7152 * g + 0.0722 * b
+
+            def css(rgba):
+                dark = relative_luminance(rgba) < text_color_threshold
+                text_color = '#f1f1f1' if dark else '#000000'
+                return 'background-color: {b};color: {c};'.format(
+                    b=colors.rgb2hex(rgba), c=text_color
+                )
+
+            if s.ndim == 1:
+                return [css(rgba) for rgba in rgbas]
+            else:
+                return pd.DataFrame(
+                    [[css(rgba) for rgba in row] for row in rgbas],
+                    index=s.index, columns=s.columns
+                )
 
     def set_properties(self, subset=None, **kwargs):
         """

pandas/tests/io/formats/test_style.py

@@ -1056,6 +1056,34 @@ def test_text_color_threshold_raises(self, text_color_threshold):
             df.style.background_gradient(
                 text_color_threshold=text_color_threshold)._compute()
 
+    @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', 'color: #000000']
+        high = ['background-color: #08306b', 'color: #f1f1f1']
+        mid = ['background-color: #abd0e6', 'color: #000000']
+        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