DiscreteSystem (#986)

* DiscreteSystem added

* Update test/test_epidemiological_compartmental.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

* Update test/test_epidemiological_compartmental.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

* Update test/test_epidemiological_compartmental.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

* Update test/test_epidemiological_compartmental.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

* DiscreteSystem doesn't need two files

* Rename test_epidemiological_compartmental.jl to discretesystem.jl

* Update runtests.jl

* Update discretesystem.jl

* Update discrete_system.jl

* Update discretesystem.jl

* Update discretesystem.jl

* Update discretesystem.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

Author: ovile27

src/ModelingToolkit.jl

@@ -130,6 +130,8 @@ include("systems/pde/pdesystem.jl")
 include("systems/reaction/reactionsystem.jl")
 include("systems/dependency_graphs.jl")
 
+include("systems/discrete_system/discrete_system.jl")
+
 include("systems/systemstructure.jl")
 using .SystemStructures
 
@@ -162,6 +164,7 @@ export Term, Sym
 export SymScope, LocalScope, ParentScope, GlobalScope
 export independent_variable, states, parameters, equations, controls, observed, structure
 export structural_simplify
+export DiscreteSystem, DiscreteProblem
 
 export calculate_jacobian, generate_jacobian, generate_function
 export calculate_tgrad, generate_tgrad

src/systems/discrete_system/discrete_system.jl

@@ -0,0 +1,107 @@
+"""
+$(TYPEDEF)
+
+A system of difference equations.
+
+# Fields
+$(FIELDS)
+
+# Example
+
+```
+using ModelingToolkit
+
+@parameters t σ ρ β
+@variables x(t) y(t) z(t) next_x(t) next_y(t) next_z(t)
+
+eqs = [next_x ~ σ*(y-x),
+       next_y ~ x*(ρ-z)-y,
+       next_z ~ x*y - β*z]
+
+de = DiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β])
+```
+"""
+struct DiscreteSystem <: AbstractSystem
+    """The differential equations defining the discrete system."""
+    eqs::Vector{Equation}
+    """Independent variable."""
+    iv::Sym
+    """Dependent (state) variables."""
+    states::Vector
+    """Parameter variables."""
+    ps::Vector
+    observed::Vector{Equation}
+    """
+    Name: the name of the system
+    """
+    name::Symbol
+    """
+    systems: The internal systems. These are required to have unique names.
+    """
+    systems::Vector{DiscreteSystem}
+    """
+    default_u0: The default initial conditions to use when initial conditions
+    are not supplied in `DiscreteSystem`.
+    """
+    default_u0::Dict
+    """
+    default_p: The default parameters to use when parameters are not supplied
+    in `DiscreteSystem`.
+    """
+    default_p::Dict
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Constructs a DiscreteSystem.
+"""
+function DiscreteSystem(
+                   discreteEqs::AbstractVector{<:Equation}, iv, dvs, ps;
+                   observed = Num[],
+                   systems = DiscreteSystem[],
+                   name=gensym(:DiscreteSystem),
+                   default_u0=Dict(),
+                   default_p=Dict(),
+                  )
+    iv′ = value(iv)
+    dvs′ = value.(dvs)
+    ps′ = value.(ps)
+
+    default_u0 isa Dict || (default_u0 = Dict(default_u0))
+    default_p isa Dict || (default_p = Dict(default_p))
+    default_u0 = Dict(value(k) => value(default_u0[k]) for k in keys(default_u0))
+    default_p = Dict(value(k) => value(default_p[k]) for k in keys(default_p))
+
+    sysnames = nameof.(systems)
+    if length(unique(sysnames)) != length(sysnames)
+        throw(ArgumentError("System names must be unique."))
+    end
+    DiscreteSystem(discreteEqs, iv′, dvs′, ps′, observed, name, systems, default_u0, default_p)
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Generates an DiscreteProblem from an DiscreteSystem.
+"""
+function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
+                                    parammap=DiffEqBase.NullParameters();
+                                    eval_module = @__MODULE__,
+                                    eval_expression = true,
+                                    kwargs...)
+    dvs = states(sys)
+    ps = parameters(sys)
+    eqs = equations(sys)
+    # defs = defaults(sys)
+    t = get_iv(sys)
+    u0 = varmap_to_vars(u0map,dvs)
+    rhss = [eq.rhs for eq in eqs]
+    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(u,p,t) = f_oop(u,p,t)
+    DiscreteProblem(f,u0,tspan,p;kwargs...)
+end

test/discretesystem.jl

@@ -0,0 +1,57 @@
+# Example: Compartmental models in epidemiology
+#=
+- 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
+
+@inline function rate_to_proportion(r,t)
+    1-exp(-r*t)
+end;
+
+# Independent and dependent variables and parameters
+@parameters t c nsteps δt β γ
+@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]
+
+# System
+sys = DiscreteSystem(eqs,t,[S,I,R],[c,nsteps,δt,β,γ])
+
+# Problem
+u0 = [S => 990.0, I => 10.0, R => 0.0]
+p = [β => 0.05, c => 10.0, γ => 0.25, δt => 0.1, nsteps => 400]
+tspan = (0.0,ModelingToolkit.value(substitute(nsteps,p))) # value function (from Symbolics) is used to convert a Num to Float64
+prob_map = DiscreteProblem(sys,u0,tspan,p)
+
+# Solution
+using OrdinaryDiffEq
+sol_map = solve(prob_map,FunctionMap());
+
+# Direct Implementation
+
+function sir_map!(du,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
+    end
+    nothing
+end;
+u0 = [990.0,10.0,0.0];
+p = [0.05,10.0,0.25,0.1];
+prob_map = DiscreteProblem(sir_map!,u0,tspan,p);
+sol_map2 = solve(prob_map,FunctionMap());
+
+@test Array(sol_map) ≈ Array(sol_map2)

test/runtests.jl

@@ -6,6 +6,7 @@ using SafeTestsets, Test
 @safetestset "Simplify Test" begin include("simplify.jl") end
 @safetestset "Direct Usage Test" begin include("direct.jl") end
 @safetestset "System Linearity Test" begin include("linearity.jl") end
+@safetestset "DiscreteSystem Test" begin include("discretesystem.jl") end
 @safetestset "ODESystem Test" begin include("odesystem.jl") end
 @safetestset "LabelledArrays Test" begin include("labelledarrays.jl") end
 @safetestset "Mass Matrix Test" begin include("mass_matrix.jl") end