forked from fastplotlib/fastplotlib
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsubplot.py
More file actions
127 lines (99 loc) · 4.21 KB
/
Copy pathsubplot.py
File metadata and controls
127 lines (99 loc) · 4.21 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
import pygfx
from pygfx import Scene, OrthographicCamera, PerspectiveCamera, PanZoomController, Viewport, AxesHelper, GridHelper
from pygfx.linalg import Vector3
from .graphics import *
from .defaults import camera_types, controller_types
from typing import *
from wgpu.gui.auto import WgpuCanvas
class Subplot:
def __init__(
self,
position: Tuple[int, int] = None,
parent_dims: Tuple[int, int] = None,
camera: str = '2d',
controller: Union[pygfx.PanZoomController, pygfx.OrbitOrthoController] = None,
canvas: WgpuCanvas = None,
renderer: pygfx.Renderer = None
):
self.scene: pygfx.Scene = pygfx.Scene()
if canvas is None:
canvas = WgpuCanvas()
if renderer is None:
renderer = pygfx.renderers.WgpuRenderer(canvas)
self.canvas = canvas
self.renderer = renderer
if position is None:
position = (0, 0)
self.position: Tuple[int, int] = position
if parent_dims is None:
parent_dims = (1, 1)
self.nrows, self.ncols = parent_dims
self.camera: Union[pygfx.OrthographicCamera, pygfx.PerspectiveCamera] = camera_types[camera]()
if controller is None:
controller = controller_types[camera]()
self.controller: Union[pygfx.PanZoomController, pygfx.OrbitOrthoController] = controller
# might be better as an attribute of GridPlot
# but easier to iterate when in same object as camera and scene
self.viewport: pygfx.Viewport = pygfx.Viewport(renderer)
self.controller.add_default_event_handlers(
self.viewport,
self.camera
)
self._axes: AxesHelper = AxesHelper(size=100)
for arrow in self._axes.children:
self._axes.remove(arrow)
self._axes.set_colors('r', 'g', 'b')
self._grid: GridHelper = GridHelper(size=100, thickness=1)
self._animate_funcs = list()
def _produce_rect(self, i, j, w, h):
# print(locals())
return [
((w / self.ncols) + ((j - 1) * (w / self.ncols))),
((h / self.nrows) + ((i - 1) * (h / self.nrows))),
(w / self.ncols),
(h / self.nrows)
]
def _resize(self, canvas_dims: Tuple[int, int]):
# w, h = self.canvas.get_logical_size()
self.viewport.rect = self._produce_rect(*self.position, *canvas_dims)
def animate(self, canvas_dims: Tuple[int, int] = None):
if canvas_dims is None:
canvas_dims = self.canvas.get_logical_size()
self.controller.update_camera(self.camera)
self._resize(canvas_dims)
self.viewport.render(self.scene, self.camera)
for f in self._animate_funcs:
f()
def add_animations(self, funcs: List[callable]):
self._animate_funcs += funcs
def add_graphic(self, graphic):
self.scene.add(graphic.world_object)
if isinstance(graphic, Image):
dims = graphic.data.shape
zero_pos = Vector3(dims[0] / 2, dims[1] / 2, self.camera.position.z)
delta = zero_pos.clone().sub(self.camera.position)
zoom_level = 1 / np.mean(dims)
if self.controller.zoom_value != zoom_level:
self.controller.zoom(zoom_level)
self.controller.pan(delta)
elif isinstance(graphic, Scatter):
self.camera.show_object(graphic.world_object)
# centroid = np.mean(graphic.data, axis=0).tolist()
# zero_pos = Vector3(*centroid)
# delta = zero_pos.clone().sub(self.camera.position)
# zoom_level = 1 / np.mean(graphic.data)
# if self.controller.zoom_value != zoom_level:
# self.controller.zoom(zoom_level)
# self.controller.pan(delta)
def set_axes_visibility(self, visible: bool):
if visible:
self.scene.add(self._axes)
else:
self.scene.remove(self._axes)
def set_grid_visibility(self, visible: bool):
if visible:
self.scene.add(self._grid)
else:
self.scene.remove(self._grid)
def remove_graphic(self, graphic):
self.scene.remove(graphic.world_object)