Support higher order diff (#117)

* Support higher order diff

* add tests

* Update test_calculus.jl

* Update calculus.jl

* Generalise some routines

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "QuasiArrays"
 uuid = "c4ea9172-b204-11e9-377d-29865faadc5c"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.11.9"
+version = "0.12"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/QuasiArrays.jl

@@ -3,7 +3,7 @@ using Base, LinearAlgebra, LazyArrays, ArrayLayouts, DomainSets, FillArrays, Sta
 import Base: getindex, size, axes, axes1, length, ==, isequal, iterate, CartesianIndices, LinearIndices,
                 Indices, IndexStyle, getindex, setindex!, parent, vec, convert, similar, copy, copyto!, zero,
                 map, eachindex, eltype, first, last, firstindex, lastindex, in, reshape, permutedims, all,
-                isreal, iszero, isempty, empty, isapprox, fill!, getproperty, showarg
+                isreal, iszero, isone, isempty, empty, isapprox, fill!, getproperty, showarg
 import Base: @_inline_meta, DimOrInd, OneTo, @_propagate_inbounds_meta, @_noinline_meta,
                 DimsInteger, error_if_canonical_getindex, @propagate_inbounds, _return_type,
                 safe_tail, front, tail, _getindex, _maybe_reshape, index_ndims, _unsafe_getindex,

src/calculus.jl

@@ -51,40 +51,37 @@ cumsum_size(::NTuple{N,Integer}, A, dims) where N = error("Not implemented")
 # diff
 ####
 
-@inline diff(a::AbstractQuasiArray; dims::Integer=1) = diff_layout(MemoryLayout(a), a, dims)
-function diff_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVector, dims...)
+@inline diff(a::AbstractQuasiArray, order...; dims::Integer=1) = diff_layout(MemoryLayout(a), a, order...; dims)
+function diff_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVecOrMat, order...; dims=1)
     a = arguments(LAY, V)
-    *(diff(a[1]), tail(a)...)
+    dims == 1 || throw(ArgumentError("cannot differentiate a vector along dimension $dims"))
+    *(diff(a[1], order...), tail(a)...)
 end
 
-function diff_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiMatrix, dims=1)
-    a = arguments(LAY, V)
-    @assert dims == 1 #for type stability, for now
-    # if dims == 1
-        *(diff(a[1]), tail(a)...)
-    # else
-    #     *(front(a)..., diff(a[end]; dims=dims))
-    # end
+diff_layout(::MemoryLayout, A, order...; dims...) = diff_size(size(A), A, order...; dims...)
+diff_size(sz, a; dims...) = error("diff not implemented for $(typeof(a))")
+function diff_size(sz, a, order; dims...)
+    order < 0 && throw(ArgumentError("order must be non-negative"))
+    order == 0 && return a
+    isone(order) ? diff(a) : diff(diff(a), order-1)
 end
 
-diff_layout(::MemoryLayout, A, dims...) = diff_size(size(A), A, dims...)
-diff_size(sz, a, dims...) = error("diff not implemented for $(typeof(a))")
-
 diff(x::Inclusion; dims::Integer=1) = ones(eltype(x), diffaxes(x))
-diff(c::AbstractQuasiFill{<:Any,1}; dims::Integer=1) =  zeros(eltype(c), diffaxes(axes(c,1)))
-function diff(c::AbstractQuasiFill{<:Any,2}; dims::Integer=1)
+diff(x::Inclusion, order::Int; dims::Integer=1) = fill(ifelse(isone(order), one(eltype(x)), zero(eltype(x))), diffaxes(x,order))
+diff(c::AbstractQuasiFill{<:Any,1}, order...; dims::Integer=1) =  zeros(eltype(c), diffaxes(axes(c,1),order...))
+function diff(c::AbstractQuasiFill{<:Any,2}, order...; dims::Integer=1)
     a,b = axes(c)
     if dims == 1
-        zeros(eltype(c), diffaxes(a), b)
+        zeros(eltype(c), diffaxes(a, order...), b)
     else
-        zeros(eltype(c), a, diffaxes(b))
+        zeros(eltype(c), a, diffaxes(b, order...))
     end
 end
 
 
