|
| 1 | +import numpy as np |
| 2 | +import matplotlib.pyplot as plt |
| 3 | +from matplotlib.patches import Arc |
| 4 | +from matplotlib.transforms import IdentityTransform |
| 5 | + |
| 6 | + |
| 7 | +class AngleArc(Arc): |
| 8 | + def __init__(self, xy, vec1, vec2, size=100, units="pixels", |
| 9 | + ax=None, fig=None, **kwargs): |
| 10 | + self._xydata = xy # in data coordinates |
| 11 | + self.ax = ax or plt.gca() |
| 12 | + self.fig = fig or plt.gcf() |
| 13 | + self.vec1 = vec1 # tuple or array of coordinates, relative to xy |
| 14 | + self.vec2 = vec2 # tuple or array of coordinates, relative to xy |
| 15 | + self.size = size |
| 16 | + |
| 17 | + super().__init__(self._xydata, size, size, angle=0.0, |
| 18 | + theta1=self.theta1, theta2=self.theta2, **kwargs) |
| 19 | + |
| 20 | + self.set_transform(IdentityTransform()) |
| 21 | + |
| 22 | + if units == "relative": |
| 23 | + fig.canvas.mpl_connect("resize_event", self._resize) |
| 24 | + |
| 25 | + self.ax.add_patch(self) |
| 26 | + |
| 27 | + def _resize(self, event): |
| 28 | + x0, y0 = self.ax.transAxes.transform((0, 0)) |
| 29 | + x1, y1 = self.ax.transAxes.transform((1, 1)) |
| 30 | + dx = x1 - x0 |
| 31 | + dy = y1 - y0 |
| 32 | + smallest = min(dx, dy) |
| 33 | + self.width = 0.25*dx |
| 34 | + self.height = 0.25*dx |
| 35 | + |
| 36 | + def get_center_pixels(self): |
| 37 | + """ return center in pixel coordinates """ |
| 38 | + return self.ax.transData.transform(self._xydata) |
| 39 | + |
| 40 | + def set_center(self, xy): |
| 41 | + """ set center in data coordinates """ |
| 42 | + self._xydata = xy |
| 43 | + |
| 44 | + _center = property(get_center_pixels, set_center) |
| 45 | + |
| 46 | + def get_theta(self, vec): |
| 47 | + vec_in_pixels = self.ax.transData.transform(vec) - self._center |
| 48 | + return np.rad2deg(np.arctan2(vec_in_pixels[1], vec_in_pixels[0])) |
| 49 | + |
| 50 | + def get_theta1(self): |
| 51 | + return self.get_theta(self.vec1) |
| 52 | + |
| 53 | + def get_theta2(self): |
| 54 | + return self.get_theta(self.vec2) |
| 55 | + |
| 56 | + def set_theta(self, angle): |
| 57 | + pass |
| 58 | + |
| 59 | + theta1 = property(get_theta1, set_theta) |
| 60 | + theta2 = property(get_theta2, set_theta) |
| 61 | + |
| 62 | + |
| 63 | +fig, ax = plt.subplots() |
| 64 | + |
| 65 | +ax.plot([2, .5, 1], [0, .2, 1]) |
| 66 | +am = AngleArc((0.5, 0.2), (2, 0), (1, 1), size=100, units="relative", |
| 67 | + ax=ax, fig=fig) |
| 68 | +plt.show() |
0 commit comments