handle complex numbers correctly by default

avik-pal · avik-pal · commit 133b1d59a53a · 2023-12-19T18:48:41.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.1.1"
+version = "3.1.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -71,6 +71,7 @@ MaybeInplace = "0.1.1"
 NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.1"
+OrdinaryDiffEq = "6"
 Pkg = "1"
 PrecompileTools = "1.2"
 Printf = "<0.0.1, 1"
@@ -106,6 +107,7 @@ MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -117,4 +119,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq"]
diff --git a/src/default.jl b/src/default.jl
@@ -165,8 +165,8 @@ function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwa
 end
 
 """
-    RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
-        autodiff = nothing)
+    RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
+        precs = DEFAULT_PRECS, autodiff = nothing)
 
 A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different
 globalizing techniques (trust region updates, line searches, etc.) in order to find a
@@ -176,6 +176,11 @@ Basically, if this algorithm fails, then "most" good ways of solving your proble
 you may need to think about reformulating the model (either there is an issue with the model,
 or more precision / more stable linear solver choice is required).
 
+### Arguments
+
+  - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
+    are compatible with the problem type. Defaults to `Float64`.
+
 ### Keyword Arguments
 
   - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
@@ -196,28 +201,38 @@ or more precision / more stable linear solver choice is required).
     algorithms, consult the
     [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
 """
-function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing,
-        precs = DEFAULT_PRECS, autodiff = nothing)
-    algs = (TrustRegion(; concrete_jac, linsolve, precs),
-        TrustRegion(; concrete_jac, linsolve, precs, autodiff,
-            radius_update_scheme = RadiusUpdateSchemes.Bastin),
-        NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
-            autodiff),
-        TrustRegion(; concrete_jac, linsolve, precs,
-            radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff),
-        TrustRegion(; concrete_jac, linsolve, precs,
-            radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff))
+function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
+        precs = DEFAULT_PRECS, autodiff = nothing) where {T}
+    if __is_complex(T)
+        # Let's atleast have something here for complex numbers
+        algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)
+    else
+        algs = (TrustRegion(; concrete_jac, linsolve, precs),
+            TrustRegion(; concrete_jac, linsolve, precs, autodiff,
+                radius_update_scheme = RadiusUpdateSchemes.Bastin),
+            NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
+                autodiff),
+            TrustRegion(; concrete_jac, linsolve, precs,
+                radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff),
+            TrustRegion(; concrete_jac, linsolve, precs,
+                radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff))
+    end
     return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
 end
 
 """
-    FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing,
-        precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false),
-        prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing)
+    FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing,
+        linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false),
+        prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T}
 
 A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
 for more performance and then tries more robust techniques if the faster ones fail.
 
+### Arguments
+
+  - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
+    are compatible with the problem type. Defaults to `Float64`.
+
 ### Keyword Arguments
 
   - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
@@ -238,53 +253,76 @@ for more performance and then tries more robust techniques if the faster ones fa
     algorithms, consult the
     [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
 """
-function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing,
-        precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false),
+function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing,
+        linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false),
         prefer_simplenonlinearsolve::Val{SA} = Val(false),
-        autodiff = nothing) where {JAC, SA}
+        autodiff = nothing) where {T, JAC, SA}
     if JAC
-        algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
-            NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
-                autodiff),
-            TrustRegion(; concrete_jac, linsolve, precs, autodiff),
-            TrustRegion(; concrete_jac, linsolve, precs,
-                radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
-    else
-        # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians
-        # and thus are not included in the polyalgorithm
-        if SA
-            algs = (SimpleBroyden(),
-                Broyden(; init_jacobian = Val(:true_jacobian)),
-                SimpleKlement(),
-                NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
-                NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
-                    autodiff),
-                NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
-                    autodiff),
-                TrustRegion(; concrete_jac, linsolve, precs,
-                    radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
+        if __is_complex(T)
+            algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)
         else
