gaussian_mixture_models_3D.py

# =================================================================================================================================== #
# ----------------------------------------------------------- DESCRIPTION ----------------------------------------------------------- #
# File containing the necessary class and functions to execute, view and evaluate Gaussian Mixture Models.                            #
# Filepath: Bishop/Chapter-9/gaussian_mixture_models.py                                                                               #
# =================================================================================================================================== #


# =================================================================================================================================== #
# --------------------------------------------------------- EXTERNAL IMPORTS -------------------------------------------------------- #
import numpy as np  # Scientific computing and data manipulation.                                                                     #
import matplotlib.pyplot as plt  # Data visualization.                                                                                # 
from sklearn.datasets import make_blobs  # Generate isotropic Gaussian blobs for clustering.                                          #
import pandas as pd # Data manipulation and analysis.                                                                                 #
from matplotlib.colors import ListedColormap  # Colormap for plotting.                                                                #
import scienceplots # Custom plotting styles for matplotlib.                                                                          #
import seaborn as sns # Data visualization library based on matplotlib.                                                               #
import matplotlib.patches as mpatches # For creating custom legends.                                                                  #
# =================================================================================================================================== #

 
# =================================================================================================================================== #
# ---------------------------------------------------------- PLOTTING SETTINGS ------------------------------------------------------ #
plt.style.use(["science", "no-latex"])                                                                                                #
plt.rcParams.update({                                                                                                                 #
    "font.size": 8,               # Base font size for all text in the plot.                                                          #
    "axes.titlesize": 9,          # Title size for axes.                                                                              #
    "axes.labelsize": 9,          # Axis labels size.                                                                                 #
    "xtick.labelsize": 8,         # X-axis tick labels size.                                                                          #
    "ytick.labelsize": 8,         # Y-axis tick labels size.                                                                          #
    "legend.fontsize": 8,         # Legend font size for all text in the legend.                                                      #
    "figure.titlesize": 9,        # Overall figure title size for all text in the figure.                                             #
})                                                                                                                                    #
sns.set_theme(style="whitegrid", context="paper") # Set seaborn theme for additional aesthetics and context.                          #
# =================================================================================================================================== #


