Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a40efd2

Browse files
ksundenksunden
committed
Ellipse collection containerization
1 parent ab27eb7 commit a40efd2

2 files changed

Lines changed: 129 additions & 56 deletions

File tree

‎lib/matplotlib/collections.py‎

Lines changed: 126 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import numpy as np
1919

20-
from mpl_data_containers.description import Desc
20+
from mpl_data_containers.description import Desc, desc_like
2121

2222
import matplotlib as mpl
2323
from . import (_api, _path, artist, cbook, colorizer as mcolorizer, colors as mcolors,
@@ -124,6 +124,89 @@ def query(self, graph, parent_coordinates="axes"):
124124
return d, hash
125125

126126

127+
class EllipseCollectionContainer(CollectionContainer):
128+
def __init__(
129+
self,
130+
x: np.array,
131+
y: np.array,
132+
edgecolors: np.array,
133+
facecolors: np.array,
134+
hatchcolors: np.array,
135+
widths: np.array,
136+
heights: np.array,
137+
angles: np.array,
138+
units: str,
139+
):
140+
super().__init__(x, y, edgecolors, facecolors, hatchcolors)
141+
self.widths = np.atleast_1d(widths)
142+
self.heights = np.atleast_1d(heights)
143+
self.angles = np.atleast_1d(angles)
144+
self.units = units
145+
146+
def query(self, graph, parent_coordinates="axes"):
147+
# TODO: get dpi from graph or refactor transform to be dpi independent
148+
dpi = 100.0
149+
d, hash = super().query(graph, parent_coordinates)
150+
151+
# TODO: this section is verbose and likely to be useful elsewhere
152+
# Consider moving to one or more helper methods
153+
# For reference, this was originally from FuncContainer, with modifications
154+
desc = Desc(("N",))
155+
xy = {"x": desc, "y": desc}
156+
data_lim = graph.evaluator(
157+
desc_like(xy, coordinates="data"),
158+
desc_like(xy, coordinates=parent_coordinates),
159+
).inverse
160+
161+
screen_size = graph.evaluator(
162+
desc_like(xy, coordinates=parent_coordinates),
163+
desc_like(xy, coordinates="display"),
164+
)
165+
166+
screen_dims = screen_size.evaluate({"x": [0, 1], "y": [0, 1]})
167+
xpix, ypix = np.ceil(np.abs(np.diff(screen_dims["x"]))), np.ceil(
168+
np.abs(np.diff(screen_dims["y"]))
169+
)
170+
data_dims = data_lim.evaluate({"x": [0, 1], "y": [0, 1]})
171+
xdata, ydata = np.abs(np.diff(data_dims["x"])), np.abs(np.diff(data_dims["y"]))
172+
173+
if self.units == 'xy':
174+
sc = 1
175+
elif self.units == 'x':
176+
sc = xpix / xdata
177+
elif self.units == 'y':
178+
sc = ypix / ydata
179+
elif self.units == 'inches':
180+
sc = dpi
181+
elif self.units == 'points':
182+
sc = dpi / 72.0
183+
elif self.units == 'width':
184+
sc = xpix
185+
elif self.units == 'height':
186+
sc = ypix
187+
elif self.units == 'dots':
188+
sc = 1.0
189+
else:
190+
raise ValueError(f'Unrecognized units: {self._units!r}')
191+
192+
193+
print(f"{sc=}, {self.units=}")
194+
transforms = np.zeros((len(self.widths), 3, 3))
195+
widths = self.widths * sc
196+
heights = self.heights * sc
197+
sin_angle = np.sin(self.angles)
198+
cos_angle = np.cos(self.angles)
199+
transforms[:, 0, 0] = widths * cos_angle
200+
transforms[:, 0, 1] = heights * -sin_angle
201+
transforms[:, 1, 0] = widths * sin_angle
202+
transforms[:, 1, 1] = heights * cos_angle
203+
transforms[:, 2, 2] = 1.0
204+
205+
d["transforms"] = transforms
206+
207+
return d, hash
208+
209+
127210

128211
# "color" is excluded; it is a compound setter, and its docstring differs
129212
# in LineCollection.
@@ -2228,82 +2311,72 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
22282311
self.set_widths(widths)
22292312
self.set_heights(heights)
22302313
self.set_angles(angles)
2231-
self._units = units
2314+
self._container.units = units
22322315
self.set_transform(transforms.IdentityTransform())
22332316
self._paths = [mpath.Path.unit_circle()]
22342317

2235-
def _set_transforms(self):
2236-
"""Calculate transforms immediately before drawing."""
2318+
def _init_container(self):
2319+
return EllipseCollectionContainer(
2320+
x=np.array([]),
2321+
y=np.array([]),
2322+
edgecolors=np.array([]),
2323+
facecolors=np.array([]),
2324+
hatchcolors=np.array([]),
2325+
widths=np.array([]),
2326+
heights=np.array([]),
2327+
angles=np.array([]),
2328+
units="xy"
2329+
)
22372330

2238-
ax = self.axes
2239-
fig = self.get_figure(root=False)
22402331

2241-
if self._units == 'xy':
2242-
sc = 1
2243-
elif self._units == 'x':
2244-
sc = ax.bbox.width / ax.viewLim.width
2245-
elif self._units == 'y':
2246-
sc = ax.bbox.height / ax.viewLim.height
2247-
elif self._units == 'inches':
2248-
sc = fig.dpi
2249-
elif self._units == 'points':
2250-
sc = fig.dpi / 72.0
2251-
elif self._units == 'width':
2252-
sc = ax.bbox.width
2253-
elif self._units == 'height':
2254-
sc = ax.bbox.height
2255-
elif self._units == 'dots':
2256-
sc = 1.0
2257-
else:
2258-
raise ValueError(f'Unrecognized units: {self._units!r}')
2259-
2260-
self._container.transforms = np.zeros((len(self._widths), 3, 3))
2261-
widths = self._widths * sc
2262-
heights = self._heights * sc
2263-
sin_angle = np.sin(self._angles)
2264-
cos_angle = np.cos(self._angles)
2265-
self._container.transforms[:, 0, 0] = widths * cos_angle
2266-
self._container.transforms[:, 0, 1] = heights * -sin_angle
2267-
self._container.transforms[:, 1, 0] = widths * sin_angle
2268-
self._container.transforms[:, 1, 1] = heights * cos_angle
2269-
self._container.transforms[:, 2, 2] = 1.0
2270-
2271-
_affine = transforms.Affine2D
2272-
if self._units == 'xy':
2273-
m = ax.transData.get_affine().get_matrix().copy()
2274-
m[:2, 2:] = 0
2275-
self.set_transform(_affine(m))
2332+
def set_angles(self, angles):
2333+
"""Set the angles of the first axes, degrees CCW from the x-axis."""
2334+
if not isinstance(self._container, EllipseCollectionContainer):
2335+
raise TypeError("Cannot use 'set_angles' on custom container types")
2336+
self._container.angles = np.deg2rad(angles).ravel()
2337+
self.stale = True
22762338

22772339
def set_widths(self, widths):
22782340
"""Set the lengths of the first axes (e.g., major axis)."""
2279-
self._widths = 0.5 * np.asarray(widths).ravel()
2341+
if not isinstance(self._container, EllipseCollectionContainer):
2342+
raise TypeError("Cannot use 'set_widths' on custom container types")
2343+
self._container.widths = 0.5 * np.asarray(widths).ravel()
22802344
self.stale = True
22812345

22822346
def set_heights(self, heights):
22832347
"""Set the lengths of second axes (e.g., minor axes)."""
2284-
self._heights = 0.5 * np.asarray(heights).ravel()
2285-
self.stale = True
2286-
2287-
def set_angles(self, angles):
2288-
"""Set the angles of the first axes, degrees CCW from the x-axis."""
2289-
self._angles = np.deg2rad(angles).ravel()
2348+
if not isinstance(self._container, EllipseCollectionContainer):
2349+
raise TypeError("Cannot use 'set_heights' on custom container types")
2350+
self._container.heights = 0.5 * np.asarray(heights).ravel()
22902351
self.stale = True
22912352

22922353
def get_widths(self):
22932354
"""Get the lengths of the first axes (e.g., major axis)."""
2294-
return self._widths * 2
2355+
if not isinstance(self._container, EllipseCollectionContainer):
2356+
raise TypeError("Cannot use 'get_widths' on custom container types")
2357+
return self._container.widths * 2
22952358

22962359
def get_heights(self):
2297-
"""Set the lengths of second axes (e.g., minor axes)."""
2298-
return self._heights * 2
2360+
"""Get the lengths of second axes (e.g., minor axes)."""
2361+
if not isinstance(self._container, EllipseCollectionContainer):
2362+
raise TypeError("Cannot use 'get_heights' on custom container types")
2363+
return self._container.heights * 2
22992364

23002365
def get_angles(self):
23012366
"""Get the angles of the first axes, degrees CCW from the x-axis."""
2302-
return np.rad2deg(self._angles)
2367+
if not isinstance(self._container, EllipseCollectionContainer):
2368+
raise TypeError("Cannot use 'get_angles' on custom container types")
2369+
return np.rad2deg(self._container.angles)
23032370

23042371
@artist.allow_rasterization
23052372
def draw(self, renderer):
2306-
self._set_transforms()
2373+
if (
2374+
isinstance(self._container, EllipseCollectionContainer)
2375+
and self._container.units == "xy"
2376+
):
2377+
m = self.axes.transData.get_affine().get_matrix().copy()
2378+
m[:2, 2:] = 0
2379+
self.set_transform(transforms.Affine2D(m))
23072380
super().draw(renderer)
23082381

23092382

‎lib/matplotlib/tests/test_collections.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -430,9 +430,9 @@ def test_EllipseCollection_setter_getter():
430430
offset_transform=ax.transData,
431431
)
432432

433-
assert_array_almost_equal(ec._widths, np.array(widths).ravel() * 0.5)
434-
assert_array_almost_equal(ec._heights, np.array(heights).ravel() * 0.5)
435-
assert_array_almost_equal(ec._angles, np.deg2rad(angles).ravel())
433+
assert_array_almost_equal(ec._container.widths, np.array(widths).ravel() * 0.5)
434+
assert_array_almost_equal(ec._container.heights, np.array(heights).ravel() * 0.5)
435+
assert_array_almost_equal(ec._container.angles, np.deg2rad(angles).ravel())
436436

437437
assert_array_almost_equal(ec.get_widths(), widths)
438438
assert_array_almost_equal(ec.get_heights(), heights)

0 commit comments

Comments
 (0)