Unify PDMat and PDSparseMat + move SparseArrays support to an extension

Project.toml

@@ -1,6 +1,6 @@
 name = "PDMats"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.11.29"
+version = "0.12.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -10,10 +10,18 @@ SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 [compat]
 julia = "1"
 
+[extensions]
+PDMatsSparseArraysExt = "SparseArrays"
+
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BandedMatrices", "StaticArrays", "Test"]
+test = ["BandedMatrices", "SparseArrays", "StaticArrays", "Test"]
+
+[weakdeps]
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

README.md

@@ -27,16 +27,16 @@ Elemenent types are in princple all Real types, but in practice this is limited
 * `PDMat`: full covariance matrix, defined as
 
 ```julia
-struct PDMat{T<:Real,S<:AbstractMatrix} <: AbstractPDMat{T}
+struct PDMat{T<:Real,S<:AbstractMatrix{T},F<:Factorization} <: AbstractPDMat{T}
     mat::S                      # input matrix
-    chol::Cholesky{T,S}         # Cholesky factorization of mat
+    fact::F                     # factorization of mat
 end
 
 # Constructors
 
 PDMat(mat, chol)    # with both the input matrix and a Cholesky factorization
 
-PDMat(mat)          # with the input matrix, of type Matrix or Symmetric
+PDMat(mat)          # with the input matrix
                     # Remarks: the Cholesky factorization will be computed
                     # upon construction.
 
@@ -73,29 +73,6 @@ end
 ScalMat(d, v)        # with dimension d and diagonal value v
 ```
 
-
-* `PDSparseMat`: sparse covariance matrix, defined as
-
-```julia
-struct PDSparseMat{T<:Real,S<:AbstractSparseMatrix} <: AbstractPDMat{T}
-    mat::SparseMatrixCSC           # input matrix
-    chol::CholTypeSparse           # Cholesky factorization of mat
-end
-
-# Constructors
-
-PDSparseMat(mat, chol)    # with both the input matrix and a Cholesky factorization
-
-PDSparseMat(mat)          # with the sparse input matrix, of type SparseMatrixCSC
-                          # Remarks: the Cholesky factorization will be computed
-                          # upon construction.
-
-PDSparseMat(chol)         # with the Cholesky factorization
-                          # Remarks: the sparse matrix 'mat' will be computed upon
-                          # construction.
-```
-
-
 ## Common interface
 
 All subtypes of `AbstractPDMat` share the same API, *i.e.* with the same set of methods to operate on their instances. These methods are introduced below, where `a` is an instance of a subtype of `AbstractPDMat` to represent a positive definite matrix, `x` be a column vector or a matrix with `size(x,1) == size(a, 1) == size(a, 2)`, and `c` be a positive real value.

ext/PDMatsSparseArraysExt/PDMatsSparseArraysExt.jl

@@ -0,0 +1,25 @@
+module PDMatsSparseArraysExt
+
+using PDMats
+using SparseArrays
+
+using PDMats.LinearAlgebra
+
+if isdefined(Base, :get_extension)
+    const HAVE_CHOLMOD = isdefined(SparseArrays, :CHOLMOD)
+else
+    import SuiteSparse
+    const HAVE_CHOLMOD = isdefined(SuiteSparse, :CHOLMOD)
+end
+
+# https://github.com/JuliaLang/julia/pull/29749
+if VERSION < v"1.1.0-DEV.792"
+    eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2))
+end
+
+if HAVE_CHOLMOD
+    include("chol.jl")
+    include("pdsparsemat.jl")
+end
+
+end # module

ext/PDMatsSparseArraysExt/chol.jl

@@ -0,0 +1,9 @@
+if isdefined(Base, :get_extension)
+    const CholTypeSparse{T} = SparseArrays.CHOLMOD.Factor{T}
+else
+    const CholTypeSparse{T} = SuiteSparse.CHOLMOD.Factor{T}
+end
+
+# Take into account pivoting!
+PDMats.chol_lower(cf::CholTypeSparse) = cf.PtL
+PDMats.chol_upper(cf::CholTypeSparse) = cf.UP

ext/PDMatsSparseArraysExt/pdsparsemat.jl

