{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards       #-}

-- |
-- Module: Epidemic.Model.InhomogeneousBDSCOD
-- Copyright: (c) 2021 Alexander E. Zarebski
-- License: MIT
--
-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>
-- Stability: unstable
-- Portability: ghc
--
-- This module defines a birth-death model with continuous time sampling and
-- scheduled sampling and rates that are piece-wise constant in time.
--
-- __Example:__ we will run a simulation for one unit of time and require that
-- there be at least two sequenced samples.
--
-- >>> simDuration = TimeDelta 1.0
-- >>> atLeastTwoSequences = True
--
-- The rates can change through time so we need to specify the times at which
-- they change. In this example the birth rate starts at 1.0 and then drops down
-- to 0.5. The other rates stay at their initial values.
--
-- >>> birthRateSpec = [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 0.5, 0.5)]
-- >>> deathRateSpec = [(AbsoluteTime 0.0, 0.2)]
-- >>> sampRateSpec = [(AbsoluteTime 0.0, 0.1)]
-- >>> occRateSpec = [(AbsoluteTime 0.0, 0.1)]
--
-- There are a couple of scheduled samples with probabilities specified for
-- them, ie there will be a scheduled sample at time 0.9 where each lineage is
-- removed and sequenced individually with probability 0.1 and at times 0.5 and
-- 0.75 there is a scheduled sample where individuals are removed but /not/
-- sequenced with probabilities 0.4 and 0.5 respectively.
--
-- >>> seqSched = [(AbsoluteTime 0.9, 0.1)]
-- >>> unseqSched = [(AbsoluteTime 0.5, 0.4), (AbsoluteTime 0.75, 0.5)]
--
-- This is enough to define a 'SimulationConfiguration'. We will ignore the
-- possibility of using a termination handler for this example.
--
-- >>> ratesAndProbs = (birthRateSpec,deathRateSpec,sampRateSpec,seqSched,occRateSpec,unseqSched)
-- >>> (Just simConfig) = configuration simDuration atLeastTwoSequences Nothing ratesAndProbs
--
-- Then we can use this to generated a list of epidemic events in the simulation
--
-- >>> myEpidemicEvents = simulationWithSystem simConfig (allEvents randomEvent)
--
-- and from this we can extract the observations
--
-- >>> myObservedEvents = do
-- >>>   simState <- myEpidemicEvents
-- >>>   case simState of
-- >>>     Right es -> return $ observedEvents es
-- >>>     Left _ -> return $ Left "simulation terminated early"
--

module Epidemic.Model.InhomogeneousBDSCOD
  ( configuration
  , randomEvent
  , InhomBDSCODRates(..)
  , InhomBDSCODPop(..)
  , getNumRemovedByDeath
  , getNumRemovedBySampling
  , getNumRemovedByCatastrophe
  , getNumRemovedByOccurrence
  , getNumRemovedByDisaster
  ) where

import           Data.List                       as List
import           Data.Maybe                      (fromJust)
import qualified Data.Vector                     as V
import qualified Data.Vector.Generic             as G
import           Epidemic
import           Epidemic.Types.Events           (EpidemicEvent (..))
import           Epidemic.Types.Parameter
import           Epidemic.Types.Population
import           Epidemic.Types.Simulation       (SimulationConfiguration (..),
                                                  SimulationRandEvent (..),
                                                  TerminationHandler (..))
import           Epidemic.Types.Time             (AbsoluteTime (..),
                                                  TimeDelta (..), Timed (..),
                                                  allTimes, asTimed,
                                                  cadlagValue, maybeNextTimed)
import           Epidemic.Utility
import           System.Random.MWC
import           System.Random.MWC.Distributions (bernoulli, categorical)

data InhomBDSCODRates =
  InhomBDSCODRates
    { InhomBDSCODRates -> Timed Rate
irBirthRate       :: Timed Rate
    , InhomBDSCODRates -> Timed Rate
irDeathRate       :: Timed Rate
    , InhomBDSCODRates -> Timed Rate
irSamplingRate    :: Timed Rate
    , InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: Timed Probability
    , InhomBDSCODRates -> Timed Rate
irOccurrenceRate  :: Timed Rate
    , InhomBDSCODRates -> Timed Rate
irDisasterSpec    :: Timed Probability
    }
  deriving (Int -> InhomBDSCODRates -> ShowS
[InhomBDSCODRates] -> ShowS
InhomBDSCODRates -> String
(Int -> InhomBDSCODRates -> ShowS)
-> (InhomBDSCODRates -> String)
-> ([InhomBDSCODRates] -> ShowS)
-> Show InhomBDSCODRates
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [InhomBDSCODRates] -> ShowS
$cshowList :: [InhomBDSCODRates] -> ShowS
show :: InhomBDSCODRates -> String
$cshow :: InhomBDSCODRates -> String
showsPrec :: Int -> InhomBDSCODRates -> ShowS
$cshowsPrec :: Int -> InhomBDSCODRates -> ShowS
Show, InhomBDSCODRates -> InhomBDSCODRates -> Bool
(InhomBDSCODRates -> InhomBDSCODRates -> Bool)
-> (InhomBDSCODRates -> InhomBDSCODRates -> Bool)
-> Eq InhomBDSCODRates
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: InhomBDSCODRates -> InhomBDSCODRates -> Bool
$c/= :: InhomBDSCODRates -> InhomBDSCODRates -> Bool
== :: InhomBDSCODRates -> InhomBDSCODRates -> Bool
$c== :: InhomBDSCODRates -> InhomBDSCODRates -> Bool
Eq)

