diff --git a/doc/users/next_whats_new/add_Quiver_setters.rst b/doc/users/next_whats_new/add_Quiver_setters.rst new file mode 100644 index 000000000000..5265d3a332fc --- /dev/null +++ b/doc/users/next_whats_new/add_Quiver_setters.rst @@ -0,0 +1,36 @@ +Add ``U``, ``V`` and ``C`` setter to ``Quiver`` +----------------------------------------------- + +The ``U``, ``V`` and ``C`` values of the `~matplotlib.quiver.Quiver` +can now be changed after the collection has been created. + +.. plot:: + :include-source: true + + import matplotlib.pyplot as plt + from matplotlib.quiver import Quiver + import numpy as np + + fig, ax = plt.subplots() + X = np.arange(-10, 10, 1) + Y = np.arange(-10, 10, 1) + U, V = np.meshgrid(X, Y) + C = np.hypot(U, V) + # When X and Y are 1D and U, V are 2D, X, Y are expanded to 2D + # using X, Y = np.meshgrid(X, Y) + qc = ax.quiver(X, Y, U, V, C) + + qc.set_U(U/5) + + # The number of arrows can also be changed. + + # Get new X, Y, U, V, C + X = np.arange(-10, 10, 2) + Y = np.arange(-10, 10, 2) + U, V = np.meshgrid(X, Y) + C = np.hypot(U, V) + # Use 2D X, Y coordinate (X, Y will not be expanded to 2D) + X, Y = np.meshgrid(X, Y) + + # Set new values + qc.set_XYUVC(X, Y, U, V, C) diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py index 8fa1962d6321..20a61fe532dc 100644 --- a/lib/matplotlib/quiver.py +++ b/lib/matplotlib/quiver.py @@ -15,6 +15,7 @@ """ import math +from numbers import Number import numpy as np from numpy import ma @@ -324,7 +325,7 @@ def _init(self): self._set_transform() with cbook._setattr_cm(self.Q, pivot=self.pivot[self.labelpos], # Hack: save and restore the Umask - Umask=ma.nomask): + _Umask=ma.nomask): u = self.U * np.cos(np.radians(self.angle)) v = self.U * np.sin(np.radians(self.angle)) self.verts = self.Q._make_verts([[0., 0.]], @@ -417,21 +418,30 @@ def _parse_args(*args, caller_name='function'): else: raise _api.nargs_error(caller_name, takes="from 2 to 5", given=nargs) - nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape + nr, nc = _extract_nr_nc(U) - 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 + # Size validation for U, V, C is left to the set_XYUVC method. + 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): @@ -444,11 +454,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 API methods are set_XYUVC(), set_X(), set_Y(), set_U() and set_V(), + which can be used to change the size, orientation, and color of the + arrows; their locations are fixed when the class is instantiated. + Possibly these methods will be useful in animations. Much of the work in this class is done in the draw() method so that as much information as possible is available @@ -459,6 +468,7 @@ class Quiver(mcollections.PolyCollection): """ _PIVOT_VALS = ('tail', 'middle', 'tip') + Umask = _api.deprecate_privatize_attribute("3.9") @_docstring.Substitution(_quiver_doc) def __init__(self, ax, *args, @@ -472,11 +482,7 @@ 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) + self.scale = scale self.headwidth = headwidth self.headlength = float(headlength) @@ -496,10 +502,16 @@ 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, - closed=False, **kwargs) + super().__init__( + [], offset_transform=self.transform, closed=False, **kwargs + ) self.polykw = kwargs - self.set_UVC(U, V, C) + + self._U = self._V = self._C = None + X, Y, U, V, C, self._nr, self._nc = _parse_args( + *args, caller_name='quiver()' + ) + self.set_XYUVC(X=X, Y=Y, U=U, V=V, C=C, check_shape=True) self._dpi_at_last_init = None def _init(self): @@ -513,21 +525,79 @@ def _init(self): trans = self._set_transform() self.span = trans.inverted().transform_bbox(self.axes.bbox).width if self.width is None: - sn = np.clip(math.sqrt(self.N), 8, 25) + sn = np.clip(math.sqrt(len(self.get_offsets())), 8, 25) self.width = 0.06 * self.span / sn # _make_verts sets self.scale if not already specified if (self._dpi_at_last_init != self.axes.figure.dpi and self.scale is None): - self._make_verts(self.XY, self.U, self.V, self.angles) + self._make_verts(self.get_offsets(), self._U, self._V, self.angles) self._dpi_at_last_init = self.axes.figure.dpi + @property + def N(self): + _api.warn_deprecated("3.9", alternative="get_X().size") + return len(self.get_X()) + + @property + def X(self): + _api.warn_deprecated("3.9", alternative="get_X") + return self.get_X() + + @X.setter + def X(self, value): + _api.warn_deprecated("3.9", alternative="set_X") + return self.set_X(value) + + @property + def Y(self): + _api.warn_deprecated("3.9", alternative="get_Y") + return self.get_Y() + + @Y.setter + def Y(self, value): + _api.warn_deprecated("3.9", alternative="set_Y") + return self.set_Y(value) + + @property + def U(self): + _api.warn_deprecated("3.9", alternative="get_U") + return self.get_U() + + @U.setter + def U(self, value): + _api.warn_deprecated("3.9", alternative="set_U") + return self.set_U(value) + + @property + def V(self): + _api.warn_deprecated("3.9", alternative="get_V") + return self.get_V() + + @V.setter + def V(self, value): + _api.warn_deprecated("3.9", alternative="set_V") + return self.set_V(value) + + @property + def XY(self): + _api.warn_deprecated("3.9", alternative="get_offsets") + return self.get_offsets() + + @XY.setter + def XY(self, value): + _api.warn_deprecated("3.9", alternative="set_offsets") + self.set_offsets(offsets=value) + + def set_offsets(self, offsets): + self.set_XYUVC(X=offsets[:, 0], Y=offsets[:, 1]) + def get_datalim(self, transData): trans = self.get_transform() offset_trf = self.get_offset_transform() full_transform = (trans - transData) + (offset_trf - transData) - XY = full_transform.transform(self.XY) + XY = full_transform.transform(self.get_offsets()) bbox = transforms.Bbox.null() bbox.update_from_data_xy(XY, ignore=True) return bbox @@ -535,24 +605,138 @@ def get_datalim(self, transData): @martist.allow_rasterization def draw(self, renderer): self._init() - verts = self._make_verts(self.XY, self.U, self.V, self.angles) + verts = self._make_verts(self.get_offsets(), self._U, self._V, self.angles) self.set_verts(verts, closed=False) super().draw(renderer) self.stale = False - def set_UVC(self, U, V, C=None): + def set_X(self, X): + """ + Set positions in the horizontal direction. + + Parameters + ---------- + X : array-like + The size must the same as the existing Y. + """ + self.set_XYUVC(X=X) + + def get_X(self): + """Returns the positions in the horizontal direction.""" + return self.get_offsets()[..., 0] + + def set_Y(self, Y): + """ + Set positions in the vertical direction. + + Parameters + ---------- + Y : array-like + The size must the same as the existing X. + """ + self.set_XYUVC(Y=Y) + + def get_Y(self): + """Returns the positions in the vertical direction.""" + return self.get_offsets()[..., 1] + + def set_U(self, U): + """ + Set horizontal direction components. + + Parameters + ---------- + U : array-like + The size must the same as the existing X, Y, V or be one. + """ + self.set_XYUVC(U=U) + + def get_U(self): + """Returns the horizontal direction components.""" + return self._U + + def set_V(self, V): + """ + Set vertical direction components. + + Parameters + ---------- + V : array-like + The size must the same as the existing X, Y, U or be one. + """ + self.set_XYUVC(V=V) + + def get_V(self): + """Returns the vertical direction components.""" + return self._V + + def set_C(self, C): + """ + Set the arrow colors. + + Parameters + ---------- + C : array-like + The size must the same as the existing X, Y, U, V or be one. + """ + self.set_XYUVC(C=C) + + def get_C(self): + """Returns the arrow colors.""" + return self.get_array() + + def set_XYUVC(self, X=None, Y=None, U=None, V=None, C=None, check_shape=False): + """ + Set the positions (X, Y) and components (U, V) of the arrow vectors + and arrow colors (C) values of the arrows. + The size of the array must match with existing values. To change + the size, all arguments must be changed at once and their size + compatible. + + Parameters + ---------- + X, Y : array-like of float or None, optional + The arrow locations in the horizontal and vertical directions. + Any shape is valid so long as X and Y have the same size. + U, V : array-like or None, optional + The horizontal and vertical direction components of the arrows. + If None it is unchanged. + The size must the same as the existing U, V or be one. + C : array-like or None, optional + The arrow colors. The default is None. + The size must the same as the existing U, V or be one. + check_shape : bool + Whether to check if the shape of the parameters are + consistent. Default is False. + """ + + X = self.get_X() if X is None else X + Y = self.get_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() - 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') - + 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() + if check_shape: + for name, var in zip(('U', 'V', 'C'), (U, V, C)): + 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: mask = ma.mask_or(mask, C.mask, copy=False, shrink=True) @@ -560,11 +744,14 @@ def set_UVC(self, U, V, C=None): C = C.filled() else: C = ma.array(C, mask=mask, copy=False) - self.U = U.filled(1) - self.V = V.filled(1) - self.Umask = mask + self._U = U.filled(1) + self._V = V.filled(1) + self._Umask = mask if C is not None: self.set_array(C) + self._N = N + self._new_UV = True + super().set_offsets(np.column_stack([X, Y])) self.stale = True def _dots_per_unit(self, units): @@ -625,9 +812,9 @@ def _make_verts(self, XY, U, V, angles): a = np.abs(uv) if self.scale is None: - sn = max(10, math.sqrt(self.N)) - if self.Umask is not ma.nomask: - amean = a[~self.Umask].mean() + sn = max(10, math.sqrt(len(self.get_offsets()))) + if self._Umask is not ma.nomask: + amean = a[~self._Umask].mean() else: amean = a.mean() # crude auto-scaling @@ -657,9 +844,9 @@ def _make_verts(self, XY, U, V, angles): theta = theta.reshape((-1, 1)) # for broadcasting xy = (X + Y * 1j) * np.exp(1j * theta) * self.width XY = np.stack((xy.real, xy.imag), axis=2) - if self.Umask is not ma.nomask: + if self._Umask is not ma.nomask: XY = ma.array(XY) - XY[self.Umask] = ma.masked + XY[self._Umask] = ma.masked # This might be handled more efficiently with nans, given # that nans will end up in the paths anyway. @@ -928,10 +1115,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 diff --git a/lib/matplotlib/quiver.pyi b/lib/matplotlib/quiver.pyi index 2a043a92b4b5..33335446047f 100644 --- a/lib/matplotlib/quiver.pyi +++ b/lib/matplotlib/quiver.pyi @@ -54,11 +54,11 @@ class QuiverKey(martist.Artist): class Quiver(mcollections.PolyCollection): X: ArrayLike Y: ArrayLike - XY: ArrayLike U: ArrayLike V: ArrayLike + C: ArrayLike + XY: ArrayLike Umask: ArrayLike - N: int scale: float | None headwidth: float headlength: float @@ -121,9 +121,28 @@ class Quiver(mcollections.PolyCollection): pivot: Literal["tail", "mid", "middle", "tip"] = ..., **kwargs ) -> None: ... + @property + def N(self) -> int: ... def get_datalim(self, transData: Transform) -> Bbox: ... - def set_UVC( - self, U: ArrayLike, V: ArrayLike, C: ArrayLike | None = ... + def set_offsets(self, offsets: ArrayLike) -> None: ... + def set_X(self, X: ArrayLike) -> None: ... + def get_X(self) -> ArrayLike: ... + def set_Y(self, Y: ArrayLike) -> None: ... + def get_Y(self) -> ArrayLike: ... + def set_U(self, U: ArrayLike) -> None: ... + def get_U(self) -> ArrayLike: ... + def set_V(self, V: ArrayLike) -> None: ... + def get_V(self) -> ArrayLike: ... + def set_C(self, C: ArrayLike) -> None: ... + def get_C(self) -> ArrayLike: ... + def set_XYUVC( + self, + X: ArrayLike | None = ..., + Y: ArrayLike | None = ..., + U: ArrayLike | None = ..., + V: ArrayLike | None = ..., + C: ArrayLike | None = ..., + check_shape: bool = ..., ) -> None: ... class Barbs(mcollections.PolyCollection): diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py index 5baaeaa5d388..1b63413927de 100644 --- a/lib/matplotlib/tests/test_collections.py +++ b/lib/matplotlib/tests/test_collections.py @@ -13,6 +13,7 @@ import matplotlib.collections as mcollections import matplotlib.colors as mcolors import matplotlib.path as mpath +import matplotlib.quiver as mquiver import matplotlib.transforms as mtransforms from matplotlib.collections import (Collection, LineCollection, EventCollection, PolyCollection) @@ -357,6 +358,131 @@ def test_collection_log_datalim(fig_test, fig_ref): ax_ref.plot(x, y, marker="o", ls="") +def test_quiver_offsets(): + fig, ax = plt.subplots() + x = np.arange(-10, 10, 1) + y = np.arange(-10, 10, 1) + U, V = np.meshgrid(x, y) + X = U.ravel() + Y = V.ravel() + qc = mquiver.Quiver(ax, X, Y, U, V) + ax.add_collection(qc) + ax.autoscale_view() + + expected_offsets = np.column_stack([X, Y]) + np.testing.assert_allclose(expected_offsets, qc.get_offsets()) + + new_offsets = np.column_stack([(X + 10).ravel(), Y.ravel()]) + qc.set_offsets(new_offsets) + + np.testing.assert_allclose(qc.get_offsets(), new_offsets) + np.testing.assert_allclose(qc.get_X(), new_offsets[..., 0]) + np.testing.assert_allclose(qc.get_Y(), new_offsets[..., 1]) + + new_X = qc.get_X() + 5 + qc.set_X(new_X) + np.testing.assert_allclose(qc.get_X(), new_X) + + new_Y = qc.get_Y() + 5 + qc.set_Y(new_Y) + np.testing.assert_allclose(qc.get_Y(), new_Y) + + # new length + L = 2 + qc.set_XYUVC(X=new_X[:L], Y=new_Y[:L]) + np.testing.assert_allclose(qc.get_X(), new_X[:L]) + np.testing.assert_allclose(qc.get_Y(), new_Y[:L]) + + qc.set_XYUVC(X=X[:L], Y=Y[:L], U=qc.get_U()[:L], V=qc.get_V()[:L]) + np.testing.assert_allclose(qc.get_X(), X[:L]) + np.testing.assert_allclose(qc.get_Y(), Y[:L]) + np.testing.assert_allclose(qc.get_U(), U.ravel()[:L]) + np.testing.assert_allclose(qc.get_V(), V.ravel()[:L]) + + +def test_quiver_change_XYUVC(): + fig, ax = plt.subplots() + X = np.arange(-10, 10, 1) + Y = np.arange(-10, 10, 1) + U, V = np.meshgrid(X, Y) + C = np.hypot(U, V) + qc = mquiver.Quiver(ax, X, Y, U, V, C) + ax.add_collection(qc) + ax.autoscale_view() + + np.testing.assert_allclose(qc.get_U(), U.ravel()) + np.testing.assert_allclose(qc.get_V(), V.ravel()) + np.testing.assert_allclose(qc.get_C(), C.ravel()) + + qc.set(U=U/2, V=V/3) + np.testing.assert_allclose(qc.get_U(), U.ravel() / 2) + np.testing.assert_allclose(qc.get_V(), V.ravel() / 3) + + qc.set_U(U/4) + np.testing.assert_allclose(qc.get_U(), U.ravel() / 4) + + qc.set_V(V/6) + np.testing.assert_allclose(qc.get_V(), V.ravel() / 6) + + qc.set_C(C/3) + np.testing.assert_allclose(qc.get_C(), C.ravel() / 3) + + # check consistency not enable + qc.set_XYUVC(X=X[:2], Y=Y[:2]) + with pytest.raises(ValueError): + # setting only one of the two X, Y fails because X and Y needs + # to be stacked when passed to `offsets` + qc.set(Y=Y[:3]) + + qc.set() + np.testing.assert_allclose(qc.get_U(), U.ravel() / 4) + np.testing.assert_allclose(qc.get_V(), V.ravel() / 6) + + +def test_quiver_deprecated_attribute(): + fig, ax = plt.subplots() + X = np.arange(-10, 10, 1) + Y = np.arange(-10, 10, 1) + U, V = np.meshgrid(X, Y) + C = np.hypot(U, V) + qc = mquiver.Quiver(ax, X, Y, U, V, C) + ax.add_collection(qc) + ax.autoscale_view() + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + x = qc.X + with pytest.warns(mpl.MatplotlibDeprecationWarning): + qc.X = x * 2 + np.testing.assert_allclose(qc.get_X(), x * 2) + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + y = qc.Y + with pytest.warns(mpl.MatplotlibDeprecationWarning): + qc.Y = y * 2 + np.testing.assert_allclose(qc.get_Y(), y * 2) + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + np.testing.assert_allclose(qc.N, len(qc.get_offsets())) + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + u = qc.U + with pytest.warns(mpl.MatplotlibDeprecationWarning): + qc.U = u * 2 + np.testing.assert_allclose(qc.get_U(), u * 2) + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + v = qc.V + with pytest.warns(mpl.MatplotlibDeprecationWarning): + qc.V = v * 2 + np.testing.assert_allclose(qc.get_V(), v * 2) + + with pytest.warns(mpl.MatplotlibDeprecationWarning): + xy = qc.XY + with pytest.warns(mpl.MatplotlibDeprecationWarning): + qc.XY = xy * 2 + np.testing.assert_allclose(qc.get_offsets(), xy * 2) + + def test_quiver_limits(): ax = plt.axes() x, y = np.arange(8), np.arange(10) diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py index 7c5a9d343530..5b7262e3e920 100644 --- a/lib/matplotlib/tests/test_quiver.py +++ b/lib/matplotlib/tests/test_quiver.py @@ -24,7 +24,7 @@ def test_quiver_memory_leak(): fig, ax = plt.subplots() Q = draw_quiver(ax) - ttX = Q.X + ttX = Q.get_X() Q.remove() del Q @@ -133,7 +133,7 @@ def test_quiver_copy(): uv = dict(u=np.array([1.1]), v=np.array([2.0])) q0 = ax.quiver([1], [1], uv['u'], uv['v']) uv['v'][0] = 0 - assert q0.V[0] == 2.0 + assert q0.get_V()[0] == 2.0 @image_comparison(['quiver_key_pivot.png'], remove_text=True) @@ -332,4 +332,4 @@ def test_quiver_setuvc_numbers(): U = V = np.ones_like(X) q = ax.quiver(X, Y, U, V) - q.set_UVC(0, 1) + q.set_XYUVC(U=0, V=1)