Merge pull request #3390 from AayushSabharwal/as/initsys-fixes

ChrisRackauckas · web-flow · commit da48a1f3cba0 · 2025-02-20T13:28:08.000-05:00
fix: fix several bugs
diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml
@@ -37,6 +37,7 @@ jobs:
           - {user: SciML, repo: MethodOfLines.jl, group: Interface}
           - {user: SciML, repo: MethodOfLines.jl, group: 2D_Diffusion}
           - {user: SciML, repo: MethodOfLines.jl, group: DAE}
+          - {user: SciML, repo: ModelingToolkitNeuralNets.jl, group: All}
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1
diff --git a/Project.toml b/Project.toml
@@ -104,11 +104,11 @@ DynamicQuantities = "^0.11.2, 0.12, 0.13, 1"
 EnumX = "1.0.4"
 ExprTools = "0.1.10"
 Expronicon = "0.8"
+FMI = "0.14"
 FindFirstFunctions = "1"
 ForwardDiff = "0.10.3"
 FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
-FMI = "0.14"
 Graphs = "1.5.2"
 HomotopyContinuation = "2.11"
 InfiniteOpt = "0.5"
@@ -119,6 +119,7 @@ LabelledArrays = "1.3"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
 Libdl = "1"
 LinearAlgebra = "1"
+Logging = "1"
 MLStyle = "0.4.17"
 ModelingToolkitStandardLibrary = "2.19"
 NaNMath = "0.3, 1"
@@ -128,7 +129,7 @@ OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
 OrdinaryDiffEqCore = "1.15.0"
 OrdinaryDiffEqDefault = "1.2"
-OrdinaryDiffEqNonlinearSolve = "1.3.0"
+OrdinaryDiffEqNonlinearSolve = "1.5.0"
 PrecompileTools = "1"
 REPL = "1"
 RecursiveArrayTools = "3.26"
@@ -143,11 +144,11 @@ SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
-StochasticDiffEq = "6.72.1"
 StochasticDelayDiffEq = "1.8.1"
+StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.37"
 SymbolicUtils = "3.14"
-Symbolics = "6.29"
+Symbolics = "6.29.1"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
@@ -164,6 +165,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -187,4 +189,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve", "Logging"]
diff --git a/ext/MTKBifurcationKitExt.jl b/ext/MTKBifurcationKitExt.jl
@@ -97,7 +97,10 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
     @set! nsys.index_cache = nothing # force usage of a parameter vector instead of `MTKParameters`
     # Creates F and J functions.
     ofun = NonlinearFunction(nsys; jac = jac)
-    F = ofun.f
+    F = let f = ofun.f
+        _f(resid, u, p) = (f(resid, u, p); resid)
+        _f(u, p) = f(u, p)
+    end
     J = jac ? ofun.jac : nothing
 
     # Converts the input state guess.
@@ -136,6 +139,7 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
         args...;
         record_from_solution = record_from_solution,
         J = J,
+        inplace = true,
         kwargs...)
 end
 
diff --git a/ext/MTKChainRulesCoreExt.jl b/ext/MTKChainRulesCoreExt.jl
@@ -16,7 +16,12 @@ end
 
 function subset_idxs(idxs, portion, template)
     ntuple(Val(length(template))) do subi
