<reponame>artemmavrin/SLTools
"""Implementation of the k-means clustering algorithm.

See Section 14.3.6 of Hastie, Tibshirani, & Friedman (2009) and Section 9.1 of
Bishop (2006).

References
----------
<NAME>, <NAME>, and <NAME>. The Elements of
    Statistical Learning: Data Mining, Inference, and Prediction (Second).
    Springer 2009
<NAME>. Pattern Recognition and Machine Learning. Springer 2006.
"""

import itertools

import numpy as np

from ..generic import Predictor
from ..utils.validation import validate_bool
from ..utils.validation import validate_float
from ..utils.validation import validate_int
from ..utils.validation import validate_sample


class KMeansCluster(Predictor):
    """Partition data into K clusters based on minimizing the
    within-cluster-sum-of-squares loss function.

    Properties
    ----------
    k : int
        The number of clusters.
    centers : numpy.ndarray
        The centers of each cluster.
    standardize : bool
        Indicate whether the input data should be standardized.
    random_state : numpy.random.RandomState
        The random number generator.
    """
    k: int = None
    standardize: bool = None
    random_state: np.random.RandomState

    # Cluster centers in terms of possibly standardized units
    _centers: np.ndarray = None

    # Vectors of means and standard deviations of the feature matrix columns
    _x_mean: np.ndarray = None
    _x_std: np.ndarray = None

    def __init__(self, k, standardize=True, random_state=None):
        """Initialize a ClusterKMeans object.

        Parameters
        ----------
        k : int
            The number of clusters.
        standardize : bool
            Indicate whether the input data should be standardized.
        """
        # Validate parameters
        self.k = validate_int(k, "k", minimum=2)
        self.standardize = validate_bool(standardize, "standardize")

        # Seed the RNG
        if isinstance(random_state, np.random.RandomState):
            self.random_state = random_state
        else:
            self.random_state = np.random.RandomState(random_state)

    @property
    def centers(self):
        """Get the centers of each cluster."""
        # Check if the model is fitted
        if not self.fitted:
            raise self.unfitted_exception
        # Undo standardization if necessary
        if self.standardize:
            return self._x_mean + self._x_std * self._centers
        else:
            return self._centers

    def fit(self, x, tol=1e-5, iterations=None, repeats=5):
        """Fit the K-means cluster model to data.

        Parameters
        ----------
        x : array-like
            Feature matrix to determine the clustering.
        tol : float, optional
            Positive tolerance for algorithm convergence.
            If the Euclidean distance between the centers in successive
            iterations is less than `tol`, the algorithm stops.
        iterations : int, optional
            Maximum number of iterations to perform. If None, no maximum number
            is imposed.
        repeats : int, optional
            Number of times to repeat the algorithm with different initial
            center assignments. This can decrease the chance of finding only
            non-global minima of the loss function

        Returns
        -------
        This KMeansCluster instance.
        """
        # Validate parameters
        x, tol, iterations, repeats = \
            _validate_fit_params(x, self.k, tol, iterations, repeats)

        # n = number of observations, p = number of features
        n, p = x.shape

        # Initialize array of cluster centers
        centers = np.empty(shape=(repeats, self.k, p), dtype=np.float_)

        # Initialize array of within-cluster-sum-of-squares losses
        loss = np.empty(shape=repeats)

        # Standardize data if necessary
        if self.standardize:
            self._x_mean = x.mean(axis=0)
            self._x_std = x.std(axis=0, ddof=0)
            if np.any(self._x_std == 0):
                self._x_std[self._x_std == 0] = 1.0
            x = (x - self._x_mean) / self._x_std

        # Repeat the k-means clustering algorithm to decrease the chance of
        # finding only non-global minima of the loss function
        for i in range(repeats):
            # Initialize random starting cluster centers by choosing k
            # observations without replacement from the feature matrix
            ind = self.random_state.choice(n, size=self.k, replace=False)
            centers_ = x.take(ind, axis=0)

            if iterations is None:
                counter = itertools.count()
            else:
                counter = range(iterations)

            # Perform the k-means clustering algorithm (an EM algorithm)
            for _ in counter:
                # Save the current centers for later
                centers_old = centers_

                # (E step) Assign each observation to a cluster
                clusters = _assign_clusters(x, centers_)

                # (M step) Get new centers as the means of the currently
                # assigned clusters
                centers_ = np.zeros(shape=(self.k, p), dtype=np.float_)
                for j in range(self.k):
                    centers_[j, :] = np.mean(x[clusters == j, :], axis=0)

                # Check for convergence
                if np.linalg.norm(centers_ - centers_old) < tol:
                    break

            # Sort the final centers in lexicographic order
            ind = np.lexsort([centers_[:, p - i - 1] for i in range(p)])
            centers_ = centers_[ind]

            # Save the resulting centers and compute final cluster assignments
            centers[i] = centers_
            clusters = _assign_clusters(x, centers_)

            # Evaluate the clusters by their within-cluster-sum-of-squares
            loss[i] = _k_means_loss(x, clusters=clusters, centers=centers_)

        # Choose the cluster center assignment that minimizes the loss
        self._centers = centers[np.argmin(loss)]

        self.fitted = True
        return self

    def predict(self, x):
        """Assign each observation to one of K clusters.

        Parameters
        ----------
        x : array-like
            Feature matrix of shape (n, p) to cluster.

        Returns
        -------
        clusters : numpy.ndarray
            Array of shape (n, ) in which the i-th entry is the index of the
            cluster of the i-th observation x[i, :].
        """
        # Check if the model is fitted
        if not self.fitted:
            raise self.unfitted_exception

        # Validate the feature matrix
        x = validate_sample(x, n_dim=2)

        # Standardize data if necessary
        if self.standardize:
            x = (x - self._x_mean) / self._x_std

        # Return cluster assignments
        return _assign_clusters(x, self._centers)


def _k_means_loss(x: np.ndarray, clusters: np.ndarray, centers: np.ndarray):
    """Compute the within-cluster-sum-of-squares loss function.

    Parameters
    ----------
    x : numpy.ndarray
        Feature matrix of shape (n, p) to cluster.
    clusters : numpy.ndarray
        Array of shape (n, ) in which the i-th entry is the index of the cluster
        of the i-th observation x[i, :].
    centers : numpy.ndarray
        Array of shape (k, p) of the centers of each of the k clusters.

    Returns
    -------
    The within-cluster-sum-of-squares loss (cf. equation (14.31) in Hastie,
    Tibshirani, & Friedman (2009)).
    """
    loss = 0.0
    for i, center in enumerate(centers):
        n = np.sum(clusters == i)
        loss += n * np.sum(np.linalg.norm(x[clusters == i] - center, axis=1))
    return loss


def _assign_clusters(x: np.ndarray, centers: np.ndarray) -> np.ndarray:
    """Given cluster centers, assign data points to clusters.

    Parameters
    ----------
    x : numpy.ndarray
        Feature matrix of shape (n, p) to cluster.
    centers : numpy.ndarray
        Array of shape (k, p) of the centers of each of the k clusters.

    Returns
    -------
    clusters : numpy.ndarray
        Array of shape (n, ) in which the i-th entry is the index of the cluster
        of the i-th observation x[i, :].
    """
    clusters = np.empty(shape=len(x), dtype=np.int_)
    for i, obs in enumerate(x):
        clusters[i] = np.argmin(np.linalg.norm(obs - centers, axis=1))
    return clusters


def _validate_fit_params(x, k, tol, iterations, repeats):
    """Validate the parameters for KMeansCluster.fit().

    Parameters
    ----------
    See the parameter descriptions for GaussianMixture.fit().

    Returns
    -------
    The updated parameters.
    """
    x = validate_sample(x, n_dim=2)
    if len(x) <= k:
        print("There must be more observations than number of clusters.")

    tol = validate_float(tol, "tol", positive=True)

    if iterations is not None:
        iterations = validate_int(iterations, "iterations", minimum=1)

    repeats = validate_int(repeats, "repeats", minimum=1)

    return x, tol, iterations, repeats

