Skip to content

Commit 6e37525

Browse files
ErikQQYErikQQY
authored
Merge pull request #473 from SebastianM-C/smc/fix-extract-lcons-length
Fix `lcons`/`ucons` length in `__extract_lcons_ucons` for `optimize` path
2 parents e0cd15b + 0554143 commit 6e37525

14 files changed

Lines changed: 119 additions & 76 deletions

File tree

‎docs/Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,5 +48,5 @@ LinearSolve = "2.36.2, 3"
4848
OptimizationIpopt = "0.2.0, 1.0"
4949
OrdinaryDiffEq = "7"
5050
Plots = "1"
51-
SciMLBase = "2.60.0, 3.3"
51+
SciMLBase = "2.60.0, 3"
5252
SimpleBoundaryValueDiffEq = "1.1.0"

‎lib/BoundaryValueDiffEqCore/src/internal_problems.jl‎

Lines changed: 22 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -36,46 +36,31 @@
3636
return cost_fun
3737
end
3838

39-
@inline function __extract_lcons_ucons(
40-
prob::AbstractBVProblem, ::Type{T}, M, N, bcresid_prototype, f_prototype
41-
) where {T}
42-
L_f_prototype = length(f_prototype)
43-
L_bcresid_prototype = length(bcresid_prototype)
44-
lcons = if isnothing(prob.lcons)
45-
zeros(T, L_bcresid_prototype + (N - 1) * L_f_prototype)
46-
else
47-
lcons_length = length(prob.lcons)
48-
vcat(prob.lcons, zeros(T, N * M - lcons_length))
49-
end
50-
ucons = if isnothing(prob.ucons)
51-
zeros(T, L_bcresid_prototype + (N - 1) * L_f_prototype)
52-
else
53-
ucons_length = length(prob.ucons)
54-
vcat(prob.ucons, zeros(T, N * M - ucons_length))
55-
end
56-
return lcons, ucons
57-
end
39+
"""
40+
__extract_lcons_ucons(prob, T, constraint_length)
5841
42+
Build the `lcons` / `ucons` vectors for the equality-constrained optimization
43+
problem constructed from a BVP. Length must match the actual constraint vector
44+
(`resid_prototype`) passed to the solver — not a reconstruction from
45+
`(M, N, bcresid_prototype, f_prototype)`, which has a different formula per
46+
discretization (MIRK: `L_bc + M*(N-1)`; FIRK Expanded: `L_bc + M*(N-1)*(stage+1)`;
47+
MIRKN: `L_bc + M*2*(N-1)`; shooting: `L_bc`). Callers pass `length(resid_prototype)`
48+
directly.
49+
"""
5950
@inline function __extract_lcons_ucons(
60-
prob::AbstractBVProblem, ::Type{T}, M, N, bcresid_prototype, ::Nothing
51+
prob::AbstractBVProblem, ::Type{T}, constraint_length::Int
6152
) where {T}
62-
lcons = zeros(T, N * M)
63-
ucons = zeros(T, N * M)
64-
return lcons, ucons
65-
end
66-
67-
@inline function __extract_lcons_ucons(prob::AbstractBVProblem, ::Type{T}, M, N) where {T}
6853
lcons = if isnothing(prob.lcons)
69-
zeros(T, N * M)
54+
zeros(T, constraint_length)
7055
else
7156
lcons_length = length(prob.lcons)
72-
vcat(prob.lcons, zeros(T, N * M - lcons_length))
57+
vcat(prob.lcons, zeros(T, constraint_length - lcons_length))
7358
end
7459
ucons = if isnothing(prob.ucons)
75-
zeros(T, N * M)
60+
zeros(T, constraint_length)
7661
else
7762
ucons_length = length(prob.ucons)
78-
vcat(prob.ucons, zeros(T, N * M - ucons_length))
63+
vcat(prob.ucons, zeros(T, constraint_length - ucons_length))
7964
end
8065
return lcons, ucons
8166
end
@@ -129,7 +114,7 @@ function __construct_internal_problem(
129114
cons_j = jac,
130115
cons_jac_prototype = sparse(jac_prototype)
131116
)
132-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N, bcresid_prototype, f_prototype)
117+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
133118
lb, ub = __extract_lb_ub(prob, T, M, N)
134119

135120
return __internal_optimization_problem(
@@ -161,7 +146,7 @@ function __construct_internal_problem(
161146
cons_j = jac,
162147
cons_jac_prototype = sparse(jac_prototype)
163148
)
164-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N, bcresid_prototype, f_prototype)
149+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
165150
lb, ub = __extract_lb_ub(prob, T, M, N)
166151

