handle dds and t

Author: ChrisRackauckas

src/structural_transformation/StructuralTransformations.jl

@@ -23,7 +23,8 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
                        invalidate_cache!, Substitutions, get_or_construct_tearing_state,
                        filter_kwargs, lower_varname, setio, SparseMatrixCLIL,
-                       fast_substitute, get_fullvars, has_equations, observed
+                       fast_substitute, get_fullvars, has_equations, observed,
+                       Schedule
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview, complete

src/structural_transformation/symbolics_tearing.jl

@@ -555,6 +555,12 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [fast_substitute(observed(sys), obs_sub); subeqs]
     @set! sys.observed = obs
+
+    # Only makes sense for time-dependent
+    # TODO: generalize to SDE
+    if sys isa ODESystem
+        @set! sys.schedule = Schedule(var_eq_matching, dummy_sub)
+    end   
     @set! state.sys = sys
     @set! sys.tearing_state = state
     return invalidate_cache!(sys)

src/systems/abstractsystem.jl

@@ -576,6 +576,7 @@ for prop in [:eqs
              :preface
              :torn_matching
              :initializesystem
+             :schedule
              :tearing_state
              :substitutions
              :metadata

src/systems/diffeqs/abstractodesystem.jl

@@ -1,3 +1,8 @@
+struct Schedule
+    var_eq_matching
+    dummy_sub
+end
+
 function filter_kwargs(kwargs)
     kwargs = Dict(kwargs)
     for key in keys(kwargs)
@@ -326,7 +331,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_expression ?
-                   (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
+                   ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
                    f_gen
     f(u, p, t) = f_oop(u, p, t)
     f(du, u, p, t) = f_iip(du, u, p, t)
@@ -351,7 +356,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
             expression_module = eval_module,
             checkbounds = checkbounds, kwargs...)
         tgrad_oop, tgrad_iip = eval_expression ?
-                               (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in tgrad_gen) :
+                               ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in tgrad_gen) :
                                tgrad_gen
         if p isa Tuple
             __tgrad(u, p, t) = tgrad_oop(u, p..., t)
@@ -373,7 +378,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
             expression_module = eval_module,
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
+                           ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
                            jac_gen
         _jac(u, p, t) = jac_oop(u, p, t)
         _jac(J, u, p, t) = jac_iip(J, u, p, t)
