Probabilistic PCA

Overview of PCA

Principal component analysis (PCA) is a fundamental technique to analyse and visualise data. It is an unsupervised learning method mainly used for dimensionality reduction.

For example, we have a data matrix \(\mathbf{X} \in \mathbb{R}^{N \times D}\), and we would like to extract \(k \ll D\) principal components which captures most of the information from the original matrix. The goal is to understand \(\mathbf{X}\) through a lower dimensional subspace (e.g. two-dimensional subspace for visualisation convenience) spanned by the principal components.

In order to project the original data matrix into low dimensions, we need to find the principal directions where most of the variations of \(\mathbf{X}\) lie in. Traditionally, this is implemented via singular value decomposition (SVD) which provides a robust and accurate computational framework for decomposing matrix into products of rotation-scaling-rotation matrices, particularly for large datasets(see an illustration here):

\[ \mathbf{X}_{N \times D} = \mathbf{U}_{N \times r} \times \boldsymbol{\Sigma}_{r \times r} \times \mathbf{V}^T_{r \times D} \]

where \(\Sigma_{r \times r}\) contains only \(r := \operatorname{rank} \mathbf{X} \leq \min\{N,D\}\) non-zero singular values of \(\mathbf{X}\). If we pad \(\Sigma\) with zeros and add arbitrary orthonormal columns to \(\mathbf{U}\) and \(\mathbf{V}\), we obtain the more compact form:¹

\[ \mathbf{X}_{N \times D} = \mathbf{U}_{N \times N} \mathbf{\Sigma}_{N \times D} \mathbf{V}_{D \times D}^T \]

where \(\mathbf{U}\) and \(\mathbf{V}\) are unitary matrices (i.e. with orthonormal columns). Such a decomposition always exists for any matrix. Columns of \(\mathbf{V}\) are the principal directions/axes. The percentage of variations explained can be calculated using the ratios of singular values.²

Here we take a probabilistic perspective. For more details and a mathematical derivation, we recommend Bishop’s textbook (Christopher M. Bishop, Pattern Recognition and Machine Learning, 2006). The idea of proabilistic PCA is to find a latent variable \(z\) that can be used to describe the hidden structure in a dataset.³ Consider a data set \(\mathbf{X}_{D \times N}=\{x_i\}\) with \(i=1,2,...,N\) data points, where each data point \(x_i\) is \(D\)-dimensional (i.e. \(x_i \in \mathcal{R}^D\)). Note that, here we use the flipped version of the data matrix. We aim to represent the original \(n\) dimensional vector using a lower dimensional a latent variable \(z_i \in \mathcal{R}^k\).

We first assume that each latent variable \(z_i\) is normally distributed:

\[ z_i \sim \mathcal{N}(0, I) \]

and the corresponding data point is generated via projection:

\[ x_i | z_i \sim \mathcal{N}(\mathbf{W} z_i + \boldsymbol{μ}, \sigma^2 \mathbf{I}) \]

where the projection matrix \(\mathbf{W}_{D \times k}\) accommodates the principal axes. The above formula expresses \(x_i\) as a linear combination of the basis columns in the projection matrix W, where the combination coefficients sit in z_i (they are the coordinats of x_i in the new \(k\)-dimensional space.). We can also express the above formula in matrix form: \(\mathbf{X}_{D \times N} \approx \mathbf{W}_{D \times k} \mathbf{Z}_{k \times N}\). We are interested in inferring \(\mathbf{W}\), \(μ\) and \(\sigma\).

Classical PCA is the specific case of probabilistic PCA when the covariance of the noise becomes infinitesimally small, i.e. \(\sigma^2 \to 0\). Probabilistic PCA generalizes classical PCA, this can be seen by marginalizing out the the latent variable.⁴

The gene expression example

In the first example, we illustrate:

how to specify the probabilistic model and
how to perform inference on \(\mathbf{W}\), \(\boldsymbol{\mu}\) and \(\sigma\) using MCMC.

We use simulated gemnome data to demonstrate these. The simulation is inspired by biological measurement of expression of genes in cells, and each cell is characterized by different gene features. While the human genome is (mostly) identical between all the cells in the body, there exist interesting differences in gene expression in different human tissues and disease conditions. One way to investigate certain diseases is to look at differences in gene expression in cells from patients and healthy controls (usually from the same tissue).

