Merge branch 'master' into myb/op

YingboMa · YingboMa · commit b20cee53786d · 2023-06-01T19:31:55.000-04:00
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Yingbo Ma <mayingbo5@gmail.com>", "Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "8.56.0"
+version = "8.57.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -79,7 +79,7 @@ MacroTools = "0.5"
 NaNMath = "0.3, 1"
 RecursiveArrayTools = "2.3"
 Reexport = "0.2, 1"
-RuntimeGeneratedFunctions = "0.4.3, 0.5"
+RuntimeGeneratedFunctions = "0.5.9"
 SciMLBase = "1.76.1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0"
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
@@ -18,6 +18,7 @@ using JumpProcesses
 using DataStructures
 using SpecialFunctions, NaNMath
 using RuntimeGeneratedFunctions
+using RuntimeGeneratedFunctions: drop_expr
 using Base.Threads
 using DiffEqCallbacks
 using Graphs
diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
@@ -34,7 +34,8 @@ using ModelingToolkit.SystemStructures: algeqs, EquationsView
 
 using ModelingToolkit.DiffEqBase
 using ModelingToolkit.StaticArrays
-using ModelingToolkit: @RuntimeGeneratedFunction, RuntimeGeneratedFunctions
+using RuntimeGeneratedFunctions: @RuntimeGeneratedFunction, RuntimeGeneratedFunctions,
+                                 drop_expr
 
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
diff --git a/src/structural_transformation/bipartite_tearing/modia_tearing.jl b/src/structural_transformation/bipartite_tearing/modia_tearing.jl
@@ -73,7 +73,10 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
     # find them here [TODO: It would be good to have an explicit example of this.]
 
     @unpack graph, solvable_graph = structure
-    var_eq_matching = complete(maximal_matching(graph, eqfilter, varfilter, U))
+    var_eq_matching = maximal_matching(graph, eqfilter, varfilter, U)
+    var_eq_matching = complete(var_eq_matching,
+                               max(length(var_eq_matching),
+                                   maximum(x -> x isa Int ? x : 0, var_eq_matching)))
     var_sccs::Vector{Union{Vector{Int}, Int}} = find_var_sccs(graph, var_eq_matching)
     vargraph = DiCMOBiGraph{true}(graph)
     ict = IncrementalCycleTracker(vargraph; dir = :in)
