Port two ipo tests from v1

Author: Keno

Manifest.toml

@@ -119,12 +119,25 @@ function structural_transformation!(state::TransformationState)
             first = false
             continue
         end
-        
+
         return StateSelection.complete(var_eq_matching, nsrcs(state.structure.graph))
     end
 end
 
 using StateSelection: Unassigned, SelectedState, unassigned
+
+struct StateInvariant; end
+StateSelection.BipartiteGraphs.overview_label(::Type{StateInvariant}) = ('P', "State Invariant / Parameter", :red)
+
+struct InOut
+    ordinal::Int
+end
+StateSelection.BipartiteGraphs.overview_label(::Type{InOut}) = ('#', "IPO in var / out eq", :green)
+StateSelection.BipartiteGraphs.overview_label(io::InOut) = (string(io.ordinal), "IPO in var / out eq", :green)
+
+const IPOMatches = Union{Unassigned, SelectedState, StateInvariant, InOut}
+const IPOMatching = StateSelection.Matching{IPOMatches}
+
 function top_level_state_selection!(tstate)
     (; result, structure) = tstate
     # For the top-level problem, all external vars are state-invariant, and we do no other fissioning
@@ -135,16 +148,20 @@ function top_level_state_selection!(tstate)
     StateSelection.complete!(structure)
 
     ## Part 1: Perform the selection of differential states and subsequent tearing of the
-    #          non-linear problem at every time step. 
+    #          non-linear problem at every time step.
 
-    var_eq_matching = StateSelection.partial_state_selection_graph!(structure, highest_diff_max_match)
+    var_eq_matching = convert(IPOMatching, StateSelection.partial_state_selection_graph!(structure, highest_diff_max_match))
 
     diff_vars = BitSet()
     alg_vars = BitSet()
     explicit_eqs = BitSet()
 
     for (v, match) in enumerate(var_eq_matching)
-        v in param_vars && continue
+        if v in param_vars
+            @assert match === unassigned
+            var_eq_matching[v] = StateInvariant()
+            continue
+        end
         if match === SelectedState()
             push!(diff_vars, v)
         elseif match === unassigned
@@ -164,14 +181,19 @@ function top_level_state_selection!(tstate)
     varfilter(var) = varkind(result, structure, var) == Intrinsics.Continuous && !(var <= result.nexternalvars)
 
     ## Part 2: Perform the selection of differential states and subsequent tearing of the
-    #          non-linear problem at every time step. 
-    init_var_eq_matching = StateSelection.complete(StateSelection.maximal_matching(structure.graph;
+    #          non-linear problem at every time step.
+    init_var_eq_matching = StateSelection.complete(StateSelection.maximal_matching(structure.graph, IPOMatches;
         dstfilter = varfilter, srcfilter = eq->eqkind(result, structure, eq) in (Intrinsics.Always, Intrinsics.Initial)), nsrcs(structure.graph))
-    init_var_eq_matching = StateSelection.pss_graph_modia!(structure, init_var_eq_matching)
+    init_var_eq_matching = convert(IPOMatching, StateSelection.pss_graph_modia!(structure, init_var_eq_matching))
 
     init_state_vars = BitSet()
     init_explicit_eqs = BitSet()
     for (v, match) in enumerate(init_var_eq_matching)
+        if v in param_vars
+            @assert match === unassigned
+            init_var_eq_matching[v] = StateInvariant()
+            continue
+        end
         varfilter(v) || continue
         if match === unassigned
             push!(init_state_vars, v)

src/transform/state_selection.jl

@@ -55,15 +55,6 @@ end
 
 Base.IteratorSize(::Type{Compiler.UseRefIterator}) = Base.SizeUnknown()
 
-struct StateInvariant; end
-StateSelection.BipartiteGraphs.overview_label(::Type{StateInvariant}) = ('P', "State Invariant / Parameter", :red)
-
-struct InOut
-    ordinal::Int
-end
-StateSelection.BipartiteGraphs.overview_label(::Type{InOut}) = ('#', "IPO in var / out eq", :green)
-StateSelection.BipartiteGraphs.overview_label(io::InOut) = (string(io.ordinal), "IPO in var / out eq", :green)
-
 function schedule_incidence!(compact, var_eq_matching, curval, ::Type, var, line; vars=nothing, schedule_missing_var! = nothing)
     # This just needs the linear part, which is `0` in `Type`
     return (curval, nothing)
@@ -511,7 +502,7 @@ function matching_for_key(result::DAEIPOResult, key::TornCacheKey, structure = m
     may_use_eq(eq) = !(eq in explicit_eqs) && eqclassification(result, structure, eq) != External && eqkind(result, structure, eq) in (allow_init_eqs ? (Intrinsics.Initial, Intrinsics.Always) : (Intrinsics.Always,))
 
     # Max match is the (unique) tearing result given the choice of states
-    var_eq_matching = StateSelection.complete(StateSelection.maximal_matching(structure.graph, Union{Unassigned, SelectedState, StateInvariant, InOut};
+    var_eq_matching = StateSelection.complete(StateSelection.maximal_matching(structure.graph, IPOMatches;
         dstfilter = may_use_var, srcfilter = may_use_eq), nsrcs(structure.graph))
 
     if diff_states !== nothing

src/transform/tearing/schedule.jl

@@ -26,11 +26,51 @@ function twocall!()
 end
 
 twocall!()
-sol = solve(DAECProblem(twocall!, (1, 2) .=> 1.), IDA())
-@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
-@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[2, :], exp.(sol.t)))
-sol = solve(ODECProblem(twocall!, (1, 2) .=> 1.), Rodas5(autodiff=false))
-@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
-@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[2, :], exp.(sol.t)))
+dae_sol = solve(DAECProblem(twocall!, (1, 2) .=> 1.), IDA())
+ode_sol = solve(ODECProblem(twocall!, (1, 2) .=> 1.), Rodas5(autodiff=false))
+for (sol, i) in Iterators.product((dae_sol, ode_sol), 1:2)
+    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[i, :], exp.(sol.t)))
+end
+
+#= + NonLinear =#
+@noinline function sin!()
+    x = continuous()
+    always!(ddt(x) - sin(x))
+end
+function sin2!()
+    sin!(); sin!();
+    return nothing
+end
+dae_sol = solve(DAECProblem(sin2!, (1, 2) .=> 1.), IDA())
+ode_sol = solve(ODECProblem(sin2!, (1, 2) .=> 1.), Rodas5(autodiff=false))
+for (sol, i) in Iterators.product((dae_sol, ode_sol), 1:2)
+    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[i, :], 2*acot.(exp.(-sol.t).*cot(1/2))))
+end
+
+#= + SICM =#
+struct sicm!
+    arg::Float64
+end
+
+@noinline function (this::sicm!)()
+    x = continuous()
+    always!(ddt(x) - this.arg)
+end
+
+struct sicm2!
+    a::Float64
+    b::Float64
+end
+
+function (this::sicm2!)()
+    sicm!(this.a)(); sicm!(this.b)();
+    return nothing
+end
+dae_sol = solve(DAECProblem(sicm2!(1., 1.), (1, 2) .=> 1.), IDA())
+ode_sol = solve(ODECProblem(sicm2!(1., 1.), (1, 2) .=> 1.), Rodas5(autodiff=false))
+for (sol, i) in Iterators.product((dae_sol, ode_sol), 1:2)
+    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[i, :], 1. .+ sol.t))
+end
+
 
 end