Make TEQUILA autodifferentiable

[bclyons12]

We'd like to get this code working

using TEQUILA, MillerExtendedHarmonic, ForwardDiff
shotd3d = Shot(11, 11, 7, "../sample/g_chease_mxh_d3d")
bnd = shotd3d.surfaces[:,end]
dP_dψ = shotd3d.dP_dψ
F_dF_dψ = shotd3d.F_dF_dψ
Pbnd = shotd3d.Pbnd
Fbnd = shotd3d.Fbnd

function solve_from_bnd(bnd)
    # initialize TEQUILA 
    shot = Shot(11, 11, MXH(bnd); dP_dψ, F_dF_dψ, Pbnd, Fbnd);

    # solve TEQUILA equilibrium
    refill = solve(shot, 3; debug=false)

    return refill
end

@time refill = solve_from_bnd(bnd); # runs fine

ForwardDiff.jacobian(solve_from_bnd, bnd);

Currently, jacobian errors with the following, but I'm sure that's just the beginning:

MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:792
  (::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number}
   @ Base char.jl:50
  ...


Stacktrace:
  [1] _broadcast_getindex_evalf
    @ ./broadcast.jl:683 [inlined]
  [2] _broadcast_getindex
    @ ./broadcast.jl:656 [inlined]
  [3] getindex
    @ ./broadcast.jl:610 [inlined]
  [4] copy
    @ ./broadcast.jl:912 [inlined]
  [5] materialize
    @ ./broadcast.jl:873 [inlined]
  [6] surfaces_FE(ρ::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, surfaces::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}})
    @ TEQUILA ~/.julia/dev/TEQUILA/src/surfaces.jl:89
  [7] Shot(N::Int64, M::Int64, boundary::MXH{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}}}; Raxis::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}, Zaxis::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}, P::Nothing, dP_dψ::FiniteElementHermite.FE_rep{Vector{Float64}, Vector{Float64}}, F_dF_dψ::FiniteElementHermite.FE_rep{Vector{Float64}, Vector{Float64}}, Jt_R::Nothing, Jt::Nothing, Pbnd::Float64, Fbnd::Float64, Ip_target::Nothing, approximate_psi::Bool)
    @ TEQUILA ~/.julia/dev/TEQUILA/src/shot.jl:101
  [8] solve_from_bnd(bnd::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}})
    @ Main ./In[9]:3
  [9] chunk_mode_jacobian(f::typeof(solve_from_bnd), x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:183
 [10] jacobian(f::Function, x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}}}, ::Val{true})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:23
 [11] jacobian(f::Function, x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:19
 [12] jacobian(f::Function, x::Vector{Float64})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:19
 [13] top-level scope
    @ In[12]:1

@orso82 @sjkelly may be interested

[mbolognaJH]

@bclyons12 should I be running this in a specific directory? And is the sample data available? I also remember you mentioning this would be on a specific branch. Thanks!

[bclyons12]

@mbolognaJH This case is just on the master branch of TEQUILA, so no specific branch.

I thought the sample was in the repository, but it's not. You can find it here:
g_chease_mxh_d3d.txt

[mbolognaJH]

For your reference, I've added ForwardDiff to line 91 of the FUSE Makefile, then run make update and run this MWE. This is the output I get:

ForwardDiff.jacobian(solve_from_bnd, bnd);
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:790
  Float64(::IrrationalConstants.Sqrt2π)
   @ IrrationalConstants ~/.julia/packages/IrrationalConstants/vp5v4/src/macro.jl:112
  ...

Stacktrace:
  [1] convert(::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10})
    @ Base ./number.jl:7
  [2] setindex!(A::Matrix{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}, i1::Int64)
    @ Base ./array.jl:1019
  [3] setindex!
    @ ./subarray.jl:353 [inlined]
  [4] flat_coeffs!(flat::SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{…}, Int64}, true}, mxh::MXH{ForwardDiff.Dual{ForwardDiff.Tag{…}, Float64, 10}, Vector{ForwardDiff.Dual{…}}})
    @ MillerExtendedHarmonic ~/.julia/dev/MillerExtendedHarmonic/src/MXH.jl:46
  [5] Shot(N::Int64, M::Int64, boundary::MXH{…}; Raxis::ForwardDiff.Dual{…}, Zaxis::ForwardDiff.Dual{…}, P::Nothing, dP_dψ::FiniteElementHermite.FE_rep{…}, F_dF_dψ::FiniteElementHermite.FE_rep{…}, Jt_R::Nothing, Jt::Nothing, Pbnd::Float64, Fbnd::Float64, Ip_target::Nothing, approximate_psi::Bool)
    @ TEQUILA ~/.julia/dev/TEQUILA/src/shot.jl:98
  [6] solve_from_bnd(bnd::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10}})
    @ Main ./REPL[9]:3
  [7] chunk_mode_jacobian(f::typeof(solve_from_bnd), x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{…}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:183
  [8] jacobian(f::Function, x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{…}}}, ::Val{true})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:0
  [9] jacobian(f::Function, x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(solve_from_bnd), Float64}, Float64, 10, Vector{ForwardDiff.Dual{…}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:19
 [10] top-level scope
    @ REPL[18]:1
Some type information was truncated. Use `show(err)` to see complete types.

So my plan is to work on MillerExtendedHarmonic first.

[orso82]

That makes sense @mbolognaJH

[orso82]

https://github.com/dannys4/FFTA.jl

[mbolognaJH]

I've pushed commits to MillerExtendedHarmonic and FiniteElementHermite and added a branch called mdb/tequila-autodiff. The temporary test for ForwardDiff is running test/test_autodiff.jl.

Because FFTW.jl is not autodiff compatible, only Shot(...) from the MWE above is now autodiff compatible.

There are two possible solutions:

1. Use ForwardDiff PR #668 as a template for creating a Requires.jl implementation of FFTW + ForwardDiff extension. Wait for this PR to add ForwardDiff rules for FFTW.jl.
2. Not recommended: A further effort could be taken to switch to FFTA.jl (as @orso82 mentions above), but we haven't validated this as viable.

I can continue future efforts in the new year.