Более эффективный способ умножения каждого столбца 2-й матрицы на каждый срез 3-й матрицы

У меня есть массив 8x8x25000 W и массив 8 x 25000 r. Я хочу, чтобы каждый из 8×8 фрагментов W каждого столбца (8×1) из r и сохранял результат в Wres, который в конечном итоге станет матрицей 8×25000.

Я выполняю это, используя цикл for как таковой:

for i in range(0,25000): Wres[:,i] = np.matmul(W[:,:,i],res[:,i]) 

Но это медленно, и я надеюсь, что есть более быстрый способ достичь этого.

Есть идеи?

    Matmul может распространяться до тех пор, пока 2 массива имеют одну и ту же длину оси. Из документов:

    Если любой аргумент ND, N> 2, он рассматривается как стек матриц, находящихся в последних двух индексах, и передает их соответственно.

    Таким образом, вы должны выполнить 2 операции до matmul :

     import numpy as np a = np.random.rand(8,8,100) b = np.random.rand(8, 100) 
    1. транспонируйте a и b так, чтобы первая ось была 100 срезов
    2. добавьте дополнительный размер к b чтобы b.shape = (100, 8, 1)

    Затем:

      at = a.transpose(2, 0, 1) # swap to shape 100, 8, 8 bt = bT[..., None] # swap to shape 100, 8, 1 c = np.matmul(at, bt) 

    c теперь 100, 8, 1 , переформатируется обратно до 8, 100 :

      c = np.squeeze(c).swapaxes(0, 1) 

    или

      c = np.squeeze(c).T 

    И последнее, однострочный для удобства:

     c = np.squeeze(np.matmul(a.transpose(2, 0, 1), bT[..., None])).T 

    Альтернативой использованию np.matmul является np.einsum , который может быть выполнен за 1 короткую и, возможно, более приемлемую строку кода без цепочки методов.

    Примеры массивов:

     np.random.seed(123) w = np.random.rand(8,8,25000) r = np.random.rand(8,25000) wres = np.einsum('ijk,jk->ik',w,r) # a quick check on result equivalency to your loop print(np.allclose(np.matmul(w[:, :, 1], r[:, 1]), wres[:, 1])) True 

    Сроки эквивалентны решению @ Imanol, так что сделайте выбор. Оба они в 30 раз быстрее, чем цикл. Здесь einsum будет конкурентоспособным из-за размера массивов. С массивами, большими, чем эти, он, скорее всего, выиграет и проиграет для меньших массивов. См. Обсуждение для большего.

     def solution1(): return np.einsum('ijk,jk->ik',w,r) def solution2(): return np.squeeze(np.matmul(w.transpose(2, 0, 1), rT[..., None])).T def solution3(): Wres = np.empty((8, 25000)) for i in range(0,25000): Wres[:,i] = np.matmul(w[:,:,i],r[:,i]) return Wres %timeit solution1() 100 loops, best of 3: 2.51 ms per loop %timeit solution2() 100 loops, best of 3: 2.52 ms per loop %timeit solution3() 10 loops, best of 3: 64.2 ms per loop 

    Кредит : @ Дивакар