如何将结构化numpy数组转换为指定结构的xarray.Dataset？

我有一个结构化numpy数组，包含多个测量序列的采样数据：每个序列以l为自变量对m进行采样，不同序列由参数a区分。l的采样值不固定，且每个序列的样本数不同，无法直接生成m和l的二维数组。示例数据如下：

In [1]: data
Out[1]:array([(   0., 1323.,  69384.), (   0., 1344.,  73674.), (   0., 1344.,  73674.),
       (   0., 1439.,  76678.), (   0., 1538.,  79584.), (   0., 1643.,  82389.),
       (   0., 2382.,  95634.), (   0., 2439.,  96028.), (   0., 2439.,  96028.),
       (   0., 2574.,  98154.), (   0., 2795.,  99937.), (1219., 1316.,  59055.),
       (1219., 1332.,  61473.), (1219., 1350.,  63881.), (1219., 1372.,  66270.),
       (1219., 1372.,  66270.), (1219., 1491.,  69654.), (1219., 1617.,  72917.),
       (1219., 1749.,  76053.), (1219., 1885.,  79060.), (1219., 2028.,  81927.),
       (1219., 2072.,  82803.), (1219., 2118.,  83606.), (1219., 2166.,  84340.),
       (1219., 2846.,  91028.), (1219., 2911.,  91379.), (1219., 2977.,  91635.),
       (1219., 4164.,  95161.), (2438., 1313.,  52688.), (2438., 1331.,  54496.),
       (2438., 1350.,  56304.), (2438., 1368.,  58113.), (2438., 1480.,  60990.),
       (2438., 1598.,  63754.), (2438., 1720.,  66399.), (2438., 1846.,  68926.),
       (2438., 1978.,  71326.), (2438., 2757.,  79713.), (2438., 2819.,  80026.),
       (2438., 2882.,  80258.), (2438., 4155.,  84968.)],
      dtype=[('a', '<f8'), ('l', '<f8'), ('m', '<f8')])

我想要转换成xarray.Dataset，包含a和i两个坐标，其中i是每个测量序列内数据点的整数索引，这样能通过索引获取每个a对应的测量数据（比如选i=0获取第一个采样样本）。期望结果如下：

In [100]: desired_array
Out[100]: <xarray.Dataset>
         Dimensions:  (a: 3, i: 17)
         Coordinates:
           * a  (a) float64 0. 1219. 2438.
           * i  (i) int64 0 1 2 3 4 ... 15 16 17
         Data variables:
             l        (a, i) float64 1323. 1344. 1344. 1439. ... 2882. 4155.
             m        (a, i) float64 69384. 73674. 73674. ... 80258. 84968.

直接创建xarray会失败，因为需要二维数组输入，而data['m']是一维的。手动处理太繁琐，尝试转成Pandas DataFrame再转xarray的方法也没得到预期结果，得到的是一维结构，索引重复且无组内二级索引。

可以通过Pandas给每个a分组添加组内索引，再转成xarray，步骤如下：

1. 将结构化numpy数组转为Pandas DataFrame
   不需要手动指定索引，直接转换即可：

   import pandas as pd
   import xarray as xr
   
   df = pd.DataFrame(data)
   

2. 添加组内索引i
   按a分组，给每组内的行分配从0开始的整数索引：

   df['i'] = df.groupby('a').cumcount()
   

3. 转换为宽格式
   通过pivot方法将数据转为宽格式，不同长度的组会自动用NaN填充缺失值：

   df_wide = df.pivot(index='i', columns='a', values=['l', 'm'])
   # 调整维度顺序，让a成为第一维度
   df_wide = df_wide.swaplevel(axis=1).sort_index(axis=1)
   

4. 转换为xarray Dataset
   直接调用to_xarray()即可得到预期的结构：

   desired_ds = df_wide.to_xarray()
   

验证结果：

print(desired_ds)

输出会和你期望的结构一致，维度为(a:3, i:17)，可以通过desired_ds.sel(a=0., i=0)获取对应的数据点。

内容的提问来源于stack exchange，提问作者Zak