feat: add support for parameter dependencies

Author: AayushSabharwal

src/ModelingToolkit.jl

@@ -36,7 +36,7 @@ using PrecompileTools, Reexport
     using RecursiveArrayTools
 
     using SymbolicIndexingInterface
-    export independent_variables, unknowns, parameters
+    export independent_variables, unknowns, parameters, full_parameters
     import SymbolicUtils
     import SymbolicUtils: istree, arguments, operation, similarterm, promote_symtype,
                           Symbolic, isadd, ismul, ispow, issym, FnType,

src/systems/abstractsystem.jl

@@ -286,27 +286,12 @@ function SymbolicIndexingInterface.parameter_index(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)
-            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])
+        return if (idx = ParameterIndex(ic, sym)) !== nothing
+            idx
+        elseif (idx = ParameterIndex(ic, default_toterm(sym))) !== nothing
+            idx
         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
+            nothing
         end
     end
 
@@ -329,7 +314,7 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym::Sym
 end
 
 function SymbolicIndexingInterface.parameter_symbols(sys::AbstractSystem)
-    return parameters(sys)
+    return full_parameters(sys)
 end
 
 function SymbolicIndexingInterface.is_independent_variable(sys::AbstractSystem, sym)
@@ -419,6 +404,7 @@ for prop in [:eqs
              :metadata
              :gui_metadata
              :discrete_subsystems
+             :parameter_dependencies
              :solved_unknowns
              :split_idxs
              :parent
@@ -750,7 +736,29 @@ function parameters(sys::AbstractSystem)
         ps = first.(ps)
     end
     systems = get_systems(sys)
-    unique(isempty(systems) ? ps : [ps; reduce(vcat, namespace_parameters.(systems))])
+    result = unique(isempty(systems) ? ps :
+                    [ps; reduce(vcat, namespace_parameters.(systems))])
+    if has_parameter_dependencies(sys) &&
+       (pdeps = get_parameter_dependencies(sys)) !== nothing
+        filter(result) do sym
+            !haskey(pdeps, sym)
+        end
+    else
+        result
+    end
+end
+
+function dependent_parameters(sys::AbstractSystem)
+    if has_parameter_dependencies(sys) &&
+       (pdeps = get_parameter_dependencies(sys)) !== nothing
+        collect(keys(pdeps))
+    else
+        []
+    end
+end
+
+function full_parameters(sys::AbstractSystem)
+    vcat(parameters(sys), dependent_parameters(sys))
 end
 
 # required in `src/connectors.jl:437`
@@ -1612,7 +1620,7 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
         kwargs...)
     sts = unknowns(sys)
     t = get_iv(sys)
-    ps = parameters(sys)
+    ps = full_parameters(sys)
     p = reorder_parameters(sys, ps)
 
     fun = generate_function(sys, sts, ps; expression = Val{false})[1]
@@ -2123,3 +2131,17 @@ function Symbolics.substitute(sys::AbstractSystem, rules::Union{Vector{<:Pair},
         error("substituting symbols is not supported for $(typeof(sys))")
     end
 end
+
+function process_parameter_dependencies(pdeps, ps)
+    pdeps === nothing && return pdeps, ps
+    if pdeps isa Vector && eltype(pdeps) <: Pair
+        pdeps = Dict(pdeps)
+    elseif !(pdeps isa Dict)
+        error("parameter_dependencies must be a `Dict` or `Vector{<:Pair}`")
+    end
+
+    ps = filter(ps) do p
+        !haskey(pdeps, p)
+    end
+    return pdeps, ps
+end

src/systems/callbacks.jl

@@ -433,14 +433,14 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
 end
 
 function generate_rootfinding_callback(sys::AbstractODESystem, dvs = unknowns(sys),
-        ps = parameters(sys); kwargs...)
+        ps = full_parameters(sys); kwargs...)
     cbs = continuous_events(sys)
     isempty(cbs) && return nothing
     generate_rootfinding_callback(cbs, sys, dvs, ps; kwargs...)
 end
 
 function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknowns(sys),
-        ps = parameters(sys); kwargs...)
+        ps = full_parameters(sys); kwargs...)
     eqs = map(cb -> cb.eqs, cbs)
     num_eqs = length.(eqs)
     (isempty(eqs) || sum(num_eqs) == 0) && return nothing
@@ -556,7 +556,7 @@ function generate_discrete_callback(cb, sys, dvs, ps; postprocess_affect_expr! =
 end
 
 function generate_discrete_callbacks(sys::AbstractSystem, dvs = unknowns(sys),
-        ps = parameters(sys); kwargs...)
+        ps = full_parameters(sys); kwargs...)
     has_discrete_events(sys) || return nothing
     symcbs = discrete_events(sys)
     isempty(symcbs) && return nothing

