Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unify PDMat and PDSparseMat + move SparseArrays support to an extension #188

Open
wants to merge 18 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,7 @@ jobs:
fail-fast: false
matrix:
version:
- '1.0'
- '1.6'
- '1.9'
- '1'
- nightly
os:
Expand Down
15 changes: 10 additions & 5 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,19 +1,24 @@
name = "PDMats"
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
version = "0.11.22"
version = "0.12.0"

[deps]
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

[compat]
julia = "1"
julia = "1.9"

[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"
33 changes: 5 additions & 28 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,17 +27,17 @@ 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}
dim::Int # matrix dimension
mat::S # input matrix
chol::Cholesky{T,S} # Cholesky factorization of mat
struct PDMat{T<:Real,S<:AbstractMatrix,C<:Factorization} <: AbstractPDMat{T}
dim::Int # matrix dimension
mat::S # input matrix
chol::C # Cholesky 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.

Expand Down Expand Up @@ -76,29 +76,6 @@ 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}
dim::Int # matrix dimension
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.
Expand Down
15 changes: 15 additions & 0 deletions ext/PDMatsSparseArraysExt/PDMatsSparseArraysExt.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
module PDMatsSparseArraysExt

using PDMats
using SparseArrays

using PDMats.LinearAlgebra
dkarrasch marked this conversation as resolved.
Show resolved Hide resolved

const HAVE_CHOLMOD = isdefined(SparseArrays, :CHOLMOD)

if HAVE_CHOLMOD
include("chol.jl")
include("pdsparsemat.jl")
end

end # module
5 changes: 5 additions & 0 deletions ext/PDMatsSparseArraysExt/chol.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
const CholTypeSparse{T} = SparseArrays.CHOLMOD.Factor{T}

# Take into account pivoting!
PDMats.chol_lower(cf::CholTypeSparse) = cf.PtL
PDMats.chol_upper(cf::CholTypeSparse) = cf.UP
109 changes: 109 additions & 0 deletions ext/PDMatsSparseArraysExt/pdsparsemat.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
"""
Sparse positive definite matrix together with a Cholesky factorization object.
"""
const PDSparseMat{T<:Real,S<:AbstractSparseMatrix,C<:CholTypeSparse} = PDMat{T,S,C}

function PDMats.PDMat(mat::AbstractSparseMatrix, chol::CholTypeSparse)
d = size(mat, 1)
size(chol, 1) == d ||
throw(DimensionMismatch("Dimensions of mat and chol are inconsistent."))
PDMat{eltype(mat),typeof(mat),typeof(chol)}(d, mat, chol)
end

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.chol)}(a.dim, mat, a.chol)
end

### Arithmetics

Base.:\(a::PDSparseMat{T}, x::AbstractVecOrMat{T}) where {T<:Real} = convert(Array{T},a.chol \ 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.chol )

### Algebra

Base.inv(a::PDSparseMat{T}) where {T<:Real} = PDMat(inv(a.mat))
LinearAlgebra.det(a::PDSparseMat) = det(a.chol)
LinearAlgebra.logdet(a::PDSparseMat) = logdet(a.chol)
LinearAlgebra.sqrt(A::PDSparseMat) = PDMat(sqrt(Hermitian(Matrix(A))))

### whiten and unwhiten

function PDMats.whiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
# Can't use `ldiv!` due to missing support in SparseArrays
return copyto!(r, PDMats.chol_lower(a.chol) \ x)
end

function PDMats.unwhiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
# `*` is not defined for `PtL` factor components,
# so we can't use `chol_lower(a.chol) * x`
C = a.chol
PtL = sparse(C.L)[C.p, :]
# Can't use `lmul!` due to missing support in SparseArrays
return copyto!(r, PtL * x)
end


### quadratic forms

PDMats.quad(a::PDSparseMat, x::AbstractVector) = dot(x, a * x)
PDMats.invquad(a::PDSparseMat, x::AbstractVector) = dot(x, a \ x)

function PDMats.quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
PDMats.@check_argdims eachindex(r) == axes(x, 2)
for i in axes(x, 2)
r[i] = quad(a, x[:,i])
end
return r
end

function PDMats.invquad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
PDMats.@check_argdims eachindex(r) == axes(x, 2)
for i in axes(x, 2)
r[i] = invquad(a, x[:,i])
end
return r
end


### tri products

function PDMats.X_A_Xt(a::PDSparseMat, x::AbstractMatrix)
# `*` is not defined for `PtL` factor components,
# so we can't use `x * chol_lower(a.chol)`
C = a.chol
PtL = sparse(C.L)[C.p, :]
z = x * PtL
z * transpose(z)
end


function PDMats.Xt_A_X(a::PDSparseMat, x::AbstractMatrix)
# `*` is not defined for `UP` factor components,
# so we can't use `chol_upper(a.chol) * x`
# Moreover, `sparse` is only defined for `L` factor components
C = a.chol
UP = transpose(sparse(C.L))[:, C.p]
z = UP * x
transpose(z) * z
end


function PDMats.X_invA_Xt(a::PDSparseMat, x::AbstractMatrix)
z = a.chol \ collect(transpose(x))
x * z
end

function PDMats.Xt_invA_X(a::PDSparseMat, x::AbstractMatrix)
z = a.chol \ x
transpose(x) * z
end
10 changes: 1 addition & 9 deletions src/PDMats.jl
Original file line number Diff line number Diff line change
@@ -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,

Expand All @@ -35,19 +34,12 @@ module PDMats
"""
abstract type AbstractPDMat{T<:Real} <: AbstractMatrix{T} end

const HAVE_CHOLMOD = isdefined(SuiteSparse, :CHOLMOD)

# source files

include("chol.jl")
include("utils.jl")

include("pdmat.jl")

if HAVE_CHOLMOD
include("pdsparsemat.jl")
end

include("pdiagmat.jl")
include("scalmat.jl")

Expand Down
38 changes: 0 additions & 38 deletions src/addition.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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
8 changes: 0 additions & 8 deletions src/chol.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,3 @@ chol_lower(a::Matrix) = cholesky(Symmetric(a, :L)).L
# but this currently has an AutoDiff issue in Zygote.jl, and PDMat is
# 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
7 changes: 1 addition & 6 deletions src/pdiagmat.jl
Original file line number Diff line number Diff line change
Expand Up @@ -62,12 +62,7 @@ function \(a::PDiagMat, x::AbstractVecOrMat)
end
function /(x::AbstractVecOrMat, a::PDiagMat)
@check_argdims a.dim == size(x, 2)
if VERSION < v"1.9-"
# return matrix for 1-element vectors `x`, consistent with LinearAlgebra < 1.9
return reshape(x, Val(2)) ./ permutedims(a.diag) # = (a' \ x')'
else
return x ./ (x isa AbstractVector ? a.diag : a.diag')
end
return x ./ (x isa AbstractVector ? a.diag : a.diag')
end
Base.kron(A::PDiagMat, B::PDiagMat) = PDiagMat(vec(permutedims(A.diag) .* B.diag))

Expand Down
Loading