Particle Swarm Optimization (PSO) in Python

This Python implementation of Particle Swarm Optimization (PSO) uses NumPy and a small problem dictionary to define the objective, dimension, and bounds. It follows the structure of the MATLAB implementation, while keeping the optimizer in a reusable Python function.

PSO was introduced by James Kennedy and Russell Eberhart in 1995. Each particle stores a position, a velocity, and its personal best. The swarm shares a global best, and these two memories guide the next move.

How the Python code works

The code follows the standard PSO update rules. For particle $i$, with position $x_i$, velocity $v_i$, personal best $p_i$ and global best $g$:

Here $r_1$ and $r_2$ are vectors of uniform random numbers in $[0, 1]$, drawn again for every particle in every iteration. After each iteration, the inertia weight is multiplied by a damping ratio, $w \leftarrow w_{\text{damp}}\,w$. The meaning of each term and the choice of the coefficients are explained in the MATLAB post.

The main steps of the PSO function are:

1. Create the swarm. Each particle gets a random position drawn uniformly between the lower and upper bounds, a zero velocity, and its cost. Its personal best starts at its first position.
2. Set the global best to the best of these first positions.
3. In each iteration, update the velocity and the position of every particle with the rules above.
4. Clip the new position to the bounds, so the particle stays inside the search space.
5. Evaluate the cost. If it is lower than the personal best, update the personal best, and then check it against the global best.
6. Multiply $w$ by $w_{\text{damp}}$ and print the best cost of the iteration.

The code clips positions, but it does not limit the size of the velocity. The MATLAB version adds a velocity limit.

What is in the download

• pso/__init__.py: the PSO function, and two small helpers, tic and toc, that measure the run time.
• app.py: a sample script that defines a problem and runs PSO on it.

The only library the code needs is NumPy.

Defining a problem

A problem is a Python dictionary with four keys: CostFunction, nVar (the number of decision variables), VarMin and VarMax. The bounds can be single numbers, or NumPy arrays with nVar elements when each variable has its own range. The sample script minimizes the sphere function, $f(x) = \sum_{i=1}^{n} x_i^2$, in 10 dimensions:

# A Sample Cost Function
def Sphere(x):
    return sum(x**2)

# Define Optimization Problem
problem = {
        'CostFunction': Sphere, 
        'nVar': 10, 
        'VarMin': -5,   # Alternatively you can use a "numpy array" with nVar elements, instead of scalar
        'VarMax': 5,    # Alternatively you can use a "numpy array" with nVar elements, instead of scalar
    }

The problem is a minimization. To maximize a function, minimize its negative.

Parameters

The PSO function takes the problem and six optional parameters. Its defaults are the constriction coefficients of Clerc and Kennedy.

The sample script uses other values: 200 iterations, 50 particles, $c_1 = 1.5$, $c_2 = 2$, and an inertia weight that starts at 1 and is damped by 0.995 in every iteration.

How to run the code

Run app.py from the folder that contains it, for example with python app.py. The script imports the pso package from the same folder, prints the best cost of every iteration and the elapsed time, and then prints the global best.

The function returns two values: gbest, a dictionary with the best position and its cost, and pop, the final swarm. To solve your own problem, write a cost function that takes a NumPy array and returns a number, and put it in the problem dictionary.

Check the result

The supplied Sphere problem has a minimum of 0 at the origin. Inspect both the position and cost in gbest, and confirm that every coordinate lies inside its bounds. Repeat runs when comparing parameters because the initial swarm and updates are random. For a new objective, check the NumPy input shape and scalar return value before increasing the swarm size or iteration limit.

Applications

PSO is a general method for continuous optimization. It needs only the values of the cost function, not its gradient. Common uses include tuning the parameters of models and controllers, engineering design, and training or tuning machine learning models. Several application posts on Yarpiz use PSO, such as optimal inventory control and robot path planning.

Learn more

A free tutorial on PSO and its implementation is also available. It explains the theory and builds the MATLAB code line by line, and the Python code here follows the same structure.