feat: allow parameters to be unknowns in the initialization system

AayushSabharwal · AayushSabharwal · commit b2dbf919fa7a · 2024-10-10T12:47:32.000+05:30
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
@@ -736,6 +736,15 @@ function has_observed_with_lhs(sys, sym)
     end
 end
 
+function has_parameter_dependency_with_lhs(sys, sym)
+    has_parameter_dependencies(sys) || return false
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        return any(isequal(sym), ic.dependent_pars)
+    else
+        return any(isequal(sym), [eq.lhs for eq in parameter_dependencies(sys)])
+    end
+end
+
 function _all_ts_idxs!(ts_idxs, ::NotSymbolic, sys, sym)
     if is_variable(sys, sym) || is_independent_variable(sys, sym)
         push!(ts_idxs, ContinuousTimeseries())
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -357,7 +357,9 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         analytic = nothing,
         split_idxs = nothing,
         initializeprob = nothing,
+        update_initializeprob! = nothing,
         initializeprobmap = nothing,
+        initializeprobpmap = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
@@ -459,7 +461,9 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,
         initializeprob = initializeprob,
-        initializeprobmap = initializeprobmap)
+        update_initializeprob! = update_initializeprob!,
+        initializeprobmap = initializeprobmap,
+        initializeprobpmap = initializeprobpmap)
 end
 
 """
@@ -789,6 +793,45 @@ function get_u0(
     return u0, defs
 end
 
+struct GetUpdatedMTKParameters{G, S}
+    # `getu` functor which gets parameters that are unknowns during initialization
+    getpunknowns::G
+    # `setu` functor which returns a modified MTKParameters using those parameters
+    setpunknowns::S
+end
+
+function (f::GetUpdatedMTKParameters)(prob, initializesol)
+    mtkp = copy(parameter_values(prob))
+    f.setpunknowns(mtkp, f.getpunknowns(initializesol))
+    mtkp
+end
+
+struct UpdateInitializeprob{G, S}
+    # `getu` functor which gets all values from prob
+    getvals::G
+    # `setu` functor which updates initializeprob with values
+    setvals::S
+end
+
+function (f::UpdateInitializeprob)(initializeprob, prob)
+    f.setvals(initializeprob, f.getvals(prob))
+end
+
+function get_temporary_value(p)
+    stype = symtype(unwrap(p))
+    return if stype == Real
+        zero(Float64)
+    elseif stype <: AbstractArray{Real}
+        zeros(Float64, size(p))
+    elseif stype <: Real
+        zero(stype)
+    elseif stype <: AbstractArray
+        zeros(eltype(stype), size(p))
+    else
+        error("Nonnumeric parameter $p with symtype $stype cannot be solved for during initialization")
+    end
+end
+
 function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         implicit_dae = false, du0map = nothing,
         version = nothing, tgrad = false,
@@ -829,18 +872,38 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     end
 
     if eltype(parammap) <: Pair
-        parammap = Dict(unwrap(k) => v for (k, v) in todict(parammap))
+        parammap = Dict{Any, Any}(unwrap(k) => v for (k, v) in parammap)
     elseif parammap isa AbstractArray
         if isempty(parammap)
             parammap = SciMLBase.NullParameters()
         else
-            parammap = Dict(unwrap.(parameters(sys)) .=> parammap)
+            parammap = Dict{Any, Any}(unwrap.(parameters(sys)) .=> parammap)
         end
     end
-
+    defs = defaults(sys)
+    if has_guesses(sys)
+        guesses = merge(
+            ModelingToolkit.guesses(sys), isempty(guesses) ? Dict() : todict(guesses))
+        solvablepars = [p
+                        for p in parameters(sys)
+                        if is_parameter_solvable(p, parammap, defs, guesses)]
+
+        pvarmap = if parammap === nothing || parammap == SciMLBase.NullParameters() || !(eltype(parammap) <: Pair) && isempty(parammap)
+            defs
+        else
+            merge(defs, todict(parammap))
+        end
+        setparobserved = filter(keys(pvarmap)) do var
+            has_parameter_dependency_with_lhs(sys, var)
+        end
+    else
+        solvablepars = ()
+        setparobserved = ()
+    end
     # ModelingToolkit.get_tearing_state(sys) !== nothing => Requires structural_simplify first
     if sys isa ODESystem && build_initializeprob &&
-       (((implicit_dae || !isempty(missingvars) || !isempty(setobserved)) &&
+       (((implicit_dae || !isempty(missingvars) || !isempty(solvablepars) ||
+          !isempty(setobserved) || !isempty(setparobserved)) &&
          ModelingToolkit.get_tearing_state(sys) !== nothing) ||
         !isempty(initialization_equations(sys))) && t !== nothing
         if eltype(u0map) <: Number
@@ -854,14 +917,32 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             sys, t, u0map, parammap; guesses, warn_initialize_determined,
             initialization_eqs, eval_expression, eval_module, fully_determined, check_units)
         initializeprobmap = getu(initializeprob, unknowns(sys))
+        punknowns = [p
+                     for p in all_variable_symbols(initializeprob) if is_parameter(sys, p)]
+        getpunknowns = getu(initializeprob, punknowns)
+        setpunknowns = setp(sys, punknowns)
+        initializeprobpmap = GetUpdatedMTKParameters(getpunknowns, setpunknowns)
+        reqd_syms = parameter_symbols(initializeprob)
+        update_initializeprob! = UpdateInitializeprob(
+            getu(sys, reqd_syms), setu(initializeprob, reqd_syms))
 
         zerovars = Dict(setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0)
+        if parammap isa SciMLBase.NullParameters
+            parammap = Dict()
+        end
+        for p in punknowns
+            p = unwrap(p)
+            stype = symtype(p)
+            parammap[p] = get_temporary_value(p)
+        end
         trueinit = collect(merge(zerovars, eltype(u0map) <: Pair ? todict(u0map) : u0map))
         u0map isa StaticArraysCore.StaticArray &&
             (trueinit = SVector{length(trueinit)}(trueinit))
     else
         initializeprob = nothing
+        update_initializeprob! = nothing
         initializeprobmap = nothing
+        initializeprobpmap = nothing
         trueinit = u0map
     end
 
@@ -909,7 +990,9 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         sparse = sparse, eval_expression = eval_expression,
         eval_module = eval_module,
         initializeprob = initializeprob,
+        update_initializeprob! = update_initializeprob!,
         initializeprobmap = initializeprobmap,
+        initializeprobpmap = initializeprobpmap,
         kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end
@@ -1471,10 +1554,12 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         isys = get_initializesystem(sys; initialization_eqs, check_units)
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
         isys = structural_simplify(
-            generate_initializesystem(sys; initialization_eqs, check_units); fully_determined)
+            generate_initializesystem(
+                sys; initialization_eqs, check_units, pmap = parammap); fully_determined)
     else
         isys = structural_simplify(
-            generate_initializesystem(sys; u0map, initialization_eqs, check_units); fully_determined)
+            generate_initializesystem(
+                sys; u0map, initialization_eqs, check_units, pmap = parammap); fully_determined)
     end
 
     uninit = setdiff(unknowns(sys), [unknowns(isys); getfield.(observed(isys), :lhs)])
@@ -1498,14 +1583,15 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
     parammap = parammap isa DiffEqBase.NullParameters || isempty(parammap) ?
                [get_iv(sys) => t] :
                merge(todict(parammap), Dict(get_iv(sys) => t))
+    parammap = Dict(k => v for (k, v) in parammap if v !== missing)
     if isempty(u0map)
         u0map = Dict()
     end
     if isempty(guesses)
         guesses = Dict()
     end
 
-    u0map = merge(todict(guesses), todict(u0map))
+    u0map = merge(ModelingToolkit.guesses(sys), todict(guesses), todict(u0map))
     if neqs == nunknown
         NonlinearProblem(isys, u0map, parammap; kwargs...)
     else
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
@@ -5,6 +5,7 @@ Generate `NonlinearSystem` which initializes an ODE problem from specified initi
 """
 function generate_initializesystem(sys::ODESystem;
         u0map = Dict(),
+        pmap = Dict(),
         initialization_eqs = [],
         guesses = Dict(),
         default_dd_guess = 0.0,
@@ -74,12 +75,103 @@ function generate_initializesystem(sys::ODESystem;
         end
     end
 
-    pars = [parameters(sys); get_iv(sys)] # include independent variable as pseudo-parameter
-    eqs_ics = [eqs_ics; observed(sys)]
-    return NonlinearSystem(
-        eqs_ics, vars, pars;
-        defaults = defs, parameter_dependencies = parameter_dependencies(sys),
-        checks = check_units,
-        name, kwargs...
+    # 4) process parameters as initialization unknowns
+    paramsubs = Dict()
+    if pmap isa SciMLBase.NullParameters
+        pmap = Dict()
+    end
+    pmap = todict(pmap)
+    for p in parameters(sys)
+        if is_parameter_solvable(p, pmap, defs, guesses)
+            # If either of them are `missing` the parameter is an unknown
+            # But if the parameter is passed a value, use that as an additional
+            # equation in the system
+            _val1 = get(pmap, p, nothing)
+            _val2 = get(defs, p, nothing)
+            _val3 = get(guesses, p, nothing)
+            varp = tovar(p)
+            paramsubs[p] = varp
+            # Has a default of `missing`, and (either an equation using the value passed to `ODEProblem` or a guess)
+            if _val2 === missing
+                if _val1 !== nothing && _val1 !== missing
+                    push!(eqs_ics, varp ~ _val1)
+                    push!(defs, varp => _val1)
+                elseif _val3 !== nothing
+                    # assuming an equation exists (either via algebraic equations or initialization_eqs)
+                    push!(defs, varp => _val3)
+                elseif check_defguess
+                    error("Invalid setup: parameter $(p) has no default value, initial value, or guess")
+                end
+                # `missing` passed to `ODEProblem`, and (either an equation using default or a guess)
+            elseif _val1 === missing
+                if _val2 !== nothing && _val2 !== missing
+                    push!(eqs_ics, varp ~ _val2)
+                    push!(defs, varp => _val2)
+                elseif _val3 !== nothing
+                    push!(defs, varp => _val3)
+                elseif check_defguess
+                    error("Invalid setup: parameter $(p) has no default value, initial value, or guess")
+                end
+                # given a symbolic value to ODEProblem
+            elseif symbolic_type(_val1) != NotSymbolic()
+                push!(eqs_ics, varp ~ _val1)
+                push!(defs, varp => _val3)
+                # No value passed to `ODEProblem`, but a default and a guess are present
+                # _val2 !== missing is implied by it falling this far in the elseif chain
+            elseif _val1 === nothing && _val2 !== nothing
+                push!(eqs_ics, varp ~ _val2)
+                push!(defs, varp => _val3)
+            else
+                # _val1 !== missing and _val1 !== nothing, so a value was provided to ODEProblem
+                # This would mean `is_parameter_solvable` returned `false`, so we never end up
+                # here
+                error("This should never be reached")
+            end
+        end
+    end
+
+    # 5) parameter dependencies become equations, their LHS become unknowns
+    for eq in parameter_dependencies(sys)
+        varp = tovar(eq.lhs)
+        paramsubs[eq.lhs] = varp
+        push!(eqs_ics, eq)
+        guessval = get(guesses, eq.lhs, eq.rhs)
+        push!(defs, varp => guessval)
+    end
+
+    # 6) handle values provided for dependent parameters similar to values for observed variables
+    for (k, v) in merge(defaults(sys), pmap)
+        if is_variable_floatingpoint(k) && has_parameter_dependency_with_lhs(sys, k)
+            push!(eqs_ics, paramsubs[k] ~ v)
+        end
+    end
+
+    # parameters do not include ones that became initialization unknowns
+    pars = vcat(
+        [get_iv(sys)], # include independent variable as pseudo-parameter
+        [p for p in parameters(sys) if !haskey(paramsubs, p)]
     )
+
+    eqs_ics = Symbolics.substitute.([eqs_ics; observed(sys)], (paramsubs,))
+    vars = [vars; collect(values(paramsubs))]
+    for k in keys(defs)
+        defs[k] = substitute(defs[k], paramsubs)
+    end
+    return NonlinearSystem(eqs_ics,
+        vars,
+        pars;
+        defaults = defs,
+        checks = check_units,
+        name,
+        kwargs...)
+end
+
+function is_parameter_solvable(p, pmap, defs, guesses)
+    _val1 = pmap isa AbstractDict ? get(pmap, p, nothing) : nothing
+    _val2 = get(defs, p, nothing)
+    _val3 = get(guesses, p, nothing)
+    # either (missing is a default or was passed to the ODEProblem) or (nothing was passed to
+    # the ODEProblem and it has a default and a guess)
+    return ((_val1 === missing || _val2 === missing) ||
+           (_val1 === nothing && _val2 !== nothing)) && _val3 !== nothing
 end
diff --git a/src/systems/parameter_buffer.jl b/src/systems/parameter_buffer.jl
@@ -522,6 +522,9 @@ function _remake_buffer(indp, oldbuf::MTKParameters, idxs, vals; validate = true
     ic = get_index_cache(indp_to_system(indp))
     for (idx, val) in zip(idxs, vals)
         sym = nothing
+        if val === missing
+            val = get_temporary_value(idx)
+        end
         if symbolic_type(idx) == ScalarSymbolic()
             sym = idx
             idx = parameter_index(ic, sym)
diff --git a/src/utils.jl b/src/utils.jl
@@ -494,8 +494,8 @@ function collect_var!(unknowns, parameters, var, iv)
         push!(unknowns, var)
     end
     # Add also any parameters that appear only as defaults in the var
-    if hasdefault(var)
-        collect_vars!(unknowns, parameters, getdefault(var), iv)
+    if hasdefault(var) && (def = getdefault(var)) !== missing
+        collect_vars!(unknowns, parameters, def, iv)
     end
     return nothing
 end
