Nonogram solver, genetic algorithm

ga.hs

{-# LANGUAGE OverloadedStrings,ScopedTypeVariables #-}

import Data.Text (Text)
import qualified Data.Text as T
import Data.Bits (xor, shiftL, shiftR, (.&.))
import Data.List (transpose, find, sortBy)
import System.Random
import Control.Applicative
import Test.HUnit (Assertion(..), Counts(..), Test(..), assertEqual,
                    assertBool, assertFailure, runTestTT)
import Debug.Trace
import Control.Monad (unless, replicateM)

type Nonogram = ([[Int]], [[Int]]) --縦制約、横制約
type Z = Integer
type NonoAns  = (Z, Int, Int) -- 答え、vlen, wlen
--SAT
data BoolExp = RowSymbols Int Int  -- X_iの行。X_i0 = val & 1 X_i1 = val & (1<<1) ...
             | ColSymbols Int Int  -- X_iの列。X_0j = val & 1 X_1j = val & (1<<1) ...
             | Not BoolExp     -- NOT
             | And  [BoolExp]   -- And
             | Or  [BoolExp]   -- OR
  deriving (Eq, Show)

showNono :: Nonogram -> String
showNono nono@(vs,ws) = 
  (unlines $ verticalAlignment ws)  ++
  (unlines $ map ((flip (++) $ wx).showLine) vs)
  where 
    vlen = length vs
    wlen = length ws
    wx = replicate wlen '-'
    leftmargin = wlen + 1
    marginSpace = replicate leftmargin ' '
    showLine es = -- [1,2,3] -> "1,2,3"
      (flip (++) " ") $ T.unpack $ T.justifyRight (leftmargin - 1) ' ' $ T.intercalate "" $ map (T.pack.show) es
    verticalAlignment iss =
      map (\i->marginSpace ++ vmap i iss) [0..h-1]
      where 
        h = maximum $ map length iss 
        vmap _ [] = []
        vmap i (is:iss) = 
          c ++ vmap i iss
          where
            c =
              if lis + i >= h then show (is !! (lis+i-h))
              else " "
            lis = length is
       -- verticalAlignment "123", "54" -> "  1 "
       --                                  "  25"
       --                                  "  34"

{-
showNonoAns :: NonoAns -> String
showNonoAns  = 
  showNonoLine l ++ showNonoAns ls
  where
    showNonoLine [] = ""
    showNonoLine (i:is) = show i ++ " " ++ showNonoLine is
-}

nono2sat :: Nonogram -> BoolExp
nono2sat (vs, ws) = 
  let a = And $ (map vline2sat $ zip [0..] vs) ++ (map wline2sat $ zip [0..] ws) in
  a --trace (show a) a
  where 
    vlen = length vs
    wlen = length ws
   --i番目の縦行をsatに直す 
    vline2sat (i,vl::[Int]) = 
      let l = Or $ map (\c -> RowSymbols i c) andCodes in
      l -- trace (show l) l
      where  
        --andCodes = filter (isproper vl) [0..2^wlen-1]
        andCodes' :: Int -> [Int] -> [Int]
        andCodes' _ [] = [0]
        andCodes' _ [0] = [0]
        andCodes' pos is@(i:iss) = 
          if i+pos > wlen then []
          else
           (map (+ (shiftL (2^i-1) pos)) $ andCodes' (pos+i) iss) ++ andCodes' (pos+1) is 
        andCodes = andCodes' 0 vl
    wline2sat (j,wl) = 
      Or $ map (\c -> ColSymbols j c) andCodes
      where  
        andCodes' :: Int -> [Int] -> [Int]
        andCodes' _ [] = [0]
        andCodes' _ [0] = [0]
        andCodes' pos is@(i:iss) = 
          if i+pos > vlen then []
          else
           (map (+ (shiftL (2^i-1) pos)) $ andCodes' (pos+i) iss) ++ andCodes' (pos+1) is 
        andCodes = andCodes' 0 wl
 
--sat2fitness 適合率 1.0なら正解
sat2fitness :: NonoAns -> BoolExp -> Double 
sat2fitness ans@(ansVal,vl,wl) (RowSymbols i val) = 
  {- trace (show (ansVal, i, val, point, positions)) $ -} fromIntegral point / fromIntegral wl
  where
    positions = zip [0..] [wl*i..wl*(i+1)-1]
    point = (sum $ map (\(ps ,p)->1 - xor (fromIntegral (shiftR ansVal p .&. 1)) (shiftR val ps .&. 1))  positions)
sat2fitness ans@(ansVal,vl,wl) (ColSymbols j val) = 
  fromIntegral point / fromIntegral wl
  where
    positions = zip [0..] [j,j+wl..j+(vl-1)*wl]
    point = (sum $ map (\(ps, p)->1 - xor (fromIntegral (shiftR ansVal p .&. 1)) (shiftR val ps .&. 1))  positions)
sat2fitness ans (Not a) = 1.0 - sat2fitness ans a
sat2fitness ans (And xs) = sum (map (sat2fitness ans) xs) / fromIntegral (length xs)
sat2fitness ans (Or xs) = 
  maximum' $ map (sat2fitness ans) xs
  where
    maximum' (x:xs') = if x>=0.999 then 1.0 else max x (maximum' xs')
    maximum' [] = 0.0

sat2FitnessTests :: [Test]
sat2FitnessTests = map TestCase
  [ 
    assertEqual "issNonoansConvert#1" nonoExplAns (nonoans2iss $ iss2Nonoans $ nonoExplAns)
  , assertEqual "issNonoansConvert#2" houseAns (nonoans2iss $ iss2Nonoans $ houseAns)
  , assertNear "Sat2FitnessTest#1" 1.0 (sat2fitness (1,2,2) $ nono2sat nonoExpl2) 0.01
  , assertNear "Sat2FitnessTest#2" (1/2) (sat2fitness (7,2,2) $ nono2sat nonoExpl2) 0.01
  , assertNear "Sat2FitnessTest#3" 0.75 (sat2fitness (2,2,2) $ nono2sat nonoExpl2) 0.01
  , assertNear "Sat2FitnessTest#4" 0.75 (sat2fitness (0,2,2) $ nono2sat nonoExpl2) 0.01 
  , assertNear "Sat2FitnessTest#5" 0.75 (sat2fitness (4,2,2) $ nono2sat nonoExpl2) 0.01 
  , assertNear "Sat2FitnessTest#6" 0.5 (sat2fitness (10,2,2) $ nono2sat nonoExpl2) 0.01
  , assertNear "Sat2FitnessTest#7" 1.0 (sat2fitness (341,3,3) $ nono2sat nonoExpl3) 0.01
  , assertNear "Sat2FitnessTest#8" 1.0 
    (sat2fitness (iss2Nonoans $ transpose nonoExplAns) $ nono2sat nonoExpl) 0.01
  , assertNear "Sat2FitnessTest#9" 1.0 
    (sat2fitness (iss2Nonoans $ transpose houseAns) $ nono2sat nonoHouse) 0.01
  ]
  where
    assertNear :: String -> Double -> Double -> Double -> Assertion
    assertNear str a b d = unless (abs (a-b) <= d) $ assertFailure (mkMsg str a b d)
    mkMsg :: String -> Double -> Double -> Double -> String
    mkMsg str a b d = unlines
      [ str 
      , "expected: " ++ show a
      , "but got: " ++ show b
      , "out of range: " ++ show d
      ]
    nonoExplAns = [[0,1,1,0,0],[1,1,1,1,0],[0,1,0,1,1],[1,1,1,1,0],[0,1,1,0,0]]
    houseAns = [[0,0,1,1,1],[0,1,1,1,1],[1,1,1,1,1],[0,1,1,1,1],[0,0,1,1,1]]

iss2Nonoans :: [[Int]] -> NonoAns
iss2Nonoans iss = 
  (sum $ map line2Val $ zip [0,wl..] iss, vl, wl)
  where 
    vl = length $ iss
    wl = length $ iss !! 0
    line2Val (k, is) = sum $ map (elem2Val k) $ zip [0..] is
    elem2Val k (k', e) = if e /= 0 then 2^(k+k') else 0

nonoans2iss  :: NonoAns -> [[Int]]
nonoans2iss ans@(ansVal, vl, wl) = 
  map (makeLine wl) [0..vl-1]
  where
    makeLine wl i =
      map (\j->if (/= 0) $ shiftL 1 (i*wl+j) .&. ansVal then 1 else 0) [0..wl-1]

showNonoans :: NonoAns -> String
showNonoans ans@(ansVal, vl, wl) = 
   unlines $ map (makeLineString wl) [0..vl-1]
  where
    makeLineString wl i =
      map (\j->if (/= 0) $ shiftL 1 (i*wl+j) .&. ansVal then '*' else ' ') [0..wl-1]

--ここから探索

populationSize = 100
eliteSize = 30 --淘汰選択の際に適合度で厳密に選ぶ要素の数
rouletteSize = populationSize - eliteSize --ランダム要素を入れて選ぶ要素の数
crossover = truncate (0.6 * (fromIntegral populationSize) / 2.0)
mutation = 0.01 --変異確率

---GA (遺伝的アルゴリズム）
ga :: Nonogram -> [NonoAns] -> IO NonoAns
ga nono population = 
  --trace (show population) $ 
  do --population内に解がなかった
    mutations <- mapM mutate $ zip [0..] population
    pairs <- replicateM crossover $ mkPairs
    crossovers <- mapM createCrossover pairs
    let nextPopulation = mutations ++ tup2ToList crossovers
        fitnessNP      = zip (map (flip sat2fitness' $ sat) nextPopulation) nextPopulation
    if (fst $ head $ fitnessNP) >= 0.999 then return $ snd $ head $ fitnessNP --解が見つかった 
    else do
      let (elites, rest) = splitAt eliteSize $ sortBy (flip compare) fitnessNP
      rouletteEliminations <- selectRouletteEliminations rest rouletteSize
      let nextSurvival = elites ++ rouletteEliminations
      putStrLn $ "Best:" ++ show (fst $ head nextSurvival)
      putStrLn $ showNonoans $ snd $ head $ nextSurvival
      ga nono (map snd nextSurvival)
  where 
    nonoSize = let (_, vl, wl) = head population in vl*wl
    sat = nono2sat nono
    sat2fitness' ans sat = 
      let f = sat2fitness ans sat in f -- trace (show f) f
    --二点交叉
    createCrossover :: (NonoAns, NonoAns) -> IO (NonoAns, NonoAns) --次世代の要素の作成
    createCrossover (ans1@(ansVal1, vlen1, wlen1), ans2@(ansVal2, vlen2, wlen2)) = do --2点交差で新しい個体を作成
      point1 <- getStdRandom $ randomR (0, nonoSize-3) 
      point2 <- getStdRandom $ randomR (min (point1+1) (nonoSize-1), nonoSize-1)
      let p1mask = map (shiftL 1) [0..point1-1]
          p2mask = map (shiftL (shiftL 1 point1)) [0..(point2-point1)-1]
          p3mask = map (shiftL (shiftL 1 point2)) [0..(nonoSize-point2)-1]
          (mask1, mask2) = (sum p1mask + sum p3mask, sum p2mask)
      return $ ( ((.&.) ansVal1 mask1 + (.&.) ansVal2 mask2, vlen1, wlen1) ,
                 ((.&.) ansVal2 mask1 + (.&.) ansVal1 mask2, vlen2, wlen2) )
    --突然変異
    mutate :: (Int, NonoAns) -> IO NonoAns
    mutate (rank, (ansVal, vl, wl)) = do
      mutationBits <- getMutationBits mutation -- (mutation * fromIntegral (rank*5+populationSize) / fromIntegral populationSize)
      return (ansVal `xor` mutationBits, vl, wl)
      where
        bitMutate :: Double -> IO Z
        bitMutate mutationRate = do
          (r::Double) <- getStdRandom $ randomR (0.0, 1.0)
          if r < mutationRate then return 1 else return 0
        getMutationBits :: Double -> IO Z
        getMutationBits mutationRate = do
          mb <- replicateM nonoSize $ bitMutate mutationRate
          return $ sum $ map (\(p,b)->shiftL b p) $ zip [0..] mb
    mkPairs :: IO (NonoAns, NonoAns) --交差するための遺伝子ペアを選ぶ
    mkPairs = do
      l1 <- getStdRandom $ randomR (0, populationSize-1)
      l2 <- getStdRandom $ randomR (0, populationSize-1)
      return (population !! l1, population !! l2)
    tup2ToList [] = []
    tup2ToList ((a,b):rest) = a:b:tup2ToList rest
    --ランキングにしたがって割り当てられた生存率と乱数で生存個体を選ぶ(毎ターン、全体数-現在の順位 に比例して選ばれる)
    --あまり早くないけどまあボトルネックにはならないでしょう・・・
    selectRouletteEliminations :: [a] -> Int -> IO [a]
    selectRouletteEliminations _ 0 = return []
    selectRouletteEliminations xs num = do
      let s = sum $ [1..length xs]
          rsx = zip3 [0..] (tail $ scanl (+) 0 $ reverse [1..length xs]) xs
      r <- getStdRandom $ randomR (0, s-1)
      let c = findNum rsx r
          (h,l) = splitAt c xs
      rest <- selectRouletteEliminations (h ++ (tail l)) (num-1)
      return $ head l : rest 
      where 
        findNum rxs r = fst3 $ head $ dropWhile (\x->snd3 x<r) rxs
        fst3 (a,b,c) = a
        snd3 (a,b,c) = b

--問題例
nonoExpl = ([[1,1],[5],[2,2],[3],[1]], [[2],[4],[1,2],[4],[2]])
nonoExpl2 = ([[1],[]], [[1],[]])
nonoExpl3 = ([[1,1],[1],[1,1]], [[1,1],[1],[1,1]])
--ハウス(5*5)
nonoHouse = ([[1],[3],[5],[5],[5]], [[3],[4],[5],[4],[3]])
--ネ申(30*30)
nonoGod   = ([
               [0],[3],[5,4],[5,4],[4,3],[2,2],[2],[2,2],[5,2,3],[8,10],[11,16],[9,9,4],
               [4,3,4,10],[2,4,14],[6,9,4],[7,3,3,3],[9,12],[6,3,10],
               [3,2,2,6,2],[4,2,1,2],[4,2,2],[2,2],[3,2],[3,2],[3,3],[2,4],[2,4],[3],[2],[0]
               ]
            ,[
               [0],[1,1],[4,2],[4,3],[3,4],[4,4],[1,3,4],[2,9,3],[2,19],[3,19],[4,5,4],
               [3,4,5],[1,2,4],[3,2],[5],[7],[9],[2,7],[2,3,2,3,2],[3,10,4],
               [28],[27],[2,2,3,2],[2,2,2],[7],[11],[9],[6],[3],[0]
               ])
--ニコニコ(10*10)
nonoNiconico = ([
                 [0],[3,4],[1],[1],[5,4],[0],[3,4],[1],[1],[5,4] 
                  ]
               ,[
                 [1,1],[1,1,1,1],[1,1,1,1],[1,1,1,1],[1,1],[0],[1,1,1,1],[1,1,1,1],[1,1,1,1],[4,4] 
                  ])
--ニコニコテレビちゃん（14*14, 複数解）
nonoNicotv = ([
               [2,2],[2,2],[2],[14],[1,1],[1,1],[1,1,1],[1,1,1],[1,1],[1,2,1],[1,4,1],[1,1],[14],[1,1]
                ]
             ,[
               [10],[1,2],[1,1,1,1],[1,1,1],[1,1,1],[1,1,1,1],[2,2,1],[2,2,1],[1,1,1,1],[1,1,1],[1,1,1],[1,1,1,1],[1,2],[10]
                ])
--雨 http://www.minicgi.net/logic/logic.html?num=1871
nonoRain = ([
             [1,4,1],[4,3],[3,6],[3,2,2],[1,2,3,1],[1,3,2,1,1],[2,1,1,2,2],[3,1,1],[5,2,1],[1,4,2]
            ,[2,2,1],[1,5,1],[2,1,1,1],[1,1,3],[2,2],[1,1],[2,2],[5,1],[3]
              ]
           ,[
             [2],[2,1,1],[1,1,1,3],[2,1,4],[1,1,2],[2,2,3],[1,3,3],[1,4,3],[2,2,1,3],[1,2,1,3]
            ,[1,2,3],[1,1,1],[2,4],[2,4],[2,3,1],[1,5,1,1],[2,1],[3,2],[4,3],[1,6]
              ])


--空の解作成（初期値用-ほんとは空じゃないほうがいいですね）
blankNonoAns n = (0, n, n)

--一定確率で埋まった解を返す、乱数なのでIO..
randomNonoAns :: Int -> Double -> IO NonoAns
randomNonoAns n ratio = do
  mb <- replicateM nonoSize randBit
  return $ (sum $ map (\(p,b)->shiftL b p) $ zip [0..] mb, n, n)
  where
    randBit :: IO Z
    randBit = do
      (r::Double) <- getStdRandom $ randomR (0.0, 1.0)
      if r < ratio then return 1 else return 0
    nonoSize = n*n

--main = putStrLn $ showNono $ nonoExpl
--main = print $ nono2sat $ nonoExpl2
--main = print =<< ga nonoExpl2 [[0,0],[0,0]]
--main = print =<< (randomWalk nonoExpl3 $ blankNonoAns 3)

--main = putStrLn.("Answer is\n" ++).showNonoans =<< (ga nonoHouse $ replicate populationSize $ blankNonoAns 5)
main = putStrLn.("Answer is\n" ++).showNonoans =<< ga nonoNiconico =<< replicateM populationSize (randomNonoAns 10 0.2)
--main = putStrLn.("Answer is\n" ++).showNonoans =<< ga nonoNicotv =<< replicateM populationSize (randomNonoAns 14 0.2)
--main = putStrLn.("Answer is\n" ++).showNonoans =<< ga nonoRain =<< replicateM populationSize (randomNonoAns 20 0.2)
--main = putStrLn.("Answer is\n" ++).showNonoans =<< ga nonoGod =<< replicateM populationSize (randomNonoAns 30 0.3)

runTests :: [Test] -> IO Counts
runTests ts = runTestTT $ TestList ts

--TEST
--main = runTests $ sat2FitnessTests

↑10_10は解ける、14_14のパズルは解けなかった

http://www.slideshare.net/dekosuke/ga-nonogram