aeolus.plot.mpl — aeolus 0.4.9+9.g5a71939 documentation

# 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))

    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)