Add basic complex functions, fix 2-arg atan (JuliaDiff#64)

* add rules for complex util functions

* fix accidentally deleted line

* fix test cases

Author: simeonschaub

src/rulesets/Base/base.jl

@@ -83,14 +83,20 @@
 
 @scalar_rule(hypot(x, y), (x / Ω, y / Ω))
 @scalar_rule(sincos(x), @setup((sinx, cosx) = Ω), cosx, -sinx)
-@scalar_rule(atan(x, y), @setup(u = x^2 + y^2), (y / u, -x / u))
+@scalar_rule(atan(y, x), @setup(u = x^2 + y^2), (x / u, -y / u))
 
 @scalar_rule(max(x, y), @setup(gt = x > y), (gt, !gt))
 @scalar_rule(min(x, y), @setup(gt = x > y), (!gt, gt))
 @scalar_rule(mod(x, y), @setup((u, nan) = promote(x / y, NaN16)),
              (ifelse(isint, nan, one(u)), ifelse(isint, nan, -floor(u))))
 @scalar_rule(rem(x, y), @setup((u, nan) = promote(x / y, NaN16)),
              (ifelse(isint, nan, one(u)), ifelse(isint, nan, -trunc(u))))
+@scalar_rule(angle(x::Complex), @setup(u = abs2(x)), Wirtinger(-im//2 * x' / u, im//2 * x / u))
+@scalar_rule(angle(x::Real), Zero())
+@scalar_rule(real(x::Complex), Wirtinger(1//2, 1//2))
+@scalar_rule(real(x::Real), One())
+@scalar_rule(imag(x::Complex), Wirtinger(-im//2, im//2))
+@scalar_rule(imag(x::Real), Zero())
 
 # product rule requires special care for arguments where `mul` is non-commutative
 

test/rulesets/Base/base.jl

@@ -99,6 +99,19 @@
         end
     end
 
+    @testset "Unary complex functions" begin
+        for x in (-6, rand.((Float32, Float64, Complex{Float32}, Complex{Float64}))...)
+            rtol = x isa Complex{Float32} ? 1e-6 : 1e-9
+            test_scalar(real, x; rtol=rtol)
+            test_scalar(imag, x; rtol=rtol)
+            # TODO: implement correct complex derivative
+            x isa Real && test_scalar(abs, x; rtol=rtol)
+            test_scalar(angle, x; rtol=rtol)
+            test_scalar(abs2, x; rtol=rtol)
+            test_scalar(conj, x; rtol=rtol)
+        end
+    end
+
     @testset "*(x, y)" begin
         x, y = rand(3, 2), rand(2, 5)
         z, (dx, dy) = rrule(*, x, y)

test/runtests.jl

@@ -12,7 +12,7 @@ using Test
 # For testing purposes we use a lot of
 using ChainRulesCore: cast, extern, accumulate, accumulate!, store!, @scalar_rule,
     Wirtinger, wirtinger_primal, wirtinger_conjugate,
-    Zero, One, Casted, DNE, Thunk, DNERule
+    Zero, One, Casted, DNE, Thunk, DNERule, AbstractDifferential
 
 include("test_util.jl")