Merge pull request #2563 from AayushSabharwal/as/fix-param-defaults

feat: allow parameter defaults to depend on initial values of unknowns

Author: ChrisRackauckas

docs/src/basics/FAQ.md

@@ -91,7 +91,7 @@ and the `MTKParameters` constructor for parameters. For example:
 
 ```julia
 unew = varmap_to_vars([x => 1.0, y => 2.0, z => 3.0], unknowns(sys))
-pnew = ModelingToolkit.MTKParameters(sys, [β => 3.0, c => 10.0, γ => 2.0])
+pnew = ModelingToolkit.MTKParameters(sys, [β => 3.0, c => 10.0, γ => 2.0], unew)
 ```
 
 ## How do I handle `if` statements in my symbolic forms?

src/systems/abstractsystem.jl

@@ -1708,7 +1708,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     u0, _p, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
     ps = parameters(sys)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, p)
+        p = MTKParameters(sys, p, u0)
     else
         p = _p
         p, split_idxs = split_parameters_by_type(p)
@@ -2011,7 +2011,7 @@ function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives =
         elseif p isa Vector
             p = merge(Dict(parameters(sys) .=> p), op)
         end
-        p2 = MTKParameters(sys, p)
+        p2 = MTKParameters(sys, p, Dict(unknowns(sys) .=> u0))
     end
     linres = lin_fun(u0, p2, t)
     f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u = linres

src/systems/diffeqs/abstractodesystem.jl

@@ -889,7 +889,8 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
 
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         u0, defs = get_u0(sys, trueinit, parammap; symbolic_u0)
-        p = MTKParameters(sys, parammap)
+        check_eqs_u0(eqs, dvs, u0; kwargs...)
+        p = MTKParameters(sys, parammap, trueinit)
     else
         u0, p, defs = get_u0_p(sys,
             trueinit,

src/systems/jumps/jumpsystem.jl

@@ -338,7 +338,7 @@ function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,
 
     u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, parammap)
+        p = MTKParameters(sys, parammap, u0map)
     else
         p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
     end
@@ -395,7 +395,7 @@ function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, No
 
     u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, parammap)
+        p = MTKParameters(sys, parammap, u0map)
     else
         p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
     end

src/systems/nonlinear/nonlinearsystem.jl

@@ -363,11 +363,12 @@ function process_NonlinearProblem(constructor, sys::NonlinearSystem, u0map, para
 
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         u0, defs = get_u0(sys, u0map, parammap)
-        p = MTKParameters(sys, parammap)
+        check_eqs_u0(eqs, dvs, u0; kwargs...)
+        p = MTKParameters(sys, parammap, u0map)
     else
         u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
+        check_eqs_u0(eqs, dvs, u0; kwargs...)
     end
-    check_eqs_u0(eqs, dvs, u0; kwargs...)
 
     f = constructor(sys, dvs, ps, u0; jac = jac, checkbounds = checkbounds,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,

src/systems/optimization/optimizationsystem.jl

@@ -280,7 +280,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
     if parammap isa MTKParameters
         p = parammap
     elseif has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, parammap)
+        p = MTKParameters(sys, parammap, u0map)
     else
         p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
     end
@@ -516,7 +516,7 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0map,
 
     u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, parammap)
+        p = MTKParameters(sys, parammap, u0map)
     else
         p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
     end

src/systems/parameter_buffer.jl

@@ -10,7 +10,8 @@ struct MTKParameters{T, D, C, E, N, F, G}
     dependent_update_oop::G
 end
 
-function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = false)
+function MTKParameters(
+        sys::AbstractSystem, p, u0 = Dict(); tofloat = false, use_union = false)
     ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
         get_index_cache(sys)
     else
@@ -23,21 +24,27 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
         length(p) == length(ps) || error("Invalid parameters")
         p = ps .=> p
     end
+    if p isa SciMLBase.NullParameters || isempty(p)
+        p = Dict()
+    end
+    p = todict(p)
     defs = Dict(default_toterm(unwrap(k)) => v
     for (k, v) in defaults(sys)
     if unwrap(k) in all_ps || default_toterm(unwrap(k)) in all_ps)
-    if p isa SciMLBase.NullParameters
-        p = defs
-    else
-        extra_params = Dict(unwrap(k) => v
-        for (k, v) in p if !in(unwrap(k), all_ps) && !in(default_toterm(unwrap(k)), all_ps))
-        p = merge(defs,
-            Dict(default_toterm(unwrap(k)) => v
-            for (k, v) in p if unwrap(k) in all_ps || default_toterm(unwrap(k)) in all_ps))
-        p = Dict(k => fixpoint_sub(v, extra_params)
-        for (k, v) in p if !haskey(extra_params, unwrap(k)))
+    if eltype(u0) <: Pair
+        u0 = todict(u0)
+    elseif u0 isa AbstractArray && !isempty(u0)
+        u0 = Dict(unknowns(sys) .=> vec(u0))
+    elseif u0 === nothing || isempty(u0)
+        u0 = Dict()
     end
-
+    defs = merge(defs, u0)
+    defs = merge(defs, Dict(eq.lhs => eq.rhs for eq in observed(sys)))
+    p = merge(defs, p)
+    p = merge(Dict(unwrap(k) => v for (k, v) in p),
+        Dict(default_toterm(unwrap(k)) => v for (k, v) in p))
+    p = Dict(k => fixpoint_sub(v, p)
+    for (k, v) in p if k in all_ps || default_toterm(k) in all_ps)
     for (sym, _) in p
         if istree(sym) && operation(sym) === getindex &&
            first(arguments(sym)) in all_ps

src/variables.jl

@@ -215,54 +215,6 @@ end
     throw(ArgumentError("$vars are missing from the variable map."))
 end
 
-"""
-$(SIGNATURES)
-
-Intercept the call to `process_p_u0_symbolic` and process symbolic maps of `p` and/or `u0` if the
-user has `ModelingToolkit` loaded.
-"""
-function SciMLBase.process_p_u0_symbolic(
-        prob::Union{SciMLBase.AbstractDEProblem,
-            NonlinearProblem, OptimizationProblem,
-            SciMLBase.AbstractOptimizationCache},
-        p,
-        u0)
-    # check if a symbolic remake is possible
-    if p isa Vector && !(eltype(p) <: Pair)
-        error("Parameter values must be specified as a `Dict` or `Vector{<:Pair}`")
-    end
-    if eltype(p) <: Pair
-        hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
-            throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
-                                " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
-    end
-    if eltype(u0) <: Pair
-        hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :unknowns) ||
-            throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
-                                " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to unknown variable order."))
-    end
-
-    sys = prob.f.sys
-    defs = defaults(sys)
-    ps = parameters(sys)
-    if has_split_idxs(sys) && (split_idxs = get_split_idxs(sys)) !== nothing
-        for (i, idxs) in enumerate(split_idxs)
-            defs = mergedefaults(defs, prob.p[i], ps[idxs])
-        end
-    else
-        # assemble defaults
-        defs = defaults(sys)
-        defs = mergedefaults(defs, prob.p, ps)
-    end
-    defs = mergedefaults(defs, p, ps)
-    sts = unknowns(sys)
-    defs = mergedefaults(defs, prob.u0, sts)
-    defs = mergedefaults(defs, u0, sts)
-    u0, _, defs = get_u0_p(sys, defs)
-    p = MTKParameters(sys, p)
-    return p, u0
-end
-
 struct IsHistory end
 ishistory(x) = ishistory(unwrap(x))
 ishistory(x::Symbolic) = getmetadata(x, IsHistory, false)