gaussian_mixture.py

# coding:utf-8

import random

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import multivariate_normal

from mla.base import BaseEstimator
from mla.kmeans import KMeans


class GaussianMixture(BaseEstimator):
    """Gaussian Mixture Model: clusters with Gaussian prior.

    Finds clusters by repeatedly performing Expectation–Maximization (EM) algorithm
    on the dataset. GMM assumes the datasets is distributed in multivariate Gaussian,
    and tries to find the underlying structure of the Gaussian, i.e. mean and covariance.
    E-step computes the "responsibility" of the data to each cluster, given the mean
    and covariance; M-step computes the mean, covariance and weights (prior of each
    cluster), given the responsibilities. It iterates until the total likelihood
    changes less than the tolerance.


    Parameters
    ----------

    K : int
        The number of clusters into which the dataset is partitioned.
    max_iters: int
        The maximum iterations of assigning points to the perform EM.
        Short-circuited by the assignments converging on their own.
    init: str, default 'random'
        The name of the method used to initialize the first clustering.

        'random' - Randomly select values from the dataset as the K centroids.
        'kmeans' - Initialize the centroids, covariances, weights with KMeams's clusters.
    tolerance: float, default 1e-3
        The tolerance of difference of the two latest likelihood for convergence.
    """

    y_required = False

    def __init__(self, K=4, init="random", max_iters=500, tolerance=1e-3):
        self.K = K
        self.max_iters = max_iters
        self.init = init
        self.assignments = None
        self.likelihood = []
        self.tolerance = tolerance

    def fit(self, X, y=None):
        """Perform Expectation–Maximization (EM) until converged."""
        self._setup_input(X, y)
        self._initialize()
        for _ in range(self.max_iters):
            self._E_step()
            self._M_step()
            if self._is_converged():
                break

    def _initialize(self):
        """Set the initial weights, means and covs (with full covariance matrix).

        weights: the prior of the clusters (what percentage of data does a cluster have)
        means: the mean points of the clusters
        covs: the covariance matrix of the clusters
        """
        self.weights = np.ones(self.K)
        if self.init == "random":
            self.means = [self.X[x] for x in random.sample(range(self.n_samples), self.K)]
            self.covs = [np.cov(self.X.T) for _ in range(self.K)]

        elif self.init == "kmeans":
            kmeans = KMeans(K=self.K, max_iters=self.max_iters // 3, init="++")
            kmeans.fit(self.X)
            self.assignments = kmeans.predict()
            self.means = kmeans.centroids
            self.covs = []
            for i in np.unique(self.assignments):
                self.weights[int(i)] = (self.assignments == i).sum()
                self.covs.append(np.cov(self.X[self.assignments == i].T))
        else:
            raise ValueError("Unknown type of init parameter")
        self.weights /= self.weights.sum()

    def _E_step(self):
        """Expectation(E-step) for Gaussian Mixture."""
        likelihoods = self._get_likelihood(self.X)
        self.likelihood.append(likelihoods.sum())
        weighted_likelihoods = self._get_weighted_likelihood(likelihoods)
        self.assignments = weighted_likelihoods.argmax(axis=1)
        weighted_likelihoods /= weighted_likelihoods.sum(axis=1)[:, np.newaxis]
        self.responsibilities = weighted_likelihoods

    def _M_step(self):
        """Maximization (M-step) for Gaussian Mixture."""
        weights = self.responsibilities.sum(axis=0)
        for assignment in range(self.K):
            resp = self.responsibilities[:, assignment][:, np.newaxis]
            self.means[assignment] = (resp * self.X).sum(axis=0) / resp.sum()
            self.covs[assignment] = (self.X - self.means[assignment]).T.dot(
                (self.X - self.means[assignment]) * resp
            ) / weights[assignment]
        self.weights = weights / weights.sum()

    def _is_converged(self):
        """Check if the difference of the latest two likelihood is less than the tolerance."""
        if (len(self.likelihood) > 1) and (self.likelihood[-1] - self.likelihood[-2] <= self.tolerance):
            return True
        return False

    def _predict(self, X):
        """Get the assignments for X with GMM clusters."""
        if not X.shape:
            return self.assignments
        likelihoods = self._get_likelihood(X)
        weighted_likelihoods = self._get_weighted_likelihood(likelihoods)
        assignments = weighted_likelihoods.argmax(axis=1)
        return assignments

    def _get_likelihood(self, data):
        n_data = data.shape[0]
        likelihoods = np.zeros([n_data, self.K])
        for c in range(self.K):
            likelihoods[:, c] = multivariate_normal.pdf(data, self.means[c], self.covs[c])
        return likelihoods

    def _get_weighted_likelihood(self, likelihood):
        return self.weights * likelihood

    def plot(self, data=None, ax=None, holdon=False):
        """Plot contour for 2D data."""
        if not (len(self.X.shape) == 2 and self.X.shape[1] == 2):
            raise AttributeError("Only support for visualizing 2D data.")

        if ax is None:
            _, ax = plt.subplots()

        if data is None:
            data = self.X
            assignments = self.assignments
        else:
            assignments = self.predict(data)

        COLOR = "bgrcmyk"
        cmap = lambda assignment: COLOR[int(assignment) % len(COLOR)]

        # generate grid
        delta = 0.025
        margin = 0.2
        xmax, ymax = self.X.max(axis=0) + margin
        xmin, ymin = self.X.min(axis=0) - margin
        axis_X, axis_Y = np.meshgrid(np.arange(xmin, xmax, delta), np.arange(ymin, ymax, delta))

        def grid_gaussian_pdf(mean, cov):
            grid_array = np.array(list(zip(axis_X.flatten(), axis_Y.flatten())))
            return multivariate_normal.pdf(grid_array, mean, cov).reshape(axis_X.shape)

        # plot scatters
        if assignments is None:
            c = None
        else:
            c = [cmap(assignment) for assignment in assignments]
        ax.scatter(data[:, 0], data[:, 1], c=c)

        # plot contours
        for assignment in range(self.K):
            ax.contour(
                axis_X,
                axis_Y,
                grid_gaussian_pdf(self.means[assignment], self.covs[assignment]),
                colors=cmap(assignment),
            )

        if not holdon:
            plt.show()