feat: update clock inference codegen, fix bugs

Author: AayushSabharwal

src/discretedomain.jl

@@ -28,7 +28,7 @@ struct Shift <: Operator
     Shift(t, steps = 1) = new(value(t), steps)
 end
 Shift(steps::Int) = new(nothing, steps)
-normalize_to_differential(s::Shift) = Differential(s.t)^s.steps
+normalize_to_differential(s::Shift) = Differential(s.t)^abs(s.steps)
 function (D::Shift)(x, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
     Term{symtype(x)}(D, Any[x])
@@ -114,7 +114,7 @@ Base.hash(D::Sample, u::UInt) = hash(D.clock, xor(u, 0x055640d6d952f101))
 
 Returns true if the expression or equation `O` contains [`Sample`](@ref) terms.
 """
-hassample(O) = recursive_hasoperator(Sample, O)
+hassample(O) = recursive_hasoperator(Sample, unwrap(O))
 
 # Hold
 
@@ -140,7 +140,7 @@ Hold(x) = Hold()(x)
 
 Returns true if the expression or equation `O` contains [`Hold`](@ref) terms.
 """
-hashold(O) = recursive_hasoperator(Hold, O)
+hashold(O) = recursive_hasoperator(Hold, unwrap(O))
 
 # ShiftIndex
 

src/systems/clock_inference.jl

@@ -139,15 +139,15 @@ function split_system(ci::ClockInference{S}) where {S}
     return tss, inputs, continuous_id, id_to_clock
 end
 
-function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
+function generate_discrete_affect(osys::AbstractODESystem, syss, inputs, continuous_id, id_to_clock;
         checkbounds = true,
         eval_module = @__MODULE__, eval_expression = true)
     @static if VERSION < v"1.7"
         error("The `generate_discrete_affect` function requires at least Julia 1.7")
     end
     out = Sym{Any}(:out)
     appended_parameters = parameters(syss[continuous_id])
-    param_to_idx = Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
+    param_to_idx = Dict{Any, ParameterIndex}(p => parameter_index(osys, p) for p in appended_parameters)
     offset = length(appended_parameters)
     affect_funs = []
     init_funs = []
@@ -172,11 +172,13 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
             vv = arguments(v)[1]
             if vv in fullvars
                 push!(needed_disc_to_cont_obs, vv)
-                push!(disc_to_cont_idxs, param_to_idx[v])
+                # @show param_to_idx[v] v
+                # @assert param_to_idx[v].portion isa SciMLStructures.Discrete # TOOD: remove
+                push!(disc_to_cont_idxs, param_to_idx[v].idx)
             end
         end
         append!(appended_parameters, input, unknowns(sys))
-        cont_to_disc_obs = build_explicit_observed_function(syss[continuous_id],
+        cont_to_disc_obs = build_explicit_observed_function(osys,
             needed_cont_to_disc_obs,
             throw = false,
             expression = true,
@@ -185,30 +187,31 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
         disc_to_cont_obs = build_explicit_observed_function(sys, needed_disc_to_cont_obs,
             throw = false,
             expression = true,
-            output_type = SVector)
-        ni = length(input)
-        ns = length(unknowns(sys))
+            output_type = SVector,
+            ps = reorder_parameters(osys, parameters(sys)))
         disc = Func([
                 out,
-                DestructuredArgs(unknowns(sys)),
-                DestructuredArgs(appended_parameters),
+                DestructuredArgs(unknowns(osys)),
+                DestructuredArgs.(reorder_parameters(osys, parameters(osys)))...,
+                # DestructuredArgs(appended_parameters),
                 get_iv(sys),
             ], [],
             let_block)
-        cont_to_disc_idxs = (offset + 1):(offset += ni)
-        input_offset = offset
-        disc_range = (offset + 1):(offset += ns)
+        cont_to_disc_idxs = [parameter_index(osys, sym).idx for sym in input]
+        disc_range = [parameter_index(osys, sym).idx for sym in unknowns(sys)]
         save_vec = Expr(:ref, :Float64)
-        for i in 1:ns
-            push!(save_vec.args, :(p[$(input_offset + i)]))
+        for unk in unknowns(sys)
+            idx = parameter_index(osys, unk).idx
+            push!(save_vec.args, :(discretes[$idx]))
         end
         empty_disc = isempty(disc_range)
-
         disc_init = :(function (p, t)
             d2c_obs = $disc_to_cont_obs
-            d2c_view = view(p, $disc_to_cont_idxs)
-            disc_state = view(p, $disc_range)
-            copyto!(d2c_view, d2c_obs(disc_state, p, t))
+            discretes, repack, _ = $(SciMLStructures.canonicalize)($(SciMLStructures.Discrete()), p)
+            d2c_view = view(discretes, $disc_to_cont_idxs)
+            disc_state = view(discretes, $disc_range)
+            copyto!(d2c_view, d2c_obs(disc_state, p..., t))
+            repack(discretes)
         end)
 
         # @show disc_to_cont_idxs
@@ -220,10 +223,11 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
             c2d_obs = $cont_to_disc_obs
             d2c_obs = $disc_to_cont_obs
             # Like Sample
-            c2d_view = view(p, $cont_to_disc_idxs)
+            discretes, repack, _ = $(SciMLStructures.canonicalize)($(SciMLStructures.Discrete()), p)
+            c2d_view = view(discretes, $cont_to_disc_idxs)
             # Like Hold
-            d2c_view = view(p, $disc_to_cont_idxs)
-            disc_unknowns = view(p, $disc_range)
+            d2c_view = view(discretes, $disc_to_cont_idxs)
+            disc_unknowns = view(discretes, $disc_range)
             disc = $disc
 
             push!(saved_values.t, t)
@@ -238,12 +242,13 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
             # d2c comes last
             # @show t
             # @show "incoming", p
-            copyto!(c2d_view, c2d_obs(integrator.u, p, t))
+            copyto!(c2d_view, c2d_obs(integrator.u, p..., t))
             # @show "after c2d", p
-            $empty_disc || disc(disc_unknowns, disc_unknowns, p, t)
+            $empty_disc || disc(disc_unknowns, integrator.u, p..., t)
             # @show "after state update", p
-            copyto!(d2c_view, d2c_obs(disc_unknowns, p, t))
+            copyto!(d2c_view, d2c_obs(disc_unknowns, p..., t))
             # @show "after d2c", p
+            repack(discretes)
         end)
         sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)

