How to calculate sum of a matrix by row and store in pre-allocated diagnoal matrix

basically what I want to do is calculate the sum of a matrix, say u, by column and store it in a pre-allocated diagonal matrix, say du, in a more efficient way than

vi = sum(u, dims=2)
for i = 1:length(vi)
    du[i, i] = vi[i]

the whole idea is that i need to do an operation:
usg* A1*u + Deg*A2*u
A1 and A2 do the first and second derivation. This is a typical operation in the simulation of the convection-diffusion phenomenon. If usg is a diagonal matrix or a constant, the equation can be calculated with inplace operation (with an intermediate vector du )
mul!( lmul!(usg, mul!( du, A1, u)), A2, u, Deg, true)

In my case, Deg can be seen as constant. But usg is no constant and calculated from u by
usg = sum(u, dims=2)*const ./v1.
As v1 is a constant vector, first I can make a diagnoal matrix
rv1 = Diagnoal(1 ./v1)
usg = sum(u, dims=2)*const ./v1 = rv1*sum(u, dims=2)*const
and can be calculated using inplace operator
mul!(usg, rv1, lmul!(const, du) (take du = sum(u, dims=2) is the intermediate vector)
But now usg is a vector because sum(u, dims=2) is an vector.

So is there a way to make sum(u, dims=2) to store in a diagonal matrix? or is there a simpler way to realize the whole operation ( usg* A1*u + Deg*A2*u ) ?

thanks for reading and giving advice.

If A is a matrix and D is a pre-allocated Diagonal matrix of an appropriate size, Base.mapreducedim!(identity, Base.add_sum, transpose(D.diag .= 0), A) should work.

Or you can just write a nested loop, of course, which is probably clearer than using undocumented internals.

1 Like

You can just use sum!:

julia> D = Diagonal(rand(3));

julia> u = reshape(1:9, 3, 3);

julia> sum!(D.diag, u);

julia> D
3×3 Diagonal{Float64, Vector{Float64}}:
 12.0    ⋅     ⋅ 
   ⋅   15.0    ⋅ 
   ⋅     ⋅   18.0

julia> vi = sum(u, dims=2)
3×1 Matrix{Int64}:

julia> Base.mapreducedim!(identity, Base.add_sum, transpose(D.diag .= 0), u)
1×3 transpose(::Vector{Float64}) with eltype Float64:
 6.0  15.0  24.0

Here’s a solution that doesn’t access the internal fields of D, in case you care about this:

sum!(view(D, diagind(D)), u)

It is significantly slower than sum!(D.diag, u), though. I cannot seem to find an efficient view of a Diagonal matrix’es diagonal (aside from accessing the diag field.)

Doh, I forgot we had that! I wonder why we don’t have a (documented) mapreduce! function?


In BandedMatrices.jl there’s view(A, band(0)) but it may not be that efficient

Thanks a lot for the discussion. Really helpful!