Operations on Arrays of CartesianIndex

I came across a problem that I found surprisingly hard, given Julias excellent support for indexing multidimensional Arrays.
I lilke to apply some function along one axis on parts of an Array. The parts are defined by some predicate:

julia> size(A)
(1272, 32, 14, 120)
julia> ix = findall(predicate, A)
8904-element Vector{CartesianIndex{4}}:
 CartesianIndex(1, 1, 5, 17)
 CartesianIndex(5, 1, 5, 17)
 ...

Now I want to find the triplets of the three final indices for which the predicate is true. Note that these triplets are not necessarily permutations of single indices, so that filtering the indices out one by one would give me to many.

julia> ix2 = unique([i[2] for i in ix]);
julia> ix3 = unique([i[3] for i in ix]);
julia> ix4 = unique([i[4] for i in ix]);
julia> prod(length.([ix2,ix3,ix4]))
540

I solved it by doing this:

Base.getindex(index::CartesianIndex, r::UnitRange{Int}) = index.I[r]

Then, I can do:

ix234 = unique(getindex.(ix, [2:4]))
30-element Vector{Tuple{Int64, Int64, Int64}}:
 (1, 5, 17)
 (2, 5, 17)
 ...
for (ix2, ix3, ix4) in ix234
   ix1 = getindex.(filter(x -> x[2:4] == (ix2, ix3, ix4), ix), 1)
   f(A[ix1, ix2, ix3, ix4])
end

My guess is that this is not an uncommon task, e.g. having one Array describing the contents of another Array with the same dimensions.
I suspect there is a more elegant way to accomplish this, but I can’t figure out how. Also, before posting I searched through other posts dealing with CartesianIndex and I understood that the Julia crew wants to avoid encouraging users to broadcast over the components of CartesianIndex. Does this solution violate this ?

So, my question is: Is this something that should go into Base ? Maybe accompanied with (which I also found handy):

Base.getindex(index::CartesianIndex, t::NTuple{N, Int}) where {N} = CartesianIndex([index[i] for i in t]...)
Base.lastindex(index::CartesianIndex) = lastindex(index.I)

I think what you’re writing here is a group-by operation, which several packages support, in slightly different ways, see #32331 for discussion.

julia> A = reshape(1:12, 3,2,2,1);

julia> ix = findall(iseven, A)
6-element Vector{CartesianIndex{4}}:
 CartesianIndex(2, 1, 1, 1)
 CartesianIndex(1, 2, 1, 1)
 CartesianIndex(3, 2, 1, 1)
 CartesianIndex(2, 1, 2, 1)
 CartesianIndex(1, 2, 2, 1)
 CartesianIndex(3, 2, 2, 1)

julia> unique(I -> (I[2], I[3]), ix)
4-element Vector{CartesianIndex{4}}:
 CartesianIndex(2, 1, 1, 1)
 CartesianIndex(1, 2, 1, 1)
 CartesianIndex(2, 1, 2, 1)
 CartesianIndex(1, 2, 2, 1)

julia> using SplitApplyCombine

julia> group(I -> I.I[2:end], ix).values
4-element Vector{Vector{CartesianIndex{4}}}:
 [CartesianIndex(2, 1, 1, 1)]
 [CartesianIndex(1, 2, 1, 1), CartesianIndex(3, 2, 1, 1)]
 [CartesianIndex(2, 1, 2, 1)]
 [CartesianIndex(1, 2, 2, 1), CartesianIndex(3, 2, 2, 1)]
1 Like

Thanks @mcabbot. I didn’t know of that form of unique invocation.