fix: remove CSE hack

Author: AayushSabharwal

src/structural_transformation/symbolics_tearing.jl

@@ -708,7 +708,7 @@ Update the system equations, unknowns, and observables after simplification.
 """
 function update_simplified_system!(
         state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns;
-        cse_hack = true, array_hack = true)
+        array_hack = true)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
     diff_to_var = invview(var_to_diff)
 
@@ -732,8 +732,8 @@ function update_simplified_system!(
     unknowns = [unknowns; extra_unknowns]
     @set! sys.unknowns = unknowns
 
-    obs, subeqs, deps = cse_and_array_hacks(
-        sys, obs, solved_eqs, unknowns, neweqs; cse = cse_hack, array = array_hack)
+    obs, subeqs, deps = array_var_hack(
+        sys, obs, solved_eqs, unknowns, neweqs; array = array_hack)
 
     @set! sys.eqs = neweqs
     @set! sys.observed = obs
@@ -775,7 +775,7 @@ appear in the system. Algebraic variables are variables that are not
 differential variables.
 """
 function tearing_reassemble(state::TearingState, var_eq_matching,
-        full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
+        full_var_eq_matching = nothing; simplify = false, mm = nothing, array_hack = true)
     extra_vars = Int[]
     if full_var_eq_matching !== nothing
         for v in 𝑑vertices(state.structure.graph)
@@ -811,21 +811,14 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
 
     sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching,
-        extra_unknowns; cse_hack, array_hack)
+        extra_unknowns; array_hack)
 
     @set! state.sys = sys
     @set! sys.tearing_state = state
     return invalidate_cache!(sys)
 end
 
 """
-# HACK 1
-
-Since we don't support array equations, any equation of the sort `x[1:n] ~ f(...)[1:n]`
-gets turned into `x[1] ~ f(...)[1], x[2] ~ f(...)[2]`. Repeatedly calling `f` gets
-_very_ expensive. this hack performs a limited form of CSE specifically for this case to
-avoid the unnecessary cost. This and the below hack are implemented simultaneously
-
 # HACK 2
 
 Add equations for array observed variables. If `p[i] ~ (...)` are equations, add an
@@ -834,12 +827,7 @@ if all `p[i]` are present and the unscalarized form is used in any equation (obs
 not) we first count the number of times the scalarized form of each observed variable
 occurs in observed equations (and unknowns if it's split).
 """
-function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, array = true)
-    # HACK 1
-    # mapping of rhs to temporary CSE variable
-    # `f(...) => tmpvar` in above example
-    rhs_to_tempvar = Dict()
-
+function array_var_hack(sys, obs, subeqs, unknowns, neweqs; array = true)
     # HACK 2
     # map of array observed variable (unscalarized) to number of its
     # scalarized terms that appear in observed equations
@@ -851,36 +839,6 @@ function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, arr
         rhs = eq.rhs
         vars!(all_vars, rhs)
 
-        # 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)
-                tempvar = unwrap(Symbolics.variable(
-                    tempvar; T = Symbolics.symtype(rhs_arr)))
-                tempvar = setmetadata(
-                    tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr))
-                tempeq = tempvar ~ rhs_arr
-                rhs_to_tempvar[rhs_arr] = tempvar
-                push!(obs, tempeq)
-                push!(subeqs, tempeq)
-            end
-
-            # getindex_wrapper is used because `observed2graph` treats `x` and `x[i]` as different,
-            # so it doesn't find a dependency between this equation and `tempvar ~ rhs_arr`
-            # which fails the topological sort
-            neweq = lhs ~ getindex_wrapper(
-                rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end]))
-            obs[i] = neweq
-            subeqi = findfirst(isequal(eq), subeqs)
-            if subeqi !== nothing
-                subeqs[subeqi] = neweq
-            end
-        end
-        # end HACK 1
-
         array || continue
         iscall(lhs) || continue
         operation(lhs) === getindex || continue
@@ -891,33 +849,6 @@ function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, arr
         continue
     end
 
-    # Also do CSE for `equations(sys)`
-    if cse
-        for (i, eq) in enumerate(neweqs)
-            (; lhs, rhs) = eq
-            is_getindexed_array(rhs) || continue
-            rhs_arr = arguments(rhs)[1]
-            if !haskey(rhs_to_tempvar, rhs_arr)
-                tempvar = gensym(Symbol(lhs))
-                N = length(rhs_arr)
-                tempvar = unwrap(Symbolics.variable(
-                    tempvar; T = Symbolics.symtype(rhs_arr)))
-                tempvar = setmetadata(
-                    tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr))
-                vars!(all_vars, rhs_arr)
-                tempeq = tempvar ~ rhs_arr
-                rhs_to_tempvar[rhs_arr] = tempvar
-                push!(obs, tempeq)
-                push!(subeqs, tempeq)
-            end
-            # don't need getindex_wrapper, but do it anyway to know that this
-            # hack took place
-            neweq = lhs ~ getindex_wrapper(
-                rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end]))
-            neweqs[i] = neweq
-        end
-    end
-
     # count variables in unknowns if they are scalarized forms of variables
     # also present as observed. e.g. if `x[1]` is an unknown and `x[2] ~ (..)`
     # is an observed equation.
@@ -1007,10 +938,10 @@ new residual equations after tearing. End users are encouraged to call [`structu
 instead, which calls this function internally.
 """
 function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing,
-        simplify = false, cse_hack = true, array_hack = true, kwargs...)
+        simplify = false, array_hack = true, kwargs...)
     var_eq_matching, full_var_eq_matching = tearing(state)
     invalidate_cache!(tearing_reassemble(
-        state, var_eq_matching, full_var_eq_matching; mm, simplify, cse_hack, array_hack))
+        state, var_eq_matching, full_var_eq_matching; mm, simplify, array_hack))
 end
 
 """
@@ -1032,7 +963,7 @@ Perform index reduction and use the dummy derivative technique to ensure that
 the system is balanced.
 """
 function dummy_derivative(sys, state = TearingState(sys); simplify = false,
-        mm = nothing, cse_hack = true, array_hack = true, kwargs...)
+        mm = nothing, array_hack = true, kwargs...)
     jac = let state = state
         (eqs, vars) -> begin
             symeqs = EquationsView(state)[eqs]
@@ -1056,5 +987,5 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false,
     end
     var_eq_matching = dummy_derivative_graph!(state, jac; state_priority,
         kwargs...)
-    tearing_reassemble(state, var_eq_matching; simplify, mm, cse_hack, array_hack)
+    tearing_reassemble(state, var_eq_matching; simplify, mm, array_hack)
 end

test/structural_transformation/utils.jl

@@ -121,7 +121,7 @@ end
         @named sys = ODESystem(
             [D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t)
 
-        sys1 = structural_simplify(sys; cse_hack = false)
+        sys1 = structural_simplify(sys)
         @test length(observed(sys1)) == 6
         @test !any(observed(sys1)) do eq
             iscall(eq.rhs) &&
@@ -142,7 +142,7 @@ end
         @named sys = ODESystem(
             [D(x) ~ z[1] + z[2] + foo(z)[1] + w, y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t)
 
-        sys1 = structural_simplify(sys; cse_hack = false, fully_determined = false)
+        sys1 = structural_simplify(sys; fully_determined = false)
         @test length(observed(sys1)) == 6
         @test !any(observed(sys1)) do eq
             iscall(eq.rhs) &&