Add base class for collection with width, height and angle to be used as parent class of EllipseCollection and RectangleCollection

ericpre · ericpre · commit abbb891e599e · 2024-04-02T17:57:31.000+01:00
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1720,8 +1720,11 @@ def __init__(self, sizes, **kwargs):
         self._paths = [mpath.Path.unit_circle()]
 
 
-class EllipseCollection(Collection):
-    """A collection of ellipses, drawn using splines."""
+class _CollectionWithWidthHeightAngle(Collection):
+    """
+    Base class for collections that have an array of widths, heights and angles
+    """
+    _factor = 1
 
     def __init__(self, widths, heights, angles, *, units='points', **kwargs):
         """
@@ -1751,7 +1754,7 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
         self._units = units
         self.set_transform(transforms.IdentityTransform())
         self._transforms = np.empty((0, 3, 3))
-        self._paths = [mpath.Path.unit_circle()]
+        self._paths = [self._path_generator()]
 
     def _set_transforms(self):
         """Calculate transforms immediately before drawing."""
@@ -1797,12 +1800,12 @@ def _set_transforms(self):
 
     def set_widths(self, widths):
         """Set the lengths of the first axes (e.g., major axis)."""
-        self._widths = 0.5 * np.asarray(widths).ravel()
+        self._widths = self._factor * np.asarray(widths).ravel()
         self.stale = True
 
     def set_heights(self, heights):
         """Set the lengths of second axes (e.g., minor axes)."""
-        self._heights = 0.5 * np.asarray(heights).ravel()
+        self._heights = self._factor * np.asarray(heights).ravel()
         self.stale = True
 
     def set_angles(self, angles):
@@ -1812,11 +1815,11 @@ def set_angles(self, angles):
 
     def get_widths(self):
         """Get the lengths of the first axes (e.g., major axis)."""
-        return self._widths * 2
+        return self._widths / self._factor
 
     def get_heights(self):
         """Set the lengths of second axes (e.g., minor axes)."""
-        return self._heights * 2
+        return self._heights / self._factor
 
     def get_angles(self):
         """Get the angles of the first axes, degrees CCW from the x-axis."""
@@ -1828,6 +1831,60 @@ def draw(self, renderer):
         super().draw(renderer)
 
 
+class EllipseCollection(_CollectionWithWidthHeightAngle):
+    """
+    A collection of ellipses, drawn using splines.
+
+    Parameters
+    ----------
+    widths : array-like
+        The lengths of the first axes (e.g., major axis lengths).
+    heights : array-like
+        The lengths of second axes.
+    angles : array-like
+        The angles of the first axes, degrees CCW from the x-axis.
+    units : {'points', 'inches', 'dots', 'width', 'height', 'x', 'y', 'xy'}
+        The units in which majors and minors are given; 'width' and
+        'height' refer to the dimensions of the axes, while 'x' and 'y'
+        refer to the *offsets* data units. 'xy' differs from all others in
+        that the angle as plotted varies with the aspect ratio, and equals
+        the specified angle only when the aspect ratio is unity.  Hence
+        it behaves the same as the `~.patches.Ellipse` with
+        ``axes.transData`` as its transform.
+    **kwargs
+        Forwarded to `Collection`.
+    """
+    _path_generator = mpath.Path.unit_circle
+    _factor = 0.5
+
+
+class RectangleCollection(_CollectionWithWidthHeightAngle):
+    """
+    A collection of rectangles, drawn using splines.
+
+    Parameters
+    ----------
+    widths : array-like
+        The lengths of the first axes (e.g., major axis lengths).
+    heights : array-like
+        The lengths of second axes.
+    angles : array-like
+        The angles of the first axes, degrees CCW from the x-axis.
+    units : {'points', 'inches', 'dots', 'width', 'height', 'x', 'y', 'xy'}
+        The units in which majors and minors are given; 'width' and
+        'height' refer to the dimensions of the axes, while 'x' and 'y'
+        refer to the *offsets* data units. 'xy' differs from all others in
+        that the angle as plotted varies with the aspect ratio, and equals
+        the specified angle only when the aspect ratio is unity.  Hence
+        it behaves the same as the `~.patches.Ellipse` with
+        ``axes.transData`` as its transform.
+    **kwargs
+        Forwarded to `Collection`.
+
+    """
+    _path_generator = mpath.Path.unit_rectangle
+
+
 class PatchCollection(Collection):
     """
     A generic collection of patches.
