//
// PSO optimization for Rastrigin function
//
function y=J(x)
  y=sum(x.*x-cos(2*%pi*x))+length(x)
//  y=sum(x.*x)
endfunction
////////////////////////////////////////////////////
//
// 
clf();clear all;
//
xmin=-5;xmax=5;N=100;
x=xmin:((xmax-xmin)/(N-1)):xmax;
y=x;
z=zeros(N,N);
for i=1:N
for j=1:N
z(i,j)=J([x(i),y(j)]);
end
end
//
n=evstr(x_dialog('parameter number of the rastrigin function','2')); 
Npop=evstr(x_dialog('population number','10'));     
Ngen=evstr(x_dialog('generation number','50')); 
Vmax=evstr(x_dialog('maximal speed','1')); 
//
// PSO parameters
//
w=0.7;c1=0.7;c2=0.7;
xset('window',0);
xset('window',2);
xset('window',3);

//
//initialisation
//
PSO=zeros(Npop,3*n+1);       
PSO(:,1:n)=xmin+(xmax-xmin)*rand(Npop,n);
PSO(:,(n+1:2*n))=-Vmax+2*Vmax*rand(Npop,n);
PSO(:,(2*n+1:3*n))=PSO(:,1:n);
PSO(:,3*n+1)=1E10;
pg=zeros(1,n+1);pg(n+1)=1E10;pgval=[];Traj=[];
for i=1:Ngen
  val=[];
  for j=1:Npop
    valj=J(PSO(j,1:n))
    val=[val;valj] // the evaluation of the swarm
    // computation of the new pi
    if (valj<PSO(j,3*n+1)) then
      PSO(j,2*n+1:3*n)=PSO(j,1:n);
      PSO(j,3*n+1)=valj;
    end 
  end
  // computation of pg
  [a,b]=gsort(PSO(:,3*n+1));ind=b($);
  pg=[PSO(ind,2*n+1:3*n),PSO(ind,3*n+1)];
  pgval=[pgval,pg(n+1)];
  // new velocitiy
PSO(:,(n+1:2*n))=w*PSO(:,(n+1:2*n))+c1*rand(Npop,n).*(PSO(:,2*n+1:3*n)-PSO(:,1:n))+c2*rand(Npop,n).*(ones(Npop,1)*pg(1:n)-PSO(:,1:n));
  // new value of x
  PSO(:,1:n)=PSO(:,1:n)+PSO(:,(n+1:2*n)); 
//
//-----  dyanmic display if n=2 ----------------------------
   if (n==2) then
   xset('window',2);
   drawlater
   clf()
   Traj=[Traj;pg(1),pg(2)];
   contour2d(x,y,z,[0:0.01:0.1,0.2:1,1:10],rect=[-5.12,-5.12,5.12,5.12]);
   plot2d(PSO(:,1),PSO(:,2),-1,rect=[-5.12,-5.12,5.12,5.12]);  
   plot2d(Traj(:,1),Traj(:,2),5,rect=[-5.12,-5.12,5.12,5.12]);
   xtitle('trajectory display');
   drawnow
   xset('window',3);
   drawlater
   clf()
   contour2d(x,y,z,[0:0.01:0.1,0.2:1,1:10],rect=[-0.5,-0.5,0.5,0.5]);
   plot2d(PSO(:,1),PSO(:,2),-1,rect=[-0.5,-0.5,0.5,0.5]);   
   plot2d(Traj(:,1),Traj(:,2),5,rect=[-0.5,-0.5,0.5,0.5]);   
   xtitle('trajectory display');
   drawnow
   end
//
end
xset('window',1)
clf()
disp('best elemet found:');disp(pg(1:n))
plot2d(Npop*(1:Ngen),pgval)

