(refactor) General Cleaning and Tests

Author: kanav99

src/jacobian.jl

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

src/scalar.jl

@@ -25,14 +25,14 @@ function DiffEqBase.solve(prob::NonlinearProblem, ::Bisection, args...; maxiters
   fl, fr = f(left), f(right)
 
   if iszero(fl)
-    return fl
+    return BracketingSolution(left, right, :ExactSolutionAtLeft)
   end
 
   i = 1
   if !iszero(fr)
     while i < maxiters
       mid = (left + right) / 2
-      (mid == left || mid == right) && return left
+      (mid == left || mid == right) && return BracketingSolution(left, right, :FloatingPointLimit)
       fm = f(mid)
       if iszero(fm)
         right = mid
@@ -51,7 +51,7 @@ function DiffEqBase.solve(prob::NonlinearProblem, ::Bisection, args...; maxiters
 
   while i < maxiters
     mid = (left + right) / 2
-    (mid == left || mid == right) && return left
+    (mid == left || mid == right) && return BracketingSolution(left, right, :FloatingPointLimit)
     fm = f(mid)
     if iszero(fm)
       right = mid
@@ -63,5 +63,5 @@ function DiffEqBase.solve(prob::NonlinearProblem, ::Bisection, args...; maxiters
     i += 1
   end
 
-  return left
+  return BracketingSolution(left, right, :MaxitersExceeded)
 end

src/solve.jl

@@ -10,14 +10,18 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractBracke
     alias_u0 = false,
     maxiters = 1000,
     # bracketing algorithms only solve scalar problems
-    immutable = (prob.u0 isa StaticArray || prob.u0 isa Number),
+    immutable = (eltype(prob.u0) <: Number),
     kwargs...
   ) where {uType, iip}
 
   if !(prob.u0 isa Tuple)
     error("You need to pass a tuple of u0 in bracketing algorithms.")
   end
 
+  if eltype(prob.u0) isa AbstractArray
+    error("Bracketing Algorithms work for scalar arguments only")
+  end
+
   if alias_u0
     left, right = prob.u0
   else
@@ -28,13 +32,11 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractBracke
   fl = f(left, p)
   fr = f(right, p)
 
-  cache = alg_cache(alg, left, right, p, Val(iip))
-
-  sol = build_solution(left, Val(iip))
   if immutable
-    return BracketingImmutableSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
+    return BracketingImmutableSolver(1, f, alg, left, right, fl, fr, p, false, maxiters, :Default)
   else
-    return BracketingSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default, sol)
+    cache = alg_cache(alg, left, right, p, Val(iip))
+    return BracketingSolver(1, f, alg, left, right, fl, fr, p, cache, false, maxiters, :Default)
   end
 end
 
@@ -61,38 +63,32 @@ function DiffEqBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractNewton
     fu = f(u, p)
   end
 
-
-  sol = build_newton_solution(u, Val(iip))
   if immutable
-    return NewtonImmutableSolver(1, f, alg, u, fu, p, nothing, false, maxiters, internalnorm, :Default, tol, sol)
+    return NewtonImmutableSolver(1, f, alg, u, fu, p, false, maxiters, internalnorm, :Default, tol)
   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)
+    return NewtonSolver(1, f, alg, u, fu, p, cache, false, maxiters, internalnorm, :Default, tol)
   end
 end
 
 function DiffEqBase.solve!(solver::AbstractNonlinearSolver)
-  # sync_residuals!(solver)
   mic_check!(solver)
   while !solver.force_stop && solver.iter < solver.maxiters
     perform_step!(solver, solver.alg, solver.cache)
     solver.iter += 1
-    # sync_residuals!(solver)
   end
   if solver.iter == solver.maxiters
     solver.retcode = :MaxitersExceeded
   end
-  solver = set_solution(solver)
-  return solver.sol
+  sol = get_solution(solver)
+  return 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
@@ -101,6 +97,12 @@ function DiffEqBase.solve!(solver::AbstractImmutableNonlinearSolver)
   return sol
 end
 