"""Implementation of the k-means clustering algorithm.

See Section 14.3.6 of Hastie, Tibshirani, & Friedman (2009) and Section 9.1 of
Bishop (2006).

References
----------
<NAME>, <NAME>, and <NAME>. The Elements of
    Statistical Learning: Data Mining, Inference, and Prediction (Second).
    Springer 2009
<NAME>. Pattern Recognition and Machine Learning. Springer 2006.
"""

import itertools

import numpy as np

from ..generic import Predictor
from ..utils.validation import validate_bool
from ..utils.validation import validate_float
from ..utils.validation import validate_int
from ..utils.validation import validate_sample


class KMeansCluster(Predictor):
    """Partition data into K clusters based on minimizing the
    within-cluster-sum-of-squares loss function.

    Properties
    ----------
    k : int
        The number of clusters.
    centers : numpy.ndarray
        The centers of each cluster.
    standardize : bool
        Indicate whether the input data should be standardized.
    random_state : numpy.random.RandomState
        The random number generator.
    """
    k: int = None
    standardize: bool = None
    random_state: np.random.RandomState

    # Cluster centers in terms of possibly standardized units
    _centers: np.ndarray = None

    # Vectors of means and standard deviations of the feature matrix columns
    _x_mean: np.ndarray = None
    _x_std: np.ndarray = None

    def __init__(self, k, standardize=True, random_state=None):
        """Initialize a ClusterKMeans object.

        Parameters
        ----------
        k : int
            The number of clusters.
        standardize : bool
            Indicate whether the input data should be standardized.
        """
        # Validate parameters
        self.k = validate_int(k, "k", minimum=2)
        self.standardize = validate_bool(standardize, "standardize")

        # Seed the RNG
        if isinstance(random_state, np.random.RandomState):
            self.random_state = random_state
        else:
            self.random_state = np.random.RandomState(random_state)

    @property
    def centers(self):
        """Get the centers of each cluster."""
        # Check if the model is fitted
        if not self.fitted:
            raise self.unfitted_exception
        # Undo standardization if necessary
        if self.standardize:
            return self._x_mean + self._x_std * self._centers
        else:
            return self._centers

    def fit(self, x, tol=1e-5, iterations=None, repeats=5):
        """Fit the K-means cluster model to data.

        Parameters
        ----------
        x : array-like
            Feature matrix to determine the clustering.
        tol : float, optional
            Positive tolerance for algorithm convergence.
            If the Euclidean distance between the centers in successive
            iterations is less than `tol`, the algorithm stops.
        iterations : int, optional
            Maximum number of iterations to perform. If None, no maximum number
            is imposed.
        repeats : int, optional
            Number of times to repeat the algorithm with different initial
            center assignments. This can decrease the chance of finding only
            non-global minima of the loss function

        Returns
        -------
        This KMeansCluster instance.
        """
        # Validate parameters
        x, tol, iterations, repeats = \
            _validate_fit_params(x, self.k, tol, iterations, repeats)

        # n = number of observations, p = number of features
        n, p = x.shape

        # Initialize array of cluster centers
        centers = np.empty(shape=(repeats, self.k, p), dtype=np.float_)

        # Initialize array of within-cluster-sum-of-squares losses
        loss = np.empty(shape=repeats)

        # Standardize data if necessary
        if self.standardize:
            self._x_mean = x.mean(axis=0)
            self._x_std = x.std(axis=0, ddof=0)
            if np.any(self._x_std == 0):
                self._x_std[self._x_std == 0] = 1.0
            x = (x - self._x_mean) / self._x_std

        # Repeat the k-means clustering algorithm to decrease the chance of
        # finding only non-global minima of the loss function
        for i in range(repeats):
            # Initialize random starting cluster centers by choosing k
            # observations without replacement from the feature matrix
            ind = self.random_state.choice(n, size=self.k, replace=False)
            centers_ = x.take(ind, axis=0)

            if iterations is None:
                counter = itertools.count()
            else:
                counter = range(iterations)

            # Perform the k-means clustering algorithm (an EM algorithm)
            for _ in counter:
                # Save the current centers for later
                centers_old = centers_

                # (E step) Assign each observation to a cluster
                clusters = _assign_clusters(x, centers_)

                # (M step) Get new centers as the means of the currently
                # assigned clusters
                centers_ = np.zeros(shape=(self.k, p), dtype=np.float_)
                for j in range(self.k):
                    centers_[j, :] = np.mean(x[clusters == j, :], axis=0)

                # Check for convergence
                if np.linalg.norm(centers_ - centers_old) < tol:
                    break

            # Sort the final centers in lexicographic order
            ind = np.lexsort([centers_[:, p - i - 1] for i in range(p)])
            centers_ = centers_[ind]

            # Save the resulting centers and compute final cluster assignments
            centers[i] = centers_
            clusters = _assign_clusters(x, centers_)

            # Evaluate the clusters by their within-cluster-sum-of-squares
            loss[i] = _k_means_loss(x, clusters=clusters, centers=centers_)

        # Choose the cluster center assignment that minimizes the loss
        self._centers = centers[np.argmin(loss)]

        self.fitted = True
        return self

    def predict(self, x):
        """Assign each observation to one of K clusters.

        Parameters
        ----------
        x : array-like
            Feature matrix of shape (n, p) to cluster.

        Returns
        -------
        clusters : numpy.ndarray
            Array of shape (n, ) in which the i-th entry is the index of the
            cluster of the i-th observation x[i, :].
        """
        # Check if the model is fitted
        if not self.fitted:
            raise self.unfitted_exception

        # Validate the feature matrix
        x = validate_sample(x, n_dim=2)

        # Standardize data if necessary
        if self.standardize:
            x = (x - self._x_mean) / self._x_std

        # Return cluster assignments
        return _assign_clusters(x, self._centers)


