Usind 3d array to index 4d array - PullRequest
Новая
       15

Usind 3d array to index 4d array

2 голосов
/ 05 июня 2019

У меня есть 4d-xarray с размерами I[t, z, y, x] и 3d-xarray с размерами Z[t, y, x].Z - это индексы, которые мне нужны для измерения z в I.Теперь я хочу получить значения I[t, Z[t,y,x], y, x] и записать их в новый xarray размером O[t, y, x].

В принципе, я решил проблему с помощью for-циклов, но это слишком медленно (многои большие массивы).Есть ли способ сделать это без for-loop?

Пример кода делает то, что я хочу, но медленно: 

def get_field_at_levels(array, levels):
    shape = array.shape
    array_out = np.zeros_like(levels)
    for t in range(shape[0]):
        for x in range(shape[2]):
            for y in range(shape[3]):
                if np.isnan(levels[t, x, y]):
                    array_out[t, x, y]==float('nan')
                else:
                    array_out[t, x, y] = array[t, int(levels[t, x, y]), x, y]
    return array_out

1 Ответ

0 голосов
/ 06 июня 2019

Я преобразовал все свои xarrays в numpy.arrays, используя numpy.asarray.Теперь программа достаточно быстрая с использованием циклов.Ниже приведен пример сценария с использованием случайных чисел.В моих фактических данных у меня есть индексы вне диапазона (-1).В этом случае я хочу NaN как результат. 

import numpy as np
import time

tsize = 1
xsize = 40
ysize = 240
zsize = 260

def val_at_lev(data, Ind):
    sh=data.shape
    data2=np.empty([sh[0],sh[1]+1,sh[2],sh[3]])
    data2[:,0:sh[1],:,:]=data
    data2[:,sh[1],:,:]='nan'
    out=np.asarray(np.zeros_like(Ind))
    erg=np.asarray([data2[t,Ind[t,0,j,k],j,k] for t in range(sh[0]) for j in range(sh[2]) for k in range(sh[3])])
    out = erg.reshape(tsize,1,ysize,zsize)
    return out

# Main program
Ind=np.random.randint(-1,xsize,[tsize,1,ysize,zsize])
data=np.random.uniform(0,100,[tsize,xsize,ysize,zsize])
start_time = time.time()
erg=val_at_lev(data,Ind)
print("--- %s seconds ---" % (time.time() - start_time))
...