fix edge cases in pantiledes (#2137)

Fix edge cases in pantiledes where there are edge cases like self loops in the stem. These occured in systems like
x->D(x)

or
x->y
D(x)->D(y)
x->D(y)

and would lead to the `@assert stem′ !== nothing` failing.

Author: Oscar Smith

src/structural_transformation/pantelides.jl

@@ -126,10 +126,20 @@ function computed_highest_diff_variables(structure, ag::Union{AliasGraph, Nothin
                     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
+                        dstem = var_to_diff[stem]
+                        @assert dstem !== nothing
+                        varwhitelist[dstem] = true
+                        break
+                    end
                     # Invariant from alias elimination: Stem is chosen to have
                     # the highest differentiation order.
                     @assert stem′ !== nothing
-                    if !haskey(ag, var′) || ag[var′][2] != stem′
+                    if avar′ != stem′
                         varwhitelist[stem] = true
                         break
                     end

test/pantelides.jl

@@ -31,7 +31,7 @@ begin
     varwhitelist = computed_highest_diff_variables(structure, ag)
 
     # Correct answer is: ẋ
-    @assert varwhitelist == Bool[0, 1, 0, 0]
+    @test varwhitelist == Bool[0, 1, 0, 0]
 end
 
 begin
@@ -63,5 +63,36 @@ begin
     varwhitelist = computed_highest_diff_variables(structure, ag)
 
     # Correct answer is: ẋ
-    @assert varwhitelist == Bool[0, 1, 0, 0, 0, 0]
+    @test varwhitelist == Bool[0, 1, 0, 0, 0, 0]
+end
+
+begin
+    """
+       Vars: x, y, z
+       Eqs: 0 = f(x,y)
+       Alias: y = -z, ẏ = z
+    """
+    n_vars = 4
+    ag = AliasGraph(n_vars)
+
+    # Alias: z = 1 * ẏ
+    ag[3] = 1 => 2
+    # Alias: z = -1 * y
+    ag[4] = -1 => 2
+
+    # 0 = f(x,y)
+    graph = complete(BipartiteGraph([Int[], Int[], Int[1, 2]], n_vars))
+
+    # [x, ẋ, y, ẏ]
+    var_to_diff = DiffGraph([nothing, 4, nothing, nothing], # primal_to_diff
+                            [nothing, nothing, nothing, 2]) # diff_to_primal
+
+    # [f(x)]
+    eq_to_diff = DiffGraph([nothing, nothing, nothing], # primal_to_diff
+                           [nothing, nothing, 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[1, 0, 0, 1]
 end