add nicer discrete sys -> prob constructor (#1268)

Author: anandijain

src/systems/discrete_system/discrete_system.jl

@@ -11,14 +11,16 @@ $(FIELDS)
 ```
 using ModelingToolkit
 
-@parameters σ ρ β
-@variables t x(t) y(t) z(t) next_x(t) next_y(t) next_z(t)
+@parameters σ=28.0 ρ=10.0 β=8/3 δt=0.1
+@variables t x(t)=1.0 y(t)=0.0 z(t)=0.0
+D = Difference(t; dt=δt)
 
-eqs = [next_x ~ σ*(y-x),
-       next_y ~ x*(ρ-z)-y,
-       next_z ~ x*y - β*z]
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
 
-@named de = DiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β])
+@named de = DiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]) # or 
+@named de = DiscreteSystem(eqs)
 ```
 """
 struct DiscreteSystem <: AbstractTimeDependentSystem
@@ -45,26 +47,21 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     systems::Vector{DiscreteSystem}
     """
-    default_u0: The default initial conditions to use when initial conditions
-    are not supplied in `DiscreteSystem`.
+    defaults: The default values to use when initial conditions and/or
+    parameters are not supplied in `DiscreteProblem`.
     """
-    default_u0::Dict
-    """
-    default_p: The default parameters to use when parameters are not supplied
-    in `DiscreteSystem`.
-    """
-    default_p::Dict
+    defaults::Dict
     """
     type: type of the system
     """
     connection_type::Any
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p, connection_type; checks::Bool = true)
+    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, connection_type; checks::Bool = true)
         if checks
             check_variables(dvs, iv)
             check_parameters(ps, iv)
             all_dimensionless([dvs;ps;iv;ctrls]) ||check_units(discreteEqs)
         end
-        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p, connection_type)
+        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, connection_type)
     end
 end
 
@@ -93,7 +90,7 @@ function DiscreteSystem(
     ctrl′ = value.(controls)
 
     if !(isempty(default_u0) && isempty(default_p))
-        Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :ODESystem, force=true)
+        Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :DiscreteSystem, force=true)
     end
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
@@ -106,7 +103,46 @@ function DiscreteSystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, default_u0, default_p, connection_type, kwargs...)
+    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, defaults, connection_type, kwargs...)
+end
+
+
+function DiscreteSystem(eqs, iv=nothing; kwargs...)
+    eqs = collect(eqs)
+    # NOTE: this assumes that the order of algebric equations doesn't matter
+    diffvars = OrderedSet()
+    allstates = OrderedSet()
+    ps = OrderedSet()
+    # reorder equations such that it is in the form of `diffeq, algeeq`
+    diffeq = Equation[]
+    algeeq = Equation[]
+    # initial loop for finding `iv`
+    if iv === nothing
+        for eq in eqs
+            if !(eq.lhs isa Number) # assume eq.lhs is either Differential or Number
+                iv = iv_from_nested_difference(eq.lhs)
+                break
+            end
+        end
+    end
+    iv = value(iv)
+    iv === nothing && throw(ArgumentError("Please pass in independent variables."))
+    for eq in eqs
+        collect_vars_difference!(allstates, ps, eq.lhs, iv)
+        collect_vars_difference!(allstates, ps, eq.rhs, iv)
+        if isdifferenceeq(eq)
+            diffvar, _ = var_from_nested_difference(eq.lhs)
+            isequal(iv, iv_from_nested_difference(eq.lhs)) || throw(ArgumentError("A DiscreteSystem can only have one independent variable."))
+            diffvar in diffvars && throw(ArgumentError("The difference variable $diffvar is not unique in the system of equations."))
+            push!(diffvars, diffvar)
+            push!(diffeq, eq)
+        else
+            push!(algeeq, eq)
+        end
+    end
+    algevars = setdiff(allstates, diffvars)
+    # the orders here are very important!
+    return DiscreteSystem(append!(diffeq, algeeq), iv, vcat(collect(diffvars), collect(algevars)), ps; kwargs...)
 end
 
 """
@@ -123,16 +159,37 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     ps = parameters(sys)
     eqs = equations(sys)
     eqs = linearize_eqs(sys, eqs)
-    # defs = defaults(sys)
-    t = get_iv(sys)
-    u0 = varmap_to_vars(u0map,dvs)
+    defs = defaults(sys)
+    iv = get_iv(sys)
+
+    if parammap isa Dict
+        u0defs = merge(parammap, defs)
+    elseif eltype(parammap) <: Pair
+        u0defs = merge(Dict(parammap), defs)
+    elseif eltype(parammap) <: Number
+        u0defs = merge(Dict(zip(ps, parammap)), defs)
+    else
+        u0defs = defs
+    end
+    if u0map isa Dict
+        pdefs = merge(u0map, defs)
+    elseif eltype(u0map) <: Pair
+        pdefs = merge(Dict(u0map), defs)
+    elseif eltype(u0map) <: Number
+        pdefs = merge(Dict(zip(dvs, u0map)), defs)
+    else
+        pdefs = defs
+    end
+
+    u0 = varmap_to_vars(u0map,dvs; defaults=u0defs)
+    
     rhss = [eq.rhs for eq in eqs]
     u = dvs
-    p = varmap_to_vars(parammap,ps)
+    p = varmap_to_vars(parammap,ps; defaults=pdefs)
 
     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)
+    f(u,p,iv) = f_oop(u,p,iv)
     DiscreteProblem(f,u0,tspan,p;kwargs...)
 end
 

src/utils.jl

@@ -225,6 +225,15 @@ isdiffeq(eq) = isdifferential(eq.lhs)
 isdifference(expr) = istree(expr) && operation(expr) isa Difference
 isdifferenceeq(eq) = isdifference(eq.lhs)
 
+iv_from_nested_difference(x::Term) = operation(x) isa Difference ? iv_from_nested_difference(arguments(x)[1]) : arguments(x)[1]
+iv_from_nested_difference(x::Sym) = x
+iv_from_nested_difference(x) = missing
+
+var_from_nested_difference(x, i=0) = (missing, missing)
+var_from_nested_difference(x::Term,i=0) = operation(x) isa Difference ? var_from_nested_difference(arguments(x)[1], i + 1) : (x, i)
+var_from_nested_difference(x::Sym,i=0) = (x, i)
+
+
 isvariable(x::Num) = isvariable(value(x))
 function isvariable(x)
     x isa Symbolic || return false
@@ -305,6 +314,20 @@ function collect_vars!(states, parameters, expr, iv)
     return nothing
 end
 
+function collect_vars_difference!(states, parameters, expr, iv)
+    if expr isa Sym
+        collect_var!(states, parameters, expr, iv)
+    else
+        for var in vars(expr)
+            if istree(var) && operation(var) isa Difference
+                var, _ = var_from_nested_difference(var)
+            end
+            collect_var!(states, parameters, var, iv)
+        end
+    end
+    return nothing
+end
+
 function collect_var!(states, parameters, var, iv)
     isequal(var, iv) && return nothing
     if isparameter(var) || (istree(var) && isparameter(operation(var)))

test/discretesystem.jl

@@ -12,7 +12,7 @@ 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)
+@variables S(t) I(t) R(t)
 
 infection = rate_to_proportion(β*c*I/(S+I+R),δt)*S
 recovery = rate_to_proportion(γ,δt)*I
@@ -35,6 +35,27 @@ prob_map = DiscreteProblem(sys,u0,tspan,p)
 using OrdinaryDiffEq
 sol_map = solve(prob_map,FunctionMap());
 
+# Using defaults constructor
+@parameters t c=10.0 nsteps=400 δt=0.1 β=0.05 γ=0.25
+Diff = Difference(t; dt=0.1)
+@variables S(t)=990.0 I(t)=10.0 R(t)=0.0
+
+infection2 = rate_to_proportion(β*c*I/(S+I+R),δt)*S
+recovery2 = rate_to_proportion(γ,δt)*I
+
+eqs2 = [D(S) ~ S-infection2,
+       D(I) ~ I+infection2-recovery2,
+       D(R) ~ R+recovery2]
+
+@named sys = DiscreteSystem(eqs2; controls = [β, γ])
+@test ModelingToolkit.defaults(sys) != Dict()
+
+prob_map2 = DiscreteProblem(sys,[],tspan)
+sol_map2 = solve(prob_map,FunctionMap());
+
+@test sol_map.u == sol_map2.u
+@test sol_map.prob.p == sol_map2.prob.p
+
 # Direct Implementation
 
 function sir_map!(u_diff,u,p,t)