# This software is open source software available under the BSD-3 license.
#
# Copyright (c) 2019 Triad National Security, LLC. All rights reserved.
# Copyright (c) 2019 Lawrence Livermore National Security, LLC. All rights
# reserved.
# Copyright (c) 2019 UT-Battelle, LLC. All rights reserved.
#
# Additional copyright and license information can be found in the LICENSE file
# distributed with this code, or at
# https://raw.githubusercontent.com/MPAS-Dev/MPAS-Analysis/master/LICENSE
"""
Utilities for handling color maps and color bars
"""
# Authors
# -------
# Xylar Asay-Davis, Milena Veneziani, Luke Van Roekel, Greg Streletz
from __future__ import absolute_import, division, print_function, \
    unicode_literals
import matplotlib.pyplot as plt
import matplotlib.colors as cols
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
import xml.etree.ElementTree as ET
from six.moves import configparser
import cmocean
import pkg_resources
from six import string_types
[docs]def setup_colormap(config, configSectionName, suffix=''):
    '''
    Set up a colormap from the registry
    Parameters
    ----------
    config : instance of ConfigParser
        the configuration, containing a [plot] section with options that
        control plotting
    configSectionName : str
        name of config section
    suffix: str, optional
        suffix of colormap related options
    Returns
    -------
    colormapDict : dict
        A dictionary of colormap information.
        'colormap' specifies the name of the new colormap
        'norm' is a matplotlib norm object used to normalize the colormap
        'levels' is an array of contour levels or ``None`` if not using indexed
        color map
        'ticks' is an array of values where ticks should be placed
        'contours' is an array of contour values to plot or ``None`` if none
        have been specified
        'lineWidth' is the width of contour lines or ``None`` if not specified
        'lineColor' is the color of contour lines or ``None`` if not specified
    '''
    # Authors
    # -------
    # Xylar Asay-Davis, Milena Veneziani, Greg Streletz
    register_custom_colormaps()
    colormapType = config.get(configSectionName,
                              'colormapType{}'.format(suffix))
    if colormapType == 'indexed':
        (colormap, norm, levels, ticks) = _setup_indexed_colormap(
            config, configSectionName, suffix=suffix)
    elif colormapType == 'continuous':
        (colormap, norm, ticks) = _setup_colormap_and_norm(
            config, configSectionName, suffix=suffix)
        levels = None
    else:
        raise ValueError('config section {} option colormapType{} is not '
                         '"indexed" or "continuous"'.format(
                             configSectionName, suffix))
    option = 'contourLevels{}'.format(suffix)
    if config.has_option(configSectionName, option):
        contours = config.getExpression(configSectionName,
                                        option,
                                        usenumpyfunc=True)
        if isinstance(contours, string_types) and contours == 'none':
            contours = None
    else:
        contours = None
    option = 'contourThickness{}'.format(suffix)
    if config.has_option(configSectionName, option):
        lineWidth = config.getfloat(configSectionName, option)
    else:
        lineWidth = None
    option = 'contourColor{}'.format(suffix)
    if config.has_option(configSectionName, option):
        lineColor = config.get(configSectionName, option)
    else:
        lineColor = None
    return {'colormap': colormap, 'norm': norm, 'levels': levels,
            'ticks': ticks, 'contours': contours, 'lineWidth': lineWidth,
            'lineColor': lineColor} 
