Unify generate_jacobian

Author: HarrisonGrodin

src/ModelingToolkit.jl

@@ -9,7 +9,6 @@ import MacroTools: splitdef, combinedef
 abstract type Expression <: Number end
 abstract type AbstractOperation <: Expression end
 abstract type AbstractComponent <: Expression end
-abstract type AbstractSystem end
 
 include("variables.jl")
 
@@ -25,6 +24,7 @@ function caclulate_jacobian end
 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")

src/systems/diffeqs/diffeqsystem.jl

@@ -78,15 +78,9 @@ function calculate_jacobian(sys::DiffEqSystem, simplify=true)
     return jac
 end
 
-function generate_ode_jacobian(sys::DiffEqSystem, simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    jac = calculate_jacobian(sys, simplify)
-    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
-    block = expr_arr_to_block(exprs)
-    :((J,u,p,t)->$(block))
-end
+system_vars(sys::DiffEqSystem) = sys.dvs
+system_params(sys::DiffEqSystem) = sys.ps
+
 
 function generate_ode_iW(sys::DiffEqSystem, simplify=true)
     var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]

src/systems/nonlinear/nonlinear_system.jl

@@ -34,15 +34,8 @@ function calculate_jacobian(sys::NonlinearSystem,simplify=true)
     sys_exprs
 end
 
-function generate_nlsys_jacobian(sys::NonlinearSystem,simplify=true)
-    var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
-    jac = calculate_jacobian(sys,simplify)
-    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
-    block = expr_arr_to_block(exprs)
-    :((J,u,p,t)->$(block))
-end
+system_vars(sys::NonlinearSystem) = sys.vs
+system_params(sys::NonlinearSystem) = sys.ps
 
 export NonlinearSystem
 export generate_nlsys_function

src/systems/systems.jl

@@ -0,0 +1,19 @@
+export generate_jacobian
+
+
+abstract type AbstractSystem end
+
+
+function system_vars end
+function system_params end
+
+function generate_jacobian(sys::AbstractSystem, simplify = true)
+    vs, ps = system_vars(sys), system_params(sys)
+    var_exprs = [:($(vs[i].name) = u[$i]) for i in eachindex(vs)]
+    param_exprs = [:($(ps[i].name) = p[$i]) for i in eachindex(ps)]
+    jac = calculate_jacobian(sys, simplify)
+    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
+    exprs = vcat(var_exprs, param_exprs, vec(jac_exprs))
+    block = expr_arr_to_block(exprs)
+    :((J,u,p,t) -> $(block))
+end

test/system_construction.jl

@@ -14,7 +14,7 @@ eqs = [D(x) ~ σ*(y-x),
 de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
 ModelingToolkit.generate_ode_function(de;version=ModelingToolkit.SArrayFunction)
-jac_expr = ModelingToolkit.generate_ode_jacobian(de)
+jac_expr = generate_jacobian(de)
 jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)
 ModelingToolkit.generate_ode_iW(de)
@@ -121,7 +121,7 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs)
 nlsys_func = ModelingToolkit.generate_nlsys_function(ns)
-jac = ModelingToolkit.generate_nlsys_jacobian(ns)
+jac = generate_jacobian(ns)
 f = @eval eval(nlsys_func)
 
 # Intermediate calculations
@@ -133,4 +133,4 @@ eqs = [a ~ y-x,
 ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
 nlsys_func = ModelingToolkit.generate_nlsys_function(ns)
 jac = ModelingToolkit.calculate_jacobian(ns)
-jac = ModelingToolkit.generate_nlsys_jacobian(ns)
+jac = generate_jacobian(ns)