Collect default values of variables into defaults

Author: YingboMa

src/systems/diffeqs/odesystem.jl

@@ -101,7 +101,6 @@ function ODESystem(
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
                    connection_type=nothing,
                   )
-    
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
 
     iv′ = value(scalarize(iv))
@@ -115,6 +114,13 @@ function ODESystem(
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
+    iv′ = value(scalarize(iv))
+    dvs′ = value.(scalarize(dvs))
+    ps′ = value.(scalarize(ps))
+
+    collect_defaults!(defaults, dvs′)
+    collect_defaults!(defaults, ps′)
+
     tgrad = RefValue(Vector{Num}(undef, 0))
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
     ctrl_jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
@@ -363,4 +369,4 @@ function convert_system(::Type{<:ODESystem}, sys, t; name=nameof(sys))
     neweqs = map(sub, equations(sys))
     defs = Dict(sub(k) => sub(v) for (k, v) in defaults(sys))
     return ODESystem(neweqs, t, newsts, parameters(sys); defaults=defs, name=name)
-end
+end

src/utils.jl

@@ -116,7 +116,7 @@ function check_variables(dvs, iv)
         isequal(iv, dv) && throw(ArgumentError("Independent variable $iv not allowed in dependent variables."))
         isequal(iv, iv_from_nested_derivative(dv)) || throw(ArgumentError("Variable $dv is not a function of independent variable $iv."))
     end
-end 
+end
 
 "Get all the independent variables with respect to which differentials are taken."
 function collect_differentials(eqs)
@@ -155,3 +155,14 @@ end
 iv_from_nested_derivative(x::Term) = operation(x) isa Differential ? iv_from_nested_derivative(arguments(x)[1]) : arguments(x)[1]
 iv_from_nested_derivative(x::Sym) = x
 iv_from_nested_derivative(x) = missing
+
+hasdefault(v) = hasmetadata(v, Symbolics.VariableDefaultValue)
+getdefault(v) = value(getmetadata(v, Symbolics.VariableDefaultValue))
+setdefault(v, val) = val === nothing ? v : setmetadata(v, Symbolics.VariableDefaultValue, value(val))
+
+function collect_defaults!(defs, vars)
+    for v in vars; (haskey(defs, v) || !hasdefault(v)) && continue
+        defs[v] = getdefault(v)
+    end
+    return defs
+end