Source code for aeolus.plot.mpl

"""Matplotlib-related utilities."""
import matplotlib.colors as mcolors
from matplotlib.lines import Line2D

import numpy as np

from ..exceptions import _warn


__all__ = ("MidpointNormalize", "add_custom_legend")


class MidpointNormalize(mcolors.Normalize):
    """Normalise data around a midpoint."""

    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):  # noqa
        _warn(
            "MidpointNormalize is deprecated and will be removed in the next release. "
            "Use matplotlib.colors.CenteredNorm instead."
        )
        self.midpoint = midpoint
        mcolors.Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):  # noqa
        # Ignoring masked values and all kinds of edge cases
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y))


[docs]def add_custom_legend(ax, styles_and_labels, **leg_kw): """ Add a custom legend to matplotlib axes. Parameters ---------- ax: matplotlib.axes._subplots.AxesSubplot Axes where to put the legend. styles_and_labels: dict Dictionary with labels as keys and a dictionary of plot keywords as values. leg_kw: dict, optional Keyword arguments passed to `legend()` function. Example ------- >>> import matplotlib.pyplot as plt >>> ax = plt.axes() >>> my_dict = dict(foo=dict(color='C0', marker="X"), bar=dict(color='C1', marker="o")) >>> add_custom_legend(ax, my_dict, loc=2, title="blah") """ lines = [Line2D([0], [0], **style) for style in styles_and_labels.values()] leg = ax.legend(lines, styles_and_labels.keys(), **leg_kw) if ax.legend_ is not None: ax.add_artist(leg)