diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py index d57abf9a529c..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 @@ -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): @@ -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``, + 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 +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) @@ -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. @@ -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 : ??? + """ + 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: @@ -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): @@ -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 @@ -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 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') pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy