0

I have three two matrices A and B, the matrix product I want is diagonal(A.B.A^T), where A^T is transpose of a matrix. The dimensions of the matrices are as follows

A - (2^n, n) 
B - (n, n)

where is n is any natural number.

I want first row slice of matrix A to be multiplied with matrix B and their product is multiplied with first column of matrix A^T. I don't want the full product of A.B.A^T because I only want the diagonal slice.

It seems to me this can be achieved using eisum.

A related question on codereview

papabiceps
  • 988
  • 2
  • 18
  • 33
  • maybe just full product and use numpy.diagonal() – CuCaRot Jul 26 '19 at 09:10
  • Too slow, I'm doing way more calculations than I need to. `2^n` grows very big for `n > 10`. Since there is repeated pattern here, I'm hoping that we can use einsum and get rid of the for loop. This product is the core to my experiments and if I have save time, I would save money. – papabiceps Jul 26 '19 at 09:13
  • You can simply get `A.B` trough dot product. Then, since you only want the diagonal, multiply the previous result elementwise by `A` and do a sum reduction: `(A@B*A).sum(1)` – Brenlla Jul 26 '19 at 09:30
  • Use linked dup target one for a faster one. – Divakar Jul 26 '19 at 13:39

1 Answers1

3

Here is how to do it with einsum

np.random.seed(1)

A = np.random.randint(0,10,(8,4))
B = np.random.randint(0,10,(4,4))

# brute force for reference
np.diag(A@B@A.T)
# array([3830,  233, 2835,  958, 3706, 1273, 5478,  934])

# more economical
np.einsum('ij,jk,ik->i',A,B,A)
# array([3830,  233, 2835,  958, 3706, 1273, 5478,  934])
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • As the size grows larger the speedup is growing larger, seeing 10x-30x speedups. Thank you. Paul and `einsum` FTW. – papabiceps Jul 26 '19 at 09:44