direct calculus

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -88,6 +88,7 @@ include("equations.jl")
 include("function_registration.jl")
 include("simplify.jl")
 include("utils.jl")
+include("direct.jl")
 include("systems/diffeqs/diffeqsystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/nonlinear/nonlinear_system.jl")

src/direct.jl

@@ -0,0 +1,30 @@
+function gradient(O::Operation, vars::AbstractVector{Operation}; simplify = true)
+    out = [expand_derivatives(Differential(v)(O)) for v in vars]
+    simplify ? simplify_constants.(out) : out
+end
+
+function jacobian(ops::AbstractVector{Operation}, vars::AbstractVector{Operation}; simplify = true)
+    out = [expand_derivatives(Differential(v)(O)) for O in ops, v in vars]
+    simplify ? simplify_constants.(out) : out
+end
+
+function hessian(O::Operation, vars::AbstractVector{Operation}; simplify = true)
+    out = [expand_derivatives(Differential(v2)(Differential(v1)(O))) for v1 in vars, v2 in vars]
+    simplify ? simplify_constants.(out) : out
+end
+
+function simplified_expr(O::Operation)
+  if O isa Constant
+    return O.value
+  elseif isa(O.op, Differential)
+    return :(derivative($(simplified_expr(O.args[1])),$(simplified_expr(O.op.x))))
+  elseif isa(O.op, Variable)
+    isempty(O.args) && return O.op.name
+    return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
+  end
+  return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
+end
+
+function simplified_expr(eq::Equation)
+    Expr(:(=), simplified_expr(eq.lhs), simplified_expr(eq.rhs))
+end

src/simplify.jl

@@ -75,19 +75,3 @@ function _simplify_constants(O::Operation, shorten_tree)
 end
 _simplify_constants(x, shorten_tree) = x
 _simplify_constants(x) = _simplify_constants(x, true)
-
-function simplified_expr(O::Operation)
-  if O isa Constant
-    return O.value
-  elseif isa(O.op, Differential)
-    return :(derivative($(simplified_expr(O.args[1])),$(simplified_expr(O.op.x))))
-  elseif isa(O.op, Variable)
-    isempty(O.args) && return O.op.name
-    return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
-  end
-  return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
-end
-
-function simplified_expr(eq::Equation)
-    Expr(:(=), simplified_expr(eq.lhs), simplified_expr(eq.rhs))
-end

src/systems/nonlinear/nonlinear_system.jl

@@ -1,6 +1,5 @@
 export NonlinearSystem
 
-
 struct NLEq
     rhs::Expression
 end

test/direct.jl

@@ -0,0 +1,19 @@
+using ModelingToolkit, StaticArrays, LinearAlgebra
+using DiffEqBase
+using Test
+
+# Define some variables
+@parameters t σ ρ β
+@variables x y z
+
+eqs = [σ*(y-x),
+       x*(ρ-z)-y,
+       x*y - β*z]
+
+∂ = ModelingToolkit.jacobian(eqs,[x,y,z])
+for i in 1:3
+    ∇ = ModelingToolkit.gradient(eqs[i],[x,y,z])
+    @test isequal(∂[i,:],∇)
+end
+@test all(isequal.(ModelingToolkit.gradient(eqs[1],[x,y,z]),[σ * -1,σ,0]))
+@test all(isequal.(ModelingToolkit.hessian(eqs[1],[x,y,z]),0))