python-xarray - 从坐标增加 xarray 的维数

Question

假设我有以下二维数组

>>> import numpy as np
>>> budgets = np.array([
       [np.nan, 450.],
       [500.  , 100.],
       [np.nan, 900.],
    ])

它的值是这样定位的

>>> coords = [
        ('name' , ['Jack_teen' , 'John_adult', 'John_teen']), # over rows
        ('hobby', ['books', 'bicyle']),                       # over columns
    ]

使用 xarray我可以创建一个二维标记数组，做

>>> import xarray as xr
>>> x = xr.DataArray(budgets, coords=coords)

所以约翰十几岁的时候，他不喜欢书，如果当时有预算的话，这是可以看得出来的

>>> x.sel(name='John_teen', hobby='books')
<xarray.DataArray ()>
array(nan)
Coordinates:
    name     |S10 'John_teen'
    hobby    |S6 'books'

随着年龄的增长发生了什么变化

>>> x.sel(name='John_adult', hobby='books')
<xarray.DataArray ()>
array(500.0)
Coordinates:
    name     |S10 'John_adult'
    hobby    |S6 'books'

---

我的问题:

你如何将这个 2dl 数组转换成一个 3dl 数组，它考虑一个名为 age 的新维度(因此其坐标为 ['adult','teen'] ) 同时简化维度 name 的坐标？

请注意，name 的坐标始终 由下划线分隔，我的意思是 NAME_AGE。当然，您开始执行此操作的对象是 x。

是否有 xarray-builtin 方式来做到这一点？或者至少最快/最便宜的方法是什么？

Answer 1

因为我们最终想要一个维度'name'，我将把当前的'name'重命名为'name_age':

In [5]: x = x.rename({'name': 'name_age'})

我们可以直接从坐标值构造一个MultiIndex，并将其分配为堆叠的DataArray坐标:

In [6]: x.coords['name_age'] = pd.MultiIndex.from_tuples(
   ...:     [tuple(s.split('_')) for s in x.coords['name_age'].values],
   ...:     names=['name', 'age'])

In [7]: x
Out[7]:
<xarray.DataArray (name_age: 3, hobby: 2)>
array([[  nan,  450.],
       [ 500.,  100.],
       [  nan,  900.]])
Coordinates:
  * name_age  (name_age) MultiIndex
  - name      (name_age) object 'Jack' 'John' 'John'
  - age       (name_age) object 'teen' 'adult' 'teen'
  * hobby     (hobby) |S6 'books' 'bicyle'

如果您随后取消堆栈 'name_age'，您将获得所需的 3-D DataArray:

In [8]: x.unstack('name_age')
Out[8]:
<xarray.DataArray (hobby: 2, name: 2, age: 2)>
array([[[  nan,   nan],
        [ 500.,   nan]],

       [[  nan,  450.],
        [ 100.,  900.]]])
Coordinates:
  * hobby    (hobby) |S6 'books' 'bicyle'
  * name     (name) object 'Jack' 'John'
  * age      (age) object 'adult' 'teen'