python - 了解 xarray.apply_ufunc

Question

我有一个具有多个时间维度的 xarray slow_time、fast_time 一个表示不同对象的维度 object 和一个反射(reflect)每个对象位置的维度在每个时间点坐标。

现在的目标是使用 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

引用

首先，一个有用的引用是 this documentation page on unvectorized 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 中指定的所有维度上执行。