src/systems/diffeqs/abstractodesystem.jl

@@ -800,7 +800,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         u0 = u0_constructor(u0)
     end
 
-    p = MTKParameters(sys, parammap; toterm = default_toterm)
+    p = MTKParameters(sys, parammap)
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
@@ -931,7 +931,7 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     cbs = process_events(sys; callback, kwargs...)
     inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -1035,7 +1035,7 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     cbs = process_events(sys; callback, kwargs...)
     inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -1097,7 +1097,7 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     cbs = process_events(sys; callback, kwargs...)
     inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)

src/systems/diffeqs/odesystem.jl

@@ -172,8 +172,8 @@ 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, gui_metadata,
-            tearing_state, substitutions, complete, index_cache, discrete_subsystems,
-            solved_unknowns, split_idxs, parent)
+            tearing_state, substitutions, complete, index_cache,
+            discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
 end
 

src/systems/index_cache.jl

@@ -28,9 +28,8 @@ function IndexCache(sys::AbstractSystem)
     unks = solved_unknowns(sys)
     unk_idxs = Dict{UInt, Int}()
     for (i, sym) in enumerate(unks)
-        h = hash(unwrap(sym))
+        h = getsymbolhash(sym)
         unk_idxs[h] = i
-        setmetadata(sym, SymbolHash, h)
     end
 
     disc_buffers = Dict{DataType, Set{BasicSymbolic}}()
@@ -58,6 +57,13 @@ function IndexCache(sys::AbstractSystem)
             end
         end
     end
