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

AayushSabharwal · AayushSabharwal · commit d78b814c5f9f · 2025-05-23T15:05:49.000+05:30
diff --git a/src/problems/bvproblem.jl b/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...)
diff --git a/src/problems/daeproblem.jl b/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,
diff --git a/src/problems/ddeproblem.jl b/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...)
diff --git a/src/problems/discreteproblem.jl b/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...)
 
diff --git a/src/problems/implicitdiscreteproblem.jl b/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...)
 
diff --git a/src/problems/initializationproblem.jl b/src/problems/initializationproblem.jl
@@ -119,7 +119,7 @@ initial conditions for the given DAE.
 
     filter_missing_values!(u0map)
     filter_missing_values!(parammap)
-    u0map = merge(ModelingToolkit.guesses(sys), todict(guesses), u0map)
+    op = merge(ModelingToolkit.guesses(sys), todict(guesses), u0map, parammap)
 
     TProb = if neqs == nunknown && isempty(unassigned_vars)
         if use_scc && neqs > 0
@@ -135,8 +135,7 @@ initial conditions for the given DAE.
     else
         NonlinearLeastSquaresProblem
     end
-    TProb{iip}(isys, u0map, parammap; kwargs...,
-        build_initializeprob = false, is_initializeprob = true)
+    TProb{iip}(isys, op; kwargs..., build_initializeprob = false, is_initializeprob = true)
 end
 
 const INCOMPLETE_INITIALIZATION_MESSAGE = """
diff --git a/src/problems/intervalnonlinearproblem.jl b/src/problems/intervalnonlinearproblem.jl
@@ -33,7 +33,9 @@ function SciMLBase.IntervalNonlinearProblem(
     check_compatibility && check_compatible_system(IntervalNonlinearProblem, sys)
 
     u0map = unknowns(sys) .=> uspan[1]
-    f, u0, p = process_SciMLProblem(IntervalNonlinearFunction, sys, u0map, parammap;
+    op = anydict([unknowns(sys)[1] => uspan[1]])
+    merge!(op, to_varmap(parammap, parameters(sys)))
+    f, u0, p = process_SciMLProblem(IntervalNonlinearFunction, sys, op;
         check_compatibility, expression, kwargs...)
 
     kwargs = process_kwargs(sys; kwargs...)
diff --git a/src/problems/jumpproblem.jl b/src/problems/jumpproblem.jl
@@ -1,5 +1,5 @@
 @fallback_iip_specialize function JumpProcesses.JumpProblem{iip, spec}(
-        sys::System, u0map, tspan::Union{Tuple, Nothing}, pmap = SciMLBase.NullParameters();
+        sys::System, op, tspan::Union{Tuple, Nothing};
         check_compatibility = true, eval_expression = false, eval_module = @__MODULE__,
         checkbounds = false, cse = true, aggregator = JumpProcesses.NullAggregator(),
         callback = nothing, rng = nothing, kwargs...) where {iip, spec}
@@ -13,16 +13,16 @@
     if (has_vrjs || has_eqs)
         if has_eqs && has_noise
             prob = SDEProblem{iip, spec}(
-                sys, u0map, tspan, pmap; check_compatibility = false,
+                sys, op, tspan; check_compatibility = false,
                 build_initializeprob = false, checkbounds, cse, check_length = false,
                 kwargs...)
         elseif has_eqs
             prob = ODEProblem{iip, spec}(
-                sys, u0map, tspan, pmap; check_compatibility = false,
+                sys, op, tspan; check_compatibility = false,
                 build_initializeprob = false, checkbounds, cse, check_length = false,
                 kwargs...)
         else
-            _, u0, p = process_SciMLProblem(EmptySciMLFunction{iip}, sys, u0map, pmap;
+            _, u0, p = process_SciMLProblem(EmptySciMLFunction{iip}, sys, op;
                 t = tspan === nothing ? nothing : tspan[1], tofloat = false,
                 check_length = false, build_initializeprob = false)
             observedfun = ObservedFunctionCache(sys; eval_expression, eval_module,
@@ -32,7 +32,7 @@
             prob = ODEProblem{true}(df, u0, tspan, p; kwargs...)
         end
     else
-        _f, u0, p = process_SciMLProblem(EmptySciMLFunction{iip}, sys, u0map, pmap;
+        _f, u0, p = process_SciMLProblem(EmptySciMLFunction{iip}, sys, op;
             t = tspan === nothing ? nothing : tspan[1], tofloat = false, check_length = false, build_initializeprob = false, cse)
         f = DiffEqBase.DISCRETE_INPLACE_DEFAULT
 
diff --git a/src/problems/nonlinearproblem.jl b/src/problems/nonlinearproblem.jl
@@ -57,15 +57,15 @@
 end
 
 @fallback_iip_specialize function SciMLBase.NonlinearProblem{iip, spec}(
-        sys::System, u0map, parammap = SciMLBase.NullParameters(); expression = Val{false},
+        sys::System, op; expression = Val{false},
         check_length = true, check_compatibility = true, kwargs...) where {iip, spec}
     check_complete(sys, NonlinearProblem)
     if is_time_dependent(sys)
         sys = NonlinearSystem(sys)
     end
     check_compatibility && check_compatible_system(NonlinearProblem, sys)
 
-    f, u0, p = process_SciMLProblem(NonlinearFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(NonlinearFunction{iip, spec}, sys, op;
         check_length, check_compatibility, expression, kwargs...)
 
     kwargs = process_kwargs(sys; kwargs...)
@@ -75,12 +75,12 @@ end
 end
 
 @fallback_iip_specialize function SciMLBase.NonlinearLeastSquaresProblem{iip, spec}(
-        sys::System, u0map, parammap = DiffEqBase.NullParameters(); check_length = false,
+        sys::System, op; check_length = false,
         check_compatibility = true, expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, NonlinearLeastSquaresProblem)
     check_compatibility && check_compatible_system(NonlinearLeastSquaresProblem, sys)
 
-    f, u0, p = process_SciMLProblem(NonlinearFunction{iip}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(NonlinearFunction{iip}, sys, op;
         check_length, expression, kwargs...)
 
     kwargs = process_kwargs(sys; kwargs...)
diff --git a/src/problems/odeproblem.jl b/src/problems/odeproblem.jl
@@ -65,14 +65,14 @@
 end
 
 @fallback_iip_specialize function SciMLBase.ODEProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         callback = nothing, check_length = true, eval_expression = false,
         expression = Val{false}, eval_module = @__MODULE__, check_compatibility = true,
         kwargs...) where {iip, spec}
     check_complete(sys, ODEProblem)
     check_compatibility && check_compatible_system(ODEProblem, sys)
 
-    f, u0, p = process_SciMLProblem(ODEFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(ODEFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, expression, check_compatibility, kwargs...)
 
@@ -98,12 +98,12 @@ Generates an SteadyStateProblem from a `System` of ODEs and allows for automatic
 symbolically calculating numerical enhancements.
 """
 @fallback_iip_specialize function DiffEqBase.SteadyStateProblem{iip, spec}(
-        sys::System, u0map, parammap; check_length = true, check_compatibility = true,
+        sys::System, op; check_length = true, check_compatibility = true,
         expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, SteadyStateProblem)
     check_compatibility && check_compatible_system(SteadyStateProblem, sys)
 
