Add DiscreteFunctionExpr, DiscreteProblemExpr and docs

Signed-off-by: ErikQQY <2283984853@qq.com>

Author: ErikQQY

docs/pages.jl

@@ -28,6 +28,7 @@ pages = [
                           "systems/JumpSystem.md",
                           "systems/NonlinearSystem.md",
                           "systems/OptimizationSystem.md",
+                          "systems/DiscreteSystem.md",
                           "systems/PDESystem.md"],
     "comparison.md",
     "internals.md",

docs/src/systems/DiscreteSystem.md

@@ -0,0 +1,33 @@
+# DiscreteSystem
+
+## System Constructors
+
+```@docs
+DiscreteSystem
+```
+
+## Composition and Accessor Functions
+
+  - `get_eqs(sys)` or `equations(sys)`: The equations that define the Discrete System.
+  - `get_delay_val(sys)`: The delay of the Discrete System.
+  - `get_iv(sys)`: The independent variable of the Discrete System.
+
+## Transformations
+
+## Analyses
+
+## Applicable Calculation and Generation Functions
+
+## Standard Problem Constructors
+
+```@docs
+DiscreteFunction(sys::ModelingToolkit.DiscreteSystem, args...)
+DiscreteProblem(sys::ModelingToolkit.DiscreteSystem, args...)
+```
+
+## Expression Constructors
+
+```@docs
+DiscreteFunctionExpr
+DiscreteProblemExpr
+```

src/ModelingToolkit.jl

@@ -192,7 +192,8 @@ export SymScope, LocalScope, ParentScope, DelayParentScope, GlobalScope
 export independent_variable, equations, controls,
        observed, structure, full_equations
 export structural_simplify, expand_connections, linearize, linearization_function
-export DiscreteSystem, DiscreteProblem, DiscreteFunction
+export DiscreteSystem, DiscreteProblem, DiscreteProblemExpr, DiscreteFunction,
+       DiscreteFunctionExpr
 
 export calculate_jacobian, generate_jacobian, generate_function
 export calculate_control_jacobian, generate_control_jacobian

src/systems/discrete_system/discrete_system.jl

@@ -361,3 +361,96 @@ function SciMLBase.DiscreteFunction{iip, specialize}(sys::DiscreteSystem,
                                       observed = observedfun,
                                       analytic = analytic)
 end
+
+"""
+```julia
+    DiscreteFunctionExpr{iip}(sys::DiscreteSystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing,
+                                     kwargs...) where {iip}
+```
+
+Create a Julia expression for an `DiscreteFunction` from the [`DiscreteSystem`](@ref).
+The arguments `dvs` and `ps` are used to set the order of the dependent
+variable and parameter vectors, respectively.
+"""
+struct DiscreteFunctionExpr{iip} end
+struct DiscreteFunctionClosure{O, I} <: Function
+    f_oop::O
+    f_iip::I
+end
+(f::DiscreteFunctionClosure)(u, p, t) = f.f_oop(u, p, t)
+(f::DiscreteFunctionClosure)(du, u, p, t) = f.f_iip(du, u, p, t)
+
+function DiscreteFunctionExpr{iip}(sys::DiscreteSystem, dvs = states(sys),
+                                   ps = parameters(sys), u0 = nothing;
+                                   version = nothing, p = nothing,
+                                   linenumbers = false,
+                                   simplify = false,
+                                   kwargs...) where {iip}
+    f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
+
+    fsym = gensym(:f)
+    _f = :($fsym = $DiscreteFunctionClosure($f_oop, $f_iip))
+
+    ex = quote
+        $_f
+        DiscreteFunction{$iip}($fsym,
+                               syms = $(Symbol.(states(sys))),
+                               indepsym = $(QuoteNode(Symbol(get_iv(sys)))),
+                               paramsyms = $(Symbol.(parameters(sys))))
+    end
+    !linenumbers ? striplines(ex) : ex
+end
+
+function DiscreteFunctionExpr(sys::DiscreteSystem, args...; kwargs...)
+    DiscreteFunctionExpr{true}(sys, args...; kwargs...)
+end
+
+function process_DiscreteProblem(constructor, sys::DiscreteSystem, u0map, parammap;
+                                 version = nothing,
+                                 linenumbers = true, parallel = SerialForm(),
+                                 eval_expression = true,
+                                 use_union = false,
+                                 kwargs...)
+    eqs = equations(sys)
+    dvs = states(sys)
+    ps = parameters(sys)
+
+    defs = defaults(sys)
+    defs = mergedefaults(defs, parammap, ps)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
+    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = !use_union, use_union)
+
+    check_eqs_u0(eqs, dvs, u0; kwargs...)
+
+    f = constructor(sys, dvs, ps, u0;
+                    linenumbers = linenumbers, parallel = parallel,
+                    syms = Symbol.(dvs), paramsyms = Symbol.(ps),
+                    eval_expression = eval_expression, kwargs...)
+    return f, u0, p
+end
+
+function DiscreteProblemExpr(sys::DiscreteSystem, args...; kwargs...)
+    DiscreteProblemExpr{true}(sys, args...; kwargs...)
+end
+
+function DiscreteProblemExpr{iip}(sys::DiscreteSystem, u0map, tspan,
+                                  parammap = DiffEqBase.NullParameters();
+                                  check_length = true,
+                                  kwargs...) where {iip}
+    f, u0, p = process_DiscreteProblem(DiscreteFunctionExpr{iip}, sys, u0map, parammap;
+                                       check_length, kwargs...)
+    linenumbers = get(kwargs, :linenumbers, true)
+
+    ex = quote
+        f = $f
+        u0 = $u0
+        p = $p
+        tspan = $tspan
+        DiscreteProblem(f, u0, tspan, p; $(filter_kwargs(kwargs)...))
+    end
+    !linenumbers ? striplines(ex) : ex
+end

src/systems/jumps/jumpsystem.jl

@@ -340,10 +340,12 @@ tspan = (0.0, 250.0)
 dprob = DiscreteProblem(js, u₀map, tspan, parammap)
 ```
 """
-function DiscreteProblemExpr(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
-                             parammap = DiffEqBase.NullParameters();
-                             use_union = false,
-                             kwargs...)
+struct DiscreteProblemExpr{iip} end
+
+function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
+                                  parammap = DiffEqBase.NullParameters();
+                                  use_union = false,
+                                  kwargs...) where {iip}
     dvs = states(sys)
     ps = parameters(sys)
     defs = defaults(sys)

test/discretesystem.jl

@@ -28,17 +28,22 @@ eqs = [D(S) ~ S - infection * h,
 @named sys = DiscreteSystem(eqs, t, [S, I, R], [c, nsteps, δt, β, γ]; controls = [β, γ])
 syss = structural_simplify(sys)
 @test syss == syss
-df = DiscreteFunction(sys)
 
-# iip
-du = zeros(3)
-u = collect(1:3)
-p = collect(1:5)
-df.f(du, u, p, 0)
-@test du ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
+for df in [
+    DiscreteFunction(sys, [S, I, R], [c, nsteps, δt, β, γ]),
+    eval(DiscreteFunctionExpr(sys, [S, I, R], [c, nsteps, δt, β, γ])),
+]
+
+    # iip
+    du = zeros(3)
+    u = collect(1:3)
+    p = collect(1:5)
+    df.f(du, u, p, 0)
+    @test du ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
 
-# oop
-@test df.f(u, p, 0) ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
+    # oop
+    @test df.f(u, p, 0) ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
+end
 
 # Problem
 u0 = [S => 990.0, I => 10.0, R => 0.0]