From 2376418cca94aac8e6fb7fa1baa03d07bc0aaea2 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 25 May 2024 13:11:27 -0700 Subject: [PATCH] Add support for BigFloat --- Project.toml | 4 ++-- src/abstract_types.jl | 7 ++++--- src/algorithms/lbroyden.jl | 6 +++--- src/default.jl | 6 +----- src/globalization/line_search.jl | 4 ++-- src/internal/approximate_initialization.jl | 2 +- src/internal/helpers.jl | 4 ++-- src/internal/jacobian.jl | 5 +++-- src/internal/operators.jl | 16 ++++++++-------- src/utils.jl | 13 +++++++++++++ 10 files changed, 39 insertions(+), 28 deletions(-) diff --git a/Project.toml b/Project.toml index 6c656d419..35fdb7046 100644 --- a/Project.toml +++ b/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,7 +79,7 @@ LineSearches = "7.2" LinearAlgebra = "1.10" LinearSolve = "2.30" MINPACK = "1.2" -MaybeInplace = "0.1.1" +MaybeInplace = "0.1.3" ModelingToolkit = "9.13.0" NLSolvers = "0.5" NLsolve = "4.5" diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 20b491e19..abb603868 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -229,13 +229,14 @@ function __show_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0) __show_algorithm(io, cache.alg, (" "^(indent + 4)) * "alg = " * string(get_name(cache.alg)), indent + 4) - ustr = sprint(show, get_u(cache); context=(:compact=>true, :limit=>true)) + 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)) + 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)) + 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, ",") diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 596172ead..bead1a057 100644 --- a/src/algorithms/lbroyden.jl +++ b/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 diff --git a/src/default.jl b/src/default.jl index 660caab41..ab35b7a55 100644 --- a/src/default.jl +++ b/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] diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 038cc119a..e09a8d188 100644 --- a/src/globalization/line_search.jl +++ b/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) diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 4f17cd72f..5ce348ae3 100644 --- a/src/internal/approximate_initialization.jl +++ b/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( diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index c5241481a..83b169d7f 100644 --- a/src/internal/helpers.jl +++ b/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 diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 2abd89336..6a549721d 100644 --- a/src/internal/jacobian.jl +++ b/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 diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 8c3c8922c..c7d8f586e 100644 --- a/src/internal/operators.jl +++ b/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,7 +162,7 @@ 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 @@ -170,7 +170,7 @@ function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} end else if iip - res = similar(op.output_cache) + res = __similar(op.output_cache) op.jvp_op(res, v, u, p) return res else diff --git a/src/utils.jl b/src/utils.jl index d13555259..50faa710c 100644 --- a/src/utils.jl +++ b/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