Python vectorise les boucles for imbriquées

Question

J'aimerais avoir de l'aide pour trouver et comprendre un moyen pythonique d'optimiser les manipulations de tableaux suivantes dans des boucles for 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 tous deux ndarray les types radius est un int

Answer 1

Approche n° 1

Voici une approche vectorielle -

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 numexpr module -

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 créer réellement les mailles comme cela a été fait précédemment avec np.mgrid . L'utilisation de broadcasting pour des raisons 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 les éléments suivants 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 une folie presque 10,000x par rapport au code original et plus de 10x mieux que de créer des mailles en utilisant des opérations diffusées à la volée !

Answer 2

Disons que vous construisez d'abord un xyzy le tableau :

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 plus petit que le rayon.

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