假设我有一个这样的数据集
nx1, nx2, nx3 = 5, 3, 20
ds = xray.Dataset()
ds.coords.update({'x1': ('x1', range(nx1)),
'x2': ('x2', range(nx2)),
'x3': ('x3', range(nx3))})
ds['A'] = (['x1', 'x2', 'x3'], np.random.randn(nx1, nx2, nx3))
ds['B'] = (['x1', 'x2', 'x3'], np.random.randn(nx1, nx2, nx3))
以及一个函数func
,它接受输入变量 A 和 B,但它沿x3
维工作,只接受形状 (nx3,)
的数组,并输出形状(nx3,)
数组。 如果我想将此函数应用于上述数据集并将结果保存为名为 C 的新变量,请遵循以下方法,
required_shape = (len(ds.coords['x1']),
len(ds.coords['x2']),
len(ds.coords['x3']))
if 'C' not in ds:
ds.update({'C': (['x1', 'x2', 'x3'], np.zeros(required_shape))})
for ix1, x1 in enumerate(ds.coords['x1']):
for ix2, x2 in enumerate(ds.coords['x2']):
args = dict(x1=ix1, x2=ix2)
a = ds['A'][args]
b = ds['B'][args]
c = func(a.values, b.values)
ds['C'][args] = c
通过在数据集中初始化一个新数组并在其他维度上使用 for 循环?
我对 pandas 不感兴趣,但其他数据类型的一般解决方案也是使用推导,并摆脱嵌套循环和初始化步骤。
required_shape = (len(ds.coords['x1']),
len(ds.coords['x2']),
len(ds.coords['x3']))
ds['C'] = (['x1', 'x2', 'x3'], np.array([
func(ds['A'][args].values, ds['B'][args].values)
for ix1, x1 in enumerate(ds.coords['x1'])
for ix2, x2 in enumerate(ds.coords['x2'])
for args in (dict(x1=ix1, x2=ix2),)]).reshape(required_shape))
编辑:顺便说一下,
ds['C'] = func(ds['A'], ds['B'])
对于一个简单的函数,似乎工作得很好,例如:
def func(a, b):
return a + b