| 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 | |
|---|
| 31 | if base is not None: |
|---|
| 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 range(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 |
|---|
| 68 | |
|---|
| 69 | |
|---|