@@ -541,7 +546,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_expression ?
-                   (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
+                   ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
                    f_gen
     f(du, u, p, t) = f_oop(du, u, p, t)
     f(du, u, p::MTKParameters, t) = f_oop(du, u, p..., t)
@@ -555,7 +560,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
             expression_module = eval_module,
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_expression ?
-                           (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
+                           ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
                            jac_gen
         _jac(du, u, p, ˍ₋gamma, t) = jac_oop(du, u, p, ˍ₋gamma, t)
         _jac(du, u, p::MTKParameters, ˍ₋gamma, t) = jac_oop(du, u, p..., ˍ₋gamma, t)
@@ -624,7 +629,7 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         expression = Val{true},
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
-    f_oop, f_iip = (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
+    f_oop, f_iip = ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
     f(u, h, p, t) = f_oop(u, h, p, t)
     f(u, h, p::MTKParameters, t) = f_oop(u, h, p..., t)
     f(du, u, h, p, t) = f_iip(du, u, h, p, t)
@@ -649,10 +654,10 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
         expression = Val{true},
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
-    f_oop, f_iip = (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
+    f_oop, f_iip = ((@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
     g_gen = generate_diffusion_function(sys, dvs, ps; expression = Val{true},
         isdde = true, kwargs...)
-    g_oop, g_iip = (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in g_gen)
+    g_oop, g_iip = ((@RuntimeGeneratedFunction(ex)) for ex in g_gen)
     f(u, h, p, t) = f_oop(u, h, p, t)
     f(u, h, p::MTKParameters, t) = f_oop(u, h, p..., t)
     f(du, u, h, p, t) = f_iip(du, u, h, p, t)
@@ -849,6 +854,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         symbolic_u0 = false,
         u0_constructor = identity,
         guesses = Dict(),
+        t = nothing,
         warn_initialize_determined = true,
         kwargs...)
     eqs = equations(sys)
@@ -870,7 +876,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             u0map = unknowns(sys) .=> u0map
         end
         initializeprob = ModelingToolkit.InitializationProblem(
-            sys, u0map, parammap; guesses, warn_initialize_determined)
+            sys, t, u0map, parammap; guesses, warn_initialize_determined)
         initializeprobmap = getu(initializeprob, unknowns(sys))
 
         zerovars = setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0
@@ -1101,6 +1107,7 @@ function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem, du0map, u0map, tspan
     end
     f, du0, u0, p = process_DEProblem(DAEFunction{iip}, sys, u0map, parammap;
         implicit_dae = true, du0map = du0map, check_length,
+        t = tspan !== nothing ? tspan[1] : tspan,
         warn_initialize_determined, kwargs...)
     diffvars = collect_differential_variables(sys)
     sts = unknowns(sys)
@@ -1277,6 +1284,7 @@ function ODEProblemExpr{iip}(sys::AbstractODESystem, u0map, tspan,
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `ODEProblemExpr`")
     end
     f, u0, p = process_DEProblem(ODEFunctionExpr{iip}, sys, u0map, parammap; check_length,
+        t = tspan !== nothing ? tspan[1] : tspan,
         kwargs...)
     linenumbers = get(kwargs, :linenumbers, true)
     kwargs = filter_kwargs(kwargs)
@@ -1322,6 +1330,7 @@ function DAEProblemExpr{iip}(sys::AbstractODESystem, du0map, u0map, tspan,
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEProblemExpr`")
     end
     f, du0, u0, p = process_DEProblem(DAEFunctionExpr{iip}, sys, u0map, parammap;
+        t = tspan !== nothing ? tspan[1] : tspan,
         implicit_dae = true, du0map = du0map, check_length,
         kwargs...)
     linenumbers = get(kwargs, :linenumbers, true)
@@ -1505,11 +1514,11 @@ function InitializationProblem(sys::AbstractODESystem, args...; kwargs...)
     InitializationProblem{true}(sys, args...; kwargs...)
 end
 
-function InitializationProblem(sys::AbstractODESystem,
+function InitializationProblem(sys::AbstractODESystem, t,
         u0map::StaticArray,
         args...;
         kwargs...)
-    InitializationProblem{false, SciMLBase.FullSpecialize}(sys, u0map, args...; kwargs...)
+    InitializationProblem{false, SciMLBase.FullSpecialize}(sys, t, u0map, args...; kwargs...)
 end
 
 function InitializationProblem{true}(sys::AbstractODESystem, args...; kwargs...)
@@ -1520,7 +1529,8 @@ function InitializationProblem{false}(sys::AbstractODESystem, args...; kwargs...
     InitializationProblem{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 
-function InitializationProblem{iip, specialize}(sys::AbstractODESystem, u0map = [],
+function InitializationProblem{iip, specialize}(sys::AbstractODESystem, 
+        t, u0map = [],
         parammap = DiffEqBase.NullParameters();
         guesses = [],
         check_length = true,
@@ -1530,6 +1540,8 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEProblem`")
     end
 
+    @show u0map
+
     if isempty(u0map) && get_initializesystem(sys) !== nothing
         isys = get_initializesystem(sys)
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
@@ -1548,6 +1560,8 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if warn_initialize_determined && neqs < nunknown
         @warn "Initialization system is underdetermined. $neqs equations for $nunknown unknowns. Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false."
     end
+    
+    parammap isa DiffEqBase.NullParameters ? [independent_variable(sys) => t] : merge(todict(parammap), Dict(independent_variable(sys) => t))
 
     if neqs == nunknown
         NonlinearProblem(isys, guesses, parammap)

src/systems/diffeqs/odesystem.jl

@@ -101,6 +101,10 @@ struct ODESystem <: AbstractODESystem
     """
     initializesystem::Union{Nothing, NonlinearSystem}
     """
+    The schedule for the code generation process.
+    """
+    schedule::Any
+    """
     Type of the system.
     """
     connector_type::Any
@@ -165,11 +169,13 @@ struct ODESystem <: AbstractODESystem
     """
     parent::Any
 
+
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
             jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses,
-            torn_matching, initializesystem, connector_type, preface, cevents,
-            devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
-            tearing_state = nothing,
+            torn_matching, initializesystem, schedule, connector_type, preface, cevents,
+            devents, parameter_dependencies,
+            metadata = nothing, gui_metadata = nothing,
+            tearing_state = nothing, 
             substitutions = nothing, complete = false, index_cache = nothing,
             discrete_subsystems = nothing, solved_unknowns = nothing,
             split_idxs = nothing, parent = nothing; checks::Union{Bool, Int} = true)
@@ -185,7 +191,8 @@ struct ODESystem <: AbstractODESystem
         end
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, torn_matching,
-            initializesystem, connector_type, preface, cevents, devents, parameter_dependencies, metadata,
+            initializesystem, schedule, connector_type, preface, cevents, devents, parameter_dependencies, 
+            metadata,
             gui_metadata, tearing_state, substitutions, complete, index_cache,
             discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
@@ -202,12 +209,13 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         defaults = _merge(Dict(default_u0), Dict(default_p)),
         guesses = Dict(),
         initializesystem = nothing,
+        schedule = nothing,
         connector_type = nothing,
         preface = nothing,
         continuous_events = nothing,
         discrete_events = nothing,
         parameter_dependencies = nothing,
-        checks = true,
+        checks = true,        
         metadata = nothing,
         gui_metadata = nothing)
     name === nothing &&
@@ -253,7 +261,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, nothing, initializesystem,
-        connector_type, preface, cont_callbacks, disc_callbacks, parameter_dependencies,
+        schedule, connector_type, preface, cont_callbacks, disc_callbacks, parameter_dependencies,
         metadata, gui_metadata, checks = checks)
 end
 

src/systems/nonlinear/initializesystem.jl

@@ -6,7 +6,9 @@ Generate `NonlinearSystem` which initializes an ODE problem from specified initi
 function generate_initializesystem(sys::ODESystem;
         u0map = Dict(),
         name = nameof(sys),
-        guesses = Dict(), check_defguess = false, kwargs...)
+        guesses = Dict(), check_defguess = false, 
+        default_dd_value = 0.0,
+        kwargs...)
     sts, eqs = unknowns(sys), equations(sys)
     idxs_diff = isdiffeq.(eqs)
     idxs_alge = .!idxs_diff
@@ -18,11 +20,18 @@ function generate_initializesystem(sys::ODESystem;
     defs = merge(defaults(sys), todict(u0map))
 
     full_states = [sts; getfield.((observed(sys)), :lhs)]
+    set_full_states = Set(full_states)
+    guesses = todict(guesses)
+    schedule = getfield(sys, :schedule)
 
-    # Refactor to ODESystem construction
-    # should be ModelingToolkit.guesses(sys)
+    dd_guess = if schedule !== nothing
+        guessmap = [x[2]=>get(guesses, x[1], default_dd_value) for x in schedule.dummy_sub]
+        Dict(filter(x->!isnothing(x[1]) && x[1]∈set_full_states,guessmap))
+    else
+        Dict()
+    end
 
-    guesses = merge(get_guesses(sys), todict(guesses))
+    guesses = merge(get_guesses(sys), todict(guesses), dd_guess)
 
     for st in full_states
         if st ∈ keys(defs)
@@ -44,13 +53,13 @@ function generate_initializesystem(sys::ODESystem;
         end
     end
 
-    pars = parameters(sys)
+    pars = [parameters(sys); independent_variable(sys)]
     nleqs = [eqs_ics; observed(sys)]
-
+    
     sys_nl = NonlinearSystem(nleqs,
         full_states,
         pars;
-        defaults = merge(ModelingToolkit.defaults(sys), todict(u0)),
+        defaults = merge(ModelingToolkit.defaults(sys), todict(u0), dd_guess),
         name,
         kwargs...)
 

src/systems/systems.jl

@@ -26,7 +26,13 @@ function structural_simplify(
     else
         newsys = newsys′
     end
-    @set! newsys.parent = complete(sys; split)
+    if newsys isa ODESystem 
+        schedule = newsys.schedule
+        @set! newsys.parent = complete(sys; split)
+        @set! newsys.schedule = schedule
+    else
+        @set! newsys.parent = complete(sys; split)
+    end
     newsys = complete(newsys; split)
     if newsys′ isa Tuple
         idxs = [parameter_index(newsys, i) for i in io[1]]