refactor: store parameters from different clock partitions separately

Author: AayushSabharwal

src/systems/index_cache.jl

@@ -24,12 +24,13 @@ const UnknownIndexMap = Dict{
 
 struct IndexCache
     unknown_idx::UnknownIndexMap
+    discrete_clocks::Dict{Union{Symbol, BasicSymbolic}, Int}
     discrete_idx::ParamIndexMap
     tunable_idx::ParamIndexMap
     constant_idx::ParamIndexMap
     dependent_idx::ParamIndexMap
     nonnumeric_idx::ParamIndexMap
-    discrete_buffer_sizes::Vector{BufferTemplate}
+    discrete_buffer_sizes::Vector{Vector{BufferTemplate}}
     tunable_buffer_sizes::Vector{BufferTemplate}
     constant_buffer_sizes::Vector{BufferTemplate}
     dependent_buffer_sizes::Vector{BufferTemplate}
@@ -71,7 +72,8 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
-    disc_buffers = Dict{Any, Set{BasicSymbolic}}()
+    disc_buffers = Dict{Int, Dict{Any, Set{BasicSymbolic}}}()
+    disc_clocks = Dict{Union{Symbol, BasicSymbolic}, Int}()
     tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()
     dependent_buffers = Dict{Any, Set{BasicSymbolic}}()
@@ -84,27 +86,107 @@ function IndexCache(sys::AbstractSystem)
         push!(buf, sym)
     end
 
+    if has_discrete_subsystems(sys) && get_discrete_subsystems(sys) !== nothing
+        syss, inputs, continuous_id, _ = get_discrete_subsystems(sys)
+
+        for (i, (inps, disc_sys)) in enumerate(zip(inputs, syss))
+            i == continuous_id && continue
+            disc_buffers[i - 1] = Dict{Any, Set{BasicSymbolic}}()
+
+            for inp in inps
+                inp = unwrap(inp)
+                is_parameter(sys, inp) ||
+                    error("Discrete subsystem $i input $inp is not a parameter")
+                disc_clocks[inp] = i - 1
+                disc_clocks[default_toterm(inp)] = i - 1
+                if hasname(inp) && (!istree(inp) || operation(inp) !== getindex)
+                    disc_clocks[getname(inp)] = i - 1
+                    disc_clocks[default_toterm(inp)] = i - 1
+                end
+                insert_by_type!(disc_buffers[i - 1], inp)
+            end
+
+            for sym in unknowns(disc_sys)
+                sym = unwrap(sym)
+                is_parameter(sys, sym) ||
+                    error("Discrete subsystem $i unknown $sym is not a parameter")
+                disc_clocks[sym] = i - 1
+                disc_clocks[default_toterm(sym)] = i - 1
+                if hasname(sym) && (!istree(sym) || operation(sym) !== getindex)
+                    disc_clocks[getname(sym)] = i - 1
+                    disc_clocks[getname(default_toterm(sym))] = i - 1
+                end
+                insert_by_type!(disc_buffers[i - 1], sym)
+            end
+            t = get_iv(sys)
+            for eq in observed(disc_sys)
+                # TODO: Is this a valid check
+                # FIXME: This shouldn't be necessary
+                eq.rhs === -0.0 && continue
+                sym = eq.lhs
+                if istree(sym) && operation(sym) == Shift(t, 1)
+                    sym = only(arguments(sym))
+                end
+                # is_parameter(sys, sym) || is_parameter(sys, Hold(sym)) || continue
+                disc_clocks[sym] = i - 1
+                disc_clocks[sym] = i - 1
+                disc_clocks[default_toterm(sym)] = i - 1
+                if hasname(sym) && (!istree(sym) || operation(sym) !== getindex)
+                    disc_clocks[getname(sym)] = i - 1
+                    disc_clocks[getname(default_toterm(sym))] = i - 1
+                end
+            end
+        end
+
+        for par in inputs[continuous_id]
+            is_parameter(sys, par) || error("Discrete subsystem input is not a parameter")
+            istree(par) && operation(par) isa Hold ||
+                error("Continuous subsystem input is not a Hold")
+            if haskey(disc_clocks, par)
+                sym = par
+            else
+                sym = first(arguments(par))
+            end
+            haskey(disc_clocks, sym) ||
+                error("Variable $par not part of a discrete subsystem")
+            disc_clocks[par] = disc_clocks[sym]
+            insert_by_type!(disc_buffers[disc_clocks[sym]], par)
+        end
+    end
+
     affs = vcat(affects(continuous_events(sys)), affects(discrete_events(sys)))
+    user_affect_clock = maximum(values(disc_clocks); init = 1)
     for affect in affs
         if affect isa Equation
             is_parameter(sys, affect.lhs) || continue
-            insert_by_type!(disc_buffers, affect.lhs)
+
+            disc_clocks[affect.lhs] = user_affect_clock
+            disc_clocks[default_toterm(affect.lhs)] = user_affect_clock
+            if hasname(affect.lhs) &&
+               (!istree(affect.lhs) || operation(affect.lhs) !== getindex)
+                disc_clocks[getname(affect.lhs)] = user_affect_clock
+                disc_clocks[getname(default_toterm(affect.lhs))] = user_affect_clock
+            end
+            buffer = get!(disc_buffers, user_affect_clock, Dict{Any, Set{BasicSymbolic}}())
+            insert_by_type!(buffer, affect.lhs)
         else
             discs = discretes(affect)
             for disc in discs
                 is_parameter(sys, disc) ||
                     error("Expected discrete variable $disc in callback to be a parameter")
-                insert_by_type!(disc_buffers, disc)
+                disc = unwrap(disc)
+                disc_clocks[disc] = user_affect_clock
+                disc_clocks[default_toterm(disc)] = user_affect_clock
+                if hasname(disc) && (!istree(disc) || operation(disc) !== getindex)
+                    disc_clocks[getname(disc)] = user_affect_clock
+                    disc_clocks[getname(default_toterm(disc))] = user_affect_clock
+                end
+                buffer = get!(
+                    disc_buffers, user_affect_clock, Dict{Any, Set{BasicSymbolic}}())
+                insert_by_type!(buffer, disc)
             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 subsystem input is not a parameter")
-            insert_by_type!(disc_buffers, par)
-        end
-    end
 
     if has_parameter_dependencies(sys) &&
        (pdeps = get_parameter_dependencies(sys)) !== nothing
@@ -117,13 +199,11 @@ function IndexCache(sys::AbstractSystem)
     for p in parameters(sys)
         p = unwrap(p)
         ctype = concrete_symtype(p)
-        haskey(disc_buffers, ctype) && p in disc_buffers[ctype] && continue
+        haskey(disc_clocks, p) && continue
         haskey(dependent_buffers, ctype) && p in dependent_buffers[ctype] && continue
         insert_by_type!(
             if ctype <: Real || ctype <: AbstractArray{<:Real}
-                if is_discrete_domain(p)
-                    disc_buffers
-                elseif istunable(p, true) && Symbolics.shape(p) !== Symbolics.Unknown()
+                if istunable(p, true) && Symbolics.shape(p) !== Symbolics.Unknown()
                     tunable_buffers
                 else
                     constant_buffers
@@ -135,6 +215,31 @@ function IndexCache(sys::AbstractSystem)
         )
     end
 
+    disc_idxs = ParamIndexMap()
+    disc_buffer_sizes = [BufferTemplate[] for _ in 1:length(disc_buffers)]
+    disc_buffer_types = Set()
+    for buffer in values(disc_buffers)
+        union!(disc_buffer_types, keys(buffer))
+    end
+
+    for (clockidx, buffer) in disc_buffers
+        for (i, btype) in enumerate(disc_buffer_types)
+            if !haskey(buffer, btype)
+                push!(disc_buffer_sizes[clockidx], BufferTemplate(btype, 0))
+                continue
+            end
+            push!(disc_buffer_sizes[clockidx], BufferTemplate(btype, length(buffer[btype])))
+            for (j, sym) in enumerate(buffer[btype])
+                disc_idxs[sym] = (i, j)
+                disc_idxs[default_toterm(sym)] = (i, j)
+                if hasname(sym) && (!istree(sym) || operation(sym) !== getindex)
+                    disc_idxs[getname(sym)] = (i, j)
+                    disc_idxs[getname(default_toterm(sym))] = (i, j)
+                end
+            end
+        end
+    end
+
     function get_buffer_sizes_and_idxs(buffers::Dict{Any, Set{BasicSymbolic}})
         idxs = ParamIndexMap()
         buffer_sizes = BufferTemplate[]
@@ -152,20 +257,20 @@ function IndexCache(sys::AbstractSystem)
         return idxs, buffer_sizes
     end
 
-    disc_idxs, discrete_buffer_sizes = get_buffer_sizes_and_idxs(disc_buffers)
     tunable_idxs, tunable_buffer_sizes = get_buffer_sizes_and_idxs(tunable_buffers)
     const_idxs, const_buffer_sizes = get_buffer_sizes_and_idxs(constant_buffers)
     dependent_idxs, dependent_buffer_sizes = get_buffer_sizes_and_idxs(dependent_buffers)
     nonnumeric_idxs, nonnumeric_buffer_sizes = get_buffer_sizes_and_idxs(nonnumeric_buffers)
 
     return IndexCache(
         unk_idxs,
+        disc_clocks,
         disc_idxs,
         tunable_idxs,
         const_idxs,
         dependent_idxs,
         nonnumeric_idxs,
-        discrete_buffer_sizes,
+        disc_buffer_sizes,
         tunable_buffer_sizes,
         const_buffer_sizes,
         dependent_buffer_sizes,
@@ -193,7 +298,8 @@ function SymbolicIndexingInterface.parameter_index(ic::IndexCache, sym)
     return if (idx = check_index_map(ic.tunable_idx, sym)) !== nothing
         ParameterIndex(SciMLStructures.Tunable(), idx)
     elseif (idx = check_index_map(ic.discrete_idx, sym)) !== nothing
-        ParameterIndex(SciMLStructures.Discrete(), idx)
+        ParameterIndex(
+            SciMLStructures.Discrete(), (check_index_map(ic.discrete_clocks, sym), idx...))
     elseif (idx = check_index_map(ic.constant_idx, sym)) !== nothing
         ParameterIndex(SciMLStructures.Constants(), idx)
     elseif (idx = check_index_map(ic.nonnumeric_idx, sym)) !== nothing
@@ -205,6 +311,18 @@ function SymbolicIndexingInterface.parameter_index(ic::IndexCache, sym)
     end
 end
 
+function SymbolicIndexingInterface.is_timeseries_parameter(ic::IndexCache, sym)
+    return check_index_map(ic.discrete_clocks, sym) !== nothing
+end
+
+function SymbolicIndexingInterface.timeseries_parameter_index(ic::IndexCache, sym)
+    clockid = check_index_map(ic.discrete_clocks, sym)
+    clockid === nothing && return nothing
+    partitionid = check_index_map(ic.discrete_idx, sym)
+    partitionid === nothing && return nothing
+    return ParameterTimeseriesIndex(clockid, partitionid)
+end
+
 function check_index_map(idxmap, sym)
     if (idx = get(idxmap, sym, nothing)) !== nothing
         return idx
@@ -229,7 +347,8 @@ function ParameterIndex(ic::IndexCache, p, sub_idx = ())
     return if haskey(ic.tunable_idx, p)
         ParameterIndex(SciMLStructures.Tunable(), (ic.tunable_idx[p]..., sub_idx...))
     elseif haskey(ic.discrete_idx, p)
-        ParameterIndex(SciMLStructures.Discrete(), (ic.discrete_idx[p]..., sub_idx...))
+        ParameterIndex(SciMLStructures.Discrete(),
+            (ic.discrete_clocks[p], ic.discrete_idx[p]..., sub_idx...))
     elseif haskey(ic.constant_idx, p)
         ParameterIndex(SciMLStructures.Constants(), (ic.constant_idx[p]..., sub_idx...))
     elseif haskey(ic.dependent_idx, p)
@@ -247,10 +366,14 @@ end
 function discrete_linear_index(ic::IndexCache, idx::ParameterIndex)
     idx.portion isa SciMLStructures.Discrete || error("Discrete variable index expected")
     ind = sum(temp.length for temp in ic.tunable_buffer_sizes; init = 0)
+    for clockbuftemps in Iterators.take(ic.discrete_buffer_sizes, idx.idx[1] - 1)
+        ind += sum(temp.length for temp in clockbuftemps; init = 0)
+    end
     ind += sum(
-        temp.length for temp in Iterators.take(ic.discrete_buffer_sizes, idx.idx[1] - 1);
+        temp.length
+        for temp in Iterators.take(ic.discrete_buffer_sizes[idx.idx[1]], idx.idx[2] - 1);
         init = 0)
-    ind += idx.idx[2]
+    ind += idx.idx[3]
     return ind
 end
 
@@ -269,30 +392,32 @@ function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
     param_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.tunable_buffer_sizes)
     disc_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
-    for temp in ic.discrete_buffer_sizes)
+    for temp in Iterators.flatten(ic.discrete_buffer_sizes))
     const_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.constant_buffer_sizes)
     dep_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.dependent_buffer_sizes)
     nonnumeric_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.nonnumeric_buffer_sizes)
-
     for p in ps
+        p = unwrap(p)
         if haskey(ic.discrete_idx, p)
-            i, j = ic.discrete_idx[p]
-            disc_buf[i][j] = unwrap(p)
+            disc_offset = length(first(ic.discrete_buffer_sizes))
+            i = ic.discrete_clocks[p]
+            j, k = ic.discrete_idx[p]
+            disc_buf[(i - 1) * disc_offset + j][k] = p
         elseif haskey(ic.tunable_idx, p)
             i, j = ic.tunable_idx[p]
-            param_buf[i][j] = unwrap(p)
+            param_buf[i][j] = p
         elseif haskey(ic.constant_idx, p)
             i, j = ic.constant_idx[p]
-            const_buf[i][j] = unwrap(p)
+            const_buf[i][j] = p
         elseif haskey(ic.dependent_idx, p)
             i, j = ic.dependent_idx[p]
-            dep_buf[i][j] = unwrap(p)
+            dep_buf[i][j] = p
         elseif haskey(ic.nonnumeric_idx, p)
             i, j = ic.nonnumeric_idx[p]
-            nonnumeric_buf[i][j] = unwrap(p)
+            nonnumeric_buf[i][j] = p
         else
             error("Invalid parameter $p")
         end