Efficient linear algebra for block-diagonal matrices in compressed format

2
votes

I have a linear system in which all the matrices are block diagonal. They have N blocks identical in shape.

Matrices are stored in compressed format as numpy arrays with shape (N, n, m), while the shape of the vectors is (N, m).

I currently implemented the matrix-vector product as

import numpy as np

def mvdot(m, v):
    return (m * np.expand_dims(v, -2)).sum(-1)

Thanks to broadcasting rules, if all the blocks of a matrix are the same I have to store it only once in an array with shape (1, n, m): the product with a vector (N, m) still gives the correct (N, n) vector.

My questions are:

  • how to implement an efficient matrix-matrix product that yields the matrix with shape (N, n, m) from two matrices with shapes (N, n, p) and (N, p, m)?
  • is there a way to perform these operations with a numpy built-in (possibly faster) function? Functions like np.linalg.inv make me think that numpy was designed to support this compressed format for block diagonal matrices.
2
pythonnumpymultidimensional-arrayvectorizationlinear-algebra
Sounds like you can also use standard matrix multiplication with sparse representation for a block diagonal system. - Bort
What do you mean? using the dot method? - dPol
Yes, instead of working with a tensor (N,n,m) as a compressed format you can keep them as 2d matrices (Nn,Nm) with sparse representation and use sparse matrix matrix multiplication or matrix vector multiplication. - Bort
Which spare matrix class do you have in mind? - dPol

2 Answers

1
votes

If I understand your question correctly, you have two arrays of shape (N,n,p) and (N,p,m), respectively, and their product should be of shape (N,n,m) where element [i,:,:] is the matrix product of M1[i,:,:] and M2[i,:,:]. This can be achieved using numpy.einsum:

import numpy as np
N = 7
n,p,m = 3,4,5
M1 = np.random.rand(N,n,p)
M2 = np.random.rand(N,p,m)
Mprod = np.einsum('ijk,ikl->ijl',M1,M2)

# check if all the submatrices are what we expect
all([np.allclose(np.dot(M1[k,...],M2[k,...]),Mprod[k,...]) for k in range(N)])
# True

Numpy's einsum is an incredibly versatile construction for complicated linear operations, and it's usually pretty efficient with two operands. The idea is to rewrite your operation in an indexed way: what you need is to multiply M1[i,j,k] with M2[i,k,l] for each i,j,l, and sum over k. This is exactly what the above call to einsum does: it collapses the index k, and performs the necessary products and assignments along the remaining dimensions in the given order.

The matrix-vector product can be done similarly:

M = np.random.rand(N,n,m)
v = np.random.rand(N,m)
Mvprod = np.einsum('ijk,ik->ij',M,v)

It's possible that numpy.dot can be coerced with the proper transposes and dimension tricks to directly do what you want, but I couldn't make that work.

Both of the above operations can be done in the same function call by allowing an implicit number of dimensions within einsum:

def mvdot(M1,M2):
    return np.einsum('ijk,ik...->ij...',M1,M2)

Mprod = mvdot(M1,M2)
Mvprod = mvdot(M,v)

In case the input argument M2 is a block matrix, there will be a leading dimension appended to the result, creating a block matrix. In case M2 is a "block vector", the result will be a block vector.

1
votes

Since Python 3.5 and above, the previous example can be simplified using the matrix multiplication operator @ (numpy.matmul) which treats this case as a stack of matrices residing in the last two indexes and broadcast accordingly:

import numpy as np
N = 7
n,p,m = 3,4,5
M1 = np.random.rand(N,n,p)
M2 = np.random.rand(N,p,m)
Mprod = M1 @ M2  # similar to np.matmul(M1, M2)

all([np.allclose(np.dot(M1[k,...],M2[k,...]),Mprod[k,...]) for k in range(N)])
#True