Merge pull request #1126 from sharanry/sy/difference_ODESys

Enable usage of difference operator in ODESystem

Author: YingboMa

Project.toml

@@ -9,6 +9,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqJump = "c894b116-72e5-5b58-be3c-e6d8d4ac2b12"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -48,6 +49,7 @@ ArrayInterface = "2.8, 3.0"
 ConstructionBase = "1"
 DataStructures = "0.17, 0.18"
 DiffEqBase = "6.54.0"
+DiffEqCallbacks = "2.16"
 DiffEqJump = "6.7.5"
 DiffRules = "0.1, 1.0"
 Distributions = "0.23, 0.24, 0.25"

src/ModelingToolkit.jl

@@ -17,6 +17,7 @@ using DataStructures
 using SpecialFunctions, NaNMath
 using RuntimeGeneratedFunctions
 using Base.Threads
+using DiffEqCallbacks
 import MacroTools: splitdef, combinedef, postwalk, striplines
 import Libdl
 using DocStringExtensions
@@ -182,6 +183,7 @@ export calculate_hessian, generate_hessian
 export calculate_massmatrix, generate_diffusion_function
 export stochastic_integral_transform
 export initialize_system_structure
+export generate_difference_cb
 
 export BipartiteGraph, equation_dependencies, variable_dependencies
 export eqeq_dependencies, varvar_dependencies

src/systems/diffeqs/abstractodesystem.jl

@@ -90,7 +90,7 @@ function generate_function(
     #obsvars = map(eq->eq.lhs, observed(sys))
     #fulldvs = [dvs; obsvars]
 
-    eqs = equations(sys)
+    eqs = [eq for eq in equations(sys) if !isdifferenceeq(eq)]
     foreach(check_derivative_variables, eqs)
     # substitute x(t) by just x
     rhss = implicit_dae ? [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs] :
@@ -109,6 +109,46 @@ function generate_function(
     end
 end
 
+@inline function allequal(x)
+    length(x) < 2 && return true
+    e1 = first(x)
+    i = 2
+    @inbounds for i=2:length(x)
+        x[i] == e1 || return false
+    end
+    return true
+end
+
+function generate_difference_cb(sys::ODESystem, dvs = states(sys), ps = parameters(sys);
+    kwargs...)
+    eqs = equations(sys)
+    foreach(check_difference_variables, eqs)
+
+    rhss = [ 
+        begin
+            ind = findfirst(eq -> isdifference(eq.lhs) && isequal(arguments(eq.lhs)[1], s), eqs)
+            ind === nothing ? 0 : eqs[ind].rhs
+        end
+        for s in dvs ]
+    
+    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)
+
+    f_oop, f_iip = build_function(rhss, u, p, t; kwargs...)
+
+    f = @RuntimeGeneratedFunction(@__MODULE__, f_oop)
+
+    function cb_affect!(int) 
+        int.u += f(int.u, int.p, int.t) 
+    end
+
+    dts = [ operation(eq.lhs).dt for eq in eqs if isdifferenceeq(eq)] 
+    allequal(dts) || error("All difference variables should have same time steps.")
+
+    PeriodicCallback(cb_affect!, first(dts))
+end
+
 function time_varying_as_func(x, sys)
     # if something is not x(t) (the current state)
     # but is `x(t-1)` or something like that, pass in `x` as a callable function rather
@@ -124,7 +164,7 @@ function time_varying_as_func(x, sys)
 end
 
 function calculate_massmatrix(sys::AbstractODESystem; simplify=false)
-    eqs = equations(sys)
+    eqs = [eq for eq in equations(sys) if !isdifferenceeq(eq)]
     dvs = states(sys)
     M = zeros(length(eqs),length(eqs))
     state2idx = Dict(s => i for (i, s) in enumerate(dvs))
@@ -536,7 +576,12 @@ symbolically calculating numerical enhancements.
 function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     parammap=DiffEqBase.NullParameters();kwargs...) where iip
     f, u0, p = process_DEProblem(ODEFunction{iip}, sys, u0map, parammap; kwargs...)
