Make alias elimination work with NonlinearSystem

Author: YingboMa

src/systems/nonlinear/nonlinearsystem.jl

@@ -59,6 +59,8 @@ function NonlinearSystem(eqs, states, ps;
     NonlinearSystem(eqs, value.(states), value.(ps), value.(pins), observed, name, systems, default_u0, default_p)
 end
 
+independent_variable(::NonlinearSystem) = nothing
+
 function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)
     rhs = [eq.rhs for eq ∈ equations(sys)]
     vals = [dv for dv in states(sys)]
@@ -77,9 +79,19 @@ function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = paramete
                           conv = AbstractSysToExpr(sys), kwargs...)
 end
 
-function generate_function(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys); kwargs...)
-    rhss = [eq.rhs for eq ∈ sys.eqs]
-    return build_function(rhss, vs, ps;
+function generate_function(sys::NonlinearSystem, dvs = states(sys), ps = parameters(sys); kwargs...)
+    obsvars = map(eq->eq.lhs, observed(sys))
+    fulldvs = [dvs; obsvars]
+    fulldvs′ = makesym.(value.(fulldvs))
+
+    sub = Dict(fulldvs .=> fulldvs′)
+    # substitute x(t) by just x
+    rhss = [substitute(deq.rhs, sub) for deq ∈ equations(sys)]
+    obss = [makesym(value(eq.lhs)) ~ substitute(eq.rhs, sub) for eq ∈ observed(sys)]
+
+    dvs′ = fulldvs′[1:length(dvs)]
+    ps′ = makesym.(value.(ps), states=())
+    return build_function(Let(obss, rhss), dvs′, ps′;
                           conv = AbstractSysToExpr(sys), kwargs...)
 end
 

src/systems/reduction.jl

@@ -84,9 +84,9 @@ function maybe_alias(lhs, rhs, diff_vars, iv, conservative)
     end
 end
 
-function alias_elimination(sys::ODESystem; conservative=true)
+function alias_elimination(sys; conservative=true)
     iv = independent_variable(sys)
-    eqs = vcat(equations(sys), observed(sys))
+    eqs = equations(sys)
     diff_vars = filter(!isnothing, map(eqs) do eq
             if isdiffeq(eq)
                 arguments(eq.lhs)[1]
@@ -131,11 +131,15 @@ function alias_elimination(sys::ODESystem; conservative=true)
     end
 
     alias_vars = first.(subs)
-    sys_states = states(sys)
-    alias_eqs = topsort_equations(alias_vars .~ last.(subs), sys_states)
-
-    newstates = setdiff(sys_states, alias_vars)
-    ODESystem(neweqs, sys.iv, newstates, parameters(sys), observed=alias_eqs)
+    sts = states(sys)
+    fullsts = vcat(map(eq->eq.lhs, observed(sys)), sts)
+    alias_eqs = topsort_equations(alias_vars .~ last.(subs), fullsts)
+    newstates = setdiff(sts, alias_vars)
+
+    @set! sys.eqs = neweqs
+    @set! sys.states = newstates
+    @set! sys.observed = [observed(sys); alias_eqs]
+    return sys
 end
 
 """

test/reduction.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, OrdinaryDiffEq, Test
+using ModelingToolkit, OrdinaryDiffEq, Test, NonlinearSolve
 using ModelingToolkit: topsort_equations
 
 @variables t x(t) y(t) z(t) k(t)
@@ -179,3 +179,30 @@ let
               ]
     @test ref_eqs == equations(reduced_sys)
 end
+
+# NonlinearSystem
+@parameters t
+@variables u1(t) u2(t) u3(t) u4(t) u5(t)
+eqs = [
+       2u1 ~ 3u5
+       u2 ~ u1
+       u3 ~ 2u1 - u2
+       u4 ~ u2 + u3^2
+       u5 ~ u4^2 - u1
+      ]
+sys = NonlinearSystem(eqs, [u1, u2, u3, u4, u5], [])
+reducedsys = alias_elimination(sys)
+@test observed(reducedsys) == [u1 ~ 3/2 * u5]
+
+u0 = [
+      u1 => 1
+      u2 => 1
+      u3 => 0.3
+      u4 => 0.6
+      u5 => 2/3
+     ]
+nlprob = NonlinearProblem(reducedsys, u0)
+reducedsol = solve(nlprob, NewtonRaphson())
+residual = fill(100.0, 4)
+nlprob.f(residual, reducedsol.u, nothing)
+@test all(x->abs(x) < 1e-5, residual)