push simplify through derivative expansion

Author: ChrisRackauckas

src/differentials.jl

@@ -36,8 +36,8 @@ $(SIGNATURES)
 
 TODO
 """
-function expand_derivatives(O::Operation)
-    @. O.args = expand_derivatives(O.args)
+function expand_derivatives(O::Operation,simplify=true)
+    @. O.args = expand_derivatives(O.args,simplify)
 
     if isa(O.op, Differential)
         (D, o) = (O.op, O.args[1])
@@ -47,14 +47,16 @@ function expand_derivatives(O::Operation)
         isa(o, Operation)   || return O
         isa(o.op, Variable) && return O
 
-        return sum(1:length(o.args)) do i
-            derivative(o, i) * expand_derivatives(D(o.args[i]))
-        end |> simplify
+        x = sum(1:length(o.args)) do i
+            derivative(o, i) * expand_derivatives(D(o.args[i]),simplify)
+        end
+
+        return simplify ? ModelingToolkit.simplify(x) : x
     end
 
-    return simplify(O)
+    return simplify ? ModelingToolkit.simplify(O) : O
 end
-expand_derivatives(x) = x
+expand_derivatives(x,args...) = x
 
 # Don't specialize on the function here
 """

src/direct.jl

@@ -7,8 +7,7 @@ A helper function for computing the gradient of an expression with respect to
 an array of variable expressions.
 """
 function gradient(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v)(O)) for v in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v)(O),simplify) for v in vars]
 end
 
 """
@@ -20,8 +19,7 @@ A helper function for computing the Jacobian of an array of expressions with res
 an array of variable expressions.
 """
 function jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v)(O)) for O in ops, v in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v)(O),simplify) for O in ops, v in vars]
 end
 
 """
@@ -33,8 +31,7 @@ A helper function for computing the Hessian of an expression with respect to
 an array of variable expressions.
 """
 function hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v2)(Differential(v1)(O))) for v1 in vars, v2 in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v2)(Differential(v1)(O)),simplify) for v1 in vars, v2 in vars]
 end
 
 function simplified_expr(O::Operation)

test/bigsystem.jl

@@ -45,13 +45,31 @@ end
 
 f(du,u,nothing,0.0)
 
-multithreadedf = eval(ModelingToolkit.build_function(du,u,parallel=MultithreadedForm())[2])
+multithreadedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.MultithreadedForm())[2])
 _du = rand(N,N,3)
 _u = rand(N,N,3)
 multithreadedf(_du,_u)
 
+using Distributed
+addprocs(4)
+distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DistributedForm())[2])
+
 jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
-multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=MultithreadedForm())[2])
+serialjac = eval(ModelingToolkit.build_function(vec(jac),u)[2])
+multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.MultithreadedForm())[2])
+distributedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DistributedForm())[2])
+
+#=
+MyA = zeros(N,N)
+AMx = zeros(N,N)
+DA  = zeros(N,N)
+using BenchmarkTools
+@btime f(_du,_u,nothing,0.0)
+@btime multithreadedf(_du,_u)
+@btime distributedf(_du,_u)
 
-#_jac = similar(jac,Float64)
-#multithreadedjac(_jac,_u)
+_jac = similar(jac,Float64)
+@btime serialjac(_jac,_u)
+@btime multithreadedjac(_jac,_u)
+@btime distributedjac(_jac,_u)
+=#