Linearizing delayed equation systems (#1151)

Author: sharanry

src/systems/discrete_system/discrete_system.jl

@@ -113,6 +113,7 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     dvs = states(sys)
     ps = parameters(sys)
     eqs = equations(sys)
+    eqs = linearize_eqs(sys, eqs)
     # defs = defaults(sys)
     t = get_iv(sys)
     u0 = varmap_to_vars(u0map,dvs)
@@ -126,6 +127,55 @@ function DiffEqBase.DiscreteProblem(sys::DiscreteSystem,u0map,tspan,
     DiscreteProblem(f,u0,tspan,p;kwargs...)
 end
 
+function linearize_eqs(sys, eqs=get_eqs(sys); return_max_delay=false)
+    unique_states = unique(operation.(states(sys)))
+    max_delay = Dict(v=>0.0 for v in unique_states)
+
+    r = @rule ~t::(t -> istree(t) && any(isequal(operation(t)), operation.(states(sys))) && is_delay_var(get_iv(sys), t)) => begin
+        delay = get_delay_val(get_iv(sys), first(arguments(~t)))
+        if delay > max_delay[operation(~t)]
+            max_delay[operation(~t)] = delay
+        end
+        nothing
+    end
+    SymbolicUtils.Postwalk(r).(rhss(eqs))
+
+    if any(values(max_delay) .> 0)
+
+        dts = Dict(v=>Any[] for v in unique_states)
+        state_ops = Dict(v=>Any[] for v in unique_states)
+        for v in unique_states
+            for eq in eqs
+                if isdifferenceeq(eq) && istree(arguments(eq.lhs)[1]) && isequal(v, operation(arguments(eq.lhs)[1]))
+                    append!(dts[v], [operation(eq.lhs).dt])
+                    append!(state_ops[v], [operation(eq.lhs)])
+                end
+            end
+        end
+
+        all(length.(unique.(values(state_ops))) .<= 1) || error("Each state should be used with single difference operator.")
+        
+        dts_gcd = Dict()
+        for v in keys(dts)
+            dts_gcd[v] = (length(dts[v]) > 0) ? first(dts[v]) : nothing
+        end
+
+        lin_eqs = [
+            v(get_iv(sys) - (t)) ~ v(get_iv(sys) - (t-dts_gcd[v]))
+            for v in unique_states if max_delay[v] > 0 && dts_gcd[v]!==nothing for t in collect(max_delay[v]:(-dts_gcd[v]):0)[1:end-1] 
+        ]
+        eqs = vcat(eqs, lin_eqs)
+    end
+    if return_max_delay return eqs, max_delay end
+    eqs
+end
+
+function get_delay_val(iv, x)
+    delay = x - iv
+    isequal(delay > 0, true) && error("Forward delay not permitted")
+    return -delay
+end
+
 check_difference_variables(eq) = check_operator_variables(eq, Difference)
 
 function generate_function(

src/utils.jl

@@ -105,10 +105,25 @@ function check_parameters(ps, iv)
     end
 end
 
+function is_delay_var(iv, var)
+    args = nothing
+    try
+        args = arguments(var)
+    catch
+        return false
+    end
+    length(args) > 1 && return false
+    isequal(first(args), iv) && return false
+    delay = iv - first(args)
+    delay isa Integer || 
+    delay isa AbstractFloat ||
+    (delay isa Num && isreal(value(delay))) 
+end
+
 function check_variables(dvs, iv)
     for dv in dvs
         isequal(iv, dv) && throw(ArgumentError("Independent variable $iv not allowed in dependent variables."))
-        occursin(iv, iv_from_nested_derivative(dv)) || throw(ArgumentError("Variable $dv is not a function of independent variable $iv."))
+        (is_delay_var(iv, dv) || occursin(iv, iv_from_nested_derivative(dv))) || throw(ArgumentError("Variable $dv is not a function of independent variable $iv."))
     end
 end
 

test/discretesystem.jl

@@ -56,3 +56,37 @@ prob_map = DiscreteProblem(sir_map!,u0,tspan,p);
 sol_map2 = solve(prob_map,FunctionMap());
 
 @test Array(sol_map) ≈ Array(sol_map2)
+
+# Delayed difference equation
+@parameters t
+@variables x(..) y(..) z(t)
+D1 = Difference(t; dt=1.5)
+D2 = Difference(t; dt=2)
+
+@test ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(x(t-2)))
+@test ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(y(t-1)))
+@test !ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(z))
+@test_throws ErrorException ModelingToolkit.get_delay_val(Symbolics.value(t), Symbolics.arguments(Symbolics.value(x(t+2)))[1])
+@test_throws ErrorException z(t)
+
+# Equations
+eqs = [
+    D1(x(t)) ~ 0.4x(t) + 0.3x(t-1.5) + 0.1x(t-3),
+    D2(y(t)) ~ 0.3y(t) + 0.7y(t-2) + 0.1z,
+]
+
+# System
+@named sys = DiscreteSystem(eqs,t,[x(t),x(t-1.5),x(t-3),y(t),y(t-2),z],[])
+
+eqs2, max_delay = ModelingToolkit.linearize_eqs(sys; return_max_delay=true)
+
+@test max_delay[Symbolics.operation(Symbolics.value(x(t)))] ≈ 3
+@test max_delay[Symbolics.operation(Symbolics.value(y(t)))] ≈ 2
+
+linearized_eqs = [
+    eqs
+    x(t - 3.0) ~ x(t - 1.5)
+    x(t - 1.5) ~ x(t)
+    y(t - 2.0) ~ y(t)
+]
+@test all(eqs2 .== linearized_eqs)