matplotlib_support: Quiver plot doesn't work with mixed units

[mtryan83]

It seems like matplotlib's quiver plot isn't supported in the matplotlib_support context manager when x and y have mixed units. I'm not really sure this a unyt problem (because of how quiver works), so feel free to close this, but I think it should be at least mentioned in the documentation.

Example

# Basic example from Strogatz's Nonlinear Dynamics and Chaos (Example 6.5.2)
x = np.linspace(-2,2) * meter
y = np.linspace(-2,2), * meter/second
xg,yg = np.meshgrid(x,y)
dt = 1 * second
xdot = yg 
ydot = ((1/second) * xg - (1/second/meter**2) * xg**3)
u = xdot * dt
v = ydot * dt
# u and x have the same units and v and y have the same units to match angles='xy', scale_units='xy'
fig,ax = plt.subplots()
with matplotlib_support:
    ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=10)

produces the error:

---------------------------------------------------------------------------
UnitInconsistencyError                    Traceback (most recent call last)
Cell In[1270], line 9
      7 fig,ax = plt.subplots()
      8 with matplotlib_support:
----> 9     ax.quiver(x,y,xdot,ydot,angles='xy',scale_units='xy',scale=10)

File ~/.conda/envs/example/lib/python3.13/site-packages/matplotlib/__init__.py:1521, in _preprocess_data.<locals>.inner(ax, data, *args, **kwargs)
   1518 @functools.wraps(func)
   1519 def inner(ax, *args, data=None, **kwargs):
   1520     if data is None:
-> 1521         return func(
   1522             ax,
   1523             *map(cbook.sanitize_sequence, args),
   1524             **{k: cbook.sanitize_sequence(v) for k, v in kwargs.items()})
   1526     bound = new_sig.bind(ax, *args, **kwargs)
   1527     auto_label = (bound.arguments.get(label_namer)
   1528                   or bound.kwargs.get(label_namer))

File ~/.conda/envs/example/lib/python3.13/site-packages/matplotlib/axes/_axes.py:5501, in Axes.quiver(self, *args, **kwargs)
   5499 # Make sure units are handled for x and y values
   5500 args = self._quiver_units(args, kwargs)
-> 5501 q = mquiver.Quiver(self, *args, **kwargs)
   5502 self.add_collection(q, autolim=True)
   5503 self._request_autoscale_view()

File ~/.conda/envs/example/lib/python3.13/site-packages/matplotlib/quiver.py:524, in Quiver.__init__(self, ax, scale, headwidth, headlength, headaxislength, minshaft, minlength, units, scale_units, angles, width, color, pivot, *args, **kwargs)
    522 self.X = X
    523 self.Y = Y
--> 524 self.XY = np.column_stack((X, Y))
    525 self.N = len(X)
    526 self.scale = scale

File ~/.conda/envs/example/lib/python3.13/site-packages/unyt/array.py:2045, in unyt_array.__array_function__(self, func, types, args, kwargs)
   2043 if not all(issubclass(t, unyt_array) or t is np.ndarray for t in types):
   2044     return NotImplemented
-> 2045 return _HANDLED_FUNCTIONS[func](*args, **kwargs)

File ~/.conda/envs/example/lib/python3.13/site-packages/unyt/_array_functions.py:320, in column_stack(tup)
    318 @implements(np.column_stack)
    319 def column_stack(tup):
--> 320     ret_units = _validate_units_consistency(tup)
    321     return np.column_stack._implementation([np.asarray(_) for _ in tup]) * ret_units

File ~/.conda/envs/example/lib/python3.13/site-packages/unyt/_array_functions.py:231, in _validate_units_consistency(objs)
    229     return units[0]
    230 else:
--> 231     raise UnitInconsistencyError(*units)

UnitInconsistencyError: Expected all unyt_array arguments to have identical units. Received mixed units (m, m/s)

Note the primary problem appears to be the xy column stacking (the self.xy = np.column_stack((X, Y)), but additional problems may show up if for example x and xdot have different units. This isn't surprising but a better error message may be helpful.

Workarounds

Changing the last lines to

with matplotlib_support:
    ax.quiver(x.v, y.v, u.v, v.v, angles='xy', scale_units='xy', scale=10)

or using (artificially) consistent units (e.g. changing y, u, v units to meters) produces the expected plot:

(Though note the y-axis units would be m/s if this wasn't an issue.)

I'd be fine with using either of these workarounds but it means I can't combine a quiver plot with other types of plots that do work fine with mixed units (plot, contour, etc).

Versions:

unyt: 3.0.3
matplotlib: 3.10.0
numpy: 2.2.0

Solution/Comments

Like I said at the top, I don't really have a solution due to how quiver is implemented, but I think at minimum adding something to the matplotlib_support documentation would be helpful.

Let me know if you need any more details.

[neutrinoceros]

Thanks for reporting this.
At first glance, I would think of this as a bug in matplotlib: calling np.column_stack((X, Y)) doesn't seem to make sense in this context and I believe unyt is correct in raising an exception when X and Y do not have the same unit. Unfortunately I don't see how unyt's error message could be improved here, but I'm opened to suggestions.

Additionally, I would also assume that it could easily be fixed in matplotlib by stripping out units in a generic way, doing something like

self.X = X = np.asarray(X)
self.Y = Y = np.asarray(Y)
self.XY = np.column_stack((X, Y))

as you mentionned, other issues might pop up later down the line, but it seems worth to give it a shot first. I don't see any actionable for unyt at this point.