如何将 xarray DataArray 与长度为 1 的维度与更大的数组对齐？

Question

我想获取一个时间维度为 1 的 xarray 数据集，并简单地复制数据以将时间维度从 1 增加到 N。最有效的方法是什么？我尝试了几种方法，例如 expand_dims 和 stack，但这些方法似乎都没有达到我想要的效果。

最终我希望能够做到 moc10_H11 - moc_ctrl_clim 结果将具有与 moc10_H11 (35) 相同的尺寸。现在，当我这样做时，输出的时间维度只有 1。

为了清楚起见， moc_ctrl_clim ：

Dimensions:
time: 1, lat_aux_grid: 395, moc_z: 61
Coordinates: time (time) object 0001-01-01 00:00:00
lat_aux_grid (lat_aux_grid) float32 -79.49 -78.95 -78.42 ... 89.47 90.0
moc_z (moc_z) float32 0.0 1e+03 ... 5.25e+05 5.5e+05
Data variables:
MOC (time, moc_z, lat_aux_grid) float64
dask.array<chunksize=(1, 61, 395), meta=np.ndarray>

而 moc10_H11 有：

Dimensions:

time: 35, lat_aux_grid: 395, moc_z: 61
Coordinates: time (time) object 0001-01-01 00:00:00
lat_aux_grid (lat_aux_grid) float32 -79.49 -78.95 -78.42 ... 89.47 90.0
moc_z (moc_z) float32 0.0 1e+03 ... 5.25e+05 5.5e+05
Data variables:
MOC (time, moc_z, lat_aux_grid) float64
dask.array<chunksize=(1, 61, 395), meta=np.ndarray>

Answer 1

简短的回答，压缩数据以便 xarray 的自动对齐规则生效：

da = da.squeeze(dim='time', drop=True)

现在，您可以与按时间索引的数组配对，数据将自动广播。

说明

这背后的原因在于基于形状的numpy's broadcasting和xarray's broadcasting by dimension name之间的区别。

Numpy 按形状广播

来自numpy docs：

当对两个数组进行操作时，NumPy 会逐元素比较它们的形状。它从尾随（即最右边）尺寸开始，然后向左工作。 两个维度兼容时

1. 它们是相等的，或者
2. 其中一个是 1

例如，如果第一个维度对齐，您可以在列向量和数组之间执行元素相加：

In [3]: col_vector = np.ones(shape=(3, 1))

In [4]: col_vector
Out[4]:
array([[1.],
       [1.],
       [1.]])

In [5]: array = np.arange(12).reshape(3, 4)

In [6]: array
Out[6]:
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

In [7]: col_vector + array
Out[7]:
array([[ 1.,  2.,  3.,  4.],
       [ 5.,  6.,  7.,  8.],
       [ 9., 10., 11., 12.]])

当col_vector 添加到array 时，numpy 识别出col_vector 沿轴1 的长度为1，而array 的长度为4，因此col_vector 应沿轴1 广播（平铺）以具有添加前的长度为 4。

xarray 按维度名称广播

来自xarray docs on computation：

DataArray 对象通过维度名称而不是轴顺序自动对齐自身（numpy 术语中的“广播”）。使用 xarray，您不需要转置数组或插入长度为 1 的维度来使数组操作正常工作，这在 numpy 中通常使用numpy.reshape() 或numpy.newaxis 完成。

除此之外，xarray docs on automatic alignment:

Xarray 强制在二元运算中使用的对象的 index Coordinates（即与维度同名的坐标，用 * 标记）之间对齐。 [...] 如果任一参数中缺少维度的坐标值，则所有匹配的维度必须具有相同的大小。

改编上面的例子不仅需要指定名称和坐标维度，还需要从列向量中删除第二个维度：

In [2]: vector = xr.DataArray(np.ones(shape=3), dims=['x'], coords=[[0, 1, 2]])

In [3]: vector
Out[3]:
<xarray.DataArray (x: 3)>
array([1., 1., 1.])
Coordinates:
  * x        (x) int64 0 1 2

In [4]: arr = xr.DataArray(
   ...:     np.arange(12).reshape(3, 4),
   ...:     dims=['x', 'time'],
   ...:     coords=[[0, 1, 2], pd.date_range('2020-01-01', periods=4, freq='D')],
   ...: )

In [5]: arr
Out[5]:
<xarray.DataArray (x: 3, time: 4)>
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
Coordinates:
  * x        (x) int64 0 1 2
  * time     (time) datetime64[ns] 2020-01-01 2020-01-02 2020-01-03 2020-01-04

In [6]: vector + arr
Out[6]:
<xarray.DataArray (x: 3, time: 4)>
array([[ 1.,  2.,  3.,  4.],
       [ 5.,  6.,  7.,  8.],
       [ 9., 10., 11., 12.]])
Coordinates:
  * x        (x) int64 0 1 2
  * time     (time) datetime64[ns] 2020-01-01 2020-01-02 2020-01-03 2020-01-04

将 length-1 维度广播到更长的维度

在您的问题中，您有一个沿时间维度长度为 1 的数组，您想针对另一个具有较长时间坐标的数组进行广播。在上面的示例中，这等效于在时间维度上具有长度为 1 的“向量”：

In [7]: vector = xr.DataArray(
   ...:     np.ones(shape=(3, 1)),
   ...:     dims=['x', 'time'],
   ...:     coords=[[0, 1, 2], pd.date_range('2020-01-01', periods=1, freq='D')],
   ...: )

当针对arr（时间维度长度为4）进行广播时，只保留交集：

In [8]: vector + arr
Out[8]:
<xarray.DataArray (x: 3, time: 1)>
array([[1.],
       [5.],
       [9.]])
Coordinates:
  * time     (time) datetime64[ns] 2020-01-01
  * x        (x) int64 0 1 2

数据可以通过先用da.squeeze压缩和丢掉时间dim来按时间广播：

In [9]: vector.squeeze('time', drop=True)  + arr
Out[9]:
<xarray.DataArray (x: 3, time: 4)>
array([[ 1.,  2.,  3.,  4.],
       [ 5.,  6.,  7.,  8.],
       [ 9., 10., 11., 12.]])
Coordinates:
  * x        (x) int64 0 1 2
  * time     (time) datetime64[ns] 2020-01-01 2020-01-02 2020-01-03 2020-01-04

请注意，这种方法会忽略time 坐标中第一个数组中的信息，而是假设该信息适用于第二个数组中time 的所有元素。如果这是您正在寻找的内容，那么如图所示的挤压和放下就是要走的路。