feat: add support for SciMLStructures.Constants portion, and dependent parameters

AayushSabharwal · AayushSabharwal · commit dff1b5f00f28 · 2024-02-12T11:38:46.000+05:30
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
@@ -255,11 +255,14 @@ function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         ic = get_index_cache(sys)
         h = getsymbolhash(sym)
-        return if haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h)
+        return if haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) ||
+                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h)
             true
         else
             h = getsymbolhash(default_toterm(sym))
-            haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) || hasname(sym) && is_parameter(sys, getname(sym))
+            haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) ||
+                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h) ||
+                hasname(sym) && is_parameter(sys, getname(sym))
         end
     end
     return any(isequal(sym), parameter_symbols(sys)) ||
@@ -284,12 +287,20 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym)
             ParameterIndex(SciMLStructures.Tunable(), ic.param_idx[h])
         elseif haskey(ic.discrete_idx, h)
             ParameterIndex(SciMLStructures.Discrete(), ic.discrete_idx[h])
+        elseif haskey(ic.constant_idx, h)
+            ParameterIndex(SciMLStructures.Constants(), ic.constant_idx[h])
+        elseif haskey(ic.dependent_idx, h)
+            ParameterIndex(nothing, ic.dependent_idx[h])
         else
             h = getsymbolhash(default_toterm(sym))
             if haskey(ic.param_idx, h)
                 ParameterIndex(SciMLStructures.Tunable(), ic.param_idx[h])
             elseif haskey(ic.discrete_idx, h)
                 ParameterIndex(SciMLStructures.Discrete(), ic.discrete_idx[h])
+            elseif haskey(ic.constant_idx, h)
+                ParameterIndex(SciMLStructures.Constants(), ic.constant_idx[h])
+            elseif haskey(ic.dependent_idx, h)
+                ParameterIndex(nothing, ic.dependent_idx[h])
             else
                 nothing
             end
diff --git a/src/systems/index_cache.jl b/src/systems/index_cache.jl
@@ -16,8 +16,12 @@ struct IndexCache
     unknown_idx::Dict{UInt, Int}
     discrete_idx::Dict{UInt, Int}
     param_idx::Dict{UInt, Int}
+    constant_idx::Dict{UInt, Int}
+    dependent_idx::Dict{UInt, Int}
     discrete_buffer_sizes::Vector{BufferTemplate}
     param_buffer_sizes::Vector{BufferTemplate}
+    constant_buffer_sizes::Vector{BufferTemplate}
+    dependent_buffer_sizes::Vector{BufferTemplate}
 end
 
 function IndexCache(sys::AbstractSystem)
@@ -31,6 +35,8 @@ function IndexCache(sys::AbstractSystem)
 
     disc_buffers = Dict{DataType, Set{BasicSymbolic}}()
     tunable_buffers = Dict{DataType, Set{BasicSymbolic}}()
+    constant_buffers = Dict{DataType, Set{BasicSymbolic}}()
+    dependent_buffers = Dict{DataType, Set{BasicSymbolic}}()
 
     function insert_by_type!(buffers::Dict{DataType, Set{BasicSymbolic}}, sym)
         sym = unwrap(sym)
@@ -53,40 +59,64 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
+    all_ps = Set(unwrap.(parameters(sys)))
+    for (sym, value) in defaults(sys)
+        sym = unwrap(sym)
+        if sym in all_ps && symbolic_type(unwrap(value)) !== NotSymbolic()
+            insert_by_type!(dependent_buffers, sym)
+        end
+    end
+
     for p in parameters(sys)
         p = unwrap(p)
         ctype = concrete_symtype(p)
         haskey(disc_buffers, ctype) && p in disc_buffers[ctype] && continue
-        
-        insert_by_type!(is_discrete_domain(p) ? disc_buffers : tunable_buffers, p)
+        haskey(dependent_buffers, ctype) && p in dependent_buffers[ctype] && continue
+
+        insert_by_type!(
+            if is_discrete_domain(p)
+                disc_buffers
+            elseif istunable(p, true)
+                tunable_buffers
+            else
+                constant_buffers
+            end,
+            p
+        )
     end
 
