Merge pull request #340 from SciML/qqy/fit_parameters

Add unknown parameters estimation

Author: ChrisRackauckas

docs/pages.jl

@@ -2,8 +2,8 @@
 
 pages = ["index.md",
     "Getting Started with BVP solving in Julia" => "tutorials/getting_started.md",
-    "Tutorials" => Any["tutorials/continuation.md",
-        "tutorials/solve_nlls_bvp.md", "tutorials/extremum.md"],
+    "Tutorials" => Any["tutorials/continuation.md", "tutorials/solve_nlls_bvp.md",
+        "tutorials/unknown_parameters.md", "tutorials/extremum.md"],
     "Basics" => Any["basics/bvp_problem.md", "basics/bvp_functions.md",
         "basics/solve.md", "basics/autodiff.md", "basics/error_control.md"],
     "Solver Summaries and Recommendations" => Any[

docs/src/tutorials/unknown_parameters.md

@@ -0,0 +1,61 @@
+# Estimate Unknown Parameters in BVP
+
+When there are unknown parameters in boundary value problems, we can estimate the unknown parameters by solving the BVP, it is quite useful in practical applications in dynamical optimizations and inverse problems. This approach allows us to incorporate both the governing differential equations and boundary conditions to infer parameters that may not be directly measurable.
+
+Let's walk through this functionality with an intuitive example. In the following tutorial, we use the [Mathieu equation](https://en.wikipedia.org/wiki/Mathieu_wavelet) which is a second-order differential equation:
+
+```math
+y''+(\lambda-2q\cos(2x))y=0
+```
+
+where $\lambda$ is the unknown parameter we wish to estimate, `q` is a known real-valued parameter, with boundary conditions when `q=5`:
+
+```math
+y'(0)=0,\ y'(\pi)=0
+```
+
+The second-order BVP can be transformed into a first-order system of BVP:
+
+```math
+\begin{cases}
+y_1'=y_2\\
+y_2'=-(\lambda-2q\cos(2x))y_1
+\end{cases}
+```
+
+with boundary conditions of
+
+```math
+y_2(0)=0,\ y_2(\pi)=0
+```
+
+It is worthnoting that in this system, while we have two differetial equations, it isn't enough to estimate the unknown parameters and gurantee a unique numerical solution with only two given boundary conditions. While under the hood, the parameters are estimated simultanously with the numerical solution, it makes the boundary value problem an underconstrained BVP if the number of constraints are equal to the number of states, which may result in more than one solution. So we should provide additional constraint $y(0)=1$ from the original equation to make sure unique numerical solution and the estimated parameters are we actually wanted.
+
+With BoundaryValueDiffEq.jl, it's easy to solve boundary value problems with unknown parameters, we can just specify `fit_parameters=true` when constructing the BVP and provide the guess of the unknown parameters in `prob.p`, for example, to estimate the unknown parameters in the above BVP system:
+
+```@example unknown
+using BoundaryValueDiffEq, Plots
+tspan = (0.0, pi)
+function f!(du, u, p, t)
+    du[1] = u[2]
+    du[2] = -(p[1] - 10 * cos(2 * t)) * u[1]
+end
+function bca!(res, u, p)
+    res[1] = u[2]
+    res[2] = u[1] - 1.0
+end
+function bcb!(res, u, p)
+    res[1] = u[2]
+end
+guess(p, t) = [cos(4t); -4sin(4t)]
+bvp = TwoPointBVProblem(f!, (bca!, bcb!), guess, tspan, [15.0],
+    bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
+sol = solve(bvp, MIRK4(), dt = 0.05)
+plot(sol)
+```
+
+after solving the boundary value problem, the estimated unknown parameters can be accessed with
+
+```@example unknown
+sol.prob.p
+```

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -262,38 +262,57 @@ function __extract_problem_details(prob, u0::AbstractVectorOfArray; kwargs...)
     _u0 = first(u0.u)
     return Val(true), eltype(_u0), length(_u0), (length(u0.u) - 1), _u0
 end
-function __extract_problem_details(
-        prob, u0::AbstractArray; dt = 0.0, check_positive_dt::Bool = false)
+function __extract_problem_details(prob, u0::AbstractArray; dt = 0.0,
+        check_positive_dt::Bool = false, fit_parameters::Bool = false)
     # Problem does not have Initial Guess
     check_positive_dt && dt ≤ 0 && throw(ArgumentError("dt must be positive"))
     t₀, t₁ = prob.tspan
+    if fit_parameters
+        prob.p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+        new_u = vcat(u0, prob.p)
+        return Val(false), eltype(new_u), length(new_u), Int(cld(t₁ - t₀, dt)), new_u
+    end
     return Val(false), eltype(u0), length(u0), Int(cld(t₁ - t₀, dt)), prob.u0
 end
-function __extract_problem_details(
-        prob, f::F; dt = 0.0, check_positive_dt::Bool = false) where {F <: Function}
+function __extract_problem_details(prob, f::F; dt = 0.0, check_positive_dt::Bool = false,
+        fit_parameters::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])
+
+    u0 = __initial_guess(f, prob.p, prob.tspan[1]; fit_parameters = fit_parameters)
     t₀, t₁ = prob.tspan
     return Val(true), eltype(u0), length(u0), Int(cld(t₁ - t₀, dt)), u0
 end
 
-function __extract_problem_details(
-        prob, u0::SciMLBase.ODESolution; dt = 0.0, check_positive_dt::Bool = false)
+function __extract_problem_details(prob, u0::SciMLBase.ODESolution; dt = 0.0,
+        check_positive_dt::Bool = false, fit_parameters::Bool = false)
     # Problem passes in a initial guess function
     _u0 = first(u0.u)
     _t = u0.t
+    if fit_parameters
+        prob.p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+        new_u = vcat(_u0, prob.p)
+        return Val(false), eltype(new_u), length(new_u), Int(cld(t₁ - t₀, dt)), new_u
+    end
     return Val(true), eltype(_u0), length(_u0), (length(_t) - 1), _u0
 end
 
-function __initial_guess(f::F, p::P, t::T) where {F, P, T}
+function __initial_guess(f::F, p::P, t::T; fit_parameters = false) where {F, P, T}
     if hasmethod(f, Tuple{P, T})
+        p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+        fit_parameters && return vcat(f(p, t), p)
         return f(p, t)
     elseif hasmethod(f, Tuple{T})
         Base.depwarn("initial guess function must take 2 inputs `(p, t)` instead of just \
                      `t`. The single argument version has been deprecated and will be \
                      removed in the next major release of SciMLBase.",
             :__initial_guess)
+        p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+        fit_parameters && return vcat(f(t), p)
         return f(t)
     else
         throw(ArgumentError("`initial_guess` must be a function of the form `f(p, t)`"))
@@ -535,7 +554,7 @@ end
     return VectorOfArray([vec(__initial_guess(u₀, p, t)) for t in mesh])
 end
 @inline function __initial_guess_on_mesh(
-        prob::SecondOrderBVProblem, u₀::AbstractArray, Nig, p, alias_u0::Bool)
+        prob::SecondOrderBVProblem, u₀::AbstractArray, Nig, p)
     return VectorOfArray([copy(vec(u₀)) for _ in 1:(2 * (Nig + 1))])
 end
 

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -1,4 +1,5 @@
-@concrete struct FIRKCacheNested{iip, T, diffcache} <: AbstractBoundaryValueDiffEqCache
+@concrete struct FIRKCacheNested{iip, T, diffcache, fit_parameters} <:
+                 AbstractBoundaryValueDiffEqCache
     order::Int                 # The order of FIRK method
     stage::Int                 # The state of FIRK method
     M::Int                     # The number of equations
