How to sum elements raised to a power

Let’s assume that X is a matrix. I want to efficiently compute sum(X.^2) but without computing X.^2 which will allocate some memory for nothing. So I’m trying to do something like sum(^(2),X) which throws a method error. I thought this could work because sum(>(0.2),X) works for instance.
I know that I could easily create an auxillary function f(x)=x^2 and then compute sum(f,X). But it’s not the best in terms of readability. So my question is, is there a way to properly use ^(x, y) to achieve my purpose?

Use an anonymous function.

sum(x -> x^2, X)
5 Likes

You could also do

sum(abs2, X)

which will handle both real- and complex-valued matrices, and should have the same performance as x->x^2.

3 Likes

I use abs2 simply for ease in this type of situation

You can have a faster sum (not as safe as the default sum, which takes some care to not overflow things), using a parallel sum with @tturbo:

julia> function mysum(f,x)
           s = zero(eltype(f(x[begin])))
           @tturbo for i in eachindex(x)
               s += f(x[i])
           end
           return s
       end
mysum (generic function with 2 methods)

julia> x = rand(1000,1000);

julia> @btime mysum(abs2,$x) # 4 threads (julia -t4)
  83.500 μs (0 allocations: 0 bytes)
333115.0216837497

julia> #vs
       @btime sum(abs2,$x)
  327.982 μs (0 allocations: 0 bytes)
333115.0216837495

If that is within a peformance critical part of the code and the numbers are well behaved, it is a good alternative.

2 Likes

While the abs2 solution is good enough for this case, I think generalizing the ^(n) would be a nice addition to Julia, the same as we do sqrt.

Note that abs2 is actually different from x -> x^2 for complex numbers.

julia> X = rand(ComplexF64, 2, 2)
2×2 Matrix{ComplexF64}:
 0.681033+0.323086im   0.756258+0.454746im
 0.704345+0.603132im  0.0489073+0.595638im

julia> sum(X.^2)
0.5044940375620476 + 2.035762251252853im

julia> sum(abs2, X)
2.563956079161798

julia> sum(x -> x^2, X)
0.5044940375620476 + 2.035762251252853im

Yeah, that’s sort of what I meant by saying that it also handles complex values. But if the OP actually wants to sum squared complex values, my suggestion won’t fly.

what about LinearAlgerbra.norm?

I don’t think you need eltype here, btw.

1 Like