custom_pso.py

'''
    --------------------------------------------------------------
    Custom PSO algorithm for continuous domains
    --------------------------------------------------------------
    Autor: Rodrigo Araya 
    Email: [email protected]
'''

import math
import numpy as np
import pandas as pd
from F_newton import F_newton
from F_coulomb import F_coulomb
import plotly.graph_objects as go
import plotly.express as px


def Random_(problem, x_, auto_reduce_search_space=False):
    """ This solver generates a random solution for each particle. Additionally, it is also possible to 
        apply an exploitation method where the search space is reduced on each iteration.

    Args:
        problem (dic): Dictionary that include: objective function, lower bound and upper bound of each variable and optimal solution
        x_ (array(size=())): Current position of each particle
        auto_reduce_search_space (bool, optional): If True, it update the lower and upper bounds in a way to reduce the search space. Defaults to False.

    Returns:
        Stemp, S_f: New positions for each particle and their performance after being evaluated with the objective function.
    """
    
    
    nPart = x_.shape[0]
    nVar = x_.shape[1]
    
    if auto_reduce_search_space:
        Lo = np.min(x_, axis=0)
        Up = np.max(x_, axis=0)
    else:
        Lo = np.zeros(nVar)
        Up = np.ones(nVar)

    new_x = np.random.uniform(low=Lo, high=Up, size=(nPart, nVar))
    
    Stemp = np.zeros((nPart,nVar))
    for k in range(nPart):
        for i in range(nVar):
            Stemp[k][i] = new_x[k][i]
            if Stemp[k][i] > Up[i]:
                Stemp[k][i] = Up[i]
            elif Stemp[k][i] < Lo[i]:
                Stemp[k][i] = Lo[i]
    f,S_r,maximize = mp_evaluator(problem, Stemp)

    S_f = np.zeros((nPart,1))
    
    for i in range(len(S_r)):
        S_f[i] = f[i]
    return Stemp, S_f
    
def F_newton_(problem, x_, m, Lo, Up, G=0.00001, keep_best=True, weight_method='Linear', t=1.0, max_mass=100.0):
    _, _, new_x = F_newton(x_, m, minimize=True, G=G, keep_best=keep_best, weight_method=weight_method, t=t, max_mass=max_mass)
    nPart=x_.shape[0]
    nVar=x_.shape[1]
    Stemp = np.zeros((nPart,nVar))
    for k in range(nPart):
        for i in range(nVar):
            Stemp[k][i] = new_x[k][i]
            if Stemp[k][i] > Up[i]:
                Stemp[k][i] = Up[i]
            elif Stemp[k][i] < Lo[i]:
                Stemp[k][i] = Lo[i]
    f,S_r,maximize = mp_evaluator(problem, Stemp)

    S_f = np.zeros((nPart,1))
    
    for i in range(len(S_r)):
        S_f[i] = f[i]
    return Stemp, S_f
    
def F_coulomb_(problem, x_, v, q, P, Lo, Up, k=0.00001, weight_method='Linear', t=1.0, max_q=100.0):
    _, _, v, new_x = F_coulomb(x_, v, q, P=P, minimize=True, k=k, weight_method=weight_method, t=t, max_q=max_q)
    nPart=x_.shape[0]
    nVar=x_.shape[1]
    Stemp = np.zeros((nPart,nVar))
    for k in range(nPart):
        for i in range(nVar):
            Stemp[k][i] = new_x[k][i]
            if Stemp[k][i] > Up[i]:
                Stemp[k][i] = Up[i]
            elif Stemp[k][i] < Lo[i]:
                Stemp[k][i] = Lo[i]
    f,S_r,maximize = mp_evaluator(problem, Stemp)

    S_f = np.zeros((nPart,1))
    
    for i in range(len(S_r)):
        S_f[i] = f[i]
    return Stemp, v, S_f

