diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -13,3 +13,22 @@ simulation models already provided in the `Epidemic.Model` submodules. For the most part, writing a new epidemic model revolves around definining the `randomEvent` function. + +### Time + +For absolute times there is the `AbsoluteTime` type and for differences between +times there is the `TimeDelta` type. For quantities that vary across time the +`Timed` type is a way to represent piecewise constant functions and there are +several helper functions to query these objects. + +There is the `TimeStamp` type class for things that have an absolute time +associated with them. + +### Simulation + +The `TerminationHandler` type is used to control early termination of a +simulation. The use case of this is that you can terminate a simulation once it +reaches a state specified by a predicate. If this occurs then the simulation is +terminated and a summary function is applied to the simulation and this value is +returned. This is particularly useful if you want to terminate simulations that +have exploded and threaten to eat up all the memory on your machine. diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,50 @@ # Changelog for epi-sim +## 0.7.0 + +- Simulation functions (eg `simulationWithGenIO`) return either the simulated + events of the summary produced by the termination handler if the simulation + terminated early. +- The `InhomBDSCODPop` now records the number of individuals that have been + removed and there are getter functions exported to access this information. +- Add `TerminationHandler` to provide better control over early termination of + simulations. For example, this makes it possible to terminate the simulation + early if certain stopping conditions are met and to call a function that + summarises why the simulation is being terminated. This will break old code in + that the `configuration` functions provided by the example models all now have + an additional parameter of type `Maybe TerminationHandler`. + +## 0.6.0 + +- Improve documentation. +- Allow a flexible start time of the simulation so it does not assume a start + time of zero. +- Rename a bunch of the simulation functions so their use case is clearer. + +## 0.5.2 + +- Include helper functions in `Epidemic.Types.Simulation` to make it easier to + create PRNG with or without a fixed seed. + +## 0.5.1 + +- Bug fix and tidy up some code. + +## 0.5.0 + +- Add absolute times to the extinction and stopping time events to provide a + consistent interface. +- Add the `aggregated` function to help aggregated individual level samples into + population level samples. This is tested in `aggregationTests`. +- Add `TimeStamp` type class to abstract working with types that have an + absolute time associated with them. +- Add `TimeInterval` type for working with intervals of time, there are also + some helper functions to make it easier to work with intervals. +- Add the `maybeNextTimed` helper function and clean up some code in the `Time` + module. +- Extend `ModelParameters` class to have an `eventWeights` to provide a vector + of event weights for computing which event actually occurred. + ## 0.4.2 - Include `simulationWithGenIO` and add `scRequireCherry` to the diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,13 +1,21 @@ # epi-sim -A library for simulating epidemics, with a focus on phylodynamics and -observation models. +A library for simulating stochastic epidemic models, with a focus on +phylodynamics and observation models. + +## Epidemic events + +A realisation of the model is represented by a list of `EpidemicEvent`s which +describe the events that occurred. + ## Available models Although this package supports the definition of new models there are some that are implemented already in the `Epidemic.Model` module. Implemented models include: -1. Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see `Epidemic.Model.BDSCOD`) -2. Inhomogeneous Birth-Death-Sampling (see `Epidemic.Model.InhomogeneousBDS`) -3. Logistic Birth-Death-Sampling-Disaster (see `Epidemic.Model.LogisticBDSD`) +1. Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see `BDSCOD.hs`) +2. Inhomogeneous Birth-Death-Sampling (see `InhomogeneousBDS.hs`) +3. Inhomogeneous Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see + `InhomogeneousBDSCOD.hs`) +4. Logistic Birth-Death-Sampling-Disaster (see `LogisticBDSD.hs`) diff --git a/epi-sim.cabal b/epi-sim.cabal --- a/epi-sim.cabal +++ b/epi-sim.cabal @@ -1,6 +1,6 @@ cabal-version: 1.22 name: epi-sim -version: 0.4.2 +version: 0.7.0 synopsis: A library for simulating epidemics as birth-death processes. @@ -14,8 +14,11 @@ . 1. Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see `Epidemic.Model.BDSCOD`) 2. Inhomogeneous Birth-Death-Sampling (see `Epidemic.Model.InhomogeneousBDS`) - 3. Logistic Birth-Death-Sampling-Disaster (see `Epidemic.Model.LogisticBDSD`) + 3. Inhomogeneous BDSCOD (see `Epidemic.Model.InhomogeneousBDSCOD`) + 4. Logistic Birth-Death-Sampling-Disaster (see `Epidemic.Model.LogisticBDSD`) . + There are more details in the documentation of the "Epidemic" module. + . homepage: https://github.com/aezarebski/epi-sim#readme bug-reports: https://github.com/aezarebski/epi-sim/issues @@ -27,9 +30,9 @@ build-type: Simple category: Simulation extra-source-files: + ARCHITECTURE.md ChangeLog.md README.md - ARCHITECTURE.md source-repository head type: git @@ -40,6 +43,7 @@ Epidemic Epidemic.Model.BDSCOD Epidemic.Model.InhomogeneousBDS + Epidemic.Model.InhomogeneousBDSCOD Epidemic.Model.LogisticBDSD Epidemic.Types.Events Epidemic.Types.Newick diff --git a/src/Epidemic.hs b/src/Epidemic.hs --- a/src/Epidemic.hs +++ b/src/Epidemic.hs @@ -1,38 +1,57 @@ {-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE DeriveAnyClass #-} -{-# LANGUAGE OverloadedStrings #-} +-- | +-- Module: Epidemic +-- Copyright: (c) 2021 Alexander E. Zarebski +-- License: MIT +-- +-- Maintainer: Alexander E. Zarebski +-- Stability: unstable +-- Portability: ghc +-- +-- This package provides functionality for simulating stochastic epidemic +-- models, in particular those that are of interest in phylodynamics. There are +-- several models provided by the package, eg @Epidemic.Model.BDSCOD@, however +-- there should be the basic functionality to implement a wide range of models +-- available. Each of the models included in the package provide a +-- @configuration@ function which can be used to get a 'SimulationConfiguration' +-- and a @randomEvent@ function which returns a 'SimulationRandEvent'. With +-- these you can then use 'allEvents' to get all of the events in a simulation. +-- +-- This package also provides some functionality for working with observation +-- models, both epidemiological and phylogenetic. 'Observation' values are used +-- to describe the possible observation of an 'EpidemicEvent'. +-- +-- There is an example of how to use this package in the documentation for +-- "Epidemic.Model.InhomogeneousBDSCOD". + 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 +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 --- | The number of people added or removed in an event. +-- | The number of people added or removed in an event. In the case of an +-- extinction event the number of people removed is arbitrarily set to zero +-- because this information is available from the prior event in the sequence. eventPopDelta :: EpidemicEvent -> Integer eventPopDelta e = case e of - Infection {} -> 1 - Removal {} -> -1 - IndividualSample {} -> -1 + Infection {} -> 1 + Removal {} -> -1 + IndividualSample {} -> -1 PopulationSample {..} -> fromIntegral $ numPeople popSampPeople - StoppingTime -> 0 + StoppingTime {} -> 0 + Extinction {} -> 0 -- | The first scheduled event after a given time. firstScheduled :: @@ -44,7 +63,19 @@ prob' <- diracDeltaValue timedProb time' return (time', prob') --- | Predicate for whether there is a scheduled event during an interval. +-- | Predicate for whether there is a scheduled event during an interval. NOTE +-- that this does not consider events that happen at the start of the interval +-- as occurring between the times. +-- +-- >>> tA = AbsoluteTime 1.0 +-- >>> tB = AbsoluteTime 2.0 +-- >>> noScheduledEvent tA tB <$> asTimed [(AbsoluteTime 1.5, 0.5)] +-- Just False +-- >>> noScheduledEvent tA tB <$> asTimed [(AbsoluteTime 2.5, 0.5)] +-- Just True +-- >>> noScheduledEvent tA tB <$> asTimed [(tA, 0.5)] +-- Just True +-- noScheduledEvent :: AbsoluteTime -- ^ Start time for interval -> AbsoluteTime -- ^ End time for interval @@ -61,13 +92,11 @@ case e of Infection _ p1 p2 -> [p1, p2] Removal _ p -> [p] - (IndividualSample {..}) -> [indSampPerson] - (PopulationSample {..}) -> - V.toList personVec - where - (People personVec) = popSampPeople - Extinction -> [] - StoppingTime -> [] + IndividualSample {..} -> [indSampPerson] + PopulationSample {..} -> V.toList personVec + where (People personVec) = popSampPeople + Extinction {} -> [] + StoppingTime {} -> [] peopleInEvents :: [EpidemicEvent] -> People peopleInEvents events = @@ -82,7 +111,7 @@ infected p1 p2 e = case e of (Infection _ infector infectee) -> infector == p1 && infectee == p2 - _ -> False + _ -> False -- | The people infected by a particular person in a list of events. infectedBy :: @@ -98,38 +127,47 @@ else infectedBy person es (_:es) -> infectedBy person es --- | Run the simulation and return a @SimulationState@ which holds the history --- of the simulation. +-- | Run the simulation until the specified stopping time 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 + -> AbsoluteTime -- ^ time at which to stop the simulation + -> Maybe (TerminationHandler b c) + -> SimulationState b c -- ^ the initial/current state of the simulation -> 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 + -> IO (SimulationState b c) +allEvents _ _ _ _ ts@(TerminatedSimulation _) _ = return ts +allEvents (SimulationRandEvent randEvent) modelParams maxTime maybeTermHandler (SimulationState (currTime, currEvents, currPop, currId)) gen = + let isNotTerminated = case maybeTermHandler of + Nothing -> const True + Just (TerminationHandler hasTerminated _) -> not . hasTerminated + in if isNotTerminated currPop + then if isInfected currPop then do (newTime, event, newPop, newId) <- randEvent modelParams currTime currPop currId gen if newTime < maxTime then allEvents - simRandEvent + (SimulationRandEvent randEvent) modelParams maxTime - maybePopPredicate + maybeTermHandler (SimulationState (newTime, event : currEvents, newPop, newId)) gen else return $ SimulationState - (maxTime, StoppingTime : currEvents, currPop, currId) + ( maxTime + , StoppingTime maxTime : currEvents + , currPop + , currId) else return $ SimulationState - (currTime, Extinction : currEvents, currPop, currId) - else return TerminatedSimulation + ( currTime + , Extinction currTime : currEvents + , currPop + , currId) + else return . TerminatedSimulation $ do TerminationHandler _ termSummary <- maybeTermHandler + return $ termSummary currEvents diff --git a/src/Epidemic/Model/BDSCOD.hs b/src/Epidemic/Model/BDSCOD.hs --- a/src/Epidemic/Model/BDSCOD.hs +++ b/src/Epidemic/Model/BDSCOD.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE MultiParamTypeClasses #-} module Epidemic.Model.BDSCOD @@ -8,50 +7,45 @@ , BDSCODPopulation(..) ) where -import Data.List (nub) -import Data.Maybe (fromJust, isJust, isNothing) import qualified Data.Vector as V import qualified Data.Vector.Generic as G import Epidemic -import Epidemic.Types.Events - ( EpidemicEvent(..) - , maybeEpidemicTree - ) +import Epidemic.Types.Events (EpidemicEvent(..)) import Epidemic.Types.Parameter import Epidemic.Types.Population import Epidemic.Types.Simulation ( SimulationConfiguration(..) - , SimulationRandEvent(..) + , SimulationRandEvent(..), TerminationHandler(..) ) import Epidemic.Types.Time ( AbsoluteTime(..) , TimeDelta(..) , Timed(..) - , cadlagValue - , diracDeltaValue - , nextTime - , timeAfterDelta - , allTimes , asTimed + , maybeNextTimed + , timeAfterDelta ) import Epidemic.Utility import System.Random.MWC import System.Random.MWC.Distributions (bernoulli, categorical, exponential) -- | birth rate, death rate, sampling rate, catastrophe specification, occurrence rate and disaster specification -data BDSCODParameters = BDSCODParameters Rate Rate Rate (Timed Probability) Rate (Timed Probability) +data BDSCODParameters = + BDSCODParameters Rate Rate Rate (Timed Probability) Rate (Timed Probability) data BDSCODPopulation = BDSCODPopulation People deriving (Show) instance ModelParameters BDSCODParameters BDSCODPopulation where - rNaught _ (BDSCODParameters birthRate deathRate samplingRate _ occurrenceRate _) _ = - Just $ birthRate / (deathRate + samplingRate + occurrenceRate) - eventRate _ (BDSCODParameters birthRate deathRate samplingRate _ occurrenceRate _) _ = - Just $ birthRate + deathRate + samplingRate + occurrenceRate - birthProb _ (BDSCODParameters birthRate deathRate samplingRate _ occurrenceRate _) _ = - Just $ birthRate / (birthRate + deathRate + samplingRate + occurrenceRate) + rNaught _ (BDSCODParameters br dr sRate _ occRate _) _ = + Just $ br / (dr + sRate + occRate) + eventRate _ (BDSCODParameters br dr sRate _ occRate _) _ = + Just $ br + dr + sRate + occRate + birthProb _ (BDSCODParameters br dr sRate _ occRate _) _ = + Just $ br / (br + dr + sRate + occRate) + eventWeights _ (BDSCODParameters br dr sRate _ occRate _) _ = + Just $ V.fromList [br, dr, sRate, occRate] instance Population BDSCODPopulation where susceptiblePeople _ = Nothing @@ -63,14 +57,15 @@ configuration :: TimeDelta -- ^ Duration of the simulation -> Bool -- ^ condition upon at least two sequenced samples. + -> Maybe (BDSCODPopulation -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling. -> ( Rate , Rate , Rate , [(AbsoluteTime, Probability)] , Rate , [(AbsoluteTime, Probability)]) -- ^ Birth, Death, Sampling, Catastrophe probability, Occurrence rates and Disaster probabilities - -> Maybe (SimulationConfiguration BDSCODParameters BDSCODPopulation) -configuration maxTime atLeastCherry (birthRate, deathRate, samplingRate, catastropheSpec, occurrenceRate, disasterSpec) = do + -> Maybe (SimulationConfiguration BDSCODParameters BDSCODPopulation s) +configuration maxTime atLeastCherry maybeTHFuncs (birthRate, deathRate, samplingRate, catastropheSpec, occurrenceRate, disasterSpec) = do catastropheSpec' <- asTimed catastropheSpec disasterSpec' <- asTimed disasterSpec let bdscodParams = @@ -83,6 +78,8 @@ disasterSpec' (seedPerson, newId) = newPerson initialIdentifier bdscodPop = BDSCODPopulation (People $ V.singleton seedPerson) + termHandler = do (f1, f2) <- maybeTHFuncs + return $ TerminationHandler f1 f2 in return $ SimulationConfiguration bdscodParams @@ -90,7 +87,7 @@ newId (AbsoluteTime 0) maxTime - Nothing + termHandler atLeastCherry -- | The way in which random events are generated in this model. @@ -98,53 +95,73 @@ randomEvent = SimulationRandEvent randomEvent' -- | Return a random event from the BDSCOD-process given the current state of the process. -randomEvent' :: BDSCODParameters -- ^ Parameters of the process - -> AbsoluteTime -- ^ The current time within the process - -> BDSCODPopulation -- ^ The current state of the populaion - -> Identifier -- ^ The current state of the identifier generator - -> GenIO -- ^ The current state of the PRNG - -> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier) -randomEvent' params@(BDSCODParameters br dr sr catastInfo occr disastInfo) currTime currPop@(BDSCODPopulation currPeople) currId gen = - let netEventRate = fromJust $ eventRate currPop params currTime - eventWeights = V.fromList [br, dr, sr, occr] - in do delay <- exponential (fromIntegral (numPeople currPeople) * netEventRate) gen +randomEvent' :: + BDSCODParameters -- ^ Parameters of the process + -> AbsoluteTime -- ^ The current time within the process + -> BDSCODPopulation -- ^ The current state of the populaion + -> Identifier -- ^ The current state of the identifier generator + -> GenIO -- ^ The current state of the PRNG + -> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier) +randomEvent' params@(BDSCODParameters _ _ _ catastInfo _ disastInfo) currTime currPop@(BDSCODPopulation currPeople) currId gen = + let (Just netEventRate) = eventRate currPop params currTime + (Just weightVec) = eventWeights currPop params currTime + in do delay <- + exponential (fromIntegral (numPeople currPeople) * netEventRate) gen let newEventTime = timeAfterDelta currTime (TimeDelta delay) if noScheduledEvent currTime newEventTime (catastInfo <> disastInfo) - then do eventIx <- categorical eventWeights gen - (selectedPerson, unselectedPeople) <- randomPerson currPeople gen - return $ case eventIx of - 0 -> let (birthedPerson, newId) = newPerson currId - infEvent = Infection newEventTime selectedPerson birthedPerson - in ( newEventTime - , infEvent - , BDSCODPopulation (addPerson birthedPerson currPeople) - , newId) - 1 -> (newEventTime, Removal newEventTime selectedPerson, BDSCODPopulation unselectedPeople, currId) - 2 -> (newEventTime, IndividualSample newEventTime selectedPerson True, BDSCODPopulation unselectedPeople, currId) - 3 -> (newEventTime, IndividualSample newEventTime selectedPerson False, BDSCODPopulation unselectedPeople, currId) - _ -> error "no birth, death, sampling, occurrence event selected." - - else if noScheduledEvent currTime newEventTime catastInfo - then let (Just (disastTime,disastProb)) = firstScheduled currTime disastInfo - in do (disastEvent,postDisastPop) <- randomDisasterEvent (disastTime,disastProb) currPop gen - return (disastTime,disastEvent,postDisastPop,currId) - else if noScheduledEvent currTime newEventTime disastInfo - then let (Just (catastTime,catastProb)) = firstScheduled currTime catastInfo - in do (catastEvent,postCatastPop) <- randomCatastropheEvent (catastTime,catastProb) currPop gen - return (catastTime,catastEvent,postCatastPop,currId) - else let (Just (catastTime,catastProb)) = firstScheduled currTime catastInfo - (Just (disastTime,disastProb)) = firstScheduled currTime disastInfo - in do (scheduledEvent,postEventPop) <- if catastTime < disastTime then - randomCatastropheEvent (catastTime,catastProb) currPop gen else - randomDisasterEvent (disastTime,disastProb) currPop gen - return (min catastTime disastTime,scheduledEvent,postEventPop,currId) - + then do + eventIx <- categorical weightVec gen + (selectedPerson, unselectedPeople) <- randomPerson currPeople gen + return $ + case eventIx of + 0 -> + let (birthedPerson, newId) = newPerson currId + infEvent = + Infection newEventTime selectedPerson birthedPerson + in ( newEventTime + , infEvent + , BDSCODPopulation (addPerson birthedPerson currPeople) + , newId) + 1 -> + ( newEventTime + , Removal newEventTime selectedPerson + , BDSCODPopulation unselectedPeople + , currId) + 2 -> + ( newEventTime + , IndividualSample newEventTime selectedPerson True + , BDSCODPopulation unselectedPeople + , currId) + 3 -> + ( newEventTime + , IndividualSample newEventTime selectedPerson False + , BDSCODPopulation unselectedPeople + , currId) + _ -> + error "no birth, death, sampling, occurrence event selected." + else case maybeNextTimed catastInfo disastInfo currTime of + Just (disastTime, Right disastProb) -> + do (disastEvent, postDisastPop) <- + randomDisasterEvent + (disastTime, disastProb) + currPop + gen + return (disastTime, disastEvent, postDisastPop, currId) + Just (catastTime, Left catastProb) -> + do (catastEvent, postCatastPop) <- + randomCatastropheEvent + (catastTime, catastProb) + currPop + gen + return (catastTime, catastEvent, postCatastPop, currId) + Nothing -> error "Missing a next scheduled event when there should be one." -- | Return a randomly sampled Catastrophe event -randomCatastropheEvent :: (AbsoluteTime,Probability) -- ^ Time and probability of sampling in the catastrophe - -> BDSCODPopulation -- ^ The state of the population prior to the catastrophe - -> GenIO - -> IO (EpidemicEvent,BDSCODPopulation) +randomCatastropheEvent :: + (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the catastrophe + -> BDSCODPopulation -- ^ The state of the population prior to the catastrophe + -> GenIO + -> IO (EpidemicEvent, BDSCODPopulation) randomCatastropheEvent (catastTime, rhoProb) (BDSCODPopulation (People currPeople)) gen = do rhoBernoullis <- G.replicateM (V.length currPeople) (bernoulli rhoProb gen) let filterZip predicate a b = fst . V.unzip . V.filter predicate $ V.zip a b @@ -156,10 +173,11 @@ -- | Return a randomly sampled Disaster event -- TODO Move this into the epidemic module to keep things DRY. -randomDisasterEvent :: (AbsoluteTime,Probability) -- ^ Time and probability of sampling in the disaster - -> BDSCODPopulation -- ^ The state of the population prior to the disaster - -> GenIO - -> IO (EpidemicEvent,BDSCODPopulation) +randomDisasterEvent :: + (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the disaster + -> BDSCODPopulation -- ^ The state of the population prior to the disaster + -> GenIO + -> IO (EpidemicEvent, BDSCODPopulation) randomDisasterEvent (disastTime, nuProb) (BDSCODPopulation (People currPeople)) gen = do nuBernoullis <- G.replicateM (V.length currPeople) (bernoulli nuProb gen) let filterZip predicate a b = fst . V.unzip . V.filter predicate $ V.zip a b diff --git a/src/Epidemic/Model/InhomogeneousBDS.hs b/src/Epidemic/Model/InhomogeneousBDS.hs --- a/src/Epidemic/Model/InhomogeneousBDS.hs +++ b/src/Epidemic/Model/InhomogeneousBDS.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE MultiParamTypeClasses #-} module Epidemic.Model.InhomogeneousBDS @@ -15,30 +14,22 @@ , TimeDelta(..) , allTimes , asTimed - , diracDeltaValue - , nextTime , cadlagValue - , timeAfterDelta ) -import Control.Monad (liftM) -import Data.Maybe (fromJust, isJust, isNothing) +import Data.Maybe (fromJust) import qualified Data.Vector as V -import Epidemic import Epidemic.Types.Events ( EpidemicEvent(..) - , maybeEpidemicTree ) import Epidemic.Types.Parameter import Epidemic.Types.Population -import Epidemic.Types.Observations import Epidemic.Types.Simulation ( SimulationConfiguration(..) - , SimulationRandEvent(..) - , SimulationState(..) + , SimulationRandEvent(..), TerminationHandler(..) ) import Epidemic.Utility import System.Random.MWC -import System.Random.MWC.Distributions (categorical, exponential) +import System.Random.MWC.Distributions (categorical) data InhomBDSRates = InhomBDSRates (Timed Rate) Rate Rate @@ -50,13 +41,15 @@ instance ModelParameters InhomBDSRates InhomBDSPop where rNaught _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = let birthRate = cadlagValue timedBirthRate time - in liftM (/ (deathRate + sampleRate)) birthRate + in (/ (deathRate + sampleRate)) <$> birthRate eventRate _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = let birthRate = cadlagValue timedBirthRate time - in liftM (+ (deathRate + sampleRate)) birthRate + in (+ (deathRate + sampleRate)) <$> birthRate birthProb _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = - liftM (\br -> br / (br + deathRate + sampleRate)) $ + (\br -> br / (br + deathRate + sampleRate)) <$> cadlagValue timedBirthRate time + eventWeights _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = + Just $ V.fromList [fromJust (cadlagValue timedBirthRate time), deathRate, sampleRate] instance Population InhomBDSPop where susceptiblePeople _ = Nothing @@ -85,11 +78,14 @@ configuration :: TimeDelta -- ^ Duration of the simulation after starting at time 0. -> Bool -- ^ condition upon at least two sequenced samples. + -> Maybe (InhomBDSPop -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling. -> ([(AbsoluteTime, Rate)], Rate, Rate) -- ^ Birth, Death and Sampling rates - -> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop) -configuration maxTime atLeastCherry (tBrPairs, deathRate, sampleRate) = + -> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop s) +configuration maxTime atLeastCherry maybeTHFuncs (tBrPairs, deathRate, sampleRate) = let (seedPerson, newId) = newPerson initialIdentifier bdsPop = InhomBDSPop (People $ V.singleton seedPerson) + termHandler = do (f1, f2) <- maybeTHFuncs + return $ TerminationHandler f1 f2 in do timedBirthRate <- asTimed tBrPairs maybeIBDSRates <- inhomBDSRates timedBirthRate deathRate sampleRate if maxTime > TimeDelta 0 @@ -100,7 +96,7 @@ newId (AbsoluteTime 0) maxTime - Nothing + termHandler atLeastCherry) else Nothing @@ -115,9 +111,10 @@ -> Identifier -- ^ current identifier -> GenIO -- ^ PRNG -> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier) -randomEvent' inhomRates@(InhomBDSRates brts dr sr) currTime pop@(InhomBDSPop (people@(People peopleVec))) currId gen = +randomEvent' inhomRates@(InhomBDSRates brts _ _) currTime pop@(InhomBDSPop people) currId gen = let popSize = fromIntegral $ numPeople people :: Double - eventWeights t = V.fromList [fromJust (cadlagValue brts t), dr, sr] + --weightVecFunc :: AbsoluteTime -> Maybe (Vector Double) + weightVecFunc = eventWeights pop inhomRates -- we need a new step function to account for the population size. (Just stepFunction) = asTimed @@ -125,7 +122,7 @@ | t <- allTimes brts ] in do (Just newEventTime) <- inhomExponential stepFunction currTime gen - eventIx <- categorical (eventWeights newEventTime) gen + eventIx <- categorical (fromJust $ weightVecFunc newEventTime) gen (selectedPerson, unselectedPeople) <- randomPerson people gen return $ case eventIx of diff --git a/src/Epidemic/Model/InhomogeneousBDSCOD.hs b/src/Epidemic/Model/InhomogeneousBDSCOD.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Model/InhomogeneousBDSCOD.hs @@ -0,0 +1,315 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RecordWildCards #-} + +-- | +-- Module: Epidemic.Model.InhomogeneousBDSCOD +-- Copyright: (c) 2021 Alexander E. Zarebski +-- License: MIT +-- +-- Maintainer: Alexander E. Zarebski +-- Stability: unstable +-- Portability: ghc +-- +-- This module defines a birth-death model with continuous time sampling and +-- scheduled sampling and rates that are piece-wise constant in time. +-- +-- __Example:__ we will run a simulation for one unit of time and require that +-- there be at least two sequenced samples. +-- +-- >>> simDuration = TimeDelta 1.0 +-- >>> atLeastTwoSequences = True +-- +-- The rates can change through time so we need to specify the times at which +-- they change. In this example the birth rate starts at 1.0 and then drops down +-- to 0.5. The other rates stay at their initial values. +-- +-- >>> birthRateSpec = [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 0.5, 0.5)] +-- >>> deathRateSpec = [(AbsoluteTime 0.0, 0.2)] +-- >>> sampRateSpec = [(AbsoluteTime 0.0, 0.1)] +-- >>> occRateSpec = [(AbsoluteTime 0.0, 0.1)] +-- +-- There are a couple of scheduled samples with probabilities specified for +-- them, ie there will be a scheduled sample at time 0.9 where each lineage is +-- removed and sequenced individually with probability 0.1 and at times 0.5 and +-- 0.75 there is a scheduled sample where individuals are removed but /not/ +-- sequenced with probabilities 0.4 and 0.5 respectively. +-- +-- >>> seqSched = [(AbsoluteTime 0.9, 0.1)] +-- >>> unseqSched = [(AbsoluteTime 0.5, 0.4), (AbsoluteTime 0.75, 0.5)] +-- +-- This is enough to define a 'SimulationConfiguration'. We will ignore the +-- possibility of using a termination handler for this example. +-- +-- >>> ratesAndProbs = (birthRateSpec,deathRateSpec,sampRateSpec,seqSched,occRateSpec,unseqSched) +-- >>> (Just simConfig) = configuration simDuration atLeastTwoSequences Nothing ratesAndProbs +-- +-- Then we can use this to generated a list of epidemic events in the simulation +-- +-- >>> myEpidemicEvents = simulationWithSystem simConfig (allEvents randomEvent) +-- +-- and from this we can extract the observations +-- +-- >>> myObservedEvents = do +-- >>> simState <- myEpidemicEvents +-- >>> case simState of +-- >>> Right es -> return $ observedEvents es +-- >>> Left _ -> return $ Left "simulation terminated early" +-- + +module Epidemic.Model.InhomogeneousBDSCOD + ( configuration + , randomEvent + , InhomBDSCODRates(..) + , InhomBDSCODPop(..) + , getNumRemovedByDeath + , getNumRemovedBySampling + , getNumRemovedByCatastrophe + , getNumRemovedByOccurrence + , getNumRemovedByDisaster + ) where + +import Data.List as List +import Data.Maybe (fromJust) +import qualified Data.Vector as V +import qualified Data.Vector.Generic as G +import Epidemic +import Epidemic.Types.Events (EpidemicEvent (..)) +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Simulation (SimulationConfiguration (..), + SimulationRandEvent (..), + TerminationHandler (..)) +import Epidemic.Types.Time (AbsoluteTime (..), + TimeDelta (..), Timed (..), + allTimes, asTimed, + cadlagValue, maybeNextTimed) +import Epidemic.Utility +import System.Random.MWC +import System.Random.MWC.Distributions (bernoulli, categorical) + +data InhomBDSCODRates = + InhomBDSCODRates + { irBirthRate :: Timed Rate + , irDeathRate :: Timed Rate + , irSamplingRate :: Timed Rate + , irCatastropheSpec :: Timed Probability + , irOccurrenceRate :: Timed Rate + , irDisasterSpec :: Timed Probability + } + deriving (Show, Eq) + +-- | The population in which the epidemic occurs. This includes information +-- about the number of people that have previously been infected and +-- subsequently removed. +data InhomBDSCODPop = + InhomBDSCODPop + { ipInfectedPeople :: People + , ipNumRemovedByDeath :: Int + , ipNumRemovedBySampling :: Int + , ipNumRemovedByCatastrophe :: Int + , ipNumRemovedByOccurrence :: Int + , ipNumRemovedByDisaster :: Int + } deriving (Show) + +getNumRemovedByDeath :: InhomBDSCODPop -> Int +getNumRemovedByDeath = ipNumRemovedByDeath + +getNumRemovedBySampling :: InhomBDSCODPop -> Int +getNumRemovedBySampling = ipNumRemovedBySampling + +getNumRemovedByCatastrophe :: InhomBDSCODPop -> Int +getNumRemovedByCatastrophe = ipNumRemovedByCatastrophe + +getNumRemovedByOccurrence :: InhomBDSCODPop -> Int +getNumRemovedByOccurrence = ipNumRemovedByOccurrence + +getNumRemovedByDisaster :: InhomBDSCODPop -> Int +getNumRemovedByDisaster = ipNumRemovedByDisaster + +instance ModelParameters InhomBDSCODRates InhomBDSCODPop where + rNaught _ InhomBDSCODRates {..} time = + do + birthRate <- cadlagValue irBirthRate time + deathRate <- cadlagValue irDeathRate time + sampleRate <- cadlagValue irSamplingRate time + occurrenceRate <- cadlagValue irOccurrenceRate time + Just $ birthRate / (deathRate + sampleRate + occurrenceRate) + eventRate _ InhomBDSCODRates {..} time = + do + birthRate <- cadlagValue irBirthRate time + deathRate <- cadlagValue irDeathRate time + sampleRate <- cadlagValue irSamplingRate time + occurrenceRate <- cadlagValue irOccurrenceRate time + Just $ birthRate + deathRate + sampleRate + occurrenceRate + birthProb p inhomRates@InhomBDSCODRates {..} time = + do + birthRate <- cadlagValue irBirthRate time + totalEventRate <- eventRate p inhomRates time + Just $ birthRate / totalEventRate + eventWeights _ InhomBDSCODRates{..} time = + do + birthRate <- cadlagValue irBirthRate time + deathRate <- cadlagValue irDeathRate time + sampleRate <- cadlagValue irSamplingRate time + occurrenceRate <- cadlagValue irOccurrenceRate time + Just $ V.fromList [birthRate, deathRate, sampleRate, occurrenceRate] + +instance Population InhomBDSCODPop where + susceptiblePeople _ = Nothing + infectiousPeople = pure . ipInfectedPeople + removedPeople _ = Nothing + isInfected = not . nullPeople . ipInfectedPeople + +-- | Configuration for the simulation of the inhomogeneous rates BDSCOD process. +configuration :: + TimeDelta -- ^ Duration of the simulation after starting at time 0. + -> Bool -- ^ condition upon at least two sequenced samples. + -> Maybe (InhomBDSCODPop -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling. + -> ( [(AbsoluteTime, Rate)] + , [(AbsoluteTime, Rate)] + , [(AbsoluteTime, Rate)] + , [(AbsoluteTime, Probability)] + , [(AbsoluteTime, Rate)] + , [(AbsoluteTime, Probability)]) + -> Maybe (SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s) +configuration maxTime atLeastCherry maybeTHFuncs (tBirthRate, tDeathRate, tSampleRate, cSpec, tOccurrenceRate, dSpec) = + let (seedPerson, newId) = newPerson initialIdentifier + bdscodPop = InhomBDSCODPop { ipInfectedPeople = asPeople [seedPerson] + , ipNumRemovedByDeath = 0 + , ipNumRemovedBySampling = 0 + , ipNumRemovedByCatastrophe = 0 + , ipNumRemovedByOccurrence = 0 + , ipNumRemovedByDisaster = 0 } + in do timedBirthRate <- asTimed tBirthRate + timedDeathRate <- asTimed tDeathRate + timedSamplingRate <- asTimed tSampleRate + catastropheSpec <- asTimed cSpec + timedOccurrenceRate <- asTimed tOccurrenceRate + disasterSpec <- asTimed dSpec + let irVal = + InhomBDSCODRates + timedBirthRate + timedDeathRate + timedSamplingRate + catastropheSpec + timedOccurrenceRate + disasterSpec + termHandler = do (f1, f2) <- maybeTHFuncs + return $ TerminationHandler f1 f2 + if maxTime > TimeDelta 0 + then Just + (SimulationConfiguration + irVal + bdscodPop + newId + (AbsoluteTime 0) + maxTime + termHandler + atLeastCherry) + else Nothing + +-- | A random event and the state afterwards +randomEvent :: SimulationRandEvent InhomBDSCODRates InhomBDSCODPop +randomEvent = SimulationRandEvent randomEvent' + +randomEvent' :: + InhomBDSCODRates + -> AbsoluteTime -- ^ the current time + -> InhomBDSCODPop -- ^ the population + -> Identifier -- ^ current identifier + -> GenIO -- ^ PRNG + -> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier) +randomEvent' inhomRates@InhomBDSCODRates {..} currTime currPop currId gen = + let (Just people) = infectiousPeople currPop + popSize = fromIntegral $ numPeople people :: Double + weightVecFunc = eventWeights currPop inhomRates + -- we need a new step function to account for the population size. + (Just stepFunction) = + asTimed + [ (t, popSize * fromJust (eventRate currPop inhomRates t)) + | t <- List.sort $ concatMap allTimes [irBirthRate, irDeathRate, irSamplingRate, irOccurrenceRate] + ] + in do (Just newEventTime) <- inhomExponential stepFunction currTime gen + if noScheduledEvent currTime newEventTime (irCatastropheSpec <> irDisasterSpec) + then do + eventIx <- categorical (fromJust $ weightVecFunc newEventTime) gen + (selectedPerson, unselectedPeople) <- randomPerson people gen + return $ + case eventIx of + 0 -> + ( newEventTime + , Infection newEventTime selectedPerson birthedPerson + , currPop { ipInfectedPeople = addPerson birthedPerson people} + , newId) + where (birthedPerson, newId) = newPerson currId + 1 -> let currNumDeaths = ipNumRemovedByDeath currPop + in ( newEventTime + , Removal newEventTime selectedPerson + , currPop { ipInfectedPeople = unselectedPeople + , ipNumRemovedByDeath = currNumDeaths + 1 } + , currId ) + 2 -> let currNumSampled = ipNumRemovedBySampling currPop + in ( newEventTime + , IndividualSample newEventTime selectedPerson True + , currPop { ipInfectedPeople = unselectedPeople + , ipNumRemovedBySampling = currNumSampled + 1 } + , currId) + 3 -> let currNumOccurrence = ipNumRemovedByOccurrence currPop + in ( newEventTime + , IndividualSample newEventTime selectedPerson False + , currPop { ipInfectedPeople = unselectedPeople + , ipNumRemovedByOccurrence = currNumOccurrence + 1} + , currId) + _ -> error "no birth, death, sampling, or occurrence event selected." + else case maybeNextTimed irCatastropheSpec irDisasterSpec currTime of + Just (disastTime, Right disastProb) -> + do (disastEvent, postDisastPop) <- + randomDisasterEvent + (disastTime, disastProb) + currPop + gen + return (disastTime, disastEvent, postDisastPop, currId) + Just (catastTime, Left catastProb) -> + do (catastEvent, postCatastPop) <- + randomCatastropheEvent + (catastTime, catastProb) + currPop + gen + return (catastTime, catastEvent, postCatastPop, currId) + Nothing -> error "Missing a next scheduled event when there should be one." + +-- | Return a randomly sampled Catastrophe event and the population after that +-- event has occurred. +randomCatastropheEvent :: + (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the catastrophe + -> InhomBDSCODPop -- ^ The state of the population prior to the catastrophe + -> GenIO + -> IO (EpidemicEvent, InhomBDSCODPop) +randomCatastropheEvent (catastTime, rhoProb) currPop gen = + let (Just (People currPeople)) = infectiousPeople currPop + in do rhoBernoullis <- G.replicateM (V.length currPeople) (bernoulli rhoProb gen) + let filterZip predicate a b = fst . V.unzip . V.filter predicate $ V.zip a b + sampledPeople = People $ filterZip snd currPeople rhoBernoullis + unsampledPeople = People $ filterZip (not . snd) currPeople rhoBernoullis + currNumCatastrophe = ipNumRemovedByCatastrophe currPop + in return ( PopulationSample catastTime sampledPeople True + , currPop { ipInfectedPeople = unsampledPeople + , ipNumRemovedByCatastrophe = currNumCatastrophe + numPeople sampledPeople }) + +-- | Return a randomly sampled Disaster event and the population after that +-- event has occurred. +randomDisasterEvent :: + (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the disaster + -> InhomBDSCODPop -- ^ The state of the population prior to the disaster + -> GenIO + -> IO (EpidemicEvent, InhomBDSCODPop) +randomDisasterEvent (disastTime, nuProb) currPop gen = do + let (Just (People currPeople)) = infectiousPeople currPop + nuBernoullis <- G.replicateM (V.length currPeople) (bernoulli nuProb gen) + let filterZip predicate a b = fst . V.unzip . V.filter predicate $ V.zip a b + sampledPeople = People $ filterZip snd currPeople nuBernoullis + unsampledPeople = People $ filterZip (not . snd) currPeople nuBernoullis + currNumDisaster = ipNumRemovedByDisaster currPop + in return ( PopulationSample disastTime sampledPeople False + , currPop { ipInfectedPeople = unsampledPeople + , ipNumRemovedByDisaster = currNumDisaster + numPeople sampledPeople }) diff --git a/src/Epidemic/Model/LogisticBDSD.hs b/src/Epidemic/Model/LogisticBDSD.hs --- a/src/Epidemic/Model/LogisticBDSD.hs +++ b/src/Epidemic/Model/LogisticBDSD.hs @@ -18,10 +18,6 @@ , Timed(..) , TimeDelta(..) , asTimed - , allTimes - , diracDeltaValue - , nextTime - , cadlagValue , timeAfterDelta ) import Epidemic.Types.Parameter @@ -39,8 +35,7 @@ ) import Epidemic.Types.Simulation ( SimulationConfiguration(..) - , SimulationRandEvent(..) - , SimulationState(..) + , SimulationRandEvent(..), TerminationHandler(..) ) import Epidemic.Utility ( initialIdentifier @@ -79,11 +74,14 @@ let propCapcity = fromIntegral (numPeople pop) / fromIntegral paramsCapacity br = paramsBirthRate * (1.0 - propCapcity) in Just $ br + paramsDeathRate + paramsSamplingRate - birthProb lpop lparam@LogisticBDSDParameters {..} absTime = do + birthProb lpop lparam absTime = do er <- eventRate lpop lparam absTime Just $ br / er where br = logisticBirthRate lparam lpop + eventWeights currPop params@LogisticBDSDParameters {..} _ = + let logisticBR = logisticBirthRate params currPop + in Just $ V.fromList [logisticBR, paramsDeathRate, paramsSamplingRate] instance Population LogisticBDSDPopulation where susceptiblePeople _ = Nothing @@ -96,9 +94,10 @@ configuration :: TimeDelta -> Bool -- ^ condition upon at least two sequenced samples. + -> Maybe (LogisticBDSDPopulation -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling. -> (Rate, Int, Rate, Rate, [(AbsoluteTime, Probability)]) - -> Either String (SimulationConfiguration LogisticBDSDParameters LogisticBDSDPopulation) -configuration simDuration atLeastCherry (birthRate, capacity, deathRate, samplingRate, disasterSpec) + -> Either String (SimulationConfiguration LogisticBDSDParameters LogisticBDSDPopulation s) +configuration simDuration atLeastCherry maybeTHFuncs (birthRate, capacity, deathRate, samplingRate, disasterSpec) | minimum [birthRate, deathRate, samplingRate] < 0 = Left "negative rate provided" | capacity < 1 = Left "insufficient population capacity" @@ -116,6 +115,8 @@ disasterTP (seedPerson, newId) = newPerson initialIdentifier logBDSDPop = LogisticBDSDPopulation (People $ V.singleton seedPerson) + termHandler = do (f1, f2) <- maybeTHFuncs + return $ TerminationHandler f1 f2 in return $ SimulationConfiguration logBDSDParams @@ -123,7 +124,7 @@ newId (AbsoluteTime 0) simDuration - Nothing + termHandler atLeastCherry -- | Defines how a single random event is simulated in this model. @@ -140,14 +141,12 @@ randEvent' params@LogisticBDSDParameters {..} currTime currPop@(LogisticBDSDPopulation currPpl) currId gen = let netEventRate = (fromJust $ eventRate currPop params currTime) popSizeDouble = fromIntegral $ numPeople currPpl - logisticBR = logisticBirthRate params currPop - eventWeights = - V.fromList [logisticBR, paramsDeathRate, paramsSamplingRate] + (Just weightsVec) = eventWeights currPop params currTime in do delay <- exponential (netEventRate * popSizeDouble) gen let newEventTime = timeAfterDelta currTime (TimeDelta delay) if noScheduledEvent currTime newEventTime paramsDisasters then do - eventIx <- categorical eventWeights gen + eventIx <- categorical weightsVec gen (randPerson, otherPeople) <- randomPerson currPpl gen return $ case eventIx of diff --git a/src/Epidemic/Types/Events.hs b/src/Epidemic/Types/Events.hs --- a/src/Epidemic/Types/Events.hs +++ b/src/Epidemic/Types/Events.hs @@ -1,6 +1,6 @@ -{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE RecordWildCards #-} module Epidemic.Types.Events ( EpidemicEvent(Infection, Removal, IndividualSample, @@ -14,71 +14,68 @@ , EpidemicTree(Branch, Leaf, Shoot) , maybeEpidemicTree , isExtinctionOrStopping - , eventTime + , isIndividualSample , derivedFrom ) where -import qualified Data.Aeson as Json -import qualified Data.ByteString as B -import qualified Data.ByteString.Builder as BBuilder -import qualified Data.List as List -import qualified Data.Vector as V -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Types.Time (AbsoluteTime(..), TimeDelta(..), timeDelta) -import GHC.Generics +import qualified Data.Aeson as Json +import Epidemic.Types.Population +import Epidemic.Types.Time (AbsoluteTime (..), TimeStamp (..)) +import GHC.Generics -- | Events that can occur in an epidemic with their absolute time. data EpidemicEvent = Infection AbsoluteTime Person Person -- ^ absolute time; infector; infectee | Removal AbsoluteTime Person | IndividualSample - { indSampTime :: AbsoluteTime + { indSampTime :: AbsoluteTime , indSampPerson :: Person - , indSampSeq :: Bool + , indSampSeq :: Bool } | PopulationSample - { popSampTime :: AbsoluteTime + { popSampTime :: AbsoluteTime , popSampPeople :: People - , popSampSeq :: Bool + , popSampSeq :: Bool } - | Extinction -- ^ epidemic went extinct time time can be recovered from the preceeding removal - | StoppingTime -- ^ the simulation reached the stopping time + | Extinction AbsoluteTime -- ^ epidemic went extinct + | StoppingTime AbsoluteTime -- ^ the simulation reached the stopping time deriving (Show, Generic, Eq) instance Json.FromJSON EpidemicEvent instance Json.ToJSON EpidemicEvent +instance TimeStamp EpidemicEvent where + absTime ee = + case ee of + Infection absT _ _ -> absT + Removal absT _ -> absT + IndividualSample {..} -> indSampTime + PopulationSample {..} -> popSampTime + StoppingTime absT -> absT + Extinction absT -> absT + +-- | Predicate for the event being an individual sample event. +isIndividualSample :: EpidemicEvent -> Bool +isIndividualSample ee = + case ee of + IndividualSample {} -> True + _ -> False + -- | Predicate for whether an @EpidemicEvent@ is one of the terminal events of -- extinction or the stopping time having been reached. isExtinctionOrStopping :: EpidemicEvent -> Bool isExtinctionOrStopping e = case e of - Extinction -> True - StoppingTime -> True - _ -> False + Extinction {} -> True + StoppingTime {} -> True + _ -> False -- | Epidemic Events are ordered based on which occurred first. Since -- 'Extinction' and 'StoppingTime' events are there as placeholders they are -- placed as the end of the order. instance Ord EpidemicEvent where - Extinction <= Extinction = True - Extinction <= StoppingTime = True - Extinction <= _ = False - StoppingTime <= Extinction = False - StoppingTime <= StoppingTime = True - StoppingTime <= _ = False - e1 <= e2 = eventTime e1 <= eventTime e2 - --- | The absolute time an event occurred. -eventTime :: EpidemicEvent -> AbsoluteTime -eventTime e = - case e of - Infection time _ _ -> time - Removal time _ -> time - IndividualSample {..} -> indSampTime - PopulationSample {..} -> popSampTime + e1 <= e2 = absTime e1 <= absTime e2 -- | The events that occurred as a result of the existance of the given person. derivedFrom :: @@ -115,8 +112,8 @@ in if haveCommonPeople people popSampPeople then e : derivedEvents else derivedEvents - Extinction -> derivedFromPeople people es - StoppingTime -> derivedFromPeople people es + Extinction {} -> derivedFromPeople people es + StoppingTime {} -> derivedFromPeople people es -- | The whole transmission tree including the unobserved leaves. Lineages that -- are still extant are modelled as /shoots/ and contain a 'Person' as their @@ -143,8 +140,10 @@ if nullPeople popSampPeople then Left "The last event is a PopulationSample with no people sampled" else Right (Leaf e) - Extinction -> Left "Extinction event encountered. It should have been removed" - StoppingTime -> Left "Stopping time encountered. It should have been removed" + Extinction {} -> + Left "Extinction event encountered. It should have been removed" + StoppingTime {} -> + Left "Stopping time encountered. It should have been removed" maybeEpidemicTree (e:es) = case e of Infection _ p1 p2 -> @@ -165,5 +164,7 @@ if nullPeople popSampPeople then maybeEpidemicTree es else Right (Leaf e) - Extinction -> Left "Extinction event encountered. It should have been removed" - StoppingTime -> Left "Stopping time encountered. It should have been removed" + Extinction {} -> + Left "Extinction event encountered. It should have been removed" + StoppingTime {} -> + Left "Stopping time encountered. It should have been removed" diff --git a/src/Epidemic/Types/Newick.hs b/src/Epidemic/Types/Newick.hs --- a/src/Epidemic/Types/Newick.hs +++ b/src/Epidemic/Types/Newick.hs @@ -2,17 +2,13 @@ module Epidemic.Types.Newick where -import qualified Data.Aeson as Json -import qualified Data.ByteString as B import qualified Data.ByteString.Builder as BBuilder import qualified Data.List as List import qualified Data.Vector as V -import Epidemic.Types.Parameter import Epidemic.Types.Observations import Epidemic.Types.Events import Epidemic.Types.Population import Epidemic.Types.Time -import GHC.Generics -- | Class of types that can be expressed in Newick format. class Newick t @@ -51,7 +47,7 @@ IndividualSample {..} -> if indSampSeq then Just - ( (personByteString indSampPerson) <> + ( personByteString indSampPerson <> colonBuilder <> branchLength t indSampTime , [e]) else Nothing diff --git a/src/Epidemic/Types/Observations.hs b/src/Epidemic/Types/Observations.hs --- a/src/Epidemic/Types/Observations.hs +++ b/src/Epidemic/Types/Observations.hs @@ -1,5 +1,5 @@ +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE DeriveGeneric #-} module Epidemic.Types.Observations ( Observation(..) @@ -9,21 +9,18 @@ , pointProcessEvents , reconstructedTreeEvents , observedEvents + , aggregated ) where -import Control.Monad (liftM) -import qualified Data.Aeson as Json -import qualified Data.ByteString.Builder as BBuilder -import qualified Data.List as List -import qualified Data.Vector as V -import Epidemic.Types.Events - ( EpidemicEvent(..) - , EpidemicTree(..) - , maybeEpidemicTree - ) -import Epidemic.Types.Time (TimeDelta(..), timeDelta) -import Epidemic.Types.Population (People(..), personByteString) -import GHC.Generics +import qualified Data.Aeson as Json +import qualified Data.List as List +import Epidemic.Types.Events (EpidemicEvent (..), + EpidemicTree (..), + maybeEpidemicTree) +import Epidemic.Types.Population (asPeople) +import Epidemic.Types.Time (TimeInterval (..), TimeStamp (..), + inInterval) +import GHC.Generics -- | A wrapper for an 'EpidemicEvent' to indicate that this is an even that was -- observed rather than just an event of the epidemic process. @@ -35,6 +32,9 @@ instance Json.ToJSON Observation +instance TimeStamp Observation where + absTime (Observation ee) = absTime ee + -- | A representation of the events that can be observed in an epidemic but -- which are not included in the reconstructed tree, ie the unsequenced -- observations. @@ -47,8 +47,10 @@ pointProcessEvents Shoot {} = PointProcessEvents [] pointProcessEvents (Leaf e) = case e of - IndividualSample {..} -> PointProcessEvents $ if not indSampSeq then [Observation e] else [] - PopulationSample {..} -> PointProcessEvents $ if not popSampSeq then [Observation e] else [] + IndividualSample {..} -> + PointProcessEvents [Observation e | not indSampSeq] + PopulationSample {..} -> + PointProcessEvents [Observation e | not popSampSeq] _ -> PointProcessEvents [] pointProcessEvents (Branch _ lt rt) = let (PointProcessEvents lEs) = pointProcessEvents lt @@ -71,12 +73,14 @@ maybeReconstructedTree Shoot {} = Left "EpidemicTree is only a Shoot" maybeReconstructedTree (Leaf e) = case e of - IndividualSample {..} -> if indSampSeq - then Right $ RLeaf (Observation e) - else Left "Leaf with non-sequenced event individual sample" - PopulationSample {..} -> if popSampSeq - then Right $ RLeaf (Observation e) - else Left "Leaf with non-sequenced event population sample" + IndividualSample {..} -> + if indSampSeq + then Right $ RLeaf (Observation e) + else Left "Leaf with non-sequenced event individual sample" + PopulationSample {..} -> + if popSampSeq + then Right $ RLeaf (Observation e) + else Left "Leaf with non-sequenced event population sample" _ -> Left "Bad leaf in the EpidemicTree" maybeReconstructedTree (Branch e@Infection {} lt rt) | hasSequencedLeaf lt && hasSequencedLeaf rt = do @@ -96,7 +100,7 @@ case e of IndividualSample {..} -> indSampSeq PopulationSample {..} -> popSampSeq - _ -> False + _ -> False hasSequencedLeaf (Branch _ lt rt) = hasSequencedLeaf lt || hasSequencedLeaf rt -- | The events that were observed during the epidemic, ie those in the @@ -108,7 +112,7 @@ let (PointProcessEvents unseqObss) = pointProcessEvents epiTree reconTreeEvents <- if hasSequencedLeaf epiTree - then (liftM reconstructedTreeEvents) $ maybeReconstructedTree epiTree + then reconstructedTreeEvents <$> maybeReconstructedTree epiTree else Right [] return $ List.sort . List.nub $ unseqObss ++ reconTreeEvents @@ -120,3 +124,40 @@ List.sort $ obs : (reconstructedTreeEvents rtl ++ reconstructedTreeEvents rtr) RLeaf obs -> [obs] + +-- | Aggregate the sequenced and unsequenced individual level samples +aggregated :: [TimeInterval] -> [TimeInterval] -> [Observation] -> [Observation] +aggregated seqAggInts unseqAggInts = List.sort . aggUnsequenced . aggSequenced + where + aggUnsequenced = _aggregate unseqAggInts False + aggSequenced = _aggregate seqAggInts True + +-- | Aggregate observations in each of the intervals given the correct +-- sequencing status. +_aggregate :: [TimeInterval] -> Bool -> [Observation] -> [Observation] +_aggregate intervals onlySequenced obs = List.foldl' f obs intervals + where + f os i = _aggregateInInterval i onlySequenced os + +-- | Aggregate all the observations that fall in the interval and have the +-- correct sequencing status. +_aggregateInInterval :: TimeInterval -> Bool -> [Observation] -> [Observation] +_aggregateInInterval interval@TimeInterval {..} onlySequenced obs = + let asPopulationSample os absT = + Observation $ + PopulationSample + absT + (asPeople [indSampPerson ee | Observation ee <- os]) + onlySequenced + (_, aggTime) = timeIntEndPoints + toBeAggregated o = + case o of + Observation (IndividualSample {..}) -> + inInterval interval o && + (if onlySequenced + then indSampSeq + else not indSampSeq) + _ -> False + (obs2Agg, otherObs) = List.partition toBeAggregated obs + newPopSample = asPopulationSample obs2Agg aggTime + in newPopSample : otherObs diff --git a/src/Epidemic/Types/Parameter.hs b/src/Epidemic/Types/Parameter.hs --- a/src/Epidemic/Types/Parameter.hs +++ b/src/Epidemic/Types/Parameter.hs @@ -1,18 +1,45 @@ -{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module: Epidemic.Types.Parameter +-- Copyright: (c) 2021 Alexander E. Zarebski +-- License: MIT +-- +-- Maintainer: Alexander E. Zarebski +-- Stability: unstable +-- Portability: ghc +-- +-- This module defines some types and functions for working with parameters of models. +-- + module Epidemic.Types.Parameter where +import Data.Vector (Vector) import Epidemic.Types.Population (Population(..)) import Epidemic.Types.Time (AbsoluteTime(..)) -- | Class of types that can be considered parameterisations of a epidemic -- model. class (Population p) => ModelParameters a p where + + -- | The basic reproduction number. + -- + -- __NOTE__ that this is not the /effective/ reproduction number the + -- population is included in case there is structure other than immunity that + -- needs to be accounted for. rNaught :: p -> a -> AbsoluteTime -> Maybe Double + + -- | The total event rate at a particular point in time. eventRate :: p -> a -> AbsoluteTime -> Maybe Rate + + -- | The probability that an event will result in an infection. birthProb :: p -> a -> AbsoluteTime -> Maybe Probability + -- | The __unnormalised__ distribution across the possible events. + eventWeights :: p -> a -> AbsoluteTime -> Maybe (Vector Double) + +-- | The rate at which an event occurs type Rate = Double +-- | A probability type Probability = Double diff --git a/src/Epidemic/Types/Population.hs b/src/Epidemic/Types/Population.hs --- a/src/Epidemic/Types/Population.hs +++ b/src/Epidemic/Types/Population.hs @@ -1,6 +1,23 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE DeriveGeneric #-} +-- | +-- Module: Epidemic.Types.Population +-- Copyright: (c) 2021 Alexander E. Zarebski +-- License: MIT +-- +-- Maintainer: Alexander E. Zarebski +-- Stability: unstable +-- Portability: ghc +-- +-- This module defines some types and functions for working with identifiers, +-- persons, people, and populations. +-- +-- * An 'Identifier' is used as a unique label for a 'Person', +-- * a 'Person' is a single individual, +-- * a group of 'People' is a collection of persons, +-- * and 'Population' is a typeclass for working with people that have some structure. +-- + module Epidemic.Types.Population ( Person(Person) , People(People) @@ -16,12 +33,10 @@ , personByteString ) where -import qualified Data.Aeson as Json -import qualified Data.ByteString as B +import qualified Data.Aeson as Json import qualified Data.ByteString.Builder as BBuilder -import Data.ByteString.Internal (c2w) -import qualified Data.Vector as V -import GHC.Generics +import qualified Data.Vector as V +import GHC.Generics -- | Class of types that can represent populations in an epidemic simulation. class Population a where @@ -67,7 +82,7 @@ -- | Predicate for whether two sets of people have any members in common. haveCommonPeople :: People -> People -> Bool -haveCommonPeople (People ps1) (People ps2) = V.any (\p -> V.elem p ps2) ps1 +haveCommonPeople (People ps1) (People ps2) = V.any (`V.elem` ps2) ps1 -- | Predicate for whether there are any people nullPeople :: People -> Bool diff --git a/src/Epidemic/Types/Simulation.hs b/src/Epidemic/Types/Simulation.hs --- a/src/Epidemic/Types/Simulation.hs +++ b/src/Epidemic/Types/Simulation.hs @@ -4,38 +4,46 @@ ( SimulationConfiguration(..) , SimulationState(..) , SimulationRandEvent(..) + , TerminationHandler(..) + , genIOFromFixed + , genIOFromWord32 + , genIOFromSystem ) where -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Types.Time (AbsoluteTime(..), TimeDelta(..), timeDelta) -import System.Random.MWC +import qualified Data.Vector.Unboxed as Unboxed +import Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Time (AbsoluteTime (..), TimeDelta (..)) +import GHC.Word (Word32) +import System.Random.MWC (GenIO, create, createSystemRandom, + initialize) -data SimulationConfiguration r p = +data SimulationConfiguration r p s = SimulationConfiguration { -- | The event rates - scRates :: r + scRates :: r -- | The population - , scPopulation :: p + , scPopulation :: p -- | A new identifier - , scNewIdentifier :: Identifier + , scNewIdentifier :: Identifier -- | The absolute time at which the simulation starts - , scStartTime :: AbsoluteTime + , scStartTime :: AbsoluteTime -- | The duration of the simulation until it stops - , scSimDuration :: TimeDelta + , scSimDuration :: TimeDelta -- | The simulation terminates if this predicate is not satisfied - , scValidPopulation :: Maybe (p -> Bool) + , scTerminationHandler :: Maybe (TerminationHandler p s) -- | The simulation requires at least two sequenced samples - , scRequireCherry :: Bool + , scRequireCherry :: Bool } -- | Either there is a valid simulation state which contains a sequence of --- epidemic events of there is a terminated simulation which indicates that --- the simulation has been rejected. -data SimulationState b +-- epidemic events along with the time and population or, if the simulation has +-- terminated early there is another value to indicate that along with a value +-- which can be used to indicate why the simulation was terminated early. +data SimulationState b c = SimulationState (AbsoluteTime, [EpidemicEvent], b, Identifier) - | TerminatedSimulation + | TerminatedSimulation (Maybe c) deriving (Eq, Show) data SimulationRandEvent a b where @@ -48,3 +56,28 @@ -> GenIO -> IO (AbsoluteTime, EpidemicEvent, b, Identifier)) -> SimulationRandEvent a b + +-- | Check if a simulation should be terminated and if it should be terminated, +-- then compute a summary explaining why. The first function is used to +-- determine whether the population has entered a state which requires the +-- simulation to terminate early and the second can be use to write a summary of +-- the events that led to the termination. +data TerminationHandler b c where + TerminationHandler + :: Population b + => (b -> Bool) + -> ([EpidemicEvent] -> c) + -> TerminationHandler b c + +-- | A PRNG seed based on the given number. This is the best choice for +-- reproducible simulations. +genIOFromWord32 :: Word32 -> IO GenIO +genIOFromWord32 seed = initialize (Unboxed.fromList [seed]) + +-- | A PRNG seed generated by the system's random number generator. +genIOFromSystem :: IO GenIO +genIOFromSystem = createSystemRandom + +-- | A PRNG seed which is hard coded into @mwc-random@. +genIOFromFixed :: IO GenIO +genIOFromFixed = create diff --git a/src/Epidemic/Types/Time.hs b/src/Epidemic/Types/Time.hs --- a/src/Epidemic/Types/Time.hs +++ b/src/Epidemic/Types/Time.hs @@ -1,25 +1,33 @@ -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE RecordWildCards #-} module Epidemic.Types.Time ( AbsoluteTime(..) , TimeDelta(..) + , TimeInterval(..) , Timed(..) - , timeDelta - , diracDeltaValue - , timeAfterDelta - , nextTime - , cadlagValue - , isAscending - , hasTime + , TimeStamp(..) , allTimes + , allValues + , asConsecutiveIntervals1 , asTimed + , cadlagValue + , diracDeltaValue + , hasTime + , inInterval + , isAscending + , maybeNextTimed + , nextTime + , timeAfterDelta + , timeDelta + , timeInterval1 + , timeInterval2 ) where -import qualified Data.Aeson as Json -import qualified Data.List as List -import qualified Data.Maybe as Maybe -import GHC.Generics +import qualified Data.Aeson as Json +import qualified Data.List as List +import qualified Data.Maybe as Maybe +import GHC.Generics -- | Absolute time. newtype AbsoluteTime = @@ -30,10 +38,25 @@ instance Json.ToJSON AbsoluteTime --- | Predicate for an infinite absolute time -isInfiniteAbsoluteTime :: AbsoluteTime -> Bool -isInfiniteAbsoluteTime (AbsoluteTime t) = isInfinite t +-- | A type that has an absolute time associated with it and can be treated as +-- having a temporal ordering. +-- +-- > a = AbsoluteTime 1 +-- > b = AbsoluteTime 2 +-- > a `isBefore` b +-- +class TimeStamp a where + absTime :: a -> AbsoluteTime + isAfter :: a -> a -> Bool + isAfter x y = absTime x > absTime y + + isBefore :: a -> a -> Bool + isBefore x y = absTime x < absTime y + +instance TimeStamp AbsoluteTime where + absTime = id + -- | Duration of time between two absolute times. newtype TimeDelta = TimeDelta Double @@ -43,6 +66,18 @@ instance Json.ToJSON TimeDelta +-- | An interval of time +data TimeInterval = + TimeInterval + { timeIntEndPoints :: (AbsoluteTime, AbsoluteTime) + , timeIntDuration :: TimeDelta + } + deriving (Generic, Eq, Show) + +instance Json.FromJSON TimeInterval + +instance Json.ToJSON TimeInterval + -- | The duration of time between two absolute times -- -- >>> timeDelta (AbsoluteTime 1) (AbsoluteTime 2.5) @@ -62,6 +97,28 @@ timeAfterDelta :: AbsoluteTime -> TimeDelta -> AbsoluteTime timeAfterDelta (AbsoluteTime t0) (TimeDelta d) = AbsoluteTime (t0 + d) +-- | Construct a 'TimeInterval' from the end points. +timeInterval1 :: AbsoluteTime -> AbsoluteTime -> TimeInterval +timeInterval1 start end = TimeInterval (start, end) (timeDelta start end) + +-- | Construct a 'TimeInterval' from the start time and the duration. +timeInterval2 :: AbsoluteTime -> TimeDelta -> TimeInterval +timeInterval2 start duration = + TimeInterval (start, timeAfterDelta start duration) duration + +-- | Check if an 'AbsoluteTime' sits within a 'TimeInterval'. +inInterval :: TimeStamp a => TimeInterval -> a -> Bool +inInterval TimeInterval {..} x = + let (start, end) = timeIntEndPoints + absT = absTime x + in start <= absT && absT <= end + +-- | Construct a list of consecutive intervals divided by the given absolute +-- times. +asConsecutiveIntervals1 :: [AbsoluteTime] -> [TimeInterval] +asConsecutiveIntervals1 absTimes = + zipWith timeInterval1 (init absTimes) (tail absTimes) + -- | Type containing values at times. The times are increasing as required by -- @asTimed@. newtype Timed a = @@ -89,8 +146,8 @@ isAscending :: Ord a => [a] -> Bool isAscending xs = case xs of - [] -> True - [_] -> True + [] -> True + [_] -> True (x:y:xs') -> x <= y && isAscending (y : xs') -- | Evaluate the timed object treating it as a cadlag function @@ -122,13 +179,7 @@ -- | Check if there exists a pair with a particular time index. hasTime :: Timed a -> AbsoluteTime -> Bool -hasTime (Timed txs) = hasTime' txs - -hasTime' :: [(AbsoluteTime, a)] -> AbsoluteTime -> Bool -hasTime' txs q = - case txs of - ((t, _):txs') -> t == q || hasTime' txs' q - [] -> False +hasTime tx absT = elem absT $ allTimes tx -- | Return the value of the next time if possible or an exact match if it -- exists. @@ -153,3 +204,55 @@ -- allTimes :: Timed a -> [AbsoluteTime] allTimes (Timed txs) = [t | (t, _) <- txs, not $ isInfiniteAbsoluteTime t] + +-- | The values that the timed variable takes. NOTE that it is safe to use +-- 'fromJust' here because 'allTimes' only returns times for which there is a +-- cadlag value anyway. +-- +-- >>> (Just tx) = asTimed [(AbsoluteTime 1,2),(AbsoluteTime 1.5,1)] +-- >>> allValues tx +-- [2,1] +-- +allValues :: Timed a -> [a] +allValues timed = Maybe.fromJust . cadlagValue timed <$> allTimes timed + +-- | Predicate for an infinite absolute time +isInfiniteAbsoluteTime :: AbsoluteTime -> Bool +isInfiniteAbsoluteTime (AbsoluteTime t) = isInfinite t + +-- | Look at both of the timed objects and, if possible, return the time that +-- the first one changes along with the value it changes to. +-- +-- >>> (Just tA) = asTimed [(AbsoluteTime 1, (1.1 :: Double)), (AbsoluteTime 3, 2.3)] +-- >>> (Just tB) = asTimed [(AbsoluteTime 2, (1 :: Int))] +-- >>> maybeNextTimed tA tB (AbsoluteTime 0.5) +-- Just (AbsoluteTime 1.0,Left 1.1) +-- >>> maybeNextTimed tA tB (AbsoluteTime 1.5) +-- Just (AbsoluteTime 2.0,Right 1) +-- >>> maybeNextTimed tA tB (AbsoluteTime 3.5) +-- Nothing +-- +maybeNextTimed :: Timed a + -> Timed b + -> AbsoluteTime + -> Maybe (AbsoluteTime, Either a b) +maybeNextTimed timedA timedB absT = + let f = flip nextTime absT + g1 timed at = do -- two functions are needed for the different types. + v <- diracDeltaValue timed at + if isInfiniteAbsoluteTime at + then Nothing + else Just (at, Left v) + g2 timed at = do + v <- diracDeltaValue timed at + if isInfiniteAbsoluteTime at + then Nothing + else Just (at, Right v) + in case (f timedA, f timedB) of + (Just tA, Just tB) -> + if tA < tB + then g1 timedA tA + else g2 timedB tB + (Just tA, Nothing) -> g1 timedA tA + (Nothing, Just tB) -> g2 timedB tB + (Nothing, Nothing) -> Nothing diff --git a/src/Epidemic/Utility.hs b/src/Epidemic/Utility.hs --- a/src/Epidemic/Utility.hs +++ b/src/Epidemic/Utility.hs @@ -1,42 +1,42 @@ -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} -module Epidemic.Utility where +module Epidemic.Utility ( initialIdentifier + , inhomExponential + , randomPerson + , maybeToRight + , newPerson + , isReconTreeLeaf + , simulationWithSystem + , simulationWithFixedSeed + , simulationWithGenIO + ) where -import Control.Applicative -import Control.Monad (liftM) -import Control.Monad.Primitive (PrimMonad, PrimState) -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as Char8 -import qualified Data.List as List -import qualified Data.Maybe as Maybe -import qualified Data.Vector as V -import Epidemic -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Types.Simulation -import Epidemic.Types.Time - ( AbsoluteTime(..) - , Timed(..) - , TimeDelta(..) - , diracDeltaValue - , nextTime - , cadlagValue - , timeAfterDelta - ) -import GHC.Generics (Generic) -import System.Random.MWC -import System.Random.MWC.Distributions (exponential) +import Control.Monad.Primitive (PrimMonad, PrimState) +import qualified Data.List as List +import qualified Data.Maybe as Maybe +import qualified Data.Vector as V +import Epidemic +import Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Simulation +import Epidemic.Types.Time (AbsoluteTime (..), + TimeDelta (..), Timed (..), + cadlagValue, nextTime, + timeAfterDelta) +import System.Random.MWC +import System.Random.MWC.Distributions (exponential) initialIdentifier :: Identifier initialIdentifier = Identifier 1 +-- | A new person constructed from the given identifier and a new identifier. newPerson :: Identifier -> (Person, Identifier) newPerson idntty@(Identifier idInt) = (Person idntty, Identifier (idInt + 1)) +-- | An element of a vector and the vector with that element removed. selectElem :: V.Vector a -> Int -> (a, V.Vector a) selectElem v n | n == 0 = (V.head v, V.tail v) @@ -44,6 +44,8 @@ let (foo, bar) = V.splitAt n v in (V.head bar, foo V.++ (V.tail bar)) +-- | A random person and the remaining group of people after they have been +-- sampled with removal. randomPerson :: People -> GenIO -> IO (Person, People) randomPerson people@(People persons) gen = do u <- uniform gen @@ -61,7 +63,7 @@ instance Show NBranch where show (NBranch st (Just l)) = show st ++ ":" ++ show l - show (NBranch st Nothing) = show st + show (NBranch st Nothing) = show st data NBranchSet = NBranchSet [NBranch] @@ -77,8 +79,8 @@ instance Show NSubtree where show (NLeaf (Just n)) = n - show (NLeaf Nothing) = "" - show (NInternal bs) = show bs + show (NLeaf Nothing) = "" + show (NInternal bs) = show bs data NTree = NTree [NBranch] @@ -87,106 +89,102 @@ instance Show NTree where show (NTree bs) = show (NBranchSet bs) ++ ";" --- | Example run --- > (Success foo) = parseString newickTree mempty "((foo:1.1,bar:1.2):1.3,baz:1.4);" --- > (Success bar) = parseString newickTree mempty $ show foo --- > foo == bar --- True -sort :: Ord a => [a] -> [a] -sort = List.sort - +-- | The number of elements of the list that map to @True@ under the predicate. count' :: (a -> Bool) -> [a] -> Int -count' p = go 0 - where - go n [] = n - go n (x:xs) - | p x = go (n + 1) xs - | otherwise = go n xs +count' p xs = sum [if p x then 1 else 0 | x <- xs] --- | Run a simulation described by a configuration object with the provided --- PRNG. +-- | Run a simulation described by a configuration object and the model's +-- @allEvents@ style function (see the example in +-- "Epidemic.Model.InhomogeneousBDSCOD") using the provided PRNG. simulationWithGenIO :: (ModelParameters a b, Population b) - => SimulationConfiguration a b - -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) + => SimulationConfiguration a b c + -> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c)) -> GenIO - -> IO [EpidemicEvent] + -> IO (Either (Maybe c) [EpidemicEvent]) simulationWithGenIO config@SimulationConfiguration {..} allEventsFunc gen = if scRequireCherry - then do - simulation' config allEventsFunc gen + then + simulationAtLeastCherry config allEventsFunc gen else do - SimulationState (_, events, _, _) <- + simState <- allEventsFunc scRates (timeAfterDelta scStartTime scSimDuration) - scValidPopulation - (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + scTerminationHandler + (SimulationState (scStartTime, [], scPopulation, scNewIdentifier)) gen - return $ sort events + return $ case simState of + SimulationState (_, events, _, _) -> Right $ List.sort events + TerminatedSimulation maybeSummary -> Left maybeSummary --- | Run a simulation described by a configuration object using the fixed PRNG --- that is hardcoded in the @mwc-random@ package. -simulation :: +-- | Run a simulation using a fixed PRNG random seed. +simulationWithFixedSeed :: (ModelParameters a b, Population b) - => SimulationConfiguration a b - -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) - -> IO [EpidemicEvent] -simulation config allEventsFunc = do - gen <- System.Random.MWC.create :: IO GenIO + => SimulationConfiguration a b c + -> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c)) + -> IO (Either (Maybe c) [EpidemicEvent]) +simulationWithFixedSeed config allEventsFunc = do + gen <- genIOFromFixed simulationWithGenIO config allEventsFunc gen --- | Predicate for whether an epidemic event will appear as a leaf in the --- reconstructed tree. -isReconTreeLeaf :: EpidemicEvent -> Bool -isReconTreeLeaf e = - case e of - IndividualSample {..} -> indSampSeq - PopulationSample {..} -> popSampSeq - _ -> False - -- | Simulation conditioned upon there being at least two sequenced samples. --- NOTE This function is deprecated and will be removed in future versions. -simulation' :: +simulationAtLeastCherry :: (ModelParameters a b, Population b) - => SimulationConfiguration a b - -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) + => SimulationConfiguration a b c + -> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c)) -> GenIO - -> IO [EpidemicEvent] -simulation' config@SimulationConfiguration {..} allEventsFunc gen = do - SimulationState (_, events, _, _) <- + -> IO (Either (Maybe c) [EpidemicEvent]) +simulationAtLeastCherry config@SimulationConfiguration {..} allEventsFunc gen = do + simState <- allEventsFunc scRates (timeAfterDelta scStartTime scSimDuration) - scValidPopulation - (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + scTerminationHandler + (SimulationState (scStartTime, [], scPopulation, scNewIdentifier)) gen - if count' isReconTreeLeaf events >= 2 - then return $ sort events - else simulation' config allEventsFunc gen + case simState of + SimulationState (_, events, _, _) -> + if count' isReconTreeLeaf events >= 2 + then return $ Right $ List.sort events + else simulationAtLeastCherry config allEventsFunc gen + TerminatedSimulation maybeSummary -> return $ Left maybeSummary -- | Run a simulation described by a configuration object but using a random -- seed generated by the system rather than a seed -simulationWithSystemRandom :: +simulationWithSystem :: (ModelParameters a b, Population b) - => SimulationConfiguration a b - -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) - -> IO [EpidemicEvent] -simulationWithSystemRandom config@SimulationConfiguration {..} allEventsFunc = do - SimulationState (_, events, _, _) <- + => SimulationConfiguration a b c + -> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c)) + -> IO (Either (Maybe c) [EpidemicEvent]) +simulationWithSystem config@SimulationConfiguration {..} allEventsFunc = do + simState <- withSystemRandom $ \g -> allEventsFunc scRates (timeAfterDelta scStartTime scSimDuration) - scValidPopulation - (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + scTerminationHandler + (SimulationState (scStartTime, [], scPopulation, scNewIdentifier)) g - if scRequireCherry - then (if count' isReconTreeLeaf events >= 2 - then return $ sort events - else simulationWithSystemRandom config allEventsFunc) - else return $ sort events + case simState of + SimulationState (_, events, _, _) -> + if scRequireCherry + then (if count' isReconTreeLeaf events >= 2 + then return $ Right $ List.sort events + else simulationWithSystem config allEventsFunc) + else return $ Right $ List.sort events + TerminatedSimulation maybeSummary -> return $ Left maybeSummary +-- | Predicate for whether an epidemic event will appear as a leaf in the +-- reconstructed tree. For scheduled sequenced samples this will only return +-- true if there was at least one lineage observed. +isReconTreeLeaf :: EpidemicEvent -> Bool +isReconTreeLeaf e = + case e of + IndividualSample {..} -> indSampSeq + PopulationSample {..} -> popSampSeq && not (nullPeople popSampPeople) + _ -> False + -- | The number of lineages at the end of a simulation. finalSize :: [EpidemicEvent] -- ^ The events from the simulation @@ -235,4 +233,4 @@ maybeToRight a maybeB = case maybeB of (Just b) -> Right b - Nothing -> Left a + Nothing -> Left a diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,31 +1,38 @@ {-# LANGUAGE OverloadedStrings #-} -import Control.Exception (evaluate) -import Control.Monad -import qualified Data.Aeson as Json -import qualified Data.ByteString as B -import qualified Data.ByteString.Builder as BBuilder -import Data.Either (isRight) -import Data.Maybe (fromJust, isJust, isNothing) -import qualified Data.Vector as V -import Epidemic -import qualified Epidemic.Model.BDSCOD as BDSCOD -import qualified Epidemic.Model.InhomogeneousBDS as InhomBDS -import Epidemic.Types.Events -import Epidemic.Types.Observations -import Epidemic.Types.Time -import Epidemic.Types.Newick -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Utility -import Statistics.Sample -import qualified System.Random.MWC as MWC -import Test.Hspec +import Control.Exception (evaluate) +import Control.Monad +import qualified Data.Aeson as Json +import qualified Data.ByteString as B +import qualified Data.ByteString.Builder as BBuilder +import Data.Either (isRight) +import Data.Maybe (fromJust, isJust, + isNothing) +import qualified Data.Vector as V +import Epidemic +import qualified Epidemic.Model.BDSCOD as BDSCOD +import qualified Epidemic.Model.InhomogeneousBDS as InhomBDS +import qualified Epidemic.Model.InhomogeneousBDSCOD as InhomBDSCOD +import Epidemic.Types.Events +import Epidemic.Types.Newick +import Epidemic.Types.Observations +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Simulation (SimulationState (..), + TerminationHandler (..), + genIOFromFixed, + genIOFromSystem, + genIOFromWord32) +import Epidemic.Types.Time +import Epidemic.Utility +import Statistics.Sample +import qualified System.Random.MWC as MWC +import Test.Hspec -- | Helper function for converting from Either to Maybe monad. either2Maybe x = case x of Right v -> Just v - Left _ -> Nothing + Left _ -> Nothing -- | y is within n% of x from x. withinNPercent n x y = x - d < y && y < x + d @@ -214,12 +221,13 @@ (observedEvents demoFullEvents04)) `shouldBe` True it "Disasters can be simulated" $ do - demoSim <- - simulation + (Right demoSim) <- + simulationWithFixedSeed (fromJust (BDSCOD.configuration (TimeDelta 4) False + Nothing ( 1.3 , 0.1 , 0.1 @@ -239,7 +247,7 @@ , indSampPerson = p1 , indSampSeq = False } - , StoppingTime + , StoppingTime (AbsoluteTime 6.0) ] expectedObs = [ Observation @@ -317,6 +325,14 @@ isJust (nextTime demoTimed (AbsoluteTime 2.0)) `shouldBe` True isJust (nextTime demoTimed (AbsoluteTime 2.1)) `shouldBe` True isJust (nextTime demoTimed (AbsoluteTime 10.0)) `shouldBe` True + it "the maybeNextTimed function works as expected" $ do + let (Just tA) = asTimed [(AbsoluteTime 1, (1.1 :: Double)), (AbsoluteTime 3, 2.3)] + (Just tB) = asTimed [(AbsoluteTime 2, (1 :: Int))] + maybeNextTimed tA tB (AbsoluteTime 0.5) == Just (AbsoluteTime 1.0,Left 1.1) `shouldBe` True + maybeNextTimed tA tB (AbsoluteTime 1.5) == Just (AbsoluteTime 2.0,Right 1) `shouldBe` True + maybeNextTimed tA tB (AbsoluteTime 2.5) == Just (AbsoluteTime 3.0,Left 2.3) `shouldBe` True + isNothing (maybeNextTimed tA tB (AbsoluteTime 3.5)) `shouldBe` True + it "shifted times work" $ let sf = fromJust $ @@ -337,6 +353,61 @@ asTimed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, -1.0)] (isJust $ InhomBDS.inhomBDSRates timedBirthRate 0.5 0.5) `shouldBe` False +simTypeTests = + describe "Test Types.Simulation PRNG helpers" $ + do it "check genIOFromFixed always gives same result" $ + do g1 <- genIOFromFixed + u11 <- MWC.uniform g1 :: IO Double + u12 <- MWC.uniform g1 :: IO Double + + g2 <- genIOFromFixed + u21 <- MWC.uniform g2 :: IO Double + u22 <- MWC.uniform g2 :: IO Double + + u11 == u21 `shouldBe` True + u11 /= u22 `shouldBe` True + u12 == u22 `shouldBe` True + + it "check genIOFromSystem always gives different results" $ + do g1 <- genIOFromSystem + u11 <- MWC.uniform g1 :: IO Double + u12 <- MWC.uniform g1 :: IO Double + + g2 <- genIOFromSystem + u21 <- MWC.uniform g2 :: IO Double + u22 <- MWC.uniform g2 :: IO Double + + u11 /= u12 `shouldBe` True + u11 /= u21 `shouldBe` True + u11 /= u22 `shouldBe` True + u12 /= u21 `shouldBe` True + u12 /= u22 `shouldBe` True + u21 /= u22 `shouldBe` True + + it "check genIOFromWord32 works as expected" $ + do g1 <- genIOFromWord32 1 + u11 <- MWC.uniform g1 :: IO Double + u12 <- MWC.uniform g1 :: IO Double + + g2 <- genIOFromWord32 1 + u21 <- MWC.uniform g2 :: IO Double + u22 <- MWC.uniform g2 :: IO Double + + g3 <- genIOFromWord32 2 + u31 <- MWC.uniform g3 :: IO Double + u32 <- MWC.uniform g3 :: IO Double + + u11 == u21 `shouldBe` True + u11 /= u22 `shouldBe` True + u12 == u22 `shouldBe` True + + u11 /= u12 `shouldBe` True + u11 /= u31 `shouldBe` True + u11 /= u32 `shouldBe` True + u12 /= u31 `shouldBe` True + u12 /= u32 `shouldBe` True + u31 /= u32 `shouldBe` True + inhomExpTests = describe "Test the inhomogeneous exponential variate generator" $ let rate1 = 2.0 @@ -347,9 +418,8 @@ fromJust $ asTimed [(AbsoluteTime 0, 1e-10), (AbsoluteTime 1, rate1)] mean2 = 1 / rate1 + 1 var2 = var1 - genAction = MWC.createSystemRandom in do it "check we can get a positive variate out" $ do - gen <- genAction + gen <- genIOFromSystem u1 <- MWC.uniform gen :: IO Double (u1 > 0) `shouldBe` True (Just x1) <- inhomExponential sF1 (AbsoluteTime 0) gen @@ -357,14 +427,14 @@ (x1 < AbsoluteTime 100) `shouldBe` True True `shouldBe` True it "check the mean and variance look sensible" $ do - gen <- genAction + gen <- genIOFromSystem xBoxed <- V.replicateM 20000 (inhomExponential sF1 (AbsoluteTime 0) gen) let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed withinNPercent 5 (mean x) mean1 `shouldBe` True withinNPercent 5 (variance x) var1 `shouldBe` True it "check the mean and variance look sensible with delay" $ do - gen <- genAction + gen <- genIOFromSystem xBoxed <- V.replicateM 20000 (inhomExponential sF2 (AbsoluteTime 0) gen) let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed @@ -390,15 +460,82 @@ , [(simRhoTime, simRho)] , simOmega , [(simNuTime, simNu)]) - simConfig = BDSCOD.configuration simDuration True simParams - in it "stress testing the observed events function" $ do - null (observedEvents []) `shouldBe` True - simEvents <- - simulation (fromJust simConfig) (allEvents BDSCOD.randomEvent) - any isReconTreeLeaf simEvents `shouldBe` True - let (Right oes) = observedEvents simEvents - (length oes > 1) `shouldBe` True + simConfig = BDSCOD.configuration simDuration True Nothing simParams + in do it "stress testing the observed events function" $ do + null (observedEvents []) `shouldBe` True + (Right simEvents) <- + simulationWithFixedSeed (fromJust simConfig) (allEvents BDSCOD.randomEvent) + any isReconTreeLeaf simEvents `shouldBe` True + let (Right oes) = observedEvents simEvents + (length oes > 1) `shouldBe` True + it "check leaves are recognised correctly" $ do + let (absT, person) = (AbsoluteTime 1.0, Person (Identifier 1)) + infEvent = Infection absT person person + remEvent = Removal absT person + sampIndv = IndividualSample absT person + popSampEmpty = PopulationSample absT $ asPeople [] + popSampPerson = PopulationSample absT $ asPeople [person] + isReconTreeLeaf infEvent `shouldBe` False + isReconTreeLeaf remEvent `shouldBe` False + isReconTreeLeaf (sampIndv True) `shouldBe` True + isReconTreeLeaf (sampIndv False) `shouldBe` False + isReconTreeLeaf (popSampEmpty True) `shouldBe` False + isReconTreeLeaf (popSampEmpty False) `shouldBe` False + isReconTreeLeaf (popSampPerson True) `shouldBe` True + isReconTreeLeaf (popSampPerson False) `shouldBe` False + +resultAA = demoSampleEvents01 + +resultAB = + [ Infection (AbsoluteTime 1) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p4 + , IndividualSample (AbsoluteTime 6) p4 True + , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False + , IndividualSample (AbsoluteTime 12) p5 True + , PopulationSample (AbsoluteTime 17.0) (asPeople []) False + ] + +resultBA = + [ Infection (AbsoluteTime 1) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p4 + , IndividualSample (AbsoluteTime 6) p4 True + , IndividualSample (AbsoluteTime 8) p2 False + , IndividualSample (AbsoluteTime 11) p6 False + , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True + ] + +resultBB = + [ Infection (AbsoluteTime 1.0) p1 p2 + , IndividualSample (AbsoluteTime 3.0) (Person (Identifier 1)) True + , Infection (AbsoluteTime 4.0) p2 p4 + , IndividualSample (AbsoluteTime 6.0) (Person (Identifier 4)) True + , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False + , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True + , PopulationSample (AbsoluteTime 17) (asPeople []) False + ] + +aggregationTests = + describe "Aggregation functionality tests" $ do + it "check it does nothing unless it needs to" $ + let demoObs1 = [Observation ee | ee <- demoFullEvents01] + demoObs2 = [Observation ee | ee <- demoSampleEvents01] + in do + (aggregated [] [] demoObs1) == demoObs1 `shouldBe` True + (aggregated [] [] demoObs2) == demoObs2 `shouldBe` True + it "check relevant intervals are processed correctly" $ + let demoObs = [Observation ee | ee <- demoSampleEvents01] + demoSeqInts = asConsecutiveIntervals1 [AbsoluteTime 10, AbsoluteTime 13] + demoUnseqInts = asConsecutiveIntervals1 [AbsoluteTime 7, AbsoluteTime 12, AbsoluteTime 17] + in do + aggregated [] [] demoObs == (map Observation resultAA) `shouldBe` True + aggregated [] demoUnseqInts demoObs == (map Observation resultAB) `shouldBe` True + aggregated demoSeqInts [] demoObs == (map Observation resultBA) `shouldBe` True + aggregated demoSeqInts demoUnseqInts demoObs == (map Observation resultBB) `shouldBe` True + + inhomogeneousBDSTest = describe "InhomogeneousBDS module tests" $ do it "Check the observedEvents filters out removals" $ @@ -677,19 +814,6 @@ (length demoEvents == 4) `shouldBe` True (maybeEpidemicTree demoEvents == maybeEpidemicTree (tail demoEvents)) `shouldBe` True - -- it "asNewickString works for EpidemicTree" $ do - -- let trickyEvents = [ - -- Infection (AbsoluteTime 0.3) (Person (Identifier 1)) (Person (Identifier 2)), - -- Infection (AbsoluteTime 0.4) (Person (Identifier 2)) (Person (Identifier 3)), - -- IndividualSample (AbsoluteTime 0.6) (Person (Identifier 3)) True, - -- IndividualSample (AbsoluteTime 0.7) (Person (Identifier 1)) True] - -- let maybeNewickPair = asNewickString (AbsoluteTime 0, Person (Identifier 1)) =<< maybeEpidemicTree trickyEvents - -- let newickTarget = BBuilder.stringUtf8 "(1:0.39999999999999997,(2:Infinity,3:0.19999999999999996):0.10000000000000003):0.3" - -- let maybeReconTree = maybeReconstructedTree =<< maybeEpidemicTree trickyEvents - -- isJust maybeNewickPair `shouldBe` True - -- [IndividualSample (AbsoluteTime 0.6) (Person (Identifier 3)) True, IndividualSample (AbsoluteTime 0.7) (Person (Identifier 1)) True] == snd (fromJust maybeNewickPair) `shouldBe` True - -- equalBuilders newickTarget (fst $ fromJust maybeNewickPair) `shouldBe` True - -- isJust maybeReconTree `shouldBe` True it "asNewickString works for ReconstructedTree" $ do isJust (asNewickString @@ -750,3 +874,62 @@ helperTypeTests jsonTests newickTests + aggregationTests + simTypeTests + terminationTests1 + +terminationTests1 = + describe "Termination handling tests: InhomogeneousBDSCOD" $ do + let duration = TimeDelta 2.0 + birthRateSpec = [(AbsoluteTime 0.0, 1.5), (AbsoluteTime 0.5, 0.5)] + deathRateSpec = [(AbsoluteTime 0.0, 0.4)] + sampRateSpec = [(AbsoluteTime 0.0, 0.1)] + occRateSpec = [(AbsoluteTime 0.0, 0.1)] + seqSched = [(AbsoluteTime 0.9, 0.1)] + unseqSched = [(AbsoluteTime 0.5, 0.4), (AbsoluteTime 0.75, 0.5)] + ratesAndProbs = (birthRateSpec,deathRateSpec,sampRateSpec,seqSched,occRateSpec,unseqSched) + conf maybeTH = fromJust $ InhomBDSCOD.configuration duration True maybeTH ratesAndProbs + -- We need one simulation configuration for each of the termination + -- handlers that we want to test. + simConfigNothing = conf Nothing + simConfigNever = conf (Just (const False, const ())) + simConfigAlways = conf (Just (const True, const ())) + numDeadThreshold = 3 + simConfigSometimes = conf (Just ((>numDeadThreshold) . InhomBDSCOD.getNumRemovedByDeath, + \es -> length [() | Removal _ _ <- es])) + allEventsFunc = allEvents InhomBDSCOD.randomEvent + it "test simulation works without hander" $ + do + (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc + -- There should always be at least one event in the simulation. + (length esNothing > 0) `shouldBe` True + it "test simulation works with handler that does not trigger" $ + do + (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc + (Right esNever) <- simulationWithFixedSeed simConfigNever allEventsFunc + -- If the handler never triggers this should look the same as not having + -- the event handler. + (all id $ zipWith (==) esNothing esNever) `shouldBe` True + it "test simulation works with handler that always triggers" $ + do + -- If the handler always triggers this should always return the summary. + replicateM_ 30 + (do esAlways <- simulationWithSystem simConfigAlways allEventsFunc + esAlways == Left (Just ()) `shouldBe` True + ) + it "test simulation works with handler that sometimes triggers" $ + do + -- If the handler only sometimes triggers then we need to test both + -- branches. + replicateM_ 30 + (do esSometimes <- simulationWithSystem simConfigSometimes allEventsFunc + case esSometimes of + -- If the termination handler did not trigger the number of + -- removals should not exceed the threshold allowed by the + -- termination handler + (Right es) -> do length [() | Removal _ _ <- es] <= numDeadThreshold `shouldBe` True + -- If the termination handler did trigger then we should know + -- exactly how many removals there was. + (Left (Just n)) -> do n == numDeadThreshold + 1 `shouldBe` True + (Left Nothing) -> True `shouldBe` False -- this branch should not be reached. + )