Skip to content

Commit c2e7dc9

Browse files
Merge pull request #122 from avik-pal/ap/nls
Overcontrained / Underconstrained BVPs
2 parents e100313 + 347f22a commit c2e7dc9

14 files changed

+395
-45
lines changed

Project.toml

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "BoundaryValueDiffEq"
22
uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
3-
version = "5.3.0"
3+
version = "5.4.0"
44

55
[deps]
66
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -11,6 +11,7 @@ ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
1111
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
1212
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
1313
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
14+
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
1415
NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
1516
PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
1617
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -41,8 +42,9 @@ BandedMatrices = "1"
4142
ConcreteStructs = "0.2"
4243
DiffEqBase = "6.135"
4344
ForwardDiff = "0.10"
44-
LinearAlgebra = "1.6"
45-
NonlinearSolve = "2.5"
45+
LinearAlgebra = "1.9"
46+
LinearSolve = "2"
47+
NonlinearSolve = "2.6.1"
4648
ODEInterface = "0.5"
4749
OrdinaryDiffEq = "6"
4850
PreallocationTools = "0.4"
@@ -52,7 +54,7 @@ RecursiveArrayTools = "2.38.10"
5254
Reexport = "0.2, 1.0"
5355
SciMLBase = "2.5"
5456
Setfield = "1"
55-
SparseArrays = "1.6"
57+
SparseArrays = "1.9"
5658
SparseDiffTools = "2.9"
5759
TruncatedStacktraces = "1"
5860
UnPack = "1"
@@ -61,6 +63,7 @@ julia = "1.9"
6163
[extras]
6264
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
6365
DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
66+
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
6467
ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
6568
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
6669
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -69,4 +72,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
6972
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
7073

7174
[targets]
72-
test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua"]
75+
test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua", "LinearSolve"]

README.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,9 @@ Precompilation can be controlled via `Preferences.jl`
4848
- `PrecompileMIRK` -- Precompile the MIRK2 - MIRK6 algorithms (default: `true`).
4949
- `PrecompileShooting` -- Precompile the single shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
5050
- `PrecompileMultipleShooting` -- Precompile the multiple shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
51+
- `PrecompileMIRKNLLS` -- Precompile the MIRK2 - MIRK6 algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above).
52+
- `PrecompileShootingNLLS` -- Precompile the single shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.
53+
- `PrecompileMultipleShootingNLLS` -- Precompile the multiple shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.
5154

