本文基于 优化算法笔记（三）粒子群算法（1） - 简书 (jianshu.com) 进行实现，建议先看原理。

输出结果如下

PSO.gif

实现代码如下

# 粒子群算法
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from PIL import Image
import shutil
import os
import glob 

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def plot_jpg(start, end, point_g, point_best, err, m, n, lower, upper, W, title):
    plt.figure(figsize=(8, 12))
    gs = gridspec.GridSpec(3, 2)
    
    ax1 = plt.subplot(gs[:2, :2])
    ax1.scatter(start[:, 0], start[:, 1], alpha=0.3, color='green', s=20, label='当前位置')  # 当前位置
    ax1.scatter(point_g[0], point_g[1], alpha=1, color='blue', s=20, label='当前最优点')  # 全局最优点
    ax1.scatter(point_best[0], point_best[1], alpha=0.3, color='red', label='目标点')  # 最优点

    for i in range(n):
        ax1.text(start[i][0]-5, start[i][1], f'{i}', alpha=0.3, fontsize=10, color='red')
        ax1.plot([start[i][0], end[i][0]], [start[i][1], end[i][1]], alpha=0.3, color='gray')
    
    ax1.grid(True, color='gray', linestyle='-.', linewidth=0.5)
    ax1.set_xlim(lower[0]*1.2, upper[0]*1.2)
    ax1.set_ylim(lower[1]*1.2, upper[1]*1.2)
    ax1.set_xlabel(f'iter:{m}  W: {W:.8f}  dist: {err[-1]:.8f}'  )
    ax1.set_title(title)
    ax1.legend(loc='lower right', bbox_to_anchor=(1, 0), ncol=1)
    
    ax2 = plt.subplot(gs[2, :])
    ax2.plot(range(len(err)), err, marker='o', markersize=5)
    ax2.grid(True, color='gray', linestyle='-.', linewidth=0.5)
    ax2.set_xlim(0, max_iter)
    ax2.set_ylim(0, np.ceil(max(err)))
    ax2.set_xticks(range(0, max_iter, 5))
    
    plt.savefig(rf'./tmp/tmp_{m:04}.png')
    plt.close()


# 目标函数
def target(point):
    return (point[0]-a)**2 + (point[1]-b)**2

# 速度限制
def limit_speed(speed, maxV):
    rate = ((speed[:, 0]**2 + speed[:, 1]**2)**0.5)/maxV
    rate = np.where(rate > 1, rate, 1)
    for t in range(d):
        speed[:, t] = speed[:, t]/rate
    return speed

def PSO(C1, C2, W, maxV):
    # 初始化粒子位置
    start = np.random.random(size=(n, d))
    for _ in range(d):
        start[:, _] = start[:, _]*(upper_lim[_]-lower_lim[_])+lower_lim[_]
    # 初始化粒子速度
    speed = np.random.random(size=(n, d))
    for _ in range(d):
        speed[:, _] = speed[:, _]*(upper_lim[_]-lower_lim[_])+lower_lim[_]
    # 速度限制
    speed = limit_speed(speed, maxV)  

    if os.path.exists(tmp_path):
        shutil.rmtree(tmp_path) 
    os.makedirs(tmp_path, exist_ok=True)

    errors = [target(min(start.tolist(), key=target))**0.5]
    for _ in range(max_iter):
        # 下次目标位置
        r1, r2 = np.random.random(2)
        # 全局最优点
        point_g = min(start.tolist(), key=target)
        # 各粒子历史最优点
        if _ == 0:
            point_p = start.copy()
        else:
            point_p = np.array([min([point_p[i], end[i]], key=target) for i in range(n)])
        # 更新速度
        speed = W*speed + r1*C1*(point_p - start) + r2*C2*(point_g - start)
        speed = limit_speed(speed, maxV)  # 速度限制 
        # 更新位置
        end = start + speed
        
        title = f'PSO\nC1:{C1} C2:{C2} n:{n} maxV:{maxV} max_iter:{max_iter}'
        plot_jpg(start, end, point_g, point_best, errors,  _, n, lower_lim, upper_lim, W, title)
        # 更新开始位置
        start = end
        errors.append(target(min(start.tolist(), key=target))**0.5)
        # 惯性衰减
        # W *= 0.8
        W -= step
    plot_jpg(start, end, point_g, point_best, errors,  max_iter, n, lower_lim, upper_lim, W, title)
    return errors


C1 = 2  # C1:自我学习因子
C2 = 2  # C2:全局学习因子
W = 0.5  # W:惯性系数
maxV = 20  # 最大速率

n = 10  # 粒子数量
d = 2  # 粒子维度
max_iter = 50  # 迭代次数

step = W/(max_iter)  # 惯性系数衰减

# 搜索区间 
lower_lim = [-100, -100]
upper_lim = [100, 100]

# 目标点
a, b = 0, 0
point_best = (a, b)

# 临时文件路径
tmp_path = r'./tmp/'

err = PSO(C1, C2, W, maxV)

# png 转 gif
images = [Image.open(png) for png in glob.glob(os.path.join(tmp_path, '*.png'))]
im = images.pop(0)
im.save(r"./PSO.gif", save_all=True, append_images=images, duration=500)

im = Image.open(r"./PSO.gif")
im.show()
im.close()