@@ -31,7 +32,8 @@ end
 
 Base.eltype(::FIRKCacheNested{iip, T}) where {iip, T} = T
 
-@concrete struct FIRKCacheExpand{iip, T, diffcache} <: AbstractBoundaryValueDiffEqCache
+@concrete struct FIRKCacheExpand{iip, T, diffcache, fit_parameters} <:
+                 AbstractBoundaryValueDiffEqCache
     order::Int                 # The order of FIRK method
     stage::Int                 # The state of FIRK method
     M::Int                     # The number of equations
@@ -103,9 +105,14 @@ function init_nested(
         error("Algorithm doesn't support adaptivity. Please choose a higher order algorithm.")
     end
     diffcache = __cache_trait(alg.jac_alg)
+    fit_parameters = haskey(prob.kwargs, :fit_parameters)
 
     t₀, t₁ = prob.tspan
-    ig, T, M, Nig, X = __extract_problem_details(prob; dt, check_positive_dt = true)
+    ig, T,
+    M,
+    Nig,
+    X = __extract_problem_details(
+        prob; dt, check_positive_dt = true, fit_parameters = fit_parameters)
     mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
     mesh_dt = diff(mesh)
 
@@ -116,7 +123,7 @@ function init_nested(
     fᵢ₂_cache = vec(zero(X))
 
     # Don't flatten this here, since we need to expand it later if needed
-    y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p)
+    y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
 
     y = __alloc.(copy.(y₀.u))
     TU, ITU = constructRK(alg, T)
@@ -143,7 +150,17 @@ function init_nested(
     bcresid_prototype = __vec(bcresid_prototype)
     f,
     bc = if X isa AbstractVector
-        prob.f, prob.f.bc
+        if fit_parameters == true
+            l_parameters = length(prob.p)
+            vecf! = function (du, u, p, t)
+                prob.f(du, u, @view(u[(end - l_parameters + 1):end]), t)
+                du[(end - l_parameters + 1):end] .= 0
+            end
+            vecbc! = prob.f.bc
+            vecf!, vecbc!
+        else
+            prob.f, prob.f.bc
+        end
     elseif iip
         vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
         vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
@@ -169,7 +186,8 @@ function init_nested(
 
     prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
 
-    K0 = __K0_on_u0(prob.u0, stage) # Somewhat arbitrary initialization of K
+    # Somewhat arbitrary initialization of K
+    K0 = __K0_on_u0(prob, stage; fit_parameters = fit_parameters)
 
     nestprob_p = zeros(T, M + 2)
 
@@ -181,7 +199,7 @@ function init_nested(
             (K, p) -> FIRK_nlsolve(K, p, f, TU, prob.p), K0, nestprob_p)
     end
 
-    return FIRKCacheNested{iip, T, typeof(diffcache)}(
+    return FIRKCacheNested{iip, T, typeof(diffcache), fit_parameters}(
         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,
         y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, nestprob, resid₁_size,
@@ -197,11 +215,16 @@ function init_expanded(
         error("Algorithm $(alg) doesn't support adaptivity. Please choose a higher order algorithm.")
     end
     diffcache = __cache_trait(alg.jac_alg)
+    fit_parameters = haskey(prob.kwargs, :fit_parameters)
 
     iip = isinplace(prob)
 
     t₀, t₁ = prob.tspan
-    ig, T, M, Nig, X = __extract_problem_details(prob; dt, check_positive_dt = true)
+    ig, T,
+    M,
+    Nig,
+    X = __extract_problem_details(
+        prob; dt, check_positive_dt = true, fit_parameters = fit_parameters)
     mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
     mesh_dt = diff(mesh)
 
@@ -215,7 +238,7 @@ function init_expanded(
     fᵢ₂_cache = vec(zero(X))
 
     # Don't flatten this here, since we need to expand it later if needed
-    _y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p)
+    _y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
     y₀ = extend_y(_y₀, Nig + 1, stage)
     y = __alloc.(copy.(y₀.u)) # Runtime dispatch
 
@@ -240,7 +263,17 @@ function init_expanded(
     bcresid_prototype = __vec(bcresid_prototype)
     f,
     bc = if X isa AbstractVector
-        prob.f, prob.f.bc
+        if fit_parameters == true
+            l_parameters = length(prob.p)
+            vecf! = function (du, u, p, t)
+                prob.f(du, u, @view(u[(end - l_parameters + 1):end]), t)
+                du[(end - l_parameters + 1):end] .= 0
+            end
+            vecbc! = prob.f.bc
+            vecf!, vecbc!
+        else
+            prob.f, prob.f.bc
+        end
     elseif iip
         vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
         vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
@@ -266,7 +299,7 @@ function init_expanded(
 
     prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
 
-    return FIRKCacheExpand{iip, T, typeof(diffcache)}(
+    return FIRKCacheExpand{iip, T, typeof(diffcache), fit_parameters}(
         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,
         y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, resid₁_size,
@@ -299,9 +332,12 @@ function __expand_cache!(cache::FIRKCacheNested)
     return cache
 end
 
-function SciMLBase.solve!(cache::FIRKCacheExpand{iip, T}) where {iip, T}
+function SciMLBase.solve!(cache::FIRKCacheExpand{
+        iip, T, diffcache, fit_parameters}) where {iip, T, diffcache, fit_parameters}
     (abstol, adaptive, _), kwargs = __split_kwargs(; cache.kwargs...)
     info::ReturnCode.T = ReturnCode.Success
+    prob = cache.prob
+    length_u = cache.in_size
 
     # We do the first iteration outside the loop to preserve type-stability of the
     # `original` field of the solution
@@ -314,19 +350,28 @@ function SciMLBase.solve!(cache::FIRKCacheExpand{iip, T}) where {iip, T}
         end
     end
 
-    u = shrink_y(
-        [reshape(y, cache.in_size) for y in cache.y₀], length(cache.mesh), cache.stage)
+    # Parameter estimation, put the estimated parameters to sol.prob.p
+    if fit_parameters
+        length_u = cache.M - length(prob.p)
+        prob = remake(prob; p = first(cache.y₀)[(length_u + 1):end])
+        map(x -> resize!(x, length_u), cache.y₀)
+        resize!(cache.fᵢ₂_cache, length_u)
+    end
+
+    u = shrink_y([reshape(y, length_u) for y in cache.y₀], length(cache.mesh), cache.stage)
 
     interpolation = __build_interpolation(cache, u)
 
     odesol = DiffEqBase.build_solution(
-        cache.prob, cache.alg, cache.mesh, u; interp = interpolation, retcode = info)
-    return __build_solution(cache.prob, odesol, sol_nlprob)
+        prob, cache.alg, cache.mesh, u; interp = interpolation, retcode = info)
+    return __build_solution(prob, odesol, sol_nlprob)
 end
 
-function SciMLBase.solve!(cache::FIRKCacheNested{iip, T}) where {iip, T}
+function SciMLBase.solve!(cache::FIRKCacheNested{
+        iip, T, diffcache, fit_parameters}) where {iip, T, diffcache, fit_parameters}
     (abstol, adaptive, _), kwargs = __split_kwargs(; cache.kwargs...)
     info::ReturnCode.T = ReturnCode.Success
+    prob = cache.prob
 
     # We do the first iteration outside the loop to preserve type-stability of the
     # `original` field of the solution
@@ -339,13 +384,21 @@ function SciMLBase.solve!(cache::FIRKCacheNested{iip, T}) where {iip, T}
         end
     end
 
+    # Parameter estimation, put the estimated parameters to sol.prob.p
+    if fit_parameters
+        length_u = cache.M - length(prob.p)
+        prob = remake(prob; p = first(cache.y₀)[(length_u + 1):end])
+        map(x -> resize!(x, length_u), cache.y₀)
+        resize!(cache.fᵢ₂_cache, length_u)
+    end
+
     u = recursivecopy(cache.y₀)
 
     interpolation = __build_interpolation(cache, u.u)
 
     odesol = DiffEqBase.build_solution(
-        cache.prob, cache.alg, cache.mesh, u.u; interp = interpolation, retcode = info)
-    return __build_solution(cache.prob, odesol, sol_nlprob)
+        prob, cache.alg, cache.mesh, u.u; interp = interpolation, retcode = info)
+    return __build_solution(prob, odesol, sol_nlprob)
 end
 
 function __perform_firk_iteration(