src/systems/clock_inference.jl

@@ -198,15 +198,15 @@ function generate_discrete_affect(
             throw = false,
             expression = true,
             output_type = SVector,
-            ps = reorder_parameters(osys, parameters(sys)))
+            ps = reorder_parameters(osys, full_parameters(sys)))
         ni = length(input)
         ns = length(unknowns(sys))
         disc = Func(
             [
                 out,
                 DestructuredArgs(unknowns(osys)),
                 if use_index_cache
-                    DestructuredArgs.(reorder_parameters(osys, parameters(osys)))
+                    DestructuredArgs.(reorder_parameters(osys, full_parameters(osys)))
                 else
                     (DestructuredArgs(appended_parameters),)
                 end...,

src/systems/diffeqs/abstractodesystem.jl

@@ -80,7 +80,8 @@ function calculate_control_jacobian(sys::AbstractODESystem;
     return jac
 end
 
-function generate_tgrad(sys::AbstractODESystem, dvs = unknowns(sys), ps = parameters(sys);
+function generate_tgrad(
+        sys::AbstractODESystem, dvs = unknowns(sys), ps = full_parameters(sys);
         simplify = false, kwargs...)
     tgrad = calculate_tgrad(sys, simplify = simplify)
     pre = get_preprocess_constants(tgrad)
@@ -100,7 +101,7 @@ function generate_tgrad(sys::AbstractODESystem, dvs = unknowns(sys), ps = parame
 end
 
 function generate_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
-        ps = parameters(sys);
+        ps = full_parameters(sys);
         simplify = false, sparse = false, kwargs...)
     jac = calculate_jacobian(sys; simplify = simplify, sparse = sparse)
     pre = get_preprocess_constants(jac)
@@ -118,7 +119,7 @@ function generate_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
 end
 
 function generate_control_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
-        ps = parameters(sys);
+        ps = full_parameters(sys);
         simplify = false, sparse = false, kwargs...)
     jac = calculate_control_jacobian(sys; simplify = simplify, sparse = sparse)
     p = reorder_parameters(sys, ps)
@@ -146,7 +147,7 @@ function generate_dae_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
 end
 
 function generate_function(sys::AbstractODESystem, dvs = unknowns(sys),
-        ps = parameters(sys);
+        ps = full_parameters(sys);
         implicit_dae = false,
         ddvs = implicit_dae ? map(Differential(get_iv(sys)), dvs) :
                nothing,
@@ -314,7 +315,7 @@ end
 
 function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         dvs = unknowns(sys),
-        ps = parameters(sys), u0 = nothing;
+        ps = full_parameters(sys), u0 = nothing;
         version = nothing, tgrad = false,
         jac = false, p = nothing,
         t = nothing,
@@ -830,7 +831,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         kwargs...)
     eqs = equations(sys)
     dvs = unknowns(sys)
-    ps = parameters(sys)
+    ps = full_parameters(sys)
     iv = get_iv(sys)
 
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
@@ -845,7 +846,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             symbolic_u0)
         p, split_idxs = split_parameters_by_type(p)
         if p isa Tuple
-            ps = Base.Fix1(getindex, parameters(sys)).(split_idxs)
+            ps = Base.Fix1(getindex, full_parameters(sys)).(split_idxs)
             ps = (ps...,) #if p is Tuple, ps should be Tuple
         end
     end
@@ -997,7 +998,7 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
                 cbs = CallbackSet(discrete_cbs...)
             end
         else
-            cbs = CallbackSet(cbs, discrete_cbs)
+            cbs = CallbackSet(cbs, discrete_cbs...)
         end
     else
         svs = nothing
@@ -1060,7 +1061,7 @@ function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem, du0map, u0map, tspan
 end
 
 function generate_history(sys::AbstractODESystem, u0; kwargs...)
-    p = reorder_parameters(sys, parameters(sys))
+    p = reorder_parameters(sys, full_parameters(sys))
     build_function(u0, p..., get_iv(sys); expression = Val{false}, kwargs...)
 end
 

src/systems/diffeqs/odesystem.jl

@@ -111,6 +111,11 @@ struct ODESystem <: AbstractODESystem
     """
     discrete_events::Vector{SymbolicDiscreteCallback}
     """
+    A mapping from dependent parameters to expressions describing how they are calculated from
+    other parameters.
+    """
+    parameter_dependencies::Union{Nothing, Dict}
+    """
     Metadata for the system, to be used by downstream packages.
     """
     metadata::Any
@@ -154,7 +159,7 @@ struct ODESystem <: AbstractODESystem
     function ODESystem(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 = nothing, gui_metadata = nothing,
+            devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing,
             substitutions = nothing, complete = false, index_cache = nothing,
             discrete_subsystems = nothing, solved_unknowns = nothing,
@@ -171,8 +176,8 @@ struct ODESystem <: AbstractODESystem
         end
         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,
+            connector_type, preface, cevents, devents, parameter_dependencies, metadata,
+            gui_metadata, tearing_state, substitutions, complete, index_cache,
             discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
 end
@@ -190,6 +195,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         preface = nothing,
         continuous_events = nothing,
         discrete_events = nothing,
+        parameter_dependencies = nothing,
         checks = true,
         metadata = nothing,
         gui_metadata = nothing)
@@ -225,10 +231,12 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     end
     cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
     disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
+    parameter_dependencies, ps′ = process_parameter_dependencies(
+        parameter_dependencies, ps′)
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing,
-        connector_type, preface, cont_callbacks, disc_callbacks,
+        connector_type, preface, cont_callbacks, disc_callbacks, parameter_dependencies,
         metadata, gui_metadata, checks = checks)
 end
 
@@ -323,7 +331,7 @@ function build_explicit_observed_function(sys, ts;
         output_type = Array,
         checkbounds = true,
         drop_expr = drop_expr,
-        ps = parameters(sys),
+        ps = full_parameters(sys),
         throw = true)
     if (isscalar = !(ts isa AbstractVector))
         ts = [ts]

src/systems/diffeqs/sdesystem.jl

@@ -104,6 +104,11 @@ struct SDESystem <: AbstractODESystem
     """
     discrete_events::Vector{SymbolicDiscreteCallback}
     """
+    A mapping from dependent parameters to expressions describing how they are calculated from
+    other parameters.
+    """
+    parameter_dependencies::Union{Nothing, Dict}
+    """
     Metadata for the system, to be used by downstream packages.
     """
     metadata::Any
@@ -128,7 +133,7 @@ struct SDESystem <: AbstractODESystem
             tgrad,
             jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-            cevents, devents, metadata = nothing, gui_metadata = nothing,
+            cevents, devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
             complete = false, index_cache = nothing, parent = nothing;
             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
@@ -144,7 +149,7 @@ struct SDESystem <: AbstractODESystem
         new(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac,
             Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents,
-            metadata, gui_metadata, complete, index_cache, parent)
+            parameter_dependencies, metadata, gui_metadata, complete, index_cache, parent)
     end
 end
 
@@ -161,6 +166,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
         checks = true,
         continuous_events = nothing,
         discrete_events = nothing,
+        parameter_dependencies = nothing,
         metadata = nothing,
         gui_metadata = nothing)
     name === nothing &&
