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

<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()
</code>

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

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