Usually, we can assume that the changes in gene expression only affect a subset of all genes (and these can be linked to diseases in some way). One of the challenges for this kind of data is to explore the underlying structure, e.g. to make the connection between a certain state (healthy/disease) and gene expression. This becomes difficult when the dimensions is very large (up to 20000 genes across 1000s of cells). So in order to find structure in this data, it is useful to project the data into a lower dimensional space.

Regardless of the biological background, the more abstract problem formulation is to project the data living in high-dimensional space onto a representation in lower-dimensional space where most of the variation is concentrated in the first few dimensions. We use PCA to explore underlying structure or pattern which may not necessarily be obvious from looking at the raw data itself.

Step 1: configuration of dependencies

First, we load the dependencies used.

using Turing
using Mooncake
using LinearAlgebra, FillArrays

# Packages for visualization
using DataFrames, StatsPlots, Measures

# Set a seed for reproducibility.
using Random
Random.seed!(1789);

Precompiling Turing...
    848.0 ms  ? OptimizationBase
   1482.3 ms  ? Optimization
   2168.4 ms  ? OptimizationOptimJL
Info Given Turing was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
   5749.9 ms  ? Turing
   5894.3 ms  ? Turing → TuringOptimExt
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
Precompiling Optimization...
    810.8 ms  ? OptimizationBase
