Looking for elegant cartesian grid array assignment

I’m building an array of cube centers for a 3-D cartesion grid, something like (ignore for the moment, for simplicity, that the cube centers are actually located at an offset of dx/2, dy/2, dz/2).

for i=1:nx
  for j=1:ny
    for k = 1:nz
        a[i,j,k,1] = i*dx # i=1:nx
        a[i,j,k,2] = j*dy # j=1:ny
        a[i,j,k,3] = k*dz #  k=1:nz
    end
  end
end

realizing that the sequences are repeated, i.e. you can do

xs = dx * range(0,nx-1,nx)
for j=1:ny
  for k=1:nz
    a[:,j,k,1] .= xs
  end
end

and then you need two more loops for the 2nd and 3rd indexes.

Works great, but it’s not clever enough :laughing: . It seems like this is amenable to a one-liner hack. I am seeking that one-liner hack, just for fun. Even the brute force 3-loop calculation of the values for a 64M array takes LESS than 2 seconds, so not worried about efficiency at all.

I’m just figuring some julia hacker out there can give me a valuable lesson in broadcasting :slight_smile:

You are iterating the loops in an order suboptimal for the memory layout. Iterate the first dimension in the tightest loop, i.e, swap the loop order in the first snippet.

2 Likes

Doesn’t

a[:,:, :,1] .= xs

work?

1 Like

yowzer. a factor of 3 speed-up. I really should have known that as many discussions as i’ve seen about the column-major-ness of julia.

yes it does! How ?! More impressively if the first and second dimensions are the same it still works… However

ys = dy * range(0, ny-1, ny)

a1[:,:,:,2] .= ys

does not. But if you make ys a 1xny vector then i think it does work, but then i couldn’t figure out how to make things work for the 3rd index.

“Julia provides broadcast, which expands singleton dimensions in array arguments to match the corresponding dimension in the other array without using extra memory, and applies the given function elementwise”
https://docs.julialang.org/en/v1/manual/arrays/#Broadcasting-1

1 Like

You can do this, just putting the vector along the 3rd axis:

zsteps = reshape(range(0, nz-1, step=1), 1,1,:) 
a1[:,:,:,3] .= zsteps .* dz
1 Like

Now I get it! LOL. Thank you both for your very helpful answers!

Doesn’t everyone like benchmarks ? :wink:

Remarkably, the element-by-element assign and the precomputed vector assigns seem to be about the same speed. I’ll just include one result here, but i ran several different dimensions, as well as reversing the call order, and they track each other very closely.

  14.889 ms (3 allocations: 42.06 MiB)
[0.0005, 0.0005, 0.0005]
[0.2995, 0.17450000000000002, 0.0345]
(300, 175, 35, 3)
1837500
  13.552 ms (5 allocations: 42.06 MiB)
[0.0005, 0.0005, 0.0005]
[0.2995, 0.17450000000000002, 0.0345]
(300, 175, 35, 3)
1837500
using BenchmarkTools

function f0(nx, ny, nz, dx, dy, dz)
    a1 = Array{Float64,4}(undef, nx, ny, nz, 3)
    
    for k=1:nz
        for j=1:ny
            for i=1:nx
                a1[i,j,k,1] = (i-0.5)*dx
                a1[i,j,k,2] = (j-0.5)*dy
                a1[i,j,k,3] = (k-0.5)*dz
            end
        end
    end

    a1
end

                
function f1(nx, ny, nz, dx, dy, dz)
    
    a1 = Array{Float64,4}(undef, nx, ny, nz, 3)

    # println(a)

    # sequences of values that need only be calculated once
    
    xs = dx * (0.5 .+ range(0, nx-1, step=1))
    ys = reshape(dy * (0.5 .+ range(0, ny-1, step=1)), 1, :, 1)
    zs = reshape(dz * (0.5 .+ range(0, nz-1, step=1)), 1, 1, :)

    # assign x, y, z values in sequence
    
    a1[:,:,:,1] .= xs
    a1[:,:,:,2] .= ys
    a1[:,:,:,3] .= zs

    a1
end



function main()
    # electric fields are established at the node i,j,k exactly
    # the magnetic fields are offset 1/2 step

    nx = 300
    ny = 175
    nz = 35

    dx = 1E-3
    dy = 1E-3
    dz = 1E-3

    a = @btime f0($nx, $ny, $nz, $dx, $dy, $dz)
    println(a[1,1,1,:])
    println(a[end,end,end,:])
    println(size(a))
    println(nx*ny*nz)

    a = @btime f1($nx, $ny, $nz, $dx, $dy, $dz)
    println(a[1,1,1,:])
    println(a[end,end,end,:])
    println(size(a))
    println(nx*ny*nz)
end

main()