2

Is there an expression (perhaps using np.tensordot) that most expressively captures a sparse matrix vector multplicatiuon between a block matrix of diagonal blocks and a vector of the corresponding size?

I have a working implementation that performs the exact operations I want, but I used two python loops (see below) instead of an appropriate numpy command, which probably exists.

For example:

import numpy as np

outer_size = 2
inner_size = 5
shape = (outer_size, outer_size, inner_size)
diag_block = np.arange(np.prod(shape)).reshape(shape)
true_diag = np.bmat([[np.diag(diag_block[i,j]) for j in range(shape[1])] for i in range(shape[0])]).A
x = np.arange(shape[1] * shape[2])

def sparse_prod(diags, x):
    outer_size = diags.shape[0]
    return np.hstack(sum(diags[i] * x.reshape(outer_size, -1)) for i in range(outer_size))

print(true_diag.dot(x))
print(sparse_prod(diag_block, x))
VF1
  • 1,594
  • 2
  • 11
  • 31

1 Answers1

3

Approach #1 : You could use np.einsum -

np.einsum('ijk,jk->ik', diag_block, x.reshape(outer_size, -1)).ravel()

Basically, we are keeping the last axis aligned between the two inputs and sum-reducing the second and first axes respectively.

Approach #2 : Depending on the shapes of the input arrays, you might want to use np.dot/np.tensordot in a loop as discussed in some detail in this post.

Here's such an approach with a loop -

m,_,n = diag_block.shape
x2D = x.reshape(outer_size, -1)
out = np.empty((m,n))
for i in range(n):
    out[:,i] = np.dot(diag_block[...,i], x2D[:,i])
out.shape = out.size  # Flatten
Community
  • 1
  • 1
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Right, I agree einsum would work, but, as your link points out, it's far slower than the alternatives. Tensordot is equivalent, but doen't require knowledge of the Einstein summation subscript syntax and needs no subscript compliation. Would you mind adding the tensordot equivalent for completeness (I couldn't figure it out). – VF1 Jan 13 '17 at 17:30
  • If you'd written that you want to `dot` a (n, m, k) array with a (m, k) one, the `einsum` notation is trivial. Your `block_diagonal` framing obscures that, to some degree. Why are you concerned with the `subscript compilation`? That's done in C code. – hpaulj Jan 13 '17 at 17:50
  • Even when I increase the sizes to 20 and 50, `einsum` wins. `true_diag.dot(x)` is slow even without including the `bmat` construction step. And much slower with that step. The `dot` operates on a 1000x1000 space, `einsum` on a 20x20x50. – hpaulj Jan 13 '17 at 18:07
  • @hpaulj Are you talking to me? Or to author of the answer? I don't like einsum as much because it is not widely known, and its documentation basically says "generalize the behavior of this function from examples". On the other hand, `tensordot` has a formal, precise definition in the docs as a tensor expansion/contraction along selected axes. Also, you're confusing the fact that I'm checking the answer with a huge operation that I'd never do in real life simply to validate that my `sparse_prod` works. – VF1 Jan 13 '17 at 18:41
  • @Divakar right, sorry for making you go through the trouble again - I can see how to do it with `dot`, I was wondering in particular about `tensordot`, without the external loop. – VF1 Jan 13 '17 at 18:41
  • @VF1 Since, we need to keep the last axis aligned between those two inputs, we can't use `np.tensordot` on the entire tensor. – Divakar Jan 13 '17 at 18:44
  • `tensordot` just transposes and reshapes the inputs, and then does one `dot` - followed by more transposing and reshaping if needed. – hpaulj Jan 13 '17 at 18:45
  • @Divakar OK, I accepted in that case. Would you mind putting that `np.tensordot` is impossible in your answer? – VF1 Jan 13 '17 at 19:03
  • @VF1 Maybe you got the concept of `tensordot` twisted. `Tensordot` is for `3D` or higher dimensional arrays. Within a loop, for this problem, we don't have such higher dim arrays, we only have 2D and 1D arrays. Do you want me to add this statement to my answer post or is this enough to clarify your doubts? – Divakar Jan 13 '17 at 19:06
  • No, I was suggesting to use `tensordot` **outside** of the loop, but couldn't figure out how. You said it was impossible, I was wondering how. – VF1 Jan 13 '17 at 21:06
  • @VF1 I am confused - Outside of which loop? The loop I meant was the loop I used in approach#2 : `for i in range(n): out[:,i] = ....` and I meant it doesn't make sense to use `np.tensordot` within that loop. – Divakar Jan 13 '17 at 21:19
  • @Divakar I agree. I'm suggesting that there might be a one-line statement using `np.tensordot` with **no python loop at all** equivalent to the `np.einsum`, but I was hoping you could confirm – VF1 Jan 13 '17 at 21:44
  • @VF1 I see. There isn't any, sorry :) – Divakar Jan 13 '17 at 21:45