WIP: use the preface API from Symbolics

src/inputoutput.jl

@@ -231,15 +231,12 @@
  p = map(x -> time_varying_as_func(value(x), sys), ps)
  t = get_iv(sys)
 
- # pre = has_difference ? (ex -> ex) : get_postprocess_fbody(sys)
-
  args = (u, inputs, p, t)
  if implicit_dae
  ddvs = map(Differential(get_iv(sys)), dvs)
  args = (ddvs, args...)
  end
- process = get_postprocess_fbody(sys)
- f = build_function(rhss, args...; postprocess_fbody = process,
+ f = build_function(rhss, args...; preface = get_preface_vec(sys),
  expression = Val{false}, kwargs...)
  (; f, dvs, ps, io_sys = sys)
 end

src/structural_transformation/StructuralTransformations.jl

@@ -19,7 +19,7 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
  has_tearing_state, defaults, InvalidSystemException,
  ExtraEquationsSystemException,
  ExtraVariablesSystemException,
- get_postprocess_fbody, vars!,
+ get_preface_vec, vars!,
  IncrementalCycleTracker, add_edge_checked!, topological_sort,
  invalidate_cache!, Substitutions, get_or_construct_tearing_state,
  filter_kwargs, lower_varname, setio, SparseMatrixCLIL,

src/structural_transformation/codegen.jl

@@ -197,9 +197,9 @@
  f = Func([DestructuredArgs(vars, inbounds = !checkbounds)
  DestructuredArgs(params, inbounds = !checkbounds)],
  [],
- pre(Let(needed_assignments[inner_idxs],
+ Let(vcat(pre, needed_assignments[inner_idxs]),
  funex,
- false))) |> SymbolicUtils.Code.toexpr
+ false)) |> SymbolicUtils.Code.toexpr
 
  # solver call contains code to call the root-finding solver on the function f
  solver_call = LiteralExpr(quote
@@ -301,9 +301,9 @@
  @set! sys.unknown_states = states
  syms = map(Symbol, states)
 
- pre = get_postprocess_fbody(sys)
+ pre = get_preface_vec(sys)
  cpre = get_preprocess_constants(rhss)
- pre2 = x -> pre(cpre(x))
+ pre2 = vcat(pre, cpre)
 
  expr = SymbolicUtils.Code.toexpr(Func([out
  DestructuredArgs(states,
@@ -312,10 +312,10 @@
  inbounds = !checkbounds)
  independent_variables(sys)],
  [],
- pre2(Let([torn_expr;
+ Let([pre2; torn_expr;
  assignments[is_not_prepended_assignment]],
  funbody,
- false))),
+ false)),
  sol_states)
  if expression
  expr, states
@@ -490,17 +490,17 @@
  pre = get_postprocess_fbody(sys)
  cpre = get_preprocess_constants([obs[1:maxidx];
  isscalar ? ts[1] : MakeArray(ts, output_type)])
- pre2 = x -> pre(cpre(x))
+ pre2 = vcat(pre, cpre)
  ex = Code.toexpr(Func([DestructuredArgs(unknown_states, inbounds = !checkbounds)
  DestructuredArgs(parameters(sys), inbounds = !checkbounds)
  independent_variables(sys)],
  [],
- pre2(Let([collect(Iterators.flatten(solves))
+ Let(vcat(pre2, [collect(Iterators.flatten(solves))
  assignments[is_not_prepended_assignment]
  map(eq -> eq.lhs ← eq.rhs, obs[1:maxidx])
- subs],
+ subs]),
  isscalar ? ts[1] : MakeArray(ts, output_type),
- false))), sol_states)
+ false)), sol_states)
 
  expression ? ex : drop_expr(@RuntimeGeneratedFunction(ex))
 end

src/systems/callbacks.jl

@@ -349,7 +349,7 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
  rf_oop, rf_ip = build_function(rhss, u, p, t; expression = getexpr,
  wrap_code = add_integrator_header(integ, outvar),
  outputidxs = update_inds,
- postprocess_fbody = pre,
+ preface = pre,
  kwargs...)
  # applied user-provided function to the generated expression
  if postprocess_affect_expr! !== nothing
