Parallelization of outer product on multiple vectors using cython - numpy

Question

Suppose to have a numpy array m of shape (N, M) and I want to compute
res = np.zeros((M, M))
for i in range(N):
res += np.outer(m[i], m[i])
This loop can be made more efficient using einsum, i.e.
res = np.sum(np.einsum('ij,ik->ijk', m , m, axis=0)
but this requires the storage of a N x M x M matrix, which can be (and in my case is) very demanding.
I thought to build this function in cython, using parallelzation
import numpy as np
cimport numpy as np
from cython.parallel import prange
def get_s(double[:,:] m):
cdef Py_ssize_t i = 0
cdef int n = m.shape[1]
res = 0.
for i in prange(n, nogil=True):
res += np.outer(m[i], m[i])
return res
The idea of this code is that
The run of this code produces a lot of errors since I am using python objects, disallowed operations, and I don't know how to properly initialize res.

Answer 1

You cannot use numpy functions (np.outer) is the nogil context. So you just spell it out with loops.
Futhermore, your res variable seems to be an array, so you'll need to declare one and initialize it.
Last, you want the loops to compile to C, thus use typed memoryviews. It's easiest to use numpy arrays for memory management and take memoryviews of them. Taking it all together,
%%cython -a
cimport cython
import numpy as np
#cython.boundscheck(False)
#cython.wraparound(False)
def m_outer(double[:, ::1] a):
n, m = a.shape[0], a.shape[1]
cdef double[:, ::1] resm = np.zeros((m, m))
for i in range(a.shape[0]):
for j in range(a.shape[1]):
for k in range(a.shape[1]):
resm[j, k] += a[i, j] * a[i, k]
return np.asarray(resm)
A way to write these things (maybe the way) is to write it in python (nevermind the speed), validate the output on a small example (I use 3-by-4), then cythonize.
When cythonizing, use %cython -a and examine the generated C code.
Now, there are two obvious opportunities here: reorder the loops to lift loop-constants and use prange. Both are left as an exercise to the reader.
And the very last note. Unless it's an educational exercise, note that what you are really computing is a matrix product A.T # A.

Answer 2

Your iteration:
In [139]: res = np.zeros((4,4))
In [140]: for i in range(3): res += np.outer(m[i],m[i])
In [141]: res
Out[141]:
array([[ 80., 92., 104., 116.],
[ 92., 107., 122., 137.],
[104., 122., 140., 158.],
[116., 137., 158., 179.]])
We can do the same outer with broadcasting:
In [142]: np.sum(m[:,:,None]*m[:,None,:], axis=0)
Out[142]:
array([[ 80, 92, 104, 116],
[ 92, 107, 122, 137],
[104, 122, 140, 158],
[116, 137, 158, 179]])
(yes, this does make a temporary (N,M,M) array)
The suggested single step einsum:
In [143]: np.einsum('ij,ik->jk',m,m)
Out[143]:
array([[ 80, 92, 104, 116],
[ 92, 107, 122, 137],
[104, 122, 140, 158],
[116, 137, 158, 179]])
Which is just a simple dot product (with the appropriate transpose):
In [144]: m.T.dot(m)
Out[144]:
array([[ 80, 92, 104, 116],
[ 92, 107, 122, 137],
[104, 122, 140, 158],
[116, 137, 158, 179]])
In [145]: m.T#m
Out[145]:
array([[ 80, 92, 104, 116],
[ 92, 107, 122, 137],
[104, 122, 140, 158],
[116, 137, 158, 179]])
Since numpy dot's use fast BLAS code, I doubt if you can improve on it with cython.