def register_custom_colormaps():
    name = 'ferret'
    backgroundColor = (0.9, 0.9, 0.9)
    red = np.array([[0, 0.6],
                    [0.15, 1],
                    [0.35, 1],
                    [0.65, 0],
                    [0.8, 0],
                    [1, 0.75]])
    green = np.array([[0, 0],
                      [0.1, 0],
                      [0.35, 1],
                      [1, 0]])
    blue = np.array([[0, 0],
                     [0.5, 0],
                     [0.9, 0.9],
                     [1, 0.9]])
    colorCount = 21
    colorList = np.ones((colorCount, 4), float)
    colorList[:, 0] = np.interp(np.linspace(0, 1, colorCount),
                                red[:, 0], red[:, 1])
    colorList[:, 1] = np.interp(np.linspace(0, 1, colorCount),
                                green[:, 0], green[:, 1])
    colorList[:, 2] = np.interp(np.linspace(0, 1, colorCount),
                                blue[:, 0], blue[:, 1])
    colorList = colorList[::-1, :]
    colorMap = cols.LinearSegmentedColormap.from_list(
        name, colorList, N=255)
    colorMap.set_bad(backgroundColor)
    _register_colormap_and_reverse(name, colorMap)
    name = 'erdc_iceFire_H'
    colorArray = np.array([
        [-1, 4.05432e-07, 0, 5.90122e-06],
        [-0.87451, 0, 0.120401, 0.302675],
        [-0.74902, 0, 0.216583, 0.524574],
        [-0.623529, 0.0552475, 0.345025, 0.6595],
        [-0.498039, 0.128047, 0.492588, 0.720288],
        [-0.372549, 0.188955, 0.641309, 0.792092],
        [-0.247059, 0.327673, 0.784935, 0.873434],
        [-0.121569, 0.60824, 0.892164, 0.935547],
        [0.00392157, 0.881371, 0.912178, 0.818099],
        [0.129412, 0.951407, 0.835621, 0.449279],
        [0.254902, 0.904481, 0.690489, 0],
        [0.380392, 0.85407, 0.510864, 0],
        [0.505882, 0.777093, 0.33018, 0.00088199],
        [0.631373, 0.672862, 0.139087, 0.00269398],
        [0.756863, 0.508815, 0, 0],
        [0.882353, 0.299417, 0.000366289, 0.000547829],
        [1, 0.0157519, 0.00332021, 4.55569e-08]], float)
    colorCount = 255
    colorList = np.ones((colorCount, 4), float)
    x = colorArray[:, 0]
    for cIndex in range(3):
        colorList[:, cIndex] = np.interp(
            np.linspace(-1., 1., colorCount),
            x, colorArray[:, cIndex + 1])
    colorMap = cols.LinearSegmentedColormap.from_list(
        name, colorList, N=255)
    _register_colormap_and_reverse(name, colorMap)
    name = 'erdc_iceFire_L'
    colorArray = np.array([
        [-1, 0.870485, 0.913768, 0.832905],
        [-0.87451, 0.586919, 0.887865, 0.934003],
        [-0.74902, 0.31583, 0.776442, 0.867858],
        [-0.623529, 0.18302, 0.632034, 0.787722],
        [-0.498039, 0.117909, 0.484134, 0.713825],
        [-0.372549, 0.0507239, 0.335979, 0.654741],
        [-0.247059, 0, 0.209874, 0.511832],
        [-0.121569, 0, 0.114689, 0.28935],
        [0.00392157, 0.0157519, 0.00332021, 4.55569e-08],
        [0.129412, 0.312914, 0, 0],
        [0.254902, 0.520865, 0, 0],
        [0.380392, 0.680105, 0.15255, 0.0025996],
        [0.505882, 0.785109, 0.339479, 0.000797922],
        [0.631373, 0.857354, 0.522494, 0],
        [0.756863, 0.910974, 0.699774, 0],
        [0.882353, 0.951921, 0.842817, 0.478545],
        [1, 0.881371, 0.912178, 0.818099]], float)
    colorCount = 255
    colorList = np.ones((colorCount, 4), float)
    x = colorArray[:, 0]
    for cIndex in range(3):
        colorList[:, cIndex] = np.interp(
            np.linspace(-1., 1., colorCount),
            x, colorArray[:, cIndex + 1])
    colorMap = cols.LinearSegmentedColormap.from_list(
        name, colorList, N=255)
    _register_colormap_and_reverse(name, colorMap)
    name = 'BuOr'
    colors1 = plt.cm.PuOr(np.linspace(0., 1, 256))
    colors2 = plt.cm.RdBu(np.linspace(0, 1, 256))
    # combine them and build a new colormap, just the orange from the first
    # and the blue from the second
    colorList = np.vstack((colors1[0:128, :], colors2[128:256, :]))
    # reverse the order
    colorList = colorList[::-1, :]
    colorMap = cols.LinearSegmentedColormap.from_list(name, colorList)
    _register_colormap_and_reverse(name, colorMap)
    name = 'Maximenko'
    colorArray = np.array([
        [-1, 0., 0.45882352941, 0.76470588235],
        [-0.666667, 0., 0.70196078431, 0.90588235294],
        [-0.333333, 0.3294117647, 0.87058823529, 1.],
        [0., 0.76470588235, 0.94509803921, 0.98039215686],
        [0.333333, 1., 1., 0.],
        [0.666667, 1., 0.29411764705, 0.],
        [1, 1., 0., 0.]], float)
    colorCount = 255
    colorList = np.ones((colorCount, 4), float)
    x = colorArray[:, 0]
    for cIndex in range(3):
        colorList[:, cIndex] = np.interp(
            np.linspace(-1., 1., colorCount),
            x, colorArray[:, cIndex + 1])
    colorMap = cols.LinearSegmentedColormap.from_list(
        name, colorList, N=255)
    _register_colormap_and_reverse(name, colorMap)
    # add the cmocean color maps
    mapNames = list(cmocean.cm.cmapnames)
    # don't bother with gray (already exists, I think)
    mapNames.pop(mapNames.index('gray'))
    for mapName in mapNames:
        _register_colormap_and_reverse(mapName, getattr(cmocean.cm, mapName))
    # add Scientific Colour-Maps 3.0 from
    # http://www.fabiocrameri.ch/colourmaps.php
    for mapName in ['berlin', 'bilbao', 'broc', 'cork', 'davos', 'devon',
                    'grayC', 'lajolla', 'lapaz', 'lisbon', 'oleron', 'oslo',
                    'roma', 'tofino', 'tokyo', 'turku', 'vik']:
        xmlFile = pkg_resources.resource_filename(
            __name__, 'ColourMapSuite3/{}/{}.xml'.format(mapName, mapName))
        _read_xml_colormap(xmlFile, mapName)
    # add SciVisColor colormaps from
    # https://sciviscolor.org/home/colormaps/
    for mapName in ['3wave-yellow-grey-blue', '3Wbgy5',
                    '4wave-grey-red-green-mgreen', '5wave-yellow-brown-blue',
                    'blue-1', 'blue-3', 'blue-6', 'blue-8', 'blue-orange-div',
                    'brown-2', 'brown-5', 'brown-8', 'green-1', 'green-4',
                    'green-7', 'green-8', 'orange-5', 'orange-6',
                    'orange-green-blue-gray', 'purple-7', 'purple-8', 'red-1',
                    'red-3', 'red-4', 'yellow-1', 'yellow-7']:
        xmlFile = pkg_resources.resource_filename(
            __name__, 'SciVisColorColormaps/{}.xml'.format(mapName))
        _read_xml_colormap(xmlFile, mapName)
    name = 'white_cmo_deep'
    # modify cmo.deep to start at white
    colors2 = plt.cm.get_cmap('cmo.deep')(np.linspace(0, 1, 224))
    colorCount = 32
    colors1 = np.ones((colorCount, 4), float)
    x = np.linspace(0., 1., colorCount+1)[0:-1]
    white = [1., 1., 1., 1.]
    for cIndex in range(4):
        colors1[:, cIndex] = np.interp(x, [0., 1.],
                                       [white[cIndex], colors2[0, cIndex]])
    colors = np.vstack((colors1, colors2))
    # generating a smoothly-varying LinearSegmentedColormap
    cmap = LinearSegmentedColormap.from_list(name, colors)
    _register_colormap_and_reverse(name, cmap)
