Skip to content

Commit c1fa892

Browse files
avik-palavik-pal
authored
Merge pull request #438 from SciML/ap/bigfloat
Add support for BigFloat
2 parents 128e353 + 301f874 commit c1fa892

File tree

12 files changed

+80
-36
lines changed

12 files changed

+80
-36
lines changed

‎Project.toml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ CUDA = "5.2"
6666
ConcreteStructs = "0.2.3"
6767
DiffEqBase = "6.149.0"
6868
Enzyme = "0.12"
69-
ExplicitImports = "1.4.4"
69+
ExplicitImports = "1.5"
7070
FastBroadcast = "0.2.8, 0.3"
7171
FastClosures = "0.3.2"
7272
FastLevenbergMarquardt = "0.1"
@@ -79,8 +79,8 @@ LineSearches = "7.2"
7979
LinearAlgebra = "1.10"
8080
LinearSolve = "2.30"
8181
MINPACK = "1.2"
82-
MaybeInplace = "0.1.1"
83-
ModelingToolkit = "9.13.0"
82+
MaybeInplace = "0.1.3"
83+
ModelingToolkit = "9.15.0"
8484
NLSolvers = "0.5"
8585
NLsolve = "4.5"
8686
NaNMath = "1"

‎src/abstract_types.jl‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -202,19 +202,17 @@ Abstract Type for all NonlinearSolve.jl Caches.
202202
abstract type AbstractNonlinearSolveCache{iip, timeit} end
203203

204204
function SymbolicIndexingInterface.symbolic_container(cache::AbstractNonlinearSolveCache)
205-
cache.prob
205+
return cache.prob
206206
end
207207
function SymbolicIndexingInterface.parameter_values(cache::AbstractNonlinearSolveCache)
208-
parameter_values(symbolic_container(cache))
208+
return parameter_values(symbolic_container(cache))
209209
end
210210
function SymbolicIndexingInterface.state_values(cache::AbstractNonlinearSolveCache)
211-
state_values(symbolic_container(cache))
211+
return state_values(symbolic_container(cache))
212212
end
213213

214214
function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol)
215-
if sym == :ps
216-
return ParameterIndexingProxy(cache)
217-
end
215+
sym == :ps && return ParameterIndexingProxy(cache)
218216
return getfield(cache, sym)
219217
end
220218

@@ -230,10 +228,17 @@ function __show_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0)
230228
println(io, "$(nameof(typeof(cache)))(")
231229
__show_algorithm(io, cache.alg,
232230
(" "^(indent + 4)) * "alg = " * string(get_name(cache.alg)), indent + 4)
233-
println(io, ",")
234-
println(io, (" "^(indent + 4)) * "u = ", get_u(cache), ",")
235-
println(io, (" "^(indent + 4)) * "residual = ", get_fu(cache), ",")
236-
println(io, (" "^(indent + 4)) * "inf-norm(residual) = ", norm(get_fu(cache), Inf), ",")
231+
232+
ustr = sprint(show, get_u(cache); context = (:compact => true, :limit => true))
233+
println(io, ",\n" * (" "^(indent + 4)) * "u = $(ustr),")
234+
235+
residstr = sprint(show, get_fu(cache); context = (:compact => true, :limit => true))
236+
println(io, (" "^(indent + 4)) * "residual = $(residstr),")
237+
238+
normstr = sprint(
239+
show, norm(get_fu(cache), Inf); context = (:compact => true, :limit => true))
240+
println(io, (" "^(indent + 4)) * "inf-norm(residual) = $(normstr),")
241+
237242
println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",")
238243
println(io, " "^(indent + 4) * "retcode = ", cache.retcode)
239244
print(io, " "^(indent) * ")")

‎src/algorithms/lbroyden.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,9 @@ end
117117

118118
function BroydenLowRankJacobian(fu, u; threshold::Int = 10, alpha = true)
119119
T = promote_type(eltype(u), eltype(fu))
120-
U = similar(fu, T, length(fu), threshold)
121-
Vᵀ = similar(u, T, length(u), threshold)
122-
cache = similar(u, T, threshold)
120+
U = __similar(fu, T, length(fu), threshold)
121+
Vᵀ = __similar(u, T, length(u), threshold)
122+
cache = __similar(u, T, threshold)
123123
return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache, T(alpha))
124124
end
125125

‎src/default.jl‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -266,11 +266,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
266266
alias_u0 = false # If immutable don't care about aliasing
267267
end
268268
u0 = prob.u0
269-
if alias_u0
270-
u0_aliased = similar(u0)
271-
else
272-
u0_aliased = u0 # Irrelevant
273-
end
269+
u0_aliased = alias_u0 ? __similar(u0) : u0
274270
end]
275271
for i in 1:N
276272
cur_sol = sol_syms[i]

