Skip to content

Commit 2692074

Browse files
ErikQQYErikQQY
authored
Merge pull request #460 from SciML/qqy/fix_initial_guess
2 parents 2d549af + f5d177c commit 2692074

9 files changed

Lines changed: 275 additions & 99 deletions

File tree

‎lib/BoundaryValueDiffEqCore/Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ Adapt = "4.1.1"
3333
Aqua = "0.8"
3434
ArrayInterface = "7.18"
3535
ConcreteStructs = "0.2.3"
36-
DiffEqBase = "6.183"
36+
DiffEqBase = "6.213"
3737
ForwardDiff = "0.10.38, 1"
3838
Integrals = "4.7.1, 5"
3939
InteractiveUtils = "<0.0.1, 1"

‎lib/BoundaryValueDiffEqCore/src/utils.jl‎

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -323,8 +323,6 @@ function __extract_problem_details(
323323
_u0 = first(u0.u)
324324
_t = u0.t
325325
if tune_parameters
326-
prob.p isa SciMLBase.NullParameters &&
327-
throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
328326
new_u = vcat(_u0, __tunable_part(prob.p))
329327
return Val(false), eltype(new_u), length(new_u), Int(cld(t₁ - t₀, dt)), new_u
330328
end
@@ -333,11 +331,7 @@ end
333331

334332
function __initial_guess(f::F, p::P, t::T; tune_parameters = false) where {F, P, T}
335333
if hasmethod(f, Tuple{P, T})
336-
if tune_parameters
337-
p isa SciMLBase.NullParameters &&
338-
throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
339-
return vcat(f(p, t), __tunable_part(p))
340-
end
334+
tune_parameters && return vcat(f(p, t), __tunable_part(p))
341335
return f(p, t)
342336
elseif hasmethod(f, Tuple{T})
343337
Base.depwarn(
@@ -346,11 +340,7 @@ function __initial_guess(f::F, p::P, t::T; tune_parameters = false) where {F, P,
346340
removed in the next major release of SciMLBase.",
347341
:__initial_guess
348342
)
349-
if tune_parameters
350-
p isa SciMLBase.NullParameters &&
351-
throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
352-
return vcat(f(t), __tunable_part(p))
353-
end
343+
tune_parameters && return vcat(f(t), __tunable_part(p))
354344
return f(t)
355345
else
356346
throw(ArgumentError("`initial_guess` must be a function of the form `f(p, t)`"))
@@ -621,27 +611,46 @@ initial guess, it returns `vec(u₀)`.
621611
Returns the initial guess on the mesh. For `DiffEqArray` assumes that the mesh is the same
622612
as the mesh of the `DiffEqArray`.
623613
"""
624-
@inline function __initial_guess_on_mesh(u₀::AbstractVector{<:AbstractArray}, _, p)
614+
@inline function __initial_guess_on_mesh(u₀::AbstractVector{<:AbstractArray}, mesh, p; tune_parameters = false)
615+
tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀])
625616
return VectorOfArray([copy(vec(u)) for u in u₀])
626617
end
627-
@inline function __initial_guess_on_mesh(u₀::VectorOfArray, _, p)
618+
@inline function __initial_guess_on_mesh(u₀::VectorOfArray, mesh, p; tune_parameters = false)
619+
tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
628620
return copy(u₀)
629621
end
630-
@inline function __initial_guess_on_mesh(u₀::DiffEqArray, mesh, p)
622+
@inline function __initial_guess_on_mesh(u₀::DiffEqArray, mesh, p; tune_parameters = false)
623+
tune_parameters && return DiffEqArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
631624
return copy(u₀)
632625
end
633-
@inline function __initial_guess_on_mesh(u₀::SciMLBase.ODESolution, mesh, p)
626+
@inline function __initial_guess_on_mesh(u₀::SciMLBase.ODESolution, mesh, p; tune_parameters = false)
627+
tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
634628
return copy(VectorOfArray(u₀.u))
635629
end
636-
@inline function __initial_guess_on_mesh(u₀::AbstractArray, mesh, p)
630+
@inline function __initial_guess_on_mesh(u₀::AbstractArray, mesh, p; tune_parameters = false)
631+
tune_parameters && return VectorOfArray([vcat(vec(u₀), __tunable_part(p)) for _ in mesh])
637632
return VectorOfArray([copy(vec(u₀)) for _ in mesh])
638633
end
639-
@inline function __initial_guess_on_mesh(u₀::F, mesh, p) where {F}
634+
@inline function __initial_guess_on_mesh(u₀::F, mesh, p; tune_parameters = false) where {F}
635+
tune_parameters && return VectorOfArray([vcat(vec(__initial_guess(u₀, p, t)), __tunable_part(p)) for t in mesh])
640636
return VectorOfArray([vec(__initial_guess(u₀, p, t)) for t in mesh])
641637
end
638+
@inline function __initial_guess_on_mesh(u₀::Number, mesh, p; tune_parameters = false)
639+
tune_parameters && return VectorOfArray([vcat([u₀], __tunable_part(p)) for _ in mesh])
640+
return VectorOfArray([copy([u₀]) for _ in mesh])
641+
end
642642
@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::AbstractArray, Nig, p)
643643
return VectorOfArray([copy(vec(u₀)) for _ in 1:(2 * (Nig + 1))])
644644
end
645+
@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::AbstractVector{<:AbstractVector}, _, p)
646+
return VectorOfArray(vcat([copy(vec(u)) for u in u₀], [copy(vec(u)) for u in u₀]))
647+
end
648+
@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::VectorOfArray, _, p)
649+
return VectorOfArray(vcat(copy(u₀.u), copy(u₀.u)))
650+
end
651+
@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::SciMLBase.ODESolution, Nig, p)
652+
return VectorOfArray(vcat(copy(VectorOfArray(u₀.u)), copy(VectorOfArray(u₀.u))))
653+
end
645654

646655
# Construct BVP Solution
647656
function __build_solution(prob::AbstractBVProblem, odesol, nlsol::SciMLBase.NonlinearSolution)

‎lib/BoundaryValueDiffEqFIRK/src/firk.jl‎

Lines changed: 47 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,11 @@ function init_nested(
126126
end
127127
diffcache = __cache_trait(alg.jac_alg)
128128
tune_parameters = haskey(prob.kwargs, :tune_parameters)
129+
if tune_parameters
130+
prob.p isa SciMLBase.NullParameters &&
131+
throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
132+
end
133+
129134
constraint = (!isnothing(prob.f.inequality)) ||
130135
(!isnothing(prob.f.equality)) ||
131136
(!isnothing(prob.lb)) ||
@@ -135,18 +140,18 @@ function init_nested(
135140
ig, T,
136141
M,
137142
Nig,
138-
X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
143+
u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
139144
mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
140145
mesh_dt = diff(mesh)
141146

142147
chunksize = pickchunksize(M * (Nig - 1))
143148
__alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
144149

145-
fᵢ_cache = __alloc(zero(X))
146-
fᵢ₂_cache = vec(zero(X))
150+
fᵢ_cache = __alloc(zero(u0))
151+
fᵢ₂_cache = vec(zero(u0))
147152

148153
# Don't flatten this here, since we need to expand it later if needed
149-
y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
154+
y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p; tune_parameters = tune_parameters)
150155

151156
y = __alloc.(copy.(y₀.u))
152157
TU, ITU = constructRK(alg, T)
@@ -156,17 +161,17 @@ function init_nested(
156161

157162
k_discrete = if !constraint
158163
[
159-
__maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
164+
__maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
160165
for _ in 1:Nig
161166
]
162167
else
163168
[
164-
__maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
169+
__maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
165170
for _ in 1:Nig
166171
]
167172
end
168173

169-
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, X)
174+
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, u0)
170175

171176
residual = if iip
172177
if !constraint
@@ -192,12 +197,12 @@ function init_nested(
192197
nothing
193198
end
194199

195-
defect = VectorOfArray([safe_similar(X, ifelse(adaptive, M, 0)) for _ in 1:Nig])
200+
defect = VectorOfArray([safe_similar(u0, ifelse(adaptive, M, 0)) for _ in 1:Nig])
196201

197202
# Transform the functions to handle non-vector inputs
198203
bcresid_prototype = __vec(bcresid_prototype)
199204
f,
200-
bc = if X isa AbstractVector
205+
bc = if u0 isa AbstractVector
201206
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
202207
tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
203208
l_parameters = length(tunable_part)
@@ -220,39 +225,39 @@ function init_nested(
220225
prob.f, prob.f.bc
221226
end
222227
elseif iip
223-
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
228+
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
224229
vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
225-
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(X))
230+
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
226231
else
227232
(
228233
@closure(
229234
(
230235
r, u,
231236
p,
232-
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(X))
237+
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(u0))
233238
),
234239
@closure(
235240
(
236241
r, u, p,
237-
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(X))
242+
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(u0))
238243
),
239244
)
240245
end
241246
vecf!, vecbc!
242247
else
243-
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
248+
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
244249
vecbc = if !(prob.problem_type isa TwoPointBVProblem)
245-
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
250+
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
246251
else
247252
(
248-
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(X))),
249-
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(X))),
253+
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
254+
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
250255
)
251256
end
252257
vecf, vecbc
253258
end
254259

255-
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
260+
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = u0) : prob
256261

257262
# Somewhat arbitrary initialization of K
258263
K0 = __K0_on_u0(prob, stage; tune_parameters = tune_parameters)
@@ -266,7 +271,7 @@ function init_nested(
266271
end
267272

268273
return FIRKCacheNested{iip, T, typeof(diffcache), tune_parameters}(
269-
alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p,
274+
alg_order(alg), stage, M, size(u0), f, bc, prob_, prob.problem_type, prob.p,
270275
alg, TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete,
271276
y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, nestprob, resid₁_size, prob.singular_term,
272277
nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...), verbose_spec
@@ -287,6 +292,10 @@ function init_expanded(
287292
end
288293
diffcache = __cache_trait(alg.jac_alg)
289294
tune_parameters = haskey(prob.kwargs, :tune_parameters)
295+
if tune_parameters
296+
prob.p isa SciMLBase.NullParameters &&
297+
throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
298+
end
290299
constraint = (!isnothing(prob.f.inequality)) ||
291300
(!isnothing(prob.f.equality)) ||
292301
(!isnothing(prob.lb)) ||
@@ -296,7 +305,7 @@ function init_expanded(
296305
ig, T,
297306
M,
298307
Nig,
299-
X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
308+
u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
300309
mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
301310
mesh_dt = diff(mesh)
302311

@@ -308,27 +317,27 @@ function init_expanded(
308317
chunksize = pickchunksize(M + M * Nig * (stage + 1))
309318
__alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
310319

311-
fᵢ_cache = __alloc(zero(X)) # Runtime dispatch
312-
fᵢ₂_cache = vec(zero(X))
320+
fᵢ_cache = __alloc(zero(u0)) # Runtime dispatch
321+
fᵢ₂_cache = vec(zero(u0))
313322

314323
# Don't flatten this here, since we need to expand it later if needed
315-
_y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
324+
_y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p; tune_parameters = tune_parameters)
316325
y₀ = extend_y(_y₀, Nig + 1, stage)
317326
y = __alloc.(copy.(y₀.u)) # Runtime dispatch
318327

319328
k_discrete = if !constraint
320329
[
321-
__maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
330+
__maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
322331
for _ in 1:Nig
323332
] # Runtime dispatch
324333
else
325334
[
326-
__maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
335+
__maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
327336
for _ in 1:Nig
328337
] # Runtime dispatch
329338
end
330339

331-
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, X)
340+
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, u0)
332341

333342
residual = if iip
334343
if !constraint
@@ -354,12 +363,12 @@ function init_expanded(
354363
nothing
355364
end
356365

357-
defect = VectorOfArray([similar(X, ifelse(adaptive, M, 0)) for _ in 1:Nig])
366+
defect = VectorOfArray([similar(u0, ifelse(adaptive, M, 0)) for _ in 1:Nig])
358367

359368
# Transform the functions to handle non-vector inputs
360369
bcresid_prototype = __vec(bcresid_prototype)
361370
f,
362-
bc = if X isa AbstractVector
371+
bc = if u0 isa AbstractVector
363372
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
364373
tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
365374
l_parameters = length(tunable_part)
@@ -382,43 +391,43 @@ function init_expanded(
382391
prob.f, prob.f.bc
383392
end
384393
elseif iip
385-
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
394+
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
386395
vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
387-
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(X))
396+
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
388397
else
389398
(
390399
@closure(
391400
(
392401
r, u,
393402
p,
394-
) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(X))
403+
) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(u0))
395404
),
396405
@closure (
397406
(
398407
r, u,
399408
p,
400-
) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(X))
409+
) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(u0))
401410
)
402411
)
403412
end
404413
vecf!, vecbc!
405414
else
406-
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
415+
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
407416
vecbc = if !(prob.problem_type isa TwoPointBVProblem)
408-
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
417+
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
409418
else
410419
(
411-
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(X))),
412-
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(X))),
420+
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
421+
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
413422
)
414423
end
415424
vecf, vecbc
416425
end
417426

418-
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
427+
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = u0) : prob
419428

420429
return FIRKCacheExpand{iip, T, typeof(diffcache), tune_parameters}(
421-
alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p,
430+
alg_order(alg), stage, M, size(u0), f, bc, prob_, prob.problem_type, prob.p,
422431
alg, TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete,
423432
y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, resid₁_size, prob.singular_term, nlsolve_kwargs,
424433
optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...), verbose_spec

‎lib/BoundaryValueDiffEqFIRK/test/expanded/firk_basic_tests.jl‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -525,3 +525,46 @@ end
525525
@test sol_struct.prob.p.params ≈ [17.09658] atol = 1.0e-5
526526
@test sol_struct.prob.p.params ≈ sol_vec.prob.p atol = 1.0e-10
527527
end
528+
529+
#=
530+
# The initial guess for expanded FIRK just stall the CI, need to find out why.
531+
@testitem "Test initial guess" begin
532+
tspan = (0.0, 1.0)
533+
function f!(du, u, p, t)
534+
cond = 0.002
535+
vol_heat = 0.2
536+
du[1] = -u[2] / cond
537+
du[2] = vol_heat
538+
du[3] = 0.0
539+
end
540+
function bca!(res_a, u_a, p)
541+
res_a[1] = u_a[2]
542+
res_a[2] = u_a[1] - 100.0
543+
end
544+
function bcb!(res_b, u_b, p)
545+
tref = 20.0
546+
res_b[1] = u_b[3] * (u_b[1] - tref) - u_b[2]
547+
end
548+
u_guess = [
549+
[100.0, 0.0, 0.006666666666666668],
550+
[99.5, 0.020000000000000004, 0.006666666666666668],
551+
[98.0, 0.04000000000000001, 0.006666666666666668],
552+
[95.5, 0.060000000000000005, 0.006666666666666668],
553+
[92.0, 0.08000000000000002, 0.006666666666666668],
554+
[87.5, 0.1, 0.006666666666666668],
555+
[82.0, 0.12000000000000001, 0.006666666666666668],
556+
[75.5, 0.14, 0.006666666666666668],
557+
[68.0, 0.16000000000000003, 0.006666666666666668],
558+
[59.49999999999999, 0.18000000000000002, 0.006666666666666668],
559+
[50.0, 0.2, 0.006666666666666668],
560+
]
561+
562+
bvp1 = TwoPointBVProblem(f!, (bca!, bcb!), u_guess, tspan; bcresid_prototype = (zeros(2), zeros(1)))
563+
sol1 = solve(bvp1, LobattoIIIa3(), dt = 0.1, adaptive = false, nlsolve_kwargs = (; maxiters = 0))
564+
@test sol1.u == u_guess
565+
566+
bvp2 = TwoPointBVProblem(f!, (bca!, bcb!), sol1, tspan; bcresid_prototype = (zeros(2), zeros(1)))
567+
sol2 = solve(bvp2, LobattoIIIa3(), dt = 0.1, adaptive = false, nlsolve_kwargs = (; maxiters = 0))
568+
@test sol2.u == u_guess
569+
end
570+
=#

0 commit comments

Comments
 (0)