python 粒子群算法实现

import matplotlib.pyplot as plt
from random import uniform
from random import random
from numpy import array, sin, cos, sqrt

class Particle(object):
    def __init__(self, position):
        self.x = position
        self.v = array([0., 0.])
        self.best = position

def fit(p):
    (x, y) = p
    return (x**2+y**2)/32-cos(x)-cos(y)

def PSO():
    c1 = 2
    c2 = 2
    N = 100
    # 初始化
    particle = [Particle(array([uniform(-10, 10), uniform(-10, 10)]))
                for i in range(N)]
    # 初始化全局最优位置
    gbest = particle[0].x
    for p in particle:
        if fit(p.x) < fit(gbest):
            gbest = p.x
    plt.ion()
    for t in range(100):
        plt.clf()
        # 对于所有粒子
        for p in particle:
            # 更新速度和位置
            p.v += c1*random()*(p.best-p.x)+c2*random()*(gbest-p.x)
            v = sqrt(p.v[0]**2+p.v[1]**2)
            if v > 1:
                p.v /= v
            p.x += p.v
            # 更新最优位置
            if fit(p.x) < fit(p.best):
                p.best = p.x
            if fit(p.best) < fit(gbest):
                gbest = p.best
        plt.scatter(*array([p.x for p in particle]).T)
        plt.xlim([-10, 10])
        plt.ylim([-10, 10])
        plt.pause(0.2)
    plt.ioff()
    plt.show()

PSO()