Merge pull request #3466 from AayushSabharwal/as/better-nlsys-init

feat: handle `Initial(x)` initialization_eqs in time-independent systems

Author: ChrisRackauckas

Project.toml

@@ -92,7 +92,7 @@ DataInterpolations = "6.4"
 DataStructures = "0.17, 0.18"
 DeepDiffs = "1"
 DelayDiffEq = "5.50"
-DiffEqBase = "6.157"
+DiffEqBase = "6.165.1"
 DiffEqCallbacks = "2.16, 3, 4"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"

src/doc

@@ -676,23 +676,27 @@ function SymbolicUtils.maketerm(::Type{<:BasicSymbolic}, ::Initial, args, meta)
     return metadata(val, meta)
 end
 
+supports_initialization(sys::AbstractSystem) = true
+
 function add_initialization_parameters(sys::AbstractSystem)
     @assert !has_systems(sys) || isempty(get_systems(sys))
+    supports_initialization(sys) || return sys
+    is_initializesystem(sys) && return sys
+
     all_initialvars = Set{BasicSymbolic}()
     # time-independent systems don't initialize unknowns
-    if is_time_dependent(sys)
-        eqs = equations(sys)
-        if !(eqs isa Vector{Equation})
-            eqs = Equation[x for x in eqs if x isa Equation]
-        end
-        obs, eqs = unhack_observed(observed(sys), eqs)
-        for x in Iterators.flatten((unknowns(sys), Iterators.map(eq -> eq.lhs, obs)))
-            x = unwrap(x)
-            if iscall(x) && operation(x) == getindex
-                push!(all_initialvars, arguments(x)[1])
-            else
-                push!(all_initialvars, x)
-            end
+    # but may initialize parameters using guesses for unknowns
+    eqs = equations(sys)
+    if !(eqs isa Vector{Equation})
+        eqs = Equation[x for x in eqs if x isa Equation]
+    end
+    obs, eqs = unhack_observed(observed(sys), eqs)
+    for x in Iterators.flatten((unknowns(sys), Iterators.map(eq -> eq.lhs, obs)))
+        x = unwrap(x)
+        if iscall(x) && operation(x) == getindex
+            push!(all_initialvars, arguments(x)[1])
+        else
+            push!(all_initialvars, x)
         end
     end
     for eq in parameter_dependencies(sys)
@@ -722,15 +726,8 @@ Returns true if the parameter `p` is of the form `Initial(x)`.
 """
 function isinitial(p)
     p = unwrap(p)
-    if iscall(p)
-        operation(p) isa Initial && return true
-        if operation(p) === getindex
-            operation(arguments(p)[1]) isa Initial && return true
-        end
-    else
-        return false
-    end
-    return false
+    return iscall(p) && (operation(p) isa Initial ||
+            operation(p) === getindex && isinitial(arguments(p)[1]))
 end
 
 """
@@ -744,7 +741,8 @@ the global structure of the system.
 One property to note is that 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; split = true, flatten = true)
+function complete(
+        sys::AbstractSystem; split = true, flatten = true, add_initial_parameters = true)
     newunknowns = OrderedSet()
     newparams = OrderedSet()
     iv = has_iv(sys) ? get_iv(sys) : nothing
@@ -765,7 +763,9 @@ function complete(sys::AbstractSystem; split = true, flatten = true)
             @set! newsys.parent = complete(sys; split = false, flatten = false)
         end
         sys = newsys
-        sys = add_initialization_parameters(sys)
+        if add_initial_parameters
+            sys = add_initialization_parameters(sys)
+        end
     end
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
@@ -1345,20 +1345,8 @@ function parameters(sys::AbstractSystem; initial_parameters = false)
     systems = get_systems(sys)
     result = unique(isempty(systems) ? ps :
                     [ps; reduce(vcat, namespace_parameters.(systems))])
-    if !initial_parameters
-        if is_time_dependent(sys)
-            # time-dependent systems have `Initial` parameters for all their
-            # unknowns/pdeps, all of which should be hidden.
-            filter!(x -> !iscall(x) || !isa(operation(x), Initial), result)
-        else
-            # time-independent systems only have `Initial` parameters for
-            # pdeps. Any other `Initial` parameters should be kept (e.g. initialization
-            # systems)
-            filter!(
-                x -> !iscall(x) || !isa(operation(x), Initial) ||
-                         !has_parameter_dependency_with_lhs(sys, only(arguments(x))),
-                result)
-        end
+    if !initial_parameters && !is_initializesystem(sys)
+        filter!(x -> !iscall(x) || !isa(operation(x), Initial), result)
     end
     return result
 end

src/systems/abstractsystem.jl

@@ -421,3 +421,5 @@ end
 function DiscreteFunctionExpr(sys::DiscreteSystem, args...; kwargs...)
     DiscreteFunctionExpr{true}(sys, args...; kwargs...)
 end
+
+supports_initialization(::DiscreteSystem) = false

src/systems/discrete_system/discrete_system.jl

@@ -333,10 +333,8 @@ function SciMLBase.ImplicitDiscreteProblem(
 
     u0map = to_varmap(u0map, dvs)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
-    @show u0map
     f, u0, p = process_SciMLProblem(
         ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, kwargs...)
-    @show u0
 
     kwargs = filter_kwargs(kwargs)
     ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)

src/systems/discrete_system/implicit_discrete_system.jl

@@ -507,7 +507,7 @@ function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothi
             systems = get_systems(sys), defaults = defaults(sys), guesses = guesses(sys),
             parameter_dependencies = parameter_dependencies(sys),
             metadata = get_metadata(sys), gui_metadata = get_gui_metadata(sys))
-        osys = complete(osys)
+        osys = complete(osys; add_initial_parameters = false)
         return ODEProblem(osys, u0map, tspan, parammap; check_length = false,
             build_initializeprob = false, kwargs...)
     else
@@ -685,3 +685,5 @@ function (ratemap::JumpSysMajParamMapper{U, V, W})(maj::MassActionJump, newparam
     scale_rates && JumpProcesses.scalerates!(maj.scaled_rates, maj.reactant_stoch)
     nothing
 end
+
+supports_initialization(::JumpSystem) = false