I have recently built a class that is an implementation of kMeans from scratch. I believe there is room for improvement and I would happily receive some feedback. The project can be found at: https://github.com/EmpanS/kMeans-From-Scratch

All code for the class is found below:

# Import useful libraries
import numpy as np

class kMeans:
    """A class used to perform the k-means clustering algorithm on a data set. The maximum number of
     iterations is set by the user, if it converges to a solution, it stops iterating. 

    AVAILABLE_DIST_F : (list of str)
        Contains the available distance functions. L1_norm is the Manhattan distance, L2_norm is the
        ordinary Euclidean distance.

    k : (int)
        Represents the number of clusters

    X : (numpy array)
        The data to cluster, must be an (m x n)-numpy array with m observations and n features.

    verbose : (boolean)
        A boolean representing if printing should be done while training the model.

    h_params : (dictionary)
        Contains two hyper-parameters, number of iterations (n_iter) and distance function (dist_f)

    random_state : (int)
        Optional setting for the random state. The k-means algorithm does not guarantee finding a 
        global minimum, but the final clusters depends on the initial random cluters.
    labels : (numpy array)
        Contains the predicted label for each observation, i.e., what cluster it belongs to.
    cluster_centers (numpy array)
        Contains the n-dimensional coordinates for each cluster.
    update_h_params(self, h_params)
        Updates hyper parameters.
    fit(self, X=None)
        Performs the k-means algorithm on the passed data X or, if no data is passed, on self.X 

    __calculate_distances(X, centers)
        Calculates the distances between all observations in X and all centers of clusters. Uses the
        distance function already specified as a hyper-parameter.    

    __validate_param(h_param, setting)
        Validate new hyper-parameter settings.

    def __init__(self, k, X, verbose=True, h_params=None, random_state=None):
        self.AVAILABLE_DIST_F = ["L1_norm", "L2_norm"]
        self.k = k
        self.X = X
        self.verbose = verbose
        self.h_params = {'n_iter':100, 'dist_f':'L2_norm'}

        if h_params != None:

        self.labels = np.full((X.shape[0], 1), np.nan)
        self.cluster_centers = np.full((self.k, 1), np.nan)

    def update_h_params(self, h_params):
        """Updates the hyper parameters.

        h_params : (dict)
            Dictionary containing the hyper parameter/s and its updated setting/s.


        if type(h_params) != dict:
            raise TypeError('The argument must be a dictionary.')
        for h_param, setting in h_params.items():
            self.__validate_param(h_param, setting)
            self.h_params[h_param] = setting 

    def fit(self, X=None):
        """Performs the k-means algorithm. First, all observations in X gets randomly assigned a 
        label. Then, the function iterates until a solution is found (converged) or the maximum 
        number of iterations is reached. Each iteration performs the following:
            - Update labels
            - Check convergence
            - Calculate new cluster centers

        X : (numpy array)
            The data to cluster, must be an (m x n)-numpy array with m observations and n features.

        wss : (numpy array)
            A numpy array that saves the within cluster sum of squares for each iteration.    
        Labels : (numpy array)
            A numpy array containing all new labels for the observations in X.

        if X == None:
            X = self.X

        if (n := self.random_state) != None:

        # Initiate array to save within cluster sum of squares
        wss = np.zeros((1, self.h_params['n_iter']))

        # Randomly draw k observations and set them as the initial cluster centers 
        center_index = np.random.choice(X.shape[0], size=self.k, replace=False)
        cluster_centers = X[center_index]
        old_labels = None

        for iter in range(self.h_params['n_iter']):
            # Label the observations using the updated cluster centers
            distances  = self.__calculate_distances(X, cluster_centers)
            labels = np.argmin(distances, axis=1)

            # Calculate the within-sum-of-squares
            wss[0,iter] = sum(np.min(distances, axis=1))

            # Check convergence
            if np.all(labels == old_labels):
                if self.verbose:
                    print(f"Converged to a solution after {iter} iterations!")
                return(wss[0,:(iter)], labels)
                old_labels = labels

            # Calculate new cluster centers
            for i in range(self.k):
                cluster_centers[i] = np.sum(X[labels==i],axis=0)/(X[labels==i].shape[0])

        if self.verbose:
            print(f"Did not converged, reached max iterations. Completed {iter+1} iterations.")
        return(wss[0,:], labels)        

    def __calculate_distances(self, X, centers):
        Calculates the distances between all observations in X and all cluster centers. The already
        specified distance function (found in self.h_params) is used to calculate the distances.

        X : (numpy array)
            A matrix (m x n) containing all observations.

        centers : (numpy array)
            A matrix (k x n) where k is the number of clusters, containing all cluster centers. 

        labels : (numpy array)
            A numpy array containing all new labels for the observations in X.

        # Initiate a distance matrix
        distance_m = np.tile(centers.flatten(), (X.shape[0],1))

        # Duplicate data matrix to same dimension as distance matrix
        X_m = np.tile(X, (centers.shape[0]))

        if self.h_params["dist_f"] == "L2_norm":
            # Complete the distance matrix using the L2-norm
            distance_m = np.reshape(distance_m - X_m, (X.shape[0]*centers.shape[0], X.shape[1]))
            distance_m = np.sum(np.square(distance_m),axis=1, keepdims=True)

            # Reshape distance matrix
            distance_m = np.sqrt(np.reshape(distance_m, (X.shape[0], len(centers))))

        elif self.h_params["dist_f"] == "L1_norm":
            # Complete the distance matrix using the L1-norm
            distance_m = np.reshape(distance_m - X_m, (X.shape[0]*centers.shape[0], X.shape[1]))
            distance_m = np.sum(np.abs(distance_m),axis=1, keepdims=True)

            # Reshape distance matrix
            distance_m = np.reshape(distance_m, (X.shape[0], len(centers)))

            raise ValueError('Could not calculate distance, no distance function found.')

    def __validate_param(self, h_param, setting):
        Validates a given hyper-parameter update. The update must must have a valid key and value.

        h_param : (str)
            The hyper parameter to update 

        setting : (int) or (str)
            The new setting of the hyper parameter

        None - (Throws an error if not valid.)

        if h_param not in self.h_params.keys():
            raise KeyError("No hyper parameter is named " + str(h_param) + ", it is a wrong value of key. Must be either 'n_iter' or 'dist_f'.")

        if h_param == "n_iter":
            if type(setting) != int or setting <= 0:
                raise ValueError("n_iter must be a positive integer which " + str(setting) + " is not.")
        else: # Setting for the distance function
            if setting not in self.AVAILABLE_DIST_F:
                raise ValueError(str(setting) + " is not an available distance function. Available functions are: " + str(self.AVAILABLE_DIST_F))

Unnecessary type checking

In update_h_params, you write

if type(h_params) != dict:
            raise TypeError('The argument must be a dictionary.')

In Python, we don't care about the actual types of objects, only the interfaces they offer. This ideology is often referred to as Duck Typing, stemming from the notion that if something "walks like a duck, and quacks like a duck, its probably a duck." In this method you don't actually care whether or not the user passed a dict, only that whatever they gave you implements an items() method that returns an iterator of tuples. Your end users may very well have a reason to be using a different flavor of dictionary down the road, such as an OrderedDict or a defaultdict, which would be needlessly rejected by the current implementation.

Using __methods as "private" methods

The double-underscore syntax for method names is not intended for creating "private" methods. Rather, it's a tool to add namescrambling to method names so that subclasses don't accidentally overwrite a critical method. This doesn't appear to apply to your class.

To mark an attribute or method as "private", a single underscore will suffice.

Re-inventing logging

Each time you write

if self.verbose:

could be more idiomatically replaced by using the standard logging module. This has the added benefit of being much more configurable and will lead to fewer surprises when other developers try to use your code and receive unexpected writes to stdout.

Duplicate constants

I'd advise replace the attribute

self.AVAILABLE_DIST_F = ["L1_norm", "L2_norm"]

with a single class attribute

class kMeans:
    AVAILABLE_DIST_F = ["L1_norm", "L2_norm"]

As currently implemented, a new list will be created for each kMeans instance you create, which is a minor, but unnecessary, memory overhead. This also would more clearly convey that AVAILABLE_DIST_F does not change on an instance-by-instance basis.

Naming Conventions

Per PEP-8, it is recommend that class names use the CapsWords convention. Therefore, you may wish to rename kMeans to KMeans.

TMI in Docstrings

In your docstring for kMeans, you document each of the public methods your class provides, along with a short synopsis of their functionality. It is worth noting that standard Python documentation generating tools, such as sphinx, will create blocks such as this automatically, making this section somewhat unnecessary. In addition, by documenting methods in two separate locations, you make your class liable to have its methods change without the associated documentation being updated along with them. I'd advise documenting your method in their own respective docstrings and restricting the class-level docstring to the high-level overview of your class' interface.

Using Type Hints

You currently only note the expect types of various method parameters in docstrings. This information can be more effectively conveys by adding type hints to the methods themselves. For example, you update_h_params method could be re-written as

from typing import Dict, Any

def update_h_params(self, h_params: Dict[str, Any]):

Using this feature helps many editors and static type checking tools analyze your program more appropriately. Both PyCharm and mypy fully support these annotations.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.