def evaluator(problem, x):
    # Dado que se trabaja bajo el dominio [0, 1] para cada variable, se debe crear una funcion que regrese a los valores a los dominios originales
    # Ejemplo (ndim=2): si se tiene el dominio original [-5, 5] y [-2, 2] para la variables x1, x2 y se cambio a [0, 1] y [0, 1] se vuelve al original tras:
    # y1=ax+b -> -5 = a*0+b -> b=-5
    # y1=ax+b -> 5 = a*1+-5 -> a=10
    # y1=10*x+-5
    # y2=ax+b -> -2 = a*0+b -> b=-2
    # y2=ax+b -> 2 = a*1+-2 -> a=4
    # y2=4*x+-2
    # Luego se aplica y1(x1) e y2(x2) respectivamente. Extendible a n dimensiones.
    # Dado que [0, 1] no cambia, se generaliza la formula b=lb y a=ub-lb -> y_{i} = (ub-lb)_{i}*x + lb_{i}
    lb = [var[0] for var in problem['bounds']]
    ub = [var[1] for var in problem['bounds']]
    x = [(ub[ind]-lb[ind])*i+lb[ind] for ind, i in enumerate(x)]
    # calculate fitness
    f = problem['f'](x)
    fitness = dict(Obj=f)
    return fitness

def mp_evaluator(problem, x):
    results = [evaluator(problem, c) for c in x]
    f = [r['Obj'] for r in results]
    # maximization or minimization problem
    maximize = False
    return (f, [r for r in results],maximize)

def correct_x(problem, x_):
    lb = [var[0] for var in problem['bounds']]
    ub = [var[1] for var in problem['bounds']]
    x = np.empty(x_.shape)
    for ind, row in enumerate(x_):
        corr_row = np.array([(ub[ind]-lb[ind])*i+lb[ind] for ind, i in enumerate(row)])
        x[ind]=corr_row
    return x



