-
Notifications
You must be signed in to change notification settings - Fork 64
Expand file tree
/
Copy path_scatter.py
More file actions
620 lines (496 loc) · 20.6 KB
/
_scatter.py
File metadata and controls
620 lines (496 loc) · 20.6 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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
from typing import Sequence
import numpy as np
import pygfx
from ._base import (
GraphicFeature,
BufferManager,
GraphicFeatureEvent,
block_reentrance,
)
marker_names = {
# MPL
"o": "circle",
"s": "square",
"D": "diamond",
"+": "plus",
"x": "cross",
"^": "triangle_up",
"<": "triangle_left",
">": "triangle_right",
"v": "triangle_down",
"*": "asterisk6",
# Unicode
"●": "circle",
"○": "ring",
"■": "square",
"♦": "diamond",
"♥": "heart",
"♠": "spade",
"♣": "club",
"✳": "asterisk6",
"▲": "triangle_up",
"▼": "triangle_down",
"◀": "triangle_left",
"▶": "triangle_right",
# Emojis (these may look like their plaintext variants in your editor)
"❤️": "heart",
"♠️": "spade",
"♣️": "club",
"♦️": "diamond",
"💎": "diamond",
"💍": "ring",
"✳️": "asterisk6",
"📍": "pin",
}
def user_input_to_marker(name):
resolved_name = marker_names.get(name, name).lower()
if resolved_name not in pygfx.MarkerShape:
raise ValueError(
f"markers must be a string in: {list(pygfx.MarkerShape) + list(marker_names.keys())}, not {name!r}"
)
return resolved_name
def validate_user_markers_array(markers):
# make sure all markers are valid
# need to validate before converting to ints because
# we can't use control flow in the vectorized function
unique_values = np.unique(markers)
for m in unique_values:
user_input_to_marker(m)
# fast vectorized function to convert array of user markers to the standardized strings
# TODO: can probably use search-sorted for this too
vectorized_user_markers_to_std_markers = np.vectorize(marker_names.get, otypes=["<U14"])
# maps the human-readable marker name to the integers stored in the buffer
marker_int_mapping = dict(pygfx.MarkerInt.__members__)
# search sorted is the fastest way to map an array of str -> array of int
# see: https://github.com/pygfx/pygfx/issues/1215
# Prepare for searchsorted
def init_searchsorted(markers_mapping):
keys = np.array(list(markers_mapping.keys()))
vals = np.array(list(markers_mapping.values()))
order = np.argsort(keys)
keys = keys[order]
vals = vals[order]
return keys, vals
marker_int_searchsorted_keys, marker_int_searchsorted_vals = init_searchsorted(
marker_int_mapping
)
def searchsorted_markers_to_int_array(markers_str_array: np.ndarray[str]):
# Vectorized lookup
indices = np.searchsorted(marker_int_searchsorted_keys, markers_str_array)
return marker_int_searchsorted_vals[indices]
def parse_markers_init(markers: str | Sequence[str] | np.ndarray, n_datapoints: int):
# first validate then allocate buffers
if isinstance(markers, str):
markers = user_input_to_marker(markers)
elif isinstance(markers, (tuple, list, np.ndarray)):
validate_user_markers_array(markers)
# allocate buffers
markers_int_array = np.zeros(n_datapoints, dtype=np.int32)
marker_str_length = max(map(len, list(pygfx.MarkerShape)))
markers_readable_array = np.empty(n_datapoints, dtype=f"<U{marker_str_length}")
if isinstance(markers, str):
# all markers in the array are identical, so set the entire array
markers_readable_array[:] = markers
markers_int_array[:] = marker_int_mapping[markers]
elif isinstance(markers, (np.ndarray, tuple, list)):
# distinct marker for each point
# first vectorized map from user marker strings to "standard" marker strings
markers_readable_array = vectorized_user_markers_to_std_markers(markers)
# map standard marker strings to integer array
markers_int_array[:] = searchsorted_markers_to_int_array(markers_readable_array)
return markers_int_array, markers_readable_array
class VertexMarkers(BufferManager):
event_info_spec = [
{
"dict key": "key",
"type": "slice, index (int) or numpy-like fancy index",
"description": "key at which markers were indexed/sliced",
},
{
"dict key": "value",
"type": "str | np.ndarray[str]",
"description": "new marker values for points that were changed",
},
]
def __init__(
self,
markers: str | Sequence[str] | np.ndarray,
n_datapoints: int,
property_name: str = "markers",
):
"""
Manages the markers buffer for the scatter points. Supports fancy indexing.
"""
markers_int_array, self._markers_readable_array = parse_markers_init(
markers, n_datapoints
)
super().__init__(markers_int_array, property_name=property_name)
@property
def value(self) -> np.ndarray[str]:
"""numpy array of per-vertex marker shapes in human-readable form"""
return self._markers_readable_array
@property
def value_int(self) -> np.ndarray[np.int32]:
"""numpy array of the actual int32 buffer that represents per-vertex marker shapes on the GPU"""
return self.buffer.data
def _set_markers_arrays(self, key, value, n_markers):
if isinstance(value, str):
# set markers at these indices to this value
m = user_input_to_marker(value)
self._markers_readable_array[key] = m
self.value_int[key] = marker_int_mapping[m]
elif isinstance(value, (np.ndarray, list, tuple)):
if n_markers != len(value):
raise IndexError(
f"Must provide one marker value, or an array/list/tuple of marker values with the same length "
f"as the slice. You have provided the slice: {key}, which refers to {n_markers} markers, but "
f"provided {len(value)} new marker values. You must provide 1 or {n_markers} values."
)
validate_user_markers_array(value)
new_markers_human_readable = vectorized_user_markers_to_std_markers(value)
new_markers_int = searchsorted_markers_to_int_array(
new_markers_human_readable
)
self._markers_readable_array[key] = new_markers_human_readable
self.value_int[key] = new_markers_int
else:
raise TypeError(
"new markers value must be a str, Sequence or np.ndarray of new marker values"
)
def set_value(self, graphic, value):
"""set all the markers, create new buffer if necessary"""
if isinstance(value, (np.ndarray, list, tuple)):
if self.buffer.data.shape[0] != len(value):
# need to create a new buffer
markers_int_array, self._markers_readable_array = parse_markers_init(
value, len(value)
)
# create the new buffer, old buffer should get dereferenced
self._fpl_buffer = pygfx.Buffer(markers_int_array)
graphic.world_object.geometry.markers = self._fpl_buffer
self._emit_event(self._property_name, key=slice(None), value=value)
return
self[:] = value
@block_reentrance
def __setitem__(
self,
key: int | slice | list[int | bool] | np.ndarray[int | bool],
value: str | Sequence[str] | np.ndarray[str],
):
if isinstance(key, int):
if key >= self.value.size:
raise IndexError(f"index : {key} out of bounds: {self.value.size}")
if not isinstance(value, str):
# only a single marker should be provided if changing one at one index
raise TypeError(
f"you must provide a <str> marker value if providing a single <int> index, "
f"you have passed index: {key} and value: {value}"
)
m = user_input_to_marker(value)
self._markers_readable_array[key] = m
self.value_int[key] = marker_int_mapping[m]
elif isinstance(key, slice):
# find the number of new markers by converting slice to range and then parse markers
start, stop, step = key.indices(self.value.size)
n_markers = len(range(start, stop, step))
self._set_markers_arrays(key, value, n_markers)
elif isinstance(key, (list, np.ndarray)):
key = np.asarray(key) # convert to array if list
if key.dtype == bool:
# make sure len is same
if not key.size == self.buffer.data.shape[0]:
raise IndexError(
f"Length of array for fancy indexing must match number of datapoints.\n"
f"There are {len(self.buffer.data.shape[0])} datapoints and you have passed "
f"a bool array of size: {key.size}"
)
n_markers = np.count_nonzero(key)
self._set_markers_arrays(key, value, n_markers)
# if it's an array of int
elif np.issubdtype(key.dtype, np.integer):
if key.size > self.buffer.data.shape[0]:
raise IndexError(
f"Length of array for fancy indexing must be <= n_datapoints. "
f"There are: {self.buffer.data.shape[0]} datapoints, you have passed an "
f"integer array for fancy indexing of size: {key.size}"
)
n_markers = key.size
self._set_markers_arrays(key, value, n_markers)
else:
# fancy indexing doesn't make sense with non-integer types
raise TypeError(
f"can only using integer or booleans arrays for fancy indexing, your array is of type: {key.dtype}"
)
else:
raise TypeError(
f"Can only set markers by slicing/indexing using the one of the following types: "
f"int | slice | list[int | bool] | np.ndarray[int | bool], you have passed"
f"sliced using the following type: {type(key)}"
)
# _update_range handles parsing the key to
# determine offset and size for GPU upload
self._update_range(key)
self._emit_event(self._property_name, key, value)
def __len__(self):
return len(self.buffer.data)
class UniformMarker(GraphicFeature):
event_info_spec = [
{
"dict key": "value",
"type": "str | None",
"description": "new marker value",
},
]
def __init__(self, marker: str, property_name: str = "markers"):
"""Manages evented uniform buffer for scatter marker"""
self._value = user_input_to_marker(marker)
super().__init__(property_name=property_name)
@property
def value(self) -> str:
return self._value
@block_reentrance
def set_value(self, graphic, value: str):
value = user_input_to_marker(value)
graphic.world_object.material.marker = value
self._value = value
event = GraphicFeatureEvent(type=self._property_name, info={"value": value})
self._call_event_handlers(event)
class UniformEdgeColor(GraphicFeature):
event_info_spec = [
{
"dict key": "value",
"type": "str | np.ndarray | pygfx.Color | Sequence[float]",
"description": "new edge_color",
},
]
def __init__(
self,
edge_color: str | np.ndarray | pygfx.Color | Sequence[float],
property_name: str = "edge_colors",
):
"""Manages evented uniform buffer for scatter marker edge_color"""
self._value = pygfx.Color(edge_color)
super().__init__(property_name=property_name)
@property
def value(self) -> pygfx.Color:
return self._value
@block_reentrance
def set_value(
self, graphic, value: str | np.ndarray | pygfx.Color | Sequence[float]
):
graphic.world_object.material.edge_color = pygfx.Color(value)
self._value = value
event = GraphicFeatureEvent(type=self._property_name, info={"value": value})
self._call_event_handlers(event)
class EdgeWidth(GraphicFeature):
event_info_spec = [
{
"dict key": "value",
"type": "float",
"description": "new edge_width",
},
]
def __init__(self, edge_width: float, property_name: str = "edge_width"):
"""Manages evented uniform buffer for scatter marker edge_width"""
self._value = edge_width
super().__init__(property_name=property_name)
@property
def value(self) -> float:
return self._value
@block_reentrance
def set_value(self, graphic, value: float):
graphic.world_object.material.edge_width = value
self._value = value
event = GraphicFeatureEvent(type=self._property_name, info={"value": value})
self._call_event_handlers(event)
class UniformRotations(GraphicFeature):
event_info_spec = [
{
"dict key": "value",
"type": "float",
"description": "new edge_width",
},
]
def __init__(self, edge_width: float, property_name: str = "point_rotations"):
"""Manages evented uniform buffer for scatter marker rotation"""
self._value = edge_width
super().__init__(property_name=property_name)
@property
def value(self) -> float:
return self._value
@block_reentrance
def set_value(self, graphic, value: float):
graphic.world_object.material.rotations = value
self._value = value
event = GraphicFeatureEvent(type=self._property_name, info={"value": value})
self._call_event_handlers(event)
class VertexRotations(BufferManager):
event_info_spec = [
{
"dict key": "key",
"type": "slice, index (int) or numpy-like fancy index",
"description": "key at which point rotations were indexed/sliced",
},
{
"dict key": "value",
"type": "int | float | array-like",
"description": "new rotation values for points that were changed",
},
]
def __init__(
self,
rotations: int | float | np.ndarray | Sequence[int | float],
n_datapoints: int,
property_name: str = "point_rotations",
):
"""
Manages rotations buffer of scatter points.
"""
sizes = self._fix_rotations(rotations, n_datapoints)
super().__init__(data=sizes, property_name=property_name)
def _fix_rotations(
self,
sizes: int | float | np.ndarray | Sequence[int | float],
n_datapoints: int,
):
if np.issubdtype(type(sizes), np.number):
# single value given
sizes = np.full(
n_datapoints, sizes, dtype=np.float32
) # force it into a float to avoid weird gpu errors
elif isinstance(
sizes, (np.ndarray, tuple, list)
): # if it's not a ndarray already, make it one
sizes = np.asarray(sizes, dtype=np.float32) # read it in as a numpy.float32
if (sizes.ndim != 1) or (sizes.size != n_datapoints):
raise ValueError(
f"sequence of `rotations` must be 1 dimensional with "
f"the same length as the number of datapoints"
)
else:
raise TypeError(
"`rotations` must be a single <int>, <float>, or a sequence (array, list, tuple) of int"
"or float with the length equal to the number of datapoints"
)
return sizes
def set_value(self, graphic, value):
"""set all rotations, create new buffer if necessary"""
if isinstance(value, (np.ndarray, list, tuple)):
if self.buffer.data.shape[0] != value.shape[0]:
# need to create a new buffer
value = self._fix_rotations(value, len(value))
data = np.empty(shape=(len(value),), dtype=np.float32)
# create the new buffer, old buffer should get dereferenced
self._fpl_buffer = pygfx.Buffer(data)
graphic.world_object.geometry.rotations = self._fpl_buffer
self._emit_event(self._property_name, key=slice(None), value=value)
return
self[:] = value
@block_reentrance
def __setitem__(
self,
key: int | slice | np.ndarray[int | bool] | list[int | bool],
value: int | float | np.ndarray | Sequence[int | float],
):
# this is a very simple 1D buffer, no parsing required, directly set buffer
self.buffer.data[key] = value
self._update_range(key)
self._emit_event(self._property_name, key, value)
def __len__(self):
return len(self.buffer.data)
class VertexPointSizes(BufferManager):
event_info_spec = [
{
"dict key": "key",
"type": "slice, index (int) or numpy-like fancy index",
"description": "key at which point sizes were indexed/sliced",
},
{
"dict key": "value",
"type": "int | float | array-like",
"description": "new size values for points that were changed",
},
]
def __init__(
self,
sizes: int | float | np.ndarray | Sequence[int | float],
n_datapoints: int,
property_name: str = "sizes",
):
"""
Manages sizes buffer of scatter points.
"""
sizes = self._fix_sizes(sizes, n_datapoints)
super().__init__(data=sizes, property_name=property_name)
def _fix_sizes(
self,
sizes: int | float | np.ndarray | Sequence[int | float],
n_datapoints: int,
):
if np.issubdtype(type(sizes), np.number):
# single value given
sizes = np.full(
n_datapoints, sizes, dtype=np.float32
) # force it into a float to avoid weird gpu errors
elif isinstance(
sizes, (np.ndarray, tuple, list)
): # if it's not a ndarray already, make it one
sizes = np.asarray(sizes, dtype=np.float32) # read it in as a numpy.float32
if (sizes.ndim != 1) or (sizes.size != n_datapoints):
raise ValueError(
f"sequence of `sizes` must be 1 dimensional with "
f"the same length as the number of datapoints"
)
else:
raise TypeError(
"sizes must be a single <int>, <float>, or a sequence (array, list, tuple) of int"
"or float with the length equal to the number of datapoints"
)
if np.count_nonzero(sizes < 0) > 1:
raise ValueError(
"All sizes must be positive numbers greater than or equal to 0.0."
)
return sizes
def set_value(self, graphic, value):
"""set all sizes, create new buffer if necessary"""
if isinstance(value, (np.ndarray, list, tuple)):
if self.buffer.data.shape[0] != len(value):
# create new buffer
value = self._fix_sizes(value, len(value))
data = np.empty(shape=(len(value),), dtype=np.float32)
# create the new buffer, old buffer should get dereferenced
self._fpl_buffer = pygfx.Buffer(data)
graphic.world_object.geometry.sizes = self._fpl_buffer
self._emit_event(self._property_name, key=slice(None), value=value)
return
self[:] = value
@block_reentrance
def __setitem__(
self,
key: int | slice | np.ndarray[int | bool] | list[int | bool],
value: int | float | np.ndarray | Sequence[int | float],
):
# this is a very simple 1D buffer, no parsing required, directly set buffer
self.buffer.data[key] = value
self._update_range(key)
self._emit_event(self._property_name, key, value)
def __len__(self):
return len(self.buffer.data)
class UniformSize(GraphicFeature):
event_info_spec = [
{"dict key": "value", "type": "float", "description": "new size value"},
]
def __init__(self, value: int | float, property_name: str = "sizes"):
"""Manages uniform size for scatter material"""
self._value = float(value)
super().__init__(property_name=property_name)
@property
def value(self) -> float:
return self._value
@block_reentrance
def set_value(self, graphic, value: float | int):
value = float(value)
graphic.world_object.material.size = value
self._value = value
event = GraphicFeatureEvent(type=self._property_name, info={"value": value})
self._call_event_handlers(event)