| 89 | |
| 90 | |
| 91 | class Ellipse: |
| 92 | def __init__(self, x, y, width, height, theta): |
| 93 | self.x = x |
| 94 | self.y = y |
| 95 | self.width = width |
| 96 | self.height = height |
| 97 | self.theta = theta # in radians |
| 98 | self._geometry = None |
| 99 | |
| 100 | @property |
| 101 | def parameters(self): |
| 102 | return self.x, self.y, self.width, self.height, self.theta |
| 103 | |
| 104 | @property |
| 105 | def aspect_ratio(self): |
| 106 | return max(self.width, self.height) / min(self.width, self.height) |
| 107 | |
| 108 | def calc_similarity_with(self, other_ellipse): |
| 109 | max_dist = max(self.height, self.width, other_ellipse.height, other_ellipse.width) |
| 110 | dist = math.sqrt((self.x - other_ellipse.x) ** 2 + (self.y - other_ellipse.y) ** 2) |
| 111 | |
| 112 | if max_dist == 0: |
| 113 | max_dist = 1 |
| 114 | |
| 115 | cost1 = 1 - min(dist / max_dist, 1) |
| 116 | cost2 = abs(math.cos(self.theta - other_ellipse.theta)) |
| 117 | return 0.8 * cost1 + 0.2 * cost2 * cost1 |
| 118 | |
| 119 | def contains_points(self, xy, tol=0.1): |
| 120 | ca = math.cos(self.theta) |
| 121 | sa = math.sin(self.theta) |
| 122 | x_demean = xy[:, 0] - self.x |
| 123 | y_demean = xy[:, 1] - self.y |
| 124 | return ( |
| 125 | ((ca * x_demean + sa * y_demean) ** 2 / (0.5 * self.width) ** 2) |
| 126 | + ((sa * x_demean - ca * y_demean) ** 2 / (0.5 * self.height) ** 2) |
| 127 | ) <= 1 + tol |
| 128 | |
| 129 | def draw(self, show_axes=True, ax=None, **kwargs): |
| 130 | import matplotlib.pyplot as plt |
| 131 | from matplotlib.lines import Line2D |
| 132 | from matplotlib.transforms import Affine2D |
| 133 | |
| 134 | if ax is None: |
| 135 | ax = plt.subplot(111, aspect="equal") |
| 136 | el = patches.Ellipse( |
| 137 | xy=(self.x, self.y), |
| 138 | width=self.width, |
| 139 | height=self.height, |
| 140 | angle=np.rad2deg(self.theta), |
| 141 | **kwargs, |
| 142 | ) |
| 143 | ax.add_patch(el) |
| 144 | if show_axes: |
| 145 | major = Line2D([-self.width / 2, self.width / 2], [0, 0], lw=3, zorder=3) |
| 146 | minor = Line2D([0, 0], [-self.height / 2, self.height / 2], lw=3, zorder=3) |
| 147 | trans = Affine2D().rotate(self.theta).translate(self.x, self.y) + ax.transData |
| 148 | major.set_transform(trans) |