Gaussian Mixture Model (GMM) Usage, Analysis and Visualization¶

In [1]:

import torch
import numpy as np
import matplotlib.pyplot as plt
import os

os.chdir('../')

import tgmm
import importlib
importlib.reload(tgmm)

from tgmm import GaussianMixture, GMMInitializer, dynamic_figsize, plot_gmm, generate_gmm_data

device = 'cuda' if torch.cuda.is_available() else 'cpu'
random_state = 0
np.random.seed(random_state)
torch.manual_seed(random_state)

if device == 'cuda':
    torch.cuda.manual_seed(random_state)
    print('CUDA version:', torch.version.cuda)
    print('Device:', torch.cuda.get_device_name(0))
else:
    print('Using CPU')

import torch import numpy as np import matplotlib.pyplot as plt import os os.chdir('../') import tgmm import importlib importlib.reload(tgmm) from tgmm import GaussianMixture, GMMInitializer, dynamic_figsize, plot_gmm, generate_gmm_data device = 'cuda' if torch.cuda.is_available() else 'cpu' random_state = 0 np.random.seed(random_state) torch.manual_seed(random_state) if device == 'cuda': torch.cuda.manual_seed(random_state) print('CUDA version:', torch.version.cuda) print('Device:', torch.cuda.get_device_name(0)) else: print('Using CPU')

Using CPU

Synthetic Data Generation¶

The synthetic dataset is generated by combining four Gaussian components

In [2]:

n_samples = [800, 200, 1000, 1000]
centers = [np.array([0, 2]),
           np.array([2, -2]),
           np.array([0, 0]),
           np.array([2, 2])]
covs = [
    1.0 * np.eye(2),                    # spherical covariance
    0.5 * np.eye(2),                    # spherical covariance, fewer points
    np.array([[2, 0], [0, 0.5]]),       # diagonal covariance
    np.array([[0.2, 0.5], [0.5, 2]])    # full covariance
]

# Generate synthetic data using the new function
X_tensor, y_tensor = generate_gmm_data(centers, covs, n_samples, device=device, random_state=random_state)

# Convert to numpy for plotting
X = X_tensor.cpu().numpy()
labels = y_tensor.cpu().numpy()
legend_labels = [f'Component {i+1}' for i in range(len(n_samples))]

n_features = X.shape[1]
n_components = len(n_samples)

plot_gmm(X=X, true_labels=labels, title='Original Data', legend_labels=legend_labels)
plt.show()

