Search code examples
foreachjuliareduction

Clean way to loop over a masked list in Julia


In Julia, I have a list of neighbors of a location stored in all_neighbors[loc]. This allows me to quickly loop over these neighbors conveniently with the syntax for neighbor in all_neighbors[loc]. This leads to readable code such as the following:

active_neighbors = 0
for neighbor in all_neighbors[loc]
    if cube[neighbor] == ACTIVE
       active_neighbors += 1
    end
end

Astute readers will see that this is nothing more than a reduction. Because I'm just counting active neighbors, I figured I could do this in a one-liner using the count function. However,

# This does not work
active_neighbors = count(x->x==ACTIVE, cube[all_neighbors[loc]])

does not work because the all_neighbors mask doesn't get interpreted correctly as simply a mask over the cube array. Does anyone know the cleanest way to write this reduction? An alternative solution I came up with is:

active_neighbors = count(x->x==ACTIVE, [cube[all_neighbors[loc][k]] for k = 1:length(all_neighbors[loc])])

but I really don't like this because it's even less readable than what I started with. Thanks for any advice!


Solution

  • This should work:

    count(x -> cube[x] == ACTIVE, all_neighbors[loc])