def _k_means_loss(x: np.ndarray, clusters: np.ndarray, centers: np.ndarray):
    """Compute the within-cluster-sum-of-squares loss function.

    Parameters
    ----------
    x : numpy.ndarray
        Feature matrix of shape (n, p) to cluster.
    clusters : numpy.ndarray
        Array of shape (n, ) in which the i-th entry is the index of the cluster
        of the i-th observation x[i, :].
    centers : numpy.ndarray
        Array of shape (k, p) of the centers of each of the k clusters.

    Returns
    -------
    The within-cluster-sum-of-squares loss (cf. equation (14.31) in Hastie,
    Tibshirani, & Friedman (2009)).
    """
    loss = 0.0
    for i, center in enumerate(centers):
        n = np.sum(clusters == i)
        loss += n * np.sum(np.linalg.norm(x[clusters == i] - center, axis=1))
    return loss


def _assign_clusters(x: np.ndarray, centers: np.ndarray) -> np.ndarray:
    """Given cluster centers, assign data points to clusters.

    Parameters
    ----------
    x : numpy.ndarray
        Feature matrix of shape (n, p) to cluster.
    centers : numpy.ndarray
        Array of shape (k, p) of the centers of each of the k clusters.

    Returns
    -------
    clusters : numpy.ndarray
        Array of shape (n, ) in which the i-th entry is the index of the cluster
        of the i-th observation x[i, :].
    """
    clusters = np.empty(shape=len(x), dtype=np.int_)
    for i, obs in enumerate(x):
        clusters[i] = np.argmin(np.linalg.norm(obs - centers, axis=1))
    return clusters


def _validate_fit_params(x, k, tol, iterations, repeats):
    """Validate the parameters for KMeansCluster.fit().

    Parameters
    ----------
    See the parameter descriptions for GaussianMixture.fit().

    Returns
    -------
    The updated parameters.
    """
    x = validate_sample(x, n_dim=2)
    if len(x) <= k:
        print("There must be more observations than number of clusters.")

    tol = validate_float(tol, "tol", positive=True)

    if iterations is not None:
        iterations = validate_int(iterations, "iterations", minimum=1)

    repeats = validate_int(repeats, "repeats", minimum=1)

    return x, tol, iterations, repeats