def custom_pso(problem, nPart, nVar, maxiter, G=0.00001, k=0.0001, keep_best=True, weight_method='Linear', t=1.0, seed=0, max_mass=100.0, solver='Newton', P=3, max_q=100.0, auto_reduce_search_space=False, dinamic=False):
    """The procedure of this algorithm is as follows:
    1. Generate a file of size nSize=nPart.
    2. Generate nPart initial solutions. Each solution contains nVar variables.
    3. Evaluate the nPart solutions and add them to the file ordered from best to worst solution.
    4. While the termination condition is not met (iterations > maxiter):
        4.1 Generate new solutions using some solver. Options = ['Random', 'Newton', 'Coulomb']
        4.2 Add new solutions to the file.
        4.3 Remove duplicate solutions.
        4.4 Evaluate solutions and sort from best to worst solution.
        4.5 Keep in the file the nPar best solutions.
        4.6 Save iteration results in a history (dataframe).
        4.7 Evaluate termination condition. If negative, return to step 4.1.
    
    Args:
        problem (dic): Dictionary that include: objective function, lower bound and upper bound of each variable and optimal solution
        nPart (_type_): Quantity of particles
        nVar (_type_): Quantity of variables
        maxiter (_type_): Maximum number of iterations.
        seed (int, optional): set the generation of random numbers. Defaults to 0.
        solver (str, optional): solver to apply. Defaults to 'Newton'. Options=['Random', 'Newton', 'Coulomb'].
        
        Random Solver:
            auto_reduce_search_space (bool, optional): If True, it update the lower and upper bounds in a way to reduce the search space. Defaults to False.

        Newton Solver:
            G (float, optional): Gravitational constant. Defaults to 0.00001.
            keep_best (bool, optional): It will keep the best value obtained in each iteration. Defaults to True.
            weight_method (str, optional): method to reassign particle mass. Defaults to 'Linear'. Options=['Linear', 'Quadratic', 'Exponential'].
            t (float, optional): time. Defaults to 1.0.
            max_mass (float, optional): upper bound of mass to assing. Defaults to 100.0.
            dinamic (bool, optional): It will change the max_mass value depending on the current iteration and difference between best and worst value obteined. Defaults to False.
            
        Coulomb Solver:
            P (int): quantity of electric fields (the best P particles become electric fields)
            max_q (float, optional): upper bound of electric charge to assing. Defaults to 100.0.
            weight_method (str, optional): method to reassign electric charges. Defaults to 'Linear'. Options=['Linear', 'Quadratic', 'Exponential'].
            k (float, optional): electric constant. Defaults to 0.0001.
            t (float, optional): time. Defaults to 1.0.
            dinamic (bool, optional): It will change the max_q value depending on the current iteration and difference between best and worst value obteined. Defaults to False.
            
        
    Returns:
        df, best_var, best_sol: dataframe contening data from all iterations, the best variables and the best value obtained
    """
    
    # number of variables
    parameters_v = [f'x{i+1}' for i in range(nVar)]
    
    # number of variables
    nVar = len(parameters_v)

    # size of solution archive
    nSize = nPart

    # number of Particules
    nPart = nPart

    # maximum iterations
    maxiter = maxiter

    # bounds of variables
    Up = [1]*nVar
    Lo = [0]*nVar
    
    # dinamic q
    din_max_q=max_q
    
    # dinamic mass
    din_max_mass = max_mass

    # initilize matrices
    S = np.zeros((nSize,nVar))
    S_f = np.zeros((nSize,1))

    # initial velocity
    v = np.zeros((nSize,nVar))

    # history
    columns_ = ['iter']
    for i in parameters_v: columns_.append(i)
    columns_.append('f')
    df = pd.DataFrame(columns=columns_)
    
    # generate first random solution
    #np.random.seed(seed)
    Srand = np.random.uniform(low=0,high=1,size=(nPart, nVar))
    f,S_r,maximize = mp_evaluator(problem, Srand)


    for i in range(len(S_r)):
        S_f[i] = f[i]
        
    # add responses and "fitness" column to solution
    S = np.hstack((Srand, v, S_f))
    # sort according to fitness (last column)
    S = sorted(S, key=lambda row: row[-1],reverse = maximize)
    S = np.array(S)

    # save each iteration
    iter_n = np.full((nPart, 1), 0.0)
    x_ = S[:, 0:nVar]
    x = correct_x(problem, x_)
    res = np.reshape(S[:, -1], (nPart, 1))
    rows=np.hstack((iter_n, x, res))
    df_new = pd.DataFrame(rows, columns=columns_)
    df = pd.concat([df, df_new])

    iterations=1
    
    # iterations
    while True:
        # get a new solution
        if solver == "Random":
            Stemp, S_f = Random_(problem, S[:, :nVar], auto_reduce_search_space=auto_reduce_search_space)
        
        elif solver=="Newton":
            din_mass=0
            if dinamic: 
                din_mass = ((maxiter-iterations)/(maxiter))*(np.max(S_f)-np.min(S_f))
            else:
                din_mass=max_mass
            
            m = np.reshape(S[:, -1], (nPart, 1))
            Stemp, S_f = F_newton_(problem, S[:, :nVar], m, Lo, Up, G=G, keep_best=keep_best, weight_method=weight_method, t=t, max_mass=din_mass)

        elif solver == "Coulomb":
            din_q=0
            if dinamic:
                din_q = ((maxiter-iterations)/(maxiter))*(np.max(S_f)-np.min(S_f))
            else:
                din_q=max_q
            q = np.reshape(S[:, -1], (nPart, 1))
            Stemp, v, S_f = F_coulomb_(problem, S[:, :nVar], v, q, P, Lo, Up, k=k, weight_method=weight_method, t=t, max_q=din_q)
        
        # add responses and "fitness" column to solution
        Ssample = np.hstack((Stemp, v, S_f))

        # add new solutions in the solutions table
        Solution_temp = np.vstack((S,Ssample))

        # delate duplicated rows
        Solution_temp = np.unique(Solution_temp, axis=0)

        # sort according to "fitness"
        Solution_temp = sorted(Solution_temp, key=lambda row: row[-1],reverse = maximize)
        Solution_temp = np.array(Solution_temp)

        # keep best solutions
        S = Solution_temp[:nSize][:]

        # save each iteration
        iter_n = np.full((nPart, 1), iterations)
        x_ = S[:, 0:nVar]
        x = correct_x(problem, x_)
        res = np.reshape(S[:, -1], (nPart, 1))
        rows=np.hstack((iter_n, x, res))
        df_new = pd.DataFrame(rows, columns=columns_)
        df = pd.concat([df, df_new])

        iterations += 1
        if iterations > maxiter:
            break

    best_sol = np.min(df['f'])
    ind_best_sol = np.argmin(df['f'])
    best_var = df.iloc[ind_best_sol, 1:len(parameters_v)+1]
    return df, best_var, best_sol

# Test Functions.
# Adapted from "https://www.sfu.ca/~ssurjano/optimization.html"

def Ackley(x, a=20.0, b=0.2, c=2.0*np.pi):
    d = len(x)
    sum1 = np.sum(np.square(x))
    sum2 = np.sum(np.array([np.cos(c*i) for i in x]))
    term1 = -a * np.exp(-b * np.sqrt(sum1 / d))
    term2 = -np.exp(sum2 / d)
    return term1 + term2 + a + np.exp(1)

def sixth_Bukin(x):
    return 100 * np.sqrt(np.abs(x[1] - 0.01*x[0]**2)) + 0.01*np.abs(x[0] + 10)

