Allow expansion of matrices in diff and sum (#160)

* Allow expansion of matrices in diff and sum

* Update bases.jl

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "ContinuumArrays"
 uuid = "7ae1f121-cc2c-504b-ac30-9b923412ae5c"
-version = "0.15.1"
+version = "0.15.2"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/ContinuumArrays.jl

@@ -104,9 +104,9 @@ include("plotting.jl")
 # sum/dot
 ###
 
-sum_size(::Tuple{InfiniteCardinal{1}}, a, dims) = _sum(expand(a), dims)
+sum_size(::Tuple{InfiniteCardinal{1}, Vararg{Integer}}, a, dims) = _sum(expand(a), dims)
 dot_size(::InfiniteCardinal{1}, a, b) = dot(expand(a), expand(b))
-diff_size(::Tuple{InfiniteCardinal{1}}, a, dims) = diff(expand(a); dims=dims)
+diff_size(::Tuple{InfiniteCardinal{1}, Vararg{Integer}}, a, dims) = diff(expand(a); dims=dims)
 function copy(d::Dot{<:ExpansionLayout,<:ExpansionLayout,<:AbstractQuasiArray,<:AbstractQuasiArray})
     a,b = d.A,d.B
     P,c = basis(a),coefficients(a)

src/bases/bases.jl

@@ -567,6 +567,8 @@ end
 ####
 # sum
 ####
+sum_layout(::AbstractBasisLayout, Vm, dims) = error("Overload _sum(::$(typeof(Vm)), ::$(typeof(dims)))")
+
 function sum_layout(::SubBasisLayout, Vm, dims)
     dims == 1 || error("not implemented")
     sum(parent(Vm); dims=dims)[:,parentindices(Vm)[2]]
@@ -588,19 +590,21 @@ cumsum_layout(::ExpansionLayout, A, dims) = cumsum_layout(ApplyLayout{typeof(*)}
 # diff
 ###
 
-function diff_layout(::SubBasisLayout, Vm, dims::Integer)
+diff_layout(::AbstractBasisLayout, Vm, dims...) = error("Overload diff(::$(typeof(Vm)))")
+
+function diff_layout(::SubBasisLayout, Vm, dims::Integer=1)
     dims == 1 || error("not implemented")
     diff(parent(Vm); dims=dims)[:,parentindices(Vm)[2]]
 end
 
-function diff_layout(::WeightedBasisLayout{<:SubBasisLayout}, Vm, dims::Integer)
+function diff_layout(::WeightedBasisLayout{SubBasisLayout}, Vm, dims::Integer=1)
     dims == 1 || error("not implemented")
     w = weight(Vm)
     V = unweighted(Vm)
     view(diff(w .* parent(V)), parentindices(V)...)
 end
 
-function diff_layout(::MappedBasisLayouts, V, dims)
+function diff_layout(::MappedBasisLayouts, V, dims::Integer=1)
     kr = basismap(V)
     @assert kr isa AbstractAffineQuasiVector
     D = diff(demap(V); dims=dims)

test/test_chebyshev.jl

@@ -168,7 +168,11 @@ Base.:(==)(::FooBasis, ::FooBasis) = true
         @test dot(exp.(x), x) ≈ 2/ℯ
         @test diff(exp.(x))[0.1] ≈ exp(0.1)
 
+        @test sum(exp.(x .* (1:2)'); dims=1) ≈ [ℯ - 1/ℯ (ℯ^2 - 1/ℯ^2)/2]
+        @test diff(exp.(x .* (1:2)'); dims=1)[0.1,:] ≈ [exp(0.1), exp(0.2)]
+
         @test_throws ErrorException diff(wT[:,1:3])
+        @test_throws ErrorException sum(wT[:,1:3]; dims=1)
         @test_throws ErrorException cumsum(x)
     end