Fix NonlinearSystem

Author: HarrisonGrodin

README.md

@@ -88,7 +88,7 @@ derivatives are zero. We use (unknown) variables for our nonlinear system.
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs)
+ns = NonlinearSystem(eqs, [x,y,z])
 nlsys_func = generate_function(ns)
 ```
 
@@ -272,7 +272,7 @@ a = y - x
 eqs = [0 ~ σ*a,
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
+ns = NonlinearSystem(eqs, [x,y,z])
 nlsys_func = generate_function(ns)
 ```
 

src/equations.jl

@@ -10,32 +10,3 @@ Base.:(==)(a::Equation, b::Equation) = isequal((a.lhs, a.rhs), (b.lhs, b.rhs))
 Base.:~(lhs::Expression, rhs::Expression) = Equation(lhs, rhs)
 Base.:~(lhs::Expression, rhs::Number    ) = Equation(lhs, rhs)
 Base.:~(lhs::Number    , rhs::Expression) = Equation(lhs, rhs)
-
-_is_known(O::Operation) = O.op.known
-_is_unknown(O::Operation) = !O.op.known
-
-function extract_elements(eqs, predicates)
-    result = [Variable[] for p ∈ predicates]
-    vars = foldl(vars!, eqs; init=Set{Variable}())
-
-    for var ∈ vars
-        for (i, p) ∈ enumerate(predicates)
-            p(var) && (push!(result[i], var); break)
-        end
-    end
-
-    return result
-end
-
-vars(exprs) = foldl(vars!, exprs; init = Set{Variable}())
-function vars!(vars, O)
-    isa(O, Operation) || return vars
-    for arg ∈ O.args
-        if isa(arg, Operation)
-            isa(arg.op, Variable) && push!(vars, arg.op)
-            vars!(vars, arg)
-        end
-    end
-
-    return vars
-end

src/systems/nonlinear/nonlinear_system.jl

@@ -9,17 +9,26 @@ function Base.convert(::Type{NLEq}, eq::Equation)
     return NLEq(eq.rhs)
 end
 Base.:(==)(a::NLEq, b::NLEq) = a.rhs == b.rhs
-get_args(eq::NLEq) = Expression[eq.rhs]
 
 struct NonlinearSystem <: AbstractSystem
     eqs::Vector{NLEq}
-    vs::Vector{Variable}
+    vs::Vector{Expression}
     ps::Vector{Variable}
+    function NonlinearSystem(eqs, vs)
+        rhss = [eq.rhs for eq ∈ eqs]
+        ps = reduce(∪, map(_find_params(vs), rhss); init = vnil())
+        new(eqs, vs, collect(ps))
+    end
 end
 
-function NonlinearSystem(eqs)
-    vs, ps = extract_elements(eqs, [_is_unknown, _is_known])
-    NonlinearSystem(eqs, vs, ps)
+vnil() = Set{Variable}()
+_find_params(vs) = Base.Fix2(_find_params, vs)
+function _find_params(O, vs)
+    isa(O, Operation) || return vnil()
+    any(isequal(O), vs) && return vnil()
+    ps = reduce(∪, map(_find_params(vs), O.args); init = vnil())
+    isa(O.op, Variable) && push!(ps, O.op)
+    return ps
 end
 
 
@@ -31,10 +40,12 @@ end
 
 function generate_jacobian(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction)
     jac = calculate_jacobian(sys)
-    return build_function(jac, sys.vs, sys.ps; version = version)
+    return build_function(jac, clean.(sys.vs), sys.ps; version = version)
 end
 
-function generate_function(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction)
+function generate_function(sys::NonlinearSystem, vs, ps; version::FunctionVersion = ArrayFunction)
     rhss = [eq.rhs for eq ∈ sys.eqs]
