Merge pull request #3360 from AayushSabharwal/as/refactor-codegen

refactor: centralize all code generation

Author: ChrisRackauckas

Project.toml

@@ -147,7 +147,7 @@ StochasticDiffEq = "6.72.1"
 StochasticDelayDiffEq = "1.8.1"
 SymbolicIndexingInterface = "0.3.37"
 SymbolicUtils = "3.10.1"
-Symbolics = "6.25"
+Symbolics = "6.27"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"

src/ModelingToolkit.jl

@@ -150,6 +150,7 @@ include("systems/connectors.jl")
 include("systems/analysis_points.jl")
 include("systems/imperative_affect.jl")
 include("systems/callbacks.jl")
+include("systems/codegen_utils.jl")
 include("systems/problem_utils.jl")
 
 include("systems/nonlinear/nonlinearsystem.jl")

src/inputoutput.jl

@@ -249,17 +249,8 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
         ddvs = map(Differential(get_iv(sys)), dvs)
         args = (ddvs, args...)
     end
-    process = get_postprocess_fbody(sys)
-    wrapped_arrays_vars = disturbance_argument ?
-                          wrap_array_vars(
-        sys, rhss; dvs, ps, inputs, extra_args = (disturbance_inputs,)) :
-                          wrap_array_vars(sys, rhss; dvs, ps, inputs)
-    f = build_function(rhss, args...; postprocess_fbody = process,
-        expression = Val{true}, wrap_code = wrap_mtkparameters(
-            sys, false, 3, Int(disturbance_argument) + 1) .∘
-                                            wrapped_arrays_vars .∘
-                                            wrap_parameter_dependencies(sys, false),
-        kwargs...)
+    f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
+        p_end = length(p) + 2 + implicit_dae)
     f = eval_or_rgf.(f; eval_expression, eval_module)
     (; f, dvs, ps, io_sys = sys)
 end

src/systems/abstractsystem.jl

