How to implement macro for assign array of arbitrary dims

General Usage
1

Hello, I would like to do something like d[:, :, 1] = some_array:

function foo(d, x)
    d[:,:,1] = x
end

Then I could write the following macro:

macro gen_foo(index)
    fun = gensym()
    quote
        function $fun(d, x)
            d[:,:,$index] = x
        end
    end
end

newfoo = @gen_foo(2)
d = rand(3, 4, 5)
x = rand(3, 4)
newfoo(d, x)
d[:,:,2] == x # true

However, the number of : is hard-coded here and I would like to generalize it that I could do d[:, :, ..., :] = some_array given how many : I need by trying this:

macro gen_foo2(ndims, index)
    fun = gensym()
    quote
        function $fun(d, x)
            eval(quote
                Expr(
                    :(=),
                    Expr(:ref, $$d, [:(:) for _ in 1: $$ndims]..., $$index),
                    $$x,
                )
            end)
        end
    end
end # not work

I couldn’t make it quite work and not sure how to reference arguments in $fun. Does anyone have a clue for me?

Thanks!

2

Not a full answer, but you may be looking for this neat package: https://github.com/ChrisRackauckas/EllipsisNotation.jl

2 Likes
3

Thanks. I took a look but not with help. On the other hand, I just implemented by adding another macro:

macro assign(d, x, ndims, index)
    Expr(:(=), Expr(:ref, d, [:(:) for _ in 1: ndims]..., index), x)
end

macro gen_foo2(ndims, index)
    fun = gensym()
    quote
        function $(fun)(d, x)
            @assign(d, x, $ndims, $index)
        end
    end
end


d2 = rand(2, 2, 2, 2, 2)
x2 = rand(2, 2, 2, 2)

newfoo2 = @gen_foo2(4, 2)

newfoo2(d2, x2)

d2[:,:,:,:,2] == x2 # true
4

I’m not sure if this is any better but there is CartesianIndices.

julia> a = rand(2,2,2,2,2);

julia> b= zeros(2,2,2,2);

julia> CartesianIndices((1:2,1:2, 2))
2×2×1 CartesianIndices{3,Tuple{UnitRange{Int64},UnitRange{Int64},UnitRange{Int64}}}:
[:, :, 1] =
 CartesianIndex(1, 1, 2)  CartesianIndex(1, 2, 2)
 CartesianIndex(2, 1, 2)  CartesianIndex(2, 2, 2)

julia> idxs = CartesianIndices(( Base.OneTo.(size(b))... , 2))
2×2×2×2×1 CartesianIndices{5,NTuple{5,UnitRange{Int64}}}:
[:, :, 1, 1, 1] =
 CartesianIndex(1, 1, 1, 1, 2)  CartesianIndex(1, 2, 1, 1, 2)
 CartesianIndex(2, 1, 1, 1, 2)  CartesianIndex(2, 2, 1, 1, 2)

[:, :, 2, 1, 1] =
 CartesianIndex(1, 1, 2, 1, 2)  CartesianIndex(1, 2, 2, 1, 2)
 CartesianIndex(2, 1, 2, 1, 2)  CartesianIndex(2, 2, 2, 1, 2)

[:, :, 1, 2, 1] =
 CartesianIndex(1, 1, 1, 2, 2)  CartesianIndex(1, 2, 1, 2, 2)
 CartesianIndex(2, 1, 1, 2, 2)  CartesianIndex(2, 2, 1, 2, 2)

[:, :, 2, 2, 1] =
 CartesianIndex(1, 1, 2, 2, 2)  CartesianIndex(1, 2, 2, 2, 2)
 CartesianIndex(2, 1, 2, 2, 2)  CartesianIndex(2, 2, 2, 2, 2)

julia> a[idxs] .= b;

julia> a
2×2×2×2×2 Array{Float64,5}:
[:, :, 1, 1, 1] =
 0.204811  0.81601 
 0.287767  0.474275

[:, :, 2, 1, 1] =
 0.345389  0.119216
 0.524387  0.574756

[:, :, 1, 2, 1] =
 0.444125  0.589582
 0.869282  0.675045

[:, :, 2, 2, 1] =
 0.595017  0.00766057
 0.394598  0.780346  

[:, :, 1, 1, 2] =
 0.0  0.0
 0.0  0.0

[:, :, 2, 1, 2] =
 0.0  0.0
 0.0  0.0

[:, :, 1, 2, 2] =
 0.0  0.0
 0.0  0.0

[:, :, 2, 2, 2] =
 0.0  0.0
 0.0  0.0
5

Thank you. The cartesian index never came to my mind until you mentioned it, I implement and benchmark them:

macro assign(dest, src, ndim, dest_index, src_index)
    Expr(
        :(=),
        Expr(:ref, esc(dest), [:(:) for _ in 1: ndim]..., esc(dest_index)),
        Expr(:call, :selectdim, esc(src), ndim+1, esc(src_index))
    )
end

macro assign2(dest, src, dest_index, src_index)
    quote
        didx = CartesianIndices((Base.OneTo.(size($(esc(dest)))[1:end-1])..., $(esc(dest_index))))
        sidx = CartesianIndices((Base.OneTo.(size($(esc(src)))[1:end-1])..., $(esc(src_index))))
        $(esc(dest))[didx] = selectdim($(esc(src)), ndims($(esc(src))), $(esc(src_index)))
    end
end

macro gen_foo(ndims, dest_index, src_index)
    fun = gensym()
    quote
        function $(fun)(d, s)
            @assign(d, s, $ndims, $(esc(dest_index)), $(esc(src_index)))
        end
    end
end

macro gen_foo2(dest_index, src_index)
    fun = gensym()
    quote
        function $(fun)(d, s)
            @assign2(d, s, $(esc(dest_index)), $(esc(src_index)))
        end
    end
end

d1 = rand(100, 100, 100, 100)
x1 = rand(100, 100, 100, 100)
foo1 = @gen_foo(3, 3, 2)
foo1(d1, x1)
d1[:,:,:,3] == x1[:,:,:,2] # true

d2 = rand(100, 100, 100, 100)
x2 = rand(100, 100, 100, 100)
foo2 = @gen_foo2(3, 2)
foo2(d2, x2)
d2[:,:,:,3] == x2[:,:,:,2] # true

using BenchmarkTools

@btime foo1($d1, $x1) # 1.510 ms (9 allocations: 256 bytes)
@btime foo2($d2, $x2) # 4.134 ms (25 allocations: 1.00 KiB)

Looks like the handy :, :, : is faster than use CartesianIndex.

6

I’m not sure why you would need an auxiliary macro @assign here…

How about the following?

macro genfun(ndims, index1, index2)
    fun = gensym()
    quote
        function $(fun)(d, x)
            d[$((esc(:(:)) for _ in 1:ndims)...), $index1] = selectdim(x, $(ndims+1), $index2)
        end
    end
end

This seems to generate the same code as your @gen_foo implementation.

2 Likes
7

Thank you! Your solution is much clearer than mine!