diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,184 @@ +# Changelog for epi-sim + +## 0.2.2.0 + +- Clean up for a release candidate. + +## 0.2.1.0 + +- Update the stack resolver to `lts-16.17` and add bounds on the package + 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. + +## 0.2.0.1 + +- Improve documentation in `Epidemic.BDSCOD`. +- Upgrade to `epi-types-0.2.1.2` for the bug fix. + +## 0.2.0.0 + +- Predicates such as `isSampling` and `isOccurrence` have been replaced with + clearer named alternatives: `isNonReconTreeObservation` and `isReconTreeLeaf` + to avoid uncertainty about whether scheduled events are included. +- Upgrade to `epi-types-0.2.1.0` which adds substantial amounts of Newick + functionality meaning this can be removed. + +## 0.1.9.0 + +- Upgrade to `epi-types-0.2.0.0` which makes a lot of functions redundant so + they have been removed. + +## 0.1.8.8 + +- Upgrade to `epi-types-0.1.1.1` and remove `joinTimed` since it is no longer + needed. + +## 0.1.8.7 + +- Upgrade to `epi-types-0.1.1.0` and start using the `Timed *` class. + +## 0.1.8.6 + +- Use `-j2` to slightly improve compilation time. +- Include a `observedEvents` function which was missing from the + `InhomogeneousBDS` module. + +## 0.1.8.5 + +- Start using cabal directly without hpack. +- Move time dependent parameters into a `epi-sim` package so they are easier to + share, and import them here. + +## 0.1.8.4 + +- Improve documentation of construction of `inhomBDSRates` and include a test to + check this conforms to expectations. +- Implement a birth-death-sampling model with time-dependent birth rates and + include an example of what this returns in the README. + +## 0.1.8.3 + +- Improve error handling in `birth-death-lines.R`. +- Refactor some of the plotting code in `birth-death-lines.R` to make it easier + to maintain. +- Use the `future` and `furrr` packages in R to generate the visualisation of a + simulation in parallel since the serial version is frustratingly slow. Both + packages are on CRAN so it does not seem unreasonable. + +## 0.1.8.2 + +- Include an additional parameter in the `simulation` and + `simulationWithSystemRandom` to determine whether the simulation should be + conditioned on having at least two sampled leaves in the reconstructed tree. + It does this by repeatedly simulating until such a tree is acheved. + +## 0.1.8.1 + +- Implement an example of the inhomogeneous birth-death process and include it + in the `README`. +- Fix broken example using the BDSCOD model and organise examples in `examples/` + rather than just the single one in `app/`. + +## 0.1.8.0 + +- Define an inhomogeneous birth-death process in `InhomogeneousBD`. +- Define an inhomogeneous exponential random sampler using `Timed Rate`s. +- Extend the `ModelParameters` class to have a `birthProb` function. +- Include a `Timed` type for time varying values. +- Include a test that the final size of a birth-death simulation is + approximately correct. +- Format the changelog to make it easier to read outside of a browser. +- Fix naming convention and exports in the `BirthDeath`, `BirthDeathSampling` + and `BirthDeathSamplingOccurrence`modules to conform to the pattern used in + the `BDSCOD` module and move `simulation` into the `Utility` module since it + is common between all models +- Adjust the type of the `ModelParameters` interface to include potential time + dependence. + +## 0.1.7.2 + +- Make `Event` an instance of `FromRecord` to match with `ToRecord`. + +## 0.1.7.1 + +- Provide `vis/birth-death-lines.R` which visualises a simulation. + +## 0.1.7.0 + +- Provide `vis/ott.R` to visualise observations through time and make + `vis/ltt.R` ignore disasters. +- Replace some Catastrophe specific functions with equivalents for handling + scheduled events more generally + + `noCatastrophe` becomes `noScheduledEvent` + + `firstCatastrophe` becomes `firstScheduled` +- Change the naming of the BDSCO model to be consistent with the new BDSCOD + model. +- Implement a birth-death-sampling-catastrophe-occurrence-disaster model in + `Epidemic.BDSCOD`. + +## 0.1.6.2 + +- Update the visualisation code in `vis/ltt.R` such that it parses the + catastrophe events to draw the LTT. +- BUG FIX: In the observed events catastrophes are replicated when they should + not be. + +## 0.1.6.1 + +- Include a visualisation script for the LTT of reconstructed trees: `vis/ltt.R` + which is described in the README. + +## 0.1.6.0 + +- Move the `ToField` declarations into the `Epidemic` module to avoid orphaned + instance and select a representation for multiple people in the CSV output. +- Implement a birth-death-sampling-catastrophe-occurrence model in + `Epidemic.BirthDeathSamplingCatastropheOccurrence`; it's a shame about the + length of the name. +- Include a contructor for catastrophe events in which multiple people can be + sampled simultaneously. + +## 0.1.5.1 + +- Include a `.gitignore`. +- Improved documentation with `haddock`. + +## 0.1.5.0 + +- Include some testing with `hspec`. +- The `Epidemic.BirthDeathSamplingOccurrence` module now has a function to + extract the observable events. +- Change the simulation functions to return events in the order they occurred. +- Change the configuration functions to take tuples of model parameters. + +## 0.1.4.0 + +- Make `Event` an instance of the `Ord` typeclass. +- Implement types and conversion functions for a `TransmissionTree` and a + `SampleTree`. +- Implement a parser and printer for Newick in `Epidemic.Utility`. + +## 0.1.3.0 + +- Implement a birth-death-sampling-occurrence model in + `Epidemic.BirthDeathSamplingOccurrence`. +- Implement a visualisation of the birth-death-sampling LTT in + `vis/birthDeathSamplingVisualisation.R`. +- Implement a birth-death-sampling model in `Epidemic.BirthDeathSampling`. + +## 0.1.2.0 + +- Implement a visualisation of the infection network in + `vis/infection-tree.vg.json`. +- Implement `Epidemic.Utility.eventsAsJsonTree`. +- Move the `Epidemic.Simulation` module to `Epidemic.Utility`. + +## 0.1.1.0 + +- Write the output to CSV using `cassava` + +## 0.1.0.0 + +- Basic prototype to simulate a constant rate birth-death process. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Alexander E. Zarebski + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,24 @@ +# epi-sim +A library for simulating epidemics, with a focus on phylodynamics and +observation models. + +## 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 + +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. diff --git a/epi-sim.cabal b/epi-sim.cabal new file mode 100644 --- /dev/null +++ b/epi-sim.cabal @@ -0,0 +1,106 @@ +cabal-version: 1.22 +name: epi-sim +version: 0.2.2.0 +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/ + . + 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. + . +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 + +source-repository head + 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 + 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 + default-language: Haskell2010 + 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 + build-depends: + aeson, + base, + bytestring, + epi-sim, + primitive, + vector, + cassava, + hspec, + mwc-random, + statistics, + trifecta + default-language: Haskell2010 diff --git a/src/Epidemic.hs b/src/Epidemic.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic.hs @@ -0,0 +1,206 @@ +{-# 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 qualified Data.Vector as V +import Data.Word +import Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import GHC.Generics (Generic) + + + +-- | 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 + +-- | 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 time timedProb = do + time' <- nextTime timedProb time + prob' <- diracDeltaValue timedProb time' + return (time',prob') + +-- | Predicate for whether there is a scheduled event during an interval. +noScheduledEvent :: Time -- ^ Start time for interval + -> Time -- ^ 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) + +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 + +peopleInEvents :: [EpidemicEvent] -> People +peopleInEvents events = + People . V.fromList . nub . concat $ map personsInEvent events + + +-- | Predicate for whether the first person infected the second in the given event +infected :: Person -- ^ Potential infector + -> Person -- ^ Potential infectee + -> EpidemicEvent -- ^ Given event + -> Bool +infected p1 p2 e = + case e of + (Infection _ infector infectee) -> infector == p1 && infectee == p2 + _ -> False + + +-- | The people infected by a particular person in a list of events. +infectedBy :: Person -- ^ Potential infector + -> [EpidemicEvent] -- ^ Events + -> People +infectedBy person events = + case events of + [] -> People V.empty + (Infection _ infector infectee :es) -> + if infector == person + then addPerson infectee $ infectedBy person es + else infectedBy person es + (_:es) -> infectedBy person es + + +-- | 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' diff --git a/src/Epidemic/BDSCOD.hs b/src/Epidemic/BDSCOD.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/BDSCOD.hs @@ -0,0 +1,180 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/BirthDeath.hs @@ -0,0 +1,97 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/BirthDeathSampling.hs @@ -0,0 +1,101 @@ +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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/BirthDeathSamplingCatastropheOccurrence.hs @@ -0,0 +1,136 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/BirthDeathSamplingOccurrence.hs @@ -0,0 +1,119 @@ +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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/InhomogeneousBD.hs @@ -0,0 +1,115 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/InhomogeneousBDS.hs @@ -0,0 +1,135 @@ +{-# 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/Types/Events.hs b/src/Epidemic/Types/Events.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Events.hs @@ -0,0 +1,347 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DeriveGeneric #-} + +module Epidemic.Types.Events + ( EpidemicEvent(Infection, Removal, Sampling, Catastrophe, + Occurrence, Disaster) + , EpidemicTree(Branch, Leaf, Shoot) + , maybeEpidemicTree + , 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 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 + 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 + +-- | Epidemic Events are ordered based on which occurred first. +instance Ord EpidemicEvent where + e1 <= e2 = eventTime e1 <= eventTime e2 + +-- | The absolute time an event occurred. +eventTime :: EpidemicEvent -> Time +eventTime e = + case e of + Infection time _ _ -> time + Removal time _ -> time + Sampling time _ -> time + Catastrophe time _ -> time + Occurrence time _ -> time + Disaster time _ -> time + +-- | The events that occurred as a result of the existance of the given person. +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 _ [] = [] +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 + +{-| 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. +-} +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) + + +-- | 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) = + case e of + Infection _ p1 p2 -> + let infectorEvents = derivedFrom p1 (es : ess) + infecteeEvents = derivedFrom p2 (es : ess) + in do leftTree <- + if null infectorEvents + then Just (Shoot p1) + else maybeEpidemicTree infectorEvents + rightTree <- + if null infecteeEvents + then Just (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 diff --git a/src/Epidemic/Types/Observations.hs b/src/Epidemic/Types/Observations.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Observations.hs @@ -0,0 +1,2 @@ + +module Epidemic.Types.Observations where diff --git a/src/Epidemic/Types/Parameter.hs b/src/Epidemic/Types/Parameter.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Parameter.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE DeriveGeneric #-} + +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) + +instance Semigroup (Timed a) where + (Timed x) <> (Timed y) = Timed $ List.sortOn fst (x ++ y) + +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 new file mode 100644 --- /dev/null +++ b/src/Epidemic/Types/Population.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DeriveGeneric #-} + +module Epidemic.Types.Population + ( Person(Person) + , People(People) + , asPeople + , includesPerson + , haveCommonPeople + , nullPeople + , numPeople + , addPerson + , removePerson + , personByteString + ) where + +import qualified Data.Aeson as Json +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 + +newtype Person = + Person Integer + 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) + +newtype People = + People (V.Vector Person) + deriving (Show, Eq, Generic) + +instance Json.FromJSON People + +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 + +-- | Predicate for whether a person is one of the people +includesPerson :: People -> Person -> Bool +includesPerson (People persons) person = V.elem person persons + +-- | Predicate for whether two sets of people have any members in common. +haveCommonPeople :: People -> People -> Bool +haveCommonPeople (People ps1) (People ps2) = V.any (\p -> V.elem p ps2) ps1 + +-- | Predicate for whether there are any people +nullPeople :: People -> Bool +nullPeople (People persons) = V.null persons + +-- | The number of people +numPeople :: People -> Int +numPeople (People persons) = V.length persons + +-- | Add a person to a group of people +addPerson :: Person -> People -> People +addPerson person (People persons) = People $ V.cons person persons + +-- | Remove a person from a group of people +removePerson :: Person -> People -> People +removePerson person (People persons) = People $ V.filter (/= person) persons + +-- | A bytestring builder for a person +personByteString :: Person -> BBuilder.Builder +personByteString (Person n) = BBuilder.integerDec n diff --git a/src/Epidemic/Utility.hs b/src/Epidemic/Utility.hs new file mode 100644 --- /dev/null +++ b/src/Epidemic/Utility.hs @@ -0,0 +1,234 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +module Epidemic.Utility where + +import Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Control.Monad (liftM) +import qualified Data.List as List +import qualified Data.Maybe as Maybe +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as Char8 +import GHC.Generics (Generic) +import qualified Data.Vector as V +import System.Random.MWC +import System.Random.MWC.Distributions +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Applicative +import Text.Trifecta + +import Epidemic + +data SimulationConfiguration r p = + SimulationConfiguration + { rates :: r + , population :: p + , newIdentifier :: Integer + , timeLimit :: Time + } + +initialIdentifier :: Integer +initialIdentifier = 1 + +newPerson :: Integer -> (Person, Integer) +newPerson identifier = (Person identifier, identifier + 1) + +selectElem :: V.Vector a -> Int -> (a, V.Vector a) +selectElem v n + | n == 0 = (V.head v, V.tail v) + | otherwise = + 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 + u <- uniform gen + return $ selectElem persons (floor (u * numPersons)) + where + numPersons = fromIntegral $ V.length persons :: Double + + +type NName = Maybe String + +type NLength = Maybe Double + +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) + +instance Show NBranchSet where + show (NBranchSet bs) = "(" ++ (List.intercalate "," (map show bs)) ++ ")" + +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) + +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 + +-- | 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 + +-- | 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 + +-- | 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 + + +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 + if count' isReconTreeLeaf events >= 2 + then return $ sort events + else simulation' config allEvents 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, _, _) <- + withSystemRandom $ \g -> + allEvents rates timeLimit (0, [], population, newIdentifier) g + if atLeastCherry + then (if count' isReconTreeLeaf events >= 2 + then return $ sort events + else simulationWithSystemRandom True config allEvents) + else return $ sort events + + +-- | The number of lineages at the end of a simulation. +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. +randInhomExp :: PrimMonad m + => Double -- ^ Timer + -> Timed Double -- ^ Step function + -> Gen (PrimState m) -- ^ Generator. + -> m (Maybe Double) +randInhomExp crrT stepFunc gen = + let crrR = cadlagValue stepFunc crrT + nxtT = nextTime stepFunc crrT + 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) + else return Nothing diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,508 @@ +{-# LANGUAGE OverloadedStrings #-} + +import Control.Exception (evaluate) +import Control.Monad +import qualified Data.Aeson as Json +import qualified Data.ByteString as B +import qualified Data.ByteString.Builder as BBuilder +import Data.Csv +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 Epidemic.Types.Events +import Epidemic.Types.Parameter +import Epidemic.Types.Population +import Epidemic.Utility +import Statistics.Sample +import qualified System.Random.MWC as MWC +import Test.Hspec + +-- | 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 + +p2 = Person 2 + +p3 = Person 3 + +p4 = Person 4 + +p5 = Person 5 + +p6 = Person 6 + +p7 = Person 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 + ] + +demoSampleEvents01 = + [ Infection 1 p1 p2 + , Sampling 3 p1 + , Infection 4 p2 p4 + , Sampling 6 p4 + , Occurrence 8 p2 + , Occurrence 11 p6 + , Sampling 12 p5 + ] + +-- | 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 + ] + +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 + ] + +-- | 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]) + ] + +demoSampleEvents03 = + [ Infection 1 p1 p4 + , Infection 2 p1 p2 + , Sampling 3 p1 + , Infection 4 p2 p3 + , Catastrophe 6 (asPeople [p2, p3, p4]) + ] + +-- | 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]) + ] + +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]) + ] + +eventHandlingTests = do + describe "Post-simulation processing" $ do + it "Extracting observed events" $ do + (demoSampleEvents01 == BDSO.observedEvents demoFullEvents01) `shouldBe` + True + (demoSampleEvents02 == BDSO.observedEvents demoFullEvents02) `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 + 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 + 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` + 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 + describe "Disaster definitions" $ do + it "Disasters are handled correctly" $ do + (demoSampleEvents04 == fromJust (BDSCOD.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 + 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 + +helperFuncTests = do + describe "Helpers in Utility" $ do + it "the isAscending function works" $ do + (isAscending ([] :: [Time])) `shouldBe` True + (isAscending [-1.0]) `shouldBe` True + (isAscending [1.0]) `shouldBe` True + (isAscending [1.0, 2.0]) `shouldBe` True + (isAscending [1.0, 2.0, 3.0]) `shouldBe` True + (isAscending [1.0, -2.0]) `shouldBe` False + (isAscending [1.0, -2.0, 3.0]) `shouldBe` False + (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)] + 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 + 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 + 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 + 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 + 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 + 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 + 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 + + +inhomExpTests = + describe "Test the inhomogeneous exponential variate generator" $ + let rate1 = 2.0 + sF1 = fromJust $ asTimed [(0, rate1)] + mean1 = 1 / rate1 + var1 = 1 / (rate1 ** 2.0) + sF2 = fromJust $ asTimed [(0, 1e-10),(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 + + +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 + + + + + +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 + + +helperTypeTests = do + describe "Helpers for working with the types" $ do + it "the isAscending function works" $ do + (isAscending ([] :: [Time])) `shouldBe` True + (isAscending [-1.0]) `shouldBe` True + (isAscending [1.0]) `shouldBe` True + (isAscending [1.0, 2.0]) `shouldBe` True + (isAscending [1.0, 2.0, 3.0]) `shouldBe` True + (isAscending [1.0, -2.0]) `shouldBe` False + (isAscending [1.0, -2.0, 3.0]) `shouldBe` False + (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)] + 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 + 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 + 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 + 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 + 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 + +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 + (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 + encodedObj == demoJson `shouldBe` True + isJust decodedJson `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 + +main :: IO () +main = + hspec $ do + eventHandlingTests + birthDeathTests + helperFuncTests + readwriteTests + inhomExpTests + illFormedTreeTest + inhomogeneousBDSTest + helperTypeTests + jsonTests + newickTests