Non-mutable Iterator

Author: kanav99

Project.toml

@@ -8,4 +8,13 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+
+[extras]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["BenchmarkTools", "Test"]

src/NonlinearSolve.jl

@@ -4,11 +4,14 @@ module NonlinearSolve
   @reexport using DiffEqBase
   using UnPack: @unpack
   using FiniteDiff, ForwardDiff
+  using Setfield
+  using StaticArrays
 
   abstract type AbstractNonlinearSolveAlgorithm end
   abstract type AbstractBracketingAlgorithm <: AbstractNonlinearSolveAlgorithm end
   abstract type AbstractNewtonAlgorithm{CS,AD} <: AbstractNonlinearSolveAlgorithm end
   abstract type AbstractNonlinearSolver end
+  abstract type AbstractImmutableNonlinearSolver <: AbstractNonlinearSolver end
 
   include("jacobian.jl")
   include("types.jl")
@@ -25,4 +28,6 @@ module NonlinearSolve
   # DiffEq styled algorithms
   export Bisection, Falsi, NewtonRaphson
   export ScalarBisection, ScalarNewton
+
+  export reinit!
 end # module

src/jacobian.jl

@@ -6,6 +6,13 @@ end
 (uf::JacobianWrapper)(u) = uf.f(u, uf.p)
 (uf::JacobianWrapper)(res, u) = uf.f(res, u, uf.p)
 
+mutable struct ImmutableJacobianWrapper{fType, pType}
+    f::fType
+    p::pType
+end
+
+(uf::ImmutableJacobianWrapper)(u) = uf.f(u, uf.p)
+
 function calc_J(solver, cache)
     @unpack u, f, p, alg = solver
     @unpack uf = cache
@@ -15,9 +22,17 @@ function calc_J(solver, cache)
     return J
 end
 
+function calc_J(solver, uf::ImmutableJacobianWrapper)
+    @unpack u, f, p, alg = solver
+    @set! uf.f = f
+    @set! 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)
+        J = ForwardDiff.jacobian(f, Ref(x)[])
     else
         J = FiniteDiff.finite_difference_derivative(f, x, solver.alg.diff_type, eltype(x))
     end

src/raphson.jl

@@ -15,7 +15,7 @@ mutable struct NewtonRaphsonCache{ufType, L, jType, uType, JC}
     jac_config::JC
 end
 
-mutable struct NewtonRaphsonConstantCache{ufType}
+struct NewtonRaphsonConstantCache{ufType}
     uf::ufType
 end
 
@@ -48,7 +48,7 @@ function perform_step!(solver, alg::NewtonRaphson, cache::NewtonRaphsonConstantC
     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
+    if iszero(solver.fu) || solver.internalnorm(solver.fu) < solver.tol
         solver.force_stop = true
     end
 end
@@ -65,3 +65,15 @@ function perform_step!(solver, alg::NewtonRaphson, cache::NewtonRaphsonCache)
         solver.force_stop = true
     end
 end
+
+function perform_step(solver, alg::NewtonRaphson)
+    @unpack u, fu, f, p = solver
+    J = calc_J(solver, ImmutableJacobianWrapper(f, p))
+    @set! solver.u = u - J \ fu
+    fu = f(solver.u, p)
+    @set! solver.fu = fu
+    if iszero(solver.fu) || solver.internalnorm(solver.fu) < solver.tol
+        @set! solver.force_stop = true
+    end
+    return solver
+end

src/solve.jl

@@ -9,6 +9,7 @@ end
 function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractBracketingAlgorithm, args...;
     alias_u0 = false,
     maxiters = 1000,
+    immutable = (prob.u0 isa StaticArray || prob.u0 isa Number),
     kwargs...
   ) where {uType, iip}
 
@@ -29,12 +30,17 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractBracke
   cache = alg_cache(alg, left, right, p, Val(iip))
 
   sol = build_solution(left, Val(iip))
-  return BracketingSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
+  if immutable
+    return BracketingImmutableSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
+  else
+    return BracketingSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
+  end
 end
 
 function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewtonAlgorithm, args...;
     alias_u0 = false,
     maxiters = 1000,
+    immutable = (prob.u0 isa StaticArray || prob.u0 isa Number),
     tol = 1e-6,
     internalnorm = Base.Fix2(DiffEqBase.ODE_DEFAULT_NORM, nothing),
     kwargs...
@@ -54,10 +60,14 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewton
     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, internalnorm, :Default, tol, sol)
+  if immutable
+    return NewtonImmutableSolver(1, f, alg, u, fu, p, nothing, false, maxiters, internalnorm, :Default, tol, sol)
+  else
+    cache = alg_cache(alg, f, u, p, Val(iip))
+    return NewtonSolver(1, f, alg, u, fu, p, cache, false, maxiters, internalnorm, :Default, tol, sol)
+  end
 end
 
 function DiffEqBase.solve!(solver::AbstractNonlinearSolver)
