Merge pull request #4533 from SciML/as/perfect-aliases-arr-eqs

fix: properly handle aliasing elements of unscalarized array variables

Author: AayushSabharwal

src/systems/alias_elimination.jl

@@ -145,6 +145,7 @@ function find_perfect_aliases!(
         c1 == c2 && push!(eqs_to_rm, ieq)
     end
 
+    eqs_to_substitute = Int[]
     for (root, group_vars) in alias_sets
         target = group_target[root]
         state.always_present[target] = true
@@ -165,8 +166,7 @@ function find_perfect_aliases!(
             dnbors = copy(𝑑neighbors(graph, v))
             for e in dnbors
                 snbors = 𝑠neighbors(graph, e)
-                eqs[e] = substitute(eqs[e], subs)
-                original_eqs[e] = substitute(original_eqs[e], subs)
+                push!(eqs_to_substitute, e)
                 Graphs.rem_edge!(graph, e, v)
                 Graphs.add_edge!(graph, e, target)
             end
@@ -185,8 +185,7 @@ function find_perfect_aliases!(
                 dnbors = copy(𝑑neighbors(graph, dv))
                 for e in dnbors
                     snbors = 𝑠neighbors(graph, e)
-                    eqs[e] = substitute(eqs[e], subs)
-                    original_eqs[e] = substitute(original_eqs[e], subs)
+                    push!(eqs_to_substitute, e)
                     Graphs.rem_edge!(graph, e, dv)
                     Graphs.add_edge!(graph, e, dtarget)
                 end
@@ -197,6 +196,23 @@ function find_perfect_aliases!(
         end
     end
 
+    # We need to handle unscalarized array variables
+    for k in keys(subs)
+        k, isarr = split_indexed_var(k)
+        isarr || continue
+        haskey(subs, k) && continue
+        subs[k] = Symbolics.SConst(collect(k))
+    end
+
+    ir = get_irstructure(sys)
+    subber = SU.IRSubstituter{false}(ir, subs)
+    for e in eqs_to_substitute
+        # Double substitute to handle unscalarized array variables. First one
+        # substitutes `x` to `[x[1], x[2]]`. The second substitutes `x[1]` and `x[2]`.
+        eqs[e] = subber(subber(eqs[e]))
+        original_eqs[e] = subber(subber(original_eqs[e]))
+    end
+
     # After substitution, alias equations that connected a sticky non-target
     # variable to a zero-priority variable become structurally identical to the
     # direct alias between the sticky variable and the target (since the

test/reduction.jl

@@ -1,4 +1,5 @@
 using ModelingToolkit, OrdinaryDiffEq, Test, NonlinearSolve, LinearAlgebra
+using Symbolics
 using OrdinaryDiffEqRosenbrock
 using ModelingToolkit: topsort_equations, t_nounits as t, D_nounits as D, unwrap
 
@@ -374,3 +375,12 @@ end
 
     @test Set(unknowns(sys)) == Set([e, c, d])
 end
+
+@testset "`eliminate_perfect_aliases!` correctly handles unscalarized arrays" begin
+    @variables x(t)[1:2] y(t)[1:2]
+    @mtkcompile sys = System([D(x) ~ x, y[1] ~ y[2], dot(x, y) ~ 1], t)
+    @test any(equations(sys)) do eq
+        isequal(eq, 0 ~ 1 - dot(x, Symbolics.SConst([y[2], y[2]]))) ||
+        isequal(eq, 0 ~ 1 - dot(x, Symbolics.SConst([y[1], y[1]])))
+    end
+end