Constructor for PDMat with Cholesky pre- and post- multiplied by diagonal matrix

[sethaxen]

Suppose a user has a Cholesky decomposition of a matrix R and a vector of positive entries σ and wishes to create Σ = PDMat(cholesky(σ .* Matrix(R) .* σ')). This is common in Bayesian inference. e.g. a Cholesky factor of a correlation matrix is drawn from a LKJCholesky distribution, while a vector of standard deviations is drawn from some other distribution, and one wants to construct the Cholesky factorization of the covariance matrix for use in MvNormal.

This can be done much more efficiently with Σ = PDMat(Cholesky(R.uplo == 'U' ? R.U * Diagonal(σ) : Diagonal(σ) * R.L, R.uplo, 0)). But this is still quite long. It would be ideal to have either a constructor for PDMat or some other utility function like the following:

function PDMat(fac::Cholesky, scale::AbstractVector)
    factors_scaled = fac.uplo == 'U' ? fac.factors .* scale' : scale .* fac.factors
    return PDMat(Cholesky(factors_scaled, fac.uplo, 0))
end
function PDMat(fac::Cholesky, scale::UniformScaling)
    return PDMat(Cholesky(fac.factors * scale, fac.uplo, 0))
end
PDMat(fac::Cholesky, scale::Number) = PDMat(fac, scale * I)

See also discussions in JuliaStats/Distributions.jl#1336 and TuringLang/Turing.jl#1629 (comment).

[sethaxen]

This kind of transform is called a congruence transformation and is already computed by the API function X_A_Xt. So all we would have to do is provide the following overloads to ensure that for cases where we can guarantee the result is PD of a certain structure and where we can avoid recomputing a factorization or using excess storage, we implement the necessary overloads to return the appropriate AbstractPDMat subtype:

import PDMats: X_A_Xt

X_A_Xt(A::PDiagMat, B::PDiagMat) = PDiagMat(abs2.(B.diag) .* A.diag)
X_A_Xt(A::PDiagMat, B::ScalMat) = ScalMat(B.dim, abs2(B.value)) * A
function X_A_Xt(A::ScalMat, B::PDiagMat)
    A.dim == size(B, 1) || throw(DimensionMismatch("Inconsistent argument dimensions."))
    return PDiagMat(A.value .* abs2.(B.diag))
end
X_A_Xt(A::ScalMar, B::ScalMat) = B * A * B'
function X_A_Xt(A::PDMat, B::ScalMat)
    B.dim == size(A, 1) || throw(DimensionMismatch("Inconsistent argument dimensions."))
    b2 = abs2(B.value)
    mat = b2 * A.mat
    uplo = A.chol.uplo
    if uplo === 'U'
        factors = A.chol.factors * B.value'
    else
        factors = B.value * A.chol.factors
    end
    chol = Cholesky(factors, uplo, A.chol.info)
    return PDMat(mat, chol)
end
function X_A_Xt(A::PDMat, B::PDiagMat)
    b = B.diag
    mat = A.mat .* (b .* b')
    uplo = A.chol.uplo
    if uplo === 'U'
        factors = A.chol.factors .* b'
    else
        factors = b .* A.chol.factors
    end
    chol = Cholesky(factors, uplo, A.chol.info)
    return PDMat(mat, chol)
end

Similar overloads could in some cases be provided for Xt_A_X, X_invA_Xt, and Xt_invA_X.

Then in the future a user could do something like this:

julia> R = rand(LKJPDMat(10, 1.0));  # correlation matrix

julia> s = rand(truncated(TDist(7); lower=0), 10);  # standard deviations

julia> Σ = X_A_Xt(R, PDiagMat(s));  # covariance matrix

[sethaxen]

@devmotion what do you think of this solution?

[devmotion]

I like this proposal.