-- | The population in which the epidemic occurs. This includes information
-- about the number of people that have previously been infected and
-- subsequently removed.
data InhomBDSCODPop =
  InhomBDSCODPop
  { InhomBDSCODPop -> People
ipInfectedPeople          :: People
  , InhomBDSCODPop -> Int
ipNumRemovedByDeath       :: Int
  , InhomBDSCODPop -> Int
ipNumRemovedBySampling    :: Int
  , InhomBDSCODPop -> Int
ipNumRemovedByCatastrophe :: Int
  , InhomBDSCODPop -> Int
ipNumRemovedByOccurrence  :: Int
  , InhomBDSCODPop -> Int
ipNumRemovedByDisaster    :: Int
  } deriving (Int -> InhomBDSCODPop -> ShowS
[InhomBDSCODPop] -> ShowS
InhomBDSCODPop -> String
(Int -> InhomBDSCODPop -> ShowS)
-> (InhomBDSCODPop -> String)
-> ([InhomBDSCODPop] -> ShowS)
-> Show InhomBDSCODPop
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [InhomBDSCODPop] -> ShowS
$cshowList :: [InhomBDSCODPop] -> ShowS
show :: InhomBDSCODPop -> String
$cshow :: InhomBDSCODPop -> String
showsPrec :: Int -> InhomBDSCODPop -> ShowS
$cshowsPrec :: Int -> InhomBDSCODPop -> ShowS
Show)

getNumRemovedByDeath :: InhomBDSCODPop -> Int
getNumRemovedByDeath :: InhomBDSCODPop -> Int
getNumRemovedByDeath = InhomBDSCODPop -> Int
ipNumRemovedByDeath

getNumRemovedBySampling :: InhomBDSCODPop -> Int
getNumRemovedBySampling :: InhomBDSCODPop -> Int
getNumRemovedBySampling = InhomBDSCODPop -> Int
ipNumRemovedBySampling

getNumRemovedByCatastrophe :: InhomBDSCODPop -> Int
getNumRemovedByCatastrophe :: InhomBDSCODPop -> Int
getNumRemovedByCatastrophe = InhomBDSCODPop -> Int
ipNumRemovedByCatastrophe

getNumRemovedByOccurrence :: InhomBDSCODPop -> Int
getNumRemovedByOccurrence :: InhomBDSCODPop -> Int
getNumRemovedByOccurrence = InhomBDSCODPop -> Int
ipNumRemovedByOccurrence

getNumRemovedByDisaster :: InhomBDSCODPop -> Int
getNumRemovedByDisaster :: InhomBDSCODPop -> Int
getNumRemovedByDisaster = InhomBDSCODPop -> Int
ipNumRemovedByDisaster

instance ModelParameters InhomBDSCODRates InhomBDSCODPop where
  rNaught :: InhomBDSCODPop -> InhomBDSCODRates -> AbsoluteTime -> Maybe Rate
rNaught InhomBDSCODPop
_ InhomBDSCODRates {Timed Rate
irDisasterSpec :: Timed Rate
irOccurrenceRate :: Timed Rate
irCatastropheSpec :: Timed Rate
irSamplingRate :: Timed Rate
irDeathRate :: Timed Rate
irBirthRate :: Timed Rate
irDisasterSpec :: InhomBDSCODRates -> Timed Rate
irOccurrenceRate :: InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: InhomBDSCODRates -> Timed Rate
irSamplingRate :: InhomBDSCODRates -> Timed Rate
irDeathRate :: InhomBDSCODRates -> Timed Rate
irBirthRate :: InhomBDSCODRates -> Timed Rate
..} AbsoluteTime
time =
    do
      Rate
birthRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irBirthRate AbsoluteTime
time
      Rate
deathRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irDeathRate AbsoluteTime
time
      Rate
sampleRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irSamplingRate AbsoluteTime
time
      Rate
occurrenceRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irOccurrenceRate AbsoluteTime
time
      Rate -> Maybe Rate
forall a. a -> Maybe a
Just (Rate -> Maybe Rate) -> Rate -> Maybe Rate
forall a b. (a -> b) -> a -> b
$ Rate
birthRate Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ (Rate
deathRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
sampleRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
occurrenceRate)
  eventRate :: InhomBDSCODPop -> InhomBDSCODRates -> AbsoluteTime -> Maybe Rate
eventRate InhomBDSCODPop
_ InhomBDSCODRates {Timed Rate
irDisasterSpec :: Timed Rate
irOccurrenceRate :: Timed Rate
irCatastropheSpec :: Timed Rate
irSamplingRate :: Timed Rate
irDeathRate :: Timed Rate
irBirthRate :: Timed Rate
irDisasterSpec :: InhomBDSCODRates -> Timed Rate
irOccurrenceRate :: InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: InhomBDSCODRates -> Timed Rate
irSamplingRate :: InhomBDSCODRates -> Timed Rate
irDeathRate :: InhomBDSCODRates -> Timed Rate
irBirthRate :: InhomBDSCODRates -> Timed Rate
..} AbsoluteTime
time =
    do
      Rate
birthRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irBirthRate AbsoluteTime
time
      Rate
deathRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irDeathRate AbsoluteTime
time
      Rate
sampleRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irSamplingRate AbsoluteTime
time
      Rate
occurrenceRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irOccurrenceRate AbsoluteTime
time
      Rate -> Maybe Rate
forall a. a -> Maybe a
Just (Rate -> Maybe Rate) -> Rate -> Maybe Rate
forall a b. (a -> b) -> a -> b
$ Rate
birthRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
deathRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
sampleRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
occurrenceRate
  birthProb :: InhomBDSCODPop -> InhomBDSCODRates -> AbsoluteTime -> Maybe Rate
birthProb InhomBDSCODPop
p inhomRates :: InhomBDSCODRates
inhomRates@InhomBDSCODRates {Timed Rate
irDisasterSpec :: Timed Rate
irOccurrenceRate :: Timed Rate
irCatastropheSpec :: Timed Rate
irSamplingRate :: Timed Rate
irDeathRate :: Timed Rate
irBirthRate :: Timed Rate
irDisasterSpec :: InhomBDSCODRates -> Timed Rate
irOccurrenceRate :: InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: InhomBDSCODRates -> Timed Rate
irSamplingRate :: InhomBDSCODRates -> Timed Rate
irDeathRate :: InhomBDSCODRates -> Timed Rate
irBirthRate :: InhomBDSCODRates -> Timed Rate
..} AbsoluteTime
time =
    do
      Rate
birthRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irBirthRate AbsoluteTime
time
      Rate
totalEventRate <- InhomBDSCODPop -> InhomBDSCODRates -> AbsoluteTime -> Maybe Rate
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Rate
eventRate InhomBDSCODPop
p InhomBDSCODRates
inhomRates AbsoluteTime
time
      Rate -> Maybe Rate
forall a. a -> Maybe a
Just (Rate -> Maybe Rate) -> Rate -> Maybe Rate
forall a b. (a -> b) -> a -> b
$ Rate
birthRate Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ Rate
totalEventRate
  eventWeights :: InhomBDSCODPop
-> InhomBDSCODRates -> AbsoluteTime -> Maybe (Vector Rate)
eventWeights InhomBDSCODPop
_ InhomBDSCODRates{Timed Rate
irDisasterSpec :: Timed Rate
irOccurrenceRate :: Timed Rate
irCatastropheSpec :: Timed Rate
irSamplingRate :: Timed Rate
irDeathRate :: Timed Rate
irBirthRate :: Timed Rate
irDisasterSpec :: InhomBDSCODRates -> Timed Rate
irOccurrenceRate :: InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: InhomBDSCODRates -> Timed Rate
irSamplingRate :: InhomBDSCODRates -> Timed Rate
irDeathRate :: InhomBDSCODRates -> Timed Rate
irBirthRate :: InhomBDSCODRates -> Timed Rate
..} AbsoluteTime
time =
    do
      Rate
birthRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irBirthRate AbsoluteTime
time
      Rate
deathRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irDeathRate AbsoluteTime
time
      Rate
sampleRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irSamplingRate AbsoluteTime
time
      Rate
occurrenceRate <- Timed Rate -> AbsoluteTime -> Maybe Rate
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Rate
irOccurrenceRate AbsoluteTime
time
      Vector Rate -> Maybe (Vector Rate)
