Various cleanup

Author: HarrisonGrodin

src/ModelingToolkit.jl

@@ -7,30 +7,33 @@ using MacroTools
 import MacroTools: splitdef, combinedef
 
 abstract type Expression <: Number end
-abstract type AbstractOperation <: Expression end
-abstract type AbstractComponent <: Expression end
+abstract type AbstractComponent end
+abstract type AbstractSystem end
 
-include("variables.jl")
-
-Base.promote_rule(::Type{T},::Type{T2}) where {T<:Number,T2<:Expression} = Expression
+Base.promote_rule(::Type{<:Number},::Type{<:Expression}) = Expression
 Base.zero(::Type{<:Expression}) = Constant(0)
 Base.one(::Type{<:Expression}) = Constant(1)
 
-function caclulate_jacobian end
+function calculate_jacobian end
+function generate_jacobian end
+function generate_function end
 
 @enum FunctionVersion ArrayFunction=1 SArrayFunction=2
 
+include("variables.jl")
 include("operations.jl")
 include("differentials.jl")
 include("equations.jl")
-include("systems/systems.jl")
 include("systems/diffeqs/diffeqsystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/nonlinear/nonlinear_system.jl")
 include("function_registration.jl")
 include("simplify.jl")
 include("utils.jl")
 
-export Operation, Expression, AbstractComponent, AbstractDomain
+export Operation, Expression, AbstractComponent
+export calculate_jacobian, generate_jacobian, generate_function
+export ArrayFunction, SArrayFunction
 export @register
+
 end # module

src/operations.jl

@@ -1,5 +1,4 @@
-# Parameterize by T so that way it can be Vector{Expression} which is defined after
-struct Operation <: AbstractOperation
+struct Operation <: Expression
     op::Function
     args::Vector{Expression}
 end

src/systems/systems.jl

@@ -30,6 +30,30 @@ function flatten_expr!(x)
     x
 end
 
+function build_function(rhss, vs, ps, args = (); version::FunctionVersion)
+    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
+    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
+    (ls, rs) = zip(var_pairs..., param_pairs...)
+
+    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
+
+    if version === ArrayFunction
+        X = gensym()
+        sys_exprs = [:($X[$i] = $(convert(Expr, rhs))) for (i, rhs) ∈ enumerate(rhss)]
+        let_expr = Expr(:let, var_eqs, build_expr(:block, sys_exprs))
+        :(($X,u,p,$(args...)) -> $let_expr)
+    elseif version === SArrayFunction
+        sys_expr = build_expr(:tuple, [convert(Expr, rhs) for rhs ∈ rhss])
+        let_expr = Expr(:let, var_eqs, sys_expr)
+        :((u,p,$(args...)) -> begin
+            X = $let_expr
+            T = StaticArrays.similar_type(typeof(u), eltype(X))
+            T(X)
+        end)
+    end
+end
+
+
 is_constant(::Constant) = true
 is_constant(::Any) = false
 

src/utils.jl

@@ -25,7 +25,7 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 sys = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
-jac = ModelingToolkit.calculate_jacobian(sys)
+jac = calculate_jacobian(sys)
 @test jac[1,1] == σ*-1
 @test jac[1,2] == σ
 @test jac[1,3] == 0

test/derivatives.jl

@@ -14,7 +14,7 @@ de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 generate_function(de)
 generate_function(de;version=ModelingToolkit.SArrayFunction)
 jac_expr = generate_jacobian(de)
-jac = ModelingToolkit.calculate_jacobian(de)
+jac = calculate_jacobian(de)
 f = ODEFunction(de)
 ModelingToolkit.generate_ode_iW(de)
 
@@ -67,7 +67,7 @@ eqs = [D(x) ~ σ*a,
        D(z) ~ x*y - β*z]
 de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 generate_function(de)
-jac = ModelingToolkit.calculate_jacobian(de)
+jac = calculate_jacobian(de)
 f = ODEFunction(de)
 
 # Define a nonlinear system
@@ -108,7 +108,7 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
-jac = ModelingToolkit.calculate_jacobian(ns)
+jac = calculate_jacobian(ns)
 @testset "nlsys jacobian" begin
     @test jac[1,1] == σ * -1
     @test jac[1,2] == σ
@@ -131,5 +131,5 @@ eqs = [0 ~ σ*a,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
 nlsys_func = generate_function(ns)
-jac = ModelingToolkit.calculate_jacobian(ns)
+jac = calculate_jacobian(ns)
 jac = generate_jacobian(ns)