Найти последнее не NaN-значение вдоль оси в отсортированном многомерном массиве numpy - PullRequest
1 голос
/ 10 марта 2020

Я смотрю на некоторые трехмерные данные о температуре океана (время, глубина, долгота, широта) и хотел бы извлечь значение на самой низкой глубине, чтобы создать двухмерную карту температуры на дне океана.

Дно океана - это маска, которая создает некий отсортированный массив вдоль оси глубины со всеми значениями NaN, сконцентрированными в конце оси 1.

Некоторые примеры кода, чтобы повторить это:

import numpy as np

A=np.random.rand(6,50,300,360)*100
A.ravel()[np.random.choice(A.size, 10000000, replace=False)] = np.nan
A.sort(axis=1)

Затем, следуя Получая последний не-нановый индекс отсортированной матрицы numpy или pandas dataframe , я получаю массив, который содержит индекс конечного не-NaN элемент вдоль оси 1:

lv=(~np.isnan(A)).sum(axis=1)-1

Теперь сложная часть извлекает значения из оси 1 из A , используя lv (массив элементов, которые я хочу извлечь ). Пока что мой лучший метод (который работает) - создать пустой массив соответствующего размера и заполнить его поэлементно:

B=np.zeros(lv.shape,dtype=np.float32)
for i in range(t):
    for j in range(y):
        for k in range(x):
            B[i,j,k]=A[i,lv[i,j,k],j,k]

Однако это довольно медленно; неоправданно так для объема данных, на котором я собираюсь использовать это (много ТБ).

Любые идеи о том, как упростить этот последний этап (как в Pandas найти последнее значение не NAN , а для numpy)? Я думаю что-то вроде (хотя я понимаю, что это даже не имеет смысла):

B=A[:,lv[:],:,:]

Я также пробовал варианты np.take, np.take_along_axis и np.choose безуспешно.

Заранее спасибо за любые предложения!

1 Ответ

0 голосов
/ 10 марта 2020

Numpy s take_along_axis должно работать для этого. Последний шаг можно выразить следующим образом:

B = np.take_along_axis(A, lv[:,None,:,:], axis=1).squeeze()
...