refactor: account for new parameter dependencies

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -295,7 +295,7 @@ function has_parameter_dependency_with_lhs(sys, sym)
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
         return haskey(ic.dependent_pars_to_timeseries, unwrap(sym))
     else
-        return any(isequal(sym), [eq.lhs for eq in parameter_dependencies(sys)])
+        return any(isequal(sym), [eq.lhs for eq in get_parameter_dependencies(sys)])
     end
 end
 
@@ -565,7 +565,7 @@ function add_initialization_parameters(sys::AbstractSystem; split = true)
         D = Differential(get_iv(sys))
         union!(all_initialvars, [D(v) for v in all_initialvars if iscall(v)])
     end
-    for eq in parameter_dependencies(sys)
+    for eq in get_parameter_dependencies(sys)
         is_variable_floatingpoint(eq.lhs) || continue
         push!(all_initialvars, eq.lhs)
     end
@@ -1314,8 +1314,15 @@ function parameter_dependencies(sys::AbstractSystem)
     get_parameter_dependencies(sys)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Return all of the parameters of the system, including hidden initial parameters and ones
+eliminated via `parameter_dependencies`.
+"""
 function full_parameters(sys::AbstractSystem)
-    vcat(parameters(sys; initial_parameters = true), dependent_parameters(sys))
+    dep_ps = [eq.lhs for eq in get_parameter_dependencies(sys)]
+    vcat(parameters(sys; initial_parameters = true), dep_ps)
 end
 
 """
