Support for u0 initial guess function

Author: avik-pal

Project.toml

@@ -13,6 +13,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 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"
@@ -22,6 +23,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
@@ -56,6 +58,7 @@ SciMLBase = "2.5"
 Setfield = "1"
 SparseArrays = "1.9"
 SparseDiffTools = "2.9"
+Tricks = "0.1"
 TruncatedStacktraces = "1"
 UnPack = "1"
 julia = "1.9"

src/BoundaryValueDiffEq.jl

@@ -5,7 +5,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
 @recompile_invalidations begin
     using ADTypes, Adapt, BandedMatrices, DiffEqBase, ForwardDiff, LinearAlgebra,
         NonlinearSolve, PreallocationTools, Preferences, RecursiveArrayTools, Reexport,
-        SciMLBase, Setfield, SparseArrays, SparseDiffTools
+        SciMLBase, Setfield, SparseArrays, SparseDiffTools, Tricks
 
     import ADTypes: AbstractADType
     import ArrayInterface: matrix_colors,
@@ -22,6 +22,31 @@ end
 
 @reexport using ADTypes, DiffEqBase, NonlinearSolve, SparseDiffTools, SciMLBase
 
+# TODO: Upstream
+# For BVPs we want to propagate even a function u0
+function DiffEqBase.get_concrete_u0(prob::BVProblem, isadapt, t0, kwargs)
+    if haskey(kwargs, :u0)
+        u0 = kwargs[:u0]
+    else
+        u0 = prob.u0
+    end
+
+    isadapt && eltype(u0) <: Integer && (u0 = float.(u0))
+
+    _u0 = DiffEqBase.handle_distribution_u0(u0)
+
+    if isinplace(prob) && (_u0 isa Number || _u0 isa DiffEqBase.SArray)
+        throw(DiffEqBase.IncompatibleInitialConditionError())
+    end
+
+    if _u0 isa Tuple
+        throw(DiffEqBase.TupleStateError())
+    end
+
+    return _u0
+end
+# End of Upstream
+
 include("types.jl")
 include("utils.jl")
 include("algorithms.jl")

src/solve/mirk.jl

@@ -35,18 +35,20 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0,
         abstol = 1e-3, adaptive = true, kwargs...)
     @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
     iip = isinplace(prob)
+
     _, T, M, n, X = __extract_problem_details(prob; dt, check_positive_dt = true)
+    # NOTE: Assumes the user provided initial guess is on a uniform mesh
+    mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1))
+
+    mesh_dt = diff(mesh)
+
     chunksize = pickchunksize(M * (n + 1))
 
     __alloc = x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
 
     fᵢ_cache = __alloc(similar(X))
     fᵢ₂_cache = vec(similar(X))
 
-    # NOTE: Assumes the user provided initial guess is on a uniform mesh
-    mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1))
-    mesh_dt = diff(mesh)
-
     defect_threshold = T(0.1)  # TODO: Allow user to specify these
     MxNsub = 3000              # TODO: Allow user to specify these
 
@@ -100,7 +102,9 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0,
         vecf, vecbc
     end
 
-    return MIRKCache{iip, T}(alg_order(alg), stage, M, size(X), f, bc, prob,
+    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
+
+    return MIRKCache{iip, T}(alg_order(alg), stage, M, size(X), f, bc, prob_,
         prob.problem_type, prob.p, alg, TU, ITU, bcresid_prototype, mesh, mesh_dt,
         k_discrete, k_interp, y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, new_stages,
         resid₁_size, (; defect_threshold, MxNsub, abstol, dt, adaptive, kwargs...))

src/solve/multiple_shooting.jl

@@ -49,8 +49,8 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
                 ode_cache_loss_fn; kwargs..., verbose, odesolve_kwargs...)
         else
             u_at_nodes = __multiple_shooting_initialize!(nodes, u_at_nodes, prob, alg,
-                cur_nshoot, all_nshoots[i - 1], ig, ode_cache_loss_fn; kwargs..., verbose,
-                odesolve_kwargs...)
+                cur_nshoot, all_nshoots[i - 1], ig, ode_cache_loss_fn, u0; kwargs...,
+                verbose, odesolve_kwargs...)
         end
 
         if prob.problem_type isa TwoPointBVProblem
@@ -362,9 +362,13 @@ end
     resize!(nodes, nshoots + 1)
     nodes .= range(tspan[1], tspan[2]; length = nshoots + 1)
 
