Merge remote-tracking branch 'Codra/develop-2.1_new-colormaps' into develop

Author: PierreRaybaut

colormaps/matplotlib_cmaps.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import copy
+import sys
+from typing import Callable
+
+import matplotlib.colors as pltc
+import matplotlib.pyplot as plt
+import numpy as np
+
+from plotpy.mathutils.colormap import (
+    DEFAULT_COLORMAPS,
+    DEFAULT_COLORMAPS_PATH,
+    CmapDictType,
+    EditableColormap,
+    save_colormaps,
+)
+
+PERCEPTUALLY_UNIFORM_CMAPS = ["viridis", "plasma", "inferno", "magma", "cividis"]
+SEQUENTIAL_CMAPS = [
+    "Greys",
+    "Purples",
+    "Blues",
+    "Greens",
+    "Oranges",
+    "Reds",
+    "YlOrBr",
+    "YlOrRd",
+    "OrRd",
+    "PuRd",
+    "RdPu",
+    "BuPu",
+    "GnBu",
+    "PuBu",
+    "YlGnBu",
+    "PuBuGn",
+    "BuGn",
+    "YlGn",
+]
+SEQUENTIAL_CMAPS2 = [
+    "binary",
+    "gist_yarg",
+    "gist_gray",
+    "gray",
+    "bone",
+    "pink",
+    "spring",
+    "summer",
+    "autumn",
+    "winter",
+    "cool",
+    "Wistia",
+    "hot",
+    "afmhot",
+    "gist_heat",
+    "copper",
+]
+DIVERGING_CMAPS = [
+    "PiYG",
+    "PRGn",
+    "BrBG",
+    "PuOr",
+    "RdGy",
+    "RdBu",
+    "RdYlBu",
+    "RdYlGn",
+    "Spectral",
+    "coolwarm",
+    "bwr",
+    "seismic",
+]
+CYCLIC_CMAPS = ["twilight", "twilight_shifted", "hsv"]
+QUALITATIVE_CMAPS = [
+    "Pastel1",
+    "Pastel2",
+    "Paired",
+    "Accent",
+    "Dark2",
+    "Set1",
+    "Set2",
+    "Set3",
+    "tab10",
+    "tab20",
+    "tab20b",
+    "tab20c",
+]
+MISCELLANEOUS_CMAPS = [
+    "flag",
+    "prism",
+    "ocean",
+    "gist_earth",
+    "terrain",
+    "gist_stern",
+    "gnuplot",
+    "gnuplot2",
+    "CMRmap",
+    "cubehelix",
+    "brg",
+    "gist_rainbow",
+    "rainbow",
+    "jet",
+    "turbo",
+    "nipy_spectral",
+    "gist_ncar",
+]
+
+SORTED_MATPLOTLIB_COLORMAPS: list[str] = [
+    *PERCEPTUALLY_UNIFORM_CMAPS,
+    *SEQUENTIAL_CMAPS,
+    *SEQUENTIAL_CMAPS2,
+    *DIVERGING_CMAPS,
+    *CYCLIC_CMAPS,
+    *QUALITATIVE_CMAPS,
+    *MISCELLANEOUS_CMAPS,
+]
+
+
+def rgb_colors_to_hex_list(
+    colors: list[tuple[int, int, int]]
+) -> list[tuple[float, str]]:
+    """Convert a list of RGB colors to a list of tuples with the position of the color
+    and the color in hex format. Positions evenly distributed between 0 and 1.
+
+    Args:
+        colors: list of RGB colors
+
+    Returns:
+        list of tuples with the position of the color and the color in hex format
+    """
+    return [(i / len(colors), pltc.to_hex(color)) for i, color in enumerate(colors)]
+
+
+def _interpolate(
+    val: float, vmin: tuple[float, float, float], vmax: tuple[float, float, float]
+):
+    """Interpolate between two level of a color.
+
+    Args:
+        val: value to interpolate
+        vmin: R, G or B tuple from a matplotlib segmented colormap
+        vmax: R, G or B tuple from matplotlib segmented colormap
+
+    Returns:
+        The interpolated R, G or B component
+    """
+    interp = (val - vmin[0]) / (vmax[0] - vmin[0])
+    return (1 - interp) * vmin[1] + interp * vmax[2]
+
+
+def std_segmented_cmap_to_hex_list(cmdata: dict[str, list[tuple[float, float, float]]]):
+    """Convert a matplotlib segmented colormap to a list of tuples with the position of
+    the color and the color in hex format.
+
+    Args:
+        cmdata: segmented colormap data
+
+    Returns:
+        list of tuples with the position of the color and the color in hex format
+    """
+    colors: list[tuple[float, str]] = []
+    red = np.array(cmdata["red"])
+    green = np.array(cmdata["green"])
+    blue = np.array(cmdata["blue"])
+    indices = sorted(set(red[:, 0]) | set(green[:, 0]) | set(blue[:, 0]))
+    for i in indices:
+        idxr = red[:, 0].searchsorted(i)
+        idxg = green[:, 0].searchsorted(i)
+        idxb = blue[:, 0].searchsorted(i)
+        compr = _interpolate(i, red[idxr - 1], red[idxr])
+        compg = _interpolate(i, green[idxg - 1], green[idxg])
+        compb = _interpolate(i, blue[idxb - 1], blue[idxb])
+        colors.append((i, pltc.to_hex((compr, compg, compb))))
+    return colors
+
+
+InterpFuncT = Callable[[np.ndarray], np.ndarray]
+
+
+def func_segmented_cmap_to_hex_list(
+    n: int,
+    cmap: pltc.LinearSegmentedColormap,
+) -> list[tuple[float, str]]:
+    """Convert a matplotlib segmented colormap to a list of tuples with the position of
+    the color and the color in hex format. The input colormap contains function for each
+    color RGB component instead of a list of colors.
+
+    Args:
+        n: number of colors to generate
+        cmap: segmented colormap
+
+    Returns:
+        list of tuples with the position of the color and the color in hex format
+    """
+    colors = []
+    arr = np.linspace(0, 1, n, dtype=float)
+    colors = [(i, pltc.to_hex(rgba)) for i, rgba in zip(arr, cmap(arr))]
+    return colors
+
+
+def continuous_to_descrete_cmap(cmap: EditableColormap) -> EditableColormap:
+    """Convert a continuous colormap to a descrete one.
+
+    Args:
+        cmap: colormap to convert
+
+    Returns:
+        descrete colormap
+    """
+    raw_cmap: tuple[tuple[float, str], ...] = cmap.to_tuples()
+    new_raw_cmap: list[tuple[float, str]] = [raw_cmap[0]]
+    n = len(raw_cmap)
+    coeff = (n - 1) / n
+    for i, (pos, color) in enumerate(raw_cmap[1:]):
+        prev_pos, prev_color = raw_cmap[i]
+        curr_pos, curr_color = pos, color
+        new_pos = curr_pos * coeff
+        new_raw_cmap.append((new_pos, prev_color))
+        new_raw_cmap.append((new_pos, curr_color))
+    new_raw_cmap.append(raw_cmap[-1])
+
+    return EditableColormap.from_iterable(new_raw_cmap, name=cmap.name)
+
+
+def sort_mpl_colormaps(colormaps: CmapDictType) -> CmapDictType:
+    """Filter and sort input colormaps to follow the same order (by category) as in the
+     matplotlib colormaps documentation. Colormaps not found in the matplotlib
+     are filtered out.
+
+    Args:
+        colormaps: Dictionnary of colormaps to extract and order
+
+    Returns:
+        Filtered and sorted colormaps dictionnary
+    """
+    ordered_colormaps: CmapDictType = {}
+    lower_cmap_names = [cm.lower() for cm in SORTED_MATPLOTLIB_COLORMAPS]
+    for lower_name in lower_cmap_names:
+        cmap = colormaps.get(lower_name, None)
+        if lower_name.endswith("_r"):
+            continue
+        if cmap is None:
+            print(f"Colormap {lower_name} not found in input colormaps.")
+            continue
+        ordered_colormaps[lower_name] = cmap
+    return ordered_colormaps
+
+
+def append_non_mpl_colormaps(mpl_colormaps: CmapDictType, colormaps: CmapDictType):
+    """Append colormaps not found in the matplotlib colormaps to the input colormaps.
+     Mutate the input in place.
+
+    Args:
+        mpl_colormaps: dictionnary of matplotlib colormaps. Mutated in place.
+        colormaps: dictionnary of colormaps to append to the matplotlib colormaps
+    """
+    colormap_names = set(SORTED_MATPLOTLIB_COLORMAPS)
+    for colormap in colormaps.values():
+        if colormap.name not in colormap_names:
+            print(f"{colormap} not in matplotlib colormaps.")
+            mpl_colormaps[colormap.name.lower()] = colormap
+
+
+def main(cmaps: CmapDictType, out_json_path: str = DEFAULT_COLORMAPS_PATH):
+
+    new_cmaps: dict[str, list[tuple[float, str]]] = {}
+
+    # Uniform colormaps with a .colors attribute that return a list of RGB colors
+    cmaps_with_colors = [
+        "magma",
+        "viridis",
+        "inferno",
+        "plasma",
+        "cividis",
+    ]
+
+    # Discrete colormaps, same as uniform colormaps but the colormap must be post
+    # processed to become descrete
+    descrete_cmaps = [
+        "Pastel1",
+        "Pastel2",
+        "Paired",
+        "Accent",
+        "Dark2",
+        "Set1",
+        "Set2",
+        "Set3",
+    ]
+
+    cmaps_with_colors.extend(descrete_cmaps)
+
+    # Colormaps with a _segmented_data attribute that contains the R, G and B components
+    # as lists of tuples
+    segmented_cmaps = [
+        "coolwarm",
+        "bwr",
+        "seismic",
+    ]
+
+    # Colormaps with a _segmentdata attribute that contains the R, G and B components as
+    # functions that return the color for a given position
+    interp_cmaps = ["gnuplot2", "CMRmap", "rainbow", "turbo", "afmhot"]
+
+    for cm_name in cmaps_with_colors:
+        cmap = plt.get_cmap(cm_name)
+        new_cmaps[cm_name] = rgb_colors_to_hex_list(cmap.colors)
+
+    for cm_name in descrete_cmaps:
+        cmap = EditableColormap.from_iterable(new_cmaps[cm_name], name=cm_name)
+        new_cmaps[cm_name] = list(continuous_to_descrete_cmap(cmap).to_tuples())
+
+    for cm_name in segmented_cmaps:
+        cmap = plt.get_cmap(cm_name)
+        new_cmaps[cm_name] = std_segmented_cmap_to_hex_list(cmap._segmentdata)
+
+    n = 128
+
+    for cm_name in interp_cmaps:
+        cmap = plt.get_cmap(cm_name)
+        new_cmaps[cm_name] = func_segmented_cmap_to_hex_list(n, cmap)
+
+    for name, raw_cm in new_cmaps.items():
+        cmaps[name.lower()] = EditableColormap.from_iterable(raw_cm, name=name)
+
+    ordered_cmaps = sort_mpl_colormaps(cmaps)
+    append_non_mpl_colormaps(ordered_cmaps, cmaps)
+    save_colormaps(out_json_path, ordered_cmaps)
+
+
+if __name__ == "__main__":
+    cmaps = copy.deepcopy(DEFAULT_COLORMAPS)
+    out_json = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_COLORMAPS_PATH
+    main(cmaps, out_json)