Skip to content

Commit

Permalink
Merge pull request #122 from avik-pal/ap/nls
Browse files Browse the repository at this point in the history 
Overcontrained / Underconstrained BVPs
  • Loading branch information
@ChrisRackauckas
ChrisRackauckas authored Nov 4, 2023
2 parents e100313 + 347f22a commit c2e7dc9
Show file tree
Hide file tree
Showing 14 changed files with 395 additions and 45 deletions.
13 changes: 8 additions & 5 deletions Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "BoundaryValueDiffEq"
uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
version = "5.3.0"
version = "5.4.0"

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

[targets]
test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua"]
test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua", "LinearSolve"]
3 changes: 3 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,9 @@ Precompilation can be controlled via `Preferences.jl`
- `PrecompileMIRK` -- Precompile the MIRK2 - MIRK6 algorithms (default: `true`).
- `PrecompileShooting` -- Precompile the single shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
- `PrecompileMultipleShooting` -- Precompile the multiple shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
- `PrecompileMIRKNLLS` -- Precompile the MIRK2 - MIRK6 algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above).
- `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.
- `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.

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):

Expand Down
75 changes: 75 additions & 0 deletions ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,81 @@ end
solve(prob, alg)
end
end

function f1_nlls!(du, u, p, t)
du[1] = u[2]
du[2] = -u[1]
end

f1_nlls(u, p, t) = [u[2], -u[1]]

function bc1_nlls!(resid, sol, p, t)
solₜ₁ = sol(0.0)
solₜ₂ = sol(100.0)
resid[1] = solₜ₁[1]
resid[2] = solₜ₂[1] - 1
resid[3] = solₜ₂[2] + 1.729109
return nothing
end
bc1_nlls(sol, p, t) = [sol(0.0)[1], sol(100.0)[1] - 1, sol(100.0)[2] + 1.729109]

bc1_nlls_a!(resid, ua, p) = (resid[1] = ua[1])
bc1_nlls_b!(resid, ub, p) = (resid[1] = ub[1] - 1; resid[2] = ub[2] + 1.729109)

bc1_nlls_a(ua, p) = [ua[1]]
bc1_nlls_b(ub, p) = [ub[1] - 1, ub[2] + 1.729109]

tspan = (0.0, 100.0)
u0 = [0.0, 1.0]
bcresid_prototype1 = Array{Float64}(undef, 3)
bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))

probs = [
BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
u0, tspan),
BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
u0, tspan),
TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
bcresid_prototype = bcresid_prototype2),
TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
bcresid_prototype = bcresid_prototype2),
]

algs = []

if @load_preference("PrecompileShootingNLLS", VERSION≥v"1.10-")
append!(algs,
[
Shooting(Tsit5();
nlsolve = LevenbergMarquardt(;
autodiff = AutoForwardDiff(chunksize = 2))),
Shooting(Tsit5();
nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2))),
])
end

if @load_preference("PrecompileMultipleShootingNLLS", VERSION≥v"1.10-")
append!(algs,
[
MultipleShooting(10, Tsit5();
nlsolve = LevenbergMarquardt(;
autodiff = AutoForwardDiff(chunksize = 2)),
jac_alg = BVPJacobianAlgorithm(;
bc_diffmode = AutoForwardDiff(; chunksize = 2),
nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
MultipleShooting(10, Tsit5();
nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2)),
jac_alg = BVPJacobianAlgorithm(;
bc_diffmode = AutoForwardDiff(; chunksize = 2),
nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
])
end

@compile_workload begin
for prob in probs, alg in algs
solve(prob, alg)
end
end
end

end
62 changes: 62 additions & 0 deletions src/BoundaryValueDiffEq.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,68 @@ end
solve(prob, alg; dt = 0.2)
end
end

function f1_nlls!(du, u, p, t)
du[1] = u[2]
du[2] = -u[1]
end

f1_nlls(u, p, t) = [u[2], -u[1]]

function bc1_nlls!(resid, sol, p, t)
solₜ₁ = sol[1]
solₜ₂ = sol[end]
resid[1] = solₜ₁[1]
resid[2] = solₜ₂[1] - 1
resid[3] = solₜ₂[2] + 1.729109
return nothing
end
bc1_nlls(sol, p, t) = [sol[1][1], sol[end][1] - 1, sol[end][2] + 1.729109]

bc1_nlls_a!(resid, ua, p) = (resid[1] = ua[1])
bc1_nlls_b!(resid, ub, p) = (resid[1] = ub[1] - 1; resid[2] = ub[2] + 1.729109)

bc1_nlls_a(ua, p) = [ua[1]]
bc1_nlls_b(ub, p) = [ub[1] - 1, ub[2] + 1.729109]

tspan = (0.0, 100.0)
u0 = [0.0, 1.0]
bcresid_prototype1 = Array{Float64}(undef, 3)
bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))

probs = [
BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
u0, tspan),
BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
u0, tspan),
TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
bcresid_prototype = bcresid_prototype2),
TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
bcresid_prototype = bcresid_prototype2),
]

jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))

nlsolvers = [LevenbergMarquardt(), GaussNewton()]

algs = []

if Preferences.@load_preference("PrecompileMIRKNLLS", VERSION≥v"1.10-")
for nlsolve in nlsolvers
append!(algs,
[
MIRK2(; jac_alg, nlsolve), MIRK3(; jac_alg, nlsolve),
MIRK4(; jac_alg, nlsolve), MIRK5(; jac_alg, nlsolve),
MIRK6(; jac_alg, nlsolve),
])
end
end

@compile_workload begin
for prob in probs, alg in algs
solve(prob, alg; dt = 0.2)
end
end
end

