Question

我想沿着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和{{1 }}，具体取决于这些尺寸。数据集如下所示：

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 ...计算vec1和vec2的{{3}}，它是矩阵{ {1}}。

使用Numpy，可以如下计算：

i1

现在outer[i1,i2] = vec1[i1] * vec2[i2]沿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不能很好地处理重复的维度-实际上，这可能是合理的，尽管它使得处理矩阵值的数据更加麻烦。）

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

Answer 1

由于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]

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

1 个答案: