Repository: JuliaStats/MixedModels.jl (julia)

Issue:
Be Diagonal where you can
We should allow the λ field in an `ReMat` in the `MixedModels` package to be `LowerTriangular{T,Union{Diagonal{T},Matrix{T}}}` and check if the `inds` make the result diagonal.  For example, `zerocorr` applied to an `ReTerm` results in a diagonal λ.  A surprisingly large amount of the time in `updateL!` is spent in the mulitplications by Λ, especially for a `UniformBlockDiagonal` (1, 1) block.

_Originally posted by @dmbates in https://github.com/RePsychLing/MixedModelsSim.jl/pull/42#issuecomment-801161864_

Pull request:
Closes #487 .

I had to loosen some type constraints in various places to make this work and I didn't add in any additional `mul!`  methods, so I don't know if we're fully taking advantages of this. @dmbates Feel free to add those in if you want, but those wouldn't be breaking changes.

I also changed `rownormalize!` to be non mutating. This made specializing for `Diagonal` easier and actually cut out an extra allocation.

Patch:
diff --git a/.github/workflows/Future.yml b/.github/workflows/Future.yml
index 2a86b4c24..270237de5 100644
--- a/.github/workflows/Future.yml
+++ b/.github/workflows/Future.yml
@@ -1,12 +1,13 @@
 name: Future
 on:
-  pull_request:
-    branches:
-      - master
-      - v3.x
-    paths-ignore:
-      - 'LICENSE.md'
-      - 'README.md'
+  # pull_request:
+  #   branches:
+  #     - master
+  #     - v3.x
+  #   paths-ignore:
+  #     - 'LICENSE.md'
+  #     - 'README.md'
+  workflow_dispatch:
 jobs:
   ci:
     runs-on: ${{ matrix.os }}
