How to use CartesianIndices to replace two nested loops

General Usage
,
1

I’m trying to make a scatter plot where the x,y pairs form a (not necessarily rectangular) grid. The color of each point depends in a rather complicated way on the index of the point. Given two vectors, a1 and a2, my points are n*a1+m*a2, where n and m each span some integer range.

I made this work by using two nested loops and computing a single index for each point, but I can’t help but think that CartesianIndices is perfect for this use case. Despite spending a couple of hours, I couldn’t figure out how to use CartesianIndices. Here’s a minimal working example.

using Plots
nx, ny = 4,5
x = zeros(nx*ny)
y = zeros(nx*ny)
c = zeros(nx*ny)
a1, a2 = [1.1,.5], [-.2,1.0]
b = hcat(a1,a2)
for i in 0:nx-1
    for j in 0:ny-1
        idx = j+i*ny+1
        x[idx],y[idx] = b*[i,j]
        c[idx] = mod(i,7)/7+mod(j,3)/3
    end
end
scatter(x,y,marker_z=c,colormap=:thermal)

image
image1176×790 19.2 KB

While this is a passable solution, it seems rather like I’m still writing in Fortran…any suggestions on how to make this more Julian?

2

One way:

CI = CartesianIndices((0:ny-1, 0:nx-1))
LI = LinearIndices(CI)
for (I, L) in zip(CI, LI)
    x[L], y[L] = b * [I[2], I[1]]
    c[L] = mod(I[2], 7)/7 + mod(I[1], 3)/3
end
1 Like
3

Do you really need CartesianIndices? A simple solution without them:

using StructArrays

xyc = map(Iterators.product(1:nx, 1:ny)) do (i, j)
    x, y = b * [i, j]
    c = mod(i,7)/7+mod(j,3)/3
    (; x, y, c)
end |> StructArray

# access x, y, c arrays as xyc.x and so on
2 Likes
4

Without any additional packages one could also write this problem as a comprehension:

xyc = vec([[b*[i,j]..., mod(i,7)/7+mod(j,3)/3] for i in 0:nx-1, j in 0:ny-1])
x, y, c = [getindex.(xyc, i) for i in 1:3]

It should be possible to write this even simpler.

OP’s nested loops could also be written using the simplified syntax:

for i in 0:nx-1, j in 0:ny-1
..
end
1 Like
5

I learned a lot looking at your solution. I haven’t used Iterators or StructArrays before so this really helpful in expanding my Julia repertoire. I also like that you didn’t need to pre-allocate the arrays. Thanks for the suggestions.

6

This answered my question directly. But your later suggestion, and @aplavin’s, are even better suggestions (I think). I like that neither one needs preallocated arrays.