Go back to <: Real

JuliaDiff · May 22, 2024 · 6dedcac · 6dedcac
1 parent 3770f8b
commit 6dedcac
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 67 deletions.
diff --git a/docs/make.jl b/docs/make.jl
@@ -15,7 +15,7 @@ makedocs(;
         assets = String[]),
     pages = [
         "Home" => "index.md",
-        "API" => "api.md",
+        "API" => "api.md"
     ])
 
 deploydocs(;

diff --git a/examples/halley.jl b/examples/halley.jl
@@ -4,7 +4,7 @@ using TaylorDiff
 using LinearAlgebra
 using LinearSolve
 
-function newton(f, x0, p; tol=1e-10, maxiter=100)
+function newton(f, x0, p; tol = 1e-10, maxiter = 100)
     x = x0
     for i in 1:maxiter
         fx = f(x, p)
@@ -22,7 +22,7 @@ function newton(f, x0, p; tol=1e-10, maxiter=100)
     return x
 end
 
-function halley(f, x0, p; tol=1e-10, maxiter=100)
+function halley(f, x0, p; tol = 1e-10, maxiter = 100)
     x = x0
     for i in 1:maxiter
         fx = f(x, p)

diff --git a/examples/integration.jl b/examples/integration.jl
@@ -17,10 +17,11 @@ eqsdiff: RHS
 t: constructed by TaylorScalar{T, N}(t0, 1), which means unit perturbation
 x0: initial value
 """
-function jetcoeffs_taylordiff(eqsdiff::Function, t::TaylorScalar{T, N}, x0::U, params) where
-    {T<:Real, U<:Number, N}
+function jetcoeffs_taylordiff(eqsdiff::Function, t::TaylorScalar{T, N}, x0::U,
+        params) where
+        {T <: Real, U <: Number, N}
     x = TaylorScalar{U, N}(x0) # x.values[1] is defined, others are 0
-    for index in 1:N-1 # computes x.values[index + 1]
+    for index in 1:(N - 1) # computes x.values[index + 1]
         f = eqsdiff(x, params, t)
         df = TaylorDiff.extract_derivative(f, index)
         x = update_coefficient(x, index + 1, df)
@@ -35,11 +36,13 @@ eqsdiff!: RHS, in non-allocation form
 t: constructed by TaylorScalar{T, N}(t0, 1), which means unit perturbation
 x0: initial value
 """
-function jetcoeffs_array_taylordiff(eqsdiff!::Function, t::TaylorScalar{T, N}, x0::AbstractArray{U, D}, params) where
-    {T<:Real, U<:Number, N, D}
+function jetcoeffs_array_taylordiff(
+        eqsdiff!::Function, t::TaylorScalar{T, N}, x0::AbstractArray{U, D},
+        params) where
+        {T <: Real, U <: Number, N, D}
     x = map(TaylorScalar{U, N}, x0) # x.values[1] is defined, others are 0
     f = similar(x)
-    for index in 1:N-1 # computes x.values[index + 1]
+    for index in 1:(N - 1) # computes x.values[index + 1]
         eqsdiff!(f, x, params, t)
         df = TaylorDiff.extract_derivative.(f, index)
         x = update_coefficient.(x, index + 1, df)

diff --git a/src/derivative.jl b/src/derivative.jl
@@ -42,11 +42,11 @@ function derivatives end
 # Core APIs
 
 # Added to help Zygote infer types
-@inline function make_seed(x::T, l::S, ::Val{N}) where {T <: TN, S <: TN, N}
+@inline function make_seed(x::T, l::S, ::Val{N}) where {T <: Real, S <: Real, N}
     TaylorScalar{T, N}(x, convert(T, l))
 end
 
-@inline function make_seed(x::AbstractArray{T}, l, vN::Val{N}) where {T <: TN, N}
+@inline function make_seed(x::AbstractArray{T}, l, vN::Val{N}) where {T <: Real, N}
     broadcast(make_seed, x, l, vN)
 end
 

diff --git a/src/primitive.jl b/src/primitive.jl
@@ -74,9 +74,13 @@ end
 
 # Binary
 
+const AMBIGUOUS_TYPES = (AbstractFloat, Irrational, Integer, Rational, Real, RoundingMode)
+
 for op in [:>, :<, :(==), :(>=), :(<=)]
-    @eval @inline $op(a::Number, b::TaylorScalar) = $op(a, value(b)[1])
-    @eval @inline $op(a::TaylorScalar, b::Number) = $op(value(a)[1], b)
+    for R in AMBIGUOUS_TYPES
+        @eval @inline $op(a::TaylorScalar, b::$R) = $op(value(a)[1], b)
+        @eval @inline $op(a::$R, b::TaylorScalar) = $op(a, value(b)[1])
+    end
     @eval @inline $op(a::TaylorScalar, b::TaylorScalar) = $op(value(a)[1], value(b)[1])
 end
 
@@ -110,41 +114,43 @@ end
     return :(@inbounds $ex)
 end
 
-@generated function ^(t::TaylorScalar{T, N}, n::S) where {S <: Number, T, N}
-    ex = quote
-        v = value(t)
-        w11 = 1
-        u1 = ^(v[1], n)
-    end
-    for k in 1:N
+for R in (Integer, Real)
+    @eval @generated function ^(t::TaylorScalar{T, N}, n::S) where {S <: $R, T, N}
         ex = quote
-            $ex
-            $(Symbol('p', k)) = ^(v[1], n - $(k - 1))
+            v = value(t)
+            w11 = 1
+            u1 = ^(v[1], n)
         end
-    end
-    for i in 2:N
-        subex = quote
-            $(Symbol('w', i, 1)) = 0
+        for k in 1:N
+            ex = quote
+                $ex
+                $(Symbol('p', k)) = ^(v[1], n - $(k - 1))
+            end
         end
-        for k in 2:i
+        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
                 $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)]...))
