xarray 坐标相关计算

Question

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

mean_values = sum(values)/len(values)

但后者为

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

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

给定一个数据集 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 友好的方式编写函数，然后在您的数据上调用它。即

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 与数据结构最一致。似乎均值和方差是您希望在相同索引上对齐的两个不同数组，因此数据集是理想的

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 的引用是我所追求的属性。