|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# |
| 3 | +# Licensed under the terms of the BSD 3-Clause |
| 4 | +# (see plotpy/LICENSE for details) |
| 5 | + |
| 6 | +# pylint: disable=C0103 |
| 7 | + |
| 8 | +""" |
| 9 | +Colormap utilities |
| 10 | +------------------ |
| 11 | +
|
| 12 | +""" |
| 13 | + |
| 14 | +from __future__ import annotations |
| 15 | + |
| 16 | +import _cm |
| 17 | +import numpy as np |
| 18 | +from qtpy import QtGui as QG |
| 19 | +from qwt import QwtInterval |
| 20 | + |
| 21 | +from plotpy.mathutils.colormaps import ( |
| 22 | + DEFAULT_COLORMAPS, |
| 23 | + DEFAULT_COLORMAPS_PATH, |
| 24 | + save_colormaps, |
| 25 | +) |
| 26 | +from plotpy.widgets.colormap.widget import EditableColormap # Reuse matplotlib data |
| 27 | + |
| 28 | +# usefull to obtain a full color map |
| 29 | +FULLRANGE = QwtInterval(0.0, 1.0) |
| 30 | + |
| 31 | +COLORMAPS = {} |
| 32 | +EXTRA_COLORMAPS = [] # custom build colormaps |
| 33 | +ICON_CACHE = {} |
| 34 | + |
| 35 | + |
| 36 | +def _interpolate(val, vmin, vmax): |
| 37 | + """Interpolate a color component between to values as provided |
| 38 | + by matplotlib colormaps |
| 39 | + """ |
| 40 | + interp = (val - vmin[0]) / (vmax[0] - vmin[0]) |
| 41 | + return (1 - interp) * vmin[1] + interp * vmax[2] |
| 42 | + |
| 43 | + |
| 44 | +def _setup_colormap(cmap: EditableColormap, cmdata): |
| 45 | + """Setup a CustomQwtLinearColorMap according to |
| 46 | + matplotlib's data |
| 47 | + """ |
| 48 | + red = np.array(cmdata["red"]) |
| 49 | + green = np.array(cmdata["green"]) |
| 50 | + blue = np.array(cmdata["blue"]) |
| 51 | + qmin = QG.QColor() |
| 52 | + qmin.setRgbF(red[0, 2], green[0, 2], blue[0, 2]) |
| 53 | + qmax = QG.QColor() |
| 54 | + qmax.setRgbF(red[-1, 2], green[-1, 2], blue[-1, 2]) |
| 55 | + cmap.setColorInterval(qmin, qmax) |
| 56 | + indices = sorted(set(red[:, 0]) | set(green[:, 0]) | set(blue[:, 0])) |
| 57 | + for i in indices[1:-1]: |
| 58 | + idxr = red[:, 0].searchsorted(i) |
| 59 | + idxg = green[:, 0].searchsorted(i) |
| 60 | + idxb = blue[:, 0].searchsorted(i) |
| 61 | + compr = _interpolate(i, red[idxr - 1], red[idxr]) |
| 62 | + compg = _interpolate(i, green[idxg - 1], green[idxg]) |
| 63 | + compb = _interpolate(i, blue[idxb - 1], blue[idxb]) |
| 64 | + col = QG.QColor() |
| 65 | + col.setRgbF(compr, compg, compb) |
| 66 | + cmap.addColorStop(i, col) |
| 67 | + |
| 68 | + |
| 69 | +def get_cmap(name: str) -> EditableColormap: |
| 70 | + """Get a colormap from its name |
| 71 | +
|
| 72 | + Args: |
| 73 | + name: colormap name |
| 74 | +
|
| 75 | + Returns: |
| 76 | + Return a QwtColormap based on matplotlib's colormap of the same name |
| 77 | + We avoid rebuilding the cmap by keeping it in cache |
| 78 | + """ |
| 79 | + if name in COLORMAPS: |
| 80 | + return COLORMAPS[name] |
| 81 | + |
| 82 | + colormap = EditableColormap() |
| 83 | + COLORMAPS[name] = colormap |
| 84 | + COLORMAPS[colormap] = name |
| 85 | + data = getattr(_cm, "_" + name + "_data") |
| 86 | + _setup_colormap(colormap, data) |
| 87 | + return colormap |
| 88 | + |
| 89 | + |
| 90 | +def get_colormap_list() -> list[str]: |
| 91 | + """Builds a list of available colormaps by introspection of the _cm module |
| 92 | +
|
| 93 | + Returns: |
| 94 | + list of colormap names |
| 95 | + """ |
| 96 | + cmlist = [] |
| 97 | + cmlist += EXTRA_COLORMAPS |
| 98 | + for name in dir(_cm): |
| 99 | + if name.endswith("_data"): |
| 100 | + obj = getattr(_cm, name) |
| 101 | + if isinstance(obj, dict): |
| 102 | + cmlist.append(name[1:-5]) |
| 103 | + return cmlist |
| 104 | + |
| 105 | + |
| 106 | +if __name__ == "__main__": |
| 107 | + for cmap in get_colormap_list(): |
| 108 | + DEFAULT_COLORMAPS[cmap] = get_cmap(cmap) |
| 109 | + save_colormaps(DEFAULT_COLORMAPS_PATH, DEFAULT_COLORMAPS) |
0 commit comments