def Cross_in_Tray(x):
    return -0.0001 * math.pow(np.abs(np.sin(x[0]) * np.sin(x[1]) * np.exp(np.abs(100.0 - np.sqrt(x[0]**2 + x[1]**2)/np.pi)))+1.0, 0.1)

def Drop_Wave(x):
    return -1.0*(1.0 + np.cos(12.0 * np.sqrt(x[0]**2 + x[1]**2))) / (0.5 * (x[0]**2 + x[1]**2) + 2.0)

def Eggholder(x):
    return -(x[1] + 47.0) * np.sin(np.sqrt(np.abs(x[1] + x[0]/2 + 47.0))) - x[0] * np.sin(np.sqrt(np.abs(x[0] - (x[1] + 47.0))))

def Griewank(x):
    sum_part=0.0
    prod_part=1.0
    for i, xi in enumerate(x):
        sum_part += xi/4000.0
        prod_part *= np.cos(xi/np.sqrt(i+1))
    return sum_part - prod_part + 1.0

def Holder_Table(x):
    return -np.abs(np.sin(x[0])*np.cos(x[1])*np.exp(np.abs(1.0-(np.sqrt(x[0]**2 + x[1]**2)/np.pi))))

def Levy(x):
    # d dimensions
    w1 = 1.0 + (x[0]-1.0)/4.0
    wd = 1.0 + (x[-1]-1.0)/4.0
    sum_part = 0.0
    for i, xi in enumerate(x):
        wi = 1.0 + (xi-1.0)/4.0
        sum_part += ((wi-1.0)**2)*(1.0 + 10.0*math.pow(np.sin(np.pi*wi+1.0), 2))
    return math.pow(np.sin(np.pi*w1), 2) + sum_part + ((wd-1.0)**2)*(1.0 + math.pow(np.sin(2*np.pi*wd), 2))

def Rastrigin(x):
    # d dimensions
    d = len(x)
    sum_part = 0.0
    for i, xi in enumerate(x):
        sum_part += (xi**2) - 10.0*np.cos(2*np.pi*xi)
    return 10*d + sum_part

def second_Schaffe(x):
    return 0.5 + (math.pow(np.sin(x[0]**2 + x[1]**2), 2)-0.5)/(1.0 + 0.001*(x[0]**2 + x[1]**2))**2

def fourth_Schaffer(x):
    return 0.5 + (math.pow(np.cos(np.abs(x[0]**2 - x[1]**2)), 2)-0.5)/(1.0 + 0.001*(x[0]**2 + x[1]**2))**2

def Schwefel(x):
    # d dimensions
    d = len(x)
    sum_part = 0.0
    for xi in x:
        sum_part += xi*np.sin(np.sqrt(np.abs(xi)))
    return 418.9829*d - sum_part

def Shubert(x):
    sum_1 = 0.0
    sum_2 = 0.0
    for i in np.arange(5):
        i = float(i + 1)
        sum_1 += i*np.cos((i+1)*x[0] + i)
        sum_2 += i*np.cos((i+1)*x[1] + i)
    return sum_1*sum_2

def Styblinski_Tang(x):
    f = (sum([math.pow(i,4)-16*math.pow(i,2)+5*i for i in x])/2)
    return f

def Easom(x):
    f = np.cos(x[0])*np.cos(x[1])*np.exp(-(math.pow(x[0]-np.pi, 2) + math.pow(x[1]-np.pi, 2)))
    return f

def Bohachevsky(x):
    f = x[0]**2 + 2*x[1]**2 -0.3*np.cos(3*np.pi*x[0]) - 0.4*np.cos(4*np.pi*x[1]) + 0.7
    return f

def Perm_d_beta(x, d=2, beta=4):
    f = 0.0
    for i in range(d):
        for j in range(d):
            f += ((j+1) + beta)*(x[j]**(i+1) - (1/(j+1)**(i+1)))**2
    return f
    
def Rotated_Hyper_Ellipsoid(x, d=2):
    f = 0.0
    for i in range(d):
        for j in range(i+1):
            f += x[j]**2
    return f
            
def Sum_of_Different_Powers(x, d=2):
    f = 0.0
    for i in range(d):
        f += np.abs(x[i])**i+1+1
    return f
        
def SUM_SQUARES(x, d=2):    
    f = 0.0
    for i in range(d):
        f += (i+1)*x[i]**2
    return f