@@ -75,6 +85,21 @@ function DiffEqBase.solve!(solver::AbstractNonlinearSolver)
   return solver.sol
 end
 
+function DiffEqBase.solve!(solver::AbstractImmutableNonlinearSolver)
+  # sync_residuals!(solver)
+  solver = mic_check(solver)
+  while !solver.force_stop && solver.iter < solver.maxiters
+    solver = perform_step(solver, solver.alg)
+    @set! solver.iter += 1
+    # sync_residuals!(solver)
+  end
+  if solver.iter == solver.maxiters
+    @set! solver.retcode = :MaxitersExceeded
+  end
+  sol = get_solution(solver)
+  return sol
+end
+
 function mic_check!(solver::BracketingSolver)
   @unpack f, fl, fr = solver
   flr = fl * fr
@@ -90,10 +115,29 @@ function mic_check!(solver::BracketingSolver)
   nothing
 end
 
+function mic_check(solver::BracketingImmutableSolver)
+  @unpack f, fl, fr = solver
+  flr = fl * fr
+  fzero = zero(flr)
+  (flr > fzero) && error("Non bracketing interval passed in bracketing method.")
+  if fl == fzero
+    @set! solver.force_stop = true
+    @set! solver.retcode = :ExactSolutionAtLeft
+  elseif fr == fzero
+    @set! solver.force_stop = true
+    @set! solver.retcode = :ExactSolutionAtRight
+  end
+  solver
+end
+
 function mic_check!(solver::NewtonSolver)
   nothing
 end
 
+function mic_check(solver::NewtonImmutableSolver)
+  solver
+end
+
 function check_for_exact_solution!(solver::BracketingSolver)
   @unpack fl, fr = solver
   fzero = zero(fl)
@@ -113,7 +157,27 @@ function set_solution!(solver::BracketingSolver)
   solver.sol.retcode = solver.retcode
 end
 
+function get_solution(solver::BracketingImmutableSolver)
+  return (left = solver.left, right = solver.right, retcode = solver.retcode)
+end
+
 function set_solution!(solver::NewtonSolver)
   solver.sol.u = solver.u
   solver.sol.retcode = solver.retcode
-end
+end
+
+function get_solution(solver::NewtonImmutableSolver)
+  return (u = solver.u, retcode = solver.retcode)
+end
+
+function reinit!(solver::NewtonSolver, prob::NonlinearProblem{uType, true}) where {uType}
+  @. solver.u = prob.u0
+  solver.iter = 1
+  solver.force_stop = false
+end
+
+function reinit!(solver::NewtonSolver, prob::NonlinearProblem{uType, false}) where {uType}
+  solver.u = prob.u0
+  solver.iter = 1
+  solver.force_stop = false
+end

src/types.jl

@@ -14,6 +14,23 @@ mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType
     sol::solType
 end
 
+struct BracketingImmutableSolver{fType, algType, uType, resType, pType, cacheType, solType} <: AbstractImmutableNonlinearSolver
+    iter::Int
+    f::fType
+    alg::algType
+    left::uType
+    right::uType
+    fl::resType
+    fr::resType
+    p::pType
+    cache::cacheType
+    force_stop::Bool
+    maxiters::Int
+    retcode::Symbol
+    sol::solType
+end
+
+
 mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, INType, tolType, solType} <: AbstractNonlinearSolver
     iter::Int
     f::fType
@@ -30,6 +47,22 @@ mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, IN
     sol::solType
 end
 
+struct NewtonImmutableSolver{fType, algType, uType, resType, pType, cacheType, INType, tolType, solType} <: AbstractImmutableNonlinearSolver
+    iter::Int
+    f::fType
+    alg::algType
+    u::uType
+    fu::resType
+    p::pType
+    cache::cacheType
+    force_stop::Bool
+    maxiters::Int
+    internalnorm::INType
+    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)

test/runtests.jl

@@ -0,0 +1,28 @@
+using NonlinearSolve
+using StaticArrays
+using BenchmarkTools
+using Test
+
+function benchmark_immutable()
+    probN = NonlinearProblem((u,p) -> u .* u .- 2, @SVector[1.0, 1.0])
+    solver = init(probN, NewtonRaphson(), immutable = true, tol = 1e-9)
+    sol = @btime solve!($solver)
+    @test all(sol.u .* sol.u .- 2 .< 1e-9)
+end
+
+function benchmark_mutable()
+    probN = NonlinearProblem((u,p) -> u .* u .- 2, @SVector[1.0, 1.0])
+    solver = init(probN, NewtonRaphson(), immutable = false, tol = 1e-9)
+    sol = @btime (reinit!($solver, $probN); solve!($solver))
+    @test all(sol.u .* sol.u .- 2 .< 1e-9)
+end
+
+function benchmark_scalar()
+    probN = NonlinearProblem((u,p) -> u .* u .- 2, 1.0)
+    sol = @btime (solve($probN, ScalarNewton()))
+    @test sol * sol - 2 < 1e-9
+end
+
+benchmark_immutable()
+benchmark_mutable()
+benchmark_scalar()