Is there an elegant way to sum columns in a Vector of Static arrays in Julia?

user_13952307 0
Viewed 228

I store velocities for particles as a Vector of SVectors. Each row is an SVector with that indices x, y, z velocity.

SVectors are great for their speed in arithmetic, but they are a bit... difficult to manipulate. How would I add up all of the x, y, z squared velocities in my Vector in a way that is more elegant than the following

using StaticArrays
n = 5
v = [SVector{3}(rand(), rand(), rand())  for i = 1:n]
x, y, z = 0.0, 0.0, 0.0
for i=1:n
    x += v[i][1]^2
    y += v[i][2]^2
    z += v[i][3]^2
end
sumv = SVector{3}(x, y, z)

If I just wanted to sum the x, y, z in my velocities, without squaring, Julia is simple, just sum(v) will give me a vector of the summed columns.

One solution I have is

sum([v[i].^2 for i=1:n])

but there must be a simpler solution that doesn't require a comprehension?

julia arrays
2 Answers

The sum function can take a higher order function that transforms each element before it adds them.

julia> sum(x->x.^2, v)
3-element SArray{Tuple{3},Float64,1,3} with indices SOneTo(3):
 2.104874346619557
 1.2512923674436118
 1.5781536056190726

This gives the same answer as your comprehension:

julia> sum(x->x.^2, v) == sum([v[i].^2 for i=1:n])
true
user_1760710

What about SMatrices?

n=5
mat = SMatrix{3, n}([rand() for i = 1:n*3])

sum(mat .^ 2, dims=2)
user_104592330
Related