Merge pull request #1119 from sharanry/sy/use_difference_for_discreteSys

Use Difference operator for discrete system

Author: YingboMa

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

@@ -116,8 +116,29 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     u = dvs
     p = varmap_to_vars(parammap,ps)
 
-    f_gen = build_function(rhss, dvs, ps, t; expression=Val{eval_expression}, expression_module=eval_module)
-    f_oop,f_iip = (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen)
+    f_gen = generate_function(sys; expression=Val{eval_expression}, expression_module=eval_module)
+    f_oop, _ = (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen)
     f(u,p,t) = f_oop(u,p,t)
     DiscreteProblem(f,u0,tspan,p;kwargs...)
 end
+
+isdifference(expr) = istree(expr) && operation(expr) isa Difference
+isdifferenceeq(eq) = isdifference(eq.lhs)
+
+check_difference_variables(eq) = check_operator_variables(eq, Difference)
+
+function generate_function(
+        sys::DiscreteSystem, dvs = states(sys), ps = parameters(sys);
+        kwargs...
+    )
+    eqs = equations(sys)
+    foreach(check_difference_variables, eqs)
+    # substitute x(t) by just x
+    rhss = [eq.rhs for eq in eqs]
+
+    u = map(x->time_varying_as_func(value(x), sys), dvs)
+    p = map(x->time_varying_as_func(value(x), sys), ps)
+    t = get_iv(sys)
+    
+    build_function(rhss, u, p, t; kwargs...)
+end

src/utils.jl

@@ -166,3 +166,24 @@ function collect_defaults!(defs, vars)
     end
     return defs
 end
+
+"Throw error when difference/derivative operation occurs in the R.H.S."
+@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
+
+"Check if difference/derivative operation occurs in the R.H.S. of an equation"
+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

test/discretesystem.jl

@@ -3,23 +3,24 @@
 - https://github.com/epirecipes/sir-julia/blob/master/markdown/function_map/function_map.md
 - https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#Deterministic_versus_stochastic_epidemic_models
 =#
-using ModelingToolkit
+using ModelingToolkit, Test
 
 @inline function rate_to_proportion(r,t)
     1-exp(-r*t)
 end;
 
 # Independent and dependent variables and parameters
 @parameters t c nsteps δt β γ
+D = Difference(t; dt=0.1)
 @variables S(t) I(t) R(t) next_S(t) next_I(t) next_R(t)
 
 infection = rate_to_proportion(β*c*I/(S+I+R),δt)*S
 recovery = rate_to_proportion(γ,δt)*I
 
 # Equations
-eqs = [next_S ~ S-infection,
-       next_I ~ I+infection-recovery,
-       next_R ~ R+recovery]
+eqs = [D(S) ~ S-infection,
+       D(I) ~ I+infection-recovery,
+       D(R) ~ R+recovery]
 
 # System
 sys = DiscreteSystem(eqs,t,[S,I,R],[c,nsteps,δt,β,γ]; controls = [β, γ])
@@ -36,16 +37,16 @@ sol_map = solve(prob_map,FunctionMap());
 
 # Direct Implementation
 
-function sir_map!(du,u,p,t)
+function sir_map!(u_diff,u,p,t)
     (S,I,R) = u
     (β,c,γ,δt) = p
     N = S+I+R
     infection = rate_to_proportion(β*c*I/N,δt)*S
     recovery = rate_to_proportion(γ,δt)*I
     @inbounds begin
-        du[1] = S-infection
-        du[2] = I+infection-recovery
-        du[3] = R+recovery
+        u_diff[1] = S-infection
+        u_diff[2] = I+infection-recovery
+        u_diff[3] = R+recovery
     end
     nothing
 end;