xarray：具有低开销轴坐标变换的极坐标pcolormesh

问题描述

我正在尝试绘制二维xarray DataArray，该二维数组表示以极坐标表示的参数变量。 重要：theta坐标以度为单位，而不是弧度。以下代码段创建了一个示例数据集：

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"),)

我想将此数据绘制为极坐标pcolormesh。我还想依靠xarray的绘图例程从尽可能多的功能（构面，绘图自定义等）中受益。 Matplotlib的极坐标投影假定theta角以弧度给出：如果我寻求简单的解决方案，我首先必须将theta坐标转换为弧度，但是我不想修改数组到位。除了复制数组并转换副本的theta，没有比找到更好的方法了，例如：

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
    )

这将产生所需的图：

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

实际问题

但是我想避免重复数据：我的实际数据集比那大得多。我进行了一些研究，发现了Matplotlib的转换管道，我有种感觉，可以用它在绘图例程中动态地插入此转换，但是到目前为止，我什么都还不能正常工作。有人知道我该如何进行吗？

解决方法

由于@kmuehlbauer的建议和对xarray.DataArray.assign_coords() docs的仔细检查，我设法产生了我想要的东西。

首先，我将测试数据修改为还包含单位元数据：

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"

然后，我改进了kwargs处理以自动进行单位转换，并创建了一个与theta维度相关联的额外坐标：

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,)

这里非常重要的一点是assign_coords()返回它从中调用的数据数组的副本，并且该副本的值实际上引用了原始数组，因此除了创建额外的坐标外，不增加内存开销。按照@kmuehlbauer的建议就地修改数据数组很简单（只需将da_plot = da.assign_coords(...)替换为da = da.assign_coords(...)）。

然后我们得到相同的图（没有轴标签，因为我更改了默认值以便隐藏它们）：

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