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

module Epidemic.Model.LogisticBDSD
  ( configuration
  , randomEvent
  , LogisticBDSDParameters(..)
  , LogisticBDSDPopulation(..)
  ) where

import Data.Maybe (fromJust)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import Epidemic (firstScheduled, noScheduledEvent)
import Epidemic.Types.Events (EpidemicEvent(..))
import Epidemic.Types.Time
  ( AbsoluteTime(..)
  , Timed(..)
  , TimeDelta(..)
  , asTimed
  , timeAfterDelta
  )
import Epidemic.Types.Parameter
  (  ModelParameters(..)
  , Probability
  , Rate
  )
import Epidemic.Types.Population
  ( Identifier(..)
  , People(..)
  , Population(..)
  , addPerson
  , nullPeople
  , numPeople
  )
import Epidemic.Types.Simulation
  ( SimulationConfiguration(..)
  , SimulationRandEvent(..), TerminationHandler(..)
  )
import Epidemic.Utility
  ( initialIdentifier
  , maybeToRight
  , newPerson
  , randomPerson
  )
import System.Random.MWC (GenIO)
import System.Random.MWC.Distributions (bernoulli, categorical, exponential)

-- | The parameters of the logistic-BDSD process. This process allows for
-- infections, removals, sampling and disasters.
data LogisticBDSDParameters =
  LogisticBDSDParameters
    { LogisticBDSDParameters -> Rate
paramsBirthRate :: Rate
    , LogisticBDSDParameters -> Int
paramsCapacity :: Int
    , LogisticBDSDParameters -> Rate
paramsDeathRate :: Rate
    , LogisticBDSDParameters -> Rate
paramsSamplingRate :: Rate
    , LogisticBDSDParameters -> Timed Rate
paramsDisasters :: Timed Probability
    }
  deriving (Int -> LogisticBDSDParameters -> ShowS
[LogisticBDSDParameters] -> ShowS
LogisticBDSDParameters -> String
(Int -> LogisticBDSDParameters -> ShowS)
-> (LogisticBDSDParameters -> String)
-> ([LogisticBDSDParameters] -> ShowS)
-> Show LogisticBDSDParameters
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LogisticBDSDParameters] -> ShowS
$cshowList :: [LogisticBDSDParameters] -> ShowS
show :: LogisticBDSDParameters -> String
$cshow :: LogisticBDSDParameters -> String
showsPrec :: Int -> LogisticBDSDParameters -> ShowS
$cshowsPrec :: Int -> LogisticBDSDParameters -> ShowS
Show)

newtype LogisticBDSDPopulation =
  LogisticBDSDPopulation People
  deriving (Int -> LogisticBDSDPopulation -> ShowS
[LogisticBDSDPopulation] -> ShowS
LogisticBDSDPopulation -> String
(Int -> LogisticBDSDPopulation -> ShowS)
-> (LogisticBDSDPopulation -> String)
-> ([LogisticBDSDPopulation] -> ShowS)
-> Show LogisticBDSDPopulation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LogisticBDSDPopulation] -> ShowS
$cshowList :: [LogisticBDSDPopulation] -> ShowS
show :: LogisticBDSDPopulation -> String
$cshow :: LogisticBDSDPopulation -> String
showsPrec :: Int -> LogisticBDSDPopulation -> ShowS
$cshowsPrec :: Int -> LogisticBDSDPopulation -> ShowS
Show)

-- | The per lineage birth rate accounting for the population size.
logisticBirthRate :: LogisticBDSDParameters -> LogisticBDSDPopulation -> Rate
logisticBirthRate :: LogisticBDSDParameters -> LogisticBDSDPopulation -> Rate
logisticBirthRate LogisticBDSDParameters {Rate
Int
Timed Rate
paramsDisasters :: Timed Rate
paramsSamplingRate :: Rate
paramsDeathRate :: Rate
paramsCapacity :: Int
paramsBirthRate :: Rate
paramsDisasters :: LogisticBDSDParameters -> Timed Rate
paramsSamplingRate :: LogisticBDSDParameters -> Rate
paramsDeathRate :: LogisticBDSDParameters -> Rate
paramsCapacity :: LogisticBDSDParameters -> Int
paramsBirthRate :: LogisticBDSDParameters -> Rate
..} (LogisticBDSDPopulation People
pop) =
  let propCapacity :: Rate
