Splitting up BoundaryValueDiffEq

Signed-off-by: ErikQQY <2283984853@qq.com>

.github/workflows/Tests.yml

@@ -34,6 +34,21 @@ jobs:
           - "FIRK(EXPANDED)"
           - "FIRK(NESTED)"
           - "WRAPPERS"
+
+          - "BoundaryValueDiffEqFIRK"
+          - "BoundaryValueDiffEqMIRK"
+          - "BoundaryValueDiffEqShooting"
+
+    steps:
+      # Explicitly develop the libraries first before running the tests for now.
+      # This is necessary since the tests are likely to fail otherwise, given that all
+      # the libs haven't been registered yet.
+      - name: "Develop the libraries since they haven't been registered yet"
+        run: |
+          julia --project=. -e '
+              using Pkg;
+              Pkg.develop(map(path ->Pkg.PackageSpec.(;path="$(@__DIR__)/lib/$(path)"), readdir("./lib")));
+          '
     uses: "SciML/.github/.github/workflows/tests.yml@v1"
     with:
       group: "${{ matrix.group }}"

lib/BoundaryValueDiffEqCore/Project.toml

@@ -0,0 +1,60 @@
+name = "BoundaryValueDiffEqCore"
+uuid = "8bd40a78-d78a-4856-9888-3324b15bdffe"
+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"
+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"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+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"
+ConcreteStructs = "0.2.3"
+DiffEqBase = "6.146"
+DiffEqDevTools = "2.44"
+FastAlmostBandedMatrices = "0.1.1"
+FastClosures = "0.3"
+ForwardDiff = "0.10.36"
+JET = "0.8"
+LinearAlgebra = "1.10"
+LinearSolve = "2.21"
+Logging = "1.10"
+NonlinearSolve = "3.8.1"
+OrdinaryDiffEq = "6.89.0"
+PreallocationTools = "0.4.24"
+PrecompileTools = "1.2"
+Preferences = "1.4"
+Random = "1.10"
+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"

lib/BoundaryValueDiffEqCore/src/BoundaryValueDiffEqCore.jl

@@ -0,0 +1,39 @@
+module BoundaryValueDiffEqCore
+
+import PrecompileTools: @compile_workload, @setup_workload
+
+using ADTypes, Adapt, ArrayInterface, DiffEqBase, ForwardDiff, LinearAlgebra,
+      NonlinearSolve, OrdinaryDiffEq, Preferences, RecursiveArrayTools, Reexport, SciMLBase,
+      Setfield, SparseDiffTools
+
+using PreallocationTools: PreallocationTools, DiffCache
+
+# Special Matrix Types
+using BandedMatrices, FastAlmostBandedMatrices, SparseArrays
+
+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, StandardBVProblem, __solve, _unwrap_val
+
+@reexport using ADTypes, DiffEqBase, NonlinearSolve, OrdinaryDiffEq, SparseDiffTools,
+                SciMLBase
+
+include("types.jl")
+include("utils.jl")
+include("algorithms.jl")
+include("alg_utils.jl")
+include("default_nlsolve.jl")
+include("sparse_jacobians.jl")
+
+function __solve(prob::BVProblem, alg::BoundaryValueDiffEqAlgorithm, args...; kwargs...)
+    cache = init(prob, alg, args...; kwargs...)
+    return solve!(cache)
+end
+
+end

lib/BoundaryValueDiffEqCore/src/alg_utils.jl

@@ -0,0 +1,3 @@
+SciMLBase.isautodifferentiable(::BoundaryValueDiffEqAlgorithm) = true
+SciMLBase.allows_arbitrary_number_types(::BoundaryValueDiffEqAlgorithm) = true
+SciMLBase.allowscomplex(alg::BoundaryValueDiffEqAlgorithm) = true

lib/BoundaryValueDiffEqCore/src/algorithms.jl

