Particle Swarm Optimization: runnable experiments

Companion notebook to the PSO for filter design topic. We watch a swarm escape a local minimum that traps gradient descent, design a biquad with PSO, and look honestly at run-to-run variability. Everything imports from the clean module pso.py.

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal

from pso import pso, gradient_descent, design_iir_pso, iir_is_stable

def rastrigin(x):
    x = np.asarray(x)
    return 10 * len(x) + np.sum(x**2 - 10 * np.cos(2 * np.pi * x))

1. Swarm vs gradient on a rugged surface

The Rastrigin function has one global minimum at the origin buried in a lattice of local minima. Gradient descent from a poor start stalls; the swarm does not.

bounds = [(-5.12, 5.12)] * 2
gbest, gf, hist_pso = pso(rastrigin, bounds, n_particles=40, n_iters=150,
                          rng=np.random.default_rng(3))
x_gd, f_gd, hist_gd = gradient_descent(rastrigin, [4.0, -3.0], bounds, n_iters=150)
print('PSO  -> error %.3f at %s' % (gf, np.round(gbest, 2)))
print('GD   -> error %.3f at %s (trapped)' % (f_gd, np.round(x_gd, 2)))

g = np.linspace(-5.12, 5.12, 300)
X, Y = np.meshgrid(g, g)
Z = 20 + X**2 - 10*np.cos(2*np.pi*X) + Y**2 - 10*np.cos(2*np.pi*Y)
fig, ax = plt.subplots(figsize=(6, 5))
ax.contourf(X, Y, Z, levels=30, cmap='viridis')
ax.plot(*x_gd, 'X', color='red', ms=13, label='gradient descent')
ax.plot(*gbest, '*', color='cyan', ms=18, label='PSO')
ax.set_title('Rastrigin: swarm escapes, gradient stalls'); ax.legend(fontsize=8)
fig.tight_layout(); plt.show()

plt.figure(figsize=(9, 3.4))
plt.plot(hist_pso, label='PSO', lw=2)
plt.plot(hist_gd, label='gradient descent', lw=2)
plt.xlabel('iteration'); plt.ylabel('best error so far')
plt.title('Convergence'); plt.legend(); plt.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

2. Designing a biquad with PSO

A particle is a coefficient vector; the fitness is magnitude-response error against a target, with a stability wall. Here we match a Butterworth bandpass.

fs = 8000
freqs = np.linspace(0, fs/2, 200)
sos_target = signal.butter(2, [800, 1600], btype='band', fs=fs, output='sos')
_, h_target = signal.sosfreqz(sos_target, worN=freqs, fs=fs)
target = np.abs(h_target)

sos, err, hist = design_iir_pso(freqs, target, fs, rng=np.random.default_rng(2))
_, h_pso = signal.sosfreqz(sos, worN=freqs, fs=fs)
print('MSE %.2e, stable=%s' % (err, iir_is_stable(sos[0, 4], sos[0, 5])))

plt.figure(figsize=(9, 3.6))
plt.plot(freqs, target, 'k--', lw=2, label='target')
plt.plot(freqs, np.abs(h_pso), color='C1', lw=1.8, label='PSO design')
plt.xlabel('Frequency [Hz]'); plt.ylabel('|H(f)|')
plt.title('PSO-designed bandpass biquad'); plt.legend(); plt.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

3. Honesty: PSO is stochastic

A single run can be lucky or unlucky. The right way to report a stochastic optimiser is a distribution over seeds, not one cherry-picked result.

finals = []
for seed in range(30):
    _, f, _ = pso(rastrigin, bounds, n_particles=40, n_iters=120,
                  rng=np.random.default_rng(seed))
    finals.append(f)
finals = np.array(finals)
print('over 30 seeds: best %.3f, median %.3f, worst %.3f' %
      (finals.min(), np.median(finals), finals.max()))

plt.figure(figsize=(8, 3.2))
plt.hist(finals, bins=15, color='C0', alpha=0.8)
plt.xlabel('final error'); plt.ylabel('count')
plt.title('Distribution of PSO outcomes over 30 seeds')
plt.tight_layout(); plt.show()

Most runs reach the global basin, but not all, which is exactly why PSO carries no convergence guarantee. See the embedded page for running the swarm on a microcontroller for online adaptation.