-diffaxes(a::Inclusion{<:Any,<:AbstractVector}) = Inclusion(a.domain[1:end-1])
-diffaxes(a::OneTo) = oneto(length(a)-1)
-diffaxes(a) = a # default is differentiation does not change axes
+diffaxes(a::Inclusion{<:Any,<:AbstractVector}, order=1) = Inclusion(a.domain[1:end-order])
+diffaxes(a::OneTo, order=1) = oneto(length(a)-order)
+diffaxes(a, order...) = a # default is differentiation does not change axes
 
 diff(b::QuasiVector; dims::Integer=1) = QuasiVector(diff(b.parent) ./ diff(b.axes[1]), (diffaxes(axes(b,1)),))
 function diff(A::QuasiMatrix; dims::Integer=1)

src/quasibroadcast.jl

@@ -186,13 +186,13 @@ LazyArrays._broadcast_mul_mul((A,b)::Tuple{AbstractQuasiMatrix,AbstractQuasiVect
 # support (A .* B) * y
 _broadcasted_mul(a::Tuple{Number,Vararg{Any}}, b::AbstractQuasiVector) = (first(a)*sum(b), _broadcasted_mul(tail(a), b)...)
 _broadcasted_mul(a::Tuple{Number,Vararg{Any}}, B::AbstractQuasiMatrix) = (first(a)*sum(B; dims=1), _broadcasted_mul(tail(a), B)...)
-_broadcasted_mul(a::Tuple{AbstractQuasiVector,Vararg{Any}}, b::AbstractQuasiVector) = (first(a)*sum(b), _broadcasted_mul(tail(a), b)...)
-_broadcasted_mul(a::Tuple{AbstractQuasiVector,Vararg{Any}}, B::AbstractQuasiMatrix) = (first(a)*sum(B; dims=1), _broadcasted_mul(tail(a), B)...)
-_broadcasted_mul(A::Tuple{AbstractQuasiMatrix,Vararg{Any}}, b::AbstractQuasiVector) = (axes(first(A),2) == Base.OneTo(1) ? first(A)*sum(b) : (first(A)*b), _broadcasted_mul(tail(A), b)...)
-_broadcasted_mul(A::Tuple{AbstractQuasiMatrix,Vararg{Any}}, B::AbstractQuasiMatrix) = (axes(first(A),2) == Base.OneTo(1) ? first(A)*sum(B; dims=1) : (first(A)*B), _broadcasted_mul(tail(A), B)...)
+_broadcasted_mul(a::Tuple{AbstractQuasiVector,Vararg{Any}}, b::AbstractQuasiOrVector) = (first(a)*sum(b), _broadcasted_mul(tail(a), b)...)
+_broadcasted_mul(a::Tuple{AbstractQuasiVector,Vararg{Any}}, B::AbstractQuasiOrMatrix) = (first(a)*sum(B; dims=1), _broadcasted_mul(tail(a), B)...)
+_broadcasted_mul(A::Tuple{AbstractQuasiMatrix,Vararg{Any}}, b::AbstractQuasiOrVector) = (axes(first(A),2) == Base.OneTo(1) ? first(A)*sum(b) : (first(A)*b), _broadcasted_mul(tail(A), b)...)
+_broadcasted_mul(A::Tuple{AbstractQuasiMatrix,Vararg{Any}}, B::AbstractQuasiOrMatrix) = (axes(first(A),2) == Base.OneTo(1) ? first(A)*sum(B; dims=1) : (first(A)*B), _broadcasted_mul(tail(A), B)...)
 _broadcasted_mul(A::AbstractQuasiMatrix, b::Tuple{Number,Vararg{Any}}) = (sum(A; dims=2)*first(b)[1], _broadcasted_mul(A, tail(b))...)
-_broadcasted_mul(A::AbstractQuasiMatrix, b::Tuple{Union{AbstractVector,AbstractQuasiVector},Vararg{Any}}) = (size(first(b),1) == 1 ? (sum(A; dims=2)*first(b)[1]) : (A*first(b)), _broadcasted_mul(A, tail(b))...)
-_broadcasted_mul(A::AbstractQuasiMatrix, B::Tuple{Union{AbstractMatrix,AbstractQuasiMatrix},Vararg{Any}}) = (size(first(B),1) == 1 ? (sum(A; dims=2) * first(B)) : (A * first(B)), _broadcasted_mul(A, tail(B))...)
+_broadcasted_mul(A::AbstractQuasiMatrix, b::Tuple{AbstractQuasiOrVector,Vararg{Any}}) = (size(first(b),1) == 1 ? (sum(A; dims=2)*first(b)[1]) : (A*first(b)), _broadcasted_mul(A, tail(b))...)
+_broadcasted_mul(A::AbstractQuasiMatrix, B::Tuple{AbstractQuasiOrMatrix,Vararg{Any}}) = (size(first(B),1) == 1 ? (sum(A; dims=2) * first(B)) : (A * first(B)), _broadcasted_mul(A, tail(B))...)
 
 
 

test/test_calculus.jl

@@ -30,19 +30,25 @@ using QuasiArrays, IntervalSets, Test
 
     @testset "Diff" begin
         x = range(0, 1; length=10_000)
-        @test diff(Inclusion(x)) == ones(Inclusion(x[1:end-1]))
-        @test diff(ones(Inclusion(x))) == zeros(Inclusion(x[1:end-1]))
+        @test diff(Inclusion(x)) == diff(Inclusion(x),1) == ones(Inclusion(x[1:end-1]))
+        @test diff(Inclusion(x),2) == diff(diff(Inclusion(x))) == zeros(Inclusion(x[1:end-2]))
+        @test diff(ones(Inclusion(x))) == diff(ones(Inclusion(x)),1) == zeros(Inclusion(x[1:end-1]))
+        @test diff(ones(Inclusion(x)),2) == diff(diff(ones(Inclusion(x)))) == zeros(Inclusion(x[1:end-2]))
 
         @test diff(ones(Inclusion(x), Inclusion(x))) == zeros(Inclusion(x[1:end-1]), Inclusion(x))
+        @test diff(ones(Inclusion(x), Inclusion(x)), 2) == zeros(Inclusion(x[1:end-2]), Inclusion(x))
         @test diff(ones(Inclusion(x), Inclusion(x)); dims=2) == zeros(Inclusion(x), Inclusion(x[1:end-1]))
+        @test diff(ones(Inclusion(x), Inclusion(x)), 2; dims=2) == zeros(Inclusion(x), Inclusion(x[1:end-2]))
 
         b = QuasiVector(exp.(x), x)
 
         @test diff(b) ≈ b[Inclusion(x[1:end-1])] atol=1E-2
+        @test diff(b,2) ≈ b[Inclusion(x[1:end-2])] atol=1E-1
 
 
         A = QuasiArray(randn(3,2), (1:0.5:2,0:0.5:0.5))
         @test diff(A; dims=1)[:,0] == diff(A[:,0])
+        @test diff(A,2; dims=1)[:,0] == diff(diff(A[:,0]))
         @test diff(A; dims=2)[1,:] == diff(A[1,:])
 
         @testset "* diff" begin
@@ -57,10 +63,18 @@ using QuasiArrays, IntervalSets, Test
 
     @testset "Interval" begin
         @test diff(Inclusion(0.0..1)) ≡ ones(Inclusion(0.0..1))
-        @test diff(ones(Inclusion(0.0..1))) ≡ zeros(Inclusion(0.0..1))
-        @test diff(ones(Inclusion(0.0..1), Base.OneTo(3))) ≡ zeros(Inclusion(0.0..1), Base.OneTo(3))
+        @test diff(Inclusion(0.0..1),1) ≡ fill(1.0,Inclusion(0.0..1))
+        @test diff(Inclusion(0.0..1),2) ≡ fill(0.0,Inclusion(0.0..1))
+        @test diff(ones(Inclusion(0.0..1))) ≡ diff(ones(Inclusion(0.0..1)),1) ≡ diff(ones(Inclusion(0.0..1)),2) ≡ zeros(Inclusion(0.0..1))
+        @test diff(ones(Inclusion(0.0..1), Base.OneTo(3))) ≡ diff(ones(Inclusion(0.0..1), Base.OneTo(3)),2) ≡ zeros(Inclusion(0.0..1), Base.OneTo(3))
         @test diff(ones(Inclusion(0.0..1), Base.OneTo(3)); dims=2) ≡ zeros(Inclusion(0.0..1), Base.OneTo(2))
         @test diff(ones(Base.OneTo(3), Inclusion(0.0..1))) ≡ zeros(Base.OneTo(2), Inclusion(0.0..1))
         @test diff(ones(Base.OneTo(3), Inclusion(0.0..1)); dims=2) ≡ zeros(Base.OneTo(3), Inclusion(0.0..1))
     end
+
+    @testset "Incomplete" begin
+        struct IncompleteQuasiArray <: AbstractQuasiVector{Int} end
+        Base.axes(::IncompleteQuasiArray) = (Base.OneTo(3),)
+        @test_throws ErrorException diff(IncompleteQuasiArray())
+    end
 end