Многомерное логическое индексирование XArray - PullRequest
Новая
       10

Многомерное логическое индексирование XArray

0 голосов
/ 21 сентября 2019

Рассмотрим XArray DataArray: 

DA=xr.DataArray(np.random.rand(3,4,5))

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

mask=DA.mean(dim='dim_1')>.5 (маска 3x5)

Теперь я хотел бы извлечь все элементы в маске в простой вектор, сделать что-то с этим вектором и вставить результат обратно в массив данных.

DA[mask],что кажется интуитивно понятным, return IndexError: 2-dimensional boolean indexing is not supported.

xr.where позволяет установить для невыбранных элементов значение nan, но это не то же самое, что извлечение вектора.

numpy многомерное логическое индексирование (DA.values[mask.values]) здесь не сработает, потому что размеры не совпадают.

Единственное решение, которое я нашел, - это расширить маску таким неуклюжим образом: 

def multidim_masked_fun(DA,mask,fun):
    expanded_mask=DA.where(xr.ufuncs.logical_not(mask),'masked')=='masked'
    DB=DA.copy()
    DB.values[expanded_mask.values]=fun(DA.values[expanded_mask.values])
    return DB

(Iиспользовал строку «маскированный» вместо «нан», чтобы избежать взаимодействия с потенциальными значениями «нан», существующими в DA до маскирования)

Есть ли более естественный способ добиться этого?

...