forall a. a -> Maybe a
Just (Vector Rate -> Maybe (Vector Rate))
-> Vector Rate -> Maybe (Vector Rate)
forall a b. (a -> b) -> a -> b
$ [Rate] -> Vector Rate
forall a. [a] -> Vector a
V.fromList [Rate
birthRate, Rate
deathRate, Rate
sampleRate, Rate
occurrenceRate]

instance Population InhomBDSCODPop where
  susceptiblePeople :: InhomBDSCODPop -> Maybe People
susceptiblePeople InhomBDSCODPop
_ = Maybe People
forall a. Maybe a
Nothing
  infectiousPeople :: InhomBDSCODPop -> Maybe People
infectiousPeople = People -> Maybe People
forall (f :: * -> *) a. Applicative f => a -> f a
pure (People -> Maybe People)
-> (InhomBDSCODPop -> People) -> InhomBDSCODPop -> Maybe People
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InhomBDSCODPop -> People
ipInfectedPeople
  removedPeople :: InhomBDSCODPop -> Maybe People
removedPeople InhomBDSCODPop
_ = Maybe People
forall a. Maybe a
Nothing
  isInfected :: InhomBDSCODPop -> Bool
isInfected = Bool -> Bool
not (Bool -> Bool)
-> (InhomBDSCODPop -> Bool) -> InhomBDSCODPop -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. People -> Bool
nullPeople (People -> Bool)
-> (InhomBDSCODPop -> People) -> InhomBDSCODPop -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InhomBDSCODPop -> People
ipInfectedPeople

