Merge pull request #2016 from SciML/myb/solver_states

Add `unknown_states`

Author: YingboMa

src/structural_transformation/codegen.jl

@@ -297,7 +297,8 @@ function build_torn_function(sys;
                        out,
                        rhss)
 
-    states = fullvars[states_idxs]
+    states = Any[fullvars[i] for i in states_idxs]
+    @set! sys.unknown_states = states
     syms = map(Symbol, states)
 
     pre = get_postprocess_fbody(sys)
@@ -402,7 +403,7 @@ function build_observed_function(state, ts, var_eq_matching, var_sccs,
 
     fullvars = state.fullvars
     s = state.structure
-    solver_states = fullvars[is_solver_state_idxs]
+    unknown_states = fullvars[is_solver_state_idxs]
     algvars = fullvars[.!is_solver_state_idxs]
 
     required_algvars = Set(intersect(algvars, vars))
@@ -489,7 +490,7 @@ function build_observed_function(state, ts, var_eq_matching, var_sccs,
     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)
+    ex = Code.toexpr(Func([DestructuredArgs(unknown_states, inbounds = !checkbounds)
                            DestructuredArgs(parameters(sys), inbounds = !checkbounds)
                            independent_variables(sys)],
                           [],

src/systems/abstractsystem.jl

@@ -213,7 +213,8 @@ for prop in [:eqs
              :tearing_state
              :substitutions
              :metadata
-             :discrete_subsystems]
+             :discrete_subsystems
+             :unknown_states]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
     @eval begin
@@ -466,7 +467,7 @@ function namespace_expr(O, sys, n = nameof(sys))
     end
 end
 
-function SymbolicIndexingInterface.states(sys::AbstractSystem)
+function states(sys::AbstractSystem)
     sts = get_states(sys)
     systems = get_systems(sys)
     unique(isempty(systems) ?
@@ -580,8 +581,21 @@ end
 
 SymbolicIndexingInterface.is_indep_sym(sys::AbstractSystem, sym) = isequal(sym, get_iv(sys))
 
+"""
+$(SIGNATURES)
+
+Return a list of actual states needed to be solved by solvers.
+"""
+function unknown_states(sys::AbstractSystem)
+    sts = states(sys)
+    if has_unknown_states(sys)
+        sts = something(get_unknown_states(sys), sts)
+    end
+    return sts
+end
+
 function SymbolicIndexingInterface.state_sym_to_index(sys::AbstractSystem, sym)
-    findfirst(isequal(sym), SymbolicIndexingInterface.states(sys))
+    findfirst(isequal(sym), unknown_states(sys))
 end
 function SymbolicIndexingInterface.is_state_sym(sys::AbstractSystem, sym)
     !isnothing(SymbolicIndexingInterface.state_sym_to_index(sys, sym))

src/systems/diffeqs/odesystem.jl

@@ -127,16 +127,22 @@ struct ODESystem <: AbstractODESystem
     """
     complete::Bool
     """
-    discrete_subsystems: a list of discrete subsystems
+    discrete_subsystems: a list of discrete subsystems.
     """
     discrete_subsystems::Any
+    """
+    unknown_states: a list of actual states needed to be solved by solvers. Only
+    used for ODAEProblem.
+    """
+    unknown_states::Union{Nothing, Vector{Any}}
 
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
                        jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
                        torn_matching, connector_type, preface, cevents,
                        devents, metadata = nothing, tearing_state = nothing,
                        substitutions = nothing, complete = false,
-                       discrete_subsystems = nothing; checks::Union{Bool, Int} = true)
+                       discrete_subsystems = nothing, unknown_states = nothing;
+                       checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
             check_parameters(ps, iv)
@@ -149,7 +155,7 @@ struct ODESystem <: AbstractODESystem
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
             connector_type, preface, cevents, devents, metadata, tearing_state,
-            substitutions, complete, discrete_subsystems)
+            substitutions, complete, discrete_subsystems, unknown_states)
     end
 end