Skip to content

Add set_XY and set_data to Quiver #22407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

Add set_XY and set_data to Quiver #22407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
56498dd
add set_XY method to quiver and barb
ianhi Feb 5, 2022
f563d82
create set_data method for quiver
ianhi Feb 5, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 105 additions & 34 deletions lib/matplotlib/quiver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"""

import math
from numbers import Number
import weakref

import numpy as np
Expand Down Expand Up @@ -426,19 +427,28 @@ def _parse_args(*args, caller_name='function'):

nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape

if X is not None:
X = X.ravel()
Y = Y.ravel()
if len(X) == nc and len(Y) == nr:
X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
elif len(X) != len(Y):
raise ValueError('X and Y must be the same size, but '
f'X.size is {X.size} and Y.size is {Y.size}.')
else:
if X is None:
indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
X, Y = [np.ravel(a) for a in indexgrid]
# Size validation for U, V, C is left to the set_UVC method.
return X, Y, U, V, C
return X, Y, U, V, C, nr, nc


def _process_XY(X, Y, nc, nr):
X = X.ravel()
Y = Y.ravel()
if len(X) == nc and len(Y) == nr:
X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
elif len(X) != len(Y):
raise ValueError(
'X and Y must be the same size, but '
f'X.size is {X.size} and Y.size is {Y.size}.'
)
return X, Y


def _extract_nr_nc(U):
return (1, U.shape[0]) if U.ndim == 1 else U.shape


def _check_consistent_shapes(*arrays):
Expand All @@ -451,11 +461,10 @@ 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
arrows; their locations are fixed when the class is
instantiated. Possibly this method will be useful
in animations.
The only API methods are ``set_UVC()``, ``set_XY``,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove "only"? (Does three count as only?)

Single back-ticks will link, right? I guess one should be consistent with either using () for all or none of the methods.

and ``set_data``, which can be used
to change the size, orientation, color and locations
of the arrows.

Much of the work in this class is done in the draw()
method so that as much information as possible is available
Expand All @@ -479,11 +488,9 @@ def __init__(self, ax, *args,
%s
"""
self._axes = ax # The attr actually set by the Artist.axes property.
X, Y, U, V, C = _parse_args(*args, caller_name='quiver()')
self.X = X
self.Y = Y
self.XY = np.column_stack((X, Y))
self.N = len(X)
X, Y, U, V, C, self._nr, self._nc = _parse_args(
*args, caller_name='quiver()'
)
self.scale = scale
self.headwidth = headwidth
self.headlength = float(headlength)
Expand All @@ -503,10 +510,14 @@ def __init__(self, ax, *args,
self.transform = kwargs.pop('transform', ax.transData)
kwargs.setdefault('facecolors', color)
kwargs.setdefault('linewidths', (0,))
super().__init__([], offsets=self.XY, offset_transform=self.transform,
super().__init__([], offset_transform=self.transform,
closed=False, **kwargs)
self.polykw = kwargs
self.set_UVC(U, V, C)

self.X = self.Y = self.U = self.V = self.C = None
self.set_data(X, Y, U, V, C)
# self.U = self.V = self.C = None
# self.set_UVC(U, V, C)
self._initialized = False

weak_self = weakref.ref(self) # Prevent closure over the real self.
Expand Down Expand Up @@ -567,17 +578,58 @@ def draw(self, renderer):
self.stale = False

def set_UVC(self, U, V, C=None):
self.set_data(U=U, V=V, C=C)

def set_XY(self, X, Y):
"""
Update the locations of the arrows.

Parameters
----------
X, Y : arraylike of float
The arrow locations, any shape is valid so long
as X and Y have the same size.
"""
self.set_data(X=X, Y=Y)

def set_data(self, X=None, Y=None, U=None, V=None, C=None):
"""
Update the locations and/or rotation and color of the arrows.

Parameters
----------
X, Y : arraylike of float
The arrow locations, any shape is valid so long
as X and Y have the same size.
U, V : ???
C : ???
Comment on lines +604 to +605
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note to self - this is part of the "cleaning up"

"""
X = self.X if X is None else X
Y = self.Y if Y is None else Y
if U is None or isinstance(U, Number):
nr, nc = (self._nr, self._nc)
else:
nr, nc = _extract_nr_nc(U)
X, Y = _process_XY(X, Y, nc, nr)
N = len(X)

# 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()
V = ma.masked_invalid(V, copy=True).ravel()
if C is not None:
C = ma.masked_invalid(C, copy=True).ravel()
U = ma.masked_invalid(self.U if U is None else U, copy=True).ravel()
V = ma.masked_invalid(self.V if V is None else V, copy=True).ravel()
if C is not None or self.C is not None:
C = ma.masked_invalid(
self.C if C is None else C, copy=True
).ravel()
for name, var in zip(('U', 'V', 'C'), (U, V, C)):
if not (var is None or var.size == self.N or var.size == 1):
raise ValueError(f'Argument {name} has a size {var.size}'
f' which does not match {self.N},'
' the number of arrow positions')
if not (var is None or var.size == N or var.size == 1):
raise ValueError(
f'Argument {name} has a size {var.size}'
f' which does not match {N},'
' the number of arrow positions'
)

# now shapes are validated and we can start assigning things

mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
if C is not None:
Expand All @@ -591,7 +643,12 @@ def set_UVC(self, U, V, C=None):
self.Umask = mask
if C is not None:
self.set_array(C)
self.X = X
self.Y = Y
self.XY = np.column_stack([X, Y])
self.N = N
self._new_UV = True
self.set_offsets(self.XY)
self.stale = True

def _dots_per_unit(self, units):
Expand Down Expand Up @@ -963,10 +1020,11 @@ def __init__(self, ax, *args,
kwargs['linewidth'] = 1

# Parse out the data arrays from the various configurations supported
x, y, u, v, c = _parse_args(*args, caller_name='barbs()')
self.x = x
self.y = y
xy = np.column_stack((x, y))
x, y, u, v, c, self._nr, self._nc = _parse_args(
*args, caller_name='barbs()'
)
self.x, self.y = _process_XY(x, y, self._nr, self._nc)
xy = np.column_stack([self.x, self.y])

# Make a collection
barb_size = self._length ** 2 / 4 # Empirically determined
Expand Down Expand Up @@ -1198,6 +1256,19 @@ def set_UVC(self, U, V, C=None):
self._offsets = xy
self.stale = True

def set_XY(self, X, Y):
"""
Update the locations of the arrows.

Parameters
----------
X, Y : arraylike of float
The arrow locations, any shape is valid so long
as X and Y have the same size.
"""
self.X, self.Y = _process_XY(X, Y, self._nr, self._nc)
self.set_offsets(np.column_stack([X, Y]))

def set_offsets(self, xy):
"""
Set the offsets for the barb polygons. This saves the offsets passed
Expand Down
3 changes: 2 additions & 1 deletion lib/matplotlib/tests/test_quiver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,10 +167,11 @@ def test_quiver_key_xy():
for ax, angle_str in zip(axs, ('uv', 'xy')):
ax.set_xlim(-1, 8)
ax.set_ylim(-0.2, 0.2)
q = ax.quiver(X, Y, U, V, pivot='middle',
q = ax.quiver(X+1, Y+2, U, V, pivot='middle',
units='xy', width=0.05,
scale=2, scale_units='xy',
angles=angle_str)
q.set_XY(X, Y)
for x, angle in zip((0.2, 0.5, 0.8), (0, 45, 90)):
ax.quiverkey(q, X=x, Y=0.8, U=1, angle=angle, label='', color='b')

Expand Down