web-dev-qa-db-fra.com

Python vectorisation imbriquée pour les boucles

J'apprécierais un peu d'aide pour trouver et comprendre une manière Pythonique d'optimiser les manipulations de tableaux suivantes dans les boucles imbriquées:

def _func(a, b, radius):
    "Return 0 if a>b, otherwise return 1"
    if distance.euclidean(a, b) < radius:
        return 1
    else:
        return 0

def _make_mask(volume, roi, radius):
    mask = numpy.zeros(volume.shape)
    for x in range(volume.shape[0]):
        for y in range(volume.shape[1]):
            for z in range(volume.shape[2]):
                mask[x, y, z] = _func((x, y, z), roi, radius)
    return mask

Où volume.shape (182, 218, 200) et roi.shape (3,) sont les deux types ndarray; et radius est un int

pythonnumpyfor-loopoptimizationvectorization
34
Kambiz

Approche n ° 1

Voici une approche vectorisée -

m,n,r = volume.shape
x,y,z = np.mgrid[0:m,0:n,0:r]
X = x - roi[0]
Y = y - roi[1]
Z = z - roi[2]
mask = X**2 + Y**2 + Z**2 < radius**2

Amélioration possible: nous pouvons probablement accélérer la dernière étape avec le module numexpr -

import numexpr as ne

mask = ne.evaluate('X**2 + Y**2 + Z**2 < radius**2')

Approche n ° 2

Nous pouvons également construire progressivement les trois plages correspondant aux paramètres de forme et effectuer la soustraction contre les trois éléments de roi à la volée sans réellement créer les maillages comme précédemment avec np.mgrid. Cela bénéficierait de l'utilisation de broadcasting à des fins d'efficacité. L'implémentation ressemblerait à ceci -

m,n,r = volume.shape
vals = ((np.arange(m)-roi[0])**2)[:,None,None] + \
       ((np.arange(n)-roi[1])**2)[:,None] + ((np.arange(r)-roi[2])**2)
mask = vals < radius**2

Version simplifiée: Merci à @Bi Rico d'avoir suggéré une amélioration ici car nous pouvons utiliser np.ogrid pour effectuer ces opérations de manière un peu plus concise, comme ceci -

m,n,r = volume.shape    
x,y,z = np.ogrid[0:m,0:n,0:r]-roi
mask = (x**2+y**2+z**2) < radius**2

Test d'exécution

Définitions des fonctions -

def vectorized_app1(volume, roi, radius):
    m,n,r = volume.shape
    x,y,z = np.mgrid[0:m,0:n,0:r]
    X = x - roi[0]
    Y = y - roi[1]
    Z = z - roi[2]
    return X**2 + Y**2 + Z**2 < radius**2

def vectorized_app1_improved(volume, roi, radius):
    m,n,r = volume.shape
    x,y,z = np.mgrid[0:m,0:n,0:r]
    X = x - roi[0]
    Y = y - roi[1]
    Z = z - roi[2]
    return ne.evaluate('X**2 + Y**2 + Z**2 < radius**2')

def vectorized_app2(volume, roi, radius):
    m,n,r = volume.shape
    vals = ((np.arange(m)-roi[0])**2)[:,None,None] + \
           ((np.arange(n)-roi[1])**2)[:,None] + ((np.arange(r)-roi[2])**2)
    return vals < radius**2

def vectorized_app2_simplified(volume, roi, radius):
    m,n,r = volume.shape    
    x,y,z = np.ogrid[0:m,0:n,0:r]-roi
    return (x**2+y**2+z**2) < radius**2

Horaires -

In [106]: # Setup input arrays  
     ...: volume = np.random.Rand(90,110,100) # Half of original input sizes 
     ...: roi = np.random.Rand(3)
     ...: radius = 3.4
     ...: 

In [107]: %timeit _make_mask(volume, roi, radius)
1 loops, best of 3: 41.4 s per loop

In [108]: %timeit vectorized_app1(volume, roi, radius)
10 loops, best of 3: 62.3 ms per loop

In [109]: %timeit vectorized_app1_improved(volume, roi, radius)
10 loops, best of 3: 47 ms per loop

In [110]: %timeit vectorized_app2(volume, roi, radius)
100 loops, best of 3: 4.26 ms per loop

In [139]: %timeit vectorized_app2_simplified(volume, roi, radius)
100 loops, best of 3: 4.36 ms per loop

Donc, comme toujours broadcasting montrant sa magie pour un fou presque 10,000x accélération sur le code d'origine et plus de 10x mieux que de créer des maillages en utilisant des opérations diffusées à la volée!

61
Divakar

Supposons que vous créez d'abord un tableau xyzy:

import itertools

xyz = [np.array(p) for p in itertools.product(range(volume.shape[0]), range(volume.shape[1]), range(volume.shape[2]))]

Maintenant, en utilisant numpy.linalg.norm ,

np.linalg.norm(xyz - roi, axis=1) < radius

vérifie si la distance pour chaque Tuple de roi est inférieure au rayon.

Enfin, juste reshape le résultat aux dimensions dont vous avez besoin.

7
Ami Tavory