-            algs = (Broyden(),
-                Broyden(; init_jacobian = Val(:true_jacobian)),
-                Klement(; linsolve, precs),
-                NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
+            algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
                 NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
                     autodiff),
                 TrustRegion(; concrete_jac, linsolve, precs, autodiff),
                 TrustRegion(; concrete_jac, linsolve, precs,
                     radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
         end
+    else
+        # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians
+        # and thus are not included in the polyalgorithm
+        if SA
+            if __is_complex(T)
+                algs = (SimpleBroyden(),
+                    Broyden(; init_jacobian = Val(:true_jacobian)),
+                    SimpleKlement(),
+                    NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
+            else
+                algs = (SimpleBroyden(),
+                    Broyden(; init_jacobian = Val(:true_jacobian)),
+                    SimpleKlement(),
+                    NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
+                    NewtonRaphson(; concrete_jac, linsolve, precs,
+                        linesearch = BackTracking(), autodiff),
+                    NewtonRaphson(; concrete_jac, linsolve, precs,
+                        linesearch = BackTracking(), autodiff),
+                    TrustRegion(; concrete_jac, linsolve, precs,
+                        radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
+            end
+        else
+            if __is_complex(T)
+                algs = (Broyden(),
+                    Broyden(; init_jacobian = Val(:true_jacobian)),
+                    Klement(; linsolve, precs),
+                    NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
+            else
+                algs = (Broyden(),
+                    Broyden(; init_jacobian = Val(:true_jacobian)),
+                    Klement(; linsolve, precs),
+                    NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
+                    NewtonRaphson(; concrete_jac, linsolve, precs,
+                        linesearch = BackTracking(), autodiff),
+                    TrustRegion(; concrete_jac, linsolve, precs, autodiff),
+                    TrustRegion(; concrete_jac, linsolve, precs,
+                        radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
+            end
+        end
     end
     return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
 end
 
 """
-    FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
-                              precs = DEFAULT_PRECS, kwargs...)
+    FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
+        precs = DEFAULT_PRECS, kwargs...)
 
 A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
 for more performance and then tries more robust techniques if the faster ones fail.
 
+### Arguments
+
+  - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
+    are compatible with the problem type. Defaults to `Float64`.
+
 ### Keyword Arguments
 
   - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
@@ -305,42 +343,58 @@ for more performance and then tries more robust techniques if the faster ones fa
     algorithms, consult the
     [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
 """
-function FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
-        precs = DEFAULT_PRECS, kwargs...)
-    algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),
-        GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
-            kwargs...),
-        LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
+function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing,
+        linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T}
+    if __is_complex(T)
+        algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),
+            LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
+    else
+        algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),
+            TrustRegion(; concrete_jac, linsolve, precs, kwargs...),
+            GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
+                kwargs...),
+            TrustRegion(; concrete_jac, linsolve, precs,
+                radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...),
+            LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
+    end
     return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS))
 end
 
 ## Defaults
 
 ## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify
 ##       the conditioning of the Jacobian and such
+
+## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support
+##       complex numbers. We should use the `DiffEqBase` trait for this once all of the
+##       NonlinearSolve algorithms support it. For now we just do a check and remove the
+##       unsupported ones from default
+
 ## Defaults to a fast and robust poly algorithm in most cases. If the user went through
 ## the trouble of specifying a custom jacobian function, we should use algorithms that
 ## can use that!
-
 function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
     must_use_jacobian = Val(prob.f.jac !== nothing)
-    return SciMLBase.__init(prob, FastShortcutNonlinearPolyalg(; must_use_jacobian),
+    return SciMLBase.__init(prob,
+        FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian),
         args...; kwargs...)
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
     must_use_jacobian = Val(prob.f.jac !== nothing)
     prefer_simplenonlinearsolve = Val(prob.u0 isa SArray)
     return SciMLBase.__solve(prob,
-        FastShortcutNonlinearPolyalg(; must_use_jacobian,
+        FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian,
             prefer_simplenonlinearsolve), args...; kwargs...)
 end
 
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...)
-    return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
+    return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...;
+        kwargs...)
 end
 
 function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...;
         kwargs...)
-    return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
+    return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...;
+        kwargs...)
 end
diff --git a/src/utils.jl b/src/utils.jl
@@ -494,3 +494,8 @@ end
 @inline __diag(x::AbstractMatrix) = diag(x)
 @inline __diag(x::AbstractVector) = x
 @inline __diag(x::Number) = x
+
+@inline __is_complex(::Type{ComplexF64}) = true
+@inline __is_complex(::Type{ComplexF32}) = true
+@inline __is_complex(::Type{Complex}) = true
+@inline __is_complex(::Type{T}) where {T} = false
diff --git a/test/polyalgs.jl b/test/polyalgs.jl
@@ -1,4 +1,4 @@
-using NonlinearSolve, Test, NaNMath
+using NonlinearSolve, Test, NaNMath, OrdinaryDiffEq
 
 f(u, p) = u .* u .- 2
 u0 = [1.0, 1.0]
@@ -58,3 +58,28 @@ p = 2.0
 prob = NonlinearProblem(ff_interval, u0, p)
 sol = solve(prob; abstol = 1e-9)
 @test SciMLBase.successful_retcode(sol)
+
+# Shooting Problem: Taken from BoundaryValueDiffEq.jl
+# Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the
+# algorithms which dont support those.
+function ode_func!(du, u, p, t)
+    du[1] = u[2]
+    du[2] = -u[1]
+    return nothing
+end
+
+function objective_function!(resid, u0, p)
+    odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p)
+    sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false)
+    resid[1] = sol(0.0)[1]
+    resid[2] = sol(100.0)[1] - 1.0
+    return nothing
+end
+
+prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im)
+sol = solve(prob; abstol = 1e-10)
+@test SciMLBase.successful_retcode(sol)
+# This test is not meant to return success but test that all the default solvers can handle
+# complex valued problems
+@test_nowarn solve(prob; abstol = 1e-19, maxiters = 10)
+@test_nowarn solve(prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, maxiters = 10)