+"""
+  mic_check(solver::AbstractImmutableNonlinearSolver)
+  mic_check!(solver::AbstractNonlinearSolver)
+
+Checks before running main solving iterations.
+"""
 function mic_check!(solver::BracketingSolver)
   @unpack f, fl, fr = solver
   flr = fl * fr
@@ -139,52 +141,35 @@ function mic_check(solver::NewtonImmutableSolver)
   solver
 end
 
-function check_for_exact_solution!(solver::BracketingSolver)
-  @unpack fl, fr = solver
-  fzero = zero(fl)
-  if fl == fzero
-    solver.retcode = :ExactSolutionAtLeft
-    return true
-  elseif fr == fzero
-    solver.retcode = :ExactSolutionAtRight
-    return true
-  end
-  return false
-end
-
-function set_solution(solver::BracketingSolver)
-  sol = solver.sol
-  @set! sol.left = solver.left
-  @set! sol.right = solver.right
-  @set! sol.retcode = solver.retcode
-  @set! solver.sol = sol
-  return solver
-end
+"""
+  get_solution(solver::Union{BracketingImmutableSolver, BracketingSolver})
+  get_solution(solver::Union{NewtonImmutableSolver, NewtonSolver})
 
-function get_solution(solver::BracketingImmutableSolver)
-  return (left = solver.left, right = solver.right, retcode = solver.retcode)
+Form solution object from solver types
+"""
+function get_solution(solver::Union{BracketingImmutableSolver, BracketingSolver})
+  return BracketingSolution(solver.left, solver.right, solver.retcode)
 end
 
-function set_solution(solver::NewtonSolver)
-  sol = solver.sol
-  @set! sol.u = solver.u
-  @set! sol.retcode = solver.retcode
-  @set! solver.sol = sol
-  return solver
+function get_solution(solver::Union{NewtonImmutableSolver, NewtonSolver})
+  return NewtonSolution(solver.u, solver.retcode)
 end
 
-function get_solution(solver::NewtonImmutableSolver)
-  return (u = solver.u, retcode = solver.retcode)
-end
+"""
+  reinit!(solver, prob)
 
+Reinitialize solver to the original starting conditions
+"""
 function reinit!(solver::NewtonSolver, prob::NonlinearProblem{uType, true}) where {uType}
   @. solver.u = prob.u0
   solver.iter = 1
   solver.force_stop = false
+  return solver
 end
 
 function reinit!(solver::NewtonSolver, prob::NonlinearProblem{uType, false}) where {uType}
   solver.u = prob.u0
   solver.iter = 1
   solver.force_stop = false
+  return solver
 end

src/types.jl

@@ -1,4 +1,4 @@
-mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType, solType} <: AbstractNonlinearSolver
+mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType} <: AbstractNonlinearSolver
     iter::Int
     f::fType
     alg::algType