diff --git a/src/structural_transformation/codegen.jl b/src/structural_transformation/codegen.jl
@@ -219,7 +219,7 @@ function gen_nlsolve!(is_not_prepended_assignment, eqs, vars, u0map::AbstractDic
     end
 
     nlsolve_expr = Assignment[preassignments
-                              fname ← @RuntimeGeneratedFunction(f)
+                              fname ← drop_expr(@RuntimeGeneratedFunction(f))
                               DestructuredArgs(vars, inbounds = !checkbounds) ← solver_call]
 
     nlsolve_expr
@@ -345,7 +345,7 @@ function build_torn_function(sys;
             end
         end
 
-        ODEFunction{true, SciMLBase.AutoSpecialize}(@RuntimeGeneratedFunction(expr),
+        ODEFunction{true, SciMLBase.AutoSpecialize}(drop_expr(@RuntimeGeneratedFunction(expr)),
                                                     sparsity = jacobian_sparsity ?
                                                                torn_system_with_nlsolve_jacobian_sparsity(state,
                                                                                                           var_eq_matching,
@@ -501,7 +501,7 @@ function build_observed_function(state, ts, var_eq_matching, var_sccs,
                                    isscalar ? ts[1] : MakeArray(ts, output_type),
                                    false))), sol_states)
 
-    expression ? ex : @RuntimeGeneratedFunction(ex)
+    expression ? ex : drop_expr(@RuntimeGeneratedFunction(ex))
 end
 
 """
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
@@ -583,10 +583,10 @@ Tear the nonlinear equations in system. When `simplify=true`, we simplify the
 new residual equations after tearing. End users are encouraged to call [`structural_simplify`](@ref)
 instead, which calls this function internally.
 """
-function tearing(sys::AbstractSystem; simplify = false)
-    state = TearingState(sys)
+function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing,
+                 simplify = false, kwargs...)
     var_eq_matching = tearing(state)
-    invalidate_cache!(tearing_reassemble(state, var_eq_matching; simplify = simplify))
+    invalidate_cache!(tearing_reassemble(state, var_eq_matching; mm, simplify))
 end
 
 """
diff --git a/src/systems/alias_elimination.jl b/src/systems/alias_elimination.jl
@@ -468,10 +468,10 @@ function observed2graph(eqs, states)
 end
 
 function fixpoint_sub(x, dict)
-    y = substitute(x, dict)
+    y = fast_substitute(x, dict)
     while !isequal(x, y)
         y = x
-        x = substitute(y, dict)
+        x = fast_substitute(y, dict)
     end
 
     return x
diff --git a/src/systems/callbacks.jl b/src/systems/callbacks.jl
@@ -354,7 +354,8 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
         # applied user-provided function to the generated expression
         if postprocess_affect_expr! !== nothing
             postprocess_affect_expr!(rf_ip, integ)
-            (expression == Val{false}) && (return @RuntimeGeneratedFunction(rf_ip))
+            (expression == Val{false}) &&
+                (return drop_expr(@RuntimeGeneratedFunction(rf_ip)))
         end
         rf_ip
     end
diff --git a/src/systems/clock_inference.jl b/src/systems/clock_inference.jl
@@ -225,7 +225,7 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
     end
     if eval_expression
         affects = map(affect_funs) do a
-            @RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a)))
+            drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
         end
     else
         affects = map(a -> toexpr(LiteralExpr(a)), affect_funs)
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -281,7 +281,8 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = s
                               expression_module = eval_module, checkbounds = checkbounds,
                               kwargs...)
     f_oop, f_iip = eval_expression ?
-                   (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen) : f_gen
+                   (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
+                   f_gen
     f(u, p, t) = f_oop(u, p, t)
     f(du, u, p, t) = f_iip(du, u, p, t)
 
@@ -299,7 +300,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = s
                                    expression_module = eval_module,
                                    checkbounds = checkbounds, kwargs...)
         tgrad_oop, tgrad_iip = eval_expression ?
-                               (@RuntimeGeneratedFunction(eval_module, ex) for ex in tgrad_gen) :
+                               (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in tgrad_gen) :
                                tgrad_gen
         _tgrad(u, p, t) = tgrad_oop(u, p, t)
         _tgrad(J, u, p, t) = tgrad_iip(J, u, p, t)
@@ -314,7 +315,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = s
                                     expression_module = eval_module,
                                     checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (@RuntimeGeneratedFunction(eval_module, ex) for ex in jac_gen) :
+                           (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
                            jac_gen
         _jac(u, p, t) = jac_oop(u, p, t)
         _jac(J, u, p, t) = jac_iip(J, u, p, t)
@@ -423,7 +424,8 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                               expression_module = eval_module, checkbounds = checkbounds,
                               kwargs...)
     f_oop, f_iip = eval_expression ?
-                   (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen) : f_gen
+                   (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
+                   f_gen
     f(du, u, p, t) = f_oop(du, u, p, t)
     f(out, du, u, p, t) = f_iip(out, du, u, p, t)
 
@@ -434,7 +436,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                         expression_module = eval_module,
                                         checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (@RuntimeGeneratedFunction(eval_module, ex) for ex in jac_gen) :
+                           (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
                            jac_gen
         _jac(du, u, p, ˍ₋gamma, t) = jac_oop(du, u, p, ˍ₋gamma, t)
 
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -387,7 +387,7 @@ function build_explicit_observed_function(sys, ts;
               pre(Let(obsexprs,
                       isscalar ? ts[1] : MakeArray(ts, output_type),
                       false))) |> toexpr
-    expression ? ex : @RuntimeGeneratedFunction(ex)
+    expression ? ex : drop_expr(@RuntimeGeneratedFunction(ex))
 end
 
 function _eq_unordered(a, b)
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
@@ -385,10 +385,12 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
     ps = scalarize.(ps)
 
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression}, kwargs...)
-    f_oop, f_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in f_gen) : f_gen
+    f_oop, f_iip = eval_expression ?
+                   (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in f_gen) : f_gen
     g_gen = generate_diffusion_function(sys, dvs, ps; expression = Val{eval_expression},
                                         kwargs...)
-    g_oop, g_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in g_gen) : g_gen
+    g_oop, g_iip = eval_expression ?
+                   (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in g_gen) : g_gen
 
     f(u, p, t) = f_oop(u, p, t)
     f(du, u, p, t) = f_iip(du, u, p, t)
@@ -399,7 +401,7 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
         tgrad_gen = generate_tgrad(sys, dvs, ps; expression = Val{eval_expression},
                                    kwargs...)
         tgrad_oop, tgrad_iip = eval_expression ?
-                               (@RuntimeGeneratedFunction(ex) for ex in tgrad_gen) :
+                               (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in tgrad_gen) :
                                tgrad_gen
         _tgrad(u, p, t) = tgrad_oop(u, p, t)
         _tgrad(J, u, p, t) = tgrad_iip(J, u, p, t)
@@ -411,7 +413,8 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
         jac_gen = generate_jacobian(sys, dvs, ps; expression = Val{eval_expression},
                                     sparse = sparse, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (@RuntimeGeneratedFunction(ex) for ex in jac_gen) : jac_gen
+                           (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in jac_gen) :
+                           jac_gen
         _jac(u, p, t) = jac_oop(u, p, t)
         _jac(J, u, p, t) = jac_iip(J, u, p, t)
     else
@@ -422,10 +425,10 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
         tmp_Wfact, tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true;
                                                        expression = Val{true}, kwargs...)
         Wfact_oop, Wfact_iip = eval_expression ?
-                               (@RuntimeGeneratedFunction(ex) for ex in tmp_Wfact) :
+                               (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in tmp_Wfact) :
                                tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = eval_expression ?
-                                   (@RuntimeGeneratedFunction(ex) for ex in tmp_Wfact_t) :
+                                   (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in tmp_Wfact_t) :
                                    tmp_Wfact_t
         _Wfact(u, p, dtgamma, t) = Wfact_oop(u, p, dtgamma, t)
         _Wfact(W, u, p, dtgamma, t) = Wfact_iip(W, u, p, dtgamma, t)
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
@@ -227,7 +227,7 @@ function SciMLBase.DiscreteProblem(sys::DiscreteSystem, u0map = [], tspan = get_
 
     f_gen = generate_function(sys; expression = Val{eval_expression},
                               expression_module = eval_module)
-    f_oop, _ = (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen)
+    f_oop, _ = (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
     f(u, p, iv) = f_oop(u, p, iv)
     fd = DiscreteFunction(f; syms = Symbol.(dvs), indepsym = Symbol(iv),
                           paramsyms = Symbol.(ps), sys = sys)
@@ -339,7 +339,8 @@ function SciMLBase.DiscreteFunction{iip, specialize}(sys::DiscreteSystem,
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression},
                               expression_module = eval_module, kwargs...)
     f_oop, f_iip = eval_expression ?
-                   (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen) : f_gen
+                   (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
+                   f_gen
     f(u, p, t) = f_oop(u, p, t)
     f(du, u, p, t) = f_iip(du, u, p, t)
 
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
@@ -196,11 +196,11 @@ function generate_affect_function(js::JumpSystem, affect, outputidxs)
 end
 
 function assemble_vrj(js, vrj, statetoid)
-    rate = @RuntimeGeneratedFunction(generate_rate_function(js, vrj.rate))
+    rate = drop_expr(@RuntimeGeneratedFunction(generate_rate_function(js, vrj.rate)))
     outputvars = (value(affect.lhs) for affect in vrj.affect!)
     outputidxs = [statetoid[var] for var in outputvars]
-    affect = @RuntimeGeneratedFunction(generate_affect_function(js, vrj.affect!,
-                                                                outputidxs))
+    affect = drop_expr(@RuntimeGeneratedFunction(generate_affect_function(js, vrj.affect!,
+                                                                          outputidxs)))
     VariableRateJump(rate, affect)
 end
 
@@ -217,11 +217,11 @@ function assemble_vrj_expr(js, vrj, statetoid)
 end
 
 function assemble_crj(js, crj, statetoid)
-    rate = @RuntimeGeneratedFunction(generate_rate_function(js, crj.rate))
+    rate = drop_expr(@RuntimeGeneratedFunction(generate_rate_function(js, crj.rate)))
     outputvars = (value(affect.lhs) for affect in crj.affect!)
     outputidxs = [statetoid[var] for var in outputvars]
-    affect = @RuntimeGeneratedFunction(generate_affect_function(js, crj.affect!,
-                                                                outputidxs))
+    affect = drop_expr(@RuntimeGeneratedFunction(generate_affect_function(js, crj.affect!,
+                                                                          outputidxs)))
     ConstantRateJump(rate, affect)
 end
 
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
@@ -228,7 +228,8 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys
                                           sparse = false, simplify = false,
                                           kwargs...) where {iip}
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression}, kwargs...)
-    f_oop, f_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in f_gen) : f_gen
+    f_oop, f_iip = eval_expression ?
+                   (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in f_gen) : f_gen
     f(u, p) = f_oop(u, p)
     f(du, u, p) = f_iip(du, u, p)
 
@@ -237,7 +238,8 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys
                                     simplify = simplify, sparse = sparse,
                                     expression = Val{eval_expression}, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (@RuntimeGeneratedFunction(ex) for ex in jac_gen) : jac_gen
+                           (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in jac_gen) :
+                           jac_gen
         _jac(u, p) = jac_oop(u, p)
         _jac(J, u, p) = jac_iip(J, u, p)
     else
diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
@@ -609,18 +609,19 @@ function structural_simplify(sys::OptimizationSystem; kwargs...)
         end
     end
     nlsys = NonlinearSystem(econs, states(sys), parameters(sys); name = :___tmp_nlsystem)
-    snlsys = structural_simplify(nlsys; check_consistency = false, kwargs...)
+    snlsys = structural_simplify(nlsys; fully_determined = false, kwargs...)
     obs = observed(snlsys)
     subs = Dict(eq.lhs => eq.rhs for eq in observed(snlsys))
     seqs = equations(snlsys)
-    cons_simplified = Array{eltype(cons), 1}(undef, length(icons) + length(seqs))
+    cons_simplified = similar(cons, length(icons) + length(seqs))
     for (i, eq) in enumerate(Iterators.flatten((seqs, icons)))
-        cons_simplified[i] = substitute(eq, subs)
+        cons_simplified[i] = fixpoint_sub(eq, subs)
     end
     newsts = setdiff(states(sys), keys(subs))
     @set! sys.constraints = cons_simplified
     @set! sys.observed = [observed(sys); obs]
-    @set! sys.op = substitute(equations(sys), subs)
+    neweqs = fixpoint_sub.(equations(sys), (subs,))
+    @set! sys.op = length(neweqs) == 1 ? first(neweqs) : neweqs
     @set! sys.states = newsts
     return sys
 end
diff --git a/src/systems/pde/pdesystem.jl b/src/systems/pde/pdesystem.jl
@@ -114,7 +114,7 @@ struct PDESystem <: ModelingToolkit.AbstractMultivariateSystem
                     p = ps isa SciMLBase.NullParameters ? [] : map(a -> a.first, ps)
                     args = vcat(DestructuredArgs(p), args)
                     ex = Func(args, [], eq.rhs) |> toexpr
-                    eq.lhs => @RuntimeGeneratedFunction(ex)
+                    eq.lhs => drop_expr(@RuntimeGeneratedFunction(ex))
                 end
             end
         end
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
@@ -554,14 +554,15 @@ function merge_io(io, inputs)
 end
 
 function structural_simplify!(state::TearingState, io = nothing; simplify = false,
-                              check_consistency = true, kwargs...)
+                              check_consistency = true, fully_determined = true,
+                              kwargs...)
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
         ModelingToolkit.infer_clocks!(ci)
         tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
         cont_io = merge_io(io, inputs[continuous_id])
         sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify,
-                                                check_consistency,
+                                                check_consistency, fully_determined,
                                                 kwargs...)
         if length(tss) > 1
             # TODO: rename it to something else
@@ -576,7 +577,7 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
                 end
                 dist_io = merge_io(io, inputs[i])
                 ss, = _structural_simplify!(state, dist_io; simplify, check_consistency,
-                                            kwargs...)
+                                            fully_determined, kwargs...)
                 append!(appended_parameters, inputs[i], states(ss))
                 discrete_subsystems[i] = ss
             end
@@ -588,14 +589,16 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
         end
     else
         sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,
-                                                kwargs...)
+                                                fully_determined, kwargs...)
     end
     has_io = io !== nothing
     return has_io ? (sys, input_idxs) : sys
 end
 
 function _structural_simplify!(state::TearingState, io; simplify = false,
-                               check_consistency = true, kwargs...)
+                               check_consistency = true, fully_determined = true,
+                               kwargs...)
+    check_consistency &= fully_determined
     has_io = io !== nothing
     orig_inputs = Set()
     if has_io
@@ -606,7 +609,11 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
     if check_consistency
         ModelingToolkit.check_consistency(state, orig_inputs)
     end
-    sys = ModelingToolkit.dummy_derivative(sys, state; simplify, mm)
+    if fully_determined
+        sys = ModelingToolkit.dummy_derivative(sys, state; simplify, mm, check_consistency)
+    else
+        sys = ModelingToolkit.tearing(sys, state; simplify, mm, check_consistency)
+    end
     fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = ModelingToolkit.topsort_equations(observed(sys), fullstates)
     ModelingToolkit.invalidate_cache!(sys), input_idxs
diff --git a/src/utils.jl b/src/utils.jl
