xarray：重塑数据，拆分维

Question

I have a Dataset in xarray with the following dimensions:

Dimensions:      (subject: 30, session: 5, time: 45000)
Coordinates:
  * subject      (subject) object '110' '112' '114' '117' ...
  * session      (session) object 'week1' 'week2' 'week3' ...
  * time         (time) timedelta64[ns] 00:00:00 00:00:00.040000 ...

I want to split each trial (subject/session combo) into smaller time segments, for example into 3 segments of 15000 values each, The resulting dimensions may look as follows:

(subject: 30, session: 5, segment: 3, time: 15000)

I've searched and tried a lot of things but have not succeeded, how can this be done?

One of the things that I've been trying, that seems to be close, is creating a new MultiIndex and unstacking it.

segment_data = np.repeat(range(3),len(ds.time)//3)
segment = xr.Variable(dims='time',data=segment_data)
newtime_data = np.tile(ds.time[:len(ds.time)//3],3)
newtime = xr.Variable(dims='time',data=newtime_data)
dsr = ds.assign_coords(segment=segment,newtime=newtime)
dsr = dsr.set_index(segment='segment',newtime='newtime')
dsr = dsr.stack(fragment=['segment','newtime'])

However that last line takes a huge amount of memory and seems to create a dimension fragment: len(ds.time)**2 , which doesn't seem right. I'm also no sure what I would have to do after this ( unstack('fragment') ?).

edit: Some more attempts have brought me here:

x = np.repeat(range(3),15000)
y = np.tile(ds.time[:len(ds.time)//3],3)
dsr = (ds.assign_coords(segment=x,time2=y)
      .set_index(fragment=['segment','time2'])
      .unstack('fragment'))

Which gives this:

(subject: 30, segment: 3, session: 5, time: 45000, time2: 15000)

This seems close but it's not quite there since every time2 point now has 45000 values while it should be a single value:

dsr.isel(subject=0,segment=0,session=0,time2=0)
# (time: 45000)

edit: I finally found a way to do it, see my answer. Futher suggestions welcome!

Answer 1

First make sure you have the labels for the two new dimensions. In this case as follows:

x = range(3) # 3 segments
y = ds.time[:len(ds.time)//3] # the first 1/3rd of the time labels

Then create a pandas MultiIndex from these labels*.

ind = pd.MultiIndex.from_product((x,y),names=('segment','new_time'))

Finally, replace the time index in the Dataset by this new index, and then unstack its levels to create the two required dimensions.

dsr = ds.assign(time=ind).unstack('time')

You may want to use rename to rename the new dimension:

dsr = dsr.rename({'new_time':'time'})

Resulting dimensions:

(subject: 30, segment: 3, session: 5, time: 15000)

The only thing that's off now is the order of the dimensions (ideally segment and session should be swapped). I thought transpose would help here but "although the order of dimensions on each array will change, the dataset dimensions themselves will remain in fixed (sorted) order." ** So I'll probably live with it like this.

_{* Note you won't be able to use the name of the dimension you want to split, so we have 'new_time' here. An unnecessary limitation of assign ?}

_{** Another limitation that I can't explain.}