<reponame>cffbots/ESMValTool<filename>esmvaltool/utils/color_tables/show_color_tables.py<gh_stars>100-1000
"""Utility script for inspecting and converting ncl color tables."""
# pylint: disable=import-outside-toplevel
import logging

logger = logging.getLogger(__name__)

NCL_SCRIPT = """
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"

begin
  print("Autogenerated NCL script is running.")
  wks = gsn_open_wks("pdf","{{ outdir }}/available_colormaps_for_ncl")
  {% for n in list_of_snippets %}{{n}} {% endfor %}
end
"""

COLOR_SNIPPET = """
  cmap := read_colormap_file("{{ path }}")
  opt = True
  opt@Frame = False
  draw_color_palette(wks, cmap, opt)
  gsn_text_ndc(wks, "{{ name }}", 0.5, 0.9, True)
  frame(wks)
"""


def load_ncl_color_map(name, colorpath):
    """Load ncl color map to a list that is returned."""
    def _format(content):
        out = []
        for item in content.split("\n"):
            item = item.strip()
            if item and not ('ncolors' in item or item.startswith('#')
                             or item.startswith(';')):
                out.append([int(elem) / 256
                            for elem in item.split()[0:3]] + [1])
        return out

    filename = "{0}/{1}.rgb".format(colorpath, name)
    import os
    if not os.path.exists(filename):
        raise ValueError("Path {0} does not exist.".format(filename))
    with open(filename, 'r') as ncl_color_map:
        return _format(ncl_color_map.read())


def get_color_map(name, colorpath):
    """Convert colormap from ncl to python.

    Parameters
    ----------
    name: str
        Name of ncl color map
    Returns
    -------
    matplotlib.colors.ListedColorMap object
    """
    import matplotlib
    import yaml
    colors = load_ncl_color_map(name, colorpath)
    logger.debug("RGB values for '%s':\n%s", name, yaml.dump(colors))
    return matplotlib.colors.ListedColormap(colors, name=name, N=None)


def list_ncl_color_maps(colorpath):
    """Get list of all available ncl color maps."""
    import os

    def _format(name):
        return os.path.splitext(os.path.basename(name))[0]

    out = []
    for (_, _, filenames) in os.walk(colorpath):
        out.extend([
            _format(filename) for filename in filenames
            if 'rgb' in filename.split('.')
        ])
    return out


def plot_example_for_colormap(name, colorpath, outdir='./'):
    """Create plots of given color map using python."""
    logger.info("Plotting example for '%s'", name)
    import os
    import matplotlib
    matplotlib.use("Agg")  # noqa
    import matplotlib.pyplot as plt
    import numpy as np
    fig = plt.figure(1)
    axis = fig.add_axes([0.1, 0.3, 0.5, 0.5])
    np.random.seed(12345678)
    data = np.random.randn(30, 30)
    psm = axis.pcolormesh(data,
                          cmap=get_color_map(name, colorpath),
                          rasterized=True,
                          vmin=-4,
                          vmax=4)
    fig.colorbar(psm, ax=axis)
    plt.savefig(os.path.join(outdir, "{0}.png".format(name)))
    plt.close()


def main_plot_python_cm(colorpath, outpath):
    """Execute functions for python plots."""
    for name in list_ncl_color_maps(colorpath):
        plot_example_for_colormap(name, colorpath, outdir=outpath)


def main_plot_ncl_cm(colorpath, outpath):
    """Execute functions for ncl plots."""
    from jinja2 import Template
    t_color_snippet = Template(COLOR_SNIPPET)
    template = Template(NCL_SCRIPT)
    list_of_snippets = []
    import glob
    import subprocess
    import tempfile
    import os
    for path in glob.glob(colorpath + "/*rgb"):
        _, tail = os.path.split(path)
        list_of_snippets.append(t_color_snippet.render(path=path, name=tail))
    with tempfile.NamedTemporaryFile(mode='w', suffix='ncl') as fname:
        fname.write(
            template.render(list_of_snippets=sorted(list_of_snippets),
                            outdir=outpath))
        subprocess.check_call(["ncl", fname.name])


