epi-sim 0.4.2 → 0.7.0

raw patch · 18 files changed

+1275/−435 lines, 18 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• ARCHITECTURE.md +19/−0
• ChangeLog.md +45/−0
• README.md +13/−5
• epi-sim.cabal +7/−3
• src/Epidemic.hs +89/−51
• src/Epidemic/Model/BDSCOD.hs +88/−70
• src/Epidemic/Model/InhomogeneousBDS.hs +18/−21
• src/Epidemic/Model/InhomogeneousBDSCOD.hs +315/−0
• src/Epidemic/Model/LogisticBDSD.hs +13/−14
• src/Epidemic/Types/Events.hs +44/−43
• src/Epidemic/Types/Newick.hs +1/−5
• src/Epidemic/Types/Observations.hs +65/−24
• src/Epidemic/Types/Parameter.hs +28/−1
• src/Epidemic/Types/Population.hs +22/−7
• src/Epidemic/Types/Simulation.hs +50/−17
• src/Epidemic/Types/Time.hs +128/−25
• src/Epidemic/Utility.hs +97/−99
• test/Spec.hs +233/−50

@@ -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.

@@ -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

@@ -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`)

@@ -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

@@ -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 <aezarebski@gmail.com>+-- 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

@@ -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

@@ -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

@@ -0,0 +1,315 @@+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE RecordWildCards       #-}++-- |+-- Module: Epidemic.Model.InhomogeneousBDSCOD+-- Copyright: (c) 2021 Alexander E. Zarebski+-- License: MIT+--+-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>+-- 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 })

@@ -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

@@ -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"

@@ -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

@@ -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

@@ -1,18 +1,45 @@-{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE MultiParamTypeClasses #-} +-- |+-- Module: Epidemic.Types.Parameter+-- Copyright: (c) 2021 Alexander E. Zarebski+-- License: MIT+--+-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>+-- 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

@@ -1,6 +1,23 @@-{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE DeriveGeneric #-} +-- |+-- Module: Epidemic.Types.Population+-- Copyright: (c) 2021 Alexander E. Zarebski+-- License: MIT+--+-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>+-- 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

@@ -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

@@ -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

@@ -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

@@ -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.+                )