Coverage for src / tracekit / visualization / colors.py: 76%

159 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-11 23:04 +0000

1"""Color palette selection and accessibility utilities. 

2 

3This module provides intelligent color palette selection based on data 

4characteristics and accessibility requirements with WCAG contrast checking. 

5 

6 

7Example: 

8 >>> from tracekit.visualization.colors import select_optimal_palette 

9 >>> colors = select_optimal_palette(n_channels=3, palette_type="qualitative") 

10 

11References: 

12 WCAG 2.1 contrast guidelines 

13 Colorblind-safe palette design (Brettel 1997) 

14 ColorBrewer schemes 

15""" 

16 

17from __future__ import annotations 

18 

19from typing import Literal 

20 

21import numpy as np 

22 

23# Predefined colorblind-safe palettes 

24COLORBLIND_SAFE_QUALITATIVE = [ 

25 "#0173B2", # Blue 

26 "#DE8F05", # Orange 

27 "#029E73", # Green 

28 "#CC78BC", # Purple 

29 "#CA9161", # Brown 

30 "#949494", # Gray 

31 "#ECE133", # Yellow 

32 "#56B4E9", # Light blue 

33] 

34 

35SEQUENTIAL_VIRIDIS = [ 

36 "#440154", 

37 "#481567", 

38 "#482677", 

39 "#453781", 

40 "#404788", 

41 "#39568C", 

42 "#33638D", 

43 "#2D708E", 

44 "#287D8E", 

45 "#238A8D", 

46 "#1F968B", 

47 "#20A387", 

48 "#29AF7F", 

49 "#3CBB75", 

50 "#55C667", 

51 "#73D055", 

52 "#95D840", 

53 "#B8DE29", 

54 "#DCE319", 

55 "#FDE724", 

56] 

57 

58DIVERGING_COOLWARM = [ 

59 "#3B4CC0", 

60 "#5977E3", 

61 "#7D9EF2", 

62 "#A2C0F9", 

63 "#C7DDFA", 

64 "#E8F0FC", 

65 "#F9EBE5", 

66 "#F6CFBB", 

67 "#F0AD8E", 

68 "#E68462", 

69 "#D8583E", 

70 "#C52A1E", 

71 "#B40426", 

72] 

73 

74 

75def select_optimal_palette( 

76 n_colors: int, 

77 *, 

78 palette_type: Literal["sequential", "diverging", "qualitative"] | None = None, 

79 data_range: tuple[float, float] | None = None, 

80 colorblind_safe: bool = True, 

81 background_color: str = "#FFFFFF", 

82 min_contrast_ratio: float = 4.5, 

83) -> list[str]: 

84 """Select optimal color palette based on data characteristics. 

85 

86 : Automatically select optimal color palettes based on 

87 data characteristics, plot type, and accessibility requirements. 

88 

89 Args: 

90 n_colors: Number of colors needed 

91 palette_type: Type of palette ("sequential", "diverging", "qualitative") 

92 If None, auto-select based on n_colors and data_range 

93 data_range: Data range (min, max) for auto-detecting bipolar signals 

94 colorblind_safe: Ensure colorblind-safe palette (default: True) 

95 background_color: Background color for contrast checking (default: white) 

96 min_contrast_ratio: Minimum WCAG contrast ratio (default: 4.5 for AA) 

97 

98 Returns: 

99 List of color hex codes 

100 

101 Raises: 

102 ValueError: If n_colors is invalid or palette cannot meet requirements 

103 

104 Example: 

105 >>> # Auto-select for 3 channels 

106 >>> colors = select_optimal_palette(3) 

107 >>> # Diverging palette for bipolar data 

108 >>> colors = select_optimal_palette(10, palette_type="diverging") 

109 

110 References: 

111 VIS-023: Data-Driven Color Palette 

112 WCAG 2.1 contrast ratio guidelines (AA: 4.5:1, AAA: 7:1) 

113 ColorBrewer sequential/diverging schemes 

114 """ 

115 if n_colors < 1: 

116 raise ValueError("n_colors must be >= 1") 

