Further allocation reduction in power computations

ext/TaylorSeriesIAExt.jl

@@ -41,7 +41,7 @@ function ^(a::Taylor1{Interval{T}}, r::S) where {T<:Real, S<:Real}
     c_order = l0 == 0 ? a.order : min(a.order, trunc(Int,r*a.order))
     c = Taylor1(zero(aux), c_order)
     for k = 0:c_order
-        TS.pow!(c, aa, r, k)
+        TS.pow!(c, aa, c[0], r, k)
     end
 
     return c
@@ -66,7 +66,7 @@ function ^(a::TaylorN{Interval{T}}, r::S) where {T<:Real, S<:Real}
 
     c = TaylorN( a0r, a.order)
     for ord in 1:a.order
-        TS.pow!(c, aa, r, ord)
+        TS.pow!(c, aa, c, r, ord)
     end
 
     return c

src/arithmetic.jl

@@ -162,7 +162,9 @@ for T in (:Taylor1, :HomogeneousPolynomial, :TaylorN)
 end
 
 for T in (:Taylor1, :TaylorN)
-    @eval zero(a::$T) = $T(zero.(a.coeffs))
+    @eval function zero(a::$T)
+        return $T(zero.(a.coeffs))
+    end
     @eval function one(a::$T)
         b = zero(a)
         b[0] = one(b[0])
@@ -572,6 +574,42 @@ for T in (:Taylor1, :TaylorN)
     end
 end
 
+# in-place product: `a` <- `a*b`
+# this method computes the product `a*b` and saves it back into `a`
+# assumes `a` and `b` are of same order
+function mul!(a::TaylorN{T}, b::TaylorN{T}) where {T<:Number}
+    @inbounds for k in reverse(eachindex(a))
+        mul!(a, a, b[0][1], k)
+        for l in 1:k
+            mul!(a[k], a[k-l], b[l])
+        end
+    end
+    return nothing
+end
+function mul!(a::Taylor1{T}, b::Taylor1{T}) where {T<:Number}
+    @inbounds for k in reverse(eachindex(a))
+        # a[k] <- a[k]*b[0]
+        mul!(a, a, b[0], k)
+        for l in 1:k
+            # a[k] <- a[k] + a[k-l] * b[l]
+            a[k] += a[k-l] * b[l]
+        end
+    end
+    return nothing
+end
+function mul!(a::Taylor1{TaylorN{T}}, b::Taylor1{TaylorN{T}}) where {T<:NumberNotSeries}
+    @inbounds for k in reverse(eachindex(a))
+        mul!(a, a, b[0], k)
+        for l in 1:k
+            # a[k] += a[k-l] * b[l]
+            for m in eachindex(a[k])
+                mul!(a[k], a[k-l], b[l], m)
+            end
+        end
+    end
+    return nothing
+end
+
 function mul!(res::Taylor1{TaylorN{T}}, a::Taylor1{TaylorN{T}}, b::Taylor1{TaylorN{T}},
         ordT::Int) where {T<:NumberNotSeries}
     # Sanity

src/auxiliary.jl

@@ -18,9 +18,8 @@ function resize_coeffs1!(coeffs::Array{T,1}, order::Int) where {T<:Number}
     lencoef = length(coeffs)
     resize!(coeffs, order+1)
     if order > lencoef-1
-        z = zero(coeffs[1])
         @simd for ord in lencoef+1:order+1
-            @inbounds coeffs[ord] = z
+            @inbounds coeffs[ord] = zero(coeffs[1])
         end
     end
     return nothing
@@ -39,9 +38,8 @@ function resize_coeffsHP!(coeffs::Array{T,1}, order::Int) where {T<:Number}
     @assert order ≤ get_order() && lencoef ≤ num_coeffs
     num_coeffs == lencoef && return nothing
     resize!(coeffs, num_coeffs)
-    z = zero(coeffs[1])
     @simd for ord in lencoef+1:num_coeffs
-        @inbounds coeffs[ord] = z
+        @inbounds coeffs[ord] = zero(coeffs[1])
     end
     return nothing
 end

