Merge pull request #1487 from ValentinKaisermayer/patch-process_DEProblem

ChrisRackauckas · web-flow · commit dc60fe8cdf25 · 2022-03-19T17:38:50.000-04:00
Patch process_DEproblem
diff --git a/src/structural_transformation/codegen.jl b/src/structural_transformation/codegen.jl
@@ -495,6 +495,8 @@ function ODAEProblem{iip}(
     ps = parameters(sys)
     defs = defaults(sys)
 
+    defs = ModelingToolkit.mergedefaults(defs,parammap,ps)
+    defs = ModelingToolkit.mergedefaults(defs,u0map,dvs)
     u0 = ModelingToolkit.varmap_to_vars(u0map, dvs; defaults=defs)
     p = ModelingToolkit.varmap_to_vars(parammap, ps; defaults=defs)
 
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -557,36 +557,22 @@ function process_DEProblem(constructor, sys::AbstractODESystem,u0map,parammap;
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
-    defs = defaults(sys)
     iv = get_iv(sys)
-    if parammap isa Dict
-        u0defs = merge(parammap, defs)
-    elseif eltype(parammap) <: Pair
-        u0defs = merge(Dict(parammap), defs)
-    elseif eltype(parammap) <: Number
-        u0defs = merge(Dict(zip(ps, parammap)), defs)
-    else
-        u0defs = defs
-    end
-    if u0map isa Dict
-        pdefs = merge(u0map, defs)
-    elseif eltype(u0map) <: Pair
-        pdefs = merge(Dict(u0map), defs)
-    elseif eltype(u0map) <: Number
-        pdefs = merge(Dict(zip(dvs, u0map)), defs)
-    else
-        pdefs = defs
-    end
-
-    u0 = varmap_to_vars(u0map,dvs; defaults=u0defs)
+    
+    defs = defaults(sys)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs)
+    
+    u0 = varmap_to_vars(u0map,dvs; defaults=defs)
+    p = varmap_to_vars(parammap,ps; defaults=defs)
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
-        du0 = varmap_to_vars(du0map, ddvs; defaults=defaults, toterm=identity)
+        defs = mergedefaults(defs,du0map, ddvs)
+        du0 = varmap_to_vars(du0map,ddvs; defaults=defs, toterm=identity)
     else
         du0 = nothing
         ddvs = nothing
     end
-    p = varmap_to_vars(parammap,ps; defaults=pdefs)
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)
 
