CODEWITHSUNDEEP
pythonnumpydaskpython-xarray
shaunakde
shaunakde

Reputation: 3106

Move dimensions to original order after apply_ufunc in xarray?

My dataset has 3 dimensions in the order (time, y, x) and I use apply_ufunc to apply a computation along the time dimension. This rearranges the order of the dimensions as (y, x, time). I need to restructure the xarray so its in the (time, y, x) order as the original dataset. How would I go along doing this?

Here is a visual description of what's happening:

Before:

dcube before

Then I apply my function:

dcube = xr.apply_ufunc(
        bc.clip_and_normalize_percentile, 
        dcube, 
        input_core_dims=[["time"]], 
        output_core_dims=[["time"]], 
        dask = 'allowed',
        vectorize=True
    )

as expected time is moved to the last dimension:

after

How do I rearrange this so that its in the order of the original array? Are there parameters that prevent apply_ufunc from moving the dims?

Upvotes: 2

Views: 996

Answers (2)

Val
Val

Reputation: 7023

The docs say that

Core dimensions are automatically moved to the last axes of input variables before applying func, which facilitates using NumPy style generalized ufuncs

so it's unlikely that there's a way (or any parameters) to prevent that.

What I've been doing is simply call .transpose afterwards to restore the initial order.

In your example, that would look like:

dcube = dcube.transpose("time", ...)

fixing time to be the first dimension and shifting all other ones behind using ....

Upvotes: 3

NNN
NNN

Reputation: 591

Would np.swapaxes help?

import numpy as np
aa = np.arange(2*3*4).reshape(2,3,4)
bb = aa.swapaxes(2,0)
print(bb.shape)
print(aa[0,1,2])
print(bb[2,1,0])

seems np.einsum can work too

import numpy as np
aa = np.arange(2*3*4).reshape(2,3,4)
bb = np.einsum('ijk->kji',aa)
print(bb.shape)
print(aa[0,1,2])
print(bb[2,1,0])

Upvotes: 1

Related Questions