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

Question

我试图绘制一个二维 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 的转换管道，我觉得我可以用它在绘图例程中动态插入这个转换，但到目前为止我无法正常工作。 有没有人知道我该如何继续？

Answer 1

感谢@kmuehlbauer 的建议和对xarray.DataArray.assign_coords()文档的仔细检查，我设法产生了我想要的。

首先，我修改了我的测试数据以包含单元元数据：

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