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

ericpre · ericpre · commit 78de816de7a1 · 2024-04-02T19:02:33.000+01:00
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1720,8 +1720,10 @@ 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
+    """
 
     def __init__(self, widths, heights, angles, *, units='points', **kwargs):
         """
@@ -1751,7 +1753,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 +1799,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 = 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 = np.asarray(heights).ravel()
         self.stale = True
 
     def set_angles(self, angles):
@@ -1812,11 +1814,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
 
     def get_heights(self):
         """Set the lengths of second axes (e.g., minor axes)."""
-        return self._heights * 2
+        return self._heights
 
     def get_angles(self):
         """Get the angles of the first axes, degrees CCW from the x-axis."""
@@ -1828,6 +1830,59 @@ 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.half_unit_circle
+
+
+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/path.py b/lib/matplotlib/path.py
@@ -799,6 +799,21 @@ def unit_regular_asterisk(cls, numVertices):
         """
         return cls.unit_regular_star(numVertices, 0.0)
 
+    _half_unit_circle = None
+
+    @classmethod
+    def half_unit_circle(cls):
+        """
+        Return the readonly :class:`Path` of the half unit circle.
+
+        For most cases, :func:`Path.circle` will be what you want.
+        """
+        if cls._half_unit_circle is None:
+            cls._half_unit_circle = cls.circle(
+                center=(0, 0), radius=0.5, readonly=True
+                )
+        return cls._half_unit_circle
+
     _unit_circle = None
 
     @classmethod
diff --git a/lib/matplotlib/path.pyi b/lib/matplotlib/path.pyi
@@ -107,6 +107,8 @@ class Path:
     @classmethod
     def unit_regular_asterisk(cls, numVertices: int) -> Path: ...
     @classmethod
+    def half_unit_circle(cls) -> Path: ...
+    @classmethod
     def unit_circle(cls) -> Path: ...
     @classmethod
     def circle(
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,30 @@ 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', [mcollections.EllipseCollection, mcollections.RectangleCollection]
+    )
+def test_WidthHeightAngleCollection_setter_getter(Class):
     # Test widths, heights and angle setter
     rng = np.random.default_rng(0)
 
@@ -421,7 +444,7 @@ def test_EllipseCollection_setter_getter():
 
     fig, ax = plt.subplots()
 
-    ec = mcollections.EllipseCollection(
+    ec = Class(
         widths=widths,
         heights=heights,
         angles=angles,
@@ -430,8 +453,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())
+    assert_array_almost_equal(ec._heights, np.array(heights).ravel())
     assert_array_almost_equal(ec._angles, np.deg2rad(angles).ravel())
 
     assert_array_almost_equal(ec.get_widths(), widths)
