ENH: Allow elementwise coloring in background_gradient with axis=None

[dhirschfeld]

At present the axis argument to background_gradient only accepts 0 or 1 so that the resulting "heatmap" can only be applied columnwise or rowwise rather than to the whole table.

Signature: df.style.background_gradient(cmap='PuBu', low=0, high=0, axis=0, subset=None)

Parameters
----------
cmap: str or colormap
    matplotlib colormap
low, high: float
    compress the range by these values.
axis: int or str
    1 or 'columns' for colunwise, 0 or 'index' for rowwise
subset: IndexSlice
    a valid slice for ``data`` to limit the style application to

I propose that axis=None be allowed to mean applying a color gradient elementwise.

[dhirschfeld]

df = pd.DataFrame(np.sqrt((np.mgrid[0:11, 0:11]**2).sum(axis=0)))

Example of axis=0 coloring whereas I would like the color gradient to be radially increasing like the numbers

[dhirschfeld]

Note: This would be consistent with the axis argument to apply:

[dhirschfeld]

My workaround:

def background_gradient(self, cvals=None, cmin=None, cmax=None, cmap='viridis', **css):
    """For use with `DataFrame.style.apply` this function will apply a heatmap
    color gradient *elementwise* to the calling DataFrame

    Parameters
    ----------
    self : pd.DataFrame
        The calling DataFrame. This argument is automatically passed in by the
        `DataFrame.style.apply` method

    cvals : pd.DataFrame
        If specified this DataFrame is used to determine the color gradient

    cmin : float
        If specified, any values below this will be clipped to the bottom of
        the cmap

    cmax : float
        If specified, any values above this will be clipped to the top of
        the cmap

    cmap : colormap or str
        The colormap to use

    css : dict
        Any extra inline css key/value pars to pass to the styler

    Returns
    -------
    pd.DataFrame
        The css styles to apply

    """
    from matplotlib import cm
    from matplotlib.colors import rgb2hex
    if cvals is None:
        cvals = self.values.ravel().copy()
    else:
        assert cvals.shape == self.shape
        cvals = cvals.values.ravel().copy()
    cvals -= cmin or cvals.min()
    cvals /= cmax or cvals.max()
    cvals = cvals.clip(0, 1)
    styles = []
    for rgb in cm.viridis_r(cvals):
        style = [
            "{}: {}".format(key.replace('_', '-'), value)
            for key, value in css.items()
        ]
        style.append("background-color: {}".format(rgb2hex(rgb)))
        styles.append('; '.join(style))
    styles = np.asarray(styles).reshape(self.shape)
    return pd.DataFrame(styles, index=self.index, columns=self.columns)

[jorisvandenbossche]

Sounds logical, PR certainly welcome I think!
cc @TomAugspurger

[TomAugspurger]

Yep, that's be great if you could submit a PR.

[jreback]

this should be fairly generalized i think, iow where we accept axis

[jorisvandenbossche]

@dhirschfeld Interested to do a PR for this? (I also stumbled into a case where I would need this)

BTW, your ability to specify a vmin/vmax is also a nice feature I think.

[dhirschfeld]

@jorisvandenbossche - I'd like to finish it up by submitting a PR but it's unlikely I'll have time in at least the next couple of weeks so I'd be happy for anyone else interested to pick it up. In the interim, my implementation posted above is available for anyone who finds it...

I had a need to explicitly specify the cmin/cmax values so I implemented that directly rather than the high/low compression. I also thought the **css made for a very usable interactive api but concede that a case could be made that it was too magic