Indexing array with a matrix and scalar

General Usage
, ,
1

A bit rusty after a while away from julia… How would I fill a 3D array with N values at positions given as follows:

  • 1st index given as a scalar (from a for loop)
    • the other two indices are provided as vectors

I’m using nested for loops but this feels very inelegant. Preferred built-in solutions however, as the script should have no dependencies.

u = [1 3 4 5 6 8]
up = [2 3 4 5 6 9]
values = [1.0im 3.0im 2.0 1.0 5.4im 6.0im+2]
qmax = 9
Nl = 4
dummy = zeros(ComplexF64,(Nl,qmax,qmax))
for i in 1:Nl
    for iu in 1:length(u)
        dummy[i,u[iu],up[iu]] = values[iu]
    end
end

dummy
2

One easy simplification:

for iu in 1:length(u)
    dummy[1:Nl, u[iu], up[iu]] .= values[iu]
end
3

Thanks, that loop over i I’m actually OK with (it makes the code easier to understand in my case). I don’t know how to index dummy with a matrix to inject all the values at once or this i-th slice.

4

Don’t you want dummy[1, u, up] .= values ?

5

I’d like that kind of syntax, yes, but it returns a very different result

6

Ah right, broadcast assignment replicates the values in all the positions given in vectors, while you want no replication, so try this instead:

for i in 1:Nl
    setindex!.(Ref(dummy), values, i, u, up)
end

(though I find the solution with two loops no less elegant)

7

If efficiency is of interest, here is the comparative timing:

u =  [1 3 4 5 6 8]
up = [2 3 4 5 6 9]
values = [1.0im 3.0im 2.0 1.0 5.4im 6.0im+2]
qmax = 9
Nl = 4
dummy = zeros(ComplexF64,(Nl,qmax,qmax))

function f1!(mat, Nl, u, up, values)
    for i in 1:Nl
        for iu in 1:length(u)
            mat[i,u[iu],up[iu]] = values[iu]
        end
    end
    return mat
end

function f2!(mat, Nl, u, up, values)
    for iu in 1:length(u) 
        mat[1:Nl, u[iu], up[iu]] .= values[iu]
    end
    return mat
end

function f3!(mat, Nl, u, up, values)
    for i in 1:Nl
        setindex!.(Ref(mat), values, i, u, up)
    end
    return mat
end

using BenchmarkTools

m1 = @btime f1!($dummy, $Nl, $u, $up, $values); # 41.070 ns (0 allocations: 0 bytes)
m2 = @btime f2!($dummy, $Nl, $u, $up, $values); # 27.182 ns (0 allocations: 0 bytes)
m3 = @btime f3!($dummy, $Nl, $u, $up, $values); # 271.935 ns (4 allocations: 384 bytes)
m1 == m2 == m3 # true

f1! can be made as fast as f2! by putting @inbounds in front of the outer loop. But eliding bounds checking like this can be dangerous.

1 Like
8

A little bit safer version of f2!:

function f2new!(mat, Nl, u, up, values)
    for iu in eachindex(u, up, values) 
        mat[1:Nl, u[iu], up[iu]] .= values[iu]
    end
    return mat
end
m2new = @btime f2new!($dummy, $Nl, $u, $up, $values); # 26.506 ns (0 allocations: 0 bytes)
9

While I don’t really recommend loading a package to do this once, if you have many such things Tullio.jl can write the loops for you, from a 1-line expression:

julia> using Tullio

julia> dummy2 = zero(dummy);

julia> dummy == @tullio dummy2[i,u[iu],up[iu]] = values[iu]
true

julia> f4!(dummy, _, u, up, values) = @tullio dummy[i,u[iu],up[iu]] = values[iu];

julia> m4 = @btime f4!($dummy, $Nl, $u, $up, $values);  # similar to f1!
  35.204 ns (0 allocations: 0 bytes)
2 Likes
10

Thanks for the suggestions and timings. I was still curious to find an option with matrix indexing. In matlab I would probably use sub2ind to find the linear indices of the matrix where the values should go. I found a discussion on its Julia equivalent being removed, and suggestion to use CartesianIndices or LinearIndices instead. I don’t fully get how they work, but this seems reasonable:

dummy2 = zeros(ComplexF64,(Nl,qmax,qmax))
for i in 1:Nl
  dummy2[CartesianIndex.(i, u, up)] .= values
end