Coverage for pyrc \ core \ visualization \ vpython.py: 38%

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# ------------------------------------------------------------------------------ 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11 

12from pyrc.core.visualization.color.vpython import cycle_palette_vpython 

13from pyrc.tools.functions import get_nested_attribute 

14 

15if TYPE_CHECKING: 

16 from pyrc.core.components.templates import Cell 

17 

18 

19def color_cells_by_attribute(nodes: list[Cell], attribute: str, opacity: float | int = 0.3, *args, **kwargs): 

20 """ 

21 Color nodes based on a specified attribute value. 

22 

23 Parameters 

24 ---------- 

25 nodes : list[Cell] 

26 List of Cell objects to color 

27 attribute : str 

28 Attribute path to use for coloring (e.g., 'material.name' or 'x') 

29 opacity : float, optional 

30 Opacity value that is set to every vbox. 

31 args, kwargs : 

32 Passed to cycle_palette_vpython 

33 Can be used to choose the color map like: cmap="managua50" 

34 """ 

35 color_generator = cycle_palette_vpython(*args, **kwargs) 

36 

37 # Map to track attribute values and their assigned colors 

38 attribute_colors = {} 

39 

40 # Color each node 

41 node: Cell 

42 for node in nodes: 

43 value = get_nested_attribute(node, attribute) 

44 

45 if value not in attribute_colors: 

46 # get new color from generator 

47 attribute_colors[value] = next(color_generator) 

48 

49 # Apply color vector to vbox 

50 node.vbox.color = attribute_colors[value] 

51 node.vbox.opacity = opacity