Extraction des indices d'une entité à partir d'un tableau multidimensionnel

Question

J'ai un tableau représentant les valeurs de la concentration d'eau dans les nuages ​​dans un espace tridimensionnel. Aux endroits où la concentration de l'eau des nuages ​​est au-dessus d'un certain seuil, je dis que j'ai un nuage (voir la section ci-dessous). La majeure partie du domaine est sèche, mais il y a un nuage de stratocumulus sur la majeure partie du domaine avec une base à environ 400 mètres.

Ce que je veux faire est extrait (x, y, z) les coordonnées de la base des nuages ​​et du sommet des nuages. Ensuite, je veux utiliser ces coordonnées sur un tableau tridimensionnel différent représentant la composante verticale de la vitesse du vent pour obtenir le courant ascendant à la base du nuage. Ce que je fais maintenant fonctionne mais il est lent. Je me sens comme il doit y avoir un moyen de tirer parti de NumPy pour l'accélérer.

C'est ce que je fais maintenant:

# 3d array representing cloud water at a particular timestep t 
qc = QC(t) 

# get the coordinates where there is cloud 
cloud_coords = argwhere(qc > qc_thresh) 

# Arrays to hold the z values of cloud base (cb) and cloud top (ct) 
zcb = zeros((nx,ny)) 
zct = zeros((nx,ny)) 

# Since each coordinate (x,y) will in general have multiple z values 
# for cloud I have to loop over all (x,y) and 
# pull out max and min height for each point (x,y) 
for x in range(nx): 
    # Pull out all the coordinates with a given x value 
    xslice = cloud_coords[ where(cloud_coords[:,0] == x) ] 

    for y in range(ny):  
     # for the given x value select a particular y value 
     column = xslice[ where(xslice[:,1] == y) ] 

     try: 
      zcb[x,y] = min(column[:,2]) 
      zct[x,y] = max(column[:,2]) 
     except: 
      # Because there may not be any cloud at all 
      # (a "hole") we fill the array with an average value 
      zcb[x,y] = mean(zcb[zcb.nonzero()]) 
      zct[x,y] = mean(zct[zct.nonzero()]) 


# Because I intend to use these as indices I need them to be ints 
zcb = array(zcb, dtype='int') 
zct = array(zct, dtype='int') 

La sortie est un tableau à deux dimensions contenant les coordonnées z de la base des nuages ​​(et haut)

J'utilise ensuite ces indices sur un autre tableau pour obtenir des variables comme la vitesse du vent à la base du nuage:

wind = W(t) 
j,i = meshgrid(arange(ny),arange(nx)) 
wind_base = wind[i,j,zcb] 

Je le fais pour plusieurs timesteps dans la simulation et la partie la plus lente est la boucle python sur toutes les coordonnées (x, y). Toute aide sur l'utilisation de NumPy pour extraire ces valeurs plus rapidement serait grandement appréciée!

Answer 1

Votre soupçon que numpy peut être mis à bon usage avec votre problème est correct. En fait, il y a plusieurs inefficacités que vous faites, par exemple en créant explicitement un nouveau tableau en utilisant np.array() à la fin, et une avec un dtype de int qui est un objet complexe en python 3.

Vous pouvez faire la plupart des travailler dans quelques lignes vectorisées de numpy. L'idée est qu'il suffit de trouver les indices (le long de l'axe des hauteurs) où un nuage apparaît, ou à la fin d'un nuage. Nous pouvons le faire de façon vectorisée en utilisant numpy.argmax. C'est vraiment le cœur d'une solution efficace:

import numpy as np 
import matplotlib.pyplot as plt 

# generate dummy data 
qc_thresh = 0.6 
nx,ny,nz = 400,400,100 
qc = np.zeros((nx,ny,nz)) 
# insert random cloud layer 
qc[...,50:80] = np.random.rand(nx,ny,30) 
# insert holes in clouds for completeness 
qc[np.random.randint(nx,size=2*nx),np.random.randint(ny,size=2*nx),:] = 0 

def compute_cloud_boundaries(): 
    cloud_arr = qc > qc_thresh 

    # find boundaries by making use of np.argmax returning first maximum 
    zcb = np.argmax(cloud_arr,axis=-1) 
    zct = nz - 1 - np.argmax(cloud_arr[...,::-1],axis=-1) 

    # logical (nx,ny)-shaped array where there's a cloud 
    cloud_inds = (zcb | (zct!=nz-1)).astype(bool) 
    # this is short for `(zcb==0) | (zct!=nz-1)` 

    # fill the rest with the mean 
    zcb[np.logical_not(cloud_inds)] = zcb[cloud_inds].mean() 
    zct[np.logical_not(cloud_inds)] = zct[cloud_inds].mean() 

    return zcb,zct 

j'ai vérifié ci-dessus (complet avec le petit correspondant par exemple) par rapport à votre approche, et il donne exactement le même résultat. Comme je l'ai dit, l'idée est que cloud_arr = qc > qc_thresh est un tableau logique nous indiquant où l'humidité est assez grande pour se qualifier pour un nuage. Ensuite, nous regardons le maximum de cette matrice (essentiellement binaire!) Le long du dernier axe (hauteur). L'appel np.argmax nous dira la première valeur de hauteur (la plus basse) pour chaque index 2d planaire. Afin d'obtenir le sommet des nuages, nous devons inverser notre tableau et faire la même chose de l'autre côté (en prenant soin de transformer les indices qui en résultent). Inverser le tableau crée une vue plutôt qu'une copie, donc c'est aussi efficace. Enfin, nous corrigeons les points où il n'y a pas de nuages; au lieu d'une meilleure contrainte, nous vérifions où l'indice le plus élevé renvoyé par argmax correspond au point de contour. Compte tenu des données météorologiques réalistes, nous pouvons être sûrs que les mesures les plus basses et les plus élevées ne correspondent pas aux nuages, donc ce devrait être un critère de sécurité.

Voici une section transversale des données fictives pour le spectacle:

horaires non représentatifs pour le 400x400x100 cas ci-dessus:

In [24]: %timeit compute_cloud_boundaries() 
10 loops, best of 3: 29.1 ms per loop 

In [25]: %timeit orig() # original loopy version from the question 
1 loop, best of 3: 9.37 s per loop 

qui semble être plus d'un 300 -plus augmentation de la vitesse. Bien sûr, votre cas d'utilisation réel sera le bon test de cette approche, mais ça devrait aller. En ce qui concerne l'étape d'indexation, vous pouvez économiser de la mémoire en utilisant une grille ouverte pour les index et en utilisant la diffusion en réseau. Ne pas avoir à allouer (nx,ny) tableaux supplémentaires pourraient accélérer en forme de cette étape un peu trop:

wind = W(t) 
i,j = np.ogrid[:nx,:ny] 
wind_base = wind[i,j,zcb] 

Comme vous pouvez le voir, np.ogrid crée une grille ouverte de forme (nx,1) et (1,ny) qui diffusent ensemble à ce qui est équivalent à la meshgrid appel.