Comment optimiser une boucle numpy qui somme des valeurs d'un tableau qui est indexé par un autre tableau où les valeurs sont égales à l'index de boucle

Question

J'ai ce morceau de code qui est appelé plusieurs fois pendant l'exécution de l'application. Il prend un tableau de nombres qui représentent des valeurs (tableau_valeurs). Elles doivent être additionnées dans les zones définies dans zone_array. Zone_ids représente une liste de toutes les zones possibles dans zone_array.

Son essentiellement quelque chose dans les lignes de: J'ai une carte raster de population et je veux savoir combien de personnes vivent dans chaque zone de la carte de la zone.

le code:

values = np.zeros(len(zone_ids)) 
for i in zone_ids: 
    values[i] = round(np.nansum(value_array[zone_array == i]), 2) 
return values 

Le coupable semble être la boucle, mais je n'ai pas trouvé un moyen d'éliminer et avoir les mêmes résultats.

Je l'ai essayé avec bincount mais je n'ai pas réussi. L'utilisation de numba jit n'a également aucun effet. Je voudrais rester loin de Cython car ce code sera utilisé dans un plugin Qgis qui n'a pas de support Cython.

code de test:

import numpy as np 


def fill_values(zone_array, value_array, zone_ids): 
    values = np.zeros(len(zone_ids)) 
    for i in zone_ids: 
     values[i] = round(np.nansum(value_array[zone_array == i]), 2) 
    return values 


def run(): 
    # 300 different zones 
    zone_ids = range(300) 
    # zone map with 300 zones 
    zone_array = (np.random.rand(2000, 2000) * 300).astype(int) 
    # value map from which we want the sum of values per zone (real map can have NaN values) 
    value_array = (np.random.rand(2000, 2000) * 10.) 
    value_array[5, 5] = np.NAN 
    fill_values(zone_array, value_array, zone_ids) 


if __name__ == '__main__': 
    run() 

1,92 s ± 17,5 ms par boucle (moyenne ± std dev 7 pistes, 1 en boucle chacun..)

Avec la mise en oeuvre de bincount comme suggéré par Divakar:

203 ms ± 15,2 ms par boucle (moyenne ± std. dev. 7 pistes, 1 boucle chacun)

Answer 1

Avec une utilisation directe de bincount, vous auriez NaNs dans les sommations. Ainsi, vous pouvez simplement remplacer le NaNs par zeros et utiliser bincount. Cela devrait être beaucoup plus rapide, étant une solution vectorisée.

Par conséquent, la mise en œuvre serait -

val_nonan = np.where(np.isnan(value_array), 0, value_array) 
out = np.round(np.bincount(zone_array.ravel(), val_nonan.ravel()),2)