refactor: support reverting to old behavior by passing split=false

Author: AayushSabharwal

NEWS.md

@@ -30,6 +30,7 @@
   `@parameters p::Int`. Array-valued parameters must be array symbolics; `@parameters p = [1.0, 2.0]`
   is now invalid and must be changed to `@parameters p[1:2] = [1.0, 2.0]`. The index of a parameter
   in the system is also not guaranteed to be an `Int`, and will instead be a custom undocumented type.
-  - To restore the old behavior, use `ModelingToolkit.@set sys.index_cache = nothing` before creating
-    a problem, and after calling `structural_simplify`.
+  To restore the old behavior:
+  - Pass the `split = false` keyword to `structural_simplify`. E.g. `ss = structural_simplify(sys; split = false)`.
+  - Pass `split = false` to `@mtkbuild`. E.g. `@mtkbuild sys = ODESystem(...) split = false`.
 - Discrete-time system using `Difference` are unsupported. Instead, use the new `Clock`-based syntax.

src/systems/abstractsystem.jl

@@ -373,8 +373,8 @@ $(TYPEDSIGNATURES)
 Mark a system as completed. If a system is complete, the system will no longer
 namespace its subsystems or variables, i.e. `isequal(complete(sys).v.i, v.i)`.
 """
-function complete(sys::AbstractSystem)
-    if has_index_cache(sys)
+function complete(sys::AbstractSystem; split = true)
+    if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
     end
     isdefined(sys, :complete) ? (@set! sys.complete = true) : sys
@@ -1380,12 +1380,19 @@ macro component(expr)
     esc(component_post_processing(expr, false))
 end
 
-macro mtkbuild(expr)
+macro mtkbuild(exprs...)
+    expr = exprs[1]
     named_expr = ModelingToolkit.named_expr(expr)
     name = named_expr.args[1]
+    kwargs = if length(exprs) > 1
+        NamedTuple{Tuple(ex.args[1] for ex in Base.tail(exprs))}(Tuple(ex.args[2] for ex in Base.tail(exprs)))
+    else
+        (;)
+    end
+        @show kwargs
     esc(quote
         $named_expr
-        $name = $structural_simplify($name)
+        $name = $structural_simplify($name; $(kwargs)...)
     end)
 end
 

src/systems/diffeqs/abstractodesystem.jl

@@ -793,7 +793,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    u0, _, defs = get_u0_p(sys,
+    u0, _p, defs = get_u0_p(sys,
         u0map,
         parammap;
         tofloat,
@@ -803,7 +803,11 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         u0 = u0_constructor(u0)
     end
 
-    p = MTKParameters(sys, parammap)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p = MTKParameters(sys, parammap)
+    else
+        p = _p
+    end
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)

src/systems/nonlinear/nonlinearsystem.jl

@@ -353,7 +353,9 @@ function process_NonlinearProblem(constructor, sys::NonlinearSystem, u0map, para
     ps = parameters(sys)
 
     u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
-    p = MTKParameters(sys, parammap)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p = MTKParameters(sys, parammap)
+    end
     check_eqs_u0(eqs, dvs, u0; kwargs...)
 
     f = constructor(sys, dvs, ps, u0; jac = jac, checkbounds = checkbounds,

src/systems/optimization/optimizationsystem.jl

@@ -623,7 +623,7 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0map,
     end
 end
 
-function structural_simplify(sys::OptimizationSystem; kwargs...)
+function structural_simplify(sys::OptimizationSystem; split = true, kwargs...)
     sys = flatten(sys)
     cons = constraints(sys)
     econs = Equation[]
@@ -650,6 +650,6 @@ function structural_simplify(sys::OptimizationSystem; kwargs...)
     neweqs = fixpoint_sub.(equations(sys), (subs,))
     @set! sys.op = length(neweqs) == 1 ? first(neweqs) : neweqs
     @set! sys.unknowns = newsts
-    sys = complete(sys)
+    sys = complete(sys; split)
     return sys
 end

src/systems/systems.jl

@@ -16,7 +16,7 @@ The optional argument `io` may take a tuple `(inputs, outputs)`.
 This will convert all `inputs` to parameters and allow them to be unconnected, i.e.,
 simplification will allow models where `n_unknowns = n_equations - n_inputs`.
 """
-function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
+function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false, split = true,
         kwargs...)
     newsys′ = __structural_simplify(sys, io; simplify, kwargs...)
     if newsys′ isa Tuple
@@ -25,8 +25,8 @@ function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false
     else
         newsys = newsys′
     end
-    @set! newsys.parent = complete(sys)
-    newsys = complete(newsys)
+    @set! newsys.parent = complete(sys; split)
+    newsys = complete(newsys; split)
     if newsys′ isa Tuple
         idxs = [parameter_index(newsys, i) for i in io[1]]
         return newsys, idxs