stylix/palette-generator/Ai/Evolutionary.hs

{-# LANGUAGE MultiParamTypeClasses #-}

module Ai.Evolutionary ( EvolutionConfig(..), Species(..), evolve ) where

import Control.Applicative ( liftA2 )
import Data.Bifunctor ( first, second )
import Data.List ( mapAccumR, sortBy )
import Data.Ord ( Down(Down, getDown), comparing )
import System.Random ( RandomGen, mkStdGen, randomR )
import Text.Printf ( printf )

{- |
Find every possible combination of two values, with the first value
coming from one list and the second value coming from a different list.
-}
cartesianProduct :: [a] -> [b] -> [(a, b)]
cartesianProduct = liftA2 (,)

{- |
Find every possible combination of two values, with both values coming
from the same list. Values are allowed to be paired with themself.
-}
cartesianSquare :: [a] -> [(a, a)]
cartesianSquare as = as `cartesianProduct` as

-- | Pick a random element from a list using a random generator.
randomFromList :: (RandomGen r) => r -> [a] -> (a, r)
randomFromList generator list
  = let (index, generator') = randomR (0, length list - 1) generator
     in (list !! index, generator')

{- |
Map over a list, passing a random generator into the mapped
function each time it is called. A random generator is returned
along with the new list.
-}
mapWithGen :: (r -> a -> (r, b)) -> (r, [a]) -> (r, [b])
mapWithGen = uncurry . mapAccumR

unfoldWithGen :: (r -> (r, a)) -> Int -> r -> (r, [a])
unfoldWithGen _ 0 generator = (generator, [])
unfoldWithGen f size generator =
  let (generator', as) = unfoldWithGen f (size - 1) generator
      (generator'', a) = f generator'
   in (generator'', a:as)

{- |
A genotype is a value which is generated by the genetic algorithm.

The environment is used to specify the problem for which
we are trying to find the optimal genotype.
-}
class Species environment genotype where
  -- | Generate a new genotype at random.
  generate :: (RandomGen r) => environment -> r -> (r, genotype)

  -- | Randomly combine two genotypes.
  crossover :: (RandomGen r) => environment -> r -> genotype -> genotype -> (r, genotype)

  -- | Randomly mutate a genotype using the given environment.
  mutate :: (RandomGen r) => environment -> r -> genotype -> (r, genotype)

  -- | Score a genotype. Higher numbers are better.
  fitness :: environment -> genotype -> Double

-- | Parameters for the genetic algorithm.
data EvolutionConfig = EvolutionConfig
  { -- | The number of genotypes processed on each pass.
    populationSize :: Int,
    -- | How many genotypes make it through to the next pass.
    survivors :: Int,
    -- | The chance of a genotype being randomly changed
    --   before crossover. Between 0 and 1.
    mutationProbability :: Double,
    -- | When the fitness score improves by less than this percentage,
    --   the algorithm will stop.
    changeThreshold :: Double
  }

{- |
Randomly mutate the given genotype, if the mutation probability
from the 'EvolutionConfig' says yes.
-}
randomMutation :: (RandomGen r, Species e g)
               => e -- ^ Environment
               -> EvolutionConfig
               -> r -- ^ Random generator
               -> g -- ^ Genotype to mutate
               -> (r, g)
randomMutation environment config generator chromosome
  = let (r, generator') = randomR (0.0, 1.0) generator
     in if r <= mutationProbability config
        then mutate environment generator' chromosome
        else (generator', chromosome)

{- |
Select the fittest survivors from a population,
to be moved to the next pass of the algorithm.
-}
naturalSelection :: (Species e g)
                 => e -- ^ Environment
                 -> EvolutionConfig
                 -> [g] -- ^ Original population
                 -> [(Double, g)] -- ^ Survivors with fitness scores
naturalSelection environment config
  = take (survivors config)
  . map (first getDown)
  . sortBy (comparing fst)
  -- Avoid computing fitness multiple times during sorting
  -- Down reverses the sort order so that the best fitness comes first
  . map (\genotype -> (Down $ fitness environment genotype, genotype))

-- | Run one pass of the genetic algorithm over a given population.
evolveGeneration :: (RandomGen r, Species e g)
                 => e -- ^ Environment
                 -> EvolutionConfig
                 -> (r, [g]) -- ^ Random generator, population from previous generation
                 -> (r, Double, [g]) -- ^ New random generator, maximum fitness, new population
evolveGeneration environment config (generator, population)
  = (newGenerator, maximum fitnesses, newPopulation)
  where
    (fitnesses, newPopulation) = unzip newPopulationWithFitness

    (newGenerator, newPopulationWithFitness) =
      second (naturalSelection environment config)
      $ mapWithGen (randomMutation environment config)
      $ unfoldWithGen randomCrossover (populationSize config) generator

    randomCrossover gen = let (pair, gen') = randomFromList gen pairs
                           in (uncurry $ crossover environment gen') pair

    pairs = cartesianSquare population

evolveUntilThreshold :: (RandomGen r, Species e g)
                     => e -- ^ Environment
                     -> EvolutionConfig
                     -> [Double] -- ^ Fitnesses of previous generations
                     -> (r, [g]) -- ^ Random generator, population from previous generation
                     -> IO (r, [g]) -- ^ New random generator, final population
evolveUntilThreshold environment config fitnesses (generator, population) =
  do
    let (generator', fitness, population') =
          evolveGeneration environment config (generator, population)

        -- Begins at 0 on the first iteration
        generationNumber = length fitnesses

        fitnesses' = fitness : fitnesses
        recentFitnesses = take 5 fitnesses'

        {-
           On the first iteration there is only one recent fitness, so the
           improvement would be calculated as 0%. To prevent the algorithm
           stopping immediately, we fall back to 100% in this case.
        -}
        change =
          if generationNumber < 1
          then 1
          else 1 - (head recentFitnesses / last recentFitnesses);

    printf "Generation: %3i  Fitness: %7.1f Improvement: %5.1f%%\n"
      generationNumber fitness (change * 100)

    if change < changeThreshold config
      then return (generator', population')
      else evolveUntilThreshold environment config fitnesses' (generator', population')

{- |
Create the initial population, to be fed into the first
pass of the genetic algorithm.
-}
initialGeneration :: (RandomGen r, Species e g)
                  => e -- ^ Environment
                  -> EvolutionConfig
                  -> r -- ^ Random generator
                 -> (r, [g]) -- ^ New random generator, population
initialGeneration environment config
  = unfoldWithGen (generate environment) (survivors config)

-- | Run the full genetic algorithm.
evolve :: Species e g
       => e -- ^ Environment
       -> EvolutionConfig
       -> IO g -- ^ Optimal genotype
evolve environment config = do
  (_, population) <-
    evolveUntilThreshold environment config []
    $ initialGeneration environment config
    $ mkStdGen 0 -- Fixed seed for determinism

  return $ head population