Derivatives using SVector{Dual}

[dbeach24]

@andyferris Thanks for the offer to help. To avoid sabotaging #25, I'm putting my "Jacobian" code in this ticket.

Let's begin with a small, differentiable vector-valued function:

function cart_to_sphere(cart)
    x, y, z = cart
    r = sqrt(x^2 + y^2 + z^2)
    θ = acos(z / r)
    ϕ = atan2(y, x)
    SVector(r, θ, ϕ)
end

Suppose I want to evaluate this function and its derivative simultaneously, using AD forward mode with dual numbers. Because this is a low-dimensional space, I'd like to use SVector / SMatrix wherever possible.

@inbounds function eval_jacobian{N,T}(f::Function, x::SVector{N,T})
    dualtype = ForwardDiff.Dual{N,T}
    partials = make_partials(GetPartials{N,T}())
    xp = SVector{N,dualtype}(dualtype(x[i], partials[i]) for i=1:N)
    yp = f(xp)
    y = map(ForwardDiff.value, yp)
    P = length(yp)
    J = SMatrix{P,N}((ForwardDiff.partials(yp[i])[j] for i=1:P, j=1:N) ...) # slow line?
    (y, J)
end

Because the "diagonal" partials can be precomputed and re-used for any low-dimension space, I have moved these to a generated function. The function returns an MVector to avoid copying the result on each invocation.

# just a type to invoke make_partials with
immutable GetPartials{N,T} end

@generated function make_partials{N,T}(::GetPartials{N,T})
    partialtype = ForwardDiff.Partials{N,T}
    partials = MVector{N,partialtype}(
        partialtype(ntuple(N) do x x == i ? one(T) : zero(T) end)
        for i=1:N
    )
    return :($partials)
end

Evaluating make_partials is essentially free:

julia> @code_native make_partials(GetPartials{9,Float64}())
    .section    __TEXT,__text,regular,pure_instructions
Filename: REPL[28]
    pushq   %rbp
    movq    %rsp, %rbp
Source line: 2
    movabsq $4569288656, %rax       ## imm = 0x11059CFD0
    popq    %rbp
    retq

Based on timing tests, calling eval_jacobian(cart_to_sphere, SVector(1.,2.,3.)) is far more expensive than it could be:

julia> @time for i=1:10^5 eval_jacobian(cart_to_sphere, SVector(1.,2.,3.)) end
  0.684567 seconds (4.21 M allocations: 227.805 MB, 5.21% gc time)

If I change eval_jacobian to return yp (skipping the formation of y and J), then the timing is substantially better:

julia> @time for i=1:10^5 eval_jacobian(cart_to_sphere, SVector(1.,2.,3.)) end
  0.074104 seconds (1.40 M allocations: 57.983 MB, 9.93% gc time)

I think the cost of constructing the Jacobian matrix is the single biggest problem with my code right now.

Ideally, I'd like to see a version with no dynamic allocations (once compiled). Perhaps this will require crafting my own dual-number implementation using SVector to hold the partials.

Thoughts?

[dbeach24]

Made some other changes to my code, and I think I can confidently state that my only struggle is getting an efficient conversion from SVector{Dual} -> SMatrix.

For comparison:

vector_from_duals{P,N,T}(yp::SVector{P, ForwardDiff.Dual{N,T}}) = map(ForwardDiff.value, yp)

@inbounds function jacobian_from_duals{P,N,T}(yp::SVector{P, ForwardDiff.Dual{N,T}})
    SMatrix{P,N}(
        ntuple(P*N) do n
            i, j = divrem(n-1, N)
            ForwardDiff.partials(yp[i+1])[j+1]
        end ...
    )
end

(Note, I've tried several variations of jacobian_from_duals, including using a generator expression, and both with and without the trailing ellipsis. The fastest version I have so far is the one shown above.)

When I time these:

julia> @time for i=1:10^5 vector_from_duals(yp) end
  0.002037 seconds (100.00 k allocations: 3.052 MB)

julia> @time for i=1:10^5 jacobian_from_duals(yp) end
  0.360131 seconds (2.20 M allocations: 51.880 MB, 1.33% gc time)

So, jacobian_from_duals is about 100x slower than vector_from_duals (in the 3x3 case), but I would expect it to only be 3x slower (since it is copying 3x as much data).

I suppose I could write a @generated function that unrolls this copy. Is there a better way?

[dbeach24]

This seems to work well:

@generated function jacobian_from_duals{P,N,T}(yp::SVector{P, ForwardDiff.Dual{N,T}})
    result = :(SMatrix{P,N}())
    for n=1:P*N
        i, j = divrem(n-1, N)
        p, q = i+1, j+1
        term = :(yp[$p].partials[$q])
        push!(result.args, term)
    end
    result
end

[andyferris]

Yes a generated version works well. Your version might not be optimal for large P, N because of Julia compiler details. E.g. conversion to Tuple using Varargs (i.e. slurping with ...) is sometimes slow, so it is safer to wrap the tuple yourself, and inlining is also important for optimal speed (and getting around that varags problem and other potential issues for larger P, N). So my own implementations tend to go:

@generated function jacobian_from_duals{P,N,T}(yp::SVector{P, ForwardDiff.Dual{N,T}})
    type = SMatrix{P,N}    
    exprs = [:(yp[$p].partials[$q]) for p = 1:P, n = 1:N]
    return quote
        $(Expr(:meta, :inline)
        @inbounds return $type($(Expr(:tuple, exprs...)))    
    end
end

This should be safe for any P or N. Note that there is a shortcut for making a tuple (I forgot for now maybe it's just $((exprs...)) or something - you can find it with trial and error). The inbounds might be important for speed (SVector normally does OK without it since inference can reason about the size of the tuple and indexing with compile-time constants but if there are issues I'll have to add extra code to getindex).

[c42f]

@dbeach24 by the way, did you see CoordinateTransformations.jl? In there, we have some of the same issues of differential geometry and AD... See JuliaGeometry/CoordinateTransformations.jl#6 for what might be the start of exactly the same conversation.

[andyferris]

Yes, I was secretly excited when I saw what you were working on, @dbeach24. :)

[andyferris]

@dbeach24 did you get what you want here?

[dbeach24]

@andyferris Yes, I did. Thanks for the help.