fix: undo the hack in generate_initializesystem

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -12,7 +12,9 @@ function generate_initializesystem(sys::ODESystem;
         algebraic_only = false,
         check_units = true, check_defguess = false,
         name = nameof(sys), kwargs...)
-    vars = unique([unknowns(sys); getfield.((observed(sys)), :lhs)])
+    trueobs = unhack_observed(observed(sys))
+    @show trueobs
+    vars = unique([unknowns(sys); getfield.(trueobs, :lhs)])
     vars_set = Set(vars) # for efficient in-lookup
 
     eqs = equations(sys)
@@ -24,7 +26,7 @@ function generate_initializesystem(sys::ODESystem;
     D = Differential(get_iv(sys))
     diffmap = merge(
         Dict(eq.lhs => eq.rhs for eq in eqs_diff),
-        Dict(D(eq.lhs) => D(eq.rhs) for eq in observed(sys))
+        Dict(D(eq.lhs) => D(eq.rhs) for eq in trueobs)
     )
 
     # 1) process dummy derivatives and u0map into initialization system
@@ -166,15 +168,14 @@ function generate_initializesystem(sys::ODESystem;
     )
 
     # 7) use observed equations for guesses of observed variables if not provided
-    obseqs = observed(sys)
-    for eq in obseqs
+    for eq in trueobs
         haskey(defs, eq.lhs) && continue
         any(x -> isequal(default_toterm(x), eq.lhs), keys(defs)) && continue
 
         defs[eq.lhs] = eq.rhs
     end
 
-    eqs_ics = Symbolics.substitute.([eqs_ics; obseqs], (paramsubs,))
+    eqs_ics = Symbolics.substitute.([eqs_ics; trueobs], (paramsubs,))
     vars = [vars; collect(values(paramsubs))]
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
@@ -324,3 +325,37 @@ function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
         return nothing, nothing, nothing, nothing
     end
 end
+
+"""
+Counteracts the CSE/array variable hacks in `symbolics_tearing.jl` so it works with
+initialization.
+"""
+function unhack_observed(eqs::Vector{Equation})
+    subs = Dict()
+    tempvars = Set()
+    rm_idxs = Int[]
+    for (i, eq) in enumerate(eqs)
+        iscall(eq.rhs) || continue
+        if operation(eq.rhs) == StructuralTransformations.change_origin
+            push!(rm_idxs, i)
+            continue
+        end
+        if operation(eq.rhs) == StructuralTransformations.getindex_wrapper
+            var, idxs = arguments(eq.rhs)
+            subs[eq.rhs] = var[idxs...]
+            push!(tempvars, var)
+        end
+    end
+
+    for (i, eq) in enumerate(eqs)
+        if eq.lhs in tempvars
+            subs[eq.lhs] = eq.rhs
+            push!(rm_idxs, i)
+        end
+    end
+
+    eqs = eqs[setdiff(eachindex(eqs), rm_idxs)]
+    return map(eqs) do eq
+        fixpoint_sub(eq.lhs, subs) ~ fixpoint_sub(eq.rhs, subs)
+    end
+end

test/structural_transformation/utils.jl

@@ -53,6 +53,14 @@ end
     @test any(eq -> isequal(eq.lhs, z), observed(sys))
     prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [foo => _tmp_fn])
     @test_nowarn prob.f(prob.u0, prob.p, 0.0)
+
+    isys = ModelingToolkit.generate_initializesystem(sys)
+    @test length(unknowns(isys)) == 5
+    @test length(equations(isys)) == 4
+    @test !any(equations(isys)) do eq
+        iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper,
+            StructuralTransformations.change_origin]
+    end
 end
 
 @testset "scalarized array observed calling same function multiple times" begin
@@ -69,4 +77,12 @@ end
     prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [foo => _tmp_fn2])
     @test_nowarn prob.f(prob.u0, prob.p, 0.0)
     @test val[] == 1
+
+    isys = ModelingToolkit.generate_initializesystem(sys)
+    @test length(unknowns(isys)) == 3
+    @test length(equations(isys)) == 2
+    @test !any(equations(isys)) do eq
+        iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper,
+            StructuralTransformations.change_origin]
+    end
 end