Python Xarray: how to convert a 3-d DataArray to a 2-d stacked Pandas dataframe

Question

I have a 3d xarray DataArray volume of data of time series data for multiple runs of a model. So the rows are indexed by the simulation timestep, the columns are just a variety of variables captured about the model, and then the depth coordinate represents the individual simulation run, since I run the entire simulation multiple time.

My goal is to take this 3d xarray DataArray and convert it to a 2d pandas dataframe so that I can export it to a CSV file. I order to do that, I need to stack each of the simulation runs on top of each other, so that the 3d array is converted to a 2d array.

I have some code to generate some test data, but I am not familiar enough with Xarray to know how to do this kind of stacking.

So here is some code to develop test data.

import xarray as xr
import pandas as pd
import numpy as np
from tqdm import tqdm

results_matrix = np.zeros([5, 7, 4])
simulation_matrix = xr.DataArray(results_matrix,
                                      coords={'simdata': ['val1', 'val2','val3','val4'],
                                              'run': range(5),
                                              'year': range(7)},
                                      dims=('run', 'year', 'simdata'))

itercount = 0
for i in tqdm(range(5)):
    simulation_matrix[i, :, :] = i
    itercount += 1

This code will generate a DataArray that looks like

<xarray.DataArray (run: 5, year: 7, simdata: 4)>
array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]],
 ... Additional arrays truncated

I want this converted to a 2d Pandas dataframe something like

        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.]]]

Any suggestions?

UPDATED:

Based upon comments from @rahlf23 and @DSM, I had some luck with simulation_matrix.to_dataframe('fred').unstack().

        fred
simdata val1    val2    val3    val4
run year                
0  0    0   0.0 0.0 0.0 0.0
   1    0.0 0.0 0.0 0.0
   2    0.0 0.0 0.0 0.0
   3    0.0 0.0 0.0 0.0
   4    0.0 0.0 0.0 0.0
   5    0.0 0.0 0.0 0.0
   6    0.0 0.0 0.0 0.0
1   0   1.0 1.0 1.0 1.0
   1    1.0 1.0 1.0 1.0
   2    1.0 1.0 1.0 1.0
   3    1.0 1.0 1.0 1.0
   4    1.0 1.0 1.0 1.0
   5    1.0 1.0 1.0 1.0
   6    1.0 1.0 1.0 1.0

Answer 1

Using your test data, you can use to_pandas() and pd.concat():

df = pd.concat([simulation_matrix.loc[i,:,:].to_pandas() for i in range(simulation_matrix.shape[2])])

Yields:

simdata  val1  val2  val3  val4
year                           
0         0.0   0.0   0.0   0.0
1         0.0   0.0   0.0   0.0
2         0.0   0.0   0.0   0.0
3         0.0   0.0   0.0   0.0
4         0.0   0.0   0.0   0.0
5         0.0   0.0   0.0   0.0
6         0.0   0.0   0.0   0.0
0         1.0   1.0   1.0   1.0
1         1.0   1.0   1.0   1.0
2         1.0   1.0   1.0   1.0
3         1.0   1.0   1.0   1.0
4         1.0   1.0   1.0   1.0
5         1.0   1.0   1.0   1.0
6         1.0   1.0   1.0   1.0
0         2.0   2.0   2.0   2.0
1         2.0   2.0   2.0   2.0
2         2.0   2.0   2.0   2.0
3         2.0   2.0   2.0   2.0
4         2.0   2.0   2.0   2.0
5         2.0   2.0   2.0   2.0
6         2.0   2.0   2.0   2.0
0         3.0   3.0   3.0   3.0
1         3.0   3.0   3.0   3.0
2         3.0   3.0   3.0   3.0
3         3.0   3.0   3.0   3.0
4         3.0   3.0   3.0   3.0
5         3.0   3.0   3.0   3.0
6         3.0   3.0   3.0   3.0

Answer 2

You can use .to_dataframe and then unstack, you just need to pass a name to attach to the dataset (which becomes a column containing that value):

In [41]: simulation_matrix.to_dataframe("results").unstack()
Out[41]: 
         results               
simdata     val1 val2 val3 val4
run year                       
0   0        0.0  0.0  0.0  0.0
    1        0.0  0.0  0.0  0.0
    2        0.0  0.0  0.0  0.0
    3        0.0  0.0  0.0  0.0
    4        0.0  0.0  0.0  0.0
    5        0.0  0.0  0.0  0.0
    6        0.0  0.0  0.0  0.0
1   0        1.0  1.0  1.0  1.0
    1        1.0  1.0  1.0  1.0
    2        1.0  1.0  1.0  1.0
    3        1.0  1.0  1.0  1.0
    4        1.0  1.0  1.0  1.0
    5        1.0  1.0  1.0  1.0
    6        1.0  1.0  1.0  1.0
2   0        2.0  2.0  2.0  2.0
    1        2.0  2.0  2.0  2.0
    2        2.0  2.0  2.0  2.0
    3        2.0  2.0  2.0  2.0
    4        2.0  2.0  2.0  2.0
    5        2.0  2.0  2.0  2.0
    6        2.0  2.0  2.0  2.0
3   0        3.0  3.0  3.0  3.0
    1        3.0  3.0  3.0  3.0
    2        3.0  3.0  3.0  3.0
    3        3.0  3.0  3.0  3.0
    4        3.0  3.0  3.0  3.0
    5        3.0  3.0  3.0  3.0
    6        3.0  3.0  3.0  3.0
4   0        4.0  4.0  4.0  4.0
    1        4.0  4.0  4.0  4.0
    2        4.0  4.0  4.0  4.0
    3        4.0  4.0  4.0  4.0
    4        4.0  4.0  4.0  4.0
    5        4.0  4.0  4.0  4.0
    6        4.0  4.0  4.0  4.0

All the "run" values are there even though the default representation only shows the first in a repeated group for conciseness:

In [50]: df = simulation_matrix.to_dataframe("results").unstack()

In [51]: df.reset_index().head()
Out[51]: 
        run year results               
simdata             val1 val2 val3 val4
0         0    0     0.0  0.0  0.0  0.0
1         0    1     0.0  0.0  0.0  0.0
2         0    2     0.0  0.0  0.0  0.0
3         0    3     0.0  0.0  0.0  0.0
4         0    4     0.0  0.0  0.0  0.0