export Shooting, MultipleShooting
Expand Down
50 changes: 28 additions & 22 deletions src/solve/mirk.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -274,11 +274,12 @@ function __mirk_loss_collocation(u, p, y, mesh, residual, cache)
end

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

resid_bc = cache.bcresid_prototype
L = length(resid_bc)
resid_collocation = similar(y, cache.M * (N - 1))

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

jac = if iip
(J, u, p) -> __mirk_mpoint_jacobian!(J, u, p, jac_alg.bc_diffmode,
(J, u, p) -> __mirk_mpoint_jacobian!(J, u, jac_alg.bc_diffmode,
jac_alg.nonbc_diffmode, cache_bc, cache_collocation, loss_bcₚ,
loss_collocationₚ, resid_bc, resid_collocation, cache.M)
loss_collocationₚ, resid_bc, resid_collocation, L)
else
(u, p) -> __mirk_mpoint_jacobian(u, p, jac_prototype, jac_alg.bc_diffmode,
(u, p) -> __mirk_mpoint_jacobian(jac_prototype, u, jac_alg.bc_diffmode,
jac_alg.nonbc_diffmode, cache_bc, cache_collocation, loss_bcₚ,
loss_collocationₚ, cache.M)
loss_collocationₚ, L)
end

return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p)
nlf = NonlinearFunction{iip}(loss; resid_prototype = vcat(resid_bc, resid_collocation),
jac, jac_prototype)
return (L == cache.M ? NonlinearProblem : NonlinearLeastSquaresProblem)(nlf, y, cache.p)
end

function __mirk_mpoint_jacobian!(J, x, p, bc_diffmode, nonbc_diffmode, bc_diffcache,
function __mirk_mpoint_jacobian!(J, x, bc_diffmode, nonbc_diffmode, bc_diffcache,
nonbc_diffcache, loss_bc::BC, loss_collocation::C, resid_bc, resid_collocation,
M::Int) where {BC, C}
sparse_jacobian!(@view(J[1:M, :]), bc_diffmode, bc_diffcache, loss_bc, resid_bc, x)
sparse_jacobian!(@view(J[(M + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
L::Int) where {BC, C}
sparse_jacobian!(@view(J[1:L, :]), bc_diffmode, bc_diffcache, loss_bc, resid_bc, x)
sparse_jacobian!(@view(J[(L + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
loss_collocation, resid_collocation, x)
return nothing
end

function __mirk_mpoint_jacobian(x, p, J, bc_diffmode, nonbc_diffmode, bc_diffcache,
nonbc_diffcache, loss_bc::BC, loss_collocation::C, M::Int) where {BC, C}
sparse_jacobian!(@view(J[1:M, :]), bc_diffmode, bc_diffcache, loss_bc, x)
sparse_jacobian!(@view(J[(M + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
function __mirk_mpoint_jacobian(J, x, bc_diffmode, nonbc_diffmode, bc_diffcache,
nonbc_diffcache, loss_bc::BC, loss_collocation::C, L::Int) where {BC, C}
sparse_jacobian!(@view(J[1:L, :]), bc_diffmode, bc_diffcache, loss_bc, x)
sparse_jacobian!(@view(J[(L + 1):end, :]), nonbc_diffmode, nonbc_diffcache,
loss_collocation, x)
return J
end

function __construct_nlproblem(cache::MIRKCache{iip}, y, loss_bc::BC, loss_collocation::C,
loss::L, ::TwoPointBVProblem) where {iip, BC, C, L}
loss::LF, ::TwoPointBVProblem) where {iip, BC, C, LF}
@unpack nlsolve, jac_alg = cache.alg
N = length(cache.mesh)

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

resid = vcat(cache.bcresid_prototype[1:prod(cache.resid_size[1])],
resid = vcat(@view(cache.bcresid_prototype[1:prod(cache.resid_size[1])]),
similar(y, cache.M * (N - 1)),
cache.bcresid_prototype[(prod(cache.resid_size[1]) + 1):end])
@view(cache.bcresid_prototype[(prod(cache.resid_size[1]) + 1):end]))
L = length(cache.bcresid_prototype)

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

jac = if iip
(J, u, p) -> __mirk_2point_jacobian!(J, u, p, jac_alg.diffmode, diffcache, lossₚ,
(J, u, p) -> __mirk_2point_jacobian!(J, u, jac_alg.diffmode, diffcache, lossₚ,
resid)
else
(u, p) -> __mirk_2point_jacobian(u, p, jac_prototype, jac_alg.diffmode, diffcache,
(u, p) -> __mirk_2point_jacobian(u, jac_prototype, jac_alg.diffmode, diffcache,
lossₚ)
end

return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p)
nlf = NonlinearFunction{iip}(loss; resid_prototype = copy(resid), jac, jac_prototype)

return (L == cache.M ? NonlinearProblem : NonlinearLeastSquaresProblem)(nlf, y, cache.p)
end

function __mirk_2point_jacobian!(J, x, p, diffmode, diffcache, loss_fn::L, resid) where {L}
function __mirk_2point_jacobian!(J, x, diffmode, diffcache, loss_fn::L, resid) where {L}
sparse_jacobian!(J, diffmode, diffcache, loss_fn, resid, x)
return J
end

function __mirk_2point_jacobian(x, p, J, diffmode, diffcache, loss_fn::L) where {L}
function __mirk_2point_jacobian(x, J, diffmode, diffcache, loss_fn::L) where {L}
sparse_jacobian!(J, diffmode, diffcache, loss_fn, x)
return J
end
Loading

0 comments on commit c2e7dc9

Please sign in to comment.