-    ODEProblem{iip}(f,u0,tspan,p;kwargs...)
+    if any(isdifferenceeq.(equations(sys)))
+        ODEProblem{iip}(f,u0,tspan,p;difference_cb=generate_difference_cb(sys),kwargs...)
+    else
+        ODEProblem{iip}(f,u0,tspan,p;kwargs...)
+    end
+    
 end
 
 """
@@ -563,7 +608,12 @@ function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem,du0map,u0map,tspan,
     diffvars = collect_differential_variables(sys)
     sts = states(sys)
     differential_vars = map(Base.Fix2(in, diffvars), sts)
-    DAEProblem{iip}(f,du0,u0,tspan,p;differential_vars=differential_vars,kwargs...)
+    if any(isdifferenceeq.(equations(sys)))
+        DAEProblem{iip}(f,du0,u0,tspan,p;difference_cb=generate_difference_cb(sys),differential_vars=differential_vars,kwargs...)
+    else
+        DAEProblem{iip}(f,du0,u0,tspan,p;differential_vars=differential_vars,kwargs...)    
+    end
+    
 end
 
 """
@@ -713,6 +763,3 @@ end
 function SteadyStateProblemExpr(sys::AbstractODESystem, args...; kwargs...)
     SteadyStateProblemExpr{true}(sys, args...; kwargs...)
 end
-
-isdifferential(expr) = istree(expr) && operation(expr) isa Differential
-isdiffeq(eq) = isdifferential(eq.lhs)

src/systems/diffeqs/odesystem.jl

