Cette boucle dans une fonction numba peut-elle être optimisée pour s'exécuter plus rapidement?

Il semble que vous traitez de nombreux doublons dans votre exemple.

Dans cette version, je ne recalcule aucune valeur pour un «n» que nous avons déjà vu.

Je ne sais pas si c'est correct dans votre cas ou non, mais cela me fait gagner environ 0,4 seconde.

#!/usr/bin/env python


import numba as nb
import numpy as np
import time

temp = np.ones(25000000)

@nb.jit(fastmath=True)
def func(temp):
    return [n for n in temp]

time1 = time.time()
result = func(temp)
print("Took: ", time.time()-time1, " seconds for ", len(temp), " items")

Original: Hashtable / p>

Une boucle simple de 25 millions d'éléments utilisant np.ones:

#!/usr/bin/env python


import numba as nb
import numpy as np
import time
#######create some synthetic data for illustration purpose##################
size = 5000
eps = 0.2
theta_c = 0.4
temp = np.ones(size)
neighbour = np.random.randint(size, size=(size, 3))
coschi = np.random.random_sample((size))
theta = np.random.random_sample((size))*np.pi/2
pwr = np.cos(theta)
###################end of dummy data##########################

###################-----main loop------###############
@nb.jit(fastmath=True)
def func(theta, pwr, neighbour, coschi, temp):

    hashtable = {}

    for k in range(np.argmax(pwr), 5000*(pwr.size)):
        n = k % pwr.size

        if not hashtable.get(n, False):
            hashtable[n] = 1

            #taking into account regions with different super wind direction
            if (np.abs(theta[n]-np.pi/2.) < np.abs(theta_c)):
                    adj = neighbour[n, 1]
            else:
                    adj = neighbour[n, 0]
            psi_diff = np.abs(np.arccos(coschi[adj])-np.arccos(coschi[n]))
            temp5 = temp[adj]**5
            e_temp = 1. - np.exp(-temp5*psi_diff/np.abs(eps))
            retval = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)

            temp[n] = retval


    return temp


#check time
time1 = time.time()
result = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, " 

Numba est vraiment génial. Mais vous êtes désespéré, n'oubliez pas que vous pouvez toujours écrire en C (youtube). Sur mon propre problème, j'ai obtenu un gain de performance de 30% par rapport à numba simplement en traduisant ligne par ligne en C.

Si vous souhaitez dépenser cet effort, je vous suggère d'utiliser eigen pour les opérations vectorielles (avec une taille de vecteur connue au moment de la compilation) et pybind11 car il se traduit de manière native entre numpy et eigen. Gardez votre boucle principale en Python, bien sûr. Assurez-vous d'utiliser les indicateurs de compilation appropriés (comme -O3 -march = native , -mtune = native , -ffast-math ) et essayez différents compilateurs (pour moi, la sortie de gcc était deux fois plus rapide que celle de clang , mais des collègues ont signalé le contraire).

Si vous ne connaissez aucun C ++, il peut être plus judicieux de vous limiter au C pur et à aucune bibliothèque (car cela réduit la complexité). Mais vous allez traiter directement les API Python et numpy C (pas si compliqué, mais beaucoup plus de code, et vous apprendrez tout sur les composants internes de Python).

Permettez-moi de commencer par quelques commentaires généraux:

• Si vous utilisez numba et que vous vous souciez vraiment des performances, vous devez éviter toute possibilité que numba crée du code en mode objet. Cela signifie que vous devez utiliser numba.njit (...) ou numba.jit (nopython = True, ...) au lieu de numba.jit (. ..) .

  Cela ne fait aucune différence dans votre cas, mais cela rend l'intention plus claire et lance des exceptions dès que quelque chose n'est pas pris en charge en mode nopython (rapide).

• Vous devez faire attention à ce que vous chronométrez et comment. Le premier appel à une fonction numba-jitted (qui n'est pas compilée à l'avance) inclura le coût de compilation. Vous devez donc l'exécuter une fois avant de le chronométrer pour obtenir des horaires précis. Et pour des horaires plus précis, vous devez appeler la fonction plus d'une fois. J'aime IPythons % timeit dans Jupyter Notebooks / Lab pour avoir une idée approximative des performances.

  Je vais donc utiliser:

  499 ms ± 4.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func4 with fastmath on func4_inner 
  

  De cette façon, j'utilise le premier appel (qui inclut le temps de compilation) avec une assertion pour m'assurer qu'il produit vraiment (presque) la même sortie, puis chronométrer la fonction en utilisant une méthode de synchronisation plus robuste (par rapport au temps ).

Levage de np.arccos

Commençons maintenant par quelques optimisations de performances réelles: une évidence est que vous pouvez hisser certains "invariants", par exemple le np.arccos (coschi [...]) est calculé beaucoup plus souvent qu'il n'y a d'éléments réels dans le coschi . Vous parcourez chaque élément de coschi environ 5000 fois et cela fait deux np.arccos par boucle! Calculons donc une fois le arccos du coschi et stockons-le dans un tableau intermédiaire pour pouvoir y accéder à l'intérieur de la boucle:

1.79 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # original
844 ms ± 41.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func2
707 ms ± 31.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func3
550 ms ± 4.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func4

• a préalablement calculé les indices en utilisant np.tile et en découpant au lieu de l'approche range avec % .
• Utilisé NumPy (en dehors de numba) pour calculer le np.arccos .

Timings finaux et résumé

@nb.njit(fastmath=True)
def func2(theta, pwr, neighbour, coschi, temp):
    arccos_coschi = np.arccos(coschi)
    for k in range(np.argmax(pwr), 5000 * pwr.size): 
        n = k % pwr.size
        if np.abs(theta[n] - np.pi / 2.) < np.abs(theta_c):
            adj = neighbour[n, 1]
        else:
            adj = neighbour[n, 0]
        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
        temp5 = temp[adj]**5;
        e_temp = 1. - np.exp(-temp5 * psi_diff / np.abs(eps))
        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
    return temp

Donc, au final, la dernière approche est environ 3 fois plus rapide (sans fastmath ) que l'approche originale. Si vous êtes sûr de vouloir utiliser fastmath , appliquez simplement fastmath = True sur le func4_inner et ce sera encore plus rapide:

res1 = func(theta, pwr, neighbour, coschi, np.ones(size))
res2 = # other approach

np.testing.assert_allclose(res1, res2)

%timeit func(theta, pwr, neighbour, coschi, np.ones(size))
%timeit # other approach

Cependant, comme je l'ai déjà dit, fastmath peut ne pas être approprié si vous voulez des résultats exacts (ou du moins pas trop inexacts).

Ici aussi, plusieurs optimisations dépendent beaucoup du matériel disponible et des caches de processeur (en particulier pour les parties du code limitées en bande passante mémoire). Vous devez vérifier comment ces approches fonctionnent les unes par rapport aux autres sur votre ordinateur.