+    if has_discrete_subsystems(sys) && get_discrete_subsystems(sys) !== nothing
+        _, inputs, continuous_id, _ = get_discrete_subsystems(sys)
+        for par in inputs[continuous_id]
+            is_parameter(sys, par) || error("Discrete subsytem input is not a parameter")
+            insert_by_type!(disc_buffers, par)
+        end
+    end
 
     all_ps = Set(unwrap.(parameters(sys)))
     for (sym, value) in defaults(sys)
@@ -91,8 +97,9 @@ function IndexCache(sys::AbstractSystem)
         idx = 1
         for (T, buf) in buffers
             for p in buf
-                h = hash(p)
-                setmetadata(p, SymbolHash, h)
+                h = getsymbolhash(p)
+                idxs[h] = idx
+                h = getsymbolhash(default_toterm(p))
                 idxs[h] = idx
                 idx += 1
             end

src/systems/parameter_buffer.jl

@@ -6,23 +6,26 @@ struct MTKParameters{T, D, C, E, F}
     dependent_update::F
 end
 
-function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat = false, use_union = false)
+function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = false)
     ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
         get_index_cache(sys)
     else
         error("Cannot create MTKParameters if system does not have index_cache")
     end
     all_ps = Set(unwrap.(parameters(sys)))
+    union!(all_ps, default_toterm.(unwrap.(parameters(sys))))
     if p isa Vector && !(eltype(p) <: Pair) && !isempty(p)
         ps = parameters(sys)
         length(p) == length(ps) || error("Invalid parameters")
         p = ps .=> p
     end
-    defs = Dict(unwrap(k) => v for (k, v) in defaults(sys) if unwrap(k) in all_ps)
+    defs = Dict(default_toterm(unwrap(k)) => v for (k, v) in defaults(sys) if unwrap(k) in all_ps)
     if p isa SciMLBase.NullParameters
         p = defs
     else
-        p = merge(defs, Dict(unwrap(k) => v for (k, v) in p))
+        extra_params = Dict(unwrap(k) => v for (k, v) in p if !in(unwrap(k), all_ps))
+        p = merge(defs, Dict(default_toterm(unwrap(k)) => v for (k, v) in p))
+        p = Dict(k => fixpoint_sub(v, extra_params) for (k, v) in p if !haskey(extra_params, unwrap(k)))
     end
 
     tunable_buffer = ArrayPartition((Vector{temp.type}(undef, temp.length) for temp in ic.param_buffer_sizes)...)
@@ -41,6 +44,10 @@ function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat
         elseif haskey(ic.dependent_idx, h)
             dep_buffer[ic.dependent_idx[h]] = val
             dependencies[wrap(sym)] = wrap(p[sym])
+        elseif !isequal(default_toterm(sym), sym)
+            set_value(default_toterm(sym), val)
+        else
+            error("Symbol $sym does not have an index")
         end
     end
 

src/systems/systemstructure.jl

@@ -558,6 +558,7 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
         ModelingToolkit.infer_clocks!(ci)
+        time_domains = merge(Dict(state.fullvars .=> ci.var_domain), Dict(default_toterm.(state.fullvars) .=> ci.var_domain))
         tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
         cont_io = merge_io(io, inputs[continuous_id])
         sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify,
@@ -586,6 +587,8 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
             @set! sys.defaults = merge(ModelingToolkit.defaults(sys),
                 Dict(v => 0.0 for v in Iterators.flatten(inputs)))
         end
+        ps = [setmetadata(sym, TimeDomain, get(time_domains, sym, Continuous())) for sym in get_ps(sys)]
+        @set! sys.ps = ps
     else
         sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,
             fully_determined, kwargs...)

test/clock.jl

@@ -116,7 +116,7 @@ ss = structural_simplify(sys);
 Tf = 1.0
 prob = ODEProblem(ss, [x => 0.0, y => 0.0], (0.0, Tf),
     [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
-@test sort(prob.p) == [0, 1.0, 2.0, 3.0, 4.0] # yd, kp, z(k+1), z(k), ud
+@test sort(vcat(prob.p...)) == [0, 1.0, 2.0, 3.0, 4.0] # yd, kp, z(k+1), z(k), ud
 sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 # For all inputs in parameters, just initialize them to 0.0, and then set them
 # in the callback.