refactor: support clock systems without index caches

AayushSabharwal · ChrisRackauckas · commit f9f7ae192c60 · 2024-02-19T07:50:53.000-05:00
diff --git a/src/systems/clock_inference.jl b/src/systems/clock_inference.jl
@@ -146,12 +146,15 @@ function generate_discrete_affect(
     @static if VERSION < v"1.7"
         error("The `generate_discrete_affect` function requires at least Julia 1.7")
     end
-    has_index_cache(osys) && get_index_cache(osys) !== nothing || error("System must have index_cache for clock support")
+    use_index_cache = has_index_cache(osys) && get_index_cache(osys) !== nothing
     out = Sym{Any}(:out)
     appended_parameters = parameters(syss[continuous_id])
-    param_to_idx = Dict{Any, ParameterIndex}(p => parameter_index(osys, p)
-    for p in appended_parameters)
     offset = length(appended_parameters)
+    param_to_idx = if use_index_cache
+        Dict{Any, ParameterIndex}(p => parameter_index(osys, p) for p in appended_parameters)
+    else
+        Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
+    end
     affect_funs = []
     init_funs = []
     svs = []
@@ -170,18 +173,20 @@ function generate_discrete_affect(
             push!(fullvars, s)
         end
         needed_disc_to_cont_obs = []
-        disc_to_cont_idxs = ParameterIndex[]
+        if use_index_cache
+            disc_to_cont_idxs = ParameterIndex[]
+        else
+            disc_to_cont_idxs = Int[]
+        end
         for v in inputs[continuous_id]
             vv = arguments(v)[1]
             if vv in fullvars
                 push!(needed_disc_to_cont_obs, vv)
-                # @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])
             end
         end
         append!(appended_parameters, input, unknowns(sys))
-        cont_to_disc_obs = build_explicit_observed_function(osys,
+        cont_to_disc_obs = build_explicit_observed_function(use_index_cache ? osys : syss[continuous_id],
             needed_cont_to_disc_obs,
             throw = false,
             expression = true,
@@ -192,36 +197,62 @@ function generate_discrete_affect(
             expression = true,
             output_type = SVector,
             ps = reorder_parameters(osys, parameters(sys)))
+        ni = length(input)
+        ns = length(unknowns(sys))
         disc = Func(
             [
                 out,
                 DestructuredArgs(unknowns(osys)),
-                DestructuredArgs.(reorder_parameters(osys, parameters(osys)))...,
-                # DestructuredArgs(appended_parameters),
+                if use_index_cache
+                    DestructuredArgs.(reorder_parameters(osys, parameters(osys)))
+                else
+                    (DestructuredArgs(appended_parameters),)
+                end...,
                 get_iv(sys)
             ],
             [],
             let_block)
-        cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
-        disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
+        if use_index_cache
+            cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
+            disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
+        else
+            cont_to_disc_idxs = (offset + 1):(offset += ni)
+            input_offset = offset
+            disc_range = (offset + 1):(offset += ns)
+        end
         save_vec = Expr(:ref, :Float64)
-        for unk in unknowns(sys)
-            idx = parameter_index(osys, unk)
-            push!(save_vec.args, :($(parameter_values)(p, $idx)))
+        if use_index_cache
+            for unk in unknowns(sys)
+                idx = parameter_index(osys, unk)
+                push!(save_vec.args, :($(parameter_values)(p, $idx)))
+            end
+        else
+            for i in 1:ns
+                push!(save_vec.args, :(p[$(input_offset + i)]))
+            end
         end
         empty_disc = isempty(disc_range)
-        disc_init = :(function (p, t)
-            d2c_obs = $disc_to_cont_obs
-            disc_state = Tuple($(parameter_values)(p, i) for i in $disc_range)
-            result = d2c_obs(disc_state, p..., t)
-            for (val, i) in zip(result, $disc_to_cont_idxs)
-                # prevent multiple updates to dependents
-                _set_parameter_unchecked!(p, val, i; update_dependent = false)
-            end
-            discretes, repack, _ = $(SciMLStructures.canonicalize)(
-                $(SciMLStructures.Discrete()), p)
-            repack(discretes) # to force recalculation of dependents
-        end)
+        disc_init = if use_index_cache
+            :(function (p, t)
+                d2c_obs = $disc_to_cont_obs
+                disc_state = Tuple($(parameter_values)(p, i) for i in $disc_range)
+                result = d2c_obs(disc_state, p..., t)
+                for (val, i) in zip(result, $disc_to_cont_idxs)
+                    # prevent multiple updates to dependents
+                    _set_parameter_unchecked!(p, val, i; update_dependent = false)
+                end
+                discretes, repack, _ = $(SciMLStructures.canonicalize)(
+                    $(SciMLStructures.Discrete()), p)
+                repack(discretes) # to force recalculation of dependents
+            end)
+        else
+            :(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))
+            end)
+        end
 
         # @show disc_to_cont_idxs
         # @show cont_to_disc_idxs
@@ -230,8 +261,18 @@ function generate_discrete_affect(
             @unpack u, p, t = integrator
             c2d_obs = $cont_to_disc_obs
             d2c_obs = $disc_to_cont_obs
+            $(
+                if use_index_cache
+                    :(disc_unknowns = [$(parameter_values)(p, i) for i in $disc_range])
+                else
+                    quote
+                        c2d_view = view(p, $cont_to_disc_idxs)
+                        d2c_view = view(p, $disc_to_cont_idxs)
+                        disc_unknowns = view(p, $disc_range)
+                    end
+                end
+            )
             # TODO: find a way to do this without allocating
-            disc_unknowns = [$(parameter_values)(p, i) for i in $disc_range]
             disc = $disc
 
             push!(saved_values.t, t)
@@ -246,26 +287,56 @@ function generate_discrete_affect(
             # d2c comes last
             # @show t
             # @show "incoming", p
-            result = c2d_obs(integrator.u, p..., t)
-            for (val, i) in zip(result, $cont_to_disc_idxs)
-                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
-            end
+            $(
+                if use_index_cache
+                    quote
+                        result = c2d_obs(integrator.u, p..., t)
+                        for (val, i) in zip(result, $cont_to_disc_idxs)
+                            $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+                        end
+                    end
+                else
+                    :(copyto!(c2d_view, c2d_obs(integrator.u, p, t)))
+                end
+            )
             # @show "after c2d", p
-            if !$empty_disc
-                disc(disc_unknowns, integrator.u, p..., t)
-                for (val, i) in zip(disc_unknowns, $disc_range)
-                    $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+            $(
+                if use_index_cache
+                    quote
+                        if !$empty_disc
+                            disc(disc_unknowns, integrator.u, p..., t)
+                            for (val, i) in zip(disc_unknowns, $disc_range)
+                                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+                            end
+                        end
+                    end
+                else
+                    :($empty_disc || disc(disc_unknowns, disc_unknowns, p, t))
                 end
-            end
+            )
             # @show "after state update", p
-            result = d2c_obs(disc_unknowns, p..., t)
-            for (val, i) in zip(result, $disc_to_cont_idxs)
-                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
-            end
+            $(
+                if use_index_cache
+                    quote
+                        result = d2c_obs(disc_unknowns, p..., t)
+                        for (val, i) in zip(result, $disc_to_cont_idxs)
+                            $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+                        end
+                    end
+                else
+                    :(copyto!(d2c_view, d2c_obs(disc_unknowns, p, t)))
+                end
+            )
             # @show "after d2c", p
-            discretes, repack, _ = $(SciMLStructures.canonicalize)(
-                $(SciMLStructures.Discrete()), p)
-            repack(discretes)
+            $(
+                if use_index_cache
+                    quote
+                        discretes, repack, _ = $(SciMLStructures.canonicalize)(
+                            $(SciMLStructures.Discrete()), p)
+                        repack(discretes)
+                    end
+                end
+            )
         end)
         sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)
diff --git a/test/clock.jl b/test/clock.jl
@@ -118,6 +118,12 @@ 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(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)
+
+ss_nosplit = structural_simplify(sys; split = false)
+prob_nosplit = ODEProblem(ss_nosplit, [x => 0.0, y => 0.0], (0.0, Tf),
+    [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
+@test sort(prob_nosplit.p) == [0, 1.0, 2.0, 3.0, 4.0] # yd, kp, z(k+1), z(k), ud
+sol_nosplit = solve(prob_nosplit, Tsit5(), kwargshandle = KeywordArgSilent)
 # For all inputs in parameters, just initialize them to 0.0, and then set them
 # in the callback.
 
@@ -154,8 +160,11 @@ prob = ODEProblem(foo!, [0.0], (0.0, Tf), [1.0, 4.0, 2.0, 3.0], callback = cb)
 #                               ud initializes to kp * (r - yd) + z = 1 * (1 - 0) + 3 = 4
 sol2 = solve(prob, Tsit5())
 @test sol.u == sol2.u
+@test sol_nosplit.u == sol2.u
 @test saved_values.t == sol.prob.kwargs[:disc_saved_values][1].t
+@test saved_values.t == sol_nosplit.prob.kwargs[:disc_saved_values][1].t
 @test saved_values.saveval == sol.prob.kwargs[:disc_saved_values][1].saveval
+@test saved_values.saveval == sol_nosplit.prob.kwargs[:disc_saved_values][1].saveval
 
 @info "Testing multi-rate hybrid system"
 dt = 0.1
@@ -280,10 +289,13 @@ ci, varmap = infer_clocks(cl)
 @test varmap[u] == Continuous()
 
 ss = structural_simplify(cl)
+ss_nosplit = structural_simplify(cl; split = false)
 
 if VERSION >= v"1.7"
     prob = ODEProblem(ss, [x => 0.0], (0.0, 1.0), [kp => 1.0])
+    prob_nosplit = ODEProblem(ss_nosplit, [x => 0.0], (0.0, 1.0), [kp => 1.0])
     sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
+    sol_nosplit = solve(prob_nosplit, Tsit5(), kwargshandle = KeywordArgSilent)
 
     function foo!(dx, x, p, t)
         kp, ud1, ud2 = p
@@ -314,6 +326,7 @@ if VERSION >= v"1.7"
     sol2 = solve(prob, Tsit5())
 
     @test sol.u≈sol2.u atol=1e-6
+    @test sol_nosplit.u≈sol2.u atol=1e-6
 end
 
 ##