-    return build_function(rhss, sys.vs, sys.ps; version = version)
+    vs′ = [clean(v) for v ∈ vs]
+    ps′ = [clean(p) for p ∈ ps]
+    return build_function(rhss, vs′, ps′; version = version)
 end

src/utils.jl

@@ -69,3 +69,17 @@ Base.occursin(t::Expression, x::Expression) = isequal(x, t)
 
 clean(x::Variable) = x
 clean(O::Operation) = isa(O.op, Variable) ? O.op : throw(ArgumentError("invalid variable: $(O.op)"))
+
+
+vars(exprs) = foldl(vars!, exprs; init = Set{Variable}())
+function vars!(vars, O)
+    isa(O, Operation) || return vars
+    for arg ∈ O.args
+        if isa(arg, Operation)
+            isa(arg.op, Variable) && push!(vars, arg.op)
+            vars!(vars, arg)
+        end
+    end
+
+    return vars
+end

test/derivatives.jl

@@ -31,11 +31,10 @@ d2 = D(sin(t)*cos(t))
 @test isequal(expand_derivatives(d1), t*cos(t)+sin(t))
 @test isequal(expand_derivatives(d2), simplify_constants(cos(t)*cos(t)+sin(t)*(-1*sin(t))))
 
-@test_broken begin
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-sys = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
+sys = NonlinearSystem(eqs, [x,y,z])
 jac = calculate_jacobian(sys)
 @test isequal(jac[1,1], σ*-1)
 @test isequal(jac[1,2], σ)
@@ -46,7 +45,6 @@ jac = calculate_jacobian(sys)
 @test isequal(jac[3,1], y)
 @test isequal(jac[3,2], x)
 @test isequal(jac[3,3], -1*β)
-end
 
 # Variable dependence checking in differentiation
 @variables a(t) b(a)

test/system_construction.jl

@@ -13,6 +13,12 @@ function test_diffeq_inference(name, de, iv, dvs, ps)
         @test Set([p.name  for p  ∈ de.ps ]) == Set(ps)
     end
 end
+function test_nlsys_inference(name, de, vs, ps)
+    @testset "NonlinearSystem construction: $name" begin
+        @test Set(vs) == Set(vs)
+        @test Set([p.name for p ∈ de.ps]) == Set(ps)
+    end
+end
 
 # Define a differential equation
 eqs = [D(x) ~ σ*(y-x),
@@ -96,21 +102,14 @@ jac = calculate_jacobian(de)
     f = ODEFunction(de)
 end
 
-@test_broken begin
 # Define a nonlinear system
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs,[x,y,z],[t,σ,ρ,β])
-ns2 = NonlinearSystem(eqs)
-for el in (:vs, :ps)
-    names2 = sort(collect(var.name for var in getfield(ns2,el)))
-    names = sort(collect(var.name for var in getfield(ns,el)))
-    @test names2 == names
-end
+ns = NonlinearSystem(eqs, [x,y,z])
+test_nlsys_inference("standard", ns, (x, y, z), (:σ, :ρ, :β))
 
 generate_function(ns, [x,y,z], [σ,ρ,β])
-end
 
 @derivatives D'~t
 @parameters A() B() C()
@@ -125,16 +124,15 @@ de = ODESystem(eqs)
     du ≈ [-1, -1/3]
 end
 
-@test_broken begin
 # Now nonlinear system with only variables
-@variables x y z
+@variables x() y() z()
 @parameters σ() ρ() β()
 
 # Define a nonlinear system
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
+ns = NonlinearSystem(eqs, [x,y,z])
 jac = calculate_jacobian(ns)
 @testset "nlsys jacobian" begin
     @test isequal(jac[1,1], σ * -1)
@@ -156,8 +154,7 @@ f = @eval eval(nlsys_func)
 eqs = [0 ~ σ*a,
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
+ns = NonlinearSystem(eqs, [x,y,z])
 nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])
 jac = calculate_jacobian(ns)
 jac = generate_jacobian(ns)
-end