From 86e67209904b860e974b67a0c6b734ccf425a04c Mon Sep 17 00:00:00 2001 From: ErikQQY <2283984853@qq.com> Date: Mon, 19 Aug 2024 15:25:35 +0800 Subject: [PATCH] Better BigFloat support Signed-off-by: ErikQQY <2283984853@qq.com> --- src/default.jl | 2 +- src/globalization/line_search.jl | 8 ++++---- src/internal/helpers.jl | 4 ++-- src/internal/jacobian.jl | 2 +- src/internal/operators.jl | 16 ++++++++-------- src/utils.jl | 10 +--------- 6 files changed, 17 insertions(+), 25 deletions(-) diff --git a/src/default.jl b/src/default.jl index ab35b7a55..b4613a6f7 100644 --- a/src/default.jl +++ b/src/default.jl @@ -266,7 +266,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb alias_u0 = false # If immutable don't care about aliasing end u0 = prob.u0 - u0_aliased = alias_u0 ? __similar(u0) : u0 + u0_aliased = alias_u0 ? zero(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 4fca295e7..5de4610b6 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -125,7 +125,7 @@ function __internal_init( deriv_op = nothing elseif SciMLBase.has_jvp(f) if isinplace(prob) - jvp_cache = __similar(fu) + jvp_cache = zero(fu) deriv_op = @closure (du, u, fu, p) -> begin f.jvp(jvp_cache, du, u, p) dot(fu, jvp_cache) @@ -135,7 +135,7 @@ function __internal_init( end elseif SciMLBase.has_vjp(f) if isinplace(prob) - vjp_cache = __similar(u) + vjp_cache = zero(u) deriv_op = @closure (du, u, fu, p) -> begin f.vjp(vjp_cache, fu, u, p) dot(du, vjp_cache) @@ -149,7 +149,7 @@ function __internal_init( alg.autodiff, prob; check_reverse_mode = true) vjp_op = VecJacOperator(prob, fu, u; autodiff) if isinplace(prob) - vjp_cache = __similar(u) + vjp_cache = zero(u) deriv_op = @closure (du, u, fu, p) -> dot(du, vjp_op(vjp_cache, fu, u, p)) else deriv_op = @closure (du, u, fu, p) -> dot(du, vjp_op(fu, u, p)) @@ -159,7 +159,7 @@ function __internal_init( alg.autodiff, prob; check_forward_mode = true) jvp_op = JacVecOperator(prob, fu, u; autodiff) if isinplace(prob) - jvp_cache = __similar(fu) + jvp_cache = zero(fu) deriv_op = @closure (du, u, fu, p) -> dot(fu, jvp_op(jvp_cache, du, u, p)) else deriv_op = @closure (du, u, fu, p) -> dot(fu, jvp_op(du, u, p)) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index db88a7d23..abf6f1099 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 ? zero(u) : promote_type(eltype(u), eltype(f.resid_prototype)).(f.resid_prototype) f(fu, u, p) else @@ -156,7 +156,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(zero(_resid)) @closure u -> begin 𝐟(du, u) return du diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 6a549721d..b712b93e2 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -85,7 +85,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) : copy(f.jac_prototype) elseif f.jac_prototype === nothing - __init_bigfloat_array!!(init_jacobian( + zero(init_jacobian( jac_cache; preserve_immutable = Val(true))) else f.jac_prototype diff --git a/src/internal/operators.jl b/src/internal/operators.jl index c7d8f586e..5bbfcb0bf 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 = zero(fu) + cache2 = zero(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}.(zero(u), ForwardDiff.Partials.(tuple.(u))) cache2 = Dual{typeof(ForwardDiff.Tag(uf, eltype(fu))), eltype(fu), - 1}.(__similar(fu), ForwardDiff.Partials.(tuple.(fu))) + 1}.(zero(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 = zero(fu) + cache2 = zero(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 = zero(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 = zero(op.output_cache) op.jvp_op(res, v, u, p) return res else diff --git a/src/utils.jl b/src/utils.jl index 50faa710c..64399a349 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -163,13 +163,5 @@ 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 + return zero(y) end