@@ -2095,7 +2102,7 @@ function Base.show(
     end
 
     # Print parameter dependencies
-    npdeps = has_parameter_dependencies(sys) ? length(parameter_dependencies(sys)) : 0
+    npdeps = has_parameter_dependencies(sys) ? length(get_parameter_dependencies(sys)) : 0
     npdeps > 0 && printstyled(io, "\nParameter dependencies ($npdeps):"; bold)
     npdeps > 0 && hint && print(io, " see parameter_dependencies($name)")
 
@@ -2604,15 +2611,15 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
     eqs = union(get_eqs(basesys), get_eqs(sys))
     sts = union(get_unknowns(basesys), get_unknowns(sys))
     ps = union(get_ps(basesys), get_ps(sys))
-    dep_ps = union(parameter_dependencies(basesys), parameter_dependencies(sys))
+    dep_ps = union(get_parameter_dependencies(basesys), get_parameter_dependencies(sys))
     obs = union(get_observed(basesys), get_observed(sys))
     cevs = union(get_continuous_events(basesys), get_continuous_events(sys))
     devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
     meta = merge(get_metadata(basesys), get_metadata(sys))
     syss = union(get_systems(basesys), get_systems(sys))
     args = length(ivs) == 0 ? (eqs, sts, ps) : (eqs, ivs[1], sts, ps)
-    kwargs = (parameter_dependencies = dep_ps, observed = obs, continuous_events = cevs,
+    kwargs = (observed = obs, continuous_events = cevs,
         discrete_events = devs, defaults = defs, systems = syss, metadata = meta,
         name = name, description = description, gui_metadata = gui_metadata)
 
@@ -2626,7 +2633,10 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
             kwargs, (; assertions = merge(get_assertions(basesys), get_assertions(sys))))
     end
 
-    return T(args...; kwargs...)
+    newsys = T(args...; kwargs...)
+    @set! newsys.parameter_dependencies = dep_ps
+
+    return newsys
 end
 
 """
@@ -2768,9 +2778,10 @@ function Symbolics.substitute(sys::AbstractSystem, rules::Union{Vector{<:Pair},
         initialization_eqs = fast_substitute(get_initialization_eqs(sys), rules)
         cstrs = fast_substitute(get_constraints(sys), rules)
         subsys = map(s -> substitute(s, rules), get_systems(sys))
-        System(eqs, get_iv(sys); name = nameof(sys), defaults = defs,
-            guesses = guess, parameter_dependencies = pdeps, systems = subsys, noise_eqs,
+        newsys = System(eqs, get_iv(sys); name = nameof(sys), defaults = defs,
+            guesses = guess, systems = subsys, noise_eqs,
             observed, initialization_eqs, constraints = cstrs)
+        @set! newsys.parameter_dependencies = pdeps
     else
         error("substituting symbols is not supported for $(typeof(sys))")
     end
@@ -2846,7 +2857,7 @@ See also: [`ModelingToolkit.dump_variable_metadata`](@ref), [`ModelingToolkit.du
 """
 function dump_parameters(sys::AbstractSystem)
     defs = defaults(sys)
-    pdeps = parameter_dependencies(sys)
+    pdeps = get_parameter_dependencies(sys)
     metas = map(dump_variable_metadata.(parameters(sys))) do meta
         if haskey(defs, meta.var)
             meta = merge(meta, (; default = defs[meta.var]))

src/systems/codegen_utils.jl

@@ -246,7 +246,7 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
         p_start += 1
         p_end += 1
     end
-    pdeps = parameter_dependencies(sys)
+    pdeps = get_parameter_dependencies(sys)
 
     # only get the necessary observed equations, avoiding extra computation
     if add_observed && !isempty(obs)

src/systems/diffeqs/basic_transformations.jl

@@ -223,12 +223,13 @@ function change_independent_variable(
         wasflat = isempty(systems)
         sys = typeof(sys)( # recreate system with transformed fields
             eqs, iv2, unknowns, ps; observed, initialization_eqs,
-            parameter_dependencies, defaults, guesses, connector_type,
+            defaults, guesses, connector_type,
             assertions, name = nameof(sys), description = description(sys)
         )
         sys = compose(sys, systems) # rebuild hierarchical system
         if wascomplete
             sys = complete(sys; split = wassplit, flatten = wasflat) # complete output if input was complete
+            @set! sys.parameter_dependencies = parameter_dependencies
         end
         return sys
     end

src/systems/index_cache.jl

@@ -315,7 +315,7 @@ function IndexCache(sys::AbstractSystem)
     dependent_pars_to_timeseries = Dict{
         Union{BasicSymbolic, CallWithMetadata}, TimeseriesSetType}()
 
-    for eq in parameter_dependencies(sys)
+    for eq in get_parameter_dependencies(sys)
         sym = eq.lhs
         vs = vars(eq.rhs)
         timeseries = TimeseriesSetType()

src/systems/nonlinear/initializesystem.jl

@@ -167,15 +167,15 @@ function generate_initializesystem_timevarying(sys::AbstractSystem;
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
     end
-    return System(eqs_ics,
+    isys = System(eqs_ics,
         vars,
         pars;
         defaults = defs,
         checks = check_units,
-        parameter_dependencies = new_parameter_deps,
         name,
         is_initializesystem = true,
         kwargs...)
+    @set isys.parameter_dependencies = new_parameter_deps
 end
 
 """
@@ -280,15 +280,15 @@ function generate_initializesystem_timeindependent(sys::AbstractSystem;
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
     end
-    return System(eqs_ics,
+    isys = System(eqs_ics,
         vars,
         pars;
         defaults = defs,
         checks = check_units,
-        parameter_dependencies = new_parameter_deps,
         name,
         is_initializesystem = true,
         kwargs...)
+    @set isys.parameter_dependencies = new_parameter_deps
 end
 
 """

src/systems/system.jl

@@ -723,8 +723,8 @@ function flatten(sys::System, noeqs = false)
         parameters(sys; initial_parameters = true), brownians(sys);
         jumps = jumps(sys), constraints = constraints(sys), costs = costs,
         consolidate = default_consolidate, observed = observed(sys),
-        parameter_dependencies = parameter_dependencies(sys), defaults = defaults(sys),
-        guesses = guesses(sys), continuous_events = continuous_events(sys),
+        defaults = defaults(sys), guesses = guesses(sys),
+        continuous_events = continuous_events(sys),
         discrete_events = discrete_events(sys), assertions = assertions(sys),
         is_dde = is_dde(sys), tstops = symbolic_tstops(sys),
         initialization_eqs = initialization_equations(sys),
@@ -924,12 +924,12 @@ function NonlinearSystem(sys::System)
         fast_substitute(eq, subrules)
     end
     nsys = System(eqs, unknowns(sys), [parameters(sys); get_iv(sys)];
-        parameter_dependencies = parameter_dependencies(sys),
         defaults = merge(defaults(sys), Dict(get_iv(sys) => Inf)), guesses = guesses(sys),
         initialization_eqs = initialization_equations(sys), name = nameof(sys),
-        observed = obs)
+        observed = obs, systems = map(NonlinearSystem, get_systems(sys)))
     if iscomplete(sys)
         nsys = complete(nsys; split = is_split(sys))
+        @set! nsys.parameter_dependencies = get_parameter_dependencies(sys)
     end
     return nsys
 end
@@ -983,25 +983,29 @@ end
 
 Construct a system of equations with associated noise terms.
 """
-function SDESystem(eqs::Vector{Equation}, noise, iv; is_scalar_noise = false, kwargs...)
+function SDESystem(eqs::Vector{Equation}, noise, iv; is_scalar_noise = false,
+        parameter_dependencies = Equation[], kwargs...)
     if is_scalar_noise
         if !(noise isa Vector)
             throw(ArgumentError("Expected noise to be a vector if `is_scalar_noise`"))
         end
         noise = repeat(reshape(noise, (1, :)), length(eqs))
     end
-    return System(eqs, iv; noise_eqs = noise, kwargs...)
+    sys = System(eqs, iv; noise_eqs = noise, kwargs...)
+    @set sys.parameter_dependencies = parameter_dependencies
 end
 
 function SDESystem(
-        eqs::Vector{Equation}, noise, iv, dvs, ps; is_scalar_noise = false, kwargs...)
+        eqs::Vector{Equation}, noise, iv, dvs, ps; is_scalar_noise = false,
+        parameter_dependencies = Equation[], kwargs...)
     if is_scalar_noise
         if !(noise isa Vector)
             throw(ArgumentError("Expected noise to be a vector if `is_scalar_noise`"))
         end
         noise = repeat(reshape(noise, (1, :)), length(eqs))
     end
-    return System(eqs, iv, dvs, ps; noise_eqs = noise, kwargs...)
+    sys = System(eqs, iv, dvs, ps; noise_eqs = noise, kwargs...)
+    @set sys.parameter_dependencies = parameter_dependencies
 end
 
 function SDESystem(sys::System, noise; kwargs...)

src/systems/systems.jl

@@ -159,10 +159,12 @@ function __mtkcompile(sys::AbstractSystem; simplify = false,
         ssys = System(Vector{Equation}(full_equations(ode_sys)),
             get_iv(ode_sys), unknowns(ode_sys), parameters(ode_sys); noise_eqs,
             name = nameof(ode_sys), observed = observed(ode_sys), defaults = defaults(sys),
-            parameter_dependencies = parameter_dependencies(sys), assertions = assertions(sys),
+            assertions = assertions(sys),
             guesses = guesses(sys), initialization_eqs = initialization_equations(sys),
             continuous_events = continuous_events(sys),
             discrete_events = discrete_events(sys))
+        @set! ssys.parameter_dependencies = get_parameter_dependencies(sys)
+        return ssys
     end
 end
 

src/utils.jl

@@ -512,15 +512,6 @@ function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Dif
             collect_vars!(unknowns, parameters, eq, iv; depth, op)
         end
     end
-    if has_parameter_dependencies(sys)
-        for eq in parameter_dependencies(sys)
-            if eq isa Pair
-                collect_vars!(unknowns, parameters, eq, iv; depth, op)
-            else
-                collect_vars!(unknowns, parameters, eq, iv; depth, op)
-            end
-        end
-    end
     if has_constraints(sys)
         for eq in constraints(sys)
             eqtype_supports_collect_vars(eq) || continue