trace(A'*A)

Given a (dense) matrix A, what would be a fast, idiomatic way of calculating \text{tr}(A'A)? I can solve it as

A = Float64.(reshape(1:20, :, 2))
trace(A*A')                     # = 2870, but obviously wasteful
sum(i -> (a = @view A[i, :]; dot(a, a)), 1:size(A, 1)) # works, convoluted

I imagine there must be a simple solution that is compiled to some neat BLAS function, but I can’t come up with anything.

It is Monday so I could be horribly wrong here but isn’t tr(A'A) the sum of the squares of the elements of A? If I am correct then you just need

sum(abs2, A)
5 Likes

Thanks. It is pretty clear that I should stop programming for today :smile:

1 Like

Just a note for clarity in case anyone else is reading this, the trace is the sum of the diagonal elements of A, not all of the elements.

Interestingly, for a reason that is not obvious to me, your solution does in fact give the correct trace though, whereas if we defined some square matrix

B = A*A'

then

trace(B)

is not equal to

sum(identity,B)

so I don’t see why

sum(abs2,A)

does give the right answer.

Edit:
I should have known better than to comment late at night on linear algebra. I completely forgot the product identities

Check sources such as Wikipedia to confirm that the trace of a product X*Y.' can be written as sum(X .* Y). The difference is, you are using sum(identity, B) where B = X*Y.' already. Then, you need to use sum(identity, diag(B)). In the solution provided, he is using A not A*A.', avoiding matrix multiplication and vectorizing the operations.

The trace of a matrix is simply the sum of the diagonal elements. The squaring is a result of the multiplication A’A.

Yikes. Now its me who needs to stop programming and go to bed. Forgot about the product identities and mistyped the definition of trace.