Enumerate multi dimensional array

Is there a way to use enumerate of a multi-dimensional array, something like:

function new_matrix(matrix)
    new_matrix = copy(matrix)
    for (i,j,element) in enumerate(matrix)
        new_matrix[i,j] = i*j*matrix[i,j]
    end
    return new_matrix
end

?

1 Like

You may look for
axes(matrix)

axes doesn’t include the value, unlike enumerate. And actually, it’s not enumerate you want, even for one-dimensional vectors. You should use pairs, since enumerate always starts counting at 1 no matter what sort of array you have (even 0-indexed arrays).

julia> a = randn(3,2)
3×2 Array{Float64,2}:
  0.446886    0.168215
 -0.0590402   1.89755
  0.0420936  -1.62462

julia> for (ind, val) in pairs(a)
       println(ind, ": ", val)
       end
CartesianIndex(1, 1): 0.4468857733868482
CartesianIndex(2, 1): -0.0590401519075785
CartesianIndex(3, 1): 0.04209357154257153
CartesianIndex(1, 2): 0.16821454509635914
CartesianIndex(2, 2): 1.8975549461332841
CartesianIndex(3, 2): -1.624623544140236

You can get (i, j) from calling Tuple(ind), but you should rather use the cartesian index directly:

function new_matrix(matrix)
    new_matrix = copy(matrix)
    for (ind, element) in pairs(matrix)
        (i, j) = Tuple(ind)
        new_matrix[ind] = i * j * matrix[ind]  # or ind[1]*ind[2] * matrix[ind]
    end
    return new_matrix
end
6 Likes

Maybe:

function new_matrix(matrix::AbstractMatrix)
    new_matrix = similar(matrix)
    for ij in CartesianIndices(matrix)
        new_matrix[ij] = ij[1]*ij[2]*matrix[ij]
    end
    return new_matrix
end

You could do for (ij, val) in zip(CartesianIndices(matrix), matrix) to be closer to enumerate, but it’s probably not worth enumerate is mainly useful for non-indexable collections. Update: there is also pairs, as mentioned by @DNF above.

3 Likes

BTW, there are some not-so-nice things here. Don’t give your output matrix the same name as the function: new_matrix, that’s probably a bad idea. And a tip: use similar(matrix) instead of copy(matrix), it can be a lot faster.

Here’s some code that actually uses the pairs feature better:

function new_matrix(matrix)
    matrix_ = similar(matrix)
    for (ind, element) in pairs(matrix)
        matrix_[ind] = ind[1] * ind[2] * element 
    end
    return matrix_
end

And this code works for all array dimensionalities from 0 to N

function new_arr(arr)
    arr_ = similar(arr)
    for (ind, element) in pairs(arr)
        arr_[ind] = prod(Tuple(ind)) * element
    end
    return arr_
end
4 Likes

Even shorter:

new_matrix(matrix) = [ij[1]*ij[2]*matrix[ij] for ij in CartesianIndices(matrix)]

which has the advantage of more flexibly computing the return type. (e.g. if you have a Matrix{Bool} it will return a Matrix{Int} thanks to promotion.)

4 Likes

Surprisingly pairs can’t be used with map, otherwise this might be another neat way:

julia> map(enumerate(rand(Bool, 2,10))) do (i,v)
         i * v
       end
2×10 Matrix{Int64}:
 0  0  5  0   9  11   0  0   0  0
 2  0  0  8  10   0  14  0  18  0

julia> map(pairs(rand(Bool, 2,10))) do (i,v)
         prod(i.I) * v
       end
ERROR: map is not defined on dictionaries
1 Like

And for arbitrary dimensions:

new_arr(x) = [prod(Tuple(ind)) * element for (ind, element) in pairs(x)]
2 Likes

Thanks everyone, this is really helpful!

It is indeed weird that map deliberately does not support pairs. Last year I asked a question about this in a slightly different context of NamedTuples: Map over namedtuple keys and values: confusing error.

Thanks, I hadn’t seen the thread from last year. I made https://github.com/JuliaLang/julia/pull/38150 and we will see what people think.

1 Like

Looking at the code above and without taking credit from the authors of all the clever solutions, for the lazy the fast and intuitive TensorCast is a friend, it trivially enumerates and provides elementwise access:

# N-dimensional very fast solution by @DNF
function new_arr(arr)
    arr_ = similar(arr)
    for (ind, element) in pairs(arr)
        arr_[ind] = prod(Tuple(ind)) * element
    end
    return arr_
end

a = rand(4,3,2)
b = new_arr(a)

# Lazy solution using TensorCast (not as fast but versatile and easy to remember):
using TensorCast
@cast c[i,j,k] := i*j*k*a[i,j,k]

c == b  # true

Merry Christmas! :christmas_tree: