Better indexing into a KronExpansion (#207)

* Better indexing into a KronExpansion

* Update bases.jl

* increase coverage

* v0.19.6

Author: dlfivefifty

Project.toml

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

src/bases/bases.jl

@@ -581,15 +581,34 @@ end
 
 
 # we represent as a Mul with a banded matrix
-# sublayout(::AbstractBasisLayout, ::Type{<:Tuple{<:Inclusion,<:Integer}}) = SubBasisLayout()
-sublayout(::AbstractBasisLayout, ::Type{<:Tuple{<:Inclusion,<:AbstractVector}}) = SubBasisLayout()
-sublayout(::AbstractBasisLayout, ::Type{<:Tuple{<:AbstractAffineQuasiVector,<:AbstractVector}}) = MappedBasisLayout()
-sublayout(::WeightedBasisLayouts, ::Type{<:Tuple{<:AbstractAffineQuasiVector,<:AbstractVector}}) = MappedWeightedBasisLayout()
-sublayout(::WeightedBasisLayout, ::Type{<:Tuple{<:Inclusion,<:AbstractVector}}) = SubWeightedBasisLayout()
-sublayout(::MappedWeightedBasisLayout, ::Type{<:Tuple{<:Inclusion,<:AbstractVector}}) = MappedWeightedBasisLayout()
+# sublayout(::AbstractBasisLayout, ::Type{Tuple{Inclusion,Integer}}) = SubBasisLayout()
+sublayout(::AbstractBasisLayout, ::Type{<:Tuple{Inclusion,AbstractVector}}) = SubBasisLayout()
+sublayout(::AbstractBasisLayout, ::Type{<:Tuple{AbstractAffineQuasiVector,AbstractVector}}) = MappedBasisLayout()
+sublayout(::WeightedBasisLayouts, ::Type{<:Tuple{AbstractAffineQuasiVector,AbstractVector}}) = MappedWeightedBasisLayout()
+sublayout(::WeightedBasisLayout, ::Type{<:Tuple{Inclusion,AbstractVector}}) = SubWeightedBasisLayout()
+# sublayout(::MappedWeightedBasisLayout, ::Type{<:Tuple{Inclusion,AbstractVector}}) = MappedWeightedBasisLayout() # not used
+sublayout(lay::ExpansionLayout, ::Type{<:Tuple{Inclusion,Integer}}) = lay
+sublayout(lay::ExpansionLayout, ::Type{<:Tuple{Inclusion,AbstractVector}}) = lay
+
+
+sub_basis_layout(_, P, j) = basis(P) # TODO: restrict to ExpansionLayout?
+function basis(V::SubQuasiArray{<:Any, N, <:Any, <:Tuple{Inclusion,Any}}) where N
+    P = parent(V)
+    _,j = parentindices(V)
+    sub_basis_layout(MemoryLayout(P), P, j)
+end
+
+
+sub_coefficients_layout(_, P, j) = coefficients(P)[:,j] # TODO: restrict to ExpansionLayout?
+function coefficients(V::SubQuasiArray{<:Any, N, <:Any, <:Tuple{Inclusion,Any}}) where N
+     P = parent(V)
+    _,j = parentindices(V)
+    sub_coefficients_layout(MemoryLayout(P), P, j)
+end
 
 @inline sub_materialize(::AbstractBasisLayout, V::AbstractQuasiArray) = V
 @inline sub_materialize(::AbstractBasisLayout, V::AbstractArray) = V
+@inline sub_materialize(::ExpansionLayout, V::AbstractQuasiArray) = basis(V) * coefficients(V)
 
 demap(x) = x
 demap(x::BroadcastQuasiArray) = BroadcastQuasiArray(x.f, map(demap, arguments(x))...)

src/bases/basiskron.jl

@@ -9,5 +9,17 @@ is a MemoryLayout corresponding to a quasi-matrix corresponding to the 2D expans
 """
 struct KronExpansionLayout{LayA, LayB} <: AbstractLazyLayout end
 applylayout(::Type{typeof(*)}, ::LayA, ::CoefficientLayouts, ::AdjointBasisLayout{LayB}) where {LayA <: AbstractBasisLayout, LayB <: AbstractBasisLayout} = KronExpansionLayout{LayA,LayB}()
+
 sublayout(::KronExpansionLayout, inds) = sublayout(ApplyLayout{typeof(*)}(), inds)
+sublayout(::KronExpansionLayout{LayA, LayB}, inds::Type{<:Tuple{Inclusion,AbstractVector}}) where {LayA,LayB} = ExpansionLayout{LayA}()
+
+
+sub_basis_layout(::KronExpansionLayout, P, j) = first(arguments(P))
+
+
+function sub_coefficients_layout(::KronExpansionLayout, P, j)
+   _,X,Bt = arguments(P)
+    X * Bt[:,j]
+end
+
 sum_layout(::KronExpansionLayout, F, dims...) = sum_layout(ApplyLayout{typeof(*)}(), F, dims...)

test/test_basiskron.jl

@@ -18,4 +18,8 @@ end
     @test sum(F; dims=2)[0.1,1] ≈ 7.3
 
     @test F[0.1,:][0.2] ≈ F[:,0.2][0.1] ≈ F[0.1,0.2]
+
+    @test F[:,[0.1,0.2]][0.3,:] ≈ F[0.3,[0.1,0.2]] ≈ F[0.3,:][[0.1,0.2]]
+    @test F[[0.1,0.2],:][:,0.3] ≈ F[[0.1,0.2],0.3] ≈ F[:,0.3][[0.1,0.2]]
+    @test F[[0.1,0.2],[0.3,0.4]] ≈ F[[0.1,0.2],:][:,[0.3,0.4]] ≈ F[:,[0.3,0.4]][[0.1,0.2],:]
 end

test/test_chebyshev.jl

@@ -117,6 +117,7 @@ Base.:(==)(::FooBasis, ::FooBasis) = true
         @test MemoryLayout(w[y] .* T[y,:]) isa MappedWeightedBasisLayout
         @test wT[y,:][[0.1,0.2],1:5] == (w[y] .* T[y,:])[[0.1,0.2],1:5] == (w .* T[:,1:5])[y,:][[0.1,0.2],:]
         @test MemoryLayout(wT[y,1:3]) isa MappedWeightedBasisLayout
+        @test MemoryLayout(wT[y,1:3][:,1:2]) isa MappedWeightedBasisLayout
         @test wT[y,1:3][[0.1,0.2],1:2] == wT[y[[0.1,0.2]],1:2]
 
         @test T[y,:]'T[y,:] ≈ grammatrix(T[y,:]) ≈ (T'T)/2

test/test_splines.jl

@@ -655,9 +655,20 @@ import ContinuumArrays: basis, AdjointBasisLayout, ExpansionLayout, BasisLayout,
         @test expand(L,f)[0.1] ≈ 1
     end
 
-    @testset "vec" begin
-        L = LinearSpline([-1,0,1])
-        F = L * randn(3,1)
-        @test vec(F)[0.1] == F[0.1,1]
+    @testset "matrix coefficients" begin
+        @testset "vec" begin
+            L = LinearSpline([-1,0,1])
+            F = L * randn(3,1)
+            @test vec(F)[0.1] == F[0.1,1]
+        end
+
+        @testset "sub" begin
+            L = LinearSpline([-1,0,1])
+            F = L * randn(3,2)
+            @test MemoryLayout(F[:,1]) isa ExpansionLayout
+            @test MemoryLayout(F[:,1:2]) isa ExpansionLayout
+            @test F[:,1][0.1] ≈ F[0.1,1]
+            @test F[:,1:2][0.1,:] ≈ F[0.1,1:2]
+        end
     end
 end