def _setup_colormap_and_norm(config, configSectionName, suffix=''):
    '''
    Set up a colormap from the registry
    Parameters
    ----------
    config : instance of ConfigParser
        the configuration, containing a [plot] section with options that
        control plotting
    configSectionName : str
        name of config section
    suffix: str, optional
        suffix of colormap related options
    Returns
    -------
    colormap : srt
        new colormap
    norm : ``mapplotlib.colors.Normalize``
        the norm used to normalize the colormap
    ticks : array of float
        the tick marks on the colormap
    '''
    # Authors
    # -------
    # Xylar Asay-Davis
    register_custom_colormaps()
    colormap = plt.get_cmap(config.get(configSectionName,
                                       'colormapName{}'.format(suffix)))
    normType = config.get(configSectionName, 'normType{}'.format(suffix))
    kwargs = config.getExpression(configSectionName,
                                  'normArgs{}'.format(suffix))
    if normType == 'symLog':
        norm = cols.SymLogNorm(**kwargs)
    elif normType == 'log':
        norm = cols.LogNorm(**kwargs)
    elif normType == 'linear':
        norm = cols.Normalize(**kwargs)
    else:
        raise ValueError('Unsupported norm type {} in section {}'.format(
            normType, configSectionName))
    try:
        ticks = config.getExpression(
            configSectionName, 'colorbarTicks{}'.format(suffix),
            usenumpyfunc=True)
    except(configparser.NoOptionError):
        ticks = None
    return (colormap, norm, ticks)
