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: 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.