Allow TriangularRFP for diagonal blocks of L

Project.toml

@@ -24,6 +24,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RectangularFullPacked = "27983f2f-6524-42ba-a408-2b5a31c238e4"
 RegressionFormulae = "545c379f-4ec2-4339-9aea-38f2fb6a8ba2"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -71,6 +72,7 @@ PrecompileTools = "1"
 Printf = "1"
 ProgressMeter = "1.7"
 Random = "1"
+RectangularFullPacked = "0.2.1"
 RegressionFormulae = "0.1.3"
 SparseArrays = "1"
 StableRNGs = "0.1, 1"

ext/MixedModelsForwardDiffExt.jl

@@ -4,7 +4,6 @@ using MixedModels
 using MixedModels: AbstractReMat,
     block,
     BlockedSparse,
-    cholUnblocked!,
     copyscaleinflate!,
     kp1choose2,
     LD,
@@ -140,7 +139,9 @@ function MixedModels.fd_updateL!(A::Vector, L::Vector, reterms::Vector)
     for j in eachindex(reterms) # pre- and post-multiply by Λ, add I to diagonal
         cj = reterms[j]
         diagind = kp1choose2(j)
-        copyscaleinflate!(L[diagind], A[diagind], cj)
+        LdH = L[diagind]
+        LdH = isa(LdH, LowerTriangular) ? Hermitian(LdH.data, :L) : Hermitian(LdH, :L)
+        copyscaleinflate!(LdH, A[diagind], cj)
         for i in (j + 1):(k + 1)     # postmultiply column by Λ
             bij = block(i, j)
             rmulΛ!(copyto!(L[bij], A[bij]), cj)
@@ -151,16 +152,11 @@ function MixedModels.fd_updateL!(A::Vector, L::Vector, reterms::Vector)
     end
     for j in 1:(k + 1)             # blocked Cholesky
         Ljj = L[kp1choose2(j)]
+        LjjH = isa(Ljj, LowerTriangular) ? Hermitian(Ljj.data, :L) : Hermitian(Ljj, :L)
         for jj in 1:(j - 1)
-            fd_rankUpdate!(
-                Hermitian(Ljj, :L),
-                L[block(j, jj)],
-                -one(eltype(Ljj)),
-                one(eltype(Ljj)),
-            )
+            fd_rankUpdate!(LjjH, L[block(j, jj)], -one(eltype(Ljj)), one(eltype(Ljj)))
         end
-        fd_cholUnblocked!(Ljj, Val{:L})
-        LjjT = isa(Ljj, Diagonal) ? Ljj : LowerTriangular(Ljj)
+        fd_cholUnblocked!(LjjH)
         for i in (j + 1):(k + 1)
             Lij = L[block(i, j)]
             for jj in 1:(j - 1)
@@ -172,7 +168,7 @@ function MixedModels.fd_updateL!(A::Vector, L::Vector, reterms::Vector)
                     one(eltype(Lij)),
                 )
             end
