Enable usage of difference operator in ODESystem

Author: sharanry

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
@@ -180,6 +181,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

@@ -94,7 +94,7 @@ function generate_function(
     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] :
-                          [eq.rhs for eq in eqs]
+                          [eq.rhs for eq in eqs if isdiffeq(eq)]
     #rhss = Let(obss, rhss)
 
     # TODO: add an optional check on the ordering of observed equations
@@ -109,6 +109,38 @@ function generate_function(
     end
 end
 
+function generate_difference_cb(sys::ODESystem, dvs = states(sys), ps = parameters(sys);
+    kwargs...)
+    eqs = equations(sys)
+    foreach(check_difference_variables, eqs)
+    # substitute x(t) by just x
+
+    # map(x -> isempty(x) ? Val{0} : first(x), 
+    rhss = [ 
+        begin
+            ind = findfirst(eq -> isdifference(eq.lhs) && isequal(arguments(eq.lhs)[1], s), eqs)
+            ind === nothing ? Val{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)] 
+    all(dts .== dts[1]) || error("All difference variables should have same time steps.")
+
+    PeriodicCallback(cb_affect!, dts[1])
+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
@@ -713,6 +745,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

@@ -133,42 +133,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...)
     # NOTE: this assumes that the order of algebric equations doesn't matter
     diffvars = OrderedSet()
@@ -206,30 +170,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)
@@ -317,21 +257,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

@@ -121,9 +121,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