-    f, u0, p = process_SciMLProblem(ODEFunction{iip}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(ODEFunction{iip}, sys, op;
         steady_state = true, check_length, check_compatibility, expression,
         force_initialization_time_independent = true, kwargs...)
 
diff --git a/src/problems/optimizationproblem.jl b/src/problems/optimizationproblem.jl
@@ -87,13 +87,13 @@ function SciMLBase.OptimizationProblem(sys::System, args...; kwargs...)
 end
 
 function SciMLBase.OptimizationProblem{iip}(
-        sys::System, u0map, parammap = SciMLBase.NullParameters(); lb = nothing,
+        sys::System, op; lb = nothing,
         ub = nothing, check_compatibility = true, expression = Val{false},
         kwargs...) where {iip}
     check_complete(sys, OptimizationProblem)
     check_compatibility && check_compatible_system(OptimizationProblem, sys)
 
-    f, u0, p = process_SciMLProblem(OptimizationFunction{iip}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(OptimizationFunction{iip}, sys, op;
         check_compatibility, tofloat = false, check_length = false, expression, kwargs...)
 
     dvs = unknowns(sys)
@@ -114,7 +114,7 @@ function SciMLBase.OptimizationProblem{iip}(
     end
 
     ps = parameters(sys)
-    defs = merge(defaults(sys), to_varmap(parammap, ps), to_varmap(u0map, dvs))
+    defs = merge(defaults(sys), to_varmap(op, dvs))
     lb = varmap_to_vars(dvs .=> lb, dvs; defaults = defs, tofloat = false)
     ub = varmap_to_vars(dvs .=> ub, dvs; defaults = defs, tofloat = false)
 
diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl
@@ -69,9 +69,8 @@ function SciMLBase.SCCNonlinearProblem(sys::System, args...; kwargs...)
     SCCNonlinearProblem{true}(sys, args...; kwargs...)
 end
 
-function SciMLBase.SCCNonlinearProblem{iip}(sys::System, u0map,
-        parammap = SciMLBase.NullParameters(); eval_expression = false, eval_module = @__MODULE__,
-        cse = true, kwargs...) where {iip}
+function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = false,
+        eval_module = @__MODULE__, cse = true, kwargs...) where {iip}
     if !iscomplete(sys) || get_tearing_state(sys) === nothing
         error("A simplified `System` is required. Call `mtkcompile` on the system before creating an `SCCNonlinearProblem`.")
     end
@@ -85,7 +84,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, u0map,
 
     if length(var_sccs) == 1
         return NonlinearProblem{iip}(
-            sys, u0map, parammap; eval_expression, eval_module, kwargs...)
+            sys, op; eval_expression, eval_module, kwargs...)
     end
 
     condensed_graph = MatchedCondensationGraph(
@@ -102,7 +101,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, u0map,
     obs = observed(sys)
 
     _, u0, p = process_SciMLProblem(
-        EmptySciMLFunction{iip}, sys, u0map, parammap; eval_expression, eval_module, kwargs...)
+        EmptySciMLFunction{iip}, sys, op; eval_expression, eval_module, kwargs...)
 
     explicitfuns = []
     nlfuns = []
diff --git a/src/problems/sddeproblem.jl b/src/problems/sddeproblem.jl
@@ -43,15 +43,15 @@
 end
 
 @fallback_iip_specialize function SciMLBase.SDDEProblem{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, sparse = false, sparsenoise = sparse,
         expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, SDDEProblem)
     check_compatibility && check_compatible_system(SDDEProblem, sys)
 
-    f, u0, p = process_SciMLProblem(SDDEFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(SDDEFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, cse, checkbounds,
         eval_expression, eval_module, check_compatibility, sparse, symbolic_u0 = true,
         expression, u0_constructor, kwargs...)
diff --git a/src/problems/sdeproblem.jl b/src/problems/sdeproblem.jl
@@ -67,14 +67,14 @@
 end
 
 @fallback_iip_specialize function SciMLBase.SDEProblem{iip, spec}(
-        sys::System, u0map, tspan, parammap = SciMLBase.NullParameters();
+        sys::System, op, tspan;
         callback = nothing, check_length = true, eval_expression = false,
         eval_module = @__MODULE__, check_compatibility = true, sparse = false,
         sparsenoise = sparse, expression = Val{false}, kwargs...) where {iip, spec}
     check_complete(sys, SDEProblem)
     check_compatibility && check_compatible_system(SDEProblem, sys)
 
-    f, u0, p = process_SciMLProblem(SDEFunction{iip, spec}, sys, u0map, parammap;
+    f, u0, p = process_SciMLProblem(SDEFunction{iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, check_compatibility, sparse, expression, kwargs...)
 
diff --git a/src/systems/nonlinear/homotopy_continuation.jl b/src/systems/nonlinear/homotopy_continuation.jl
@@ -549,13 +549,13 @@ function HomotopyContinuationProblem{false}(sys::System, args...; kwargs...)
 end
 
 function HomotopyContinuationProblem{iip, spec}(
-        sys::System, u0map, pmap = SciMLBase.NullParameters();
+        sys::System, op;
         kwargs...) where {iip, spec}
     if !iscomplete(sys)
         error("A completed `System` is required. Call `complete` or `mtkcompile` on the system before creating a `HomotopyContinuationProblem`")
     end
     f, u0, p = process_SciMLProblem(
-        HomotopyNonlinearFunction{iip, spec}, sys, u0map, pmap; kwargs...)
+        HomotopyNonlinearFunction{iip, spec}, sys, op; kwargs...)
 
     kwargs = filter_kwargs(kwargs)
     return NonlinearProblem{iip}(f, u0, p; kwargs...)
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
@@ -594,7 +594,11 @@ function SciMLBase.remake_initialization_data(
     filter_missing_values!(u0map)
     filter_missing_values!(pmap)
 
-    op, missing_unknowns, missing_pars = build_operating_point!(sys,
+    merge!(u0map, pmap)
+    op = u0map
+    u0map = anydict()
+    pmap = anydict()
+    missing_unknowns, missing_pars = build_operating_point!(sys, op,
         u0map, pmap, defs, dvs, ps)
     floatT = float_type_from_varmap(op)
     u0_constructor = p_constructor = identity
diff --git a/src/systems/parameter_buffer.jl b/src/systems/parameter_buffer.jl
@@ -27,7 +27,7 @@ This requires that `complete` has been called on the system (usually via
 the default behavior).
 """
 function MTKParameters(
-        sys::AbstractSystem, p, u0 = Dict(); tofloat = false,
+        sys::AbstractSystem, op; tofloat = false,
         t0 = nothing, substitution_limit = 1000, floatT = nothing,
         p_constructor = identity)
     ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
@@ -40,17 +40,17 @@ function MTKParameters(
 
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
-    u0 = to_varmap(u0, dvs)
-    symbols_to_symbolics!(sys, u0)
-    p = to_varmap(p, ps)
-    symbols_to_symbolics!(sys, p)
+    op = to_varmap(op, ps)
+    symbols_to_symbolics!(sys, op)
     defs = add_toterms(recursive_unwrap(defaults(sys)))
 
-    is_time_dependent(sys) && add_observed!(sys, u0)
-    add_parameter_dependencies!(sys, p)
+    is_time_dependent(sys) && add_observed!(sys, op)
+    add_parameter_dependencies!(sys, op)
 
-    op, missing_unknowns, missing_pars = build_operating_point!(sys,
-        u0, p, defs, dvs, ps)
+    u0map = anydict()
+    pmap = anydict()
+    missing_unknowns, missing_pars = build_operating_point!(sys, op,
+        u0map, pmap, defs, dvs, ps)
 
     if t0 !== nothing
         op[get_iv(sys)] = t0
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
