xarray 坐标相关计算

Question

我将 xarray 与我有测量值和错误的数据一起使用。 我将这些沿着维度矩存储在数据集中，坐标值和方差。 例如，当我计算一个维度的平均值时，我需要对值和方差进行不同的处理，因为前者应该被组合为

mean_values = sum(values)/len(values)

但后者作为

mean_variance = sum(variances**2)/len(variances) 。

目前我正在通过形成两个新数据集并将它们连接起来来做到这一点。 这非常难看，令人费解，不适合更复杂的计算。 我希望能够一步完成这种操作，也许通过定义一个以值和方差作为输入的函数，然后将数据集维度时刻广播到它上面。

给定具有维度moment, time, position的数据集q_lp ：

q_lp_av = q_lp.sel(moment='value').mean(dim='time')
q_lp_var = q_lp.sel(moment='variance').reduce(average_of_squares, dim='time')
q_lp = xr.concat([q_lp_common_av, q_lp_common_var], dim='moment')

其中average_of_squares定义为

def average_of_squares(data, axis=None):
    sums = np.sum(data**2, axis=axis)
    if axis:
        return sums/np.shape(data)[axis]**2
    return sums/len(data)**2

• 有什么更好的方法来处理这个问题？
• 是否可以使用xr.apply_ufunc和my_average函数一步就地完成此操作？
• 我不应该将论文放在一个数据集中吗？ q_lp 稍后会与其他量（也具有维度 moment、pos 和 time）组合到 DataSet 中。

我感谢讨论、想法、技巧和示例链接。

编辑：澄清一下，我不喜欢拆分 DataArray，单独处理每个时刻并再次连接它们。 我更愿意执行以下操作（用于说明的未经测试的伪代码）：

def multi_moment_average(mean, variance):
    mean = np.average(mean)
    variance = np.sum(variance**2)/len(variance)
    return mean, variance

q_lp.reduce(multi_moment_average, broadcast='moment', dim='time')

最小工作示例：

import numpy as np
import xarray as xr


def average_of_squares(data, axis=None):
    sums = np.sum(data**2, axis=axis)
    if axis:
        return sums/np.shape(data)[axis]**2
    return sums/len(data)**2


times = np.arange(10)
positions = np.array([1, 3, 5])
values = np.ones((len(times), len(positions))) * (2 + np.random.rand())
variance = np.ones((len(times), len(positions))) * np.random.rand()

q_lp = xr.DataArray(np.array([values, variance]),
                    coords=[['value', 'variance'], times, positions],
                    dims=['moment', 'time', 'position'])

q_lp_av = q_lp.sel(moment='value').mean(dim='time')
q_lp_var = q_lp.sel(moment='variance').reduce(average_of_squares, dim='time')
q_lp = xr.concat([q_lp_av, q_lp_var], dim='moment')

Answer 1

我认为您可以以 xarray 友好的方式编写您的函数，然后在您的数据上调用它。 IE

def average_of_squares(data, dim=None):
    sums = (data ** 2).sum(dim)
    return sums/data.count(dim)**2

q_lp_var = q_lp.sel(moment='variance').pipe(average_of_squares, dim='time')

将它们连接在同一个DataArray中很好； 不过，它可能更适合Dataset上的项目。

这是否回答你的问题？

编辑：重新编辑的问题，我认为将项目保存在 Dataset 而不是 DataArray 中与数据结构最一致。 似乎均值和方差是您想要在相同索引上对齐的两个不同数组，因此 Dataset 是理想的

Answer 2

我找到了一个适合我需求的解决方案，但仍然感谢更多建议：

groupby 可以沿指定的维度分离 Dataset 或 DataArray，其列表创建 (key, value) 元组，其 dict 本质上具有关键字字典的形式。 见http://xarray.pydata.org/en/stable/groupby.html

因此，我当前的解决方案如下所示：

import xarray as xr

def function_applier(data, function, split_dimension=None, **function_kwargs):
    return xr.concat(
                function(
                    **dict(list(data.groupby(split_dimension))),
                    **function_kwargs),
                dim=split_dimension)

现在我可以定义将特定坐标作为输入的函数，这些函数可以编写为也适用于例如 numpy 数组。 （MWE 在此处使用我原始问题的具体示例）

import numpy as np

def average_of_gaussians(val, var, dim=None): 
    return val.mean(dim), (var ** 2).sum(dim)/var.count(dim)

val = np.random.rand(12).reshape(2,6)
var = 0.1*np.random.rand(12).reshape(2,6)

da = xr.DataArray([val, var],
                  dims=['moment','time','position'],
                  coords=[['val','var'],
                          np.arange(6),
                          ['a','b']])

>>>da
<xarray.DataArray (moment: 2, position: 2, time: 6)>
array([[[0.66233728, 0.71419351, 0.96758741, 0.96949021, 0.94594299,
         0.05080628],
        [0.44005458, 0.64616657, 0.69865189, 0.84970553, 0.19561433,
         0.8529829 ]],

       [[0.02209967, 0.02152369, 0.09181031, 0.00223527, 0.01448938,
         0.01484197],
        [0.05651841, 0.04942305, 0.08250529, 0.04258035, 0.00184209,
         0.0957248 ]]])
Coordinates:
  * moment    (moment) <U3 'val' 'var'
  * position  (position) <U1 'a' 'b'
  * time      (time) int32 0 1 2 3 4 5

>>>function_applier(da,
                 average_of_gaussians,
                 split_dimension='moment',
                 dim='time')
<xarray.DataArray (moment: 2, position: 2)>
array([[0.71839295, 0.61386263],
       [0.001636  , 0.00390397]])
Coordinates:
  * position  (position) <U1 'a' 'b'
  * moment    (moment) object 'val' 'var'

注意输入名称等于 average_of_gaussians 的坐标。 对一个函数中每个变量的不同操作以及其中缺少对 xarray 的引用是我所追求的属性。