colors.py 1.63 KB
Newer Older
1
2
3
4
5
6
7
"""
colors.py
Connor Hainje (connor.hainje@pnnl.gov)

Defines the colors to be used for the detectors and particles.
"""

8
from matplotlib.colors import LinearSegmentedColormap
9
10
import seaborn as sns

11
_det_cmap = sns.cubehelix_palette(
12
13
14
15
    5, start=1.5, rot=1.5, dark=0.3, light=0.8, reverse=True
)

DETECTORS = {
16
    "SVD": "gray",
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
    "CDC": _det_cmap[0],
    "TOP": _det_cmap[1],
    "ARICH": _det_cmap[2],
    "KLM": _det_cmap[3],
    "ECL": _det_cmap[4],
}

PARTICLES = {
    "e": f"C{0}",
    "mu": f"C{1}",
    "pi": f"C{2}",
    "K": f"C{3}",
    "p": f"C{4}",
    "d": f"C{5}",
}
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52


def fill_color_list(color_list, detectors=True):
    from . import const

    if color_list is not None:
        return color_list

    if detectors:
        return [DETECTORS[d] for d in const.DETECTORS]
    else:
        return [PARTICLES[h] for h in const.PARTICLES]


def make_colormap(color_list, detectors=True):
    color_list = fill_color_list(color_list, detectors=detectors)
    cmap_name = "detectors" if detectors else "particles"
    cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N=len(color_list))
    return cmap


53
def make_linear_colormap(color, cmap_name="custom_cmap", low=0.3, high=0.9):
54
55
56
57
58
59
60
61
62
63
64
65
66
    from matplotlib.colors import to_rgba, rgb_to_hsv, hsv_to_rgb

    hsv = rgb_to_hsv(to_rgba(color, alpha=1)[:3])
    lo = hsv_to_rgb((hsv[0], hsv[1], low))
    hi = hsv_to_rgb((hsv[0], hsv[1], high))

    cdict = {
        "red": ((0.0, lo[0], lo[0]), (1.0, hi[0], hi[0])),
        "green": ((0.0, lo[1], lo[1]), (1.0, hi[1], hi[1])),
        "blue": ((0.0, lo[2], lo[2]), (1.0, hi[2], hi[2])),
    }

    return LinearSegmentedColormap(cmap_name, cdict)