I have a simple algebric problem and I would like to solve it with numpy (of course that I could solve it easily with numba, but that is not the point).

Let us consider a first random matrix A with size (m x n), with n a big value, and a second random matrix B with size (n x n).

A = np.random.random((1E6, 1E2))
B = np.random.random((1E2, 1E2))

We want to compute the following expression:

np.diag(np.dot(np.dot(A,B),B.T))

The problem is that the entire matrix is loaded to the memory and only then is extracted the diagonal. Is it possible to do this operation in a more efficient way?

CodePudding user response：

Yes.

N = 1E6

A = np.random.random((N, 1E2))
B = np.random.random((1E2, 1E2))

result = 0;
for i in range(N):
  result  = np.dot(np.dot(A[i,:], B[i,:])[i, :], B.T[i, :])

  # Replacing B.T[i, :] with B[:, i].T might be a little more efficient

Explanation:

Say we have: K = np.dot(np.dot(A,B),B.T).

Then, K[0,0] = (A[0, :] * B[:,0])[0, :] * B.T[:])

• Let X = (A[0, :] * B[:,0]), which is the [0, 0] element of np.dot(A,B)
• Then X[0, :] * B.T[:, 0] is the [0, 0] element of np.dot(np.dot(A,B),B.T)
• Then X[0, :] * B.T[:, 0] = (A[0, :] * B[:,0])[0, :] * B.T[:])

We can also generalize this result to: K[i,i] = (A[i, :] * B[:,i])[i, :] * B.T[:, i])

CodePudding user response：

1. Changed (A*B)*B.T to A*(B*B.T)
2. Multiplied only this part of A (A[:B.shape[0]]) that would result in the diagonal part of the matrix

import numpy as np
import time

A = np.random.random((1000_000, 100))
B = np.random.random((100, 100))

start_time = time.time()
result = np.diag(np.dot(np.dot(A, B), B.T))
print('Baseline: ', time.time() - start_time)

start_time = time.time()
for i in range(100):
    result2 = np.diag(np.dot(A[:B.shape[0]], np.dot(B, B.T)))
print('Optimized: ', (time.time() - start_time) / 100)
stop = 1

assert np.allclose(result, result2)

Baseline:  1.7957241535186768
Optimized:  0.00016015291213989258

CodePudding user response：

This is how I would approach it from your starting expression

np.diag(np.dot(np.dot(A,B),B.T))

You can start by grouping terms:

np.diag(np.dot(A, np.dot(B,B.T)))

then only use the first relevant (square) part of A:

np.diag(np.dot(A[:B.shape[0], :], np.dot(B,B.T)))

and then avoid the extra multiplications (that will fall out of the diagonal), by doing the element-wise multiplications yourself:

np.sum( np.multiply(A[:B.shape[0], :].T, np.dot(B,B.T)), 0)