Skip to content

get the same 'hue' behaviour in .quiver compared to the standard .plot() #10397

Open
@ACHMartin

Description

@ACHMartin

Is your feature request related to a problem?

I have a 3 dimensions arrows, I would like to plot on top of each other.
I tried to use the usual hue attribute within xr.quiver, but it doesn't behave similarly to xr.plot(), where hue is used to unpack the extra dimension.

Below a minimal example of my issue:

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

x_size = 6; y_size = 4; w_size = 2
xv = np.arange(0,x_size)
yv = np.arange(0,y_size)
wv =  np.arange(0,w_size)
x, y, w = np.meshgrid(xv, yv, wv)
u = np.random.uniform(size=(y_size, x_size, w_size))
v = np.random.uniform(size=(y_size, x_size, w_size))

myds = xr.Dataset({
    'x': (['x'], xv),
    'y': (['y'], yv),
    'w': (['w'], wv),
    'u': (['y','x','w'], u),
    'v': (['y','x','w'], v),
})

If I unpack my data using col, it works fine:

myds.plot.quiver(x='x', y='y', u='u', v='v', col='w', 
                 scale=10, add_guide=False,)

Image

If I add the hue option, I got the colours as expected:

myds.plot.quiver(x='x', y='y', u='u', v='v', col='w', 
                 hue='w',
                 scale=10, add_guide=False,)

Image

but if I only use hue and not col, I get an error, because quiver doesn't unpack along the remaining dimension:

myds.plot.quiver(x='x', y='y', u='u', v='v', 
                 hue='w',
                 scale=10, add_guide=False,)

with the following error:

File /env/lib/python3.11/site-packages/matplotlib/quiver.py:420, in _parse_args(caller_name, *args)
    417 else:
    418     raise _api.nargs_error(caller_name, takes="from 2 to 5", given=nargs)
--> 420 nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape
    422 if X is not None:
    423     X = X.ravel()

ValueError: too many values to unpack (expected 2)

Describe the solution you'd like

What I would expect is to have a behaviour similar to what we usually have with xr.plot, with automatic unpacking along the dimension with the hue option.

Describe alternatives you've considered

We can do it by reshaping the data, but we loose xarray simplicity.

xc = x.reshape((x_size, y_size * w_size))
yc = y.reshape((x_size, y_size * w_size))
wc = w.reshape((x_size, y_size * w_size))
uc = u.reshape((x_size, y_size * w_size))
vc = v.reshape((x_size, y_size * w_size))

plt.quiver(xc, yc, uc, vc, wc, scale=10)

Image

I tried to 'reshape' my xarray dataset with .stack, but I didn't find a suitable solution to work with xr.quiver.

Additional context

No response

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions