xarray:具有低开销轴坐标转换的极坐标 pcolormesh
xarray: polar pcolormesh with low-overhead axis coordinate transformation
我正在尝试绘制一个二维 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()
returns 调用它的数据数组的一个副本,这个副本的值实际上引用了原始数组,因此除了创建额外的坐标。按照@kmuehlbauer 的建议修改数据数组 in-place 很简单(只需将 da_plot = da.assign_coords(...)
替换为 da = da.assign_coords(...)
)。
然后我们得到相同的图(没有轴标签,因为我更改了默认值以隐藏它们):
pcolormesh_polar_cheap(my_da, r="r", theta="theta")
我正在尝试绘制一个二维 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()
returns 调用它的数据数组的一个副本,这个副本的值实际上引用了原始数组,因此除了创建额外的坐标。按照@kmuehlbauer 的建议修改数据数组 in-place 很简单(只需将 da_plot = da.assign_coords(...)
替换为 da = da.assign_coords(...)
)。
然后我们得到相同的图(没有轴标签,因为我更改了默认值以隐藏它们):
pcolormesh_polar_cheap(my_da, r="r", theta="theta")