Fix power (^) function gives NaN at 0

Project.toml

@@ -7,6 +7,7 @@ version = "0.2.1"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesOverloadGeneration = "f51149dc-2911-5acf-81fc-2076a2a81d4f"
+IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 SliceMap = "82cb661a-3f19-5665-9e27-df437c7e54c8"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
@@ -18,6 +19,7 @@ ChainRulesCore = "1"
 ChainRulesOverloadGeneration = "0.1"
 SliceMap = "0.2"
 SpecialFunctions = "2"
+IrrationalConstants = "0.2"
 SymbolicUtils = "1"
 Zygote = "0.6.55"
 julia = "1.6"

src/codegen.jl

@@ -1,5 +1,6 @@
 using ChainRulesCore
 using SpecialFunctions
+using IrrationalConstants: sqrtπ
 using SymbolicUtils, SymbolicUtils.Code
 using SymbolicUtils: BasicSymbolic, Pow
 
@@ -26,7 +27,7 @@ using SymbolicUtils: BasicSymbolic, Pow
 @scalar_rule asech(x::BasicSymbolic) inv(x * -sqrt(1 - x^2))
 @scalar_rule asinh(x::BasicSymbolic) inv(sqrt(x^2 + 1))
 @scalar_rule atanh(x::BasicSymbolic) inv(1 - x^2)
-@scalar_rule erf(x::BasicSymbolic) exp(-x^2)*(2 / sqrt(pi))
+@scalar_rule erf(x::BasicSymbolic) exp(-x^2) * 2/sqrtπ
 
 dummy = (NoTangent(), 1)
 @syms t₁

src/primitive.jl

@@ -110,37 +110,37 @@
 @generated function ^(t::TaylorScalar{T, N}, n::S) where {S <: Number, T, N}
     ex = quote
         v = value(t)
-        v1 = ^(v[1], n)
+        w11 = 1
+        u1 = ^(v[1], n)
     end
-    for i in 2:N
+    for k in 1:N
         ex = quote
             $ex
-            $(Symbol('v', i)) = +($([:((n * $(binomial(i - 2, j - 1)) -
-                                        $(binomial(i - 2, j - 2))) * $(Symbol('v', j)) *
-                                       v[$(i + 1 - j)])
-                                     for j in 1:(i - 1)]...)) / v[1]
+            $(Symbol('p', k)) = ^(v[1], n - $(k - 1))
         end
     end
-    ex = :($ex; TaylorScalar($([Symbol('v', i) for i in 1:N]...)))
-    return :(@inbounds $ex)
-end
-
-@generated function ^(t::TaylorScalar{T, N}, n::S) where {S <: Integer, T, N}
-    # TODO: optimize for small powers
-    ex = quote
-        v = value(t)
-        v1 = ^(v[1], n)
-    end
     for i in 2:N
+        subex = quote
+            $(Symbol('w', i, 1)) = 0
+        end
+        for k in 2:i
+            subex = quote
+                $subex
+                $(Symbol('w', i, k)) = +($([:((n * $(binomial(i - 2, j - 1)) -
+                                               $(binomial(i - 2, j - 2))) *
+                                              $(Symbol('w', j, k - 1)) *
+                                              v[$(i + 1 - j)])
+                                            for j in (k - 1):(i - 1)]...))
+            end
+        end
         ex = quote
             $ex
-            $(Symbol('v', i)) = +($([:((n * $(binomial(i - 2, j - 1)) -
-                                        $(binomial(i - 2, j - 2))) * $(Symbol('v', j)) *
-                                       v[$(i + 1 - j)])
-                                     for j in 1:(i - 1)]...)) / v[1]
+            $subex
+            $(Symbol('u', i)) = +($([:($(Symbol('w', i, k)) * $(Symbol('p', k)))
+                                     for k in 2:i]...))
         end
     end
-    ex = :($ex; TaylorScalar($([Symbol('v', i) for i in 1:N]...)))
+    ex = :($ex; TaylorScalar($([Symbol('u', i) for i in 1:N]...)))
     return :(@inbounds $ex)
 end
 

test/primitive.jl

@@ -30,14 +30,18 @@ end
 end
 
 @testset "Binary functions" begin
-    some_number, another_number = 1.9, 2.6
+    some_number, another_number = 1.9, 5.6
     for f in (*, /), order in (1, 4)
         fdm = central_fdm(12, order)
         closure = x -> exp(f(x, another_number))
         @test derivative(closure, some_number, order)≈fdm(closure, some_number) rtol=1e-6
     end
-    for f in (x -> x^7, x -> x^another_number), order in (2, 4)
+    for f in (x -> x^7, x -> x^another_number), order in (1, 2, 4)
         fdm = central_fdm(12, order)
         @test derivative(f, some_number, order)≈fdm(f, some_number) rtol=1e-6
     end
+    for f in (x -> x^7, x -> x^another_number), order in (1, 2)
+        fdm = forward_fdm(12, order)
+        @test derivative(f, 0, order)≈fdm(f, 0) atol=1e-6
+    end
 end