module MCSP.Algorithms.PSO (
    -- * Updaters
    Weight,
    Updater,
    randomVelocity,
    globalGuideDirection,
    localGuideDirection,
    previousVelocity,

    -- * Data structures
    PSOGuide (..),
    Particle (..),
    Swarm (..),
    particleSwarmOptimization,
) where

import Control.Applicative (pure)
import Control.Monad (fmap, mapM, replicateM, (>>=))
import Data.Eq (Eq, (==))
import Data.Function (($))
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.List (intercalate, (++))
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty.Extra (maximum1)
import Data.Maybe (Maybe (..))
import Data.Ord (Ord (..))
import Data.Vector.Unboxed (Unbox, Vector, length)
import GHC.Err (error)
import GHC.Exts (fromListN, toList)
import GHC.Num ((+))
import GHC.Stack (HasCallStack)
import Numeric (showFFloat)
import Text.Show (Show, show, showListWith)

import MCSP.Algorithms.Vector (Default, sortLike, uniformSN, zeros, (.+), (.-))
import MCSP.System.Random (Random, iterateR)

-- -------------------------------------------------------------
-- Based on https://github.com/brianshourd/haskell-Calypso
-- -------------------------------------------------------------

-- | Represents an evaluation of how good a solution is.
type Grade = Int

-- | Element of a vector used to sort values of a permutation problem.
type Weight = Default

-- ----- --
-- Guide --
-- ----- --

-- | Information about a specific position (weights) and
-- the value of the objective function at that point.
data PSOGuide a = PsoGuide
    { -- | Position (weights).
      forall a. PSOGuide a -> Vector Weight
guideWeights :: Vector Weight,
      -- | Grade of the position.
      forall a. PSOGuide a -> Grade
guideGrade :: Grade,
      -- | Values sorted by weights.
      forall a. PSOGuide a -> Vector a
sortedValues :: Vector a
    }

instance Show (PSOGuide a) where
    show :: PSOGuide a -> String
show PSOGuide a
g = (Weight -> ShowS) -> [Weight] -> ShowS
forall a. (a -> ShowS) -> [a] -> ShowS
showListWith Weight -> ShowS
showF (Vector Weight -> [Item (Vector Weight)]
forall l. IsList l => l -> [Item l]
toList (Vector Weight -> [Item (Vector Weight)])
-> Vector Weight -> [Item (Vector Weight)]
forall a b. (a -> b) -> a -> b
$ PSOGuide a -> Vector Weight
forall a. PSOGuide a -> Vector Weight
guideWeights PSOGuide a
g) String
" -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Grade -> String
forall a. Show a => a -> String
show (PSOGuide a -> Grade
forall a. PSOGuide a -> Grade
guideGrade PSOGuide a
g)
      where
        showF :: Weight -> ShowS
showF = Maybe Grade -> Weight -> ShowS
forall a. RealFloat a => Maybe Grade -> a -> ShowS
showFFloat (Grade -> Maybe Grade
forall a. a -> Maybe a
Just Grade
3)

instance Eq (PSOGuide a) where
    PsoGuide Vector Weight
_ Grade
x Vector a
_ == :: PSOGuide a -> PSOGuide a -> Bool
== PsoGuide Vector Weight
_ Grade
y Vector a
_ = Grade
x Grade -> Grade -> Bool
forall a. Eq a => a -> a -> Bool
== Grade
y

instance Ord (PSOGuide a) where
    PsoGuide Vector Weight
_ Grade
x Vector a
_ <= :: PSOGuide a -> PSOGuide a -> Bool
<= PsoGuide Vector Weight
_ Grade
y Vector a
_ = Grade
x Grade -> Grade -> Bool
forall a. Ord a => a -> a -> Bool
<= Grade
y

-- -------- --
-- Particle --
-- -------- --