# =================================================================================================================================== #
# ----------------------------------------------------------- FUNCTIONS ------------------------------------------------------------- #
def get_meshgrid_3D(x: np.ndarray, y: np.ndarray, z: np.ndarray, nx: int, ny: int, nz: int, margin: float = 0.1) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Generate a 3D meshgrid for plotting.

    Parameters:
        x (np.ndarray): X-coordinates.
        y (np.ndarray): Y-coordinates.
        z (np.ndarray): Z-coordinates.
        nx (int): Number of points along the x-axis.
        ny (int): Number of points along the y-axis.
        nz (int): Number of points along the z-axis.
        margin (float): Margin for the grid.

    Returns:
        tuple[np.ndarray, np.ndarray, np.ndarray]: 3D meshgrid arrays.
    """
    x_min, x_max = (1 + margin) * x.min() - margin * x.max(), (1 + margin) * x.max() - margin * x.min()
    y_min, y_max = (1 + margin) * y.min() - margin * y.max(), (1 + margin) * y.max() - margin * y.min()
    z_min, z_max = (1 + margin) * z.min() - margin * z.max(), (1 + margin) * z.max() - margin * z.min()
    
    xx, yy, zz = np.meshgrid(
        np.linspace(x_min, x_max, nx),
        np.linspace(y_min, y_max, ny),
        np.linspace(z_min, z_max, nz)
    )
    return xx, yy, zz


def plot_predicted_label_3D(ax: plt.Axes, clf, xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, X: np.ndarray, y: np.ndarray) -> None:
    """
    Plot the predicted labels in 3D.

    Parameters:
        ax (plt.Axes): Matplotlib 3D Axes object.
        clf: Classifier object with a predict method.
        xx, yy, zz (np.ndarray): Meshgrid coordinates.
        X (np.ndarray): Data points.
        y (np.ndarray): True labels.
    """
    # Use LaTeX for plot typography
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')
    
    # Create a sample grid for visualization (3D grids are too dense to visualize fully)
    sample_pts = 15  # Points per dimension
    x_sample = np.linspace(xx.min(), xx.max(), sample_pts)
    y_sample = np.linspace(yy.min(), yy.max(), sample_pts)
    z_sample = np.linspace(zz.min(), zz.max(), sample_pts)
    
    # Create sample points for all 3 dimensions
    X_pred = []
    for x in x_sample:
        for y in y_sample:
            for z in z_sample:
                X_pred.append([x, y, z])
    X_pred = np.array(X_pred)
    
    # Predict labels
    pred_labels = clf.predict(X_pred)
    
    # Plot the sample points with predicted labels
    # Plot each prediction class separately
    for label in np.unique(pred_labels):
        mask = (pred_labels == label)
        ax.scatter(X_pred[mask, 0], X_pred[mask, 1], X_pred[mask, 2],
                   color=f'C{label}', alpha=0.05, s=20)
    
    # Plot the original data points - class by class
    for label in np.unique(y):
        mask = (y == label)
        ax.scatter(X[mask, 0], X[mask, 1], X[mask, 2],
                   color=f'C{label}', edgecolors="k", linewidth=0.15, s=50,
                   label=f"Class {label}")
    
    ax.set_xlabel(r"$\cos(\phi)$")
    ax.set_ylabel(r"$\sin(\phi)$")
    ax.set_zlabel("RSSI [dBm]")
    ax.set_title("Predicted Labels")
    ax.legend()


def plot_prob_density_3D(ax: plt.Axes, model, xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, X: np.ndarray, t: np.ndarray) -> None:
    """
    Plot the probability density in 3D.

    Parameters:
        ax (plt.Axes): Matplotlib 3D Axes object.
        model: Model object with calculate_probability_density method.
        xx, yy, zz (np.ndarray): Meshgrid coordinates.
        X (np.ndarray): Data points.
        t (np.ndarray): True labels.
    """
    # Define fixed colors for clusters
    cluster_colors = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9']
    
    # Plot clusters as ellipsoids
    for k in range(model.K):
        # Get the mean and covariance for this cluster
        mean = model.mu[k]
        cov = model.sigma[k]
        
        # Generate points on a unit sphere
        u = np.linspace(0, 2 * np.pi, 20)
        v = np.linspace(0, np.pi, 20)
        x = np.outer(np.cos(u), np.sin(v))
        y = np.outer(np.sin(u), np.sin(v))
        z = np.outer(np.ones_like(u), np.cos(v))
        
        # Stack the sphere points
        sphere_points = np.vstack((x.flatten(), y.flatten(), z.flatten())).T
        
        # Calculate eigenvectors and eigenvalues of the covariance matrix
        eigvals, eigvecs = np.linalg.eigh(cov)
        
        # Make sure the eigenvalues are positive (for numerical stability)
        eigvals = np.maximum(eigvals, 1e-10)
        
        # Scale sphere points by the square root of eigenvalues
        scaled_sphere = sphere_points @ np.diag(np.sqrt(eigvals)) @ eigvecs.T
        
        # Add the cluster mean
        ellipsoid_points = scaled_sphere + mean
        
        # Reshape for plotting
        x_ellipsoid = ellipsoid_points[:, 0].reshape(x.shape)
        y_ellipsoid = ellipsoid_points[:, 1].reshape(y.shape)
        z_ellipsoid = ellipsoid_points[:, 2].reshape(z.shape)
        
        # Plot the ellipsoid with fixed color
        ax.plot_surface(x_ellipsoid, y_ellipsoid, z_ellipsoid, 
                       color=cluster_colors[k % len(cluster_colors)], alpha=0.2)
    
    # Plot the original data points - class by class
    for label in np.unique(t):
        mask = (t == label)
        ax.scatter(X[mask, 0], X[mask, 1], X[mask, 2],
                  color=cluster_colors[int(label) % len(cluster_colors)], edgecolor='k', linewidth=0.15, 
                  alpha=0.25, s=50, label=f"Class {label}")
    
    ax.set_xlabel(r"$\cos(\phi)$")
    ax.set_ylabel(r"$\sin(\phi)$")
    ax.set_zlabel("RSSI [dBm]")
    ax.set_title("Probability Density")
    ax.legend()


def plot_result_3D(model, xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, X: np.ndarray, y: np.ndarray) -> None:
    """
    Plot the 3D GMM results with two views.

    Parameters:
        model: Model object with predict and calculate_probability_density methods.
        xx, yy, zz (np.ndarray): Meshgrid coordinates.
        X (np.ndarray): Data points.
        y (np.ndarray): True labels.
    """
    fig = plt.figure(figsize=(15, 7))
    
    # First view - show predicted labels
    ax1 = fig.add_subplot(121, projection='3d')
    plot_predicted_label_3D(ax1, model, xx, yy, zz, X, y)
    
    # Second view - show probability density
    ax2 = fig.add_subplot(122, projection='3d')
    plot_prob_density_3D(ax2, model, xx, yy, zz, X, y)
    
    plt.tight_layout()
    plt.show()


def plot_result_with_point_3D(model, xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, X: np.ndarray, y: np.ndarray, point: tuple, cluster_names: list = None) -> None:
    """
    Plot the 3D GMM results with a specific test point highlighted.

    Parameters:
        model: Model object with predict and calculate_probability_density methods.
        xx, yy, zz (np.ndarray): Meshgrid coordinates.
        X (np.ndarray): Data points.
        y (np.ndarray): True labels.
        point (tuple): Point (x, y, z) to be highlighted.
        cluster_names (list): List of cluster names.
    """
    fig = plt.figure(figsize=(15, 7))
    
    # First view - show predicted labels
    ax1 = fig.add_subplot(121, projection='3d')
    plot_predicted_label_3D(ax1, model, xx, yy, zz, X, y)
    
    # Add the test point
    ax1.scatter([point[0]], [point[1]], [point[2]], 
               color='lime', s=200, edgecolors='white', linewidth=1.5,
               label='Test Point')
    
    # Second view - show probability density
    ax2 = fig.add_subplot(122, projection='3d')
    plot_prob_density_3D(ax2, model, xx, yy, zz, X, y)
    
    # Add the test point
    ax2.scatter([point[0]], [point[1]], [point[2]], 
               color='lime', s=200, edgecolors='white', linewidth=1.5,
               label='Test Point')
    
    # Add cluster labels near means
    for k in range(model.K):
        mean = model.mu[k]
        cluster_name = cluster_names[k] if cluster_names and k < len(cluster_names) else f'Cluster {k}'
        
        ax1.text(mean[0], mean[1], mean[2], cluster_name, 
                fontsize=10, ha='center', va='center')
        ax2.text(mean[0], mean[1], mean[2], cluster_name, 
                fontsize=10, ha='center', va='center')
    
    ax1.legend()
    ax2.legend()
    plt.tight_layout()
    plt.show()
# =================================================================================================================================== #


# =================================================================================================================================== #
# ------------------------------------------------------------ CLASSES -------------------------------------------------------------- #
class GaussianMixtureModel:
    """
    Gaussian Mixture Model class.

    Attributes:
        K (int): Number of clusters.
        mu (np.ndarray): Means of the clusters.
        sigma (np.ndarray): Covariances of the clusters.
        pi (np.ndarray): Mixing coefficients.
    """

    def __init__(self, K: int) -> None:
        """
        Initialize the Gaussian Mixture Model.

        Parameters:
            K (int): Number of clusters.
        """
        self.K = K
        self.mu = None
        self.sigma = None
        self.pi = None

    def init_params(self, X: np.ndarray, random_state: int = None) -> None:
        """
        Initialize the parameters of the model.

        Parameters:
            X (np.ndarray): Data points.
            random_state (int): Random state for reproducibility.
        """
        n_samples, n_features = np.shape(X)
        random = np.random.RandomState(seed=random_state)

        self.pi = np.ones(self.K) / self.K
        self.mu = X[random.choice(n_samples, size=self.K, replace=False)]
        self.sigma = np.tile(np.diag(np.var(X, axis=0)), (self.K, 1, 1))

    def calc_nmat(self, X: np.ndarray) -> np.ndarray:
        """
        Calculate the normal matrix.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            np.ndarray: Normal matrix.
        """
        n_samples, n_features = np.shape(X)

        difference = np.reshape(X, (n_samples, 1, n_features)) - np.reshape(self.mu, (1, self.K, n_features))
        L = np.linalg.inv(self.sigma)
        exponent = np.einsum("nkj, nkj -> nk", np.einsum("nki, kij -> nkj", difference, L), difference)
        return np.exp(-0.5 * exponent) / np.sqrt(np.linalg.det(self.sigma)) / (2 * np.pi) ** (n_features / 2)

    def E_step(self, X: np.ndarray) -> np.ndarray:
        """
        Perform the E-step of the EM algorithm.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            np.ndarray: Responsibilities matrix.
        """
        n_samples, n_features = np.shape(X)

        nmat = self.calc_nmat(X)
        tmp = nmat * self.pi
        gamma = tmp / np.reshape(np.sum(tmp, axis=1), (n_samples, 1))
        return gamma

    def M_step(self, X: np.ndarray, gamma: np.ndarray) -> None:
        """
        Perform the M-step of the EM algorithm.

        Parameters:
            X (np.ndarray): Data points.
            gamma (np.ndarray): Responsibilities matrix.
        """
        n_samples, n_features = np.shape(X)

        difference = np.reshape(X, (n_samples, 1, n_features)) - np.reshape(self.mu, (1, self.K, n_features))
        Nk = np.sum(gamma, axis=0)
        self.pi = Nk / n_samples
        self.mu = gamma.T @ X / np.reshape(Nk, (self.K, 1))
        self.sigma = np.einsum("nki, nkj -> kij", np.einsum("nk, nki -> nki", gamma, difference), difference) / np.reshape(Nk, (self.K, 1, 1))

    def calculate_probability_density(self, X: np.ndarray) -> np.ndarray:
        """
        Calculate the probability density.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            np.ndarray: Probability density.
        """
        return self.calc_nmat(X) @ self.pi

    def calculate_log_likelihood(self, X: np.ndarray) -> float:
        """
        Calculate the log likelihood.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            float: Log likelihood.
        """
        return np.sum(np.log(self.calculate_probability_density(X)))

    def fit(self, X: np.ndarray, max_iter: int = 100, tol: float = 1e-6, display_message: bool = False, random_state: int = None) -> None:
        """
        Fit the model to the data.

        Parameters:
            X (np.ndarray): Data points.
            max_iter (int): Maximum number of iterations.
            tol (float): Tolerance for convergence.
            display_message (bool): Whether to display messages during fitting.
            random_state (int): Random state for reproducibility.
        """
        self.init_params(X, random_state=random_state)
        log_likelihood = -np.float64("inf")

        for i in range(max_iter):
            gamma = self.E_step(X)
            self.M_step(X, gamma)
            previous_log_likelihood = log_likelihood
            log_likelihood = self.calculate_log_likelihood(X)

            if (log_likelihood - previous_log_likelihood) < tol:
                break

            if display_message:
                print(f"Iteration {i}: {log_likelihood}")

    def predict_probability(self, X: np.ndarray) -> np.ndarray:
        """
        Predict the probabilities for each cluster.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            np.ndarray: Probabilities for each cluster.
        """
        return self.E_step(X)

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict the cluster labels.

        Parameters:
            X (np.ndarray): Data points.

        Returns:
            np.ndarray: Predicted cluster labels.
        """
        return np.argmax(self.predict_probability(X), axis=1)
# =================================================================================================================================== #