Merge pull request #438 from SciML/ap/bigfloat

Add support for BigFloat

Author: avik-pal

Project.toml

@@ -66,7 +66,7 @@ CUDA = "5.2"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.149.0"
 Enzyme = "0.12"
-ExplicitImports = "1.4.4"
+ExplicitImports = "1.5"
 FastBroadcast = "0.2.8, 0.3"
 FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
@@ -79,8 +79,8 @@ LineSearches = "7.2"
 LinearAlgebra = "1.10"
 LinearSolve = "2.30"
 MINPACK = "1.2"
-MaybeInplace = "0.1.1"
-ModelingToolkit = "9.13.0"
+MaybeInplace = "0.1.3"
+ModelingToolkit = "9.15.0"
 NLSolvers = "0.5"
 NLsolve = "4.5"
 NaNMath = "1"

src/abstract_types.jl

@@ -202,19 +202,17 @@ Abstract Type for all NonlinearSolve.jl Caches.
 abstract type AbstractNonlinearSolveCache{iip, timeit} end
 
 function SymbolicIndexingInterface.symbolic_container(cache::AbstractNonlinearSolveCache)
-    cache.prob
+    return cache.prob
 end
 function SymbolicIndexingInterface.parameter_values(cache::AbstractNonlinearSolveCache)
-    parameter_values(symbolic_container(cache))
+    return parameter_values(symbolic_container(cache))
 end
 function SymbolicIndexingInterface.state_values(cache::AbstractNonlinearSolveCache)
-    state_values(symbolic_container(cache))
+    return state_values(symbolic_container(cache))
 end
 
 function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol)
-    if sym == :ps
-        return ParameterIndexingProxy(cache)
-    end
+    sym == :ps && return ParameterIndexingProxy(cache)
     return getfield(cache, sym)
 end
 
@@ -230,10 +228,17 @@ function __show_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0)
     println(io, "$(nameof(typeof(cache)))(")
     __show_algorithm(io, cache.alg,
         (" "^(indent + 4)) * "alg = " * string(get_name(cache.alg)), indent + 4)
-    println(io, ",")
-    println(io, (" "^(indent + 4)) * "u = ", get_u(cache), ",")
-    println(io, (" "^(indent + 4)) * "residual = ", get_fu(cache), ",")
-    println(io, (" "^(indent + 4)) * "inf-norm(residual) = ", norm(get_fu(cache), Inf), ",")
+
+    ustr = sprint(show, get_u(cache); context = (:compact => true, :limit => true))
+    println(io, ",\n" * (" "^(indent + 4)) * "u = $(ustr),")
+
+    residstr = sprint(show, get_fu(cache); context = (:compact => true, :limit => true))
+    println(io, (" "^(indent + 4)) * "residual = $(residstr),")
+
+    normstr = sprint(
+        show, norm(get_fu(cache), Inf); context = (:compact => true, :limit => true))
+    println(io, (" "^(indent + 4)) * "inf-norm(residual) = $(normstr),")
+
     println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",")
     println(io, " "^(indent + 4) * "retcode = ", cache.retcode)
     print(io, " "^(indent) * ")")

src/algorithms/lbroyden.jl

@@ -117,9 +117,9 @@ end
 
 function BroydenLowRankJacobian(fu, u; threshold::Int = 10, alpha = true)
     T = promote_type(eltype(u), eltype(fu))
-    U = similar(fu, T, length(fu), threshold)
-    Vᵀ = similar(u, T, length(u), threshold)
-    cache = similar(u, T, threshold)
+    U = __similar(fu, T, length(fu), threshold)
+    Vᵀ = __similar(u, T, length(u), threshold)
+    cache = __similar(u, T, threshold)
     return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache, T(alpha))
 end
 

src/default.jl

@@ -266,11 +266,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
                     alias_u0 = false  # If immutable don't care about aliasing
                 end
                 u0 = prob.u0
-                if alias_u0
-                    u0_aliased = similar(u0)
-                else
-                    u0_aliased = u0  # Irrelevant
-                end
+                u0_aliased = alias_u0 ? __similar(u0) : u0
             end]
             for i in 1:N
                 cur_sol = sol_syms[i]

src/globalization/line_search.jl

