-
Notifications
You must be signed in to change notification settings - Fork 5
/
Main.hs
154 lines (125 loc) · 4.94 KB
/
Main.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
{-# LANGUAGE Arrows #-}
module Main where
import System.IO
import Data.List
import Data.Maybe
import Control.Monad.Random
import Control.Monad.Reader
--import Control.Monad.Trans.MSF
import Data.MonadicStreamFunction
import SIR
type Time = Double
type DTime = Double
type SIRAgentMSF g = MSF (ReaderT DTime (Rand g)) [SIRState] SIRState
contactRate :: Double
contactRate = 5.0
infectionProb :: Double
infectionProb = 0.05
illnessDuration :: Double
illnessDuration = 15.0
agentCount :: Int
agentCount = 100
infectedCount :: Int
infectedCount = 10
rngSeed :: Int
rngSeed = 42
dt :: DTime
dt = 1.0
t :: Time
t = 150
main :: IO ()
main = do
hSetBuffering stdout NoBuffering
let g = mkStdGen rngSeed
let as = initAgents agentCount infectedCount
let ass = runSimulationUntil g t dt as
let dyns = aggregateAllStates ass
let fileName = "SIR_DUNAI_DYNAMICS_" ++ show agentCount ++ "agents.m"
writeAggregatesToFile fileName dyns
runSimulationUntil :: RandomGen g => g -> Time -> DTime -> [SIRState] -> [[SIRState]]
runSimulationUntil g t dt as = evalRand (runReaderT ass dt) g -- runReader (evalRand ass g) dt -- evalRand (runReader ass dt) g
where
steps = floor $ t / dt
ticks = replicate steps ()
msfs = map sirAgent as
ass = embed (sirSimulation msfs as) ticks
sirSimulation :: RandomGen g => [SIRAgentMSF g] -> [SIRState] -> MSF (ReaderT DTime (Rand g)) () [SIRState]
sirSimulation msfs as = MSF $ \_a -> do
(as', msfs') <- foldM (sirSimulationAux as) ([], []) msfs
return (as, sirSimulation msfs' as')
where
sirSimulationAux :: [SIRState]
-> ([SIRState], [SIRAgentMSF g])
-> SIRAgentMSF g
-> ReaderT DTime (Rand g) ([SIRState], [SIRAgentMSF g])
sirSimulationAux as (accStates, accSfs) sf = do
(s, sf') <- unMSF sf as
return (s : accStates, sf' : accSfs)
sirAgent :: RandomGen g => SIRState -> SIRAgentMSF g
sirAgent Susceptible = susceptibleAgent
sirAgent Infected = infectedAgent
sirAgent Recovered = recoveredAgent
susceptibleAgent :: RandomGen g => SIRAgentMSF g
susceptibleAgent = switch susceptibleAgentInfectedEvent (const infectedAgent)
where
susceptibleAgentInfectedEvent :: RandomGen g => MSF (ReaderT DTime (Rand g)) [SIRState] (SIRState, Maybe ())
susceptibleAgentInfectedEvent = arrM susceptibleAgentInfectedEventAux
where
susceptibleAgentInfectedEventAux :: RandomGen g => [SIRState] -> ReaderT DTime (Rand g) (SIRState, Maybe ())
susceptibleAgentInfectedEventAux as = do
randContactCount <- lift $ randomExpM (1 / contactRate)
aInfs <- lift $ doTimes (floor randContactCount) (susceptibleAgentAux as)
let mayInf = find (Infected==) aInfs
if isJust mayInf
then return (Infected, Just ())
else return (Susceptible, Nothing)
susceptibleAgentAux :: RandomGen g => [SIRState] -> Rand g SIRState
susceptibleAgentAux as = do
randContact <- randomElem as
if (Infected == randContact)
then infect
else return Susceptible
infect :: RandomGen g => Rand g SIRState
infect = do
doInfect <- randomBoolM infectionProb
if doInfect
then return Infected
else return Susceptible
infectedAgent :: RandomGen g => SIRAgentMSF g
infectedAgent = switch infectedAgentRecoveredEvent (const recoveredAgent)
where
infectedAgentRecoveredEvent :: RandomGen g => MSF (ReaderT DTime (Rand g)) [SIRState] (SIRState, Maybe ())
infectedAgentRecoveredEvent = proc _ -> do
recEvt <- occasionally illnessDuration () -< ()
let a = maybe Infected (const Recovered) recEvt
returnA -< (a, recEvt)
recoveredAgent :: RandomGen g => SIRAgentMSF g
recoveredAgent = arr (const Recovered)
initAgents :: Int -> Int -> [SIRState]
initAgents n i = sus ++ inf
where
sus = replicate (n - i) Susceptible
inf = replicate i Infected
doTimes :: (Monad m) => Int -> m a -> m [a]
doTimes n f = forM [0..n - 1] (\_ -> f)
-- NOTE: is in spirit of the Yampa implementation
occasionally :: RandomGen g => Time -> b -> MSF (ReaderT DTime (Rand g)) a (Maybe b)
occasionally t_avg b
| t_avg > 0 = MSF (const tf)
| otherwise = error "AFRP: occasionally: Non-positive average interval."
where
-- Generally, if events occur with an average frequency of f, the
-- probability of at least one event occurring in an interval of t
-- is given by (1 - exp (-f*t)). The goal in the following is to
-- decide whether at least one event occurred in the interval of size
-- dt preceding the current sample point. For the first point,
-- we can think of the preceding interval as being 0, implying
-- no probability of an event occurring.
tf = do
dt <- ask
r <- lift $ getRandomR (0, 1)
let p = 1 - exp (-(dt / t_avg))
let evt = if r < p
then Just b
else Nothing
return (evt, MSF (const tf))