src/constructors.jl

@@ -156,7 +156,7 @@ struct TaylorN{T<:Number} <: AbstractSeries{T}
     order   :: Int
 
     function TaylorN{T}(v::Array{HomogeneousPolynomial{T},1}, order::Int) where T<:Number
-        coeffs = zeros(HomogeneousPolynomial{T}, order)
+        coeffs = isempty(v) ? zeros(HomogeneousPolynomial{T}, order) : zeros(v[1], order)
         @inbounds for i in eachindex(v)
             ord = v[i].order
             if ord ≤ order

src/dictmutfunct.jl

@@ -60,7 +60,7 @@ const _dict_binary_ops = Dict(
     :- => (:subst!, (:_res, :_arg1, :_arg2, :_k), :(_res = _arg1 - _arg2)),
     :* => (:mul!, (:_res, :_arg1, :_arg2, :_k), :(_res = _arg1 * _arg2)),
     :/ => (:div!, (:_res, :_arg1, :_arg2, :_k), :(_res = _arg1 / _arg2)),
-    :^ => (:pow!, (:_res, :_arg1, :_arg2, :_k), :(_res = _arg1 ^ float(_arg2))),
+    :^ => (:pow!, (:_res, :_arg1, :_arg2, :_arg3, :_k), :(_res = _arg1 ^ float(_arg3))),
 );
 
 """
@@ -196,4 +196,3 @@ internal mutating functions.
 Evaluating the entries generates symbols that represent
 the actual calls to the internal mutating functions.
 """ _dict_binary_calls
-

src/evaluate.jl

