feat: add IndexCache to ODESystem

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -313,6 +313,9 @@ Mark a system as completed. If a system is complete, the system will no longer
 namespace its subsystems or variables, i.e. `isequal(complete(sys).v.i, v.i)`.
 """
 function complete(sys::AbstractSystem)
+    if has_index_cache(sys)
+        @set! sys.index_cache = IndexCache(sys)
+    end
     isdefined(sys, :complete) ? (@set! sys.complete = true) : sys
 end
 
@@ -354,7 +357,8 @@ for prop in [:eqs
     :discrete_subsystems
     :solved_unknowns
     :split_idxs
-    :parent]
+    :parent
+    :index_cache]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
     @eval begin

src/systems/diffeqs/odesystem.jl

@@ -131,6 +131,10 @@ struct ODESystem <: AbstractODESystem
     """
     complete::Bool
     """
+    Cached data for fast symbolic indexing.
+    """
+    index_cache::Union{Nothing, IndexCache}
+    """
     A list of discrete subsystems.
     """
     discrete_subsystems::Any
@@ -152,7 +156,7 @@ struct ODESystem <: AbstractODESystem
             torn_matching, connector_type, preface, cevents,
             devents, metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing,
-            substitutions = nothing, complete = false,
+            substitutions = nothing, complete = false, index_cache = nothing,
             discrete_subsystems = nothing, solved_unknowns = nothing,
             split_idxs = nothing, parent = nothing; checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
@@ -168,7 +172,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, gui_metadata,
-            tearing_state, substitutions, complete, discrete_subsystems,
+            tearing_state, substitutions, complete, index_cache, discrete_subsystems,
             solved_unknowns, split_idxs, parent)
     end
 end

src/systems/index_cache.jl

@@ -0,0 +1,77 @@
+abstract type SymbolHash end
+
+struct IndexCache
+    unknown_idx::Dict{UInt, Int}
+    discrete_idx::Dict{UInt, Int}
+    param_idx::Dict{UInt, Int}
+    discrete_buffer_type_and_size::Vector{Tuple{DataType, Int}}
+    param_buffer_type_and_size::Vector{Tuple{DataType, Int}}
+end
+
+function IndexCache(sys::AbstractSystem)
+    unks = solved_unknowns(sys)
+    unk_idxs = Dict{UInt, Int}()
+    for (i, sym) in enumerate(unks)
+        h = hash(unwrap(sym))
+        unk_idxs[h] = i
+        setmetadata(sym, SymbolHash, h)
+    end
+
+    # split parameters, also by type
+    discrete_params = Dict{DataType, Any}()
+    tunable_params = Dict{DataType, Any}()
+
+    for p in parameters(sys)
+        T = symtype(p)
+        buf = get!(is_discrete_domain(p) ? discrete_params : tunable_params, T, [])
+        push!(buf, unwrap(p))
+    end
+    
+    disc_idxs = Dict{UInt, Int}()
+    discrete_buffer_type_and_size = Tuple{DataType, Int}[]
+    didx = 1
+
+    for (T, ps) in discrete_params
+        push!(discrete_buffer_type_and_size, (T, length(ps)))
+        for p in ps
+            h = hash(p)
+            disc_idxs[h] = didx
+            didx += 1
+            setmetadata(p, SymbolHash, h)
+        end
+    end
+
+    param_idxs = Dict{UInt, Int}()
+    param_buffer_type_and_size = Tuple{DataType, Int}[]
+    pidx = 1
+
+    for (T, ps) in tunable_params
+        push!(param_buffer_type_and_size, (T, length(ps)))
+        for p in ps
+            h = hash(p)
+            param_idxs[h] = pidx
+            pidx += 1
+            setmetadata(p, SymbolHash, h)
+        end
+    end
+
+    return IndexCache(unk_idxs, disc_idxs, param_idxs, discrete_buffer_type_and_size, param_buffer_type_and_size)
+end
+
+function reorder_parameters(ic::IndexCache, ps)
+    param_bufs = ArrayPartition((Vector{BasicSymbolic{T}}(undef, sz) for (T, sz) in ic.param_buffer_type_and_size)...)
+    disc_bufs = ArrayPartition((Vector{BasicSymbolic{T}}(undef, sz) for (T, sz) in ic.discrete_buffer_type_and_size)...)
+
+    for p in ps
+        h = hasmetadata(p, SymbolHash) ? getmetadata(p, SymbolHash) : hash(unwrap(p))
+        if haskey(ic.discrete_idx, h)
+            disc_bufs[ic.discrete_idx[h]] = unwrap(p)
+        elseif haskey(ic.param_idx, h)
+            param_bufs[ic.param_idx[h]] = unwrap(p)
+        else
+            error("Invalid parameter $p")
+        end
+    end
+
+    return (param_bufs.x..., disc_bufs.x...)
+end