Understanding xarray.apply_ufunc

Question

I have an xarray with multiple time dimensions slow_time, fast_time a dimension representing different objects object and a dimension reflecting the position of each object at each point in time coords.

The goal is now to apply a rotation using scipy.spatial.transform.Rotation to each position in this array, for every point in time.

I'm struggling to figure out how to use xarray.apply_ufunc to do what I want, mainly because the concept of input_core_dimensions isn't really clear to me.

The code below shows what I'm trying to do:

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, ...??)

The behaviour I'm looking for is essentially like four nested for loops that apply the rotations to each point like this pseudocode

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

But I'm unsure how to proceed.

Using this

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!

Answer 1

Reference

First off, a helpful reference would be this documentation page on unvectorized ufunc

Solution

As I understand your question you want to apply a rotation to the position vector of each object at every time. The way you have set up your data already puts the coordinates as the final dimension of the array.

Translating your pseudocode to generate a reference dataset rotatedPositions yields:

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])

where I hard-coded the object dimension size.

In essence the apply_rotation function takes 1 3-vector (1D array of size 3) and 2 scalars and returns a 1D array of size 3 (3 vector).

Following the documentation mentioned above I arrive at the following call to apply_ufunc:

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

Testing via

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

indicates success.

Explanation

As I understand the documentation cited above apply_ufunc will take the function to be applied as first argument, then all positional arguments in order.

Next one has to provide the dimension labels of each dataset that will correspond to the data that will be core to apply_rotation working. Here this is coords, as we manipulate coordinates. Since neither theta nor phi have this dimension we do not specify anything for them.

Next we have to specify the dimensions the output data will have, since we just transform the output data we keep output_core_dims=[['coords']]. Leaving this out would lead apply_ufunc to assume output data to be 0-dimensional (a scalar).

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