-        [Base.tail(idx.idx) for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+        result = [Base.tail(idx.idx)
+                  for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+        if isempty(result)
+            result = []
+        end
+        result
     end
 end
 
diff --git a/src/linearization.jl b/src/linearization.jl
@@ -170,8 +170,8 @@ function (linfun::LinearizationFunction)(u, p, t)
     if u !== nothing # Handle systems without unknowns
         linfun.num_states == length(u) ||
             error("Number of unknown variables ($(linfun.num_states)) does not match the number of input unknowns ($(length(u)))")
-        integ_cache = linfun.caches
-        integ = MockIntegrator{true}(u, p, t, integ_cache)
+        integ_cache = (linfun.caches,)
+        integ = MockIntegrator{true}(u, p, t, integ_cache, nothing)
         u, p, success = SciMLBase.get_initial_values(
             linfun.prob, integ, fun, linfun.initializealg, Val(true);
             linfun.initialize_kwargs...)
@@ -218,7 +218,7 @@ Mock `DEIntegrator` to allow using `CheckInit` without having to create a new in
 
 $(TYPEDFIELDS)
 """
-struct MockIntegrator{iip, U, P, T, C} <: SciMLBase.DEIntegrator{Nothing, iip, U, T}
+struct MockIntegrator{iip, U, P, T, C, O} <: SciMLBase.DEIntegrator{Nothing, iip, U, T}
     """
     The state vector.
     """
@@ -235,10 +235,14 @@ struct MockIntegrator{iip, U, P, T, C} <: SciMLBase.DEIntegrator{Nothing, iip, U
     The integrator cache.
     """
     cache::C
+    """
+    Integrator "options" for `CheckInit`.
+    """
+    opts::O
 end
 
-function MockIntegrator{iip}(u::U, p::P, t::T, cache::C) where {iip, U, P, T, C}
-    return MockIntegrator{iip, U, P, T, C}(u, p, t, cache)
+function MockIntegrator{iip}(u::U, p::P, t::T, cache::C, opts::O) where {iip, U, P, T, C, O}
+    return MockIntegrator{iip, U, P, T, C, O}(u, p, t, cache, opts)
 end
 
 SymbolicIndexingInterface.state_values(integ::MockIntegrator) = integ.u
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
@@ -647,6 +647,7 @@ function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, arr
         # HACK 1
         if cse && is_getindexed_array(rhs)
             rhs_arr = arguments(rhs)[1]
+            iscall(rhs_arr) && operation(rhs_arr) isa Symbolics.Operator && continue
             if !haskey(rhs_to_tempvar, rhs_arr)
                 tempvar = gensym(Symbol(lhs))
                 N = length(rhs_arr)
@@ -719,6 +720,7 @@ function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, arr
         Symbolics.shape(sym) != Symbolics.Unknown() || continue
         arg1 = arguments(sym)[1]
         cnt = get(arr_obs_occurrences, arg1, 0)
+        cnt == 0 && continue
         arr_obs_occurrences[arg1] = cnt + 1
     end
     for eq in neweqs
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
@@ -669,7 +669,11 @@ end
 
 # This is required so `fast_substitute` works
 function SymbolicUtils.maketerm(::Type{<:BasicSymbolic}, ::Initial, args, meta)
-    return metadata(Initial()(args...), meta)
+    val = Initial()(args...)
+    if symbolic_type(val) == NotSymbolic()
+        return val
+    end
+    return metadata(val, meta)
 end
 
 function add_initialization_parameters(sys::AbstractSystem)
diff --git a/src/systems/codegen_utils.jl b/src/systems/codegen_utils.jl
@@ -12,8 +12,14 @@ end
 
 Given the arguments to `build_function_wrapper`, return a list of assignments which
 reconstruct array variables if they are present scalarized in `args`.
+
+# Keyword Arguments
+
+- `argument_name` a function of the form `(::Int) -> Symbol` which takes the index of
+  an argument to the generated function and returns the name of the argument in the
+  generated function.
 """
-function array_variable_assignments(args...)
+function array_variable_assignments(args...; argument_name = generated_argument_name)
     # map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
     var_to_arridxs = Dict{BasicSymbolic, Array{Tuple{Int, Int}}}()
     for (i, arg) in enumerate(args)
@@ -60,12 +66,12 @@ function array_variable_assignments(args...)
             end
             # view and reshape
 
-            expr = term(reshape, term(view, generated_argument_name(buffer_idx), idxs),
+            expr = term(reshape, term(view, argument_name(buffer_idx), idxs),
                 size(arrvar))
         else
             elems = map(idxs) do idx
                 i, j = idx
-                term(getindex, generated_argument_name(i), j)
+                term(getindex, argument_name(i), j)
             end
             # use `MakeArray` syntax and generate a stack-allocated array
             expr = term(SymbolicUtils.Code.create_array, SArray, nothing,
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -1291,7 +1291,9 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         guesses = Dict()
     end
 
-    u0map = merge(ModelingToolkit.guesses(sys), todict(guesses), todict(u0map))
+    filter_missing_values!(u0map)
+    filter_missing_values!(parammap)
+    u0map = merge(ModelingToolkit.guesses(sys), todict(guesses), u0map)
 
     fullmap = merge(u0map, parammap)
     u0T = Union{}
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -476,11 +476,16 @@ function build_explicit_observed_function(sys, ts;
         end
     end
     allsyms = Set(all_symbols(sys))
+    iv = has_iv(sys) ? get_iv(sys) : nothing
     for var in vs
         var = unwrap(var)
         newvar = get(ns_map, var, nothing)
         if newvar !== nothing
             namespace_subs[var] = newvar
+            var = newvar
+        end
+        if throw && !var_in_varlist(var, allsyms, iv)
+            Base.throw(ArgumentError("Symbol $var is not present in the system."))
         end
     end
     ts = fast_substitute(ts, namespace_subs)
@@ -522,12 +527,14 @@ function build_explicit_observed_function(sys, ts;
     if fns isa Tuple
         oop, iip = eval_or_rgf.(fns; eval_expression, eval_module)
         f = GeneratedFunctionWrapper{(
-            p_start, length(args) - length(ps) + 1, is_split(sys))}(oop, iip)
+            p_start + is_dde(sys), length(args) - length(ps) + 1 + is_dde(sys), is_split(sys))}(
+            oop, iip)
         return return_inplace ? (f, f) : f
     else
         f = eval_or_rgf(fns; eval_expression, eval_module)
         f = GeneratedFunctionWrapper{(
-            p_start, length(args) - length(ps) + 1, is_split(sys))}(f, nothing)
+            p_start + is_dde(sys), length(args) - length(ps) + 1 + is_dde(sys), is_split(sys))}(
+            f, nothing)
         return f
     end
 end
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
@@ -322,7 +322,8 @@ function (rip::ReconstructInitializeprob)(srcvalp, dstvalp)
     if state_values(dstvalp) === nothing
         return nothing, newp
     end
-    T = eltype(state_values(srcvalp))
+    srcu0 = state_values(srcvalp)
+    T = srcu0 === nothing || isempty(srcu0) ? Union{} : eltype(srcu0)
     if parameter_values(dstvalp) isa MTKParameters
         if !isempty(newp.tunable)
             T = promote_type(eltype(newp.tunable), T)
@@ -332,7 +333,7 @@ function (rip::ReconstructInitializeprob)(srcvalp, dstvalp)
     end
     if T == eltype(state_values(dstvalp))
         u0 = state_values(dstvalp)
-    else
+    elseif T != Union{}
         u0 = T.(state_values(dstvalp))
     end
     buf, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Tunable(), newp)
@@ -511,9 +512,10 @@ end
 
 function SciMLBase.late_binding_update_u0_p(
         prob, sys::AbstractSystem, u0, p, t0, newu0, newp)
-    u0 === missing && return newu0, newp
-    eltype(u0) <: Pair || return newu0, newp
+    u0 === missing && return newu0, (p === missing ? copy(newp) : newp)
+    eltype(u0) <: Pair || return newu0, (p === missing ? copy(newp) : newp)
 
+    newp = p === missing ? copy(newp) : newp
     newu0 = DiffEqBase.promote_u0(newu0, newp, t0)
     tunables, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Tunable(), newp)
     tunables = DiffEqBase.promote_u0(tunables, newu0, t0)
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
@@ -578,11 +578,19 @@ function CacheWriter(sys::AbstractSystem, buffer_types::Vector{TypeT},
         Symbol(:tmp, i) ← SetArray(true, :(out[$i]), get(exprs, T, []))
     end
 
+    function argument_name(i::Int)
+        if i <= length(solsyms)
+            return :($(generated_argument_name(1))[$i])
+        end
+        return generated_argument_name(i - length(solsyms))
+    end
+    array_assignments = array_variable_assignments(solsyms...; argument_name)
     fn = build_function_wrapper(
-        sys, nothing, :out, DestructuredArgs(DestructuredArgs.(solsyms)),
-        DestructuredArgs.(rps)...; p_start = 3, p_end = length(rps) + 2,
+        sys, nothing, :out,
+        DestructuredArgs(DestructuredArgs.(solsyms), generated_argument_name(1)),
+        rps...; p_start = 3, p_end = length(rps) + 2,
         expression = Val{true}, add_observed = false,
-        extra_assignments = [obs_assigns; body])
+        extra_assignments = [array_assignments; obs_assigns; body])
     fn = eval_or_rgf(fn; eval_expression, eval_module)
     fn = GeneratedFunctionWrapper{(3, 3, is_split(sys))}(fn, nothing)
     return CacheWriter(fn)
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
@@ -117,6 +117,7 @@ Variables as they are specified in `vars` will take priority over their `toterm`
 function add_fallbacks!(
         varmap::AnyDict, vars::Vector, fallbacks::Dict; toterm = default_toterm)
     missingvars = Set()
+    arrvars = Set()
     for var in vars
         haskey(varmap, var) && continue
         ttvar = toterm(var)
@@ -157,6 +158,7 @@ function add_fallbacks!(
                     fallbacks, arrvar, nothing) get(fallbacks, ttarrvar, nothing) Some(nothing)
                 if val !== nothing
                     val = val[idxs...]
+                    is_sized_array_symbolic(arrvar) && push!(arrvars, arrvar)
                 end
             else
                 val = nothing
@@ -170,6 +172,10 @@ function add_fallbacks!(
         end
     end
 
+    for arrvar in arrvars
+        varmap[arrvar] = collect(arrvar)
+    end
+
     return missingvars
 end
 
@@ -269,9 +275,9 @@ entry for `eq.lhs`, insert the reverse mapping if `eq.rhs` is not a number.
 """
 function add_observed_equations!(varmap::AbstractDict, eqs)
     for eq in eqs
-        if haskey(varmap, eq.lhs)
+        if var_in_varlist(eq.lhs, keys(varmap), nothing)
             eq.rhs isa Number && continue
-            haskey(varmap, eq.rhs) && continue
+            var_in_varlist(eq.rhs, keys(varmap), nothing) && continue
             !iscall(eq.rhs) || issym(operation(eq.rhs)) || continue
             varmap[eq.rhs] = eq.lhs
         else
diff --git a/src/utils.jl b/src/utils.jl
@@ -1267,3 +1267,22 @@ function symbol_to_symbolic(sys::AbstractSystem, sym; allsyms = all_symbols(sys)
     end
     return sym
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Check if `var` is present in `varlist`. `iv` is the independent variable of the system,
+and should be `nothing` if not applicable.
+"""
+function var_in_varlist(var, varlist::AbstractSet, iv)
+    var = unwrap(var)
+    # simple case
+    return var in varlist ||
+           # indexed array symbolic, unscalarized array present
+           (iscall(var) && operation(var) === getindex && arguments(var)[1] in varlist) ||
+           # unscalarized sized array symbolic, all scalarized elements present
+           (symbolic_type(var) == ArraySymbolic() && is_sized_array_symbolic(var) &&
+            all(x -> x in varlist, collect(var))) ||
+           # delayed variables
+           (isdelay(var, iv) && var_in_varlist(operation(var)(iv), varlist, iv))
+end
diff --git a/test/debugging.jl b/test/debugging.jl
@@ -1,4 +1,5 @@
 using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq, SymbolicIndexingInterface
+import Logging
 using ModelingToolkit: t_nounits as t, D_nounits as D, ASSERTION_LOG_VARIABLE
 
 @variables x(t)
@@ -25,11 +26,11 @@ end
         dsys = debug_system(sys; functions = [])
         @test is_parameter(dsys, ASSERTION_LOG_VARIABLE)
         prob = Problem(dsys, [x => 0.1], (0.0, 5.0))
-        sol = solve(prob, alg)
-        @test !SciMLBase.successful_retcode(sol)
-        prob.ps[ASSERTION_LOG_VARIABLE] = true
         sol = @test_logs (:error, r"ohno") match_mode=:any solve(prob, alg)
         @test !SciMLBase.successful_retcode(sol)
+        prob.ps[ASSERTION_LOG_VARIABLE] = false
+        sol = @test_logs min_level=Logging.Error solve(prob, alg)
+        @test !SciMLBase.successful_retcode(sol)
     end
 end
 
@@ -41,10 +42,10 @@ end
         dsys = debug_system(outer; functions = [])
         @test is_parameter(dsys, ASSERTION_LOG_VARIABLE)
         prob = Problem(dsys, [inner.x => 0.1], (0.0, 5.0))
-        sol = solve(prob, alg)
-        @test !SciMLBase.successful_retcode(sol)
-        prob.ps[ASSERTION_LOG_VARIABLE] = true
         sol = @test_logs (:error, r"ohno") match_mode=:any solve(prob, alg)
         @test !SciMLBase.successful_retcode(sol)
+        prob.ps[ASSERTION_LOG_VARIABLE] = false
+        sol = @test_logs min_level=Logging.Error solve(prob, alg)
+        @test !SciMLBase.successful_retcode(sol)
     end
 end
diff --git a/test/downstream/inversemodel.jl b/test/downstream/inversemodel.jl
diff --git a/test/extensions/ad.jl b/test/extensions/ad.jl
diff --git a/test/fmi/fmi.jl b/test/fmi/fmi.jl
diff --git a/test/initial_values.jl b/test/initial_values.jl
diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
