Вычисление внешнего произведения двух переменных в наборе данных xarray - PullRequest
Новая
       7

Вычисление внешнего произведения двух переменных в наборе данных xarray

0 голосов
/ 03 июля 2018

Я хотел бы вычислить внешнее произведение двух переменных по общему измерению в пределах 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 и 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 относительно индекса i1, который является матрицей outer[i1,i2] = vec1[i1] * vec2[i2].

Используя Numpy, это можно рассчитать следующим образом: 

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

Вопрос: Существует ли удобный способ вычисления внешнего продукта, подобного этому, с использованием функциональности из xarray (которая в идеале должна быть совместима с массивами dask и способной выполнять вычисления параллельно), не возвращаясь к вызову NumPy функции на значения?

1 Ответ

0 голосов
/ 03 июля 2018

Поскольку xarray передает массивы на основе имен их измерений, этот внешний продукт может быть вычислен следующим образом: 

In [2]: ds['vec1'] * ds['vec2'].rename(i1='i2')

Это эквивалентно следующему скриптовому сценарию 

v1 = ds['vec1'].values
v2 = ds['vec2'].values    
v1[:, :, np.newaxis] * v2[:, np.newaxis]
...