Merge pull request #1191 from SciML/myb/fixpreface

Support preface in more codegen functions

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -19,7 +19,8 @@ using ModelingToolkit: ODESystem, AbstractSystem,var_from_nested_derivative, Dif
                        isdiffeq, isdifferential, get_structure, get_iv, independent_variables,
                        get_structure, defaults, InvalidSystemException,
                        ExtraEquationsSystemException,
-                       ExtraVariablesSystemException
+                       ExtraVariablesSystemException,
+                       get_postprocess_fbody
 
 using ModelingToolkit.BipartiteGraphs
 using LightGraphs

src/structural_transformation/codegen.jl

@@ -201,6 +201,7 @@ function build_torn_function(
     s = structure(sys)
     states = map(i->s.fullvars[i], diffvars_range(s))
     syms = map(Symbol, states)
+    pre = get_postprocess_fbody(sys)
 
     expr = SymbolicUtils.Code.toexpr(
         Func(
@@ -211,10 +212,10 @@ function build_torn_function(
               independent_variables(sys)
              ],
              [],
-             Let(
+             pre(Let(
                  collect(Iterators.flatten(get_torn_eqs_vars(sys, checkbounds=checkbounds))),
                  odefunbody
-                )
+                ))
             )
     )
     if expression
@@ -307,6 +308,7 @@ function build_observed_function(
             obs[observed_idx[sym]].rhs
         end
     end
+    pre = get_postprocess_fbody(sys)
 
     ex = Func(
         [
@@ -315,13 +317,13 @@ function build_observed_function(
          independent_variables(sys)
         ],
         [],
-        Let(
+        pre(Let(
             [
              collect(Iterators.flatten(solves))
              map(eq -> eq.lhs←eq.rhs, obs[1:maxidx])
             ],
             isscalar ? output[1] : MakeArray(output, output_type)
-           )
+           ))
     ) |> Code.toexpr
 
     expression ? ex : @RuntimeGeneratedFunction(ex)

src/systems/diffeqs/abstractodesystem.jl

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

src/systems/diffeqs/odesystem.jl

@@ -244,13 +244,14 @@ function build_explicit_observed_function(
     ps = DestructuredArgs(parameters(sys), inbounds=!checkbounds)
     ivs = independent_variables(sys)
     args = [dvs, ps, ivs...]
+    pre = get_postprocess_fbody(sys)
 
     ex = Func(
         args, [],
-        Let(
+        pre(Let(
             map(eq -> eq.lhs←eq.rhs, obs[1:maxidx]),
             isscalar ? output[1] : MakeArray(output, output_type)
-           )
+           ))
     ) |> toexpr
 
     expression ? ex : @RuntimeGeneratedFunction(ex)

src/utils.jl

@@ -293,3 +293,15 @@ function collect_var!(states, parameters, var, iv)
     end
     return nothing
 end
+
+
+function get_postprocess_fbody(sys)
+    if has_preface(sys) && (pre = preface(sys); pre !== nothing)
+        pre_ = let pre=pre
+            ex -> Let(pre, ex)
+        end
+    else
+        pre_ = ex -> ex
+    end
+    return pre_
+end