How to use nested list comprehension in julia

user_16021234 0
Viewed 187

How do I make the following code into nested list comprehension?

node_x = 5
node_y = 5
node_z = 5
xyz = Matrix(undef, node_x*node_y*node_z,3)
ii = 0
dx = 1.0
    for k in 0:node_z-1
        for j in 0:node_y-1
            for i in 0:node_x-1
                x = i * dx
                y = j * dx
                z = k * dx
                ii += 1
                #println([x, y, z])
                xyz[ii, 1] = x
                xyz[ii, 2] = y
                xyz[ii, 3] = z
            end
        end
    end

In python and numpy, I can write such as following codes.

xyz = np.array([[i*dx, j*dx, k*dx] for k in range(node_z) for j in range(node_y) for i in range(node_x)])
julia
1 Answers

Comprehensions can be nested just the same, it's just range that is a bit different, but in your case there is the start:end syntactic sugar:

julia> [[i*dx, j*dx, k*dx] for k in 1:node_z for j in 1:node_y for i in 1:node_x]
125-element Vector{Vector{Float64}}:
 [1.0, 1.0, 1.0]
 [2.0, 1.0, 1.0]
 ⋮
 [4.0, 5.0, 5.0]
 [5.0, 5.0, 5.0]

To get the same array as your Python example, you'd have to permute the dimensions of the 3-element vectors and concatenate the list:

julia> vcat(([i*dx j*dx k*dx] for k in 1:node_z for j in 1:node_y for i in 1:node_x)...)
125×3 Matrix{Float64}:
 1.0  1.0  1.0
 2.0  1.0  1.0
 ⋮
 4.0  5.0  5.0
 5.0  5.0  5.0
user_56021490
Related