Fix tests

Author: YingboMa

src/structural_transformation/codegen.jl

@@ -249,7 +249,7 @@ function build_torn_function(
     states_idxs = collect(diffvars_range(s))
     mass_matrix_diag = ones(length(states_idxs))
 
-    assignments, deps, bf_states = tearing_assignments(sys)
+    assignments, deps, sol_states = tearing_assignments(sys)
     invdeps = map(_->BitSet(), deps)
     for (i, d) in enumerate(deps)
         for a in d
@@ -310,16 +310,16 @@ function build_torn_function(
                  funbody
                 ))
             ),
-        bf_states
+        sol_states
     )
     if expression
         expr, states
     else
-        observedfun = let sys=sys, dict=Dict(), assignments=assignments, deps=(deps, invdeps), bf_states=bf_states, var2assignment=var2assignment
+        observedfun = let sys=sys, dict=Dict(), assignments=assignments, deps=(deps, invdeps), sol_states=sol_states, var2assignment=var2assignment
             function generated_observed(obsvar, u, p, t)
                 obs = get!(dict, value(obsvar)) do
                     build_observed_function(sys, obsvar, var_eq_matching, var_sccs,
-                                            assignments, deps, bf_states, var2assignment,
+                                            assignments, deps, sol_states, var2assignment,
                                             checkbounds=checkbounds,
                                            )
                 end
@@ -358,7 +358,7 @@ function build_observed_function(
         sys, ts, var_eq_matching, var_sccs,
         assignments,
         deps,
-        bf_states,
+        sol_states,
         var2assignment;
         expression=false,
         output_type=Array,
@@ -445,7 +445,7 @@ function build_observed_function(
             ],
             isscalar ? ts[1] : MakeArray(ts, output_type)
            ))
-       ), bf_states)
+       ), sol_states)
 
     expression ? ex : @RuntimeGeneratedFunction(ex)
 end

src/structural_transformation/symbolics_tearing.jl

@@ -62,13 +62,13 @@ function tearing_assignments(sys::AbstractSystem)
     if empty_substitutions(sys)
         assignments = []
         deps = Int[]
-        bf_states = Code.LazyState()
+        sol_states = Code.LazyState()
     else
         subs, deps = 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))
+        sol_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
     end
-    return assignments, deps, bf_states
+    return assignments, deps, sol_states
 end
 
 function solve_equation(eq, var, simplify)

src/systems/diffeqs/abstractodesystem.jl

@@ -99,24 +99,12 @@ function generate_function(
     p = map(x->time_varying_as_func(value(x), sys), ps)
     t = get_iv(sys)
 
-    if empty_substitutions(sys)
-        bf_states = Code.LazyState()
-        pre = has_difference ? (ex -> ex) : get_postprocess_fbody(sys)
-    else
-        subs, = get_substitutions(sys)
-        bf_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
-        if has_difference
-            pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs], ex)
-        else
-            process = get_postprocess_fbody(sys)
-            pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs], process(ex))
-        end
-    end
+    pre, sol_states = get_substitutions_and_solved_states(sys, no_postprocess = has_difference)
 
     if implicit_dae
-        build_function(rhss, ddvs, u, p, t; postprocess_fbody=pre, states=bf_states, kwargs...)
+        build_function(rhss, ddvs, u, p, t; postprocess_fbody=pre, states=sol_states, kwargs...)
     else
-        build_function(rhss, u, p, t; postprocess_fbody=pre, states=bf_states, kwargs...)
+        build_function(rhss, u, p, t; postprocess_fbody=pre, states=sol_states, kwargs...)
     end
 end
 

src/systems/discrete_system/discrete_system.jl

@@ -19,7 +19,7 @@ eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
 
