Main.hs

{-
  | = CLI application

  Generates an expression approximating an unknown
  multivariable function f(x1,x2,x3,...,x_8) using data from an external file.

  File format: comma-separated values with the function evaluations
  in the last column, i.e.

  x1 x2 x3 ... x_8 f(x1,x2,x3,...,x_8)

  For example, run 2000 iterations expressions of maximal length 30
  with mutation probability 0.2 and the population of 400 chromosomes.

  $ stack exec hmep -- -f Data.csv -t 2000 --var 0.2 -l 30 -m 0.2 -p 400

-}

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ExtendedDefaultRules #-}

module Main where

-- CLI options from `optparse-applicative`
import           Options.Applicative as Op
import           Data.Semigroup ( (<>) )
import qualified Data.Vector as V
-- CSV files from `cassava`
import           Data.Csv ( decode
                          , HasHeader (..)
                          )
import           Data.List ( foldl' )
import qualified Data.ByteString.Lazy as BS
import           Control.Monad ( foldM )
import           Text.Printf ( printf )
-- Random operations from `probable`
import           Math.Probable.Random ( double )

import           AI.MEP

-- | (Mean) square error
dist :: Double -> Double -> Double
dist x y = (x - y)^2

-- Functions available to genetically produced programs.
-- Modify this to your needs.
ops = V.fromList [
   ('*', (*)),
   ('+', (+)),
   -- Avoid division by zero
   ('/', \x y -> x / (y + 1e-6)),
   ('-', (-)),
   ('s', \x -> sin. const x),
   ('f', \x -> fromIntegral. floor. const x),
   -- Power; invalid operation results in zero
   ('^', \x y -> let z = x**y in if isNaN z || isInfinite z then 0 else z),
   ('e', \x _ -> let z = exp x in if isNaN z || isInfinite z then 0 else z),
   ('a', \x -> abs. const x)
   -- ('n', min),
   -- ('x', max)
 ]

-- | CLI options are parsed to this data structure
data ProgOptions = ProgOptions
  { _inputFile :: FilePath
  , _chromosomeLength :: Int
  , _mutationProb :: Double
  , _varProb :: Double
  , _constProb :: Double
  , _popSize :: Int
  , _maxIter :: Int
  }

-- | CLI options parser
progOptions :: Parser ProgOptions
progOptions = ProgOptions
  -- Mandatory field (no default value)
  <$> strOption
      ( short 'f'
      <> metavar "<input file>"
      <> help "Input file path. Format: comma-separated, two columns." )
  <*> option auto
      ( long "length"
      <> short 'l'
      <> metavar "30"
      <> value 30
      <> help "Chromosome length" )
  <*> option auto
      ( long "mutation"
      <> short 'm'
      <> metavar "0.05"
      <> value 0.05
      <> help "Mutation probability" )
  <*> option auto
      ( long "var"
      <> short 'r'
      <> metavar "0.1"
      <> value 0.1
      <> help "Probability to generate a new variable gene" )
  <*> option auto
      ( long "const"
      <> short 'c'
      <> metavar "0.05"
      <> value 0.05
      <> help "Probability to generate a new constant gene" )
  <*> option auto
      ( long "population"
      <> short 'p'
      <> metavar "200"
      <> value 200
      <> help "Population size" )
  <*> option auto
      ( long "total"
      <> short 't'
      <> metavar "200"
      <> value 200
      <> help "Total number of iterations" )

main :: IO ()
main = Op.execParser opts >>= run
  where
    opts = Op.info (Op.helper <*> progOptions)
      ( fullDesc
     <> header "A CLI interface to Haskell multi expression programming" )

last9 (_, _, _, _, _, _, _, _, y) = y
init9 (x1,x2,x3,x4,x5,x6,x7,x8,_) = V.fromList [x1,x2,x3,x4,x5,x6,x7,x8]

run :: ProgOptions -> IO ()
run arg = do
  let pVar = _varProb arg
      pConst = _constProb arg
      pMut = _mutationProb arg
      config0 = defaultConfig
        { c'ops = ops
        , c'length = _chromosomeLength arg
        , p'mutation = pMut
        , p'var = pVar
        , p'const = pConst
        , c'popSize = _popSize arg
        }
      maxIter = _maxIter arg

  let f = _inputFile arg

  putStrLn $ printf "Chromosome length: %d" (_chromosomeLength arg)
  putStrLn $ printf "Population size: %d" (_popSize arg)
  putStrLn $ "Mutation probability: " ++ show pMut
  putStrLn $ "Probability to generate a new variable gene: " ++ show pVar
  putStrLn $ "Probability to generate a new constant gene: " ++ show pConst
  let pNew = 1 - pConst - pVar
  putStrLn $ "Probability to generate a new operator: " ++ show pNew
  putStrLn ""

  putStrLn $ "Reading file " ++ f
  bs <- BS.readFile f
  -- TODO: Improve parsing with Cassava
  let result = decode NoHeader bs :: Either String (V.Vector (Double, Double, Double, Double, Double, Double, Double, Double, Double))
  case result of
    Left err -> error err
    Right parsed -> do
      let dtaY = V.map last9 parsed  -- Last column
          dtaX = V.map init9 parsed
          dataset = (dtaX, dtaY) :: (V.Vector (V.Vector Double), V.Vector Double)
      let datasetSize = V.length dtaX
          dim = V.length (V.head dtaX)  -- Input dimensionality
          config = config0 { c'vars = dim }
      putStrLn $ printf "Fetched %d records" datasetSize
      putStrLn $ printf "Input dimensions: %d\n" dim

      -- Randomly create a population of chromosomes
      pop <- runRandIO $ initialize config

      let loss = regressionLoss dist dataset

      -- Evaluate the initial population
      let popEvaluated = evaluatePopulation loss pop
          norm = fromIntegral datasetSize

      putStrLn $ "Average loss in the initial population " ++ show (avgLoss popEvaluated / norm)

      -- Declare how to produce the new generation
      let nextGeneration = evolve config loss (mutation3 config) crossover binaryTournament

      -- Specify the I/O loop, which logs every 5 generation
          runIO pop i = do
            newPop <- runRandIO $ foldM (\xg _ -> nextGeneration xg) pop [1..generations]
            putStrLn $ "Population " ++ show (i * generations) ++ ": average loss " ++ show (avgLoss newPop / norm)
            return newPop
              where generations = 5

      -- The final population
      final <- foldM runIO popEvaluated [1..(maxIter `div` 5)]

      putStrLn $ show $ best final

      putStrLn "Interpreted expression:"
      putStrLn $ generateCode (best final)