julia differential-equations fixed-point-iteration

Wrong function call when solving for steady state using julia's NonlinearSolve

I am using NonlinearSolve to solve a steady state problem. When I do

prob = SteadyStateProblem(eoms!, σ₀, p)
sol = solve(
    prob,
    SSRootfind(),
    abstol = slv.abstol,
    reltol = slv.reltol,
)

I see that eoms! is called correctly for the first time. Here is the function shape:

function eoms!(
    du::Vector{Float64},
    u::Vector{Float64},
    p::Tuple{Any,Any,Any,Any,Any},
    t::Float64 = 0.0,
)

When it is called for the second time, however, it is called with incorrect arguments, i.e., for example for du:

ForwardDiff.Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}, Float64, 8}[Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0), Dual{ForwardDiff.Tag{SciMLBase.JacobianWrapper{true, NonlinearFunction{true, SciMLBase.FullSpecialize, SciMLBase.var"#41#50"{ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.Solve.eoms!), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, Tuple{Main.Solve.Instance, Main.SymmetricLattice.Instance, Main.Laser.Instance, Float64, Int64}}, Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0)]

Does anyone know what I am doing wrong? Note: When I do a time evolution (i.e., when I declare an ODEProblem instead of a SteadyStateProblem), the solver manages to find a reasonable solution.

Solution

ForwardDiff.Dual <: Real (see src). So if you want to specify that du and u must be vectors of reals, then to workaround the error, one way is to specify type Vector{R} where {R <: Real). Then the function will accept either Vector{Float64} or Vector{ForwardDiff.Dual}.

Error-reproducing example:

function f!(
    du::Vector{Float64},
    u::Vector{Float64},
    p::Tuple{Any,Any,Any,Any,Any},
    t::Float64 = 0.0,
)
    # added: made-up function body
    (a, b, c, d, e) = p
    du .= [a*sin(u[1]), b*cos(u[2]), c*exp(u[3]), d*tan(u[4]), e*cot(u[5])]
end

Running solver

using SciMLBase: SteadyStateProblem, solve
using SteadyStateDiffEq: SSRootfind

p = (1.0, 2.0, 3.0, 4.0, 5.0)
σ₀ = rand(5)
slv = (; abstol = 1e-10, reltol = 1e-10)

prob = SteadyStateProblem(f!, σ₀, p)
sol = solve(
    prob,
    SSRootfind(),
    abstol = slv.abstol,
    reltol = slv.reltol,
)

produces error

ERROR: MethodError: no method matching f!(::Vector{ForwardDiff.Dual{…}}, ::Vector{ForwardDiff.Dual{…}}, ::NTuple{5, Float64}, ::Float64)

Closest candidates are:
  f!(::Vector{Float64}, ::Vector{Float64}, ::NTuple{5, Any}, ::Float64)
   @ Main REPL[3]:1
  f!(::Vector{Float64}, ::Vector{Float64}, ::NTuple{5, Any})
   @ Main REPL[3]:1

Working example:

function f!(
    du::Vector{R},  # was: Vector{Float64}
    u::Vector{R},   # was: Vector{Float64}
    p::Tuple{Any,Any,Any,Any,Any},
    t::Float64 = 0.0,
) where {R <: Real} # added: where {R <: Real}
    # made-up function body
    (a, b, c, d, e) = p
    du .= [a*sin(u[1]), b*cos(u[2]), c*exp(u[3]), d*tan(u[4]), e*cot(u[5])]
end

Running solver

using SciMLBase: SteadyStateProblem, solve
using SteadyStateDiffEq: SSRootfind

p = (1.0, 2.0, 3.0, 4.0, 5.0)
σ₀ = rand(5)
slv = (; abstol = 1e-10, reltol = 1e-10)

prob = SteadyStateProblem(f!, σ₀, p)
sol = solve(
    prob,
    SSRootfind(),
    abstol = slv.abstol,
    reltol = slv.reltol,
)

produces solution

retcode: Success
u: 5-element Vector{Float64}:
   0.0
  51.83627878423159
 -24.893051092705488
   0.0
   1.5707963267948966