+                $(Symbol('u', i)) = +($([:($(Symbol('w', i, k)) * $(Symbol('p', k)))
+                                         for k in 2:i]...))
             end
         end
-        ex = quote
-            $ex
-            $subex
-            $(Symbol('u', i)) = +($([:($(Symbol('w', i, k)) * $(Symbol('p', k)))
-                                     for k in 2:i]...))
-        end
+        ex = :($ex; TaylorScalar($([Symbol('u', i) for i in 1:N]...)))
+        return :(@inbounds $ex)
     end
-    ex = :($ex; TaylorScalar($([Symbol('u', i) for i in 1:N]...)))
-    return :(@inbounds $ex)
 end
 
 @generated function raise(f::T, df::TaylorScalar{T, M},

diff --git a/src/scalar.jl b/src/scalar.jl
@@ -4,6 +4,20 @@ import Base: convert, promote_rule
 
 export TaylorScalar
 
+"""
+    TaylorDiff.can_taylor(V::Type)
+
+Determines whether the type V is allowed as the scalar type in a
+Dual. By default, only `<:Real` types are allowed.
+"""
+can_taylorize(::Type{<:Real}) = true
+can_taylorize(::Type) = false
+
+@noinline function throw_cannot_taylorize(V::Type)
+    throw(ArgumentError("Cannot create a Taylor polynomial over scalar type $V." *
+                        " If the type behaves as a scalar, define TaylorDiff.can_taylorize(::Type{$V}) = true."))
+end
+
 """
     TaylorScalar{T, N}
 
@@ -13,20 +27,23 @@ Representation of Taylor polynomials.
 
 - `value::NTuple{N, T}`: i-th element of this stores the (i-1)-th derivative
 """
-struct TaylorScalar{T, N}
+struct TaylorScalar{T, N} <: Real
     value::NTuple{N, T}
+    function TaylorScalar{T, N}(value::NTuple{N, T}) where {T, N}
+        can_taylorize(T) || throw_cannot_taylorize(T)
+        new{T, N}(value)
+    end
 end
 
-TN = Union{TaylorScalar, Number}
-
-@inline TaylorScalar(xs::Vararg{T, N}) where {T, N} = TaylorScalar(xs)
+TaylorScalar(value::NTuple{N, T}) where {T, N} = TaylorScalar{T, N}(value)
+TaylorScalar(value::Vararg{T, N}) where {T, N} = TaylorScalar{T, N}(value)
 
 """
     TaylorScalar{T, N}(x::T) where {T, N}
 
 Construct a Taylor polynomial with zeroth order coefficient.
 """
-@generated function TaylorScalar{T, N}(x::T) where {T, N}
+@generated function TaylorScalar{T, N}(x::S) where {T, S <: Real, N}
     return quote
         $(Expr(:meta, :inline))
         TaylorScalar((T(x), $(zeros(T, N - 1)...)))
@@ -38,7 +55,7 @@ end
 
 Construct a Taylor polynomial with zeroth and first order coefficient, acting as a seed.
 """
-@generated function TaylorScalar{T, N}(x::T, d::T) where {T, N}
+@generated function TaylorScalar{T, N}(x::S, d::S) where {T, S <: Real, N}
     return quote
         $(Expr(:meta, :inline))
         TaylorScalar((T(x), T(d), $(zeros(T, N - 2)...)))
@@ -68,31 +85,11 @@ end
 end
 @inline primal(t::TaylorScalar) = extract_derivative(t, 1)
 
-@inline zero(::Type{TaylorScalar{T, N}}) where {T, N} = TaylorScalar{T, N}(zero(T))
-@inline one(::Type{TaylorScalar{T, N}}) where {T, N} = TaylorScalar{T, N}(one(T))
-@inline zero(::TaylorScalar{T, N}) where {T, N} = zero(TaylorScalar{T, N})
-@inline one(::TaylorScalar{T, N}) where {T, N} = one(TaylorScalar{T, N})
-
-adjoint(t::TaylorScalar) = t
-conj(t::TaylorScalar) = t
-
 function promote_rule(::Type{TaylorScalar{T, N}},
         ::Type{S}) where {T, S, N}
     TaylorScalar{promote_type(T, S), N}
 end
 
-# Number-like convention (I patched them after removing <: Number)
-
-convert(::Type{TaylorScalar{T, N}}, x::TaylorScalar{T, N}) where {T, N} = x
-function convert(::Type{TaylorScalar{T, N}}, x::S) where {T, S, N}
-    TaylorScalar{T, N}(convert(T, x))
-end
-for op in (:+, :-, :*, :/)
-    @eval @inline $op(a::TaylorScalar, b::Number) = $op(promote(a, b)...)
-    @eval @inline $op(a::Number, b::TaylorScalar) = $op(promote(a, b)...)
-end
-transpose(t::TaylorScalar) = t
-
 function Base.AbstractFloat(x::TaylorScalar{T, N}) where {T, N}
     TaylorScalar{Float64, N}(convert(NTuple{N, Float64}, x.value))
 end
diff --git a/test/primitive.jl b/test/primitive.jl
@@ -14,7 +14,8 @@ end
 
 @testset "Unary functions" begin
     some_number = 3.7
-    for f in (x -> exp(x^2), expm1, exp2, exp10, x -> sin(x^2), x -> cos(x^2), sinpi, cospi,
+    for f in (
+            x -> exp(x^2), expm1, exp2, exp10, x -> sin(x^2), x -> cos(x^2), sinpi, cospi,
             sqrt, cbrt,
             inv), order in (1, 4)
         fdm = central_fdm(12, order)