-- | Configuration for the simulation of the inhomogeneous rates BDSCOD process.
configuration ::
     TimeDelta -- ^ Duration of the simulation after starting at time 0.
  -> Bool -- ^ condition upon at least two sequenced samples.
  -> Maybe (InhomBDSCODPop -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling.
  -> ( [(AbsoluteTime, Rate)]
     , [(AbsoluteTime, Rate)]
     , [(AbsoluteTime, Rate)]
     , [(AbsoluteTime, Probability)]
     , [(AbsoluteTime, Rate)]
     , [(AbsoluteTime, Probability)])
  -> Maybe (SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s)
configuration :: TimeDelta
-> Bool
-> Maybe (InhomBDSCODPop -> Bool, [EpidemicEvent] -> s)
-> ([(AbsoluteTime, Rate)], [(AbsoluteTime, Rate)],
    [(AbsoluteTime, Rate)], [(AbsoluteTime, Rate)],
    [(AbsoluteTime, Rate)], [(AbsoluteTime, Rate)])
-> Maybe
     (SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s)
configuration TimeDelta
maxTime Bool
atLeastCherry Maybe (InhomBDSCODPop -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs ([(AbsoluteTime, Rate)]
tBirthRate, [(AbsoluteTime, Rate)]
tDeathRate, [(AbsoluteTime, Rate)]
tSampleRate, [(AbsoluteTime, Rate)]
cSpec, [(AbsoluteTime, Rate)]
tOccurrenceRate, [(AbsoluteTime, Rate)]
dSpec) =
  let (Person
seedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
initialIdentifier
      bdscodPop :: InhomBDSCODPop
bdscodPop = InhomBDSCODPop :: People -> Int -> Int -> Int -> Int -> Int -> InhomBDSCODPop
InhomBDSCODPop { ipInfectedPeople :: People
ipInfectedPeople = [Person] -> People
asPeople [Person
seedPerson]
                                 , ipNumRemovedByDeath :: Int
ipNumRemovedByDeath = Int
0
                                 , ipNumRemovedBySampling :: Int
ipNumRemovedBySampling = Int
0
                                 , ipNumRemovedByCatastrophe :: Int
ipNumRemovedByCatastrophe = Int
0
                                 , ipNumRemovedByOccurrence :: Int
ipNumRemovedByOccurrence = Int
0
                                 , ipNumRemovedByDisaster :: Int
ipNumRemovedByDisaster = Int
0 }
   in do Timed Rate
timedBirthRate <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
tBirthRate
         Timed Rate
timedDeathRate <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
tDeathRate
         Timed Rate
timedSamplingRate <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
tSampleRate
         Timed Rate
catastropheSpec <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
cSpec
         Timed Rate
timedOccurrenceRate <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
tOccurrenceRate
         Timed Rate
disasterSpec <- [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
dSpec
         let irVal :: InhomBDSCODRates
irVal =
               Timed Rate
-> Timed Rate
-> Timed Rate
-> Timed Rate
-> Timed Rate
-> Timed Rate
-> InhomBDSCODRates
InhomBDSCODRates
                 Timed Rate
timedBirthRate
                 Timed Rate
timedDeathRate
                 Timed Rate
timedSamplingRate
                 Timed Rate
catastropheSpec
                 Timed Rate
timedOccurrenceRate
                 Timed Rate
disasterSpec
             termHandler :: Maybe (TerminationHandler InhomBDSCODPop s)
termHandler = do (InhomBDSCODPop -> Bool
f1, [EpidemicEvent] -> s
f2) <- Maybe (InhomBDSCODPop -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs
                              TerminationHandler InhomBDSCODPop s
-> Maybe (TerminationHandler InhomBDSCODPop s)
forall (m :: * -> *) a. Monad m => a -> m a
return (TerminationHandler InhomBDSCODPop s
 -> Maybe (TerminationHandler InhomBDSCODPop s))
-> TerminationHandler InhomBDSCODPop s
-> Maybe (TerminationHandler InhomBDSCODPop s)
forall a b. (a -> b) -> a -> b
$ (InhomBDSCODPop -> Bool)
-> ([EpidemicEvent] -> s) -> TerminationHandler InhomBDSCODPop s
forall b c.
Population b =>
(b -> Bool) -> ([EpidemicEvent] -> c) -> TerminationHandler b c
TerminationHandler InhomBDSCODPop -> Bool
f1 [EpidemicEvent] -> s
f2
         if TimeDelta
maxTime TimeDelta -> TimeDelta -> Bool
forall a. Ord a => a -> a -> Bool
> Rate -> TimeDelta
TimeDelta Rate
0
           then SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s
-> Maybe
     (SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s)
forall a. a -> Maybe a
Just
                  (InhomBDSCODRates
-> InhomBDSCODPop
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler InhomBDSCODPop s)
-> Bool
-> SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s
forall r p s.
r
-> p
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler p s)
-> Bool
-> SimulationConfiguration r p s
SimulationConfiguration
                     InhomBDSCODRates
irVal
                     InhomBDSCODPop
bdscodPop
                     Identifier
newId
                     (Rate -> AbsoluteTime
AbsoluteTime Rate
0)
                     TimeDelta
maxTime
                     Maybe (TerminationHandler InhomBDSCODPop s)
termHandler
                     Bool
atLeastCherry)
           else Maybe (SimulationConfiguration InhomBDSCODRates InhomBDSCODPop s)
forall a. Maybe a
Nothing

-- | A random event and the state afterwards
randomEvent :: SimulationRandEvent InhomBDSCODRates InhomBDSCODPop
randomEvent :: SimulationRandEvent InhomBDSCODRates InhomBDSCODPop
randomEvent = (InhomBDSCODRates
 -> AbsoluteTime
 -> InhomBDSCODPop
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier))
-> SimulationRandEvent InhomBDSCODRates InhomBDSCODPop
forall a b.
(ModelParameters a b, Population b) =>
(a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent InhomBDSCODRates
-> AbsoluteTime
-> InhomBDSCODPop
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
randomEvent'

randomEvent' ::
     InhomBDSCODRates
  -> AbsoluteTime -- ^ the current time
  -> InhomBDSCODPop -- ^ the population
  -> Identifier -- ^ current identifier
  -> GenIO -- ^ PRNG
  -> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
randomEvent' :: InhomBDSCODRates
-> AbsoluteTime
-> InhomBDSCODPop
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
randomEvent' inhomRates :: InhomBDSCODRates
inhomRates@InhomBDSCODRates {Timed Rate
irDisasterSpec :: Timed Rate
irOccurrenceRate :: Timed Rate
irCatastropheSpec :: Timed Rate
irSamplingRate :: Timed Rate
irDeathRate :: Timed Rate
irBirthRate :: Timed Rate
irDisasterSpec :: InhomBDSCODRates -> Timed Rate
irOccurrenceRate :: InhomBDSCODRates -> Timed Rate
irCatastropheSpec :: InhomBDSCODRates -> Timed Rate
irSamplingRate :: InhomBDSCODRates -> Timed Rate
irDeathRate :: InhomBDSCODRates -> Timed Rate
irBirthRate :: InhomBDSCODRates -> Timed Rate
..} AbsoluteTime
currTime InhomBDSCODPop
currPop Identifier
currId GenIO
gen =
  let (Just People
people) = InhomBDSCODPop -> Maybe People
forall a. Population a => a -> Maybe People
infectiousPeople InhomBDSCODPop
currPop
      popSize :: Rate
popSize = Int -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Rate) -> Int -> Rate
forall a b. (a -> b) -> a -> b
$ People -> Int
numPeople People
people :: Double
      weightVecFunc :: AbsoluteTime -> Maybe (Vector Rate)
weightVecFunc = InhomBDSCODPop
-> InhomBDSCODRates -> AbsoluteTime -> Maybe (Vector Rate)
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe (Vector Rate)
eventWeights InhomBDSCODPop
currPop InhomBDSCODRates
inhomRates
      -- we need a new step function to account for the population size.
      (Just Timed Rate
stepFunction) =
        [(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed
          [ (AbsoluteTime
t, Rate
popSize Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
* Maybe Rate -> Rate
forall a. HasCallStack => Maybe a -> a
fromJust (InhomBDSCODPop -> InhomBDSCODRates -> AbsoluteTime -> Maybe Rate
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Rate
eventRate InhomBDSCODPop
currPop InhomBDSCODRates
inhomRates AbsoluteTime
t))
          | AbsoluteTime
t <- [AbsoluteTime] -> [AbsoluteTime]
forall a. Ord a => [a] -> [a]
List.sort ([AbsoluteTime] -> [AbsoluteTime])
-> [AbsoluteTime] -> [AbsoluteTime]
forall a b. (a -> b) -> a -> b
$ (Timed Rate -> [AbsoluteTime]) -> [Timed Rate] -> [AbsoluteTime]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Timed Rate -> [AbsoluteTime]
forall a. Timed a -> [AbsoluteTime]
allTimes [Timed Rate
irBirthRate, Timed Rate
irDeathRate, Timed Rate
irSamplingRate, Timed Rate
irOccurrenceRate]
          ]
   in do (Just AbsoluteTime
newEventTime) <- Timed Rate -> AbsoluteTime -> GenIO -> IO (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
Timed Rate
-> AbsoluteTime -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
inhomExponential Timed Rate
stepFunction AbsoluteTime
currTime GenIO
gen
         if AbsoluteTime -> AbsoluteTime -> Timed Rate -> Bool
noScheduledEvent AbsoluteTime
currTime AbsoluteTime
newEventTime (Timed Rate
irCatastropheSpec Timed Rate -> Timed Rate -> Timed Rate
forall a. Semigroup a => a -> a -> a
<> Timed Rate
irDisasterSpec)
           then do
             Int
eventIx <- Vector Rate -> Gen RealWorld -> IO Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Rate) =>
v Rate -> g -> m Int
categorical (Maybe (Vector Rate) -> Vector Rate
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Vector Rate) -> Vector Rate)
-> Maybe (Vector Rate) -> Vector Rate
forall a b. (a -> b) -> a -> b
$ AbsoluteTime -> Maybe (Vector Rate)
weightVecFunc AbsoluteTime
newEventTime) Gen RealWorld
GenIO
gen
             (Person
selectedPerson, People
unselectedPeople) <- People -> GenIO -> IO (Person, People)
randomPerson People
people GenIO
gen
             (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return ((AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
 -> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier))
-> (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall a b. (a -> b) -> a -> b
$
               case Int
eventIx of
                 Int
0 ->
                   ( AbsoluteTime
newEventTime
                   , AbsoluteTime -> Person -> Person -> EpidemicEvent
Infection AbsoluteTime
newEventTime Person
selectedPerson Person
birthedPerson
                   , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = Person -> People -> People
addPerson Person
birthedPerson People
people}
                   , Identifier
newId)
                   where (Person
birthedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
currId
                 Int
1 -> let currNumDeaths :: Int
currNumDeaths = InhomBDSCODPop -> Int
ipNumRemovedByDeath InhomBDSCODPop
currPop
                      in ( AbsoluteTime
newEventTime
                         , AbsoluteTime -> Person -> EpidemicEvent
Removal AbsoluteTime
newEventTime Person
selectedPerson
                         , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = People
unselectedPeople
                                   , ipNumRemovedByDeath :: Int
ipNumRemovedByDeath = Int
currNumDeaths Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 }
                         , Identifier
currId )
                 Int
2 -> let currNumSampled :: Int
currNumSampled = InhomBDSCODPop -> Int
ipNumRemovedBySampling InhomBDSCODPop
currPop
                      in ( AbsoluteTime
newEventTime
                         , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
selectedPerson Bool
True
                         , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = People
unselectedPeople
                                   , ipNumRemovedBySampling :: Int
ipNumRemovedBySampling = Int
currNumSampled Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 }
                         , Identifier
currId)
                 Int
3 -> let currNumOccurrence :: Int
currNumOccurrence = InhomBDSCODPop -> Int
ipNumRemovedByOccurrence InhomBDSCODPop
currPop
                      in ( AbsoluteTime
newEventTime
                         , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
selectedPerson Bool
False
                         , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = People
unselectedPeople
                                   , ipNumRemovedByOccurrence :: Int
ipNumRemovedByOccurrence = Int
currNumOccurrence Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}
                         , Identifier
currId)
                 Int
_ -> String -> (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall a. HasCallStack => String -> a
error String
"no birth, death, sampling, or occurrence event selected."
           else case Timed Rate
-> Timed Rate
-> AbsoluteTime
-> Maybe (AbsoluteTime, Either Rate Rate)
forall a b.
Timed a
-> Timed b -> AbsoluteTime -> Maybe (AbsoluteTime, Either a b)
maybeNextTimed Timed Rate
irCatastropheSpec Timed Rate
irDisasterSpec AbsoluteTime
currTime of
                  Just (AbsoluteTime
disastTime, Right Rate
disastProb) ->
                    do (EpidemicEvent
disastEvent, InhomBDSCODPop
postDisastPop) <-
                         (AbsoluteTime, Rate)
-> InhomBDSCODPop -> GenIO -> IO (EpidemicEvent, InhomBDSCODPop)
randomDisasterEvent
                         (AbsoluteTime
disastTime, Rate
disastProb)
                         InhomBDSCODPop
currPop
                         GenIO
gen
                       (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
disastTime, EpidemicEvent
disastEvent, InhomBDSCODPop
postDisastPop, Identifier
currId)
                  Just (AbsoluteTime
catastTime, Left Rate
catastProb) ->
                    do (EpidemicEvent
catastEvent, InhomBDSCODPop
postCatastPop) <-
                         (AbsoluteTime, Rate)
-> InhomBDSCODPop -> GenIO -> IO (EpidemicEvent, InhomBDSCODPop)
randomCatastropheEvent
                         (AbsoluteTime
catastTime, Rate
catastProb)
                         InhomBDSCODPop
currPop
                         GenIO
gen
                       (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
catastTime, EpidemicEvent
catastEvent, InhomBDSCODPop
postCatastPop, Identifier
currId)
                  Maybe (AbsoluteTime, Either Rate Rate)
Nothing -> String
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSCODPop, Identifier)
forall a. HasCallStack => String -> a
error String
"Missing a next scheduled event when there should be one."

-- | Return a randomly sampled Catastrophe event and the population after that
-- event has occurred.
randomCatastropheEvent ::
     (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the catastrophe
  -> InhomBDSCODPop -- ^ The state of the population prior to the catastrophe
  -> GenIO
  -> IO (EpidemicEvent, InhomBDSCODPop)
randomCatastropheEvent :: (AbsoluteTime, Rate)
-> InhomBDSCODPop -> GenIO -> IO (EpidemicEvent, InhomBDSCODPop)
randomCatastropheEvent (AbsoluteTime
catastTime, Rate
rhoProb) InhomBDSCODPop
currPop GenIO
gen =
  let (Just (People Vector Person
currPeople)) = InhomBDSCODPop -> Maybe People
forall a. Population a => a -> Maybe People
infectiousPeople InhomBDSCODPop
currPop
  in do Vector Bool
rhoBernoullis <- Int -> IO Bool -> IO (Vector Bool)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> m a -> m (v a)
G.replicateM (Vector Person -> Int
forall a. Vector a -> Int
V.length Vector Person
currPeople) (Rate -> Gen RealWorld -> IO Bool
forall g (m :: * -> *). StatefulGen g m => Rate -> g -> m Bool
bernoulli Rate
rhoProb Gen RealWorld
GenIO
gen)
        let filterZip :: ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (a, b) -> Bool
predicate Vector a
a Vector b
b = (Vector a, Vector b) -> Vector a
forall a b. (a, b) -> a
fst ((Vector a, Vector b) -> Vector a)
-> (Vector (a, b) -> (Vector a, Vector b))
-> Vector (a, b)
-> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (a, b) -> (Vector a, Vector b)
forall a b. Vector (a, b) -> (Vector a, Vector b)
V.unzip (Vector (a, b) -> (Vector a, Vector b))
-> (Vector (a, b) -> Vector (a, b))
-> Vector (a, b)
-> (Vector a, Vector b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> Bool) -> Vector (a, b) -> Vector (a, b)
forall a. (a -> Bool) -> Vector a -> Vector a
V.filter (a, b) -> Bool
predicate (Vector (a, b) -> Vector a) -> Vector (a, b) -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector b -> Vector (a, b)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector a
a Vector b
b
            sampledPeople :: People
sampledPeople = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ ((Person, Bool) -> Bool)
-> Vector Person -> Vector Bool -> Vector Person
forall a b. ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (Person, Bool) -> Bool
forall a b. (a, b) -> b
snd Vector Person
currPeople Vector Bool
rhoBernoullis
            unsampledPeople :: People
unsampledPeople = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ ((Person, Bool) -> Bool)
-> Vector Person -> Vector Bool -> Vector Person
forall a b. ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (Bool -> Bool
not (Bool -> Bool)
-> ((Person, Bool) -> Bool) -> (Person, Bool) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Person, Bool) -> Bool
forall a b. (a, b) -> b
snd) Vector Person
currPeople Vector Bool
rhoBernoullis
            currNumCatastrophe :: Int
currNumCatastrophe = InhomBDSCODPop -> Int
ipNumRemovedByCatastrophe InhomBDSCODPop
currPop
         in (EpidemicEvent, InhomBDSCODPop)
-> IO (EpidemicEvent, InhomBDSCODPop)
forall (m :: * -> *) a. Monad m => a -> m a
return ( AbsoluteTime -> People -> Bool -> EpidemicEvent
PopulationSample AbsoluteTime
catastTime People
sampledPeople Bool
True
                   , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = People
unsampledPeople
                             , ipNumRemovedByCatastrophe :: Int
ipNumRemovedByCatastrophe = Int
currNumCatastrophe Int -> Int -> Int
forall a. Num a => a -> a -> a
+ People -> Int
numPeople People
sampledPeople })

