Add diagtrav function with dims keyword (#151)

* Add diagtrav function with dims keyword

* add tests

* Update test_blockkron.jl

* remove moved code

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.11.6"
+version = "0.11.7"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -28,7 +28,7 @@ BlockArrays = "1.0"
 BlockBandedMatrices = "0.13"
 FillArrays = "1.0"
 InfiniteArrays = "0.15"
-LazyArrays = "2.2.3"
+LazyArrays = "2.8"
 MatrixFactorizations = "3.0"
 StaticArrays = "1.0"
 julia = "1.10"

src/LazyBandedMatrices.jl

@@ -38,6 +38,8 @@ const LazyBlockBandedLayouts = LazyArraysBlockBandedMatricesExt.LazyBlockBandedL
 
 export DiagTrav, KronTrav, blockkron, BlockKron, BlockBroadcastArray, BlockVcat, BlockHcat, BlockHvcat, unitblocks
 
+## TODO: export diagtrav, invdiagtrav
+
 include("tridiag.jl")
 include("bidiag.jl")
 include("special.jl")

src/blockconcat.jl

@@ -492,30 +492,7 @@ resize!(c::BlockBroadcastVector{T,typeof(vcat)}, N::Block{1}) where T = BlockBro
 # BlockVec
 ####
 
-# support LazyArrays v2.8 where BlockVec is moved
-if isdefined(LazyBandedMatrices.LazyArraysBlockArraysExt, :BlockVec)
-    const BlockVec = LazyBandedMatrices.LazyArraysBlockArraysExt.BlockVec
-else
-    const BlockVec{T, M<:AbstractMatrix{T}} = ApplyVector{T, typeof(blockvec), <:Tuple{M}}
-
-    BlockVec{T}(M::AbstractMatrix{T}) where T = ApplyVector{T}(blockvec, M)
-    BlockVec(M::AbstractMatrix{T}) where T = BlockVec{T}(M)
-    axes(b::BlockVec) = (blockedrange(Fill(size(b.args[1])...)),)
-    size(b::BlockVec) = (length(b.args[1]),)
-
-    view(b::BlockVec, K::Block{1}) = view(b.args[1], :, Int(K))
-    Base.@propagate_inbounds getindex(b::BlockVec, k::Int) = b.args[1][k]
-    Base.@propagate_inbounds setindex!(b::BlockVec, v, k::Int) = setindex!(b.args[1], v, k)
-
-    _resize!(A::AbstractMatrix, m, n) = A[1:m, 1:n]
-    _resize!(At::Transpose, m, n) = transpose(transpose(At)[1:n, 1:m])
-    _resize!(Ac::Adjoint, m, n) = (Ac')[1:n, 1:m]'
-    resize!(b::BlockVec, K::Block{1}) = BlockVec(_resize!(b.args[1], size(b.args[1],1), Int(K)))
-
-    applylayout(::Type{typeof(blockvec)}, ::AbstractPaddedLayout) = PaddedColumns{ApplyLayout{typeof(blockvec)}}()
-    paddeddata(b::BlockVec) = BlockVec(paddeddata(b.args[1]))
-end
-
+const BlockVec = LazyBandedMatrices.LazyArraysBlockArraysExt.BlockVec
 
 ####
 # summary

src/blockkron.jl

@@ -237,7 +237,22 @@ function getindex(A::InvDiagTrav{T}, k::Int, j::Int)  where T
     end
 end
 
-invdiagtrav(a) = InvDiagTrav(a)
+diagtrav(a::AbstractMatrix) = DiagTrav(a)
+function diagtrav(a::AbstractArray{T,3}; dims=1:3) where T
+    if dims == 1:3
+        DiagTrav(a)
+    else
+        @assert dims == 1:2
+        ret = BlockedMatrix{T}(undef, (_krontrav_axes(axes(a,1), axes(a,2)), axes(a,3)))
+        for ℓ in axes(a,3)
+            ret[:,ℓ] = DiagTrav(view(a,:,:,ℓ))
+        end
+        ret
+    end
+end
+
+diagtrav(a::InvDiagTrav) = a.vector
+invdiagtrav(a::AbstractVector) = InvDiagTrav(a)
 invdiagtrav(a::DiagTrav) = a.array
 
 -(A::DiagTrav) = DiagTrav(-A.array)

test/test_blockkron.jl

@@ -4,7 +4,7 @@ using LazyBandedMatrices, FillArrays, BandedMatrices, BlockBandedMatrices, Block
 
 using LinearAlgebra
 import BlockBandedMatrices: isbandedblockbanded, isbanded, BandedBlockBandedStyle, BandedLayout, _BandedBlockBandedMatrix
-import LazyBandedMatrices: KronTravBandedBlockBandedLayout, BroadcastBandedLayout, BroadcastBandedBlockBandedLayout, arguments, call, blockcolsupport, InvDiagTrav, invdiagtrav, pad, krontrav
+import LazyBandedMatrices: KronTravBandedBlockBandedLayout, BroadcastBandedLayout, BroadcastBandedBlockBandedLayout, arguments, call, blockcolsupport, InvDiagTrav, invdiagtrav, pad, krontrav, diagtrav
 import ArrayLayouts: FillLayout, OnesLayout
 import LazyArrays: resizedata!, FillLayout, arguments, colsupport, call, LazyArrayStyle
 import BandedMatrices: BandedColumns
@@ -24,7 +24,7 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
 @testset "Kron" begin
     @testset "DiagTrav" begin
         A = [1 2 3; 4 5 6; 7 8 9]
-        @test DiagTrav(A) == Vector(DiagTrav(A)) == [1, 4, 2, 7, 5, 3]
+        @test DiagTrav(A) == Vector(DiagTrav(A)) == diagtrav(A) == [1, 4, 2, 7, 5, 3]
         @test resize!(DiagTrav(A), Block(2)) == [1, 4,2]
         @test maximum(abs, DiagTrav(A)) == 7
         @test copy(DiagTrav(A)) == DiagTrav(A)
@@ -51,14 +51,18 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
         @test resize!(DiagTrav(A), Block(2)) == [1, 3,2]
 
         A = DiagTrav(randn(3,3,3))
+        @test A == diagtrav(A.array)
         @test A[Block(1)] == A[1:1,1,1]
         @test A[Block(2)] == [A.array[2,1,1], A.array[1,2,1], A.array[1,1,2]]
         @test A[Block(3)] == [A.array[3,1,1], A.array[2,2,1], A.array[2,1,2],
                             A.array[1,3,1], A.array[1,2,2], A.array[1,1,3]]
         @test A == [A[Block(1)]; A[Block(2)]; A[Block(3)]]
 
+        A = randn(3,3,3)
+        @test diagtrav(A; dims=1:2) == [diagtrav(A[:,:,1]) diagtrav(A[:,:,2]) diagtrav(A[:,:,3])]
+
         A = reshape(1:9,3,3)'
-        @test DiagTrav(A) == Vector(DiagTrav(A)) == [1, 4, 2, 7, 5, 3]
+        @test DiagTrav(A) == Vector(DiagTrav(A)) == diagtrav(A) == [1, 4, 2, 7, 5, 3]
         A = reshape(1:12,3,4)'
         @test DiagTrav(A) == [1, 4, 2, 7, 5, 3, 10, 8, 6]
         A = reshape(1:12,3,4)
@@ -84,6 +88,7 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
     @testset "InvDiagTrav" begin
         A = [1 2 3; 4 5 6; 7 8 9]
         @test invdiagtrav(BlockedVector(DiagTrav(A))) == invdiagtrav(DiagTrav(A)) == [1 2 3; 4 5 0; 7 0 0]
+        @test diagtrav(invdiagtrav(diagtrav(A))) == diagtrav(invdiagtrav(BlockedVector(diagtrav(A))))  == diagtrav(A)
     end
 
     @testset "BlockKron" begin