-    N = length(first(u0))
-    u_at_nodes = similar(first(u0), (nshoots + 1) * N)
-    recursive_flatten!(u_at_nodes, u0)
+    # NOTE: We don't check `u0 isa Function` since `u0` in-principle can be a callable
+    #       struct
+    u0_ = u0 isa AbstractArray ? u0 : [__initial_guess(u0, prob.p, t) for t in nodes]
+
+    N = length(first(u0_))
+    u_at_nodes = similar(first(u0_), (nshoots + 1) * N)
+    recursive_flatten!(u_at_nodes, u0_)
 
     return u_at_nodes
 end
@@ -401,7 +405,8 @@ end
         end
     else
         @warn "Initialization using odesolve failed. Initializing using 0s. It is \
-               recommended to provide an `initial_guess` in this case."
+               recommended to provide an initial guess function via \
+               `u0 = <function>(p, t)` or `u0 = <function>(t)` in this case."
         fill!(u_at_nodes, 0)
     end
 
@@ -410,16 +415,16 @@ end
 
 # Grid coarsening
 @views function __multiple_shooting_initialize!(nodes, u_at_nodes_prev, prob, alg,
-        nshoots, old_nshoots, ig, odecache_; kwargs...)
-    @unpack f, u0, tspan, p = prob
+        nshoots, old_nshoots, ig, odecache_, u0; kwargs...)
+    @unpack f, tspan, p = prob
     prev_nodes = copy(nodes)
     odecache = odecache_ isa Vector ? first(odecache_) : odecache_
 
     resize!(nodes, nshoots + 1)
     nodes .= range(tspan[1], tspan[2]; length = nshoots + 1)
-    N = _unwrap_val(ig) ? length(first(u0)) : length(u0)
+    N = length(u0)
 
-    u_at_nodes = similar(_unwrap_val(ig) ? first(u0) : u0, N + nshoots * N)
+    u_at_nodes = similar(u0, N + nshoots * N)
     u_at_nodes[1:N] .= u_at_nodes_prev[1:N]
     u_at_nodes[(end - N + 1):end] .= u_at_nodes_prev[(end - N + 1):end]
 

src/utils.jl

@@ -139,17 +139,38 @@ function __extract_problem_details(prob, u0::AbstractArray; dt = 0.0,
     t₀, t₁ = prob.tspan
     return Val(false), eltype(u0), length(u0), Int(cld(t₁ - t₀, dt)), prob.u0
 end
-function __extract_problem_details(prob, ::F; kwargs...) where {F <: Function}
-    throw(ArgumentError("passing `u0` as a function is not supported yet. Curently we only \
-                         support AbstractArray or Vector of AbstractArrays as input! \
-                         Use the latter format for passing in initial guess!"))
+function __extract_problem_details(prob, f::F; dt = 0.0,
+        check_positive_dt::Bool = false) where {F <: Function}
+    # Problem passes in a initial guess function
+    check_positive_dt && dt ≤ 0 && throw(ArgumentError("dt must be positive"))
+    u0 = __initial_guess(f, prob.p, prob.tspan[1])
+    t₀, t₁ = prob.tspan
+    return Val(true), eltype(u0), length(u0), Int(cld(t₁ - t₀, dt)), u0
 end
 
-__initial_state_from_prob(prob::BVProblem, mesh) = __initial_state_from_prob(prob.u0, mesh)
-__initial_state_from_prob(u0::AbstractArray, mesh) = [copy(vec(u0)) for _ in mesh]
-function __initial_state_from_prob(u0::AbstractVector{<:AbstractVector}, _)
+function __initial_guess(f::F, p::P, t::T) where {F, P, T}
+    if static_hasmethod(f, Tuple{P, T})
+        return f(p, t)
+    elseif static_hasmethod(f, Tuple{T})
+        return f(t)
+    else
+        throw(ArgumentError("`initial_guess` must be a function of the form `f(p, t)` or \
+                             `f(t)`"))
+    end
+end
+
+function __initial_state_from_prob(prob::BVProblem, mesh)
+    return __initial_state_from_prob(prob, prob.u0, mesh)
+end
+function __initial_state_from_prob(::BVProblem, u0::AbstractArray, mesh)
+    return [copy(vec(u0)) for _ in mesh]
+end
+function __initial_state_from_prob(::BVProblem, u0::AbstractVector{<:AbstractVector}, _)
     return [copy(vec(u)) for u in u0]
 end
+function __initial_state_from_prob(prob::BVProblem, f::F, mesh) where {F}
+    return [__initial_guess(f, prob.p, t) for t in mesh]
+end
 
 function __get_bcresid_prototype(prob::BVProblem, u)
     return __get_bcresid_prototype(prob.problem_type, prob, u)