def TRID(x, d=2):
    sum_1 = 0.0
    sum_2 = 0.0
    for i in range(d):
        sum_1 += (x[i] - 1)**2
    for i in range(d-1):
        sum_2 += x[i+1]*x[i]   
    f = sum_1 + sum_2
    return f 

def BOOTH(x):
    f = (x[0] + 2*x[1] -7)**2 + (2*x[0] + x[1] - 5)**2
    return f

def Matyas(x):
    f = 0.26*(x[0]**2 + x[1]**2) - 0.48*x[0]*x[1]
    return f

def MCCORMICK(x):
    f = np.sin(x[0] + x[1]) + (x[0] - x[1])**2 - 1.5*x[0] + 2.5*x[1] + 1
    return f

def Power_Sum(x, d=2, b=[8, 18, 44, 114]):
    f = 0.0
    for i in range(d):
        sum_1 = 0.0
        for j in range(d):
            sum_1 += x[j]**(i+1)
        f += (sum_1 - b[i])**2
    return f
            
def Zakharov(x, d=2):
    f = 0.0
    sum_1 = 0.0
    sum_2 = 0.0
    sum_3 = 0.0
    for i in range(d):
        sum_1 += x[i]**2
        sum_2 += 0.5*(i+1)*x[i]
        sum_3 += 0.5*(i+1)*x[i]
    f = sum_1 + sum_2**2 + sum_3**4
    return f

def THREE_HUMP_CAMEL(x):
    f = 2*x[0]**2 - 1.05*x[0]**4 + (x[0]**6/6) + x[0]*x[1] + x[1]**2
    return f

def SIX_HUMP_CAMEL(x):
    f = (4 - 2.1*x[0] + (x[0]**4)/3)*x[0]**2 + x[0]*x[1] + (-4 + 4*x[1]**2)*x[1]**2
    return f

def DIXON_PRICE(x, d=2):
    sum_1 = 0.0
    for i in range(d-1):
        i = i + 1
        sum_1 += (i+1)*(2*x[i]**2 - x[i-1])**2
    f = (x[0] - 1)**2 + sum_1
    return f 
    
def ROSENBROCK(x, d=2):
    f = 0.0 
    for i in range(d-1):
        f += 100*(x[i+1] - x[i]**2)**2 + (x[i] - 1)**2
    return f    

def DE_JONG(x):
    a = [[-32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16, 32], 
         [-32, -32, -32, -32, -32, -16, -16, -16, -16, -16, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 32, 32, 32, 32, 32]]
    sum_1 = 0.0
    for i in range(25):
        sum_1 += 1/((i+1) + (x[0] - a[0][i])**6 + (x[1] - a[1][i])**6)
    f = (0.002 + sum_1)**(-1)
    return f

def MICHALEWICZ(x, d=2, m=10):
    f = 0.0
    for i in range(d):
        f += np.sin(x[i])*np.sin(((i+1)*x[i]**2)/np.pi)**(2*m) 
    f = -f
    return f
  
def BEALE(x):
    f = (1.5 - x[0] + x[0]*x[1])**2 + (2.25 - x[0] + x[0]*x[1]**2)**2 + (2.625 - x[0] + x[0]*x[1]**3)**2
    return f
    
def BRANIN(x):
    a=1
    b=5.1/(4*(np.pi)**2)
    c = 5/np.pi
    r = 6
    s = 10
    t = 1/(8*np.pi)
    f = a*(x[1] - b*x[0]**2 + c*x[0] - r)**2 + s*(1 - t)*np.cos(x[0]) + s
    return f
    
def GOLDSTEIN_PRICE(x):
    f = (1 + ((x[0] + x[1] + 1)**2)*(19 - 14*x[0] + 3*x[0]**2 - 14*x[1] + 6*x[0]*x[1] + 3*x[1]**2))*(30 + ((2*x[0] - 3*x[1])**2)*(18 - 32*x[0] + 12*x[0]**2 + 48*x[1] - 36*x[0]*x[1] + 27*x[1]**2))    
    return f

def PERM_D_BETA(x, d=2, beta=1):
    f = 0.0
    for i in range(d):
        sum_1 = 0
        for j in range(d):
            sum_1 += ((j+1)**(i+1) + beta)*((x[j]/(j+1))**(i+1) - 1)
        f += sum_1
    return f