‎src/globalization/line_search.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ function __internal_init(
115115
else
116116
if SciMLBase.has_jvp(f)
117117
if isinplace(prob)
118-
g_cache = similar(u)
118+
g_cache = __similar(u)
119119
grad_op = @closure (u, fu, p) -> f.vjp(g_cache, fu, u, p)
120120
else
121121
grad_op = @closure (u, fu, p) -> f.vjp(fu, u, p)
@@ -125,7 +125,7 @@ function __internal_init(
125125
alg.autodiff, prob; check_reverse_mode = true)
126126
vjp_op = VecJacOperator(prob, fu, u; autodiff)
127127
if isinplace(prob)
128-
g_cache = similar(u)
128+
g_cache = __similar(u)
129129
grad_op = @closure (u, fu, p) -> vjp_op(g_cache, fu, u, p)
130130
else
131131
grad_op = @closure (u, fu, p) -> vjp_op(fu, u, p)

‎src/internal/approximate_initialization.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ function __internal_init(
9797
@assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!"
9898
J = one.(_vec(fu)) .* α
9999
else
100-
J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
100+
J_ = __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
101101
J = alg.structure(__make_identity!!(J_, α); alias = true)
102102
end
103103
return InitializedApproximateJacobianCache(

‎src/internal/helpers.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, iip}
33
(; f, u0, p) = prob
44
if iip
5-
fu = f.resid_prototype === nothing ? similar(u) :
5+
fu = f.resid_prototype === nothing ? __similar(u) :
66
promote_type(eltype(u), eltype(f.resid_prototype)).(f.resid_prototype)
77
f(fu, u, p)
88
else
@@ -154,7 +154,7 @@ function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool
154154

155155
𝐅 = if force_oop === True && applicable(𝐟, u0, u0)
156156
_resid = resid isa Number ? [resid] : _vec(resid)
157-
du = _vec(similar(_resid))
157+
du = _vec(__similar(_resid))
158158
@closure u -> begin
159159
𝐟(du, u)
160160
return du

‎src/internal/jacobian.jl‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,10 +82,11 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
8282
else
8383
if has_analytic_jac
8484
f.jac_prototype === nothing ?
85-
similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
85+
__similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
8686
copy(f.jac_prototype)
8787
elseif f.jac_prototype === nothing
88-
init_jacobian(jac_cache; preserve_immutable = Val(true))
88+
__init_bigfloat_array!!(init_jacobian(
89+
jac_cache; preserve_immutable = Val(true)))
8990
else
9091
f.jac_prototype
9192
end

‎src/internal/operators.jl‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,8 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
7474
@closure (v, u, p) -> auto_vecjac(uf, u, v)
7575
elseif vjp_autodiff isa AutoFiniteDiff
7676
if iip
77-
cache1 = similar(fu)
78-
cache2 = similar(fu)
77+
cache1 = __similar(fu)
78+
cache2 = __similar(fu)
7979
@closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2)
8080
else
8181
@closure (v, u, p) -> num_vecjac(uf, __mutable(u), v)
@@ -106,17 +106,17 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
106106
if iip
107107
# FIXME: Technically we should propagate the tag but ignoring that for now
108108
cache1 = Dual{typeof(ForwardDiff.Tag(uf, eltype(u))), eltype(u),
109-
1}.(similar(u), ForwardDiff.Partials.(tuple.(u)))
109+
1}.(__similar(u), ForwardDiff.Partials.(tuple.(u)))
110110
cache2 = Dual{typeof(ForwardDiff.Tag(uf, eltype(fu))), eltype(fu),
111-
1}.(similar(fu), ForwardDiff.Partials.(tuple.(fu)))
111+
1}.(__similar(fu), ForwardDiff.Partials.(tuple.(fu)))
112112
@closure (Jv, v, u, p) -> auto_jacvec!(Jv, uf, u, v, cache1, cache2)
113113
else
114114
@closure (v, u, p) -> auto_jacvec(uf, u, v)
115115
end
116116
elseif jvp_autodiff isa AutoFiniteDiff
117117
if iip
118-
cache1 = similar(fu)
119-
cache2 = similar(u)
118+
cache1 = __similar(fu)
119+
cache2 = __similar(u)
120120
@closure (Jv, v, u, p) -> num_jacvec!(Jv, uf, u, v, cache1, cache2)
121121
else
122122
@closure (v, u, p) -> num_jacvec(uf, u, v)
@@ -162,15 +162,15 @@ end
162162
function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip}
163163
if vjp
164164
if iip
165-
res = similar(op.output_cache)
165+
res = __similar(op.output_cache)
166166
op.vjp_op(res, v, u, p)
167167
return res
168168
else
169169
return op.vjp_op(v, u, p)
170170
end
171171
else
172172
if iip
173-
res = similar(op.output_cache)
173+
res = __similar(op.output_cache)
174174
op.jvp_op(res, v, u, p)
175175
return res
176176
else

‎src/utils.jl‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,3 +160,16 @@ function __reinit_internal!(stats::NLStats)
160160
stats.njacs = 0
161161
stats.nsolve = 0
162162
end
163+
164+
function __similar(x, args...; kwargs...)
165+
y = similar(x, args...; kwargs...)
166+
return __init_bigfloat_array!!(y)
167+
end
168+
169+
function __init_bigfloat_array!!(x)
170+
if ArrayInterface.can_setindex(x)
171+
eltype(x) <: BigFloat && fill!(x, BigFloat(0))
172+
return x
173+
end
174+
return x
175+
end

0 commit comments

Comments
 (0)