Info Given Optimization was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
   1493.7 ms  ? Optimization
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
Precompiling OptimizationBase...
Info Given OptimizationBase was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
    806.1 ms  ? OptimizationBase
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
┌ Warning: Replacing docs for `CommonSolve.solve :: Tuple{SciMLBase.OptimizationProblem, Any, Vararg{Any}}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
┌ Warning: Replacing docs for `CommonSolve.init :: Tuple{SciMLBase.OptimizationProblem, Any, Vararg{Any}}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
┌ Warning: Replacing docs for `CommonSolve.solve! :: Tuple{SciMLBase.AbstractOptimizationCache}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
Precompiling OptimizationOptimJL...
    819.7 ms  ? OptimizationBase
   1039.0 ms  ? Optimization
Info Given OptimizationOptimJL was explicitly requested, output will be shown live 
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-d70f-644f316ff045 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
   1181.7 ms  ? OptimizationOptimJL
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-d70f-644f316ff045 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
Precompiling TuringOptimExt...
    815.3 ms  ? OptimizationBase
   1016.4 ms  ? Optimization
   1151.8 ms  ? OptimizationOptimJL
   3880.0 ms  ? Turing
Info Given TuringOptimExt was explicitly requested, output will be shown live 
┌ Warning: Module Turing with build ID ffffffff-ffff-ffff-334a-b015a4bc4d53 is missing from the cache.
│ This may mean Turing [fce5fe82-541a-59a6-adf8-730c64b5f9a0] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
    638.6 ms  ? Turing → TuringOptimExt
┌ Warning: Module Turing with build ID ffffffff-ffff-ffff-334a-b015a4bc4d53 is missing from the cache.
│ This may mean Turing [fce5fe82-541a-59a6-adf8-730c64b5f9a0] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541

All packages used in this tutorial are listed here. You can install them via using Pkg; Pkg.add("package_name").

Package usages:

We use DataFrames for instantiating matrices, LinearAlgebra and FillArrays to perform matrix operations; Turing for model specification and MCMC sampling, Mooncake for automatic differentiation when sampling. StatsPlots for visualising the resutls. , Measures for setting plot margin units. As all examples involve sampling, for reproducibility we set a fixed seed using the Random standard library.

Step 2: Data generation

Here, we simulate the biological gene expression problem described earlier. We simulate 60 cells, each cell has 9 gene features. This is a simplified problem with only a few cells and genes for demonstration purpose, which is not comparable to the complexity in real-life (e.g. thousands of features for each individual). Even so, spotting the structures or patterns in a 9-feature space would be a challenging task; it would be nice to reduce the dimentionality using p-PCA.

By design, we mannually divide the 60 cells into two groups. the first 3 gene features of the first 30 cells have mean 10, while those of the last 30 cells have mean 10. These two groups of cells differ in the expression of genes.

n_genes = 9 # D
n_cells = 60 # N

# create a diagonal block like expression matrix, with some non-informative genes;
# not all features/genes are informative, some might just not differ very much between cells)
mat_exp = randn(n_genes, n_cells)
mat_exp[1:(n_genes ÷ 3), 1:(n_cells ÷ 2)] .+= 10
mat_exp[(2 * (n_genes ÷ 3) + 1):end, (n_cells ÷ 2 + 1):end] .+= 10

3×30 view(::Matrix{Float64}, 7:9, 31:60) with eltype Float64:
 11.7413   10.5735   11.3817    8.50923  …   7.38716   8.7734   11.4717
  9.28533  11.1225    9.43421  10.8904      11.6846   10.7264    9.64063
  9.92113   8.42122   9.59885   9.90799      9.40715   8.40956  10.2522

To visualize the \((D=9) \times (N=60)\) data matrix mat_exp, we use the heatmap plot.

heatmap(
    mat_exp;
    c=:summer,
    colors=:value,
    xlabel="cell number",
    yflip=true,
    ylabel="gene feature",
    yticks=1:9,
    colorbar_title="expression",
)

Note that:

We have made distinct feature differences between these two groups of cells (it is fairly obvious from looking at the raw data), in practice and with large enough data sets, it is often impossible to spot the differences from the raw data alone.
If you have some patience and compute resources you can increase the size of the dataset, or play around with the noise levels to make the problem increasingly harder.

Step 3: Create the pPCA model

Here we construct the probabilistic model pPCA(). As per the p-PCA formula, we think of each row (i.e. each gene feature) following a \(N=60\) dimensional multivariate normal distribution centered around the corresponding row of \(\mathbf{W}_{D \times k} \times \mathbf{Z}_{k \times N} + \boldsymbol{\mu}_{D \times N}\).

@model function pPCA(X::AbstractMatrix{<:Real}, k::Int)
    # retrieve the dimension of input matrix X.
    N, D = size(X)

    # weights/loadings W
    W ~ filldist(Normal(), D, k)

    # latent variable z
    Z ~ filldist(Normal(), k, N)

    # mean offset
    μ ~ MvNormal(Eye(D))
    genes_mean = W * Z .+ reshape(μ, n_genes, 1)
    return X ~ arraydist([MvNormal(m, Eye(N)) for m in eachcol(genes_mean')])
end;

The function pPCA() accepts:

an data array \(\mathbf{X}\) (with no. of instances x dimension no. of features, NB: it is a transpose of the original data matrix);
an integer \(k\) which indicates the dimension of the latent space (the space the original feature matrix is projected onto).

Specifically:

it first extracts the dimension \(D\) and number of instances \(N\) of the input matrix;
draw samples of each entries of the projection matrix \(\mathbf{W}\) from a standard normal;
draw samples of the latent variable \(\mathbf{Z}_{k \times N}\) from an MND;
draw samples of the offset \(\boldsymbol{\mu}\) from an MND, assuming uniform offset for all instances;
Finally, we iterate through each gene dimension in \(\mathbf{X}\), and define an MND for the sampling distribution (i.e. likelihood).

Step 4: Sampling-based inference of the pPCA model

Here we aim to perform MCMC sampling to infer the projection matrix \(\mathbf{W}_{D \times k}\), the latent variable matrix \(\mathbf{Z}_{k \times N}\), and the offsets \(\boldsymbol{\mu}_{N \times 1}\).

We run the inference using the NUTS sampler. By default, sample samples a single chain (in this case with 500 samples). You can also use different samplers if you wish.

setprogress!(false)

k = 2 # k is the dimension of the projected space, i.e. the number of principal components/axes of choice
ppca = pPCA(mat_exp', k) # instantiate the probabilistic model
chain_ppca = sample(ppca, NUTS(; adtype=AutoMooncake()), 500);

┌ Info: Found initial step size
└   ϵ = 0.4

The samples are saved in chain_ppca, which is an MCMCChains.Chains object. We can check its shape:

size(chain_ppca) # (no. of iterations, no. of vars, no. of chains) = (500, 159, 1)

(500, 161, 1)

Sampling statistics such as R-hat, ESS, mean estimates, and so on can also be obtained from this:

describe(chain_ppca)

Chains MCMC chain (500×161×1 Array{Float64, 3}):

Iterations        = 251:1:750
Number of chains  = 1
Samples per chain = 500
Wall duration     = 124.77 seconds
Compute duration  = 124.77 seconds
parameters        = W[1, 1], W[2, 1], W[3, 1], W[4, 1], W[5, 1], W[6, 1], W[7, 1], W[8, 1], W[9, 1], W[1, 2], W[2, 2], W[3, 2], W[4, 2], W[5, 2], W[6, 2], W[7, 2], W[8, 2], W[9, 2], Z[1, 1], Z[2, 1], Z[1, 2], Z[2, 2], Z[1, 3], Z[2, 3], Z[1, 4], Z[2, 4], Z[1, 5], Z[2, 5], Z[1, 6], Z[2, 6], Z[1, 7], Z[2, 7], Z[1, 8], Z[2, 8], Z[1, 9], Z[2, 9], Z[1, 10], Z[2, 10], Z[1, 11], Z[2, 11], Z[1, 12], Z[2, 12], Z[1, 13], Z[2, 13], Z[1, 14], Z[2, 14], Z[1, 15], Z[2, 15], Z[1, 16], Z[2, 16], Z[1, 17], Z[2, 17], Z[1, 18], Z[2, 18], Z[1, 19], Z[2, 19], Z[1, 20], Z[2, 20], Z[1, 21], Z[2, 21], Z[1, 22], Z[2, 22], Z[1, 23], Z[2, 23], Z[1, 24], Z[2, 24], Z[1, 25], Z[2, 25], Z[1, 26], Z[2, 26], Z[1, 27], Z[2, 27], Z[1, 28], Z[2, 28], Z[1, 29], Z[2, 29], Z[1, 30], Z[2, 30], Z[1, 31], Z[2, 31], Z[1, 32], Z[2, 32], Z[1, 33], Z[2, 33], Z[1, 34], Z[2, 34], Z[1, 35], Z[2, 35], Z[1, 36], Z[2, 36], Z[1, 37], Z[2, 37], Z[1, 38], Z[2, 38], Z[1, 39], Z[2, 39], Z[1, 40], Z[2, 40], Z[1, 41], Z[2, 41], Z[1, 42], Z[2, 42], Z[1, 43], Z[2, 43], Z[1, 44], Z[2, 44], Z[1, 45], Z[2, 45], Z[1, 46], Z[2, 46], Z[1, 47], Z[2, 47], Z[1, 48], Z[2, 48], Z[1, 49], Z[2, 49], Z[1, 50], Z[2, 50], Z[1, 51], Z[2, 51], Z[1, 52], Z[2, 52], Z[1, 53], Z[2, 53], Z[1, 54], Z[2, 54], Z[1, 55], Z[2, 55], Z[1, 56], Z[2, 56], Z[1, 57], Z[2, 57], Z[1, 58], Z[2, 58], Z[1, 59], Z[2, 59], Z[1, 60], Z[2, 60], μ[1], μ[2], μ[3], μ[4], μ[5], μ[6], μ[7], μ[8], μ[9]
internals         = n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size, lp, logprior, loglikelihood

Summary Statistics
  parameters      mean       std      mcse   ess_bulk   ess_tail      rhat   e ⋯
      Symbol   Float64   Float64   Float64    Float64    Float64   Float64     ⋯

     W[1, 1]   -2.8296    0.5668    0.1343    16.5461    63.6634    1.0044     ⋯
     W[2, 1]   -3.0151    0.5332    0.1159    19.4966    48.7369    1.0016     ⋯
     W[3, 1]   -2.6595    0.6334    0.1547    15.6084    58.9459    1.0071     ⋯
     W[4, 1]   -0.1128    0.1660    0.0297    30.0485    66.6391    1.0016     ⋯
     W[5, 1]   -0.3150    0.2174    0.0450    23.3737    41.2575    1.0020     ⋯
     W[6, 1]   -0.0492    0.1458    0.0093   245.8926   263.5520    1.0016     ⋯
     W[7, 1]    3.0339    0.4369    0.0850    25.2554    56.5235    0.9983     ⋯
     W[8, 1]    2.7789    0.6417    0.1575    15.5515    53.9110    1.0053     ⋯
     W[9, 1]    2.9742    0.5518    0.1245    17.8949    55.6637    1.0012     ⋯
     W[1, 2]    1.6836    0.9658    0.2732    13.0954    29.7267    1.0007     ⋯
     W[2, 2]    1.5033    1.0130    0.2841    13.0654    31.8178    1.0021     ⋯
     W[3, 2]    1.8885    0.9115    0.2557    13.1707    31.9881    1.0002     ⋯
     W[4, 2]   -0.3311    0.1643    0.0106   256.6495   304.3756    0.9994     ⋯
     W[5, 2]   -0.4704    0.2004    0.0262    56.4601   211.6755    1.0097     ⋯
     W[6, 2]    0.0644    0.1878    0.0132   200.6478   365.6208    1.0108     ⋯
     W[7, 2]   -1.2011    1.0344    0.2965    12.6951    28.9230    0.9995     ⋯
     W[8, 2]   -1.9229    0.9556    0.2687    13.0725    29.1472    1.0016     ⋯
      ⋮           ⋮         ⋮         ⋮         ⋮          ⋮          ⋮        ⋱
                                                   1 column and 130 rows omitted

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

     W[1, 1]   -3.6662   -3.2431   -2.8870   -2.5186   -1.4836
     W[2, 1]   -3.7522   -3.3926   -3.1054   -2.7600   -1.6272
     W[3, 1]   -3.5615   -3.1479   -2.7264   -2.3402   -1.1093
     W[4, 1]   -0.4423   -0.2106   -0.1122    0.0036    0.1810
     W[5, 1]   -0.7532   -0.4556   -0.3179   -0.1726    0.0993
     W[6, 1]   -0.3441   -0.1328   -0.0563    0.0394    0.2419
     W[7, 1]    1.8604    2.8300    3.0945    3.3175    3.7213
     W[8, 1]    1.1907    2.4429    2.8725    3.2570    3.6932
     W[9, 1]    1.4473    2.6918    3.0668    3.3650    3.8001
     W[1, 2]   -0.6491    1.1287    1.8429    2.3714    3.1602
     W[2, 2]   -1.0512    0.9783    1.6632    2.2130    3.1345
     W[3, 2]   -0.4531    1.3831    2.0505    2.5013    3.3082
     W[4, 2]   -0.6342   -0.4499   -0.3291   -0.2179   -0.0312
     W[5, 2]   -0.8574   -0.6090   -0.4818   -0.3432   -0.0480
     W[6, 2]   -0.3074   -0.0520    0.0688    0.1882    0.4098
     W[7, 2]   -2.7925   -1.9070   -1.3899   -0.5715    1.2784
     W[8, 2]   -3.3290   -2.6040   -2.1132   -1.4094    0.6130
      ⋮           ⋮         ⋮         ⋮         ⋮         ⋮
                                                130 rows omitted

Step 5: posterior predictive checks

We try to reconstruct the input data using the posterior mean as parameter estimates. We first retrieve the samples for the projection matrix W from chain_ppca. This can be done using the Julia group(chain, parameter_name) function. Then we calculate the mean value for each element in \(W\), averaging over the whole chain of samples.

# Extract parameter estimates for predicting x - mean of posterior
W = reshape(mean(group(chain_ppca, :W))[:, 2], (n_genes, k))
Z = reshape(mean(group(chain_ppca, :Z))[:, 2], (k, n_cells))
μ = mean(group(chain_ppca, :μ))[:, 2]

mat_rec = W * Z .+ repeat(μ; inner=(1, n_cells))

9×60 Matrix{Float64}:
  9.31871     9.5518     10.0025     …   0.64444    0.8444      0.223469
  9.68383     9.83608    10.4575         1.16466    0.996148    0.092868
  9.26425     9.58451     9.85939        0.349686   0.939309    0.610091
 -0.236833   -0.359212   -0.147743       0.408592  -0.106151   -0.461941
  0.123233   -0.0683059   0.301473       0.794934  -0.0461873  -0.662084
  0.0485938   0.0642583   0.0550314  …  -0.189659  -0.140087   -0.121693
  0.63398     0.583789   -0.19291        8.43082    9.00929    10.1772
  0.405128    0.0871515  -0.224566       9.59464    9.04471     9.43062
  0.267097    0.0837768  -0.480546       8.9202     8.95487     9.76258

heatmap(
    mat_rec;
    c=:summer,
    colors=:value,
    xlabel="cell number",
    yflip=true,
    ylabel="gene feature",
    yticks=1:9,
    colorbar_title="expression",
)

We can quantitatively check the absolute magnitudes of the column average of the gap between mat_exp and mat_rec:

diff_matrix = mat_exp .- mat_rec
for col in 4:6
    @assert abs(mean(diff_matrix[:, col])) <= 0.5
end

We observe that, using posterior mean, the recovered data matrix mat_rec has values align with the original data matrix - particularly the same pattern in the first and last 3 gene features are captured, which implies the inference and p-PCA decomposition are successful. This is satisfying as we have just projected the original 9-dimensional space onto a 2-dimensional space - some info has been cut off in the projection process, but we haven’t lost any important info, e.g. the key differences between the two groups. The is the desirable property of PCA: it picks up the principal axes along which most of the (original) data variations cluster, and remove those less relevant. If we choose the reduced space dimension \(k\) to be exactly \(D\) (the original data dimension), we would recover exactly the same original data matrix mat_exp, i.e. all information will be preserved.

Now we have represented the original high-dimensional data in two dimensions, without lossing the key information about the two groups of cells in the input data. Finally, the benefits of performing PCA is to analyse and visualise the dimension-reduced data in the projected, low-dimensional space. we save the dimension-reduced matrix \(\mathbf{Z}\) as a DataFrame, rename the columns and visualise the first two dimensions.

df_pca = DataFrame(Z', :auto)
rename!(df_pca, Symbol.(["z" * string(i) for i in collect(1:k)]))
df_pca[!, :type] = repeat([1, 2]; inner=n_cells ÷ 2)

scatter(df_pca[:, :z1], df_pca[:, :z2]; xlabel="z1", ylabel="z2", group=df_pca[:, :type])

We see the two groups are well separated in this 2-D space. As an unsupervised learning method, performing PCA on this dataset gives membership for each cell instance. Another way to put it: 2 dimensions is enough to capture the main structure of the data.

Further extension: automatic choice of the number of principal components with ARD

A direct question arises from above practice is: how many principal components do we want to keep, in order to sufficiently represent the latent structure in the data? This is a very central question for all latent factor models, i.e. how many dimensions are needed to represent that data in the latent space. In the case of PCA, there exist a lot of heuristics to make that choice. For example, We can tune the number of principal components using empirical methods such as cross-validation based on some criteria such as MSE between the posterior predicted (e.g. mean predictions) data matrix and the original data matrix or the percentage of variation explained ⁵.

For p-PCA, this can be done in an elegant and principled way, using a technique called Automatic Relevance Determination (ARD). ARD can help pick the correct number of principal directions by regularizing the solution space using a parameterized, data-dependent prior distribution that effectively prunes away redundant or superfluous features ⁶. Essentially, we are using a specific prior over the factor loadings \(\mathbf{W}\) that allows us to prune away dimensions in the latent space. The prior is determined by a precision hyperparameter \(\alpha\). Here, smaller values of \(\alpha\) correspond to more important components. You can find more details about this in, for example, Bishop (2006) ⁷.

@model function pPCA_ARD(X)
    # Dimensionality of the problem.
    N, D = size(X)

    # latent variable Z
    Z ~ filldist(Normal(), D, N)

    # weights/loadings w with Automatic Relevance Determination part
    α ~ filldist(Gamma(1.0, 1.0), D)
    W ~ filldist(MvNormal(zeros(D), 1.0 ./ sqrt.(α)), D)

    mu = (W' * Z)'

    tau ~ Gamma(1.0, 1.0)
    return X ~ arraydist([MvNormal(m, 1.0 / sqrt(tau)) for m in eachcol(mu)])
end;

Instead of drawing samples of each entry in \(\mathbf{W}\) from a standard normal, this time we repeatedly draw \(D\) samples from the \(D\)-dimensional MND, forming a \(D \times D\) matrix \(\mathbf{W}\). This matrix is a function of \(\alpha\) as the samples are drawn from the MND parameterized by \(\alpha\). We also introduce a hyper-parameter \(\tau\) which is the precision in the sampling distribution. We also re-paramterise the sampling distribution, i.e. each dimension across all instances is a 60-dimensional multivariate normal distribution. Re-parameterisation can sometimes accelrate the sampling process.

We instantiate the model and ask Turing to sample from it using NUTS sampler. The sample trajectories of \(\alpha\) is plotted using the plot function from the package StatsPlots.

ppca_ARD = pPCA_ARD(mat_exp') # instantiate the probabilistic model
chain_ppcaARD = sample(ppca_ARD, NUTS(; adtype=AutoMooncake()), 500) # sampling
plot(group(chain_ppcaARD, :α); margin=6.0mm)

┌ Info: Found initial step size
└   ϵ = 0.0125

Again, we do some inference diagnostics. Here we look at the convergence of the chains for the \(α\) parameter. This parameter determines the relevance of individual components. We see that the chains have converged and the posterior of the \(\alpha\) parameters is centered around much smaller values in two instances. In the following, we will use the mean of the small values to select the relevant dimensions (remember that, smaller values of \(\alpha\) correspond to more important components.). We can clearly see from the values of \(\alpha\) that there should be two dimensions (corresponding to \(\bar{\alpha}_3=\bar{\alpha}_5≈0.05\)) for this dataset.

# Extract parameter mean estimates of the posterior
W = permutedims(reshape(mean(group(chain_ppcaARD, :W))[:, 2], (n_genes, n_genes)))
Z = permutedims(reshape(mean(group(chain_ppcaARD, :Z))[:, 2], (n_genes, n_cells)))'
α = mean(group(chain_ppcaARD, :α))[:, 2]
plot(α; label="α")

We can inspect α to see which elements are small (i.e. high relevance). To do this, we first sort α using sortperm() (in ascending order by default), and record the indices of the first two smallest values (among the \(D=9\) \(\alpha\) values). After picking the desired principal directions, we extract the corresponding subset loading vectors from \(\mathbf{W}\), and the corresponding dimensions of \(\mathbf{Z}\). We obtain a posterior predicted matrix \(\mathbf{X} \in \mathbb{R}^{2 \times 60}\) as the product of the two sub-matrices, and compare the recovered info with the original matrix.

α_indices = sortperm(α)[1:2]
k = size(α_indices)[1]
X_rec = W[:, α_indices] * Z[α_indices, :]

df_rec = DataFrame(X_rec', :auto)
heatmap(
    X_rec;
    c=:summer,
    colors=:value,
    xlabel="cell number",
    yflip=true,
    ylabel="gene feature",
    yticks=1:9,
    colorbar_title="expression",
)

We observe that, the data in the original space is recovered with key information, the distinct feature values in the first and last three genes for the two cell groups, are preserved. We can also examine the data in the dimension-reduced space, i.e. the selected components (rows) in \(\mathbf{Z}\).

df_pro = DataFrame(Z[α_indices, :]', :auto)
rename!(df_pro, Symbol.(["z" * string(i) for i in collect(1:k)]))
df_pro[!, :type] = repeat([1, 2]; inner=n_cells ÷ 2)
scatter(
    df_pro[:, 1], df_pro[:, 2]; xlabel="z1", ylabel="z2", color=df_pro[:, "type"], label=""
)

This plot is very similar to the low-dimensional plot above, with the relevant dimensions chosen based on the values of \(α\) via ARD. When you are in doubt about the number of dimensions to project onto, ARD might provide an answer to that question.

Final comments.

p-PCA is a linear map which linearly transforms the data between the original and projected spaces. It can also thought as a matrix factorisation method, in which \(\mathbf{X}=(\mathbf{W} \times \mathbf{Z})^T\). The projection matrix can be understood as a new basis in the projected space, and \(\mathbf{Z}\) are the new coordinates.

Footnotes

Gilbert Strang, Introduction to Linear Algebra, 5th Ed., Wellesley-Cambridge Press, 2016.↩︎
Gareth M. James, Daniela Witten, Trevor Hastie, Robert Tibshirani, An Introduction to Statistical Learning, Springer, 2013.↩︎
Probabilistic PCA by TensorFlow, “https://www.tensorflow.org/probability/examples/Probabilistic_PCA”.↩︎
Probabilistic PCA by TensorFlow, “https://www.tensorflow.org/probability/examples/Probabilistic_PCA”.↩︎
Gareth M. James, Daniela Witten, Trevor Hastie, Robert Tibshirani, An Introduction to Statistical Learning, Springer, 2013.↩︎
David Wipf, Srikantan Nagarajan, A New View of Automatic Relevance Determination, NIPS 2007.↩︎
Christopher Bishop, Pattern Recognition and Machine Learning, Springer, 2006.↩︎