Particle Swarm Optimization in Python: Full Implementation and Results

PSO Formulas

v_i and x_i denote the velocity and position of particle i in each dimension, pbest_i the best position found by the particle, and gbest the best position found by the whole swarm.

Code

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['font.family'] = ['SimHei']  # set default font
mpl.rcParams['axes.unicode_minus'] = False  # fix minus sign display

def fitness_func(X):
    """Calculate particle fitness (objective function). X shape is size*2."""
    A = 10
    pi = np.pi
    x = X[:, 0]
    y = X[:, 1]
    return 2 * A + x ** 2 - A * np.cos(2 * pi * x) + y ** 2 - A * np.cos(2 * pi * y)

def velocity_update(V, X, pbest, gbest, c1, c2, w, max_val):
    """
    Update particle velocities according to the PSO formula.
    V: current velocity matrix (size*2)
    X: current position matrix (size*2)
    pbest: best positions of each particle (size*2)
    gbest: best position of the swarm (1*2)
    """
    size = X.shape[0]
    r1 = np.random.random((size, 1))
    r2 = np.random.random((size, 1))
    V = w * V + c1 * r1 * (pbest - X) + c2 * r2 * (gbest - X)
    V[V < -max_val] = -max_val
    V[V > max_val] = max_val
    return V

def position_update(X, V):
    """
    Update particle positions according to the PSO formula.
    X: current position matrix (size*2)
    V: current velocity matrix (size*2)
    """
    return X + V

def pso():
    # PSO parameters
    w = 1          # inertia weight
    c1 = 2         # cognitive coefficient
    c2 = 2         # social coefficient
    dim = 2
    size = 20
    iter_num = 1000
    max_val = 0.5
    best_fitness = float(9e10)

    # initialize swarm
    X = np.random.uniform(-5, 5, size=(size, dim))
    V = np.random.uniform(-0.5, 0.5, size=(size, dim))
    p_fitness = fitness_func(X)
    g_fitness = p_fitness.min()
    fitneess_value_list = [g_fitness]
    pbest = X
    gbest = X[p_fitness.argmin()]

    for i in range(1, iter_num):
        V = velocity_update(V, X, pbest, gbest, c1, c2, w, max_val)
        X = position_update(X, V)
        p_fitness2 = fitness_func(X)
        g_fitness2 = p_fitness2.min()
        for j in range(size):
            if p_fitness[j] > p_fitness2[j]:
                pbest[j] = X[j]
                p_fitness[j] = p_fitness2[j]
        if g_fitness > g_fitness2:
            gbest = X[p_fitness2.argmin()]
            g_fitness = g_fitness2
        fitneess_value_list.append(g_fitness)

    print("Best value: %.5f" % fitneess_value_list[-1])
    print(f"Best solution: x={round(gbest[0],5)}, y={round(gbest[1],5)}")
    plt.plot(fitneess_value_list, color='r')
    plt.title('Iteration Process')
    plt.show()

Result: Best value is 0.00002, best solution is x = -0.00029, y = -0.00019.

Reference - “Python Optimization Algorithms in Practice” (Su Zhenyu).