Multiply array with diagonal matrix stored as vector - python

I have a 1D array A = [a, b, c...] (length N_A) and a 3D array T of shape (N_A, N_B, N_A). A is meant to represent a diagonal N_A by N_A matrix.
I'd like to perform contractions of A with T without having to promote A to dense storage. In particular, I'd like to do
np.einsum('ij, ikl', A, T)
and
np.einsum('ikl, lm', T, A)
is it possible to do such things while keeping A sparse?
Note this question is similar to
dot product with diagonal matrix, without creating it full matrix
but not identical, since it's not clear to me how one generalizes to more complicated index patterns.

np.einsum('ij, ikl', np.diag(a), t) is equivalent to (a * t.T).T.
np.einsum('ikl, lm', t, np.diag(a)) is equivalent to a * t.
(found by trial-and-error)

Related

Is there a numpy/scipy dot product for sparse matrix, calculating only the diagonal entries of the result?

Imagine having 2 sparse matrix:
> A, A.shape = (n,m)
> B, B.shape = (m,n)
I would like to compute the dot product A*B, but then only keep the diagonal. The matrices being big, I actually don't want to compute other values than the ones in the diagonal.
This is a variant of the question Is there a numpy/scipy dot product, calculating only the diagonal entries of the result?
Where the most relevant answer seems to be to use np.einsum:
np.einsum('ij,ji->i', A, B)
However this does not work:
ValueError: einstein sum subscripts string contains too many subscripts for operand 0
The solution is to use todense(), but it increases a lot the memory usage: np.einsum('ij,ji->i', A.todense(), B.todense())
The other solution, that I currently use, is to iterate over all the rows of A and compute each product in the loop :
for i in range(len_A):
result = np.float32(A[i].dot(B[:, i])[0, 0])
...
None of these solutions seems perfect. Is there an equivalent to np.einsum that could work with sparse matrices ?
[sum(A[i]*B.T[i]) for i in range(min(A.shape[0], B.shape[1]))]
otherwise this is faster:
l = min(A.shape[0], B.shape[1])
(A[np.arange(l)]*B.T[np.arange(l)]).sum(axis=1)
In general you shouldn't try to use numpy functions on the scipy.sparse arrays. In your case I'd first make sure both arrays actually have a compatible shape, that is
A, A.shape = (r,m)
B, B.shape = (m,r)
where r = min(n, p). Then we can compute the diagonal of the matrix product using
d = (A.multiply(B.T)).sum(axis=1)
Here we compute the entry wise row-column products, and manually sum them up. This avoids all the unnecessary computations you'd get using dot/#/*. (Note that unlike in numpy, both * and # perform matrix multiplication.)

Broadcasting - 3D field of coefficients to 3D field of matrices given matrix basis

I have a (large) 4D array, consisting of the 5 coefficients in a given basis for a matrix field. Given the 5 basis matrices, I want to efficiently calculate the matrix field.
The coefficient field c[x,y,z,i] being the value of i-th coefficient at position x,y,z
And the matrix field M[x,y,z,a,b] being the (3,3) matrix at position x,y,z
And the basis matrices T_1,...T_5, being the (3,3) basis matrices
I could loop over each position in space:
M[x,y,z,:,:] = T_1[:,:]*c[x,y,z,0] + T_2[:,:]*c[x,y,z,1]...T_5[:,:]*c[x,y,z,4]
But this is very inefficient. My attempts at using np.multiply,np.sum result in broadcasting errors due to the ambiguity of the desired product being a field of 3x3 matrices.
Keep in mind that to numpy, these 4 and 5d arrays are just that, not 3d arrays containing 2d matrices, etc.
Let's try to write your calculation in a way that clarifies dimensions:
M[x,y,z] = T_1*c[x,y,z,0] + T_2*c[x,y,z,1]...T_5*c[x,y,z,4]
M[x,y,z,:,:] = T_1[:,:]*c[x,y,z,0] + T_2[:,:]*c[x,y,z,1]...T_5[:,:]*c[x,y,z,4]
c[x,y,z,i] is a coefficient, right? So M is a weighted sum of the T_n arrays?
One way of expressing this is:
T = np.stack([T_1, T_2, ...T_5], axis=0) # 3d (nab)
M = np.einsum('nab,xyzn->xyzab', T, c)
We could alternatively stack T_i on a new last axis
T = np.stack([T_1, T_2 ...T_5], axis=2) # (abn)
M = np.einsum('abn,xyzn->xyzab', T, c)
or as broadcasted multiplication plus sum:
M = (T[None,None,None,:,:,:] * c[:,:,:,None,None,:]).sum(axis=-1)
I'm writing this code without testing, so there may be errors, but I think the basic outline is right.
It could also be written as a dot, if I can put the n dimension last in one argument, and 2nd to the last in the other. Or with tensordot. But there's less control over broadcasting of the other dimensions.
For test calculations you could also reshape these arrays so that the x,y,z are rolled into one, and the a,b into another, e.g
M[xyz,:] = T_n[ab]*c[xyz,n] # etc