@@ -161,47 +161,27 @@ time-independent systems. If `split=true` (the default) was passed to [`complete
 object.
 """
 function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys),
-        ps = parameters(sys); wrap_code = nothing, postprocess_fbody = nothing, states = nothing,
+        ps = parameters(sys);
         expression = Val{true}, eval_expression = false, eval_module = @__MODULE__,
         cachesyms::Tuple = (), kwargs...)
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system.")
     end
     p = (reorder_parameters(sys, unwrap.(ps))..., cachesyms...)
     isscalar = !(exprs isa AbstractArray)
-    if wrap_code === nothing
-        wrap_code = isscalar ? identity : (identity, identity)
-    end
-    pre, sol_states = get_substitutions_and_solved_unknowns(sys, isscalar ? [exprs] : exprs)
-    if postprocess_fbody === nothing
-        postprocess_fbody = pre
-    end
-    if states === nothing
-        states = sol_states
-    end
     fnexpr = if is_time_dependent(sys)
-        build_function(exprs,
+        build_function_wrapper(sys, exprs,
             dvs,
             p...,
             get_iv(sys);
             kwargs...,
-            postprocess_fbody,
-            states,
-            wrap_code = wrap_code .∘ wrap_mtkparameters(sys, isscalar) .∘
-                        wrap_array_vars(sys, exprs; dvs, cachesyms) .∘
-                        wrap_parameter_dependencies(sys, isscalar),
             expression = Val{true}
         )
     else
-        build_function(exprs,
+        build_function_wrapper(sys, exprs,
             dvs,
             p...;
             kwargs...,
-            postprocess_fbody,
-            states,
-            wrap_code = wrap_code .∘ wrap_mtkparameters(sys, isscalar) .∘
-                        wrap_array_vars(sys, exprs; dvs, cachesyms) .∘
-                        wrap_parameter_dependencies(sys, isscalar),
             expression = Val{true}
         )
     end
@@ -844,7 +824,7 @@ end
 
 function SymbolicIndexingInterface.all_symbols(sys::AbstractSystem)
     syms = all_variable_symbols(sys)
-    for other in (parameter_symbols(sys), independent_variable_symbols(sys))
+    for other in (full_parameters(sys), independent_variable_symbols(sys))
         isempty(other) || (syms = vcat(syms, other))
     end
     return syms
@@ -2578,7 +2558,7 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
 
     fun_expr = generate_function(sys, sts, ps; expression = Val{true})[1]
     fun = eval_or_rgf(fun_expr; eval_expression, eval_module)
-    dx = fun(sts, p..., t)
+    dx = fun(sts, p, t)
 
     h = build_explicit_observed_function(sys, outputs; eval_expression, eval_module)
     y = h(sts, p, t)

src/systems/callbacks.jl

@@ -519,70 +519,19 @@ end
 # handles ensuring that affect! functions work with integrator arguments
 function add_integrator_header(
         sys::AbstractSystem, integrator = gensym(:MTKIntegrator), out = :u)
-    if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        function (expr)
-            p = gensym(:p)
-            Func(
-                [
-                    DestructuredArgs([expr.args[1], p, expr.args[end]],
-                    integrator, inds = [:u, :p, :t])
-                ],
-                [],
-                Let(
-                    [DestructuredArgs([arg.name for arg in expr.args[2:(end - 1)]], p),
-                        expr.args[2:(end - 1)]...],
-                    expr.body,
-                    false)
-            )
-        end,
-        function (expr)
-            p = gensym(:p)
-            Func(
-                [
-                    DestructuredArgs([expr.args[1], expr.args[2], p, expr.args[end]],
-                    integrator, inds = [out, :u, :p, :t])
-                ],
-                [],
-                Let(
-                    [DestructuredArgs([arg.name for arg in expr.args[3:(end - 1)]], p),
-                        expr.args[3:(end - 1)]...],
-                    expr.body,
-                    false)
-            )
-        end
-    else
-        expr -> Func([DestructuredArgs(expr.args, integrator, inds = [:u, :p, :t])], [],
-            expr.body),
-        expr -> Func(
-            [DestructuredArgs(expr.args, integrator, inds = [out, :u, :p, :t])], [],
-            expr.body)
-    end
+    expr -> Func([DestructuredArgs(expr.args, integrator, inds = [:u, :p, :t])], [],
+        expr.body),
+    expr -> Func(
+        [DestructuredArgs(expr.args, integrator, inds = [out, :u, :p, :t])], [],
+        expr.body)
 end
 
 function condition_header(sys::AbstractSystem, integrator = gensym(:MTKIntegrator))
-    if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        function (expr)
-            p = gensym(:p)
-            res = Func(
-                [expr.args[1], expr.args[2],
-                    DestructuredArgs([p], integrator, inds = [:p])],
-                [],
-                Let(
-                    [
-                        DestructuredArgs([arg.name for arg in expr.args[3:end]], p),
-                        expr.args[3:end]...
-                    ], expr.body, false
-                )
-            )
-            return res
-        end
-    else
-        expr -> Func(
-            [expr.args[1], expr.args[2],
-                DestructuredArgs(expr.args[3:end], integrator, inds = [:p])],
-            [],
-            expr.body)
-    end
+    expr -> Func(
+        [expr.args[1], expr.args[2],
+            DestructuredArgs(expr.args[3:end], integrator, inds = [:p])],
+        [],
+        expr.body)
 end
 
 function callback_save_header(sys::AbstractSystem, cb)
@@ -628,11 +577,10 @@ function compile_condition(cb::SymbolicDiscreteCallback, sys, dvs, ps;
         cmap = map(x -> x => getdefault(x), cs)
         condit = substitute(condit, cmap)
     end
-    expr = build_function(
+    expr = build_function_wrapper(sys,
         condit, u, t, p...; expression = Val{true},
-        wrap_code = condition_header(sys) .∘
-                    wrap_array_vars(sys, condit; dvs, ps, inputs = true) .∘
-                    wrap_parameter_dependencies(sys, !(condit isa AbstractArray)),
+        p_start = 3, p_end = length(p) + 2,
+        wrap_code = condition_header(sys),
         kwargs...)
     if expression == Val{true}
         return expr
@@ -715,14 +663,12 @@ function compile_affect(eqs::Vector{Equation}, cb, sys, dvs, ps; outputidxs = no
         end
         t = get_iv(sys)
         integ = gensym(:MTKIntegrator)
-        pre = get_preprocess_constants(rhss)
-        rf_oop, rf_ip = build_function(rhss, u, p..., t; expression = Val{true},
+        rf_oop, rf_ip = build_function_wrapper(
+            sys, rhss, u, p..., t; expression = Val{true},
             wrap_code = callback_save_header(sys, cb) .∘
-                        add_integrator_header(sys, integ, outvar) .∘
-                        wrap_array_vars(sys, rhss; dvs, ps = _ps) .∘
-                        wrap_parameter_dependencies(sys, false),
+                        add_integrator_header(sys, integ, outvar),
             outputidxs = update_inds,
-            postprocess_fbody = pre,
+            create_bindings = false,
             kwargs...)
         # applied user-provided function to the generated expression
         if postprocess_affect_expr! !== nothing