propCapacity = Int -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (People -> Int
numPeople People
pop) Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ Int -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
paramsCapacity
   in Rate
paramsBirthRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
* (Rate
1.0 Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
- Rate
propCapacity)

instance ModelParameters LogisticBDSDParameters LogisticBDSDPopulation where
  rNaught :: LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe Rate
rNaught LogisticBDSDPopulation
_ LogisticBDSDParameters
_ AbsoluteTime
_ = Maybe Rate
forall a. Maybe a
Nothing
  eventRate :: LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe Rate
eventRate (LogisticBDSDPopulation People
pop) LogisticBDSDParameters {Rate
Int
Timed Rate
paramsDisasters :: Timed Rate
paramsSamplingRate :: Rate
paramsDeathRate :: Rate
paramsCapacity :: Int
paramsBirthRate :: Rate
paramsDisasters :: LogisticBDSDParameters -> Timed Rate
paramsSamplingRate :: LogisticBDSDParameters -> Rate
paramsDeathRate :: LogisticBDSDParameters -> Rate
paramsCapacity :: LogisticBDSDParameters -> Int
paramsBirthRate :: LogisticBDSDParameters -> Rate
..} AbsoluteTime
_ =
    let propCapcity :: Rate
propCapcity = Int -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (People -> Int
numPeople People
pop) Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ Int -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
paramsCapacity
        br :: Rate
br = Rate
paramsBirthRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
* (Rate
1.0 Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
- Rate
propCapcity)
     in Rate -> Maybe Rate
forall a. a -> Maybe a
Just (Rate -> Maybe Rate) -> Rate -> Maybe Rate
forall a b. (a -> b) -> a -> b
$ Rate
br Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
paramsDeathRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
+ Rate
paramsSamplingRate
  birthProb :: LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe Rate
birthProb LogisticBDSDPopulation
lpop LogisticBDSDParameters
lparam AbsoluteTime
absTime = do
    Rate
er <- LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe Rate
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Rate
eventRate LogisticBDSDPopulation
lpop LogisticBDSDParameters
lparam AbsoluteTime
absTime
    Rate -> Maybe Rate
forall a. a -> Maybe a
Just (Rate -> Maybe Rate) -> Rate -> Maybe Rate
forall a b. (a -> b) -> a -> b
$ Rate
br Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ Rate
er
    where
      br :: Rate
br = LogisticBDSDParameters -> LogisticBDSDPopulation -> Rate
logisticBirthRate LogisticBDSDParameters
lparam LogisticBDSDPopulation
lpop
  eventWeights :: LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe (Vector Rate)
eventWeights LogisticBDSDPopulation
currPop params :: LogisticBDSDParameters
params@LogisticBDSDParameters {Rate
Int
Timed Rate
paramsDisasters :: Timed Rate
paramsSamplingRate :: Rate
paramsDeathRate :: Rate
paramsCapacity :: Int
paramsBirthRate :: Rate
paramsDisasters :: LogisticBDSDParameters -> Timed Rate
paramsSamplingRate :: LogisticBDSDParameters -> Rate
paramsDeathRate :: LogisticBDSDParameters -> Rate
paramsCapacity :: LogisticBDSDParameters -> Int
paramsBirthRate :: LogisticBDSDParameters -> Rate
..} AbsoluteTime
_ =
    let logisticBR :: Rate
logisticBR = LogisticBDSDParameters -> LogisticBDSDPopulation -> Rate
logisticBirthRate LogisticBDSDParameters
params LogisticBDSDPopulation
currPop
        in 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
logisticBR, Rate
paramsDeathRate, Rate
paramsSamplingRate]

instance Population LogisticBDSDPopulation where
  susceptiblePeople :: LogisticBDSDPopulation -> Maybe People
susceptiblePeople LogisticBDSDPopulation
_ = Maybe People
forall a. Maybe a
Nothing
  infectiousPeople :: LogisticBDSDPopulation -> Maybe People
infectiousPeople (LogisticBDSDPopulation People
people) = People -> Maybe People
forall a. a -> Maybe a
Just People
people
  removedPeople :: LogisticBDSDPopulation -> Maybe People
removedPeople LogisticBDSDPopulation
_ = Maybe People
forall a. Maybe a
Nothing
  isInfected :: LogisticBDSDPopulation -> Bool
isInfected (LogisticBDSDPopulation People
people) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ People -> Bool
nullPeople People
people