-    disc_idxs = Dict{UInt, Int}()
-    discrete_buffer_sizes = BufferTemplate[]
-    didx = 1
-    for (T, buf) in disc_buffers
-        for p in buf
-            h = hash(p)
-            setmetadata(p, SymbolHash, h)
-            disc_idxs[h] = didx
-            didx += 1
-        end
-        push!(discrete_buffer_sizes, BufferTemplate(T, length(buf)))
-    end
-    param_idxs = Dict{UInt, Int}()
-    param_buffer_sizes = BufferTemplate[]
-    pidx = 1
-    for (T, buf) in tunable_buffers
-        for p in buf
-            h = hash(p)
-            setmetadata(p, SymbolHash, h)
-            param_idxs[h] = pidx
-            pidx += 1
+    function get_buffer_sizes_and_idxs(buffers::Dict{DataType, Set{BasicSymbolic}})
+        idxs = Dict{UInt, Int}()
+        buffer_sizes = BufferTemplate[]
+        idx = 1
+        for (T, buf) in buffers
+            for p in buf
+                h = hash(p)
+                setmetadata(p, SymbolHash, h)
+                idxs[h] = idx
+                idx += 1
+            end
+            push!(buffer_sizes, BufferTemplate(T, length(buf)))
         end
-        push!(param_buffer_sizes, BufferTemplate(T, length(buf)))
+        return idxs, buffer_sizes
     end
 
-    return IndexCache(unk_idxs, disc_idxs, param_idxs, discrete_buffer_sizes, param_buffer_sizes)
+    disc_idxs, discrete_buffer_sizes = get_buffer_sizes_and_idxs(disc_buffers)
+    param_idxs, param_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)
+
+    return IndexCache(
+        unk_idxs,
+        disc_idxs,
+        param_idxs,
+        const_idxs,
+        dependent_idxs,
+        discrete_buffer_sizes,
+        param_buffer_sizes,
+        const_buffer_sizes,
+        dependent_buffer_sizes,
+    )
 end
 
 function reorder_parameters(sys::AbstractSystem, ps; kwargs...)
@@ -102,19 +132,25 @@ end
 function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
     param_buf = ArrayPartition((fill(variable(:DEF), temp.length) for temp in ic.param_buffer_sizes)...)
     disc_buf = ArrayPartition((fill(variable(:DEF), temp.length) for temp in ic.discrete_buffer_sizes)...)
+    const_buf = ArrayPartition((fill(variable(:DEF), temp.length) for temp in ic.constant_buffer_sizes)...)
+    dep_buf = ArrayPartition((fill(variable(:DEF), temp.length) for temp in ic.dependent_buffer_sizes)...)
 
     for p in ps
         h = getsymbolhash(p)
         if haskey(ic.discrete_idx, h)
             disc_buf[ic.discrete_idx[h]] = unwrap(p)
         elseif haskey(ic.param_idx, h)
             param_buf[ic.param_idx[h]] = unwrap(p)
+        elseif haskey(ic.constant_idx, h)
+            const_buf[ic.constant_idx[h]] = unwrap(p)
+        elseif haskey(ic.dependent_idx, h)
+            dep_buf[ic.dependent_idx[h]] = unwrap(p)
         else
             error("Invalid parameter $p")
         end
     end
 
-    result = broadcast.(unwrap, (param_buf.x..., disc_buf.x...))
+    result = broadcast.(unwrap, (param_buf.x..., disc_buf.x..., const_buf.x..., dep_buf.x...))
     if drop_missing
         result = map(result) do buf
             filter(buf) do sym
diff --git a/src/systems/parameter_buffer.jl b/src/systems/parameter_buffer.jl
@@ -1,6 +1,9 @@
-struct MTKParameters{T, D}
+struct MTKParameters{T, D, C, E, F}
     tunable::T
     discrete::D
+    constant::C
+    dependent::E
+    dependent_update::F
 end
 
 function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat = false, use_union = false)
@@ -10,7 +13,7 @@ function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat
         error("Cannot create MTKParameters if system does not have index_cache")
     end
     all_ps = Set(unwrap.(parameters(sys)))
-    if p isa Vector && !(eltype(p) <: Pair)
+    if p isa Vector && !(eltype(p) <: Pair) && !isempty(p)
         ps = parameters(sys)
         length(p) == length(ps) || error("Invalid parameters")
         p = ps .=> p
