xarray广播如何与numpy功能一起工作？

Question

我正在尝试使用我的xarray坐标进行坐标转换

我有一个像DataArray：

d = xr.DataArray(np.zeros((10, 10, 1)), dims=['x', 'y', 'z']

并且我正在执行以下操作：

r = np.sqrt(d.x**2 + d.y**2 + d.z**2)
theta = np.arctan2(np.sqrt(d.x**2 + d.y**2), d.z)
phi = np.arctan2(d.y, d.x)

我的形状是：

In [212]: r.shape
Out[212]: (10, 10, 1)

In [214]: theta.shape
Out[214]: (10, 10)

In [216]: phi.shape
Out[216]: (10,)

我希望看到我为所有人获得的(10, 10, 1)形状。

看起来xarray的奇特广播只能用于基本算术运算，而numpy函数则天真地播放。

这是对的吗？

是否有批准的解决方法？我可以编写自己的xarray感知arctan2函数，它可以正常播放，但我希望我没有。

Answer 1

不幸的是，NumPy函数通常不会在xarray对象上进行适当的广播。但是xarray确实带有许多这些函数的包装版本，包括xarray.ufuncs module中的arctan2。所以xarray.ufuncs.arctan2应该完全符合您的要求。