-- | Create an simulation configuration or return an error message if this is
-- not possible.
configuration ::
     TimeDelta
  -> Bool -- ^ condition upon at least two sequenced samples.
  -> Maybe (LogisticBDSDPopulation -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling.
  -> (Rate, Int, Rate, Rate, [(AbsoluteTime, Probability)])
  -> Either String (SimulationConfiguration LogisticBDSDParameters LogisticBDSDPopulation s)
configuration :: TimeDelta
-> Bool
-> Maybe (LogisticBDSDPopulation -> Bool, [EpidemicEvent] -> s)
-> (Rate, Int, Rate, Rate, [(AbsoluteTime, Rate)])
-> Either
     String
     (SimulationConfiguration
        LogisticBDSDParameters LogisticBDSDPopulation s)
configuration TimeDelta
simDuration Bool
atLeastCherry Maybe (LogisticBDSDPopulation -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs (Rate
birthRate, Int
capacity, Rate
deathRate, Rate
samplingRate, [(AbsoluteTime, Rate)]
disasterSpec)
  | [Rate] -> Rate
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Rate
birthRate, Rate
deathRate, Rate
samplingRate] Rate -> Rate -> Bool
forall a. Ord a => a -> a -> Bool
< Rate
0 =
    String
-> Either
     String
     (SimulationConfiguration
        LogisticBDSDParameters LogisticBDSDPopulation s)
forall a b. a -> Either a b
Left String
"negative rate provided"
  | Int
capacity Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String
-> Either
     String
     (SimulationConfiguration
        LogisticBDSDParameters LogisticBDSDPopulation s)
forall a b. a -> Either a b
Left String
"insufficient population capacity"
  | Bool
otherwise = do
    Timed Rate
disasterTP <-
      String -> Maybe (Timed Rate) -> Either String (Timed Rate)
forall a b. a -> Maybe b -> Either a b
maybeToRight
        String
"could not construct timed probability"
        ([(AbsoluteTime, Rate)] -> Maybe (Timed Rate)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Rate)]
disasterSpec)
    let logBDSDParams :: LogisticBDSDParameters
logBDSDParams =
          Rate -> Int -> Rate -> Rate -> Timed Rate -> LogisticBDSDParameters
LogisticBDSDParameters
            Rate
birthRate
            Int
capacity
            Rate
deathRate
            Rate
samplingRate
            Timed Rate
disasterTP
        (Person
seedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
initialIdentifier
        logBDSDPop :: LogisticBDSDPopulation
logBDSDPop = People -> LogisticBDSDPopulation
LogisticBDSDPopulation (Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ Person -> Vector Person
forall a. a -> Vector a
V.singleton Person
seedPerson)
        termHandler :: Maybe (TerminationHandler LogisticBDSDPopulation s)
termHandler = do (LogisticBDSDPopulation -> Bool
f1, [EpidemicEvent] -> s
f2) <- Maybe (LogisticBDSDPopulation -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs
                         TerminationHandler LogisticBDSDPopulation s
-> Maybe (TerminationHandler LogisticBDSDPopulation s)
forall (m :: * -> *) a. Monad m => a -> m a
return (TerminationHandler LogisticBDSDPopulation s
 -> Maybe (TerminationHandler LogisticBDSDPopulation s))
-> TerminationHandler LogisticBDSDPopulation s
-> Maybe (TerminationHandler LogisticBDSDPopulation s)
forall a b. (a -> b) -> a -> b
$ (LogisticBDSDPopulation -> Bool)
-> ([EpidemicEvent] -> s)
-> TerminationHandler LogisticBDSDPopulation s
forall b c.
Population b =>
(b -> Bool) -> ([EpidemicEvent] -> c) -> TerminationHandler b c
TerminationHandler LogisticBDSDPopulation -> Bool
f1 [EpidemicEvent] -> s
f2
     in SimulationConfiguration
  LogisticBDSDParameters LogisticBDSDPopulation s
-> Either
     String
     (SimulationConfiguration
        LogisticBDSDParameters LogisticBDSDPopulation s)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationConfiguration
   LogisticBDSDParameters LogisticBDSDPopulation s
 -> Either
      String
      (SimulationConfiguration
         LogisticBDSDParameters LogisticBDSDPopulation s))
-> SimulationConfiguration
     LogisticBDSDParameters LogisticBDSDPopulation s
-> Either
     String
     (SimulationConfiguration
        LogisticBDSDParameters LogisticBDSDPopulation s)
forall a b. (a -> b) -> a -> b
$
        LogisticBDSDParameters
-> LogisticBDSDPopulation
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler LogisticBDSDPopulation s)
-> Bool
-> SimulationConfiguration
     LogisticBDSDParameters LogisticBDSDPopulation s