5255
To set these preferences before loading the package, do the following (replacing `PrecompileShooting` with the preference you want to set, or pass in multiple pairs to set them together):
5356

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,81 @@ end
6363
solve(prob, alg)
6464
end
6565
end
66+
67+
function f1_nlls!(du, u, p, t)
68+
du[1] = u[2]
69+
du[2] = -u[1]
70+
end
71+
72+
f1_nlls(u, p, t) = [u[2], -u[1]]
73+
74+
function bc1_nlls!(resid, sol, p, t)
75+
solₜ₁ = sol(0.0)
76+
solₜ₂ = sol(100.0)
77+
resid[1] = solₜ₁[1]
78+
resid[2] = solₜ₂[1] - 1
79+
resid[3] = solₜ₂[2] + 1.729109
80+
return nothing
81+
end
82+
bc1_nlls(sol, p, t) = [sol(0.0)[1], sol(100.0)[1] - 1, sol(100.0)[2] + 1.729109]
83+
84+
bc1_nlls_a!(resid, ua, p) = (resid[1] = ua[1])
85+
bc1_nlls_b!(resid, ub, p) = (resid[1] = ub[1] - 1; resid[2] = ub[2] + 1.729109)
86+
87+
bc1_nlls_a(ua, p) = [ua[1]]
88+
bc1_nlls_b(ub, p) = [ub[1] - 1, ub[2] + 1.729109]
89+
90+
tspan = (0.0, 100.0)
91+
u0 = [0.0, 1.0]
92+
bcresid_prototype1 = Array{Float64}(undef, 3)
93+
bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))
94+
95+
probs = [
96+
BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
97+
u0, tspan),
98+
BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
99+
u0, tspan),
100+
TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
101+
bcresid_prototype = bcresid_prototype2),
102+
TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
103+
bcresid_prototype = bcresid_prototype2),
104+
]
105+
106+
algs = []
107+
108+
if @load_preference("PrecompileShootingNLLS", VERSIONv"1.10-")
109+
append!(algs,
110+
[
111+
Shooting(Tsit5();
112+
nlsolve = LevenbergMarquardt(;
113+
autodiff = AutoForwardDiff(chunksize = 2))),
114+
Shooting(Tsit5();
115+
nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2))),
116+
])
117+
end
118+
119+
if @load_preference("PrecompileMultipleShootingNLLS", VERSIONv"1.10-")
120+
append!(algs,
121+
[
122+
MultipleShooting(10, Tsit5();
123+
nlsolve = LevenbergMarquardt(;
124+
autodiff = AutoForwardDiff(chunksize = 2)),
125+
jac_alg = BVPJacobianAlgorithm(;
126+
bc_diffmode = AutoForwardDiff(; chunksize = 2),
127+
nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
128+
MultipleShooting(10, Tsit5();
129+
nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2)),
130+
jac_alg = BVPJacobianAlgorithm(;
131+
bc_diffmode = AutoForwardDiff(; chunksize = 2),
132+
nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
133+
])
134+
end
135+
136+
@compile_workload begin
137+
for prob in probs, alg in algs
138+
solve(prob, alg)
139+
end
140+
end
66141
end
67142

68143
end

src/BoundaryValueDiffEq.jl

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,68 @@ end
8989
solve(prob, alg; dt = 0.2)
9090
end
9191
end
92+
93+
function f1_nlls!(du, u, p, t)
94+
du[1] = u[2]
95+
du[2] = -u[1]
96+
end
97+
98+
f1_nlls(u, p, t) = [u[2], -u[1]]
99+
100+
function bc1_nlls!(resid, sol, p, t)
101+
solₜ₁ = sol[1]
102+
solₜ₂ = sol[end]
103+
resid[1] = solₜ₁[1]
104+
resid[2] = solₜ₂[1] - 1
105+
resid[3] = solₜ₂[2] + 1.729109
106+
return nothing
107+
end
108+
bc1_nlls(sol, p, t) = [sol[1][1], sol[end][1] - 1, sol[end][2] + 1.729109]
109+
110+
bc1_nlls_a!(resid, ua, p) = (resid[1] = ua[1])
111+
bc1_nlls_b!(resid, ub, p) = (resid[1] = ub[1] - 1; resid[2] = ub[2] + 1.729109)
112+
113+
bc1_nlls_a(ua, p) = [ua[1]]
114+
bc1_nlls_b(ub, p) = [ub[1] - 1, ub[2] + 1.729109]
115+
116+
tspan = (0.0, 100.0)
117+
u0 = [0.0, 1.0]
118+
bcresid_prototype1 = Array{Float64}(undef, 3)
119+
bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))
120+
121+
probs = [
122+
BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
123+
u0, tspan),
124+
BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
125+
u0, tspan),
126+
TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
127+
bcresid_prototype = bcresid_prototype2),
128+
TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
129+
bcresid_prototype = bcresid_prototype2),
130+
]
131+
132+
jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))
133+
134+
nlsolvers = [LevenbergMarquardt(), GaussNewton()]
135+
136+
algs = []
137+
138+
if Preferences.@load_preference("PrecompileMIRKNLLS", VERSIONv"1.10-")
139+
for nlsolve in nlsolvers
140+
append!(algs,
141+
[
142+
MIRK2(; jac_alg, nlsolve), MIRK3(; jac_alg, nlsolve),
143+
MIRK4(; jac_alg, nlsolve), MIRK5(; jac_alg, nlsolve),
144+
MIRK6(; jac_alg, nlsolve),
145+
])
146+
end
147+
end
148+
149+
@compile_workload begin
150+
for prob in probs, alg in algs
151+
solve(prob, alg; dt = 0.2)
152+
end
153+
end
92154
end
93155