diff --git a/NEWS.md b/NEWS.md
index 37b976fa7..5aa63403d 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -14,6 +14,10 @@ MixedModels v4.0.0 Release Notes
   rank of `X` by `statsrank`. [#479]
 * Artifacts are now loaded lazily: the test data loaded via `dataset` is
   downloaded on first use [#486]
+* `ReMat` and `PCA` now support covariance factors (`λ`) that are `LowerTriangular`
+  or `Diagonal`. This representation is both more memory efficient and
+  enables additional computational optimizations for particular covariance
+  structures.[#489]
 
 Run-time formula syntax
 -----------------------
@@ -193,3 +197,4 @@ Package dependencies
 [#482]: https://github.com/JuliaStats/MixedModels.jl/issues/482
 [#484]: https://github.com/JuliaStats/MixedModels.jl/issues/484
 [#486]: https://github.com/JuliaStats/MixedModels.jl/issues/486
+[#489]: https://github.com/JuliaStats/MixedModels.jl/issues/489
diff --git a/src/pca.jl b/src/pca.jl
index 8b82df06b..710581c16 100644
--- a/src/pca.jl
+++ b/src/pca.jl
@@ -11,8 +11,8 @@ Principal Components Analysis
 * `corr` is this a correlation matrix?
 """
 struct PCA{T<:AbstractFloat}
-    covcor::Symmetric{T,Matrix{T}}
-    sv::SVD{T,T,Matrix{T}}
+    covcor::Symmetric{T,<:AbstractMatrix{T}}
+    sv::SVD{T,T,<:AbstractMatrix{T}}
     rnames::Union{Vector{String}, Missing}
     corr::Bool
 end
@@ -30,7 +30,7 @@ If `corr=true`, then the covariance is first standardized to the correlation sca
 """
 
 function PCA(covfac::AbstractMatrix, rnames=missing; corr::Bool=true)
-    covf = corr ? rownormalize!(copy(covfac)) : copy(covfac)
+    covf = corr ? rownormalize(covfac) : covfac
     PCA(Symmetric(covf*covf', :L), svd(covf), rnames, corr)
 end
 
@@ -108,13 +108,12 @@ function Base.show(io::IO, ::MIME"text/plain", pca::PCA;
     if loadings
         println(io, "\nComponent loadings")
         printmat = round.(pca.loadings, digits=ndigitsmat)
-
         if pca.rnames !== missing
             pclabs = [Text(""); Text.( "PC$i" for i in 1:length(pca.rnames))]
             pclabs = reshape(pclabs, 1, :)
             # this hurts type stability,
             # but this show method shouldn't be a bottleneck
-            printmat = [pclabs; Text.(pca.rnames) printmat]
+            printmat = [pclabs; Text.(pca.rnames) Matrix(printmat)]
         end
 
         Base.print_matrix(io, printmat)
diff --git a/src/remat.jl b/src/remat.jl
index 419964c3a..5f21c4569 100644
--- a/src/remat.jl
+++ b/src/remat.jl
@@ -22,7 +22,7 @@ mutable struct ReMat{T,S} <: AbstractReMat{T}
     cnames::Vector{String}
     z::Matrix{T}
     wtz::Matrix{T}
-    λ::LowerTriangular{T,Matrix{T}}
+    λ::Union{LowerTriangular{T,Matrix{T}},Diagonal{T,Vector{T}}}
     inds::Vector{Int}
     adjA::SparseMatrixCSC{T,Int32}
     scratch::Matrix{T}
@@ -64,8 +64,11 @@ function _amalgamate(reterms::Vector, T::Type)
                 offset += sz
             end
             inds = (1:abs2(Snew))[vec(btemp)]
-
-            λ = LowerTriangular(Matrix{T}(I, Snew, Snew))
+            if inds == diagind(btemp)
+                λ = Diagonal{T}(I(Snew))
+            else
+                λ = LowerTriangular(Matrix{T}(I, Snew, Snew))
+            end
             scratch =  foldl(vcat, rr.scratch for rr in trms)
 
             push!(
@@ -127,7 +130,7 @@ Overwrite `v` with the elements of the blocks in the lower triangle of `A.Λ` (c
 """
 function getθ!(v::AbstractVector{T}, A::ReMat{T}) where {T}
     length(v) == length(A.inds) || throw(DimensionMismatch("length(v) ≠ length(A.inds)"))
-    m = A.λ.data
+    m = A.λ
     @inbounds for (j, ind) in enumerate(A.inds)
         v[j] = m[ind]
     end
@@ -173,8 +176,9 @@ Return the vector of lower bounds on the parameters, `θ` associated with `A`
 These are the elements in the lower triangle of `A.λ` in column-major ordering.
 Diagonals have a lower bound of `0`.  Off-diagonals have a lower-bound of `-Inf`.
 """
-lowerbd(A::ReMat{T}) where {T} =
-    T[x ∈ diagind(A.λ.data) ? zero(T) : T(-Inf) for x in A.inds]
+function lowerbd(A::ReMat{T}) where {T}
+    T[x ∈ diagind(A.λ) ? zero(T) : T(-Inf) for x in A.inds]
+end
 
 """
     isnested(A::ReMat, B::ReMat)
@@ -451,8 +455,8 @@ end
 rowlengths(A::ReMat{T,1}) where {T} = vec(abs.(A.λ.data))
 
 function rowlengths(A::ReMat)
-    ld = A.λ.data
-    [norm(view(ld, i, 1:i)) for i in 1:size(ld, 1)]
+    ld = A.λ
+    isa(ld, Diagonal) ? abs.(ld.diag) : [norm(view(ld, i, 1:i)) for i in 1:size(ld, 1)]
 end
 
 """
@@ -538,15 +542,24 @@ function ρ(i, λ::AbstractMatrix{T}, im::Matrix{Bool}, indpairs, σs, sc::T)::T
     end
 end
 
-function σρs(A::ReMat{T}, sc::T) where {T}
-    λ = A.λ.data
+function _σρs(λ::LowerTriangular{T}, sc::T, im::Matrix{Bool}, cnms::Vector{Symbol}) where {T}
+    λ = λ.data
     k = size(λ, 1)
-    im = indmat(A)
     indpairs = checkindprsk(k)
-    σs = NamedTuple{(Symbol.(A.cnames)...,)}(ntuple(i -> sc*norm(view(λ,i,1:i)), k))
+    σs = NamedTuple{(cnms...,)}(ntuple(i -> sc*norm(view(λ,i,1:i)), k))
     NamedTuple{(:σ,:ρ)}((σs, ntuple(i -> ρ(i,λ,im,indpairs,σs,sc), (k * (k - 1)) >> 1)))
 end
 
+function _σρs(λ::Diagonal{T}, sc::T, im::Matrix{Bool}, cnms::Vector{Symbol}) where {T}
+    dsc = sc .* λ.diag
+    k = length(dsc)
+    σs = NamedTuple{(cnms...,)}(NTuple{k, T}(dsc))
+    NamedTuple{(:σ,:ρ)}((σs, ntuple(i -> -zero(T), (k * (k - 1)) >> 1)))
+end
+
+function σρs(A::ReMat{T}, sc::T) where {T}
+    _σρs(A.λ, sc, indmat(A), Symbol.(A.cnames))
+end
 """
     corrmat(A::ReMat)
 
@@ -589,9 +602,7 @@ end
 vsize(A::ReMat{T,S}) where {T,S} = S
 
 function zerocorr!(A::ReMat{T}) where {T}
-    λ = A.λ
-    # zero out all entries not on the diagonal
-    λ[setdiff(A.inds, diagind(λ))] .= 0
+    λ = A.λ = Diagonal(A.λ)
     A.inds = intersect(A.inds, diagind(λ))
     A
 end
diff --git a/src/utilities.jl b/src/utilities.jl
index 9ef6877c2..06afa8824 100644
--- a/src/utilities.jl
+++ b/src/utilities.jl
@@ -81,7 +81,8 @@ function Base.sum!(s::AbstractVector{T}, a::RaggedArray{T}) where {T}
     s
 end
 
-function rownormalize!(A::AbstractMatrix)
+function rownormalize(A::AbstractMatrix)
+    A = copy(A)
     for r in eachrow(A)
         # all zeros arise in zerocorr situations
         if !iszero(r)
@@ -91,6 +92,10 @@ function rownormalize!(A::AbstractMatrix)
     A
 end
 
+function rownormalize(A::LowerTriangular{T, Diagonal{T, Vector{T}}}) where T
+    one(T) * I(size(A,1))
+end
+
 """
     replicate(f::Function, n::Integer; use_threads=false)
 


Test patch:
diff --git a/test/pls.jl b/test/pls.jl
index 0163f3a5f..ef99450f7 100644
--- a/test/pls.jl
+++ b/test/pls.jl
@@ -382,6 +382,16 @@ end
     # amalgamate should set these to -0.0 to indicate structural zeros
     @test all(ρs_intercept .=== -0.0)
 
+    @testset "diagonal λ in zerocorr" begin
+        # explicit zerocorr
+        fmzc = models(:sleepstudy)[2]
+        λ = first(fmzc.reterms).λ
+        @test λ isa Diagonal{Float64, Vector{Float64}}
+        # implicit zerocorr via almagation
+        fmnc = models(:sleepstudy)[3]
+        λ = first(fmnc.reterms).λ
+        @test λ isa Diagonal{Float64, Vector{Float64}}
+    end
 
     show(io, BlockDescription(first(models(:sleepstudy))))
     @test countlines(seekstart(io)) == 3