WIP

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -38,6 +38,7 @@ using SparseArrays
 using NonlinearSolve
 
 export tearing, dae_index_lowering, check_consistency
+export tearing_assignments, tearing_substitution
 export build_torn_function, build_observed_function, ODAEProblem
 export sorted_incidence_matrix
 

src/structural_transformation/codegen.jl

@@ -96,7 +96,7 @@ function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_s
 end
 
 """
-    exprs = gen_nlsolve(eqs::Vector{Equation}, vars::Vector, u0map::Dict; checkbounds = true)
+    exprs = gen_nlsolve(eqs::Vector{Equation}, vars::Vector, u0map::Dict; checkbounds = true, assignments)
 
 Generate `SymbolicUtils` expressions for a root-finding function based on `eqs`,
 as well as a call to the root-finding solver.
@@ -112,13 +112,24 @@ exprs = [fname = f, numerical_nlsolve(fname, ...)]
 - `u0map`: A `Dict` which maps variables in `eqs` to values, e.g., `defaults(sys)` if `eqs = equations(sys)`.
 - `checkbounds`: Apply bounds checking in the generated code.
 """
-function gen_nlsolve(eqs, vars, u0map::AbstractDict; checkbounds=true)
+function gen_nlsolve(eqs, vars, u0map::AbstractDict; checkbounds=true, assignments)
     isempty(vars) && throw(ArgumentError("vars may not be empty"))
     length(eqs) == length(vars) || throw(ArgumentError("vars must be of the same length as the number of equations to find the roots of"))
     rhss = map(x->x.rhs, eqs)
     # We use `vars` instead of `graph` to capture parameters, too.
-    allvars = unique(collect(Iterators.flatten(map(ModelingToolkit.vars, rhss))))
-    params = setdiff(allvars, vars) # these are not the subject of the root finding
+    allvars = Set(Iterators.flatten(ModelingToolkit.vars(r) for r in rhss))
+    vars_set = Set(vars)
+    params = setdiff(allvars, vars_set) # these are not the subject of the root finding
+    #needed_assignments = filter(a->a.lhs in params, assignments)
+    needed_assignments = assignments[1:findlast(a->a.lhs in params, assignments)]
+    params = setdiff(params, [a.lhs for a in needed_assignments])
+    @show needed_assignments, params
+    for a in needed_assignments
+        ModelingToolkit.vars!(params, a.rhs)
+    end
+    params = setdiff(params, vars_set) # these are not the subject of the root finding
+    @show params
+    # inductor1₊v, inductor2₊v
 
     # splatting to tighten the type
     u0 = []
@@ -141,10 +152,10 @@ function gen_nlsolve(eqs, vars, u0map::AbstractDict; checkbounds=true)
     f = Func(
         [
          DestructuredArgs(vars, inbounds=!checkbounds)
-         DestructuredArgs(params, inbounds=!checkbounds)
+         DestructuredArgs(collect(params), inbounds=!checkbounds)
         ],
         [],
-        isscalar ? rhss[1] : MakeArray(rhss, SVector)
+        Let(needed_assignments, isscalar ? rhss[1] : MakeArray(rhss, SVector))
     ) |> SymbolicUtils.Code.toexpr
 
     # solver call contains code to call the root-finding solver on the function f
