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

Skip to content

Commit 47f94dd

Browse files
ksundenksunden
committed
Ellipse collection containerization
1 parent b6e4ca4 commit 47f94dd

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.
@@ -2212,82 +2295,72 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
22122295
self.set_widths(widths)
22132296
self.set_heights(heights)
22142297
self.set_angles(angles)
2215-
self._units = units
2298+
self._container.units = units
22162299
self.set_transform(transforms.IdentityTransform())
22172300
self._paths = [mpath.Path.unit_circle()]
22182301

2219-
def _set_transforms(self):
2220-
"""Calculate transforms immediately before drawing."""
2302+
def _init_container(self):
2303+
return EllipseCollectionContainer(
2304+
x=np.array([]),
2305+
y=np.array([]),
2306+
edgecolors=np.array([]),
2307+
facecolors=np.array([]),
2308+
hatchcolors=np.array([]),
2309+
widths=np.array([]),
2310+
heights=np.array([]),
2311+
angles=np.array([]),
2312+
units="xy"
2313+
)
22212314

2222-
ax = self.axes
2223-
fig = self.get_figure(root=False)
22242315

2225-
if self._units == 'xy':
2226-
sc = 1
2227-
elif self._units == 'x':
2228-
sc = ax.bbox.width / ax.viewLim.width
2229-
elif self._units == 'y':
2230-
sc = ax.bbox.height / ax.viewLim.height
2231-
elif self._units == 'inches':
2232-
sc = fig.dpi
2233-
elif self._units == 'points':
2234-
sc = fig.dpi / 72.0
2235-
elif self._units == 'width':
2236-
sc = ax.bbox.width
2237-
elif self._units == 'height':
2238-
sc = ax.bbox.height
2239-
elif self._units == 'dots':
2240-
sc = 1.0
2241-
else:
2242-
raise ValueError(f'Unrecognized units: {self._units!r}')
2243-
2244-
self._container.transforms = np.zeros((len(self._widths), 3, 3))
2245-
widths = self._widths * sc
2246-
heights = self._heights * sc
2247-
sin_angle = np.sin(self._angles)
2248-
cos_angle = np.cos(self._angles)
2249-
self._container.transforms[:, 0, 0] = widths * cos_angle
2250-
self._container.transforms[:, 0, 1] = heights * -sin_angle
2251-
self._container.transforms[:, 1, 0] = widths * sin_angle
2252-
self._container.transforms[:, 1, 1] = heights * cos_angle
2253-
self._container.transforms[:, 2, 2] = 1.0
2254-
2255-
_affine = transforms.Affine2D
2256-
if self._units == 'xy':
2257-
m = ax.transData.get_affine().get_matrix().copy()
2258-
m[:2, 2:] = 0
2259-
self.set_transform(_affine(m))
2316+
def set_angles(self, angles):
2317+
"""Set the angles of the first axes, degrees CCW from the x-axis."""
2318+
if not isinstance(self._container, EllipseCollectionContainer):
2319+
raise TypeError("Cannot use 'set_angles' on custom container types")
2320+
self._container.angles = np.deg2rad(angles).ravel()
2321+
self.stale = True
22602322

22612323
def set_widths(self, widths):
22622324
"""Set the lengths of the first axes (e.g., major axis)."""
2263-
self._widths = 0.5 * np.asarray(widths).ravel()
2325+
if not isinstance(self._container, EllipseCollectionContainer):
2326+
raise TypeError("Cannot use 'set_widths' on custom container types")
2327+
self._container.widths = 0.5 * np.asarray(widths).ravel()
22642328
self.stale = True
22652329

22662330
def set_heights(self, heights):
22672331
"""Set the lengths of second axes (e.g., minor axes)."""
2268-
self._heights = 0.5 * np.asarray(heights).ravel()
2269-
self.stale = True
2270-
2271-
def set_angles(self, angles):
2272-
"""Set the angles of the first axes, degrees CCW from the x-axis."""
2273-
self._angles = np.deg2rad(angles).ravel()
2332+
if not isinstance(self._container, EllipseCollectionContainer):
2333+
raise TypeError("Cannot use 'set_heights' on custom container types")
2334+
self._container.heights = 0.5 * np.asarray(heights).ravel()
22742335
self.stale = True
22752336

22762337
def get_widths(self):
22772338
"""Get the lengths of the first axes (e.g., major axis)."""
2278-
return self._widths * 2
2339+
if not isinstance(self._container, EllipseCollectionContainer):
2340+
raise TypeError("Cannot use 'get_widths' on custom container types")
2341+
return self._container.widths * 2
22792342

22802343
def get_heights(self):
2281-
"""Set the lengths of second axes (e.g., minor axes)."""
2282-
return self._heights * 2
2344+
"""Get the lengths of second axes (e.g., minor axes)."""
2345+
if not isinstance(self._container, EllipseCollectionContainer):
2346+
raise TypeError("Cannot use 'get_heights' on custom container types")
2347+
return self._container.heights * 2
22832348

22842349
def get_angles(self):
22852350
"""Get the angles of the first axes, degrees CCW from the x-axis."""
2286-
return np.rad2deg(self._angles)
2351+
if not isinstance(self._container, EllipseCollectionContainer):
2352+
raise TypeError("Cannot use 'get_angles' on custom container types")
2353+
return np.rad2deg(self._container.angles)
22872354

22882355
@artist.allow_rasterization
22892356
def draw(self, renderer):
2290-
self._set_transforms()
2357+
if (
2358+
isinstance(self._container, EllipseCollectionContainer)
2359+
and self._container.units == "xy"
2360+
):
2361+
m = self.axes.transData.get_affine().get_matrix().copy()
2362+
m[:2, 2:] = 0
2363+
self.set_transform(transforms.Affine2D(m))
22912364
super().draw(renderer)
22922365

22932366

‎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)