@@ -0,0 +1,142 @@
+"""
+Sparse positive definite matrix together with a Cholesky factorization object.
+"""
+const PDSparseMat{T<:Real,S<:AbstractSparseMatrix{T},C<:CholTypeSparse} = PDMat{T,S,C}
+
+function PDMats.PDMat(mat::AbstractSparseMatrix, chol::CholTypeSparse)
+    PDMat{eltype(mat),typeof(mat),typeof(chol)}(mat, chol)
+end
+Base.@deprecate PDMat{T,S}(d::Int, m::AbstractSparseMatrix{T}, c::CholTypeSparse) where {T,S} PDSparseMat{T,S,typeof(c)}(m, c)
+
+PDMats.PDMat(mat::SparseMatrixCSC) = PDMat(mat, cholesky(mat))
+PDMats.PDMat(fac::CholTypeSparse) = PDMat(sparse(fac), fac)
+
+PDMats.AbstractPDMat(A::CholTypeSparse) = PDMat(A)
+
+### Conversion
+function Base.convert(::Type{PDMat{T}}, a::PDSparseMat) where {T<:Real}
+    # CholTypeSparse only supports Float64 and ComplexF64 type parameters!
+    # So there is no point in recomputing `cholesky(mat)` and we just reuse
+    # the existing Cholesky factorization
+    mat = convert(AbstractMatrix{T}, a.mat)
+    return PDMat{T,typeof(mat),typeof(a.fact)}(mat, a.fact)
+end
+
+### Arithmetics
+
+Base.:\(a::PDSparseMat{T}, x::AbstractVecOrMat{T}) where {T<:Real} = convert(Array{T},a.fact \ convert(Array{Float64},x)) #to avoid limitations in sparse factorization library CHOLMOD, see e.g., julia issue #14076
+Base.:/(x::AbstractVecOrMat{T}, a::PDSparseMat{T}) where {T<:Real} = convert(Array{T},convert(Array{Float64},x) / a.fact )
+
+### Algebra
+
+Base.inv(a::PDSparseMat) = PDMat(inv(a.mat))
+LinearAlgebra.cholesky(a::PDSparseMat) = a.fact
+Base.sqrt(A::PDSparseMat) = PDMat(sqrt(Hermitian(Matrix(A))))
+
+### whiten and unwhiten
+
+function PDMats.whiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims axes(r) == axes(x)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    # Can't use `ldiv!` due to missing support in SparseArrays
+    return copyto!(r, PDMats.chol_lower(cholesky(a)) \ x)
+end
+
+function PDMats.unwhiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims axes(r) == axes(x)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    # `*` is not defined for `PtL` factor components,
+    # so we can't use `chol_lower(cholesky(a)) * x`
+    C = cholesky(a)
+    PtL = sparse(C.L)[C.p, :]
+    return copyto!(r, PtL * x)
+end
+
+function PDMats.whiten(a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    return PDMats.chol_lower(cholesky(a)) \ x
+end
+
+function PDMats.unwhiten(a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    # `*` is not defined for `PtL` factor components,
+    # so we can't use `chol_lower(cholesky(a)) * x`
+    C = cholesky(a)
+    PtL = sparse(C.L)[C.p, :]
+    return PtL * x
+end
+
+### quadratic forms
+
+function PDMats.quad(a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    # https://github.com/JuliaLang/julia/commit/2425ae760fb5151c5c7dd0554e87c5fc9e24de73
+    if VERSION < v"1.4.0-DEV.92"
+        z = a.mat * x
+        return x isa AbstractVector ? dot(x, z) : map(dot, eachcol(x), eachcol(z))
+    else
+        return x isa AbstractVector ? dot(x, a.mat, x) : map(Base.Fix1(quad, a), eachcol(x))
+    end
+end
+
+function PDMats.quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
+    PDMats.@check_argdims eachindex(r) == axes(x, 2)
+    @inbounds for i in axes(x, 2)
+        xi = view(x, :, i)
+        # https://github.com/JuliaLang/julia/commit/2425ae760fb5151c5c7dd0554e87c5fc9e24de73
+        if VERSION < v"1.4.0-DEV.92"
+            # Can't use `lmul!` with buffer due to missing support in SparseArrays
+            r[i] = dot(xi, a.mat * xi)
+        else
+            r[i] = dot(xi, a.mat, xi)
+        end
+    end
+    return r
+end
+
+function PDMats.invquad(a::PDSparseMat, x::AbstractVecOrMat)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    z = cholesky(a) \ x
+    return x isa AbstractVector ? dot(x, z) : map(dot, eachcol(x), eachcol(z))
+end
+
+function PDMats.invquad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
+    PDMats.@check_argdims eachindex(r) == axes(x, 2)
+    PDMats.@check_argdims a.dim == size(x, 1)
+    # Can't use `ldiv!` with buffer due to missing support in SparseArrays
+    C = cholesky(a)
+    @inbounds for i in axes(x, 2)
+        xi = view(x, :, i)
+        r[i] = dot(xi, C \ xi)
+    end
+    return r
+end
+
+
+### tri products
+
+function PDMats.X_A_Xt(a::PDSparseMat, x::AbstractMatrix{<:Real})
+    PDMats.@check_argdims a.dim == size(x, 2)
+    z = a.mat * transpose(x)
+    return Symmetric(x * z)
+end
+
+
+function PDMats.Xt_A_X(a::PDSparseMat, x::AbstractMatrix{<:Real})
+    PDMats.@check_argdims a.dim == size(x, 1)
+    z = a.mat * x
+    return Symmetric(transpose(x) * z)
+end
+
+
+function PDMats.X_invA_Xt(a::PDSparseMat, x::AbstractMatrix{<:Real})
+    PDMats.@check_argdims a.dim == size(x, 2)
+    z = cholesky(a) \ collect(transpose(x))
+    return Symmetric(x * z)
+end
+
+function PDMats.Xt_invA_X(a::PDSparseMat, x::AbstractMatrix{<:Real})
+    PDMats.@check_argdims a.dim == size(x, 1)
+    z = cholesky(a) \ x
+    return Symmetric(transpose(x) * z)
+end

src/PDMats.jl

@@ -1,14 +1,13 @@
 module PDMats
 
-    using LinearAlgebra, SparseArrays, SuiteSparse
+    using LinearAlgebra
 
     import Base: +, *, \, /, ==, convert, inv, Matrix, kron
 
     export
         # Types
         AbstractPDMat,
         PDMat,
-        PDSparseMat,
         PDiagMat,
         ScalMat,
 
@@ -35,19 +34,12 @@ module PDMats
     """
     abstract type AbstractPDMat{T<:Real} <: AbstractMatrix{T} end
 
-    const HAVE_CHOLMOD = isdefined(SuiteSparse, :CHOLMOD)
-
     # source files
 
     include("utils.jl")
     include("chol.jl")
 
     include("pdmat.jl")
-
-    if HAVE_CHOLMOD
-        include("pdsparsemat.jl")
-    end
-
     include("pdiagmat.jl")
     include("scalmat.jl")
 
@@ -56,4 +48,8 @@ module PDMats
 
     include("deprecates.jl")
 
+    # Support for sparse arrays
+    if !isdefined(Base, :get_extension)
+        include("../ext/PDMatsSparseArraysExt/PDMatsSparseArraysExt.jl")
+    end
 end # module

src/addition.jl

@@ -4,38 +4,18 @@
 +(a::PDMat, b::AbstractPDMat) = PDMat(a.mat + Matrix(b))
 +(a::PDiagMat, b::AbstractPDMat) = PDMat(_adddiag!(Matrix(b), a.diag, true))
 +(a::ScalMat, b::AbstractPDMat) = PDMat(_adddiag!(Matrix(b), a.value))
-if HAVE_CHOLMOD
-    +(a::PDSparseMat, b::AbstractPDMat) = PDMat(a.mat + Matrix(b))
-end
 
 +(a::PDMat, b::PDMat) = PDMat(a.mat + b.mat)
 +(a::PDMat, b::PDiagMat) = PDMat(_adddiag(a.mat, b.diag))
 +(a::PDMat, b::ScalMat) = PDMat(_adddiag(a.mat, b.value))
-if HAVE_CHOLMOD
-    +(a::PDMat, b::PDSparseMat) = PDMat(a.mat + b.mat)
-end
 
 +(a::PDiagMat, b::PDMat) = PDMat(_adddiag(b.mat, a.diag))
 +(a::PDiagMat, b::PDiagMat) = PDiagMat(a.diag + b.diag)
 +(a::PDiagMat, b::ScalMat) = PDiagMat(a.diag .+ b.value)
-if HAVE_CHOLMOD
-    +(a::PDiagMat, b::PDSparseMat) = PDSparseMat(_adddiag(b.mat, a.diag))
-end
 
 +(a::ScalMat, b::PDMat) = PDMat(_adddiag(b.mat, a.value))
 +(a::ScalMat, b::PDiagMat) = PDiagMat(a.value .+ b.diag)
 +(a::ScalMat, b::ScalMat) = ScalMat(a.dim, a.value + b.value)
-if HAVE_CHOLMOD
-    +(a::ScalMat, b::PDSparseMat) = PDSparseMat(_adddiag(b.mat, a.value))
-end
-
-if HAVE_CHOLMOD
-    +(a::PDSparseMat, b::PDMat) = PDMat(Matrix(a) + b.mat)
-    +(a::PDSparseMat, b::PDiagMat) = PDSparseMat(_adddiag(a.mat, b.diag))
-    +(a::PDSparseMat, b::ScalMat) = PDSparseMat(_adddiag(a.mat, b.value))
-    +(a::PDSparseMat, b::PDSparseMat) = PDSparseMat(a.mat + b.mat)
-end
-
 
 # between pdmat and uniformscaling (multiple of identity)
 
@@ -45,34 +25,16 @@ end
 pdadd(a::PDMat, b::AbstractPDMat, c::Real) = PDMat(a.mat + Matrix(b * c))
 pdadd(a::PDiagMat, b::AbstractPDMat, c::Real) = PDMat(_adddiag!(Matrix(b * c), a.diag, one(c)))
 pdadd(a::ScalMat, b::AbstractPDMat, c::Real) = PDMat(_adddiag!(Matrix(b * c), a.value))
-if HAVE_CHOLMOD
-    pdadd(a::PDSparseMat, b::AbstractPDMat, c::Real) = PDMat(a.mat + Matrix(b * c))
-end
 
 pdadd(a::PDMat, b::PDMat, c::Real) = PDMat(a.mat + b.mat * c)
 pdadd(a::PDMat, b::PDiagMat, c::Real) = PDMat(_adddiag(a.mat, b.diag, c))
 pdadd(a::PDMat, b::ScalMat, c::Real) = PDMat(_adddiag(a.mat, b.value * c))
-if HAVE_CHOLMOD
-    pdadd(a::PDMat, b::PDSparseMat, c::Real) = PDMat(a.mat + b.mat * c)
-end
 
 pdadd(a::PDiagMat, b::PDMat, c::Real) = PDMat(_adddiag!(b.mat * c, a.diag, one(c)))
 pdadd(a::PDiagMat, b::PDiagMat, c::Real) = PDiagMat(a.diag + b.diag * c)
 pdadd(a::PDiagMat, b::ScalMat, c::Real) = PDiagMat(a.diag .+ b.value * c)
-if HAVE_CHOLMOD
-    pdadd(a::PDiagMat, b::PDSparseMat, c::Real) = PDSparseMat(_adddiag!(b.mat * c, a.diag, one(c)))
-end
 
 pdadd(a::ScalMat, b::PDMat, c::Real) = PDMat(_adddiag!(b.mat * c, a.value))
 pdadd(a::ScalMat, b::PDiagMat, c::Real) = PDiagMat(a.value .+ b.diag * c)
 pdadd(a::ScalMat, b::ScalMat, c::Real) = ScalMat(a.dim, a.value + b.value * c)
-if HAVE_CHOLMOD
-    pdadd(a::ScalMat, b::PDSparseMat, c::Real) = PDSparseMat(_adddiag!(b.mat * c, a.value))
-end
 
-if HAVE_CHOLMOD
-    pdadd(a::PDSparseMat, b::PDMat, c::Real) = PDMat(a.mat + b.mat * c)
-    pdadd(a::PDSparseMat, b::PDiagMat, c::Real) = PDSparseMat(_adddiag(a.mat, b.diag, c))
-    pdadd(a::PDSparseMat, b::ScalMat, c::Real) = PDSparseMat(_adddiag(a.mat, b.value * c))
-    pdadd(a::PDSparseMat, b::PDSparseMat, c::Real) = PDSparseMat(a.mat + b.mat * c)
-end

src/chol.jl

@@ -17,14 +17,6 @@ chol_lower(a::Matrix) = cholesky(Symmetric(a, :L)).L
 # type-restricted to be Real, so they are equivalent.
 chol_upper(a::Matrix) = cholesky(Symmetric(a, :U)).U
 
-if HAVE_CHOLMOD
-    CholTypeSparse{T} = SuiteSparse.CHOLMOD.Factor{T}
-
-    # Take into account pivoting!
-    chol_lower(cf::CholTypeSparse) = cf.PtL
-    chol_upper(cf::CholTypeSparse) = cf.UP
-end
-
 # Interface for `Cholesky`
 
 dim(A::Cholesky) = LinearAlgebra.checksquare(A)

src/deprecates.jl

@@ -17,7 +17,7 @@ using Base: @deprecate
 
 @deprecate dim(a::AbstractMatrix) LinearAlgebra.checksquare(a)
 
-@deprecate PDMat{T,S}(d::Int, m::AbstractMatrix{T}, c::Cholesky{T,S}) where {T,S} PDMat{T,S}(m, c)
+@deprecate PDMat{T,S}(d::Int, m::AbstractMatrix{T}, c::Cholesky{T,S}) where {T,S} PDMat{T,S,Cholesky{T,S}}(m, c)
 
 @deprecate PDiagMat(dim::Int, diag::AbstractVector{<:Real}) PDiagMat(diag)
 @deprecate PDiagMat{T,V}(dim, diag) where {T<:Real, V<:AbstractVector{T}} PDiagMat{T,V}(diag)