[wIP] Allow TriangularRFP for diagonal blocks of L [ci skip]

Project.toml

@@ -19,6 +19,7 @@ PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RectangularFullPacked = "27983f2f-6524-42ba-a408-2b5a31c238e4"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -57,6 +58,7 @@ PooledArrays = "0.5, 1"
 PrecompileTools = "1"
 ProgressMeter = "1.7"
 Random = "1"
+RectangularFullPacked = "0.2.1"
 SparseArrays = "1"
 StableRNGs = "0.1, 1"
 StaticArrays = "0.11, 0.12, 1"

src/MixedModels.jl

@@ -25,6 +25,7 @@ using PooledArrays: PooledArrays, PooledArray
 using PrecompileTools: PrecompileTools, @setup_workload, @compile_workload
 using ProgressMeter: ProgressMeter, Progress, ProgressUnknown, finish!, next!
 using Random: Random, AbstractRNG, randn!
+using RectangularFullPacked: HermitianRFP, TriangularRFP
 using SparseArrays: SparseArrays, SparseMatrixCSC, SparseVector, dropzeros!, nnz
 using SparseArrays: nonzeros, nzrange, rowvals, sparse
 using StaticArrays: StaticArrays, SVector
@@ -214,7 +215,6 @@ include("profile/profile.jl")
 include("nlopt.jl")
 include("prima.jl")
 
-
 # aliases with non-unicode function names
 const settheta! = setθ!
 const profilesigma = profileσ

src/blockdescription.jl

@@ -32,11 +32,17 @@ end
 
 BlockDescription(m::GeneralizedLinearMixedModel) = BlockDescription(m.LMM)
 
-shorttype(::UniformBlockDiagonal, ::UniformBlockDiagonal) = "BlkDiag"
-shorttype(::UniformBlockDiagonal, ::Matrix) = "BlkDiag/Dense"
+function shorttype(
+    ::UniformBlockDiagonal, ::LowerTriangular{T,UniformBlockDiagonal{T}}
+) where {T}
+    return "BlkDiag"
+end
+shorttype(::UniformBlockDiagonal, ::LowerTriangular) = "BlkDiag/Dense"
 shorttype(::SparseMatrixCSC, ::BlockedSparse) = "Sparse"
 shorttype(::Diagonal, ::Diagonal) = "Diagonal"
 shorttype(::Diagonal, ::Matrix) = "Diag/Dense"
+shorttype(::Diagonal, ::LowerTriangular) = "Diag/Dense"
+shorttype(::Diagonal, ::TriangularRFP) = "Diag/TrRFP"
 shorttype(::Matrix, ::Matrix) = "Dense"
 shorttype(::SparseMatrixCSC, ::SparseMatrixCSC) = "Sparse"
 shorttype(::SparseMatrixCSC, ::Matrix) = "Sparse/Dense"

src/linalg.jl