diff --git a/lib/matplotlib/collections.pyi b/lib/matplotlib/collections.pyi
@@ -168,7 +168,7 @@ class EventCollection(LineCollection):
 class CircleCollection(_CollectionWithSizes):
     def __init__(self, sizes: float | ArrayLike, **kwargs) -> None: ...
 
-class EllipseCollection(Collection):
+class _CollectionWithWidthHeightAngle(Collection):
     def __init__(
         self,
         widths: ArrayLike,
@@ -187,6 +187,9 @@ class EllipseCollection(Collection):
     def get_heights(self) -> ArrayLike: ...
     def get_angles(self) -> ArrayLike: ...
 
+class EllipseCollection(_CollectionWithWidthHeightAngle): ...
+class RectangleCollection(_CollectionWithWidthHeightAngle): ...
+
 class PatchCollection(Collection):
     def __init__(
         self, patches: Iterable[Patch], *, match_original: bool = ..., **kwargs
diff --git a/lib/matplotlib/tests/baseline_images/test_collections/RectangleCollection_test_image.png b/lib/matplotlib/tests/baseline_images/test_collections/RectangleCollection_test_image.png
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -410,7 +410,31 @@ def test_EllipseCollection():
     ax.autoscale_view()
 
 
-def test_EllipseCollection_setter_getter():
+@image_comparison(['RectangleCollection_test_image.png'], remove_text=True)
+def test_RectangleCollection():
+    # Test basic functionality
+    fig, ax = plt.subplots()
+    x = np.arange(4)
+    y = np.arange(3)
+    X, Y = np.meshgrid(x, y)
+    XY = np.vstack((X.ravel(), Y.ravel())).T
+
+    ww = X / x[-1]
+    hh = Y / y[-1]
+    aa = np.ones_like(ww) * 20  # first axis is 20 degrees CCW from x axis
+
+    ec = mcollections.RectangleCollection(
+        ww, hh, aa, units='x', offsets=XY, offset_transform=ax.transData,
+        facecolors='none')
+    ax.add_collection(ec)
+    ax.autoscale_view()
+
+
+@pytest.mark.parametrize(
+    'Class,scale',
+    [(mcollections.EllipseCollection, 0.5), (mcollections.RectangleCollection, 1)]
+    )
+def test_WidthHeightAngleCollection_setter_getter(Class, scale):
     # Test widths, heights and angle setter
     rng = np.random.default_rng(0)
 
@@ -421,7 +445,7 @@ def test_EllipseCollection_setter_getter():
 
     fig, ax = plt.subplots()
 
-    ec = mcollections.EllipseCollection(
+    ec = Class(
         widths=widths,
         heights=heights,
         angles=angles,
@@ -430,8 +454,8 @@ def test_EllipseCollection_setter_getter():
         offset_transform=ax.transData,
         )
 
-    assert_array_almost_equal(ec._widths, np.array(widths).ravel() * 0.5)
-    assert_array_almost_equal(ec._heights, np.array(heights).ravel() * 0.5)
+    assert_array_almost_equal(ec._widths, np.array(widths).ravel() * scale)
+    assert_array_almost_equal(ec._heights, np.array(heights).ravel() * scale)
     assert_array_almost_equal(ec._angles, np.deg2rad(angles).ravel())
 
     assert_array_almost_equal(ec.get_widths(), widths)
