diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,15 @@ +# ARCHITECTURE + +This document describes the high-level structure of the `epi-sim` package to +make it easier to navigate the source code. + +## Modules + +The `Epidemic` and `Epidemic.Utilty` modules provide basic functions for +simulating epidemics and working with the resulting data sets, eg simulate a +birth-death process filter out the observed cases and write them to CSV. The +submodules under `Epidemic.Types` provide combinators and functionality for +working these these data and writing your own simulations. There are some basic +simulation models already provided in the `Epidemic.Model` submodules. For the +most part, writing a new epidemic model revolves around definining the +`randomEvent` function. diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,120 @@ # Changelog for epi-sim +## 0.4.2 + +- Include `simulationWithGenIO` and add `scRequireCherry` to the + `SimulationConfiguration` record type to make it easier to control how + simulations are conditioned upon particular observations. +- Documentation. +- Bug fix in edge case of no sequenced samples. + +## 0.4.1 + +- Update the `simulationWithSystemRandom` function so that this works again. + +## 0.4.0.0 + +- The following changes to the `EpidemicEvent` type will be the real sticking + point in moving from `0.3.0.0` to `0.4.0.0`: + +``` +-- | Events that can occur in an epidemic with their absolute time. +data EpidemicEvent + = Infection AbsoluteTime Person Person -- ^ infection time, infector, infectee + | Removal AbsoluteTime Person -- ^ removal without observation + | Sampling AbsoluteTime Person -- ^ removal and inclusion in phylogeny + | Catastrophe AbsoluteTime People -- ^ scheduled sampling of lineages + | Occurrence AbsoluteTime Person -- ^ removal and observed by not in phylogeny + | Disaster AbsoluteTime People -- ^ scheduled occurrence of lineages + | Extinction -- ^ epidemic went extinct time time can be recovered from the preceeding removal + | StoppingTime -- ^ the simulation reached the stopping time + deriving (Show, Generic, Eq) +``` + +becomes + +``` +-- | Events that can occur in an epidemic with their absolute time. +data EpidemicEvent + = Infection AbsoluteTime Infector Infectee + | Removal AbsoluteTime Person + | IndividualSample + { indSampTime :: AbsoluteTime + , indSampPerson :: Person + , indSampSeq :: Bool + } + | PopulationSample + { popSampTime :: AbsoluteTime + , popSampPeople :: People + , popSampSeq :: Bool + } + | Extinction -- ^ epidemic went extinct time time can be recovered from the preceeding removal + | StoppingTime -- ^ the simulation reached the stopping time + deriving (Show, Generic, Eq) +``` + +- Remove CSV export, now there is only JSON export. If you want to include an + orphan instance for working with `cassava` the following might be useful + +``` +instance Csv.ToRecord EpidemicEvent where + toRecord e = + case e of + (Infection time person1 person2) -> + Csv.record + [ "infection" + , Csv.toField time + , Csv.toField person1 + , Csv.toField person2 + ] + (Removal time person) -> + Csv.record ["removal", Csv.toField time, Csv.toField person, "NA"] + (Sampling time person) -> + Csv.record ["sampling", Csv.toField time, Csv.toField person, "NA"] + (Catastrophe time people) -> + Csv.record ["catastrophe", Csv.toField time, Csv.toField people, "NA"] + (Occurrence time person) -> + Csv.record ["occurrence", Csv.toField time, Csv.toField person, "NA"] + (Disaster time people) -> + Csv.record ["disaster", Csv.toField time, Csv.toField people, "NA"] + Extinction -> Csv.record ["extinction", "NA", "NA", "NA"] + StoppingTime -> Csv.record ["stop", "NA", "NA", "NA"] +``` + +- Move the Newick material into `Epidemic.Types.Newick`. +- Remvoe the `TransmissionTree` and `SampleTree` data types because there is + already the `EpidemicTree` and `ReconstructedTree` in the + `Epidemic.Types.Events` which should be used preferentially. +- Remove dependency on `trifecta` since this functionality is not necessary. + +## 0.3.0.0 + +- Add an `ARCHITECTURE.md` file to outline the structure of this package. +- Move all the temporal types into their own submodule `Epidemic.Types.Time` so + they are easier to isolate. +- Provide a `Epidemic.Types.Observations` module to provide the functionality + surrounding extracting the observed events from a simulation. +- Add `Extinction` and `StoppingTime` constructors for the `EpidemicEvent` type + so that we can encode why the simulation finished in the events. As a result + of this change, every simulation that was not terminated early should end with + an `Extinction` event or a `StoppingTime` event. +- Update the resolver to 17.2 +- Extend the `ModelParameters` class to include a `Population` parameter type + since this is needed to compute event rates in for the logistic model. +- Create `Epidemic.Type.Simulation` module for types relating to running generic + simulations from the models to avoid confusion as to where these are defined. +- Move class definitions into corresponding `Epidemic.Type.X` modules. +- Add a `Epidemic.Model.LogisticBDSD` module implementing a logistic birth-death + process with unscheduled sampling and scheduled unsequenced sampling. +- Move the models into a new `Epidemic.Model` module so that it is clearer that + these are really just examples of putting together functionality provided by + the rest of the library. +- Use a new type `Identifier` to represent identities of people rather than a + raw integer this way it is clearer what it really is. +- Replace `Time` with `AbsoluteTime` and `TimeDelta` types to make it explicit + what is being represented. +- Reduce the number of models that are included to the more interesting subset. + ## 0.2.2.0 - Clean up for a release candidate. @@ -10,7 +125,7 @@ versions to use. - Remove dependency upon `epi-types` by moving its modules into this package and include the tests from that package. -- Remove unsed `Setup.hs` file. +- Remove unused `Setup.hs` file. ## 0.2.0.1 diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -4,21 +4,10 @@ ## Available models -1. Birth-Death (see `Epidemic.BirthDeath`) -2. Birth-Death-Sampling (see `Epidemic.BirthDeathSampling`) -3. Birth-Death-Sampling-Occurrence (see `Epidemic.BirthDeathSamplingOccurrence`) -4. Birth-Death-Sampling-Catastrophe-Occurrence (see `Epidemic.BirthDeathSamplingCatastropheOccurrence`) -5. Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see `Epidemic.BDSCOD`) -6. Inhomogeneous Birth-Death (see `Epidemic.InhomogeneousBD`) -7. Inhomogeneous Birth-Death-Sampling (see `Epidemic.InhomogeneousBDS`) - -## Output +Although this package supports the definition of new models there are some that +are implemented already in the `Epidemic.Model` module. Implemented models +include: -The output is a CSV with a header encoding which events occurred when and to -whom: `event,time,primaryPerson,secondaryPerson`. The *primary person* is either -the infecting person or the person who has been removed in some manner, the -*secondary person* is the person who was infected, or this is a missing value. -There are functions to assist in extracting observations from a full simulation: -`birthDeathSamplingOccurrenceObservedEvents`. In the case of a catastrophe event -where multiple individuals may be removed, they are represented as a colon -separated list of identifiers in the `primaryPerson` field. +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`) diff --git a/epi-sim.cabal b/epi-sim.cabal --- a/epi-sim.cabal +++ b/epi-sim.cabal @@ -1,106 +1,85 @@ -cabal-version: 1.22 -name: epi-sim -version: 0.2.2.0 -synopsis: A library for simulating epidemics as birth-death processes. +cabal-version: 1.22 +name: epi-sim +version: 0.4.2 +synopsis: + A library for simulating epidemics as birth-death processes. + description: A library for simulating epidemics, with a focus on phylodynamics and observation models. . - /Available models/ - . - * Birth-Death (see `Epidemic.BirthDeath`) - . - * Birth-Death-Sampling (see `Epidemic.BirthDeathSampling`) - . - * Birth-Death-Sampling-Occurrence (see `Epidemic.BirthDeathSamplingOccurrence`) - . - * Birth-Death-Sampling-Catastrophe-Occurrence (see `Epidemic.BirthDeathSamplingCatastropheOccurrence`) - . - * Birth-Death-Sampling-Catastrophe-Occurrence-Disaster (see `Epidemic.BDSCOD`) - . - * Inhomogeneous Birth-Death (see `Epidemic.InhomogeneousBD`) - . - * Inhomogeneous Birth-Death-Sampling (see `Epidemic.InhomogeneousBDS`) - . - /Output format/ + Although this package supports the definition of new models there are some that + are implemented already in the `Epidemic.Model` module. Implemented models + include: . - The output is a CSV with a header encoding which events occurred when and to - whom: @event,time,primaryPerson,secondaryPerson@. The @primary person@ is either - the infecting person or the person who has been removed in some manner, the - @secondary person@ is the person who was infected, or this is a missing value. - There are functions to assist in extracting observations from a full simulation: - @birthDeathSamplingOccurrenceObservedEvents@. In the case of a catastrophe event - where multiple individuals may be removed, they are represented as a colon - separated list of identifiers in the @primaryPerson@ field. + 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`) . -homepage: https://github.com/aezarebski/epi-sim#readme -bug-reports: https://github.com/aezarebski/epi-sim/issues -author: Alexander Zarebski -maintainer: aezarebski@gmail.com -copyright: 2020 Alexander Zarebski -license: MIT -license-file: LICENSE -build-type: Simple -category: Simulation + +homepage: https://github.com/aezarebski/epi-sim#readme +bug-reports: https://github.com/aezarebski/epi-sim/issues +author: Alexander Zarebski +maintainer: aezarebski@gmail.com +copyright: 2020 Alexander Zarebski +license: MIT +license-file: LICENSE +build-type: Simple +category: Simulation extra-source-files: - README.md - ChangeLog.md + ChangeLog.md + README.md + ARCHITECTURE.md source-repository head - type: git + type: git location: https://github.com/aezarebski/epi-sim library exposed-modules: - Epidemic - Epidemic.BDSCOD - Epidemic.BirthDeath - Epidemic.BirthDeathSampling - Epidemic.BirthDeathSamplingCatastropheOccurrence - Epidemic.BirthDeathSamplingOccurrence - Epidemic.InhomogeneousBD - Epidemic.InhomogeneousBDS - Epidemic.Types.Events - Epidemic.Types.Observations - Epidemic.Types.Parameter - Epidemic.Types.Population - Epidemic.Utility - other-modules: - Paths_epi_sim - hs-source-dirs: - src + Epidemic + Epidemic.Model.BDSCOD + Epidemic.Model.InhomogeneousBDS + Epidemic.Model.LogisticBDSD + Epidemic.Types.Events + Epidemic.Types.Newick + Epidemic.Types.Observations + Epidemic.Types.Parameter + Epidemic.Types.Population + Epidemic.Types.Simulation + Epidemic.Types.Time + Epidemic.Utility + + other-modules: Paths_epi_sim + hs-source-dirs: src build-depends: - aeson >= 1.4.0 && < 1.5, - base >= 4.8.2 && < 4.14, - bytestring >= 0.10.6 && < 0.11, - primitive >= 0.6.1 && < 0.8, - vector >= 0.11.0 && < 0.13, - cassava >= 0.5.2 && < 0.6, - hspec >= 2.7.4 && < 2.8, - mwc-random >= 0.14.0 && < 0.15, - statistics >= 0.15.0 && < 0.16, - trifecta >= 2.1 && < 2.2 + aeson >=1.4.7 && <1.6 + , base >=4.14.1 && <4.15 + , bytestring >=0.10.10 && <0.11 + , hspec >=2.7.6 && <2.8 + , mwc-random >=0.14.0 && <0.16 + , primitive >=0.7.0 && <0.8 + , statistics >=0.15.2 && <0.16 + , vector >=0.12.1 && <0.13 + default-language: Haskell2010 - ghc-options: -Wincomplete-patterns + ghc-options: -Wincomplete-patterns test-suite epi-sim-test - type: exitcode-stdio-1.0 - main-is: Spec.hs - other-modules: - Paths_epi_sim - hs-source-dirs: - test - ghc-options: -threaded -rtsopts -with-rtsopts=-N + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: Paths_epi_sim + hs-source-dirs: test + ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: - aeson, - base, - bytestring, - epi-sim, - primitive, - vector, - cassava, - hspec, - mwc-random, - statistics, - trifecta + aeson + , base + , bytestring + , epi-sim + , hspec + , mwc-random + , primitive + , statistics + , vector + default-language: Haskell2010 diff --git a/src/Epidemic.hs b/src/Epidemic.hs --- a/src/Epidemic.hs +++ b/src/Epidemic.hs @@ -1,206 +1,135 @@ +{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveAnyClass #-} {-# LANGUAGE OverloadedStrings #-} + module Epidemic where import Control.Monad import qualified Data.ByteString as B import Data.ByteString.Internal (c2w) -import Data.Csv import Data.List (nub) +import Data.Maybe (fromJust, isJust, isNothing) import qualified Data.Vector as V import Data.Word import Epidemic.Types.Events import Epidemic.Types.Parameter import Epidemic.Types.Population +import Epidemic.Types.Simulation + ( SimulationConfiguration(..) + , SimulationRandEvent(..) + , SimulationState(..) + ) +import Epidemic.Types.Time (AbsoluteTime(..), Timed(..), diracDeltaValue, nextTime) import GHC.Generics (Generic) - - +import System.Random.MWC -- | The number of people added or removed in an event. eventPopDelta :: EpidemicEvent -> Integer -eventPopDelta e = case e of - Infection{} -> 1 - Removal _ _ -> -1 - Sampling _ _ -> -1 - Catastrophe _ people -> fromIntegral $ numPeople people - Occurrence _ _ -> -1 - Disaster _ people -> fromIntegral $ numPeople people +eventPopDelta e = + case e of + Infection {} -> 1 + Removal {} -> -1 + IndividualSample {} -> -1 + PopulationSample {..} -> fromIntegral $ numPeople popSampPeople + StoppingTime -> 0 -- | The first scheduled event after a given time. -firstScheduled :: Time -- ^ The given time - -> Timed Probability -- ^ The information about all scheduled events - -> Maybe (Time,Probability) +firstScheduled :: + AbsoluteTime -- ^ The given time + -> Timed Probability -- ^ The information about all scheduled events + -> Maybe (AbsoluteTime, Probability) firstScheduled time timedProb = do time' <- nextTime timedProb time prob' <- diracDeltaValue timedProb time' - return (time',prob') + return (time', prob') -- | Predicate for whether there is a scheduled event during an interval. -noScheduledEvent :: Time -- ^ Start time for interval - -> Time -- ^ End time for interval - -> Timed Probability -- ^ Information about all scheduled events - -> Bool +noScheduledEvent :: + AbsoluteTime -- ^ Start time for interval + -> AbsoluteTime -- ^ End time for interval + -> Timed Probability -- ^ Information about all scheduled events + -> Bool noScheduledEvent _ _ (Timed []) = True noScheduledEvent a b (Timed ((shedTime, _):scheduledEvents)) = - not (a < shedTime && shedTime <= b) && noScheduledEvent a b (Timed scheduledEvents) + not (a < shedTime && shedTime <= b) && + noScheduledEvent a b (Timed scheduledEvents) +-- | A list of the people involved in an 'EpidemicEvent'. personsInEvent :: EpidemicEvent -> [Person] -personsInEvent e = case e of - (Infection _ p1 p2) -> [p1,p2] - (Removal _ p) -> [p] - (Sampling _ p) -> [p] - (Catastrophe _ (People persons)) -> V.toList persons - (Occurrence _ p) -> [p] - (Disaster _ (People persons)) -> V.toList persons +personsInEvent e = + case e of + Infection _ p1 p2 -> [p1, p2] + Removal _ p -> [p] + (IndividualSample {..}) -> [indSampPerson] + (PopulationSample {..}) -> + V.toList personVec + where + (People personVec) = popSampPeople + Extinction -> [] + StoppingTime -> [] peopleInEvents :: [EpidemicEvent] -> People peopleInEvents events = People . V.fromList . nub . concat $ map personsInEvent events - -- | Predicate for whether the first person infected the second in the given event -infected :: Person -- ^ Potential infector - -> Person -- ^ Potential infectee - -> EpidemicEvent -- ^ Given event - -> Bool +infected :: + Person -- ^ Potential infector + -> Person -- ^ Potential infectee + -> EpidemicEvent -- ^ Given event + -> Bool infected p1 p2 e = case e of (Infection _ infector infectee) -> infector == p1 && infectee == p2 _ -> False - -- | The people infected by a particular person in a list of events. -infectedBy :: Person -- ^ Potential infector - -> [EpidemicEvent] -- ^ Events - -> People +infectedBy :: + Person -- ^ Potential infector + -> [EpidemicEvent] -- ^ Events + -> People infectedBy person events = case events of [] -> People V.empty - (Infection _ infector infectee :es) -> + (Infection _ infector infectee:es) -> if infector == person then addPerson infectee $ infectedBy person es else infectedBy person es (_:es) -> infectedBy person es - --- | Predicate for whether a person or one of their descendents satisfies a --- predicate -hasDescendentWhich :: [EpidemicEvent] - -> (Person -> Bool) - -> Person - -> Bool -hasDescendentWhich events predicate person = - predicate person || - any (hasDescendentWhich events predicate) (V.toList descendents) - where - (People descendents) = infectedBy person events - -hasSampledDescendent :: [EpidemicEvent] -> Person -> Bool -hasSampledDescendent events = hasDescendentWhich events (wasSampled events) - --- | Predicate for whether a person was sampled in the given events -wasSampled :: [EpidemicEvent] -- ^ The given events - -> Person -- ^ The person of interest - -> Bool -wasSampled events person = - case events of - (Sampling _ sampledPerson:es) -> - sampledPerson == person || wasSampled es person - (Catastrophe _ (People sampledPeople):es) -> - person `V.elem` sampledPeople || wasSampled es person - (_:es) -> wasSampled es person - [] -> False - --- | Return the sampling event of a person who was sampled. -samplingEvent :: [EpidemicEvent] -> Person -> EpidemicEvent -samplingEvent events person = - case events of - (se@(Sampling _ sampledPerson):remainingEvents) -> - if sampledPerson == person - then se - else samplingEvent remainingEvents person - (se@(Catastrophe _ (People sampledPeople)):remainingEvents) -> - if person `V.elem` sampledPeople - then se - else samplingEvent remainingEvents person - _:remainingEvents -> samplingEvent remainingEvents person - _ -> error "person does not appear to have been sampled." - - -class ModelParameters a where - rNaught :: a -> Time -> Maybe Double - eventRate :: a -> Time -> Maybe Rate - birthProb :: a -> Time -> Maybe Probability - -class Population a where - susceptiblePeople :: a -> Maybe People - infectiousPeople :: a -> Maybe People - removedPeople :: a -> Maybe People - isInfected :: a -> Bool - - -data TransmissionTree - = TTUnresolved Person - | TTDeath People EpidemicEvent - | TTBirth Person EpidemicEvent (TransmissionTree, TransmissionTree) - deriving (Show) - --- | A transmission tree of all the events starting from a given person -transmissionTree :: [EpidemicEvent] -> Person -> TransmissionTree -transmissionTree (e@(Infection _ p1 p2):es) person - | p1 == person = TTBirth person e (transmissionTree es p1,transmissionTree es p2) - | null es = TTUnresolved person - | otherwise = transmissionTree es person -transmissionTree (e@(Removal _ p1):es) person - | p1 == person = TTDeath (peopleInEvents [e]) e - | otherwise = transmissionTree es person -transmissionTree (e@(Sampling _ p1):es) person - | p1 == person = TTDeath (peopleInEvents [e]) e - | otherwise = transmissionTree es person -transmissionTree (e@(Catastrophe _ (People people)):es) person - | person `V.elem` people = TTDeath (People people) e - | otherwise = transmissionTree es person -transmissionTree (e@(Occurrence _ p1):es) person - | p1 == person = TTDeath (peopleInEvents [e]) e - | otherwise = transmissionTree es person -transmissionTree (e@(Disaster _ (People people)):es) person - | person `V.elem` people = TTDeath (People people) e - | otherwise = transmissionTree es person -transmissionTree [] person = TTUnresolved person - --- | A predicate for whether there is a sampled leaf in the transmission tree -hasSampledLeaf :: TransmissionTree -> Bool -hasSampledLeaf t = case t of - (TTUnresolved _) -> False - (TTDeath _ (Sampling _ _)) -> True - (TTDeath _ (Catastrophe _ _)) -> True - (TTDeath _ _) -> False - (TTBirth _ _ (t1,t2)) -> hasSampledLeaf t1 || hasSampledLeaf t2 - -data SampleTree - = STBirth EpidemicEvent (SampleTree,SampleTree) - | STDeath EpidemicEvent - deriving (Show) - --- | A transmission tree with all non-sampling leaves removed -sampleTree :: TransmissionTree -> SampleTree -sampleTree transTree = case transTree of - (TTBirth _ e@Infection {} (t1,t2)) - | hasSampledLeaf t1 && hasSampledLeaf t2 -> STBirth e (sampleTree t1,sampleTree t2) - | hasSampledLeaf t1 -> sampleTree t1 - | hasSampledLeaf t2 -> sampleTree t2 - (TTDeath _ e@(Sampling _ _)) -> STDeath e - (TTDeath _ e@(Catastrophe _ _)) -> STDeath e - _ -> error "ill-formed transmission tree" - --- | Recurse through the tree and extract all birth and death events. -sampleTreeEvents' :: SampleTree -> [EpidemicEvent] -sampleTreeEvents' sTree = - case sTree of - (STDeath e) -> [e] - (STBirth e (s1, s2)) -> e : sampleTreeEvents s1 ++ sampleTreeEvents s2 - --- | The unique events in a sample tree. -sampleTreeEvents :: SampleTree -> [EpidemicEvent] -sampleTreeEvents = nub . sampleTreeEvents' +-- | Run the simulation and return a @SimulationState@ which holds the history +-- of the simulation. +allEvents :: + (ModelParameters a b, Population b) + => SimulationRandEvent a b + -> a + -> AbsoluteTime + -> Maybe (b -> Bool) -- ^ predicate for a valid population + -> SimulationState b + -> GenIO + -> IO (SimulationState b) +allEvents _ _ _ _ TerminatedSimulation _ = return TerminatedSimulation +allEvents simRandEvent@(SimulationRandEvent randEvent) modelParams maxTime maybePopPredicate (SimulationState (currTime, currEvents, currPop, currId)) gen = + if isNothing maybePopPredicate || + (isJust maybePopPredicate && fromJust maybePopPredicate currPop) + then if isInfected currPop + then do + (newTime, event, newPop, newId) <- + randEvent modelParams currTime currPop currId gen + if newTime < maxTime + then allEvents + simRandEvent + modelParams + maxTime + maybePopPredicate + (SimulationState + (newTime, event : currEvents, newPop, newId)) + gen + else return $ + SimulationState + (maxTime, StoppingTime : currEvents, currPop, currId) + else return $ + SimulationState + (currTime, Extinction : currEvents, currPop, currId) + else return TerminatedSimulation diff --git a/src/Epidemic/BDSCOD.hs b/src/Epidemic/BDSCOD.hs deleted file mode 100644 --- a/src/Epidemic/BDSCOD.hs +++ /dev/null @@ -1,180 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -module Epidemic.BDSCOD - ( configuration - , allEvents - , observedEvents - ) where - -import Data.List (nub) -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import qualified Data.Vector.Generic as G -import Epidemic -import Epidemic.Types.Events - ( EpidemicEvent(..) - , PointProcessEvents(..) - , ReconstructedTree(..) - , maybeEpidemicTree - , maybeReconstructedTree - , pointProcessEvents - ) -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Utility -import System.Random.MWC -import System.Random.MWC.Distributions (bernoulli, categorical, exponential) - - -data BDSCODParameters - -- | birth rate, death rate, sampling rate, catastrophe specification, occurrence rate and disaster specification - = BDSCODParameters Rate Rate Rate (Timed Probability) Rate (Timed Probability) - -instance ModelParameters BDSCODParameters 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) - -newtype BDSCODPopulation = - BDSCODPopulation People - deriving (Show) - -instance Population BDSCODPopulation where - susceptiblePeople _ = Nothing - infectiousPeople (BDSCODPopulation people) = Just people - removedPeople _ = Nothing - isInfected (BDSCODPopulation (People people)) = not $ V.null people - --- | Configuration of a birth-death-sampling-occurrence simulation -configuration :: Time -- ^ Duration of the simulation - -> (Rate,Rate,Rate,[(Time,Probability)],Rate,[(Time,Probability)]) -- ^ Birth, Death, Sampling, Catastrophe probability and Occurrence rates - -> Maybe (SimulationConfiguration BDSCODParameters BDSCODPopulation) -configuration maxTime (birthRate, deathRate, samplingRate, catastropheSpec, occurrenceRate, disasterSpec) = - do catastropheSpec' <- asTimed catastropheSpec - disasterSpec' <- asTimed disasterSpec - let bdscodParams = - BDSCODParameters - birthRate - deathRate - samplingRate - catastropheSpec' - occurrenceRate - disasterSpec' - (seedPerson, newId) = newPerson initialIdentifier - bdscodPop = BDSCODPopulation (People $ V.singleton seedPerson) - in return $ SimulationConfiguration bdscodParams bdscodPop newId maxTime - --- | Return a random event from the BDSCOD-process given the current state of the process. -randomEvent :: BDSCODParameters -- ^ Parameters of the process - -> Time -- ^ The current time within the process - -> BDSCODPopulation -- ^ The current state of the populaion - -> Integer -- ^ The current state of the identifier generator - -> GenIO -- ^ The current state of the PRNG - -> IO (Time, EpidemicEvent, BDSCODPopulation, Integer) -randomEvent params@(BDSCODParameters br dr sr catastInfo occr disastInfo) currTime currPop@(BDSCODPopulation (People currPeople)) currId gen = - let netEventRate = fromJust $ eventRate params currTime - eventWeights = V.fromList [br, dr, sr, occr] - in do delay <- exponential (fromIntegral (V.length currPeople) * netEventRate) gen - nextTime <- pure $ currTime + delay - if noScheduledEvent currTime nextTime (catastInfo <> disastInfo) - then do eventIx <- categorical eventWeights gen - (selectedPerson, unselectedPeople) <- randomPerson currPeople gen - return $ case eventIx of - 0 -> let (birthedPerson, newId) = newPerson currId - event = Infection nextTime selectedPerson birthedPerson - in ( nextTime - , event - , BDSCODPopulation (People $ V.cons birthedPerson currPeople) - , newId) - 1 -> (nextTime, Removal nextTime selectedPerson, BDSCODPopulation (People unselectedPeople), currId) - 2 -> (nextTime, Sampling nextTime selectedPerson, BDSCODPopulation (People unselectedPeople), currId) - 3 -> (nextTime, Occurrence nextTime selectedPerson, BDSCODPopulation (People unselectedPeople), currId) - _ -> error "no birth, death, sampling, occurrence event selected." - - else if noScheduledEvent currTime nextTime 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 nextTime 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) - - --- | Return a randomly sampled Catastrophe event -randomCatastropheEvent :: (Time,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 - sampledPeople = filterZip snd currPeople rhoBernoullis - unsampledPeople = filterZip (not . snd) currPeople rhoBernoullis - in return - ( Catastrophe catastTime (People sampledPeople) - , BDSCODPopulation (People unsampledPeople)) - --- | Return a randomly sampled Disaster event -randomDisasterEvent :: (Time,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 - sampledPeople = filterZip snd currPeople nuBernoullis - unsampledPeople = filterZip (not . snd) currPeople nuBernoullis - in return - ( Disaster disastTime (People sampledPeople) - , BDSCODPopulation (People unsampledPeople)) - -allEvents :: - BDSCODParameters - -> Time - -> (Time, [EpidemicEvent], BDSCODPopulation, Integer) - -> GenIO - -> IO (Time, [EpidemicEvent], BDSCODPopulation, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - --- | The events from the nodes of a reconstructed tree __not__ in time sorted --- order. -reconstructedTreeEvents :: ReconstructedTree -> [EpidemicEvent] -reconstructedTreeEvents node = case node of - (RBranch e lt rt) -> e:(reconstructedTreeEvents lt ++ reconstructedTreeEvents rt) - (RLeaf e) -> [e] - --- | Just the observable events from a list of all the events that occurred in a --- simulation of the BDSCOD-process. These events are the result of extracting --- the events from the reconstructed tree and getting the point process events --- that make up the unsequenced samples (see `pointProcessEvents` for details on --- this latter data.) -observedEvents :: [EpidemicEvent] -- ^ All of the simulation events - -> Maybe [EpidemicEvent] -observedEvents eEvents = do - epiTree <- maybeEpidemicTree eEvents - reconTree <- maybeReconstructedTree epiTree - let (PointProcessEvents nonReconTreeEvents) = pointProcessEvents epiTree - let reconTreeEvents = reconstructedTreeEvents reconTree - return . sort . nub $ nonReconTreeEvents ++ reconTreeEvents diff --git a/src/Epidemic/BirthDeath.hs b/src/Epidemic/BirthDeath.hs deleted file mode 100644 --- a/src/Epidemic/BirthDeath.hs +++ /dev/null @@ -1,97 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -module Epidemic.BirthDeath - ( configuration - , allEvents - ) where - -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Types.Events -import qualified Data.Vector as V -import System.Random.MWC -import System.Random.MWC.Distributions (bernoulli, exponential) - -import Epidemic -import Epidemic.Utility - -data BDRates = - BDRates Rate Rate - -instance ModelParameters BDRates where - rNaught (BDRates birthRate deathRate) _ = Just $ birthRate / deathRate - eventRate (BDRates birthRate deathRate) _ = Just $ birthRate + deathRate - birthProb (BDRates birthRate deathRate) _ = Just $ birthRate / (birthRate + deathRate) - -newtype BDPopulation = - BDPopulation People - deriving (Show) - -instance Population BDPopulation where - susceptiblePeople _ = Nothing - infectiousPeople (BDPopulation people) = Just people - removedPeople _ = Nothing - isInfected (BDPopulation (People people)) = not $ V.null people - --- | Return a BD-process parameters object -birthDeathRates :: Rate -- ^ birth rate - -> Rate -- ^ death rate - -> Maybe BDRates -birthDeathRates birthRate deathRate - | birthRate >= 0 && deathRate >= 0 = Just $ BDRates birthRate deathRate - | otherwise = Nothing - --- | Configuration of a birth-death simulation. -configuration :: Time -- ^ Duration of the simulation - -> (Rate, Rate) -- ^ Birth and Death rates - -> Maybe (SimulationConfiguration BDRates BDPopulation) -configuration maxTime (birthRate, deathRate) = - let (seedPerson, newId) = newPerson initialIdentifier - bdPop = BDPopulation (People $ V.singleton seedPerson) - in do maybeBDRates <- birthDeathRates birthRate deathRate - if maxTime > 0 then Just (SimulationConfiguration maybeBDRates bdPop newId maxTime) else Nothing - -randomBirthDeathEvent :: - BDRates - -> Time - -> BDPopulation - -> Integer - -> GenIO - -> IO (Time, EpidemicEvent, BDPopulation, Integer) -randomBirthDeathEvent (BDRates br dr) currTime (BDPopulation (People currPeople)) currId gen = do - delay <- exponential (fromIntegral (V.length currPeople) * (br + dr)) gen - isBirth <- bernoulli (br / (br + dr)) gen - (selectedPerson, unselectedPeople) <- randomPerson currPeople gen - return $ - if isBirth - then let newTime = currTime + delay - (birthedPerson, newId) = newPerson currId - event = Infection newTime selectedPerson birthedPerson - in ( newTime - , event - , BDPopulation (People $ V.cons birthedPerson currPeople) - , newId) - else let newTime = currTime + delay - event = Removal newTime selectedPerson - in (newTime, event, BDPopulation (People unselectedPeople), currId) - -allEvents :: - BDRates - -> Time - -> (Time, [EpidemicEvent], BDPopulation, Integer) - -> GenIO - -> IO (Time, [EpidemicEvent], BDPopulation, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomBirthDeathEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - diff --git a/src/Epidemic/BirthDeathSampling.hs b/src/Epidemic/BirthDeathSampling.hs deleted file mode 100644 --- a/src/Epidemic/BirthDeathSampling.hs +++ /dev/null @@ -1,101 +0,0 @@ -module Epidemic.BirthDeathSampling - ( configuration - , allEvents - ) where - - -import Epidemic.Types.Parameter -import Epidemic.Types.Population -import Epidemic.Types.Events -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import System.Random.MWC -import System.Random.MWC.Distributions (categorical, exponential) - -import Epidemic -import Epidemic.Utility - -data BDSRates = - BDSRates Rate Rate Rate - -instance ModelParameters BDSRates where - rNaught (BDSRates birthRate deathRate samplingRate) _ = - Just $ birthRate / (deathRate + samplingRate) - eventRate (BDSRates birthRate deathRate samplingRate) _ = - Just $ birthRate + deathRate + samplingRate - birthProb (BDSRates birthRate deathRate samplingRate) _ = - Just $ birthRate / (birthRate + deathRate + samplingRate) - -newtype BDSPopulation = - BDSPopulation People - deriving (Show) - -instance Population BDSPopulation where - susceptiblePeople _ = Nothing - infectiousPeople (BDSPopulation people) = Just people - removedPeople _ = Nothing - isInfected (BDSPopulation (People people)) = not $ V.null people - -birthDeathSamplingRates :: Rate -> Rate -> Rate -> BDSRates -birthDeathSamplingRates = BDSRates -- birthRate deathRate samplingRate - --- | Configuration of a birth-death-sampling simulation. -configuration :: Time -- ^ Duration of the simulation - -> (Rate,Rate,Rate) -- ^ Birth, Death and Sampling rates - -> SimulationConfiguration BDSRates BDSPopulation -configuration maxTime (birthRate, deathRate, samplingRate) = - let bdsRates = birthDeathSamplingRates birthRate deathRate samplingRate - (seedPerson, newId) = newPerson initialIdentifier - bdsPop = BDSPopulation (People $ V.singleton seedPerson) - in SimulationConfiguration bdsRates bdsPop newId maxTime - -randomBirthDeathSamplingEvent :: - BDSRates - -> Time - -> BDSPopulation - -> Integer - -> GenIO - -> IO (Time, EpidemicEvent, BDSPopulation, Integer) -randomBirthDeathSamplingEvent bdsRates@(BDSRates br dr sr) currTime (BDSPopulation (People currPeople)) currId gen = - let netEventRate = fromJust $ eventRate bdsRates currTime - eventWeights = V.fromList [br,dr,sr] - in - do delay <- exponential (fromIntegral (V.length currPeople) * netEventRate) gen - eventIx <- categorical eventWeights gen - (selectedPerson, unselectedPeople) <- randomPerson currPeople gen - return $ case eventIx of - 0 -> let newTime = currTime + delay - (birthedPerson, newId) = newPerson currId - event = Infection newTime selectedPerson birthedPerson - in ( newTime - , event - , BDSPopulation (People $ V.cons birthedPerson currPeople) - , newId) - 1 -> let newTime = currTime + delay - event = Removal newTime selectedPerson - in (newTime, event, BDSPopulation (People unselectedPeople), currId) - 2 -> let newTime = currTime + delay - event = Sampling newTime selectedPerson - in (newTime, event, BDSPopulation (People unselectedPeople), currId) - _ -> error "no birth-death-sampling event selected." - -allEvents :: - BDSRates - -> Time - -> (Time, [EpidemicEvent], BDSPopulation, Integer) - -> GenIO - -> IO (Time, [EpidemicEvent], BDSPopulation, Integer) -allEvents bdsRates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomBirthDeathSamplingEvent bdsRates currTime currPop currId gen - if newTime < maxTime - then allEvents - bdsRates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - diff --git a/src/Epidemic/BirthDeathSamplingCatastropheOccurrence.hs b/src/Epidemic/BirthDeathSamplingCatastropheOccurrence.hs deleted file mode 100644 --- a/src/Epidemic/BirthDeathSamplingCatastropheOccurrence.hs +++ /dev/null @@ -1,136 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -module Epidemic.BirthDeathSamplingCatastropheOccurrence - ( configuration - , allEvents - , observedEvents - ) where - -import Epidemic.Types.Population -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import qualified Data.Vector.Generic as G -import System.Random.MWC -import System.Random.MWC.Distributions (categorical, exponential, bernoulli) - -import Epidemic -import Epidemic.Utility - - -data BDSCOParameters - -- | birth rate, death rate, sampling rate, catastrophe probability and occurrence rate. - = BDSCOParameters Rate Rate Rate (Timed Probability) Rate - -instance ModelParameters BDSCOParameters where - rNaught (BDSCOParameters birthRate deathRate samplingRate _ occurrenceRate) _ = - Just $ birthRate / (deathRate + samplingRate + occurrenceRate) - eventRate (BDSCOParameters birthRate deathRate samplingRate _ occurrenceRate) _ = - Just $ birthRate + deathRate + samplingRate + occurrenceRate - birthProb (BDSCOParameters birthRate deathRate samplingRate _ occurrenceRate) _ = - Just $ birthRate / (birthRate + deathRate + samplingRate + occurrenceRate) - -newtype BDSCOPopulation = - BDSCOPopulation People - deriving (Show) - -instance Population BDSCOPopulation where - susceptiblePeople _ = Nothing - infectiousPeople (BDSCOPopulation people) = Just people - removedPeople _ = Nothing - isInfected (BDSCOPopulation (People people)) = not $ V.null people - --- | Configuration of a birth-death-sampling-occurrence simulation -configuration :: Time -- ^ Duration of the simulation - -> (Rate,Rate,Rate,[(Time,Probability)],Rate) -- ^ Birth, Death, Sampling, Catastrophe probability and Occurrence rates - -> Maybe (SimulationConfiguration BDSCOParameters BDSCOPopulation) -configuration maxTime (birthRate, deathRate, samplingRate, catastropheProb, occurrenceRate) = do - catastropheTimedProb <- asTimed catastropheProb - let bdscoParams = - BDSCOParameters - birthRate - deathRate - samplingRate - catastropheTimedProb - occurrenceRate - (seedPerson, newId) = newPerson initialIdentifier - bdscoPop = BDSCOPopulation (People $ V.singleton seedPerson) - in Just $ SimulationConfiguration bdscoParams bdscoPop newId maxTime - --- | Return a random event from the BDSCO-process given the current state of the process. -randomBdscoEvent :: - BDSCOParameters -- ^ Parameters of the process - -> Time -- ^ The current time within the process - -> BDSCOPopulation -- ^ The current state of the populaion - -> Integer -- ^ The current state of the identifier generator - -> GenIO -- ^ The current state of the PRNG - -> IO (Time, EpidemicEvent, BDSCOPopulation, Integer) -randomBdscoEvent params@(BDSCOParameters br dr sr catastInfo occr) currTime currPop@(BDSCOPopulation (People people)) currId gen = - let netEventRate = fromJust $ eventRate params currTime - eventWeights = V.fromList [br, dr, sr, occr] - in - do delay <- exponential (fromIntegral (V.length people) * netEventRate) gen - nextEventTime <- pure $ currTime + delay - if noScheduledEvent currTime nextEventTime catastInfo - then do eventIx <- categorical eventWeights gen - (selectedPerson, unselectedPeople) <- randomPerson people gen - return $ case eventIx of - 0 -> let (birthedPerson, newId) = newPerson currId - event = Infection nextEventTime selectedPerson birthedPerson - in ( nextEventTime - , event - , BDSCOPopulation (People $ V.cons birthedPerson people) - , newId) - 1 -> (nextEventTime, Removal nextEventTime selectedPerson, BDSCOPopulation (People unselectedPeople), currId) - 2 -> (nextEventTime, Sampling nextEventTime selectedPerson, BDSCOPopulation (People unselectedPeople), currId) - 3 -> (nextEventTime, Occurrence nextEventTime selectedPerson, BDSCOPopulation (People unselectedPeople), currId) - _ -> error "no birth, death, sampling, occurrence event selected." - else let (Just (catastTime,catastProb)) = firstScheduled currTime catastInfo - in do (catastEvent,postCatastPop) <- randomCatastropheEvent (catastTime,catastProb) currPop gen - return (catastTime,catastEvent,postCatastPop,currId) - - --- | Return a randomly sampled Catastrophe event -randomCatastropheEvent :: (Time,Probability) -- ^ Time and probability of sampling in the catastrophe - -> BDSCOPopulation -- ^ The state of the population prior to the catastrophe - -> GenIO - -> IO (EpidemicEvent,BDSCOPopulation) -randomCatastropheEvent (catastTime, rhoProb) (BDSCOPopulation (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 - sampledPeople = filterZip snd currPeople rhoBernoullis - unsampledPeople = filterZip (not . snd) currPeople rhoBernoullis - in return - ( Catastrophe catastTime (People sampledPeople) - , BDSCOPopulation (People unsampledPeople)) - -allEvents :: - BDSCOParameters - -> Time - -> (Time, [EpidemicEvent], BDSCOPopulation, Integer) - -> GenIO - -> IO (Time, [EpidemicEvent], BDSCOPopulation, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomBdscoEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - - --- | Just the observable events from a list of all the events in a simulation. -observedEvents :: [EpidemicEvent] -- ^ All of the simulation events - -> [EpidemicEvent] -observedEvents events = sort $ occurrenceEvents ++ sampleTreeEvents'' - where - occurrenceEvents = filter isNonReconTreeObservation events - sampleTreeEvents'' = - sampleTreeEvents . sampleTree $ transmissionTree events (Person 1) diff --git a/src/Epidemic/BirthDeathSamplingOccurrence.hs b/src/Epidemic/BirthDeathSamplingOccurrence.hs deleted file mode 100644 --- a/src/Epidemic/BirthDeathSamplingOccurrence.hs +++ /dev/null @@ -1,119 +0,0 @@ -module Epidemic.BirthDeathSamplingOccurrence - ( configuration - , allEvents - , observedEvents - ) where - -import Epidemic.Types.Population -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import System.Random.MWC -import System.Random.MWC.Distributions (categorical, exponential) - -import Epidemic -import Epidemic.Utility - -data BDSORates = - BDSORates Rate Rate Rate Rate - -instance ModelParameters BDSORates where - rNaught (BDSORates birthRate deathRate samplingRate occurrenceRate) _ = - Just $ birthRate / (deathRate + samplingRate + occurrenceRate) - eventRate (BDSORates birthRate deathRate samplingRate occurrenceRate) _ = - Just $ birthRate + deathRate + samplingRate + occurrenceRate - birthProb (BDSORates birthRate deathRate samplingRate occurrenceRate) _ = - Just $ birthRate / (birthRate + deathRate + samplingRate + occurrenceRate) - -newtype BDSOPopulation = - BDSOPopulation People - deriving (Show) - -instance Population BDSOPopulation where - susceptiblePeople _ = Nothing - infectiousPeople (BDSOPopulation people) = Just people - removedPeople _ = Nothing - isInfected (BDSOPopulation (People people)) = not $ V.null people - -birthDeathSamplingOccurrenceRates :: Rate -> Rate -> Rate -> Rate -> BDSORates -birthDeathSamplingOccurrenceRates = BDSORates -- birthRate deathRate samplingRate occurrenceRate - --- | Configuration of a birth-death-sampling-occurrence simulation -configuration :: Time -- ^ Duration of the simulation - -> (Rate,Rate,Rate,Rate) -- ^ Birth, Death, Sampling and Occurrence rates - -> SimulationConfiguration BDSORates BDSOPopulation -configuration maxTime (birthRate, deathRate, samplingRate, occurrenceRate) = - let bdsoRates = - birthDeathSamplingOccurrenceRates - birthRate - deathRate - samplingRate - occurrenceRate - (seedPerson, newId) = newPerson initialIdentifier - bdsoPop = BDSOPopulation (People $ V.singleton seedPerson) - in SimulationConfiguration bdsoRates bdsoPop newId maxTime - -randomBirthDeathSamplingOccurrenceEvent :: - BDSORates - -> Time - -> BDSOPopulation - -> Integer - -> GenIO - -> IO (Time, EpidemicEvent, BDSOPopulation, Integer) -randomBirthDeathSamplingOccurrenceEvent rates@(BDSORates br dr sr ocr) currTime (BDSOPopulation (People currPeople)) currId gen = - let netEventRate = fromJust $ eventRate rates currTime - eventWeights = V.fromList [br,dr,sr,ocr] - in - do delay <- exponential (fromIntegral (V.length currPeople) * netEventRate) gen - eventIx <- categorical eventWeights gen - (selectedPerson, unselectedPeople) <- randomPerson currPeople gen - return $ case eventIx of - 0 -> let newTime = currTime + delay - (birthedPerson, newId) = newPerson currId - event = Infection newTime selectedPerson birthedPerson - in ( newTime - , event - , BDSOPopulation (People $ V.cons birthedPerson currPeople) - , newId) - 1 -> let newTime = currTime + delay - event = Removal newTime selectedPerson - in (newTime, event, BDSOPopulation (People unselectedPeople), currId) - 2 -> let newTime = currTime + delay - event = Sampling newTime selectedPerson - in (newTime, event, BDSOPopulation (People unselectedPeople), currId) - 3 -> let newTime = currTime + delay - event = Occurrence newTime selectedPerson - in (newTime, event, BDSOPopulation (People unselectedPeople), currId) - _ -> error "no birth-death-sampling-occurrence event selected." - -allEvents :: - BDSORates - -> Time - -> (Time, [EpidemicEvent], BDSOPopulation, Integer) - -> GenIO - -> IO (Time, [EpidemicEvent], BDSOPopulation, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomBirthDeathSamplingOccurrenceEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - - --- | Just the observable events from a list of all the events in a simulation. -observedEvents :: [EpidemicEvent] -- ^ All of the simulation events - -> [EpidemicEvent] -observedEvents events = - sort $ occurrenceEvents ++ sampleTreeEvents'' - where - occurrenceEvents = filter isNonReconTreeObservation events - sampleTreeEvents'' = - sampleTreeEvents . sampleTree $ transmissionTree events (Person 1) diff --git a/src/Epidemic/InhomogeneousBD.hs b/src/Epidemic/InhomogeneousBD.hs deleted file mode 100644 --- a/src/Epidemic/InhomogeneousBD.hs +++ /dev/null @@ -1,115 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -module Epidemic.InhomogeneousBD - ( configuration - , allEvents - ) where - -import Epidemic.Types.Population -import Epidemic.Types.Parameter -import Epidemic.Types.Events -import Control.Monad (liftM) -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import System.Random.MWC -import System.Random.MWC.Distributions (bernoulli, exponential) -import Epidemic -import Epidemic.Utility - -data InhomBDRates = - InhomBDRates (Timed Rate) Rate - -instance ModelParameters InhomBDRates where - rNaught (InhomBDRates timedBirthRate deathRate) time = - let birthRate = cadlagValue timedBirthRate time - in liftM (/ deathRate) birthRate - eventRate (InhomBDRates timedBirthRate deathRate) time = - let birthRate = cadlagValue timedBirthRate time - in liftM (+ deathRate) birthRate - birthProb (InhomBDRates timedBirthRate deathRate) time = - liftM (\br -> br / (br + deathRate)) $ cadlagValue timedBirthRate time - -newtype InhomBDPop = - InhomBDPop People - deriving (Show) - -instance Population InhomBDPop where - susceptiblePeople _ = Nothing - infectiousPeople (InhomBDPop people) = Just people - removedPeople _ = Nothing - isInfected (InhomBDPop people) = not $ nullPeople people - --- | Return a BD-process parameters object -inhomBirthDeathRates :: [(Time, Rate)] -- ^ birth rate - -> Rate -- ^ death rate - -> Maybe InhomBDRates -inhomBirthDeathRates tBrPairs deathRate - | all (\x -> 0 < snd x) tBrPairs && deathRate >= 0 = liftM (\tbr -> InhomBDRates tbr deathRate) $ asTimed tBrPairs - | otherwise = Nothing - --- | Configuration of a inhomogeneous birth-death simulation. -configuration :: Time -- ^ Duration of the simulation - -> ([(Time,Rate)], Rate) -- ^ Birth and Death rates - -> Maybe (SimulationConfiguration InhomBDRates InhomBDPop) -configuration maxTime (tBrPairs, deathRate) = - let (seedPerson, newId) = newPerson initialIdentifier - bdPop = InhomBDPop (People $ V.singleton seedPerson) - in do maybeIBDRates <- inhomBirthDeathRates tBrPairs deathRate - if maxTime > 0 - then Just (SimulationConfiguration maybeIBDRates bdPop newId maxTime) - else Nothing - --- | A random event and the state afterwards -randomEvent :: - InhomBDRates -- ^ model parameters - -> Time -- ^ the current time - -> InhomBDPop -- ^ the population - -> Integer -- ^ current identifier - -> GenIO -- ^ PRNG - -> IO (Time, EpidemicEvent, InhomBDPop, Integer) -randomEvent inhomRates@(InhomBDRates brts@(Timed brts') dr) currTime (InhomBDPop (people@(People peopleVec))) currId gen = - let popSize = fromIntegral $ numPeople people :: Double - stepTimes = map fst brts' - stepFunction = fromJust $ asTimed [(t-currTime,popSize * fromJust (eventRate inhomRates t)) | t <- stepTimes] - in do delay <- inhomExponential stepFunction gen - isBirth <- bernoulli (fromJust (birthProb inhomRates (currTime + delay))) gen - (selectedPerson, unselectedPeople) <- randomPerson peopleVec gen - return $ - if isBirth - then let newTime = currTime + delay - (birthedPerson, newId) = newPerson currId - event = - Infection newTime selectedPerson birthedPerson - in ( newTime - , event - , InhomBDPop (addPerson birthedPerson people) - , newId) - else let newTime = currTime + delay - event = Removal newTime selectedPerson - in ( newTime - , event - , InhomBDPop (People unselectedPeople) - , currId) - --- | The state of the simulation at the time of the last event prior to the --- stopping time. -allEvents :: - InhomBDRates -- ^ model parameters - -> Time -- ^ stopping time - -> (Time, [EpidemicEvent], InhomBDPop, Integer) -- ^ simulation state - -> GenIO -- ^ PRNG - -> IO (Time, [EpidemicEvent], InhomBDPop, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - diff --git a/src/Epidemic/InhomogeneousBDS.hs b/src/Epidemic/InhomogeneousBDS.hs deleted file mode 100644 --- a/src/Epidemic/InhomogeneousBDS.hs +++ /dev/null @@ -1,135 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -module Epidemic.InhomogeneousBDS - ( configuration - , allEvents - , observedEvents - , inhomBDSRates - ) where - -import Epidemic.Types.Population -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Control.Monad (liftM) -import Data.Maybe (fromJust) -import qualified Data.Vector as V -import System.Random.MWC -import System.Random.MWC.Distributions (categorical, exponential) -import Epidemic -import Epidemic.Utility - -data InhomBDSRates = - InhomBDSRates (Timed Rate) Rate Rate - -instance ModelParameters InhomBDSRates where - rNaught (InhomBDSRates timedBirthRate deathRate sampleRate) time = - let birthRate = cadlagValue timedBirthRate time - in liftM (/ (deathRate + sampleRate)) birthRate - eventRate (InhomBDSRates timedBirthRate deathRate sampleRate) time = - let birthRate = cadlagValue timedBirthRate time - in liftM (+ (deathRate + sampleRate)) birthRate - birthProb (InhomBDSRates timedBirthRate deathRate sampleRate) time = - liftM (\br -> br / (br + deathRate + sampleRate)) $ cadlagValue timedBirthRate time - -newtype InhomBDSPop = - InhomBDSPop People - deriving (Show) - -instance Population InhomBDSPop where - susceptiblePeople _ = Nothing - infectiousPeople (InhomBDSPop people) = Just people - removedPeople _ = Nothing - isInfected (InhomBDSPop people) = not $ nullPeople people - --- | Return a BDS-process parameters object --- --- Note that this requires that the rates are all positive, if they are not it --- will return @Nothing@. -inhomBDSRates :: [(Time, Rate)] -- ^ birth rate - -> Rate -- ^ death rate - -> Rate -- ^ sample rate - -> Maybe InhomBDSRates -inhomBDSRates tBrPairs deathRate sampleRate - | all (\x -> 0 < snd x) tBrPairs && deathRate >= 0 && sampleRate >= 0 = - (\tbr -> InhomBDSRates tbr deathRate sampleRate) <$> asTimed tBrPairs - | otherwise = Nothing - --- | Configuration of a inhomogeneous birth-death-sampling simulation. --- --- Note that this requires that the timed rates are all positive, if they are --- not it will return @Nothing@ which can lead to cryptic bugs. -configuration :: Time -- ^ Duration of the simulation - -> ([(Time,Rate)], Rate, Rate) -- ^ Birth, Death and Sampling rates - -> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop) -configuration maxTime (tBrPairs, deathRate, sampleRate) = - let (seedPerson, newId) = newPerson initialIdentifier - bdsPop = InhomBDSPop (People $ V.singleton seedPerson) - in do maybeIBDSRates <- inhomBDSRates tBrPairs deathRate sampleRate - if maxTime > 0 - then Just - (SimulationConfiguration maybeIBDSRates bdsPop newId maxTime) - else Nothing - --- | A random event and the state afterwards -randomEvent :: - InhomBDSRates -- ^ model parameters - -> Time -- ^ the current time - -> InhomBDSPop -- ^ the population - -> Integer -- ^ current identifier - -> GenIO -- ^ PRNG - -> IO (Time, EpidemicEvent, InhomBDSPop, Integer) -randomEvent inhomRates@(InhomBDSRates brts@(Timed brts') dr sr) currTime (InhomBDSPop (people@(People peopleVec))) currId gen = - let popSize = fromIntegral $ numPeople people :: Double - stepTimes = map fst brts' - stepFunction = fromJust $ asTimed [(t-currTime,popSize * fromJust (eventRate inhomRates t)) | t <- stepTimes] - eventWeights t = V.fromList [fromJust (cadlagValue brts t), dr, sr] - in do delay <- inhomExponential stepFunction gen - eventIx <- categorical (eventWeights (currTime + delay)) gen - (selectedPerson, unselectedPeople) <- randomPerson peopleVec gen - return $ case eventIx of - 0 -> let newTime = currTime + delay - (birthedPerson, newId) = newPerson currId - event = Infection newTime selectedPerson birthedPerson - in ( newTime - , event - , InhomBDSPop (addPerson birthedPerson people) - , newId) - 1 -> let newTime = currTime + delay - event = Removal newTime selectedPerson - in (newTime, event, InhomBDSPop (People unselectedPeople), currId) - 2 -> let newTime = currTime + delay - event = Sampling newTime selectedPerson - in (newTime, event, InhomBDSPop (People unselectedPeople), currId) - _ -> error "no birth-death-sampling event selected." - --- | The state of the simulation at the time of the last event prior to the --- stopping time. -allEvents :: - InhomBDSRates -- ^ model parameters - -> Time -- ^ stopping time - -> (Time, [EpidemicEvent], InhomBDSPop, Integer) -- ^ simulation state - -> GenIO -- ^ PRNG - -> IO (Time, [EpidemicEvent], InhomBDSPop, Integer) -allEvents rates maxTime currState@(currTime, currEvents, currPop, currId) gen = - if isInfected currPop - then do - (newTime, event, newPop, newId) <- - randomEvent rates currTime currPop currId gen - if newTime < maxTime - then allEvents - rates - maxTime - (newTime, event : currEvents, newPop, newId) - gen - else return currState - else return currState - - --- | Just the observable events from a list of all the events in a simulation. -observedEvents :: [EpidemicEvent] -- ^ All of the simulation events - -> [EpidemicEvent] -observedEvents [] = [] -observedEvents events = sort $ sampleTreeEvents'' - where - sampleTreeEvents'' = - sampleTreeEvents . sampleTree $ transmissionTree events (Person 1) diff --git a/src/Epidemic/Model/BDSCOD.hs b/src/Epidemic/Model/BDSCOD.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Model/BDSCOD.hs @@ -0,0 +1,170 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE MultiParamTypeClasses #-} + +module Epidemic.Model.BDSCOD + ( configuration + , randomEvent + , BDSCODParameters(..) + , 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.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Simulation + ( SimulationConfiguration(..) + , SimulationRandEvent(..) + ) +import Epidemic.Types.Time + ( AbsoluteTime(..) + , TimeDelta(..) + , Timed(..) + , cadlagValue + , diracDeltaValue + , nextTime + , timeAfterDelta + , allTimes + , asTimed + ) +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 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) + +instance Population BDSCODPopulation where + susceptiblePeople _ = Nothing + infectiousPeople (BDSCODPopulation people) = Just people + removedPeople _ = Nothing + isInfected (BDSCODPopulation (People people)) = not $ V.null people + +-- | Configuration of a birth-death-sampling-occurrence-disaster simulation +configuration :: + TimeDelta -- ^ Duration of the simulation + -> Bool -- ^ condition upon at least two sequenced samples. + -> ( 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 + catastropheSpec' <- asTimed catastropheSpec + disasterSpec' <- asTimed disasterSpec + let bdscodParams = + BDSCODParameters + birthRate + deathRate + samplingRate + catastropheSpec' + occurrenceRate + disasterSpec' + (seedPerson, newId) = newPerson initialIdentifier + bdscodPop = BDSCODPopulation (People $ V.singleton seedPerson) + in return $ + SimulationConfiguration + bdscodParams + bdscodPop + newId + (AbsoluteTime 0) + maxTime + Nothing + atLeastCherry + +-- | The way in which random events are generated in this model. +randomEvent :: SimulationRandEvent BDSCODParameters BDSCODPopulation +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 + 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) + + +-- | 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 (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 + sampledPeople = filterZip snd currPeople rhoBernoullis + unsampledPeople = filterZip (not . snd) currPeople rhoBernoullis + in return + ( PopulationSample catastTime (People sampledPeople) True + , BDSCODPopulation (People unsampledPeople)) + +-- | 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 (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 + sampledPeople = filterZip snd currPeople nuBernoullis + unsampledPeople = filterZip (not . snd) currPeople nuBernoullis + in return + ( PopulationSample disastTime (People sampledPeople) False + , BDSCODPopulation (People unsampledPeople)) diff --git a/src/Epidemic/Model/InhomogeneousBDS.hs b/src/Epidemic/Model/InhomogeneousBDS.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Model/InhomogeneousBDS.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE MultiParamTypeClasses #-} + +module Epidemic.Model.InhomogeneousBDS + ( configuration + , randomEvent + , inhomBDSRates + , InhomBDSRates(..) + , InhomBDSPop(..) + ) where + +import Epidemic.Types.Time + ( AbsoluteTime(..) + , Timed(..) + , TimeDelta(..) + , allTimes + , asTimed + , diracDeltaValue + , nextTime + , cadlagValue + , timeAfterDelta + ) +import Control.Monad (liftM) +import Data.Maybe (fromJust, isJust, isNothing) +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(..) + ) +import Epidemic.Utility +import System.Random.MWC +import System.Random.MWC.Distributions (categorical, exponential) + +data InhomBDSRates = + InhomBDSRates (Timed Rate) Rate Rate + +data InhomBDSPop = + InhomBDSPop People + deriving (Show) + +instance ModelParameters InhomBDSRates InhomBDSPop where + rNaught _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = + let birthRate = cadlagValue timedBirthRate time + in liftM (/ (deathRate + sampleRate)) birthRate + eventRate _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = + let birthRate = cadlagValue timedBirthRate time + in liftM (+ (deathRate + sampleRate)) birthRate + birthProb _ (InhomBDSRates timedBirthRate deathRate sampleRate) time = + liftM (\br -> br / (br + deathRate + sampleRate)) $ + cadlagValue timedBirthRate time + +instance Population InhomBDSPop where + susceptiblePeople _ = Nothing + infectiousPeople (InhomBDSPop people) = Just people + removedPeople _ = Nothing + isInfected (InhomBDSPop people) = not $ nullPeople people + +-- | Return a BDS-process parameters object +-- +-- Note that this requires that the rates are all positive, if they are not it +-- will return @Nothing@. +inhomBDSRates :: + Timed Rate -- ^ birth rate + -> Rate -- ^ death rate + -> Rate -- ^ sample rate + -> Maybe InhomBDSRates +inhomBDSRates timedBirthRate@(Timed tBrPairs) deathRate sampleRate + | all (\x -> 0 < snd x) tBrPairs && deathRate >= 0 && sampleRate >= 0 = + Just $ InhomBDSRates timedBirthRate deathRate sampleRate + | otherwise = Nothing + +-- | Configuration of a inhomogeneous birth-death-sampling simulation. +-- +-- Note that this requires that the timed rates are all positive, if they are +-- not it will return @Nothing@ which can lead to cryptic bugs. +configuration :: + TimeDelta -- ^ Duration of the simulation after starting at time 0. + -> Bool -- ^ condition upon at least two sequenced samples. + -> ([(AbsoluteTime, Rate)], Rate, Rate) -- ^ Birth, Death and Sampling rates + -> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop) +configuration maxTime atLeastCherry (tBrPairs, deathRate, sampleRate) = + let (seedPerson, newId) = newPerson initialIdentifier + bdsPop = InhomBDSPop (People $ V.singleton seedPerson) + in do timedBirthRate <- asTimed tBrPairs + maybeIBDSRates <- inhomBDSRates timedBirthRate deathRate sampleRate + if maxTime > TimeDelta 0 + then Just + (SimulationConfiguration + maybeIBDSRates + bdsPop + newId + (AbsoluteTime 0) + maxTime + Nothing + atLeastCherry) + else Nothing + +randomEvent :: SimulationRandEvent InhomBDSRates InhomBDSPop +randomEvent = SimulationRandEvent randomEvent' + +-- | A random event and the state afterwards +randomEvent' :: + InhomBDSRates -- ^ model parameters + -> AbsoluteTime -- ^ the current time + -> InhomBDSPop -- ^ the population + -> Identifier -- ^ current identifier + -> GenIO -- ^ PRNG + -> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier) +randomEvent' inhomRates@(InhomBDSRates brts dr sr) currTime pop@(InhomBDSPop (people@(People peopleVec))) currId gen = + let popSize = fromIntegral $ numPeople people :: Double + eventWeights t = V.fromList [fromJust (cadlagValue brts t), dr, sr] + -- we need a new step function to account for the population size. + (Just stepFunction) = + asTimed + [ (t, popSize * fromJust (eventRate pop inhomRates t)) + | t <- allTimes brts + ] + in do (Just newEventTime) <- inhomExponential stepFunction currTime gen + eventIx <- categorical (eventWeights newEventTime) gen + (selectedPerson, unselectedPeople) <- randomPerson people gen + return $ + case eventIx of + 0 -> + ( newEventTime + , Infection newEventTime selectedPerson birthedPerson + , InhomBDSPop (addPerson birthedPerson people) + , newId) + where (birthedPerson, newId) = newPerson currId + 1 -> + ( newEventTime + , Removal newEventTime selectedPerson + , InhomBDSPop unselectedPeople + , currId) + 2 -> + ( newEventTime + , IndividualSample newEventTime selectedPerson True + , InhomBDSPop unselectedPeople + , currId) + _ -> error "no birth-death-sampling event selected." diff --git a/src/Epidemic/Model/LogisticBDSD.hs b/src/Epidemic/Model/LogisticBDSD.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Model/LogisticBDSD.hs @@ -0,0 +1,193 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RecordWildCards #-} + +module Epidemic.Model.LogisticBDSD + ( configuration + , randomEvent + , LogisticBDSDParameters(..) + , LogisticBDSDPopulation(..) + ) where + +import Data.Maybe (fromJust) +import qualified Data.Vector as V +import qualified Data.Vector.Generic as G +import Epidemic (firstScheduled, noScheduledEvent) +import Epidemic.Types.Events (EpidemicEvent(..)) +import Epidemic.Types.Time + ( AbsoluteTime(..) + , Timed(..) + , TimeDelta(..) + , asTimed + , allTimes + , diracDeltaValue + , nextTime + , cadlagValue + , timeAfterDelta + ) +import Epidemic.Types.Parameter + ( ModelParameters(..) + , Probability + , Rate + ) +import Epidemic.Types.Population + ( Identifier(..) + , People(..) + , Population(..) + , addPerson + , nullPeople + , numPeople + ) +import Epidemic.Types.Simulation + ( SimulationConfiguration(..) + , SimulationRandEvent(..) + , SimulationState(..) + ) +import Epidemic.Utility + ( initialIdentifier + , maybeToRight + , newPerson + , randomPerson + ) +import System.Random.MWC (GenIO) +import System.Random.MWC.Distributions (bernoulli, categorical, exponential) + +-- | The parameters of the logistic-BDSD process. This process allows for +-- infections, removals, sampling and disasters. +data LogisticBDSDParameters = + LogisticBDSDParameters + { paramsBirthRate :: Rate + , paramsCapacity :: Int + , paramsDeathRate :: Rate + , paramsSamplingRate :: Rate + , paramsDisasters :: Timed Probability + } + deriving (Show) + +newtype LogisticBDSDPopulation = + LogisticBDSDPopulation People + deriving (Show) + +-- | The per lineage birth rate accounting for the population size. +logisticBirthRate :: LogisticBDSDParameters -> LogisticBDSDPopulation -> Rate +logisticBirthRate LogisticBDSDParameters {..} (LogisticBDSDPopulation pop) = + let propCapacity = fromIntegral (numPeople pop) / fromIntegral paramsCapacity + in paramsBirthRate * (1.0 - propCapacity) + +instance ModelParameters LogisticBDSDParameters LogisticBDSDPopulation where + rNaught _ _ _ = Nothing + eventRate (LogisticBDSDPopulation pop) LogisticBDSDParameters {..} _ = + let propCapcity = fromIntegral (numPeople pop) / fromIntegral paramsCapacity + br = paramsBirthRate * (1.0 - propCapcity) + in Just $ br + paramsDeathRate + paramsSamplingRate + birthProb lpop lparam@LogisticBDSDParameters {..} absTime = do + er <- eventRate lpop lparam absTime + Just $ br / er + where + br = logisticBirthRate lparam lpop + +instance Population LogisticBDSDPopulation where + susceptiblePeople _ = Nothing + infectiousPeople (LogisticBDSDPopulation people) = Just people + removedPeople _ = Nothing + isInfected (LogisticBDSDPopulation people) = not $ nullPeople people + +-- | Create an simulation configuration or return an error message if this is +-- not possible. +configuration :: + TimeDelta + -> Bool -- ^ condition upon at least two sequenced samples. + -> (Rate, Int, Rate, Rate, [(AbsoluteTime, Probability)]) + -> Either String (SimulationConfiguration LogisticBDSDParameters LogisticBDSDPopulation) +configuration simDuration atLeastCherry (birthRate, capacity, deathRate, samplingRate, disasterSpec) + | minimum [birthRate, deathRate, samplingRate] < 0 = + Left "negative rate provided" + | capacity < 1 = Left "insufficient population capacity" + | otherwise = do + disasterTP <- + maybeToRight + "could not construct timed probability" + (asTimed disasterSpec) + let logBDSDParams = + LogisticBDSDParameters + birthRate + capacity + deathRate + samplingRate + disasterTP + (seedPerson, newId) = newPerson initialIdentifier + logBDSDPop = LogisticBDSDPopulation (People $ V.singleton seedPerson) + in return $ + SimulationConfiguration + logBDSDParams + logBDSDPop + newId + (AbsoluteTime 0) + simDuration + Nothing + atLeastCherry + +-- | Defines how a single random event is simulated in this model. +randomEvent :: SimulationRandEvent LogisticBDSDParameters LogisticBDSDPopulation +randomEvent = SimulationRandEvent randEvent' + +randEvent' :: + LogisticBDSDParameters + -> AbsoluteTime + -> LogisticBDSDPopulation + -> Identifier + -> GenIO + -> IO (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier) +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] + in do delay <- exponential (netEventRate * popSizeDouble) gen + let newEventTime = timeAfterDelta currTime (TimeDelta delay) + if noScheduledEvent currTime newEventTime paramsDisasters + then do + eventIx <- categorical eventWeights gen + (randPerson, otherPeople) <- randomPerson currPpl gen + return $ + case eventIx of + 0 -> + let (infectedPerson, newId) = newPerson currId + infEvent = + Infection newEventTime randPerson infectedPerson + newPop = + LogisticBDSDPopulation + (addPerson infectedPerson currPpl) + in (newEventTime, infEvent, newPop, newId) + 1 -> + ( newEventTime + , Removal newEventTime randPerson + , LogisticBDSDPopulation otherPeople + , currId) + 2 -> + ( newEventTime + , IndividualSample newEventTime randPerson True + , LogisticBDSDPopulation otherPeople + , currId) + _ -> error "do not recognise the type of event index." + else let (Just dsstr@(dsstrTime, _)) = + firstScheduled currTime paramsDisasters + in do (schdEvent, postEventPpl) <- + randomDisasterEvent dsstr currPop gen + return (dsstrTime, schdEvent, postEventPpl, currId) + +-- | 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 + -> LogisticBDSDPopulation -- ^ The state of the population prior to the disaster + -> GenIO + -> IO (EpidemicEvent, LogisticBDSDPopulation) +randomDisasterEvent (dsstrTime, dsstrProb) (LogisticBDSDPopulation (People currPpl)) gen = do + randBernoullis <- G.replicateM (V.length currPpl) (bernoulli dsstrProb gen) + let filterZip predicate a b = fst . V.unzip . V.filter predicate $ V.zip a b + sampledPeople = filterZip snd currPpl randBernoullis + unsampledPeople = filterZip (not . snd) currPpl randBernoullis + in return + ( PopulationSample dsstrTime (People sampledPeople) False + , LogisticBDSDPopulation (People unsampledPeople)) diff --git a/src/Epidemic/Types/Events.hs b/src/Epidemic/Types/Events.hs --- a/src/Epidemic/Types/Events.hs +++ b/src/Epidemic/Types/Events.hs @@ -1,347 +1,169 @@ +{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE DeriveGeneric #-} module Epidemic.Types.Events - ( EpidemicEvent(Infection, Removal, Sampling, Catastrophe, - Occurrence, Disaster) + ( EpidemicEvent(Infection, Removal, IndividualSample, + PopulationSample, StoppingTime, Extinction) + , popSampPeople + , popSampSeq + , popSampTime + , indSampPerson + , indSampSeq + , indSampTime , EpidemicTree(Branch, Leaf, Shoot) , maybeEpidemicTree + , isExtinctionOrStopping , eventTime - , ReconstructedTree(RBranch, RLeaf) - , maybeReconstructedTree - , PointProcessEvents(PointProcessEvents) - , pointProcessEvents , derivedFrom - , Newick - , asNewickString ) where import qualified Data.Aeson as Json import qualified Data.ByteString as B import qualified Data.ByteString.Builder as BBuilder -import qualified Data.Csv as Csv 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 -- | Events that can occur in an epidemic with their absolute time. data EpidemicEvent - = Infection Time Person Person -- ^ infection time, infector, infectee - | Removal Time Person -- ^ removal without observation - | Sampling Time Person -- ^ removal and inclusion in phylogeny - | Catastrophe Time People -- ^ scheduled sampling of lineages - | Occurrence Time Person -- ^ removal and observed by not in phylogeny - | Disaster Time People -- ^ scheduled occurrence of lineages + = Infection AbsoluteTime Person Person -- ^ absolute time; infector; infectee + | Removal AbsoluteTime Person + | IndividualSample + { indSampTime :: AbsoluteTime + , indSampPerson :: Person + , indSampSeq :: Bool + } + | PopulationSample + { popSampTime :: AbsoluteTime + , popSampPeople :: People + , popSampSeq :: Bool + } + | Extinction -- ^ epidemic went extinct time time can be recovered from the preceeding removal + | StoppingTime -- ^ the simulation reached the stopping time deriving (Show, Generic, Eq) instance Json.FromJSON EpidemicEvent instance Json.ToJSON EpidemicEvent -instance Csv.ToRecord EpidemicEvent where - toRecord e = - case e of - (Infection time person1 person2) -> - Csv.record - [ "infection" - , Csv.toField time - , Csv.toField person1 - , Csv.toField person2 - ] - (Removal time person) -> - Csv.record ["removal", Csv.toField time, Csv.toField person, "NA"] - (Sampling time person) -> - Csv.record ["sampling", Csv.toField time, Csv.toField person, "NA"] - (Catastrophe time people) -> - Csv.record ["catastrophe", Csv.toField time, Csv.toField people, "NA"] - (Occurrence time person) -> - Csv.record ["occurrence", Csv.toField time, Csv.toField person, "NA"] - (Disaster time people) -> - Csv.record ["disaster", Csv.toField time, Csv.toField people, "NA"] - -et :: B.ByteString -> Csv.Record -> Bool -et bs r = (== bs) . head $ V.toList r - -instance Csv.FromRecord EpidemicEvent where - parseRecord r - | et "infection" r = - Infection <$> (r Csv..! 1) <*> (Person <$> (r Csv..! 2)) <*> - (Person <$> (r Csv..! 3)) - | et "removal" r = Removal <$> (r Csv..! 1) <*> (Person <$> (r Csv..! 2)) - | et "sampling" r = Sampling <$> (r Csv..! 1) <*> (Person <$> (r Csv..! 2)) - | et "catastrophe" r = Catastrophe <$> (r Csv..! 1) <*> (r Csv..! 2) - | et "occurrence" r = - Occurrence <$> (r Csv..! 1) <*> (Person <$> (r Csv..! 2)) - | et "disaster" r = Disaster <$> (r Csv..! 1) <*> (r Csv..! 2) - | otherwise = undefined +-- | 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 --- | Epidemic Events are ordered based on which occurred first. +-- | 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 -> Time +eventTime :: EpidemicEvent -> AbsoluteTime eventTime e = case e of Infection time _ _ -> time Removal time _ -> time - Sampling time _ -> time - Catastrophe time _ -> time - Occurrence time _ -> time - Disaster time _ -> time + IndividualSample {..} -> indSampTime + PopulationSample {..} -> popSampTime -- | The events that occurred as a result of the existance of the given person. -derivedFrom :: Person - -> [EpidemicEvent] -- ^ ordered epidemic events - -> [EpidemicEvent] +derivedFrom :: + Person + -> [EpidemicEvent] -- ^ ordered epidemic events + -> [EpidemicEvent] derivedFrom person = derivedFromPeople (asPeople [person]) -- | The events that occurred as a result of the existance of a group of people -derivedFromPeople :: People - -> [EpidemicEvent] -- ^ ordered epidemic events - -> [EpidemicEvent] +derivedFromPeople :: + People + -> [EpidemicEvent] -- ^ ordered epidemic events + -> [EpidemicEvent] derivedFromPeople _ [] = [] -derivedFromPeople people (e:es) = case e of - Infection _ p1 p2 -> if includesPerson people p1 || includesPerson people p2 - then let people' = addPerson p2 (addPerson p1 people) - in e : derivedFromPeople people' es - else derivedFromPeople people es - Removal _ p -> let derivedEvents = derivedFromPeople people es - in if includesPerson people p - then e:derivedEvents - else derivedEvents - Sampling _ p -> let derivedEvents = derivedFromPeople people es - in if includesPerson people p - then e:derivedEvents - else derivedEvents - Catastrophe _ ps -> let derivedEvents = derivedFromPeople people es - in if haveCommonPeople people ps - then e:derivedEvents - else derivedEvents - Occurrence _ p -> let derivedEvents = derivedFromPeople people es - in if includesPerson people p - then e:derivedEvents - else derivedEvents - Disaster _ ps -> let derivedEvents = derivedFromPeople people es - in if haveCommonPeople people ps - then e:derivedEvents - else derivedEvents +derivedFromPeople people (e:es) = + case e of + Infection _ p1 p2 -> + if includesPerson people p1 || includesPerson people p2 + then let people' = addPerson p2 (addPerson p1 people) + in e : derivedFromPeople people' es + else derivedFromPeople people es + Removal _ p -> + let derivedEvents = derivedFromPeople people es + in if includesPerson people p + then e : derivedEvents + else derivedEvents + IndividualSample {..} -> + let derivedEvents = derivedFromPeople people es + in if includesPerson people indSampPerson + then e : derivedEvents + else derivedEvents + PopulationSample {..} -> + let derivedEvents = derivedFromPeople people es + in if haveCommonPeople people popSampPeople + then e : derivedEvents + else derivedEvents + Extinction -> derivedFromPeople people es + StoppingTime -> derivedFromPeople people es -{-| A representation of the whole transmission tree in a realisation of an -epidemic including the unobserved leaves. Lineages that are still extant are -modelled as shoots and contain a `Person` as their data rather than an event. --} +-- | The whole transmission tree including the unobserved leaves. Lineages that +-- are still extant are modelled as /shoots/ and contain a 'Person' as their +-- data rather than an event. data EpidemicTree - = Branch EpidemicEvent EpidemicTree EpidemicTree -- ^ Internal node representing infection event - | Leaf EpidemicEvent -- ^ External node representing removal event - | Shoot Person -- ^ External node representing extant lineages - deriving (Show,Eq) - + = Branch EpidemicEvent EpidemicTree EpidemicTree + | Leaf EpidemicEvent + | Shoot Person + deriving (Show, Eq) --- | A tree representation of the epidemic events. -maybeEpidemicTree :: [EpidemicEvent] -- ^ ordered epidemic events - -> Maybe EpidemicTree -maybeEpidemicTree [] = Nothing -maybeEpidemicTree [e] = case e of - Catastrophe _ people -> if nullPeople people - then Nothing - else Just (Leaf e) - Disaster _ people -> if nullPeople people - then Nothing - else Just (Leaf e) - Infection _ p1 p2 -> Just (Branch e (Shoot p1) (Shoot p2)) - _ -> Just (Leaf e) -maybeEpidemicTree (e:es:ess) = +-- | If possible return an 'EpidemicTree' describing the /sorted/ list of +-- 'EpidemicEvent'. +maybeEpidemicTree :: + [EpidemicEvent] -- ^ ordered epidemic events + -> Either String EpidemicTree +maybeEpidemicTree [] = + Left "There are no EpidemicEvent values to construct a tree with." +maybeEpidemicTree [e] = case e of + Infection _ p1 p2 -> Right (Branch e (Shoot p1) (Shoot p2)) + Removal {} -> Right (Leaf e) + IndividualSample {} -> Right (Leaf e) + PopulationSample {..} -> + 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" +maybeEpidemicTree (e:es) = + case e of Infection _ p1 p2 -> - let infectorEvents = derivedFrom p1 (es : ess) - infecteeEvents = derivedFrom p2 (es : ess) + let infectorEvents = derivedFrom p1 es + infecteeEvents = derivedFrom p2 es in do leftTree <- if null infectorEvents - then Just (Shoot p1) - else maybeEpidemicTree infectorEvents + then Right (Shoot p1) + else maybeEpidemicTree infectorEvents rightTree <- if null infecteeEvents - then Just (Shoot p2) - else maybeEpidemicTree infecteeEvents + then Right (Shoot p2) + else maybeEpidemicTree infecteeEvents return $ Branch e leftTree rightTree - Catastrophe _ people -> if nullPeople people - then maybeEpidemicTree (es:ess) - else Just (Leaf e) - Disaster _ people -> if nullPeople people - then maybeEpidemicTree (es:ess) - else Just (Leaf e) - _ -> Just (Leaf e) - -{-| A representation of the reconstructed tree which is the phylogeny connecting - all the `Sampling` and `Catastrophe` events. --} -data ReconstructedTree - = RBranch EpidemicEvent ReconstructedTree ReconstructedTree - | RLeaf EpidemicEvent - deriving (Show, Eq) - - --- | A tree representation of the reconstructed phylogeny. -maybeReconstructedTree :: EpidemicTree -> Maybe ReconstructedTree - -maybeReconstructedTree Shoot{} = Nothing - -maybeReconstructedTree (Leaf e) = case e of - Sampling{} -> Just $ RLeaf e - Catastrophe{} -> Just $ RLeaf e - _ -> Nothing - -maybeReconstructedTree (Branch e@Infection{} lt rt) - | hasSequencedLeaf lt && hasSequencedLeaf rt = - do - rlt <- maybeReconstructedTree lt - rrt <- maybeReconstructedTree rt - Just $ RBranch e rlt rrt - | hasSequencedLeaf lt = maybeReconstructedTree lt - | hasSequencedLeaf rt = maybeReconstructedTree rt - | otherwise = Nothing -maybeReconstructedTree Branch{} = Nothing - --- | Predicate for whether an `EpidemicTree` has a leaf which corresponds to a --- node in the `ReconstructedTree`. -hasSequencedLeaf :: EpidemicTree -> Bool - -hasSequencedLeaf Shoot {} = False - -hasSequencedLeaf (Leaf e) = - case e of - Sampling {} -> True - Catastrophe {} -> True - _ -> False - -hasSequencedLeaf (Branch _ lt rt) = hasSequencedLeaf lt || hasSequencedLeaf rt - - -{-| A representation of the events that can be observed in an epidemic but which - are not included in the reconstructed tree, i.e. the `Occurrence` and - `Disaster` events. --} -newtype PointProcessEvents = PointProcessEvents [EpidemicEvent] - --- | Extract the events from an epidemic tree which are observed but not part of --- the reconstructed tree. -pointProcessEvents :: EpidemicTree -> PointProcessEvents - -pointProcessEvents Shoot {} = PointProcessEvents [] - -pointProcessEvents (Leaf e) = case e of - Occurrence {} -> PointProcessEvents [e] - Disaster {} -> PointProcessEvents [e] - _ -> PointProcessEvents [] - -pointProcessEvents (Branch _ lt rt) = - let (PointProcessEvents lEs) = pointProcessEvents lt - (PointProcessEvents rEs) = pointProcessEvents rt - allEs = List.sort $ lEs ++ rEs - in PointProcessEvents allEs - -class Newick t where - -- | Return a representation of the tree in Newick format. - asNewickString :: (Time,Person) -> t -> Maybe (BBuilder.Builder, [EpidemicEvent]) - - -ampersandBuilder :: BBuilder.Builder -ampersandBuilder = BBuilder.charUtf8 '&' - -colonBuilder :: BBuilder.Builder -colonBuilder = BBuilder.charUtf8 ':' - -leftBraceBuilder :: BBuilder.Builder -leftBraceBuilder = BBuilder.charUtf8 '(' - -rightBraceBuilder :: BBuilder.Builder -rightBraceBuilder = BBuilder.charUtf8 ')' - -commaBuilder :: BBuilder.Builder -commaBuilder = BBuilder.charUtf8 ',' - -catastrophePeopleBuilder :: People -> BBuilder.Builder -catastrophePeopleBuilder (People persons) = - mconcat $ - List.intersperse ampersandBuilder [personByteString p | p <- V.toList persons] - - -instance Newick EpidemicTree where - asNewickString (_, p) (Shoot p') = - if p /= p' - then Nothing - else let identifier = personByteString p - bl = BBuilder.stringUtf8 "Infinity" - in Just (identifier <> colonBuilder <> bl, []) - - asNewickString (t, p) (Leaf e) = - let identifier = personByteString p - bl a b = BBuilder.doubleDec $ b - a - in case e of - Infection {} -> Nothing - (Removal t' p') -> - if p /= p' - then Nothing - else Just (identifier <> colonBuilder <> bl t t', [e]) - (Sampling t' p') -> - if p /= p' - then Nothing - else Just (identifier <> colonBuilder <> bl t t', [e]) - (Catastrophe t' ps) -> - if ps `includesPerson` p - then Just (identifier <> colonBuilder <> bl t t', [e]) - else Nothing - (Occurrence t' p') -> - if p /= p' - then Nothing - else Just (identifier <> colonBuilder <> bl t t', [e]) - (Disaster t' ps) -> - if ps `includesPerson` p - then Just (identifier <> colonBuilder <> bl t t', [e]) - else Nothing - - asNewickString (t, p) (Branch e lt rt) = - case e of - (Infection t' p1 p2) -> - if p /= p1 - then Nothing - else do - (leftNS, leftEs) <- asNewickString (t', p1) lt - (rightNS, rightEs) <- asNewickString (t', p2) rt - let bl = BBuilder.doubleDec $ t' - t - return - ( leftBraceBuilder <> - leftNS <> - commaBuilder <> rightNS <> rightBraceBuilder <> colonBuilder <> bl - , List.sort $ leftEs ++ rightEs) - _ -> Nothing - - - -instance Newick ReconstructedTree where - asNewickString (t, _) (RLeaf e) = - let bl a b = BBuilder.doubleDec $ b - a - in case e of - (Sampling t' p) -> Just ((personByteString p) <> colonBuilder <> bl t t', [e]) - Infection {} -> Nothing - Removal {} -> Nothing - (Catastrophe t' ps) -> Just (catastrophePeopleBuilder ps <> colonBuilder <> bl t t', [e]) - Occurrence {} -> Nothing - Disaster {} -> Nothing - - asNewickString (t, _) (RBranch e lt rt) = - case e of - (Infection t' p1 p2) -> - do - (leftNS, leftEs) <- asNewickString (t', p1) lt - (rightNS, rightEs) <- asNewickString (t', p2) rt - let bl = BBuilder.doubleDec $ t' - t - return - ( leftBraceBuilder <> - leftNS <> - commaBuilder <> rightNS <> rightBraceBuilder <> colonBuilder <> bl - , List.sort $ leftEs ++ rightEs) - _ -> Nothing + Removal {} -> Right (Leaf e) + IndividualSample {} -> Right (Leaf e) + PopulationSample {..} -> + 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" diff --git a/src/Epidemic/Types/Newick.hs b/src/Epidemic/Types/Newick.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Newick.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE RecordWildCards #-} + +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 + where + asNewickString :: + (AbsoluteTime, Person) -- ^ The person and time of the root of the tree + -> t + -> Maybe (BBuilder.Builder, [EpidemicEvent]) + +ampersandBuilder :: BBuilder.Builder +ampersandBuilder = BBuilder.charUtf8 '&' + +colonBuilder :: BBuilder.Builder +colonBuilder = BBuilder.charUtf8 ':' + +leftBraceBuilder :: BBuilder.Builder +leftBraceBuilder = BBuilder.charUtf8 '(' + +rightBraceBuilder :: BBuilder.Builder +rightBraceBuilder = BBuilder.charUtf8 ')' + +commaBuilder :: BBuilder.Builder +commaBuilder = BBuilder.charUtf8 ',' + +catastrophePeopleBuilder :: People -> BBuilder.Builder +catastrophePeopleBuilder (People persons) = + mconcat $ + List.intersperse ampersandBuilder [personByteString p | p <- V.toList persons] + +instance Newick ReconstructedTree where + asNewickString (t, _) (RLeaf (Observation e)) = + let branchLength a b = BBuilder.doubleDec td + where + (TimeDelta td) = timeDelta a b + in case e of + IndividualSample {..} -> + if indSampSeq + then Just + ( (personByteString indSampPerson) <> + colonBuilder <> branchLength t indSampTime + , [e]) + else Nothing + PopulationSample {..} -> + if popSampSeq + then Just + ( catastrophePeopleBuilder popSampPeople <> + colonBuilder <> branchLength t popSampTime + , [e]) + else Nothing + _ -> Nothing + asNewickString (t, _) (RBranch (Observation e) lt rt) = + case e of + (Infection t' p1 p2) -> do + (leftNS, leftEs) <- asNewickString (t', p1) lt + (rightNS, rightEs) <- asNewickString (t', p2) rt + let branchLength = BBuilder.doubleDec td + where + (TimeDelta td) = timeDelta t t' + return + ( leftBraceBuilder <> + leftNS <> + commaBuilder <> + rightNS <> rightBraceBuilder <> colonBuilder <> branchLength + , List.sort $ leftEs ++ rightEs) + _ -> Nothing diff --git a/src/Epidemic/Types/Observations.hs b/src/Epidemic/Types/Observations.hs --- a/src/Epidemic/Types/Observations.hs +++ b/src/Epidemic/Types/Observations.hs @@ -1,2 +1,122 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} -module Epidemic.Types.Observations where +module Epidemic.Types.Observations + ( Observation(..) + , ReconstructedTree(..) + , maybeReconstructedTree + , PointProcessEvents(..) + , pointProcessEvents + , reconstructedTreeEvents + , observedEvents + ) 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 + +-- | A wrapper for an 'EpidemicEvent' to indicate that this is an even that was +-- observed rather than just an event of the epidemic process. +newtype Observation = + Observation EpidemicEvent + deriving (Show, Ord, Eq, Generic) + +instance Json.FromJSON Observation + +instance Json.ToJSON Observation + +-- | 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. +newtype PointProcessEvents = + PointProcessEvents [Observation] + +-- | Extract the events from an epidemic tree which are observed but not part of +-- the reconstructed tree, ie the ones that are not sequenced. +pointProcessEvents :: EpidemicTree -> PointProcessEvents +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 [] + _ -> PointProcessEvents [] +pointProcessEvents (Branch _ lt rt) = + let (PointProcessEvents lEs) = pointProcessEvents lt + (PointProcessEvents rEs) = pointProcessEvents rt + allEs = List.sort $ lEs ++ rEs + in PointProcessEvents allEs + +-- | A representation of the reconstructed tree, ie the tree where the leaves +-- correspond to sequenced observations. +data ReconstructedTree + = RBranch Observation ReconstructedTree ReconstructedTree + | RLeaf Observation + deriving (Show, Eq) + +-- | The reconstructed phylogeny obtained by pruning an 'EpidemicTree' which +-- contains represents the transmission tree of the epidemic. In the case where +-- there are no sequenced samples in the epidemic then there is no tree to +-- reconstruct which is why this function is in the either monad. +maybeReconstructedTree :: EpidemicTree -> Either String ReconstructedTree +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" + _ -> Left "Bad leaf in the EpidemicTree" +maybeReconstructedTree (Branch e@Infection {} lt rt) + | hasSequencedLeaf lt && hasSequencedLeaf rt = do + rlt <- maybeReconstructedTree lt + rrt <- maybeReconstructedTree rt + Right $ RBranch (Observation e) rlt rrt + | hasSequencedLeaf lt = maybeReconstructedTree lt + | hasSequencedLeaf rt = maybeReconstructedTree rt + | otherwise = Left "Neither subtree has a sequenced leaf" +maybeReconstructedTree Branch {} = Left "EpidemicTree is a bad branch" + +-- | Predicate for whether an 'EpidemicTree' has any leaf which corresponds to a +-- sequenced observation and hence should be included in a @ReconstructedTree@. +hasSequencedLeaf :: EpidemicTree -> Bool +hasSequencedLeaf Shoot {} = False +hasSequencedLeaf (Leaf e) = + case e of + IndividualSample {..} -> indSampSeq + PopulationSample {..} -> popSampSeq + _ -> False +hasSequencedLeaf (Branch _ lt rt) = hasSequencedLeaf lt || hasSequencedLeaf rt + +-- | The events that were observed during the epidemic, ie those in the +-- reconstructed tree and any unsequenced samples. If this is not possible an +-- error message will be returned. +observedEvents :: [EpidemicEvent] -> Either String [Observation] +observedEvents epiEvents = do + epiTree <- maybeEpidemicTree epiEvents + let (PointProcessEvents unseqObss) = pointProcessEvents epiTree + reconTreeEvents <- + if hasSequencedLeaf epiTree + then (liftM reconstructedTreeEvents) $ maybeReconstructedTree epiTree + else Right [] + return $ List.sort . List.nub $ unseqObss ++ reconTreeEvents + +-- | A sorted list of all of the observations in the reconstructed tree. +reconstructedTreeEvents :: ReconstructedTree -> [Observation] +reconstructedTreeEvents rt = + case rt of + RBranch obs rtl rtr -> + List.sort $ + obs : (reconstructedTreeEvents rtl ++ reconstructedTreeEvents rtr) + RLeaf obs -> [obs] diff --git a/src/Epidemic/Types/Parameter.hs b/src/Epidemic/Types/Parameter.hs --- a/src/Epidemic/Types/Parameter.hs +++ b/src/Epidemic/Types/Parameter.hs @@ -1,84 +1,18 @@ {-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE MultiParamTypeClasses #-} module Epidemic.Types.Parameter where -import qualified Data.Aeson as Json -import qualified Data.List as List -import qualified Data.Maybe as Maybe -import GHC.Generics - -type Time = Double - --- | Type containing values at times. The times are increasing as required by --- @asTimed@. -newtype Timed a = - Timed [(Time, a)] - deriving (Generic, Eq, Show) - -instance Json.FromJSON a => Json.FromJSON (Timed a) - -instance Json.ToJSON a => Json.ToJSON (Timed a) +import Epidemic.Types.Population (Population(..)) +import Epidemic.Types.Time (AbsoluteTime(..)) -instance Semigroup (Timed a) where - (Timed x) <> (Timed y) = Timed $ List.sortOn fst (x ++ y) +-- | Class of types that can be considered parameterisations of a epidemic +-- model. +class (Population p) => ModelParameters a p where + rNaught :: p -> a -> AbsoluteTime -> Maybe Double + eventRate :: p -> a -> AbsoluteTime -> Maybe Rate + birthProb :: p -> a -> AbsoluteTime -> Maybe Probability type Rate = Double type Probability = Double - --- | Construct a timed list if possible. -asTimed :: Num a - => [(Time,a)] -- ^ list of ascending times and values - -> Maybe (Timed a) -asTimed tas = if isAscending $ map fst tas then Just (Timed $ tas ++ [(1e100,-1)]) else Nothing - --- | Predicate to check if a list of orderable objects is in ascending order. -isAscending :: Ord a => [a] -> Bool -isAscending xs = case xs of - [] -> True - [_] -> True - (x:y:xs') -> x <= y && isAscending (y:xs') - --- | Evaluate the timed object treating it as a cadlag function -cadlagValue :: Timed a -> Time -> Maybe a -cadlagValue (Timed txs) = cadlagValue' txs - - -cadlagValue' :: [(Time,a)] -> Time -> Maybe a -cadlagValue' [] _ = Nothing -cadlagValue' ((t, x):txs) q = - if q < t - then Nothing - else let nextCLV = cadlagValue' txs q - in if Maybe.isNothing nextCLV - then Just x - else nextCLV - - --- | Evaluate the timed object treating it as a direct delta function -diracDeltaValue :: Timed a -> Time -> Maybe a -diracDeltaValue (Timed txs) = diracDeltaValue' txs - -diracDeltaValue' :: [(Time,a)] -> Time -> Maybe a -diracDeltaValue' txs q = case txs of - ((t,x):txs') -> if t == q then Just x else diracDeltaValue' txs' q - [] -> Nothing - --- | Check if there exists a pair with a particular time index. -hasTime :: Timed a -> Time -> Bool -hasTime (Timed txs) = hasTime' txs - -hasTime' :: [(Time,a)] -> Time -> Bool -hasTime' txs q = case txs of - ((t,_):txs') -> t == q || hasTime' txs' q - [] -> False - --- | Return the value of the next time if possible or an exact match if it --- exists. -nextTime :: Timed a -> Time -> Maybe Time -nextTime (Timed txs) = nextTime' txs - -nextTime' :: [(Time,a)] -> Time -> Maybe Time -nextTime' txs q = case txs of - ((t,_):txs') -> if q < t then Just t else nextTime' txs' q - [] -> Nothing diff --git a/src/Epidemic/Types/Population.hs b/src/Epidemic/Types/Population.hs --- a/src/Epidemic/Types/Population.hs +++ b/src/Epidemic/Types/Population.hs @@ -4,6 +4,8 @@ module Epidemic.Types.Population ( Person(Person) , People(People) + , Population(..) + , Identifier(Identifier) , asPeople , includesPerson , haveCommonPeople @@ -18,24 +20,35 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Builder as BBuilder import Data.ByteString.Internal (c2w) -import qualified Data.Csv as Csv import qualified Data.Vector as V import GHC.Generics +-- | Class of types that can represent populations in an epidemic simulation. +class Population a where + susceptiblePeople :: a -> Maybe People + infectiousPeople :: a -> Maybe People + removedPeople :: a -> Maybe People + isInfected :: a -> Bool + +-- | A type to hold an integer which is unique to each 'Person'. +newtype Identifier = + Identifier Integer + deriving (Show, Generic, Eq) + +instance Json.FromJSON Identifier + +instance Json.ToJSON Identifier + +-- | A type to represent a single person in a group of 'People' newtype Person = - Person Integer + Person Identifier deriving (Show, Generic, Eq) instance Json.FromJSON Person instance Json.ToJSON Person -instance Csv.ToField Person where - toField (Person n) = Csv.toField n - -instance Csv.FromField Person where - parseField f = Person <$> (Csv.parseField f :: Csv.Parser Integer) - +-- | A type to represent a population. newtype People = People (V.Vector Person) deriving (Show, Eq, Generic) @@ -44,14 +57,6 @@ instance Json.ToJSON People -instance Csv.ToField People where - toField (People persons) = - B.intercalate ":" $ V.toList $ V.map Csv.toField persons - -instance Csv.FromField People where - parseField f = - (People . V.fromList) <$> (mapM Csv.parseField $ B.split (c2w ':') f) - -- | A list of persons as a people asPeople :: [Person] -> People asPeople persons = People $ V.fromList persons @@ -82,4 +87,4 @@ -- | A bytestring builder for a person personByteString :: Person -> BBuilder.Builder -personByteString (Person n) = BBuilder.integerDec n +personByteString (Person (Identifier n)) = BBuilder.integerDec n diff --git a/src/Epidemic/Types/Simulation.hs b/src/Epidemic/Types/Simulation.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Simulation.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE GADTs #-} + +module Epidemic.Types.Simulation + ( SimulationConfiguration(..) + , SimulationState(..) + , SimulationRandEvent(..) + ) where + +import Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Types.Time (AbsoluteTime(..), TimeDelta(..), timeDelta) +import System.Random.MWC + +data SimulationConfiguration r p = + SimulationConfiguration + { -- | The event rates + scRates :: r + -- | The population + , scPopulation :: p + -- | A new identifier + , scNewIdentifier :: Identifier + -- | The absolute time at which the simulation starts + , scStartTime :: AbsoluteTime + -- | The duration of the simulation until it stops + , scSimDuration :: TimeDelta + -- | The simulation terminates if this predicate is not satisfied + , scValidPopulation :: Maybe (p -> Bool) + -- | The simulation requires at least two sequenced samples + , 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 + = SimulationState (AbsoluteTime, [EpidemicEvent], b, Identifier) + | TerminatedSimulation + deriving (Eq, Show) + +data SimulationRandEvent a b where + SimulationRandEvent + :: (ModelParameters a b, Population b) + => (a + -> AbsoluteTime + -> b + -> Identifier + -> GenIO + -> IO (AbsoluteTime, EpidemicEvent, b, Identifier)) + -> SimulationRandEvent a b diff --git a/src/Epidemic/Types/Time.hs b/src/Epidemic/Types/Time.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Time.hs @@ -0,0 +1,155 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveGeneric #-} + +module Epidemic.Types.Time + ( AbsoluteTime(..) + , TimeDelta(..) + , Timed(..) + , timeDelta + , diracDeltaValue + , timeAfterDelta + , nextTime + , cadlagValue + , isAscending + , hasTime + , allTimes + , asTimed + ) where + +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 = + AbsoluteTime Double + deriving (Generic, Eq, Show, Ord) + +instance Json.FromJSON AbsoluteTime + +instance Json.ToJSON AbsoluteTime + +-- | Predicate for an infinite absolute time +isInfiniteAbsoluteTime :: AbsoluteTime -> Bool +isInfiniteAbsoluteTime (AbsoluteTime t) = isInfinite t + +-- | Duration of time between two absolute times. +newtype TimeDelta = + TimeDelta Double + deriving (Generic, Eq, Show, Ord) + +instance Json.FromJSON TimeDelta + +instance Json.ToJSON TimeDelta + +-- | The duration of time between two absolute times +-- +-- >>> timeDelta (AbsoluteTime 1) (AbsoluteTime 2.5) +-- TimeDelta 1.5 +-- +timeDelta :: + AbsoluteTime -- ^ start + -> AbsoluteTime -- ^ finish + -> TimeDelta +timeDelta (AbsoluteTime t0) (AbsoluteTime t1) = TimeDelta (t1 - t0) + +-- | The time after a given delay +-- +-- >>> timeAfterDelta (AbsoluteTime 1) (TimeDelta 2.5) +-- AbsoluteTime 3.5 +-- +timeAfterDelta :: AbsoluteTime -> TimeDelta -> AbsoluteTime +timeAfterDelta (AbsoluteTime t0) (TimeDelta d) = AbsoluteTime (t0 + d) + +-- | Type containing values at times. The times are increasing as required by +-- @asTimed@. +newtype Timed a = + Timed [(AbsoluteTime, a)] + deriving (Generic, Eq, Show) + +instance Json.FromJSON a => Json.FromJSON (Timed a) + +instance Json.ToJSON a => Json.ToJSON (Timed a) + +instance Semigroup (Timed a) where + (Timed x) <> (Timed y) = Timed $ List.sortOn fst (x ++ y) + +-- | Construct a timed list if possible. +asTimed :: + Num a + => [(AbsoluteTime, a)] -- ^ list of ascending times and values + -> Maybe (Timed a) +asTimed tas = + if isAscending $ map fst tas + then Just (Timed $ tas ++ [(AbsoluteTime (1 / 0), -1)]) + else Nothing + +-- | Predicate to check if a list of orderable objects is in ascending order. +isAscending :: Ord a => [a] -> Bool +isAscending xs = + case xs of + [] -> True + [_] -> True + (x:y:xs') -> x <= y && isAscending (y : xs') + +-- | Evaluate the timed object treating it as a cadlag function +cadlagValue :: Timed a -> AbsoluteTime -> Maybe a +cadlagValue (Timed txs) = cadlagValue' txs + +cadlagValue' :: [(AbsoluteTime, a)] -> AbsoluteTime -> Maybe a +cadlagValue' [] _ = Nothing +cadlagValue' ((t, x):txs) q = + if q < t + then Nothing + else let nextCLV = cadlagValue' txs q + in if Maybe.isNothing nextCLV + then Just x + else nextCLV + +-- | Evaluate the timed object treating it as a direct delta function +diracDeltaValue :: Timed a -> AbsoluteTime -> Maybe a +diracDeltaValue (Timed txs) = diracDeltaValue' txs + +diracDeltaValue' :: [(AbsoluteTime, a)] -> AbsoluteTime -> Maybe a +diracDeltaValue' txs q = + case txs of + ((t, x):txs') -> + if t == q + then Just x + else diracDeltaValue' txs' q + [] -> Nothing + +-- | 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 + +-- | Return the value of the next time if possible or an exact match if it +-- exists. +nextTime :: Timed a -> AbsoluteTime -> Maybe AbsoluteTime +nextTime (Timed txs) = nextTime' txs + +nextTime' :: [(AbsoluteTime, a)] -> AbsoluteTime -> Maybe AbsoluteTime +nextTime' txs q = + case txs of + ((t, _):txs') -> + if q < t + then Just t + else nextTime' txs' q + [] -> Nothing + +-- | Return a list of the (finite) absolute times that the step function changes +-- value. +-- +-- >>> let demoMaybeTimed = asTimed [(AbsoluteTime 1,2),(AbsoluteTime 1.5,1)] +-- >>> liftM allTimes demoMaybeTimed +-- Just [AbsoluteTime 1.0,AbsoluteTime 1.5] +-- +allTimes :: Timed a -> [AbsoluteTime] +allTimes (Timed txs) = [t | (t, _) <- txs, not $ isInfiniteAbsoluteTime t] diff --git a/src/Epidemic/Utility.hs b/src/Epidemic/Utility.hs --- a/src/Epidemic/Utility.hs +++ b/src/Epidemic/Utility.hs @@ -1,39 +1,41 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE OverloadedStrings #-} + module Epidemic.Utility where -import Epidemic.Types.Events -import Epidemic.Types.Parameter -import Epidemic.Types.Population +import Control.Applicative import Control.Monad (liftM) -import qualified Data.List as List -import qualified Data.Maybe as Maybe +import Control.Monad.Primitive (PrimMonad, PrimState) import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as Char8 -import GHC.Generics (Generic) +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 -import Control.Monad.Primitive (PrimMonad, PrimState) -import Control.Applicative -import Text.Trifecta +import System.Random.MWC.Distributions (exponential) -import Epidemic -data SimulationConfiguration r p = - SimulationConfiguration - { rates :: r - , population :: p - , newIdentifier :: Integer - , timeLimit :: Time - } - -initialIdentifier :: Integer -initialIdentifier = 1 +initialIdentifier :: Identifier +initialIdentifier = Identifier 1 -newPerson :: Integer -> (Person, Integer) -newPerson identifier = (Person identifier, identifier + 1) +newPerson :: Identifier -> (Person, Identifier) +newPerson idntty@(Identifier idInt) = (Person idntty, Identifier (idInt + 1)) selectElem :: V.Vector a -> Int -> (a, V.Vector a) selectElem v n @@ -42,193 +44,195 @@ let (foo, bar) = V.splitAt n v in (V.head bar, foo V.++ (V.tail bar)) -randomPerson :: V.Vector Person -> GenIO -> IO (Person, V.Vector Person) -randomPerson persons gen = do +randomPerson :: People -> GenIO -> IO (Person, People) +randomPerson people@(People persons) gen = do u <- uniform gen - return $ selectElem persons (floor (u * numPersons)) - where - numPersons = fromIntegral $ V.length persons :: Double - + let personIx = floor (u * (fromIntegral $ numPeople people :: Double)) + (person, remPeople) = selectElem persons personIx + in return (person, People remPeople) type NName = Maybe String type NLength = Maybe Double -data NBranch = NBranch NSubtree NLength deriving (Eq) +data NBranch = + NBranch NSubtree NLength + deriving (Eq) instance Show NBranch where show (NBranch st (Just l)) = show st ++ ":" ++ show l show (NBranch st Nothing) = show st -data NBranchSet = NBranchSet [NBranch] deriving (Eq) +data NBranchSet = + NBranchSet [NBranch] + deriving (Eq) instance Show NBranchSet where show (NBranchSet bs) = "(" ++ (List.intercalate "," (map show bs)) ++ ")" -data NSubtree = NLeaf NName | NInternal NBranchSet deriving (Eq) +data NSubtree + = NLeaf NName + | NInternal NBranchSet + deriving (Eq) instance Show NSubtree where show (NLeaf (Just n)) = n show (NLeaf Nothing) = "" show (NInternal bs) = show bs -data NTree = NTree [NBranch] deriving (Eq) +data NTree = + NTree [NBranch] + deriving (Eq) instance Show NTree where show (NTree bs) = show (NBranchSet bs) ++ ";" --- Name → empty | string -newickName :: (Monad f, CharParsing f) => f NName -newickName = optional (some alphaNum) >>= pure - --- Leaf → Name -newickLeaf :: (Monad f, CharParsing f) => f NSubtree -newickLeaf = do - n <- newickName - pure (NLeaf n) - --- Length → empty | ":" number -newickLength :: (TokenParsing f, Monad f, CharParsing f) => f NLength -newickLength = do - maybeLength <- optional ((symbolic ':') >> double) - pure maybeLength - --- Branch → Subtree Length -newickBranch :: (TokenParsing f, Monad f, CharParsing f) => f NBranch -newickBranch = do - st <- newickSubtree - l <- newickLength - pure (NBranch st l) - --- BranchSet → Branch | Branch "," BranchSet -newickBranchSet :: (TokenParsing f, Monad f, CharParsing f) => f NBranchSet -newickBranchSet = do - bs <- sepBy1 newickBranch comma - pure (NBranchSet bs) - --- Internal → "(" BranchSet ")" Name -newickInternal :: (TokenParsing f, Monad f, CharParsing f) => f NSubtree -newickInternal = do - bs <- parens newickBranchSet - pure (NInternal bs) - --- Subtree → Leaf | Internal -newickSubtree :: (TokenParsing f, Monad f, CharParsing f) => f NSubtree -newickSubtree = choice [newickInternal,newickLeaf] - --- Tree → Subtree ";" | Branch ";" -newickTree :: (TokenParsing f, Monad f, CharParsing f) => f NTree -newickTree = do - (NBranchSet bs) <- parens newickBranchSet - symbolic ';' - pure (NTree 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 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 + where + go n [] = n + go n (x:xs) + | p x = go (n + 1) xs + | otherwise = go n xs --- | Run a simulation described by a configuration object. -simulation :: (ModelParameters a) - => Bool -- ^ Condition upon at least two leaves in the reconstructed tree - -> SimulationConfiguration a b - -> (a -> Time -> (Time, [EpidemicEvent], b, Integer) -> GenIO -> IO (Time, [EpidemicEvent], b, Integer)) - -> IO [EpidemicEvent] -simulation True config allEvents = do - gen <- System.Random.MWC.create :: IO GenIO - simulation' config allEvents gen -simulation False SimulationConfiguration {..} allEvents = do - gen <- System.Random.MWC.create :: IO GenIO - (_, events, _, _) <- - allEvents rates timeLimit (0, [], population, newIdentifier) gen - return $ sort events +-- | Run a simulation described by a configuration object with the provided +-- PRNG. +simulationWithGenIO :: + (ModelParameters a b, Population b) + => SimulationConfiguration a b + -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) + -> GenIO + -> IO [EpidemicEvent] +simulationWithGenIO config@SimulationConfiguration {..} allEventsFunc gen = + if scRequireCherry + then do + simulation' config allEventsFunc gen + else do + SimulationState (_, events, _, _) <- + allEventsFunc + scRates + (timeAfterDelta scStartTime scSimDuration) + scValidPopulation + (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + gen + return $ sort events --- | Predicate for whether an epidemic event is either an occurrence or a --- disaaster. -isNonReconTreeObservation :: EpidemicEvent -> Bool -isNonReconTreeObservation e = case e of - Occurrence {} -> True - Disaster {} -> True - _ -> False +-- | Run a simulation described by a configuration object using the fixed PRNG +-- that is hardcoded in the @mwc-random@ package. +simulation :: + (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 + 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 - Sampling {} -> True - Catastrophe {} -> True - _ -> False - +isReconTreeLeaf e = + case e of + IndividualSample {..} -> indSampSeq + PopulationSample {..} -> popSampSeq + _ -> False -simulation' :: (ModelParameters a) => SimulationConfiguration a b - -> (a -> Time -> (Time, [EpidemicEvent], b, Integer) -> GenIO -> IO (Time, [EpidemicEvent], b, Integer)) - -> GenIO - -> IO [EpidemicEvent] -simulation' config@SimulationConfiguration {..} allEvents gen = do - (_, events, _, _) <- - allEvents rates timeLimit (0, [], population, newIdentifier) gen +-- | Simulation conditioned upon there being at least two sequenced samples. +-- NOTE This function is deprecated and will be removed in future versions. +simulation' :: + (ModelParameters a b, Population b) + => SimulationConfiguration a b + -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b)) + -> GenIO + -> IO [EpidemicEvent] +simulation' config@SimulationConfiguration {..} allEventsFunc gen = do + SimulationState (_, events, _, _) <- + allEventsFunc + scRates + (timeAfterDelta scStartTime scSimDuration) + scValidPopulation + (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + gen if count' isReconTreeLeaf events >= 2 then return $ sort events - else simulation' config allEvents gen - + else simulation' config allEventsFunc gen -- | Run a simulation described by a configuration object but using a random -- seed generated by the system rather than a seed -simulationWithSystemRandom :: (ModelParameters a) - => Bool -- ^ Condition upon at least two leaves in the reconstructed tree - -> SimulationConfiguration a b - -> (a -> Time -> (Time, [EpidemicEvent], b, Integer) -> GenIO -> IO (Time, [EpidemicEvent], b, Integer)) - -> IO [EpidemicEvent] -simulationWithSystemRandom atLeastCherry config@SimulationConfiguration {..} allEvents = do - (_, events, _, _) <- +simulationWithSystemRandom :: + (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, _, _) <- withSystemRandom $ \g -> - allEvents rates timeLimit (0, [], population, newIdentifier) g - if atLeastCherry + allEventsFunc + scRates + (timeAfterDelta scStartTime scSimDuration) + scValidPopulation + (SimulationState (AbsoluteTime 0, [], scPopulation, scNewIdentifier)) + g + if scRequireCherry then (if count' isReconTreeLeaf events >= 2 - then return $ sort events - else simulationWithSystemRandom True config allEvents) + then return $ sort events + else simulationWithSystemRandom config allEventsFunc) else return $ sort events - -- | The number of lineages at the end of a simulation. -finalSize :: [EpidemicEvent] -- ^ The events from the simulation - -> Integer +finalSize :: + [EpidemicEvent] -- ^ The events from the simulation + -> Integer finalSize = foldl (\x y -> x + eventPopDelta y) 1 --- | Generate exponentially distributed random variates with inhomogeneous rate. -inhomExponential :: PrimMonad m - => Timed Double -- ^ Step function - -> Gen (PrimState m) -- ^ Generator. - -> m Double -inhomExponential stepFunc gen = do - maybeVariate <- randInhomExp 0 stepFunc gen - if Maybe.isJust maybeVariate - then return $ Maybe.fromJust maybeVariate - else inhomExponential stepFunc gen +-- | Generate exponentially distributed random variates with inhomogeneous rate +-- starting from a particular point in time. +-- +-- Assuming the @stepFunc@ is the intensity of arrivals and @t0@ is the start +-- time this returns @t1@ the time of the next arrival. +inhomExponential :: + PrimMonad m + => Timed Double -- ^ Step function + -> AbsoluteTime -- ^ Start time + -> Gen (PrimState m) -- ^ Generator + -> m (Maybe AbsoluteTime) +inhomExponential stepFunc t0 = randInhomExp t0 stepFunc -- | Generate exponentially distributed random variates with inhomogeneous rate. -randInhomExp :: PrimMonad m - => Double -- ^ Timer - -> Timed Double -- ^ Step function - -> Gen (PrimState m) -- ^ Generator. - -> m (Maybe Double) +-- +-- __TODO__ The algorithm used here generates more variates than are needed. It +-- would be nice to use a more efficient implementation. +-- +randInhomExp :: + PrimMonad m + => AbsoluteTime -- ^ Timer + -> Timed Double -- ^ Step function + -> Gen (PrimState m) -- ^ Generator. + -> m (Maybe AbsoluteTime) randInhomExp crrT stepFunc gen = let crrR = cadlagValue stepFunc crrT nxtT = nextTime stepFunc crrT - in if (Maybe.isJust crrR && Maybe.isJust nxtT) + in if Maybe.isJust crrR && Maybe.isJust nxtT then do crrD <- exponential (Maybe.fromJust crrR) gen - if crrT + crrD < (Maybe.fromJust nxtT) - then return $ Just (crrD + crrT) - else (randInhomExp (Maybe.fromJust nxtT) stepFunc gen) + let propT = timeAfterDelta crrT (TimeDelta crrD) + if propT < Maybe.fromJust nxtT + then return $ Just propT + else randInhomExp (Maybe.fromJust nxtT) stepFunc gen else return Nothing + +-- | Helper function for converting between the Maybe monad and the Either +-- monad. +maybeToRight :: a -> Maybe b -> Either a b +maybeToRight a maybeB = + case maybeB of + (Just b) -> Right b + Nothing -> Left a diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -5,16 +5,16 @@ import qualified Data.Aeson as Json import qualified Data.ByteString as B import qualified Data.ByteString.Builder as BBuilder -import Data.Csv +import Data.Either (isRight) import Data.Maybe (fromJust, isJust, isNothing) import qualified Data.Vector as V import Epidemic -import qualified Epidemic.BDSCOD as BDSCOD -import qualified Epidemic.BirthDeath as BD -import qualified Epidemic.BirthDeathSamplingCatastropheOccurrence as BDSCO -import qualified Epidemic.BirthDeathSamplingOccurrence as BDSO -import qualified Epidemic.InhomogeneousBDS as InhomBDS +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 @@ -22,195 +22,241 @@ 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 + -- | y is within n% of x from x. withinNPercent n x y = x - d < y && y < x + d where d = n * x / 100 -p1 = Person 1 +p1 = Person (Identifier 1) -p2 = Person 2 +p2 = Person (Identifier 2) -p3 = Person 3 +p3 = Person (Identifier 3) -p4 = Person 4 +p4 = Person (Identifier 4) -p5 = Person 5 +p5 = Person (Identifier 5) -p6 = Person 6 +p6 = Person (Identifier 6) -p7 = Person 7 +p7 = Person (Identifier 7) -- | The first set of test data does not have any catastrophe events. demoFullEvents01 = - [ Infection 1 p1 p2 - , Infection 2 p1 p3 - , Sampling 3 p1 - , Infection 4 p2 p4 - , Infection 5 p2 p5 - , Sampling 6 p4 - , Infection 7 p3 p6 - , Occurrence 8 p2 - , Removal 9 p3 - , Infection 10 p5 p7 - , Occurrence 11 p6 - , Sampling 12 p5 - , Removal 13 p7 + [ Infection (AbsoluteTime 1) p1 p2 + , Infection (AbsoluteTime 2) p1 p3 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p4 + , Infection (AbsoluteTime 5) p2 p5 + , IndividualSample (AbsoluteTime 6) p4 True + , Infection (AbsoluteTime 7) p3 p6 + , IndividualSample (AbsoluteTime 8) p2 False + , Removal (AbsoluteTime 9) p3 + , Infection (AbsoluteTime 10) p5 p7 + , IndividualSample (AbsoluteTime 11) p6 False + , IndividualSample (AbsoluteTime 12) p5 True + , Removal (AbsoluteTime 13) p7 ] demoSampleEvents01 = - [ Infection 1 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p4 - , Sampling 6 p4 - , Occurrence 8 p2 - , Occurrence 11 p6 - , Sampling 12 p5 + [ 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 + , IndividualSample (AbsoluteTime 12) p5 True ] -- | The second set of test data is the same as the first but includes a -- catastrophe event. demoFullEvents02 = - [ Infection 1 p1 p2 - , Infection 2 p1 p3 - , Sampling 3 p1 - , Infection 4 p2 p4 - , Infection 5 p2 p5 - , Sampling 6 p4 - , Infection 7 p3 p6 - , Occurrence 8 p2 - , Removal 9 p3 - , Infection 10 p5 p7 - , Catastrophe 11 (asPeople [p5]) - , Occurrence 12 p6 - , Removal 13 p7 + [ Infection (AbsoluteTime 1) p1 p2 + , Infection (AbsoluteTime 2) p1 p3 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p4 + , Infection (AbsoluteTime 5) p2 p5 + , IndividualSample (AbsoluteTime 6) p4 True + , Infection (AbsoluteTime 7) p3 p6 + , IndividualSample (AbsoluteTime 8) p2 False + , Removal (AbsoluteTime 9) p3 + , Infection (AbsoluteTime 10) p5 p7 + , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True + , IndividualSample (AbsoluteTime 12) p6 False + , Removal (AbsoluteTime 13) p7 ] demoSampleEvents02 = - [ Infection 1 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p4 - , Sampling 6 p4 - , Occurrence 8 p2 - , Catastrophe 11 (asPeople [p5]) - , Occurrence 12 p6 + [ 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 + , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True + , IndividualSample (AbsoluteTime 12) p6 False ] -- | Another test set to test that catastrophes are handled correctly. demoFullEvents03 = - [ Infection 1 p1 p4 - , Infection 2 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p3 - , Infection 5 p4 p5 - , Catastrophe 6 (asPeople [p2, p3, p4]) + [ Infection (AbsoluteTime 1) p1 p4 + , Infection (AbsoluteTime 2) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p3 + , Infection (AbsoluteTime 5) p4 p5 + , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True ] demoSampleEvents03 = - [ Infection 1 p1 p4 - , Infection 2 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p3 - , Catastrophe 6 (asPeople [p2, p3, p4]) + [ Infection (AbsoluteTime 1) p1 p4 + , Infection (AbsoluteTime 2) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p3 + , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True ] -- | Another test to make sure that disasters are handled. demoFullEvents04 = - [ Infection 1 p1 p4 - , Infection 2 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p3 - , Infection 5 p4 p5 - , Catastrophe 6 (asPeople [p2, p3, p4]) - , Infection 7 p5 p6 - , Infection 8 p5 p7 - , Disaster 9 (asPeople [p5, p6]) + [ Infection (AbsoluteTime 1) p1 p4 + , Infection (AbsoluteTime 2) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p3 + , Infection (AbsoluteTime 5) p4 p5 + , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True + , Infection (AbsoluteTime 7) p5 p6 + , Infection (AbsoluteTime 8) p5 p7 + , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False ] demoSampleEvents04 = - [ Infection 1 p1 p4 - , Infection 2 p1 p2 - , Sampling 3 p1 - , Infection 4 p2 p3 - , Catastrophe 6 (asPeople [p2, p3, p4]) - , Disaster 9 (asPeople [p5, p6]) + [ Infection (AbsoluteTime 1) p1 p4 + , Infection (AbsoluteTime 2) p1 p2 + , IndividualSample (AbsoluteTime 3) p1 True + , Infection (AbsoluteTime 4) p2 p3 + , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True + , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False ] eventHandlingTests = do describe "Post-simulation processing" $ do it "Extracting observed events" $ do - (demoSampleEvents01 == BDSO.observedEvents demoFullEvents01) `shouldBe` + let demoEvents = + [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored! + , Infection (AbsoluteTime 1.0) p1 p2 + , PopulationSample (AbsoluteTime 1.5) (asPeople []) True + , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True + ] + (length demoEvents == 4) `shouldBe` True + ((length <$> observedEvents (tail demoEvents)) == (Right 2)) `shouldBe` True - (demoSampleEvents02 == BDSO.observedEvents demoFullEvents02) `shouldBe` + ((length <$> observedEvents (demoEvents)) == (Right 2)) `shouldBe` True + (observedEvents (demoEvents) == observedEvents (tail demoEvents)) `shouldBe` True - let demoEvents = [Catastrophe 0.5 (asPeople []) -- Because the first event is a null event it can be ignored! - ,Infection 1.0 p1 p2 - ,Catastrophe 1.5 (asPeople []) - ,Catastrophe 2.0 (asPeople [p1,p2])] - (length demoEvents == 4) `shouldBe` True - ((length <$> BDSCOD.observedEvents (tail demoEvents)) == (Just 2)) `shouldBe` True - ((length <$> BDSCOD.observedEvents (demoEvents)) == (Just 2)) `shouldBe` True - (BDSCOD.observedEvents (demoEvents) == BDSCOD.observedEvents (tail demoEvents)) `shouldBe` True - (maybeEpidemicTree (demoEvents) == maybeEpidemicTree (tail demoEvents)) `shouldBe` True + (maybeEpidemicTree (demoEvents) == maybeEpidemicTree (tail demoEvents)) `shouldBe` + True describe "Catastrophe definitions" $ do it "Check we can find a catastrophe" $ do - (noScheduledEvent 0 1 (Timed [])) `shouldBe` True - (noScheduledEvent 0 1 (Timed [(2, 0.5)])) `shouldBe` True - (noScheduledEvent 0 1 (Timed [(0.5, 0.5)])) `shouldBe` False - (noScheduledEvent 0 1 (Timed [(2, 0.6), (0.5, 0.5)])) `shouldBe` False + (noScheduledEvent (AbsoluteTime 0) (AbsoluteTime 1) (Timed [])) `shouldBe` + True + (noScheduledEvent + (AbsoluteTime 0) + (AbsoluteTime 1) + (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe` + True + (noScheduledEvent + (AbsoluteTime 0) + (AbsoluteTime 1) + (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe` + False + (noScheduledEvent + (AbsoluteTime 0) + (AbsoluteTime 1) + (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe` + False it "Check we can find a particular catastrophe" $ do - (firstScheduled 1 (Timed [])) `shouldBe` Nothing - (firstScheduled 1 (Timed [(2, 0.5)])) `shouldBe` Just (2, 0.5) - (firstScheduled 1 (Timed [(0.5, 0.5)])) `shouldBe` Nothing - (firstScheduled 1 (Timed [(2, 0.6), (0.5, 0.5)])) `shouldBe` Just (2, 0.6) - isNothing (asTimed [(2, 0.6 :: Rate), (0.5, 0.5), (1.5, 0.4)]) `shouldBe` True - (firstScheduled 1 (Timed [(2, 0.6), (0.5, 0.5), (1.5, 0.4)])) `shouldBe` - Just (2, 0.6) - it "Works on a very specific case it seems to not like" $ do - (noScheduledEvent 2.28 (2.28 + 0.42) (Timed [(2.3, 0.9)])) `shouldBe` False - it "Catastrophes are handled correctly" $ do - (demoSampleEvents03 == BDSCO.observedEvents demoFullEvents03) `shouldBe` + (firstScheduled (AbsoluteTime 1) (Timed [])) `shouldBe` Nothing + (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe` + Just (AbsoluteTime 2, 0.5) + (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe` + Nothing + (firstScheduled + (AbsoluteTime 1) + (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe` + Just (AbsoluteTime 2, 0.6) + isNothing + (asTimed + [ (AbsoluteTime 2, 0.6 :: Rate) + , (AbsoluteTime 0.5, 0.5) + , (AbsoluteTime 1.5, 0.4) + ]) `shouldBe` True - it "Catastrophes can be simulated" $ do - demoSim <- - simulation False - (fromJust (BDSCO.configuration 4 (1.3, 0.1, 0.1, ([(3, 0.5)]), 0.2))) - BDSCO.allEvents - length demoSim > 1 `shouldBe` True + (firstScheduled + (AbsoluteTime 1) + (Timed + [ (AbsoluteTime 2, 0.6) + , (AbsoluteTime 0.5, 0.5) + , (AbsoluteTime 1.5, 0.4) + ])) `shouldBe` + Just (AbsoluteTime 2, 0.6) + it "Works on a very specific case it seems to not like" $ do + (noScheduledEvent + (AbsoluteTime 2.28) + (AbsoluteTime (2.28 + 0.42)) + (Timed [(AbsoluteTime 2.3, 0.9)])) `shouldBe` + False describe "Disaster definitions" $ do it "Disasters are handled correctly" $ do - (demoSampleEvents04 == fromJust (BDSCOD.observedEvents demoFullEvents04)) `shouldBe` + ((Right $ [Observation e | e <- demoSampleEvents04]) == + (observedEvents demoFullEvents04)) `shouldBe` True it "Disasters can be simulated" $ do demoSim <- - simulation False - (fromJust (BDSCOD.configuration 4 (1.3, 0.1, 0.1, [(3, 0.5)], 0.2, [(3.5, 0.5)]))) - BDSCOD.allEvents + simulation + (fromJust + (BDSCOD.configuration + (TimeDelta 4) + False + ( 1.3 + , 0.1 + , 0.1 + , [(AbsoluteTime 3, 0.5)] + , 0.2 + , [(AbsoluteTime 3.5, 0.5)]))) + (allEvents BDSCOD.randomEvent) length demoSim > 1 `shouldBe` True - -birthDeathTests = do - describe "BirthDeath module tests" $ do - it "Construct a simulation configuration" $ do - (isJust (BD.configuration 1 (1, 1))) `shouldBe` True - (isJust (BD.configuration (-1) (1, 1))) `shouldBe` False - (isJust (BD.configuration 1 ((-1), 1))) `shouldBe` False - (isJust (BD.configuration 1 (1, (-1)))) `shouldBe` False - (isJust (BD.configuration 1 ((-1), (-1)))) `shouldBe` False - it "Mean behaviour is approximately correct" $ - let mean xs = fromIntegral (sum xs) / (fromIntegral $ length xs) - meanFinalSize = exp ((2.1 - 0.2) * 1.5) - randomBDEvents = - simulationWithSystemRandom False - (fromJust $ BD.configuration 1.5 (2.1, 0.2)) - BD.allEvents - numRepeats = 3000 - in do finalSizes <- replicateM numRepeats (finalSize <$> randomBDEvents) - (withinNPercent 5 (mean finalSizes) meanFinalSize) `shouldBe` True + describe "Extracting observed events" $ do + it "unseq obs still extracted if no seq obs" $ do + let noSequencedEvents = + [ Infection (AbsoluteTime 4.1) p1 p2 + , Infection (AbsoluteTime 4.3) p2 p3 + , Infection (AbsoluteTime 4.5) p2 p4 + , IndividualSample + { indSampTime = AbsoluteTime 5.3 + , indSampPerson = p1 + , indSampSeq = False + } + , StoppingTime + ] + expectedObs = + [ Observation + (IndividualSample + { indSampTime = AbsoluteTime 5.3 + , indSampPerson = p1 + , indSampSeq = False + }) + ] + isRight (observedEvents noSequencedEvents) `shouldBe` True + ((Right expectedObs) == (observedEvents noSequencedEvents)) `shouldBe` + True helperFuncTests = do describe "Helpers in Utility" $ do it "the isAscending function works" $ do - (isAscending ([] :: [Time])) `shouldBe` True + (isAscending ([] :: [AbsoluteTime])) `shouldBe` True (isAscending [-1.0]) `shouldBe` True (isAscending [1.0]) `shouldBe` True (isAscending [1.0, 2.0]) `shouldBe` True @@ -220,151 +266,160 @@ (isAscending [1.0, 2.0, -3.0]) `shouldBe` False it "the asTimed function works" $ do (isJust $ asTimed []) `shouldBe` True - (isJust $ asTimed [(0, 1)]) `shouldBe` True - (isJust $ asTimed [(0, 1), (1, 3)]) `shouldBe` True - (isJust $ asTimed [(0, 3), (1, 1)]) `shouldBe` True - (isJust $ asTimed [(1, 3), (0, 1)]) `shouldBe` False - let demoTimed = fromJust $ asTimed [(0, 1.2), (1, 3.1), (2, 2.7)] + (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True + (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe` + True + (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe` + True + (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe` + False + let demoTimed = + fromJust $ + asTimed + [ (AbsoluteTime 0, 1.2) + , (AbsoluteTime 1, 3.1) + , (AbsoluteTime 2, 2.7) + ] in do it "the cadlagValue function works" $ do - (isJust $ cadlagValue demoTimed (-1.0)) `shouldBe` False - ((== 1.2) . fromJust $ cadlagValue demoTimed 0.0) `shouldBe` True - ((== 1.2) . fromJust $ cadlagValue demoTimed 0.5) `shouldBe` True - ((== 3.1) . fromJust $ cadlagValue demoTimed 1.5) `shouldBe` True + (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe` + False + ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe` + True + ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe` + True + ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe` + True it "the diracDeltaValue function works" $ do - ((== 1.2) . fromJust $ diracDeltaValue demoTimed 0) `shouldBe` True - (isJust $ diracDeltaValue demoTimed 1) `shouldBe` True - (isJust $ diracDeltaValue demoTimed 0.9) `shouldBe` False - (isJust $ diracDeltaValue demoTimed 1.1) `shouldBe` False + ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe` + True + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe` + True + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe` + False + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe` + False it "the hasTime function works" $ do - (hasTime demoTimed 0) `shouldBe` True - (hasTime demoTimed 0.5) `shouldBe` False - (hasTime demoTimed 1) `shouldBe` True - (hasTime demoTimed 1.5) `shouldBe` False + (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True + (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False + (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True + (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False it "the nextTime function works" $ do - (0 == (fromJust $ nextTime demoTimed (-1))) `shouldBe` True - (1 == (fromJust $ nextTime demoTimed (0))) `shouldBe` True - (1 == (fromJust $ nextTime demoTimed (0.5))) `shouldBe` True + (AbsoluteTime 0 == + (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe` + True + (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe` + True + (AbsoluteTime 1 == + (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe` + True it "the nextTime function handles the last time correctly" $ do - isJust (nextTime demoTimed 1.9) `shouldBe` True - isJust (nextTime demoTimed 2.0) `shouldBe` True - isJust (nextTime demoTimed 2.1) `shouldBe` True - isJust (nextTime demoTimed 10.0) `shouldBe` True + isJust (nextTime demoTimed (AbsoluteTime 1.9)) `shouldBe` True + 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 "shifted times work" $ - let sf = fromJust $ asTimed [(-1.0,2.0),(1,3.0)] - val1 = cadlagValue sf 0 - val2 = cadlagValue sf (-2.0) - val3 = cadlagValue sf 1.5 - in do - isJust val1 `shouldBe` True - val1 == Just 2.0 `shouldBe` True - (not $ isJust val2) `shouldBe` True - isJust val3 `shouldBe` True - val3 == Just 3.0 `shouldBe` True + let sf = + fromJust $ + asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)] + val1 = cadlagValue sf (AbsoluteTime 0) + val2 = cadlagValue sf (AbsoluteTime (-2.0)) + val3 = cadlagValue sf (AbsoluteTime 1.5) + in do isJust val1 `shouldBe` True + val1 == Just 2.0 `shouldBe` True + (not $ isJust val2) `shouldBe` True + isJust val3 `shouldBe` True + val3 == Just 3.0 `shouldBe` True it "the asTimed function returns nothing as expected" $ do - (isJust $ asTimed [(0.0,-1)]) `shouldBe` True - (isJust $ asTimed [(0.0,1),(1.0,-1)]) `shouldBe` True - (isJust $ InhomBDS.inhomBDSRates [(0.0,1),(1.0,-1)] 0.5 0.5) `shouldBe` False - - - -readwriteTests = - do - describe "Change Event read/write" $ do - it "check we can writte an event" $ - let demoPerson = Person 3 - demoPersonField = toField demoPerson - demoPersonField' = "3" - demoEvent = Removal 1.0 demoPerson - demoRecord = toRecord demoEvent - demoRecord' = V.fromList ["removal", "1.0", "3", "NA"] :: Record - (Right demoEvent') = - runParser (parseRecord demoRecord) :: Either String EpidemicEvent - demoRecord2 = - toRecord (Catastrophe 1.0 (asPeople [p2, p3])) - (Right demoEvent2@(Catastrophe _ people2)) = - runParser (parseRecord demoRecord2) :: Either String EpidemicEvent - demoRecord2' = toRecord demoEvent2 - in do (demoPersonField' == demoPersonField) `shouldBe` True - (demoRecord' == demoRecord) `shouldBe` True - (demoEvent' == demoEvent) `shouldBe` True - (demoRecord2' == demoRecord2) `shouldBe` True - (numPeople people2 == 2) `shouldBe` True - + (isJust $ asTimed [(AbsoluteTime 0.0, -1)]) `shouldBe` True + (isJust $ asTimed [(AbsoluteTime 0.0, 1), (AbsoluteTime 1.0, -1)]) `shouldBe` + True + let (Just timedBirthRate) = + asTimed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, -1.0)] + (isJust $ InhomBDS.inhomBDSRates timedBirthRate 0.5 0.5) `shouldBe` False inhomExpTests = describe "Test the inhomogeneous exponential variate generator" $ let rate1 = 2.0 - sF1 = fromJust $ asTimed [(0, rate1)] + sF1 = fromJust $ asTimed [(AbsoluteTime 0, rate1)] mean1 = 1 / rate1 var1 = 1 / (rate1 ** 2.0) - sF2 = fromJust $ asTimed [(0, 1e-10),(1, rate1)] + sF2 = + 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 - u1 <- MWC.uniform gen :: IO Double - (u1 > 0) `shouldBe` True - x1 <- inhomExponential sF1 gen - (x1 > 0) `shouldBe` True - (x1 < 100) `shouldBe` True - True `shouldBe` True - it "check the mean and variance look sensible" $ - do gen <- genAction - x <- V.replicateM 20000 (inhomExponential sF1 gen) - 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 - x <- V.replicateM 20000 (inhomExponential sF2 gen) - withinNPercent 5 (mean x) mean2 `shouldBe` True - withinNPercent 5 (variance x) var2 `shouldBe` True - + in do it "check we can get a positive variate out" $ do + gen <- genAction + u1 <- MWC.uniform gen :: IO Double + (u1 > 0) `shouldBe` True + (Just x1) <- inhomExponential sF1 (AbsoluteTime 0) gen + (x1 > AbsoluteTime 0) `shouldBe` True + (x1 < AbsoluteTime 100) `shouldBe` True + True `shouldBe` True + it "check the mean and variance look sensible" $ do + gen <- genAction + 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 + xBoxed <- + V.replicateM 20000 (inhomExponential sF2 (AbsoluteTime 0) gen) + let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed + withinNPercent 5 (mean x) mean2 `shouldBe` True + withinNPercent 5 (variance x) var2 `shouldBe` True +illFormedTreeTest :: SpecWith () illFormedTreeTest = describe "Prevent the simulator returning a broken tree" $ do - let simDuration = 0.2 - simLambda = 3.2 - simMu = 0.3 - simPsi = 0.3 - simRho = 0.15 - simRhoTime = 2.6 - simOmega = 0.3 - simNu = 0.15 - simNuTime = 3.0 - simParams = (simLambda, simMu, simPsi, [(simRhoTime,simRho)], simOmega, [(simNuTime,simNu)]) - simConfig = BDSCOD.configuration simDuration simParams - in it "stress testing the observed events function" $ - do - null (BDSCOD.observedEvents []) `shouldBe` True - simEvents <- simulation True (fromJust simConfig) BDSCOD.allEvents - any isReconTreeLeaf simEvents `shouldBe` True - (length (fromJust $ BDSCOD.observedEvents simEvents) > 1) `shouldBe` True - - - - + let simDuration = TimeDelta 0.2 + simLambda = 3.2 + simMu = 0.3 + simPsi = 0.3 + simRho = 0.15 + simRhoTime = AbsoluteTime 2.6 + simOmega = 0.3 + simNu = 0.15 + simNuTime = AbsoluteTime 3.0 + simParams = + ( simLambda + , simMu + , simPsi + , [(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 inhomogeneousBDSTest = describe "InhomogeneousBDS module tests" $ do it "Check the observedEvents filters out removals" $ - let demoAllEvents = [Infection 0.1 p1 p2 - ,Sampling 0.2 p1 - ,Removal 0.3 p3 - ,Sampling 0.4 p2] - demoObsEvents = [Infection 0.1 p1 p2 - ,Sampling 0.2 p1 - ,Sampling 0.4 p2] - compObsEvents = InhomBDS.observedEvents demoAllEvents - in do - (compObsEvents == demoObsEvents) `shouldBe` True - + let demoAllEvents = + [ Infection (AbsoluteTime 0.1) p1 p2 + , IndividualSample (AbsoluteTime 0.2) p1 True + , Removal (AbsoluteTime 0.3) p3 + , IndividualSample (AbsoluteTime 0.4) p2 True + ] + demoObsEvents = + [ Infection (AbsoluteTime 0.1) p1 p2 + , IndividualSample (AbsoluteTime 0.2) p1 True + , IndividualSample (AbsoluteTime 0.4) p2 True + ] + compObsEvents = observedEvents demoAllEvents + in do (compObsEvents == (Right [Observation e | e <- demoObsEvents])) `shouldBe` True helperTypeTests = do describe "Helpers for working with the types" $ do it "the isAscending function works" $ do - (isAscending ([] :: [Time])) `shouldBe` True + (isAscending ([] :: [AbsoluteTime])) `shouldBe` True (isAscending [-1.0]) `shouldBe` True (isAscending [1.0]) `shouldBe` True (isAscending [1.0, 2.0]) `shouldBe` True @@ -374,35 +429,59 @@ (isAscending [1.0, 2.0, -3.0]) `shouldBe` False it "the asTimed function works" $ do (isJust $ asTimed []) `shouldBe` True - (isJust $ asTimed [(0, 1)]) `shouldBe` True - (isJust $ asTimed [(0, 1), (1, 3)]) `shouldBe` True - (isJust $ asTimed [(0, 3), (1, 1)]) `shouldBe` True - (isJust $ asTimed [(1, 3), (0, 1)]) `shouldBe` False - let demoTimed = fromJust $ asTimed [(0, 1.2), (1, 3.1), (2, 2.7)] + (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True + (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe` + True + (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe` + True + (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe` + False + let demoTimed = + fromJust $ + asTimed + [ (AbsoluteTime 0, 1.2) + , (AbsoluteTime 1, 3.1) + , (AbsoluteTime 2, 2.7) + ] in do it "the cadlagValue function works" $ do - (isJust $ cadlagValue demoTimed (-1.0)) `shouldBe` False - ((== 1.2) . fromJust $ cadlagValue demoTimed 0.0) `shouldBe` True - ((== 1.2) . fromJust $ cadlagValue demoTimed 0.5) `shouldBe` True - ((== 3.1) . fromJust $ cadlagValue demoTimed 1.5) `shouldBe` True + (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe` + False + ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe` + True + ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe` + True + ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe` + True it "the diracDeltaValue function works" $ do - ((== 1.2) . fromJust $ diracDeltaValue demoTimed 0) `shouldBe` True - (isJust $ diracDeltaValue demoTimed 1) `shouldBe` True - (isJust $ diracDeltaValue demoTimed 0.9) `shouldBe` False - (isJust $ diracDeltaValue demoTimed 1.1) `shouldBe` False + ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe` + True + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe` + True + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe` + False + (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe` + False it "the hasTime function works" $ do - (hasTime demoTimed 0) `shouldBe` True - (hasTime demoTimed 0.5) `shouldBe` False - (hasTime demoTimed 1) `shouldBe` True - (hasTime demoTimed 1.5) `shouldBe` False + (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True + (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False + (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True + (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False it "the nextTime function works" $ do - (0 == (fromJust $ nextTime demoTimed (-1))) `shouldBe` True - (1 == (fromJust $ nextTime demoTimed (0))) `shouldBe` True - (1 == (fromJust $ nextTime demoTimed (0.5))) `shouldBe` True + (AbsoluteTime 0 == + (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe` + True + (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe` + True + (AbsoluteTime 1 == + (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe` + True it "shifted times work" $ - let sf = fromJust $ asTimed [(-1.0, 2.0), (1, 3.0)] - val1 = cadlagValue sf 0 - val2 = cadlagValue sf (-2.0) - val3 = cadlagValue sf 1.5 + let sf = + fromJust $ + asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)] + val1 = cadlagValue sf (AbsoluteTime 0) + val2 = cadlagValue sf (AbsoluteTime (-2.0)) + val3 = cadlagValue sf (AbsoluteTime 1.5) in do isJust val1 `shouldBe` True val1 == Just 2.0 `shouldBe` True (not $ isJust val2) `shouldBe` True @@ -412,94 +491,259 @@ jsonTests = do describe "Converting to and from JSON" $ do it "Conversion of Timed Rate" $ do - let demoObj = Timed [(0.0, 1.0), (1.0, 1.0)] :: Timed Rate + let demoObj = + Timed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, 1.0)] :: Timed Rate (Timed demoVals) = demoObj demoJson = "[[0,1],[1,1]]" encodedObj = Json.encode demoObj decodedJson = Json.decode demoJson :: Maybe (Timed Rate) in do True `shouldBe` True - let (Timed foo) = demoObj in demoVals == foo `shouldBe` True + let (Timed foo) = demoObj + in demoVals == foo `shouldBe` True encodedObj == demoJson `shouldBe` True isJust decodedJson `shouldBe` True - let (Timed bar) = fromJust decodedJson in demoVals == bar `shouldBe` True + let (Timed bar) = fromJust decodedJson + in demoVals == bar `shouldBe` True equalBuilders :: BBuilder.Builder -> BBuilder.Builder -> Bool equalBuilders a b = BBuilder.toLazyByteString a == BBuilder.toLazyByteString b - newickTests = - let p1 = Person 1 - p2 = Person 2 - p3 = Person 3 - ps = asPeople [p1,p2] - maybeEpiTree = maybeEpidemicTree [Infection 1 p1 p2,Infection 2 p2 p3,Catastrophe 3 (asPeople [p1,p3]),Removal 4 p2] - maybeEpiTree' = maybeEpidemicTree [Infection 1 p1 p2,Infection 2 p2 p3,Catastrophe 3 (asPeople [p1,p3]),Sampling 4 p2] - maybeEpiTree'' = maybeEpidemicTree [Infection 1 p1 p2,Infection 2 p2 p3,Disaster 3 (asPeople [p1,p3]),Sampling 4 p2] - in - describe "Writing to Newick" $ do - it "equalBuilders works as expected" $ do - equalBuilders (BBuilder.charUtf8 ':') (BBuilder.charUtf8 ':') `shouldBe` True - equalBuilders (BBuilder.charUtf8 'a') (BBuilder.charUtf8 ':') `shouldBe` False - it "derivedFrom works as expected" $ do - let p1 = Person 1 - let p2 = Person 2 - let p3 = Person 3 - let e = [Infection 0.3 p1 p2] - derivedFrom p1 e == derivedFrom p2 e `shouldBe` True - derivedFrom p1 e /= derivedFrom p3 e `shouldBe` True - derivedFrom p1 e /= [] `shouldBe` True - null (derivedFrom p3 e) `shouldBe` True - derivedFrom p1 e == e `shouldBe` True - let foo = derivedFrom (Person 1) [Infection 0.3 (Person 1) (Person 2),Sampling 0.7 (Person 1)] - let bar = derivedFrom (Person 2) [Infection 0.3 (Person 1) (Person 2),Sampling 0.7 (Person 1)] - foo == bar `shouldBe` True - it "maybeEpidemicTree works as expected: 1" $ do - let e1 = Removal 1 (Person 1) - maybeEpidemicTree [e1] == Just (Leaf e1) `shouldBe` True - let t1 = maybeEpidemicTree [Infection 0.3 (Person 1) (Person 2),Sampling 0.6 (Person 2),Sampling 0.7 (Person 1)] - let t2 = Just (Branch (Infection 0.3 (Person 1) (Person 2)) (Leaf (Sampling 0.7 (Person 1))) (Leaf (Sampling 0.6 (Person 2)))) - t1 == t2 `shouldBe` True - maybeEpidemicTree [Infection 0.3 (Person 1) (Person 2)] == Just (Branch (Infection 0.3 (Person 1) (Person 2)) (Shoot (Person 1)) (Shoot (Person 2))) `shouldBe` True - maybeEpidemicTree [Infection 0.3 (Person 1) (Person 2),Sampling 0.7 (Person 1)] == Just (Branch (Infection 0.3 (Person 1) (Person 2)) (Leaf (Sampling 0.7 (Person 1))) (Shoot (Person 2))) `shouldBe` True - let trickyEvents = [Infection 0.3 (Person 1) (Person 2),Infection 0.4 (Person 2) (Person 3),Sampling 0.6 (Person 3),Sampling 0.7 (Person 1)] - isJust (maybeEpidemicTree trickyEvents) `shouldBe` True - it "maybeEpidemicTree works as expected: 2" $ do - let p1 = Person 1 - p2 = Person 2 - demoEvents = [Catastrophe 0.5 (asPeople []) -- Because the first event is a null event it can be ignored! - ,Infection 1.0 p1 p2 - ,Catastrophe 1.5 (asPeople []) - ,Catastrophe 2.0 (asPeople [p1,p2])] - (length demoEvents == 4) `shouldBe` True - (maybeEpidemicTree demoEvents == maybeEpidemicTree (tail demoEvents)) `shouldBe` True - it "asNewickString works for EpidemicTree" $ do - let trickyEvents = [Infection 0.3 (Person 1) (Person 2),Infection 0.4 (Person 2) (Person 3),Sampling 0.6 (Person 3),Sampling 0.7 (Person 1)] - let maybeNewickPair = asNewickString (0, Person 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 - [Sampling 0.6 (Person 3),Sampling 0.7 (Person 1)] == snd (fromJust maybeNewickPair) `shouldBe` True - equalBuilders newickTarget (fst $ fromJust maybeNewickPair) `shouldBe` True - isJust maybeReconTree `shouldBe` True - it "asNewickString works for ReconstructedTree" $ do - isJust (asNewickString (0,Person 1) (RLeaf (Sampling 1 (Person 1)))) `shouldBe` True - let trickyEvents = [Infection 0.3 (Person 1) (Person 2),Infection 0.4 (Person 2) (Person 3),Sampling 0.6 (Person 3),Sampling 0.7 (Person 1)] - let maybeNewickPair = asNewickString (0, Person 1) =<< maybeReconstructedTree =<< maybeEpidemicTree trickyEvents - let newickTarget = BBuilder.stringUtf8 "(1:0.39999999999999997,3:0.3):0.3" - isJust maybeNewickPair `shouldBe` True - [Sampling 0.6 (Person 3),Sampling 0.7 (Person 1)] == snd (fromJust maybeNewickPair) `shouldBe` True - equalBuilders newickTarget (fst $ fromJust maybeNewickPair) `shouldBe` True - let catasNewick = (asNewickString (0,Person 1) (RLeaf (Catastrophe 1 (asPeople [Person 1,Person 2])))) - let catasTarget = BBuilder.stringUtf8 "1&2:1.0" - equalBuilders catasTarget (fst $ fromJust catasNewick) `shouldBe` True + let p1 = Person (Identifier 1) + p2 = Person (Identifier 2) + p3 = Person (Identifier 3) + ps = asPeople [p1, p2] + maybeEpiTree = + maybeEpidemicTree + [ Infection (AbsoluteTime 1) p1 p2 + , Infection (AbsoluteTime 2) p2 p3 + , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True + , Removal (AbsoluteTime 4) p2 + ] + maybeEpiTree' = + maybeEpidemicTree + [ Infection (AbsoluteTime 1) p1 p2 + , Infection (AbsoluteTime 2) p2 p3 + , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True + , IndividualSample (AbsoluteTime 4) p2 True + ] + maybeEpiTree'' = + maybeEpidemicTree + [ Infection (AbsoluteTime 1) p1 p2 + , Infection (AbsoluteTime 2) p2 p3 + , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) False + , IndividualSample (AbsoluteTime 4) p2 True + ] + in describe "Writing to Newick" $ do + it "equalBuilders works as expected" $ do + equalBuilders (BBuilder.charUtf8 ':') (BBuilder.charUtf8 ':') `shouldBe` + True + equalBuilders (BBuilder.charUtf8 'a') (BBuilder.charUtf8 ':') `shouldBe` + False + it "derivedFrom works as expected" $ do + let p1 = Person (Identifier 1) + let p2 = Person (Identifier 2) + let p3 = Person (Identifier 3) + let e = [Infection (AbsoluteTime 0.3) p1 p2] + derivedFrom p1 e == derivedFrom p2 e `shouldBe` True + derivedFrom p1 e /= derivedFrom p3 e `shouldBe` True + derivedFrom p1 e /= [] `shouldBe` True + null (derivedFrom p3 e) `shouldBe` True + derivedFrom p1 e == e `shouldBe` True + let foo = + derivedFrom + (Person (Identifier 1)) + [ Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2)) + , IndividualSample + (AbsoluteTime 0.7) + (Person (Identifier 1)) + True + ] + let bar = + derivedFrom + (Person (Identifier 2)) + [ Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2)) + , IndividualSample + (AbsoluteTime 0.7) + (Person (Identifier 1)) + True + ] + foo == bar `shouldBe` True + it "maybeEpidemicTree works as expected: 1" $ do + let e1 = Removal (AbsoluteTime 1) (Person (Identifier 1)) + maybeEpidemicTree [e1] == Right (Leaf e1) `shouldBe` True + let t1 = + maybeEpidemicTree + [ Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2)) + , IndividualSample + (AbsoluteTime 0.6) + (Person (Identifier 2)) + True + , IndividualSample + (AbsoluteTime 0.7) + (Person (Identifier 1)) + True + ] + let t2 = + Right + (Branch + (Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2))) + (Leaf + (IndividualSample + (AbsoluteTime 0.7) + (Person (Identifier 1)) + True)) + (Leaf + (IndividualSample + (AbsoluteTime 0.6) + (Person (Identifier 2)) + True))) + t1 == t2 `shouldBe` True + maybeEpidemicTree + [ Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2)) + ] == + Right + (Branch + (Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2))) + (Shoot (Person (Identifier 1))) + (Shoot (Person (Identifier 2)))) `shouldBe` + True + maybeEpidemicTree + [ Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2)) + , IndividualSample (AbsoluteTime 0.7) (Person (Identifier 1)) True + ] == + Right + (Branch + (Infection + (AbsoluteTime 0.3) + (Person (Identifier 1)) + (Person (Identifier 2))) + (Leaf + (IndividualSample + (AbsoluteTime 0.7) + (Person (Identifier 1)) + True)) + (Shoot (Person (Identifier 2)))) `shouldBe` + True + 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 + ] + isRight (maybeEpidemicTree trickyEvents) `shouldBe` True + it "maybeEpidemicTree works as expected: 2" $ do + let p1 = Person (Identifier 1) + p2 = Person (Identifier 2) + demoEvents = + [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored! + , Infection (AbsoluteTime 1.0) p1 p2 + , PopulationSample (AbsoluteTime 1.5) (asPeople []) True + , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True + ] + (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 + (AbsoluteTime 0, Person (Identifier 1)) + (RLeaf + (Observation (IndividualSample + (AbsoluteTime 1) + (Person (Identifier 1)) + True)))) `shouldBe` + True + 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 et = maybeEpidemicTree trickyEvents :: Either String EpidemicTree + rt = maybeReconstructedTree =<< et :: Either String ReconstructedTree + maybeNewickPair = asNewickString (AbsoluteTime 0, Person (Identifier 1)) =<< (either2Maybe rt) + let newickTarget = + BBuilder.stringUtf8 "(1:0.39999999999999997,3:0.3):0.3" + 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 + let catasNewick = + (asNewickString + (AbsoluteTime 0, Person (Identifier 1)) + (RLeaf + (Observation (PopulationSample + (AbsoluteTime 1) + (asPeople + [Person (Identifier 1), Person (Identifier 2)]) + True)))) + let catasTarget = BBuilder.stringUtf8 "1&2:1.0" + equalBuilders catasTarget (fst $ fromJust catasNewick) `shouldBe` True main :: IO () main = hspec $ do eventHandlingTests - birthDeathTests helperFuncTests - readwriteTests inhomExpTests illFormedTreeTest inhomogeneousBDSTest