-- | A single particle of a swarm.
data Particle a = Particle
    { -- | Position of the particle.
      forall a. Particle a -> Vector Weight
particleWeights :: Vector Weight,
      -- | Velocity of the particle.
      forall a. Particle a -> Vector Weight
vel :: Vector Weight,
      -- | Best position the particle found so far.
      forall a. Particle a -> PSOGuide a
pGuide :: PSOGuide a
    }

instance Show (Particle a) where
    show :: Particle a -> String
show Particle {Vector Weight
PSOGuide a
particleWeights :: forall a. Particle a -> Vector Weight
vel :: forall a. Particle a -> Vector Weight
pGuide :: forall a. Particle a -> PSOGuide a
particleWeights :: Vector Weight
vel :: Vector Weight
pGuide :: PSOGuide a
..} =
        String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate
            String
"\n"
            [ String
Item [String]
"Particle:",
              String
"- Position: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vector Weight -> String
forall a. Show a => a -> String
show Vector Weight
particleWeights,
              String
"- Guide: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PSOGuide a -> String
forall a. Show a => a -> String
show PSOGuide a
pGuide
            ]

-- -------- --
-- Updaters --
-- -------- --

-- | Implicit parameters used in `Updater`.
type UpdaterContext a =
    ( ?particle :: Particle a,
      ?global :: PSOGuide a,
      ?iteration :: Int
    )

-- | Evaluate the new velocity of a particle using the global best and iteration number.
type Updater a = UpdaterContext a => Random (Vector Weight)

