Plot colorwheel inset in matplotlib (Python).

mpl_colorwheel.py

"""Color wheel for matplotlib plots."""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import inset_locator


def add_cwheel(ax_parent, rect, label, cmap=None, half=True, nticks=4):
    """Add color wheel to a matplotlib axes.

    Parameters
    ----------
    ax_parent : :class:`~matplotlib.axes.Axes`
        Parent axes for the colorwheel inset.
    rect : sequence of float
        The dimensions [left, bottom, width, height] of the inset axes.
        Quantities are normalised to parent axes width and height.
    label : str
        Label text for the color wheel.
    cmap : :class:`~matplotlib.colors.Colormap`, optional
        Colormap to use for the color wheel. Uses the 'hsv' colormap by
        default. Cyclic colormap recommended.
    half : bool, optional
        Draw a half-circle color fan if False (default), otherwise draw a full
        circle (NOT IMPLEMENTED YET).
    nticks : int, optional
        Number of ticks for the angle labels. Default is 4.

    """

    if cmap is None:
        cmap = plt.get_cmap('hsv')

    # Add colorwheel in 'sticky' axes for robust plot resizing.
    fig = ax_parent.get_figure()
    cwheel = fig.add_axes([0, 0, 1, 1], projection="polar", in_layout=False)
    anchored_pos = inset_locator.InsetPosition(ax_parent, rect)
    cwheel.set_axes_locator(anchored_pos)
    cwheel.set_theta_zero_location('N')
    cwheel.set_theta_direction(-1)

    if half:
        # Fiddle with tick lims and values.
        cwheel.set_thetamin(-90)
        cwheel.set_thetamax(-90)
        cwheel.set_xticks(
            np.linspace(np.deg2rad(-90), np.deg2rad(90), nticks)
        )
        thetavals = np.linspace(
            np.deg2rad(-90), np.deg2rad(90), cmap.N
        )
    else:
        raise NotImplementedError  # TODO: Support full 360 colorwheel.

    # Clean up stuff we don't need.
    cwheel.set_yticklabels([])
    cwheel.spines['polar'].set_visible(False)
    cwheel.set_facecolor((0, 0, 0, 0))  # Transparent cutout inside wheel.

    # Plot the color fan.
    rvals = np.linspace(0.5, 1, 2)
    _, thetagrid = np.meshgrid(rvals, thetavals)
    cwheel.pcolormesh(thetavals, rvals, thetagrid.T, cmap=cmap)

    if half:
        # Negative label padding to use the free space below the half circle.
        cwheel.set_xlabel(label, labelpad=-10)
    else:
        cwheel.set_xlabel(label)


def test_cwheel():
    fig, ax =plt.subplots()
    add_cwheel(ax, [0.75, 0.05, 0.15, 0.15], "Angle")
    fig.show()