Refactor code

Author: sharanry

src/systems/diffeqs/abstractodesystem.jl

@@ -78,19 +78,7 @@ function generate_control_jacobian(sys::AbstractODESystem, dvs = states(sys), ps
     return build_function(jac, dvs, ps, get_iv(sys); kwargs...)
 end
 
-@noinline function throw_invalid_derivative(dervar, eq)
-    msg = "The derivative variable must be isolated to the left-hand " *
-    "side of the equation like `$dervar ~ ...`.\n Got $eq."
-    throw(InvalidSystemException(msg))
-end
-
-function check_derivative_variables(eq, expr=eq.rhs)
-    istree(expr) || return nothing
-    if operation(expr) isa Differential
-        throw_invalid_derivative(expr, eq)
-    end
-    foreach(Base.Fix1(check_derivative_variables, eq), arguments(expr))
-end
+check_derivative_variables(eq) = check_operator_variables(eq, Differential)
 
 function generate_function(
         sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);

src/systems/discrete_system/discrete_system.jl

@@ -124,60 +124,7 @@ end
 isdifference(expr) = istree(expr) && operation(expr) isa Difference
 isdifferenceeq(eq) = isdifference(eq.lhs)
 
-difference_vars(x::Sym) = Set([x])
-difference_vars(exprs::Symbolic) = difference_vars([exprs])
-difference_vars(exprs) = foldl(difference_vars!, exprs; init = Set())
-difference_vars!(difference_vars, eq::Equation) = (difference_vars!(difference_vars, eq.lhs); difference_vars!(difference_vars, eq.rhs); difference_vars)
-function difference_vars!(difference_vars, O)
-    if isa(O, Sym)
-        return push!(difference_vars, O)
-    end
-    !istree(O) && return difference_vars
-
-    operation(O) isa Difference && return push!(difference_vars, O)
-
-    if operation(O) === (getindex) &&
-        first(arguments(O)) isa Symbolic
-
-        return push!(difference_vars, O)
-    end
-
-    symtype(operation(O)) <: FnType && push!(difference_vars, O)
-    for arg in arguments(O)
-        difference_vars!(difference_vars, arg)
-    end
-
-    return difference_vars
-end
-
-function collect_difference_variables(sys::DiscreteSystem)
-    eqs = equations(sys)
-    vars = Set()
-    difference_vars = Set()
-    for eq in eqs
-        difference_vars!(vars, eq)
-        for v in vars
-            isdifference(v) || continue
-            push!(difference_vars, arguments(v)[1])
-        end
-        empty!(vars)
-    end
-    return difference_vars
-end
-
-@noinline function throw_invalid_difference(difvar, eq)
-    msg = "The difference variable must be isolated to the left-hand " *
-    "side of the equation like `$difvar ~ ...`.\n Got $eq."
-    throw(InvalidSystemException(msg))
-end
-
-function check_difference_variables(eq, expr=eq.rhs)
-    istree(expr) || return nothing
-    if operation(expr) isa Difference
-        throw_invalid_difference(expr, eq)
-    end
-    foreach(Base.Fix1(check_difference_variables, eq), arguments(expr))
-end
+check_difference_variables(eq) = check_operator_variables(eq, Difference)
 
 function generate_function(
         sys::DiscreteSystem, dvs = states(sys), ps = parameters(sys);

src/utils.jl

@@ -166,3 +166,22 @@ function collect_defaults!(defs, vars)
     end
     return defs
 end
+
+@noinline function throw_invalid_operator(opvar, eq, op::Type)
+    if op === Difference
+        optext = "difference"
+    elseif op === Differential
+        optext="derivative" 
+    end
+    msg = "The $optext variable must be isolated to the left-hand " *
+    "side of the equation like `$opvar ~ ...`.\n Got $eq."
+    throw(InvalidSystemException(msg))
+end
+
+function check_operator_variables(eq, op::Type, expr=eq.rhs)
+    istree(expr) || return nothing
+    if operation(expr) isa op
+        throw_invalid_operator(expr, eq, op)
+    end
+    foreach(expr -> check_operator_variables(eq, op, expr), arguments(expr))
+end