@@ -134,42 +134,6 @@ function ODESystem(
     ODESystem(deqs, iv′, dvs′, ps′, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
 end
 
-vars(x::Sym) = Set([x])
-vars(exprs::Symbolic) = vars([exprs])
-vars(exprs) = foldl(vars!, exprs; init = Set())
-vars!(vars, eq::Equation) = (vars!(vars, eq.lhs); vars!(vars, eq.rhs); vars)
-function vars!(vars, O)
-    if isa(O, Sym)
-        return push!(vars, O)
-    end
-    !istree(O) && return vars
-
-    operation(O) isa Differential && return push!(vars, O)
-
-    if operation(O) === (getindex) &&
-        first(arguments(O)) isa Symbolic
-
-        return push!(vars, O)
-    end
-
-    symtype(operation(O)) <: FnType && push!(vars, O)
-    for arg in arguments(O)
-        vars!(vars, arg)
-    end
-
-    return vars
-end
-
-find_derivatives!(vars, expr::Equation, f=identity) = (find_derivatives!(vars, expr.lhs, f); find_derivatives!(vars, expr.rhs, f); vars)
-function find_derivatives!(vars, expr, f)
-    !istree(O) && return vars
-    operation(O) isa Differential && push!(vars, f(O))
-    for arg in arguments(O)
-        vars!(vars, arg)
-    end
-    return vars
-end
-
 function ODESystem(eqs, iv=nothing; kwargs...)
     eqs = collect(eqs)
     # NOTE: this assumes that the order of algebric equations doesn't matter
@@ -208,30 +172,6 @@ function ODESystem(eqs, iv=nothing; kwargs...)
     return ODESystem(append!(diffeq, algeeq), iv, vcat(collect(diffvars), collect(algevars)), ps; kwargs...)
 end
 
-function collect_vars!(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 Differential
-                var, _ = var_from_nested_derivative(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)))
-        push!(parameters, var)
-    else
-        push!(states, var)
-    end
-    return nothing
-end
-
 # NOTE: equality does not check cached Jacobian
 function Base.:(==)(sys1::ODESystem, sys2::ODESystem)
     iv1 = independent_variable(sys1)
@@ -319,21 +259,6 @@ function _eq_unordered(a, b)
     return true
 end
 
-function collect_differential_variables(sys::ODESystem)
-    eqs = equations(sys)
-    vars = Set()
-    diffvars = Set()
-    for eq in eqs
-        vars!(vars, eq)
-        for v in vars
-            isdifferential(v) || continue
-            push!(diffvars, arguments(v)[1])
-        end
-        empty!(vars)
-    end
-    return diffvars
-end
-
 # We have a stand-alone function to convert a `NonlinearSystem` or `ODESystem`
 # to an `ODESystem` to connect systems, and we later can reply on
 # `structural_simplify` to convert `ODESystem`s to `NonlinearSystem`s.

src/systems/discrete_system/discrete_system.jl

@@ -122,9 +122,6 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     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(

src/utils.jl

@@ -187,3 +187,93 @@ function check_operator_variables(eq, op::Type, expr=eq.rhs)
     end
     foreach(expr -> check_operator_variables(eq, op, expr), arguments(expr))
 end
+
+isdifferential(expr) = istree(expr) && operation(expr) isa Differential
+isdiffeq(eq) = isdifferential(eq.lhs)
+
+isdifference(expr) = istree(expr) && operation(expr) isa Difference
+isdifferenceeq(eq) = isdifference(eq.lhs)
+
+vars(x::Sym; op=Differential) = Set([x])
+vars(exprs::Symbolic; op=Differential) = vars([exprs]; op=op)
+vars(exprs; op=Differential) = foldl((x, y) -> vars!(x, y; op=op), exprs; init = Set())
+vars!(vars, eq::Equation; op=Differential) = (vars!(vars, eq.lhs; op=op); vars!(vars, eq.rhs; op=op); vars)
+function vars!(vars, O; op=Differential)
+    if isa(O, Sym)
+        return push!(vars, O)
+    end
+    !istree(O) && return vars
+
+    operation(O) isa op && return push!(vars, O)
+
+    if operation(O) === (getindex) &&
+        first(arguments(O)) isa Symbolic
+
+        return push!(vars, O)
+    end
+
+    symtype(operation(O)) <: FnType && push!(vars, O)
+    for arg in arguments(O)
+        vars!(vars, arg; op=op)
+    end
+
+    return vars
+end
+difference_vars(x::Sym) = vars(x; op=Difference)
+difference_vars(exprs::Symbolic) = vars(exprs; op=Difference)
+difference_vars(exprs) = vars(exprs; op=Difference)
+difference_vars!(vars, eq::Equation) = vars!(vars, eq; op=Difference)
+difference_vars!(vars, O) = vars!(vars, O; op=Difference)
+
+function collect_operator_variables(sys, isop::Function)
+    eqs = equations(sys)
+    vars = Set()
+    diffvars = Set()
+    for eq in eqs
+        vars!(vars, eq)
+        for v in vars
+            isop(v) || continue
+            push!(diffvars, arguments(v)[1])
+        end
+        empty!(vars)
+    end
+    return diffvars
+end
+collect_differential_variables(sys) = collect_operator_variables(sys, isdifferential)
+collect_difference_variables(sys) = collect_operator_variables(sys, isdifference)
+
+# 
+
+find_derivatives!(vars, expr::Equation, f=identity) = (find_derivatives!(vars, expr.lhs, f); find_derivatives!(vars, expr.rhs, f); vars)
+function find_derivatives!(vars, expr, f)
+    !istree(O) && return vars
+    operation(O) isa Differential && push!(vars, f(O))
+    for arg in arguments(O)
+        vars!(vars, arg)
+    end
+    return vars
+end
+
+function collect_vars!(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 Differential
+                var, _ = var_from_nested_derivative(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)))
+        push!(parameters, var)
+    else
+        push!(states, var)
+    end
+    return nothing
+end

test/odesystem.jl

@@ -387,3 +387,45 @@ let
     sys = ODESystem(D.(x) .~ x)
     @test_nowarn structural_simplify(sys)
 end
+
+# Mixed Difference Differential equations
+@parameters t a b c d
+@variables x(t) y(t)
+δ = Differential(t)
+D = Difference(t; dt=0.1)
+eqs = [
+    δ(x) ~ a*x - b*x*y,
+    δ(y) ~ -c*y + d*x*y,
+    D(x) ~ y
+]
+de = ODESystem(eqs,t,[x,y],[a,b,c,d])
+@test generate_difference_cb(de) isa ModelingToolkit.DiffEqCallbacks.DiscreteCallback
+
+# doesn't work with ODEFunction
+# prob = ODEProblem(ODEFunction{false}(de),[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
+
+prob = ODEProblem(de,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0], check_length=false)
+@test prob.kwargs[:difference_cb] isa ModelingToolkit.DiffEqCallbacks.DiscreteCallback
+
+sol = solve(prob, Tsit5(); callback=prob.kwargs[:difference_cb], tstops=prob.tspan[1]:0.1:prob.tspan[2])
+
+# Direct implementation
+function lotka(du,u,p,t)
+    x = u[1]
+    y = u[2]
+    du[1] = p[1]*x - p[2]*x*y
+    du[2] = -p[3]*y + p[4]*x*y
+end
+
+prob2 = ODEProblem(lotka,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
+function periodic_difference_affect!(int)
+    int.u += [int.u[2], 0]
+end
+
+difference_cb = ModelingToolkit.PeriodicCallback(periodic_difference_affect!, 0.1) 
+
+sol2 = solve(prob2, Tsit5(); callback=difference_cb, tstops=collect(prob.tspan[1]:0.1:prob.tspan[2])[2:end]
+)
+
+@test sol(0:0.01:1)[x] ≈ sol2(0:0.01:1)[1,:]
+@test sol(0:0.01:1)[y] ≈ sol2(0:0.01:1)[2,:]