94156
export Shooting, MultipleShooting

src/solve/mirk.jl

Lines changed: 28 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -274,11 +274,12 @@ function __mirk_loss_collocation(u, p, y, mesh, residual, cache)
274274
end
275275

276276
function __construct_nlproblem(cache::MIRKCache{iip}, y, loss_bc::BC, loss_collocation::C,
277-
loss::L, ::StandardBVProblem) where {iip, BC, C, L}
277+
loss::LF, ::StandardBVProblem) where {iip, BC, C, LF}
278278
@unpack nlsolve, jac_alg = cache.alg
279279
N = length(cache.mesh)
280280

281281
resid_bc = cache.bcresid_prototype
282+
L = length(resid_bc)
282283
resid_collocation = similar(y, cache.M * (N - 1))
283284

284285
loss_bcₚ = iip ? ((du, u) -> loss_bc(du, u, cache.p)) : (u -> loss_bc(u, cache.p))
@@ -302,45 +303,48 @@ function __construct_nlproblem(cache::MIRKCache{iip}, y, loss_bc::BC, loss_collo
302303
jac_prototype = vcat(init_jacobian(cache_bc), init_jacobian(cache_collocation))
303304

304305
jac = if iip
305-
(J, u, p) -> __mirk_mpoint_jacobian!(J, u, p, jac_alg.bc_diffmode,
306+
(J, u, p) -> __mirk_mpoint_jacobian!(J, u, jac_alg.bc_diffmode,
306307
jac_alg.nonbc_diffmode, cache_bc, cache_collocation, loss_bcₚ,
307-
loss_collocationₚ, resid_bc, resid_collocation, cache.M)
308+
loss_collocationₚ, resid_bc, resid_collocation, L)
308309
else
309-
(u, p) -> __mirk_mpoint_jacobian(u, p, jac_prototype, jac_alg.bc_diffmode,
310+
(u, p) -> __mirk_mpoint_jacobian(jac_prototype, u, jac_alg.bc_diffmode,
310311
jac_alg.nonbc_diffmode, cache_bc, cache_collocation, loss_bcₚ,
311-
loss_collocationₚ, cache.M)
312+
loss_collocationₚ, L)
312313
end
313314

314-
return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p)
315+
nlf = NonlinearFunction{iip}(loss; resid_prototype = vcat(resid_bc, resid_collocation),
316+
jac, jac_prototype)
317+
return (L == cache.M ? NonlinearProblem : NonlinearLeastSquaresProblem)(nlf, y, cache.p)
315318
end
316319