167152
return __internal_optimization_problem(
@@ -194,7 +179,7 @@ function __construct_internal_problem(
194179
cons_j = jac,
195180
cons_jac_prototype = sparse(jac_prototype)
196181
)
197-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N)
182+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
198183
lb, ub = __extract_lb_ub(prob, T, M, N)
199184

200185
return __internal_optimization_problem(
@@ -237,7 +222,7 @@ function __construct_internal_problem(
237222
cons_j = jac,
238223
cons_jac_prototype = sparse(jac_prototype)
239224
)
240-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N, bcresid_prototype, f_prototype)
225+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
241226
lb, ub = __extract_lb_ub(prob, T, M, N)
242227

243228
return __internal_optimization_problem(
@@ -268,7 +253,7 @@ function __construct_internal_problem(
268253
cons_j = jac,
269254
cons_jac_prototype = sparse(jac_prototype)
270255
)
271-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N, bcresid_prototype, f_prototype)
256+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
272257
lb, ub = __extract_lb_ub(prob, T, M, N)
273258

274259
return __internal_optimization_problem(
@@ -301,7 +286,7 @@ function __construct_internal_problem(
301286
cons_j = jac,
302287
cons_jac_prototype = sparse(jac_prototype)
303288
)
304-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N)
289+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
305290
lb, ub = __extract_lb_ub(prob, T, M, N)
306291