@@ -0,0 +1,20 @@
+# Algorithms
+abstract type BoundaryValueDiffEqAlgorithm <: SciMLBase.AbstractBVPAlgorithm end
+
+## Disable the ugly verbose printing by default
+@inline __modifier_text!(list, fieldname, field) = push!(list, "$fieldname = $(field)")
+@inline __modifier_text!(list, fieldname, ::Nothing) = list
+@inline __modifier_text!(list, fieldname, ::Missing) = list
+@inline function __modifier_text!(list, fieldname, field::SciMLBase.AbstractODEAlgorithm)
+    push!(list, "$fieldname = $(__nameof(field))()")
+end
+
+function Base.show(io::IO, alg::BoundaryValueDiffEqAlgorithm)
+    print(io, "$(__nameof(alg))(")
+    modifiers = String[]
+    for field in fieldnames(typeof(alg))
+        __modifier_text!(modifiers, field, getfield(alg, field))
+    end
+    print(io, join(modifiers, ", "))
+    print(io, ")")
+end

lib/BoundaryValueDiffEqCore/src/sparse_jacobians.jl

@@ -0,0 +1,38 @@
+# This file defines several common patterns of sparse Jacobians we see in the BVP solvers.
+function _sparse_like(I, J, x::AbstractArray, m = maximum(I), n = maximum(J))
+    I′ = adapt(parameterless_type(x), I)
+    J′ = adapt(parameterless_type(x), J)
+    V = __ones_like(x, length(I))
+    return sparse(I′, J′, V, m, n)
+end
+
+# NOTE: We don't retain the Banded Structure in non-TwoPoint BVP cases since vcat/hcat makes
+# it into a dense array. Instead we can atleast exploit sparsity!
+
+# FIXME: Fix the cases where fast_scalar_indexing is not possible
+
+# Helpers for IIP/OOP functions
+function __sparse_jacobian_cache(::Val{iip}, ad, sd, fn, fx, y) where {iip}
+    if iip
+        sparse_jacobian_cache(ad, sd, fn, fx, y)
+    else
+        sparse_jacobian_cache(ad, sd, fn, y; fx)
+    end
+end
+
+@concrete struct ColoredMatrix
+    M
+    row_colorvec
+    col_colorvec
+end
+
+Base.size(M::ColoredMatrix, args...) = size(M.M, args...)
+Base.eltype(M::ColoredMatrix) = eltype(M.M)
+
+ColoredMatrix() = ColoredMatrix(nothing, nothing, nothing)
+
+function __sparsity_detection_alg(M::ColoredMatrix)
+    return PrecomputedJacobianColorvec(;
+        jac_prototype = M.M, M.row_colorvec, M.col_colorvec)
+end
+__sparsity_detection_alg(::ColoredMatrix{Nothing}) = NoSparsityDetection()

src/types.jl → lib/BoundaryValueDiffEqCore/src/types.jl

@@ -1,58 +1,3 @@
-# MIRK Method Tableaus
-struct MIRKTableau{sType, cType, vType, bType, xType}
-    """Discrete stages of MIRK formula"""
-    s::sType
-    c::cType
-    v::vType
-    b::bType
-    x::xType
-
-    function MIRKTableau(s, c, v, b, x)
-        @assert eltype(c) == eltype(v) == eltype(b) == eltype(x)
-        return new{typeof(s), typeof(c), typeof(v), typeof(b), typeof(x)}(s, c, v, b, x)
-    end
-end
-
-struct MIRKInterpTableau{s, c, v, x, τ}
-    s_star::s
-    c_star::c
-    v_star::v
-    x_star::x
-    τ_star::τ
-
-    function MIRKInterpTableau(s_star, c_star, v_star, x_star, τ_star)
-        @assert eltype(c_star) == eltype(v_star) == eltype(x_star)
-        return new{
-            typeof(s_star), typeof(c_star), typeof(v_star), typeof(x_star), typeof(τ_star)}(
-            s_star, c_star, v_star, x_star, τ_star)
-    end
-end
-
-# FIRK Method Tableaus
-struct FIRKTableau{nested, sType, aType, cType, bType}
-    """Discrete stages of RK formula"""
-    s::sType
-    a::aType
-    c::cType
-    b::bType
-
-    function FIRKTableau(s, a, c, b, nested)
-        @assert eltype(a) == eltype(c) == eltype(b)
-        return new{nested, typeof(s), typeof(a), typeof(c), typeof(b)}(s, a, c, b)
-    end
-end
-
-struct FIRKInterpTableau{nested, c, m}
-    q_coeff::c
-    τ_star::m
-    stage::Int
-
-    function FIRKInterpTableau(q_coeff, τ_star, stage, nested::Bool)
-        @assert eltype(q_coeff) == eltype(τ_star)
-        return new{nested, typeof(q_coeff), typeof(τ_star)}(q_coeff, τ_star, stage)
-    end
-end
-
 # Sparsity Detection
 @concrete struct BVPJacobianAlgorithm
     bc_diffmode

