-
Notifications
You must be signed in to change notification settings - Fork 64
Expand file tree
/
Copy pathscatter.py
More file actions
129 lines (101 loc) · 4.26 KB
/
scatter.py
File metadata and controls
129 lines (101 loc) · 4.26 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
from typing import *
import numpy as np
import pygfx
from ._positions_base import PositionsGraphic
from ._features import PointsSizesFeature, UniformSize, SizeSpace
class ScatterGraphic(PositionsGraphic):
_features = {"data", "sizes", "colors", "cmap", "size_space"}
def __init__(
self,
data: Any,
colors: str | np.ndarray | tuple[float] | list[float] | list[str] = "w",
uniform_color: bool = False,
alpha: float = 1.0,
cmap: str = None,
cmap_transform: np.ndarray = None,
isolated_buffer: bool = True,
sizes: float | np.ndarray | Iterable[float] = 1,
uniform_size: bool = False,
size_space: str = "screen",
**kwargs,
):
"""
Create a Scatter Graphic, 2d or 3d
Parameters
----------
data: array-like
Scatter data to plot, 2D must be of shape [n_points, 2], 3D must be of shape [n_points, 3]
colors: str, array, or iterable, default "w"
specify colors as a single human readable string, a single RGBA array,
or an iterable of strings or RGBA arrays
uniform_color: bool, default False
if True, uses a uniform buffer for the scatter point colors,
basically saves GPU VRAM when the entire line has a single color
alpha: float, optional, default 1.0
alpha value for the colors
cmap: str, optional
apply a colormap to the scatter instead of assigning colors manually, this
overrides any argument passed to "colors"
cmap_transform: 1D array-like or list of numerical values, optional
if provided, these values are used to map the colors from the cmap
isolated_buffer: bool, default True
whether the buffers should be isolated from the user input array.
Generally always ``True``, ``False`` is for rare advanced use.
sizes: float or iterable of float, optional, default 1.0
size of the scatter points
uniform_size: bool, default False
if True, uses a uniform buffer for the scatter point sizes,
basically saves GPU VRAM when all scatter points are the same size
size_space: str, default "screen"
coordinate space in which the size is expressed ("screen", "world", "model")
kwargs
passed to Graphic
"""
super().__init__(
data=data,
colors=colors,
uniform_color=uniform_color,
alpha=alpha,
cmap=cmap,
cmap_transform=cmap_transform,
isolated_buffer=isolated_buffer,
size_space=size_space,
**kwargs,
)
n_datapoints = self.data.value.shape[0]
geo_kwargs = {"positions": self._data.buffer}
material_kwargs = {"pick_write": True}
self._size_space = SizeSpace(size_space)
if uniform_color:
material_kwargs["color_mode"] = "uniform"
material_kwargs["color"] = self.colors
else:
material_kwargs["color_mode"] = "vertex"
geo_kwargs["colors"] = self.colors.buffer
if uniform_size:
material_kwargs["size_mode"] = "uniform"
self._sizes = UniformSize(sizes)
material_kwargs["size"] = self.sizes
else:
material_kwargs["size_mode"] = "vertex"
self._sizes = PointsSizesFeature(sizes, n_datapoints=n_datapoints)
geo_kwargs["sizes"] = self.sizes.buffer
material_kwargs["size_space"] = self.size_space
world_object = pygfx.Points(
pygfx.Geometry(**geo_kwargs),
material=pygfx.PointsMaterial(**material_kwargs),
)
self._set_world_object(world_object)
@property
def sizes(self) -> PointsSizesFeature | float:
"""Get or set the scatter point size(s)"""
if isinstance(self._sizes, PointsSizesFeature):
return self._sizes
elif isinstance(self._sizes, UniformSize):
return self._sizes.value
@sizes.setter
def sizes(self, value):
if isinstance(self._sizes, PointsSizesFeature):
self._sizes[:] = value
elif isinstance(self._sizes, UniformSize):
self._sizes.set_value(self, value)