Created
December 8, 2016 15:47
-
-
Save Dierk/d6c3a1d69c8f2d5654798cf7f58568f9 to your computer and use it in GitHub Desktop.
AntHill simulation using Frege STM
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
module Ants where | |
import STM | |
import Control.Concurrent (forkIO, forkOS) | |
import Data.List (!!, nub, sortBy) | |
import System.Random | |
import fregefx.JavaFxAll (GraphicsContext, Color, Group, Stage, Canvas, VBox) | |
import fregefx.JavaFxUtils (<:, withStage, FregeFX, withUI) | |
data Location = L { !x :: Int | |
, !y :: Int | |
} | |
derive Eq Location | |
derive Show Location | |
data Cell = Cell { | |
--- amount of food in the cell; >= 0 | |
!food :: Int, | |
--- amount of pheromons in the cell; >= 0 | |
!pher :: Int, | |
--- Indicates whether an ant is in the cell and that this cell is its position | |
!ant :: Maybe Ant, | |
--- indicates whether this cell is in the ant hill | |
!home :: Bool | |
} where | |
takeFood :: Cell -> Cell | |
takeFood cell = cell.{ | |
food <- if isJust cell.ant then pred else id, | |
ant <- fmap _.{ food = True } | |
} | |
dropFood :: Cell -> Cell | |
dropFood cell = cell.{ | |
food <- succ, | |
ant <- fmap _.{ food = False } | |
} | |
placeAnt :: Cell -> Ant -> Cell | |
placeAnt cell ant = cell.{ ant = Just ant } | |
removeAnt :: Cell -> Cell | |
removeAnt cell = cell.{ ant = Nothing } | |
dropPheromone :: Cell -> Cell | |
dropPheromone cell = cell.{ pher <- excite, ant <- fmap calmdown} where | |
calmdown ant = ant.{excited = min excitement (ant.excited - boost)} | |
excite pher = min pherLimit (pher + 1 + boost) | |
boost = case cell.ant of | |
Nothing -> error "cannot happen!" | |
Just ant -> if ant.excited >= pherBoost then pherBoost else 0 | |
derive Show Cell | |
--- directions from North to North West going clock-wise. | |
data Direction = N | NE | E | SE | S | SW | W | NW where | |
--- clockwise direction change | |
succ :: Direction -> Direction | |
succ dir = case dir of | |
N -> NE | |
NE -> E | |
E -> SE | |
SE -> S | |
S -> SW | |
SW -> W | |
W -> NW | |
NW -> N | |
--- counter-clockwise direction change | |
pred :: Direction -> Direction | |
pred dir = case dir of | |
N -> NW | |
NE -> N | |
E -> NE | |
SE -> E | |
S -> SE | |
SW -> S | |
W -> SW | |
NW -> W | |
derive Show Direction | |
--- Translate a Direction into a pair of x and y offsets. | |
deltaDirection :: Direction -> Location | |
deltaDirection x = case x of | |
N -> L 0 (-1) | |
NE -> L 1 (-1) | |
E -> L 1 0 | |
SE -> L 1 1 | |
S -> L 0 1 | |
SW -> L (-1) 1 | |
W -> L (-1) 0 | |
NW -> L (-1) (-1) | |
data Ant = Ant { | |
!direction :: Direction, | |
!food :: Bool, | |
!excited :: Int -- extra pheros to drop; >= 0 | |
} where | |
--- turn the ant so often to the right (positive numbers) or to the left (negative) | |
turn :: Ant -> Int -> Ant | |
turn ant amount | |
| amount == 0 = ant | |
| amount > 0 = turn ant.{direction <- _.succ } (amount - 1) | |
| otherwise = turn ant.{direction <- _.pred } (amount + 1) | |
derive Show Ant | |
data AntAgent = AntAgent { | |
!world :: World, | |
!location :: Location, | |
rng :: TVar StdGen -- random number generator | |
} where | |
dice :: AntAgent -> STM Double | |
dice agent = do | |
rng <- agent.rng.read | |
let (result, rng') = Double.random rng | |
agent.rng.write rng' | |
return result | |
--- move ant from old into new cell if new cell is empty | |
move :: AntAgent -> STM AntAgent | |
move agent = do | |
let oldp = agent.world.place agent.location | |
oldc <- oldp.read | |
let ant = unJust oldc.ant | |
newloc = agent.world.deltaLocation agent.location ant.direction | |
newp = agent.world.place newloc | |
newc <- newp.read | |
case newc.ant of | |
Just _ -> return agent | |
Nothing -> do | |
if oldc.home | |
then oldp.write oldc.removeAnt | |
else oldp.write oldc.dropPheromone.removeAnt | |
newp.write $ newc.placeAnt ant | |
return agent.{ location = newloc } | |
behave :: AntAgent -> STM AntAgent | |
behave agent = do | |
let oldp = agent.world.place agent.location | |
cell <- oldp.read | |
let | |
world = agent.world | |
ant = unJust cell.ant | |
ahead = world.place (world.deltaLocation agent.location ant.direction) | |
aheadLeft = world.place (world.deltaLocation agent.location (ant.direction.pred)) | |
aheadRight = world.place (world.deltaLocation agent.location (ant.direction.succ)) | |
cellAhead <- ahead.read | |
cellAheadLeft <- aheadLeft.read | |
cellAheadRight <- aheadRight.read | |
let | |
places = [cellAhead, cellAheadLeft, cellAheadRight] | |
if ant.food | |
-- going home? | |
then | |
-- going home | |
-- at home? | |
if cell.home | |
-- at home | |
then do | |
oldp.write $ cell.dropFood.placeAnt (ant.{food=False,excited=excitement}.turn 4) | |
return agent | |
else | |
-- not at home | |
if cellAhead.home && isNothing cellAhead.ant | |
then (move agent >>= return) -- examine: why is bind return needed? | |
else | |
let | |
actions :: [STM AntAgent] | |
actions = [move agent >>= return, | |
(oldp.write (cell.placeAnt (ant.turn (-1)))) >> return agent, | |
(oldp.write (cell.placeAnt (ant.turn 1 ))) >> return agent] | |
rankAhead cell = if isJust cell.ant then 0 else cell.pher | |
rankHome cell = if cell.home then 2 else 1 | |
ranks = ranked places [ | |
[rankAhead], | |
[rankHome, _.pher], | |
[rankHome, _.pher] | |
] | |
in do | |
r <- agent.dice | |
let i = wrand r ranks | |
actions !! i | |
else | |
-- foraging | |
if cell.food > 0 && !cell.home | |
-- food in cell and not home? | |
then do | |
oldp.write $ cell.takeFood.placeAnt (ant.{food=True,excited=excitement}.turn 4) | |
return agent | |
-- no food? | |
else | |
if cellAhead.food > 0 && !cellAhead.home && isNothing cellAhead.ant | |
then (move agent >>= return) | |
else | |
let | |
actions :: [STM AntAgent] | |
actions = [move agent >>= return, | |
(oldp.write (cell.placeAnt (ant.turn (-1)))) >> return agent, | |
(oldp.write (cell.placeAnt (ant.turn 1 ))) >> return agent] | |
rankAhead c = case c.ant of | |
Just _ -> 0 | |
otherwise -> c.pher + c.food + (if c.home then 0 else 2) | |
ranks = ranked places [[rankAhead], | |
[_.food, _.pher, const 1], | |
[_.food, _.pher, const 1] | |
] | |
in do | |
r <- agent.dice | |
let i = wrand r ranks | |
actions !! i | |
act :: AntAgent -> IO Bool -> IO () | |
act agent running = do | |
Thread.sleep antSleep | |
stillRunning <- running | |
if stillRunning | |
then do agent' <- atomically $ agent.behave | |
act agent' running | |
else return () | |
--- for each cell, apply its list of ranking functions | |
ranked :: [Cell] -> [[(Cell -> Int)]] -> [Int] | |
ranked cs fss = zipWith go cs fss where | |
go c fs = (succ $ sum $ map ($ c) fs) | |
--- index selection from list of weights, given a (random) "distance" rate | |
wrand :: Double -> [Int] -> Int | |
wrand rate weights = loop 0 weights 0 where | |
loop i [] acc = i.pred | |
loop i (x:xs) acc = -- count weight indexes | |
if distance < (x + acc) -- until distance has been reached | |
then i | |
else loop i.succ xs (x + acc) | |
total = sum weights | |
distance = min (max 0 (total - 1)) (rate * total.double).long.fromIntegral | |
--- The world consists of TVars of type Cell. | |
data World = private World { !cells :: [[TVar Cell]] | |
, !dimX :: Int | |
, !dimY :: Int | |
} where | |
new :: Int -> Int -> STM World | |
new dimX dimY = do | |
cells <- mapM (const $ mapM (const $ TVar.new (Cell 0 0 Nothing False)) [0..dimY - 1]) [0..dimX - 1] | |
return $ World cells dimX dimY | |
place :: World -> Location -> TVar Cell | |
place world (L x y) = (world.cells !! x) !! y | |
deltaLocation :: World -> Location -> Direction -> Location | |
deltaLocation world loc dir = | |
case deltaDirection dir of | |
L dx dy = L ((loc.x + dx + world.dimX) `rem` world.dimX) | |
((loc.y + dy + world.dimY) `rem` world.dimY) | |
locations :: World -> [Location] | |
locations world = [L x y | x <- [0..world.dimX - 1], y <- [0..world.dimY - 1]] | |
evaporate :: World -> IO () | |
evaporate world = do | |
mapM_ evap $ concat world.cells | |
where | |
evap place = atomically $ do | |
cell <- place.read | |
place.write $ cell.{pher <- \x -> max 0 (x - 1)} | |
evaporator :: World -> IO () | |
evaporator world = do | |
Thread.sleep evaporationRate | |
evaporate world | |
createWorld :: Int -> Int -> Int -> Int -> IO (World, [AntAgent]) | |
createWorld dimX dimY nantsSqrt nfood = do | |
world <- atomically $ World.new dimX dimY | |
-- create food places | |
rng <- getStdGen | |
(prngX, prngY) = rng.split | |
(_, prngFood) = prngY.split | |
foodLocations = take nfood $ nub $ zipWith L | |
(Int.randomRs (25, 65 - 1) prngX) | |
(Int.randomRs (25, 65 - 1) prngY) | |
atomically $ zipWithM_ (setFood world) foodLocations (take nfood $ Int.randomRs foodRange prngFood) | |
-- create home | |
atomically $ mapM_ (setHome world) homeLocations | |
-- create ants and ant agents | |
rng <- getStdGen | |
rngs <- mapM (atomically • TVar.new) $ take nants $ makeRngs rng | |
turns <- sequence $ replicate nants $ Int.randomRIO (0, 7) | |
ants = zipWith Ant.turn (repeat (Ant N False 0)) turns | |
atomically $ zipWithM_ (setAnt world) ants homeLocations | |
agents = zipWith (AntAgent world) homeLocations rngs | |
return (world, agents) | |
where | |
nants = nantsSqrt * nantsSqrt | |
homeOffsetX = dimX `div` 4 | |
homeOffsetY = dimY `div` 4 | |
homeRangeX = [homeOffsetX..(homeOffsetX + nantsSqrt - 1)] | |
homeRangeY = [homeOffsetY..(homeOffsetY + nantsSqrt - 1)] | |
homeLocations = [L x y | x <- homeRangeX, y <- homeRangeY] | |
makeRngs :: StdGen -> [StdGen] | |
makeRngs rng = x : makeRngs y | |
where (x, y) = rng.split | |
--makeRnds rng = r : makeRnds g | |
-- where (r, g) = Double.random rng | |
setHome world loc = do | |
p = world.place loc | |
c <- p.read | |
p.write $ c.{home = True} | |
setFood world loc amt = do | |
p = world.place loc | |
c <- p.read | |
p.write $ c.{food=amt} | |
setAnt world ant loc = do | |
p = world.place loc | |
c <- p.read | |
p.write $ c.placeAnt ant.{excited=excitement} | |
drawWorld :: World -> GraphicsContext -> IO () | |
drawWorld world ctx = do | |
colorBg <- Color.white | |
ctx.setFill colorBg | |
ctx.fillRect 0 0 world.dimX.double world.dimY.double | |
cells <- atomically $ mapM TVar.read $ concat world.cells | |
sequence_ $ zipWith drawCell world.locations cells | |
homeColor <- Color.blue | |
ctx.setStroke homeColor | |
ctx.setLineWidth (1.0 / displayScale) | |
ctx.strokeRect homeOffsetX homeOffsetY nantsSqrt.double nantsSqrt.double | |
where | |
homeOffsetX = world.dimX.double / 4 | |
homeOffsetY = world.dimY.double / 4 | |
drawCell loc (Cell food pher ant home) = do | |
pherColor <- Color.rgb 0 255 0 ((min 1 (pher.double / pherScale))) :: IO Color | |
foodColor <- Color.rgb 255 0 0 ((min 100 food).double / 100) :: IO Color | |
fillCell loc.x loc.y pherColor | |
fillCell loc.x loc.y foodColor | |
case ant of | |
Just ant -> drawAnt loc.x.double loc.y.double ant | |
Nothing -> return () | |
where | |
fillCell x y c = do | |
ctx.setFill c | |
ctx.fillRect x.double y.double 1 1 | |
drawAnt x y (Ant dir food _) = do | |
color <- if food | |
then Color.red | |
else Color.black | |
ctx.setStroke color | |
ctx.setLineWidth (3 / displayScale) | |
ctx.strokeLine fx fy tx ty | |
where | |
fx = x + 0.5 + (lx.double * 0.5) | |
fy = y + 0.5 + (ly.double * 0.5) | |
tx = x + 0.5 -- (lx.double * 0.5) | |
ty = y + 0.5 -- (ly.double * 0.5) | |
L lx ly = deltaDirection dir | |
createCanvas :: Double -> IO Canvas | |
createCanvas scale = do | |
c <- Canvas.new width height :: IO Canvas | |
ctx <- c.getGraphicsContext2D :: IO GraphicsContext | |
ctx.scale scale scale | |
return c | |
buildUI :: TMVar (Maybe (GraphicsContext -> IO ())) -> TVar Bool -> Group -> Stage -> IO Group | |
buildUI mvar running root stage = do | |
stage.setTitle "Ant Colony" | |
canvas <- createCanvas displayScale | |
ctx <- canvas.getGraphicsContext2D :: IO GraphicsContext | |
forkOS (loop running $ animator mvar ctx) | |
root <: do | |
vbox <- VBox.new 5d :: IO VBox | |
kids <- vbox.getChildren | |
kids.add canvas | |
return vbox | |
where | |
animator :: TMVar (Maybe (GraphicsContext -> IO ())) -> GraphicsContext -> IO () | |
animator mvar ctx = do | |
mf <- atomically $ takeTMVar mvar | |
case mf of | |
Just f -> withUI $ f ctx | |
Nothing -> return () | |
displayScale = 8.0 | |
dim = 80 | |
dimX = dim | |
dimY = dim | |
width = dimX.double * displayScale | |
height = dimY.double * displayScale | |
nantsSqrt = 7 | |
nfood = 35 | |
foodRange = (30, 100) | |
antSleep = 20L | |
evaporationRate = 700L | |
pherScale = 20.0 | |
excitement = 50 | |
pherBoost = 5 | |
pherLimit = 40 | |
animate :: TMVar (Maybe (GraphicsContext -> IO ())) -> World -> IO () | |
animate pipe world = do | |
atomically $ putTMVar pipe $ Just $ drawWorld world | |
Thread.sleep 100 | |
loop :: TVar Bool -> IO () -> IO () | |
loop running f = do | |
r <- atomically $ running.read | |
if r | |
then (f >> loop running f) | |
else return () | |
reporter :: World -> IO () | |
reporter world = do | |
cells <- atomically $ mapM TVar.read $ concat world.cells | |
let nfood = (sum (map _.food cells)) + foodAnts | |
foodHome = sum $ map _.food $ filter _.home cells | |
foodWild = nfood - foodHome | |
foodAnts = fromIntegral $ length $ filter _.food ants | |
ants = map unJust $ filter isJust $ map _.ant cells | |
nants = fromIntegral $ length $ ants | |
phers = sum $ map _.pher cells | |
println ( | |
"Food: " ++ show nfood ++ " (" ++ | |
"home: " ++ show foodHome ++ ", " ++ | |
"outside: " ++ show foodWild ++ ") " ++ | |
"Ants: " ++ show nants ++ " (" ++ | |
"foraging: " ++ show (nants - foodAnts) ++ ", " ++ | |
"carring food: " ++ show foodAnts ++ ") " ++ | |
"Pheromones: " ++ show phers | |
) | |
main = do | |
running <- atomically $ TVar.new True | |
(world, agents) <- createWorld dimX dimY nantsSqrt nfood | |
mvar <- atomically $ newEmptyTMVar | |
forkOS (loop running $ animate mvar world) | |
forkOS (loop running $ evaporator world) | |
forkIO (Thread.sleep 2000 >> mapM_ (\a -> forkOS (a.act $ atomically $ running.read)) agents) | |
forkOS (loop running $ (reporter world >> Thread.sleep 2000)) | |
FregeFX.launch $ withStage (buildUI mvar running) | |
atomically $ running.write False | |
atomically $ putTMVar mvar Nothing | |
return () |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment