CODEWITHSUNDEEP
julia
Adobe
Adobe

Reputation: 13467

Vectorize function with two args of different types?

In Julia one can vectorize 2 arg function, when both args are of the same type:

function mysquare(n::Number, x::Number)
    return n * x^2
end

println(mysquare(2, 2.)) # ⇒ 8.0

@vectorize_2arg Number mysquare

println(mysquare(2, [2., 3., 4.])) # ⇒ [8.0,18.0,32.0]

Is there a way, aside from doing it manually, to vectorize function of two args with different types:

function mysquare(n::Int64, x::Float64)
    return n * x^2
end

println(mysquare(2, 2.))

@vectorize_2arg Int64 Float64 mysquare # ⇒ wrong number of args

println(mysquare(2, [2., 3., 4.]))

Upvotes: 1

Views: 171

Answers (1)

Adobe
Adobe

Reputation: 13467

I modified the code of vectorize_2arg macro. I'm not sure I fully understand all the implications here, but it works:

macro vectorize_2arg_full(S1,S2,f)
    S1 = esc(S1); S2 = esc(S2); f = esc(f); T1 = esc(:T1); T2 = esc(:T2)
    quote
        ($f){$T1<:$S1, $T2<:$S2}(x::($T1), y::AbstractArray{$T2}) =
            reshape([ ($f)(x, y[i]) for i=1:length(y) ], size(y))
        ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::($T2)) =
            reshape([ ($f)(x[i], y) for i=1:length(x) ], size(x))

        function ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::AbstractArray{$T2})
            shp = promote_shape(size(x),size(y))
            reshape([ ($f)(x[i], y[i]) for i=1:length(x) ], shp)
        end
    end
end

function mysquare(n::Int64, x::Float64)
    return n * x^2
end

println(mysquare(2, 2.))

@vectorize_2arg_full Int64 Float64 mysquare

println(mysquare(2, [2., 3., 4.]))

Can someone code-review it?

Edit

I posted that on a Julia mailing list: one could use 

@vectorize_2arg Number mysquare

for that. No need for tweaks.

Upvotes: 1

Related Questions