Как векторизовать вычисления на массивах разных измерений?

Question

У меня есть несколько больших массивов комплексных чисел, на которых мне нужно выполнить вычисления.

import numpy as np

# Reduced sizes -- real ones are orders of magnitude larger
n, d, l = 50000, 3, 1000

# Two complex matrices and a column vector
x = np.random.rand(n, d) + 1j*np.random.rand(n, d)
s = np.random.rand(n, d) + 1j*np.random.rand(n, d)
v = np.random.rand(l)[:, np.newaxis]

Функция в основном x*v*s для каждой строки x (и s), а затем это произведение суммируется по строке. Поскольку массивы имеют разные размеры, я не могу найти способ векторизации вычислений, и он слишком медленный, чтобы использовать цикл for.

Моя текущая реализация такова (~ 3,5 секунды):

h = []
for i in range(len(x)):
    h.append(np.sum(x[i,:]*v*s[i,:], axis=1))

h = np.asarray(h)

Я также пытался использовать np.apply_along_axis() с расширенной матрицей, но он только немного быстрее (~ 2,6 с) и не так удобочитаем.

def func(m, v):
    return np.sum(m[:d]*v*m[d:], axis=1)

h = np.apply_along_axis(func, 1, np.hstack([x, s]), v)

Что может быть быстрее для вычисления этого результата? Я могу использовать другие пакеты, такие как dask, если это поможет.

Answer 1

С трансляцией это должно работать:

np.sum(((x*s)[...,None]*v[:,0], axis=1)

но с вашими размерами образца я получаю ошибку памяти. «Внешний» широковещательный массив (n,d,l) форма слишком велика для моей памяти.

Я могу уменьшить использование памяти, перебирая меньшее d измерение:

res = np.zeros((n,l), dtype=x.dtype) 
for i in range(d): 
    res += (x[:,i]*s[:,i])[:,None]*v[:,0] 

Это тестирует так же, как ваш h, но я не смог завершить временные тесты. Как правило, итерации по меньшему измерению быстрее.

Я могу повторить вещи с небольшими размерами.

Это, вероятно, также может быть выражено как einsum проблема, хотя это может не помочь с этими измерениями.

---

In [1]: n, d, l = 5000, 3, 1000 
   ...:  
   ...: # Two complex matrices and a column vector 
   ...: x = np.random.rand(n, d) + 1j*np.random.rand(n, d) 
   ...: s = np.random.rand(n, d) + 1j*np.random.rand(n, d) 
   ...: v = np.random.rand(l)[:, np.newaxis]                                                           
In [2]:                                                                                                
In [2]: h = [] 
   ...: for i in range(len(x)): 
   ...:     h.append(np.sum(x[i,:]*v*s[i,:], axis=1)) 
   ...:  
   ...: h = np.asarray(h)                                                                              
In [3]: h.shape                                                                                        
Out[3]: (5000, 1000)
In [4]: res = np.zeros((n,l), dtype=x.dtype) 
   ...: for i in range(d): 
   ...:     res += (x[:,i]*s[:,i])[:,None]*v[:,0] 
   ...:                                                                                                
In [5]: res.shape                                                                                      
Out[5]: (5000, 1000)
In [6]: np.allclose(res,h)                                                                             
Out[6]: True
In [7]: %%timeit  
   ...: h = [] 
   ...: for i in range(len(x)): 
   ...:     h.append(np.sum(x[i,:]*v*s[i,:], axis=1)) 
   ...: h = np.asarray(h) 
   ...:  
   ...:                                                                                                
490 ms ± 3.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [8]: %%timeit 
   ...: res = np.zeros((n,l), dtype=x.dtype) 
   ...: for i in range(d): 
   ...:     res += (x[:,i]*s[:,i])[:,None]*v[:,0] 
   ...:                                                                                                

354 ms ± 1.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [9]:                                                                                                
In [9]: np.sum((x*s)[...,None]*v[:,0], axis=1).shape                                                   
Out[9]: (5000, 1000)
In [10]: out = np.sum((x*s)[...,None]*v[:,0], axis=1)                                                  
In [11]: np.allclose(h,out)                                                                            
Out[11]: True
In [12]: timeit out = np.sum((x*s)[...,None]*v[:,0], axis=1)                                           
310 ms ± 964 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Некоторая экономия времени, но не большая.

И einsum версия:

In [13]: np.einsum('ij,ij,k->ik',x,s,v[:,0]).shape                                                     
Out[13]: (5000, 1000)
In [14]: np.allclose(np.einsum('ij,ij,k->ik',x,s,v[:,0]),h)                                            
Out[14]: True
In [15]: timeit np.einsum('ij,ij,k->ik',x,s,v[:,0]).shape                                              
167 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Хорошая экономия времени. Но я не знаю, как это будет масштабироваться.

---

Но einsum заставил меня понять, что мы можем суммировать по d измерению раньше, перед умножением на v - и получить много времени и использования памяти:

In [16]: np.allclose(np.sum(x*s, axis=1)[:,None]*v[:,0],h)                                             
Out[16]: True
In [17]: timeit np.sum(x*s, axis=1)[:,None]*v[:,0]                                                     
68.4 ms ± 1.58 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

@ cs95 попал туда первым!

Согласно комментарию @ PaulPanzer, флаг оптимизации помогает. Вероятно, он делает тот же вывод - который мы можем суммировать на j рано:

In [18]: timeit np.einsum('ij,ij,k->ik',x,s,v[:,0],optimize=True).shape                                
91.6 ms ± 991 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)