@@ -24,12 +27,20 @@ function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat
 
     tunable_buffer = ArrayPartition((Vector{temp.type}(undef, temp.length) for temp in ic.param_buffer_sizes)...)
     disc_buffer = ArrayPartition((Vector{temp.type}(undef, temp.length) for temp in ic.discrete_buffer_sizes)...)
+    const_buffer = ArrayPartition((Vector{temp.type}(undef, temp.length) for temp in ic.constant_buffer_sizes)...)
+    dep_buffer = ArrayPartition((Vector{temp.type}(undef, temp.length) for temp in ic.dependent_buffer_sizes)...)
+    dependencies = Dict{Num, Num}()
     function set_value(sym, val)
         h = getsymbolhash(sym)
         if haskey(ic.param_idx, h)
             tunable_buffer[ic.param_idx[h]] = val
         elseif haskey(ic.discrete_idx, h)
             disc_buffer[ic.discrete_idx[h]] = val
+        elseif haskey(ic.constant_idx, h)
+            const_buffer[ic.constant_idx[h]] = val
+        elseif haskey(ic.dependent_idx, h)
+            dep_buffer[ic.dependent_idx[h]] = val
+            dependencies[wrap(sym)] = wrap(p[sym])
         end
     end
 
@@ -49,22 +60,46 @@ function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm, tofloat
         end
     end
 
+    dep_exprs = ArrayPartition((wrap.(v) for v in dep_buffer.x)...)
+    for (sym, val) in dependencies
+        h = getsymbolhash(sym)
+        idx = ic.dependent_idx[h]
+        dep_exprs[idx] = wrap(fixpoint_sub(val, dependencies))
+    end
+    p = reorder_parameters(ic, parameters(sys))[begin:end-length(dep_buffer.x)]
+    update_function = if isempty(dep_exprs.x)
+        (_...) -> ()
+    else
+        RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(build_function(dep_exprs, p...)[2])
+    end
     # everything is an ArrayPartition so it's easy to figure out how many
     # distinct vectors we have for each portion as `ArrayPartition.x`
-    if tunable_buffer isa ArrayPartition && isempty(tunable_buffer.x) || isempty(tunable_buffer)
+    if isempty(tunable_buffer.x)
         tunable_buffer = ArrayPartition(Float64[])
     end
-    if disc_buffer isa ArrayPartition && isempty(disc_buffer.x) || isempty(disc_buffer)
+    if isempty(disc_buffer.x)
         disc_buffer = ArrayPartition(Float64[])
     end
+    if isempty(const_buffer.x)
+        const_buffer = ArrayPartition(Float64[])
+    end
+    if isempty(dep_buffer.x)
+        dep_buffer = ArrayPartition(Float64[])
+    end
     if use_union
         tunable_buffer = ArrayPartition(restrict_array_to_union(tunable_buffer))
         disc_buffer = ArrayPartition(restrict_array_to_union(disc_buffer))
+        const_buffer = ArrayPartition(restrict_array_to_union(const_buffer))
+        dep_buffer = ArrayPartition(restrict_array_to_union(dep_buffer))
     elseif tofloat
         tunable_buffer = ArrayPartition(Float64.(tunable_buffer))
         disc_buffer = ArrayPartition(Float64.(disc_buffer))
+        const_buffer = ArrayPartition(Float64.(const_buffer))
+        dep_buffer = ArrayPartition(Float64.(dep_buffer))
     end
-    return MTKParameters{typeof(tunable_buffer), typeof(disc_buffer)}(tunable_buffer, disc_buffer)
+    return MTKParameters{typeof(tunable_buffer), typeof(disc_buffer), typeof(const_buffer),
+                         typeof(dep_buffer), typeof(update_function)}(tunable_buffer,
+                         disc_buffer, const_buffer, dep_buffer, update_function)
 end
 
 SciMLStructures.isscimlstructure(::MTKParameters) = true
@@ -74,20 +109,24 @@ SciMLStructures.ismutablescimlstructure(::MTKParameters) = true
 for (Portion, field) in [
     (SciMLStructures.Tunable, :tunable)
     (SciMLStructures.Discrete, :discrete)
+    (SciMLStructures.Constants, :constant)
 ]
     @eval function SciMLStructures.canonicalize(::$Portion, p::MTKParameters)
         function repack(values)
             p.$field .= values