@@ -195,11 +201,12 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
     Wfact_t = RefValue(EMPTY_JAC)
     cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
     disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
-
+    parameter_dependencies, ps′ = process_parameter_dependencies(
+        parameter_dependencies, ps′)
     SDESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, neqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-        cont_callbacks, disc_callbacks, metadata, gui_metadata; checks = checks)
+        cont_callbacks, disc_callbacks, parameter_dependencies, metadata, gui_metadata; checks = checks)
 end
 
 function SDESystem(sys::ODESystem, neqs; kwargs...)
@@ -220,7 +227,7 @@ function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = unknowns(sys),
-        ps = parameters(sys); isdde = false, kwargs...)
+        ps = full_parameters(sys); isdde = false, kwargs...)
     eqs = get_noiseeqs(sys)
     if isdde
         eqs = delay_to_function(sys, eqs)
@@ -285,7 +292,7 @@ function stochastic_integral_transform(sys::SDESystem, correction_factor)
     end
 
     SDESystem(deqs, get_noiseeqs(sys), get_iv(sys), unknowns(sys), parameters(sys),
-        name = name, checks = false)
+        name = name, parameter_dependencies = get_parameter_dependencies(sys), checks = false)
 end
 
 """
@@ -393,7 +400,8 @@ function Girsanov_transform(sys::SDESystem, u; θ0 = 1.0)
     # return modified SDE System
     SDESystem(deqs, noiseeqs, get_iv(sys), unknown_vars, parameters(sys);
         defaults = Dict(θ => θ0), observed = [weight ~ θ / θ0],
-        name = name, checks = false)
+        name = name, parameter_dependencies = get_parameter_dependencies(sys),
+        checks = false)
 end
 
 function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),