-            rdiv!(Lij, LjjT')
+            rdiv!(Lij, Ljj')
         end
     end
     return nothing
@@ -194,26 +190,25 @@ end
 ##### Cholesky factorization
 #####
 
-function MixedModels.fd_cholUnblocked!(A::StridedMatrix, ::Type{Val{:L}})
-    cholesky!(Hermitian(A, :L))
+function MixedModels.fd_cholUnblocked!(D::Hermitian{T,Diagonal{T,Vector{T}}}) where {T}
+    D.data.diag .= sqrt.(D.data.diag)
+    return D
+end
+
+function MixedModels.fd_cholUnblocked!(A::Hermitian{T,Matrix{T}}) where {T}
+    A.uplo == 'L' || throw(ArgumentError("A.uplo should be 'L'"))
+    cholesky!(A)
     return A
 end
 
-function MixedModels.fd_cholUnblocked!(D::UniformBlockDiagonal, ::Type{T}) where {T}
-    Ddat = D.data
+function MixedModels.fd_cholUnblocked!(D::Hermitian{T,UniformBlockDiagonal{T}}) where {T}
+    Ddat = D.data.data
     for k in axes(Ddat, 3)
-        fd_cholUnblocked!(view(Ddat, :, :, k), T)
+        cholesky!(Hermitian(view(Ddat, :, :, k), :L))
     end
     return D
 end
 
-function MixedModels.fd_cholUnblocked!(A::Diagonal, ::Type{T}) where {T}
-    A.diag .= sqrt.(A.diag)
-    return A
-end
-
-MixedModels.fd_cholUnblocked!(A::AbstractMatrix, ::Type{T}) where {T} = cholUnblocked!(A, T)
-
 #####
 ##### Rank update
 #####

src/MixedModels.jl

@@ -27,8 +27,10 @@ using PrecompileTools: PrecompileTools, @setup_workload, @compile_workload
 using Printf: @sprintf
 using ProgressMeter: ProgressMeter, Progress, finish!, next!
 using Random: Random, AbstractRNG, randn!
+using RectangularFullPacked: HermitianRFP, TriangularRFP
+using RectangularFullPacked.LAPACK_RFP: sfrk!
 using RegressionFormulae: fulldummy
-using SparseArrays: SparseArrays, SparseMatrixCSC, SparseVector, dropzeros!, findnz
+using SparseArrays: SparseArrays, SparseMatrixCSC, SparseVector, dropzeros!
 using SparseArrays: nnz, nonzeros, nzrange, rowvals, sparse
 using StaticArrays: StaticArrays, SVector
 using Statistics: Statistics, mean, quantile, std

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}
+    sfrk!(C.transr, C.uplo, 'N', T(α), A, T(β), C.data)
+    return C
+end
+
 """
     _columndot(rv, nz, rngi, rngj)
 
@@ -64,12 +76,24 @@ 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())
+    adim = lower ? 1 : 2
+    if size(C, 1) ≠ size(A, adim)
+        throw(
+            DimensionMismatch(
+                "size(C,1) = $(size(C, 1)) does not match size(A,$adim) = $(size(A, adim))"
+            ),
+        )
+    end
+
+    Cd, rv, nz = C.data, A.rowval, A.nzval
     isone(β) || rmul!(lower ? LowerTriangular(Cd) : UpperTriangular(Cd), β)
     if lower
         @inbounds for jj in axes(A, 2)
@@ -96,10 +120,65 @@ 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)
+    if uppercase(C.transr) ≠ 'N' || uppercase(C.uplo) ≠ 'L'
+        throw(ArgumentError("HermitianRFP C is not non-trans, lower"))
+    end
+    rv, nz = A.rowval, A.nzval
+    Cdat = C.data
+    An = size(A, 1)
+    if An ≠ size(C, 1)
+        throw(DimensionMismatch("size(A, 1) == $An ≠ $(size(C, 1)) = size(C, 1)"))
+    end
+    tall = iseven(An)        # is Cdat in the tall format or the wide format?
+    Cdn, Cdm = size(Cdat)
+    @assert (Cdn == An + tall) && (Cdm == ((An + !tall) >> 1))
+
+    isone(β) || rmul!(Cdat, β)
+    for jj in axes(A, 2)
+        rngjj = nzrange(A, jj)
+        lenrngjj = length(rngjj)
+        for (k, j) in enumerate(rngjj)
+            anzj = α * nz[j]
+            rvj = rv[j]                         # updates in rvj column of C
+            if rvj ≤ Cdm                        # in the trapezoidal part of Cdat
+                offset = (rvj - 1) * Cdn + tall
+                for i in k:lenrngjj
+                    kk = rngjj[i]
+                    linind = offset + rv[kk]
+                    Cdat[linind] = muladd(nz[kk], anzj, Cdat[linind])
+                end
+            else                                # in the transposed triangular part of Cdat
+                Cdrow = rvj - Cdm
+                for i in k:lenrngjj
+                    kk = rngjj[i]
+                    rvkk = rv[kk]
+                    @assert rvkk ≥ rvj
+                    linind = (rvkk - Cdm - tall) * Cdn + Cdrow
+                    Cdat[linind] = muladd(nz[kk], anzj, Cdat[linind])
+                end
+            end
+        end
+    end
+    return C
+end
+
 function rankUpdate!(C::HermOrSym, A::BlockedSparse, α, β)
     return rankUpdate!(C, sparse(A), α, β)
 end
 
+# HermitianRFP is not a HermOrSym, so it needs its own BlockedSparse method that
+# forwards to the HermitianRFP/SparseMatrixCSC kernel.
+function rankUpdate!(C::HermitianRFP{T}, A::BlockedSparse{T}, α, β) where {T}
+    return rankUpdate!(C, sparse(A), α, β)
+end
+
 function rankUpdate!(
     C::HermOrSym{T,Diagonal{T,Vector{T}}}, A::StridedMatrix{T}, α, β
 ) where {T}