Self-organizing map online algorithm

som.py

import numpy as np
import matplotlib.pylab as plt

import torch
from torch import nn

from sklearn.metrics.pairwise import euclidean_distances

from som_measures import topographic_error, quantization_error


def quantization_error(som=None, x=None, d=None):
    if d is None:
        if som is None or x is None:
            raise ValueError('')
        else:
            d = euclidean_distances(x, som)
    return np.mean(np.min(d, axis=1))


def topographic_error(dist_fun, som=None, x=None, d=None):
    if d is None:
        if som is None or x is None:
            raise ValueError('')
        else:
            d = euclidean_distances(x, som)

    bmus = np.argsort(d, axis=1)[:, :2]
    return np.mean(np.array([dist_fun(bmu[0], bmu[1]) > 1 for bmu in bmus]))


class SOM(nn.Module):
    def __init__(self,
                 n_units=16,
                 dim=2,
                 iters=1000,
                 lrate=(0.5, 0.01),
                 sigma=(0.5, 0.01)):
        super().__init__()
        self.n_units = n_units
        self.dim = dim
        self.iters = iters
        self.lrate_i = lrate[0]
        self.lrate_f = lrate[1]
        self.sigma_i = sigma[0]
        self.sigma_f = sigma[1]

        self.initCodebook()
        self.initLatice()

        self.t = torch.linspace(0, 1, steps=iters)
        self.lrate = self.lrate_i * (self.lrate_f / self.lrate_i)**self.t
        self.sigma = self.sigma_i * (self.sigma_f / self.sigma_i)**self.t

    def initCodebook(self):
        self.codebook = torch.ones([self.n_units*self.n_units, self.dim])
        nn.init.uniform_(self.codebook, a=0, b=0.01)

    def initLatice(self):
        line = torch.linspace(0, 1, steps=self.n_units)
        grid_x, grid_y = torch.meshgrid(line, line, indexing="ij")
        p = torch.cat((grid_x.flatten().reshape(-1, 1),
                       grid_y.flatten().reshape(-1, 1)), dim=1)
        self.dist = torch.cdist(p, p)

    def fit(self, X):
        for t in range(self.iters):
            x = X[np.random.randint(X.shape[0])]

            BMU = torch.argmin((((self.codebook - x))**2).sum(dim=-1))

            h = torch.exp(-(self.dist[BMU]/self.sigma[t])**2).unsqueeze(1)

            self.codebook -= self.lrate[t] * h * (self.codebook - x)

    def drawMap(self, X=None, ax=None, xy_lim=(0, 1)):
        wx = self.codebook[:, 0].reshape(self.n_units, self.n_units).numpy()
        wy = self.codebook[:, 1].reshape(self.n_units, self.n_units).numpy()
        if ax is None:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        ax.scatter(wx, wy, s=50, c='w', edgecolors='k', zorder=3)
        for i in range(self.n_units):
            ax.plot(wx[i, :], wy[i, :], 'k', alpha=0.85, lw=1.5, zorder=2)
            ax.plot(wx[:, i], wy[:, i], 'k', alpha=0.85, lw=1.5, zorder=2)
            ax.set_xticks([xy_lim[0], xy_lim[1]])
            ax.set_xticklabels([str(xy_lim[0]), str(xy_lim[1])],
                               fontsize=14,
                               weight='bold')
            ax.set_yticks([xy_lim[0], xy_lim[1]])
            ax.set_yticklabels([str(xy_lim[0]), str(xy_lim[1])],
                               fontsize=14,
                               weight='bold')

        if X is not None:
            ax.scatter(X[:, 0], X[:, 1],
                       s=35,
                       c='r',
                       edgecolor='r',
                       alpha=0.15,
                       zorder=0)

    def drawParameters(self):
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(self.lrate)
        ax.plot(self.sigma)


if __name__ == "__main__":
    X = torch.from_numpy(np.random.uniform(0, 1, (1000, 2)))

    fig = plt.figure(figsize=(11, 5))
    ax1 = fig.add_subplot(121)

    som = SOM(iters=3000)
    som.fit(X)
    som.drawMap(X=X, ax=ax1)
    ax1.text(0, 1.06, "A",
             ha="left",
             fontsize=18,
             weight="bold")

    def dist(k: int, l: int) -> int:
        return abs(k // 16 - l // 16) + abs(k % 16 - l % 16)

    print(topographic_error(dist, som=som.codebook, x=X))
    print(quantization_error(som=som.codebook, x=X))

    T = np.random.uniform(0, 2*np.pi, 1000)
    R = np.sqrt(np.random.uniform(0.5**2, 1, len(T)))
    X = np.c_[R * np.cos(T), R * np.sin(T)]

    ax2 = fig.add_subplot(122)

    som.fit(X)
    som.drawMap(X=X, ax=ax2, xy_lim=(-1, 1))
    ax2.text(-1, 1.11, "B",
             ha="left",
             fontsize=18,
             weight="bold")

    print(topographic_error(dist, som=som.codebook, x=X))
    print(quantization_error(som=som.codebook, x=X))

    plt.savefig("som_maps.png")
    plt.show()