Implementation of the k-means clustering algorithm.  See Section 14.3.6 of Hastie, Tibshirani, & Friedman (2009) and Section 9.1 of Bishop (2006).  References ---------- <NAME>, <NAME>, and <NAME>. The Elements of     Statistical Learning: Data Mining, Inference, and Prediction (Second).     Springer 2009 <NAME>. Pattern Recognition and Machine Learning. Springer 2006. Partition data into K clusters based on minimizing the     within-cluster-sum-of-squares loss function.      Properties     ----------     k : int         The number of clusters.     centers : numpy.ndarray         The centers of each cluster.     standardize : bool         Indicate whether the input data should be standardized.     random_state : numpy.random.RandomState         The random number generator. # Cluster centers in terms of possibly standardized units # Vectors of means and standard deviations of the feature matrix columns Initialize a ClusterKMeans object.          Parameters         ----------         k : int             The number of clusters.         standardize : bool             Indicate whether the input data should be standardized. # Validate parameters # Seed the RNG Get the centers of each cluster. # Check if the model is fitted # Undo standardization if necessary Fit the K-means cluster model to data.          Parameters         ----------         x : array-like             Feature matrix to determine the clustering.         tol : float, optional             Positive tolerance for algorithm convergence.             If the Euclidean distance between the centers in successive             iterations is less than `tol`, the algorithm stops.         iterations : int, optional             Maximum number of iterations to perform. If None, no maximum number             is imposed.         repeats : int, optional             Number of times to repeat the algorithm with different initial             center assignments. This can decrease the chance of finding only             non-global minima of the loss function          Returns         -------         This KMeansCluster instance. # Validate parameters # n = number of observations, p = number of features # Initialize array of cluster centers # Initialize array of within-cluster-sum-of-squares losses # Standardize data if necessary # Repeat the k-means clustering algorithm to decrease the chance of # finding only non-global minima of the loss function # Initialize random starting cluster centers by choosing k # observations without replacement from the feature matrix # Perform the k-means clustering algorithm (an EM algorithm) # Save the current centers for later # (E step) Assign each observation to a cluster # (M step) Get new centers as the means of the currently # assigned clusters # Check for convergence # Sort the final centers in lexicographic order # Save the resulting centers and compute final cluster assignments # Evaluate the clusters by their within-cluster-sum-of-squares # Choose the cluster center assignment that minimizes the loss Assign each observation to one of K clusters.          Parameters         ----------         x : array-like             Feature matrix of shape (n, p) to cluster.          Returns         -------         clusters : numpy.ndarray             Array of shape (n, ) in which the i-th entry is the index of the             cluster of the i-th observation x[i, :]. # Check if the model is fitted # Validate the feature matrix # Standardize data if necessary # Return cluster assignments Compute the within-cluster-sum-of-squares loss function.      Parameters     ----------     x : numpy.ndarray         Feature matrix of shape (n, p) to cluster.     clusters : numpy.ndarray         Array of shape (n, ) in which the i-th entry is the index of the cluster         of the i-th observation x[i, :].     centers : numpy.ndarray         Array of shape (k, p) of the centers of each of the k clusters.      Returns     -------     The within-cluster-sum-of-squares loss (cf. equation (14.31) in Hastie,     Tibshirani, & Friedman (2009)). Given cluster centers, assign data points to clusters.      Parameters     ----------     x : numpy.ndarray         Feature matrix of shape (n, p) to cluster.     centers : numpy.ndarray         Array of shape (k, p) of the centers of each of the k clusters.      Returns     -------     clusters : numpy.ndarray         Array of shape (n, ) in which the i-th entry is the index of the cluster         of the i-th observation x[i, :]. Validate the parameters for KMeansCluster.fit().      Parameters     ----------     See the parameter descriptions for GaussianMixture.fit().      Returns     -------     The updated parameters.