Calculate a trace of large matrix in python

Question

I have a matrix X and I need to write a function, which calculate a trace of matrix .

I wrote a next script:

import numpy as np
def test(matrix):
    return (np.dot(matrix, matrix.T)).trace()

np.random.seed(42)
matrix = np.random.uniform(size=(1000, 1))

print(test(matrix))

It works fine on small matrix, but when I try to calculate on large matrix (for example on matrix with shape (50000, 1)), it gives me a memory error.

I tried to find a solution to the problem in other questions on the site, but nothing helped me. I would be grateful for any advice!

Answer 1

The number you're trying to compute is just the sum of the squares of all entries of X. Sum the squares instead of computing a giant matrix product full of entries you don't want:

return (X**2).sum()

Or ravel the matrix and use dot, which is probably faster for contiguous X:

raveled = X.ravel()
return raveled.dot(raveled)

Actually, ravel is probably faster for non-contiguous X, too - even when ravel needs to copy, it's not doing more allocation than (X**2).sum().