@@ -46,7 +46,7 @@ function LinearAlgebra.mul!(
 end
 
 function LinearAlgebra.ldiv!(
-    A::UpperTriangular{T,<:Adjoint{T,UniformBlockDiagonal{T}}}, B::StridedVector{T}
+    A::UpperTriangular{T,Adjoint{T,UniformBlockDiagonal{T}}}, B::StridedVector{T}
 ) where {T}
     adjA = A.data
     length(B) == size(A, 2) || throw(DimensionMismatch(""))
@@ -60,7 +60,8 @@ function LinearAlgebra.ldiv!(
 end
 
 function LinearAlgebra.rdiv!(
-    A::Matrix{T}, B::UpperTriangular{T,<:Adjoint{T,UniformBlockDiagonal{T}}}
+    A::Matrix{T},
+    B::UpperTriangular{T,Adjoint{T,UniformBlockDiagonal{T}}},
 ) where {T}
     m, n = size(A)
     Bd = B.data.parent
@@ -78,7 +79,8 @@ function LinearAlgebra.rdiv!(
 end
 
 function LinearAlgebra.rdiv!(
-    A::BlockedSparse{T,S,P}, B::UpperTriangular{T,<:Adjoint{T,UniformBlockDiagonal{T}}}
+    A::BlockedSparse{T,S,P},
+    B::UpperTriangular{T,Adjoint{T,UniformBlockDiagonal{T}}},
 ) where {T,S,P}
     Bpd = B.data.parent
     Bdat = Bpd.data

src/linalg/cholUnblocked.jl

@@ -8,17 +8,21 @@ because these are part of the inner calculations in a blocked Cholesky factoriza
 """
 function cholUnblocked! end
 
-function cholUnblocked!(D::Diagonal{T}, ::Type{Val{:L}}) where {T<:AbstractFloat}
-    Ddiag = D.diag
+function cholUnblocked!(D::Hermitian{T,Diagonal{T,Vector{T}}}) where {T}
+    Ddiag = D.data.diag
     @inbounds for i in eachindex(Ddiag)
         (ddi = Ddiag[i]) ≤ zero(T) && throw(PosDefException(i))
         Ddiag[i] = sqrt(ddi)
     end
-
     return D
 end
 
-function cholUnblocked!(A::StridedMatrix{T}, ::Type{Val{:L}}) where {T<:BlasFloat}
+function cholUnblocked!(A::Hermitian{T,Matrix{T}}) where {T}
+    A.uplo == 'L' || throw(ArgumentError("A.uplo should be 'L'"))
+    return cholUnblocked!(A.data)
+end
+
+function cholUnblocked!(A::StridedMatrix{T}) where {T<:BlasFloat}
     n = LinearAlgebra.checksquare(A)
     if n == 1
         A[1] < zero(T) && throw(PosDefException(1))
@@ -36,10 +40,12 @@ function cholUnblocked!(A::StridedMatrix{T}, ::Type{Val{:L}}) where {T<:BlasFloa
     return A
 end
 
-function cholUnblocked!(D::UniformBlockDiagonal, ::Type{Val{:L}})
-    Ddat = D.data
+function cholUnblocked!(D::Hermitian{T,UniformBlockDiagonal{T}}) where {T}
+    Ddat = D.data.data
     for k in axes(Ddat, 3)
-        cholUnblocked!(view(Ddat, :, :, k), Val{:L})
+        cholUnblocked!(view(Ddat, :, :, k))
     end
     return D
 end
+
+cholUnblocked!(D::HermitianRFP) = LinearAlgebra.cholesky!(D)

src/linalg/logdet.jl

@@ -7,13 +7,19 @@ Return `log(det(tril(A)))` evaluated in place.
 """
 LD(d::Diagonal{T}) where {T<:Number} = sum(log, d.diag)
 
-function LD(d::UniformBlockDiagonal{T}) where {T}
+function LD(d::LowerTriangular{T,UniformBlockDiagonal{T}}) where {T}
     dat = d.data
     return sum(log, dat[j, j, k] for j in axes(dat, 2), k in axes(dat, 3))
 end
 
 LD(d::DenseMatrix{T}) where {T} = @inbounds sum(log, d[k] for k in diagind(d))
 
+LD(d::LowerTriangular{T,Matrix{T}}) where {T} = LD(d.data)
+
+function LD(d::TriangularRFP{T}) where {T}
+    return sum(log, d[j, j] for j in axes(d, 2))
+end
+
 """
     logdet(m::LinearMixedModel)
 

src/linalg/rankUpdate.jl

@@ -3,7 +3,7 @@
     rankUpdate!(C, A, α)
     rankUpdate!(C, A, α, β)
 
-A rank-k update, C := α*A'A + β*C, of a Hermitian (Symmetric) matrix.
+A rank-k update, C := α*A*A' + β*C, of a Hermitian (Symmetric) matrix.
 
 `α` and `β` both default to 1.0.  When `α` is -1.0 this is a downdate operation.
 The name `rankUpdate!` is borrowed from [https://github.com/andreasnoack/LinearAlgebra.jl]
@@ -39,6 +39,18 @@ function rankUpdate!(C::HermOrSym{T,S}, A::StridedMatrix{T}, α, β) where {T,S}
     return C
 end
 
+# function rankUpdate!(
+#     C::HermOrSym{T,S}, A::StridedMatrix{T}, α, β
+# ) where {T,S<:LowerTriangular}
+#     BLAS.syrk!(C.uplo, 'N', T(α), A, T(β), C.data.data)
+#     return C
+# end
+
+# function rankUpdate!(C::HermitianRFP{T}, A::StridedMatrix{T}, α, β) where {T}
+#     BLAS.syrk!('N', T(α), A, T(β), C)
+#     return C
+# end
+
 """
     _columndot(rv, nz, rngi, rngj)
 
@@ -64,12 +76,17 @@ function _columndot(rv, nz, rngi, rngj)
     return accum
 end
 
-function rankUpdate!(C::HermOrSym{T,S}, A::SparseMatrixCSC{T}, α, β) where {T,S}
+function rankUpdate!(
+    C::HermOrSym{T,S},
+    A::SparseMatrixCSC{T},
+    α,
+    β,
+) where {T,S}
     require_one_based_indexing(C, A)
     m, n = size(A)
     Cd, rv, nz = C.data, A.rowval, A.nzval
     lower = C.uplo == 'L'
-    (lower ? m : n) == size(C, 2) || throw(DimensionMismatch())
+    #    (lower ? m : n) == size(C, 2) || throw(DimensionMismatch())  # Doesn't make sense with lower.  
     isone(β) || rmul!(lower ? LowerTriangular(Cd) : UpperTriangular(Cd), β)
     if lower
         @inbounds for jj in axes(A, 2)
@@ -96,6 +113,38 @@ function rankUpdate!(C::HermOrSym{T,S}, A::SparseMatrixCSC{T}, α, β) where {T,
     return C
 end
 
+function rankUpdate!(
+    C::HermitianRFP{T},
+    A::SparseMatrixCSC{T},
+    α,
+    β,
+) where {T}
+    require_one_based_indexing(C, A)
+    Ctr = TriangularRFP(C.data, C.transr, C.uplo)  # need TriangularRFP for write access
+    rv, nz = A.rowval, A.nzval
+    if size(A, 1) ≠ size(C, 1)
+        throw(DimensionMismatch("size(A, 1) == $(size(A,1)) ≠ $(size(C, 1)) = size(C, 1)"))
+    end
+    isone(β) || rmul!(Ctr, β)
+    if C.uplo == 'L'
+        @inbounds for jj in axes(A, 2)
+            rangejj = nzrange(A, jj)
+            lenrngjj = length(rangejj)
+            for (k, j) in enumerate(rangejj)
+                anzj = α * nz[j]
+                rvj = rv[j]
+                for i in k:lenrngjj
+                    kk = rangejj[i]
+                    Ctr[rv[kk], rvj] = muladd(nz[kk], anzj, Ctr[rv[kk], rvj])
+                end
+            end
+        end
+    else
+        throw(ArgumentError("rankUpdate! requires C be stored in the lower triangle"))
+    end
+    return C
+end
+
 function rankUpdate!(C::HermOrSym, A::BlockedSparse, α, β)
     return rankUpdate!(C, sparse(A), α, β)
 end