4

Currently I use

Na = (3, 2, 4)
Nb = Na[1:]
A = np.arange(np.prod(Na)).reshape(Na)
b = np.arange(np.prod(Nb)).reshape(Nb)

I want to calculate:

r_ik = sum_j(A_ijk * b_jk)

r = np.empty((A.shape[0], A.shape[2])
for i in range(A.shape[2]):
    r[:, i] = np.dot(A[:, :, i], b[:, i])

In words: A is a "stack of 4 matrices" (shaped (3,2)), i.e. a 3d-array, b is a "stack of 4 vectors" (shaped (3,)), i.e. a 2d-array. The desired result is a "stack of 4 matrix-vector-products", i.e. a stack of vectors, i.e. again a 2d-array (shaped (3, 4)).

I had a medium-depth look to np.einsum and np.tensordot but any solution that I constructed with these was at least as long and less readable than my loop-solution.

However I think there should be a one-liner for that simple problem.

cknoll
  • 2,130
  • 4
  • 18
  • 34

1 Answers1

6

I'm not sure why np.einsum didn't work for you. This should do the trick:

r = np.einsum('ijk,jk->ik', A, b)

I consider this to be very readable, because it exactly reflects the mathematical formula you've given in your question.

Praveen
  • 6,872
  • 3
  • 43
  • 62
  • Indeed it works. I was confused by the formulation "Whenever a label is repeated, it is summed" (einsum docs). I want only one summation but `j` and `k` are repeated. Therefor I only tried `np.einsum('ijk,jl->ik', A, b)` – cknoll Aug 01 '17 at 15:23
  • 1
    A few paragraphs further down says that `output` indexing, ` allows summing to be disallowed or forced when desired`. Putting `k` in the output turns off its summation, and it just 'goes along for the ride'. – hpaulj Aug 01 '17 at 16:43