newtothis
newtothis

Reputation: 525

Alternatives to `for` loops in this context in Julia

I have a matrix in Julia which has 0,1 as elements. Suppose I consider row i (at random), then I want to go through all the elements of the row, and for column j where the element is 1, I want to performing an arithmetic operation involving elements i,j of a different list. I want this done for all possible 'j' columns in the matrix. Currently I have the following code for an NxN square matrix mat

e = 0
for j in (1:N)
     if mat[i,i] == 1 
          e += -(list[i]*list[j])
     end
end

However, I need to perform this for a reasonably large matrix, over and over again, so I would like to know if there is faster way to do this. I think I may be able to use Iterators but I don't really know how to do that.

Upvotes: 0

Views: 309

Answers (1)

ahnlabb
ahnlabb

Reputation: 2162

Julia will be able to optimise that code rather well, however we can simplify and improve performance somewhat. Since you mention iterators let's start with a solution using an iterator comprehension where we use sum to sum the values and factor out the common factor list[i]:

list[i]*sum(list[j] for j in 1:length(list) if mat[i,j] == 1)

This is more or less interpreted by julia as:

-list[i]*sum(Iterators.map(j->list[j], Iterators.filter(j -> mat[i,j] == 1, 1:length(list))))

but the comprehension is more concise and may look more familiar to you.

This is slightly more performant than the original (20%-30% time savings), however, iterators are often not the fastest solution unless memory is a bottleneck. Let's try broadcasting instead:

-list[i]*sum(list[mat[i,:] .== 1])

This increases performance further (40% savings from original in my tests).

We don't need a copy of mat[i,:] so let's use a view:

list[i]*sum(list[@view(mat[i,:]) .== 1])

Now we've shaved off more than 50%. Of course it might be simpler to just use a product (since we have 0s and 1s) and performance is similar:

-list[i]*sum(list .* mat[i,:])

Further performance gains should be possible if you can arrange your data differently, e.g. use a BitMatrix instead of numbers and use a transposed matrix so we iterate over columns instead of rows (since Julia uses column-major order).

Benchmarks:

using BenchmarkTools

function original(mat, list, i)
    N = length(list)
    e = 0
    for j in (1:N)
         if mat[i,j] == 1
              e += -(list[i]*list[j])
         end
    end
    e
end

function sol1(mat, list, i)
    -list[i]*sum(list[j] for j in 1:length(list) if mat[i,j] == 1)
end

function sol2(mat, list, i)
    -list[i]*sum(list[mat[i,:] .== 1])
end

function sol3(mat, list, i)
    -list[i]*sum(list[@view(mat[i,:]) .== 1])
end

function sol4(mat, list, i)
    -list[i]*sum(list .* mat[i,:])
end

mat = rand(0:1, 1000, 1000)
list = rand(1000)
i = rand(1:1000)

@btime original($mat, $list, $i)
@btime sol1($mat, $list, $i)
@btime sol2($mat, $list, $i)
@btime sol3($mat, $list, $i)
@btime sol4($mat, $list, $i)

Upvotes: 2

Related Questions