An implementation of Gaussian Mean Shift Procedure

gaussian_mean_shift.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
import matplotlib.animation as animation

def gaussian_kernel(x, sigma):
    return 1 / (np.sqrt(2*np.pi)*sigma) * np.exp(-(x**2)/(2*(sigma**2)))

def x_update(x, xi, sigma):
    return np.sum(gaussian_kernel(xi - x, sigma) * x) / np.sum(gaussian_kernel(xi - x, sigma))

def gaussian_mean_shift(x, max_iter=10000):
    k = gaussian_kde(x)
    sigma = k.factor * x.std(ddof=1)

    x_ = np.copy(x)
    l = x.shape[0]
    history = []
    for _ in range(max_iter):
        x_old = np.copy(x_)
        history.append(x_old) # solution history
        for xi in range(l):
            x_[xi] = x_update(x, x_[xi], sigma)
        if np.sum(np.abs(x_ - x_old)) < 1e-10:
            break
    history = np.asarray(history)
    return k, x_, history

data = np.array([])
tmp = np.random.normal(0.0, 1.0, size=100)
data = np.append(data, tmp)
tmp = np.random.normal(5.0, 1.0, size=200)
data = np.append(data, tmp)

k, max_x, x_history = gaussian_mean_shift(data)

N = 1000 # number of sample
X = np.linspace(-5, 10, N)
fig, ax1 = plt.subplots()
ax1.hist(data, bins=50, label='data')
plt.legend(loc='upper left')
ax2 = ax1.twinx()
ax2.plot(X, k(X), c='red', label='kde')
scat = ax2.scatter(x_history[0], k(x_history[0]), c='black', label='solution')
plt.legend(loc='upper right')

frames = x_history.shape[0]

def update(i):
    ii = i + 1
    x = x_history[ii]
    y = k(x_history[ii])
    scat.set_offsets(np.c_[x, y])

ani = animation.FuncAnimation(fig, update, frames=frames-1, interval=200)
plt.show()
# ani.save("mean_shift.gif", writer = 'imagemagick')