`findall` from a vector of CartesianIndices in Julia

Question

CI is the given vector. Each element is CartesianIndex.

julia> CI
5-element Vector{CartesianIndex{3}}:                                     
 CartesianIndex(1, 3, 1)                                                
 CartesianIndex(1, 1, 4)                                                 
 CartesianIndex(2, 2, 3)                                                 
 CartesianIndex(1, 2, 4)                                                 
 CartesianIndex(1, 1, 1)  

I want to find elements in CI for which the d-th element is k. Let f(CI,d,k) be such a function.

julia> f(CI,1,2)
[CartesianIndex(2, 2, 3) ]

julia> f(CI,2,1)
[CartesianIndex(1, 1, 4), CartesianIndex(1, 1, 1)]

How can I define such a function f? I guess a naive way is converting CI to a matrix array, then use findall function. But I feel there is a more clever way.

Answer 1

There are many ways in Julia to accomplish this task. Here are a few examples:

One method of the filter function takes as its first argument an arbitrary unary function that is passed each element of the collection in question and, if the function returns false, removes that element from the collection. CartesianIndex values can be indexed like Tuples, so use this to find which elements of CI have a value of value in dimension d:

julia> CI = [CartesianIndex(1,3,1), CartesianIndex(1,1,4), CartesianIndex(2,2,3), CartesianIndex(1,2,4), CartesianIndex(1,1,1)]
5-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 3, 1)
 CartesianIndex(1, 1, 4)
 CartesianIndex(2, 2, 3)
 CartesianIndex(1, 2, 4)
 CartesianIndex(1, 1, 1)

julia> CI[2]
CartesianIndex(1, 3, 1)

julia> CI[2][1]
1

julia> filter(idx -> idx[2] == 1, CI)
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)

julia> f1(ci, d, value) = filter(idx -> idx[d] == value, ci)
f1 (generic function with 1 method)

julia> f1(CI, 1, 2)
1-element Vector{CartesianIndex{3}}:
 CartesianIndex(2, 2, 3)

julia> f1(CI, 2, 1)
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)

You can even use the in-place filter! function to avoid allocating and returning a copy of the input array (assuming you do not need the array again later):

julia> CIcopy = copy(CI)
5-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 3, 1)
 CartesianIndex(1, 1, 4)
 CartesianIndex(2, 2, 3)
 CartesianIndex(1, 2, 4)
 CartesianIndex(1, 1, 1)

julia> f1!(ci, d, value) = filter!(idx -> idx[d] == value, ci)
f1! (generic function with 1 method)

julia> f1!(CIcopy, 2, 1)
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)

julia> CIcopy  # proving it was modified in place
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)

The findall function can be used like filter, but it will be slower due to the temporary array that has to be constructed in the process:

julia> findall(idx -> idx[2] == 1, CI)
2-element Vector{Int64}:
 2
 5

julia> f2(ci, d, value) = ci[findall(idx -> idx[d] == value, ci)]
f2 (generic function with 1 method)

julia> f2(CI, 1, 2)
1-element Vector{CartesianIndex{3}}:
 CartesianIndex(2, 2, 3)

julia> f2(CI, 2, 1)
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)

You can also broadcast the getindex method and the == operator and use logical indexing of vectors to the same effect:

julia> getindex(CI[1], 2)
3

julia> getindex.(CI, 2)
5-element Vector{Int64}:
 3
 1
 2
 2
 1

julia> getindex.(CI, 2) .== 1
5-element BitVector:
 0
 1
 0
 0
 1

julia> f3(ci, d, value) = ci[getindex.(ci, Ref(d)) .== value]
f3 (generic function with 1 method)

julia> f3(CI, 1, 2)
1-element Vector{CartesianIndex{3}}:
 CartesianIndex(2, 2, 3)

julia> f3(CI, 2, 1)
2-element Vector{CartesianIndex{3}}:
 CartesianIndex(1, 1, 4)
 CartesianIndex(1, 1, 1)