@@ -115,7 +115,7 @@ function __internal_init(
     else
         if SciMLBase.has_jvp(f)
             if isinplace(prob)
-                g_cache = similar(u)
+                g_cache = __similar(u)
                 grad_op = @closure (u, fu, p) -> f.vjp(g_cache, fu, u, p)
             else
                 grad_op = @closure (u, fu, p) -> f.vjp(fu, u, p)
@@ -125,7 +125,7 @@ function __internal_init(
                 alg.autodiff, prob; check_reverse_mode = true)
             vjp_op = VecJacOperator(prob, fu, u; autodiff)
             if isinplace(prob)
-                g_cache = similar(u)
+                g_cache = __similar(u)
                 grad_op = @closure (u, fu, p) -> vjp_op(g_cache, fu, u, p)
             else
                 grad_op = @closure (u, fu, p) -> vjp_op(fu, u, p)

src/internal/approximate_initialization.jl

@@ -97,7 +97,7 @@ function __internal_init(
         @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!"
         J = one.(_vec(fu)) .* α
     else
-        J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
+        J_ = __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
         J = alg.structure(__make_identity!!(J_, α); alias = true)
     end
     return InitializedApproximateJacobianCache(

src/internal/helpers.jl

@@ -2,7 +2,7 @@
 function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, iip}
     (; f, u0, p) = prob
     if iip
-        fu = f.resid_prototype === nothing ? similar(u) :
+        fu = f.resid_prototype === nothing ? __similar(u) :
              promote_type(eltype(u), eltype(f.resid_prototype)).(f.resid_prototype)
         f(fu, u, p)
     else
@@ -154,7 +154,7 @@ function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool
 
     𝐅 = if force_oop === True && applicable(𝐟, u0, u0)
         _resid = resid isa Number ? [resid] : _vec(resid)
-        du = _vec(similar(_resid))
+        du = _vec(__similar(_resid))
         @closure u -> begin
             𝐟(du, u)
             return du

src/internal/jacobian.jl

@@ -82,10 +82,11 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
     else
         if has_analytic_jac
             f.jac_prototype === nothing ?
-            similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
+            __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
             copy(f.jac_prototype)
         elseif f.jac_prototype === nothing
-            init_jacobian(jac_cache; preserve_immutable = Val(true))
+            __init_bigfloat_array!!(init_jacobian(
+                jac_cache; preserve_immutable = Val(true)))
         else
             f.jac_prototype
         end

src/internal/operators.jl

@@ -74,8 +74,8 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
             @closure (v, u, p) -> auto_vecjac(uf, u, v)
         elseif vjp_autodiff isa AutoFiniteDiff
             if iip
-                cache1 = similar(fu)
-                cache2 = similar(fu)
+                cache1 = __similar(fu)
+                cache2 = __similar(fu)
                 @closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2)
             else
                 @closure (v, u, p) -> num_vecjac(uf, __mutable(u), v)
@@ -106,17 +106,17 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
             if iip
                 # FIXME: Technically we should propagate the tag but ignoring that for now
                 cache1 = Dual{typeof(ForwardDiff.Tag(uf, eltype(u))), eltype(u),
-                    1}.(similar(u), ForwardDiff.Partials.(tuple.(u)))
+                    1}.(__similar(u), ForwardDiff.Partials.(tuple.(u)))
                 cache2 = Dual{typeof(ForwardDiff.Tag(uf, eltype(fu))), eltype(fu),
-                    1}.(similar(fu), ForwardDiff.Partials.(tuple.(fu)))
+                    1}.(__similar(fu), ForwardDiff.Partials.(tuple.(fu)))
                 @closure (Jv, v, u, p) -> auto_jacvec!(Jv, uf, u, v, cache1, cache2)
             else
                 @closure (v, u, p) -> auto_jacvec(uf, u, v)
             end
         elseif jvp_autodiff isa AutoFiniteDiff
             if iip
-                cache1 = similar(fu)
-                cache2 = similar(u)
+                cache1 = __similar(fu)
+                cache2 = __similar(u)
                 @closure (Jv, v, u, p) -> num_jacvec!(Jv, uf, u, v, cache1, cache2)
             else
                 @closure (v, u, p) -> num_jacvec(uf, u, v)
@@ -162,15 +162,15 @@ end
 function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip}
     if vjp
         if iip
-            res = similar(op.output_cache)
+            res = __similar(op.output_cache)
             op.vjp_op(res, v, u, p)
             return res
         else
             return op.vjp_op(v, u, p)
         end
     else
         if iip
-            res = similar(op.output_cache)
+            res = __similar(op.output_cache)
             op.jvp_op(res, v, u, p)
             return res
         else

src/utils.jl

@@ -160,3 +160,16 @@ function __reinit_internal!(stats::NLStats)
     stats.njacs = 0
     stats.nsolve = 0
 end
+
+function __similar(x, args...; kwargs...)
+    y = similar(x, args...; kwargs...)
+    return __init_bigfloat_array!!(y)
+end
+
+function __init_bigfloat_array!!(x)
+    if ArrayInterface.can_setindex(x)
+        eltype(x) <: BigFloat && fill!(x, BigFloat(0))
+        return x
+    end
+    return x
+end

test/misc/exotic_type_tests.jl

@@ -0,0 +1,27 @@
+# File for different types of exotic types
+@testsetup module NonlinearSolveExoticTypeTests
+using NonlinearSolve
+
+fn_iip = NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p)
+fn_oop = NonlinearFunction{false}((u, p) -> u .* u .- p)
+
+u0 = BigFloat[1.0, 1.0, 1.0]
+prob_iip_bf = NonlinearProblem{true}(fn_iip, u0, BigFloat(2))
+prob_oop_bf = NonlinearProblem{false}(fn_oop, u0, BigFloat(2))
+
+export fn_iip, fn_oop, u0, prob_iip_bf, prob_oop_bf
+end
+
+@testitem "BigFloat Support" tags=[:misc] setup=[NonlinearSolveExoticTypeTests] begin
+    using NonlinearSolve, LinearAlgebra
+
+    for alg in [NewtonRaphson(), Broyden(), Klement(), DFSane(), TrustRegion()]
+        sol = solve(prob_oop_bf, alg)
+        @test norm(sol.resid, Inf) < 1e-6
+        @test SciMLBase.successful_retcode(sol.retcode)
+
+        sol = solve(prob_iip_bf, alg)
+        @test norm(sol.resid, Inf) < 1e-6
+        @test SciMLBase.successful_retcode(sol.retcode)
+    end
+end

test/misc/qa_tests.jl

@@ -26,4 +26,6 @@ end
         skip = (NonlinearSolve, Base, Core, SimpleNonlinearSolve, SciMLBase)) === nothing
 
     @test check_no_stale_explicit_imports(NonlinearSolve) === nothing
+
+    @test check_all_qualified_accesses_via_owners(NonlinearSolve) === nothing
 end