Merge pull request #4678 from SciML/as/fix-ci

fix: fix bugs, improve type-stability

Author: AayushSabharwal

lib/ModelingToolkitBase/src/problems/nonlinearproblem.jl

@@ -90,8 +90,8 @@ function generate_nonlinear_bounds(sys::AbstractSystem, op)
     left_merge!(lbmap, op)
     lb = varmap_to_vars(lbmap, dvs; tofloat = false)
     ubmap = SymmapT()
-    for (var, b) in zip(dvs, lb)
-        b === -Inf && continue
+    for (var, b) in zip(dvs, ub)
+        b === Inf && continue
         write_possibly_indexed_array!(ubmap, var, Symbolics.SConst(b), COMMON_NOTHING)
     end
     left_merge!(ubmap, op)

lib/ModelingToolkitBase/src/systems/index_cache.jl

@@ -644,14 +644,20 @@ function reorder_parameters(ic::IndexCache, ps::Vector{SymbolicT}; drop_missing
             if i isa Int
                 param_buf[i] = p
             else
-                param_buf[i] = collect(p)
+                i = (first(i)::Int):(last(i)::Int)
+                for (buf_i, p_i) in zip(i, SU.stable_eachindex(p))
+                    param_buf[buf_i] = p[p_i]
+                end
             end
         elseif haskey(ic.initials_idx, p)
             i = ic.initials_idx[p]
             if i isa Int
                 initials_buf[i] = p
             else
-                initials_buf[i] = collect(p)
+                i = (first(i)::Int):(last(i)::Int)
+                for (buf_i, p_i) in zip(i, SU.stable_eachindex(p))
+                    initials_buf[buf_i] = p[p_i]
+                end
             end
         elseif haskey(ic.constant_idx, p)
             i, j = ic.constant_idx[p]

lib/ModelingToolkitBase/src/systems/system.jl

@@ -1176,7 +1176,7 @@ end
 
 Get the metadata associated with key `k` in system `sys` or `default` if it does not exist.
 """
-function SymbolicUtils.getmetadata(sys::AbstractSystem, k::DataType, default)
+function SymbolicUtils.getmetadata(sys::AbstractSystem, @nospecialize(k::DataType), default)
     meta = get_metadata(sys)
     return get(meta, k, default)
 end
@@ -1188,7 +1188,7 @@ Set the metadata associated with key `k` in system `sys` to value `v`. This is a
 out-of-place operation, and will return a shallow copy of `sys` with the appropriate
 metadata values.
 """
-function SymbolicUtils.setmetadata(sys::AbstractSystem, k::DataType, v)
+function SymbolicUtils.setmetadata(sys::AbstractSystem, @nospecialize(k::DataType), @nospecialize(v))
     meta = get_metadata(sys)
     meta = Base.ImmutableDict(meta, k => v)::MetadataT
     return @set sys.metadata = meta

lib/ModelingToolkitBase/src/systems/systems.jl

@@ -983,6 +983,12 @@ function observed2graph(sys::AbstractSystem, eqs::Vector{Equation}, unknowns::Ve
     return graph, assigns
 end
 
+"""
+Toggle to control whether `topsort_equations` prints the equations in the
+cycle, if present.
+"""
+TOPSORT_EQS_PRINT_CYCLE::Bool = false
+
 """
     $(TYPEDSIGNATURES)
 
@@ -1054,7 +1060,37 @@ function topsort_equations(sys::AbstractSystem, eqs::Vector{Equation}, unknowns:
         end
     end
 
-    (check && idx != neqs) && throw(ArgumentError("The equations have at least one cycle."))
+    if check && idx != neqs
+        # Build a directed eq→eq subgraph over unsorted equations, find smallest SCC.
+        if TOPSORT_EQS_PRINT_CYCLE
+            unsorted = findall(>(0), degrees)
+            unsorted_set = Set(unsorted)
+            n_unsorted = length(unsorted)
+            old_to_new = Dict(old => new for (new, old) in enumerate(unsorted))
+
+            g = SimpleDiGraph(n_unsorted)
+            for src_old in unsorted
+                for dst_old in 𝑑neighbors(graph, assigns[src_old])
+                    dst_old in unsorted_set || continue
+                    add_edge!(g, old_to_new[src_old], old_to_new[dst_old])
+                end
+            end
+
+            sccs = strongly_connected_components(g)
+            nontrivial = filter(scc -> length(scc) >= 2, sccs)
+            smallest_new = isempty(nontrivial) ? collect(1:n_unsorted) :
+                nontrivial[argmin(length.(nontrivial))]
+
+            println("=== topsort_equations: CYCLE DETECTED ===")
+            println("Smallest cycle ($(length(smallest_new)) equations):")
+            for new_idx in smallest_new
+                old_idx = unsorted[new_idx]
+                println("  LHS = $(unknowns[assigns[old_idx]])")
+                println("  EQ  = $(eqs[old_idx])")
+            end
+        end
+        throw(ArgumentError("The equations have at least one cycle."))
+    end
 
     return ordered_eqs
 end

lib/ModelingToolkitBase/test/initializationsystem.jl

@@ -326,8 +326,15 @@ end
 initprob = ModelingToolkitBase.InitializationProblem(sys, 0.0; missing_guess_value)
 conditions = getfield.(equations(initprob.f.sys), :rhs)
 
-@test initprob isa SCCNonlinearProblem
-@test initprob.probs isa Tuple{<:LinearProblem}
+# The init system as-is is linear. Since it is so severely overdetermined, MTK can sometimes
+# choose a bad simplification that results in a nonlinear system.
+# TODO: Support retaining `observed` equations in proper `mtkcompile` and generate
+# `x ~ Initial(x)` as such `observed` equations. This makes `mtkcompile`'s job easier
+# on initialization systems, but can lead to parameter-only equations.
+if !@isdefined(ModelingToolkit)
+    @test initprob isa SCCNonlinearProblem
+    @test initprob.probs isa Tuple{<:LinearProblem}
+end
 if @isdefined(ModelingToolkit)
     initsol = solve(initprob, reltol = 1.0e-12, abstol = 1.0e-12)
     @test SciMLBase.successful_retcode(initsol)