I am trying to implement a Jacobian-Free Newton Krylov (JFNK) solver using existing functions in NonlinearSolve.jl and related packages. These methods should only require the evaluation of the directional derivatives J*v for a vector v, rather than the full Jacobian matrix J.
In a situation when Jv is known analytically and given as an input to the nonlinear function via the ‘jvp=’ argument, the solver ignores this argument and attempts to use auto-differentiation to evaluate Jv rather than calling the function. I am wondering how to fix this so that an analytic formula for Jv can be input.
Here is a simple example demonstrating the issue. This example code calls the desired JFNK methods to perform a time-step of a non-linear diffusion-type equation implemented by a finite-difference method; in it F(u,p) is the non-linear function (with parameter p) to solve, and J(v,u,p) is an analytic evaluation of the Jacobian at u acting on the vector v. This code returns an error saying that an argument of type ForwardDiff.Dual cannot be passed into F, indicating that NonlinearSolve is attempting to use automatic differentiation to evaluate Jv rather than calling the input function.
function F(u :: Vector{Float64}, p :: Vector{Float64})
Δ=Tridiagonal(-ones(99), 2*ones(100),-ones(99)) ##finite-difference Laplacian w/ 100 volumes
return u + 0.1 * u .* Δ * u- p ## parameter p is previous time-step
end
function J(v :: Vector{Float64},u :: Vector{Float64},p :: Vector{Float64})
Δ=Tridiagonal(-ones(99), 2*ones(100),-ones(99))
return v + 0.1* ( u.* Δ*v + v.* Δ*u) ## Analytic formula for J*v via product rule
end
u0=ones(100)
NLP_F=NonlinearProblem{false}(NonlinearFunction{false}(F, jvp=J),u0,u0) ##u0 is both initial guess and parameter
sol=solve(NLP_F,NewtonRaphson(linsolve =KrylovJL_GMRES()))
( This particular code above can be implemented more easily by removing the Vector{Float64} specifications allowing Automatic Differentiation to work, but in more complicated situations this approach may not be viable.)
Any help or thoughts would be appreciated! Thanks.