hooks to "prepare" xarray objects for plotting

[keewis]

From xarray-contrib/pint-xarray#61 (comment)

matplotlib has a module called matplotlib.units which manages a mapping of types to hooks. This is then used to convert custom types to something matplotlib can work with, and to optionally add axis labels. For example, with pint:

In [9]: ureg = pint.UnitRegistry()
   ...: ureg.setup_matplotlib()
   ...:
   ...: t = ureg.Quantity(np.arange(10), "s")
   ...: v = ureg.Quantity(5, "m / s")
   ...: x = v * t
   ...: 
   ...: fig, ax = plt.subplots(1, 1)
   ...: ax.plot(t, x)
   ...: 
   ...: plt.show()

this will plot the data without UnitStrippedWarnings and even attach the units as labels to the axis (the format is hard-coded in pint right now).

While this is pretty neat there are some issues:

• xarray's plotting code converts to masked_array, dropping metadata on the duck array (which means matplotlib won't see the duck arrays)
• we will end up overwriting the axis labels once the variable names are added (not sure if there's a way to specify a label format?)
• it is matplotlib specific, which means we have to reimplement once we go through with the plotting entrypoints discussed in ENH: Plotting backend options #3553 and Add entrypoint for plotting backends #3640

All of this makes me wonder: should we try to maintain our own mapping of hooks which "prepare" the object based on the data's type? My initial idea would be that the hook function receives a Dataset or DataArray object and modifies it to convert the data to numpy arrays and optionally modifies the attrs.

For example for pint the hook would return the result of .pint.dequantify() but it could also be used to explicitly call .get on cupy arrays or .todense on sparse arrays.

xref #5561

[dcherian]

For example for pint the hook would return the result of .pint.dequantify() but it could also be used to explicitly call .get on cupy arrays or .todense on sparse arrays.

We discussed as_numpy or to_numpy_data here: #3245

[keewis]

that will work if we want to remove the units, but .pint.dequantify() also adds the units as string-valued attributes. In other words, I'd like the new hook to work on xarray objects instead of the wrapped duck array.

[dcherian]

wouldn't as_numpy on a pint-backed DataArray return a DataArray with numpy data and units in the attribute?

[dcherian]

Oh I see. as_numpy is using data.magnitude and dropping the units. I think we should special-case pint and add the units attribute on the DataArray like .pint.dequantify

[dcherian]

I think we should special-case pint and add the units attribute on the DataArray like .pint.dequantify

We would also want this to write pint-backed xarray datasets to netCDF.

[keewis]

that makes sense, but I think long term it might be better to find a way to delegate this to duck array extension libraries (it really should call .pint.dequantify()).