pso <- function(x0=c(-5, -8), sig=1, n=10, N=200, w.ini=0.9, w.fin=0.4, phi1=2, phi2=2) {
# A simple version of the PSO algorithm with fully connected topology.
# For optimization of the Rosenbrock function (classroom demonstration).
# Uses Gaussian initial positions of particles and zero initial weights.
# 
# x0 ... the mean value of the initial particle positions
# sig ... the standard deviation of the initial positions of particles
# n ... the number of particles
# N ... the number of iterations (no sophisticated termination rule)
# w.ini ... the initial inertia weight
# w.fin ... the final inertia weight
# phi1, phi2 ... the acceleration parameters


  # The Rosenbrock function to be minimized
  f <- function(x){(1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2}

  # Plot the image of the Rosenbrock function
  yy <- xx <- seq(-10, 10, by=0.1)
  zz <- yy %o% xx
  for(i in 1:201) {
    for(j in 1:201) {
      zz[i, j] <- f(c(xx[i], yy[j]))
    }
  }
  image(xx, yy, zz, xlab="x", ylab="y", col=topo.colors(15), breaks=exp(0:15) - 1.5)
  points(1, 1, cex=4, col="red", lwd=3)

  # Initialization
  d <- length(x0)
  v <- x <- matrix(0, ncol=d, nrow=n)
  for(i in 1:n) {
    x[i,] <- x0 + rnorm(d, sd=sig)
  }

  pb <- x
  pb.val <- rep(0, n)
  for(i in 1:n) {
    pb.val[i] <- f(pb[i,]) 
  }

  g <- x[1,]
  g.val <- f(g)
  for(i in 2:n) {
    fact <- f(x[i,])
    if(g.val > fact) {
      g <- x[i,]
      g.val <- fact
    }
  }

  w <- w.ini
  w.del <- (w.ini - w.fin) / N

  # The main cycle
  for(t in 1:N) {

    # Plot the swarm
    for(i in 1:n) {
      points(x[i, 1], x[i, 2], pch=19, cex=0.5, col="red")
    }

    # Delay the execution just to be able to see the animation
    Sys.sleep(0.05)

    for(i in 1:n) {
      points(x[i, 1], x[i, 2], pch=19, cex=0.5)
    } 

    # Update the velocities and the positions of the particles
    for(i in 1:n) {
      v[i,] <- w * v[i,] + phi1 * runif(1) * (pb[i,] - x[i,]) + phi2 * runif(1) * (g - x[i,])  # The heart of the method
      x[i,] <- x[i,] + v[i,]
      fact <- f(x[i,])
      if(pb.val[i] > fact) {
        pb[i,] <- x[i,]
        pb.val[i] <- fact
        if(g.val > fact) {
          g <- x[i,]
          g.val <- fact
        }
      }
    }
  
    w <- w - w.del
  }

  # Plot the final swarm and return the results
  for(i in 1:n) {
    points(x[i, 1], x[i, 2], pch=19, cex=0.5, col="red")
  }

  list(g=g, g.val=g.val)
}

# Demo:
pso(c(-7, 5))
pso(c(-7, -7))
pso(c(7, -5))

# A teraz sa pohrajme s ostatnymi parametrami...