Merge pull request #3648 from AayushSabharwal/as/problem-refactor-2

refactor: change problem constructors to `XProblem(sys, op[, tspan])`

Author: AayushSabharwal

ext/MTKCasADiDynamicOptExt.jl

@@ -76,8 +76,10 @@ function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap;
         steps = nothing,
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
-    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
+    _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) :
+             merge(Dict(u0map), Dict(guesses))
+    pmap = MTK.to_varmap(pmap, parameters(sys))
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap);
         t = tspan !== nothing ? tspan[1] : tspan, output_type = MX, kwargs...)
 
     pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))

ext/MTKInfiniteOptExt.jl

@@ -59,8 +59,10 @@ function MTK.JuMPDynamicOptProblem(sys::System, u0map, tspan, pmap;
         steps = nothing,
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
-    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
+    _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) :
+             merge(Dict(u0map), Dict(guesses))
+    pmap = MTK.to_varmap(pmap, parameters(sys))
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap);
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))
@@ -86,8 +88,10 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap;
         steps = nothing,
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
-    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
+    _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) :
+             merge(Dict(u0map), Dict(guesses))
+    pmap = MTK.to_varmap(pmap, parameters(sys))
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap);
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))

src/ModelingToolkit.jl

@@ -234,7 +234,7 @@ PrecompileTools.@compile_workload begin
     using ModelingToolkit
     @variables x(ModelingToolkit.t_nounits)
     @named sys = System([ModelingToolkit.D_nounits(x) ~ -x], ModelingToolkit.t_nounits)
-    prob = ODEProblem(mtkcompile(sys), [x => 30.0], (0, 100), [], jac = true)
+    prob = ODEProblem(mtkcompile(sys), [x => 30.0], (0, 100), jac = true)
     @mtkmodel __testmod__ begin
         @constants begin
             c = 1.0

src/deprecations.jl

@@ -8,4 +8,3 @@ macro mtkbuild(exprs...)
         @mtkcompile $(exprs...)
     end |> esc
 end
-

src/linearization.jl

@@ -73,7 +73,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     end
 
     prob = ODEProblem{true, SciMLBase.FullSpecialize}(
-        sys, op, (nothing, nothing), p; allow_incomplete = true,
+        sys, merge(op, anydict(p)), (nothing, nothing); allow_incomplete = true,
         algebraic_only = true, guesses)
     u0 = state_values(prob)
 

src/problems/bvproblem.jl

@@ -44,7 +44,7 @@ If the `System` has algebraic equations, like `x(t)^2 + y(t)^2`, the resulting
 `BVProblem` must be solved using BVDAE solvers, such as Ascher.
 """
 @fallback_iip_specialize function SciMLBase.BVProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         check_compatibility = true, cse = true,
         checkbounds = false, eval_expression = false, eval_module = @__MODULE__,
         expression = Val{false}, guesses = Dict(), callback = nothing,
@@ -55,22 +55,23 @@ If the `System` has algebraic equations, like `x(t)^2 + y(t)^2`, the resulting
 
     # Systems without algebraic equations should use both fixed values + guesses
     # for initialization.
-    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
+    _op = has_alg_eqs(sys) ? op : merge(Dict(op), Dict(guesses))
 
     fode, u0, p = process_SciMLProblem(
-        ODEFunction{iip, spec}, sys, _u0map, parammap; guesses,
+        ODEFunction{iip, spec}, sys, _op; guesses,
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility = false, cse,
         checkbounds, time_dependent_init = false, expression, kwargs...)
 
     dvs = unknowns(sys)
     stidxmap = Dict([v => i for (i, v) in enumerate(dvs)])
-    u0_idxs = has_alg_eqs(sys) ? collect(1:length(dvs)) : [stidxmap[k] for (k, v) in u0map]
+    u0_idxs = has_alg_eqs(sys) ? collect(1:length(dvs)) :
+              [stidxmap[k] for (k, v) in op if haskey(stidxmap, k)]
     fbc = generate_boundary_conditions(
         sys, u0, u0_idxs, tspan[1]; expression = Val{false},
         wrap_gfw = Val{true}, cse, checkbounds)
 
-    if (length(constraints(sys)) + length(u0map) > length(dvs))
-        @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
+    if (length(constraints(sys)) + length(op) > length(dvs))
+        @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
     end
 
     kwargs = process_kwargs(sys; expression, kwargs...)

src/problems/daeproblem.jl

@@ -62,15 +62,15 @@
 end
 
 @fallback_iip_specialize function SciMLBase.DAEProblem{iip, spec}(
-        sys::System, du0map, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         callback = nothing, check_length = true, eval_expression = false,
         eval_module = @__MODULE__, check_compatibility = true,
         expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, DAEProblem)
     check_compatibility && check_compatible_system(DAEProblem, sys)
 
-    f, du0, u0, p = process_SciMLProblem(DAEFunction{iip, spec}, sys, u0map, parammap;
-        du0map, t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
+    f, du0, u0, p = process_SciMLProblem(DAEFunction{iip, spec}, sys, op;
+        t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, check_compatibility, implicit_dae = true, expression, kwargs...)
 
     kwargs = process_kwargs(sys; expression, callback, eval_expression, eval_module,

src/problems/ddeproblem.jl

@@ -42,14 +42,14 @@
 end
 
 @fallback_iip_specialize function SciMLBase.DDEProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         callback = nothing, check_length = true, cse = true, checkbounds = false,
         eval_expression = false, eval_module = @__MODULE__, check_compatibility = true,
         u0_constructor = identity, expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, DDEProblem)
     check_compatibility && check_compatible_system(DDEProblem, sys)
 
-    f, u0, p = process_SciMLProblem(DDEFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(DDEFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, cse, checkbounds,
         eval_expression, eval_module, check_compatibility, symbolic_u0 = true,
         expression, u0_constructor, kwargs...)

src/problems/discreteproblem.jl

@@ -39,15 +39,15 @@
 end
 
 @fallback_iip_specialize function SciMLBase.DiscreteProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         check_compatibility = true, expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, DiscreteProblem)
     check_compatibility && check_compatible_system(DiscreteProblem, sys)
 
     dvs = unknowns(sys)
     u0map = to_varmap(u0map, dvs)
     add_toterms!(u0map; replace = true)
-    f, u0, p = process_SciMLProblem(DiscreteFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(DiscreteFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility, expression,
         kwargs...)
 

src/problems/implicitdiscreteproblem.jl

@@ -43,16 +43,16 @@
 end
 
 @fallback_iip_specialize function SciMLBase.ImplicitDiscreteProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         check_compatibility = true, expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, ImplicitDiscreteProblem)
     check_compatibility && check_compatible_system(ImplicitDiscreteProblem, sys)
 
     dvs = unknowns(sys)
-    u0map = to_varmap(u0map, dvs)
-    add_toterms!(u0map; replace = true)
+    op = to_varmap(op, dvs)
+    add_toterms!(op; replace = true)
     f, u0, p = process_SciMLProblem(
-        ImplicitDiscreteFunction{iip, spec}, sys, u0map, parammap;
+        ImplicitDiscreteFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility,
         expression, kwargs...)