粒子群优化算法-PSO

import numpy as np
import matplotlib.pyplot as plt

N = 10
dim = 2
iter_times = 100

def function(vec):
    return vec[0]**2+vec[1]**2-5

def initSwarm(N,dim):
    x = np.random.random((N,dim))
    v = np.random.random((N,dim))
    pbest = x
    fitness = [function(vec) for vec in x]
    gbest = pbest[np.argmin(fitness)]
    best = fitness[np.argmin(fitness)]
    return x,v,pbest,fitness,gbest,best

x,v,pbest,fitness,gbest,best = initSwarm(N,dim)

def searchBestParam(iter_times):
    result = []
    for time in range(iter_times):
        for index in range(N):
            k = function(x[index])
            if k < fitness[index]:
                fitness[index] = k
                pbest[index] = x[index]
        best = fitness[np.argmin(fitness)]
        gbest = x[np.argmin(fitness)]
        result.append(best)
        
        for index in range(N):
            v[index] = 0.8*v[index]+2*np.random.rand()*(pbest[index]-x[index])+2*np.random.rand()*(gbest-x[index])
            x[index] = x[index] + v[index]
    print('自变量:',gbest)
    print('因变量:',best)

searchBestParam(iter_times)