Handling matrix multiplication in log space in Python

I am implementing a Hidden Markov Model and thus am dealing with very small probabilities. I am handling the underflow by representing variables in log space (so x → log(x)) which has the side effect that multiplication is now replaced by addition and addition is handled via numpy.logaddexp or similar.
Is there an easy way to handle matrix multiplication in log space?
This is the best way I could come up with to do it.
from scipy.special import logsumexp
def log_space_product(A,B):
Astack = np.stack([A]*A.shape[0]).transpose(2,1,0)
Bstack = np.stack([B]*B.shape[1]).transpose(1,0,2)
return logsumexp(Astack+Bstack, axis=0)
The inputs A and B are the logs of the matrices A0 and B0 you want to multiply, and the functions returns the log of A0B0. The idea is that the i,j spot in log(A0B0) is the log of the dot product of the ith row of A0 and the jth column of B0. So it is the logsumexp of the ith row of A plus the jth column of B.
In the code, Astack is built so the i,j spot is a vector containing the ith row of A, and Bstack is built so the i,j spot is a vector containing the jth column of B. Thus Astack + Bstack is a 3D tensor whose i,j spot is the ith row of A plus the jth column of B. Taking logsumexp with axis = 0 then gives the desired result.
Erik's response doesn't seem to work for some non-square matrices (e.g. n*m times m*r). Here is a version that takes that into account:
def log_space_product(A,B):
Astack = np.stack([A]*B.shape[1]).transpose(1,0,2)
Bstack = np.stack([B]*A.shape[0]).transpose(0,2,1)
return logsumexp(Astack+Bstack, axis=2)
where the i, j spot of A contains the i-th row of A and i, j spot of B contains the i-th column of B.
This happens because [A] * B.shape[1] is of shape (r, n, m) which is transposed into (n, r, m), and [B] * A.shape[0] is of shape (n, m, r) which is transposed into (n, r, m). We want their first two dimensions to be (n, r) because the result matrix needs to be of shape (n, r).
Took a while to figure out myself. Hope this helps anyone implementing a HMM!

numpy dot product and matrix product