@@ -11,10 +11,9 @@ mutable struct BracketingSolver{fType, algType, uType, resType, pType, cacheType
     force_stop::Bool
     maxiters::Int
     retcode::Symbol
-    sol::solType
 end
 
-struct BracketingImmutableSolver{fType, algType, uType, resType, pType, cacheType, solType} <: AbstractImmutableNonlinearSolver
+struct BracketingImmutableSolver{fType, algType, uType, resType, pType} <: AbstractImmutableNonlinearSolver
     iter::Int
     f::fType
     alg::algType
@@ -23,15 +22,13 @@ struct BracketingImmutableSolver{fType, algType, uType, resType, pType, cacheTyp
     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
+mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, INType, tolType} <: AbstractNonlinearSolver
     iter::Int
     f::fType
     alg::algType
@@ -44,51 +41,35 @@ mutable struct NewtonSolver{fType, algType, uType, resType, pType, cacheType, IN
     internalnorm::INType
     retcode::Symbol
     tol::tolType
-    sol::solType
 end
 
-struct NewtonImmutableSolver{fType, algType, uType, resType, pType, cacheType, INType, tolType, solType} <: AbstractImmutableNonlinearSolver
+struct NewtonImmutableSolver{fType, algType, uType, resType, pType, INType, tolType} <: 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)
-    nothing
-end
-
-mutable struct BracketingSolution{uType}
+struct BracketingSolution{uType}
     left::uType
     right::uType
     retcode::Symbol
 end
 
-function build_solution(u_prototype, ::Val{true})
-    return BracketingSolution(similar(u_prototype), similar(u_prototype), :Default)
-end
-
-function build_solution(u_prototype, ::Val{false})
-    return BracketingSolution(zero(u_prototype), zero(u_prototype), :Default)
-end
-
 struct NewtonSolution{uType}
     u::uType
     retcode::Symbol
 end
 
-function build_newton_solution(u_prototype, ::Val{iip}) where iip
-    return NewtonSolution(zero(u_prototype), :Default)
+function sync_residuals!(solver::BracketingSolver)
+    solver.fl = solver.f(solver.left, solver.p)
+    solver.fr = solver.f(solver.right, solver.p)
+    nothing
 end
-

src/utils.jl

@@ -26,4 +26,4 @@ function value_derivative(f::F, x::R) where {F,R}
 end
 
 DiffEqBase.has_Wfact(f::Function) = false
-DiffEqBase.has_Wfact_t(f::Function) = false
+DiffEqBase.has_Wfact_t(f::Function) = false

test/runtests.jl

@@ -68,3 +68,68 @@ for p in 1.1:0.1:100.0
     @test g(p) ≈ sqrt(p)
     @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p))
 end
+
+# Error Checks
+
+f, u0 = (u, p) -> u .* u .- 2.0, @SVector[1.0, 1.0]
+probN = NonlinearProblem(f, u0)
+
+@test solve(probN, NewtonRaphson()).u[end] ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(); immutable = false).u[end] ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(;autodiff=false)).u[end] ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(;autodiff=false); immutable = false).u[end] ≈ sqrt(2.0)
+
+f, u0 = (u, p) -> u .* u .- 2.0, 1.0
+probN = NonlinearProblem(f, u0)
+
+@test solve(probN, NewtonRaphson()).u ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(); immutable = false).u ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(;autodiff=false)).u ≈ sqrt(2.0)
+@test solve(probN, NewtonRaphson(;autodiff=false); immutable = false).u ≈ sqrt(2.0)
+
+
+# Bisection Tests
+f, u0 = (u, p) -> u .* u .- 2.0, (1.0, 2.0)
+probB = NonlinearProblem(f, u0)
+
+# this should call the fast scalar overload
+@test solve(probB, Bisection()).left ≈ sqrt(2.0)
+
+# these should call the iterator version
+solver = init(probB, Bisection())
+@test solver isa NonlinearSolve.BracketingImmutableSolver
+# Question: Do we need BracketingImmutableSolver? We have fast scalar overload and 
+# Bracketing solvers work only for scalars.
+
+solver = init(probB, Bisection(); immutable = false)
+@test solver isa NonlinearSolve.BracketingSolver
+@test solve!(solver).left ≈ sqrt(2.0)
+
+# Garuntee Tests for Bisection
+f = function (u, p)
+    if u < 2.0
+        return u - 2.0
+    elseif u > 3.0
+        return u - 3.0
+    else
+        return 0.0
+    end
+end
+probB = NonlinearProblem(f, (0.0, 4.0))
+
+solver = init(probB, Bisection(;exact_left = true); immutable = false)
+sol = solve!(solver)
+@test f(sol.left, nothing) < 0.0
+@test f(nextfloat(sol.left), nothing) >= 0.0
+
+solver = init(probB, Bisection(;exact_right = true); immutable = false)
+sol = solve!(solver)
+@test f(sol.right, nothing) > 0.0
+@test f(prevfloat(sol.right), nothing) <= 0.0
+
+solver = init(probB, Bisection(;exact_left = true, exact_right = true); immutable = false)
+sol = solve!(solver)
+@test f(sol.left, nothing) < 0.0
+@test f(nextfloat(sol.left), nothing) >= 0.0
+@test f(sol.right, nothing) > 0.0
+@test f(prevfloat(sol.right), nothing) <= 0.0