def _setup_indexed_colormap(config, configSectionName, suffix=''):
    '''
    Set up a colormap from the registry
    Parameters
    ----------
    config : instance of ConfigParser
        the configuration, containing a [plot] section with options that
        control plotting
    configSectionName : str
        name of config section
    suffix: str, optional
        suffix of colormap related options
    colorMapType
    Returns
    -------
    colormap : srt
        new colormap
    norm : ``mapplotlib.colors.Normalize``
        the norm used to normalize the colormap
    ticks : array of float
        the tick marks on the colormap
    '''
    # Authors
    # -------
    # Xylar Asay-Davis, Milena Veneziani, Greg Streletz
    colormap = plt.get_cmap(config.get(configSectionName,
                                       'colormapName{}'.format(suffix)))
    indices = config.getExpression(configSectionName,
                                   'colormapIndices{}'.format(suffix),
                                   usenumpyfunc=True)
    try:
        levels = config.getExpression(
            configSectionName, 'colorbarLevels{}'.format(suffix),
            usenumpyfunc=True)
    except(configparser.NoOptionError):
        levels = None
    if levels is not None:
        # set under/over values based on the first/last indices in the colormap
        underColor = colormap(indices[0])
        overColor = colormap(indices[-1])
        if len(levels) + 1 == len(indices):
            # we have 2 extra values for the under/over so make the colormap
            # without these values
            indices = indices[1:-1]
        elif len(levels) - 1 != len(indices):
            # indices list must be either one element shorter
            # or one element longer than colorbarLevels list
            raise ValueError('length mismatch between indices and '
                             'colorbarLevels')
        colormap = cols.ListedColormap(colormap(indices),
                                       'colormapName{}'.format(suffix))
        colormap.set_under(underColor)
        colormap.set_over(overColor)
    norm = cols.BoundaryNorm(levels, colormap.N)
    try:
        ticks = config.getExpression(
            configSectionName, 'colorbarTicks{}'.format(suffix),
            usenumpyfunc=True)
    except(configparser.NoOptionError):
        ticks = levels
    return (colormap, norm, levels, ticks)
def _read_xml_colormap(xmlFile, mapName):
    '''Read in an XML colormap'''
    xml = ET.parse(xmlFile)
    root = xml.getroot()
    colormap = root.findall('ColorMap')
    if len(colormap) > 0:
        colormap = colormap[0]
        colorDict = {'red': [], 'green': [], 'blue': []}
        for point in colormap.findall('Point'):
            x = float(point.get('x'))
            color = [float(point.get('r')), float(point.get('g')),
                     float(point.get('b'))]
            colorDict['red'].append((x, color[0], color[0]))
            colorDict['green'].append((x, color[1], color[1]))
            colorDict['blue'].append((x, color[2], color[2]))
        cmap = LinearSegmentedColormap(mapName, colorDict, 256)
        _register_colormap_and_reverse(mapName, cmap)
def _register_colormap_and_reverse(mapName, cmap):
    plt.register_cmap(mapName, cmap)
    plt.register_cmap('{}_r'.format(mapName), cmap.reversed())
def _plot_color_gradients():
    '''from https://matplotlib.org/tutorials/colors/colormaps.html'''
    cmap_list = [m for m in plt.cm.cmap_d if not m.endswith("_r")]
    gradient = np.linspace(0, 1, 256)
    gradient = np.vstack((gradient, gradient))
    nrows = len(cmap_list)
    fig, axes = plt.subplots(figsize=(7.2, 0.25 * nrows), nrows=nrows)
    fig.subplots_adjust(top=0.99, bottom=0.01, left=0.35, right=0.99)
    for ax, name in zip(axes, cmap_list):
        ax.imshow(gradient, aspect='auto', cmap=plt.get_cmap(name))
        pos = list(ax.get_position().bounds)
        x_text = pos[0] - 0.01
        y_text = pos[1] + pos[3] / 2.
        fig.text(x_text, y_text, name, va='center', ha='right', fontsize=10)
    # Turn off *all* ticks & spines, not just the ones with colormaps.
    for ax in axes:
        ax.set_axis_off()
    plt.savefig('colormaps.png', dpi=100)
# vim: foldmethod=marker ai ts=4 sts=4 et sw=4 ft=python