Add preface for DiscreteSystem (#1429)

Author: sharanry

src/systems/discrete_system/discrete_system.jl

@@ -56,21 +56,25 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     structure::Any
     """
-    type: type of the system
+    preface: inject assignment statements before the evaluation of the RHS function.
+    """
+    preface::Any
+    """
+    connector_type: type of the system
     """
     connector_type::Any
     """
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
 
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, connector_type, substitutions=nothing; checks::Bool = true)
+    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, preface, connector_type, substitutions=nothing; 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, defaults, structure, connector_type, substitutions)
+        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, preface, connector_type, substitutions)
     end
 end
 
@@ -88,6 +92,7 @@ function DiscreteSystem(
                    default_u0=Dict(),
                    default_p=Dict(),
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
+                   preface=nothing,
                    connector_type=nothing,
                    kwargs...,
                   )
@@ -113,7 +118,7 @@ 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, defaults, nothing, connector_type, kwargs...)
+    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, defaults, nothing, preface, connector_type, kwargs...)
 end
 
 

test/discretesystem.jl

@@ -123,3 +123,60 @@ linearized_eqs = [
 RHS2 = RHS
 @unpack RHS = fol
 @test isequal(RHS, RHS2)
+
+@testset "Preface tests" begin
+    @parameters t
+    using OrdinaryDiffEq
+    using Symbolics
+    using DiffEqBase: isinplace
+    using ModelingToolkit
+    using SymbolicUtils.Code
+    using SymbolicUtils: Sym
+
+    c = [0]
+    f =  function f(c, d::Vector{Float64}, u::Vector{Float64}, p, t::Float64, dt::Float64)
+            c .= [c[1] + 1]
+            d .= randn(length(u))
+            nothing
+        end
+    
+    dummy_identity(x, _) = x
+    @register dummy_identity(x, y)
+
+    u0 = ones(5)
+    p0 = Float64[]
+    syms = [Symbol(:a, i) for i in 1:5]
+    syms_p = Symbol[]
+    dt = 0.1
+    @assert isinplace(f, 6)
+    wf = let c=c, buffer = similar(u0), u=similar(u0), p=similar(p0), dt=dt
+            t -> (f(c, buffer, u, p, t, dt); buffer)
+    end
+
+    num = hash(f) ⊻ length(u0) ⊻ length(p0)
+    buffername = Symbol(:fmi_buffer_, num)
+
+    Δ = DiscreteUpdate(t; dt=dt)
+    us = map(s->(@variables $s(t))[1], syms)
+    ps = map(s->(@variables $s(t))[1], syms_p)
+    buffer, = @variables $buffername[1:length(u0)]
+    dummy_var = Sym{Any}(:_) # this is safe because _ cannot be a rvalue in Julia
+
+    ss = Iterators.flatten((us, ps))
+    vv = Iterators.flatten((u0, p0))
+    defs = Dict{Any, Any}(s=>v for (s, v) in zip(ss, vv))
+
+    preface = [
+               Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:u)), us))
+               Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:p)), ps))
+               Assignment(buffer, term(wf, t))
+              ]
+    eqs = map(1:length(us)) do i
+        Δ(us[i]) ~ dummy_identity(buffer[i], us[i])
+    end
+
+    @named sys = DiscreteSystem(eqs, t, us, ps; defaults=defs, preface=preface)
+    prob = DiscreteProblem(sys, [], (0.0, 1.0))
+    sol = solve(prob, FunctionMap(); dt=dt)
+    @test c[1]+1 == length(sol)
+end

test/odesystem.jl

@@ -535,3 +535,61 @@ eqs = copy(eqs)
 eqs[end] = D(D(z)) ~ α*x - β*y
 @named sys = ODESystem(eqs, t, [x,y,z],[α,β])
 @test_throws Any ODEFunction(sys)
+
+
+@testset "Preface tests" begin
+    using OrdinaryDiffEq
+    using Symbolics
+    using DiffEqBase: isinplace
+    using ModelingToolkit
+    using SymbolicUtils.Code
+    using SymbolicUtils: Sym
+
+    c = [0]
+    function f(c, du::AbstractVector{Float64}, u::AbstractVector{Float64}, p, t::Float64)
+        c .= [c[1]+1]
+        du .= randn(length(u))
+        nothing
+    end
+
+    dummy_identity(x, _) = x
+    @register dummy_identity(x, y)
+
+    u0 = ones(5)
+    p0 = Float64[]
+    syms = [Symbol(:a, i) for i in 1:5]
+    syms_p = Symbol[]
+
+    @assert isinplace(f, 5)
+    wf = let buffer=similar(u0), u=similar(u0), p=similar(p0), c=c
+            t -> (f(c, buffer, u, p, t); buffer)
+        end
+
+    num = hash(f) ⊻ length(u0) ⊻ length(p0)
+    buffername = Symbol(:fmi_buffer_, num)
+
+    D = Differential(t)
+    us = map(s->(@variables $s(t))[1], syms)
+    ps = map(s->(@variables $s(t))[1], syms_p)
+    buffer, = @variables $buffername[1:length(u0)]
+    dummy_var = Sym{Any}(:_) # this is safe because _ cannot be a rvalue in Julia
+
+    ss = Iterators.flatten((us, ps))
+    vv = Iterators.flatten((u0, p0))
+    defs = Dict{Any, Any}(s=>v for (s, v) in zip(ss, vv))
+
+    preface = [
+            Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:u)), us))
+            Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:p)), ps))
+            Assignment(buffer, term(wf, t))
+            ]
+    eqs = map(1:length(us)) do i
+        D(us[i]) ~ dummy_identity(buffer[i], us[i])
+    end
+
+    @named sys = ODESystem(eqs, t, us, ps; defaults=defs, preface=preface)
+    prob = ODEProblem(sys, [], (0.0, 1.0))
+    sol = solve(prob, Euler(); dt=0.1)
+    
+    @test c[1] == length(sol)
+end