+            p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
         end
         return p.$field, repack, true
     end
 
     @eval function SciMLStructures.replace(::$Portion, p::MTKParameters, newvals)
-        @set p.$field = newvals
+        @set! p.$field = newvals
+        p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
     end
 
     @eval function SciMLStructures.replace!(::$Portion, p::MTKParameters, newvals)
         p.$field .= newvals
+        p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
         nothing
     end
 end
@@ -98,6 +137,10 @@ function SymbolicIndexingInterface.parameter_values(p::MTKParameters, i::Paramet
         return p.tunable[idx]
     elseif portion isa SciMLStructures.Discrete
         return p.discrete[idx]
+    elseif portion isa SciMLStructures.Constants
+        return p.constant[idx]
+    elseif portion === nothing
+        return p.dependent[idx]
     else
         error("Unhandled portion $portion")
     end
@@ -109,51 +152,62 @@ function SymbolicIndexingInterface.set_parameter!(p::MTKParameters, val, idx::Pa
         p.tunable[idx] = val
     elseif portion isa SciMLStructures.Discrete
         p.discrete[idx] = val
+    elseif portion isa SciMLStructures.Constants
+        p.constant[idx] = val
+    elseif portion === nothing
+        error("Cannot set value of parameter: ")
     else
         error("Unhandled portion $portion")
     end
+    p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
 end
 
 # for compiling callbacks
 # getindex indexes the vectors, setindex! linearly indexes values
 # it's inconsistent, but we need it to be this way
 function Base.getindex(buf::MTKParameters, i)
-    if i <= length(buf.tunable.x)
-        buf.tunable.x[i]
-    else
-        buf.discrete.x[i - length(buf.tunable.x)]
+    if !isempty(buf.tunable)
+        i <= length(buf.tunable.x) && return buf.tunable.x[i]
+        i -= length(buf.tunable.x)
+    end
+    if !isempty(buf.discrete)
+        i <= length(buf.discrete.x) && return buf.discrete.x[i]
+        i -= length(buf.discrete.x)
+    end
+    if !isempty(buf.constant)
+        i <= length(buf.constant.x) && return buf.constant.x[i]
+        i -= length(buf.constant.x)
     end
+    isempty(buf.dependent) || return buf.dependent.x[i]
+    throw(BoundsError(buf, i))
 end
 function Base.setindex!(buf::MTKParameters, val, i)
     if i <= length(buf.tunable)
         buf.tunable[i] = val
-    else
+    elseif i <= length(buf.tunable) + length(buf.discrete)
         buf.discrete[i - length(buf.tunable)] = val
+    else
+        buf.constant[i - length(buf.tunable) - length(buf.discrete)] = val
     end
+    buf.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
 end
 
 function Base.iterate(buf::MTKParameters, state = 1)
-    tunable = if isempty(buf.tunable)
-        ()
-    elseif buf.tunable isa ArrayPartition
-        buf.tunable.x
-    end
-    discrete = if isempty(buf.discrete)
-        ()
-    elseif buf.discrete isa ArrayPartition
-        buf.discrete.x
-    end
-    if state <= length(tunable)
-        return (tunable[state], state + 1)
-    elseif state <= length(tunable) + length(discrete)
-        return (discrete[state - length(tunable)], state + 1)
+    total_len = 0
+    isempty(buf.tunable) || (total_len += length(buf.tunable.x))
+    isempty(buf.discrete) || (total_len += length(buf.discrete.x))
+    isempty(buf.constant) || (total_len += length(buf.constant.x))
+    isempty(buf.dependent) || (total_len += length(buf.dependent.x))
+    if state <= total_len
+        return (buf[state], state + 1)
     else
         return nothing
     end
 end
 
 function Base.:(==)(a::MTKParameters, b::MTKParameters)
-    return a.tunable == b.tunable && a.discrete == b.discrete
+    return a.tunable == b.tunable && a.discrete == b.discrete &&
+        a.constant == b.constant && a.dependent == b.dependent
 end
 
 # to support linearize/linearization_function
