Simulating and Fitting a Gaussian Mixture Model

This tutorial demonstrates how to use StateSpaceDynamics.jl to create a Probabilistic PCA model and fit it using the EM algorithm.

Load Packages

using StateSpaceDynamics
using LinearAlgebra
using Random
using Plots
using StatsPlots
using StableRNGs
using Distributions

rng = StableRNG(1234);

Create a State-Space Model

D = 2
k = 2

2

define true parameters for the model to sample from

W_true = [
   -1.64   0.2;
    0.9  -2.8
]

σ²_true = 0.5
μ_true = [1.65, -1.3]

ppca = ProbabilisticPCA(W_true, σ²_true, μ_true)

ProbabilisticPCA{Float64, Matrix{Float64}, Vector{Float64}}([-1.64 0.2; 0.9 -2.8], 0.5, [1.65, -1.3], 2, 2, Matrix{Float64}(undef, 2, 0))

Sample data from the model

num_obs = 500
X, z = rand(rng, ppca, num_obs)

([-0.16292996949597693 4.4667986169103875 … -1.664030995267812 2.256120297611268; -2.8032655940478572 0.8620051665088987 … 3.015699033993531 1.8530191339070186], [0.5194945485037596 -1.6756611579095566 … 1.2133013606558642 -0.42285308575916947; 0.9051397268611271 -1.2947334444127405 … -1.1944870517261137 -1.4742335297619584])

Plot sampled data from the model

x1 = X[1, :]
x2 = X[2, :]
labels = map(i -> (abs(z[1,i]) > abs(z[2,i]) ? 1 : 2), 1:size(z,2))

p = plot()

scatter!(
    p, x1, x2;
    group      = labels,
    xlabel     = "X₁",
    ylabel     = "X₂",
    title      = "Samples grouped by dominant latent factor",
    label= ["Latent 1" "Latent 2"],
    legend     = :topright,
    markersize = 5,
)

p

Paramter recovery: Initialize a new model with default parameters and fit to the data using EM.

#define default parameters
D = 2
k = 2
W = randn(rng, D, k)
σ² = 0.5
μ_vector = randn(rng, 2)

fit_ppca = ProbabilisticPCA(W, σ², μ_vector)

ProbabilisticPCA{Float64, Matrix{Float64}, Vector{Float64}}([-0.6286783412298981 1.6477984375260266; -1.9544959850153094 2.20883182977146], 0.5, [-1.2857779269609082, 0.6481990189960374], 2, 2, Matrix{Float64}(undef, 2, 0))

Fit model using EM Algorithm

lls = fit!(fit_ppca, X)

26-element Vector{Float64}:
 -2942.9846914566488
 -2649.021975079838
 -2596.3727771298704
 -2578.2703501279584
 -2571.2629060042264
 -2568.437824572584
 -2567.265115905659
 -2566.7600331093895
 -2566.5343236584004
 -2566.430444632034
     ⋮
 -2566.337136047855
 -2566.3369917179375
 -2566.3369215564903
 -2566.336887441408
 -2566.336870850639
 -2566.3368627813325
 -2566.3368588563276
 -2566.3368569470554
 -2566.336856018277

Confirm model convergence using log likelihoods

ll_plot = plot(
    lls;
    xlabel="Iteration",
    ylabel="Log-Likelihood",
    title="EM Convergence (PPCA)",
    marker=:circle,
    label="log_likelihood",
    reuse=false,
)

ll_plot

Plot the learned lower dimension latent space over the Data

x1, x2 = X[1, :], X[2, :]
μ1, μ2  = fit_ppca.μ
W_fit = fit_ppca.W
w1      = W_fit[:, 1]
w2      = W_fit[:, 2]

P = plot()

scatter!(
    P, x1, x2;
    xlabel     = "X₁",
    ylabel     = "X₂",
    title      = "Data with PPCA loading directions",
    label      = "Data",
    alpha      = 0.5,
    markersize = 4,
)

k = size(W_fit, 2)
bases_x = repeat([μ1], k)
bases_y = repeat([μ2], k)

2-element Vector{Float64}:
 0.6481990189960374
 0.6481990189960374

Add the component arrows in both directions

quiver!(
  P, [μ1], [μ2];
  quiver      = ([ w1[1]], [ w1[2]]),
  arrow       = :arrow, lw=3,
  color       = :red,
  label       = "W₁"
)
quiver!(
  P, [μ1], [μ2];
  quiver      = ([-w1[1]], [-w1[2]]),
  arrow       = :arrow, lw=3,
  color       = :red,
  label       = ""
)

quiver!(
  P, [μ1], [μ2];
  quiver      = ([ w2[1]], [ w2[2]]),
  arrow       = :arrow, lw=3,
  color       = :green,
  label       = "W₂"
)
quiver!(
  P, [μ1], [μ2];
  quiver      = ([-w2[1]], [-w2[2]]),
  arrow       = :arrow, lw=3,
  color       = :green,
  label       = ""
)

P

This page was generated using Literate.jl.