I'm working with numpy arrays of shape (N,), (N,3) and (N,3,3) which represent sequences of scalars, vectors and matrices in 3D space. I have implemented pointwise dot product, matrix multiplication, and matrix/vector multiplication as follows:
def dot_product(v, w):
return np.einsum('ij, ij -> i', v, w)
def matrix_vector_product(M, v):
return np.einsum('ijk, ik -> ij', M, v)
def matrix_matrix_product(A, B):
return np.einsum('ijk, ikl -> ijl', A, B)
As you can see I use einsum for lack of a better solution. To my surprise I was not able to use np.dot... which seems not suitable for this need. Is there a more numpythonic way to implement these function?
In particular it would be nice if the functions could work also on the shapes (3,) and (3,3) by broadcasting the first missing axis. I think I need ellipsis, but I don't quite understand how to achieve the result.
These operations cannot be reshaped into general BLAS calls and looping BLAS calls would be quite slow for arrays of this size. As such, einsum is likely optimal for this kind of operation.
Your functions can be generalized with ellipses as follows:
def dot_product(v, w):
return np.einsum('...j,...j->...', v, w)
def matrix_vector_product(M, v):
return np.einsum('...jk,...k->...j', M, v)
def matrix_matrix_product(A, B):
return np.einsum('...jk,...kl->...jl', A, B)
Just as working notes, these 3 calculations can also be written as:
np.einsum(A,[0,1,2],B,[0,2,3],[0,1,3])
np.einsum(M,[0,1,2],v,[0,2],[0,1])
np.einsum(w,[0,1],v,[0,1],[0])
Or with Ophion's generalization
np.einsum(A,[Ellipsis,1,2], B, ...)
It shouldn't be hard to generate the [0,1,..] lists based on the dimensions of the inputs arrays.
By focusing on generalizing the einsum expressions, I missed the fact that what you are trying to reproduce is N small dot products.
np.array([np.dot(i,j) for i,j in zip(a,b)])
It's worth keeping mind that np.dot uses fast compiled code, and focuses on calculations where the arrays are large. Where as your problem is one of calculating many small dot products.
And without extra arguments that define axes, np.dot performs just 2 of the possible combinations, ones which can be expressed as:
np.einsum('i,i', v1, v2)
np.einsum('...ij,...jk->...ik', m1, m2)
An operator version of dot would face the same limitation - no extra parameters to specify how the axes are to be combined.
It may also be instructive to note what tensordot does to generalize dot:
def tensordot(a, b, axes=2):
....
newshape_a = (-1, N2)
...
newshape_b = (N2, -1)
....
at = a.transpose(newaxes_a).reshape(newshape_a)
bt = b.transpose(newaxes_b).reshape(newshape_b)
res = dot(at, bt)
return res.reshape(olda + oldb)
It can perform a dot with summation over several axes. But after the transposing and reshaping is done, the calculation becomes the standard dot with 2d arrays.
This could have been flagged as a duplicate issue. People have asking about doing multiple dot products for some time.
Matrix vector multiplication along array axes
suggests using numpy.core.umath_tests.matrix_multiply
https://stackoverflow.com/a/24174347/901925 equates:
matrix_multiply(matrices, vectors[..., None])
np.einsum('ijk,ik->ij', matrices, vectors)
The C documentation for matrix_multiply notes:
* This implements the function
* out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.
inner1d from the same directory does the same same for (N,n) vectors
inner1d(vector, vector)
np.einsum('ij,ij->i', vector, vector)
# out[n] = sum_i { in1[n, i] * in2[n, i] }
Both are UFunc, and can handle broadcasting on the right most dimensions. In numpy/core/test/test_ufunc.py these functions are used to exercise the UFunc mechanism.
matrix_multiply(np.ones((4,5,6,2,3)),np.ones((3,2)))
https://stackoverflow.com/a/16704079/901925 adds that this kind of calculation can be done with * and sum, eg
(w*v).sum(-1)
(M*v[...,None]).sum(-1)
(A*B.swapaxes(...)).sum(-1)
On further testing, I think inner1d and matrix_multiply match your dot and matrix-matrix product cases, and the matrix-vector case if you add the [...,None]. Looks like they are 2x faster than the einsum versions (on my machine and test arrays).
https://github.com/numpy/numpy/blob/master/doc/neps/return-of-revenge-of-matmul-pep.rst
is the discussion of the # infix operator on numpy. I think the numpy developers are less enthused about this PEP than the Python ones.

Numpy matrix multiplication with array of matrices

I have several numpy arrays that I would like to multiply (using dot, so matrix multiplication). I'd like to put them all into a numpy array, but I can't figure out how to do it.
E.g.
a = np.random.randn((10,2,2))
b = np.random.randn((10,2))
So I have 10 2x2 matrices (a) and 10 2x1 matrices (b). What I could do is this:
c = np.zeros((10,2))
for i in range(10):
c[i] = np.dot(a[i,:,:],b[i,:])
You get the idea.
But I feel like there's a usage of dot or tensordot or something that would do this in one line really easily. I just can't make sense of the dot and tensordot functions for >2 dimensions like this.
You could use np.einsum:
c = np.einsum('ijk,ik->ij', a, b)
einsum performs a sum of products. Since matrix multiplication is a sum of products, any matrix multiplication can be expressed using einsum. It is based on Einstein summation notation.
The first argument to einsum, ijk,ik->ij is a string of subscripts.
ijk declares that a has three axes which are denoted by i, j, and k.
ik, similarly, declares that the axes of b will be denoted i and k.
When the subscripts repeat, those axes are locked together for purposes of summation.
The part of the subscript that follows the -> shows the axes which will remain after summation.
Since the k appears on the left (of the ->) but disappears on the right, there is summation over k. It means that the sum
c_ij = sum over k ( a_ijk * b_ik )
should be computed. Since this sum can be computed for each i and j, the result is an array with subscripts i and j.

Categories