n_samples = [800, 200, 1000, 1000] centers = [np.array([0, 2]), np.array([2, -2]), np.array([0, 0]), np.array([2, 2])] covs = [ 1.0 * np.eye(2), # spherical covariance 0.5 * np.eye(2), # spherical covariance, fewer points np.array([[2, 0], [0, 0.5]]), # diagonal covariance np.array([[0.2, 0.5], [0.5, 2]]) # full covariance ] # Generate synthetic data using the new function X_tensor, y_tensor = generate_gmm_data(centers, covs, n_samples, device=device, random_state=random_state) # Convert to numpy for plotting X = X_tensor.cpu().numpy() labels = y_tensor.cpu().numpy() legend_labels = [f'Component {i+1}' for i in range(len(n_samples))] n_features = X.shape[1] n_components = len(n_samples) plot_gmm(X=X, true_labels=labels, title='Original Data', legend_labels=legend_labels) plt.show()

Fitting the GMM¶

Here we fit a GMM to the data. We choose the number of clusters n_components to be the same as the number of Gaussians we defined in the previous section. We choose a randomu initialisation, i.e. init_means='random' and fit the GMM five times, so n_init=5 in order to make it less likely to land in local minima. We set verbose=True and verbose_interval=10 to see what is happening during the fitting process.

In [3]:

# Convert to tensor (if needed for further processing)
X_tensor = torch.tensor(X, dtype=torch.float32, device=device)

n_components = len(n_samples)

# Initialize the GMM
gmm = GaussianMixture(
    n_components=n_components,
    init_means='random',
    n_init=10,
    verbose=True,
    verbose_interval=100,
)

# Fit the GMM
gmm.fit(X_tensor)

# Get predictions
y_pred = gmm.predict(X_tensor)

# Compute per-sample log-likelihoods
log_probs = gmm.score_samples(X_tensor)

# Get probabilities for each sample for each component
probs = gmm.predict_proba(X_tensor)
probs = probs.detach()

# Generate new samples
n_samples_to_generate = len(X)
gmm_samples, gmm_labels = gmm.sample(n_samples_to_generate)
gmm_samples = gmm_samples.detach().cpu().numpy()
gmm_labels = gmm_labels.detach().cpu().numpy()

# Compute probabilities for each generated sample
generated_probs = gmm.predict_proba(torch.tensor(gmm_samples, dtype=torch.float32).to(device))
generated_probs = generated_probs.detach().cpu().numpy()

print('Mean per-sample log-likelihood:                        ', gmm.score(X_tensor))
print('Mean per-sample log-likelihood (torch.mean(log_probs)):', torch.mean(log_probs).item())
print('Lower bound:                                           ', gmm.lower_bound_)
for i in range(probs.shape[0]):
    assert np.isclose(torch.sum(probs[i]).item(), 1.0), f"Probabilities for sample {i} do not sum to 1"
print('Number of iterations:                                  ', gmm.n_iter_)
print('Converged:                                             ', gmm.converged_)

# Convert to tensor (if needed for further processing) X_tensor = torch.tensor(X, dtype=torch.float32, device=device) n_components = len(n_samples) # Initialize the GMM gmm = GaussianMixture( n_components=n_components, init_means='random', n_init=10, verbose=True, verbose_interval=100, ) # Fit the GMM gmm.fit(X_tensor) # Get predictions y_pred = gmm.predict(X_tensor) # Compute per-sample log-likelihoods log_probs = gmm.score_samples(X_tensor) # Get probabilities for each sample for each component probs = gmm.predict_proba(X_tensor) probs = probs.detach() # Generate new samples n_samples_to_generate = len(X) gmm_samples, gmm_labels = gmm.sample(n_samples_to_generate) gmm_samples = gmm_samples.detach().cpu().numpy() gmm_labels = gmm_labels.detach().cpu().numpy() # Compute probabilities for each generated sample generated_probs = gmm.predict_proba(torch.tensor(gmm_samples, dtype=torch.float32).to(device)) generated_probs = generated_probs.detach().cpu().numpy() print('Mean per-sample log-likelihood: ', gmm.score(X_tensor)) print('Mean per-sample log-likelihood (torch.mean(log_probs)):', torch.mean(log_probs).item()) print('Lower bound: ', gmm.lower_bound_) for i in range(probs.shape[0]): assert np.isclose(torch.sum(probs[i]).item(), 1.0), f"Probabilities for sample {i} do not sum to 1" print('Number of iterations: ', gmm.n_iter_) print('Converged: ', gmm.converged_)

[Run 1] Iteration 0: log-likelihood=-3.83872

[Run 1] Iteration 34: log-likelihood=-3.15956, Converged!
[Run 2] Iteration 0: log-likelihood=-3.82558
[Run 2] Iteration 50: log-likelihood=-3.15923, Converged!
[Run 3] Iteration 0: log-likelihood=-3.76025
[Run 3] Iteration 67: log-likelihood=-3.15923, Converged!
[Run 4] Iteration 0: log-likelihood=-3.76986

[Run 4] Iteration 35: log-likelihood=-3.15955, Converged!
[Run 5] Iteration 0: log-likelihood=-3.77588

[Run 5] Iteration 47: log-likelihood=-3.15936, Converged!
[Run 6] Iteration 0: log-likelihood=-3.77075
[Run 6] Iteration 41: log-likelihood=-3.57570, Converged!
[Run 7] Iteration 0: log-likelihood=-3.80822
[Run 7] Iteration 34: log-likelihood=-3.15929, Converged!
[Run 8] Iteration 0: log-likelihood=-3.75723
[Run 8] Iteration 9: log-likelihood=-3.22951, Converged!
[Run 9] Iteration 0: log-likelihood=-3.56511
[Run 9] Iteration 20: log-likelihood=-3.15973, Converged!
[Run 10] Iteration 0: log-likelihood=-3.77652

[Run 10] Iteration 80: log-likelihood=-3.15922, Converged!

Mean per-sample log-likelihood:                         -3.15913987159729
Mean per-sample log-likelihood (torch.mean(log_probs)): -3.15913987159729
Lower bound:                                            -3.15913987159729
Number of iterations:                                   80
Converged:                                              True

Methods of the GMM¶

The GaussianMixture provides the following methods:

• Predicted Cluster Labels: Using the predict method.
• Per-Sample Log-Likelihoods: Using the score_samples method.
• Posterior Probabilities: For each data sample across all components.
• New Samples: Generated from the fitted GMM, along with the corresponding component indices.

Diagnostic information such as mean log-likelihood, convergence status, and number of iterations is printed to verify the model's performance.

Inspecting the Learned GMM Parameters¶

After fitting the model, we inspect the parameters learned by the GMM. For each component, the following are printed:

• Weight: The mixing proportion for the component.
• Mean: The estimated mean of the component.
• Covariance: The covariance matrix (or vector/scalar for other types) representing the spread of the component.

In [4]:

# weights, means, covariances of each component
for i in range(gmm.n_components):
    print(f'\nComponent {i+1}:')
    print('Weight:')
    print(gmm.weights_[i].item())
    print('Mean:')
    print(gmm.means_[i].detach().cpu().numpy())
    print('Covariance:')
    print(gmm.covariances_[i].detach().cpu().numpy())

plot_gmm(X=X, gmm=gmm, title='GMM', legend_labels=legend_labels)

# weights, means, covariances of each component for i in range(gmm.n_components): print(f'\nComponent {i+1}:') print('Weight:') print(gmm.weights_[i].item()) print('Mean:') print(gmm.means_[i].detach().cpu().numpy()) print('Covariance:') print(gmm.covariances_[i].detach().cpu().numpy()) plot_gmm(X=X, gmm=gmm, title='GMM', legend_labels=legend_labels)

Component 1:
Weight:
0.32896649837493896
Mean:
[-0.15161191 -0.00941133]
Covariance:
[[ 3.969354   -0.01766335]
 [-0.01766335  0.24021453]]

Component 2:
Weight:
0.33582615852355957
Mean:
[2.000501  2.0055819]
Covariance:
[[0.28738084 1.0938029 ]
 [1.0938029  4.242133  ]]

Component 3:
Weight:
0.2692916989326477
Mean:
[-0.01153989  1.971482  ]
Covariance:
[[ 0.9504869  -0.02864964]
 [-0.02864964  0.9700867 ]]

Component 4:
Weight:
0.06591557711362839
Mean:
[ 1.997651  -2.0003922]
Covariance:
[[ 0.20961574 -0.00229622]
 [-0.00229622  0.26018155]]

Out[4]:

<Axes: title={'center': 'GMM'}, xlabel='Feature 1', ylabel='Feature 2'>

Visualization of Clustering and Generated Samples¶

In this section, we visualize several key results using our custom plot_gmm function:

1. Per-Sample Log-Likelihood:
   A continuous scatter plot where each point is colored by its log-likelihood value. A colorbar is included to indicate the log-likelihood scale.

2. GMM Clustering Results:
   A cluster plot where points are colored by their predicted labels. The plot also overlays the final GMM component means and ellipses representing covariance contours.

3. Generated Samples:
   The new samples generated by the GMM are plotted along with the model's estimated parameters to evaluate how well the model reproduces the data distribution.

In [5]:

plot_gmm(X=X, title='Per-Sample Log-Likelihood', log_probs=log_probs, colorbar_label='Log-Likelihood')
plt.show()

plot_gmm(X=X, true_labels=y_pred, title='GMM Results', legend_labels=[f'Cluster {i}' for i in range(n_components)])
plt.show()

plot_gmm(X=gmm_samples, gmm=gmm, title='GMM Results (Generated Samples)')
plt.show()

plot_gmm(X=X, title='Per-Sample Log-Likelihood', log_probs=log_probs, colorbar_label='Log-Likelihood') plt.show() plot_gmm(X=X, true_labels=y_pred, title='GMM Results', legend_labels=[f'Cluster {i}' for i in range(n_components)]) plt.show() plot_gmm(X=gmm_samples, gmm=gmm, title='GMM Results (Generated Samples)') plt.show()

Visualizing Component Probability Distributions¶

Finally, we analyze the distribution of component responsibilities (i.e., the probability that each sample belongs to a particular component):

• Observed Data:
  A grid of subplots is generated where each subplot shows the probability distribution for a single GMM component across the observed data. Each plot uses a continuous color scale (e.g., 'RdBu') with an associated colorbar.

• Generated Samples:
  Similarly, we plot the component probability distributions for the generated samples, using a different colormap (e.g., 'PRGn').

These plots provide insight into how the GMM assigns probabilities to each data point and how well each component is represented in both the original and generated data.

In [6]:

nrows, ncols = 2, 2
figsize = dynamic_figsize(nrows, ncols)

# Observed Data
fig, axs = plt.subplots(nrows, ncols, figsize=figsize)

for k, ax in enumerate(axs.ravel()):
    # Assume probs is an array of shape (N, n_components)
    prob_k = probs[:, k]
    plot_gmm(X=X, ax=ax, title=f'Component {k+1} Probability', log_probs=prob_k, colormap='RdBu', colorbar_label='Probability')
plt.suptitle("Observed Data: Probability Distributions Across GMM Components")
plt.tight_layout()
plt.show()

# Generated Samples
fig, axs = plt.subplots(nrows, ncols, figsize=figsize)

for k, ax in enumerate(axs.ravel()):
    prob_k = generated_probs[:, k]
    plot_gmm(X=gmm_samples, gmm=gmm, ax=ax, show_ellipses=False, title=f'Component {k+1} Probability', log_probs=prob_k, colormap='PRGn', colorbar_label='Probability')
fig.suptitle("Generated Samples: Probability Distributions Across GMM Components")
plt.tight_layout()
plt.show()

nrows, ncols = 2, 2 figsize = dynamic_figsize(nrows, ncols) # Observed Data fig, axs = plt.subplots(nrows, ncols, figsize=figsize) for k, ax in enumerate(axs.ravel()): # Assume probs is an array of shape (N, n_components) prob_k = probs[:, k] plot_gmm(X=X, ax=ax, title=f'Component {k+1} Probability', log_probs=prob_k, colormap='RdBu', colorbar_label='Probability') plt.suptitle("Observed Data: Probability Distributions Across GMM Components") plt.tight_layout() plt.show() # Generated Samples fig, axs = plt.subplots(nrows, ncols, figsize=figsize) for k, ax in enumerate(axs.ravel()): prob_k = generated_probs[:, k] plot_gmm(X=gmm_samples, gmm=gmm, ax=ax, show_ellipses=False, title=f'Component {k+1} Probability', log_probs=prob_k, colormap='PRGn', colorbar_label='Probability') fig.suptitle("Generated Samples: Probability Distributions Across GMM Components") plt.tight_layout() plt.show()

Comparing Different Covariance Types¶

In this section, we examine how the choice of covariance type affects the clustering results of the GMM. We consider six different covariance types:

• Tied Spherical Covariance:
  All components share the same covariance matrix, which is spherical: $$ \Sigma_1 = \Sigma_2 = \cdots = \Sigma_K = \sigma^2 I_d = \sigma^2 \begin{bmatrix} 1 & 0 & \cdots & 0 \\ 0 & 1 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & 1 \end{bmatrix}. $$

• Spherical Covariance:
  Each component $ k $ has its own spherical covariance matrix: $$ \Sigma_k = \sigma_k^2 I_d = \sigma_k^2 \begin{bmatrix} 1 & 0 & \cdots & 0 \\ 0 & 1 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & 1 \end{bmatrix}, \quad k = 1,\dots,K. $$

• Tied Diagonal Covariance:
  All components share the same diagonal covariance matrix: $$ \Sigma_1 = \Sigma_2 = \cdots = \Sigma_K = \begin{bmatrix} \sigma_1^2 & 0 & \cdots & 0 \\ 0 & \sigma_2^2 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \sigma_d^2 \end{bmatrix}. $$

• Diagonal Covariance:
  Each component $ k $ has its own diagonal covariance matrix: $$ \Sigma_k = \begin{bmatrix} \sigma_{k1}^2 & 0 & \cdots & 0 \\ 0 & \sigma_{k2}^2 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \sigma_{kd}^2 \end{bmatrix}, \quad k = 1,\dots,K. $$

• Tied Full Covariance:
  All components share the same full covariance matrix: $$ \Sigma_1 = \Sigma_2 = \cdots = \Sigma_K = \Sigma = \begin{bmatrix} \sigma_{11} & \sigma_{12} & \cdots & \sigma_{1d} \\ \sigma_{12} & \sigma_{22} & \cdots & \sigma_{2d} \\ \vdots & \vdots & \ddots & \vdots \\ \sigma_{1d} & \sigma_{2d} & \cdots & \sigma_{dd} \end{bmatrix}. $$

• Full Covariance:
  Each component $ k $ has its own full covariance matrix: $$ \Sigma_k = \begin{bmatrix} \sigma_{k11} & \sigma_{k12} & \cdots & \sigma_{k1d} \\ \sigma_{k12} & \sigma_{k22} & \cdots & \sigma_{k2d} \\ \vdots & \vdots & \ddots & \vdots \\ \sigma_{k1d} & \sigma_{k2d} & \cdots & \sigma_{kdd} \end{bmatrix}, \quad k = 1,\dots,K. $$

For each covariance type, we:

1. Initialize a GMM with 4 components using k-means for mean initialization.
2. Fit the model to the dataset (converted to a PyTorch tensor, X_tensor).
3. Predict the cluster labels.
4. Visualize the clustering results with our custom plot_gmm function in a grid of subplots.

In [7]:

cov_types = ['tied_spherical', 'spherical', 'tied_diag', 'diag', 'tied_full', 'full']

torch.manual_seed(random_state)

nrows, ncols = 3, 2
figsize = dynamic_figsize(nrows, ncols)

fig, axs = plt.subplots(nrows, ncols, figsize=figsize)
axs = axs.flatten()

for ax, cov_type in zip(axs, cov_types):
    gmm = GaussianMixture(
        n_features=2,
        n_components=4,
        covariance_type=cov_type,
        init_means='kmeans',
        random_state=random_state,
        device=device,
        cem=False,
    )
    gmm.fit(X_tensor)
    labels_pred = gmm.predict(X_tensor).cpu().numpy()
    plot_gmm(X, gmm=gmm, ax=ax, title=f'Covariance: {cov_type}', legend_labels=legend_labels)

plt.suptitle("GMM Clustering with Different Covariance Types")
plt.tight_layout(rect=[0, 0, 1, 0.99])  # leave space for the suptitle
plt.show()

cov_types = ['tied_spherical', 'spherical', 'tied_diag', 'diag', 'tied_full', 'full'] torch.manual_seed(random_state) nrows, ncols = 3, 2 figsize = dynamic_figsize(nrows, ncols) fig, axs = plt.subplots(nrows, ncols, figsize=figsize) axs = axs.flatten() for ax, cov_type in zip(axs, cov_types): gmm = GaussianMixture( n_features=2, n_components=4, covariance_type=cov_type, init_means='kmeans', random_state=random_state, device=device, cem=False, ) gmm.fit(X_tensor) labels_pred = gmm.predict(X_tensor).cpu().numpy() plot_gmm(X, gmm=gmm, ax=ax, title=f'Covariance: {cov_type}', legend_labels=legend_labels) plt.suptitle("GMM Clustering with Different Covariance Types") plt.tight_layout(rect=[0, 0, 1, 0.99]) # leave space for the suptitle plt.show()

Comparing Different Initialization Methods¶

In this section, we investigate how different initialization methods affect the GMM clustering results when using a fixed covariance type (here, full). The initialization methods compared are:

• random
• points
• kpp (k-means++)
• kmeans
• maxdist

Below, we provide a brief mathematical description of each method:

1. Random Initialization:
   This method computes the empirical mean $\bar{x} $ and covariance $\Sigma$ of the dataset and then draws each initial center from a Gaussian distribution: $$ \mu_i = \bar{x} + L z, \quad z \sim \mathcal{N}(0, I_d), $$ where $ L $ is the Cholesky factor of $\Sigma$ (i.e. $\Sigma = LL^\top$).
   This yields centers that roughly follow the global data distribution.

2. Points Initialization:
   This method simply selects $ k $ points uniformly at random from the dataset. Mathematically, if $ S \subset \{1,2,\dots,N\} $ is a random subset of $ k $ indices, then: $$ \mu_i = x_{s_i}, \quad i = 1, \dots, k. $$

3. k-means++ Initialization (kpp):
   The first center is chosen uniformly at random. For each subsequent center, the probability of selecting data point $ x_j $ is proportional to the square of its distance to the nearest already-chosen center: $$ P(x_j) = \frac{D(x_j)}{\sum_{j=1}^{N} D(x_j)}, \quad \text{with } D(x_j) = \min_{l=1,\dots,i-1} \|x_j - \mu_l\|^2. $$ This procedure tends to choose centers that are spread out.

4. k-means Initialization:
   Starting from the k-means++ initialization, the standard k-means algorithm is run iteratively:

   • Assignment step:
     Assign each point $ x_j $ to the nearest center: $$ c_j = \arg\min_{i} \|x_j - \mu_i\|^2. $$
   • Update step:
     Recompute each center as the mean of its assigned points: $$ \mu_i = \frac{1}{|C_i|} \sum_{x_j \in C_i} x_j. $$ These steps repeat until convergence (i.e. the centers do not change significantly).

5. Max-Distance Initialization (maxdist):
   This is a modified k-means++ method:

   • The first center is chosen randomly.
   • For $ i = 2,\dots,k $, choose the next center as: $$ \mu_i = \arg\max_{x \in \mathcal{D}} \min_{l=1,\dots,i-1} \|x - \mu_l\|, $$ ensuring that the new center is the point with the largest minimum distance to the existing centers.
   • Finally, the first center is reselected as the point with the largest distance from the centers $ \mu_2,\dots,\mu_k $: $$ \mu_1 = \arg\max_{x \in \mathcal{D}} \min_{l=2,\dots,k} \|x - \mu_l\|. $$

For each initialization method, we obtain a set of initial means which are then passed (via the means_init parameter) to the GMM. This bypasses the model’s internal initialization logic, allowing us to compare the impact of different starting points on the final clustering.

1. Use a helper function (get_init_means) to obtain the initial means from the corresponding method.
2. Initialize a GMM with these precomputed initial means by providing them via the means_init parameter (thus bypassing the internal initialization logic).
3. Fit the GMM to the dataset (X_tensor).
4. Predict cluster labels and visualize the clustering results with our plot_gmm function, which also overlays the initial means (marked in red).

In [8]:

cov_types = ['full']
init_methods = ['random', 'points', 'kpp', 'kmeans', 'maxdist']

torch.manual_seed(random_state)

def get_init_means(method, data, k):
    """
    Returns initial means for a given method from GMMInitializer.
    """
    if method == 'random':
        return GMMInitializer.random(data, k)
    elif method == 'points':
        return GMMInitializer.points(data, k)
    elif method == 'kpp':
        return GMMInitializer.kpp(data, k)
    elif method == 'kmeans':
        return GMMInitializer.kmeans(data, k)
    elif method == 'maxdist':
        return GMMInitializer.maxdist(data, k)
    else:
        raise ValueError(f"Unknown init method: {method}")

nrows, ncols = len(init_methods), len(cov_types)
figsize = dynamic_figsize(nrows, ncols)

fig, axs = plt.subplots(nrows, ncols, figsize=figsize)

for col_idx, method in enumerate(init_methods):
    ax = axs[col_idx]

    # 1) Get initial means from our GMMInitializer
    init_means = get_init_means(method, X_tensor, k=4)

    # 2) Create GMM with 'init_means' so it uses these means
    gmm = GaussianMixture(
        n_features=2,
        n_components=4,
        covariance_type=cov_types[0],
        init_means=init_means,
        random_state=0,
        device=device,
        max_iter=1000,
        tol=1e-6,
        reg_covar=1e-10,
    )

    # 3) Fit the GMM on X_tensor
    gmm.fit(X_tensor)
    # 4) Predict cluster labels
    labels_pred = gmm.predict(X_tensor).cpu().numpy()

    # 5) Plot
    title = f'Initialisation: {method}'
    plot_gmm(X, gmm=gmm, ax=ax, title=title, show_initial_means=True, legend_labels=legend_labels)

plt.suptitle("GMM with Different Initialisation Methods")
plt.tight_layout(rect=[0, 0, 1, 0.98])  # leave space for the suptitle
plt.show()

cov_types = ['full'] init_methods = ['random', 'points', 'kpp', 'kmeans', 'maxdist'] torch.manual_seed(random_state) def get_init_means(method, data, k): """ Returns initial means for a given method from GMMInitializer. """ if method == 'random': return GMMInitializer.random(data, k) elif method == 'points': return GMMInitializer.points(data, k) elif method == 'kpp': return GMMInitializer.kpp(data, k) elif method == 'kmeans': return GMMInitializer.kmeans(data, k) elif method == 'maxdist': return GMMInitializer.maxdist(data, k) else: raise ValueError(f"Unknown init method: {method}") nrows, ncols = len(init_methods), len(cov_types) figsize = dynamic_figsize(nrows, ncols) fig, axs = plt.subplots(nrows, ncols, figsize=figsize) for col_idx, method in enumerate(init_methods): ax = axs[col_idx] # 1) Get initial means from our GMMInitializer init_means = get_init_means(method, X_tensor, k=4) # 2) Create GMM with 'init_means' so it uses these means gmm = GaussianMixture( n_features=2, n_components=4, covariance_type=cov_types[0], init_means=init_means, random_state=0, device=device, max_iter=1000, tol=1e-6, reg_covar=1e-10, ) # 3) Fit the GMM on X_tensor gmm.fit(X_tensor) # 4) Predict cluster labels labels_pred = gmm.predict(X_tensor).cpu().numpy() # 5) Plot title = f'Initialisation: {method}' plot_gmm(X, gmm=gmm, ax=ax, title=title, show_initial_means=True, legend_labels=legend_labels) plt.suptitle("GMM with Different Initialisation Methods") plt.tight_layout(rect=[0, 0, 1, 0.98]) # leave space for the suptitle plt.show()

Comparing Different Weight Initialization Methods¶

In this section, we investigate how different weight initialization methods affect the starting mixture weights (before EM optimization). The weight initialization methods compared are:

• uniform: Equal weight for every component
• random: Sampled from a symmetric Dirichlet distribution
• kmeans: Proportional to cluster sizes from a nearest-mean assignment

The initial weights assigned to each component are shown in each panel's title; marker sizes are scaled by weight (scale_size_by_weight=True).

In [9]:

init_methods_weights = ['uniform', 'random', 'kmeans']

torch.manual_seed(random_state)

nrows, ncols = 1, len(init_methods_weights)
figsize = dynamic_figsize(nrows, ncols)
fig, axs = plt.subplots(nrows, ncols, figsize=figsize)

for ax, method in zip(axs, init_methods_weights):
    gmm = GaussianMixture(
        n_features=2,
        n_components=4,
        covariance_type='full',
        init_weights=method,
        random_state=random_state,
        device=device,
        max_iter=1000,
        tol=1e-6,
        reg_covar=1e-10,
    )
    gmm.fit(X_tensor)
    weights_str = ', '.join(f'{w:.2f}' for w in gmm.initial_weights_.cpu().tolist())
    title = f"init_weights='{method}'\ninitial weights: [{weights_str}]"
    plot_gmm(X, gmm=gmm, ax=ax, title=title, legend_labels=legend_labels, scale_size_by_weight=True)

plt.suptitle("GMM with Different Weight Initialization Methods")
plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.show()

init_methods_weights = ['uniform', 'random', 'kmeans'] torch.manual_seed(random_state) nrows, ncols = 1, len(init_methods_weights) figsize = dynamic_figsize(nrows, ncols) fig, axs = plt.subplots(nrows, ncols, figsize=figsize) for ax, method in zip(axs, init_methods_weights): gmm = GaussianMixture( n_features=2, n_components=4, covariance_type='full', init_weights=method, random_state=random_state, device=device, max_iter=1000, tol=1e-6, reg_covar=1e-10, ) gmm.fit(X_tensor) weights_str = ', '.join(f'{w:.2f}' for w in gmm.initial_weights_.cpu().tolist()) title = f"init_weights='{method}'\ninitial weights: [{weights_str}]" plot_gmm(X, gmm=gmm, ax=ax, title=title, legend_labels=legend_labels, scale_size_by_weight=True) plt.suptitle("GMM with Different Weight Initialization Methods") plt.tight_layout(rect=[0, 0, 1, 0.93]) plt.show()

Comparing Different Covariance Initialization Methods¶

In this section, we investigate how different covariance initialization methods affect the GMM's starting ellipses (before EM optimization) and the resulting fit. The covariance initialization methods compared are:

• empirical: Estimated from data assigned to each component by nearest mean (default)
• eye: Identity-like matrices, scaled by 1 + reg_covar
• random: Random positive semi-definite matrices
• global: The single global covariance of the whole dataset, repeated across components

In [10]:

init_methods_cov = ['empirical', 'eye', 'random', 'global']

torch.manual_seed(random_state)

nrows, ncols = 1, len(init_methods_cov)
figsize = dynamic_figsize(nrows, ncols)
fig, axs = plt.subplots(nrows, ncols, figsize=figsize)

for ax, method in zip(axs, init_methods_cov):
    gmm = GaussianMixture(
        n_features=2,
        n_components=4,
        covariance_type='full',
        init_covariances=method,
        random_state=random_state,
        device=device,
        max_iter=1000,
        tol=1e-6,
        reg_covar=1e-10,
    )
    gmm.fit(X_tensor)
    title = f"init_covariances='{method}'"
    plot_gmm(X, gmm=gmm, ax=ax, title=title, legend_labels=legend_labels, show_initial_means=True)

plt.suptitle("GMM with Different Covariance Initialization Methods")
plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.show()

init_methods_cov = ['empirical', 'eye', 'random', 'global'] torch.manual_seed(random_state) nrows, ncols = 1, len(init_methods_cov) figsize = dynamic_figsize(nrows, ncols) fig, axs = plt.subplots(nrows, ncols, figsize=figsize) for ax, method in zip(axs, init_methods_cov): gmm = GaussianMixture( n_features=2, n_components=4, covariance_type='full', init_covariances=method, random_state=random_state, device=device, max_iter=1000, tol=1e-6, reg_covar=1e-10, ) gmm.fit(X_tensor) title = f"init_covariances='{method}'" plot_gmm(X, gmm=gmm, ax=ax, title=title, legend_labels=legend_labels, show_initial_means=True) plt.suptitle("GMM with Different Covariance Initialization Methods") plt.tight_layout(rect=[0, 0, 1, 0.93]) plt.show()

Training Controls¶

Beyond covariance type and initialization strategy, a few parameters control the EM optimization itself:

• max_iter, tol: stop after this many iterations, or once the relative log-likelihood improvement drops below tol.
• reg_covar: regularization added to the covariance diagonal, needed to avoid singular/degenerate covariances.
• n_init: run EM from this many independent random initializations and keep the best (highest log-likelihood) one; the winning seed is stored in best_random_state_.
• warm_start: reuse the previous fit's parameters as the starting point for the next call to fit, instead of re-initializing.
• random_state: seed for reproducibility. With n_init > 1, initialization i uses random_state + i.

max_iter, tol, random_state, and warm_start can also be overridden on a single call to fit(...), without changing the model's own defaults.

In [11]:

# n_init keeps the best of several random restarts, and records which seed won
gmm_multi = GaussianMixture(n_components=4, n_init=5, random_state=0, max_iter=200)
gmm_multi.fit(X_tensor)
print(f"Best log-likelihood: {gmm_multi.lower_bound_:.4f}")
print(f"Winning random_state: {gmm_multi.best_random_state_}")

# warm_start: continue optimizing an already-fitted model instead of restarting from scratch
gmm_warm = GaussianMixture(n_components=4, warm_start=True, max_iter=5, random_state=0)
gmm_warm.fit(X_tensor)
print(f"\nAfter 5 iterations: log-likelihood={gmm_warm.lower_bound_:.4f}, converged={gmm_warm.converged_}")
gmm_warm.fit(X_tensor)  # picks up where it left off instead of re-initializing
print(f"After 5 more (warm-started) iterations: log-likelihood={gmm_warm.lower_bound_:.4f}, converged={gmm_warm.converged_}")

# fit() accepts overrides for max_iter/tol/random_state/warm_start without touching the model's own defaults
gmm_override = GaussianMixture(n_components=4, max_iter=1000, random_state=0)
gmm_override.fit(X_tensor, max_iter=10, tol=1e-2)
print(f"\nOverridden fit (max_iter=10, tol=1e-2): n_iter_={gmm_override.n_iter_}, converged={gmm_override.converged_}")

# n_init keeps the best of several random restarts, and records which seed won gmm_multi = GaussianMixture(n_components=4, n_init=5, random_state=0, max_iter=200) gmm_multi.fit(X_tensor) print(f"Best log-likelihood: {gmm_multi.lower_bound_:.4f}") print(f"Winning random_state: {gmm_multi.best_random_state_}") # warm_start: continue optimizing an already-fitted model instead of restarting from scratch gmm_warm = GaussianMixture(n_components=4, warm_start=True, max_iter=5, random_state=0) gmm_warm.fit(X_tensor) print(f"\nAfter 5 iterations: log-likelihood={gmm_warm.lower_bound_:.4f}, converged={gmm_warm.converged_}") gmm_warm.fit(X_tensor) # picks up where it left off instead of re-initializing print(f"After 5 more (warm-started) iterations: log-likelihood={gmm_warm.lower_bound_:.4f}, converged={gmm_warm.converged_}") # fit() accepts overrides for max_iter/tol/random_state/warm_start without touching the model's own defaults gmm_override = GaussianMixture(n_components=4, max_iter=1000, random_state=0) gmm_override.fit(X_tensor, max_iter=10, tol=1e-2) print(f"\nOverridden fit (max_iter=10, tol=1e-2): n_iter_={gmm_override.n_iter_}, converged={gmm_override.converged_}")

/home/asp/Downloads/HeaDS/tgmm/tgmm/gmm.py:814: UserWarning: With n_init=5 and random_state=0, initializations will use random states [0, 1, ..., 4]. The best initialization's random state will be stored in best_random_state_.
  warnings.warn(

Best log-likelihood: -3.1591
Winning random_state: 2

After 5 iterations: log-likelihood=-3.2142, converged=False
After 5 more (warm-started) iterations: log-likelihood=-3.1899, converged=False

Overridden fit (max_iter=10, tol=1e-2): n_iter_=4, converged=True

/home/asp/Downloads/HeaDS/tgmm/tgmm/gmm.py:879: UserWarning: Best result from random_state=2. To reproduce this specific result: use random_state=2 with n_init=1.
  warnings.warn(
/home/asp/Downloads/HeaDS/tgmm/tgmm/gmm.py:975: UserWarning: Run 1: EM did not converge after 5 iterations.
  warnings.warn(
/home/asp/Downloads/HeaDS/tgmm/tgmm/gmm.py:887: UserWarning: EM did not converge. Consider increasing max_iter or adjusting tol.
  warnings.warn(

Custom (Tensor) Initialization¶

init_means, init_weights, and init_covariances each also accept an explicit torch.Tensor instead of a method name, letting you seed the model with domain knowledge -- for example the output of GMMInitializer's static methods (see the GMM Initializer guide), a previous run's fitted parameters, or prior expertise about the data.

In [12]:

custom_means = torch.tensor([[-1.0, 5.0], [0.0, 5.0], [1.0, 5.0], [2.0, 5.0]])
custom_weights = torch.tensor([0.4, 0.2, 0.2, 0.2])
custom_covs = torch.stack([torch.eye(2) for _ in range(4)])

gmm_custom = GaussianMixture(
    n_components=4,
    covariance_type='full',
    init_means=custom_means,
    init_weights=custom_weights,
    init_covariances=custom_covs,
    random_state=0,
)
gmm_custom.fit(X_tensor)
print(f"Converged: {gmm_custom.converged_}, log-likelihood={gmm_custom.lower_bound_:.4f}")
print(f"Initial weights (as provided): {gmm_custom.initial_weights_}")

custom_means = torch.tensor([[-1.0, 5.0], [0.0, 5.0], [1.0, 5.0], [2.0, 5.0]]) custom_weights = torch.tensor([0.4, 0.2, 0.2, 0.2]) custom_covs = torch.stack([torch.eye(2) for _ in range(4)]) gmm_custom = GaussianMixture( n_components=4, covariance_type='full', init_means=custom_means, init_weights=custom_weights, init_covariances=custom_covs, random_state=0, ) gmm_custom.fit(X_tensor) print(f"Converged: {gmm_custom.converged_}, log-likelihood={gmm_custom.lower_bound_:.4f}") print(f"Initial weights (as provided): {gmm_custom.initial_weights_}")

Converged: True, log-likelihood=-3.1592
Initial weights (as provided): tensor([0.4000, 0.2000, 0.2000, 0.2000])

Classification EM (CEM)¶

Setting cem=True swaps the soft E-step for a hard classification step between the E-step and M-step: each point is assigned entirely to its most likely component instead of a probability distribution over components. This can converge faster, at the cost of being more sensitive to initialization. See the CEM tutorial for a deeper comparison against standard EM.

In [13]:

gmm_em = GaussianMixture(n_components=4, covariance_type='full', random_state=0, max_iter=200)
gmm_em.fit(X_tensor)

gmm_cem = GaussianMixture(n_components=4, covariance_type='full', cem=True, random_state=0, max_iter=200)
gmm_cem.fit(X_tensor)

fig, axes = plt.subplots(1, 2, figsize=dynamic_figsize(1, 2))
plot_gmm(X, gmm=gmm_em, ax=axes[0], title=f'Standard EM ({gmm_em.n_iter_} iters, LL={gmm_em.lower_bound_:.3f})', legend_labels=legend_labels)
plot_gmm(X, gmm=gmm_cem, ax=axes[1], title=f'CEM ({gmm_cem.n_iter_} iters, LL={gmm_cem.lower_bound_:.3f})', legend_labels=legend_labels)
plt.tight_layout()
plt.show()

gmm_em = GaussianMixture(n_components=4, covariance_type='full', random_state=0, max_iter=200) gmm_em.fit(X_tensor) gmm_cem = GaussianMixture(n_components=4, covariance_type='full', cem=True, random_state=0, max_iter=200) gmm_cem.fit(X_tensor) fig, axes = plt.subplots(1, 2, figsize=dynamic_figsize(1, 2)) plot_gmm(X, gmm=gmm_em, ax=axes[0], title=f'Standard EM ({gmm_em.n_iter_} iters, LL={gmm_em.lower_bound_:.3f})', legend_labels=legend_labels) plot_gmm(X, gmm=gmm_cem, ax=axes[1], title=f'CEM ({gmm_cem.n_iter_} iters, LL={gmm_cem.lower_bound_:.3f})', legend_labels=legend_labels) plt.tight_layout() plt.show()

MAP Estimation with Priors¶

Each of weight_concentration_prior (Dirichlet), mean_prior + mean_precision_prior (Gaussian), and covariance_prior + degrees_of_freedom_prior (Inverse-Wishart) turns the corresponding M-step update from a maximum-likelihood estimate into a maximum a posteriori one, regularizing it toward the prior. When both the mean and covariance priors are supplied, tgmm uses their Normal-Inverse-Wishart (NIW) conjugate joint update instead of independent ones. See the Priors tutorial and NIW Priors Comparison for a deeper dive.

In [14]:

gmm_map = GaussianMixture(
    n_components=4,
    n_features=2,
    covariance_type='full',
    weight_concentration_prior=torch.ones(4) * 2.0,
    mean_prior=torch.zeros(4, 2),
    mean_precision_prior=0.1,
    covariance_prior=torch.eye(2).unsqueeze(0).repeat(4, 1, 1),
    degrees_of_freedom_prior=3.0,
    random_state=0,
)
gmm_map.fit(X_tensor)
print(f"Converged: {gmm_map.converged_}, log-likelihood={gmm_map.lower_bound_:.4f}")

fig, ax = plt.subplots()
plot_gmm(X, gmm=gmm_map, ax=ax, title='MAP Estimation (Weight + NIW Mean/Covariance Priors)', legend_labels=legend_labels)
plt.show()

gmm_map = GaussianMixture( n_components=4, n_features=2, covariance_type='full', weight_concentration_prior=torch.ones(4) * 2.0, mean_prior=torch.zeros(4, 2), mean_precision_prior=0.1, covariance_prior=torch.eye(2).unsqueeze(0).repeat(4, 1, 1), degrees_of_freedom_prior=3.0, random_state=0, ) gmm_map.fit(X_tensor) print(f"Converged: {gmm_map.converged_}, log-likelihood={gmm_map.lower_bound_:.4f}") fig, ax = plt.subplots() plot_gmm(X, gmm=gmm_map, ax=ax, title='MAP Estimation (Weight + NIW Mean/Covariance Priors)', legend_labels=legend_labels) plt.show()

Converged: True, log-likelihood=-3.1633

Constrained Sampling¶

sample() supports rejection-sampling constraints beyond plain unconstrained sampling: sampling only from a specific component, within a std_radius (Mahalanobis distance) or std_range, within a confidence/confidence_range region (converted internally via the $\chi^2$ distribution), or within center_radius of a center_point. See the Sampling tutorial for the full set of options.

In [15]:

# Sample only from component 0
samples_c0, _ = gmm.sample(200, component=0)

# Sample within the 95% confidence ellipse of each component
samples_95, labels_95 = gmm.sample(300, confidence=0.95)

# Sample "outliers": further than 3 standard deviations from their component mean
samples_outliers, labels_outliers = gmm.sample(100, std_range=(3.0, float('inf')))

fig, axes = plt.subplots(1, 3, figsize=dynamic_figsize(1, 3))
axes[0].scatter(samples_c0[:, 0].cpu(), samples_c0[:, 1].cpu(), s=10, alpha=0.6)
axes[0].set_title('Samples from component 0 only')
axes[1].scatter(samples_95[:, 0].cpu(), samples_95[:, 1].cpu(), c=labels_95.cpu(), cmap='turbo', s=10, alpha=0.6)
axes[1].set_title('Samples within 95% confidence region')
axes[2].scatter(samples_outliers[:, 0].cpu(), samples_outliers[:, 1].cpu(), c=labels_outliers.cpu(), cmap='turbo', s=10, alpha=0.6)
axes[2].set_title('Samples beyond 3 std devs ("outliers")')
plt.tight_layout()
plt.show()

# Sample only from component 0 samples_c0, _ = gmm.sample(200, component=0) # Sample within the 95% confidence ellipse of each component samples_95, labels_95 = gmm.sample(300, confidence=0.95) # Sample "outliers": further than 3 standard deviations from their component mean samples_outliers, labels_outliers = gmm.sample(100, std_range=(3.0, float('inf'))) fig, axes = plt.subplots(1, 3, figsize=dynamic_figsize(1, 3)) axes[0].scatter(samples_c0[:, 0].cpu(), samples_c0[:, 1].cpu(), s=10, alpha=0.6) axes[0].set_title('Samples from component 0 only') axes[1].scatter(samples_95[:, 0].cpu(), samples_95[:, 1].cpu(), c=labels_95.cpu(), cmap='turbo', s=10, alpha=0.6) axes[1].set_title('Samples within 95% confidence region') axes[2].scatter(samples_outliers[:, 0].cpu(), samples_outliers[:, 1].cpu(), c=labels_outliers.cpu(), cmap='turbo', s=10, alpha=0.6) axes[2].set_title('Samples beyond 3 std devs ("outliers")') plt.tight_layout() plt.show()

Model Saving and Loading¶

This section demonstrates the save/load functionality for GMM models. The GaussianMixture class provides four methods for model persistence:

1. save(filepath): Save the complete model (parameters + configuration) to a file
2. load(filepath, device=None): Class method to load a complete model from a file
3. save_state_dict(): Return a dictionary containing all model state (similar to PyTorch)
4. load_state_dict(state_dict): Load model state from a dictionary

These methods preserve all model parameters, configuration settings, training state, and prior settings, allowing for complete model reproducibility.

In [16]:

import tempfile
import os

# Use the previously fitted GMM (gmm) from earlier cells
print("Original model info:")
print(f"  Number of components: {gmm.n_components}")
print(f"  Covariance type: {gmm.covariance_type}")
print(f"  Converged: {gmm.converged_}")
print(f"  Number of iterations: {gmm.n_iter_}")
print(f"  Log-likelihood: {gmm.score(X_tensor):.4f}")
print(f"  Component weights: {gmm.weights_.detach().cpu().numpy()}")
print()

# ========================================
# Method 1: Complete model save/load
# ========================================
print("=== Method 1: Complete Model Save/Load ===")

# Save the complete model
with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f:
    model_path = f.name

gmm.save(model_path)
print(f"Model saved to: {model_path}")

# Load the complete model
loaded_gmm = GaussianMixture.load(model_path, device=device)
print("Model loaded successfully!")

# Verify the loaded model
print("Loaded model info:")
print(f"  Number of components: {loaded_gmm.n_components}")
print(f"  Covariance type: {loaded_gmm.covariance_type}")
print(f"  Converged: {loaded_gmm.converged_}")
print(f"  Number of iterations: {loaded_gmm.n_iter_}")
print(f"  Log-likelihood: {loaded_gmm.score(X_tensor):.4f}")
print(f"  Component weights: {loaded_gmm.weights_.detach().cpu().numpy()}")

# Test that predictions are identical
original_predictions = gmm.predict(X_tensor)
loaded_predictions = loaded_gmm.predict(X_tensor)
predictions_match = torch.equal(original_predictions, loaded_predictions)
print(f"  Predictions match: {predictions_match}")

# Clean up
os.unlink(model_path)
print()

# ========================================
# Method 2: State dictionary save/load
# ========================================
print("=== Method 2: State Dictionary Save/Load ===")

# Get state dictionary
state_dict = gmm.save_state_dict()
print(f"State dict contains {len(state_dict)} keys:")
for key in list(state_dict.keys())[:10]:  # Show first 10 keys
    print(f"  {key}")
if len(state_dict) > 10:
    print(f"  ... and {len(state_dict) - 10} more")
print()

# Create a new model and load the state
new_gmm = GaussianMixture(
    n_components=4,  # Must match the saved model
    n_features=2,    # Must match the saved model
    device=device
)

new_gmm.load_state_dict(state_dict)
print("State dict loaded into new model!")

# Verify the new model
print("New model info:")
print(f"  Number of components: {new_gmm.n_components}")
print(f"  Covariance type: {new_gmm.covariance_type}")
print(f"  Converged: {new_gmm.converged_}")
print(f"  Number of iterations: {new_gmm.n_iter_}")
print(f"  Log-likelihood: {new_gmm.score(X_tensor):.4f}")
print(f"  Component weights: {new_gmm.weights_.detach().cpu().numpy()}")

# Test predictions
new_predictions = new_gmm.predict(X_tensor)
predictions_match = torch.equal(original_predictions, new_predictions)
print(f"  Predictions match original: {predictions_match}")
print()

# ========================================
# Demonstration: Model persistence workflow
# ========================================
print("=== Practical Example: Training and Persistence Workflow ===")

# Train a new model on subset of data
subset_size = 1000
X_subset = X_tensor[:subset_size]

quick_gmm = GaussianMixture(
    n_components=3,
    covariance_type='diag',
    max_iter=100,
    verbose=False,
    device=device
)
quick_gmm.fit(X_subset)

print(f"Quick model trained:")
print(f"  Iterations: {quick_gmm.n_iter_}")
print(f"  Log-likelihood: {quick_gmm.score(X_subset):.4f}")

# Save using temporary file for demonstration
with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f:
    quick_model_path = f.name

quick_gmm.save(quick_model_path)

# Simulate loading in a different session
del quick_gmm  # Simulate model going out of scope

# Load and continue working
restored_gmm = GaussianMixture.load(quick_model_path, device=device)
print(f"Model restored from file:")
print(f"  Iterations: {restored_gmm.n_iter_}")
print(f"  Log-likelihood on subset: {restored_gmm.score(X_subset):.4f}")
print(f"  Log-likelihood on full data: {restored_gmm.score(X_tensor):.4f}")

# Generate samples from restored model
restored_samples, _ = restored_gmm.sample(200)
print(f"Generated {len(restored_samples)} samples from restored model")

# Clean up
os.unlink(quick_model_path)

print("\n=== Summary ===")
print("✓ Complete model save/load preserves all parameters and configuration")
print("✓ State dictionary methods provide PyTorch-style flexibility")  
print("✓ All training state (convergence, iterations, likelihood) is preserved")
print("✓ Models can be saved, loaded, and used for inference seamlessly")

import tempfile import os # Use the previously fitted GMM (gmm) from earlier cells print("Original model info:") print(f" Number of components: {gmm.n_components}") print(f" Covariance type: {gmm.covariance_type}") print(f" Converged: {gmm.converged_}") print(f" Number of iterations: {gmm.n_iter_}") print(f" Log-likelihood: {gmm.score(X_tensor):.4f}") print(f" Component weights: {gmm.weights_.detach().cpu().numpy()}") print() # ======================================== # Method 1: Complete model save/load # ======================================== print("=== Method 1: Complete Model Save/Load ===") # Save the complete model with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f: model_path = f.name gmm.save(model_path) print(f"Model saved to: {model_path}") # Load the complete model loaded_gmm = GaussianMixture.load(model_path, device=device) print("Model loaded successfully!") # Verify the loaded model print("Loaded model info:") print(f" Number of components: {loaded_gmm.n_components}") print(f" Covariance type: {loaded_gmm.covariance_type}") print(f" Converged: {loaded_gmm.converged_}") print(f" Number of iterations: {loaded_gmm.n_iter_}") print(f" Log-likelihood: {loaded_gmm.score(X_tensor):.4f}") print(f" Component weights: {loaded_gmm.weights_.detach().cpu().numpy()}") # Test that predictions are identical original_predictions = gmm.predict(X_tensor) loaded_predictions = loaded_gmm.predict(X_tensor) predictions_match = torch.equal(original_predictions, loaded_predictions) print(f" Predictions match: {predictions_match}") # Clean up os.unlink(model_path) print() # ======================================== # Method 2: State dictionary save/load # ======================================== print("=== Method 2: State Dictionary Save/Load ===") # Get state dictionary state_dict = gmm.save_state_dict() print(f"State dict contains {len(state_dict)} keys:") for key in list(state_dict.keys())[:10]: # Show first 10 keys print(f" {key}") if len(state_dict) > 10: print(f" ... and {len(state_dict) - 10} more") print() # Create a new model and load the state new_gmm = GaussianMixture( n_components=4, # Must match the saved model n_features=2, # Must match the saved model device=device ) new_gmm.load_state_dict(state_dict) print("State dict loaded into new model!") # Verify the new model print("New model info:") print(f" Number of components: {new_gmm.n_components}") print(f" Covariance type: {new_gmm.covariance_type}") print(f" Converged: {new_gmm.converged_}") print(f" Number of iterations: {new_gmm.n_iter_}") print(f" Log-likelihood: {new_gmm.score(X_tensor):.4f}") print(f" Component weights: {new_gmm.weights_.detach().cpu().numpy()}") # Test predictions new_predictions = new_gmm.predict(X_tensor) predictions_match = torch.equal(original_predictions, new_predictions) print(f" Predictions match original: {predictions_match}") print() # ======================================== # Demonstration: Model persistence workflow # ======================================== print("=== Practical Example: Training and Persistence Workflow ===") # Train a new model on subset of data subset_size = 1000 X_subset = X_tensor[:subset_size] quick_gmm = GaussianMixture( n_components=3, covariance_type='diag', max_iter=100, verbose=False, device=device ) quick_gmm.fit(X_subset) print(f"Quick model trained:") print(f" Iterations: {quick_gmm.n_iter_}") print(f" Log-likelihood: {quick_gmm.score(X_subset):.4f}") # Save using temporary file for demonstration with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f: quick_model_path = f.name quick_gmm.save(quick_model_path) # Simulate loading in a different session del quick_gmm # Simulate model going out of scope # Load and continue working restored_gmm = GaussianMixture.load(quick_model_path, device=device) print(f"Model restored from file:") print(f" Iterations: {restored_gmm.n_iter_}") print(f" Log-likelihood on subset: {restored_gmm.score(X_subset):.4f}") print(f" Log-likelihood on full data: {restored_gmm.score(X_tensor):.4f}") # Generate samples from restored model restored_samples, _ = restored_gmm.sample(200) print(f"Generated {len(restored_samples)} samples from restored model") # Clean up os.unlink(quick_model_path) print("\n=== Summary ===") print("✓ Complete model save/load preserves all parameters and configuration") print("✓ State dictionary methods provide PyTorch-style flexibility") print("✓ All training state (convergence, iterations, likelihood) is preserved") print("✓ Models can be saved, loaded, and used for inference seamlessly")

Original model info:
  Number of components: 4
  Covariance type: full
  Converged: True
  Number of iterations: 38
  Log-likelihood: -3.1591
  Component weights: [0.27020916 0.06466867 0.32912436 0.33599782]

=== Method 1: Complete Model Save/Load ===
Model saved to: /tmp/tmpx2flpj1z.pth
Model loaded successfully!
Loaded model info:
  Number of components: 4
  Covariance type: full
  Converged: True
  Number of iterations: 38
  Log-likelihood: -3.1591
  Component weights: [0.27020916 0.06466867 0.32912436 0.33599782]
  Predictions match: True

=== Method 2: State Dictionary Save/Load ===
State dict contains 33 keys:
  weights_
  means_
  covariances_
  initial_weights_
  initial_means_
  initial_covariances_
  n_components
  n_features
  covariance_type
  tol
  ... and 23 more

State dict loaded into new model!
New model info:
  Number of components: 4
  Covariance type: full
  Converged: True
  Number of iterations: 38
  Log-likelihood: -3.1591
  Component weights: [0.27020916 0.06466867 0.32912436 0.33599782]
  Predictions match original: True

=== Practical Example: Training and Persistence Workflow ===
Quick model trained:
  Iterations: 7
  Log-likelihood: -3.0156
Model restored from file:
  Iterations: 7
  Log-likelihood on subset: -3.0156
  Log-likelihood on full data: -5.1768
Generated 200 samples from restored model

=== Summary ===
✓ Complete model save/load preserves all parameters and configuration
✓ State dictionary methods provide PyTorch-style flexibility
✓ All training state (convergence, iterations, likelihood) is preserved
✓ Models can be saved, loaded, and used for inference seamlessly