diff --git a/NEWS.md b/NEWS.md index 19db53dd6..f6e3b95bb 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,9 @@ +MixedModels v5.8.0 Release Notes +============================== +- Allow for diagonal blocks of `L` to be stored in `RectangularFullPacked` (RFP) format, which saves roughly half the storage required for the block. This can increase the time required for `updateL!`, primarily in the `rankUpdate!` step, resulting in a time vs. memory tradeoff. The size threshold for RFP storage is a new optional argument `RFPthreshold`, which defaults to 1000. +- The RFP format stores a triangular matrix in two pieces: a trapezoidal part of roughly 3/4 of the elements, where linear indexing can be used for the updates, and a transposed triangular part with more complicated `getindex` and `setindex!` methods. +- A new Boolean optional argument, `sortlevels`, which defaults to `true`, allows for sorting the levels of the any grouping factors with RFP storage of their diagonal blocks. This is a heuristic to have more updates occur in the trapezoid part of the RFP block. It is not guaranteed to be optimal but it works well in examples we have tried. [#821] + MixedModels v5.7.1 Release Notes ============================== - Compat bump for MixedModelsDatasets. Note that some data values have changed in their least significant digits, which can change statistics computed from these. @@ -741,6 +747,7 @@ Package dependencies [#810]: https://github.com/JuliaStats/MixedModels.jl/issues/810 [#814]: https://github.com/JuliaStats/MixedModels.jl/issues/814 [#815]: https://github.com/JuliaStats/MixedModels.jl/issues/815 +[#821]: https://github.com/JuliaStats/MixedModels.jl/issues/821 [#823]: https://github.com/JuliaStats/MixedModels.jl/issues/823 [#825]: https://github.com/JuliaStats/MixedModels.jl/issues/825 [#828]: https://github.com/JuliaStats/MixedModels.jl/issues/828 diff --git a/Project.toml b/Project.toml index 043fb3022..63ceff005 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "MixedModels" uuid = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316" author = ["Phillip Alday ", "Douglas Bates "] -version = "5.7.1" +version = "5.8.0" [workspace] projects = ["docs"] @@ -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" diff --git a/ext/MixedModelsForwardDiffExt.jl b/ext/MixedModelsForwardDiffExt.jl index c3263e620..c9e54afab 100644 --- a/ext/MixedModelsForwardDiffExt.jl +++ b/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 ##### diff --git a/src/MixedModels.jl b/src/MixedModels.jl index 7e9e25aec..389c7faa3 100644 --- a/src/MixedModels.jl +++ b/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 diff --git a/src/blockdescription.jl b/src/blockdescription.jl index 9305d1ae2..efc9b682a 100644 --- a/src/blockdescription.jl +++ b/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" diff --git a/src/linalg.jl b/src/linalg.jl index 91e628153..64ab9e331 100644 --- a/src/linalg.jl +++ b/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 diff --git a/src/linalg/cholUnblocked.jl b/src/linalg/cholUnblocked.jl index b430e324a..36b8eb4f5 100644 --- a/src/linalg/cholUnblocked.jl +++ b/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) diff --git a/src/linalg/logdet.jl b/src/linalg/logdet.jl index 5f2c5a564..49fa19a15 100644 --- a/src/linalg/logdet.jl +++ b/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) diff --git a/src/linalg/rankUpdate.jl b/src/linalg/rankUpdate.jl index 705f87d53..a14598bfb 100644 --- a/src/linalg/rankUpdate.jl +++ b/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} diff --git a/src/linearmixedmodel.jl b/src/linearmixedmodel.jl index ef1b398de..e70039b1b 100644 --- a/src/linearmixedmodel.jl +++ b/src/linearmixedmodel.jl @@ -26,6 +26,27 @@ Linear mixed-effects model representation * `u`: random effects on the orthogonal scale, as a vector of matrices * `X`: the fixed-effects model matrix * `y`: the response vector + +## Named Arguments + +The `LinearMixedModel(f, tbl; ...)` constructor (and `fit(LinearMixedModel, f, tbl; ...)`) +accept the following keyword arguments: + +* `contrasts`: a `Dict` mapping column names to contrast codings, passed to schema application. +* `weights`: a vector of prior case weights; defaults to equal weights. +* `σ`: a fixed value for the residual standard deviation, or `nothing` (the default) to estimate it. +* `amalgamate`: whether to combine random-effects terms that share a grouping factor into a + single term (default `true`). +* `RFPthreshold`: the column-count threshold above which a diagonal block of `L` that incurs + fill-in (from non-nested crossed grouping factors) is stored in Rectangular Full Packed form + (`TriangularRFP`) rather than as a dense `LowerTriangular` (default `1000`). RFP roughly halves + the storage of such a block, but slows computation slightly (i.e. trades speed for memory); it does not change the fitted model. +* `sortlevels`: whether to reorder the levels of each grouping factor (other than the leading one) + by descending number of occurrences (default `true`). This concentrates the sparse rank-update + of a filled-in diagonal block in a compact region of storage, improving memory locality. This is most + helpful for large blocks that do not fit in cache and for `TriangularRFP` blocks. It leaves the + objective and estimates unchanged; only the internal level order (and hence the order of, e.g., + `raneftables` rows) differs. """ struct LinearMixedModel{T<:AbstractFloat} <: MixedModel{T} formula::FormulaTerm @@ -51,7 +72,7 @@ const _MISSING_RE_ERROR = ArgumentError( function LinearMixedModel( f::FormulaTerm, tbl::Tables.ColumnTable; contrasts=Dict{Symbol,Any}(), wts=nothing, weights=[], - σ=nothing, amalgamate=true, + σ=nothing, amalgamate=true, RFPthreshold=1000, sortlevels=true, ) fvars = StatsModels.termvars(f) tvars = Tables.columnnames(tbl) @@ -80,7 +101,7 @@ function LinearMixedModel( y, Xs = modelcols(form, tbl) - return LinearMixedModel(y, Xs, form, weights, σ, amalgamate) + return LinearMixedModel(y, Xs, form, weights, σ, amalgamate, RFPthreshold, sortlevels) end """ @@ -103,13 +124,16 @@ function LinearMixedModel( weights=[], σ=nothing, amalgamate=true, + RFPthreshold=1000, + sortlevels=true, ) T = promote_type(Float64, float(eltype(y))) # ensure eltype of model matrices is at least Float64 reterms, feterms = _split_re_fe_terms(Xs, form, T) isempty(reterms) && throw(_MISSING_RE_ERROR) return LinearMixedModel( - convert(Array{T}, y), only(feterms), reterms, form, weights, σ, amalgamate + convert(Array{T}, y), only(feterms), reterms, form, weights, σ, amalgamate, + RFPthreshold, sortlevels, ) end @@ -146,7 +170,7 @@ function _split_re_fe_terms(Xs::Tuple, form::FormulaTerm, ::Type{T}) where {T} end """ - LinearMixedModel(y, feterm, reterms, form, weights=[], σ=nothing; amalgamate=true) + LinearMixedModel(y, feterm, reterms, form, weights=[], σ=nothing; amalgamate=true, RFPthreshold=1000, sortlevels=true) Private constructor for a `LinearMixedModel` given already assembled fixed and random effects. @@ -155,6 +179,13 @@ model matrix and response), a vector of `AbstractReMat` (the random-effects model matrices), the formula and the weights. Everything else in the structure can be derived from these quantities. +When `sortlevels` is `true`, [`_sortlevels!`](@ref) is applied to every non-leading term (and +to a leading term that shares its grouping factor with a later term, as can happen with +`amalgamate=false`), so that all terms for a grouping factor use one frequency-descending level +order; the stored `form` is then updated via `_syncgrouping` to keep its grouping terms in sync. +`RFPthreshold` is forwarded to `createAL` to select Rectangular Full Packed storage for +sufficiently large filled-in diagonal blocks of `L`. See [`LinearMixedModel`](@ref) for details. + !!! note This method is internal and experimental and so may change or disappear in a future release without being considered a breaking change. @@ -167,6 +198,8 @@ function LinearMixedModel( weights=[], σ=nothing, amalgamate=true, + RFPthreshold=1000, + sortlevels=true, ) where {T} # detect and combine RE terms with the same grouping var if length(reterms) > 1 && amalgamate @@ -176,6 +209,17 @@ function LinearMixedModel( end sort!(reterms; by=nranef, rev=true) + if sortlevels + # only blocks after the first incur fill-in; the leading term's + # diagonal block is unaffected by level order, but sort it anyway + # when it shares its grouping factor with a later term (possible with + # amalgamate=false) so that all terms for a factor use one level order + tosort = Set(fname(reterms[i]) for i in 2:length(reterms)) + for rt in reterms + fname(rt) in tosort && _sortlevels!(rt) + end + form = _syncgrouping(form, reterms) + end Xy = FeMat(feterm, vec(y)) # Replace feterm's fullrankx field with a view into the shared Xymat storage, # eliminating the duplicate allocation for the full-rank X columns. @@ -190,7 +234,7 @@ function LinearMixedModel( sqrtwts = map!(sqrt, Vector{T}(undef, length(weights)), weights) reweight!.(reterms, Ref(sqrtwts)) reweight!(Xy, sqrtwts) - A, L = createAL(reterms, Xy) + A, L = createAL(reterms, Xy; RFPthreshold) θ = foldl(vcat, getθ(c) for c in reterms) optsum = OptSummary(θ) optsum.sigma = isnothing(σ) ? nothing : T(σ) @@ -230,8 +274,12 @@ function StatsAPI.fit(::Type{LinearMixedModel}, contrasts=Dict{Symbol,Any}(), σ=nothing, amalgamate=true, + RFPthreshold=1000, + sortlevels=true, kwargs...) - lmod = LinearMixedModel(f, tbl; contrasts, weights, wts, σ, amalgamate) + lmod = LinearMixedModel( + f, tbl; contrasts, weights, wts, σ, amalgamate, RFPthreshold, sortlevels + ) return fit!(lmod; kwargs...) end @@ -402,9 +450,13 @@ function _pushALblock!(A, L, blk) return push!(A, deepcopy(isa(blk, BlockedSparse) ? blk.cscmat : blk)) end -function createAL(reterms::Vector{<:AbstractReMat{T}}, Xy::FeMat{T}) where {T} +function createAL( + reterms::Vector{<:AbstractReMat{T}}, + Xy::FeMat{T}; + RFPthreshold::Int=1000, +) where {T} k = length(reterms) - vlen = kchoose2(k + 1) + vlen = kp1choose2(k + 1) A = sizehint!(AbstractMatrix{T}[], vlen) L = sizehint!(AbstractMatrix{T}[], vlen) for i in eachindex(reterms) @@ -412,16 +464,23 @@ function createAL(reterms::Vector{<:AbstractReMat{T}}, Xy::FeMat{T}) where {T} _pushALblock!(A, L, densify(reterms[i]' * reterms[j])) end end - for j in eachindex(reterms) # can't fold this into the previous loop b/c block order + for j in eachindex(reterms) # can't fold this into the previous loop b/c block order _pushALblock!(A, L, densify(Xy' * reterms[j])) end _pushALblock!(A, L, densify(Xy'Xy)) - for i in 2:k # check for fill-in due to non-nested grouping factors + for i in 2:k # check for fill-in due to non-nested grouping factors ci = reterms[i] for j in 1:(i - 1) cj = reterms[j] if !isnested(cj, ci) - for l in i:k + ind = kp1choose2(i) # location of i'th diagonal block + Li = collect(L[ind]) + L[ind] = if size(Li, 2) > RFPthreshold + TriangularRFP(Li, :L) + else + LowerTriangular(Li) + end + for l in (i + 1):k ind = block(l, i) L[ind] = Matrix(L[ind]) end @@ -429,7 +488,14 @@ function createAL(reterms::Vector{<:AbstractReMat{T}}, Xy::FeMat{T}) where {T} end end end - return identity.(A), identity.(L) + for k in axes(reterms, 1) # patch up the UniformBlockDiagonal in L for nested factors + diagind = kp1choose2(k) + Ldi = L[diagind] + if isa(Ldi, UniformBlockDiagonal) + L[diagind] = LowerTriangular(Ldi) + end + end + return identity.(A), identity.(L) # does anyone remember what the `identity` is for?` end function StatsAPI.cooksdistance(model::LinearMixedModel) @@ -775,8 +841,12 @@ end # use dispatch to distinguish Diagonal and UniformBlockDiagonal in first(L) _ldivB1!(B1::Diagonal{T}, rhs::AbstractVector{T}, ind) where {T} = rhs ./= B1.diag[ind] -function _ldivB1!(B1::UniformBlockDiagonal{T}, rhs::AbstractVector{T}, ind) where {T} - return ldiv!(LowerTriangular(view(B1.data, :, :, ind)), rhs) +function _ldivB1!( + B1::LowerTriangular{T,UniformBlockDiagonal{T}}, + rhs::AbstractVector{T}, + ind, +) where {T} + return ldiv!(LowerTriangular(view(B1.data.data, :, :, ind)), rhs) end """ @@ -1185,42 +1255,55 @@ end Base.show(io::IO, m::LinearMixedModel) = Base.show(io, MIME("text/plain"), m) """ - _findnz(A::AbstractMatrix) + _coord(A::AbstractMatrix) -Return the positions and values of the nonzeros in `A` as an (I, J, V) tuple - -When possible, methods for this generic pass through to methods for `SparseArrays.findnz`. -The exceptions are the `Matrix` and `LinearAlgebra.Diagonal` classes where our defining a -`findnz` method would be type piracy. +Return the positions and values of the nonzeros in `A` as a `TypedTables.Table` with columns `i`, `j`, and `v` """ - -function _findnz(A::Matrix) - m, n = size(A) - return ( - repeat(axes(A, 1); outer=n), - repeat(axes(A, 2); inner=m), - vec(A), - ) +function _coord(A::Diagonal) + return Table(; i=Int32.(axes(A, 1)), j=Int32.(axes(A, 2)), v=A.diag) end -function _findnz(A::Diagonal) - return (axes(A, 1), axes(A, 2), A.diag) +function _coord(A::UniformBlockDiagonal) + dat = A.data + r, c, k = size(dat) + blk = repeat(r .* (0:(k - 1)); inner=r * c) + return Table(; + i=Int32.(repeat(1:r; outer=c * k) .+ blk), + j=Int32.(repeat(1:c; inner=r, outer=k) .+ blk), + v=vec(dat), + ) end -_findnz(x::AbstractMatrix) = findnz(x) - -function SparseArrays.findnz(A::UniformBlockDiagonal) - dat = A.data +function _coord(A::LowerTriangular{T,UniformBlockDiagonal{T}}) where {T} + dat = A.data.data r, c, k = size(dat) blk = repeat(r .* (0:(k - 1)); inner=r * c) - return ( - repeat(1:r; outer=c * k) .+ blk, - repeat(1:c; inner=r, outer=k) .+ blk, - vec(dat), + return Table(; + i=Int32.(repeat(1:r; outer=c * k) .+ blk), + j=Int32.(repeat(1:c; inner=r, outer=k) .+ blk), + v=vec(dat), ) end -SparseArrays.findnz(A::BlockedSparse) = findnz(A.cscmat) +function _coord(A::SparseMatrixCSC{T,Int32}) where {T} + rv = rowvals(A) + cv = similar(rv) + for j in axes(A, 2), k in nzrange(A, j) + cv[k] = j + end + return Table(; i=rv, j=cv, v=nonzeros(A)) +end + +_coord(A::BlockedSparse) = _coord(A.cscmat) + +function _coord(A::Matrix) + m, n = size(A) + return Table(; + i=Int32.(repeat(axes(A, 1); outer=n)), + j=Int32.(repeat(axes(A, 2); inner=m)), + v=vec(A), + ) +end function sparsemat( typ::Symbol, m::LinearMixedModel{T}; fname::Symbol=first(fnames(m)), full::Bool=false @@ -1236,15 +1319,13 @@ function sparsemat( end blks = sblk:(length(reterms) + full) rowoffset, coloffset = 0, 0 - I = Int32[] - J = Int32[] - V = T[] + val = (i=Int32[], j=Int32[], v=T[]) for i in blks, j in first(blks):i Lblk = bmat[block(i, j)] - cblk = _findnz(Lblk) - append!(I, first(cblk) .+ Int32(rowoffset)) - append!(J, cblk[2] .+ Int32(coloffset)) - append!(V, last(cblk)) + cblk = _coord(Lblk) + append!(val.i, cblk.i .+ Int32(rowoffset)) + append!(val.j, cblk.j .+ Int32(coloffset)) + append!(val.v, cblk.v) if i == j coloffset = 0 rowoffset += size(Lblk, 1) @@ -1252,7 +1333,30 @@ function sparsemat( coloffset += size(Lblk, 2) end end - return tril!(sparse(I, J, V)) + return tril!(sparse(val.i, val.j, val.v)) +end + +function _triinds(n::Integer, uplo::Symbol) + inds = sizehint!(@NamedTuple{i::Int32, j::Int32}[], kp1choose2(n)) + ul = uplo == :U + for j in Int32.(1:n) + for i in Int32.(ul ? (1:j) : (j:n)) + push!(inds, (; i, j)) + end + end + return Table(inds) +end + +function _coord(A::LowerTriangular{T,Matrix{T}}) where {T} + tbl = _triinds(size(A, 2), :L) + v = [A[i, j] for (i, j) in tbl] + return Table(tbl; v) +end + +function _coord(A::TriangularRFP{T}) where {T} + tbl = _triinds(size(A, 2), Symbol(A.uplo)) + v = [A[i, j] for (i, j) in tbl] + return Table(tbl; v) end """ @@ -1389,14 +1493,16 @@ Update the blocked lower Cholesky factor, `m.L`, from `m.A` and `m.reterms` (use This is the crucial step in evaluating the objective, given a new parameter value. """ function updateL!(m::LinearMixedModel{T}) where {T} - A, L, reterms = m.A, m.L, m.reterms + (; A, L, reterms) = m k = length(reterms) copyto!(last(m.L), last(m.A)) # ensure the fixed-effects:response block is copied - for j in eachindex(reterms) # pre- and post-multiply by Λ, add I to diagonal + 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) - for i in (j + 1):(k + 1) # postmultiply column by Λ + 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) end @@ -1406,17 +1512,17 @@ function updateL!(m::LinearMixedModel{T}) where {T} 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) - rankUpdate!(Hermitian(Ljj, :L), L[block(j, jj)], -one(T), one(T)) + rankUpdate!(LjjH, L[block(j, jj)], -one(T), one(T)) end - cholUnblocked!(Ljj, Val{:L}) - LjjT = isa(Ljj, Diagonal) ? Ljj : LowerTriangular(Ljj) + cholUnblocked!(LjjH) for i in (j + 1):(k + 1) Lij = L[block(i, j)] for jj in 1:(j - 1) mul!(Lij, L[block(i, jj)], L[block(j, jj)]', -one(T), one(T)) end - rdiv!(Lij, LjjT') + rdiv!(Lij, Ljj') end end return m diff --git a/src/predict.jl b/src/predict.jl index 41a1b7e2d..1f3d906c7 100644 --- a/src/predict.jl +++ b/src/predict.jl @@ -177,11 +177,11 @@ function _predict(m::MixedModel{T}, newdata, β; new_re_levels) where {T} throw(ArgumentError("New level encountered in $grp")) end end + end - # we don't have to worry about the BLUP ordering within a given - # grouping variable because we are in the :error branch - blups = blupsold - elseif new_re_levels == :population + if new_re_levels in (:error, :population) + # the level *order* of the fitted model's reterms need not match the + # order in the new design (e.g. after sortlevels), so align by label blups = [Matrix{T}(undef, size(t.z, 1), nlevs(t)) for t in newre] for (idx, (B, newlvls, oldidx)) in enumerate(zip(blups, newlevels, oldlevelidx)) @@ -189,6 +189,7 @@ function _predict(m::MixedModel{T}, newdata, β; new_re_levels) where {T} oldloc = get(oldidx, ll, nothing) if oldloc === nothing # setting a BLUP to zero gives you the population value + # (unreachable for :error, which has already validated levels) B[:, lidx] .= zero(T) else B[:, lidx] .= @view blupsold[idx][:, oldloc] diff --git a/src/randomeffectsterm.jl b/src/randomeffectsterm.jl index 6cea36035..fe3cfb92c 100644 --- a/src/randomeffectsterm.jl +++ b/src/randomeffectsterm.jl @@ -1,3 +1,55 @@ +# Interface notes for defining new `AbstractReTerm` / `AbstractReMat` subtypes +# =========================================================================== +# +# There is no formally documented minimal interface for these abstract types: +# only `_sortlevels!(::AbstractReMat)` (a genuine no-op fallback in remat.jl) +# and `StatsAPI.coefnames(::AbstractReMat)` (assumes a `.cnames` field) are +# written against the abstract types. Everything else in the fit is dispatched +# on the concrete `RandomEffectsTerm` / `ReMat`, so a genuinely new subtype +# (not just a `ReMat` wrapper that forwards) must supply the whole surface +# below. The realistic path is usually to subtype and delegate to a held +# `ReMat`, overriding only the handful of methods being changed; if the goal is +# different block storage or a different λ structure, a new block type plus its +# `rmulΛ!`/`lmulΛ!`/`rankUpdate!`/`copyscaleinflate!` methods may suffice +# without a new `AbstractReMat` at all. +# +# `AbstractReTerm` (formula/term side; feeds `modelcols` -> `AbstractReMat`): +# - StatsModels.apply_schema(t, ::MultiSchema{FullRank}, ::Type{<:MixedModel}) +# - StatsModels.modelcols(t, d::NamedTuple) -> returns the AbstractReMat +# - StatsModels.termvars(t), StatsModels.terms(t) +# - StatsModels.is_matrix_term(::Type{T}) = false +# - is_randomeffectsterm(t) (defaults to true on AbstractReTerm) +# - Base.show +# +# `AbstractReMat` (matrix side, used throughout the fit): +# AbstractArray basics: Base.size, Base.getindex, SparseArrays.sparse, +# LinearAlgebra.Matrix +# Assembly (createAL cross-products), returning block-storage types that +# themselves support the updateL! kernels: +# - Base.:(*)(::Adjoint{<:AbstractReMat}, ::AbstractReMat) # Zᵢ'Zⱼ +# - Base.:(*)(::Adjoint{<:FeMat}, ::AbstractReMat) # X'Z +# updateL! hot loop: +# - copyscaleinflate!(LdH, A_jj, cj) # Λ'AΛ + I on the diagonal block +# - rmulΛ!(dest, cj) / lmulΛ!(cj', dest) # per block-storage type +# - block types must also support rankUpdate!, cholUnblocked!, mul!, rdiv! +# θ / parameters: getθ, getθ!, setθ!, nθ, lowerbd, vsize (== S), nranef +# Grouping metadata: fname, DataAPI.levels, DataAPI.refarray, DataAPI.refpool, +# DataAPI.refvalue, nlevs, StatsModels.isnested +# Weights: reweight! +# Copying: Base.copy, LinearAlgebra.copy_oftype +# Post-fit (ranef/simulate/VarCorr/PCA/condVar/coeftable): unscaledre!, +# rowlengths, corrmat, σvals, σs, σρs, PCA, indmat, zerocorr!, +# LinearAlgebra.cond, coefnames +# Optional (has a working fallback): _sortlevels! -- paired with the term- +# side _syncgrouping (see below); implement both or neither. +# +# `_sortlevels!` (AbstractReMat) and `_syncgrouping` (AbstractTerm) are a pair, +# both with no-op fallbacks. `_sortlevels!` reorders a ReMat's levels and +# rebuilds its `trm` with reordered contrasts; `_syncgrouping` substitutes the +# rebuilt `trm`s back into the stored formula so `modelcols` (hence predict / +# simulate on newdata) stays consistent. A new subtype with a non-trivial +# `_sortlevels!` that rebuilds its term must also add a `_syncgrouping` method, +# or the stored formula and the fitted ReMat diverge. abstract type AbstractReTerm <: AbstractTerm end struct RandomEffectsTerm <: AbstractReTerm @@ -210,3 +262,28 @@ StatsModels.modelcols(t::ZeroCorr, d::NamedTuple) = zerocorr!(modelcols(t.term, function Base.getproperty(x::ZeroCorr, s::Symbol) return s == :term ? getfield(x, s) : getproperty(x.term, s) end + +""" + _syncgrouping(form::FormulaTerm, reterms) + +Replace the grouping terms in `form` with the corresponding `trm`s of `reterms`. + +After [`_sortlevels!`](@ref) the `ReMat`s hold rebuilt `CategoricalTerm`s whose +contrasts reflect the new level order; substituting them into the stored +formula keeps `modelcols` on that formula consistent with the fitted model. +""" +function _syncgrouping(form::FormulaTerm, reterms) + newtrms = Dict( + rt.trm.sym => rt.trm for + rt in reterms if rt isa ReMat && rt.trm isa CategoricalTerm + ) + isempty(newtrms) && return form + return FormulaTerm(form.lhs, _syncgrouping.(form.rhs, Ref(newtrms))) +end +_syncgrouping(t::AbstractTerm, newtrms::Dict) = t +function _syncgrouping(t::RandomEffectsTerm, newtrms::Dict) + rhs = t.rhs + rhs isa CategoricalTerm || return t + return RandomEffectsTerm(t.lhs, get(newtrms, rhs.sym, rhs)) +end +_syncgrouping(t::ZeroCorr, newtrms::Dict) = ZeroCorr(_syncgrouping(t.term, newtrms)) diff --git a/src/remat.jl b/src/remat.jl index 17379d927..0c567e23a 100644 --- a/src/remat.jl +++ b/src/remat.jl @@ -104,6 +104,58 @@ function adjA(refs::AbstractVector, z::AbstractMatrix) return sparse(II, J, vec(z)) end +""" + _sortlevels!(A::AbstractReMat) + +Reorder the levels of the grouping factor by descending number of occurrences. + +For grouping factors whose diagonal block in `L` incurs fill-in, placing the +most frequently occurring levels first concentrates the sparse rank-update of +that block in a compact region of storage, which improves memory locality. +For `TriangularRFP` blocks it additionally keeps most updates in the +trapezoidal part of the packed storage, avoiding strided access to the +transposed triangular part. + +The `ContrastsMatrix` of the grouping factor's term (i.e. `A.trm`) is rebuilt +so that its level order (and hence its `invindex`) stays in sync with +the reordered levels. Interaction groupings store their level order only in +the `ReMat` itself, so there is nothing further to synchronize. + +!!! note "Only fully implemented for concrete `ReMat`" + There is a no-op `_sortlevels!(::AbstractReMat)` method that returns the original matrix, + so that everything will "just" work when a new subtype of `AbstractReMat` is added. + If a given concrete subtype will benefit from this, then you should implement an appropriate method. +""" +_sortlevels!(A::AbstractReMat) = A + +function _sortlevels!(A::ReMat{T,S}) where {T,S} + counts = zeros(Int, length(A.levels)) + for r in A.refs + counts[r] += 1 + end + issorted(counts; rev=true) && return A + perm = sortperm(counts; rev=true) + invp = invperm(perm) + map!(r -> invp[r], A.refs, A.refs) + A.levels .= @view A.levels[perm] + # permute the rows of adjA in place to match the permuted labels: + # each column holds the S rows of a single level, so the relabeled rowvals + # remain sorted within each column + rv = rowvals(A.adjA) + if isone(S) + map!(r -> invp[r], rv, rv) + else + map!(i -> (invp[fld1(i, S)] - 1) * S + mod1(i, S), rv, rv) + end + trm = A.trm + if trm isa CategoricalTerm + A.trm = CategoricalTerm( + trm.sym, StatsModels.ContrastsMatrix(trm.contrasts.contrasts, A.levels) + ) + end + return A +end + Base.size(A::ReMat) = (length(A.refs), length(A.scratch)) SparseArrays.sparse(A::ReMat) = adjoint(A.adjA) @@ -566,8 +618,12 @@ Overwrite L with `Λ'AΛ + I` """ function copyscaleinflate! end -function copyscaleinflate!(Ljj::Diagonal{T}, Ajj::Diagonal{T}, Λj::ReMat{T,1}) where {T} - Ldiag, Adiag = Ljj.diag, Ajj.diag +function copyscaleinflate!( + Ljj::Hermitian{T,Diagonal{T,Vector{T}}}, + Ajj::Diagonal{T,Vector{T}}, + Λj::ReMat{T,1}, +) where {T} + Ldiag, Adiag = Ljj.data.diag, Ajj.diag lambsq = abs2(only(Λj.λ.data)) @inbounds for i in eachindex(Ldiag, Adiag) Ldiag[i] = muladd(lambsq, Adiag[i], one(T)) @@ -575,23 +631,39 @@ function copyscaleinflate!(Ljj::Diagonal{T}, Ajj::Diagonal{T}, Λj::ReMat{T,1}) return Ljj end -function copyscaleinflate!(Ljj::Matrix{T}, Ajj::Diagonal{T}, Λj::ReMat{T,1}) where {T} - fill!(Ljj, zero(T)) +function copyscaleinflate!( + Ljj::Hermitian{T,Matrix{T}}, + Ajj::Diagonal{T}, + Λj::ReMat{T,1}, +) where {T} + Ld = Ljj.data + fill!(Ld, zero(T)) lambsq = abs2(only(Λj.λ.data)) @inbounds for (i, a) in enumerate(Ajj.diag) - Ljj[i, i] = muladd(lambsq, a, one(T)) + Ld[i, i] = muladd(lambsq, a, one(T)) + end + return Ljj +end + +function copyscaleinflate!(Ljj::HermitianRFP{T}, Ajj::Diagonal{T}, Λj::ReMat{T,1}) where {T} + fill!(Ljj.data, zero(T)) + lambsq = abs2(only(Λj.λ.data)) + LjjT = TriangularRFP(Ljj.data, Ljj.transr, Ljj.uplo) + @inbounds for (i, a) in enumerate(Ajj.diag) + LjjT[i, i] = muladd(lambsq, a, one(T)) end return Ljj end function copyscaleinflate!( - Ljj::UniformBlockDiagonal{T}, + Ljj::Hermitian{T,UniformBlockDiagonal{T}}, Ajj::UniformBlockDiagonal{T}, Λj::ReMat{T,S}, ) where {T,S} + Ljj.uplo == 'L' || throw(ArgumentError("Ljj.uplo should be 'L'")) λ = Λj.λ dind = diagind(S, S) - Ldat = copyto!(Ljj.data, Ajj.data) + Ldat = copyto!(Ljj.data.data, Ajj.data) for k in axes(Ldat, 3) f = view(Ldat, :, :, k) lmul!(λ', rmul!(f, λ)) @@ -603,24 +675,53 @@ function copyscaleinflate!( end function copyscaleinflate!( - Ljj::Matrix{T}, + Ljj::HermitianRFP{T}, + Ajj::UniformBlockDiagonal{T}, + Λj::ReMat{T,S}, +) where {T,S} # S is an integer - the size of the diagonal blocks + fill!(Ljj.data, zero(T)) # zero the off-block entries left over from a previous factorization + LjjT = TriangularRFP(Ljj.data, Ljj.transr, Ljj.uplo) + q, r = divrem(size(Ljj, 1), S) + iszero(r) || throw(DimensionMismatch("size(Ljj, 1) is not a multiple of S")) + λ = Λj.λ + tmp = Array{T}(undef, S, S) + offset = 0 + @inbounds for k in 1:q + copyto!(tmp, Ajj.data[:, :, k]) + lmul!(adjoint(λ), rmul!(tmp, λ)) + for j in 1:S + tmp[j, j] += one(T) + end + for j in 1:S + for i in j:S + LjjT[offset + i, offset + j] = tmp[i, j] + end + end + offset += S + end + return Ljj +end + +function copyscaleinflate!( + Ljj::Hermitian{T,Matrix{T}}, Ajj::UniformBlockDiagonal{T}, Λj::ReMat{T,S}, ) where {T,S} - copyto!(Ljj, Ajj) - n = LinearAlgebra.checksquare(Ljj) + LjjM = Ljj.data + copyto!(LjjM, Ajj) + n = LinearAlgebra.checksquare(LjjM) q, r = divrem(n, S) iszero(r) || throw(DimensionMismatch("size(Ljj, 1) is not a multiple of S")) λ = Λj.λ offset = 0 @inbounds for _ in 1:q inds = (offset + 1):(offset + S) - tmp = view(Ljj, inds, inds) + tmp = view(LjjM, inds, inds) lmul!(adjoint(λ), rmul!(tmp, λ)) offset += S end - for k in diagind(Ljj) - Ljj[k] += one(T) + for k in diagind(LjjM) + LjjM[k] += one(T) end return Ljj end diff --git a/test/RFP.jl b/test/RFP.jl new file mode 100644 index 000000000..bcc56123f --- /dev/null +++ b/test/RFP.jl @@ -0,0 +1,19 @@ +using LinearAlgebra +using MixedModels +using Test +using SparseArrays +using RectangularFullPacked + +# tests of the rankUpdate! method when `typeof(C)` is `HermitianRFP` +@testset "rankUpHermitianRFP" begin + A9 = float(sprand(Bool, 9, 12, 0.3)) + T = TriangularRFP(collect(A9 * A9'), :L) + C9 = HermitianRFP(T.data, T.transr, T.uplo) + rankUpdate!(C9, A9, -1.0, 1.0) + @test all(iszero, C9.data) + A8 = float(sprand(Bool, 8, 12, 0.3)) + T = TriangularRFP(collect(A8 * A8'), :L) + C8 = HermitianRFP(T.data, T.transr, T.uplo) + rankUpdate!(C8, A8, -1.0, 1.0) + @test all(iszero, C8.data) +end diff --git a/test/UniformBlockDiagonal.jl b/test/UniformBlockDiagonal.jl index 525b00c3a..7fd34181b 100644 --- a/test/UniformBlockDiagonal.jl +++ b/test/UniformBlockDiagonal.jl @@ -43,11 +43,11 @@ const LMM = LinearMixedModel end @testset "copyscaleinflate" begin - MixedModels.copyscaleinflate!(Lblk, ex22, vf1) + MixedModels.copyscaleinflate!(Hermitian(Lblk, :L), ex22, vf1) @test view(Lblk.data, :, :, 1) == [2.0 3.0; 2.0 5.0] setθ!(vf1, [1.0, 1.0, 1.0]) Λ = vf1.λ - MixedModels.copyscaleinflate!(Lblk, ex22, vf1) + MixedModels.copyscaleinflate!(Hermitian(Lblk, :L), ex22, vf1) target = Λ'view(ex22.data, :, :, 1) * Λ + I @test view(Lblk.data, :, :, 1) == target end @@ -64,9 +64,8 @@ const LMM = LinearMixedModel A11 = vf1'vf1 L11 = MixedModels.cholUnblocked!( MixedModels.copyscaleinflate!( - UniformBlockDiagonal(fill(0.0, size(A11.data))), A11, vf1 + Hermitian(UniformBlockDiagonal(fill(0.0, size(A11.data))), :L), A11, vf1 ), - Val{:L}, ) L21 = vf2'vf1 @test isa(L21, BlockedSparse) @@ -81,7 +80,7 @@ const LMM = LinearMixedModel # @test_broken L21.colblocks[1] == rdiv!(L21cb1, adjoint(LowerTriangular(L11.facevec[1]))) A22 = vf2'vf2 L22 = MixedModels.copyscaleinflate!( - UniformBlockDiagonal(fill(0.0, size(A22.data))), A22, vf2 + Hermitian(UniformBlockDiagonal(fill(0.0, size(A22.data))), :L), A22, vf2 ) end end diff --git a/test/bootstrap.jl b/test/bootstrap.jl index 4efec3e10..908f0ea3c 100644 --- a/test/bootstrap.jl +++ b/test/bootstrap.jl @@ -62,7 +62,7 @@ end progress=false, ) # fails in pirls! with fast=false gm4sim = refit!(simulate!(StableRNG(42), deepcopy(gm4)); progress=false) - @test isapprox(gm4.β, gm4sim.β; atol=norm(stderror(gm4))) + @test isapprox(gm4.β, gm4sim.β; atol=2 * norm(stderror(gm4))) # is the simulation within a 95%-ish confidence region? end @testset "Binomial" begin diff --git a/test/linalg.jl b/test/linalg.jl index 5d72296db..664c293b3 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -5,7 +5,7 @@ using Random using SparseArrays using Test -using MixedModels: rankUpdate! +using MixedModels: rankUpdate!, BlockedSparse, HermitianRFP, TriangularRFP @testset "mul!" begin for (m, p, n, q, k) in ( @@ -74,6 +74,39 @@ end rankUpdate!(Symmetric(zeros(100, 100), :U), sparse(transpose(L21)), 1.0, 1.0) end +@testset "rankUpdate! HermitianRFP" begin + rng = MersenneTwister(1234) + n = 6 + + # a symmetric, lower-stored base matrix and the packed HermitianRFP holding it + S = let M = randn(rng, n, n) + M * M' + n * I + end + rfp(S) = Hermitian(TriangularRFP(Matrix(LowerTriangular(S)), :L), :L) + + # dense and (Int32) sparse updates plus a genuine BlockedSparse block + Adense = randn(rng, n, 4) + Asparse = convert(SparseMatrixCSC{Float64,Int32}, sparse(sprand(rng, n, 20, 0.3))) + # crossed factors ⇒ an off-diagonal L block stored as BlockedSparse + d3 = MixedModels.dataset(:d3) + m = fit!( + LinearMixedModel(@formula(y ~ 1 + u + (1 + u | g) + (1 + u | h) + (1 + u | i)), d3); + progress=false, + ) + Ablk = m.L[findfirst(a -> a isa BlockedSparse, m.L)] + @test Ablk isa BlockedSparse + Sblk = zeros(size(Ablk, 1), size(Ablk, 1)) + I + + for (A, base) in ((Adense, S), (Asparse, S), (Ablk, Sblk)) + for (α, β) in ((1.0, 1.0), (-1.0, 1.0), (0.5, 2.0)) + ref = rankUpdate!(Hermitian(Matrix(base), :L), A, α, β) + got = rankUpdate!(rfp(base), A, α, β) + @test got isa HermitianRFP + @test LowerTriangular(Matrix(got)) ≈ LowerTriangular(Matrix(ref)) + end + end +end + #= I don't see this testset as meaningful b/c diagonal A does not occur after amalgamation of ReMat's for the same grouping factor - D.B. @testset "rankupdate!" begin @test ones(2, 2) == rankUpdate!(Hermitian(zeros(2, 2)), ones(2)) diff --git a/test/pls.jl b/test/pls.jl index f751eed0d..be0340ca8 100644 --- a/test/pls.jl +++ b/test/pls.jl @@ -283,7 +283,10 @@ end spL = sparseL(fm1) @test size(spL) == (4114, 4114) - @test 733090 < nnz(spL) < 733100 + # nnz(spL) counts the exact nonzeros remaining in the (densely stored) filled + # blocks, so it depends on the level ordering: with the default + # sortlevels=true the frequency-descending order incurs more fill + @test 744320 < nnz(spL) < 744340 spA = Symmetric(sparseA(fm1; full=true), :L) @test size(spA) == (4117, 4117) @@ -331,19 +334,72 @@ end @test size(fm2) == (73421, 28, 4100, 2) end +@testset "sortlevels" begin + insteval = MixedModels.dataset(:insteval) + form = @formula(y ~ 1 + service + (1 | s) + (1 | d)) + function levelcounts(rt) + c = zeros(Int, MixedModels.nlevs(rt)) + for r in rt.refs + c[r] += 1 + end + return c + end + m0 = LinearMixedModel(form, insteval; sortlevels=false) + m1 = LinearMixedModel(form, insteval) # sortlevels=true is the default + + @test !issorted(levelcounts(m0.reterms[2]); rev=true) + @test issorted(levelcounts(m1.reterms[2]); rev=true) + # the leading term keeps the data's level order + @test m0.reterms[1].levels == m1.reterms[1].levels + + # the objective is invariant under the level permutation + @test objective(updateL!(setθ!(m0, m0.optsum.initial))) ≈ + objective(updateL!(setθ!(m1, m1.optsum.initial))) + + # the term's contrasts stay in sync with the reordered levels + rt = m1.reterms[2] + @test rt.trm.contrasts.levels == rt.levels + @test all(rt.trm.contrasts.invindex[l] == i for (i, l) in enumerate(rt.levels)) + + # ... as does the stored formula, so that modelcols on it reproduces + # the fitted model's refs + retrm = only( + t for t in m1.formula.rhs if + t isa MixedModels.RandomEffectsTerm && MixedModels.fname(t.rhs) == :d + ) + @test retrm.rhs === rt.trm + remat = last(last(modelcols(m1.formula, columntable(insteval)))) + @test remat.refs == rt.refs + + # BLUPs align by level label and predict is positionally consistent + fit!(m1; progress=false) + @test predict(m1, insteval; new_re_levels=:error) ≈ fitted(m1) + + # with amalgamate=false, all terms for a factor share one level order, + # including a term in the leading (unsorted) position + mam = LinearMixedModel( + @formula(y ~ 1 + (1 | d) + (0 + service | d) + (1 | s)), insteval; + amalgamate=false, + ) + dre = filter(rt -> MixedModels.fname(rt) == :d, mam.reterms) + @test length(dre) == 2 + @test first(dre).levels == last(dre).levels + @test first(dre).refs == last(dre).refs +end + @testset "sleep" begin fm = last(models(:sleepstudy)) A11 = first(fm.A) @test isa(A11, UniformBlockDiagonal{Float64}) - @test isa(first(fm.L), UniformBlockDiagonal{Float64}) + @test isa(first(fm.L), LowerTriangular{Float64, UniformBlockDiagonal{Float64}}) @test size(A11) == (36, 36) a11 = view(A11.data, :, :, 1) @test a11 == [10.0 45.0; 45.0 285.0] @test size(A11.data, 3) == 18 λ = only(fm.λ) - b11 = LowerTriangular(view(first(fm.L).data, :, :, 1)) + b11 = LowerTriangular(view(first(fm.L).data.data, :, :, 1)) @test b11 * b11' ≈ λ'a11 * λ + I rtol = 1e-5 - @test count(!iszero, Matrix(first(fm.L))) == 18 * 4 + @test count(!iszero, Matrix(first(fm.L).data)) == 18 * 4 @test rank(fm) == 2 @test objective(fm) ≈ 1751.9393444636682 diff --git a/test/runtests.jl b/test/runtests.jl index 0631ca3a0..6c7a32ca0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -47,6 +47,7 @@ include("mime.jl") include("optsummary.jl") include("predict.jl") include("sigma.jl") +include("RFP.jl") @testset "PRIMA" include("prima.jl") @testset "ForwardDiff" include("forwarddiff.jl")