317-
function __mirk_mpoint_jacobian!(J, x, p, bc_diffmode, nonbc_diffmode, bc_diffcache,
320+
function __mirk_mpoint_jacobian!(J, x, bc_diffmode, nonbc_diffmode, bc_diffcache,
318321
nonbc_diffcache, loss_bc::BC, loss_collocation::C, resid_bc, resid_collocation,
319-
M::Int) where {BC, C}
320-
sparse_jacobian!(@view(J[1:M, :]), bc_diffmode, bc_diffcache, loss_bc, resid_bc, x)
321-
sparse_jacobian!(@view(J[(M + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
322+
L::Int) where {BC, C}
323+
sparse_jacobian!(@view(J[1:L, :]), bc_diffmode, bc_diffcache, loss_bc, resid_bc, x)
324+
sparse_jacobian!(@view(J[(L + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
322325
loss_collocation, resid_collocation, x)
323326
return nothing
324327
end
325328

326-
function __mirk_mpoint_jacobian(x, p, J, bc_diffmode, nonbc_diffmode, bc_diffcache,
327-
nonbc_diffcache, loss_bc::BC, loss_collocation::C, M::Int) where {BC, C}
328-
sparse_jacobian!(@view(J[1:M, :]), bc_diffmode, bc_diffcache, loss_bc, x)
329-
sparse_jacobian!(@view(J[(M + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
329+
function __mirk_mpoint_jacobian(J, x, bc_diffmode, nonbc_diffmode, bc_diffcache,
330+
nonbc_diffcache, loss_bc::BC, loss_collocation::C, L::Int) where {BC, C}
331+
sparse_jacobian!(@view(J[1:L, :]), bc_diffmode, bc_diffcache, loss_bc, x)
332+
sparse_jacobian!(@view(J[(L + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
330333
loss_collocation, x)
331334
return J
332335
end
333336

334337
function __construct_nlproblem(cache::MIRKCache{iip}, y, loss_bc::BC, loss_collocation::C,
335-
loss::L, ::TwoPointBVProblem) where {iip, BC, C, L}
338+
loss::LF, ::TwoPointBVProblem) where {iip, BC, C, LF}
336339
@unpack nlsolve, jac_alg = cache.alg
337340
N = length(cache.mesh)
338341

339342
lossₚ = iip ? ((du, u) -> loss(du, u, cache.p)) : (u -> loss(u, cache.p))
340343

341-
resid = vcat(cache.bcresid_prototype[1:prod(cache.resid_size[1])],
344+
resid = vcat(@view(cache.bcresid_prototype[1:prod(cache.resid_size[1])]),
342345
similar(y, cache.M * (N - 1)),
343-
cache.bcresid_prototype[(prod(cache.resid_size[1]) + 1):end])
346+
@view(cache.bcresid_prototype[(prod(cache.resid_size[1]) + 1):end]))
347+
L = length(cache.bcresid_prototype)
344348

345349
sd = if jac_alg.diffmode isa AbstractSparseADType
346350
__sparsity_detection_alg(__generate_sparse_jacobian_prototype(cache,
@@ -354,22 +358,24 @@ function __construct_nlproblem(cache::MIRKCache{iip}, y, loss_bc::BC, loss_collo
354358
jac_prototype = init_jacobian(diffcache)
355359

356360
jac = if iip
357-
(J, u, p) -> __mirk_2point_jacobian!(J, u, p, jac_alg.diffmode, diffcache, lossₚ,
361+
(J, u, p) -> __mirk_2point_jacobian!(J, u, jac_alg.diffmode, diffcache, lossₚ,
358362
resid)
359363
else
360-
(u, p) -> __mirk_2point_jacobian(u, p, jac_prototype, jac_alg.diffmode, diffcache,
364+
(u, p) -> __mirk_2point_jacobian(u, jac_prototype, jac_alg.diffmode, diffcache,
361365
lossₚ)
362366
end
363367

364-
return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p)
368+
nlf = NonlinearFunction{iip}(loss; resid_prototype = copy(resid), jac, jac_prototype)
369+
370+
return (L == cache.M ? NonlinearProblem : NonlinearLeastSquaresProblem)(nlf, y, cache.p)
365371
end
366372

367-
function __mirk_2point_jacobian!(J, x, p, diffmode, diffcache, loss_fn::L, resid) where {L}
373+
function __mirk_2point_jacobian!(J, x, diffmode, diffcache, loss_fn::L, resid) where {L}
368374
sparse_jacobian!(J, diffmode, diffcache, loss_fn, resid, x)
369375
return J
370376
end
371377

372-
function __mirk_2point_jacobian(x, p, J, diffmode, diffcache, loss_fn::L) where {L}
378+
function __mirk_2point_jacobian(x, J, diffmode, diffcache, loss_fn::L) where {L}
373379
sparse_jacobian!(J, diffmode, diffcache, loss_fn, x)
374380
return J
375381
end

0 commit comments

Comments
 (0)