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

Skip to content

Commit b268585

Browse files
ericpreericpre
committed
Add set_U, set_V and set_C method to matplotlib.quiver.Quiver 
1 parent 47c96df commit b268585

3 files changed

Lines changed: 109 additions & 3 deletions

File tree

‎lib/matplotlib/quiver.py‎

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -444,8 +444,8 @@ class Quiver(mcollections.PolyCollection):
444444
"""
445445
Specialized PolyCollection for arrows.
446446
447-
The only API method is set_UVC(), which can be used
448-
to change the size, orientation, and color of the
447+
The API methods are set_UVC(), set_U(), set_V() and set_C(), which
448+
can be used to change the size, orientation, and color of the
449449
arrows; their locations are fixed when the class is
450450
instantiated. Possibly this method will be useful
451451
in animations.
@@ -540,7 +540,59 @@ def draw(self, renderer):
540540
super().draw(renderer)
541541
self.stale = False
542542

543+
def set_U(self, U):
544+
"""
545+
Set x direction components of the arrow vectors.
546+
547+
Parameters
548+
----------
549+
U : array-like or None
550+
The size must the same as the existing U, V or be one.
551+
"""
552+
self.set_UVC(U, None, None)
553+
554+
def set_V(self, V):
555+
"""
556+
Set y direction components of the arrow vectors.
557+
558+
Parameters
559+
----------
560+
V : array-like or None
561+
The size must the same as the existing U, V or be one.
562+
"""
563+
self.set_UVC(None, V, None)
564+
565+
def set_C(self, C):
566+
"""
567+
Set the arrow colors.
568+
569+
Parameters
570+
----------
571+
C : array-like or None
572+
The size must the same as the existing U, V or be one.
573+
"""
574+
self.set_UVC(None, None, C)
575+
543576
def set_UVC(self, U, V, C=None):
577+
"""
578+
Set the U, V (x and y direction components of the arrow vectors) and
579+
C (arrow colors) values of the arrows.
580+
581+
Parameters
582+
----------
583+
U : array-like or None
584+
The x direction components of the arrows. If None it is unchanged.
585+
The size must the same as the existing U, V or be one.
586+
V : array-like or None
587+
The y direction components of the arrows. If None it is unchanged.
588+
The size must the same as the existing U, V or be one.
589+
C : array-like or None, optional
590+
The arrow colors. The default is None.
591+
"""
592+
if U is None:
593+
U = self.U
594+
if V is None:
595+
V = self.V
544596
# We need to ensure we have a copy, not a reference
545597
# to an array that might change before draw().
546598
U = ma.masked_invalid(U, copy=True).ravel()

‎lib/matplotlib/quiver.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,8 +122,11 @@ class Quiver(mcollections.PolyCollection):
122122
**kwargs
123123
) -> None: ...
124124
def get_datalim(self, transData: Transform) -> Bbox: ...
125+
def set_U(self, U: ArrayLike) -> None: ...
126+
def set_V(self, V: ArrayLike) -> None: ...
127+
def set_C(self, C: ArrayLike) -> None: ...
125128
def set_UVC(
126-
self, U: ArrayLike, V: ArrayLike, C: ArrayLike | None = ...
129+
self, U: ArrayLike | None, V: ArrayLike | None, C: ArrayLike | None = ...
127130
) -> None: ...
128131

129132
class Barbs(mcollections.PolyCollection):

‎lib/matplotlib/tests/test_collections.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import matplotlib.collections as mcollections
1414
import matplotlib.colors as mcolors
1515
import matplotlib.path as mpath
16+
import matplotlib.quiver as mquiver
1617
import matplotlib.transforms as mtransforms
1718
from matplotlib.collections import (Collection, LineCollection,
1819
EventCollection, PolyCollection)
@@ -357,6 +358,56 @@ def test_collection_log_datalim(fig_test, fig_ref):
357358
ax_ref.plot(x, y, marker="o", ls="")
358359

359360

361+
def test_quiver_offsets():
362+
fig, ax = plt.subplots()
363+
x = np.arange(-10, 10, 1)
364+
y = np.arange(-10, 10, 1)
365+
U, V = np.meshgrid(x, y)
366+
X = U.ravel()
367+
Y = V.ravel()
368+
qc = mquiver.Quiver(ax, X, Y, U, V)
369+
ax.add_collection(qc)
370+
ax.autoscale_view()
371+
372+
expected_offsets = np.column_stack([X, Y])
373+
np.testing.assert_allclose(expected_offsets, qc.get_offsets())
374+
375+
new_offsets = np.column_stack([(X + 10).ravel(), Y.ravel()])
376+
qc.set_offsets(new_offsets)
377+
378+
np.testing.assert_allclose(qc.get_offsets(), new_offsets)
379+
380+
381+
def test_quiver_UVC():
382+
fig, ax = plt.subplots()
383+
X = np.arange(-10, 10, 1)
384+
Y = np.arange(-10, 10, 1)
385+
U, V = np.meshgrid(X, Y)
386+
M = np.hypot(U, V)
387+
qc = mquiver.Quiver(
388+
ax, X, Y, U, V, M
389+
)
390+
ax.add_collection(qc)
391+
ax.autoscale_view()
392+
393+
np.testing.assert_allclose(qc.U, U.ravel())
394+
np.testing.assert_allclose(qc.V, V.ravel())
395+
np.testing.assert_allclose(qc.get_array(), M.ravel())
396+
397+
qc.set_UVC(U/2, V/3)
398+
np.testing.assert_allclose(qc.U, U.ravel() / 2)
399+
np.testing.assert_allclose(qc.V, V.ravel() / 3)
400+
401+
qc.set_U(U/4)
402+
np.testing.assert_allclose(qc.U, U.ravel() / 4)
403+
404+
qc.set_V(V/6)
405+
np.testing.assert_allclose(qc.V, V.ravel() / 6)
406+
407+
qc.set_C(M/10)
408+
np.testing.assert_allclose(qc.get_array(), M.ravel() / 10)
409+
410+
360411
def test_quiver_limits():
361412
ax = plt.axes()
362413
x, y = np.arange(8), np.arange(10)

0 commit comments

Comments
 (0)