class test_functions():
    def __init__(self) -> None:
        self.Ackley_ = {'name':'Ackley', 'f':Ackley, 'bounds':[[-32.768, 32.768], [-32.768, 32.768]], 'opt':[[0.0, 0.0], 0.0]}
        self.sixth_Bukin_ = {'name':'sixth_Bukin', 'f':sixth_Bukin, 'bounds':[[-15.0, -5.0], [-3.0, 3.0]], 'opt':[[-10.0, 1.0], 0.0]}
        self.Cross_in_Tray_ = {'name':'Cross_in_Tray', 'f':Cross_in_Tray, 'bounds':[[-10.0, 10.0], [-10.0, 10.0]], 'opt':[[[1.3491, -1.3491], [1.3491, 1.3491], [-1.3491, 1.3491], [-1.3491, -1.3491]], -2.06261]}
        self.Drop_Wave_ = {'name':'Drop_Wave', 'f':Drop_Wave, 'bounds':[[-5.12, 5.12], [-5.12, 5.12]], 'opt':[[0, 0], -1.0]}
        self.Eggholder_ = {'name':'Eggholder', 'f':Eggholder, 'bounds':[[-512.0, 512.0], [-512.0, 512.0]], 'opt':[[512.404, 512.404], -959.6407]}
        self.Griewank_ = {'name':'Griewank', 'f':Griewank, 'bounds':[[-600.0, 600.0], [-600.0, 600.0]], 'opt':[[0.0, 0.0], 0.0]}
        self.Holder_Table_ = {'name':'Holder_Table', 'f':Holder_Table, 'bounds':[[-10.0, 10.0], [-10.0, 10.0]], 'opt':[[[8.05502, 9.66459], [8.05502, -9.66459], [-8.05502, 9.66459], [-8.05502, -9.66459]], -19.2085]}
        self.Levy_ = {'name':'Levy', 'f':Levy, 'bounds':[[-10.0, 10.0], [-10.0, 10.0]], 'opt':[[1.0, 1.0], 0.0]}
        self.Rastrigin_ = {'name':'Rastrigin', 'f':Rastrigin, 'bounds':[[-5.12, 5.12], [-5.12, 5.12]], 'opt':[[0.0, 0.0], 0.0]}
        self.second_Schaffe_ = {'name':'second_Schaffe', 'f':second_Schaffe, 'bounds':[[-100.0, 100.0], [-100.0, 100.0]], 'opt':[[0.0, 0.0], 0.0]}
        self.fourth_Schaffer_ = {'name':'fourth_Schaffer', 'f':fourth_Schaffer, 'bounds':[[-100.0, 100.0], [-100.0, 100.0]], 'opt':[[0.0, 0.0], 0.0]}
        self.Schwefel_ = {'name':'Schwefel', 'f':Schwefel, 'bounds':[[-500.0, 500.0], [-500.0, 500.0]], 'opt':[[420.9687, 420.9687], 0.0]}
        self.Shubert_ = {'name':'Shubert', 'f':Shubert, 'bounds':[[-10.0, 10.0], [-10.0, 10.0]], 'opt':[[0.0, 0.0], -186.7309]}
        self.Styblinski_Tang_ = {'name':'Styblinski_Tang', 'f':Styblinski_Tang, 'bounds':[[-5, 5], [-5, 5]]}
        self.Easom_ = {'name':'Easom', 'f':Easom, 'bounds':[[-3, 3], [-3, 3]]}
        self.Bohachevsky_ = {'name':'Bohachevsky', 'f':Bohachevsky, 'bounds':[[-100.0, 100.0], [-100.0, 100.0]], 'opt':[[0, 0], 0]}
        self.Perm_d_beta_ = {'name':'Perm_d_beta', 'f':Perm_d_beta, 'bounds':[[-2.0, 2.0], [-2.0, 2.0]], 'opt':[[1, 0.5], 0]}
        self.Rotated_Hyper_Ellipsoid_ = {'name':'Rotated_Hyper_Ellipsoid', 'f':Rotated_Hyper_Ellipsoid, 'bounds':[[-65.536, 65.536], [-65.536, 65.536]], 'opt':[[0.0, 0.0], 0]}
        self.Sum_of_Different_Powers_ = {'name':'Sum_of_Different_Powers', 'f':Sum_of_Different_Powers, 'bounds':[[-1, 1], [-1, 1]], 'opt':[[0.0, 0.0], 0]}
        self.SUM_SQUARES_ = {'name':'SUM_SQUARES', 'f':SUM_SQUARES, 'bounds':[[-10, 10], [-10, 10]], 'opt':[[0.0, 0.0], 0]}
        self.TRID_ = {'name':'TRID', 'f':TRID, 'bounds':[[-4, 4], [-4, 4]], 'opt':[[2, 2], -2]}
        self.BOOTH_ = {'name':'BOOTH', 'f':BOOTH, 'bounds':[[-10, 10], [-10, 10]], 'opt':[[1, 3], 0]}
        self.Matyas_ = {'name':'Matyas', 'f':Matyas, 'bounds':[[-10, 10], [-10, 10]], 'opt':[[0, 0], 0]}
        self.MCCORMICK_ = {'name':'MCCORMICK', 'f':MCCORMICK, 'bounds':[[-1.5, 4], [-3, 4]], 'opt':[[-0.54719, -1.54719], -1.9133]}
        self.Power_Sum_ = {'name':'Power_Sum', 'f':Power_Sum, 'bounds':[[0, 2], [0, 2]]}
        self.Zakharov_ = {'name':'Zakharov', 'f':Zakharov, 'bounds':[[-5, 10], [-5, 10]], 'opt':[[0.0, 0.0], 0.0]}
        self.THREE_HUMP_CAMEL_ = {'name':'THREE_HUMP_CAMEL', 'f':THREE_HUMP_CAMEL, 'bounds':[[-5, 5], [-5, 5]], 'opt':[[0.0, 0.0], 0.0]}
        self.SIX_HUMP_CAMEL_  = {'name':'SIX_HUMP_CAMEL', 'f':SIX_HUMP_CAMEL, 'bounds':[[-3, 3], [-2, 2]], 'opt':[[0.0898, -0.7126], -1.0316]}
        self.DIXON_PRICE_ = {'name':'DIXON_PRICE', 'f':DIXON_PRICE, 'bounds':[[-10, 10], [-10, 10]], 'opt':[[1, 1/np.sqrt(2)], 0]}
        self.ROSENBROCK_ = {'name':'ROSENBROCK', 'f':ROSENBROCK, 'bounds':[[-5, 10], [-5, 10]], 'opt':[[1, 1], 0]}
        self.DE_JONG_ = {'name':'DE_JONG', 'f':DE_JONG, 'bounds':[[-65.536, 65.536], [-65.536, 65.536]]}
        self.MICHALEWICZ_ = {'name':'MICHALEWICZ', 'f':MICHALEWICZ, 'bounds':[[0, np.pi], [0, np.pi]], 'opt':[[2.2, 1.57], -1.8013]}
        self.BEALE_ = {'name':'BEALE', 'f':BEALE, 'bounds':[[-4.5, 4.5], [-4.5, 4.5]], 'opt':[[3, 0.5], 0]}
        self.BRANIN_ = {'name':'BRANIN', 'f':BRANIN, 'bounds':[[-5, 10], [0, 15]], 'opt':[[-np.pi, 12.275], 0.397887]}
        self.GOLDSTEIN_PRICE_ = {'name':'GOLDSTEIN_PRICE', 'f':GOLDSTEIN_PRICE, 'bounds':[[-2, 2], [-2, 2]], 'opt':[[0, -1], 3]}
        self.PERM_D_BETA_ = {'name':'PERM_D_BETA', 'f':PERM_D_BETA, 'bounds':[[-2, 2], [-2, 2]], 'opt':[[1, 2], 0]}
        self.dictionary = {'Ackley': self.Ackley_, 
                           'sixth_Bukin':self.sixth_Bukin_, 
                           'Cross_in_Tray':self.Cross_in_Tray_, 
                           'Drop_Wave':self.Drop_Wave_, 
                           'Eggholder':self.Eggholder_, 
                            'Griewank':self.Griewank_, 
                            'Holder_Table':self.Holder_Table_, 
                            'Levy':self.Levy_, 
                            'Rastrigin':self.Rastrigin_, 
                            'second_Schaffe':self.second_Schaffe_, 
                            'fourth_Schaffer':self.fourth_Schaffer_, 
                            'Schwefel':self.Schwefel_, 
                            'Shubert':self.Shubert_, 
                            'Styblinski_Tang':self.Styblinski_Tang_, 
                            'Easom':self.Easom_, 
                            'Bohachevsky':self.Bohachevsky_, 
                            'Perm_d_beta':self.Perm_d_beta_, 
                            'Rotated_Hyper_Ellipsoid':self.Rotated_Hyper_Ellipsoid_, 
                            'Sum_of_Different_Powers': self.Sum_of_Different_Powers_, 
                            'SUM_SQUARES':self.SUM_SQUARES_, 
                            'TRID':self.TRID_, 
                            'BOOTH': self.BOOTH_, 
                            'Matyas':self.Matyas_, 
                            'MCCORMICK': self.MCCORMICK_, 
                            'Power_Sum':self.Power_Sum_, 
                            'Zakharov':self.Zakharov_, 
                            'THREE_HUMP_CAMEL' :self.THREE_HUMP_CAMEL_, 
                            'SIX_HUMP_CAMEL': self.SIX_HUMP_CAMEL_, 
                            'DIXON_PRICE': self.DIXON_PRICE_, 
                            'ROSENBROCK_': self.ROSENBROCK_, 
                            'DE_JONG':self.DE_JONG_, 
                            'MICHALEWICZ': self.MICHALEWICZ_, 
                            'BEALE':self.BEALE_, 
                            'BRANIN': self.BRANIN_, 
                            'GOLDSTEIN_PRICE':self.GOLDSTEIN_PRICE_, 
                            'PERM_D_BETA':self.PERM_D_BETA_} 
        self.whole_list = list(self.dictionary.keys())


