diff --git a/Project.toml b/Project.toml index ff9fc4e..e7690c0 100644 --- a/Project.toml +++ b/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" diff --git a/src/codegen.jl b/src/codegen.jl index 6a093b4..5e27b04 100644 --- a/src/codegen.jl +++ b/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₁ diff --git a/src/primitive.jl b/src/primitive.jl index 48f82f0..ffc2917 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -110,37 +110,37 @@ end @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 diff --git a/test/primitive.jl b/test/primitive.jl index 6769276..46dab5e 100644 --- a/test/primitive.jl +++ b/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