Add blockkron (#123)

* Start work on blockkron

* Fix axes, size

* Add tests for matrix case

* Add tests, clean up getindex

Author: dlfivefifty

src/BlockArrays.jl

@@ -14,7 +14,7 @@ export PseudoBlockArray, PseudoBlockMatrix, PseudoBlockVector, PseudoBlockVecOrM
 
 export undef_blocks, undef, findblock, findblockindex
 
-export khatri_rao
+export khatri_rao, blockkron, BlockKron
 
 export blockappend!, blockpush!, blockpushfirst!, blockpop!, blockpopfirst!
 

src/blockproduct.jl

@@ -6,7 +6,7 @@ References
 * Khatri, C. G., and Rao, C. Radhakrishna (1968) Solutions to Some Functional Equations and Their Applications to Characterization of Probability Distributions. Sankhya: Indian J. Statistics, Series A 30, 167–180.
 """
 function khatri_rao(A::AbstractBlockMatrix, B::AbstractBlockMatrix)
-    # 
+    #
     Ablksize = blocksize(A)
     Bblksize = blocksize(B)
 
@@ -27,4 +27,90 @@ end
 
 function khatri_rao(A::AbstractMatrix, B::AbstractMatrix)
     kron(A, B)
-end
+end
+
+""""
+    BlockKron(A...)
+
+creates a lazy representation of kron(A...) with the natural
+block-structure imposed. This is a component in `blockkron(A...)`.
+"""
+struct BlockKron{T,N,ARGS<:Tuple} <: AbstractBlockArray{T,N}
+    args::ARGS
+end
+
+BlockKron{T,N}(A...) where {T,N} = BlockKron{T,N,typeof(A)}(A)
+BlockKron{T}(A::AbstractVector, B::AbstractVector, C::AbstractVector...) where {T} = BlockKron{T,1}(A, B, C...)
+BlockKron{T}(A, B, C...) where {T} = BlockKron{T,2}(A, B, C...)
+BlockKron(A, B, C...) = BlockKron{mapreduce(eltype,promote_type,(A,B,C...))}(A, B, C...)
+
+
+size(B::BlockKron) = size(Kron(B))
+
+size(K::BlockKron, j::Int) = prod(size.(K.args, j))
+size(a::BlockKron{<:Any,1}) = (size(a,1),)
+size(a::BlockKron{<:Any,2}) = (size(a,1), size(a,2))
+
+function axes(K::BlockKron{<:Any,1})
+    A,B = K.args
+    (blockedrange(fill(prod(size.(tail(K.args),1)), size(K.args[1],1))),)
+end
+
+function axes(K::BlockKron{<:Any,2})
+    A,B = K.args
+    blockedrange.((fill(prod(size.(tail(K.args),1)), size(K.args[1],1)),
+                   fill(prod(size.(tail(K.args),2)), size(K.args[1],2))))
+end
+
+kron_getindex((A,)::Tuple{AbstractVector}, k::Integer) = A[k]
+function kron_getindex((A,B)::NTuple{2,AbstractVector}, k::Integer)
+    K,κ = divrem(k-1, length(B))
+    A[K+1]*B[κ+1]
+end
+kron_getindex((A,)::Tuple{AbstractMatrix}, k::Integer, j::Integer) = A[k,j]
+function kron_getindex((A,B)::NTuple{2,AbstractVecOrMat}, k::Integer, j::Integer)
+    K,κ = divrem(k-1, size(B,1))
+    J,ξ = divrem(j-1, size(B,2))
+    A[K+1,J+1]*B[κ+1,ξ+1]
+end
+
+kron_getindex(args::Tuple, k::Integer, j::Integer) = kron_getindex(tuple(BlockKron(args[1:2]...), args[3:end]...), k, j)
+kron_getindex(args::Tuple, k::Integer) = kron_getindex(tuple(BlockKron(args[1:2]...), args[3:end]...), k)
+
+getindex(K::BlockKron{<:Any,1}, k::Integer) = kron_getindex(K.args, k)
+getindex(K::BlockKron{<:Any,2}, k::Integer, j::Integer) = kron_getindex(K.args, k, j)
+
+kron_getblock((a,b)::Tuple{Any,Any}, k::Integer) = a[k]*b
+kron_getblock(args, k::Integer) = args[1][k]*BlockKron(tail(args)...)
+
+kron_getblock((a,b)::Tuple{Any,Any}, k::Integer, j::Integer) = a[k,j]*b
+kron_getblock(args, k::Integer, j::Integer) = args[1][k,j]*BlockKron(tail(args)...)
+
+getblock(K::BlockKron{<:Any,1}, k::Integer) = kron_getblock(K.args, k)
+getblock(K::BlockKron{<:Any,2}, k::Integer, j::Integer) = kron_getblock(K.args, k, j)
+
+# const SubKron{T,M1,M2,R1,R2} = SubArray{T,2,<:BlockKron{T,M1,M2},<:Tuple{<:BlockSlice{R1},<:BlockSlice{R2}}}
+
+
+# BroadcastStyle(::Type{<:SubKron{<:Any,<:Any,B,Block1,Block1}}) where B =
+#     BroadcastStyle(B)
+
+
+# allow dispatch on memory layout
+_blockkron(_, A) = BlockArray(BlockKron(A...))
+
+
+""""
+    blockkron(A...)
+
+creates a blocked version of kron(A...) with the natural
+block-structure imposed.
+"""
+blockkron(A...) = _blockkron(map(MemoryLayout,A), A)
+
+"""
+    blockvec(A::AbstractMatrix)
+
+creates a blocked version of `vec(A)`, with the block structure used to represent the columns.
+"""
+blockvec(A::AbstractMatrix) = PseudoBlockVector(vec(A), Fill(size(A,1), size(A,2)))

test/test_blockproduct.jl

@@ -100,3 +100,42 @@ end
     @test khatri_rao(A, B) ≈ kron(A, B)
 end
 
+@testset "blockkron" begin
+    a = [1,2]
+    b = [3,4,5]
+    k̃ = BlockKron(a,b) # lazy version of blockkron
+    k = blockkron(a,b)
+    @test k == k̃ == kron(a,b)
+    @test blocksize(k) == blocksize(k̃) == size(a)
+
+    @test k[Block(1)] == k̃[Block(1)] == a[1]*b
+    @test k[Block(2)] == k̃[Block(2)] == a[2]*b
+    c = 6:8
+    k̄ = BlockKron(a,b,c)
+    @test k̄ == blockkron(a,b,c) == kron(a,b,c)
+    @test k̄[Block(1)][Block(1)] == a[1]*b[1]*c
+    @test k̄[Block(1)][Block(2)] == a[1]*b[2]*c
+    @test k̄[Block(2)][Block(3)] == a[2]*b[3]*c
+
+    A = randn(2,3)
+    B = randn(3,4)
+    K = blockkron(A,B)
+    K̃ = BlockKron(A,B)
+    @test K == K̃ == kron(A,B)
+    @test blocksize(K) == blocksize(K̃) == size(A)
+    @test K[Block(1,1)] == K̃[Block(1),Block(1)] == A[1,1]*B
+    @test K[Block(2,3)] == K̃[Block(2),Block(3)] == A[2,3]*B
+    C = randn(2,5)
+    K̄ = BlockKron(A,B,C)
+    @test K̄ == blockkron(A,B,C) == kron(A,B,C)
+    @test K̄[Block(1,1)][Block(1,1)] ≈ A[1,1]*B[1,1]*C
+    @test K̄[Block(2,3)][Block(3,4)] ≈ A[2,3]*B[3,4]*C
+
+    @test blockkron(a,B) == kron(a,B)
+    @test blockkron(A,b) == kron(A,b)
+    @test blockkron(A,b,c) == kron(A,b,c)
+    @test blockkron(A,b,C) == kron(A,b,C)
+
+    @test_throws MethodError BlockKron()
+    @test_throws MethodError BlockKron(a)
+end