xarray: polar pcolormesh with low-overhead axis coordinate transformation

Question

I'm trying to plot a two-dimensional xarray DataArray representing a variable parametrised in polar coordinates. Important: the theta coordinate is in degree, not in radian. The following snippet creates an example data set:

import numpy as np
import xarray as xr

res_theta = 20
thetas = np.arange(0, 360, res_theta)
res_r = 0.1
rs = np.arange(0, 1, res_r)
data = np.random.random((len(thetas), len(rs)))
my_da = xr.DataArray(
    data,
    coords=(thetas, rs),
    dims=("theta", "r"),
)

I would like to plot this data as a polar pcolormesh. I also would like to rely on xarray's plotting routines to benefit from as many features as possible (faceting, plot customisation, etc.). Matplotlib's polar projection assumes that the theta angle is given in radian: if I go for the straightforward solution, I first have to convert my theta coordinates to radian, but I don't want to modify the array in-place. I haven't found a better way than copying the array and converting the copy's theta, like this for instance:

def pcolormesh_polar_expensive(da, *args, **kwargs):
    da_tmp = da.copy()  # I'd like to avoid that
    
    # Get x value
    try:
        x = args[0]
    except IndexError:
        x = da_tmp.dims[0]
    
    da_tmp[x] = np.deg2rad(da_tmp[x])

    try:
        subplot_kws = kwargs["subplot_kws"]
    except KeyError:
        subplot_kws = {}
    
    return da_tmp.plot.pcolormesh(
        *args, 
        subplot_kws=dict(projection="polar"),
        **kwargs
    )

This produces the desired plot:

pcolormesh_polar_expensive(my_da, "theta", "r")

The Actual Problem

I however would like to avoid duplicating the data: my actual data sets are much larger than that. I made some research and found out about Matplotlib's transformation pipeline, and I have the feeling that I could use it to dynamically insert this transformation in plotting routines, but I couldn't get anything to work properly so far. Does anybody have an idea of how I could proceed?

Answer 1

Thanks to @kmuehlbauer's suggestion and a careful examination of the xarray.DataArray.assign_coords() docs, I managed to produce exactly what I wanted.

First, I modified my test data to also include unit metadata:

import numpy as np
import xarray as xr
import pint

ureg = pint.UnitRegistry()

res_r = 0.1
rs = np.arange(0, 1, res_r)
res_theta = 20
thetas = np.arange(0, 360, res_theta)
data = np.random.random((len(rs), len(thetas)))
my_da = xr.DataArray(
    data,
    coords=(rs, thetas),
    dims=("r", "theta"),
)
my_da.theta.attrs["units"] = "deg"

Then, I improved the kwargs processing to automate unit conversion and created an extra coordinate associated to the theta dimension:

def pcolormesh_polar_cheap(da, r=None, theta=None, add_labels=False, **kwargs):
    if r is None:
        r = da.dims[0]
    if theta is None:
        theta = da.dims[1]
    
    try:
        theta_units = ureg.Unit(da[theta].attrs["units"])
    except KeyError:
        theta_units = ureg.rad

    if theta_units != ureg.rad:
        theta_rad = f"{theta}_rad"
        theta_rad_values = ureg.Quantity(da[theta].values, theta_units).to(ureg.rad).magnitude
        da_plot = da.assign_coords(**{theta_rad: (theta, theta_rad_values)})
        da_plot[theta_rad].attrs = da[theta].attrs
        da_plot[theta_rad].attrs["units"] = "rad"
    else:
        theta_rad = theta
        da_plot = da
    
    kwargs["x"] = theta_rad
    kwargs["y"] = r
    kwargs["add_labels"] = add_labels

    try:
        subplot_kws = kwargs["subplot_kws"]
    except KeyError:
        subplot_kws = {}
    subplot_kws["projection"] = "polar"
    
    return da_plot.plot.pcolormesh(
        **kwargs,
        subplot_kws=subplot_kws,
    )

A very important point here is that assign_coords() returns a copy of the data array it's called from, and this copy's values actually reference the original array, thus adding no memory cost other than the creation of the extra coordinate. Modifying the data array in-place as suggested by @kmuehlbauer is straightforward (just replace da_plot = da.assign_coords(...) with da = da.assign_coords(...)).

We then get the same plot (without axis labels, since I changed the defaults so as to hide them):

pcolormesh_polar_cheap(my_da, r="r", theta="theta")