Skip to content

Commit

Permalink
Merge pull request #230 from ErikQQY/qqy/mirkn
Browse files Browse the repository at this point in the history
Add MIRKN methods
  • Loading branch information
ChrisRackauckas authored Nov 8, 2024
2 parents c232921 + 61cefda commit a92baed
Show file tree
Hide file tree
Showing 12 changed files with 827 additions and 0 deletions.
1 change: 1 addition & 0 deletions .github/workflows/Tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ jobs:
- "pre"
group:
- "MIRK"
- "MIRKN"
- "MISC"
- "SHOOTING"
- "FIRK(EXPANDED)"
Expand Down
82 changes: 82 additions & 0 deletions lib/BoundaryValueDiffEqMIRKN/Project.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
name = "BoundaryValueDiffEqMIRKN"
uuid = "9255f1d6-53bf-473e-b6bd-23f1ff009da4"
authors = ["Qingyu Qu <[email protected]>"]
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"]
62 changes: 62 additions & 0 deletions lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl
Original file line number Diff line number Diff line change
@@ -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
7 changes: 7 additions & 0 deletions lib/BoundaryValueDiffEqMIRKN/src/alg_utils.jl
Original file line number Diff line number Diff line change
@@ -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
55 changes: 55 additions & 0 deletions lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl
Original file line number Diff line number Diff line change
@@ -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
74 changes: 74 additions & 0 deletions lib/BoundaryValueDiffEqMIRKN/src/collocation.jl
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit a92baed

Please sign in to comment.