How do I use an array element as an iterator in julia?

Question

I'm trying to make my Julia code more understandable (to physicists) and thought it would be nice if I could use some sort of vector type iterator. I'm trying to use each element as an iterator. My solution so far is:

kcut=2
𝐌 = Array{Float64}(undef,3)
Ø = [0, 0, 0]

for 𝐌[1] in -kcut:kcut, 𝐌[2] in -kcut:kcut, 𝐌[3] in -kcut:kcut
    if norm(𝐌)^2 <= kcut^2 && 𝐌 != Ø
        println(𝐌)
    end
end
println(𝐌)

which almost does what it's supposed to. The only problem is once I've finished the for loop my M is redefined to the last configuration it took in the loop, [2,2,2] int this case. This does not seem to be the case when I iterate over a normal variable as below

x=1
for x in 1:10
    println(x)
end
println(x)

Which tells me x is still equal to 1 after the loop. I would like to have the same behavior if possible. Also is there any way to do this without defining M before I use it to iterate?

EDIT: I need M to be outputted as an array so I can do some Linear algebra on it.

Answer 1

You can maybe iterate directly on M, for example like this:

julia> using LinearAlgebra: norm
julia> kcut=2;
julia> Ø = [0, 0, 0];

julia> for 𝐌 in Iterators.product(-kcut:kcut, -kcut:kcut, -kcut:kcut)
           if norm(𝐌)^2 <= kcut^2 && 𝐌 != Ø
               # 𝐌 is a Tuple now, which should be best in most cases.
               #
               # If you want arrays instead:
               𝐍 = collect(𝐌)

               println(𝐌, "\t", 𝐍)
           end
       end
(0, 0, -2)      [0, 0, -2]
(-1, -1, -1)    [-1, -1, -1]
(0, -1, -1)     [0, -1, -1]
(1, -1, -1)     [1, -1, -1]
(-1, 0, -1)     [-1, 0, -1]
(0, 0, -1)      [0, 0, -1]
(1, 0, -1)      [1, 0, -1]
(-1, 1, -1)     [-1, 1, -1]
...

julia> println(𝐌)
ERROR: UndefVarError: 𝐌 not defined
Stacktrace:
 [1] top-level scope at REPL[5]:1

This way, M is only defined within the scope of the for loop.

To understand the behavior of your original code, you have to have a basic understanding of the code transformations ("code lowering") that Julia internally performs. In particular, the for loop is replaced following the rules of the iteration interface

All in all, a snippet like the following:

M = [42]
for M[1] in -k:k
    println(M[1])
end

behaves in the same way as:

M = [42]
next = iterate(-k:k)
while next !== nothing
    M[1] = i
    println(M[1])
    next = iterate(iter, state)
end

Where you see that M has to pre-exist so that M[1]=i does not fail, and there is no reason why what happens inside the loop body does not persist after it.

Answer 2

You can do the following with StaticArrays to get "tuples that can do linear algebra":

julia> using StaticArrays

julia> ⊗(u, v) = (i ⊗ j for i in u for j in v)
⊗ (generic function with 1 method)

julia> ⊗(x::Number, y::Number) = SVector(x, y)
⊗ (generic function with 2 methods)

julia> ⊗(x::SVector{N}, y::Number) where {N} = SVector(x..., y)
⊗ (generic function with 3 methods)

julia> collect((1:3) ⊗ (10:12) ⊗ (100:101))
18-element Array{SArray{Tuple{3},Int64,1,3},1}:
 [1, 10, 100]
 [1, 10, 101]
 [1, 11, 100]
 [1, 11, 101]
 [1, 12, 100]
 [1, 12, 101]
 [2, 10, 100]
 [2, 10, 101]
 [2, 11, 100]
 [2, 11, 101]
 [2, 12, 100]
 [2, 12, 101]
 [3, 10, 100]
 [3, 10, 101]
 [3, 11, 100]
 [3, 11, 101]
 [3, 12, 100]
 [3, 12, 101]

julia> using LinearAlgebra: norm

julia> for M in (1:3) ⊗ (10:12) ⊗ (100:101)
           println(norm(M))
       end
100.50373127401788
101.49876846543509
100.60815076324582
101.6021653312566
100.72239075796404
101.7152889196113
100.5186549850325
101.51354589413178
100.62305898749054
101.61692772368194
100.73728207570423
101.73003489628813
100.54352291420865
101.5381701627521
100.64790112068906
101.64152694642087
100.7620960480676
101.75460677532

But I'm not sure it's worth it. Ideally, the \otimes for SVectors and Numbers would construct a lazy data structure that only creates the appropriately sized SVectors when iterated (instead of splatting as here). I'm just too lazy to write that now.

A better variant (but slightly less mathy syntax) is to overload ⨂(spaces...) to do everything at once:

julia> ⨂(spaces::NTuple{N}) where {N} = (SVector{N}(t) for t in Iterators.product(spaces...))
⨂ (generic function with 1 method)

julia> ⨂(spaces...) = ⨂(spaces)
⨂ (generic function with 2 methods)

julia> collect(⨂(1:3, 10:11))
3×2 Array{SArray{Tuple{2},Int64,1,2},2}:
 [1, 10]  [1, 11]
 [2, 10]  [2, 11]
 [3, 10]  [3, 11]

julia> collect(⨂(1:3, 10:11, 100:101))
3×2×2 Array{SArray{Tuple{3},Int64,1,3},3}:
[:, :, 1] =
 [1, 10, 100]  [1, 11, 100]
 [2, 10, 100]  [2, 11, 100]
 [3, 10, 100]  [3, 11, 100]

[:, :, 2] =
 [1, 10, 101]  [1, 11, 101]
 [2, 10, 101]  [2, 11, 101]
 [3, 10, 101]  [3, 11, 101]

This collects to a different shape, although I'd consider that even more appropriate.