了解xarray.apply_ufunc

我有一個具有多個時間維度的 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"]])

I thought the function would be applied along subarrays of the object dimension, but now I get entire arrays passed into my function which doesn't work.

The information on apply_ufunc in the xarray documentation doesn't really makes things very clear.

Any help is appreciated!

uj5u.com熱心網友回復：

參考

首先，一個有用的參考是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 維（標量）。

Finally,vectorize=True ensures that the function is executed over all dimensions not specified in input_core_dims.

標籤：python numpy scipy python-xarray

上一篇：如何使用線性方程python計算資料框的新列

下一篇：如何計算資料框中所有行的專案總數