Specialise LinearAlgebra.lmul! for LowerTriangular blockdiagonal …

…matrices (#119) * Extend `lmul!` for `LowerTriangular` * Add tests for `lmul!` for `LowerTriangular` blockdiagonals * Update src/linalg.jl Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com> * copy all test inputs for safety Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com> * use `return` keyword Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com> * Add link to slow sampling issue * Bump version * Add `BenchmarkTools.jl` as a test dependency * Add allocation test for `lmul!` * Up allocation bound to 320 for Julia 1.0 on x64 Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>
JuliaArrays · Nov 3, 2022 · 86c57e3 · 86c57e3 · mjp98 · Nov 3, 2022
1 parent 10f2c9f
commit 86c57e3
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 10 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.38"
+version = "0.1.39"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -18,11 +18,12 @@ PDMats = "0.11"
 julia = "1"
 
 [extras]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesTestUtils", "Documenter", "PDMats", "Random", "Test"]
+test = ["BenchmarkTools", "ChainRulesTestUtils", "Documenter", "PDMats", "Random", "Test"]
diff --git a/src/linalg.jl b/src/linalg.jl
@@ -45,7 +45,7 @@ end
 """
     eigen_blockwise(B::BlockDiagonal, args...; kwargs...) -> values, vectors
 
-Computes the eigendecomposition for each block separately and keeps the block diagonal 
+Computes the eigendecomposition for each block separately and keeps the block diagonal
 structure in the matrix of eigenvectors. Hence any parameters given are applied to each
 eigendecomposition separately, but there is e.g. no global sorting of eigenvalues.
 """
@@ -58,16 +58,16 @@ function eigen_blockwise(B::BlockDiagonal, args...; kwargs...)
     values = promote([e.values for e in eigens]...)
     vectors = promote([e.vectors for e in eigens]...)
     vcat(values...), BlockDiagonal([vectors...])
-end 
+end
 
 ## This function never keeps the block diagonal structure
 function LinearAlgebra.eigen(B::BlockDiagonal, args...; kwargs...)
     values, vectors = eigen_blockwise(B, args...; kwargs...)
     vectors = Matrix(vectors) # always convert to avoid type instability (also it speeds up the permutation step)
     @static if VERSION > v"1.2.0-DEV.275"
         Eigen(LinearAlgebra.sorteig!(values, vectors, kwargs...)...)
-    else 
-        Eigen(values, vectors) 
+    else
+        Eigen(values, vectors)
     end
 end
 
@@ -157,6 +157,21 @@ function _mul!(C::BlockDiagonal, A::BlockDiagonal, B::BlockDiagonal, α::Number,
     return C
 end
 
+# Resolves MvNormal slow sampling issue https://github.com/invenia/BlockDiagonals.jl/issues/116
+function LinearAlgebra.lmul!(B::LowerTriangular{<:Any,<:BlockDiagonal}, vm::StridedVecOrMat)
+    # BlockDiagonals with non-square blocks
+    if !all(is_square, blocks(parent(B)))
+        return lmul!(LowerTriangular(Matrix(B)), vm) # Fallback on the generic LinearAlgebra method
+    end
+    row_i = 1
+    for block in blocks(parent(B))
+        nrow = size(block, 1)
+        @views lmul!(LowerTriangular(block), vm[row_i:(row_i + nrow - 1), :])
+        row_i += nrow
+    end
+    return vm
+end
+
 function LinearAlgebra.:\(B::BlockDiagonal, vm::AbstractVecOrMat)
     row_i = 1
     # BlockDiagonals with non-square blocks

diff --git a/test/linalg.jl b/test/linalg.jl
@@ -1,3 +1,4 @@
+using BenchmarkTools
 using BlockDiagonals
 using BlockDiagonals: svd_blockwise, eigen_blockwise
 using LinearAlgebra
@@ -78,7 +79,7 @@ end
                 evals, evecs = eigen(Matrix(B))
 
                 @test E isa Eigen
-                @test Matrix(E) ≈ B 
+                @test Matrix(E) ≈ B
 
                 # There is no test like @test eigen(B) == eigen(Matrix(B))
                 # 1. this fails in the complex case. Probably a convergence thing that leads to ever so slight differences
@@ -88,7 +89,7 @@ end
                 @static if VERSION < v"1.2"
                     # pre-v1.2 we need to sort the eigenvalues consistently
                     # Since eigenvalues may be complex here, I use this function, which works for this test.
-                    # This test is already somewhat fragile w. r. t. degenerate eigenvalues 
+                    # This test is already somewhat fragile w. r. t. degenerate eigenvalues
                     # and this just makes this a little worse.
                     perm_bd = sortperm(real.(evals_bd) + 100*imag.(evals_bd))
                     evals_bd = evals_bd[perm_bd]
@@ -131,7 +132,7 @@ end
                     E = Eigen(block_vals, blocks(vecs)[i])
                     evals_bd, evecs_bd = E
                     evals, evecs = eigen(block)
-                    
+
                     @test block ≈ Matrix(E)
 
                     @static if VERSION < v"1.2"
@@ -144,7 +145,7 @@ end
                         evals = evals[perm]
                         evecs = evecs[:, perm]
                     end
-    
+
                     @test evals_bd ≈ evals
                     @test all(min.(abs.(evecs_bd - evecs), abs.(evecs_bd + evecs)) .< 1e-13)
                 end
@@ -245,6 +246,23 @@ end
             end
         end
     end  # SVD
+    @testset "Left multiplication" begin
+        N1 = 20
+        N2 = 8
+        N3 = 5
+        N4 = N1 + N3 - N2
+        A = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2)])
+        B = BlockDiagonal([rand(rng, N1, N2), rand(rng, N3, N4)])
+        x = rand(rng, N1 + N2)
+        y = rand(rng, N2 + N4)
+
+        @testset "Lower triangular" begin
+            @test lmul!(LowerTriangular(A), copy(x)) ≈ lmul!(LowerTriangular(Matrix(A)), copy(x))
+            @test lmul!(LowerTriangular(B), copy(y)) ≈ lmul!(LowerTriangular(Matrix(B)), copy(y))
+            cx = copy(x)
+            @test 320 >= @ballocated lmul!($(LowerTriangular(A)), $cx)
+        end
+    end
     @testset "Left division" begin
         N1 = 20
         N2 = 8