了解 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

参考

首先，一个有用的参考是 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])

其中我 hard-coded object 维度大小。

本质上，apply_rotation 函数采用 1 个 3 向量（大小为 3 的一维数组）和 2 个标量，returns 大小为 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。

了解 xarray.apply_ufunc

Understanding xarray.apply_ufunc

python

numpy

scipy

python-xarray

参考

解决方案

说明