-- | Return a randomly sampled Disaster event and the population after that
-- event has occurred.
randomDisasterEvent ::
     (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the disaster
  -> InhomBDSCODPop -- ^ The state of the population prior to the disaster
  -> GenIO
  -> IO (EpidemicEvent, InhomBDSCODPop)
randomDisasterEvent :: (AbsoluteTime, Rate)
-> InhomBDSCODPop -> GenIO -> IO (EpidemicEvent, InhomBDSCODPop)
randomDisasterEvent (AbsoluteTime
disastTime, Rate
nuProb) InhomBDSCODPop
currPop GenIO
gen = do
  let (Just (People Vector Person
currPeople)) = InhomBDSCODPop -> Maybe People
forall a. Population a => a -> Maybe People
infectiousPeople InhomBDSCODPop
currPop
  Vector Bool
nuBernoullis <- Int -> IO Bool -> IO (Vector Bool)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> m a -> m (v a)
G.replicateM (Vector Person -> Int
forall a. Vector a -> Int
V.length Vector Person
currPeople) (Rate -> Gen RealWorld -> IO Bool
forall g (m :: * -> *). StatefulGen g m => Rate -> g -> m Bool
bernoulli Rate
nuProb Gen RealWorld
GenIO
gen)
  let filterZip :: ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (a, b) -> Bool
predicate Vector a
a Vector b
b = (Vector a, Vector b) -> Vector a
forall a b. (a, b) -> a
fst ((Vector a, Vector b) -> Vector a)
-> (Vector (a, b) -> (Vector a, Vector b))
-> Vector (a, b)
-> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (a, b) -> (Vector a, Vector b)
forall a b. Vector (a, b) -> (Vector a, Vector b)
V.unzip (Vector (a, b) -> (Vector a, Vector b))
-> (Vector (a, b) -> Vector (a, b))
-> Vector (a, b)
-> (Vector a, Vector b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> Bool) -> Vector (a, b) -> Vector (a, b)
forall a. (a -> Bool) -> Vector a -> Vector a
V.filter (a, b) -> Bool
predicate (Vector (a, b) -> Vector a) -> Vector (a, b) -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector b -> Vector (a, b)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector a
a Vector b
b
      sampledPeople :: People
sampledPeople = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ ((Person, Bool) -> Bool)
-> Vector Person -> Vector Bool -> Vector Person
forall a b. ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (Person, Bool) -> Bool
forall a b. (a, b) -> b
snd Vector Person
currPeople Vector Bool
nuBernoullis
      unsampledPeople :: People
unsampledPeople = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ ((Person, Bool) -> Bool)
-> Vector Person -> Vector Bool -> Vector Person
forall a b. ((a, b) -> Bool) -> Vector a -> Vector b -> Vector a
filterZip (Bool -> Bool
not (Bool -> Bool)
-> ((Person, Bool) -> Bool) -> (Person, Bool) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Person, Bool) -> Bool
forall a b. (a, b) -> b
snd) Vector Person
currPeople Vector Bool
nuBernoullis
      currNumDisaster :: Int
currNumDisaster = InhomBDSCODPop -> Int
ipNumRemovedByDisaster InhomBDSCODPop
currPop
   in (EpidemicEvent, InhomBDSCODPop)
-> IO (EpidemicEvent, InhomBDSCODPop)
forall (m :: * -> *) a. Monad m => a -> m a
return ( AbsoluteTime -> People -> Bool -> EpidemicEvent
PopulationSample AbsoluteTime
disastTime People
sampledPeople Bool
False
             , InhomBDSCODPop
currPop { ipInfectedPeople :: People
ipInfectedPeople = People
unsampledPeople
                       , ipNumRemovedByDisaster :: Int
ipNumRemovedByDisaster = Int
currNumDisaster Int -> Int -> Int
forall a. Num a => a -> a -> a
+ People -> Int
numPeople People
sampledPeople })