Generalize sparseL as sparsemat (#880)

* Generalize sparseL as sparsemat

* Update NEWS.md

* Run src files through format

* Add sparseA extractor

* Apply dropzeros! to sparseL but not sparseA

* minor version bump

Author: dmbates

NEWS.md

@@ -1,3 +1,7 @@
+MixedModels v5.3.0 Release Notes
+==============================
+- Implement `sparseL` as a specialization of `sparsemat`. Replace `_coord` utility with `_findnz` which, in most cases, falls through to `SparseArrays.findnz`. [#880]
+
 MixedModels v5.2.2 Release Notes
 ==============================
 - Small update to `show` methods to accommodate coming changes in Julia Markdown stdlib. [#876]
@@ -725,3 +729,4 @@ Package dependencies
 [#873]: https://github.com/JuliaStats/MixedModels.jl/issues/873
 [#875]: https://github.com/JuliaStats/MixedModels.jl/issues/875
 [#876]: https://github.com/JuliaStats/MixedModels.jl/issues/876
+[#880]: https://github.com/JuliaStats/MixedModels.jl/issues/880

Project.toml

@@ -1,7 +1,7 @@
 name = "MixedModels"
 uuid = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
 author = ["Phillip Alday <me@phillipalday.com>", "Douglas Bates <dmbates@gmail.com>"]
-version = "5.2.2"
+version = "5.3.0"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"

src/MixedModels.jl

@@ -27,8 +27,8 @@ using Printf: @sprintf
 using ProgressMeter: ProgressMeter, Progress, finish!, next!
 using Random: Random, AbstractRNG, randn!
 using RegressionFormulae: fulldummy
-using SparseArrays: SparseArrays, SparseMatrixCSC, SparseVector, dropzeros!, nnz
-using SparseArrays: nonzeros, nzrange, rowvals, sparse
+using SparseArrays: SparseArrays, SparseMatrixCSC, SparseVector, dropzeros!, findnz
+using SparseArrays: nnz, nonzeros, nzrange, rowvals, sparse
 using StaticArrays: StaticArrays, SVector
 using Statistics: Statistics, mean, quantile, std
 using StatsAPI: StatsAPI, aic, aicc, bic, coef, coefnames, coeftable, confint
@@ -156,6 +156,7 @@ export @formula,
     simulate,
     simulate!,
     sparse,
+    sparseA,
     sparseL,
     std,
     stderror,

src/linearmixedmodel.jl

@@ -800,8 +800,9 @@ Return the vector of _canonical_ lower bounds on the parameters, `θ`.
 Note that this method does not distinguish between constrained optimization and
 unconstrained optimization with post-fit canonicalization.
 """
-lowerbd(m::LinearMixedModel{T}) where {T} =
-    [(pm[2] == pm[3]) ? zero(T) : T(-Inf) for pm in m.parmap]
+function lowerbd(m::LinearMixedModel{T}) where {T}
+    return [(pm[2] == pm[3]) ? zero(T) : T(-Inf) for pm in m.parmap]
+end
 
 function mkparmap(reterms::Vector{<:AbstractReMat{T}}) where {T}
     parmap = NTuple{3,Int}[]
@@ -1133,82 +1134,101 @@ end
 Base.show(io::IO, m::LinearMixedModel) = Base.show(io, MIME("text/plain"), m)
 
 """
-    _coord(A::AbstractMatrix)
+    _findnz(A::AbstractMatrix)
+
+Return the positions and values of the nonzeros in `A` as an (I, J, V) tuple
 
-Return the positions and values of the nonzeros in `A` as a
-`NamedTuple{(:i, :j, :v), Tuple{Vector{Int32}, Vector{Int32}, Vector{Float64}}}`
+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.
 """
-function _coord(A::Diagonal)
-    return (i=Int32.(axes(A, 1)), j=Int32.(axes(A, 2)), v=A.diag)
-end
 
-function _coord(A::UniformBlockDiagonal)
-    dat = A.data
-    r, c, k = size(dat)
-    blk = repeat(r .* (0:(k - 1)); inner=r * c)
+function _findnz(A::Matrix)
+    m, n = size(A)
     return (
-        i=Int32.(repeat(1:r; outer=c * k) .+ blk),
-        j=Int32.(repeat(1:c; inner=r, outer=k) .+ blk),
-        v=vec(dat),
+        repeat(axes(A, 1); outer=n),
+        repeat(axes(A, 2); inner=m),
+        vec(A),
     )
 end
 
-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 (i=rv, j=cv, v=nonzeros(A))
+function _findnz(A::Diagonal)
+    return (axes(A, 1), axes(A, 2), A.diag)
 end
 
-_coord(A::BlockedSparse) = _coord(A.cscmat)
+_findnz(x::AbstractMatrix) = findnz(x)
 
-function _coord(A::Matrix)
-    m, n = size(A)
+function SparseArrays.findnz(A::UniformBlockDiagonal)
+    dat = A.data
+    r, c, k = size(dat)
+    blk = repeat(r .* (0:(k - 1)); inner=r * c)
     return (
-        i=Int32.(repeat(axes(A, 1); outer=n)),
-        j=Int32.(repeat(axes(A, 2); inner=m)),
-        v=vec(A),
+        repeat(1:r; outer=c * k) .+ blk,
+        repeat(1:c; inner=r, outer=k) .+ blk,
+        vec(dat),
     )
 end
 
-"""
-    sparseL(m::LinearMixedModel; fname::Symbol=first(fnames(m)), full::Bool=false)
-
-Return the lower Cholesky factor `L` as a `SparseMatrix{T,Int32}`.
+SparseArrays.findnz(A::BlockedSparse) = findnz(A.cscmat)
 
-`full` indicates whether the parts of `L` associated with the fixed-effects and response
-are to be included.
-
-`fname` specifies the first grouping factor to include. Blocks to the left of the block corresponding
- to `fname` are dropped. The default is the first, i.e., leftmost block and hence all blocks.
-"""
-function sparseL(
-    m::LinearMixedModel{T}; fname::Symbol=first(fnames(m)), full::Bool=false
+function sparsemat(
+    typ::Symbol, m::LinearMixedModel{T}; fname::Symbol=first(fnames(m)), full::Bool=false
 ) where {T}
-    L, reterms = m.L, m.reterms
+    reterms = m.reterms
+    if typ ∉ (:A, :L)
+        throw(ArgumentError("typ passed as $typ should be :A or :L"))
+    end
+    bmat = getproperty(m, typ)
     sblk = findfirst(isequal(fname), fnames(m))
     if isnothing(sblk)
         throw(ArgumentError("fname = $fname is not the name of a grouping factor"))
     end
     blks = sblk:(length(reterms) + full)
     rowoffset, coloffset = 0, 0
-    val = (i=Int32[], j=Int32[], v=T[])
+    I = Int32[]
+    J = Int32[]
+    V = T[]
     for i in blks, j in first(blks):i
-        Lblk = L[block(i, j)]
-        cblk = _coord(Lblk)
-        append!(val.i, cblk.i .+ Int32(rowoffset))
-        append!(val.j, cblk.j .+ Int32(coloffset))
-        append!(val.v, cblk.v)
+        Lblk = bmat[block(i, j)]
+        cblk = _findnz(Lblk)
+        append!(I, first(cblk) .+ Int32(rowoffset))
+        append!(J, cblk[2] .+ Int32(coloffset))
+        append!(V, last(cblk))
         if i == j
             coloffset = 0
             rowoffset += size(Lblk, 1)
         else
             coloffset += size(Lblk, 2)
         end
     end
-    return dropzeros!(tril!(sparse(val...)))
+    return tril!(sparse(I, J, V))
+end
+
+"""
+    sparseL(m::LinearMixedModel; fname::Symbol=first(fnames(m)), full::Bool=false)
+
+Return the lower Cholesky factor `L` as a `SparseMatrix{T,Int32}`.
+
+`full` indicates whether the parts of `L` associated with the fixed-effects and response
+are to be included.
+
+`fname` specifies the first grouping factor to include. Blocks to the left of the block corresponding
+ to `fname` are dropped. The default is the first, i.e., leftmost block and hence all blocks.
+
+ The matrix that is returned is lower-triangular but has not been wrapped in `LowerTriangular`.
+"""
+function sparseL(m::LinearMixedModel; fname::Symbol=first(fnames(m)), full::Bool=false)
+    return dropzeros!(sparsemat(:L, m; fname, full))
+end
+
+"""
+    sparseA(m::LinearMixedModel; fname::Symbol=first(fnames(m)), full::Bool=false)
+
+Same as [`sparseL`](@ref) but returning the sparse lower triangle of the symmetric `A` matrix.
+Most of the time the result is wrapped as, e.g. `Symmetric(sparseM(m; full=true), :L)`
+"""
+function sparseA(m::LinearMixedModel; fname::Symbol=first(fnames(m)), full::Bool=false)
+    return sparsemat(:A, m; fname, full)
 end
 
 """

test/pls.jl

@@ -277,14 +277,18 @@ end
 end
 
 @testset "InstEval" begin
-    fm1 = models(:insteval)[2]              # at one time this was the fist of the :insteval models
+    fm1 = models(:insteval)[2]              # at one time this was the first of the :insteval models
     @test size(fm1) == (73421, 2, 4114, 3)
     @test fm1.optsum.initial == ones(3)
 
     spL = sparseL(fm1)
     @test size(spL) == (4114, 4114)
     @test 733090 < nnz(spL) < 733100
 
+    spA = Symmetric(sparseA(fm1; full=true), :L)
+    @test size(spA) == (4117, 4117)
+    @test 132600 < nnz(spA.data) < 132650
+
     @test objective(fm1) ≈ 237721.76877450474 atol = 0.001
     ftd1 = fitted(fm1)
     @test size(ftd1) == (73421,)