(feat) NewtonRaphson

Author: kanav99

Project.toml

@@ -5,5 +5,7 @@ version = "0.1.0"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"

src/NonlinearSolve.jl

@@ -3,19 +3,24 @@ module NonlinearSolve
   using Reexport
   @reexport using DiffEqBase
   using UnPack: @unpack
+  using FiniteDiff, ForwardDiff
 
   abstract type AbstractNonlinearSolveAlgorithm end
   abstract type AbstractBracketingAlgorithm <: AbstractNonlinearSolveAlgorithm end
+  abstract type AbstractNewtonAlgorithm{CS,AD} <: AbstractNonlinearSolveAlgorithm end
+  abstract type AbstractNonlinearSolver end
 
+  include("jacobian.jl")
   include("types.jl")
   include("solve.jl")
   include("utils.jl")
   include("bisection.jl")
   include("falsi.jl")
+  include("raphson.jl")
 
   # raw methods
   export bisection, falsi
 
   # DiffEq styled algorithms
-  export Bisection, Falsi
+  export Bisection, Falsi, NewtonRaphson
 end # module

src/jacobian.jl

@@ -0,0 +1,25 @@
+mutable struct JacobianWrapper{fType, pType}
+    f::fType
+    p::pType
+end
+
+(uf::JacobianWrapper)(u) = uf.f(u, uf.p)
+(uf::JacobianWrapper)(res, u) = uf.f(res, u, uf.p)
+
+function calc_J(solver, cache)
+    @unpack u, f, p, alg = solver
+    @unpack uf = cache
+    uf.f = f
+    uf.p = p
+    J = jacobian(uf, u, solver)
+    return J
+end
+
+function jacobian(f, x, solver)
+    if alg_autodiff(solver.alg)
+        J = ForwardDiff.derivative(f, x)
+    else
+        J = FiniteDiff.finite_difference_derivative(f, x, solver.alg.diff_type, eltype(x))
+    end
+    return J
+end

src/raphson.jl

@@ -0,0 +1,31 @@
+struct NewtonRaphson{CS, AD, DT} <: AbstractNewtonAlgorithm{CS,AD} 
+    diff_type::DT
+end
+
+function NewtonRaphson(;autodiff=true,chunk_size=12,diff_type=Val{:forward})
+    NewtonRaphson{chunk_size, autodiff, typeof(diff_type)}(diff_type)
+end
+
+mutable struct NewtonRaphsonCache{ufType}
+    uf::ufType
+end
+
+function alg_cache(alg::NewtonRaphson, f, u, p, ::Val{true})
+    uf = JacobianWrapper(f,p)
+    NewtonRaphsonCache(uf)
+end
+
+function alg_cache(alg::NewtonRaphson, f, u, p, ::Val{false})
+    uf = JacobianWrapper(f,p)
+    NewtonRaphsonCache(uf)
+end
+
+function perform_step!(solver, alg::NewtonRaphson, cache)
+    @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

src/solve.jl

@@ -32,16 +32,34 @@ function DiffEqBase.__init(prob::NonlinearProblem{uType, iip}, alg::AbstractBrac
   return BracketingSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
 end
 
-function DiffEqBase.solve!(solver::BracketingSolver)
+function DiffEqBase.__init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewtonAlgorithm, args...;
+    alias_u0 = false,
+    maxiters = 1000,
+    tol = 1e-6,
+    kwargs...
+  ) where {uType, iip}
+
+  if alias_u0
+    u = prob.u0
+  else
+    u = deepcopy(prob.u0)
+  end
+  f = prob.f
+  p = prob.p
+  fu = f(u, p)
+
+  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)
+end
+
+function DiffEqBase.solve!(solver::AbstractNonlinearSolver)
   # sync_residuals!(solver)
   mic_check!(solver)
   while !solver.force_stop && solver.iter < solver.maxiters
-    if check_for_exact_solution!(solver)
-      break
-    else
-      perform_step!(solver, solver.alg, solver.cache)
-      solver.iter += 1
-    end
+    perform_step!(solver, solver.alg, solver.cache)
+    solver.iter += 1
     # sync_residuals!(solver)
   end
   if solver.iter == solver.maxiters
@@ -66,6 +84,10 @@ function mic_check!(solver::BracketingSolver)
   nothing
 end
 
+function mic_check!(solver::NewtonSolver)
+  nothing
+end
+
 function check_for_exact_solution!(solver::BracketingSolver)
   @unpack fl, fr = solver
   fzero = zero(fl)
@@ -79,8 +101,13 @@ function check_for_exact_solution!(solver::BracketingSolver)
   return false
 end
 
-function set_solution!(solver)
+function set_solution!(solver::BracketingSolver)
   solver.sol.left = solver.left
   solver.sol.right = solver.right
   solver.sol.retcode = solver.retcode
 end
+
+function set_solution!(solver::NewtonSolver)
+  solver.sol.u = solver.u
+  solver.sol.retcode = solver.retcode
+end

src/types.jl

@@ -1,4 +1,4 @@
-mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType, solType}
+mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType, solType} <: AbstractNonlinearSolver
     iter::Int
     f::fType
     alg::algType
@@ -14,6 +14,21 @@ mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType
     sol::solType
 end
 
+mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, tolType, solType} <: AbstractNonlinearSolver
+    iter::Int
+    f::fType
+    alg::algType
+    u::uType
+    fu::resType
+    p::pType
+    cache::cacheType
+    force_stop::Bool
+    maxiters::Int
+    retcode::Symbol
+    tol::tolType
+    sol::solType
+end
+
 function sync_residuals!(solver::BracketingSolver)
     solver.fl = solver.f(solver.left, solver.p)
     solver.fr = solver.f(solver.right, solver.p)
@@ -33,3 +48,13 @@ end
 function build_solution(u_prototype, ::Val{false})
     return BracketingSolution(zero(u_prototype), zero(u_prototype), :Default)
 end
+
+mutable struct NewtonSolution{uType}
+    u::uType
+    retcode::Symbol
+end
+
+function build_newton_solution(u_prototype, ::Val{iip}) where iip
+    return NewtonSolution(zero(u_prototype), :Default)
+end
+

src/utils.jl

@@ -10,3 +10,6 @@ end
 function nextfloat_tdir(x::T, x0::T, x1::T)::T where {T}
   x1 > x0 ? nextfloat(x) : prevfloat(x)
 end
+
+alg_autodiff(alg::AbstractNewtonAlgorithm{CS,AD}) where {CS,AD} = AD
+alg_autodiff(alg) = false