117 if min_contrast_ratio < 1.0: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true

118 raise ValueError("min_contrast_ratio must be >= 1.0") 

119 

120 # Auto-select palette type if not specified 

121 if palette_type is None: 

122 palette_type = _auto_select_palette_type(n_colors, data_range) 

123 

124 # Select base palette 

125 if palette_type == "qualitative": 

126 base_colors = ( 

127 COLORBLIND_SAFE_QUALITATIVE if colorblind_safe else _generate_qualitative(n_colors) 

128 ) 

129 elif palette_type == "sequential": 

130 base_colors = SEQUENTIAL_VIRIDIS 

131 elif palette_type == "diverging": 131 ↛ 134line 131 didn't jump to line 134 because the condition on line 131 was always true

132 base_colors = DIVERGING_COOLWARM 

133 else: 

134 raise ValueError(f"Unknown palette_type: {palette_type}") 

135 

136 # Sample colors if we need fewer than available 

137 if n_colors <= len(base_colors): 137 ↛ 143line 137 didn't jump to line 143 because the condition on line 137 was always true

138 # Evenly sample from palette 

139 indices = np.linspace(0, len(base_colors) - 1, n_colors).astype(int) 

140 colors = [base_colors[i] for i in indices] 

141 else: 

142 # Interpolate if we need more colors 

143 colors = _interpolate_colors(base_colors, n_colors) 

144 

145 # Check contrast ratios 

146 colors_with_contrast = [] 

147 bg_luminance = _relative_luminance(background_color) 

148 

149 for color in colors: 

150 color_luminance = _relative_luminance(color) 

151 contrast = _contrast_ratio(color_luminance, bg_luminance) 

152 

153 if contrast >= min_contrast_ratio: 

154 colors_with_contrast.append(color) 

155 else: 

156 # Adjust lightness to meet contrast requirement 

157 adjusted = _adjust_for_contrast(color, background_color, min_contrast_ratio) 

158 colors_with_contrast.append(adjusted) 

159 

160 return colors_with_contrast 

161 

162 

163def _auto_select_palette_type( 

164 n_colors: int, 

165 data_range: tuple[float, float] | None, 

166) -> Literal["sequential", "diverging", "qualitative"]: 

167 """Auto-select palette type based on data characteristics. 

168 

169 Args: 

170 n_colors: Number of colors needed 

171 data_range: Data range (min, max) 

172 

173 Returns: 

174 Palette type 

175 """ 

176 # Check for bipolar data (zero-crossing) 

177 if data_range is not None: 

178 min_val, max_val = data_range 

179 if min_val < 0 and max_val > 0: 

180 # Bipolar signal - use diverging 

181 return "diverging" 

182 

183 # Multi-channel (distinct categories) 

184 if n_colors <= 8: 

185 return "qualitative" 

186 

187 # Many colors or continuous data 

188 return "sequential" 

189 

190 

191def _relative_luminance(color: str) -> float: 

192 """Calculate relative luminance per WCAG 2.1. 

193 

194 Args: 

195 color: Hex color code 

196 

197 Returns: 

198 Relative luminance (0-1) 

199 """ 

200 # Parse hex color 

201 color = color.removeprefix("#") 

202 

203 r = int(color[0:2], 16) / 255.0 

204 g = int(color[2:4], 16) / 255.0 

205 b = int(color[4:6], 16) / 255.0 

206 

207 # Convert to linear RGB 

208 def to_linear(c: float) -> float: 

209 if c <= 0.03928: 

210 return c / 12.92 

211 else: 

212 return ((c + 0.055) / 1.055) ** 2.4 # type: ignore[no-any-return] 

213 

214 r_linear = to_linear(r) 

215 g_linear = to_linear(g) 

216 b_linear = to_linear(b) 

217 

218 # Calculate luminance 

219 return 0.2126 * r_linear + 0.7152 * g_linear + 0.0722 * b_linear 

220 

221 

222def _contrast_ratio(lum1: float, lum2: float) -> float: 

