xarray координатно-зависимое вычисление

Question

Я использую xarray с данными, для которых у меня есть измерения и ошибки. Я храню их по измерению момент в наборе данных с координатами значение и дисперсия . Когда я вычисляю, например, среднее значение по измерению, мне нужно, чтобы значения и дисперсии обрабатывались по-разному, поскольку первые должны быть объединены как

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')

, где medium_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_ufun c и my_average , чтобы сделать это за один шаг и на месте?
• Должен ли я вообще не собирать тезисы в один набор данных? Позже q_lp объединяется с другими величинами, также с моментом измерений, положением и временем, в DataSet.

Я благодарен за обсуждение, идеи, советы и ссылки на примеры.

Edit: чтобы уточнить, мне не нравится разбивать 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, однако.

Это отвечает на ваш вопрос?

Редактировать: повторно отредактированный вопрос, я думаю, удерживая элементы в наборе данных а не DataArray наиболее согласован со структурами данных. Кажется, что среднее значение и дисперсия - это два разных массива, которые вы хотите выровнять по одним и тем же индексам, поэтому набор данных идеален

Answer 2

Я нашел решение, которое удовлетворяет моим потребностям, но я все еще благодарен за дополнительные предложения:

groupby может разделить набор данных или массив данных по указанному измерению, список которых создает (ключ, значение) кортежи и слова это имеет по существу форму словаря ключевых слов. См. 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)

Теперь я могу определить функции, принимая указанные c координаты в качестве входных данных, которые можно записать также работать, например, для numpy массивов. (MWE, используя конкретный c пример моего первоначального вопроса здесь)

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'

Обратите внимание на входные имена, равные координатам для middle_of_gaussians. Различная операция над каждой переменной в одной функции и отсутствие ссылок на xarray в ней - это свойства, которые мне нужны.