format

Author: vyudu

docs/src/systems/DiscreteSystem.md

@@ -28,6 +28,7 @@ DiscreteFunction(sys::DiscreteSystem, args...)
 ```
 
 ## Discrete Domain
+
 ```@docs; canonical=false
 Shift
 Prev

docs/src/systems/ImplicitDiscreteSystem.md

@@ -28,6 +28,7 @@ ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
 ```
 
 ## Discrete Domain
+
 ```@docs; canonical=false
 Shift
 Prev

src/ModelingToolkit.jl

@@ -237,7 +237,8 @@ export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDEProblemExpr
 export SystemStructure
 export DiscreteSystem, DiscreteProblem, DiscreteFunction, DiscreteFunctionExpr
-export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction, ImplicitDiscreteFunctionExpr
+export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction,
+       ImplicitDiscreteFunctionExpr
 export JumpSystem
 export ODEProblem, SDEProblem
 export NonlinearFunction, NonlinearFunctionExpr

src/structural_transformation/utils.jl

@@ -530,7 +530,7 @@ Distribute a shift applied to a whole expression or equation.
 Shift(t, 1)(x + y) will become Shift(t, 1)(x) + Shift(t, 1)(y).
 Only shifts variables whose independent variable is the same t that appears in the Shift (i.e. constants, time-independent parameters, etc. do not get shifted).
 """
-function distribute_shift(var) 
+function distribute_shift(var)
     var = unwrap(var)
     var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
 
@@ -553,11 +553,13 @@ function _distribute_shift(expr, shift)
         args = arguments(expr)
 
         if ModelingToolkit.isvariable(expr)
-            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) : (return expr)
+            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) :
+            (return expr)
         elseif op isa Shift
             return shift(expr)
         else
-            return maketerm(typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
+            return maketerm(
+                typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
                 unwrap(expr).metadata)
         end
     else

src/systems/discrete_system/implicit_discrete_system.jl

@@ -239,7 +239,9 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
-ImplicitDiscreteSystem(eq::Equation, args...; kwargs...) = ImplicitDiscreteSystem([eq], args...; kwargs...)
+function ImplicitDiscreteSystem(eq::Equation, args...; kwargs...)
+    ImplicitDiscreteSystem([eq], args...; kwargs...)
+end
 
 function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
@@ -273,11 +275,13 @@ function generate_function(
     obs = observed(sys)
     shifted_obs = Symbolics.Equation[distribute_shift(Next(eq)) for eq in obs]
     obsidxs = observed_equations_used_by(sys, exprs; obs = shifted_obs)
-    extra_assignments = [Assignment(shifted_obs[i].lhs, shifted_obs[i].rhs) for i in obsidxs]
+    extra_assignments = [Assignment(shifted_obs[i].lhs, shifted_obs[i].rhs)
+                         for i in obsidxs]
 
     u_next = map(Next, dvs)
     u = dvs
-    build_function_wrapper(sys, exprs, u_next, u, ps..., iv; p_start = 3, extra_assignments, kwargs...)
+    build_function_wrapper(
+        sys, exprs, u_next, u, ps..., iv; p_start = 3, extra_assignments, kwargs...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
@@ -341,11 +345,13 @@ function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...
     ImplicitDiscreteFunction{true}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ImplicitDiscreteFunction{true}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+function SciMLBase.ImplicitDiscreteFunction{true}(
+        sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ImplicitDiscreteFunction{false}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+function SciMLBase.ImplicitDiscreteFunction{false}(
+        sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
@@ -360,7 +366,6 @@ function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
         eval_module = @__MODULE__,
         analytic = nothing,
         kwargs...) where {iip, specialize}
-
     if !iscomplete(sys)
         error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`")
     end
@@ -404,9 +409,12 @@ struct ImplicitDiscreteFunctionClosure{O, I} <: Function
     f_iip::I
 end
 (f::ImplicitDiscreteFunctionClosure)(u_next, u, p, t) = f.f_oop(u_next, u, p, t)
-(f::ImplicitDiscreteFunctionClosure)(resid, u_next, u, p, t) = f.f_iip(resid, u_next, u, p, t)
+function (f::ImplicitDiscreteFunctionClosure)(resid, u_next, u, p, t)
+    f.f_iip(resid, u_next, u, p, t)
+end
 
-function ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
+function ImplicitDiscreteFunctionExpr{iip}(
+        sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
         version = nothing, p = nothing,
         linenumbers = false,
@@ -427,4 +435,3 @@ end
 function ImplicitDiscreteFunctionExpr(sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunctionExpr{true}(sys, args...; kwargs...)
 end
-

test/implicit_discrete_system.jl

@@ -3,72 +3,73 @@ using ModelingToolkit: t_nounits as t
 using StableRNGs
 
 k = ShiftIndex(t)
-rng = StableRNG(22525) 
+rng = StableRNG(22525)
 
 @testset "Correct ImplicitDiscreteFunction" begin
     @variables x(t) = 1
-    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k) * x(k - 1) - 3], t)
     tspan = (0, 10)
 
     # u[2] - u_next[1]
     # -3 - u_next[2] + u_next[2]*u_next[1]
     f = ImplicitDiscreteFunction(sys)
-    u_next = [3., 1.5]
-    @test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
-    u_next = [0., 0.]
-    @test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
-    
+    u_next = [3.0, 1.5]
+    @test f(u_next, [2.0, 3.0], [], t) ≈ [0.0, 0.0]
+    u_next = [0.0, 0.0]
+    @test f(u_next, [2.0, 3.0], [], t) ≈ [3.0, -3.0]
+
     resid = rand(2)
-    f(resid, u_next, [2.,3.], [], t)
-    @test resid ≈ [3., -3.]
-    
-    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
-    @test prob.u0 == [3., 1.]
+    f(resid, u_next, [2.0, 3.0], [], t)
+    @test resid ≈ [3.0, -3.0]
+
+    prob = ImplicitDiscreteProblem(sys, [x(k - 1) => 3.0], tspan)
+    @test prob.u0 == [3.0, 1.0]
     prob = ImplicitDiscreteProblem(sys, [], tspan)
-    @test prob.u0 == [1., 1.]
+    @test prob.u0 == [1.0, 1.0]
     @variables x(t)
-    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-    @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k) * x(k - 1) - 3], t)
+    @test_throws ErrorException prob=ImplicitDiscreteProblem(sys, [], tspan)
 end
 
 @testset "System with algebraic equations" begin
     @variables x(t) y(t)
-    eqs = [x(k) ~ x(k-1) + x(k-2), 
-           x^2 ~ 1 - y^2]
+    eqs = [x(k) ~ x(k - 1) + x(k - 2),
+        x^2 ~ 1 - y^2]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
     f = ImplicitDiscreteFunction(sys)
 
-    function correct_f(u_next, u, p, t) 
+    function correct_f(u_next, u, p, t)
         [u[2] - u_next[1],
-         u[1] + u[2] - u_next[2],
-         1 - (u_next[1]+u_next[2])^2 - u_next[3]^2]
+            u[1] + u[2] - u_next[2],
+            1 - (u_next[1] + u_next[2])^2 - u_next[3]^2]
     end
 
     for _ in 1:10
         u_next = rand(rng, 3)
         u = rand(rng, 3)
-        @test correct_f(u_next, u, [], 0.) ≈ f(u_next, u, [], 0.)
+        @test correct_f(u_next, u, [], 0.0) ≈ f(u_next, u, [], 0.0)
     end
 
     # Initialization is satisfied.
-    prob = ImplicitDiscreteProblem(sys, [x(k-1) => .3, x(k-2) => .4], (0, 10), guesses = [y => 1])
+    prob = ImplicitDiscreteProblem(
+        sys, [x(k - 1) => 0.3, x(k - 2) => 0.4], (0, 10), guesses = [y => 1])
     @test (prob.u0[1] + prob.u0[2])^2 + prob.u0[3]^2 ≈ 1
 end
 
 @testset "Handle observables in function codegen" begin
     # Observable appears in differential equation
     @variables x(t) y(t) z(t)
-    eqs = [x(k) ~ x(k-1) + x(k-2),
-           y(k) ~ x(k) + x(k-2)*z(k-1), 
-           x + y + z ~ 2]
+    eqs = [x(k) ~ x(k - 1) + x(k - 2),
+        y(k) ~ x(k) + x(k - 2) * z(k - 1),
+        x + y + z ~ 2]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
     @test length(unknowns(sys)) == length(equations(sys)) == 3
     @test occursin("var\"y(t)\"", string(ImplicitDiscreteFunctionExpr(sys)))
 
     # Shifted observable that appears in algebraic equation is properly handled.
-    eqs = [z(k) ~ x(k) + sin(x(k)), 
-           y(k) ~ x(k-1) + x(k-2),
-           z(k) * x(k) ~ 3]
+    eqs = [z(k) ~ x(k) + sin(x(k)),
+        y(k) ~ x(k - 1) + x(k - 2),
+        z(k) * x(k) ~ 3]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
     @test occursin("var\"Shift(t, 1)(z(t))\"", string(ImplicitDiscreteFunctionExpr(sys)))
 end

test/structural_transformation/utils.jl

@@ -169,19 +169,21 @@ end
     k = ShiftIndex(t)
 
     # Expand shifts
-    @test isequal(ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
+    @test isequal(
+        ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
 
     expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
     @test isequal(ST.simplify_shifts(ST.distribute_shift(Shift(t, 2)(expr))),
-                  a*x + Shift(t, 4)(y) + b)
+        a * x + Shift(t, 4)(y) + b)
     @test isequal(ST.distribute_shift(Shift(t, 2)(exp(z))), exp(Shift(t, 2)(z)))
     @test isequal(ST.distribute_shift(Shift(t, 2)(exp(a) + b)), exp(a) + b)
 
-    expr = a^x - log(b*y) + z*x
-    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)), a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z)*Shift(t, -3)(x))
+    expr = a^x - log(b * y) + z * x
+    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)),
+        a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z) * Shift(t, -3)(x))
 
-    expr = x(k+1) ~ x + x(k-1)
-    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k-1) + x(k-2))
+    expr = x(k + 1) ~ x + x(k - 1)
+    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k - 1) + x(k - 2))
 end
 
 @testset "`map_variables_to_equations`" begin