multiargument chain rule

Author: ChrisRackauckas

src/operations.jl

@@ -6,7 +6,7 @@ end
 
 # Recursive ==
 function Base.:(==)(x::Operation,y::Operation)
-    x.op == y.op && all(isequal.(x.args,y.args))
+    x.op == y.op && length(x.args) == length(y.args) && all(isequal.(x.args,y.args))
 end
 Base.:(==)(x::Operation,y::Number) = false
 Base.:(==)(x::Number,y::Operation) = false

src/operators.jl

@@ -26,19 +26,9 @@ Variable(x::Variable,D::Differential) = Variable(x.name,x.value,x.value_type,
 
 function expand_derivatives(O::Operation)
     if O.op == Derivative
-        #=
-        diff_idxs = find(x->isequal(x,by.x),O.args)
-        (diff_idxs != nothing || length(diff_idxs) > 1) && error("Derivatives of multi-argument functions require matching a unique argument.")
-        idx = first(diff_idxs)
-        =#
-        i = 1
-        if typeof(O.args[1].args[i]) == typeof(O.args[2].x) && isequal(O.args[1].args[i],O.args[2].x)
-            Derivative(O.args[1],i)
-        else
-            D = Differential(O.args[2].x)
-            cr_exp = D*O.args[1].args[i]
-            Derivative(O.args[1],i) * expand_derivatives(cr_exp)
-        end
+        D = O.args[2]
+        o = O.args[1]
+        simplify_constants(sum(i->Derivative(o,i)*expand_derivatives(D*o.args[i]),1:length(o.args)))
     else
         for i in 1:length(O.args)
             O.args[i] = expand_derivatives(O.args[i])

src/simplify.jl

@@ -36,7 +36,7 @@ function _simplify_constants(O)
         idxs = find(x->typeof(x)<:Variable && isequal(x,Constant(0)),O.args)
         _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
         if isempty(_O.args)
-            return _O
+            return Constant(0)
         elseif length(_O.args) == 1
             return _O.args[1]
         else

src/variables.jl

@@ -65,6 +65,14 @@ function Base.:(==)(x::Variable,y::Variable)
     x.value_type == y.value_type && x.diff == y.diff
 end
 
+function Base.:(==)(x::Variable,y::Number)
+    x == Constant(y)
+end
+
+function Base.:(==)(x::Number,y::Variable)
+    Constant(x) == y
+end
+
 function Base.Expr(x::Variable)
     if x.diff == nothing
         return :($(x.name))

test/derivatives.jl

@@ -21,18 +21,22 @@ dpow2 = Derivative(^,[x, y],Val{2})
 @test dpow1 == y*x^(y-1)
 @test dpow2 == x^y*log(x)
 
+d1 = D*(sin(t)*t)
+d2 = D*(sin(t)*cos(t))
+@test expand_derivatives(d1) == t*cos(t)+sin(t)
+@test expand_derivatives(d2) == cos(t)*cos(t)+sin(t)*-sin(t)
 
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 sys = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
 jac = SciCompDSL.generate_nlsys_jacobian(sys)
-@test_broken jac[1,1] == -σ
-@test_broken jac[1,2] == σ
-@test_broken jac[1,3] == 0
+@test jac[1,1] == σ*-1
+@test jac[1,2] == σ
+@test jac[1,3] == 0
 @test jac[2,1] == ρ-z
-@test_broken jac[2,2] == -1
-@test_broken jac[2,3] == -x
+@test jac[2,2] == -1
+@test jac[2,3] == x*-1
 @test jac[3,1] == y
-@test_broken jac[3,2] == x
-@test_broken jac[3,3] == -β
+@test jac[3,2] == x
+@test jac[3,3] == -1*β