从数据集中计算新变量并使用新维度保存它们

Question

我有一个数据集 ds，其中包含变量 Var_1、Var_2 和 Var_3，它们都是物理量、温度、密度和压力。

<xarray.Dataset>
Dimensions:                         (Altitude: 100)
Coordinates:
  * Altitude                        (Altitude) float64
Data variables:
    Var_1                           (Altitude) float64
    Var_2                           (Altitude) float64
    Var_3                           (Altitude) float64

我计算了一个新变量 "Heat"，我想沿着维度 "Altitude" 和具有两个条目的新维度 "Energy_Type" 添加到我的数据集。无法将已经存在的变量分配给这个新维度。

<xarray.Dataset>
Dimensions:                         (Altitude: 100, Energy_Type: 2 )
Coordinates:
  * Altitude                        (Altitude) float64
  * New_Dim                         (Energy_Type) object 'Sensible' 'Latent'
Data variables:
    Var_1                           (Altitude) float64
    Var_2                           (Altitude) float64
    Var_3                           (Altitude) float64
    Heat                            (Altitude, Energy_Type) float64

我在计算中还使用了另外两个不需要保存的 xarray 数据集。

sensible = xr.Dataset()
sensible['Heat'] = ds.Var_1 * Dataset_X * ds.Var_2
sensible.coords['Energy_Type'] = "Sensible"
sensible = sensible.expand_dims('Energy_Type')

latent = xr.Dataset()
latent['Heat'] = ds.Var_1 * Dataset_Y * ds.Var_3
latent.coords['Energy_Type'] = "Latent"
latent = latent.expand_dims('Energy_Type')

ds = xr.merge([ds, sensible, latent])

虽然这很有效，但它对我来说似乎不是最佳方法。我是否缺少一个简单的 xarray 功能可以帮助我做得更好？

Answer 1

在我看来，您实际上是在将 Var_2 和 Var_3 视为沿 new_dim 的切片。你可以沿着这个暗淡的方向连接它们：

In [3]: reindexed_2_3 = xr.concat(
   ...:     [ds.Var_2, ds.Var_3],
   ...:     dim=pd.Index(['a', 'b'], name='New_Dim'))
   ...:

这使您可以像往常一样对组合数组进行数学运算：

In [4]: new_var = ds.Var_1 * reindexed_2_3

您当然可以将其添加回原始数据集中：

In [5]: ds['New_Var'] = new_var

In [6]: ds
Out[6]:
<xarray.Dataset>
Dimensions:   (Altitude: 100, New_Dim: 2)
Coordinates:
  * Altitude  (Altitude) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 ...
  * New_Dim   (New_Dim) object 'a' 'b'
Data variables:
    Var_1     (Altitude) float64 0.4059 0.4675 0.28 0.7884 0.523 0.6433 ...
    Var_2     (Altitude) float64 0.7325 0.4088 0.7504 0.2138 0.8064 0.4313 ...
    Var_3     (Altitude) float64 0.5138 0.6193 0.5576 0.6782 0.9428 0.7201 ...
    New_Var   (Altitude, New_Dim) float64 0.2973 0.2086 0.1911 0.2895 0.2101 ...

尽管如果连接的对象对其他操作更有帮助，您可以使用坐标集创建一个新数据集：

In [7]: xr.Dataset({'reindexed_2_3': reindexed_2_3, 'new_var': new_var})
Out[7]:
<xarray.Dataset>
Dimensions:        (Altitude: 100, New_Dim: 2)
Coordinates:
  * Altitude       (Altitude) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...
  * New_Dim        (New_Dim) object 'a' 'b'
Data variables:
    reindexed_2_3  (New_Dim, Altitude) float64 0.7325 0.4088 0.7504 0.2138 ...
    new_var        (Altitude, New_Dim) float64 0.2973 0.2086 0.1911 0.2895 ...

Answer 2

在@delgadom () 的回答的帮助下，我找到了一个令我满意的解决方案，并将post 放在这里。

sensible = ds.Var_1 * Dataset_X * ds.Var_2
latent = ds.Var_1 * Dataset_Y * ds.Var_3
Heat = xr.concat(
    [sensible, latent],
    dim=pd.Index(['sensible','latent'],name='Energy_type'))
ds['Heat'] = Heat

有了这个，代码比以前更短，更容易理解。

我也在http://xarray.pydata.org/en/stable/combining.html#concatenate找到了答案，但是之前看不懂

从数据集中计算新变量并使用新维度保存它们

Calculate new variables from a dataset and save them with a new dimension

python-xarray