为两个 DataArrays 中的每个时间点计算矩阵向量乘法

Question

我有包含头部位置向量 (x,y,z) 和头部旋转矩阵 (3x3) 的运动数据。对于每个时间点，我都有一个。它们存储在两个 DataArrays 中，如下所示：

pos = np.random.sample((1000, 3))
pos_data = xr.DataArray(pos, dims=('time', 'h_pos'), coords={'h_pos': ['x', 'y', 'z']})

rot = np.random.sample((1000, 3, 3))
rot_data = xr.DataArray(rot, dims=('time', 'h_rot_i', 'h_rot_j'))

如何通过将位置向量乘以旋转矩阵来计算每个时间点的旋转位置向量？

使用纯 numpy，我可以使用：

rotated = np.einsum('tij,tj->ti', rot, pos)

文档说 xarray.dot 类似于 np.einsum，但我无法让它工作。 dims 也不相同，在 einsum 中有两个 i，但是矩阵的 dims 是 h_rot_i，位置向量的 dims 是 h_pos..

Answer 1

您需要将 h_rot_j 维度重命名为 h_pos（或相反）以便名称匹配。然后您还需要在 dot 中指定 dims='time' 以明确保留时间维度。

理论上，以下应该有效，但目前由于 xarray bug:

而引发错误

>>> rot_data2 = rot_data.rename({'h_rot_j': 'h_pos'}
>>> xr.dot(rot_data2, pos_data, dims='h_pos')
ValueError: indexes along dimension 'x' are not equal

作为变通方法，您还应该将 h_pos 坐标值分配给 rot_data2，例如

>>> rot_data2 = rot_data.rename({'h_rot_j': 'h_pos'})
>>> rot_data2.coords['h_pos'] = pos_data.h_pos
>>> xr.dot(rot_data2, pos_data, dims='h_pos')
<xarray.DataArray (time: 1000, h_rot_i: 3)>
array([[1.034348, 0.655469, 1.235536],
       [0.815916, 0.632311, 0.603557],
       [0.87191 , 0.54483 , 0.586762],
       ...,
       [0.629219, 0.716169, 0.721611],
       [0.517938, 0.440981, 0.649848],
       [0.763335, 0.624545, 0.370505]])
Dimensions without coordinates: time, h_rot_i

此时，您可能还想使用 rename 和坐标赋值来切换 h_rot_i -> h_pos，但我将把它留作 reader :).

为两个 DataArrays 中的每个时间点计算矩阵向量乘法

Computing matrix vector multiplication for each time point in two DataArrays

python-xarray