-
Notifications
You must be signed in to change notification settings - Fork 64
Expand file tree
/
Copy pathscatter.py
More file actions
122 lines (95 loc) · 3.94 KB
/
scatter.py
File metadata and controls
122 lines (95 loc) · 3.94 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
from typing import *
import numpy as np
import pygfx
from ._positions_base import PositionsGraphic
from ._features import PointsSizesFeature, UniformSize
class ScatterGraphic(PositionsGraphic):
_features = {"data", "sizes", "colors", "cmap"}
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,
**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
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,
**kwargs,
)
n_datapoints = self.data.value.shape[0]
geo_kwargs = {"positions": self._data.buffer}
material_kwargs = {"pick_write": True}
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
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)