了解 xarray.apply_ufunc

Question

我有一个具有多个时间维度的 xarray slow_time，fast_time一个维度代表不同的对象object，一个维度反映每个对象在每个时间点的位置coords。

scipy.spatial.transform.Rotation现在的目标是针对每个时间点对该数组中的每个位置应用旋转。

我正在努力弄清楚如何使用来做我想做的事情，主要是因为我不太清楚xarray.apply_ufunc这个概念。input_core_dimensions

下面的代码显示了我正在尝试做的事情：

import numpy as np
import xarray as xr
from scipy.spatial.transform import Rotation

# dummy initial positions
initial_position = xr.DataArray(np.arange(6).reshape((-1,3)), dims=["object", "coords"])
# dummy velocities
velocity = xr.DataArray(np.array([[1, 0, 0], [0, 0.5, 0]]), dims=["object", "coords"])

slow_time = xr.DataArray(np.linspace(0, 1, 10, endpoint=False), dims=["slow_time"])
fast_time = xr.DataArray(np.linspace(0, 0.1, 100, endpoint=False), dims=["fast_time"])

# times where to evaluate my function
times = slow_time + fast_time

# this is the data I want to transform
positions = times * velocity + initial_position

# these are the rotation angles
theta = np.pi/20 * times
phi = np.pi/100 * times

def apply_rotation(vectors, theta, phi):
    R = Rotation.from_euler("xz", (theta, phi))
    result = R.apply(vectors)
    return result

rotated_positions = xr.apply_ufunc(apply_rotation,
    positions, theta, phi, ...??)

我正在寻找的行为本质上就像四个嵌套的 for 循环，将旋转应用到每个点，如这个伪代码

for pos, t, p in zip(positions, theta, phi):
    R = Rotation.from_euler("xz", (t, p))
    R.apply(pos)

但我不确定如何继续。

使用这个

rotated_positions = xr.apply_ufunc(apply_rotation,
    positions, theta, phi,
    input_core_dims=[["object"],[],[]], output_core_dims=[["object"]])

我认为该函数将沿着维度的子数组应用object，但现在我将整个数组传递到我的函数中，但这不起作用。

xarray 文档中的信息apply_ufunc并没有真正让事情变得非常清楚。

任何帮助表示赞赏！

Answer 1

参考

首先，一个有用的参考是关于非矢量化 ufunc 的文档页面

解决方案

据我了解您的问题，您希望每次都对每个对象的位置向量应用旋转。您设置数据的方式已经将坐标作为数组的最终维度。

翻译伪代码以生成参考数据集会产生rotatedPositions：

rotatedPositions = positions.copy()
for slowTimeIdx in range( len(slow_time)):
    for fastTimeIdx in range( len(fast_time) ):
        for obj in range(2):
            pos = rotatedPositions.data[slowTimeIdx, fastTimeIdx, obj].copy()
            rotatedPositions.data[slowTimeIdx, fastTimeIdx, obj] =  apply_rotation(pos, theta.data[slowTimeIdx, fastTimeIdx], phi.data[slowTimeIdx, fastTimeIdx])

我在其中硬编码了object尺寸大小。

本质上，该apply_rotation函数采用 1 个 3 向量（大小为 3 的一维数组）和 2 个标量，并返回大小为 3 的一维数组（3 个向量）。

根据上述文档，我收到以下调用apply_ufunc：

rotated = xr.apply_ufunc(apply_rotation, 
               positions,
                theta,
                phi,
               input_core_dims=[['coords'], [], []],
               output_core_dims=[['coords']],
               vectorize=True
              )

测试通过

np.allclose(rotatedPositions.data, rotated.data)

表示成功。

解释

据我了解，上面引用的文档apply_ufunc将把要应用的函数作为第一个参数，然后按顺序应用所有位置参数。

接下来必须提供每个数据集的维度标签，该标签将与将要core工作的数据相对应apply_rotation。这就是coords我们操纵坐标时的情况。由于既没有theta也phi没有这个维度，我们不为它们指定任何内容。

接下来，我们必须指定输出数据的维度，因为我们只是转换我们保留的输出数据output_core_dims=[['coords']]。忽略这一点将导致apply_ufunc假设输出数据为 0 维（标量）。

最后，vectorize=True确保该函数在 中未指定的所有维度上执行input_core_dims。