Assert that ODE derivatives are w.r.t. the independent variable (#1076)

* Check that all the derivatives in a time-dependent system (other than PDE) are derivatives with respect to the independent variable.

* Update src/utils.jl

Nicer way to do this.

Co-authored-by: Yingbo Ma <[email protected]>

* Revert accidental inclusion.

* Set is faster, no need to use OrderedSet.

* Style improvements.

* Fix warning about local assignment.

* Remove checks on differentials for systems that don't support differentials.

Co-authored-by: Yingbo Ma <[email protected]>

Author: lamorton

src/systems/control/controlsystem.jl

@@ -72,10 +72,12 @@ struct ControlSystem <: AbstractControlSystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
-    function ControlSystem(loss, deqs, iv, dvs, controls,ps, observed, name, systems, defaults)
-        check_variables(dvs,iv)
-        check_parameters(ps,iv)
-        new(loss, deqs, iv, dvs, controls,ps, observed, name, systems, defaults)
+    function ControlSystem(loss, deqs, iv, dvs, controls, ps, observed, name, systems, defaults)
+        check_variables(dvs, iv)
+        check_parameters(ps, iv)
+        check_equations(deqs, iv)
+        check_equations(observed, iv)
+        new(loss, deqs, iv, dvs, controls, ps, observed, name, systems, defaults)
     end
 end
 

src/systems/diffeqs/odesystem.jl

@@ -77,6 +77,7 @@ struct ODESystem <: AbstractODESystem
     function ODESystem(deqs, iv, dvs, ps, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
+        check_equations(deqs,iv)
         new(deqs, iv, dvs, ps, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
     end
 end

src/systems/diffeqs/sdesystem.jl

@@ -79,6 +79,7 @@ struct SDESystem <: AbstractODESystem
     function SDESystem(deqs, neqs, iv, dvs, ps, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
+        check_equations(deqs,iv)
         new(deqs, neqs, iv, dvs, ps, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
     end
 end

src/systems/discrete_system/discrete_system.jl

@@ -50,8 +50,8 @@ struct DiscreteSystem <: AbstractSystem
     """
     default_p::Dict
     function DiscreteSystem(discreteEqs, iv, dvs, ps, observed, name, systems, default_u0, default_p)
-        check_variables(dvs,iv)
-        check_parameters(ps,iv)
+        check_variables(dvs, iv)
+        check_parameters(ps, iv)
         new(discreteEqs, iv, dvs, ps, observed, name, systems, default_u0, default_p)
     end
 end

src/systems/jumps/jumpsystem.jl

@@ -52,8 +52,8 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     """
     connection_type::Any
     function JumpSystem{U}(ap::U, iv, states, ps, observed, name, systems, defaults, connection_type) where U <: ArrayPartition
-        check_variables(states,iv)
-        check_parameters(ps,iv)
+        check_variables(states, iv)
+        check_parameters(ps, iv)
         new{U}(ap, iv, states, ps, observed, name, systems, defaults, connection_type)
     end
 end

src/systems/reaction/reactionsystem.jl

@@ -154,8 +154,8 @@ struct ReactionSystem <: AbstractSystem
         iv′ = value(iv)
         states′ = value.(states)
         ps′ = value.(ps)
-        check_variables(states′,iv′)
-        check_parameters(ps′,iv′)
+        check_variables(states′, iv′)
+        check_parameters(ps′, iv′)
         new(eqs, iv′, states′, ps′, observed, name, systems)
     end
 end

src/utils.jl

@@ -105,19 +105,53 @@ function _readable_code(ex)
 end
 readable_code(expr) = JuliaFormatter.format_text(string(Base.remove_linenums!(_readable_code(expr))))
 
-function check_parameters(ps,iv)
+function check_parameters(ps, iv)
     for p in ps
-        isequal(iv,p) && throw(ArgumentError("Independent variable $iv not allowed in parameters."))
+        isequal(iv, p) && throw(ArgumentError("Independent variable $iv not allowed in parameters."))
     end
 end
 
-function check_variables(dvs,iv)
+function check_variables(dvs, iv)
     for dv in dvs
-        isequal(iv,dv) && throw(ArgumentError("Independent variable $iv not allowed in dependent variables."))
+        isequal(iv, dv) && throw(ArgumentError("Independent variable $iv not allowed in dependent variables."))
         isequal(iv, iv_from_nested_derivative(dv)) || throw(ArgumentError("Variable $dv is not a function of independent variable $iv."))
     end
 end 
 
+"Get all the independent variables with respect to which differentials are taken."
+function collect_differentials(eqs)
+    vars = Set()
+    ivs = Set()
+    for eq in eqs
+        vars!(vars, eq)
+        for v in vars
+            isdifferential(v) || continue
+            collect_ivs_from_nested_differential!(ivs, v)  
+        end
+        empty!(vars)
+    end
+    return ivs
+end
+
+"Assert that equations are well-formed when building ODE."
+function check_equations(eqs, iv)
+    ivs = collect_differentials(eqs)
+    display = collect(ivs)
+    length(ivs) <= 1 || throw(ArgumentError("Differential w.r.t. multiple variables $display are not allowed."))
+    if length(ivs) == 1
+        single_iv = pop!(ivs)
+        isequal(single_iv, iv) || throw(ArgumentError("Differential w.r.t. variable ($single_iv) other than the independent variable ($iv) are not allowed."))
+    end
+end
+"Get all the independent variables with respect to which differentials are taken."
+function collect_ivs_from_nested_differential!(ivs, x::Term)
+    op = operation(x)
+    if op isa Differential
+        push!(ivs, op.x)
+        collect_ivs_from_nested_differential!(ivs, arguments(x)[1])
+    end
+end
+
 iv_from_nested_derivative(x::Term) = operation(x) isa Differential ? iv_from_nested_derivative(arguments(x)[1]) : arguments(x)[1]
 iv_from_nested_derivative(x::Sym) = x
 iv_from_nested_derivative(x) = missing

test/controlsystem.jl

@@ -25,7 +25,7 @@ sol = solve(prob,BFGS())
         D(x) ~ - p[2]*x
         D(v) ~ p[1]*u^3
     ]
-    sys1 = ControlSystem(loss,eqs_short,t,[x,v],[u],p,name=:sys1)
-    sys2 = ControlSystem(loss,eqs_short,t,[x,v],[u],p,name=:sys1)
-    @test_throws ArgumentError ControlSystem(loss,[sys2.v ~ sys1.v],t, [],[],[],systems=[sys1, sys2])
+    sys1 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
+    sys2 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
+    @test_throws ArgumentError ControlSystem(loss, [sys2.v ~ sys1.v], t, [], [], [], systems = [sys1, sys2])
 end

test/jumpsystem.jl

@@ -155,7 +155,7 @@ sol = solve(jprob, SSAStepper());
 
 # issue #819
 @testset "Combined system name collisions" begin
-  sys1 = JumpSystem([maj1,maj2], t, [S], [β,γ],name=:sys1)
-  sys2 = JumpSystem([maj1,maj2], t, [S], [β,γ],name=:sys1)
-  @test_throws ArgumentError JumpSystem([sys1.γ ~ sys2.γ], t,[],[], systems=[sys1, sys2])
+  sys1 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
+  sys2 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
+  @test_throws ArgumentError JumpSystem([sys1.γ ~ sys2.γ], t, [], [], systems = [sys1, sys2])
 end

test/nonlinearsystem.jl

@@ -125,13 +125,13 @@ np = NonlinearProblem(ns, [0,0,0], [1,2,3], jac=true, sparse=true)
            @parameters a
            @variables x f
 
-           NonlinearSystem([0 ~ -a*x + f],[x,f],[a], name=name)
+           NonlinearSystem([0 ~ -a * x + f], [x,f], [a], name = name)
        end
 
        function issue819()
            sys1 = makesys(:sys1)
            sys2 = makesys(:sys1)
-           @test_throws ArgumentError NonlinearSystem([sys2.f ~ sys1.x, sys1.f ~ 0],[],[], systems=[sys1, sys2])
+           @test_throws ArgumentError NonlinearSystem([sys2.f ~ sys1.x, sys1.f ~ 0], [], [], systems = [sys1, sys2])
        end
        issue819()
 end