307292
return __internal_optimization_problem(
@@ -332,7 +317,7 @@ function __construct_internal_problem(
332317
cons_j = jac,
333318
cons_jac_prototype = sparse(jac_prototype)
334319
)
335-
lcons, ucons = __extract_lcons_ucons(prob, T, M, N)
320+
lcons, ucons = __extract_lcons_ucons(prob, T, length(resid_prototype))
336321
lb, ub = __extract_lb_ub(prob, T, M, N)
337322

338323
return __internal_optimization_problem(

‎lib/BoundaryValueDiffEqCore/test/util_tests.jl‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
@testset "__extract_lcons_ucons length" begin
2+
# Regression test: the function must return vectors matching the actual
3+
# constraint vector length (= length(resid_prototype)), not a reconstruction
4+
# from (M, N, ...) which was wrong for several solvers.
5+
using BoundaryValueDiffEqCore: __extract_lcons_ucons
6+
using SciMLBase: BVProblem
7+
8+
f!(du, u, p, t) = (du[1] = u[2]; du[2] = -u[1])
9+
bc!(res, u, p, t) = (res[1] = u(0.0)[1]; res[2] = u(1.0)[1])
10+
11+
# Fallback path (isnothing(prob.lcons)): both vectors have length == constraint_length
12+
prob = BVProblem(f!, bc!, [0.0, 0.0], (0.0, 1.0); bcresid_prototype = zeros(2))
13+
lc, uc = __extract_lcons_ucons(prob, Float64, 42)
14+
@test length(lc) == 42
15+
@test length(uc) == 42
16+
@test all(iszero, lc)
17+
@test all(iszero, uc)
18+
19+
# User-provided lcons/ucons: values preserved, padded with zeros to constraint_length
20+
prob2 = BVProblem(
21+
f!, bc!, [0.0, 0.0], (0.0, 1.0);
22+
bcresid_prototype = zeros(2),
23+
lcons = [-1.0, -2.0], ucons = [1.0, 2.0]
24+
)
25+
lc2, uc2 = __extract_lcons_ucons(prob2, Float64, 10)
26+
@test length(lc2) == 10
27+
@test length(uc2) == 10
28+
@test lc2[1:2] == [-1.0, -2.0]
29+
@test uc2[1:2] == [1.0, 2.0]
30+
@test all(iszero, lc2[3:end])
31+
@test all(iszero, uc2[3:end])
32+
end

‎lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,7 @@ Generate new mesh based on the defect.
260260

261261
info = ReturnCode.Success
262262

263-
ŝ = [maximum(abs, d) for d in defect] # Broadcasting breaks GPU Compilation
263+
ŝ = [maximum(abs, d) for d in defect.u] # Broadcasting breaks GPU Compilation
264264
ŝ .= (ŝ ./ abstol) .^ (T(1) / (order + 1))
265265
r₁ = maximum(ŝ)
266266
r₂ = sum(ŝ)

‎lib/BoundaryValueDiffEqFIRK/src/firk.jl‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -444,19 +444,19 @@ function __expand_cache!(cache::FIRKCacheExpand)
444444
Nₙ = length(cache.mesh)
445445
__resize!(cache.k_discrete, Nₙ - 1, cache.M, cache.TU)
446446
__resize!(cache.y, Nₙ, cache.M, cache.TU)
447-
__resize!(cache.y₀, Nₙ, cache.M, cache.TU)
447+
__resize!(cache.y₀.u, Nₙ, cache.M, cache.TU)
448448
__resize!(cache.residual, Nₙ, cache.M, cache.TU)
449-
__resize!(cache.defect, Nₙ - 1, cache.M)
449+
__resize!(cache.defect.u, Nₙ - 1, cache.M)
450450
return cache
451451
end
452452

453453
function __expand_cache!(cache::FIRKCacheNested)
454454
Nₙ = length(cache.mesh)
455455
__resize!(cache.k_discrete, Nₙ - 1, cache.M)
456456
__resize!(cache.y, Nₙ, cache.M)
457-
__resize!(cache.y₀, Nₙ, cache.M)
457+
__resize!(cache.y₀.u, Nₙ, cache.M)
458458
__resize!(cache.residual, Nₙ, cache.M)
459-
__resize!(cache.defect, Nₙ - 1, cache.M)
459+
__resize!(cache.defect.u, Nₙ - 1, cache.M)
460460
return cache
461461
end
462462

@@ -552,7 +552,7 @@ function SciMLBase.solve!(
552552
end
553553

554554
function __perform_firk_iteration(cache::Union{FIRKCacheExpand, FIRKCacheNested}, abstol, adaptive::Bool)
555-
nlprob = __construct_problem(cache, vec(cache.y₀), copy(cache.y₀))
555+
nlprob = __construct_problem(cache, copy(vec(cache.y₀)), copy(cache.y₀))
556556
solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
557557
kwargs = __concrete_kwargs(
558558
cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,
@@ -579,7 +579,7 @@ function __perform_firk_iteration(cache::Union{FIRKCacheExpand, FIRKCacheNested}
579579
# We construct a new mesh to equidistribute the defect
580580
mesh, mesh_dt, _, info = mesh_selector!(cache)
581581
if info == ReturnCode.Success
582-
__resize!(cache.y₀, length(cache.mesh), cache.M, cache.TU)
582+
__resize!(cache.y₀.u, length(cache.mesh), cache.M, cache.TU)
583583
for (i, m) in enumerate(cache.mesh)
584584
interp_eval!(cache.y₀.u[i], cache, m, mesh, mesh_dt)
585585
end

‎lib/BoundaryValueDiffEqFIRK/src/utils.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
function BoundaryValueDiffEqCore.__resize!(x::AbstractVector{<:AbstractArray}, n, _, TU::FIRKTableau{false})
1+
function BoundaryValueDiffEqCore.__resize!(x::AbstractVector{<:AbstractArray}, n, _, TU) #::FIRKTableau{false}) # TODO: remove the TU argument and just use s = length(TU.c)
22
(; s) = TU
33
N = (n - 1) * (s + 1) + 1 - length(x)
44
N == 0 && return x

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
residual[2] = u(1.0)[1]
1212
end
1313

14-
prob_func(prob, ctx, repeat) = remake(prob, p = [rand()])
14+
prob_func(prob, ctx) = remake(prob, p = [rand()])
1515

1616
u0 = [0.0, 1.0]
1717
tspan = (0, pi / 2)

‎lib/BoundaryValueDiffEqFIRK/test/nested/ensemble_tests.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
residual[2] = u(1.0)[1]
1212
end
1313

14-
prob_func(prob, ctx, repeat) = remake(prob, p = [rand()])
14+
prob_func(prob, ctx) = remake(prob, p = [rand()])
1515

1616
u0 = [0.0, 1.0]
1717
tspan = (0, pi / 2)

‎lib/BoundaryValueDiffEqMIRK/src/adaptivity.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ Generate new mesh based on the defect or the global error.
3636

3737
info = ReturnCode.Success
3838

39-
ŝ = [maximum(abs, d) for d in errors] # Broadcasting breaks GPU Compilation
39+
ŝ = [maximum(abs, d) for d in errors.u] # Broadcasting breaks GPU Compilation
4040
ŝ .= (ŝ ./ abstol) .^ (T(1) / (order + 1))
4141
r₁ = maximum(ŝ)
4242
r₂ = sum(ŝ)
@@ -90,7 +90,7 @@ end
9090

9191
info = ReturnCode.Success
9292

93-
ŝ = [maximum(abs, d) for d in errors]
93+
ŝ = [maximum(abs, d) for d in errors.u]
9494
ŝ .= (ŝ ./ abstol) .^ (T(1) / order)
9595
r₁ = maximum(ŝ)
9696
r₂ = sum(ŝ)
@@ -561,7 +561,7 @@ end
561561
errors.u,
562562
[
563563
ifelse(maximum(abs.(err.u[i])) >= maximum(abs.(err.u[i + 1])), err.u[i], err.u[i + 1])
564-
for i in 1:(length(err) - 1)
564+
for i in 1:(length(err.u) - 1)
565565
]
566566
)
567567
end

‎lib/BoundaryValueDiffEqMIRK/src/mirk.jl‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,6 @@ function SciMLBase.__init(
168168
if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
169169
tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
170170
l_parameters = length(tunable_part)
171-
base_f = f_wrapped
172171
f_wrapped = @closure (
173172
du,
174173
u,
@@ -177,22 +176,21 @@ function SciMLBase.__init(
177176
) -> begin
178177
@inbounds @views begin
179178
_p = repack(u[(end - l_parameters + 1):end])
180-
base_f(du, u, _p, t)
179+
prob.f(du, u, _p, t)
181180
fill!(du[(end - l_parameters + 1):end], zero(eltype(du)))
182181
end
183182
return nothing
184183
end
185184
elseif tune_parameters
186185
l_parameters = length(prob.p)
187-
base_f = f_wrapped
188186
f_wrapped = @closure (
189187
du,
190188
u,
191189
p,
192190
t,
193191
) -> begin
194192
@inbounds @views begin
195-
base_f(du, u, u[(end - l_parameters + 1):end], t)
193+
prob.f(du, u, u[(end - l_parameters + 1):end], t)
196194
fill!(du[(end - l_parameters + 1):end], zero(eltype(du)))
197195
end
198196
return nothing
@@ -251,12 +249,12 @@ match the length of the new mesh.
251249
function __expand_cache!(cache::MIRKCache{iip, T, use_both}) where {iip, T, use_both}
252250
Nₙ = length(cache.mesh)
253251
__resize!(cache.k_discrete, Nₙ - 1, cache.M)
254-
__resize!(cache.k_interp, Nₙ - 1, cache.M)
252+
__resize!(cache.k_interp.u, Nₙ - 1, cache.M)
255253
__resize!(cache.y, Nₙ, cache.M)
256-
__resize!(cache.y₀, Nₙ, cache.M)
254+
__resize!(cache.y₀.u, Nₙ, cache.M)
257255
__resize!(cache.residual, Nₙ, cache.M)
258-
__resize!(cache.errors, ifelse(use_both, 2 * (Nₙ - 1), (Nₙ - 1)), cache.M)
259-
__resize!(cache.new_stages, Nₙ - 1, cache.M)
256+
__resize!(cache.errors.u, ifelse(use_both, 2 * (Nₙ - 1), (Nₙ - 1)), cache.M)
257+
__resize!(cache.new_stages.u, Nₙ - 1, cache.M)
260258
return cache
261259
end
262260

@@ -269,6 +267,7 @@ function SciMLBase.solve!(
269267
(abstol, adaptive, controller, _), _ = __split_kwargs(; cache.kwargs...)
270268
info::ReturnCode.T = ReturnCode.Success
271269
prob = cache.prob
270+
length_u = cache.in_size
272271

273272
# We do the first iteration outside the loop to preserve type-stability of the
274273
# `original` field of the solution
@@ -308,7 +307,7 @@ function SciMLBase.solve!(
308307
end
309308

310309
function __perform_mirk_iteration(cache::MIRKCache, abstol, adaptive::Bool, controller::AbstractErrorControl)
311-
nlprob = __construct_problem(cache, vec(cache.y₀), copy(cache.y₀))
310+
nlprob = __construct_problem(cache, copy(vec(cache.y₀)), copy(cache.y₀))
312311
solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
313312
kwargs = __concrete_kwargs(
314313
cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,
@@ -337,7 +336,7 @@ function __perform_mirk_iteration(cache::MIRKCache, abstol, adaptive::Bool, cont
337336
mesh, mesh_dt, _, info = mesh_selector!(cache, controller)
338337
if info == ReturnCode.Success
339338
(length(mesh) < length(cache.mesh)) &&
340-
__resize!(cache.y₀, length(cache.mesh), cache.M)
339+
__resize!(cache.y₀.u, length(cache.mesh), cache.M)
341340
for (i, m) in enumerate(cache.mesh)
342341
interp_eval!(cache.y₀.u[i], cache, m, mesh, mesh_dt)
343342
end

0 commit comments

Comments
 (0)