src/utils.jl → lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -125,36 +125,6 @@ function __append_similar!(x::AbstractVector{<:MaybeDiffCache}, n, M)
     return x
 end
 
-function __append_similar!(x::AbstractVector{<:AbstractArray}, n, _, TU::FIRKTableau{false})
-    (; s) = TU
-    N = (n - 1) * (s + 1) + 1 - length(x)
-    N == 0 && return x
-    N < 0 && throw(ArgumentError("Cannot append a negative number of elements"))
-    append!(x, [similar(last(x)) for _ in 1:N])
-    return x
-end
-
-function __append_similar!(
-        x::AbstractVector{<:MaybeDiffCache}, n, M, TU::FIRKTableau{false})
-    (; s) = TU
-    N = (n - 1) * (s + 1) + 1 - length(x)
-    N == 0 && return x
-    N < 0 && throw(ArgumentError("Cannot append a negative number of elements"))
-    chunksize = isa(TU, FIRKTableau{false}) ? pickchunksize(M * (N + length(x) * (s + 1))) :
-                pickchunksize(M * (N + length(x)))
-    append!(x, [__maybe_allocate_diffcache(last(x), chunksize) for _ in 1:N])
-    return x
-end
-
-function __append_similar!(x::AbstractVectorOfArray, n, M, TU::FIRKTableau{false})
-    (; s) = TU
-    N = (n - 1) * (s + 1) + 1 - length(x)
-    N == 0 && return x
-    N < 0 && throw(ArgumentError("Cannot append a negative number of elements"))
-    append!(x, VectorOfArray([similar(last(x)) for _ in 1:N]))
-    return x
-end
-
 __append_similar!(::Nothing, n, _, _) = nothing
 function __append_similar!(x::AbstractVectorOfArray, n, _)
     N = n - length(x)

lib/BoundaryValueDiffEqFIRK/Project.toml

@@ -0,0 +1,76 @@
+name = "BoundaryValueDiffEqFIRK"
+uuid = "554fb651-e237-4e24-8fb7-12a5094f9613"
+authors = ["ErikQQY <2283984853@qq.com>"]
+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"
+BoundaryValueDiffEq = "764a87c0-6b3e-53db-9096-fe964310641d"
+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"
+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"
+BoundaryValueDiffEq = "5.10"
+ConcreteStructs = "0.2.3"
+DiffEqBase = "6.146"
+DiffEqDevTools = "2.44"
+FastAlmostBandedMatrices = "0.1.1"
+FastClosures = "0.3"
+ForwardDiff = "0.10.36"
+JET = "0.8"
+LinearAlgebra = "1.10"
+LinearSolve = "2.21"
+Logging = "1.10"
+NonlinearSolve = "3.8.1"
+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"
+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", "Random", "ReTestItems", "RecursiveArrayTools", "StaticArrays", "Test"]