Support dims keywords for matrix coefficients, Mutate first argument

Project.toml

@@ -1,7 +1,7 @@
 name = "RecurrenceRelationships"
 uuid = "807425ed-42ea-44d6-a357-6771516d7b2c"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.1.1"
+version = "0.2.0-dev"
 
 [weakdeps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -14,6 +14,7 @@ RecurrenceRelationshipsLazyArraysExt = "LazyArrays"
 RecurrenceRelationshipsLinearAlgebraExt = "LinearAlgebra"
 
 [compat]
+DynamicPolynomials = "0.6.1"
 FillArrays = "1"
 LazyArrays = "1, 2"
 SpecialFunctions = "1, 2"

README.md

@@ -233,4 +233,22 @@ julia> using InfiniteArrays
 julia> @time clenshaw(inv.(1:n), Fill(2, ∞), Zeros(∞), Ones(∞), x)
   0.004574 seconds (5 allocations: 7.630 MiB)
 0.8396901361362448
+```
+
+Clenshaw also supports in-place usage including a matrix of coefficients. This can be used to evaluate 2D polynomials, eg.
+$$
+f(x,y) = ∑_{k=1}^{n} ∑_{j=1}^n T_k(x) U_j(y) (k+2j)
+$$
+can be evaluated via:
+```julia
+julia> m,n = 5,6;
+
+julia> coeffs = ((1:m) .+ 2(1:n)');
+
+julia> x,y = 0.1,0.2;
+
+julia> A, B, C = Fill(2,n), Zeros{Int}(n), Ones{Int}(n+1);
+
+julia> clenshaw(vec(clenshaw(coeffs, x; dims=1)), A, B, C, y)
+9.358401024
 ```

src/clenshaw.jl

@@ -1,52 +1,45 @@
 
 
 """
-clenshaw!(c, A, B, C, x)
+clenshaw!(f::AbstractVecOrMat, c::AbstractVecOrMat, A, B, C, x::Number)
 
 evaluates the orthogonal polynomial expansion with coefficients `c` at points `x`,
 where `A`, `B`, and `C` are `AbstractVector`s containing the recurrence coefficients
 as defined in DLMF,
-overwriting `x` with the results.
+overwriting `f` with the results.
 
-If `c` is a matrix this treats each column as a separate vector of coefficients, returning a vector
-if `x` is a number and a matrix if `x` is a vector.
+If `c` is a matrix this treats each row (if `size(f,2) == 1`) or column (if `size(f,1) == 1`)` as a separate vector of coefficients.
 """
-clenshaw!(c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector) =
-    clenshaw!(c, A, B, C, x, one(eltype(x)), x)
-
-clenshaw!(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::Number, f::AbstractVector) =
-    clenshaw!(c, A, B, C, x, one(eltype(x)), f)
-
-
-clenshaw!(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector, f::AbstractMatrix) =
-    clenshaw!(c, A, B, C, x, one(eltype(x)), f)
+clenshaw!(f::AbstractVecOrMat, c::AbstractVecOrMat, A::AbstractVector, B::AbstractVector, C::AbstractVector, x) =
+    clenshaw!(f, c, A, B, C, x, one(eltype(x)))
 
+clenshaw!(x::AbstractVecOrMat, c::AbstractVecOrMat, A::AbstractVector, B::AbstractVector, C::AbstractVector) =
+    clenshaw!(x, c, A, B, C, x)
 
 """
-clenshaw!(c, A, B, C, x, ϕ₀, f)
+clenshaw!(f, c, A, B, C, x, ϕ₀)
 
-evaluates the orthogonal polynomial expansion with coefficients `c` at points `x`,
+evaluates the orthogonal polynomial expansion with coefficients `c` at point `x`,
 where `A`, `B`, and `C` are `AbstractVector`s containing the recurrence coefficients
 as defined in DLMF and ϕ₀ is the zeroth polynomial,
 overwriting `f` with the results.
 """
-function clenshaw!(c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector, ϕ₀, f::AbstractVector)
+function clenshaw!(f::AbstractVector, c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector, ϕ₀)
     f .= ϕ₀ .* clenshaw.(Ref(c), Ref(A), Ref(B), Ref(C), x)
 end
 
-
-function clenshaw!(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::Number, ϕ₀::Number, f::AbstractVector)
-    size(c,2) == length(f) || throw(DimensionMismatch("coeffients size and output length must match"))
-    @inbounds for j in axes(c,2)
-        f[j] = ϕ₀ * clenshaw(view(c,:,j), A, B, C, x)
-    end
-    f
-end
-
-function clenshaw!(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector, ϕ₀, f::AbstractMatrix)
-    (size(x,1),size(c,2)) == size(f) || throw(DimensionMismatch("coeffients size and output length must match"))
-    @inbounds for j in axes(c,2)
-        clenshaw!(view(c,:,j), A, B, C, x, ϕ₀, view(f,:,j))
+function clenshaw!(f::AbstractVecOrMat, c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::Number, ϕ₀::Number)
+    m,n = size(c)
+    if (size(f,1),size(f,2)) == (1,n) # dims = 1
+        @inbounds for j in axes(c,2)
+            f[1,j] = ϕ₀ * clenshaw(view(c,:,j), A, B, C, x)
+        end
+    elseif (size(f,1),size(f,2)) == (m,1) # dims = 2
+        @inbounds for k in axes(c,1)
+            f[k,1] = ϕ₀ * clenshaw(view(c,k,:), A, B, C, x)
+        end
+    else
+        throw(DimensionMismatch("coeffients size and output length must match"))
     end
     f
 end
@@ -58,7 +51,7 @@
 
 function check_clenshaw_recurrences(N, A, B, C)
     if length(A) < N || length(B) < N || length(C) < N+1
-        throw(ArgumentError("A, B must contain at least $N entries and C must contain at least $(N+1) entrie"))
+        throw(ArgumentError("A, B must contain at least $N entries and C must contain at least $(N+1) entries"))
     end
 end
 
@@ -92,40 +85,32 @@
     bn1
 end
 
-
 clenshaw(c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector) =
-    clenshaw!(c, A, B, C, copy(x))
-
-function clenshaw(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::Number)
-    T = promote_type(eltype(c),eltype(A),eltype(B),eltype(C),typeof(x))
-    clenshaw!(c, A, B, C, x, Vector{T}(undef, size(c,2)))
-end
+    clenshaw!(copy(x), c, A, B, C)
 
-function clenshaw(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractVector)
+function clenshaw(c::AbstractMatrix, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::Number; dims)
     T = promote_type(eltype(c),eltype(A),eltype(B),eltype(C),typeof(x))
-    clenshaw!(c, A, B, C, x, Matrix{T}(undef, size(x,1), size(c,2)))
+    if dims == 1
+        clenshaw!(Matrix{T}(undef, 1, size(c,2)), c, A, B, C, x)
+    elseif dims == 2
+        clenshaw!(Matrix{T}(undef, size(c,1), 1), c, A, B, C, x)
+    else
+        throw(ArgumentError("dims must be 1 or 2"))
+    end
 end
 
 ###
 # Chebyshev T special cases
 ###
 
-"""
-   clenshaw!(c, x)
-
-evaluates the first-kind Chebyshev (T) expansion with coefficients `c` at points `x`,
-overwriting `x` with the results.
-"""
-clenshaw!(c::AbstractVector, x::AbstractVector) = clenshaw!(c, x, x)
-
 
 """
-   clenshaw!(c, x, f)
+   clenshaw!(f, c, x)
 
 evaluates the first-kind Chebyshev (T) expansion with coefficients `c` at points `x`,
 overwriting `f` with the results.
 """
-function clenshaw!(c::AbstractVector, x::AbstractVector, f::AbstractVector)
+function clenshaw!(f::AbstractVector, c::AbstractVector, x::AbstractVector)
     f .= clenshaw.(Ref(c), x)
 end
 
@@ -153,23 +138,34 @@
     end
 end
 
-function clenshaw!(c::AbstractMatrix, x::Number, f::AbstractVector)
-    size(c,2) == length(f) || throw(DimensionMismatch("coeffients size and output length must match"))
-    @inbounds for j in axes(c,2)
-        f[j] = clenshaw(view(c,:,j), x)
+function clenshaw!(f::AbstractVecOrMat, c::AbstractMatrix, x::Number)
+    m,n = size(c)
+    if (size(f,1),size(f,2)) == (1,n) # dims = 1
+        @inbounds for j in axes(c,2)
+            f[1,j] = clenshaw(view(c,:,j), x)
+        end
+    elseif (size(f,1),size(f,2)) == (m,1) # dims = 2
+        @inbounds for k in axes(c,1)
+            f[k,1] = clenshaw(view(c,k,:), x)
+        end
+    else
+        throw(DimensionMismatch("coeffients size and output length must match"))
     end
     f
 end
 
-function clenshaw!(c::AbstractMatrix, x::AbstractVector, f::AbstractMatrix)
-    (size(x,1),size(c,2)) == size(f) || throw(DimensionMismatch("coeffients size and output length must match"))
-    @inbounds for j in axes(c,2)
-        clenshaw!(view(c,:,j), x, view(f,:,j))
+clenshaw!(x, c) = clenshaw!(x, c, x)
+
+function clenshaw(c::AbstractMatrix, x::Number; dims)
+    T = polynomialtype(eltype(c),typeof(x))
+    if dims == 1
+        clenshaw!(Matrix{T}(undef, 1, size(c,2)), c, x)
+    elseif dims == 2
+        clenshaw!(Matrix{T}(undef, size(c,1), 1), c, x)
+    else
+        throw(ArgumentError("dims must be 1 or 2"))
     end
-    f
 end
 
-clenshaw(c::AbstractVector, x::AbstractVector) = clenshaw!(c, copy(x))
-clenshaw(c::AbstractMatrix, x::Number) = clenshaw!(c, x, Vector{promote_type(eltype(c),typeof(x))}(undef, size(c,2)))
-clenshaw(c::AbstractMatrix, x::AbstractVector) = clenshaw!(c, x, Matrix{promote_type(eltype(c),eltype(x))}(undef, size(x,1), size(c,2)))
 
+clenshaw(c::AbstractVector, x::AbstractVector) = clenshaw!(copy(x), c)

test/runtests.jl

@@ -2,10 +2,6 @@ using RecurrenceRelationships, LinearAlgebra, Test
 using FillArrays, LazyArrays
 using DynamicPolynomials
 
-if !isdefined(LinearAlgebra, :NoPivot)
-    const NoPivot = Val{false}
-end
-
 
 @testset "forward" begin
     @testset "Chebyshev U" begin
@@ -78,7 +74,7 @@ end
         c = [1,2,3]
         @test @inferred(clenshaw(c,1)) ≡ 1 + 2 + 3
         @test @inferred(clenshaw(c,0)) ≡ 1 + 0 - 3
-        @test @inferred(clenshaw(c,[-1,0,1])) == clenshaw!(c,[-1,0,1]) == [2,-2,6]
+        @test @inferred(clenshaw(c,[-1,0,1])) == clenshaw!([-1,0,1],c) == [2,-2,6]
         @test @inferred(clenshaw(c,0.1)) == 1 + 2*0.1 + 3*cos(2acos(0.1))
         @test clenshaw(c,[-1,0,1]) isa Vector{Int}
 
@@ -91,27 +87,27 @@ end
             @test @inferred(clenshaw(elty[],1)) ≡ zero(elty)
 
             x = elty[1,0,0.1]
-            @test @inferred(clenshaw(c,x)) ≈ @inferred(clenshaw!(c,copy(x))) ≈
-                @inferred(clenshaw!(c,x,similar(x))) ≈
-                @inferred(clenshaw(cf,x)) ≈ @inferred(clenshaw!(cf,copy(x))) ≈
-                @inferred(clenshaw!(cf,x,similar(x))) ≈ elty[6,-2,-1.74]
+            @test @inferred(clenshaw(c,x)) ≈ @inferred(clenshaw!(copy(x),c)) ≈
+                @inferred(clenshaw!(similar(x),c,x)) ≈
+                @inferred(clenshaw(cf,x)) ≈ @inferred(clenshaw!(copy(x),cf)) ≈
+                @inferred(clenshaw!(similar(x),cf,x)) ≈ elty[6,-2,-1.74]
 
             @testset "Strided" begin
                 cv = view(cf,:)
                 xv = view(x,:)
-                @test clenshaw!(cv, xv, similar(xv)) == clenshaw!(cf,x,similar(x))
+                @test clenshaw!(similar(xv), cv, xv) == clenshaw!(similar(x), cf, x)
 
                 cv2 = view(cf,1:2:3)
-                @test clenshaw!(cv2, xv, similar(xv)) == clenshaw([1,3], x)
+                @test clenshaw!(similar(xv), cv2, xv) == clenshaw([1,3], x)
 
                 # modifies x and xv
-                @test clenshaw!(cv2, xv) == xv == x == clenshaw([1,3], elty[1,0,0.1])
+                @test clenshaw!(xv, cv2) == xv == x == clenshaw([1,3], elty[1,0,0.1])
             end
         end
         @testset "matrix coefficients" begin
             c = [1 2; 3 4; 5 6]
-            @test clenshaw(c,0.1) ≈ [clenshaw(c[:,1],0.1), clenshaw(c[:,2],0.1)]
-            @test clenshaw(c,[0.1,0.2]) ≈ [clenshaw(c[:,1], 0.1) clenshaw(c[:,2], 0.1); clenshaw(c[:,1], 0.2) clenshaw(c[:,2], 0.2)]
+            @test clenshaw(c,0.1; dims=1) ≈ [clenshaw(c[:,1],0.1), clenshaw(c[:,2],0.1)]'
+            @test clenshaw(c,0.1; dims=2) ≈ [clenshaw(c[1,:],0.1); clenshaw(c[2,:],0.1); clenshaw(c[3,:],0.1) ;;]
         end
     end
 
@@ -124,26 +120,28 @@ end
 
         @testset "matrix coefficients" begin
             c = [1 2; 3 4; 5 6]
-            @test clenshaw(c,A,B,C,0.1) ≈ [clenshaw(c[:,1],A,B,C,0.1), clenshaw(c[:,2],A,B,C,0.1)]
-            @test clenshaw(c,A,B,C,[0.1,0.2]) ≈ [clenshaw(c[:,1], A,B,C,0.1) clenshaw(c[:,2], A,B,C,0.1); clenshaw(c[:,1], A,B,C,0.2) clenshaw(c[:,2], A,B,C,0.2)]
+            @test clenshaw(c,A,B,C,0.1; dims=1) ≈ [clenshaw(c[:,1],A,B,C,0.1), clenshaw(c[:,2],A,B,C,0.1)]'
+            @test clenshaw(c,A,B,C,0.1; dims=2) ≈ [clenshaw(c[1,:],A,B,C,0.1); clenshaw(c[2,:],A,B,C,0.1); clenshaw(c[3,:],A,B,C,0.1) ;;]
         end
+
+        @test_throws ArgumentError clenshaw([1,2,3], A, B, Ones{Int}(2), 0.1)
     end
 
     @testset "Chebyshev-as-general" begin
         c, A, B, C = [1,2,3], [1,2,2], fill(0,3), fill(1,4)
         cf, Af, Bf, Cf = float(c), float(A), float(B), float(C)
         @test @inferred(clenshaw(c, A, B, C, 1)) ≡ 6
         @test @inferred(clenshaw(c, A, B, C, 0.1)) ≡ -1.74
-        @test @inferred(clenshaw([1,2,3], A, B, C, [-1,0,1])) == clenshaw!([1,2,3],A, B, C, [-1,0,1]) == [2,-2,6]
+        @test @inferred(clenshaw([1,2,3], A, B, C, [-1,0,1])) == clenshaw!([-1,0,1], [1,2,3],A, B, C) == [2,-2,6]
         @test clenshaw(c, A, B, C, [-1,0,1]) isa Vector{Int}
         @test @inferred(clenshaw(Float64[], A, B, C, 1)) ≡ 0.0
 
         x = [1,0,0.1]
-        @test @inferred(clenshaw(c, A, B, C, x)) ≈ @inferred(clenshaw!(c, A, B, C, copy(x))) ≈
-            @inferred(clenshaw!(c, A, B, C, x, one.(x), similar(x))) ≈
-            @inferred(clenshaw!(cf, Af, Bf, Cf, x, one.(x),similar(x))) ≈
+        @test @inferred(clenshaw(c, A, B, C, x)) ≈ @inferred(clenshaw!(copy(x), c, A, B, C)) ≈
+            @inferred(clenshaw!(similar(x), c, A, B, C, x, one.(x))) ≈
+            @inferred(clenshaw!(similar(x), cf, Af, Bf, Cf, x, one.(x))) ≈
             @inferred(clenshaw([1.,2,3], A, B, C, x)) ≈
-            @inferred(clenshaw!([1.,2,3], A, B, C, copy(x))) ≈ [6,-2,-1.74]
+            @inferred(clenshaw!(copy(x), [1.,2,3], A, B, C)) ≈ [6,-2,-1.74]
     end
 
     @testset "Legendre" begin
@@ -162,7 +160,7 @@ end
                 c = [zeros(j-1); 1]
                 @test clenshaw(c, A, B, C, 1) ≈ v_1[j] # Julia code
                 @test clenshaw(c, A, B, C, 0.1) ≈  v_f[j] # Julia code
-                @test clenshaw!(c, A, B, C, [1.0,0.1], [1.0,1.0], [0.0,0.0])  ≈ [v_1[j],v_f[j]] # libfasttransforms
+                @test clenshaw!([0.0,0.0], c, A, B, C, [1.0,0.1], [1.0,1.0])  ≈ [v_1[j],v_f[j]] # libfasttransforms
             end
         end
     end
@@ -189,6 +187,27 @@ end
         @test forwardrecurrence(Vcat(1, Fill(2, n-1)), Zeros(n), Ones(n), x) ≈ cos.((0:n-1) .* θ)
         @test clenshaw((1:n), Vcat(1, Fill(2, n-1)), Zeros(n), Ones(n+1), x) ≈ sum(cos(k * θ) * (k+1) for k = 0:n-1)
     end
+
+    @testset "2D" begin
+        # cheb U recurrence
+        m,n = 5,6
+        coeffs = ((1:m) .+ 2(1:n)')
+        x,y = 0.1,0.2
+        A, B, C = Fill(2,n), Zeros{Int}(n), Ones{Int}(n+1)
+
+
+        A_T, B_T, C_T = [1; fill(2,m-1)], fill(0,m), fill(1,m+1)
+        @test clenshaw(vec(clenshaw(coeffs, x; dims=1)), A, B, C, y) ≈ clenshaw(vec(clenshaw(coeffs, A, B, C, y; dims=2)), x) ≈
+                only(clenshaw!([0.0], clenshaw!(Matrix{Float64}(undef,1,n), coeffs, x), A, B, C, y)) ≈ 
+                only(clenshaw!([0.0], clenshaw!(Matrix{Float64}(undef,m,1), coeffs, A, B, C, y), x)) ≈
+                forwardrecurrence(A_T, B_T, C_T, x)'coeffs*forwardrecurrence(A, B, C, y)
+
+        @test_throws DimensionMismatch clenshaw!(Matrix{Float64}(undef,2,n), coeffs, x)
+        @test_throws DimensionMismatch clenshaw!(Matrix{Float64}(undef,2,n), coeffs, A_T, B_T, C_T, x)
+
+        @test_throws ArgumentError clenshaw(coeffs, x; dims=3)
+        @test_throws ArgumentError clenshaw(coeffs, A_T, B_T, C_T, x; dims=3)
+    end
 end
 
 @testset "olver" begin