@@ -691,7 +677,7 @@ merge_cb(x, y) = CallbackSet(x, y)
 
 """
 ```julia
-function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
+function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem,du0map,u0map,tspan,
                                     parammap=DiffEqBase.NullParameters();
                                     version = nothing, tgrad=false,
                                     jac = false,
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
@@ -174,33 +174,17 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     ps = parameters(sys)
     eqs = equations(sys)
     eqs = linearize_eqs(sys, eqs)
-    defs = defaults(sys)
     iv = get_iv(sys)
-
-    if parammap isa Dict
-        u0defs = merge(parammap, defs)
-    elseif eltype(parammap) <: Pair
-        u0defs = merge(Dict(parammap), defs)
-    elseif eltype(parammap) <: Number
-        u0defs = merge(Dict(zip(ps, parammap)), defs)
-    else
-        u0defs = defs
-    end
-    if u0map isa Dict
-        pdefs = merge(u0map, defs)
-    elseif eltype(u0map) <: Pair
-        pdefs = merge(Dict(u0map), defs)
-    elseif eltype(u0map) <: Number
-        pdefs = merge(Dict(zip(dvs, u0map)), defs)
-    else
-        pdefs = defs
-    end
-
-    u0 = varmap_to_vars(u0map,dvs; defaults=u0defs)
+    
+    defs = defaults(sys)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs)
+    
+    u0 = varmap_to_vars(u0map,dvs; defaults=defs)
+    p = varmap_to_vars(parammap,ps; defaults=defs)
 
     rhss = [eq.rhs for eq in eqs]
     u = dvs
-    p = varmap_to_vars(parammap,ps; defaults=pdefs)
 
     f_gen = generate_function(sys; expression=Val{eval_expression}, expression_module=eval_module)
     f_oop, _ = (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen)
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
@@ -220,9 +220,17 @@ dprob = DiscreteProblem(js, u₀map, tspan, parammap)
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,Nothing},
                                     parammap=DiffEqBase.NullParameters(); checkbounds=false, kwargs...)
+    
+    dvs = states(sys)
+    ps = parameters(sys)
+    
     defs = defaults(sys)
-    u0 = varmap_to_vars(u0map, states(sys); defaults=defs)
-    p  = varmap_to_vars(parammap, parameters(sys); defaults=defs)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs) 
+    
+    u0 = varmap_to_vars(u0map,dvs; defaults=defs)
+    p = varmap_to_vars(parammap,ps; defaults=defs)
+        
     f  = DiffEqBase.DISCRETE_INPLACE_DEFAULT
     
     # just taken from abstractodesystem.jl for ODEFunction def
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
@@ -264,7 +264,11 @@ function process_NonlinearProblem(constructor, sys::NonlinearSystem,u0map,paramm
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
+    
     defs = defaults(sys)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs)
+    
     u0 = varmap_to_vars(u0map,dvs; defaults=defs)
     p = varmap_to_vars(parammap,ps; defaults=defs)
 
diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
@@ -129,9 +129,9 @@ DiffEqBase.OptimizationProblem(sys::OptimizationSystem,args...;kwargs...) =
 
 """
 ```julia
-function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
+function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,u0map,
                                           parammap=DiffEqBase.NullParameters();
-                                          u0=nothing, lb=nothing, ub=nothing,
+                                          lb=nothing, ub=nothing,
                                           grad = false,
                                           hess = false, sparse = false,
                                           checkbounds = false,
@@ -142,7 +142,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
 Generates an OptimizationProblem from an OptimizationSystem and allows for automatically
 symbolically calculating numerical enhancements.
 """
-function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0,
+function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                           parammap=DiffEqBase.NullParameters();
                                           lb=nothing, ub=nothing,
                                           grad = false,
@@ -177,7 +177,10 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0,
     _f = DiffEqBase.OptimizationFunction{iip,AutoModelingToolkit,typeof(f),typeof(_grad),typeof(_hess),Nothing,Nothing,Nothing,Nothing}(f,AutoModelingToolkit(),_grad,_hess,nothing,nothing,nothing,nothing)
 
     defs = defaults(sys)
-    u0 = varmap_to_vars(u0,dvs; defaults=defs)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs)
+
+    u0 = varmap_to_vars(u0map,dvs; defaults=defs)
     p = varmap_to_vars(parammap,ps; defaults=defs)
     lb = varmap_to_vars(lb,dvs; check=false)
     ub = varmap_to_vars(ub,dvs; check=false)
@@ -233,7 +236,10 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
     end
 
     defs = defaults(sys)
-    u0 = varmap_to_vars(u0,dvs; defaults=defs)
+    defs = mergedefaults(defs,parammap,ps)
+    defs = mergedefaults(defs,u0map,dvs)
+
+    u0 = varmap_to_vars(u0map,dvs; defaults=defs)
     p = varmap_to_vars(parammap,ps; defaults=defs)
     lb = varmap_to_vars(lb,dvs)
     ub = varmap_to_vars(ub,dvs)
diff --git a/src/utils.jl b/src/utils.jl
@@ -459,3 +459,15 @@ function get_substitutions_and_solved_states(sys; no_postprocess=false)
     end
     return pre, sol_states
 end
+
+function mergedefaults(defaults, varmap, vars)
+    defs = if varmap isa Dict
+        merge(defaults, varmap)
+    elseif eltype(varmap) <: Pair
+        merge(defaults, Dict(varmap))
+    elseif eltype(varmap) <: Number
+        merge(defaults, Dict(zip(vars, varmap)))
+    else
+        defaults
+    end
+end
diff --git a/src/variables.jl b/src/variables.jl
@@ -73,23 +73,13 @@ function _varmap_to_vars(varmap::Dict, varlist; defaults=Dict(), check=false, to
     varmap = Dict(toterm(value(k))=>value(varmap[k]) for k in keys(varmap))
     # resolve symbolic parameter expressions
     for (p, v) in pairs(varmap)
-        val = varmap[p] = fixpoint_sub(v, varmap)
+        varmap[p] = fixpoint_sub(v, varmap)
     end
-    vs = values(varmap)
-    T′ = eltype(vs)
-    if Base.isconcretetype(T′)
-        T = T′
-    else
-        T = foldl((t, elem)->promote_type(t, eltype(elem)), vs; init=typeof(first(vs)))
-    end
-    out = Vector{T}(undef, length(varlist))
+    
     missingvars = setdiff(varlist, keys(varmap))
     check && (isempty(missingvars) || throw_missingvars(missingvars))
 
-    for (i, var) in enumerate(varlist)
-        out[i] = varmap[var]
-    end
-    out
+    out = [varmap[var] for var in varlist]
 end
 
 @noinline throw_missingvars(vars) = throw(ArgumentError("$vars are missing from the variable map."))
diff --git a/test/odesystem.jl b/test/odesystem.jl
@@ -591,4 +591,59 @@ eqs[end] = D(D(z)) ~ α*x - β*y
     sol = solve(prob, Euler(); dt=0.1)
     
     @test c[1] == length(sol)
+end
+                                                                      
+let
+    @parameters t
+    D = Differential(t)
+    @variables x[1:2](t) = zeros(2) 
+    @variables y(t) = 0
+    @parameters k = 1
+    eqs= [
+        D(x[1]) ~ x[2]
+        D(x[2]) ~ -x[1] - 0.5 * x[2] + k
+        y ~ 0.9 * x[1] + x[2]
+    ]
+    @named sys = ODESystem(eqs, t, vcat(x, [y]), [k])
+    sys = structural_simplify(sys)
+
+    u0 = [0.5, 0]
+    du0 = 0 .* copy(u0)
+    prob = DAEProblem(sys, du0, u0, (0, 50))
+    @test prob.u0 ≈ u0
+    @test prob.du0 ≈ du0
+    @test prob.p ≈ [1]
+    sol = solve(prob, IDA())
+    @test isapprox(sol[x[1]][end], 1, atol=1e-3)
+
+    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0], Pair[x[1] => 0.5], (0, 50))
+    @test prob.u0 ≈ [0.5, 0]
+    @test prob.du0 ≈ [0, 0]
+    @test prob.p ≈ [1]
+    sol = solve(prob, IDA())
+    @test isapprox(sol[x[1]][end], 1, atol=1e-3)
+
+    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0], Pair[x[1] => 0.5], (0, 50), [k => 2])
+    @test prob.u0 ≈ [0.5, 0]
+    @test prob.du0 ≈ [0, 0]
+    @test prob.p ≈ [2]
+    sol = solve(prob, IDA())
+    @test isapprox(sol[x[1]][end], 2, atol=1e-3)
+
+    # no initial conditions for D(x[1]) and D(x[2]) provided
+    @test_throws ArgumentError prob = DAEProblem(sys, Pair[], Pair[], (0, 50))    
+end
+
+#issue 1475 (mixed numeric type)
+let
+    @parameters k1 k2
+    @variables t, A(t)
+    D = Differential(t)
+    eqs = [D(A) ~ -k1*k2*A]
+    @named sys = ODESystem(eqs,t)
+    u0map = [A => 1.0]
+    pmap = (k1 => 1.0, k2 => 1)
+    tspan = (0.0,1.0)
+    prob = ODEProblem(sys, u0map, tspan, pmap)
+    @test prob.p === Tuple([(Dict(pmap))[k] for k in values(parameters(sys))])
 end