223 """Calculate WCAG contrast ratio between two luminances. 

224 

225 Args: 

226 lum1: First luminance (0-1) 

227 lum2: Second luminance (0-1) 

228 

229 Returns: 

230 Contrast ratio (1-21) 

231 """ 

232 lighter = max(lum1, lum2) 

233 darker = min(lum1, lum2) 

234 

235 return (lighter + 0.05) / (darker + 0.05) 

236 

237 

238def _adjust_for_contrast( 

239 color: str, 

240 background: str, 

241 target_ratio: float, 

242) -> str: 

243 """Adjust color lightness to meet contrast requirement. 

244 

245 Args: 

246 color: Color to adjust 

247 background: Background color 

248 target_ratio: Target contrast ratio 

249 

250 Returns: 

251 Adjusted color hex code 

252 """ 

253 # Parse color 

254 color_val = color.removeprefix("#") 

255 

256 r = int(color_val[0:2], 16) 

257 g = int(color_val[2:4], 16) 

258 b = int(color_val[4:6], 16) 

259 

260 # Convert to HSL for easier lightness adjustment 

261 h, s, l = _rgb_to_hsl(r, g, b) # noqa: E741 

262 

263 bg_lum = _relative_luminance(background) 

264 

265 # Binary search for appropriate lightness 

266 l_min, l_max = 0.0, 1.0 

267 iterations = 0 

268 max_iterations = 20 

269 

270 while iterations < max_iterations: 

271 # Try current lightness 

272 test_r, test_g, test_b = _hsl_to_rgb(h, s, l) 

273 test_color = f"#{test_r:02x}{test_g:02x}{test_b:02x}" 

274 test_lum = _relative_luminance(test_color) 

275 ratio = _contrast_ratio(test_lum, bg_lum) 

276 

277 if abs(ratio - target_ratio) < 0.1: 277 ↛ 278line 277 didn't jump to line 278 because the condition on line 277 was never true

278 break 

279 

280 if ratio < target_ratio: 280 ↛ 291line 280 didn't jump to line 291 because the condition on line 280 was always true

281 # Need more contrast - adjust lightness 

282 if bg_lum > 0.5: 282 ↛ 288line 282 didn't jump to line 288 because the condition on line 282 was always true

283 # Dark background - make lighter 

284 l_min = l 

285 l = (l + l_max) / 2 # noqa: E741 

286 else: 

287 # Light background - make darker 

288 l_max = l 

289 l = (l_min + l) / 2 # noqa: E741 

290 # Too much contrast - move back 

291 elif bg_lum > 0.5: 

292 l_max = l 

293 l = (l_min + l) / 2 # noqa: E741 

294 else: 

295 l_min = l 

296 l = (l + l_max) / 2 # noqa: E741 

297 

298 iterations += 1 

299 

300 final_r, final_g, final_b = _hsl_to_rgb(h, s, l) 

301 return f"#{final_r:02x}{final_g:02x}{final_b:02x}" 

302 

303 

304def _rgb_to_hsl(r: int, g: int, b: int) -> tuple[float, float, float]: 

305 """Convert RGB to HSL color space. 

306 

307 Args: 

308 r: Red value (0-255). 

309 g: Green value (0-255). 

310 b: Blue value (0-255). 

311 

312 Returns: 

313 (h, s, l) tuple where h in [0, 360), s and l in [0, 1] 

314 """ 

315 r_norm = r / 255.0 

316 g_norm = g / 255.0 

317 b_norm = b / 255.0 

318 

319 max_c = max(r_norm, g_norm, b_norm) 

320 min_c = min(r_norm, g_norm, b_norm) 

321 delta = max_c - min_c 

322 

323 # Lightness 

324 l = (max_c + min_c) / 2.0 # noqa: E741 

325 

326 if delta == 0: 

327 # Achromatic 

328 return (0.0, 0.0, l) 

329 

330 # Saturation 

331 s = delta / (max_c + min_c) if l < 0.5 else delta / (2.0 - max_c - min_c) 

332 

333 # Hue 

334 if max_c == r_norm: 

335 h = ((g_norm - b_norm) / delta) % 6 

336 elif max_c == g_norm: 

