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() 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() returns 调用它的数据数组的一个副本，这个副本的值实际上引用了原始数组，因此除了创建额外的坐标。按照@kmuehlbauer 的建议修改数据数组 in-place 很简单（只需将 da_plot = da.assign_coords(...) 替换为 da = da.assign_coords(...)）。

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

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

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

xarray: polar pcolormesh with low-overhead axis coordinate transformation

python

matplotlib

python-xarray

polar-coordinates

实际问题