From 56498dd132259a65e9a373d38d0addda4a59ebb2 Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Sat, 5 Feb 2022 01:09:28 -0500 Subject: [PATCH 1/2] add set_XY method to quiver and barb --- lib/matplotlib/quiver.py | 85 ++++++++++++++++++++++++++++------------ 1 file changed, 61 insertions(+), 24 deletions(-) diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py index d57abf9a529c..d3adb890e4e1 100644 --- a/lib/matplotlib/quiver.py +++ b/lib/matplotlib/quiver.py @@ -426,19 +426,24 @@ 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 _check_consistent_shapes(*arrays): @@ -451,11 +456,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``, + 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 @@ -479,11 +483,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) @@ -495,6 +497,10 @@ def __init__(self, ax, *args, self.angles = angles self.width = width + self.X, self.Y = _process_XY(X, Y, self._nr, self._nc) + self.XY = np.column_stack([self.X, self.Y]) + self.N = len(self.X) + if pivot.lower() == 'mid': pivot = 'middle' self.pivot = pivot.lower() @@ -594,6 +600,23 @@ def set_UVC(self, U, V, C=None): self._new_UV = True 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.X = X + self.Y = X + self.XY = np.column_stack((self.X, self.Y)) + self._offsets = self.XY + self.stale = True + def _dots_per_unit(self, units): """ Return a scale factor for converting from units to pixels @@ -963,10 +986,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 @@ -1198,6 +1222,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 From f563d82e3e5e752c188b255e441273fdc8742a11 Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Sat, 5 Feb 2022 01:50:45 -0500 Subject: [PATCH 2/2] create set_data method for quiver --- lib/matplotlib/quiver.py | 96 +++++++++++++++++++---------- lib/matplotlib/tests/test_quiver.py | 3 +- 2 files changed, 67 insertions(+), 32 deletions(-) diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py index d3adb890e4e1..48ee9b89eebf 100644 --- a/lib/matplotlib/quiver.py +++ b/lib/matplotlib/quiver.py @@ -15,6 +15,7 @@ """ import math +from numbers import Number import weakref import numpy as np @@ -446,6 +447,10 @@ def _process_XY(X, Y, nc, nr): 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): all_shapes = {a.shape for a in arrays} if len(all_shapes) != 1: @@ -497,10 +502,6 @@ def __init__(self, ax, *args, self.angles = angles self.width = width - self.X, self.Y = _process_XY(X, Y, self._nr, self._nc) - self.XY = np.column_stack([self.X, self.Y]) - self.N = len(self.X) - if pivot.lower() == 'mid': pivot = 'middle' self.pivot = pivot.lower() @@ -509,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. @@ -573,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 : ??? + """ + 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: @@ -597,24 +643,12 @@ def set_UVC(self, U, V, C=None): self.Umask = mask if C is not None: self.set_array(C) - self._new_UV = True - 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.X = X - self.Y = X - self.XY = np.column_stack((self.X, self.Y)) - self._offsets = self.XY + 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): diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py index 4c794ace91bc..161dbf556f6d 100644 --- a/lib/matplotlib/tests/test_quiver.py +++ b/lib/matplotlib/tests/test_quiver.py @@ -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')