337 h = ((b_norm - r_norm) / delta) + 2 

338 else: 

339 h = ((r_norm - g_norm) / delta) + 4 

340 

341 h = h * 60.0 

342 

343 return (h, s, l) 

344 

345 

346def _hsl_to_rgb(h: float, s: float, l: float) -> tuple[int, int, int]: # noqa: E741 

347 """Convert HSL to RGB color space. 

348 

349 Args: 

350 h: Hue in [0, 360) 

351 s: Saturation in [0, 1] 

352 l: Lightness in [0, 1] 

353 

354 Returns: 

355 (r, g, b) tuple with values in [0, 255] 

356 """ 

357 if s == 0: 

358 # Achromatic 

359 gray = int(l * 255) 

360 return (gray, gray, gray) 

361 

362 def hue_to_rgb(p: float, q: float, t: float) -> float: 

363 if t < 0: 

364 t += 1 

365 if t > 1: 

366 t -= 1 

367 if t < 1 / 6: 

368 return p + (q - p) * 6 * t 

369 if t < 1 / 2: 

370 return q 

371 if t < 2 / 3: 

372 return p + (q - p) * (2 / 3 - t) * 6 

373 return p 

374 

375 q = l * (1 + s) if l < 0.5 else l + s - l * s 

376 

377 p = 2 * l - q 

378 

379 h_norm = h / 360.0 

380 

381 r = hue_to_rgb(p, q, h_norm + 1 / 3) 

382 g = hue_to_rgb(p, q, h_norm) 

383 b = hue_to_rgb(p, q, h_norm - 1 / 3) 

384 

385 return (int(r * 255), int(g * 255), int(b * 255)) 

386 

387 

388def _generate_qualitative(n_colors: int) -> list[str]: 

389 """Generate qualitative color palette. 

390 

391 Args: 

392 n_colors: Number of colors 

393 

394 Returns: 

395 List of hex color codes 

396 """ 

397 # Generate evenly spaced hues 

398 colors = [] 

399 for i in range(n_colors): 

400 hue = (i * 360.0 / n_colors) % 360 

401 r, g, b = _hsl_to_rgb(hue, 0.7, 0.5) 

402 colors.append(f"#{r:02x}{g:02x}{b:02x}") 

403 

404 return colors 

405 

406 

407def _interpolate_colors(base_colors: list[str], n_colors: int) -> list[str]: 

408 """Interpolate between base colors to generate more colors. 

409 

410 Args: 

411 base_colors: Base color palette 

412 n_colors: Target number of colors 

413 

414 Returns: 

415 List of interpolated hex color codes 

416 """ 

417 if n_colors <= len(base_colors): 

418 return base_colors[:n_colors] 

419 

420 # Convert to RGB arrays 

421 rgb_array = np.zeros((len(base_colors), 3)) 

422 for i, color in enumerate(base_colors): 

423 color = color.removeprefix("#") 

424 rgb_array[i, 0] = int(color[0:2], 16) 

425 rgb_array[i, 1] = int(color[2:4], 16) 

426 rgb_array[i, 2] = int(color[4:6], 16) 

427 

428 # Interpolate 

429 indices = np.linspace(0, len(base_colors) - 1, n_colors) 

430 interp_rgb = np.zeros((n_colors, 3)) 

431 

432 for channel in range(3): 

433 interp_rgb[:, channel] = np.interp( 

434 indices, np.arange(len(base_colors)), rgb_array[:, channel] 

435 ) 

436 

437 # Convert back to hex 

438 colors = [] 

439 for i in range(n_colors): 

440 r = int(interp_rgb[i, 0]) 

441 g = int(interp_rgb[i, 1]) 

442 b = int(interp_rgb[i, 2]) 

443 colors.append(f"#{r:02x}{g:02x}{b:02x}") 

444 

445 return colors 

446 

447 

448__all__ = [ 

449 "COLORBLIND_SAFE_QUALITATIVE", 

450 "DIVERGING_COOLWARM", 

451 "SEQUENTIAL_VIRIDIS", 

452 "select_optimal_palette", 

453]