forall r p s.
r
-> p
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler p s)
-> Bool
-> SimulationConfiguration r p s
SimulationConfiguration
          LogisticBDSDParameters
logBDSDParams
          LogisticBDSDPopulation
logBDSDPop
          Identifier
newId
          (Rate -> AbsoluteTime
AbsoluteTime Rate
0)
          TimeDelta
simDuration
          Maybe (TerminationHandler LogisticBDSDPopulation s)
termHandler
          Bool
atLeastCherry

-- | Defines how a single random event is simulated in this model.
randomEvent :: SimulationRandEvent LogisticBDSDParameters LogisticBDSDPopulation
randomEvent :: SimulationRandEvent LogisticBDSDParameters LogisticBDSDPopulation
randomEvent = (LogisticBDSDParameters
 -> AbsoluteTime
 -> LogisticBDSDPopulation
 -> Identifier
 -> GenIO
 -> IO
      (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier))
-> SimulationRandEvent
     LogisticBDSDParameters LogisticBDSDPopulation
forall a b.
(ModelParameters a b, Population b) =>
(a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent LogisticBDSDParameters
-> AbsoluteTime
-> LogisticBDSDPopulation
-> Identifier
-> GenIO
-> IO
     (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
randEvent'

randEvent' ::
     LogisticBDSDParameters
  -> AbsoluteTime
  -> LogisticBDSDPopulation
  -> Identifier
  -> GenIO
  -> IO (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
randEvent' :: LogisticBDSDParameters
-> AbsoluteTime
-> LogisticBDSDPopulation
-> Identifier
-> GenIO
-> IO
     (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
randEvent' params :: LogisticBDSDParameters
params@LogisticBDSDParameters {Rate
Int
Timed Rate
paramsDisasters :: Timed Rate
paramsSamplingRate :: Rate
paramsDeathRate :: Rate
paramsCapacity :: Int
paramsBirthRate :: Rate
paramsDisasters :: LogisticBDSDParameters -> Timed Rate
paramsSamplingRate :: LogisticBDSDParameters -> Rate
paramsDeathRate :: LogisticBDSDParameters -> Rate
paramsCapacity :: LogisticBDSDParameters -> Int
paramsBirthRate :: LogisticBDSDParameters -> Rate
..} AbsoluteTime
currTime currPop :: LogisticBDSDPopulation
currPop@(LogisticBDSDPopulation People
currPpl) Identifier
currId GenIO
gen =
  let netEventRate :: Rate
netEventRate = (Maybe Rate -> Rate
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Rate -> Rate) -> Maybe Rate -> Rate
forall a b. (a -> b) -> a -> b
$ LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe Rate
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Rate
eventRate LogisticBDSDPopulation
currPop LogisticBDSDParameters
params AbsoluteTime
currTime)
      popSizeDouble :: Rate
popSizeDouble = 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
currPpl
      (Just Vector Rate
weightsVec) = LogisticBDSDPopulation
-> LogisticBDSDParameters -> AbsoluteTime -> Maybe (Vector Rate)
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe (Vector Rate)
eventWeights LogisticBDSDPopulation
currPop LogisticBDSDParameters
params AbsoluteTime
currTime
   in do Rate
delay <- Rate -> Gen RealWorld -> IO Rate
forall g (m :: * -> *). StatefulGen g m => Rate -> g -> m Rate
exponential (Rate
netEventRate Rate -> Rate -> Rate
forall a. Num a => a -> a -> a
* Rate
popSizeDouble) Gen RealWorld
GenIO
gen
         let newEventTime :: AbsoluteTime
newEventTime = AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
currTime (Rate -> TimeDelta
TimeDelta Rate
delay)
         if AbsoluteTime -> AbsoluteTime -> Timed Rate -> Bool
noScheduledEvent AbsoluteTime
currTime AbsoluteTime
newEventTime Timed Rate
paramsDisasters
           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 Vector Rate
weightsVec Gen RealWorld
GenIO
gen
             (Person
randPerson, People
otherPeople) <- People -> GenIO -> IO (Person, People)
randomPerson People
currPpl GenIO
gen
             (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
-> IO
     (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return ((AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
 -> IO
      (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier))
-> (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation,
    Identifier)
-> IO
     (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
forall a b. (a -> b) -> a -> b
$
               case Int
eventIx of
                 Int
0 ->
                   let (Person
infectedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
currId
                       infEvent :: EpidemicEvent
infEvent =
                         AbsoluteTime -> Person -> Person -> EpidemicEvent
Infection AbsoluteTime
newEventTime Person
randPerson Person
infectedPerson
                       newPop :: LogisticBDSDPopulation
newPop =
                         People -> LogisticBDSDPopulation
LogisticBDSDPopulation
                           (Person -> People -> People
addPerson Person
infectedPerson People
currPpl)
                    in (AbsoluteTime
newEventTime, EpidemicEvent
infEvent, LogisticBDSDPopulation
newPop, Identifier
newId)
                 Int
1 ->
                   ( AbsoluteTime
newEventTime
                   , AbsoluteTime -> Person -> EpidemicEvent
Removal AbsoluteTime
newEventTime Person
randPerson
                   , People -> LogisticBDSDPopulation
LogisticBDSDPopulation People
otherPeople
                   , Identifier
currId)
                 Int
2 ->
                   ( AbsoluteTime
newEventTime
                   , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
randPerson Bool
True
                   , People -> LogisticBDSDPopulation
LogisticBDSDPopulation People
otherPeople
                   , Identifier
currId)
                 Int
_ -> String
-> (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation,
    Identifier)
forall a. HasCallStack => String -> a
error String
"do not recognise the type of event index."
           else let (Just dsstr :: (AbsoluteTime, Rate)
dsstr@(AbsoluteTime
dsstrTime, Rate
_)) =
                      AbsoluteTime -> Timed Rate -> Maybe (AbsoluteTime, Rate)
firstScheduled AbsoluteTime
currTime Timed Rate
paramsDisasters
                 in do (EpidemicEvent
schdEvent, LogisticBDSDPopulation
postEventPpl) <-
                         (AbsoluteTime, Rate)
-> LogisticBDSDPopulation
-> GenIO
-> IO (EpidemicEvent, LogisticBDSDPopulation)
randomDisasterEvent (AbsoluteTime, Rate)
dsstr LogisticBDSDPopulation
currPop GenIO
gen
                       (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
-> IO
     (AbsoluteTime, EpidemicEvent, LogisticBDSDPopulation, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
dsstrTime, EpidemicEvent
schdEvent, LogisticBDSDPopulation
postEventPpl, Identifier
currId)

-- | Return a randomly sampled Disaster event
-- TODO Move this into the epidemic module to keep things DRY.
randomDisasterEvent ::
     (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the disaster
  -> LogisticBDSDPopulation -- ^ The state of the population prior to the disaster
  -> GenIO
  -> IO (EpidemicEvent, LogisticBDSDPopulation)
randomDisasterEvent :: (AbsoluteTime, Rate)
-> LogisticBDSDPopulation
-> GenIO
-> IO (EpidemicEvent, LogisticBDSDPopulation)
randomDisasterEvent (AbsoluteTime
dsstrTime, Rate
dsstrProb) (LogisticBDSDPopulation (People Vector Person
currPpl)) GenIO
gen = do
  Vector Bool
randBernoullis <- 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
currPpl) (Rate -> Gen RealWorld -> IO Bool
forall g (m :: * -> *). StatefulGen g m => Rate -> g -> m Bool
bernoulli Rate
dsstrProb 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 :: Vector Person
sampledPeople = ((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
currPpl Vector Bool
randBernoullis
      unsampledPeople :: Vector Person
unsampledPeople = ((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
currPpl Vector Bool
randBernoullis
   in (EpidemicEvent, LogisticBDSDPopulation)
-> IO (EpidemicEvent, LogisticBDSDPopulation)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( AbsoluteTime -> People -> Bool -> EpidemicEvent
PopulationSample AbsoluteTime
dsstrTime (Vector Person -> People
People Vector Person
sampledPeople) Bool
False
        , People -> LogisticBDSDPopulation
LogisticBDSDPopulation (Vector Person -> People
People Vector Person
unsampledPeople))