Merge pull request #26375 from ericpre/add_EllipseCollection_setter

QuLogic · web-flow · commit 4b42dbbdc100 · 2024-04-01T18:11:50.000-04:00
Add `widths`, `heights` and `angles` setter to `EllipseCollection`
diff --git a/doc/users/next_whats_new/add_EllipseCollection_setters.rst b/doc/users/next_whats_new/add_EllipseCollection_setters.rst
@@ -0,0 +1,40 @@
+Add ``widths``, ``heights`` and ``angles`` setter to ``EllipseCollection``
+--------------------------------------------------------------------------
+
+The ``widths``, ``heights`` and ``angles`` values of the `~matplotlib.collections.EllipseCollection`
+can now be changed after the collection has been created.
+
+.. plot::
+    :include-source: true
+
+    import matplotlib.pyplot as plt
+    from matplotlib.collections import EllipseCollection
+    import numpy as np
+
+    rng = np.random.default_rng(0)
+
+    widths = (2, )
+    heights = (3, )
+    angles = (45, )
+    offsets = rng.random((10, 2)) * 10
+
+    fig, ax = plt.subplots()
+
+    ec = EllipseCollection(
+        widths=widths,
+        heights=heights,
+        angles=angles,
+        offsets=offsets,
+        units='x',
+        offset_transform=ax.transData,
+        )
+
+    ax.add_collection(ec)
+    ax.set_xlim(-2, 12)
+    ax.set_ylim(-2, 12)
+
+    new_widths = rng.random((10, 2)) * 2
+    new_heights = rng.random((10, 2)) * 3
+    new_angles = rng.random((10, 2)) * 180
+
+    ec.set(widths=new_widths, heights=new_heights, angles=new_angles)
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1745,9 +1745,9 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
             Forwarded to `Collection`.
         """
         super().__init__(**kwargs)
-        self._widths = 0.5 * np.asarray(widths).ravel()
-        self._heights = 0.5 * np.asarray(heights).ravel()
-        self._angles = np.deg2rad(angles).ravel()
+        self.set_widths(widths)
+        self.set_heights(heights)
+        self.set_angles(angles)
         self._units = units
         self.set_transform(transforms.IdentityTransform())
         self._transforms = np.empty((0, 3, 3))
@@ -1795,6 +1795,33 @@ def _set_transforms(self):
             m[:2, 2:] = 0
             self.set_transform(_affine(m))
 
+    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.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.stale = True
+
+    def set_angles(self, angles):
+        """Set the angles of the first axes, degrees CCW from the x-axis."""
+        self._angles = np.deg2rad(angles).ravel()
+        self.stale = True
+
+    def get_widths(self):
+        """Get the lengths of the first axes (e.g., major axis)."""
+        return self._widths * 2
+
+    def get_heights(self):
+        """Set the lengths of second axes (e.g., minor axes)."""
+        return self._heights * 2
+
+    def get_angles(self):
+        """Get the angles of the first axes, degrees CCW from the x-axis."""
+        return np.rad2deg(self._angles)
+
     @artist.allow_rasterization
     def draw(self, renderer):
         self._set_transforms()
diff --git a/lib/matplotlib/collections.pyi b/lib/matplotlib/collections.pyi
@@ -180,6 +180,12 @@ class EllipseCollection(Collection):
         ] = ...,
         **kwargs
     ) -> None: ...
+    def set_widths(self, widths: ArrayLike) -> None: ...
+    def set_heights(self, heights: ArrayLike) -> None: ...
+    def set_angles(self, angles: ArrayLike) -> None: ...
+    def get_widths(self) -> ArrayLike: ...
+    def get_heights(self) -> ArrayLike: ...
+    def get_angles(self) -> ArrayLike: ...
 
 class PatchCollection(Collection):
     def __init__(
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -408,6 +408,49 @@ def test_EllipseCollection():
     ax.autoscale_view()
 
 
+def test_EllipseCollection_setter_getter():
+    # Test widths, heights and angle setter
+    rng = np.random.default_rng(0)
+
+    widths = (2, )
+    heights = (3, )
+    angles = (45, )
+    offsets = rng.random((10, 2)) * 10
+
+    fig, ax = plt.subplots()
+
+    ec = mcollections.EllipseCollection(
+        widths=widths,
+        heights=heights,
+        angles=angles,
+        offsets=offsets,
+        units='x',
+        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._angles, np.deg2rad(angles).ravel())
+
+    assert_array_almost_equal(ec.get_widths(), widths)
+    assert_array_almost_equal(ec.get_heights(), heights)
+    assert_array_almost_equal(ec.get_angles(), angles)
+
+    ax.add_collection(ec)
+    ax.set_xlim(-2, 12)
+    ax.set_ylim(-2, 12)
+
+    new_widths = rng.random((10, 2)) * 2
+    new_heights = rng.random((10, 2)) * 3
+    new_angles = rng.random((10, 2)) * 180
+
+    ec.set(widths=new_widths, heights=new_heights, angles=new_angles)
+
+    assert_array_almost_equal(ec.get_widths(), new_widths.ravel())
+    assert_array_almost_equal(ec.get_heights(), new_heights.ravel())
+    assert_array_almost_equal(ec.get_angles(), new_angles.ravel())
+
+
 @image_comparison(['polycollection_close.png'], remove_text=True, style='mpl20')
 def test_polycollection_close():
     from mpl_toolkits.mplot3d import Axes3D  # type: ignore