@@ -389,7 +389,7 @@ function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = states
  t = get_iv(sys)
  pre = get_preprocess_constants(rhss)
  rf_oop, rf_ip = build_function(rhss, u, p, t; expression = Val{false},
- postprocess_fbody = pre, kwargs...)
+ preface = pre, kwargs...)
 
  affect_functions = map(cbs) do cb # Keep affect function separate
  eq_aff = affects(cb)

src/systems/diffeqs/abstractodesystem.jl

@@ -84,14 +84,14 @@
  simplify = false, kwargs...)
  tgrad = calculate_tgrad(sys, simplify = simplify)
  pre = get_preprocess_constants(tgrad)
- return build_function(tgrad, dvs, ps, get_iv(sys); postprocess_fbody = pre, kwargs...)
+ return build_function(tgrad, dvs, ps, get_iv(sys); preface = pre, kwargs...)
 end
 
 function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
  simplify = false, sparse = false, kwargs...)
  jac = calculate_jacobian(sys; simplify = simplify, sparse = sparse)
  pre = get_preprocess_constants(jac)
- return build_function(jac, dvs, ps, get_iv(sys); postprocess_fbody = pre, kwargs...)
+ return build_function(jac, dvs, ps, get_iv(sys); preface = pre, kwargs...)
 end
 
 function generate_control_jacobian(sys::AbstractODESystem, dvs = states(sys),
@@ -113,7 +113,7 @@
  jac = ˍ₋gamma * jac_du + jac_u
  pre = get_preprocess_constants(jac)
  return build_function(jac, derivatives, dvs, ps, ˍ₋gamma, get_iv(sys);
- postprocess_fbody = pre, kwargs...)
+ preface = pre, kwargs...)
 end
 
 function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
@@ -148,11 +148,11 @@
  no_postprocess = has_difference)
 
  if implicit_dae
- build_function(rhss, ddvs, u, p, t; postprocess_fbody = pre,
+ build_function(rhss, ddvs, u, p, t; preface = pre,
  states = sol_states,
  kwargs...)
  else
- build_function(rhss, u, p, t; postprocess_fbody = pre, states = sol_states,
+ build_function(rhss, u, p, t; preface = pre, states = sol_states,
  kwargs...)
  end
  end
@@ -215,11 +215,11 @@
  d.update ? eq.rhs : eq.rhs + v
  end
 
- pre = get_postprocess_fbody(sys)
+ pre = get_preface_vec(sys)
  cpre = get_preprocess_constants(body)
- pre2 = x -> pre(cpre(x))
+ pre2 = vcat(pre, cpre)
  f_oop, f_iip = build_function(body, u, p, t; expression = Val{false},
- postprocess_fbody = pre2, kwargs...)
+ preface = pre2, kwargs...)
 
  cb_affect! = let f_oop = f_oop, f_iip = f_iip
  function cb_affect!(integ)

src/systems/diffeqs/odesystem.jl

@@ -397,12 +397,12 @@
  ipts = DestructuredArgs(inputs, inbounds = !checkbounds)
  args = [dvs, ipts, ps, ivs...]
  end
- pre = get_postprocess_fbody(sys)
+ pre = get_preface_vec(sys)
 
  ex = Func(args, [],
- pre(Let(obsexprs,
+  Let(vcat(pre, obsexprs),
  isscalar ? ts[1] : MakeArray(ts, output_type),
- false))) |> toexpr
+ false)) |> toexpr
  expression ? ex : drop_expr(@RuntimeGeneratedFunction(ex))
 end
 

src/systems/discrete_system/discrete_system.jl

@@ -297,7 +297,7 @@
 
  build_function(rhss, u, p, t; kwargs...)
  pre, sol_states = get_substitutions_and_solved_states(sys)