def plotly_graph(problem, df=None):
    function=problem['f']
    x_lb=problem['bounds'][0][0]
    x_ub=problem['bounds'][0][1]
    y_lb=problem['bounds'][1][0]
    y_ub=problem['bounds'][1][1]
    
    x = np.linspace(x_lb, x_ub, 100)
    y = np.linspace(y_lb, y_ub, 100)
    z = np.empty((100, 100))
    for ind_y, j in enumerate(y):
        for ind_x, i in enumerate(x):
            z[ind_y][ind_x] = function(np.array([i, j]))
    
    steps_ = int(np.max(df['iter']))

    fig1 = go.Figure(data=[go.Surface(x=x, y=y, z=z)])
    for step in range(steps_):
        points = df[df['iter']==step]
        points_x = list(points['x1'])
        points_y = list(points['x2'])
        points_z = list(points['f'])
        fig1.add_scatter3d(x=np.array(points_x), y=np.array(points_y), z=np.array(points_z), mode='markers', visible=False, marker=dict(size=5, color="white", line=dict(width=1, color="black")))
        fig1.update_layout(title=f"f = {step}")
    
    # Create figure
    fig = go.Figure(data=[go.Scatter3d(x=[], y=[], z=[],
                                mode="markers",
                                marker=dict(size=5, color="white", line=dict(width=1, color="black"))
                                ), fig1.data[0]]
                    )
        
    # Frames
    frames = [go.Frame(data=[go.Scatter3d(x=k['x'],
                                        y=k['y'],
                                        z=k['z']
                                        ), fig1.data[0]
                            ],
                    traces= [0],
                    name=f'frame{ind}'      
                    ) for ind, k  in  enumerate(fig1.data[1:])
            ]

    fig.update(frames=frames)

    def frame_args(duration):
        return {
                "frame": {"duration": duration},
                "mode": "immediate",
                "fromcurrent": True,
                "transition": {"duration": duration, "easing": "linear"},
                }


    sliders = [
        {"pad": {"b": 10, "t": 60},
        "len": 0.9,
        "x": 0.1,
        "y": 0,
        
        "steps": [
                    {"args": [[f.name], frame_args(0)],
                    "label": str(k),
                    "method": "animate",
                    } for k, f in enumerate(fig.frames)
                ]
        }
            ]

    fig.update_layout(

        updatemenus = [{"buttons":[
                        {
                            "args": [None, frame_args(150)],
                            "label": "Play", 
                            "method": "animate",
                        },
                        {
                            "args": [[None], frame_args(150)],
                            "label": "Pause", 
                            "method": "animate",
                    }],
                        
                    "direction": "left",
                    "pad": {"r": 10, "t": 70},
                    "type": "buttons",
                    "x": 0.1,
                    "y": 0,
                }
            ],
            sliders=sliders
        )

    fig.update_layout(sliders=sliders)
    fig.write_html('animation.html')
    return fig

        
problem=test_functions().PERM_D_BETA_  
df, best_var, best_sol = custom_pso(problem, 20, 2, maxiter=100, solver="Newton", k=0.000000001, G=0.00000001, t=1.0, max_q=0.01, auto_reduce_search_space=True, dinamic=True)
#plotly_graph(problem, df)

print(best_sol)