2017-08-31 1 views
1

Étant donné un tableau 3D et un tableau 2D,Multipliant par éléments sur l'axe final de deux tableaux

a = np.arange(10*4*3).reshape((10,4,3)) 
b = np.arange(30).reshape((10,3)) 

Comment puis-je utiliser par éléments multiplication à travers l'axe final de chaque, ce qui cc a la forme .shape comme a? C'est à dire.

c[0] = a[0] * b[0] 
c[1] = a[1] * b[1] 
# ... 
c[i] = a[i] * b[i] 

Répondre

2

Sans une somme de réduction en cause, un broadcasting simple serait vraiment efficace après l'extension b à 3D avec np.newaxis/None -

a*b[:,None,:] # or simply a*b[:,None] 

Test d'exécution -

In [531]: a = np.arange(10*4*3).reshape((10,4,3)) 
    ...: b = np.arange(30).reshape((10,3)) 
    ...: 

In [532]: %timeit np.einsum('ijk,ik->ijk', a, b) #@Brad Solomon's soln 
    ...: %timeit a*b[:,None] 
    ...: 
100000 loops, best of 3: 1.79 µs per loop 
1000000 loops, best of 3: 1.66 µs per loop 

In [525]: a = np.random.rand(100,100,100) 

In [526]: b = np.random.rand(100,100) 

In [527]: %timeit np.einsum('ijk,ik->ijk', a, b) 
    ...: %timeit a*b[:,None] 
    ...: 
1000 loops, best of 3: 1.53 ms per loop 
1000 loops, best of 3: 1.08 ms per loop 

In [528]: a = np.random.rand(400,400,400) 

In [529]: b = np.random.rand(400,400) 

In [530]: %timeit np.einsum('ijk,ik->ijk', a, b) 
    ...: %timeit a*b[:,None] 
    ...: 
10 loops, best of 3: 128 ms per loop 
10 loops, best of 3: 94.8 ms per loop 
+0

@BradSolomon C'est exact. Ajout de timings pour les confirmer. – Divakar

1

En utilisant np.einsum:

c = np.einsum('ijk,ik->ijk', a, b) 

Vérification rapide:

print(np.allclose(c[0], a[0] * b[0])) 
print(np.allclose(c[1], a[1] * b[1])) 
print(np.allclose(c[-1], a[-1] * b[-1])) 
# True 
# True 
# True