CODEWITHSUNDEEP
pythondataframepython-xarray
kots
kots

Reputation: 29

Best practice to xarray Multiindex concat

I have a set of 1000 (2D) pd.Dataframe (lets say, index:time, columns: run_id) and each one of them has 3 properties (let's say temperature, pressure, location). Ideally I would like to have everything in a xr.DataArray with 5 dimensions (or a xr.Dataset with 4 dimensions and having the last dimension as unique data variables).

I created one DataArray with two dims and 2+3 coords but then the xr.concat doesnt seem to work for multiple dimensions. (I followed the approach mentioned here Add 'constant' dimension to xarray Dataset. )

Example: I build DataArrays from individual dataframes and list of properties.

# Mock data:
data = {}
for i in np.arange(500):
    data[i] = pd.DataFrame(np.random.randn(1000, 8), index=pd.DatetimeIndex(start='01.01.2013',periods=1000,freq='h'),
                    columns=list('ABCDEFGH'))
df_catalogue = pd.DataFrame(np.random.choice(10,(500, 3)), columns=['temp','pre','zon'])

#Build DataArrays adding scalar coords
res_da = []
for i,v in df_catalogue.iterrows():
    i_df = data[i] # data is a dictionary of properly indexed dataframes

    da = xr.DataArray(i_df.values,
                   coords={'time':i_df.index.values,'runs':i_df.columns.values,
                           'temp':v['temp'], 'pre':v['pre'],'zon':v['zon']},
                   dims=['time','runs'])
    res_da.append(da)

But when I try all_da = xr.concat(res_da, dim=['temp','pre','zon']) I get strange results. What is the best way to achieve something like this:

<xarray.DataArray (time: 8000, runs: 50, temp:8, pre:10, zon: 5)>
array([[[ 4545.453613,  4545.453613, ...,  4545.453613,  4545.453613],
        [ 4545.453613,  4545.453613, ...,  4545.453613,  4545.453613],
        ..., 
        [ 4177.425781,  4177.425781, ...,  4177.425781,  4177.425781]]], dtype=float32)
Coordinates:
  * runs  (runs) object 'A' 'B' ...
  * time  (time) datetime64[ns] 2013-12-31T23:00:00 2014-01-01 ...
  * zon   (zon) 'zon1', 'zon2', 'zon3', ......
  * temp  (temp)  'XX' 'YY', 'ZZ' .....
  * pre   (pre) 'AAA', 'BBB', 'CCC' ....

Upvotes: 1

Views: 2995

Answers (1)

shoyer
shoyer

Reputation: 9593

xarray.concat only supports concatenating along a single dimension. But we can get around that by concatenating, setting a MultiIndex and then unstacking.

I'm altering your setup code because this only works if each combination of the new coordinates you're building (['temp','pre','zon']) is unique:

import numpy as np
import pandas as pd
import xarray as xr
import itertools

data = {}
for i in np.arange(500):
    data[i] = pd.DataFrame(np.random.randn(1000, 8),
                           index=pd.DatetimeIndex(start='01.01.2013',periods=1000,freq='h'),
                           columns=list('ABCDEFGH'))
cat_data = [(x, y, z)
            for x in range(20)
            for y in ['a', 'b', 'c', 'd', 'e']
            for z in ['A', 'B', 'C', 'D', 'E']]
df_catalogue = pd.DataFrame(cat_data, columns=['temp','pre','zon'])

#Build DataArrays adding scalar coords
res_da = []
for i,v in df_catalogue.iterrows():
    i_df = data[i] # data is a dictionary of properly indexed dataframes

    da = xr.DataArray(i_df.values,
                   coords={'time':i_df.index.values,'runs':i_df.columns.values,
                           'temp':v['temp'], 'pre':v['pre'],'zon':v['zon']},
                   dims=['time','runs'])
    res_da.append(da)

Then, we can simply write:

xr.concat(res_da, dim='prop').set_index(prop=['temp', 'pre', 'zon']).unstack('prop')

which results in the 5D array that you desire:

<xarray.DataArray (time: 1000, runs: 8, temp: 20, pre: 5, zon: 5)>
array([[[[[-0.690557, ..., -1.526415],
          ...,
          [ 0.737887, ...,  1.585335]],

         ...,

         [[ 0.99557 , ...,  0.256517],
          ...,
          [ 0.179632, ..., -1.236502]]],


        ...,


        [[[ 0.234426, ..., -0.149901],
          ...,
          [ 1.492255, ..., -0.380909]],

         ...,

         [[-0.36111 , ..., -0.451571],
          ...,
          [ 0.10457 , ...,  0.722738]]]]])
Coordinates:
  * time     (time) datetime64[ns] 2013-01-01 2013-01-01T01:00:00 ...
  * runs     (runs) object 'A' 'B' 'C' 'D' 'E' 'F' 'G' 'H'
  * temp     (temp) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
  * pre      (pre) object 'a' 'b' 'c' 'd' 'e'
  * zon      (zon) object 'A' 'B' 'C' 'D' 'E'

Upvotes: 5

Related Questions