Add preface

Author: YingboMa

src/systems/abstractsystem.jl

@@ -165,6 +165,7 @@ for prop in [
              :depvars
              :indvars
              :connection_type
+             :preface
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
@@ -357,6 +358,12 @@ function namespace_equation(eq::Equation, sys)
     _lhs ~ _rhs
 end
 
+function namespace_assignment(eq::Assignment, sys)
+    _lhs = namespace_expr(eq.lhs, sys)
+    _rhs = namespace_expr(eq.rhs, sys)
+    Assignment(_lhs, _rhs)
+end
+
 function namespace_expr(O, sys) where {T}
     iv = independent_variable(sys)
     O = unwrap(O)
@@ -435,6 +442,25 @@ function equations(sys::ModelingToolkit.AbstractSystem)
     end
 end
 
+function preface(sys::ModelingToolkit.AbstractSystem)
+    has_preface(sys) || return nothing
+    pre = get_preface(sys)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return pre
+    else
+        pres = pre === nothing ? [] : pre
+        for sys in systems
+            pre = get_preface(sys)
+            pre === nothing && continue
+            for eq in pre
+                push!(pres, namespace_assignment(eq, sys))
+            end
+        end
+        return isempty(pres) ? nothing : pres
+    end
+end
+
 function islinear(sys::AbstractSystem)
     rhs = [eq.rhs for eq ∈ equations(sys)]
 

src/systems/diffeqs/abstractodesystem.jl

@@ -102,10 +102,16 @@ function generate_function(
     p = map(x->time_varying_as_func(value(x), sys), ps)
     t = get_iv(sys)
 
+    if has_preface(sys)
+      pre = ex -> Let(preface(sys), ex)
+    else
+        pre = ex -> ex
+    end
+
     if implicit_dae
-        build_function(rhss, ddvs, u, p, t; kwargs...)
+        build_function(rhss, ddvs, u, p, t; postprocess_fbody=pre, kwargs...)
     else
-        build_function(rhss, u, p, t; kwargs...)
+        build_function(rhss, u, p, t; postprocess_fbody=pre, kwargs...)
     end
 end
 

src/systems/diffeqs/odesystem.jl

@@ -80,15 +80,19 @@ struct ODESystem <: AbstractODESystem
     """
     structure::Any
     """
-    type: type of the system
+    connection_type: type of the system
     """
     connection_type::Any
+    """
+    preface: injuect assignment statements before the evaluation of the RHS function.
+    """
+    preface::Any
 
-    function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
+    function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type, preface)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
         check_equations(deqs,iv)
-        new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
+        new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type, preface)
     end
 end
 
@@ -102,6 +106,7 @@ function ODESystem(
                    default_p=Dict(),
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
                    connection_type=nothing,
+                   preface=nothing,
                   )
     name === nothing && throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = collect(deqs)
@@ -135,7 +140,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
+    ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type, preface)
 end
 
 function ODESystem(eqs, iv=nothing; kwargs...)

src/variables.jl

@@ -50,6 +50,7 @@ function varmap_to_vars(varmap, varlist; defaults=Dict(), check=true, toterm=Sym
     elseif container_type <: Tuple
         (vals...,)
     else
+        vals = identity.(vals)
         SymbolicUtils.Code.create_array(container_type, eltype(vals), Val{1}(), Val(length(vals)), vals...)
     end
 end