Numpy dot product with 3d array

1
votes

I've got two arrays:

  • data of shape (2466, 2498, 9), where the dimensions are (asset, date, returns).
  • correlation_matrix of shape (2466, 2466) (with 0's on the diagonal)

I want to get the dot product that equates to the expected returns, which is the returns of each asset multiplied by the correlation_matrix. It should give a shape the same as data.

I've tried:

data.transpose([1, 2, 0]) @ correlation_matrix

but this just hangs my PC (been going 10 minutes and counting).

I also tried:

np.einsum('ijk,lm->ijk', data, correlation_matrix)

but I'm less familiar with einsum, and this also hangs.

What am I doing wrong?

2
numpy
Think you can just do data*correlation_matrix.sum(), assuming your einsum works. - Divakar
If things are hanging or taking too long, step back and test something smaller. Make sure the code is doing what you want with small arrays before stressing memory with something large. - hpaulj
Your einsum just sums all values of correlation_matrix and multiplies data by the resulting scalar. That probably not what you want. - hpaulj

2 Answers

1
votes

There are different ways to do this product:

# as you already suggested:
data.transpose([1, 2, 0]) @ correlation_matrix

# using einsum
np.einsum('ijk,il', data, correlation_matrix)

# using tensordot to explicitly specify the axes to sum over
np.tensordot(data, correlation_matrix, axes=(0,0))

All of them should give the same result. The timing for some small matrices was more or less the same for me. So your problem is the large amount of data, not an inefficient implementation.

A=np.arange(100*120*9).reshape((100, 120, 9))
B=np.arange(100**2).reshape((100,100))

timeit('A.transpose([1,2,0])@B', globals=globals(), number=100)
# 0.747475513999234
timeit("np.einsum('ijk,il', A, B)", globals=globals(), number=100)
# 0.4993825999990804
timeit('np.tensordot(A, B, axes=(0,0))', globals=globals(), number=100)
# 0.5872082839996438
2
votes

With your .transpose((1, 2, 0)) data, the correct form is:

"ijs,sk"  # -> ijk

Since for a tensor A and B, we can write:

C_{ijk} = Σ_s A_{ijs} * B_{sk}

If you want to avoid transposing your data beforehand, you can just permute the indices:

"sij,sk"  # -> ijk

To verify:

p, q, r = 2466, 2498, 9

a = np.random.randint(255, size=(p, q, r))
b = np.random.randint(255, size=(p, p))

c1 = a.transpose((1, 2, 0)) @ b
c2 = np.einsum("sij,sk", a, b)

>>> np.all(c1 == c2)
True

The amount of multiplications needed to compute this for (p, q, r) shaped data is p * np.prod(c.shape) == p * (q * r * p) == p**2 * q * r. In your case, that is 136_716_549_192 multiplications. You also need approximately the same number of additions, so that gives us somewhere close to 270 billion operations. If you want more speed, you could consider using a GPU for your computations via cupy.

def with_np():
    p, q, r = 2466, 2498, 9
    a = np.random.randint(255, size=(p, q, r))
    b = np.random.randint(255, size=(p, p))
    c1 = a.transpose((1, 2, 0)) @ b
    c2 = np.einsum("sij,sk", a, b)

def with_cp():
    p, q, r = 2466, 2498, 9
    a = cp.random.randint(255, size=(p, q, r))
    b = cp.random.randint(255, size=(p, p))
    c1 = a.transpose((1, 2, 0)) @ b
    c2 = cp.einsum("sij,sk", a, b)

>>> timeit(with_np, number=1)
513.066

>>> timeit(with_cp, number=1)
0.197

That's a speedup of 2600, including memory allocation, initialization, and CPU/GPU copy times! (A more realistic benchmark would give an even larger speedup.)