Dynamic "for"

Hello everyone!
How do I write the code below using @nloops or other Base.Cartesian elements? I’ve made several attempts. My goal is to generalize to n nested for’s.

function test()
	for i_1 = 0:1
		for i_2 = 0:1
			for i_3 = 0:1
				n = i_3*2^0 + i_2*2^1 + i_1*2^2
				println(n)
			end
		end
	end
end
1 Like

Something like this? There’s probably a cleaner way than Iterators.product, but it works.

f(i::NTuple{N,T}) where {N,T} = sum(((i, e),) -> i*2^e, zip(i, (N-1):-1:0))

@generated function test(::Val{N}) where {N}
    quote
        idx = collect(Iterators.product(fill(0:1, $N)...))
        Base.Cartesian.@nloops $N i idx begin
            n = f(Base.Cartesian.@nref $N idx i)
            @show n
        end
    end
end

test(N::Int) = test(Val(N))
julia> test(3)
n = 0
n = 4
n = 2
n = 6
n = 1
n = 5
n = 3
n = 7
1 Like

Here’s a slightly different approach - might not be quite as efficient but no metaprogramming

function power(tpl::NTuple{N, <: Integer}) where {N}
    ans = 0
    for i in 0:(N-1)
        ans += tpl[i+1] * 2 ^ i
    end
    ans
end

function test2(I::CartesianIndices)
    for i in I
        n = power(Tuple(i))
        println(n)
    end
end

julia> test2(CartesianIndices(ntuple(_ -> 0:1, 3)))
0
1
2
3
4
5
6
7

julia> test2(CartesianIndices(ntuple(_ -> 0:1, 4)))
0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
1 Like

Sorry, I didn’t post the desired result, which is: 0, 1, 2, …, 2^N-1, where N is the number of nested “for”. Right after getting the n, I need to use it in another function, not sure if I can use it out of order. Thanks stillyslalom.

Actually, your code is faster than @stillyslalom’s @generated example, and can be made faster and simpler still. (To benchmark, replace the println or @show statements with s += n, initialized to s=0, and return the sum s.) CartesianIndices are just as fast as @nloops if done properly, if somewhat less flexible (they are implemented using @nloops internally IIRC).

Consider:

function test3(v::Val)
    p2 = ntuple(i->2^(i-1), v)
    for idx in CartesianIndices(ntuple(i->0:1, v))
        n = sum(Tuple(idx) .* p2)
        @show n
    end
end

test3(Val(3))

Note that ntuple with a Val argument is completely unrolled and inlined by the compiler, and sum(Tuple(idx) .* p2) is also unrolled and inlined since the arguments are tuples with lengths known to the compiler.

5 Likes

Yeah, collect is a performance-killer here. Hopefully we’ll one day be able to index into an Iterators.product without first needing to collect.

I’m sorry, I misformulated the result. The indices are inverted and I don’t know how to invert them. The order in which the indices i_1, i_2, i_3 appear is important, because it represents the position in a chain of spins. The result needs to be:

function test()
	for i_1 = 0:1
		for i_2 = 0:1
			for i_3 = 0:1
				n = i_3*2^0 + i_2*2^1 + i_1*2^2
				println("$(i_1)$(i_2)$(i_3) => $n")
			end
		end
	end
end

000 => 0
001 => 1
010 => 2
011 => 3
100 => 4
101 => 5
110 => 6
111 => 7

And not

000 => 0
100 => 1
010 => 2
110 => 3
001 => 4
101 => 5
011 => 6
111 => 7

Just reverse the tuples?

function test3(v::Val{N}) where {N}
    p2 = reverse(ntuple(i->2^(i-1), v))
    for idx in CartesianIndices(ntuple(i->0:1, v))
        n = sum(reverse(Tuple(idx)) .* p2)
    end
end

I tested with N = 5 and found no performance penalty compared to @stevengj’s version.

2 Likes

Would that require a product method specialized for arguments with the HasMethod(getindex) trait?

Not inverting i_1, i_2, i_3

Correct result:

function test()
	for i_1 = 0:1
		for i_2 = 0:1
			for i_3 = 0:1
				n = i_3*2^0 + i_2*2^1 + i_1*2^2
				println("$(i_1)$(i_2)$(i_3) => $n")
			end
		end
	end
end

test()

000 => 0
001 => 1
010 => 2
011 => 3
100 => 4
101 => 5
110 => 6
111 => 7

Result still with i_1, i_2, i_3 inverted

function test2(v::Val{N}) where {N}
    p2 = reverse(ntuple(i->2^(i-1), v))
    for idx in CartesianIndices(ntuple(i->0:1, v))
        n = sum(reverse(Tuple(idx)) .* p2)
        println("$(idx) = > $n")
    end
end

test2(Val(3))

CartesianIndex(0, 0, 0) = > 0
CartesianIndex(1, 0, 0) = > 1
CartesianIndex(0, 1, 0) = > 2
CartesianIndex(1, 1, 0) = > 3
CartesianIndex(0, 0, 1) = > 4
CartesianIndex(1, 0, 1) = > 5
CartesianIndex(0, 1, 1) = > 6
CartesianIndex(1, 1, 1) = > 7

You need to use the reversed indices:

function test3(v::Val{N}) where {N}
    p2 = reverse(ntuple(i->2^(i-1), v))
    for idx in CartesianIndices(ntuple(i->0:1, v))
        ridx = reverse(Tuple(idx))
        n = sum(ridx .* p2)
        println(ridx[1], ridx[2], ridx[3], " => ", n)
    end
end
000 => 0
001 => 1
010 => 2
011 => 3
100 => 4
101 => 5
110 => 6
111 => 7
1 Like

------------ *** ---------------
Thank you very much stillyslalom. It worked here and in my code. Also, now I have several examples to help me understand Base.Cartesian. Thank you everyone very much.