Coverage for mango/tools/visualiser.py : 26%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# conda create -c conda-forge -n vispy_env python=3 vispy pyqt5 numpy
2# conda install pyopengl
3# conda install imageio imageio-ffmpeg
4# conda install pillow
5"""
6Mango particle visualiser.
8Controls:
10* 0 - reset cameras
11* 1 - toggle camera between regular 3D (turntable), first person (fly),
12 arcball and Pan Zoom
13* 2 - toggle between real particles and vector spheres
14* +/- - speed up/slow down simulation
15* ! - start/stop simulation
16* . - reset simulation speed
17* I - save image of current frame
19With fly camera:
21* WASD or arrow keys - move around
22* SPACE - brake
23* FC - move up-down
24* IJKL or mouse - look around
26TODO:
28* Change arrow size based on strength
29* speed up read in
30"""
32from sys import argv
33import numpy as np
34import imageio
36from vispy import app, scene, util
37from vispy.visuals.transforms import STTransform, MatrixTransform
38from vispy.visuals.transforms._util import as_vec4
39from vispy.visuals.visual import CompoundVisual
40from vispy.visuals.mesh import MeshVisual
41from vispy.scene import visuals
42from vispy.geometry import create_arrow
44from mango.managers import serverwrapper
45from mango import c
46from mango.debug import DBG
49def read(filename, speed):
50 """
51 Read in extended xyz file and names columns in recorded arrays.
53 Future:
55 * Error checking
56 * crash gracefully when memory gets too full
57 """
58 with open(filename) as f:
59 # Managing initial lines of xyz file
60 data_list = []
61 num = 0
63 for i, line in enumerate(f):
64 if i % 5e6 == 0 and i != 0: 64 ↛ 65line 64 didn't jump to line 65, because the condition on line 64 was never true
65 print(f'Read {int(i*5e6)} frames')
66 if i == 0:
67 no_mol = int(line)
68 elif i == 1:
69 comments = {}
70 line = line.split()
71 for comment in line:
72 com = comment.split("=")
73 try:
74 comments[com[0]] = float(com[1].split('[')[0])
75 except ValueError:
76 comments[com[0]] = com[1] == 'True'
77 except IndexError:
78 pass
79 # comments["i"] = int(comments["i"])
80 elif i % (no_mol + 2) == 1 or i % (no_mol + 2) == 0:
81 # ignoring future comments other than the first
82 continue
83 elif num % speed < no_mol: 83 ↛ 88line 83 didn't jump to line 88, because the condition on line 83 was never false
84 # skip frames
85 data_list += [line.split()]
86 num += 1
87 else:
88 num += 1
90 # Naming columns
91 dtypes = []
92 for i in range(len(data_list[0])):
93 if i == 0:
94 dtypes += [("name", "U5")]
95 elif i == 1:
96 dtypes += [('x', float)]
97 elif i == 2:
98 dtypes += [('y', float)]
99 elif i == 3:
100 dtypes += [('z', float)]
101 else:
102 dtypes += [('col{}'.format(i - 4), float)]
104 xyz_dump = np.rec.array([tuple(x) for x in data_list], dtype=dtypes)
105 size = xyz_dump.x.shape[0]
106 print('Number of frames: ', size // no_mol)
108 return {"xyz_dump": xyz_dump, "size": size, "no_mol": no_mol,
109 "mag": comments['mag'] if 'mag' in comments else False,
110 "force": comments['force'] if 'force' in comments else False,
111 "momentum": comments['momentum'] if 'momentum' in comments else False,
112 'box': comments['boxsize'] if 'boxsize' in comments else False}
115class Arrow3DVisual(CompoundVisual):
116 def __init__(self, radius=1.0, length=1.0, cols=30, rows=30, vertex_colors=None, face_colors=None,
117 color=(0.5, 0.5, 1, 1), edge_color=None, **kwargs):
119 mesh = create_arrow(rows, cols, radius=radius, length=length)
120 self._mesh = MeshVisual(vertices=mesh.get_vertices(),
121 faces=mesh.get_faces(),
122 vertex_colors=vertex_colors,
123 face_colors=face_colors, color=color)
124 if edge_color:
125 self._border = MeshVisual(vertices=mesh.get_vertices(),
126 faces=mesh.get_edges(),
127 color=edge_color, mode='lines')
128 else:
129 self._border = MeshVisual()
131 CompoundVisual.__init__(self, [self._mesh, self._border], **kwargs)
132 self.mesh.set_gl_state(polygon_offset_fill=True,
133 polygon_offset=(1, 1), depth_test=True)
135 @property
136 def mesh(self):
137 return self._mesh
139 @property
140 def border(self):
141 return self._border
144class scene_setup():
146 def __init__(self, application, radii, skip=1, freeze=False, true_size=True, fps=False, video=False, sm_box=True, loop=True):
147 self.app = application
149 # Prepare canvas
150 canvas = scene.SceneCanvas(keys='interactive', size=(800, 608), show=True, bgcolor='white')
151 if fps:
152 canvas.measure_fps()
153 canvas.events.mouse_move.connect(self.on_mouse_move)
154 canvas.events.key_press.connect(self.on_key_press)
155 self.render = canvas.render
156 self.csize = canvas.size
157 self.view = canvas.central_widget.add_view()
159 if video:
160 self.writer = imageio.get_writer('animation.mp4', quality=10, fps=20)
161 self.video = self.makevideo
162 self.canvas = canvas
163 else:
164 self.video = self._noop
166 fov = 60.
167 self.loops = 1
168 self.f = 0
169 self.inc = 1
170 self.stop_at_end = not loop
171 self.cont = False
172 self.camera(fov)
174 with np.errstate(all='raise'):
175 if not true_size:
176 self.cr_bsph = self._noop
177 self.scene(*self.get_data(radii, freeze, sm_box, skip))
179 self.timer = app.Timer('auto', connect=self.on_timer, start=False, app=self.app)
181 try:
182 self.app.run()
183 except:
184 pass
185 finally:
186 if video:
187 self.writer.close()
189 @staticmethod
190 def _noop(*args):
191 pass
193 def makevideo(self):
194 self.writer.append_data(self.canvas.render())
196 def on_timer(self, event):
197 self.f += self.inc
198 if self.f >= self.no_iters:
199 if self.stop_at_end and not self.cont:
200 self.timer.stop()
201 if self.f - self.inc < self.no_iters - 1:
202 self.f -= self.inc
203 else:
204 self.cont = True
205 print("LastFrame ", end='\r')
206 return
207 elif self.stop_at_end:
208 self.cont = False
209 self.f = 0
210 print(f"Start x{self.loops} ", end='\r')
211 self.loops += 1
213 for i in self.ispheres:
214 self.ispheres[i].transform.move_to(self.trans[i, self.f])
215 self.spheres[i].transform.move_to(self.trans[i, self.f])
216 self.move_arr(i)
218 self.video()
220 def get_data(self, radii, freeze, sm_box, skip=1):
221 a = read(argv[1], skip)
222 n = a['no_mol']
224 pradius = np.zeros(n) + radii
225 self.framefilelength = int(np.log10(a['size'] / n))
227 xyz = self.positions(a, freeze)
229 self.arrows(a)
231 self.no_iters = xyz.shape[1]
232 box = np.sum(pradius * 2) * 1.5 if sm_box or not a['box'] else a['box']
233 print('BOXSIZE', box)
234 return xyz, n, pradius, box
236 # def maxmin(self, arr1, arr2):
237 # norm1 = np.linalg.norm(arr1, axis=-1)
238 # norm2 = np.linalg.norm(arr2, axis=-1)
239 # conesize1 = norm1 / np.max(norm1, axis=-1)[..., None]
240 # conesize2 = norm2 / np.max(norm2, axis=-1)[..., None]
241 # self.fcones[i].transform.scale(self.fconesize[i, self.f], self.xyz[i, self.f])
242 # self.fcones[i].transform.scale(self.fconesize[i, self.f], xyz[i, self.f])
243 # conesize = self.maxmin(arr1, arr2)
244 # return conesize1 / conesize2
246 def positions(self, a, freeze):
247 n = a['no_mol']
248 x = a['xyz_dump'].x
249 y = a['xyz_dump'].y
250 z = a['xyz_dump'].z
252 xyz = np.array([x, y, z]).T
253 xyz = xyz.reshape(n, -1, 3, order='F')
255 if freeze:
256 xyz -= xyz[int(n / 2)][None] + 0.01
258 trans = []
259 for framenum, frame in enumerate(xyz):
260 trans += [[]]
261 for mol in frame:
262 trans[framenum] += [as_vec4(mol)]
264 self.trans = np.array(trans)
266 return xyz
268 def arrows(self, a):
269 n = a['no_mol']
270 cnum = np.arange(3, 9) if a['momentum'] else np.arange(0, 6)
271 if a['force'] and a['mag']:
272 fx = getattr(a['xyz_dump'], f'col{cnum[0]}')
273 fy = getattr(a['xyz_dump'], f'col{cnum[1]}')
274 fz = getattr(a['xyz_dump'], f'col{cnum[2]}')
275 mx = getattr(a['xyz_dump'], f'col{cnum[3]}')
276 my = getattr(a['xyz_dump'], f'col{cnum[4]}')
277 mz = getattr(a['xyz_dump'], f'col{cnum[5]}')
278 self.cr_arrs = self.cr_arr2
279 self.move_arr = self._mv_arr2
280 elif a['force']:
281 fx = getattr(a['xyz_dump'], f'col{cnum[0]}')
282 fy = getattr(a['xyz_dump'], f'col{cnum[1]}')
283 fz = getattr(a['xyz_dump'], f'col{cnum[2]}')
284 mx = my = mz = np.zeros_like(a['xyz_dump'].x)
285 self.cr_arrs = self._cr_arrf
286 self.move_arr = self._mv_arrf
287 elif a['mag']:
288 mx = getattr(a['xyz_dump'], f'col{cnum[0]}')
289 my = getattr(a['xyz_dump'], f'col{cnum[1]}')
290 mz = getattr(a['xyz_dump'], f'col{cnum[2]}')
291 fx = fy = fz = np.zeros_like(a['xyz_dump'].x)
292 self.cr_arrs = self._cr_arrm
293 self.move_arr = self._mv_arrm
294 else:
295 self.cr_arrs = self.move_arr = self._noop
297 if a['force']:
298 fxyz = np.array([fx, fy, fz]).T
299 fxyz = fxyz.reshape(n, -1, 3, order='F')
300 self.frotater, self.initf, self.initfa = self.cone_rotation(fxyz)
302 if a['mag']:
303 mxyz = np.array([mx, my, mz]).T
304 mxyz = mxyz.reshape(n, -1, 3, order='F')
305 self.mrotater, self.initm, self.initma = self.cone_rotation(mxyz)
306 if DBG:
307 print(self.initm, self.initma)
309 def cone_rotation(self, arr1):
310 if arr1.shape[1] > 1:
311 arr2 = self.spinner(arr1)
312 angle, rotax = self.angle(arr1, arr2, 'xyz, nix, niy ->niz')
313 angle = self.spinner2(angle)
314 rotax = self.spinner2(rotax)
315 else:
316 angle = rotax = None
317 initial_dir = np.zeros((arr1.shape[0], 3))
318 initial_dir[:, 1] = 1
319 if DBG:
320 print(arr1.shape)
321 init_angle, init_rot = self.angle(initial_dir, arr1[:, 0], 'xyz, ix, iy ->iz')
322 return self.rotate_array(angle, rotax, init_angle, init_rot), init_angle, init_rot
324 def rotate_array(self, angle, rotax, ia, ir):
325 matrix = self.rotdp(np.tile(np.eye(4), (ia.shape[0], 1, 1)), ia, ir)
326 if angle is not None:
327 rotater = np.zeros((angle.shape[0], angle.shape[1], matrix.shape[1], matrix.shape[2]))
328 for n2 in range(angle.shape[1]):
329 rotater[:, n2] = self.rotdp(matrix, angle[:, n2], rotax[:, n2]).copy()
330 else:
331 rotater = matrix[:, None, ...]
332 return rotater
334 @staticmethod
335 def rotdp(matrix, angle, axis):
336 for n in range(angle.shape[0]):
337 matrix[n] = np.dot(matrix[n], util.transforms.rotate(angle[n], axis[n])) if angle[n] != 0 else matrix[n]
338 return matrix
340 def angle(self, arr1, arr2, string=None):
341 arr1_ = np.linalg.norm(arr1, axis=-1)[..., None]
342 arr2_ = np.linalg.norm(arr2, axis=-1)[..., None]
343 dp = np.einsum('...j, ...j -> ...', arr1, arr2)[..., None] / (arr1_ * arr2_)
344 try:
345 cp = np.einsum(string, c.eijk, arr1, arr2) / (arr1_ * arr2_ * np.sin(dp))
346 except FloatingPointError:
347 cp = np.zeros_like(dp)
348 return (np.pi / 2 - np.arccos(dp)) * (180 / np.pi), cp
350 def scene(self, xyz, n, pradius, bsize):
351 self.spheres = {}
352 self.ispheres = {}
353 self.fcones = {}
354 self.mcones = {}
356 # Create an XYZAxis visual
357 self.axis = scene.visuals.XYZAxis(parent=self.view)
358 self.axis.transform = STTransform(translate=(50, 50), scale=(50, 50, 50, 1)).as_matrix()
360 self.box = visuals.Box(width=bsize, height=bsize, depth=bsize, color=None, edge_color='black', parent=self.view.scene)
361 self.box.transform = STTransform(translate=[0, 0, 0])
362 self.translucent(self.box)
364 for i in range(n):
365 self.ispheres[i] = visuals.Sphere(radius=pradius[i] * 0.1, method='ico', parent=self.view.scene,
366 color='black', subdivisions=2)
367 self.ispheres[i].transform = STTransform(translate=xyz[i, self.f])
368 self.cr_bsph(i, pradius[i], xyz[i, self.f])
369 self.cr_arrs(i, pradius[i])
371 def camera(self, fov):
372 # Create four cameras (Fly, Turntable, Arcball and PanZoom)
373 self.cam1 = scene.cameras.TurntableCamera(parent=self.view.scene, fov=fov, scale_factor=10,
374 name='Turntable')
375 self.cam2 = scene.cameras.FlyCamera(parent=self.view.scene, fov=fov, scale_factor=10, name='Fly')
376 self.cam3 = scene.cameras.ArcballCamera(parent=self.view.scene, fov=fov, scale_factor=10, name='Arcball')
377 self.cam4 = scene.cameras.PanZoomCamera(parent=self.view.scene, aspect=1, name='PanZoom')
378 self.cam4.zoom(10, center=(0.5, 0.5, 0))
379 self.view.camera = self.cam1 # Select turntable at first
381 def _mv_arrf(self, i):
382 self._arrow_move_to(self.fcones[i], self.frotater[i, self.f], self.trans[i, self.f])
384 def _mv_arrm(self, i):
385 self._arrow_move_to(self.mcones[i], self.mrotater[i, self.f], self.trans[i, self.f])
387 def _mv_arr2(self, i):
388 self._arrow_move_to(self.fcones[i], self.frotater[i, self.f], self.trans[i, self.f])
389 self._arrow_move_to(self.mcones[i], self.mrotater[i, self.f], self.trans[i, self.f])
391 def cr_bsph(self, i, pradius, xyz):
392 self.spheres[i] = visuals.Sphere(radius=pradius, method='ico', parent=self.view.scene,
393 color=(0.5, 0.6, 1, 0.1), edge_color=(0.5, 0.6, 1, 1), subdivisions=2)
394 self.translucent(self.spheres[i])
395 self.spheres[i].transform = STTransform(translate=xyz)
397 def _cr_arrf(self, i, pradius):
398 self.fcones[i] = visuals.Arrow3D(radius=pradius * 0.01, length=pradius * 0.3, parent=self.view.scene, cols=20,
399 rows=4, color=(0.5, 0.6, 1, 1), edge_color=(0.5, 0.5, 1, 1))
400 self.fcones[i].transform = MatrixTransform()
401 self._arrow_move(self.fcones[i], self.initf[i], self.initfa[i], self.trans[i, self.f])
403 def _cr_arrm(self, i, pradius):
404 self.mcones[i] = visuals.Arrow3D(radius=pradius * 0.01, length=pradius * 0.3, parent=self.view.scene, cols=20,
405 rows=4, color=(1.0, 0.5, 0.5, 1.0), edge_color=(1.0, 0.0, 0.0, 1.0))
406 self.mcones[i].transform = MatrixTransform()
407 self._arrow_move(self.mcones[i], self.initm[i], self.initma[i], self.trans[i, self.f])
409 def cr_arr2(self, i, pradius):
410 self._cr_arrf(i, pradius)
411 self._cr_arrm(i, pradius)
413 @staticmethod
414 def spinner(arr):
415 arr2 = arr.copy()
416 arr2[:, :-1] = arr[:, 1:].copy()
417 arr2[:, -1] = arr[:, 0].copy()
418 return arr2
420 @staticmethod
421 def spinner2(arr):
422 arr2 = arr.copy()
423 arr2[..., 1:] = arr[..., :-1].copy()
424 arr2[..., :1] = arr[..., -1:].copy()
425 return arr2
427 @staticmethod
428 def translucent(obj):
429 obj.mesh.set_gl_state('translucent', polygon_offset_fill=True,
430 polygon_offset=(1, 1), depth_test=False)
432 @staticmethod
433 def _arrow_move(obj, angle, raxis, trans):
434 if angle != 0:
435 obj.transform.rotate(angle, raxis)
436 obj.transform.translate_to(trans)
438 @staticmethod
439 def _arrow_move_to(obj, rot, trans):
440 obj.transform.rotate_to(rot)
441 obj.transform.translate_to(trans)
443 def on_mouse_move(self, event):
444 # Implement axis connection with cam1
445 if event.button == 1 and event.is_dragging:
446 self.axis.transform.reset()
448 self.axis.transform.rotate(self.cam1.roll, (0, 0, 1))
449 self.axis.transform.rotate(self.cam1.elevation, (1, 0, 0))
450 self.axis.transform.rotate(self.cam1.azimuth, (0, 1, 0))
452 self.axis.transform.scale((50, 50, 0.001))
453 self.axis.transform.translate((50., 50.))
454 self.axis.update()
456 def on_key_press(self, event):
457 # Implement key presses
458 if event.text == '1':
459 cam_toggle = {self.cam1: self.cam2, self.cam2: self.cam3,
460 self.cam3: self.cam4, self.cam4: self.cam1}
461 self.view.camera = cam_toggle.get(self.view.camera, self.cam2)
462 print(self.view.camera.name + ' camera')
463 self.axis.visible = self.view.camera is self.cam1
464 elif event.text == '0':
465 self.cam1.set_range()
466 self.cam2.set_range()
467 self.cam3.set_range()
468 elif event.text == '+':
469 if self.timer.interval > 0.01:
470 self.timer.interval -= 0.01
471 else:
472 self.inc *= 2
473 print(self.inc, self.timer.interval)
474 elif event.text == '-':
475 if self.inc >= 2:
476 self.inc //= 2
477 else:
478 self.timer.interval += 0.01
479 print(self.inc, self.timer.interval)
480 elif event.text == '.':
481 self.inc = 1
482 self.timer.interval = 1 / 60
483 elif event.text == '!':
484 if self.timer.running:
485 self.timer.stop()
486 else:
487 self.timer.start()
488 elif event.text == '2':
489 for i in self.spheres.values():
490 i.visible = not i.visible
491 elif event.text == 'I':
492 name = f'{self.f:0{self.framefilelength}d}.png'
493 imageio.imwrite(name, self.render())
494 print(f'frame {name} saved')
497def move_to(self, trans):
498 self.translate = trans
501def translate_to(self, trans):
502 self.matrix[-1] = trans
505def rotate_to(self, rot):
506 self.matrix = rot
509@serverwrapper('Error')
510def main():
511 print(__doc__)
512 helpstring = '{}{}{}'.format("Run with:\n mango_vis <xyzfile> <radii> <-s -v -m -l>\n",
513 " -s skipframes (default: 300)\n -v verbose\n",
514 " -m record video\n -l loop visualisation")
515 if '-h' in argv:
516 print(helpstring)
517 return
518 fps = ('-v' in argv)
519 video = ('-m' in argv)
520 loop = ('-l' in argv)
522 if '-s' in argv:
523 sind = argv.index('-s')
524 skip = int(argv[sind + 1])
525 del argv[sind + 1]
526 del argv[sind]
527 else:
528 skip = False
529 for v in ['-v', '-m', '-l']:
530 if v in argv:
531 del argv[argv.index(v)]
533 if len(argv) != 3:
534 print('xyz file or radii not given')
535 print(helpstring)
536 exit(1)
537 radii = np.array(argv[2:], dtype=float)
538 scene_setup(app.Application(), radii=radii, skip=skip or 300, freeze=False, fps=fps, video=video, sm_box=False, loop=loop)
541STTransform.move_to = move_to
542MatrixTransform.translate_to = translate_to
543MatrixTransform.rotate_to = rotate_to
544visuals.Arrow3D = visuals.create_visual_node(Arrow3DVisual)
546if __name__ == '__main__':
547 main()