diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 08298a58..66043e81 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -33,6 +33,7 @@ jobs: - "pre" group: - "MIRK" + - "MIRKN" - "MISC" - "SHOOTING" - "FIRK(EXPANDED)" diff --git a/lib/BoundaryValueDiffEqMIRKN/Project.toml b/lib/BoundaryValueDiffEqMIRKN/Project.toml new file mode 100644 index 00000000..a3a2c535 --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/Project.toml @@ -0,0 +1,82 @@ +name = "BoundaryValueDiffEqMIRKN" +uuid = "9255f1d6-53bf-473e-b6bd-23f1ff009da4" +authors = ["Qingyu Qu "] +version = "0.1.0" + +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +BandedMatrices = "aae01518-5342-5314-be14-df237901396f" +BoundaryValueDiffEqCore = "56b672f2-a5fe-4263-ab2d-da677488eb3a" +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e" +FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" +PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" + +[compat] +ADTypes = "1.2" +Adapt = "4" +Aqua = "0.8.7" +ArrayInterface = "7.7" +BandedMatrices = "1.4" +BoundaryValueDiffEqCore = "1.0.0" +ConcreteStructs = "0.2.3" +DiffEqBase = "6.146" +DiffEqDevTools = "2.44" +FastAlmostBandedMatrices = "0.1.1" +FastClosures = "0.3" +ForwardDiff = "0.10.36" +JET = "0.9.12" +LineSearch = "0.1.3" +LineSearches = "7.3.0" +LinearAlgebra = "1.10" +LinearSolve = "2.21" +Logging = "1.10" +NonlinearSolve = "3.15.1" +OrdinaryDiffEq = "6.89.0" +PreallocationTools = "0.4.24" +PrecompileTools = "1.2" +Preferences = "1.4" +Random = "1.10" +ReTestItems = "1.23.1" +RecursiveArrayTools = "3.27.0" +Reexport = "1.2" +SciMLBase = "2.40" +Setfield = "1" +SparseArrays = "1.10" +SparseDiffTools = "2.14" +StaticArrays = "1.8.1" +Test = "1.10" +julia = "1.10" + +[extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["Aqua", "DiffEqDevTools", "JET", "LinearSolve", "OrdinaryDiffEq", "Random", "ReTestItems", "RecursiveArrayTools", "StaticArrays", "Test"] diff --git a/lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl b/lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl new file mode 100644 index 00000000..a359a23f --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl @@ -0,0 +1,62 @@ +module BoundaryValueDiffEqMIRKN + +import PrecompileTools: @compile_workload, @setup_workload + +using ADTypes, Adapt, ArrayInterface, BoundaryValueDiffEqCore, DiffEqBase, ForwardDiff, + LinearAlgebra, NonlinearSolve, Preferences, RecursiveArrayTools, Reexport, SciMLBase, + Setfield, SparseDiffTools + +using PreallocationTools: PreallocationTools, DiffCache + +# Special Matrix Types +using BandedMatrices, FastAlmostBandedMatrices, SparseArrays + +import BoundaryValueDiffEqCore: BoundaryValueDiffEqAlgorithm, BVPJacobianAlgorithm, + recursive_flatten, recursive_flatten!, recursive_unflatten!, + __concrete_nonlinearsolve_algorithm, diff!, + __FastShortcutBVPCompatibleNonlinearPolyalg, + __FastShortcutBVPCompatibleNLLSPolyalg, eval_bc_residual, + eval_bc_residual!, get_tmp, __maybe_matmul!, + __append_similar!, __extract_problem_details, + __initial_guess, __maybe_allocate_diffcache, + __get_bcresid_prototype, __similar, __vec, __vec_f, + __vec_f!, __vec_bc, __vec_bc!, recursive_flatten_twopoint!, + __unsafe_nonlinearfunction, __internal_nlsolve_problem, + __extract_mesh, __extract_u0, __has_initial_guess, + __initial_guess_length, __initial_guess_on_mesh, + __flatten_initial_guess, __build_solution, __Fix3, + __sparse_jacobian_cache, __sparsity_detection_alg, + _sparse_like, ColoredMatrix, __default_sparse_ad, + __default_nonsparse_ad + +import ADTypes: AbstractADType +import ArrayInterface: matrix_colors, parameterless_type, undefmatrix, fast_scalar_indexing +import ConcreteStructs: @concrete +import DiffEqBase: solve +import FastClosures: @closure +import ForwardDiff: ForwardDiff, pickchunksize +import Logging +import RecursiveArrayTools: ArrayPartition, DiffEqArray +import SciMLBase: AbstractDiffEqInterpolation, AbstractBVProblem, + StandardSecondOrderBVProblem, StandardBVProblem, __solve, _unwrap_val + +@reexport using ADTypes, DiffEqBase, NonlinearSolve, SparseDiffTools, SciMLBase + +include("utils.jl") +include("types.jl") +include("algorithms.jl") +include("mirkn.jl") +include("alg_utils.jl") +include("collocation.jl") +include("mirkn_tableaus.jl") + +function __solve( + prob::AbstractBVProblem, alg::BoundaryValueDiffEqAlgorithm, args...; kwargs...) + cache = init(prob, alg, args...; kwargs...) + return solve!(cache) +end + +export MIRKN4, MIRKN6 +export BVPJacobianAlgorithm + +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/alg_utils.jl b/lib/BoundaryValueDiffEqMIRKN/src/alg_utils.jl new file mode 100644 index 00000000..f04e7615 --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/alg_utils.jl @@ -0,0 +1,7 @@ +for order in (4, 6) + alg = Symbol("MIRKN$(order)") + @eval alg_order(::$(alg)) = $order + @eval alg_stage(::$(alg)) = $(order - 1) +end + +SciMLBase.isadaptive(alg::AbstractMIRKN) = false diff --git a/lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl b/lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl new file mode 100644 index 00000000..cd19a2db --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl @@ -0,0 +1,55 @@ +# Algorithms +abstract type AbstractMIRKN <: BoundaryValueDiffEqAlgorithm end + +for order in (4, 6) + alg = Symbol("MIRKN$(order)") + + @eval begin + """ + $($alg)(; nlsolve = NewtonRaphson(), jac_alg = BVPJacobianAlgorithm(), + defect_threshold = 0.1, max_num_subintervals = 3000) + + $($order)th order Monotonic Implicit Runge Kutta Nyström method. + + ## Keyword Arguments + + - `nlsolve`: Internal Nonlinear solver. Any solver which conforms to the SciML + `NonlinearProblem` interface can be used. Note that any autodiff argument for + the solver will be ignored and a custom jacobian algorithm will be used. + - `jac_alg`: Jacobian Algorithm used for the nonlinear solver. Defaults to + `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to + use based on the input types and problem type. + - For `TwoPointBVProblem`, only `diffmode` is used (defaults to + `AutoSparse(AutoForwardDiff())` if possible else `AutoSparse(AutoFiniteDiff())`). + - For `BVProblem`, `bc_diffmode` and `nonbc_diffmode` are used. For + `nonbc_diffmode` defaults to `AutoSparse(AutoForwardDiff())` if possible else + `AutoSparse(AutoFiniteDiff())`. For `bc_diffmode`, defaults to `AutoForwardDiff` if + possible else `AutoFiniteDiff`. + - `defect_threshold`: Threshold for defect control. + - `max_num_subintervals`: Number of maximal subintervals, default as 3000. + + !!! note + For type-stability, the chunksizes for ForwardDiff ADTypes in + `BVPJacobianAlgorithm` must be provided. + + ## References + + ```bibtex + @article{Muir2001MonoImplicitRM, + title={Mono-Implicit Runge-Kutta-Nystr{\"o}m Methods with Application to Boundary Value Ordinary Differential Equations}, + author={Paul H. Muir and Mark F. Adams}, + journal={BIT Numerical Mathematics}, + year={2001}, + volume={41}, + pages={776-799} + } + ``` + """ + @kwdef struct $(alg){N, J <: BVPJacobianAlgorithm, T} <: AbstractMIRKN + nlsolve::N = nothing + jac_alg::J = BVPJacobianAlgorithm() + defect_threshold::T = 0.1 + max_num_subintervals::Int = 3000 + end + end +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/collocation.jl b/lib/BoundaryValueDiffEqMIRKN/src/collocation.jl new file mode 100644 index 00000000..b8bfec2d --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/collocation.jl @@ -0,0 +1,74 @@ +@views function Φ!(residual, cache::MIRKNCache, y, u, p = cache.p) + return Φ!(residual, cache.fᵢ_cache, cache.fᵢ₂_cache, cache.k_discrete, + cache.f, cache.TU, y, u, p, cache.mesh, cache.mesh_dt, cache.stage) +end + +@views function Φ!(residual, fᵢ_cache, fᵢ₂_cache, k_discrete, f!, + TU::MIRKNTableau, y, u, p, mesh, mesh_dt, stage::Int) + (; c, v, w, b, x, vp, bp, xp) = TU + L = length(mesh) + T = eltype(u) + tmp = get_tmp(fᵢ_cache, u) + tmpd = get_tmp(fᵢ₂_cache, u) + for i in 1:(L - 1) + dtᵢ = mesh_dt[i] + yᵢ = get_tmp(y[i], u) + yᵢ₊₁ = get_tmp(y[i + 1], u) + yₗ₊ᵢ = get_tmp(y[L + i], u) + yₗ₊ᵢ₊₁ = get_tmp(y[L + i + 1], u) + K = get_tmp(k_discrete[i], u) + for r in 1:stage + @. tmp = (1 - v[r]) * yᵢ + + v[r] * yᵢ₊₁ + + dtᵢ * ((c[r] - v[r] - w[r]) * yₗ₊ᵢ + w[r] * yₗ₊ᵢ₊₁) + @. tmpd = (1 - vp[r]) * yₗ₊ᵢ + vp[r] * yₗ₊ᵢ₊₁ + __maybe_matmul!(tmp, K[:, 1:(r - 1)], x[r, 1:(r - 1)], dtᵢ^2, T(1)) + __maybe_matmul!(tmpd, K[:, 1:(r - 1)], xp[r, 1:(r - 1)], dtᵢ, T(1)) + f!(K[:, r], tmpd, tmp, p, mesh[i] + c[r] * dtᵢ) + end + + # Update residual + @. residual[i] = yᵢ₊₁ - yᵢ - dtᵢ * yₗ₊ᵢ + __maybe_matmul!(residual[i], K[:, 1:stage], b[1:stage], -dtᵢ^2, T(1)) + @. residual[L + i - 1] = yₗ₊ᵢ₊₁ - yₗ₊ᵢ + __maybe_matmul!(residual[L + i - 1], K[:, 1:stage], bp[1:stage], -dtᵢ, T(1)) + end +end + +function Φ(cache::MIRKNCache, y, u, p = cache.p) + return Φ(cache.fᵢ_cache, cache.fᵢ₂_cache, cache.k_discrete, cache.f, + cache.TU, y, u, p, cache.mesh, cache.mesh_dt, cache.stage) +end + +@views function Φ(fᵢ_cache, fᵢ₂_cache, k_discrete, f, TU::MIRKNTableau, + y, u, p, mesh, mesh_dt, stage::Int) + (; c, v, w, b, x, vp, bp, xp) = TU + residual = [similar(yᵢ) for yᵢ in y[1:(end - 2)]] + L = length(mesh) + T = eltype(u) + tmp = get_tmp(fᵢ_cache, u) + tmpd = get_tmp(fᵢ₂_cache, u) + for i in 1:(L - 1) + dtᵢ = mesh_dt[i] + yᵢ = get_tmp(y[i], u) + yᵢ₊₁ = get_tmp(y[i + 1], u) + yₗ₊ᵢ = get_tmp(y[L + i], u) + yₗ₊ᵢ₊₁ = get_tmp(y[L + i + 1], u) + K = get_tmp(k_discrete[i], u) + for r in 1:stage + @. tmp = (1 - v[r]) * yᵢ + + v[r] * yᵢ₊₁ + + dtᵢ * ((c[r] - v[r] - w[r]) * yₗ₊ᵢ + w[r] * yₗ₊ᵢ₊₁) + @. tmpd = (1 - vp[r]) * yₗ₊ᵢ + vp[r] * yₗ₊ᵢ₊₁ + __maybe_matmul!(tmp, K[:, 1:(r - 1)], x[r, 1:(r - 1)], dtᵢ^2, T(1)) + __maybe_matmul!(tmpd, K[:, 1:(r - 1)], xp[r, 1:(r - 1)], dtᵢ, T(1)) + + K[:, r] .= f(tmpd, tmp, p, mesh[i] + c[r] * dtᵢ) + end + @. residual[i] = yᵢ₊₁ - yᵢ - dtᵢ * yₗ₊ᵢ + __maybe_matmul!(residual[i], K[:, 1:stage], b[1:stage], -dtᵢ^2, T(1)) + @. residual[L + i - 1] = yₗ₊ᵢ₊₁ - yₗ₊ᵢ + __maybe_matmul!(residual[L + i - 1], K[:, 1:stage], bp[1:stage], -dtᵢ, T(1)) + end + return residual +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl b/lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl new file mode 100644 index 00000000..0b6630af --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl @@ -0,0 +1,192 @@ +@concrete mutable struct MIRKNCache{iip, T} + order::Int # The order of MIRKN method + stage::Int # The state of MIRKN method + M::Int # The number of equations + in_size + f + bc + prob # BVProblem + problem_type # StandardBVProblem + p # Parameters + alg # MIRKN methods + TU # MIRKN Tableau + bcresid_prototype + mesh # Discrete mesh + mesh_dt + k_discrete # Stage information associated with the discrete Runge-Kutta-Nyström method + y + y₀ + residual + fᵢ_cache + fᵢ₂_cache + resid_size + kwargs +end + +Base.eltype(::MIRKNCache{iip, T}) where {iip, T} = T + +function SciMLBase.__init( + prob::SecondOrderBVProblem, alg::AbstractMIRKN; dt = 0.0, kwargs...) + @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg) + iip = isinplace(prob) + t₀, t₁ = prob.tspan + ig, T, M, Nig, X = __extract_problem_details(prob; dt, check_positive_dt = true) + mesh = __extract_mesh(prob.u0, t₀, t₁, Nig) + mesh_dt = diff(mesh) + + TU = constructMIRKN(alg, T) + + # Don't flatten this here, since we need to expand it later if needed + y₀ = __initial_guess_on_mesh(prob, prob.u0, Nig, prob.p, false) + chunksize = pickchunksize(M * (2 * Nig - 2)) + __alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg) + + y = __alloc.(copy.(y₀)) + fᵢ_cache = __alloc(similar(X)) + fᵢ₂_cache = __alloc(similar(X)) + stage = alg_stage(alg) + bcresid_prototype = zero(vcat(X, X)) + k_discrete = [__maybe_allocate_diffcache(similar(X, M, stage), chunksize, alg.jac_alg) + for _ in 1:Nig] + + residual = if iip + __alloc.(copy.(@view(y₀[1:end]))) + else + nothing + end + + resid_size = size(bcresid_prototype) + f, bc = if X isa AbstractVector + prob.f, prob.f.bc + elseif iip + vecf! = @closure (ddu, du, u, p, t) -> __vec_f!(ddu, du, u, p, t, prob.f, size(X)) + vecbc! = if !(prob.problem_type isa TwoPointSecondOrderBVProblem) + @closure (r, du, u, p, t) -> __vec_so_bc!( + r, du, u, p, t, prob.f.bc, resid_size, size(X)) + else + ( + @closure((r, du, u, p)->__vec_so_bc!( + r, du, u, p, first(prob.f.bc), resid_size[1], size(X))), + @closure((r, du, u, p)->__vec_so_bc!( + r, du, u, p, last(prob.f.bc), resid_size[2], size(X)))) + end + vecf!, vecbc! + else + vecf = @closure (du, u, p, t) -> __vec_f(du, u, p, t, prob.f, size(X)) + vecbc = if !(prob.problem_type isa TwoPointSecondOrderBVProblem) + @closure (du, u, p, t) -> __vec_so_bc(du, u, p, t, prob.f.bc, size(X)) + else + (@closure((du, u, p)->__vec_so_bc(du, u, p, first(prob.f.bc), size(X))), + @closure((du, u, p)->__vec_so_bc(du, u, p, last(prob.f.bc), size(X)))) + end + vecf, vecbc + end + + prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob + + return MIRKNCache{iip, T}( + alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, + prob.p, alg, TU, bcresid_prototype, mesh, mesh_dt, k_discrete, + y, y₀, residual, fᵢ_cache, fᵢ₂_cache, resid_size, kwargs) +end + +function SciMLBase.solve!(cache::MIRKNCache{iip, T}) where {iip, T} + (; mesh, M, p, prob, kwargs) = cache + nlprob = __construct_nlproblem(cache, vec(cache.y₀)) + nlsolve_alg = __concrete_nonlinearsolve_algorithm(nlprob, cache.alg.nlsolve) + sol_nlprob = __solve(nlprob, nlsolve_alg; kwargs..., alias_u0 = true) + recursive_unflatten!(cache.y₀, sol_nlprob.u) + solu = ArrayPartition.(cache.y₀[1:length(mesh)], cache.y₀[(length(mesh) + 1):end]) + return SciMLBase.build_solution( + prob, cache.alg, mesh, solu; retcode = sol_nlprob.retcode) +end + +function __construct_nlproblem(cache::MIRKNCache{iip}, y::AbstractVector) where {iip} + (; alg) = cache + pt = cache.problem_type + loss = if iip + @closure (du, u, p) -> __mirkn_loss!( + du, u, p, cache.y, pt, cache.bc, cache.residual, cache.mesh, cache) + else + @closure (u, p) -> __mirkn_loss(u, p, cache.y, pt, cache.bc, cache.mesh, cache) + end + + lossₚ = (iip ? __Fix3 : Base.Fix2)(loss, cache.p) + sd = alg.jac_alg.diffmode isa AutoSparse ? SymbolicsSparsityDetection() : + NoSparsityDetection() + ad = alg.jac_alg.diffmode + lz = reduce(vcat, cache.y₀) + jac_cache = __sparse_jacobian_cache(Val(iip), ad, sd, lossₚ, lz, lz) + jac_prototype = init_jacobian(jac_cache) + jac = if iip + @closure (J, u, p) -> __mirkn_mpoint_jacobian!(J, u, ad, jac_cache, lossₚ, lz) + else + @closure (u, p) -> __mirkn_mpoint_jacobian(jac_prototype, u, ad, jac_cache, lossₚ) + end + resid_prototype = zero(lz) + _nlf = __unsafe_nonlinearfunction{iip}(loss; resid_prototype, jac, jac_prototype) + nlprob::NonlinearProblem = NonlinearProblem(_nlf, lz, cache.p) + return nlprob +end + +function __mirkn_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 __mirkn_2point_jacobian(x, J, diffmode, diffcache, loss_fn::L) where {L} + sparse_jacobian!(J, diffmode, diffcache, loss_fn, x) + return J +end + +@inline function __internal_nlsolve_problem( + ::SecondOrderBVProblem{uType, tType, iip, nlls}, resid_prototype, + u0, args...; kwargs...) where {uType, tType, iip, nlls} + return NonlinearProblem(args...; kwargs...) +end + +function __mirkn_mpoint_jacobian!(J, x, diffmode, diffcache, loss, resid) + sparse_jacobian!(J, diffmode, diffcache, loss, resid, x) + return nothing +end + +function __mirkn_mpoint_jacobian(J, x, diffmode, diffcache, loss) + sparse_jacobian!(J, diffmode, diffcache, loss, x) + return J +end + +@views function __mirkn_loss!(resid, u, p, y, pt::StandardSecondOrderBVProblem, + bc::BC, residual, mesh, cache::MIRKNCache) where {BC} + y_ = recursive_unflatten!(y, u) + resids = [get_tmp(r, u) for r in residual] + eval_bc_residual!(resids[1:2], pt, bc, y_, p, mesh) + Φ!(resids[3:end], cache, y_, u, p) + recursive_flatten!(resid, resids) + return nothing +end + +@views function __mirkn_loss(u, p, y, pt::StandardSecondOrderBVProblem, + bc::BC, mesh, cache::MIRKNCache) where {BC} + y_ = recursive_unflatten!(y, u) + resid_bc = eval_bc_residual(pt, bc, y_, p, mesh) + resid_co = Φ(cache, y_, u, p) + return vcat(resid_bc, mapreduce(vec, vcat, resid_co)) +end + +@views function __mirkn_loss!(resid, u, p, y, pt::TwoPointSecondOrderBVProblem, + bc!::BC, residual, mesh, cache::MIRKNCache) where {BC} + y_ = recursive_unflatten!(y, u) + resids = [get_tmp(r, u) for r in residual] + Φ!(resids[3:end], cache, y_, u, p) + eval_bc_residual!(resids, pt, bc!, y_, p, mesh) + recursive_flatten!(resid, resids) + return nothing +end + +@views function __mirkn_loss(u, p, y, pt::TwoPointSecondOrderBVProblem, + bc!::BC, mesh, cache::MIRKNCache) where {BC} + y_ = recursive_unflatten!(y, u) + resid_co = Φ(cache, y_, u, p) + resid_bc = eval_bc_residual(pt, bc!, y_, p, mesh) + return vcat(resid_bc, mapreduce(vec, vcat, resid_co)) +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/mirkn_tableaus.jl b/lib/BoundaryValueDiffEqMIRKN/src/mirkn_tableaus.jl new file mode 100644 index 00000000..58d8f2a6 --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/mirkn_tableaus.jl @@ -0,0 +1,58 @@ +for order in (4, 6) + alg = Symbol("MIRKN$(order)") + f = Symbol("constructMIRKN$(order)") + @eval constructMIRKN(::$(alg), ::Type{T}) where {T} = $(f)(T) +end + +function constructMIRKN4(::Type{T}) where {T} + # RK coefficients tableau + s = 3 + c = [0, 1, 1 // 2, 3 // 10] + v = [0, 1, 1 // 2, 3 // 10] + w = [0, 0, -3 // 20, -1] + b = [1 // 6, 0, 1 // 3, 0] + x = [0 0 0 0 + 0 0 0 0 + 1//80 1//80 0 0 + 4//25 87//500 561//1000 0] + vp = [0, 1, 1 // 2, 2 // 5] + xp = [0 0 0 0 + 0 0 0 0 + 1//8 -1//8 0 0 + 349//3000 -281//3000 -46//375 0] + bp = [1 // 6, 1 // 6, 2 // 3, 0] + + TU = MIRKNTableau(s, T.(c), T.(v), T.(w), T.(b), T.(x), T.(vp), T.(bp), T.(xp)) + return TU +end + +function constructMIRKN6(::Type{T}) where {T} + # RK coefficients tableau + s = 5 + c = [0, 1, 1 // 5, 4 // 5, 1 // 2, 1 // 2, 2 // 5, 4 // 25] + v = [0, 1, 1 // 10, 9 // 10, 1 // 2, 9 // 25, 4 // 5, 2 // 5] + w = [0, 0, -1 // 50, -3 // 25, -1 // 5, -1 // 4, 0, 9 // 25] + b = [1 // 16, 0, 25 // 108, 25 // 432, 4 // 27, 0] + x = [0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 + -7//1500 -2//375 0 0 0 0 0 0 + -2//375 -7//1500 0 0 0 0 0 0 + 2329//61440 2329//61440 -5//12288 -5//12288 0 0 0 0 + 911//38400 23//1536 1073//13824 737//13824 137//5400 0 0 0 + -317//12500 -3//3125 -428//3375 -581//13500 0 -10448//84375 0 0 + 0 -7701259//1562500 -38497579//93750000 -745411//46875000 0 -7158056//9765625 10337749//15625000 0] + vp = [0, 1, 13 // 125, 112 // 125, 1 // 2, 12 // 25, + -17497 // 12500, -44493991 // 31250000] + xp = [0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 + 16//125 -4//125 0 0 0 0 0 0 + 4//125 -16//125 0 0 0 0 0 0 + -13//256 13//256 75//256 -75//256 0 0 0 0 + 183//6400 -167//6400 1165//6912 -1085//6912 4//675 0 0 0 + 29817//200000 17817//200000 223//320 127//320 14//25 -288//3125 0 0 + 76329639//500000000 45632679//500000000 11534329//20000000 7996921//20000000 1//2 56448//1953125 -12936//78125 0] + bp = [1 // 16, 1 // 16, 125 // 432, 125 // 432, 8 // 27, 0, 0, 0] + + TU = MIRKNTableau(s, T.(c), T.(v), T.(w), T.(b), T.(x), T.(vp), T.(bp), T.(xp)) + return TU +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/types.jl b/lib/BoundaryValueDiffEqMIRKN/src/types.jl new file mode 100644 index 00000000..7e60f03c --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/types.jl @@ -0,0 +1,26 @@ +# MIRKN Method Tableaus +struct MIRKNTableau{sType, cType, vType, wType, bType, xType, vpType, bpType, xpType} + """Discrete stages of MIRKN formula""" + s::sType + c::cType + v::vType + w::wType + b::bType + x::xType + vp::vpType + bp::bpType + xp::xpType + + function MIRKNTableau(s, c, v, w, b, x, vp, bp, xp) + @assert eltype(c) == + eltype(v) == + eltype(w) == + eltype(b) == + eltype(x) == + eltype(vp) == + eltype(bp) == + eltype(xp) + return new{typeof(s), typeof(c), typeof(v), typeof(w), typeof(b), + typeof(x), typeof(vp), typeof(bp), typeof(xp)}(s, c, v, w, b, x, vp, bp, xp) + end +end diff --git a/lib/BoundaryValueDiffEqMIRKN/src/utils.jl b/lib/BoundaryValueDiffEqMIRKN/src/utils.jl new file mode 100644 index 00000000..630df39d --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/src/utils.jl @@ -0,0 +1,120 @@ +function general_eval_bc_residual( + ::StandardSecondOrderBVProblem, bc::BC, y, p, mesh) where {BC} + M = length(y[1]) + L = length(mesh) + res_bc = bc(y[(L + 1):end], y[1:L], p, mesh) + return res_bc +end +function eval_bc_residual(::StandardSecondOrderBVProblem, bc::BC, y, p, mesh) where {BC} + L = length(mesh) + res_bc = bc(y[(L + 1):end], y[1:L], p, mesh) + return res_bc +end +function eval_bc_residual(::TwoPointSecondOrderBVProblem, bc::BC, sol, p, mesh) where {BC} + M = length(sol[1]) + L = length(mesh) + ua = sol isa AbstractVector ? sol[1] : sol(first(t))[1:M] + ub = sol isa AbstractVector ? sol[L] : sol(last(t))[1:M] + dua = sol isa AbstractVector ? sol[L + 1] : sol(first(t))[(M + 1):end] + dub = sol isa AbstractVector ? sol[end] : sol(last(t))[(M + 1):end] + return vcat(bc[1](dua, ua, p), bc[2](dub, ub, p)) +end + +function eval_bc_residual!(resid, ::StandardBVProblem, bc!::BC, sol, p, t) where {BC} + bc!(resid, sol, p, t) +end +function eval_bc_residual!( + resid, ::StandardSecondOrderBVProblem, bc!::BC, dsol, sol, p, t) where {BC} + bc!(resid, dsol, sol, p, t) +end + +function general_eval_bc_residual!( + resid, ::StandardSecondOrderBVProblem, bc!::BC, y, p, mesh) where {BC} + L = length(mesh) + bc!(resid, y[(L + 1):end], y[1:L], p, mesh) +end +function eval_bc_residual!( + resid, ::StandardSecondOrderBVProblem, bc!::BC, y, p, mesh) where {BC} + M = length(y[1]) + L = length(mesh) + res_bc = vcat(resid[1], resid[2]) + bc!(res_bc, y[(L + 1):end], y[1:L], p, mesh) + copyto!(resid[1], res_bc[1:M]) + copyto!(resid[2], res_bc[(M + 1):end]) +end +function eval_bc_residual!( + resid, ::TwoPointSecondOrderBVProblem, bc::BC, sol, p, mesh) where {BC} + M = length(sol[1]) + L = length(mesh) + ua = sol isa AbstractVector ? sol[1] : sol(first(t))[1:M] + ub = sol isa AbstractVector ? sol[L] : sol(last(t))[1:M] + dua = sol isa AbstractVector ? sol[L + 1] : sol(first(t))[(M + 1):end] + dub = sol isa AbstractVector ? sol[end] : sol(last(t))[(M + 1):end] + bc[1](resid[1], dua, ua, p) + bc[2](resid[2], dub, ub, p) +end + +function __get_bcresid_prototype(::TwoPointSecondOrderBVProblem, prob::BVProblem, u) + prototype = if prob.f.bcresid_prototype !== nothing + prob.f.bcresid_prototype.x + else + first(prob.f.bc)(u, prob.p), last(prob.f.bc)(u, prob.p) + end + return prototype, size.(prototype) +end +function __get_bcresid_prototype(::StandardSecondOrderBVProblem, prob::BVProblem, u) + prototype = prob.f.bcresid_prototype !== nothing ? prob.f.bcresid_prototype : + __zeros_like(u) + return prototype, size(prototype) +end + +# Restructure Non-Vector Inputs +function __vec_f!(ddu, du, u, p, t, f!, u_size) + f!(reshape(ddu, u_size), reshape(du, u_size), reshape(u, u_size), p, t) + return nothing +end + +__vec_f(du, u, p, t, f, u_size) = vec(f(reshape(du, u_size), reshape(u, u_size), p, t)) + +function __vec_so_bc!(resid, dsol, sol, p, t, bc!, resid_size, u_size) + bc!(reshape(resid, resid_size), __restructure_sol(dsol, u_size), + __restructure_sol(sol, u_size), p, t) + return nothing +end + +function __vec_so_bc!(resid, dsol, sol, p, bc!, resid_size, u_size) + bc!(reshape(resid, resid_size), reshape(dsol, u_size), reshape(sol, u_size), p) + return nothing +end + +function __vec_so_bc(dsol, sol, p, t, bc, u_size) + vec(bc(__restructure_sol(dsol, u_size), __restructure_sol(sol, u_size), p, t)) +end +function __vec_so_bc(dsol, sol, p, bc, u_size) + vec(bc(reshape(dsol, u_size), reshape(sol, u_size), p)) +end + +__get_non_sparse_ad(ad::AbstractADType) = ad + +@inline function __initial_guess_on_mesh( + prob::SecondOrderBVProblem, u₀::AbstractArray, Nig, p, alias_u0::Bool) + return [copy(vec(u₀)) for _ in 1:(2 * (Nig + 1))] +end + +function concrete_jacobian_algorithm( + jac_alg::BVPJacobianAlgorithm, prob::AbstractBVProblem, alg) + return concrete_jacobian_algorithm(jac_alg, prob.problem_type, prob, alg) +end + +function concrete_jacobian_algorithm( + jac_alg::BVPJacobianAlgorithm, prob_type, prob::SecondOrderBVProblem, alg) + u0 = __extract_u0(prob.u0, prob.p, first(prob.tspan)) + diffmode = jac_alg.diffmode === nothing ? __default_sparse_ad(u0) : jac_alg.diffmode + bc_diffmode = jac_alg.bc_diffmode === nothing ? + (prob_type isa TwoPointSecondOrderBVProblem ? __default_sparse_ad : + __default_nonsparse_ad)(u0) : jac_alg.bc_diffmode + nonbc_diffmode = jac_alg.nonbc_diffmode === nothing ? __default_sparse_ad(u0) : + jac_alg.nonbc_diffmode + + return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode) +end diff --git a/lib/BoundaryValueDiffEqMIRKN/test/mirkn_basic_tests.jl b/lib/BoundaryValueDiffEqMIRKN/test/mirkn_basic_tests.jl new file mode 100644 index 00000000..8f309323 --- /dev/null +++ b/lib/BoundaryValueDiffEqMIRKN/test/mirkn_basic_tests.jl @@ -0,0 +1,137 @@ +@testitem "Example problem from paper" begin + using BoundaryValueDiffEqMIRKN + + for order in (4, 6) + s = Symbol("MIRKN$(order)") + @eval mirkn_solver(::Val{$order}, args...; kwargs...) = $(s)(args...; kwargs...) + end + + function test!(ddu, du, u, p, t) + ϵ = 0.1 + ddu[1] = u[2] + ddu[2] = (-u[1] * du[2] - u[3] * du[3]) / ϵ + ddu[3] = (du[1] * u[3] - u[1] * du[3]) / ϵ + end + function test(du, u, p, t) + ϵ = 0.1 + return [u[2], (-u[1] * du[2] - u[3] * du[3]) / ϵ, (du[1] * u[3] - u[1] * du[3]) / ϵ] + end + + function bc!(res, du, u, p, t) + res[1] = u[1][1] + res[2] = u[end][1] + res[3] = u[1][3] + 1 + res[4] = u[end][3] - 1 + res[5] = du[1][1] + res[6] = du[end][1] + end + + function bc(du, u, p, t) + return [u[1][1], u[end][1], u[1][3] + 1, u[end][3] - 1, du[1][1], du[end][1]] + end + function bca!(resa, du, u, p) + resa[1] = u[1] + resa[2] = u[3] + 1 + resa[3] = du[1] + end + function bcb!(resb, du, u, p) + resb[1] = u[1] + resb[2] = u[3] - 1 + resb[3] = du[1] + end + + function bca(du, u, p) + [u[1], u[3] + 1, du[1]] + end + function bcb(du, u, p) + [u[1], u[3] - 1, du[1]] + end + + u0 = [1.0, 1.0, 1.0] + tspan = (0.0, 1.0) + + probArr = [SecondOrderBVProblem(test!, bc!, u0, tspan), + SecondOrderBVProblem(test, bc, u0, tspan), + TwoPointSecondOrderBVProblem( + test!, (bca!, bcb!), u0, tspan, bcresid_prototype = (zeros(3), zeros(3))), + TwoPointSecondOrderBVProblem( + test, (bca, bcb), u0, tspan, bcresid_prototype = (zeros(3), zeros(3)))] + + @testset "MIRKN$order" for order in (4, 6) + @testset "Problem $i" for i in 1:4 + sol = solve(probArr[i], mirkn_solver(Val(order)); dt = 0.01) + @test SciMLBase.successful_retcode(sol) + end + end +end + +@testitem "Convergence on Linear" begin + using LinearAlgebra, DiffEqDevTools + + for order in (4, 6) + s = Symbol("MIRKN$(order)") + @eval mirkn_solver(::Val{$order}, args...; kwargs...) = $(s)(args...; kwargs...) + end + + function f!(ddu, du, u, p, t) + ddu[1] = u[1] + end + function f(du, u, p, t) + return u[1] + end + function bc!(res, du, u, p, t) + res[1] = u[1][1] - 1 + res[2] = u[end][1] + end + function bc(du, u, p, t) + return [u[1][1] - 1, u[end][1]] + end + function bc_a!(res, du, u, p) + res[1] = u[1] - 1 + end + function bc_b!(res, du, u, p) + res[1] = u[1] + end + function bc_a(du, u, p) + return [u[1] - 1] + end + function bc_b(du, u, p) + return [u[1]] + end + u0 = [1.0] + tspan = (0.0, 1.0) + testTol = 0.2 + bvpf1 = DynamicalBVPFunction(f!, + bc!, + analytic = (u0, p, t) -> [ + (exp(-t) - exp(t - 2)) / (1 - exp(-2)), (-exp(-t) - exp(t - 2)) / (1 - exp(-2))]) + bvpf2 = DynamicalBVPFunction(f, + bc, + analytic = (u0, p, t) -> [ + (exp(-t) - exp(t - 2)) / (1 - exp(-2)), (-exp(-t) - exp(t - 2)) / (1 - exp(-2))]) + bvpf3 = DynamicalBVPFunction(f!, + (bc_a!, bc_b!), + analytic = (u0, p, t) -> [ + (exp(-t) - exp(t - 2)) / (1 - exp(-2)), (-exp(-t) - exp(t - 2)) / (1 - exp(-2))], + bcresid_prototype = (zeros(1), zeros(1)), + twopoint = Val(true)) + bvpf4 = DynamicalBVPFunction(f, + (bc_a, bc_b), + analytic = (u0, p, t) -> [ + (exp(-t) - exp(t - 2)) / (1 - exp(-2)), (-exp(-t) - exp(t - 2)) / (1 - exp(-2))], + bcresid_prototype = (zeros(1), zeros(1)), + twopoint = Val(true)) + probArr = [ + SecondOrderBVProblem(bvpf1, u0, tspan), SecondOrderBVProblem(bvpf2, u0, tspan), + TwoPointSecondOrderBVProblem(bvpf3, u0, tspan), + TwoPointSecondOrderBVProblem(bvpf4, u0, tspan)] + dts = 1 .// 2 .^ (3:-1:1) + @testset "Problem: $i" for i in (1, 2, 3, 4) + prob = probArr[i] + @testset "MIRKN$order" for order in (4, 6) + sim = test_convergence( + dts, prob, mirkn_solver(Val(order)); abstol = 1e-8, reltol = 1e-8) + @test sim.𝒪est[:final]≈order atol=testTol + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 08660990..fc173e06 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,6 +9,12 @@ function activate_mirk() Pkg.instantiate() end +function activate_mirkn() + Pkg.activate("../lib/BoundaryValueDiffEqMIRKN") + Pkg.develop(PackageSpec(path = dirname(@__DIR__))) + Pkg.instantiate() +end + function activate_firk() Pkg.activate("../lib/BoundaryValueDiffEqFIRK") Pkg.develop(PackageSpec(path = dirname(@__DIR__))) @@ -38,6 +44,13 @@ end end end + if GROUP == "All" || GROUP == "MIRKN" + @time "MIRKN solvers" begin + activate_mirkn() + ReTestItems.runtests("../lib/BoundaryValueDiffEqMIRKN/test/mirkn_basic_tests.jl") + end + end + if GROUP == "All" || GROUP == "FIRK(EXPANDED)" @time "FIRK Expanded solvers" begin activate_firk()