-- | Produce random velocity with components up to given limit.
randomVelocity :: UpdaterContext a => Random (Vector Weight)
randomVelocity :: forall a. UpdaterContext a => Random (Vector Weight)
randomVelocity = Grade -> Random (Vector Weight)
forall a. (Unbox a, Variate a, Num a) => Grade -> Random (Vector a)
uniformSN (Vector Weight -> Grade
forall a. Unbox a => Vector a -> Grade
length (Vector Weight -> Grade) -> Vector Weight -> Grade
forall a b. (a -> b) -> a -> b
$ Particle a -> Vector Weight
forall a. Particle a -> Vector Weight
particleWeights ?particle::Particle a
Particle a
?particle)
{-# INLINE randomVelocity #-}

-- | Produce random velocity in the direction of the current global best,
-- covering a random portion of the distance between them up to the given limit.
globalGuideDirection :: UpdaterContext a => Vector Weight
globalGuideDirection :: forall a. UpdaterContext a => Vector Weight
globalGuideDirection = PSOGuide a -> Vector Weight
forall a. PSOGuide a -> Vector Weight
guideWeights ?global::PSOGuide a
PSOGuide a
?global Vector Weight -> Vector Weight -> Vector Weight
forall a. (Unbox a, Num a) => Vector a -> Vector a -> Vector a
.- Particle a -> Vector Weight
forall a. Particle a -> Vector Weight
particleWeights ?particle::Particle a
Particle a
?particle
{-# INLINE globalGuideDirection #-}

-- | Produce random velocity in the direction of the current local best,
-- covering a random portion of the distance between them up to the given limit.
localGuideDirection :: UpdaterContext a => Vector Weight
localGuideDirection :: forall a. UpdaterContext a => Vector Weight
localGuideDirection = PSOGuide a -> Vector Weight
forall a. PSOGuide a -> Vector Weight
guideWeights (Particle a -> PSOGuide a
forall a. Particle a -> PSOGuide a
pGuide ?particle::Particle a
Particle a
?particle) Vector Weight -> Vector Weight -> Vector Weight
forall a. (Unbox a, Num a) => Vector a -> Vector a -> Vector a
.- Particle a -> Vector Weight
forall a. Particle a -> Vector Weight
particleWeights ?particle::Particle a
Particle a
?particle
{-# INLINE localGuideDirection #-}

-- | Just repeats the previous particle velocity.
previousVelocity :: UpdaterContext a => Vector Weight
previousVelocity :: forall a. UpdaterContext a => Vector Weight
previousVelocity = Particle a -> Vector Weight
forall a. Particle a -> Vector Weight
vel ?particle::Particle a
Particle a
?particle
{-# INLINE previousVelocity #-}

-- ----- --
-- Swarm --
-- ----- --

-- | A swarm for the PSO algorithm.
data Swarm a = Swarm
    { -- | Particles in the swarm.
      forall a. Swarm a -> NonEmpty (Particle a)
parts :: NonEmpty (Particle a),
      -- | Global guide.
      forall a. Swarm a -> PSOGuide a
gGuide :: PSOGuide a,
      -- | Current iteration.
      forall a. Swarm a -> Grade
iteration :: Int
    }

instance Show (Swarm a) where
    show :: Swarm a -> String
show Swarm {Grade
NonEmpty (Particle a)
PSOGuide a
parts :: forall a. Swarm a -> NonEmpty (Particle a)
gGuide :: forall a. Swarm a -> PSOGuide a
iteration :: forall a. Swarm a -> Grade
parts :: NonEmpty (Particle a)
gGuide :: PSOGuide a
iteration :: Grade
..} = Grade -> String
forall a. Show a => a -> String
show Grade
iteration String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" - " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PSOGuide a -> String
forall a. Show a => a -> String
show PSOGuide a
gGuide

-- | Implicit parameters used in iterations of the PSO algorithm.
type PSOContext a =
    ( Unbox a,
      ?eval :: Vector a -> Grade,
      ?values :: Vector a
    )

-- | Create a particle using the evaluation function, the vector of
-- original values and a specific position.
createParticle :: (HasCallStack, PSOContext a) => Vector Weight -> Particle a
createParticle :: forall a.
(HasCallStack, PSOContext a) =>
Vector Weight -> Particle a
createParticle Vector Weight
weights =
    Particle
        { particleWeights :: Vector Weight
particleWeights =
            if Vector Weight -> Grade
forall a. Unbox a => Vector a -> Grade
length Vector Weight
weights Grade -> Grade -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Grade
forall a. Unbox a => Vector a -> Grade
length ?values::Vector a
Vector a
?values
                then Vector Weight
weights
                else String -> Vector Weight
forall a. HasCallStack => String -> a
error String
"size mismatch while creating particle",
          vel :: Vector Weight
vel = Grade -> Vector Weight
forall a. (Unbox a, Num a) => Grade -> Vector a
zeros (Vector Weight -> Grade
forall a. Unbox a => Vector a -> Grade
length Vector Weight
weights),
          pGuide :: PSOGuide a
pGuide =
            let sortedValues :: Vector a
sortedValues = ?values::Vector a
Vector a
?values Vector a -> Vector Weight -> Vector a
forall a b.
(Unbox a, Unbox b, Ord b) =>
Vector a -> Vector b -> Vector a
`sortLike` Vector Weight
weights
             in PsoGuide {guideWeights :: Vector Weight
guideWeights = Vector Weight
weights, guideGrade :: Grade
guideGrade = ?eval::Vector a -> Grade
Vector a -> Grade
?eval Vector a
sortedValues, Vector a
sortedValues :: Vector a
sortedValues :: Vector a
sortedValues}
        }

-- | Creates a swarm from the evaluation function, the vector of
-- original values, the number of particles and a generator of positions.
createSwarm :: (HasCallStack, PSOContext a) => Int -> Random (Vector Weight) -> Random (Swarm a)
createSwarm :: forall a.
(HasCallStack, PSOContext a) =>
Grade -> Random (Vector Weight) -> Random (Swarm a)
createSwarm Grade
n Random (Vector Weight)
gen = do
    NonEmpty (Particle a)
parts <- Grade -> [Item (NonEmpty (Particle a))] -> NonEmpty (Particle a)
forall l. IsList l => Grade -> [Item l] -> l
fromListN Grade
n ([Particle a] -> NonEmpty (Particle a))
-> Random [Particle a] -> Random (NonEmpty (Particle a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Grade -> Random (Particle a) -> Random [Particle a]
forall (m :: * -> *) a. Applicative m => Grade -> m a -> m [a]
replicateM Grade
n (Vector Weight -> Particle a
forall a.
(HasCallStack, PSOContext a) =>
Vector Weight -> Particle a
createParticle (Vector Weight -> Particle a)
-> Random (Vector Weight) -> Random (Particle a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Random (Vector Weight)
gen)
    let gGuide :: PSOGuide a
gGuide = NonEmpty (PSOGuide a) -> PSOGuide a
forall a. Ord a => NonEmpty a -> a
maximum1 (NonEmpty (PSOGuide a) -> PSOGuide a)
-> NonEmpty (PSOGuide a) -> PSOGuide a
forall a b. (a -> b) -> a -> b
$ (Particle a -> PSOGuide a)
-> NonEmpty (Particle a) -> NonEmpty (PSOGuide a)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Particle a -> PSOGuide a
forall a. Particle a -> PSOGuide a
pGuide NonEmpty (Particle a)
parts
    Swarm a -> Random (Swarm a)
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Swarm a -> Random (Swarm a)) -> Swarm a -> Random (Swarm a)
forall a b. (a -> b) -> a -> b
$ Swarm {NonEmpty (Particle a)
parts :: NonEmpty (Particle a)
parts :: NonEmpty (Particle a)
parts, PSOGuide a
gGuide :: PSOGuide a
gGuide :: PSOGuide a
gGuide, iteration :: Grade
iteration = Grade
0}

-- ---------------- --
-- Iterating swarms --
-- ---------------- --

-- | Updates a single particle using an updater, the evaluation function,
-- the vector of original values and swarm information.
updateParticle :: PSOContext a => Updater a -> Swarm a -> Particle a -> Random (Particle a)
updateParticle :: forall a.
PSOContext a =>
Updater a -> Swarm a -> Particle a -> Random (Particle a)
updateParticle Updater a
newVel Swarm {Grade
NonEmpty (Particle a)
PSOGuide a
parts :: forall a. Swarm a -> NonEmpty (Particle a)
gGuide :: forall a. Swarm a -> PSOGuide a
iteration :: forall a. Swarm a -> Grade
parts :: NonEmpty (Particle a)
gGuide :: PSOGuide a
iteration :: Grade
..} Particle a
part = do
    Vector Weight
vel <- Random (Vector Weight)
Updater a
newVel
    let w' :: Vector Weight
w' = Particle a -> Vector Weight
forall a. Particle a -> Vector Weight
particleWeights Particle a
part Vector Weight -> Vector Weight -> Vector Weight
forall a. (Unbox a, Num a) => Vector a -> Vector a -> Vector a
.+ Vector Weight
vel
    let sortedValues :: Vector a
sortedValues = ?values::Vector a
Vector a
?values Vector a -> Vector Weight -> Vector a
forall a b.
(Unbox a, Unbox b, Ord b) =>
Vector a -> Vector b -> Vector a
`sortLike` Vector Weight
w'
    let newVal :: Grade
newVal = ?eval::Vector a -> Grade
Vector a -> Grade
?eval Vector a
sortedValues
    let oldGuide :: PSOGuide a
oldGuide = Particle a -> PSOGuide a
forall a. Particle a -> PSOGuide a
pGuide Particle a
part
    let pGuide' :: PSOGuide a
pGuide' =
            if Grade
newVal Grade -> Grade -> Bool
forall a. Ord a => a -> a -> Bool
>= PSOGuide a -> Grade
forall a. PSOGuide a -> Grade
guideGrade PSOGuide a
oldGuide
                then Vector Weight -> Grade -> Vector a -> PSOGuide a
forall a. Vector Weight -> Grade -> Vector a -> PSOGuide a
PsoGuide Vector Weight
w' Grade
newVal Vector a
sortedValues
                else PSOGuide a
oldGuide
    Particle a -> Random (Particle a)
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Particle {particleWeights :: Vector Weight
particleWeights = Vector Weight
w', Vector Weight
vel :: Vector Weight
vel :: Vector Weight
vel, pGuide :: PSOGuide a
pGuide = PSOGuide a
pGuide'}
  where
    ?particle = ?particle::Particle a
Particle a
part
    ?global = ?global::PSOGuide a
PSOGuide a
gGuide
    ?iteration = ?iteration::Grade
Grade
iteration

-- | Updates all particles in the swarm once.
updateSwarm :: PSOContext a => Updater a -> Swarm a -> Random (Swarm a)
updateSwarm :: forall a. PSOContext a => Updater a -> Swarm a -> Random (Swarm a)
updateSwarm Updater a
up Swarm a
swarm = do
    NonEmpty (Particle a)
newParts <- (Particle a -> Random (Particle a))
-> NonEmpty (Particle a) -> Random (NonEmpty (Particle a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM (Updater a -> Swarm a -> Particle a -> Random (Particle a)
forall a.
PSOContext a =>
Updater a -> Swarm a -> Particle a -> Random (Particle a)
updateParticle Random (Vector Weight)
Updater a
up Swarm a
swarm) (Swarm a -> NonEmpty (Particle a)
forall a. Swarm a -> NonEmpty (Particle a)
parts Swarm a
swarm)
    let bestGuide :: PSOGuide a
bestGuide = NonEmpty (PSOGuide a) -> PSOGuide a
forall a. Ord a => NonEmpty a -> a
maximum1 (NonEmpty (PSOGuide a) -> PSOGuide a)
-> NonEmpty (PSOGuide a) -> PSOGuide a
forall a b. (a -> b) -> a -> b
$ (Particle a -> PSOGuide a)
-> NonEmpty (Particle a) -> NonEmpty (PSOGuide a)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Particle a -> PSOGuide a
forall a. Particle a -> PSOGuide a
pGuide NonEmpty (Particle a)
newParts
    let newGuide :: PSOGuide a
newGuide = PSOGuide a -> PSOGuide a -> PSOGuide a
forall a. Ord a => a -> a -> a
max PSOGuide a
bestGuide (Swarm a -> PSOGuide a
forall a. Swarm a -> PSOGuide a
gGuide Swarm a
swarm)
    Swarm a -> Random (Swarm a)
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Swarm {parts :: NonEmpty (Particle a)
parts = NonEmpty (Particle a)
newParts, gGuide :: PSOGuide a
gGuide = PSOGuide a
newGuide, iteration :: Grade
iteration = Swarm a -> Grade
forall a. Swarm a -> Grade
iteration Swarm a
swarm Grade -> Grade -> Grade
forall a. Num a => a -> a -> a
+ Grade
1}

-- | Create iterations of swarms, trying to maximize the objective funtion.
particleSwarmOptimization ::
    PSOContext a => Updater a -> Random (Vector Weight) -> Int -> Random (NonEmpty (Swarm a))
particleSwarmOptimization :: forall a.
PSOContext a =>
Updater a
-> Random (Vector Weight) -> Grade -> Random (NonEmpty (Swarm a))
particleSwarmOptimization Updater a
update Random (Vector Weight)
weights Grade
size =
    Grade -> Random (Vector Weight) -> Random (Swarm a)
forall a.
(HasCallStack, PSOContext a) =>
Grade -> Random (Vector Weight) -> Random (Swarm a)
createSwarm Grade
size Random (Vector Weight)
weights Random (Swarm a)
-> (Swarm a -> Random (NonEmpty (Swarm a)))
-> Random (NonEmpty (Swarm a))
forall a b. Random a -> (a -> Random b) -> Random b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Swarm a -> Random (Swarm a))
-> Swarm a -> Random (NonEmpty (Swarm a))
forall a. (a -> Random a) -> a -> Random (NonEmpty a)
iterateR (Updater a -> Swarm a -> Random (Swarm a)
forall a. PSOContext a => Updater a -> Swarm a -> Random (Swarm a)
updateSwarm Random (Vector Weight)
Updater a
update)