Add set_U, set_V and set_C method to matplotlib.quiver.Quiver

ericpre · ericpre · commit cfe71c1e0801 · 2023-09-07T08:19:35.000+01:00
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -446,8 +446,8 @@ class Quiver(mcollections.PolyCollection):
     """
     Specialized PolyCollection for arrows.
 
-    The only API method is set_UVC(), which can be used
-    to change the size, orientation, and color of the
+    The API methods are set_UVC(), set_U(), set_V() and set_C(), which
+    can be used to change the size, orientation, and color of the
     arrows; their locations are fixed when the class is
     instantiated.  Possibly this method will be useful
     in animations.
@@ -542,7 +542,59 @@ def draw(self, renderer):
         super().draw(renderer)
         self.stale = False
 
+    def set_U(self, U):
+        """
+        Set x direction components of the arrow vectors.
+
+        Parameters
+        ----------
+        U : array-like or None
+            The size must the same as the existing U, V or be one.
+        """
+        self.set_UVC(U, None, None)
+
+    def set_V(self, V):
+        """
+        Set y direction components of the arrow vectors.
+
+        Parameters
+        ----------
+        V : array-like or None
+            The size must the same as the existing U, V or be one.
+        """
+        self.set_UVC(None, V, None)
+
+    def set_C(self, C):
+        """
+        Set the arrow colors.
+
+        Parameters
+        ----------
+        C : array-like or None
+            The size must the same as the existing U, V or be one.
+        """
+        self.set_UVC(None, None, C)
+
     def set_UVC(self, U, V, C=None):
+        """
+        Set the U, V (x and y direction components of the arrow vectors) and
+        C (arrow colors) values of the arrows.
+
+        Parameters
+        ----------
+        U : array-like or None
+            The x direction components of the arrows. If None it is unchanged.
+            The size must the same as the existing U, V or be one.
+        V : array-like or None
+            The y direction components of the arrows. If None it is unchanged.
+            The size must the same as the existing U, V or be one.
+        C : array-like or None, optional
+            The arrow colors. The default is None.
+        """
+        if U is None:
+            U = self.U
+        if V is None:
+            V = self.V
         # We need to ensure we have a copy, not a reference
         # to an array that might change before draw().
         U = ma.masked_invalid(U, copy=True).ravel()
diff --git a/lib/matplotlib/quiver.pyi b/lib/matplotlib/quiver.pyi
@@ -122,8 +122,11 @@ class Quiver(mcollections.PolyCollection):
         **kwargs
     ) -> None: ...
     def get_datalim(self, transData: Transform) -> Bbox: ...
+    def set_U(self, U: ArrayLike) -> None: ...
+    def set_V(self, V: ArrayLike) -> None: ...
+    def set_C(self, C: ArrayLike) -> None: ...
     def set_UVC(
-        self, U: ArrayLike, V: ArrayLike, C: ArrayLike | None = ...
+        self, U: ArrayLike | None, V: ArrayLike | None, C: ArrayLike | None = ...
     ) -> None: ...
     @property
     def quiver_doc(self) -> str: ...
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -13,6 +13,7 @@
 import matplotlib.collections as mcollections
 import matplotlib.colors as mcolors
 import matplotlib.path as mpath
+import matplotlib.quiver as mquiver
 import matplotlib.transforms as mtransforms
 from matplotlib.collections import (Collection, LineCollection,
                                     EventCollection, PolyCollection)
@@ -347,6 +348,56 @@ def test_collection_log_datalim(fig_test, fig_ref):
     ax_ref.plot(x, y, marker="o", ls="")
 
 
+def test_quiver_offsets():
+    fig, ax = plt.subplots()
+    x = np.arange(-10, 10, 1)
+    y = np.arange(-10, 10, 1)
+    U, V = np.meshgrid(x, y)
+    X = U.ravel()
+    Y = V.ravel()
+    qc = mquiver.Quiver(ax, X, Y, U, V)
+    ax.add_collection(qc)
+    ax.autoscale_view()
+
+    expected_offsets = np.column_stack([X, Y])
+    np.testing.assert_allclose(expected_offsets, qc.get_offsets())
+
+    new_offsets = np.column_stack([(X + 10).ravel(), Y.ravel()])
+    qc.set_offsets(new_offsets)
+
+    np.testing.assert_allclose(qc.get_offsets(), new_offsets)
+
+
+def test_quiver_UVC():
+    fig, ax = plt.subplots()
+    X = np.arange(-10, 10, 1)
+    Y = np.arange(-10, 10, 1)
+    U, V = np.meshgrid(X, Y)
+    M = np.hypot(U, V)
+    qc = mquiver.Quiver(
+        ax, X, Y, U, V, M
+        )
+    ax.add_collection(qc)
+    ax.autoscale_view()
+
+    np.testing.assert_allclose(qc.U, U.ravel())
+    np.testing.assert_allclose(qc.V, V.ravel())
+    np.testing.assert_allclose(qc.get_array(), M.ravel())
+
+    qc.set_UVC(U/2, V/3)
+    np.testing.assert_allclose(qc.U, U.ravel() / 2)
+    np.testing.assert_allclose(qc.V, V.ravel() / 3)
+
+    qc.set_U(U/4)
+    np.testing.assert_allclose(qc.U, U.ravel() / 4)
+
+    qc.set_V(V/6)
+    np.testing.assert_allclose(qc.V, V.ravel() / 6)
+
+    qc.set_C(M/10)
+    np.testing.assert_allclose(qc.get_array(), M.ravel() / 10)
+
+
 def test_quiver_limits():
     ax = plt.axes()
     x, y = np.arange(8), np.arange(10)
