fix all loops in alias graph rather than just first order loops (#2149)

Co-authored-by: Yingbo Ma <[email protected]>

Author: Oscar Smith

src/structural_transformation/pantelides.jl

@@ -121,25 +121,33 @@ function computed_highest_diff_variables(structure, ag::Union{AliasGraph, Nothin
             if ag !== nothing && haskey(ag, var)
                 (_, stem) = ag[var]
                 stem == 0 && continue
-                # Ascend the stem
-                while isempty(𝑑neighbors(graph, var))
-                    var′ = invview(var_to_diff)[var]
-                    var′ === nothing && break
-                    stem′ = invview(var_to_diff)[stem]
-                    avar′ = haskey(ag, var′) ? ag[var′][2] : nothing
-                    if avar′ == stem || var′ == stem
-                        # If we have a self-loop in the stem, we could have the
-                        # var′ also alias to the original stem. In that case, the
-                        # derivative of the stem is highest differentiated, because of the loop
+                # If we have a self-loop in the stem, we could have the
+                # var′ also alias to the original stem. In that case, the
+                # derivative of the stem is highest differentiated, because of the loop
+                loop_found = false
+                var′ = invview(var_to_diff)[var]
+                while var′ !== nothing
+                    if var′ == stem || (haskey(ag, var′) && ag[var′][2] == stem)
                         dstem = var_to_diff[stem]
                         @assert dstem !== nothing
                         varwhitelist[dstem] = true
+                        loop_found = true
                         break
                     end
+                    var′ = invview(var_to_diff)[var′]
+                end
+                loop_found && continue
+                # Ascend the stem
+                while isempty(𝑑neighbors(graph, var))
+                    var′ = invview(var_to_diff)[var]
+                    var′ === nothing && break
+                    loop_found = false
+                    cvar = var′
                     # Invariant from alias elimination: Stem is chosen to have
                     # the highest differentiation order.
+                    stem′ = invview(var_to_diff)[stem]
                     @assert stem′ !== nothing
-                    if avar′ != stem′
+                    if !haskey(ag, var′) || (ag[var′][2] != stem′)
                         varwhitelist[stem] = true
                         break
                     end

test/pantelides.jl

@@ -96,3 +96,32 @@ begin
     # Correct answer is: ẋ
     @test varwhitelist == Bool[1, 0, 0, 1]
 end
+
+begin
+    """
+       Vars: y
+       Eqs: 0 = f(y)
+       Alias: y = ÿ
+    """
+    n_vars = 3
+    ag = AliasGraph(n_vars)
+
+    # Alias: z = 1 * ÿ
+    ag[3] = 1 => 1
+
+    # 0 = f(y)
+    graph = complete(BipartiteGraph([Int[]], n_vars))
+
+    # [y, ẏ, ÿ]
+    var_to_diff = DiffGraph([2, 3, nothing], # primal_to_diff
+                            [nothing, 1, 2]) # diff_to_primal
+
+    # [f(x)]
+    eq_to_diff = DiffGraph([nothing], # primal_to_diff
+                           [nothing]) # diff_to_primal
+    structure = SystemStructure(var_to_diff, eq_to_diff, graph, nothing, nothing, false)
+    varwhitelist = computed_highest_diff_variables(structure, ag)
+
+    # Correct answer is: ẏ
+    @test varwhitelist == Bool[0, 1, 0]
+end