Experimental implementation of simple Particle Swarm Optimization algorithm

pso.py

import numpy as np
from numpy.random import rand, uniform

# reference: https://en.wikipedia.org/wiki/Particle_swarm_optimization#Parameter_selection

ob = lambda x: (np.power(x, 4.0) - x * x * 16 + x * 5).sum() / 2.0 # Styblinski-Tang function
lims = np.array([[-5.0, 4.0], [-3.0, 4.0]])

loop = 100
nps, ndim = 10, lims.shape[0]
h1, h2, h3 = 0.9, 0.9, 0.9

gx = ob(rand(1, ndim))
lxs = X = rand(nps, ndim)
vs = np.array([[uniform(-abs(u - l), abs(u - l)) for l, u in lims] for i in range(nps)])

for c in range(loop):
    for i in range(nps):
        x, v = X[i, :], vs[i, :]
        vs[i, :] = v = v * h1 + h2 * rand(ndim) * (lxs[i] - x) + h3 * rand(ndim) * (gx - x)
        X[i, :] = x = x + v
        if ob(x) < ob(lxs[i]):
            lxs[i] = x
        if ob(x) < ob(gx):
            gx = x
            print(c, ob(gx))