- build_function(rhss, u, p, t; postprocess_fbody = pre, states = sol_states, kwargs...)
+ build_function(rhss, u, p, t; preface = pre, states = sol_states, kwargs...)
 end
 
 """

src/systems/nonlinear/nonlinearsystem.jl

@@ -162,7 +162,7 @@
  sparse = false, simplify = false, kwargs...)
  jac = calculate_jacobian(sys, sparse = sparse, simplify = simplify)
  pre = get_preprocess_constants(jac)
- return build_function(jac, vs, ps; postprocess_fbody = pre, kwargs...)
+ return build_function(jac, vs, ps; preface = pre, kwargs...)
 end
 
 function calculate_hessian(sys::NonlinearSystem; sparse = false, simplify = false)
@@ -180,15 +180,15 @@
  sparse = false, simplify = false, kwargs...)
  hess = calculate_hessian(sys, sparse = sparse, simplify = simplify)
  pre = get_preprocess_constants(hess)
- return build_function(hess, vs, ps; postprocess_fbody = pre, kwargs...)
+ return build_function(hess, vs, ps; preface = pre, kwargs...)
 end
 
 function generate_function(sys::NonlinearSystem, dvs = states(sys), ps = parameters(sys);
  kwargs...)
  rhss = [deq.rhs for deq in equations(sys)]
  pre, sol_states = get_substitutions_and_solved_states(sys)
 
- return build_function(rhss, value.(dvs), value.(ps); postprocess_fbody = pre,
+ return build_function(rhss, value.(dvs), value.(ps); preface = pre,
  states = sol_states, kwargs...)
 end
 

src/systems/optimization/constraints_system.jl

@@ -180,7 +180,7 @@
  lhss = generate_canonical_form_lhss(sys)
  pre, sol_states = get_substitutions_and_solved_states(sys)
 
- func = build_function(lhss, value.(dvs), value.(ps); postprocess_fbody = pre,
+ func = build_function(lhss, value.(dvs), value.(ps); preface = pre,
  states = sol_states, kwargs...)
 
  cstr = constraints(sys)

src/systems/optimization/optimizationsystem.jl

@@ -123,7 +123,7 @@
  kwargs...)
  grad = calculate_gradient(sys)
  pre = get_preprocess_constants(grad)
- return build_function(grad, vs, ps; postprocess_fbody = pre,
+ return build_function(grad, vs, ps; preface = pre,
  conv = AbstractSysToExpr(sys), kwargs...)
 end
 
@@ -139,7 +139,7 @@
  hess = calculate_hessian(sys)
  end
  pre = get_preprocess_constants(hess)
- return build_function(hess, vs, ps; postprocess_fbody = pre,
+ return build_function(hess, vs, ps; preface = pre,
  conv = AbstractSysToExpr(sys), kwargs...)
 end
 

src/utils.jl

@@ -572,20 +572,11 @@
 """
 function get_preprocess_constants(eqs)
  cs = collect_constants(eqs)
- pre = ex -> Let(Assignment[Assignment(x, getdefault(x)) for x in cs],
- ex, false)
- return pre
+ Assignment[Assignment(x, getdefault(x)) for x in cs]
 end
 
-function get_postprocess_fbody(sys)
- if has_preface(sys) && (pre = preface(sys); pre !== nothing)
- pre_ = let pre = pre
- ex -> Let(pre, ex, false)
- end
- else
- pre_ = ex -> ex
- end
- return pre_
+function get_preface_vec(sys)
+ has_preface(sys) && (pre = preface(sys); pre !== nothing) ? pre : []
 end
 
 """
@@ -629,7 +620,7 @@
  cmap, cs = get_cmap(sys)
  if empty_substitutions(sys) && isempty(cs)
  sol_states = Code.LazyState()
- pre = no_postprocess ? (ex -> ex) : get_postprocess_fbody(sys)
+ assignments = no_postprocess ? [] : get_preface_vec(sys)
  else # Have to do some work
  if !empty_substitutions(sys)
  @unpack subs = get_substitutions(sys)
@@ -639,15 +630,14 @@
  subs = [cmap; subs] # The constants need to go first
  sol_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
  if no_postprocess
- pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs], ex,
- false)
+ assignments = Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs]
  else
- process = get_postprocess_fbody(sys)
- pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs],
- process(ex), false)
+ assignments = vcat(get_preface_vec(sys),
+  Assignment[Assignment(eq.lhs, eq.rhs)
+  for eq in subs])
  end
  end
- return pre, sol_states
+ return assignments, sol_states
 end
 
 function mergedefaults(defaults, varmap, vars)