@@ -193,18 +204,21 @@ function build_torn_function(
 
     states_idxs = collect(diffvars_range(s))
     mass_matrix_diag = ones(length(states_idxs))
+
+    assignments, bf_states = tearing_assignments(sys)
     torn_expr = []
+
     defs = defaults(sys)
     nlsolve_scc_idxs = Int[]
 
     needs_extending = false
-    for (i, scc) in enumerate(var_sccs)
+    @views for (i, scc) in enumerate(var_sccs)
         #torn_vars = [s.fullvars[var] for var in scc if var_eq_matching[var] !== unassigned]
         torn_vars_idxs = Int[var for var in scc if var_eq_matching[var] !== unassigned]
         torn_eqs_idxs = [var_eq_matching[var] for var in torn_vars_idxs]
         isempty(torn_eqs_idxs) && continue
         if length(torn_eqs_idxs) <= max_inlining_size
-            append!(torn_expr, gen_nlsolve(eqs[torn_eqs_idxs], s.fullvars[torn_vars_idxs], defs, checkbounds=checkbounds))
+            append!(torn_expr, gen_nlsolve(eqs[torn_eqs_idxs], s.fullvars[torn_vars_idxs], defs, checkbounds=checkbounds, assignments=assignments))
             push!(nlsolve_scc_idxs, i)
         else
             needs_extending = true
@@ -226,6 +240,7 @@ function build_torn_function(
 
     states = s.fullvars[states_idxs]
     syms = map(Symbol, states_idxs)
+
     pre = get_postprocess_fbody(sys)
 
     expr = SymbolicUtils.Code.toexpr(
@@ -238,18 +253,22 @@ function build_torn_function(
              ],
              [],
              pre(Let(
-                 torn_expr,
+                 [torn_expr; assignments],
                  funbody
                 ))
-            )
+            ),
+        bf_states
     )
     if expression
         expr, states
     else
-        observedfun = let sys = sys, dict = Dict()
+        observedfun = let sys=sys, dict=Dict(), assignments=assignments, bf_states=bf_states
             function generated_observed(obsvar, u, p, t)
                 obs = get!(dict, value(obsvar)) do
-                    build_observed_function(sys, obsvar, var_eq_matching, var_sccs, checkbounds=checkbounds)
+                    build_observed_function(sys, obsvar, var_eq_matching, var_sccs,
+                                            checkbounds=checkbounds,
+                                            assignments=assignments,
+                                            bf_states=bf_states)
                 end
                 obs(u, p, t)
             end
@@ -286,7 +305,9 @@ function build_observed_function(
         sys, ts, var_eq_matching, var_sccs;
         expression=false,
         output_type=Array,
-        checkbounds=true
+        checkbounds=true,
+        assignments,
+        bf_states,
     )
 
     if (isscalar = !(ts isa AbstractVector))
@@ -348,7 +369,7 @@ function build_observed_function(
     end
     pre = get_postprocess_fbody(sys)
 
-    ex = Func(
+    ex = Code.toexpr(Func(
         [
          DestructuredArgs(diffvars, inbounds=!checkbounds)
          DestructuredArgs(parameters(sys), inbounds=!checkbounds)
@@ -357,13 +378,14 @@ function build_observed_function(
         [],
         pre(Let(
             [
+             assignments
              collect(Iterators.flatten(solves))
              map(eq -> eq.lhs←eq.rhs, obs[1:maxidx])
              subs
             ],
             isscalar ? ts[1] : MakeArray(ts, output_type)
            ))
-    ) |> Code.toexpr
+       ), bf_states)
 
     expression ? ex : @RuntimeGeneratedFunction(ex)
 end

src/structural_transformation/symbolics_tearing.jl

@@ -58,6 +58,18 @@ function tearing_substitution(sys::AbstractSystem; simplify=false)
     @set! sys.substitutions = nothing
 end
 
+function tearing_assignments(sys::AbstractSystem)
+    if empty_substitutions(sys)
+        assignments = []
+        bf_states = Code.LazyState()
+    else
+        subs = get_substitutions(sys)
+        assignments = [Assignment(eq.lhs, eq.rhs) for eq in subs]
+        bf_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
+    end
+    return assignments, bf_states
+end
+
 function solve_equation(eq, var, simplify)
     rhs = value(solve_for(eq, var; simplify=simplify, check=false))
     occursin(var, rhs) && error("solving $rhs for [$var] failed")

test/components.jl

@@ -1,7 +1,7 @@
 using Test
 using ModelingToolkit, OrdinaryDiffEq
 using ModelingToolkit.BipartiteGraphs
-using ModelingToolkit.StructuralTransformations: tearing_substitution
+using ModelingToolkit.StructuralTransformations
 
 function check_contract(sys)
     sys = tearing_substitution(sys)