(feat) IIP Raphson

Author: kanav99

src/jacobian.jl

@@ -23,3 +23,23 @@ function jacobian(f, x, solver)
     end
     return J
 end
+
+function calc_J!(J, solver, cache)
+    @unpack f, u, p, alg = solver
+    @unpack du1, uf, jac_config = cache
+
+    uf.f = f
+    uf.p = p
+
+    jacobian!(J, uf, u, du1, solver, jac_config)
+end
+
+function jacobian!(J, f, x, fx, solver, jac_config)
+    alg = solver.alg
+    if alg_autodiff(alg)
+        ForwardDiff.jacobian!(J, f, fx, x, jac_config)
+    else
+        FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config)
+    end
+    nothing
+end

src/raphson.jl

@@ -1,31 +1,67 @@
-struct NewtonRaphson{CS, AD, DT} <: AbstractNewtonAlgorithm{CS,AD} 
+struct NewtonRaphson{CS, AD, DT, L} <: AbstractNewtonAlgorithm{CS,AD} 
     diff_type::DT
+    linsolve::L
 end
 
-function NewtonRaphson(;autodiff=true,chunk_size=12,diff_type=Val{:forward})
-    NewtonRaphson{chunk_size, autodiff, typeof(diff_type)}(diff_type)
+function NewtonRaphson(;autodiff=true,chunk_size=12,diff_type=Val{:forward},linsolve=DEFAULT_LINSOLVE)
+    NewtonRaphson{chunk_size, autodiff, typeof(diff_type), typeof(linsolve)}(diff_type, linsolve)
 end
 
-mutable struct NewtonRaphsonCache{ufType}
+mutable struct NewtonRaphsonCache{ufType, L, jType, uType, JC}
+    uf::ufType
+    linsolve::L
+    J::jType
+    du1::uType
+    jac_config::JC
+end
+
+mutable struct NewtonRaphsonConstantCache{ufType}
     uf::ufType
 end
 
 function alg_cache(alg::NewtonRaphson, f, u, p, ::Val{true})
     uf = JacobianWrapper(f,p)
-    NewtonRaphsonCache(uf)
+    linsolve = alg.linsolve(Val{:init}, f, u)
+    J = false .* u .* u'
+    du1 = zero(u)
+    tmp = zero(u)
+    if alg_autodiff(alg)
+        jac_config = ForwardDiff.JacobianConfig(uf, du1, u)
+    else
+        if alg.diff_type != Val{:complex}
+            du2 = zero(u)
+            jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg.diff_type)
+        else
+            jac_config = FiniteDiff.JacobianCache(Complex{eltype(tmp)}.(tmp),Complex{eltype(du1)}.(du1),nothing,alg.diff_type,eltype(u))
+        end
+    end
+    NewtonRaphsonCache(uf, linsolve, J, du1, jac_config)
 end
 
 function alg_cache(alg::NewtonRaphson, f, u, p, ::Val{false})
     uf = JacobianWrapper(f,p)
-    NewtonRaphsonCache(uf)
+    NewtonRaphsonConstantCache(uf)
 end
 
-function perform_step!(solver, alg::NewtonRaphson, cache)
+function perform_step!(solver, alg::NewtonRaphson, cache::NewtonRaphsonConstantCache)
     @unpack u, fu, f, p = solver
     J = calc_J(solver, cache)
     solver.u = u - J \ fu
     solver.fu = f(solver.u, p)
     if iszero(solver.fu) || abs(solver.fu) < solver.tol
         solver.force_stop = true
     end
-end
+end
+
+function perform_step!(solver, alg::NewtonRaphson, cache::NewtonRaphsonCache)
+    @unpack u, fu, f, p = solver
+    @unpack J, linsolve, du1 = cache
+    calc_J!(J, solver, cache)
+    # u = u - J \ fu
+    linsolve(du1, J, fu, true)
+    @. u = u - du1
+    f(fu, u, p)
+    if solver.internalnorm(solver.fu) < solver.tol
+        solver.force_stop = true
+    end
+end

src/solve.jl

@@ -36,6 +36,7 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewton
     alias_u0 = false,
     maxiters = 1000,
     tol = 1e-6,
+    internalnorm = Base.Fix2(DiffEqBase.ODE_DEFAULT_NORM, nothing),
     kwargs...
   ) where {uType, iip}
 
@@ -46,12 +47,17 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewton
   end
   f = prob.f
   p = prob.p
-  fu = f(u, p)
+  if iip
+    fu = zero(u)
+    f(fu, u, p)
+  else
+    fu = f(u, p)
+  end
 
   cache = alg_cache(alg, f, u, p, Val(iip))
 
   sol = build_newton_solution(u, Val(iip))
-  return NewtonSolver(1, f, alg, u, fu, p, cache, false, maxiters, :Default, tol, sol)
+  return NewtonSolver(1, f, alg, u, fu, p, cache, false, maxiters, internalnorm, :Default, tol, sol)
 end
 
 function DiffEqBase.solve!(solver::AbstractNonlinearSolver)

src/types.jl

@@ -14,7 +14,7 @@ mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType
     sol::solType
 end
 
-mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, tolType, solType} <: AbstractNonlinearSolver
+mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, INType, tolType, solType} <: AbstractNonlinearSolver
     iter::Int
     f::fType
     alg::algType
@@ -24,6 +24,7 @@ mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, to
     cache::cacheType
     force_stop::Bool
     maxiters::Int
+    internalnorm::INType
     retcode::Symbol
     tol::tolType
     sol::solType

src/utils.jl

@@ -24,3 +24,6 @@ function value_derivative(f::F, x::R) where {F,R}
     out = f(ForwardDiff.Dual{T}(x, one(x)))
     ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
 end
+
+DiffEqBase.has_Wfact(f::Function) = false
+DiffEqBase.has_Wfact_t(f::Function) = false