import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

# 生成X和Y的數(shù)據(jù)
X=np.arange(-5,5,0.1)
Y=np.arange(-5,5,0.1)
X,Y=np.meshgrid(X,Y)

# 目標(biāo)函數(shù)
Z=X**2+Y**2+X

# 繪圖
fig=plt.figure()
ax=Axes3D(fig)
surf=ax.plot_surface(X,Y,Z,cmap=cm.coolwarm)
plt.show()

# 速度
# Vi+1 = w*Vi + c1 * r1 * (pbest_i - Xi) + c2 * r2 * (gbest_i - Xi)
# 位置
# Xi+1 = Xi + Vi+1
# vi, xi 分別表示粒子第i維的速度和位置
# pbest_i, gbest_i 分別表示某個(gè)粒子最好位置第i維的值、整個(gè)種群最好位置第i維的值

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默認(rèn)字體
mpl.rcParams['axes.unicode_minus'] = False  # 解決保存圖像是負(fù)號(hào)'-'顯示為方塊的問題

def fitness_func(X):
    """計(jì)算粒子的的適應(yīng)度值，也就是目標(biāo)函數(shù)值，X 的維度是 size * 2 """
    A = 10
    pi = np.pi
    x = X[:, 0]
    y = X[:, 1]
    return x**2+y**2+x

def velocity_update(V, X, pbest, gbest, c1, c2, w, max_val):
    """
    根據(jù)速度更新公式更新每個(gè)粒子的速度
    :param V: 粒子當(dāng)前的速度矩陣，20*2 的矩陣
    :param X: 粒子當(dāng)前的位置矩陣，20*2 的矩陣
    :param pbest: 每個(gè)粒子歷史最優(yōu)位置，20*2 的矩陣
    :param gbest: 種群歷史最優(yōu)位置，1*2 的矩陣
    """
    size = X.shape[0]
    r1 = np.random.random((size, 1))
    r2 = np.random.random((size, 1))
    V = w*V+c1*r1*(pbest-X)+c2*r2*(gbest-X)
    # 防止越界處理
    V[V  -max_val] = -max_val
    V[V > -max_val] = max_val
    return V

def position_update(X, V):
    """
    根據(jù)公式更新粒子的位置
    :param X: 粒子當(dāng)前的位置矩陣，維度是 20*2
    :param V: 粒子當(dāng)前的速度舉著，維度是 20*2
    """
    return X+V

def pos():
    w = 1
    c1 = 2
    c2 = 2
    r1 = None
    r2 = None
    dim = 2
    size = 20
    iter_num = 1000
    max_val = 0.5
    best_fitness = float(9e10)
    fitness_val_list = []
    # 初始化種群各個(gè)粒子的位置
    X = np.random.uniform(-5, 5, size=(size, dim))
    # 初始化各個(gè)粒子的速度
    V = np.random.uniform(-0.5, 0.5, size=(size, dim))
    # print(X)
    p_fitness = fitness_func(X)
    g_fitness = p_fitness.min()
    fitness_val_list.append(g_fitness)

    # 初始化的個(gè)體最優(yōu)位置和種群最優(yōu)位置
    pbest = X
    gbest = X[p_fitness.argmin()]
    # 迭代計(jì)算
    for i in range(1, iter_num):
        V = velocity_update(V, X, pbest, gbest, c1, c2, w, max_val)
        X = position_update(X, V)
        p_fitness2 = fitness_func(X)
        g_fitness2 = p_fitness2.min()

        # 更新每個(gè)粒子的歷史最優(yōu)位置
        for j in range(size):
            if p_fitness[j] > p_fitness2[j]:
                pbest[j] = X[j]
                p_fitness[j] = p_fitness2[j]
            # 更新群體的最優(yōu)位置
            if g_fitness > g_fitness2:
                gbest = X[p_fitness2.argmin()]
                g_fitness = g_fitness2
            # 記錄最優(yōu)迭代記錄
            fitness_val_list.append(g_fitness)
            i += 1

    # 輸出迭代結(jié)果
    print("最優(yōu)值是：%.5f" % fitness_val_list[-1])
    print("最優(yōu)解是：x=%.5f,y=%.5f" % (gbest[0], gbest[1]))

    # 繪圖
    plt.plot(fitness_val_list, color='r')
    plt.title('迭代過程')
    plt.show()

pos()