@@ -413,7 +413,7 @@ function _evaluate!(suma::TaylorN{T}, a::HomogeneousPolynomial{T}, ind::Int,
         else
             for ordQ in eachindex(val)
                 zero!(vv, ordQ)
-                pow!(vv, val, vpow, ordQ)
+                pow!(vv, val, vv, vpow, ordQ)
             end
         end
         for ordQ in eachindex(suma)

src/functions.jl

@@ -303,6 +303,13 @@ end
     return nothing
 end
 
+@inline function zero!(a::TaylorN{Taylor1{T}}) where {T<:NumberNotSeries}
+    for k in eachindex(a)
+        zero!(a, k)
+    end
+    return nothing
+end
+
 @inline function one!(c::Taylor1{TaylorN{T}}, k::Int) where {T<:NumberNotSeries}
     zero!(c, k)
     if k == 0
@@ -316,6 +323,13 @@ end
     return nothing
 end
 
+@inline function identity!(c::HomogeneousPolynomial{Taylor1{T}}, a::HomogeneousPolynomial{Taylor1{T}}, k::Int) where {T<:NumberNotSeries}
+    @inbounds for l in eachindex(c[k])
+        identity!(c[k], a[k], l)
+    end
+    return nothing
+end
+
 for T in (:Taylor1, :TaylorN)
     @eval begin
         @inline function identity!(c::$T{T}, a::$T{T}, k::Int) where {T<:NumberNotSeries}

src/power.jl

@@ -58,37 +58,98 @@ end
 
 ^(a::Taylor1{TaylorN{T}}, r::Rational) where {T<:NumberNotSeries} = a^(r.num/r.den)
 
-
-
-# power_by_squaring; slightly modified from base/intfuncs.jl
-# Licensed under MIT "Expat"
-for T in (:Taylor1, :HomogeneousPolynomial, :TaylorN)
-    @eval function power_by_squaring(x::$T, p::Integer)
-        if p == 1
-            return copy(x)
-        elseif p == 0
-            return one(x)
-        elseif p == 2
-            return square(x)
-        end
+# in-place form of power_by_squaring
+# this method assumes `y`, `x` and `aux` are of same order
+for T in (:Taylor1, :TaylorN)
+    @eval function power_by_squaring!(y::$T{T}, x::$T{T}, aux::$T{T},
+            p::Integer) where {T<:NumberNotSeries}
         t = trailing_zeros(p) + 1
         p >>= t
+        # aux = x
+        for k in eachindex(aux)
+            identity!(aux, x, k)
+        end
         while (t -= 1) > 0
-            x = square(x)
+            # aux = square(aux)
+            for k in reverse(eachindex(aux))
+                sqr!(aux, k)
+            end
+        end
+        # y = aux
+        for k in eachindex(y)
+            identity!(y, aux, k)
         end
-        y = x
         while p > 0
             t = trailing_zeros(p) + 1
             p >>= t
             while (t -= 1) ≥ 0
+                # aux = square(aux)
+                for k in reverse(eachindex(aux))
+                    sqr!(aux, k)
+                end
+            end
+            # y = y * aux
+            mul!(y, aux)
+        end
+        return nothing
+    end
+end
+
+
+# power_by_squaring; slightly modified from base/intfuncs.jl
+# Licensed under MIT "Expat"
+for T in (:Taylor1, :HomogeneousPolynomial, :TaylorN)
+    @eval power_by_squaring(x::$T, p::Integer) = power_by_squaring(x, Val(p))
+    @eval power_by_squaring(x::$T, ::Val{0}) = one(x)
+    @eval power_by_squaring(x::$T, ::Val{1}) = copy(x)
+    @eval power_by_squaring(x::$T, ::Val{2}) = square(x)
+    @eval function power_by_squaring(x::$T, ::Val{P}) where P
+        p = P # copy static parameter `P` into local variable `p`
+        if $T != HomogeneousPolynomial
+            y = zero(x)
+            aux = zero(x)
+            power_by_squaring!(y, x, aux, p)
+        else
+            t = trailing_zeros(p) + 1
+            p >>= t
+            while (t -= 1) > 0
                 x = square(x)
             end
-            y *= x
+            y = x
+            while p > 0
+                t = trailing_zeros(p) + 1
+                p >>= t
+                while (t -= 1) ≥ 0
+                    x = square(x)
+                end
+                y *= x
+            end
         end
         return y
     end
 end
 
+power_by_squaring(x::TaylorN{Taylor1{T}}, ::Val{0}) where {T<:NumberNotSeries} = one(x)
+power_by_squaring(x::TaylorN{Taylor1{T}}, ::Val{1}) where {T<:NumberNotSeries} = copy(x)
+power_by_squaring(x::TaylorN{Taylor1{T}}, ::Val{2}) where {T<:NumberNotSeries} = square(x)
+function power_by_squaring(x::TaylorN{Taylor1{T}}, ::Val{P}) where {P, T<:NumberNotSeries}
+    p = P # copy static parameter `P` into local variable `p`
+    t = trailing_zeros(p) + 1
+    p >>= t
+    while (t -= 1) > 0
+        x = square(x)
+    end
+    y = x
+    while p > 0
+        t = trailing_zeros(p) + 1
+        p >>= t
+        while (t -= 1) ≥ 0
+            x = square(x)
+        end
+        y *= x
+    end
+    return y
+end
 
 ## Real power ##
 function ^(a::Taylor1{T}, r::S) where {T<:Number, S<:Real}
@@ -107,8 +168,9 @@ function ^(a::Taylor1{T}, r::S) where {T<:Number, S<:Real}
 
     c_order = l0 == 0 ? a.order : min(a.order, trunc(Int,r*a.order))
     c = Taylor1(zero(aux), c_order)
+    aux0 = deepcopy(c[0])
     for k in eachindex(c)
-        pow!(c, aa, r, k)
+        pow!(c, aa, aux0, r, k)
     end
 
     return c
@@ -132,8 +194,9 @@ function ^(a::TaylorN, r::S) where {S<:Real}
         in order to expand `^` around 0."""))
 
     c = TaylorN( zero(aux), a.order)
+    aux = deepcopy(c)
     for ord in eachindex(a)
-        pow!(c, aa, r, ord)
+        pow!(c, aa, aux, r, ord)
     end
 
     return c
@@ -158,8 +221,9 @@ function ^(a::Taylor1{TaylorN{T}}, r::S) where {T<:NumberNotSeries, S<:Real}
 
     c_order = l0 == 0 ? a.order : min(a.order, trunc(Int,r*a.order))
     c = Taylor1(zero(aux), c_order)
+    aux0 = deepcopy(c[0])
     for k in eachindex(c)
-        pow!(c, aa, r, k)
+        pow!(c, aa, aux0, r, k)
     end
 
     return c
@@ -184,7 +248,7 @@ exploits `k_0`, the order of the first non-zero coefficient of `a`.
 
 """ pow!
 
-@inline function pow!(c::Taylor1{T}, a::Taylor1{T}, r::S, k::Int) where
+@inline function pow!(c::Taylor1{T}, a::Taylor1{T}, ::T, r::S, k::Int) where
         {T<:Number, S<:Real}
 
     if r == 0
@@ -233,7 +297,7 @@ exploits `k_0`, the order of the first non-zero coefficient of `a`.
     return nothing
 end
 
-@inline function pow!(c::TaylorN{T}, a::TaylorN{T}, r::S, k::Int) where
+@inline function pow!(c::TaylorN{T}, a::TaylorN{T}, ::TaylorN{T}, r::S, k::Int) where
         {T<:NumberNotSeriesN, S<:Real}
 
     if r == 0
@@ -266,7 +330,7 @@ end
     return nothing
 end
 
-@inline function pow!(res::Taylor1{TaylorN{T}}, a::Taylor1{TaylorN{T}}, r::S,
+@inline function pow!(res::Taylor1{TaylorN{T}}, a::Taylor1{TaylorN{T}}, aux::TaylorN{T}, r::S,
         ordT::Int) where {T<:NumberNotSeries, S<:Real}
 
     if r == 0
@@ -300,7 +364,7 @@ end
     if ordT == lnull
         if isinteger(r)
             # TODO: get rid of allocations here
-            res[ordT] = a[l0]^round(Int,r) # uses power_by_squaring
+            power_by_squaring!(res[ordT], a[l0], aux, round(Int,r))
             return nothing
         end
 
@@ -310,7 +374,7 @@ end
             in order to expand `^` around 0."""))
 
         for ordQ in eachindex(a[l0])
-            pow!(res[ordT], a[l0], r, ordQ)
+            pow!(res[ordT], a[l0], aux, r, ordQ)
         end
 
         return nothing
@@ -339,7 +403,7 @@ Return `a^2`; see [`TaylorSeries.sqr!`](@ref).
 
 for T in (:Taylor1, :TaylorN)
     @eval function square(a::$T)
-        c = $T( zero(constant_term(a)), a.order)
+        c = zero(a)
         for k in eachindex(a)
             sqr!(c, a, k)
         end
@@ -447,6 +511,37 @@ for T = (:Taylor1, :TaylorN)
 
             return nothing
         end
+
+        # in-place squaring: given `c`, compute expansion of `c^2` and save back into `c`
+        @inline function sqr!(c::$T{T}, k::Int) where {T<:NumberNotSeries}
+            if k == 0
+                sqr_orderzero!(c, c)
+                return nothing
+            end
+
+            # Recursion formula
+            kodd = k%2
+            kend = (k - 2 + kodd) >> 1
+            if $T == Taylor1
+                (kend ≥ 0) && ( @inbounds c[k] = c[0] * c[k] )
+                @inbounds for i = 1:kend
+                    c[k] += c[i] * c[k-i]
+                end
+                @inbounds c[k] = 2 * c[k]
+                (kodd == 0) && ( @inbounds c[k] = c[k >> 1]^2 )
+            else
+                (kend ≥ 0) && ( @inbounds mul!(c, c[0][1], c, k) )
+                @inbounds for i = 1:kend
+                    mul!(c[k], c[i], c[k-i])
+                end
+                @inbounds mul!(c, 2, c, k)
+                if (kodd == 0)
+                    accsqr!(c[k], c[k >> 1])
+                end
+            end
+
+            return nothing
+        end
     end
 end