{-# LANGUAGE RecordWildCards #-}
module Epidemic where
import Data.List (nub)
import Data.Maybe (fromJust, isJust, isNothing)
import qualified Data.Vector as V
import Epidemic.Types.Events
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import Epidemic.Types.Simulation (SimulationRandEvent (..),
SimulationState (..),
TerminationHandler (..))
import Epidemic.Types.Time (AbsoluteTime (..), Timed (..),
diracDeltaValue, nextTime)
import System.Random.MWC
eventPopDelta :: EpidemicEvent -> Integer
eventPopDelta :: EpidemicEvent -> Integer
eventPopDelta EpidemicEvent
e =
case EpidemicEvent
e of
Infection {} -> Integer
1
Removal {} -> -Integer
1
IndividualSample {} -> -Integer
1
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ People -> Int
numPeople People
popSampPeople
StoppingTime {} -> Integer
0
Extinction {} -> Integer
0
firstScheduled ::
AbsoluteTime
-> Timed Probability
-> Maybe (AbsoluteTime, Probability)
firstScheduled :: AbsoluteTime
-> Timed Probability -> Maybe (AbsoluteTime, Probability)
firstScheduled AbsoluteTime
time Timed Probability
timedProb = do
AbsoluteTime
time' <- Timed Probability -> AbsoluteTime -> Maybe AbsoluteTime
forall a. Timed a -> AbsoluteTime -> Maybe AbsoluteTime
nextTime Timed Probability
timedProb AbsoluteTime
time
Probability
prob' <- Timed Probability -> AbsoluteTime -> Maybe Probability
forall a. Timed a -> AbsoluteTime -> Maybe a
diracDeltaValue Timed Probability
timedProb AbsoluteTime
time'
(AbsoluteTime, Probability) -> Maybe (AbsoluteTime, Probability)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
time', Probability
prob')
noScheduledEvent ::
AbsoluteTime
-> AbsoluteTime
-> Timed Probability
-> Bool
noScheduledEvent :: AbsoluteTime -> AbsoluteTime -> Timed Probability -> Bool
noScheduledEvent AbsoluteTime
_ AbsoluteTime
_ (Timed []) = Bool
True
noScheduledEvent AbsoluteTime
a AbsoluteTime
b (Timed ((AbsoluteTime
shedTime, Probability
_):[(AbsoluteTime, Probability)]
scheduledEvents)) =
Bool -> Bool
not (AbsoluteTime
a AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< AbsoluteTime
shedTime Bool -> Bool -> Bool
&& AbsoluteTime
shedTime AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
<= AbsoluteTime
b) Bool -> Bool -> Bool
&&
AbsoluteTime -> AbsoluteTime -> Timed Probability -> Bool
noScheduledEvent AbsoluteTime
a AbsoluteTime
b ([(AbsoluteTime, Probability)] -> Timed Probability
forall a. [(AbsoluteTime, a)] -> Timed a
Timed [(AbsoluteTime, Probability)]
scheduledEvents)
personsInEvent :: EpidemicEvent -> [Person]
personsInEvent :: EpidemicEvent -> [Person]
personsInEvent EpidemicEvent
e =
case EpidemicEvent
e of
Infection AbsoluteTime
_ Person
p1 Person
p2 -> [Person
p1, Person
p2]
Removal AbsoluteTime
_ Person
p -> [Person
p]
IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
..} -> [Person
indSampPerson]
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Vector Person -> [Person]
forall a. Vector a -> [a]
V.toList Vector Person
personVec
where (People Vector Person
personVec) = People
popSampPeople
Extinction {} -> []
StoppingTime {} -> []
peopleInEvents :: [EpidemicEvent] -> People
peopleInEvents :: [EpidemicEvent] -> People
peopleInEvents [EpidemicEvent]
events =
Vector Person -> People
People (Vector Person -> People)
-> ([[Person]] -> Vector Person) -> [[Person]] -> People
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Person] -> Vector Person
forall a. [a] -> Vector a
V.fromList ([Person] -> Vector Person)
-> ([[Person]] -> [Person]) -> [[Person]] -> Vector Person
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Person] -> [Person]
forall a. Eq a => [a] -> [a]
nub ([Person] -> [Person])
-> ([[Person]] -> [Person]) -> [[Person]] -> [Person]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Person]] -> [Person]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Person]] -> People) -> [[Person]] -> People
forall a b. (a -> b) -> a -> b
$ (EpidemicEvent -> [Person]) -> [EpidemicEvent] -> [[Person]]
forall a b. (a -> b) -> [a] -> [b]
map EpidemicEvent -> [Person]
personsInEvent [EpidemicEvent]
events
infected ::
Person
-> Person
-> EpidemicEvent
-> Bool
infected :: Person -> Person -> EpidemicEvent -> Bool
infected Person
p1 Person
p2 EpidemicEvent
e =
case EpidemicEvent
e of
(Infection AbsoluteTime
_ Person
infector Person
infectee) -> Person
infector Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
p1 Bool -> Bool -> Bool
&& Person
infectee Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
p2
EpidemicEvent
_ -> Bool
False
infectedBy ::
Person
-> [EpidemicEvent]
-> People
infectedBy :: Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
events =
case [EpidemicEvent]
events of
[] -> Vector Person -> People
People Vector Person
forall a. Vector a
V.empty
(Infection AbsoluteTime
_ Person
infector Person
infectee:[EpidemicEvent]
es) ->
if Person
infector Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
person
then Person -> People -> People
addPerson Person
infectee (People -> People) -> People -> People
forall a b. (a -> b) -> a -> b
$ Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es
else Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es
(EpidemicEvent
_:[EpidemicEvent]
es) -> Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es
allEvents ::
(ModelParameters a b, Population b)
=> SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEvents :: SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEvents SimulationRandEvent a b
_ a
_ AbsoluteTime
_ Maybe (TerminationHandler b c)
_ ts :: SimulationState b c
ts@(TerminatedSimulation Maybe c
_) GenIO
_ = SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return SimulationState b c
ts
allEvents (SimulationRandEvent a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent) a
modelParams AbsoluteTime
maxTime Maybe (TerminationHandler b c)
maybeTermHandler (SimulationState (AbsoluteTime
currTime, [EpidemicEvent]
currEvents, b
currPop, Identifier
currId)) GenIO
gen =
let isNotTerminated :: b -> Bool
isNotTerminated = case Maybe (TerminationHandler b c)
maybeTermHandler of
Maybe (TerminationHandler b c)
Nothing -> Bool -> b -> Bool
forall a b. a -> b -> a
const Bool
True
Just (TerminationHandler b -> Bool
hasTerminated [EpidemicEvent] -> c
_) -> Bool -> Bool
not (Bool -> Bool) -> (b -> Bool) -> b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
hasTerminated
in if b -> Bool
isNotTerminated b
currPop
then if b -> Bool
forall a. Population a => a -> Bool
isInfected b
currPop
then do
(AbsoluteTime
newTime, EpidemicEvent
event, b
newPop, Identifier
newId) <-
a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent a
modelParams AbsoluteTime
currTime b
currPop Identifier
currId GenIO
gen
if AbsoluteTime
newTime AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< AbsoluteTime
maxTime
then SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
forall a b c.
(ModelParameters a b, Population b) =>
SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEvents
((a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
forall a b.
(ModelParameters a b, Population b) =>
(a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent)
a
modelParams
AbsoluteTime
maxTime
Maybe (TerminationHandler b c)
maybeTermHandler
((AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
(AbsoluteTime
newTime, EpidemicEvent
event EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents, b
newPop, Identifier
newId))
GenIO
gen
else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> SimulationState b c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
( AbsoluteTime
maxTime
, AbsoluteTime -> EpidemicEvent
StoppingTime AbsoluteTime
maxTime EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents
, b
currPop
, Identifier
currId)
else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> SimulationState b c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
( AbsoluteTime
currTime
, AbsoluteTime -> EpidemicEvent
Extinction AbsoluteTime
currTime EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents
, b
currPop
, Identifier
currId)
else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> (Maybe c -> SimulationState b c)
-> Maybe c
-> IO (SimulationState b c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe c -> SimulationState b c
forall b c. Maybe c -> SimulationState b c
TerminatedSimulation (Maybe c -> IO (SimulationState b c))
-> Maybe c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$ do TerminationHandler b -> Bool
_ [EpidemicEvent] -> c
termSummary <- Maybe (TerminationHandler b c)
maybeTermHandler
c -> Maybe c
forall (m :: * -> *) a. Monad m => a -> m a
return (c -> Maybe c) -> c -> Maybe c
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> c
termSummary [EpidemicEvent]
currEvents