Coverage for pyrc \ core \ visualization \ color \ color.py: 76%
104 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-13 16:59 +0200
1# -------------------------------------------------------------------------------
2# Copyright (C) 2026 Joel Kimmich, Tim Jourdan
3# ------------------------------------------------------------------------------
4# License
5# This file is part of PyRC, distributed under GPL-3.0-or-later.
6# ------------------------------------------------------------------------------
8from importlib import resources
9import numpy as np
11from typing import Callable
13initial_uniform_color_map = "managua"
14initial_discrete_color_map = "managua10"
16RGB = tuple[float, float, float]
18# name -> either:
19# ("uniform", rgb[N,3])
20# ("explicit", t[N], rgb[N,3])
21CMAPS: dict[str, tuple] = {}
24def _parse_gradient(data: np.ndarray, name: str) -> tuple:
25 if data.shape[1] == 3:
26 return "uniform", data[:, :3]
27 elif data.shape[1] >= 4:
28 t = data[:, 0]
29 rgb = data[:, 1:4]
30 order = np.argsort(t)
31 return "explicit", t[order], rgb[order]
32 raise ValueError(f"Bad colormap file {name}: expected 3 or 4 columns")
35def _load_txt_colormaps() -> None:
36 pkg = __package__
37 root = resources.files(pkg)
39 for subfolder, kind in (("uniform", "gradient"), ("discrete", "discrete")):
40 try:
41 folder = root.joinpath(subfolder)
42 entries = list(folder.iterdir())
43 except (FileNotFoundError, NotADirectoryError):
44 continue
46 for p in entries:
47 if not (p.is_file() and p.name.lower().endswith(".txt")):
48 continue
49 cmap_name = p.name.rsplit(".", 1)[0]
50 lines = p.read_text(encoding="utf-8").splitlines()
52 def _first_numeric_line(lines: list[str]) -> int:
53 for i, line in enumerate(lines):
54 parts = line.split()
55 if parts and parts[0].lstrip("-").replace(".", "", 1).isdigit():
56 return i
57 raise ValueError("No numeric line found")
59 if kind == "discrete":
60 skip = _first_numeric_line(lines)
61 rgb_rows = []
62 for line in lines[skip:]:
63 parts = line.split()
64 try:
65 r, g, b = float(parts[0]), float(parts[1]), float(parts[2])
66 if all(v <= 255 for v in (r, g, b)):
67 rgb_rows.append([r / 255.0, g / 255.0, b / 255.0])
68 except (ValueError, IndexError):
69 continue
70 if not rgb_rows:
71 raise ValueError(f"No valid colours in discrete palette {p.name}")
72 CMAPS[cmap_name] = ("discrete", np.array(rgb_rows, dtype=np.float32))
73 else:
74 skip = _first_numeric_line(lines)
75 data = np.loadtxt(p, dtype=np.float32, skiprows=skip, encoding="utf-8")
76 if data.ndim == 1:
77 data = data[None, :]
78 CMAPS[cmap_name] = _parse_gradient(data, p.name)
81_load_txt_colormaps()
84def available(kind: str | None = None) -> tuple[str, ...]:
85 """
86 Return names of all loaded colormaps.
88 Parameters
89 ----------
90 kind : str or None
91 Filter by type: 'uniform', 'explicit', 'discrete', or None for all.
92 """
93 if kind is None:
94 return tuple(sorted(CMAPS.keys()))
95 return tuple(sorted(k for k, v in CMAPS.items() if v[0] == kind))
98def value_to_rgb(value: float | int | np.ndarray, cmap: str = initial_uniform_color_map) -> np.ndarray:
99 """
100 Map relative value(s) in [0, 1] to RGB using a loaded colormap.
102 Parameters
103 ----------
104 value : float, int, or np.ndarray
105 Scalar or array of values in [0, 1].
106 cmap : str
107 Colormap name.
109 Returns
110 -------
111 np.ndarray
112 Shape (..., 3), dtype float32.
113 """
114 if cmap not in CMAPS:
115 raise KeyError(f"Unknown colormap '{cmap}'. Available: {available()}")
117 value = np.clip(np.asarray(value, dtype=np.float32), 0.0, 1.0)
118 entry = CMAPS[cmap]
119 kind = entry[0]
121 if kind == "uniform":
122 rgb = entry[1]
123 n = rgb.shape[0]
124 if n == 1:
125 return np.broadcast_to(rgb[0], value.shape + (3,)).astype(np.float32, copy=False)
126 x = value * (n - 1)
127 i = np.floor(x).astype(np.int32)
128 j = np.minimum(i + 1, n - 1)
129 u = (x - i)[..., None]
130 return (rgb[i] * (1.0 - u) + rgb[j] * u).astype(np.float32, copy=False)
132 if kind == "explicit":
133 t, rgb = entry[1], entry[2]
134 r = np.interp(value, t, rgb[:, 0]).astype(np.float32, copy=False)
135 g = np.interp(value, t, rgb[:, 1]).astype(np.float32, copy=False)
136 b = np.interp(value, t, rgb[:, 2]).astype(np.float32, copy=False)
137 return np.stack([r, g, b], axis=-1)
139 # discrete: map value to nearest swatch index
140 rgb = entry[1]
141 n = rgb.shape[0]
142 idx = np.clip(np.floor(value * n).astype(np.int32), 0, n - 1)
143 return rgb[idx]
146def get_palette(cmap: str = initial_discrete_color_map) -> np.ndarray:
147 """
148 Return all swatches of a discrete palette as (N, 3) float32 array.
150 Parameters
151 ----------
152 cmap : str
153 Name of a discrete colormap.
155 Returns
156 -------
157 np.ndarray
158 Shape (N, 3), dtype float32.
160 Raises
161 ------
162 KeyError
163 If cmap is not found.
164 TypeError
165 If cmap is not a discrete palette.
166 """
167 if cmap not in CMAPS:
168 raise KeyError(f"Unknown colormap '{cmap}'. Available: {available()}")
169 if CMAPS[cmap][0] != "discrete":
170 raise TypeError(f"'{cmap}' is not a discrete palette.")
171 return CMAPS[cmap][1]
174def make_color_mapper(cmap: str = initial_uniform_color_map) -> Callable[[float | int | np.ndarray], np.ndarray]:
175 """Convenience: returns f(v)->rgb for one colormap name."""
176 return lambda v: value_to_rgb(v, cmap=cmap)
179def cycle_palette(cmap: str = initial_discrete_color_map) -> "Generator[np.ndarray, None, None]":
180 """
181 Yield RGB colours from a colormap in order, cycling indefinitely.
183 Parameters
184 ----------
185 cmap : str
186 Colormap name.
188 Yields
189 ------
190 np.ndarray
191 Shape (3,), dtype float32.
192 """
193 if cmap not in CMAPS:
194 raise KeyError(f"Unknown colormap '{cmap}'. Available: {available()}")
195 entry = CMAPS[cmap]
196 rgb = entry[1] if entry[0] in ("uniform", "discrete") else entry[2]
197 i = 0
198 n = rgb.shape[0]
199 while True:
200 yield rgb[i % n]
201 i += 1