Merge pull request #156 from JuliaDiffEq/simplified_expr

Improved direct usage

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -3,7 +3,8 @@ module ModelingToolkit
 export Operation, Expression
 export calculate_jacobian, generate_jacobian, generate_function
 export independent_variables, dependent_variables, parameters
-export @register
+export simplified_expr, eval_function
+export @register, @I
 export modelingtoolkitize
 
 
@@ -87,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,43 @@
+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(c::Constant)
+    c.value
+end
+
+function simplified_expr(eq::Equation)
+    Expr(:(=), simplified_expr(eq.lhs), simplified_expr(eq.rhs))
+end
+
+macro I(ex)
+    name = :ICompile
+    ret = return quote
+        macro $(esc(name))()
+            esc($ex)
+        end
+    end
+end

src/systems/nonlinear/nonlinear_system.jl

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

src/utils.jl

@@ -31,7 +31,7 @@ function flatten_expr!(x)
     x
 end
 
-function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); constructor=nothing)
+function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr; constructor=nothing)
     _vs = map(x-> x isa Operation ? x.op : x, vs)
     _ps = map(x-> x isa Operation ? x.op : x, ps)
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
@@ -48,20 +48,23 @@ function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs
 
     sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
     let_expr = Expr(:let, var_eqs, sys_expr)
+
+    fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))
     quote
-        function $fname($X,u,p,$(args...))
+        function $fname($X,$(fargs.args...))
             $ip_let_expr
             nothing
         end
-        function $fname(u,p,$(args...))
+        function $fname($(fargs.args...))
             X = $let_expr
-            T = $(constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, eltype(X)); du .= x)) : constructor)
-            T(X)
+            T = promote_type(map(typeof,X)...)
+            convert.(T,X)
+            construct = $(constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, T, $(size(rhss)...)); vec(du) .= x; du)) : constructor)
+            construct(X)
         end
     end
 end
 
-
 is_constant(::Constant) = true
 is_constant(::Any) = false
 

test/direct.jl

@@ -0,0 +1,127 @@
+using ModelingToolkit, StaticArrays, LinearAlgebra
+using DiffEqBase
+using Test
+
+# Calculus
+@parameters t σ ρ β
+@variables x y z
+
+eqs = [σ*(y-x),
+       x*(ρ-z)-y,
+       x*y - β*z]
+
+simpexpr = [
+   :(σ * (y - x))
+   :(x * (ρ - z) - y)
+   :(x * y - β * z)
+   ]
+
+for i in 1:3
+   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
+   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
+end
+
+∂ = 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))
+
+# Function building
+
+@parameters σ() ρ() β()
+@variables x y z
+eqs = [σ*(y-x),
+       x*(ρ-z)-y,
+       x*y - β*z]
+f = eval(ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β]))
+out = [1.0,2,3]
+o1 = f([1.0,2,3],[1.0,2,3])
+f(out,[1.0,2,3],[1.0,2,3])
+@test all(o1 .== out)
+
+function test_worldage()
+   @parameters σ() ρ() β()
+   @variables x y z
+   eqs = [σ*(y-x),
+          x*(ρ-z)-y,
+          x*y - β*z]
+   _f = eval(ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β]))
+   f(u,p) = ModelingToolkit.fast_invokelatest(_f,typeof(u),u,p)
+   f(du,u,p) = ModelingToolkit.fast_invokelatest(_f,Nothing,du,u,p)
+   out = [1.0,2,3]
+   o1 = f([1.0,2,3],[1.0,2,3])
+   f(out,[1.0,2,3],[1.0,2,3])
+end
+test_worldage()
+
+mac = @I begin
+   @parameters σ() ρ() β()
+   @variables x() y() z()
+
+   eqs = [σ*(y-x),
+          x*(ρ-z)-y,
+          x*y - β*z]
+   ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β])
+end
+f = @ICompile
+out = [1.0,2,3]
+o1 = f([1.0,2,3],[1.0,2,3])
+f(out,[1.0,2,3],[1.0,2,3])
+@test all(o1 .== out)
+
+mac = @I begin
+   @parameters σ ρ β
+   @variables x y z
+
+   eqs = [σ*(y-x),
+          x*(ρ-z)-y,
+          x*y - β*z]
+   ∂ = ModelingToolkit.jacobian(eqs,[x,y,z])
+   ModelingToolkit.build_function(∂,[x,y,z],[σ,ρ,β])
+end
+f = @ICompile
+out = zeros(3,3)
+o1 = f([1.0,2,3],[1.0,2,3])
+f(out,[1.0,2,3],[1.0,2,3])
+@test all(out .== o1)
+
+## No parameters
+@variables x y z
+eqs = [(y-x)^2,
+       x*(x-z)-y,
+       x*y - y*z]
+f = eval(ModelingToolkit.build_function(eqs,[x,y,z]))
+out = zeros(3)
+o1 = f([1.0,2,3])
+f(out,[1.0,2,3])
+@test all(out .== o1)
+
+function test_worldage()
+   @variables x y z
+   eqs = [(y-x)^2,
+          x*(x-z)-y,
+          x*y - y*z]
+   _f = eval(ModelingToolkit.build_function(eqs,[x,y,z]))
+   f(u) = ModelingToolkit.fast_invokelatest(_f,typeof(u),u)
+   f(du,u) = ModelingToolkit.fast_invokelatest(_f,Nothing,du,u)
+   out = zeros(3)
+   o1 = f([1.0,2,3])
+   f(out,[1.0,2,3])
+end
+test_worldage()
+
+mac = @I begin
+   @variables x y z
+   eqs = [(y-x)^2,
+          x*(x-z)-y,
+          x*y - y*z]
+   ModelingToolkit.build_function(eqs,[x,y,z])
+end
+f = @ICompile
+out = zeros(3)
+o1 = f([1.0,2,3])
+f(out,[1.0,2,3])
+@test all(out .== o1)

test/runtests.jl

@@ -3,4 +3,5 @@ using ModelingToolkit, Test
 @testset "Parsing Test" begin include("variable_parsing.jl") end
 @testset "Differentiation Test" begin include("derivatives.jl") end
 @testset "Simplify Test" begin include("simplify.jl") end
+@testset "Direct Usage Test" begin include("direct.jl") end
 @testset "System Construction Test" begin include("system_construction.jl") end

test/system_construction.jl

@@ -30,6 +30,9 @@ end
 eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
+
+ModelingToolkit.simplified_expr.(eqs)[1]
+:(derivative(x(t), t) = σ * (y(t) - x(t))).args
 de = ODESystem(eqs)
 test_diffeq_inference("standard", de, t, (x, y, z), (σ, ρ, β))
 generate_function(de, [x,y,z], [σ,ρ,β])