-@named de = DiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]) # or 
+@named de = DiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]) # or
 @named de = DiscreteSystem(eqs)
 ```
 """
@@ -59,13 +59,18 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     type: type of the system
     """
     connector_type::Any
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, connector_type; checks::Bool = true)
+    """
+    substitutions: substitutions generated by tearing.
+    """
+    substitutions::Any
+
+    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, connector_type, substitutions=nothing; checks::Bool = true)
         if checks
             check_variables(dvs, iv)
             check_parameters(ps, iv)
             all_dimensionless([dvs;ps;iv;ctrls]) || check_units(discreteEqs)
         end
-        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, connector_type)
+        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, connector_type, substitutions)
     end
 end
 
@@ -103,7 +108,7 @@ function DiscreteSystem(
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
-    
+
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
@@ -187,7 +192,7 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     end
 
     u0 = varmap_to_vars(u0map,dvs; defaults=u0defs)
-    
+
     rhss = [eq.rhs for eq in eqs]
     u = dvs
     p = varmap_to_vars(parammap,ps; defaults=pdefs)
@@ -226,15 +231,15 @@ function linearize_eqs(sys, eqs=get_eqs(sys); return_max_delay=false)
         end
 
         all(length.(unique.(values(state_ops))) .<= 1) || error("Each state should be used with single difference operator.")
-        
+
         dts_gcd = Dict()
         for v in keys(dts)
             dts_gcd[v] = (length(dts[v]) > 0) ? first(dts[v]) : nothing
         end
 
         lin_eqs = [
             v(get_iv(sys) - (t)) ~ v(get_iv(sys) - (t-dts_gcd[v]))
-            for v in unique_states if max_delay[v] > 0 && dts_gcd[v]!==nothing for t in collect(max_delay[v]:(-dts_gcd[v]):0)[1:end-1] 
+            for v in unique_states if max_delay[v] > 0 && dts_gcd[v]!==nothing for t in collect(max_delay[v]:(-dts_gcd[v]):0)[1:end-1]
         ]
         eqs = vcat(eqs, lin_eqs)
     end
@@ -256,12 +261,12 @@ function generate_function(
     )
     eqs = equations(sys)
     foreach(check_difference_variables, eqs)
-    # substitute x(t) by just x
     rhss = [eq.rhs for eq in eqs]
 
     u = map(x->time_varying_as_func(value(x), sys), dvs)
     p = map(x->time_varying_as_func(value(x), sys), ps)
     t = get_iv(sys)
-    
-    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...)
 end

src/systems/jumps/jumpsystem.jl

@@ -13,7 +13,7 @@ $(FIELDS)
 ```julia
 using ModelingToolkit
 
-@parameters β γ 
+@parameters β γ
 @variables t S(t) I(t) R(t)
 rate₁   = β*S*I
 affect₁ = [S ~ S - 1, I ~ I + 1]
@@ -301,7 +301,7 @@ function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
         vtoj = nothing; jtov = nothing; jtoj = nothing
     end
 
