Add Khatri Rao product (#106)

* Add Khatri Rao product

* change blocksize -> blockaxes
remove _

* change version

Author: pierrejoyot

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "0.12.0"
+version = "0.12.1"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/BlockArrays.jl

@@ -14,6 +14,8 @@ export PseudoBlockArray, PseudoBlockMatrix, PseudoBlockVector, PseudoBlockVecOrM
 
 export undef_blocks, undef, findblock, findblockindex
 
+export khatri_rao
+
 import Base: @propagate_inbounds, Array, to_indices, to_index,
             unsafe_indices, first, last, size, length, unsafe_length,
             unsafe_convert,
@@ -45,6 +47,6 @@ include("show.jl")
 include("blockarrayinterface.jl")
 include("blockbroadcast.jl")
 include("blocklinalg.jl")
-
+include("blockproduct.jl")
 
 end # module

src/blockproduct.jl

@@ -0,0 +1,30 @@
+"""
+    khatri_rao(A, B)
+
+References
+* Liu, Shuangzhe, and Gõtz Trenkler (2008) Hadamard, Khatri-Rao, Kronecker and Other Matrix Products. International J. Information and Systems Sciences 4, 160–177.
+* 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)
+
+    @assert Ablksize == Bblksize "A and B must have the same blocksize"
+
+    kblk = []
+    for iblk in blockaxes(A,1)
+        kblk_j = []
+        for _jblk in blockaxes(A,2)
+            Ablk = A[iblk, _jblk]
+            Bblk = B[iblk, _jblk]
+            push!(kblk_j, kron(Ablk, Bblk))
+        end
+        push!(kblk, tuple(kblk_j...))
+    end
+    mortar(kblk...)
+end
+
+function khatri_rao(A::AbstractMatrix, B::AbstractMatrix)
+    kron(A, B)
+end

test/runtests.jl

@@ -8,4 +8,5 @@ using BlockArrays, LinearAlgebra, Test
     include("test_blockarrayinterface.jl")
     include("test_blockbroadcast.jl")
     include("test_blocklinalg.jl")
+    include("test_blockproduct.jl")
 end

test/test_blockproduct.jl

@@ -0,0 +1,102 @@
+using BlockArrays, Test
+
+@testset "Khatri Rao Product (size)" begin
+    #A size
+    m = 5
+    n = 6
+
+    #A block size
+    mi = [3, 2]
+    ni = [4, 1, 1]
+    #B size
+    p = 8
+    q = 7
+    
+    #B block size
+    pi = [5, 3]
+    qi = [3, 2, 2]
+
+    @assert sum(mi) == m
+    @assert sum(ni) == n
+    @assert sum(qi) == q
+    @assert sum(pi) == p
+
+    A = BlockArray(ones(m, n), mi, ni)
+    B = BlockArray(ones(p, q), pi, qi)
+
+    AB = khatri_rao(A, B)
+    
+    @test blocksize(AB) == blocksize(A)
+    @test blocksize(AB) == blocksize(B)
+
+    #Test: Size of resulting blocks
+    for i in blockaxes(AB,1)
+        for j in blockaxes(AB,2)
+            @test size(AB[i, j]) == (mi[Int(i)]*pi[Int(i)], ni[Int(j)]*qi[Int(j)])
+        end
+    end
+end
+
+@testset "Khatri Rao Product (constant blocks)" begin
+    #A size
+    m = 5
+    n = 6
+
+    #A block size
+    mi = [3, 2]
+    ni = [4, 1, 1]
+    #B size
+    p = 8
+    q = 7
+    
+    #B block size
+    pi = [5, 3]
+    qi = [3, 2, 2]
+
+    @assert sum(mi) == m
+    @assert sum(ni) == n
+    @assert sum(qi) == q
+    @assert sum(pi) == p
+
+    A = BlockArray(ones(m, n), mi, ni)
+    B = BlockArray(ones(p, q), pi, qi)
+
+    #Test: Resulting values for a matrix of constant sub-blocks
+    for i in blockaxes(A,1)
+        for j in blockaxes(A,2)
+            A[i, j] .*= Int(i)+Int(j)
+        end
+    end
+
+    for i in blockaxes(B,1)
+        for j in blockaxes(B,2)
+            B[i, j] .*= Int(i)+Int(j)+10
+        end
+    end
+
+    AB = khatri_rao(A, B)
+
+    for i in blockaxes(AB,1)
+        for j in blockaxes(AB,2)
+            @test AB[i, j] ≈ (Int(i)+Int(j))*(Int(i)+Int(j)+10)*ones(mi[Int(i)]*pi[Int(i)], ni[Int(j)]*qi[Int(j)])
+        end
+    end
+
+end
+
+@testset "Khatri Rao Product (wrong blocksize)" begin
+
+    A = BlockArray(ones(1, 2), [1], [1,1])
+    B = BlockArray(ones(2, 2), [1,1], [1,1])
+
+    @test_throws AssertionError khatri_rao(A, B)
+end
+
+@testset "Khatri Rao Product (Matrix)" begin
+
+    A = ones(1, 2)
+    B = ones(2, 2)
+
+    @test khatri_rao(A, B) ≈ kron(A, B)
+end
+