Fast KronTrav * DiagTrav (#138)

Author: dlfivefifty

ext/LazyBandedMatricesInfiniteArraysExt.jl

@@ -4,11 +4,11 @@ using LazyBandedMatrices.BlockArrays
 using LazyBandedMatrices.ArrayLayouts
 
 import Base: BroadcastStyle, copy, OneTo, oneto
-import LazyBandedMatrices: _krontrav_axes, _block_interlace_axes, _broadcast_sub_arguments, AbstractLazyBandedBlockBandedLayout, KronTravBandedBlockBandedLayout, krontravargs
+import LazyBandedMatrices: _krontrav_axes, _block_interlace_axes, _broadcast_sub_arguments, AbstractLazyBandedBlockBandedLayout, KronTravBandedBlockBandedLayout, krontravargs, DiagTravLayout
 import InfiniteArrays: InfFill, TridiagonalToeplitzLayout, BidiagonalToeplitzLayout, LazyArrayStyle, OneToInf
 import LazyBandedMatrices.ArrayLayouts: MemoryLayout, sublayout, RangeCumsum, Mul
 import LazyBandedMatrices.BlockArrays: sizes_from_blocks, BlockedOneTo, BlockSlice1, BlockSlice
-import LazyBandedMatrices.LazyArrays: BroadcastBandedLayout
+import LazyBandedMatrices.LazyArrays: BroadcastBandedLayout, AbstractPaddedLayout
 
 const OneToInfCumsum = RangeCumsum{Int,OneToInf{Int}}
 
@@ -53,4 +53,8 @@ _block_interlace_axes(nbc::Int, ax::NTuple{2,BlockedOneTo{Int,OneToInf{Int}}}...
     (blockedrange(Fill(length(ax) ÷ nbc, ∞)),blockedrange(Fill(mod1(length(ax),nbc), ∞)))
 
 
+# KronTrav * DiagTrav
+
+copy(M::Mul{InfKronTravBandedBlockBandedLayout, Lay}) where Lay<:DiagTravLayout{<:AbstractPaddedLayout} = copy(Mul{KronTravBandedBlockBandedLayout, Lay}(M.A, M.B))
+
 end

src/LazyBandedMatrices.jl

@@ -13,7 +13,8 @@ import LinearAlgebra: transpose, adjoint, istriu, istril, isdiag, tril!, triu!,
 
 import ArrayLayouts: MemoryLayout, bidiagonallayout, bidiagonaluplo, diagonaldata, supdiagonaldata, subdiagonaldata,
                      symtridiagonallayout, tridiagonallayout, symmetriclayout,
-                     colsupport, rowsupport, sublayout, sub_materialize, _copyto!
+                     colsupport, rowsupport, sublayout, sub_materialize, _copyto!,
+                     materialize!, MulAdd, MatMulVecAdd
 import LazyArrays: ApplyLayout, AbstractPaddedLayout, PaddedLayout, PaddedColumns, BroadcastLayout, LazyArrayStyle, LazyLayout,
                    arguments, call, tuple_type_memorylayouts, paddeddata, _broadcast_sub_arguments, resizedata!,
                    _cumsum, convexunion, applylayout, AbstractLazyBandedLayout, ApplyBandedLayout, BroadcastBandedLayout, LazyBandedLayout

src/blockkron.jl

@@ -4,21 +4,67 @@
 # Block
 ###
 
+function _DiagTrav end
+
 """
     DiagTrav(A::AbstractMatrix)
 
 converts a matrix to a block vector by traversing the anti-diagonals.
 """
 struct DiagTrav{T, N, AA<:AbstractArray{T,N}} <: AbstractBlockVector{T}
     array::AA
-    function DiagTrav{T, N, AA}(array::AA) where {T, N, AA<:AbstractArray{T,N}}
+    global function _DiagTrav(array::AA) where {T, N, AA<:AbstractArray{T,N}}
         new{T,N,AA}(array)
     end
 end
+DiagTrav{T,N,AA}(A::AA) where {T, N, AA<:AbstractArray{T,N}} = _DiagTrav(zero_bottomright(A))
+
 DiagTrav{T,N}(A::AbstractArray) where {T,N} = DiagTrav{T,N,typeof(A)}(A)
 DiagTrav{T}(A::AbstractArray{<:Any,N}) where {T,N} = DiagTrav{T,N}(A)
 DiagTrav(A::AbstractArray{T}) where T = DiagTrav{T}(A)
 
+
+zero_bottomright(array) = zero_bottomright(array, axes(array))
+zero_bottomright!(array) = zero_bottomright!(array, axes(array))
+
+function zero_bottomright(X::AbstractMatrix, _)
+    m,n = size(X)
+    μ = max(m,n)
+    for j in rowsupport(X), k = (μ-j+2:m) ∩ colsupport(X,j)
+        iszero(X[k,j]) || return zero_bottomright!(copy(X))
+    end
+    X
+end
+
+function zero_bottomright!(X::AbstractMatrix{T}, _) where T
+    m,n = size(X)
+    μ = max(m,n)
+    for j in rowsupport(X), k = (μ-j+2:m) ∩ colsupport(X,j)
+        X[k,j] = zero(T)
+    end
+    X
+end
+
+function zero_bottomright(X::AbstractArray{<:Any,3}, _)
+    m,n,p = size(X)
+    @assert m == n == p
+    for ℓ = 0:n-1, j=0:n-1, k=max(0,n-(ℓ+j)):n-1
+        iszero(X[k+1,j+1,ℓ+1]) || return zero_bottomright!(copy(X))
+    end
+    X
+end
+
+function zero_bottomright!(X::AbstractArray{T,3}, _) where T
+    m,n,p = size(X)
+    @assert m == n == p
+    for ℓ = 0:n-1, j=0:n-1, k=max(0,n-(ℓ+j)):n-1
+        X[k+1,j+1,ℓ+1] = zero(T)
+    end
+    X
+end
+
+
+
 function _krontrav_axes(A, B)
     m,n = length(A), length(B)
     mn = min(m,n)
@@ -35,6 +81,8 @@ axes(A::DiagTrav) = (_krontrav_axes(axes(A.array)...),)
 
 copy(A::DiagTrav) = DiagTrav(copy(A.array))
 
+similar(A::DiagTrav, ::Type{T}) where T = DiagTrav(similar(A.array, T))
+
 struct DiagTravLayout{Lay} <: AbstractBlockLayout end
 MemoryLayout(::Type{<:DiagTrav{T, N, AA}}) where {T,N,AA} = DiagTravLayout{typeof(MemoryLayout(AA))}()
 
@@ -62,7 +110,7 @@ function _diagtravgetindex(_, A::AbstractMatrix, K::Block{1})
 end
 
 
-_diagtravgetindex(::AbstractStridedLayout, A::AbstractMatrix, K::Block{1}) = layout_getindex(DiagTrav(A), K)
+_diagtravgetindex(::AbstractStridedLayout, A::AbstractMatrix, K::Block{1}) = layout_getindex(_DiagTrav(A), K)
 
 function _diagtravview(::AbstractStridedLayout, A::AbstractMatrix, K::Block{1})
     k = Int(K)
@@ -110,6 +158,7 @@ function _diagtravgetindex(::AbstractStridedLayout, A::AbstractArray{T,3}, K::Bl
 end
 
 getindex(A::DiagTrav, k::Int) = A[findblockindex(axes(A,1), k)]
+setindex!(A::DiagTrav, v, k::Int) = A[findblockindex(axes(A,1), k)] = v
 
 function resize!(A::DiagTrav{<:Any,2}, K::Block{1})
     k = Int(K)
@@ -261,6 +310,7 @@ convert(::Type{B}, A::KronTrav{<:Any,2}) where B<:BandedBlockBandedMatrix = conv
 struct KronTravBandedBlockBandedLayout <: AbstractBandedBlockBandedLayout end
 struct KronTravLayout{M} <: AbstractBlockLayout end
 
+const KronTravLayouts = Union{KronTravBandedBlockBandedLayout, KronTravLayout}
 
 
 krontravlayout(::Vararg{Any,M}) where M = KronTravLayout{M}()
@@ -320,4 +370,27 @@ BroadcastStyle(::Type{KronTrav{T,N,AA,AXIS}}) where {T,N,AA,AXIS} =
 ###
 
 *(a::Number, b::KronTrav) = KronTrav(a*first(b.args), tail(b.args)...)
-*(a::KronTrav, b::Number) = KronTrav(first(a.args)*b, tail(a.args)...)
+*(a::KronTrav, b::Number) = KronTrav(first(a.args)*b, tail(a.args)...)
+
+
+function copy(M::Mul{<:KronTravLayouts, <:DiagTravLayout})
+    K,x = M.A,M.B
+    A,B = K.args
+    _krontrav_mul_diagtrav(K.args, invdiagtrav(x), eltype(M))
+end
+
+_krontrav_mul_diagtrav((A,B), X::AbstractMatrix, ::Type{T}) where T = DiagTrav(convert(AbstractMatrix{T}, B*X*A'))
+function _krontrav_mul_diagtrav((A,B,C), X::AbstractArray{<:Any,3}, ::Type{T}) where T
+    m,n,p = size(X)
+    @assert m == n == p
+    Y = similar(X, T)
+    Z = similar(X, T)
+    for k = 1:n, j=1:n mul!(view(Y,k,j,:),A,view(X,k,j,:)) end
+    for k = 1:n, j=1:n mul!(view(Z,k,:,j),B,view(Y,k,:,j)) end
+    for k = 1:n, j=1:n mul!(view(Y,:,k,j),C,view(Z,:,k,j)) end
+    DiagTrav(Y)
+end
+
+
+
+# C = α*B*X*A' + β*C

test/test_blockkron.jl

@@ -68,8 +68,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))) == [1 2 3; 4 5 0; 7 0 0]
-        @test invdiagtrav(DiagTrav(A)) == A
+        @test invdiagtrav(BlockedVector(DiagTrav(A))) == invdiagtrav(DiagTrav(A)) == [1 2 3; 4 5 0; 7 0 0]
     end
 
     @testset "BlockKron" begin
@@ -124,7 +123,9 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
             @test copy(K) == K
 
             X = [9 10; 11 0]
-            @test K*DiagTrav(X) == DiagTrav(B*X*A')
+            Y = [9 10; 11 12]
+            @test K*DiagTrav(X) == K*DiagTrav(Y) == DiagTrav(B*X*A')
+            @test K*DiagTrav(X) isa DiagTrav
 
             @test K[Block.(Base.OneTo(2)), Block.(Base.OneTo(2))] == K[Block.(1:2), Block.(1:2)] == K
         end
@@ -160,9 +161,11 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
             K = KronTrav(A,B,C)
 
             X = randn(n,n,n)
+            Y = copy(X)
             for ℓ = 0:n-1, j=0:n-1, k=max(0,n-(ℓ+j)):n-1
                 X[k+1,j+1,ℓ+1] = 0
             end
+            @test DiagTrav(Y).array == X
             Y = float(similar(X))
             for k = 1:n, j=1:n Y[k,j,:] = A*X[k,j,:] end
             for k = 1:n, j=1:n Y[k,:,j] = B*Y[k,:,j] end

test/test_lazybandedinf.jl

@@ -12,26 +12,28 @@ const OneToInfBlocks = InfiniteArraysBlockArraysExt.OneToInfBlocks
 const InfKronTravBandedBlockBandedLayout = LazyBandedMatricesInfiniteArraysExt.InfKronTravBandedBlockBandedLayout
 
 @testset "∞ LazyBandedMatrices" begin
-    @test MemoryLayout(LazyBandedMatrices.Tridiagonal(Fill(1,∞), Zeros(∞), Fill(3,∞))) isa TridiagonalToeplitzLayout
-    @test MemoryLayout(LazyBandedMatrices.Bidiagonal(Fill(1,∞), Zeros(∞), :U)) isa BidiagonalToeplitzLayout
-    @test MemoryLayout(LazyBandedMatrices.SymTridiagonal(Fill(1,∞), Zeros(∞))) isa TridiagonalToeplitzLayout
-
-    T = LazyBandedMatrices.Tridiagonal(Fill(1,∞), Zeros(∞), Fill(3,∞))
-    @test T[2:∞,3:∞] isa SubArray
-    @test exp.(T) isa BroadcastMatrix
-    @test exp.(T)[2:∞,3:∞][1:10,1:10] == exp.(T[2:∞,3:∞])[1:10,1:10] == exp.(T[2:11,3:12])
-    @test exp.(T)[2:∞,3:∞] isa BroadcastMatrix
-    @test exp.(T[2:∞,3:∞]) isa BroadcastMatrix
-
-    B = LazyBandedMatrices.Bidiagonal(Fill(1,∞), Zeros(∞), :U)
-    @test B[2:∞,3:∞] isa SubArray
-    @test exp.(B) isa BroadcastMatrix
-    @test exp.(B)[2:∞,3:∞][1:10,1:10] == exp.(B[2:∞,3:∞])[1:10,1:10] == exp.(B[2:11,3:12])
-    @test exp.(B)[2:∞,3:∞] isa BroadcastMatrix
-
-    @testset "Diagonal{Fill} * Bidiagonal" begin
-        A, B = Diagonal(Fill(2,∞)) , LazyBandedMatrices.Bidiagonal(exp.(1:∞), exp.(1:∞), :L)
-        @test (A*B)[1:10,1:10] ≈ (B*A)[1:10,1:10] ≈ 2B[1:10,1:10]
+    @testset "Tri/Bidiagonal" begin
+        @test MemoryLayout(LazyBandedMatrices.Tridiagonal(Fill(1,∞), Zeros(∞), Fill(3,∞))) isa TridiagonalToeplitzLayout
+        @test MemoryLayout(LazyBandedMatrices.Bidiagonal(Fill(1,∞), Zeros(∞), :U)) isa BidiagonalToeplitzLayout
+        @test MemoryLayout(LazyBandedMatrices.SymTridiagonal(Fill(1,∞), Zeros(∞))) isa TridiagonalToeplitzLayout
+
+        T = LazyBandedMatrices.Tridiagonal(Fill(1,∞), Zeros(∞), Fill(3,∞))
+        @test T[2:∞,3:∞] isa SubArray
+        @test exp.(T) isa BroadcastMatrix
+        @test exp.(T)[2:∞,3:∞][1:10,1:10] == exp.(T[2:∞,3:∞])[1:10,1:10] == exp.(T[2:11,3:12])
+        @test exp.(T)[2:∞,3:∞] isa BroadcastMatrix
+        @test exp.(T[2:∞,3:∞]) isa BroadcastMatrix
+
+        B = LazyBandedMatrices.Bidiagonal(Fill(1,∞), Zeros(∞), :U)
+        @test B[2:∞,3:∞] isa SubArray
+        @test exp.(B) isa BroadcastMatrix
+        @test exp.(B)[2:∞,3:∞][1:10,1:10] == exp.(B[2:∞,3:∞])[1:10,1:10] == exp.(B[2:11,3:12])
+        @test exp.(B)[2:∞,3:∞] isa BroadcastMatrix
+
+        @testset "Diagonal{Fill} * Bidiagonal" begin
+            A, B = Diagonal(Fill(2,∞)) , LazyBandedMatrices.Bidiagonal(exp.(1:∞), exp.(1:∞), :L)
+            @test (A*B)[1:10,1:10] ≈ (B*A)[1:10,1:10] ≈ 2B[1:10,1:10]
+        end
     end
 
     @testset "∞-unit blocks" begin
@@ -102,6 +104,12 @@ const InfKronTravBandedBlockBandedLayout = LazyBandedMatricesInfiniteArraysExt.I
 
         @test A*A isa KronTrav
         @test (A*A)[Block.(Base.OneTo(3)), Block.(Base.OneTo(3))] ≈ A[Block.(1:3), Block.(1:4)]A[Block.(1:4), Block.(1:3)]
+
+        @testset "mul" begin
+            X = zeros(∞,∞); X[1,1] = 1;
+            KR = Block.(1:10)
+            @test (A*DiagTrav(X))[KR] == ((A + 0I) * DiagTrav(X))[KR]
+        end
     end
 
     @testset "BlockHcat copyto!" begin