8000 Add `set_U`, `set_V` and `set_C` method to `matplotlib.quiver.Quiver` · matplotlib/matplotlib@cfe71c1 · GitHub
[go: up one dir, main page]
Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
8000
Appearance settings

Commit cfe71c1

Browse files
ericpreericpre
committed
Add set_U, set_V and set_C method to matplotlib.quiver.Quiver 
1 parent 8c58e42 commit cfe71c1

File tree

3 files changed

+109
-3
lines changed

3 files changed

+109
-3
lines changed

‎lib/matplotlib/quiver.py

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -446,8 +446,8 @@ class Quiver(mcollections.PolyCollection):
446446
"""
447447
Specialized PolyCollection for arrows.
448448
449-
The only API method is set_UVC(), which can be used
450-
to change the size, orientation, and color of the
449+
The API methods are set_UVC(), set_U(), set_V() and set_C(), which
450+
can be used to change the size, orientation, and color of the
451451
arrows; their locations are fixed when the class is
452452
instantiated. Possibly this method will be useful
453453
in animations.
@@ -542,7 +542,59 @@ def draw(self, renderer):
542542
super().draw(renderer)
543543
self.stale = False
544544

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

‎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)
@@ -347,6 +348,56 @@ def test_collection_log_datalim(fig_test, fig_ref):
347348
ax_ref.plot(x, y, marker="o", ls="")
348349

349350

351+
def test_quiver_offsets():
352+
fig, ax = plt.subplots()
353+
x = np.arange(-10, 10, 1)
354+
y = np.arange(-10, 10, 1)
355+
U, V = np.meshgrid(x, y)
356+
X = U.ravel()
357+
Y = V.ravel()
358+
qc = mquiver.Quiver(ax, X, Y, U, V)
359+
ax.add_collection(qc)
360+
ax.autoscale_view()
361+
362+
expected_offsets = np.column_stack([X, Y])
363+
np.testing.assert_allclose(expected_offsets, qc.get_offsets())
364+
365+
new_offsets = np.column_stack([(X + 10).ravel(), Y.ravel()])
366+
qc.set_offsets(new_offsets)
367+
368+
np.testing.assert_allclose(qc.get_offsets(), new_offsets)
369+
370+
371+
def test_quiver_UVC():
372+
fig, ax = plt.subplots()
373+
X = np.arange(-10, 10, 1)
374+
Y = np.arange(-10, 10, 1)
375+
U, V = np.meshgrid(X, Y)
376+
M = np.hypot(U, V)
377+
qc = mquiver.Quiver(
378+
ax, X, Y, U, V, M
379+
)
380+
ax.add_collection(qc)
381+
ax.autoscale_view()
382+
383+
np.testing.assert_allclose(qc.U, U.ravel())
384+
np.testing.assert_allclose(qc.V, V.ravel())
385+
np.testing.assert_allclose(qc.get_array(), M.ravel())
386+
387+
qc.set_UVC(U/2, V/3)
388+
np.testing.assert_allclose(qc.U, U.ravel() / 2)
389+
np.testing.assert_allclose(qc.V, V.ravel() / 3)
390+
391+
qc.set_U(U/4)
392+
np.testing.assert_allclose(qc.U, U.ravel() / 4)
393+
394+
qc.set_V(V/6)
395+
np.testing.assert_allclose(qc.V, V.ravel() / 6)
396+
397+
qc.set_C(M/10)
398+
np.testing.assert_allclose(qc.get_array(), M.ravel() / 10)
399+
400+
350401
def test_quiver_limits():
351402
ax = plt.axes()
352403
x, y = np.arange(8), np.arange(10)

0 commit comments

Comments
 (0)
0