Converting a DataArray to a DataFrame and preserve coordinate label order

user_5594299 0
Viewed 2877

Is there a simple way to convert an xarray DataArray to a pandas DataFrame, where I can prescribe which dimensions get turned into index/columns? For example, suppose I have a DataArray

import xarray as xr
weather = xr.DataArray(
    name='weather',
    data=[['Sunny', 'Windy'], ['Rainy', 'Foggy']],
    dims=['date', 'time'],
    coords={
        'date': ['Thursday', 'Friday'],
        'time': ['Morning', 'Afternoon'],
    }
)

which results in:

<xarray.DataArray 'weather' (date: 2, time: 2)>
array([['Sunny', 'Windy'],
       ['Rainy', 'Foggy']], dtype='<U5')
Coordinates:
  * date     (date) <U8 'Thursday' 'Friday'
  * time     (time) <U9 'Morning' 'Afternoon'

Suppose I now want to move that to a pandas DataFrame indexed by date, with columns time. I can kind of do this by using .to_dataframe() and then .unstack() on the resulting dataframe:

>>> weather.to_dataframe().unstack()
           weather        
time     Afternoon Morning
date                      
Friday       Foggy   Rainy
Thursday     Windy   Sunny

However, pandas will sort things so rather than Morning followed by Afternoon, I get Afternoon followed by Morning. I was rather hoping there would be an API like

weather.to_dataframe(index_dims=[...], column_dims=[...])

which could do this reshaping for me, without me having to re-sort my indices and columns afterwards.

python-xarray pandas
2 Answers

In xarray 0.16.1, dim_order was added to .to_dataframe. Does this do what you're looking for?

xr.DataArray.to_dataframe(
    self,
    name: Hashable = None,
    dim_order: List[Hashable] = None,
) -> pandas.core.frame.DataFrame
Docstring:
Convert this array and its coordinates into a tidy pandas.DataFrame.

The DataFrame is indexed by the Cartesian product of index coordinates
(in the form of a :py:class:`pandas.MultiIndex`).

Other coordinates are included as columns in the DataFrame.

Parameters
----------
name
    Name to give to this array (required if unnamed).
dim_order
    Hierarchical dimension order for the resulting dataframe.
    Array content is transposed to this order and then written out as flat
    vectors in contiguous order, so the last dimension in this list
    will be contiguous in the resulting DataFrame. This has a major
    influence on which operations are efficient on the resulting
    dataframe.

    If provided, must include all dimensions of this DataArray. By default,
    dimensions are sorted according to the DataArray dimensions order.
user_30647360

If you want to make sure that coordinate labels stay in a particular (non-default) order, you can explicitly set their type to be a CategoricalIndex:

In [25]: import xarray as xr, pandas as pd
    ...: weather = xr.DataArray(
    ...:     name='weather',
    ...:     data=[['Sunny', 'Windy'], ['Rainy', 'Foggy']],
    ...:     dims=['date', 'time'],
    ...:     coords={
    ...:         'date': pd.CategoricalIndex(['Thursday', 'Friday'], ordered=True),
    ...:         'time': pd.CategoricalIndex(['Morning', 'Afternoon'], ordered=True),
    ...:     }
    ...: )

This will be preserved when converting to pandas:

In [26]: weather.to_dataframe().unstack()
Out[26]:
         weather
time     Morning Afternoon
date
Thursday   Sunny     Windy
Friday     Rainy     Foggy
user_38887190
Related