Python 计算xarray数据集中两个变量的外积

我想计算两个变量沿一个共享维度的外积

具体来说：假设我有一个如下形式的xarray数据集：

import numpy as np
import xarray as xr

ts = np.linspace(0, 1, 100)
indices = range(10)
vecs1 = np.random.rand(len(ts), 10)
vecs2 = np.random.rand(len(ts), 10)

ds = xr.Dataset({'vec1': (['time', 'i1'], vecs1),
                 'vec2': (['time', 'i1'], vecs2)},
                coords={'time': ts,
                        'i1': indices,
                        'i2': indices})

在这个数据集中有一个时间维度，有100个时间步长，一个索引

i1=0，1，…，9

（

i2

是相同的，在一秒钟内就会变得相关）和变量

vec1

和

vec2

，取决于这些维度。数据集如下所示：

>>> ds
<xarray.Dataset>
Dimensions:  (i1: 10, i2: 10, time: 100)
Coordinates:
  * time     (time) float64 0.0 0.0101 0.0202 0.0303 0.0404 0.05051 0.06061 ...
  * i1       (i1) int64 0 1 2 3 4 5 6 7 8 9
  * i2       (i2) int64 0 1 2 3 4 5 6 7 8 9
Data variables:
    vec1     (time, i1) float64 0.2531 0.9019 0.2351 0.3897 0.8144 0.9502 ...
    vec2     (time, i1) float64 0.4962 0.05394 0.1622 0.6937 0.6703 0.5646 ...

现在

result['outer']

包含

vec1

和

vec2

沿

i1

维度的所需外积。（之所以将

i2

用作第二个索引，是因为

xarray

不能很好地处理重复的维度-这可能是合理的，实际上，尽管它使处理矩阵值数据变得更加复杂。）

---

问题：是否有一种方便的方法可以使用xarray的功能（理想情况下应与dask阵列兼容，并能够并行执行计算）来计算这样的外积，在不返回到对值调用numpy函数的情况下？

由于xarray根据数组的维度名称广播数组，此外部乘积可以按如下方式计算：

v1 = ds['vec1'].values
v2 = ds['vec2'].values

# Compute the outer product along the last axis, i.e., separately for each time step,
# giving outer[:, i1, i2] == vec1[:, i1] * vec2[:, i2] for all i1, i2.
outer = np.einsum("...i,...j->...ij", v1, v2)
# Now outer.shape == (100, 10, 10)

result = ds.merge({'outer': (['time', 'i1', 'i2'], outer)})

In [2]: ds['vec1'] * ds['vec2'].rename(i1='i2')

这相当于以下numpy脚本

v1 = ds['vec1'].values
v2 = ds['vec2'].values    
v1[:, :, np.newaxis] * v2[:, np.newaxis]

---

由于xarray根据数组的维度名称广播数组，因此可以按如下方式计算此外积：

v1 = ds['vec1'].values
v2 = ds['vec2'].values

# Compute the outer product along the last axis, i.e., separately for each time step,
# giving outer[:, i1, i2] == vec1[:, i1] * vec2[:, i2] for all i1, i2.
outer = np.einsum("...i,...j->...ij", v1, v2)
# Now outer.shape == (100, 10, 10)

result = ds.merge({'outer': (['time', 'i1', 'i2'], outer)})

In [2]: ds['vec1'] * ds['vec2'].rename(i1='i2')

这相当于以下numpy脚本

v1 = ds['vec1'].values
v2 = ds['vec2'].values    
v1[:, :, np.newaxis] * v2[:, np.newaxis]