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?
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)
(A*B)*B.T to A*(B*B.T)A[:B.shape[0]]) that would result in the diagonal part of the matriximport 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
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
Say we have: K = np.dot(np.dot(A,B),B.T).
Then, K[0,0] = (A[0, :] * B[:,0])[0, :] * B.T[:])
X = (A[0, :] * B[:,0]), which is the [0, 0] element of np.dot(A,B)X[0, :] * B.T[:, 0] is the [0, 0] element of np.dot(np.dot(A,B),B.T)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])