{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveAnyClass #-} {-# LANGUAGE OverloadedStrings #-} module Epidemic where import Control.Monad import qualified Data.ByteString as B import Data.ByteString.Internal (c2w) import Data.List (nub) import Data.Maybe (fromJust, isJust, isNothing) import qualified Data.Vector as V import Data.Word import Epidemic.Types.Events import Epidemic.Types.Parameter import Epidemic.Types.Population import Epidemic.Types.Simulation ( SimulationConfiguration(..) , SimulationRandEvent(..) , SimulationState(..) ) import Epidemic.Types.Time (AbsoluteTime(..), Timed(..), diracDeltaValue, nextTime) import GHC.Generics (Generic) import System.Random.MWC -- | The number of people added or removed in an event. eventPopDelta :: EpidemicEvent -> Integer eventPopDelta e = case e of Infection {} -> 1 Removal {} -> -1 IndividualSample {} -> -1 PopulationSample {..} -> fromIntegral $ numPeople popSampPeople StoppingTime -> 0 -- | The first scheduled event after a given time. firstScheduled :: AbsoluteTime -- ^ The given time -> Timed Probability -- ^ The information about all scheduled events -> Maybe (AbsoluteTime, Probability) firstScheduled time timedProb = do time' <- nextTime timedProb time prob' <- diracDeltaValue timedProb time' return (time', prob') -- | Predicate for whether there is a scheduled event during an interval. noScheduledEvent :: AbsoluteTime -- ^ Start time for interval -> AbsoluteTime -- ^ End time for interval -> Timed Probability -- ^ Information about all scheduled events -> Bool noScheduledEvent _ _ (Timed []) = True noScheduledEvent a b (Timed ((shedTime, _):scheduledEvents)) = not (a < shedTime && shedTime <= b) && noScheduledEvent a b (Timed scheduledEvents) -- | A list of the people involved in an 'EpidemicEvent'. personsInEvent :: EpidemicEvent -> [Person] personsInEvent e = case e of Infection _ p1 p2 -> [p1, p2] Removal _ p -> [p] (IndividualSample {..}) -> [indSampPerson] (PopulationSample {..}) -> V.toList personVec where (People personVec) = popSampPeople Extinction -> [] StoppingTime -> [] peopleInEvents :: [EpidemicEvent] -> People peopleInEvents events = People . V.fromList . nub . concat $ map personsInEvent events -- | Predicate for whether the first person infected the second in the given event infected :: Person -- ^ Potential infector -> Person -- ^ Potential infectee -> EpidemicEvent -- ^ Given event -> Bool infected p1 p2 e = case e of (Infection _ infector infectee) -> infector == p1 && infectee == p2 _ -> False -- | The people infected by a particular person in a list of events. infectedBy :: Person -- ^ Potential infector -> [EpidemicEvent] -- ^ Events -> People infectedBy person events = case events of [] -> People V.empty (Infection _ infector infectee:es) -> if infector == person then addPerson infectee $ infectedBy person es else infectedBy person es (_:es) -> infectedBy person es -- | Run the simulation and return a @SimulationState@ which holds the history -- of the simulation. allEvents :: (ModelParameters a b, Population b) => SimulationRandEvent a b -> a -> AbsoluteTime -> Maybe (b -> Bool) -- ^ predicate for a valid population -> SimulationState b -> GenIO -> IO (SimulationState b) allEvents _ _ _ _ TerminatedSimulation _ = return TerminatedSimulation allEvents simRandEvent@(SimulationRandEvent randEvent) modelParams maxTime maybePopPredicate (SimulationState (currTime, currEvents, currPop, currId)) gen = if isNothing maybePopPredicate || (isJust maybePopPredicate && fromJust maybePopPredicate currPop) then if isInfected currPop then do (newTime, event, newPop, newId) <- randEvent modelParams currTime currPop currId gen if newTime < maxTime then allEvents simRandEvent modelParams maxTime maybePopPredicate (SimulationState (newTime, event : currEvents, newPop, newId)) gen else return $ SimulationState (maxTime, StoppingTime : currEvents, currPop, currId) else return $ SimulationState (currTime, Extinction : currEvents, currPop, currId) else return TerminatedSimulation