Skip to content

Commit 8ac560b

Browse files
ErikQQYErikQQY
committed
Fix initial guess in collocation solvers
1 parent cd89ddb commit 8ac560b

3 files changed

Lines changed: 78 additions & 78 deletions

File tree

‎lib/BoundaryValueDiffEqFIRK/src/firk.jl‎

Lines changed: 38 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -135,18 +135,18 @@ function init_nested(
135135
ig, T,
136136
M,
137137
Nig,
138-
X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
138+
u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
139139
mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
140140
mesh_dt = diff(mesh)
141141

142142
chunksize = pickchunksize(M * (Nig - 1))
143143
__alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
144144

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

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

151151
y = __alloc.(copy.(y₀.u))
152152
TU, ITU = constructRK(alg, T)
@@ -156,17 +156,17 @@ function init_nested(
156156

157157
k_discrete = if !constraint
158158
[
159-
__maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
159+
__maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
160160
for _ in 1:Nig
161161
]
162162
else
163163
[
164-
__maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
164+
__maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
165165
for _ in 1:Nig
166166
]
167167
end
168168

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

171171
residual = if iip
172172
if !constraint
@@ -192,12 +192,12 @@ function init_nested(
192192
nothing
193193
end
194194

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

197197
# Transform the functions to handle non-vector inputs
198198
bcresid_prototype = __vec(bcresid_prototype)
199199
f,
200-
bc = if X isa AbstractVector
200+
bc = if u0 isa AbstractVector
201201
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
202202
tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
203203
l_parameters = length(tunable_part)
@@ -220,39 +220,39 @@ function init_nested(
220220
prob.f, prob.f.bc
221221
end
222222
elseif iip
223-
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
223+
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
224224
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))
225+
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
226226
else
227227
(
228228
@closure(
229229
(
230230
r, u,
231231
p,
232-
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(X))
232+
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(u0))
233233
),
234234
@closure(
235235
(
236236
r, u, p,
237-
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(X))
237+
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(u0))
238238
),
239239
)
240240
end
241241
vecf!, vecbc!
242242
else
243-
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
243+
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
244244
vecbc = if !(prob.problem_type isa TwoPointBVProblem)
245-
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
245+
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
246246
else
247247
(
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))),
248+
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
249+
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
250250
)
251251
end
252252
vecf, vecbc
253253
end
254254

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

257257
# Somewhat arbitrary initialization of K
258258
K0 = __K0_on_u0(prob, stage; tune_parameters = tune_parameters)
@@ -266,7 +266,7 @@ function init_nested(
266266
end
267267

268268
return FIRKCacheNested{iip, T, typeof(diffcache), tune_parameters}(
269-
alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p,
269+
alg_order(alg), stage, M, size(u0), f, bc, prob_, prob.problem_type, prob.p,
270270
alg, TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete,
271271
y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, nestprob, resid₁_size, prob.singular_term,
272272
nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...), verbose_spec
@@ -296,7 +296,7 @@ function init_expanded(
296296
ig, T,
297297
M,
298298
Nig,
299-
X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
299+
u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
300300
mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
301301
mesh_dt = diff(mesh)
302302

@@ -308,27 +308,27 @@ function init_expanded(
308308
chunksize = pickchunksize(M + M * Nig * (stage + 1))
309309
__alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
310310

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

314314
# Don't flatten this here, since we need to expand it later if needed
315-
_y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
315+
_y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p)
316316
y₀ = extend_y(_y₀, Nig + 1, stage)
317317
y = __alloc.(copy.(y₀.u)) # Runtime dispatch
318318

319319
k_discrete = if !constraint
320320
[
321-
__maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
321+
__maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
322322
for _ in 1:Nig
323323
] # Runtime dispatch
324324
else
325325
[
326-
__maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
326+
__maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
327327
for _ in 1:Nig
328328
] # Runtime dispatch
329329
end
330330

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

333333
residual = if iip
334334
if !constraint
@@ -354,12 +354,12 @@ function init_expanded(
354354
nothing
355355
end
356356

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

359359
# Transform the functions to handle non-vector inputs
360360
bcresid_prototype = __vec(bcresid_prototype)
361361
f,
362-
bc = if X isa AbstractVector
362+
bc = if u0 isa AbstractVector
363363
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
364364
tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
365365
l_parameters = length(tunable_part)
@@ -382,43 +382,43 @@ function init_expanded(
382382
prob.f, prob.f.bc
383383
end
384384
elseif iip
385-
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
385+
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
386386
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))
387+
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
388388
else
389389
(
390390
@closure(
391391
(
392392
r, u,
393393
p,
394-
) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(X))
394+
) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(u0))
395395
),
396396
@closure (
397397
(
398398
r, u,
399399
p,
400-
) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(X))
400+
) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(u0))
401401
)
402402
)
403403
end
404404
vecf!, vecbc!
405405
else
406-
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
406+
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
407407
vecbc = if !(prob.problem_type isa TwoPointBVProblem)
408-
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
408+
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
409409
else
410410
(
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))),
411+
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
412+
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
413413
)
414414
end
415415
vecf, vecbc
416416
end
417417

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

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

‎lib/BoundaryValueDiffEqMIRK/src/mirk.jl‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -57,18 +57,18 @@ function SciMLBase.__init(
5757
ig, T,
5858
N,
5959
Nig,
60-
X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
60+
u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
6161
mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
6262
mesh_dt = diff(mesh)
6363

6464
chunksize = pickchunksize(N * (Nig - 1))
6565
__alloc = @closure x -> __maybe_allocate_diffcache(vec(zero(x)), chunksize, alg.jac_alg)
6666

67-
fᵢ_cache = __alloc(zero(X))
68-
fᵢ₂_cache = vec(zero(X))
67+
fᵢ_cache = __alloc(zero(u0))
68+
fᵢ₂_cache = vec(zero(u0))
6969

7070
# Don't flatten this here, since we need to expand it later if needed
71-
y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
71+
y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p)
7272

7373
y = __alloc.(copy.(y₀.u))
7474
TU, ITU = constructMIRK(alg, T)
@@ -78,22 +78,22 @@ function SciMLBase.__init(
7878

7979
k_discrete = if !constraint
8080
[
81-
__maybe_allocate_diffcache(safe_similar(X, N, stage), chunksize, alg.jac_alg)
81+
__maybe_allocate_diffcache(safe_similar(u0, N, stage), chunksize, alg.jac_alg)
8282
for _ in 1:Nig
8383
]
8484
else
8585
[
86-
__maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
86+
__maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
8787
for _ in 1:Nig
8888
]
8989
end
9090
k_interp = if !constraint
91-
VectorOfArray([similar(X, N, ITU.s_star - stage) for _ in 1:Nig])
91+
VectorOfArray([similar(u0, N, ITU.s_star - stage) for _ in 1:Nig])
9292
else
93-
VectorOfArray([similar(X, L_f_prototype, ITU.s_star - stage) for _ in 1:Nig])
93+
VectorOfArray([similar(u0, L_f_prototype, ITU.s_star - stage) for _ in 1:Nig])
9494
end
9595

96-
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, X)
96+
bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, u0)
9797

9898
residual = if iip
9999
if !constraint
@@ -135,28 +135,28 @@ function SciMLBase.__init(
135135
errors = if !constraint
136136
VectorOfArray(
137137
[
138-
similar(X, ifelse(adaptive, N, 0))
138+
similar(u0, ifelse(adaptive, N, 0))
139139
for _ in 1:ifelse(use_both, 2Nig, Nig)
140140
]
141141
)
142142
else
143143
VectorOfArray(
144144
[
145-
similar(X, ifelse(adaptive, L_f_prototype, 0))
145+
similar(u0, ifelse(adaptive, L_f_prototype, 0))
146146
for _ in 1:ifelse(use_both, 2Nig, Nig)
147147
]
148148
)
149149
end
150150
new_stages = if !constraint
151-
VectorOfArray([similar(X, N) for _ in 1:Nig])
151+
VectorOfArray([similar(u0, N) for _ in 1:Nig])
152152
else
153-
VectorOfArray([similar(X, L_f_prototype) for _ in 1:Nig])
153+
VectorOfArray([similar(u0, L_f_prototype) for _ in 1:Nig])
154154
end
155155

156156
# Transform the functions to handle non-vector inputs
157157
bcresid_prototype = __vec(bcresid_prototype)
158158
f,
159-
bc = if X isa AbstractVector
159+
bc = if u0 isa AbstractVector
160160
f_wrapped = prob.f
161161
bc_wrapped = prob.f.bc
162162
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
@@ -194,42 +194,42 @@ function SciMLBase.__init(
194194
end
195195
f_wrapped, bc_wrapped
196196
elseif iip
197-
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
197+
vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
198198
vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
199-
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(X))
199+
@closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
200200
else
201201
(
202202
@closure(
203203
(
204204
r, u,
205205
p,
206-
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(X))
206+
) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(u0))
207207
),
208208
@closure(
209209
(
210210
r, u, p,
211-
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(X))
211+
) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(u0))
212212
),
213213
)
214214
end
215215
vecf!, vecbc!
216216
else
217-
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
217+
vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
218218
vecbc = if !(prob.problem_type isa TwoPointBVProblem)
219-
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
219+
@closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
220220
else
221221
(
222-
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(X))),
223-
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(X))),
222+
@closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
223+
@closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
224224
)
225225
end
226226
vecf, vecbc
227227
end
228228

229-
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
229+
prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = u0) : prob
230230

231231
return MIRKCache{iip, T, use_both, typeof(diffcache), tune_parameters}(
232-
alg_order(alg), stage, N, size(X), f, bc, prob_, prob.problem_type, prob.p, alg,
232+
alg_order(alg), stage, N, size(u0), f, bc, prob_, prob.problem_type, prob.p, alg,
233233
TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete, k_interp, y,
234234
y₀, residual, fᵢ_cache, fᵢ₂_cache, errors, new_stages, resid₁_size, prob.singular_term
235235
, nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, tune_parameters, kwargs...), verbose_spec

0 commit comments

Comments
 (0)