我想计算xarray.Dataset
中共享维度上两个变量的外积。
具体来说:假设我有以下形式的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 ...
现在我想在每个时间点计算outer product和vec1
的vec2
关于指数i1
,它是矩阵outer[i1,i2] = vec1[i1] * vec2[i2]
。
使用Numpy,可以计算如下:
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)})
现在,result['outer']
包含沿着vec1
维度的vec2
和i1
所需的外积。 (使用i2
作为第二个索引的原因是xarray
不能很好地处理重复维度 - 实际上,这可能是合理的,尽管它使矩阵值数据的处理更加令人厌烦。)
问题:有没有一种方便的方法来计算这样的外部产品,使用xarray的功能(理想情况下应该与dask数组兼容并且能够并行执行计算),而不是回退到在值上调用numpy函数?
当xarray基于其维度名称广播数组时,此外部产品可以按如下方式计算:
In [2]: ds['vec1'] * ds['vec2'].rename(i1='i2')
这相当于以下numpy脚本
v1 = ds['vec1'].values
v2 = ds['vec2'].values
v1[:, :, np.newaxis] * v2[:, np.newaxis]