[959eb01] | 1 | """ |
---|
| 2 | Module that draws multiple arrows in 3D coordinates |
---|
| 3 | """ |
---|
| 4 | |
---|
| 5 | from matplotlib.patches import FancyArrowPatch |
---|
| 6 | from mpl_toolkits.mplot3d import proj3d |
---|
| 7 | import time |
---|
| 8 | # from matplotlib.artist import allow_rasterization |
---|
| 9 | class Arrow3D(FancyArrowPatch): |
---|
| 10 | """ |
---|
| 11 | Draw 3D arrow |
---|
| 12 | """ |
---|
| 13 | |
---|
| 14 | def __init__(self, base, xs, ys, zs, colors, *args, **kwargs): |
---|
| 15 | """ |
---|
| 16 | Init |
---|
| 17 | |
---|
| 18 | :Params xs: [[x0, x0+dx0], [x1, x1+dx1], ...] |
---|
| 19 | :Params ys: [[y0, y0+dy0], [y1, y1+dy1], ...] |
---|
| 20 | :Params zs: [[z0, z0+dz0], [z1, z1+dz1], ...] |
---|
| 21 | :Params colors: [[R0, G0, B0], [R1, G1, B1], ...] |
---|
| 22 | where R, G, B ranges (0,1) |
---|
| 23 | """ |
---|
| 24 | FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs) |
---|
| 25 | self.leftdown = False |
---|
| 26 | self.t_click = 0 |
---|
| 27 | self._verts3d = xs, ys, zs |
---|
| 28 | self.colors = colors |
---|
| 29 | self.base = base |
---|
| 30 | |
---|
[7432acb] | 31 | if base is not None: |
---|
[959eb01] | 32 | # To turn the updating off during dragging |
---|
| 33 | base.canvas.mpl_connect('button_press_event', self.on_left_down) |
---|
| 34 | base.canvas.mpl_connect('button_release_event', self.on_left_up) |
---|
| 35 | |
---|
| 36 | def on_left_down(self, event): |
---|
| 37 | """ |
---|
| 38 | Mouse left-down event |
---|
| 39 | """ |
---|
| 40 | self.leftdown = True |
---|
| 41 | self.t_click = time.time() |
---|
| 42 | |
---|
| 43 | def on_left_up(self, event): |
---|
| 44 | """ |
---|
| 45 | Mouse left up event |
---|
| 46 | """ |
---|
| 47 | t_up = time.time() - self.t_click |
---|
| 48 | # Avoid just clicking |
---|
| 49 | if t_up > 0.1: |
---|
| 50 | self.leftdown = False |
---|
| 51 | self.base.canvas.draw() |
---|
| 52 | |
---|
| 53 | def draw(self, renderer, rasterized=True): |
---|
| 54 | """ |
---|
| 55 | Drawing actually happens here |
---|
| 56 | """ |
---|
| 57 | # Draws only when the dragging finished |
---|
| 58 | if self.leftdown: |
---|
| 59 | return |
---|
| 60 | xs3d, ys3d, zs3d = self._verts3d |
---|
| 61 | for i in xrange(len(xs3d)): |
---|
| 62 | xs, ys, _ = proj3d.proj_transform(xs3d[i], ys3d[i], zs3d[i], renderer.M) |
---|
| 63 | self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) |
---|
| 64 | self.set_color(self.colors[i]) |
---|
| 65 | FancyArrowPatch.draw(self, renderer) |
---|
| 66 | |
---|
| 67 | self.leftdown = False |
---|