Itertools pour simplifier les calculs scientifiques

En sciences nous sommes régulièrement amenés à résoudre des calculs pour l’ensemble des combinaisons de plusieurs paramètres: température, pression, volume, dimensions, tension, champ magnétique, etc.

Par exemple, si je considère un ensemble d’échantillons de forme cylindrique et de diamètre et de hauteur différentes, je peux calculer toutes les valeurs de volume possibles en bouclant sur le rayon et sur la hauteur.

Prenons maintenant un exemple plus physique, déjà abordé dans le billet précédent. Je veux calculer la fréquence de résonance d’un fil magnétique en fonction du champ qu’il subit (Omega) et de la force des interactions magnétiques issue de ses plus proches voisins (Heff).

Le fonction que je vais utiliser est la suivante :

def calc_Chi(Omega_H,Omega): #,porosite,champ
    Chi = (alpha*gamma*Ms*Omega*(Omega_H**2-((1+alpha**2)*Omega**2)))/((((Omega_H**2)-((1+alpha**2)*Omega**2))**2)+4*alpha**2*Omega**2*Omega_H**2)
    Chi_a = -(2*Omega**2*Ms*gamma*Omega_H*alpha)/((((Omega_H**2)-((1+alpha**2)*Omega**2))**2)+4*alpha**2*Omega**2*Omega_H**2)
    Chi_2 =Chi+Chi_a
    return Chi_2

Si je suis nouveau en python et que j’aime les boucles for, mon premier code pourrait ressembler à ça :

#Les 5 lignes ci-dessous définissent des vecteurs numpy
porosite = np.array(porosite) 
H = np.arange(0.1,10,0.01)    
Omega = H/gamma
Heff = calcHeff(porosite)
Chi_2 = np.array([])

#On boucle sur nos deux paramètres
for h_ in Heff: #D'abord sur le champ
    for o_ in Omega: #Et sur les fréquences
        Chi_2 = np.append(Chi_2,calc_Chi(h_,o_))

OK, ce calcul fonctionne mais si je le chronomètre pour seulement 521 fils, le temps total est de 340 secondes !

En plus d’être lent, ça n’est pas très élégant (imaginez si nous avons 4, ou 5 paramètres à faire varier).

for a in A:
    for b in B:
        for c in C:
            for d in D:
                for e in E:
                    #Fais quelque chose

En python, nous pouvons améliorer tout ça en deux étapes très simples:

1. En créant une liste contenant toutes les combinaisons désirées.
2. En donnant cette liste directement en entrée d’une fonction, et en profitant de la puissance de numpy.

Comme souvent en python, nous pourrions faire manuellement, mais heureusement il existe de nombreux modules pour nous simplifier la vie. Ici, c’est le module itertools qui va nous aider à créer une liste de paramètres. La fonction qui nous intéresse ici est itertools.product, qui est l’équivalent de boucles for imbriquées. Prenons un exemple simple :

A = [1,2,3]
B = [1,2,3,4]
C = itertools.product(A,B)
list(C)

Ce qui nous retourne la liste suivante :

[(1, 1),
(1, 2),
(1, 3),
(1, 4),
(2, 1),
(2, 2),
(2, 3),
(2, 4),
(3, 1),
(3, 2),
(3, 3),
(3, 4)]

Notez que c’est strictement équivalent à exécuter la ligne suivante dans un générateur : [(a,b) for a in A for b in B]

Maintenant que nous avons compris comment fonctionne itertools.product, je peux créer une liste avec itertools et la passer dans ma fonction calc_Chi : chaque colonne de “params” contient les valeurs des paramètres qui étaient auparavant passés via des boucles for.

Heff = calcHeff(porosite)
params = list(itertools.product(Heff,Omega))
params=np.array(params)
Chi2 = calc_Chi(params[:,0],params[:,1])

Si je chronomètre cette nouvelle solution, le temps de calcul est désormais de presque 0.7 secondes; soit un calcul environ 500 fois plus rapide !

Merci itertools pour la simplicité et merci numpy pour la rapidité 😉