class ColorTables():
    """Generate colormap samples for ESMValTool's default colormaps."""

    def __init__(self):
        logger.setLevel(logging.DEBUG)
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(
            logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(console_handler)
        self._colorpath = None
        self._outpath = None

    def _prepare_paths(self, colorpath, outpath):
        import os
        if colorpath is None:
            from esmvaltool.diag_scripts.shared.plot  \
                import __file__ as plot_path
            colorpath = os.path.join(os.path.dirname(plot_path), "rgb")

        if not os.path.isdir(colorpath):
            logger.warning("Path '%s' is invalid", colorpath)
            raise OSError
        self._colorpath = colorpath

        if not os.path.isdir(outpath) and not os.path.exists(outpath):
            logger.info("Creating directory '%s'", outpath)
            os.mkdir(outpath)
        self._outpath = outpath

    def python(self, colorpath=None, outpath="./"):
        """
        Generate samples for Python colormaps.

        Create a series of png images  with examples of ESMValTool's available
        Python colormaps

        Parameters
        ----------

        colorpath: str
            Folder to search for colormaps. Default is installed colormaps
        outpath: str
            Out directory. Default is the current directory
        """
        self._prepare_paths(colorpath, outpath)
        logger.info("Creating report with Python")
        main_plot_python_cm(self._colorpath, self._outpath)

    def ncl(self, colorpath=None, outpath="./"):
        """
        Generate samples for NCL colormaps.

        Create a pdf with examples of ESMValTool's available NCL colormaps

        Parameters
        ----------

        colorpath: str
            Folder to search for colormaps. Default is installed colormaps
        outpath: str
            Out directory. Default is the current directory
        """
        self._prepare_paths(colorpath, outpath)
        logger.info("Creating report with NCL")
        main_plot_ncl_cm(self._colorpath, self._outpath)

"""Utility script for inspecting and converting ncl color tables."""
# pylint: disable=import-outside-toplevel
import logging

logger = logging.getLogger(__name__)

NCL_SCRIPT = """
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"

begin
  print("Autogenerated NCL script is running.")
  wks = gsn_open_wks("pdf","{{ outdir }}/available_colormaps_for_ncl")
  {% for n in list_of_snippets %}{{n}} {% endfor %}
end
"""

COLOR_SNIPPET = """
  cmap := read_colormap_file("{{ path }}")
  opt = True
  opt@Frame = False
  draw_color_palette(wks, cmap, opt)
  gsn_text_ndc(wks, "{{ name }}", 0.5, 0.9, True)
  frame(wks)
"""


def load_ncl_color_map(name, colorpath):
    """Load ncl color map to a list that is returned."""
    def _format(content):
        out = []
        for item in content.split("\n"):
            item = item.strip()
            if item and not ('ncolors' in item or item.startswith('#')
                             or item.startswith(';')):
                out.append([int(elem) / 256
                            for elem in item.split()[0:3]] + [1])
        return out

    filename = "{0}/{1}.rgb".format(colorpath, name)
    import os
    if not os.path.exists(filename):
        raise ValueError("Path {0} does not exist.".format(filename))
    with open(filename, 'r') as ncl_color_map:
        return _format(ncl_color_map.read())


def get_color_map(name, colorpath):
    """Convert colormap from ncl to python.

    Parameters
    ----------
    name: str
        Name of ncl color map
    Returns
    -------
    matplotlib.colors.ListedColorMap object
    """
    import matplotlib
    import yaml
    colors = load_ncl_color_map(name, colorpath)
    logger.debug("RGB values for '%s':\n%s", name, yaml.dump(colors))
    return matplotlib.colors.ListedColormap(colors, name=name, N=None)


def list_ncl_color_maps(colorpath):
    """Get list of all available ncl color maps."""
    import os

    def _format(name):
        return os.path.splitext(os.path.basename(name))[0]

    out = []
    for (_, _, filenames) in os.walk(colorpath):
        out.extend([
            _format(filename) for filename in filenames
            if 'rgb' in filename.split('.')
        ])
    return out


def plot_example_for_colormap(name, colorpath, outdir='./'):
    """Create plots of given color map using python."""
    logger.info("Plotting example for '%s'", name)
    import os
    import matplotlib
    matplotlib.use("Agg")  # noqa
    import matplotlib.pyplot as plt
    import numpy as np
    fig = plt.figure(1)
    axis = fig.add_axes([0.1, 0.3, 0.5, 0.5])
    np.random.seed(12345678)
    data = np.random.randn(30, 30)
    psm = axis.pcolormesh(data,
                          cmap=get_color_map(name, colorpath),
                          rasterized=True,
                          vmin=-4,
                          vmax=4)
    fig.colorbar(psm, ax=axis)
    plt.savefig(os.path.join(outdir, "{0}.png".format(name)))
    plt.close()


def main_plot_python_cm(colorpath, outpath):
    """Execute functions for python plots."""
    for name in list_ncl_color_maps(colorpath):
        plot_example_for_colormap(name, colorpath, outdir=outpath)


def main_plot_ncl_cm(colorpath, outpath):
    """Execute functions for ncl plots."""
    from jinja2 import Template
    t_color_snippet = Template(COLOR_SNIPPET)
    template = Template(NCL_SCRIPT)
    list_of_snippets = []
    import glob
    import subprocess
    import tempfile
    import os
    for path in glob.glob(colorpath + "/*rgb"):
        _, tail = os.path.split(path)
        list_of_snippets.append(t_color_snippet.render(path=path, name=tail))
    with tempfile.NamedTemporaryFile(mode='w', suffix='ncl') as fname:
        fname.write(
            template.render(list_of_snippets=sorted(list_of_snippets),
                            outdir=outpath))
        subprocess.check_call(["ncl", fname.name])


class ColorTables():
    """Generate colormap samples for ESMValTool's default colormaps."""

    def __init__(self):
        logger.setLevel(logging.DEBUG)
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(
            logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(console_handler)
        self._colorpath = None
        self._outpath = None

    def _prepare_paths(self, colorpath, outpath):
        import os
        if colorpath is None:
            from esmvaltool.diag_scripts.shared.plot  \
                import __file__ as plot_path
            colorpath = os.path.join(os.path.dirname(plot_path), "rgb")

        if not os.path.isdir(colorpath):
            logger.warning("Path '%s' is invalid", colorpath)
            raise OSError
        self._colorpath = colorpath

        if not os.path.isdir(outpath) and not os.path.exists(outpath):
            logger.info("Creating directory '%s'", outpath)
            os.mkdir(outpath)
        self._outpath = outpath

    def python(self, colorpath=None, outpath="./"):
        """
        Generate samples for Python colormaps.

        Create a series of png images  with examples of ESMValTool's available
        Python colormaps

        Parameters
        ----------

        colorpath: str
            Folder to search for colormaps. Default is installed colormaps
        outpath: str
            Out directory. Default is the current directory
        """
        self._prepare_paths(colorpath, outpath)
        logger.info("Creating report with Python")
        main_plot_python_cm(self._colorpath, self._outpath)

    def ncl(self, colorpath=None, outpath="./"):
        """
        Generate samples for NCL colormaps.

        Create a pdf with examples of ESMValTool's available NCL colormaps

        Parameters
        ----------

        colorpath: str
            Folder to search for colormaps. Default is installed colormaps
        outpath: str
            Out directory. Default is the current directory
        """
        self._prepare_paths(colorpath, outpath)
        logger.info("Creating report with NCL")
        main_plot_ncl_cm(self._colorpath, self._outpath)
Utility script for inspecting and converting ncl color tables. # pylint: disable=import-outside-toplevel load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"  begin   print("Autogenerated NCL script is running.")   wks = gsn_open_wks("pdf","{{ outdir }}/available_colormaps_for_ncl")   {% for n in list_of_snippets %}{{n}} {% endfor %} end cmap := read_colormap_file("{{ path }}")   opt = True   opt@Frame = False   draw_color_palette(wks, cmap, opt)   gsn_text_ndc(wks, "{{ name }}", 0.5, 0.9, True)   frame(wks) Load ncl color map to a list that is returned. Convert colormap from ncl to python.      Parameters     ----------     name: str         Name of ncl color map     Returns     -------     matplotlib.colors.ListedColorMap object Get list of all available ncl color maps. Create plots of given color map using python. # noqa Execute functions for python plots. Execute functions for ncl plots. Generate colormap samples for ESMValTool's default colormaps. Generate samples for Python colormaps.          Create a series of png images  with examples of ESMValTool's available         Python colormaps          Parameters         ----------          colorpath: str             Folder to search for colormaps. Default is installed colormaps         outpath: str             Out directory. Default is the current directory Generate samples for NCL colormaps.          Create a pdf with examples of ESMValTool's available NCL colormaps          Parameters         ----------          colorpath: str             Folder to search for colormaps. Default is installed colormaps         outpath: str             Out directory. Default is the current directory