Looking for elegant cartesian grid array assignment

General Usage
1

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:

2

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
3

Doesn’t

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

work?

1 Like
4

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.

5

“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
6

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
7

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

8

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()