-    JumpProblem(prob, aggregator, jset; dep_graph=jtoj, vartojumps_map=vtoj, jumptovars_map=jtov, 
+    JumpProblem(prob, aggregator, jset; dep_graph=jtoj, vartojumps_map=vtoj, jumptovars_map=jtov,
                 scale_rates=false, nocopy=true, kwargs...)
 end
 

src/systems/nonlinear/nonlinearsystem.jl

@@ -54,11 +54,16 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
     type: type of the system
     """
     connector_type::Any
-    function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connector_type; checks::Bool = true)
+    """
+    substitutions: substitutions generated by tearing.
+    """
+    substitutions::Any
+
+    function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connector_type, substitutions=nothing; checks::Bool = true)
         if checks
             all_dimensionless([states;ps]) ||check_units(eqs)
         end
-        new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connector_type)
+        new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connector_type, substitutions)
     end
 end
 
@@ -123,19 +128,14 @@ end
 function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys);
                            sparse = false, simplify=false, kwargs...)
     jac = calculate_jacobian(sys,sparse=sparse, simplify=simplify)
-    return build_function(jac, vs, ps;
-                          conv = AbstractSysToExpr(sys), kwargs...)
+    return build_function(jac, vs, ps; kwargs...)
 end
 
 function generate_function(sys::NonlinearSystem, dvs = states(sys), ps = parameters(sys); kwargs...)
-    #obsvars = map(eq->eq.lhs, observed(sys))
-    #fulldvs = [dvs; obsvars]
-
     rhss = [deq.rhs for deq ∈ equations(sys)]
-    #rhss = Let(obss, rhss)
+    pre, sol_states = get_substitutions_and_solved_states(sys)
 
-    return build_function(rhss, value.(dvs), value.(ps);
-                          conv = AbstractSysToExpr(sys), kwargs...)
+    return build_function(rhss, value.(dvs), value.(ps); postprocess_fbody=pre, states=sol_states, kwargs...)
 end
 
 jacobian_sparsity(sys::NonlinearSystem) =

src/utils.jl

@@ -432,3 +432,20 @@ function empty_substitutions(sys)
     subs = get_substitutions(sys)
     isnothing(subs) || isempty(last(subs))
 end
+
+function get_substitutions_and_solved_states(sys; no_postprocess=false)
+    if empty_substitutions(sys)
+        sol_states = Code.LazyState()
+        pre = no_postprocess ? (ex -> ex) : get_postprocess_fbody(sys)
+    else
+        subs, = get_substitutions(sys)
+        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)
+        else
+            process = get_postprocess_fbody(sys)
+            pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs], process(ex))
+        end
+    end
+    return pre, sol_states
+end

test/reduction.jl

@@ -44,7 +44,7 @@ io = IOBuffer(); show(io, MIME("text/plain"), lorenz1_aliased); str = String(tak
 @test all(s->occursin(s, str), ["lorenz1", "States (2)", "Parameters (3)"])
 reduced_eqs = [
                D(x) ~ σ*(y - x)
-               D(y) ~ β + x*(ρ - (x - y)) - y
+               D(y) ~ β + x*(ρ - z) - y
               ]
 test_equal.(equations(lorenz1_aliased), reduced_eqs)
 @test isempty(setdiff(states(lorenz1_aliased), [x, y, z]))
@@ -84,10 +84,10 @@ __x = x
 # Reduced Flattened System
 
 reduced_system = structural_simplify(connected)
-reduced_system2 = structural_simplify(structural_simplify(structural_simplify(connected)))
+reduced_system2 = structural_simplify(tearing_substitution(structural_simplify(tearing_substitution(structural_simplify(connected)))))
 
 @test isempty(setdiff(states(reduced_system), states(reduced_system2)))
-@test isequal(equations(reduced_system), equations(reduced_system2))
+@test isequal(equations(tearing_substitution(reduced_system)), equations(reduced_system2))
 @test isequal(observed(reduced_system), observed(reduced_system2))
 @test setdiff(states(reduced_system), [
         s
@@ -155,7 +155,7 @@ let
     reduced_sys = structural_simplify(connected)
     ref_eqs = [
                D(ol.x) ~ ol.a*ol.x + ol.b*ol.u
-               0 ~ pc.k_P*(ol.c*ol.x + ol.d*ol.u) - ol.u
+               0 ~ pc.k_P*ol.y - ol.u
               ]
     @test ref_eqs == equations(reduced_sys)
 end
@@ -254,7 +254,7 @@ eq = [
 @named sys0 = ODESystem(eq, t)
 sys = structural_simplify(sys0)
 @test length(equations(sys)) == 1
-eq = equations(sys)[1]
+eq = equations(tearing_substitution(sys))[1]
 @test isequal(eq.lhs, 0)
 dv25 = ModelingToolkit.value(ModelingToolkit.derivative(eq.rhs, v25))
 ddv25 = ModelingToolkit.value(ModelingToolkit.derivative(eq.rhs, D(v25)))

test/structural_transformation/tearing.jl

@@ -122,7 +122,7 @@ end
 # solve for
 # 0 = u5 - hypot(sin(u5), hypot(cos(sin(u5)), hypot(sin(u5), cos(sin(u5)))))
 tornsys = tearing(sys)
-@test isequal(equations(tornsys), [0 ~ u5 + (-1 * hypot(hypot(cos(sin(u5)), hypot(sin(u5), cos(sin(u5)))), sin(u5)))])
+@test isequal(equations(tornsys), [0 ~ u5 - hypot(u4, u1)])
 prob = NonlinearProblem(tornsys, ones(1))
 sol = solve(prob, NewtonRaphson())
 @test norm(prob.f(sol.u, sol.prob.p)) < 1e-10
@@ -147,7 +147,7 @@ let (mm, _, _) = ModelingToolkit.aag_bareiss(nlsys)
 end
 
 newsys = tearing(nlsys)
-@test length(equations(newsys)) == 0
+@test length(equations(newsys)) <= 1
 
 ###
 ### DAE system
@@ -163,7 +163,8 @@ eqs = [
       ]
 @named daesys = ODESystem(eqs, t)
 newdaesys = tearing(daesys)
-@test equations(newdaesys) == [D(x) ~ z; 0 ~ x + sin(z) - p*t]
+@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 isequal(states(newdaesys), [x, z])
 prob = ODAEProblem(newdaesys, [x=>1.0], (0, 1.0), [p=>0.2])
 du = [0.0]; u = [1.0]; pr = 0.2; tt = 0.1