viz.py 文件源码

def mpl_palette(name, n_colors=6, extrema=False, cycle=False):
    """Return discrete colors from a matplotlib palette.
    Note that this handles the qualitative colorbrewer palettes
    properly, although if you ask for more colors than a particular
    qualitative palette can provide you will get fewer than you are
    expecting.

    Parameters
    ----------
    name : string
        Name of the palette. This should be a named matplotlib colormap.
    n_colors : int
        Number of discrete colors in the palette.
    extrema : boolean
        If True, include the extrema of the palette.
    cycle : boolean
        If True, return a itertools.cycle.

    Returns
    -------
    palette : colormap or itertools.cycle
        List-like object of colors as RGB tuples
    """
    if name in SEABORN_PALETTES:
        palette = SEABORN_PALETTES[name]
        # Always return as many colors as we asked for
        pal_cycle = itertools.cycle(palette)
        palette = [next(pal_cycle) for _ in range(n_colors)]

    elif name in dir(mpl.cm) or name[:-2] in dir(mpl.cm):
        cmap = getattr(mpl.cm, name)

        if extrema:
            bins = np.linspace(0, 1, n_colors)
        else:
            bins = np.linspace(0, 1, n_colors * 2 - 1 + 2)[1:-1:2]
        palette = list(map(tuple, cmap(bins)[:, :3]))

    else:
        raise ValueError("%s is not a valid palette name" % name)

    if cycle:
        return itertools.cycle(palette)
    else:
        return palette