Fixes & tests for codegen w/ constants.

Author: lamorton

src/structural_transformation/codegen.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra
 
-using ModelingToolkit: isdifferenceeq, process_events
+using ModelingToolkit: isdifferenceeq, process_events, get_preprocess_constants
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
@@ -187,12 +187,15 @@ function gen_nlsolve!(is_not_prepended_assignment, eqs, vars, u0map::AbstractDic
 
     fname = gensym("fun")
     # f is the function to find roots on
+    funex = isscalar ? rhss[1] : MakeArray(rhss, SVector)
+    @show funex
+    pre = get_preprocess_constants(funex)
     f = Func([DestructuredArgs(vars, inbounds = !checkbounds)
               DestructuredArgs(params, inbounds = !checkbounds)],
              [],
-             Let(needed_assignments[inner_idxs],
-                 isscalar ? rhss[1] : MakeArray(rhss, SVector),
-                 false)) |> SymbolicUtils.Code.toexpr
+              pre(Let(needed_assignments[inner_idxs],
+                  funex,
+                  false))) |> SymbolicUtils.Code.toexpr
 
     # solver call contains code to call the root-finding solver on the function f
     solver_call = LiteralExpr(quote
@@ -294,15 +297,17 @@ function build_torn_function(sys;
     syms = map(Symbol, states)
 
     pre = get_postprocess_fbody(sys)
+    cpre = get_preprocess_constants(rhss)
+    pre2 = x -> pre(cpre(x))
 
     expr = SymbolicUtils.Code.toexpr(Func([out
                                            DestructuredArgs(states,
-                                                            inbounds = !checkbounds)
+                                                 inbounds = !checkbounds)
                                            DestructuredArgs(parameters(sys),
-                                                            inbounds = !checkbounds)
+                                                 inbounds = !checkbounds)
                                            independent_variables(sys)],
                                           [],
-                                          pre(Let([torn_expr;
+                                          pre2(Let([torn_expr;
                                                    assignments[is_not_prepended_assignment]],
                                                   funbody,
                                                   false))),
@@ -469,12 +474,13 @@ function build_observed_function(state, ts, var_eq_matching, var_sccs,
         push!(subs, sym ← obs[eqidx].rhs)
     end
     pre = get_postprocess_fbody(sys)
-
+    cpre = get_preprocess_constants([obs[1:maxidx];  isscalar ? ts[1] : MakeArray(ts, output_type) ])
+    pre2 = x -> pre(cpre(x))
     ex = Code.toexpr(Func([DestructuredArgs(solver_states, inbounds = !checkbounds)
                            DestructuredArgs(parameters(sys), inbounds = !checkbounds)
                            independent_variables(sys)],
                           [],
-                          pre(Let([collect(Iterators.flatten(solves))
+                          pre2(Let([collect(Iterators.flatten(solves))
                                    assignments[is_not_prepended_assignment]
                                    map(eq -> eq.lhs ← eq.rhs, obs[1:maxidx])
                                    subs],

test/structural_transformation/tearing.jl

@@ -10,9 +10,10 @@ using UnPack
 ### Nonlinear system
 ###
 @parameters t
+@constants h = 1
 @variables u1(t) u2(t) u3(t) u4(t) u5(t)
 eqs = [
-    0 ~ u1 - sin(u5),
+    0 ~ u1 - sin(u5) * h,
     0 ~ u2 - cos(u1),
     0 ~ u3 - hypot(u1, u2),
     0 ~ u4 - hypot(u2, u3),
@@ -147,13 +148,13 @@ using ModelingToolkit, OrdinaryDiffEq, BenchmarkTools
 @parameters t p
 @variables x(t) y(t) z(t)
 D = Differential(t)
-eqs = [D(x) ~ z
+eqs = [D(x) ~ z * h
        0 ~ x - y
        0 ~ sin(z) + y - p * t]
 @named daesys = ODESystem(eqs, t)
 newdaesys = tearing(daesys)
-@test equations(newdaesys) == [D(x) ~ z; 0 ~ y + sin(z) - p * t]
-@test equations(tearing_substitution(newdaesys)) == [D(x) ~ z; 0 ~ x + sin(z) - p * t]
+@test equations(newdaesys) == [D(x) ~ h * z; 0 ~ y + sin(z) - p * t]
+@test equations(tearing_substitution(newdaesys)) == [D(x) ~ h * z; 0 ~ x + sin(z) - p * t]
 @test isequal(states(newdaesys), [x, z])
 prob = ODAEProblem(newdaesys, [x => 1.0], (0, 1.0), [p => 0.2])
 du = [0.0];