{-# LANGUAGE MultiParamTypeClasses #-}

module Epidemic.Model.BDSCOD
  ( configuration
  , randomEvent
  , BDSCODParameters(..)
  , BDSCODPopulation(..)
  ) where

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(..)
  , asTimed
  , maybeNextTimed
  , timeAfterDelta
  )
import Epidemic.Utility
import System.Random.MWC
import System.Random.MWC.Distributions (bernoulli, categorical, exponential)

-- | birth rate, death rate, sampling rate, catastrophe specification, occurrence rate and disaster specification
data BDSCODParameters =
  BDSCODParameters Rate Rate Rate (Timed Probability) Rate (Timed Probability)

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

instance ModelParameters BDSCODParameters BDSCODPopulation where
  rNaught :: BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe Double
rNaught BDSCODPopulation
_ (BDSCODParameters Double
br Double
dr Double
sRate Timed Double
_ Double
occRate Timed Double
_) AbsoluteTime
_ =
    Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
br Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
dr Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
occRate)
  eventRate :: BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe Double
eventRate BDSCODPopulation
_ (BDSCODParameters Double
br Double
dr Double
sRate Timed Double
_ Double
occRate Timed Double
_) AbsoluteTime
_ =
    Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
br Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dr Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
occRate
  birthProb :: BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe Double
birthProb BDSCODPopulation
_ (BDSCODParameters Double
br Double
dr Double
sRate Timed Double
_ Double
occRate Timed Double
_) AbsoluteTime
_ =
    Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
br Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
br Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dr Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
occRate)
  eventWeights :: BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe (Vector Double)
eventWeights BDSCODPopulation
_ (BDSCODParameters Double
br Double
dr Double
sRate Timed Double
_ Double
occRate Timed Double
_) AbsoluteTime
_ =
    Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just (Vector Double -> Maybe (Vector Double))
-> Vector Double -> Maybe (Vector Double)
forall a b. (a -> b) -> a -> b
$ [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList [Double
br, Double
dr, Double
sRate, Double
occRate]

instance Population BDSCODPopulation where
  susceptiblePeople :: BDSCODPopulation -> Maybe People
susceptiblePeople BDSCODPopulation
_ = Maybe People
forall a. Maybe a
Nothing
  infectiousPeople :: BDSCODPopulation -> Maybe People
infectiousPeople (BDSCODPopulation People
people) = People -> Maybe People
forall a. a -> Maybe a
Just People
people
  removedPeople :: BDSCODPopulation -> Maybe People
removedPeople BDSCODPopulation
_ = Maybe People
forall a. Maybe a
Nothing
  isInfected :: BDSCODPopulation -> Bool
isInfected (BDSCODPopulation (People Vector Person
people)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Vector Person -> Bool
forall a. Vector a -> Bool
V.null Vector Person
people

-- | Configuration of a birth-death-sampling-occurrence-disaster simulation
configuration ::
     TimeDelta -- ^ Duration of the simulation
  -> Bool -- ^ condition upon at least two sequenced samples.
  -> Maybe (BDSCODPopulation -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling.
  -> ( Rate
     , Rate
     , Rate
     , [(AbsoluteTime, Probability)]
     , Rate
     , [(AbsoluteTime, Probability)]) -- ^ Birth, Death, Sampling, Catastrophe probability, Occurrence rates and Disaster probabilities
  -> Maybe (SimulationConfiguration BDSCODParameters BDSCODPopulation s)
configuration :: TimeDelta
-> Bool
-> Maybe (BDSCODPopulation -> Bool, [EpidemicEvent] -> s)
-> (Double, Double, Double, [(AbsoluteTime, Double)], Double,
    [(AbsoluteTime, Double)])
-> Maybe
     (SimulationConfiguration BDSCODParameters BDSCODPopulation s)
configuration TimeDelta
maxTime Bool
atLeastCherry Maybe (BDSCODPopulation -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs (Double
birthRate, Double
deathRate, Double
samplingRate, [(AbsoluteTime, Double)]
catastropheSpec, Double
occurrenceRate, [(AbsoluteTime, Double)]
disasterSpec) = do
  Timed Double
catastropheSpec' <- [(AbsoluteTime, Double)] -> Maybe (Timed Double)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Double)]
catastropheSpec
  Timed Double
disasterSpec' <- [(AbsoluteTime, Double)] -> Maybe (Timed Double)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Double)]
disasterSpec
  let bdscodParams :: BDSCODParameters
bdscodParams =
        Double
-> Double
-> Double
-> Timed Double
-> Double
-> Timed Double
-> BDSCODParameters
BDSCODParameters
          Double
birthRate
          Double
deathRate
          Double
samplingRate
          Timed Double
catastropheSpec'
          Double
occurrenceRate
          Timed Double
disasterSpec'
      (Person
seedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
initialIdentifier
      bdscodPop :: BDSCODPopulation
bdscodPop = People -> BDSCODPopulation
BDSCODPopulation (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 BDSCODPopulation s)
termHandler = do (BDSCODPopulation -> Bool
f1, [EpidemicEvent] -> s
f2) <- Maybe (BDSCODPopulation -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs
                       TerminationHandler BDSCODPopulation s
-> Maybe (TerminationHandler BDSCODPopulation s)
forall (m :: * -> *) a. Monad m => a -> m a
return (TerminationHandler BDSCODPopulation s
 -> Maybe (TerminationHandler BDSCODPopulation s))
-> TerminationHandler BDSCODPopulation s
-> Maybe (TerminationHandler BDSCODPopulation s)
forall a b. (a -> b) -> a -> b
$ (BDSCODPopulation -> Bool)
-> ([EpidemicEvent] -> s) -> TerminationHandler BDSCODPopulation s
forall b c.
Population b =>
(b -> Bool) -> ([EpidemicEvent] -> c) -> TerminationHandler b c
TerminationHandler BDSCODPopulation -> Bool
f1 [EpidemicEvent] -> s
f2
   in SimulationConfiguration BDSCODParameters BDSCODPopulation s
-> Maybe
     (SimulationConfiguration BDSCODParameters BDSCODPopulation s)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationConfiguration BDSCODParameters BDSCODPopulation s
 -> Maybe
      (SimulationConfiguration BDSCODParameters BDSCODPopulation s))
-> SimulationConfiguration BDSCODParameters BDSCODPopulation s
-> Maybe
     (SimulationConfiguration BDSCODParameters BDSCODPopulation s)
forall a b. (a -> b) -> a -> b
$
      BDSCODParameters
-> BDSCODPopulation
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler BDSCODPopulation s)
-> Bool
-> SimulationConfiguration BDSCODParameters BDSCODPopulation s
forall r p s.
r
-> p
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler p s)
-> Bool
-> SimulationConfiguration r p s
SimulationConfiguration
        BDSCODParameters
bdscodParams
        BDSCODPopulation
bdscodPop
        Identifier
newId
        (Double -> AbsoluteTime
AbsoluteTime Double
0)
        TimeDelta
maxTime
        Maybe (TerminationHandler BDSCODPopulation s)
termHandler
        Bool
atLeastCherry

-- | The way in which random events are generated in this model.
randomEvent :: SimulationRandEvent BDSCODParameters BDSCODPopulation
randomEvent :: SimulationRandEvent BDSCODParameters BDSCODPopulation
randomEvent = (BDSCODParameters
 -> AbsoluteTime
 -> BDSCODPopulation
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier))
-> SimulationRandEvent BDSCODParameters BDSCODPopulation
forall a b.
(ModelParameters a b, Population b) =>
(a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent BDSCODParameters
-> AbsoluteTime
-> BDSCODPopulation
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
randomEvent'

-- | Return a random event from the BDSCOD-process given the current state of the process.
randomEvent' ::
     BDSCODParameters -- ^ Parameters of the process
  -> AbsoluteTime -- ^ The current time within the process
  -> BDSCODPopulation -- ^ The current state of the populaion
  -> Identifier -- ^ The current state of the identifier generator
  -> GenIO -- ^ The current state of the PRNG
  -> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
randomEvent' :: BDSCODParameters
-> AbsoluteTime
-> BDSCODPopulation
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
randomEvent' params :: BDSCODParameters
params@(BDSCODParameters Double
_ Double
_ Double
_ Timed Double
catastInfo Double
_ Timed Double
disastInfo) AbsoluteTime
currTime currPop :: BDSCODPopulation
currPop@(BDSCODPopulation People
currPeople) Identifier
currId GenIO
gen =
  let (Just Double
netEventRate) = BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe Double
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Double
eventRate BDSCODPopulation
currPop BDSCODParameters
params AbsoluteTime
currTime
      (Just Vector Double
weightVec) = BDSCODPopulation
-> BDSCODParameters -> AbsoluteTime -> Maybe (Vector Double)
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe (Vector Double)
eventWeights BDSCODPopulation
currPop BDSCODParameters
params AbsoluteTime
currTime
   in do Double
delay <-
           Double -> Gen RealWorld -> IO Double
forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Double
exponential (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (People -> Int
numPeople People
currPeople) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
netEventRate) Gen RealWorld
GenIO
gen
         let newEventTime :: AbsoluteTime
newEventTime = AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
currTime (Double -> TimeDelta
TimeDelta Double
delay)
         if AbsoluteTime -> AbsoluteTime -> Timed Double -> Bool
noScheduledEvent AbsoluteTime
currTime AbsoluteTime
newEventTime (Timed Double
catastInfo Timed Double -> Timed Double -> Timed Double
forall a. Semigroup a => a -> a -> a
<> Timed Double
disastInfo)
         then do
                Int
eventIx <- Vector Double -> Gen RealWorld -> IO Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical Vector Double
weightVec Gen RealWorld
GenIO
gen
                (Person
selectedPerson, People
unselectedPeople) <- People -> GenIO -> IO (Person, People)
randomPerson People
currPeople GenIO
gen
                (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return ((AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
 -> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier))
-> (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall a b. (a -> b) -> a -> b
$
                  case Int
eventIx of
                    Int
0 ->
                      let (Person
birthedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
currId
                          infEvent :: EpidemicEvent
infEvent =
                            AbsoluteTime -> Person -> Person -> EpidemicEvent
Infection AbsoluteTime
newEventTime Person
selectedPerson Person
birthedPerson
                      in ( AbsoluteTime
newEventTime
                         , EpidemicEvent
infEvent
                         , People -> BDSCODPopulation
BDSCODPopulation (Person -> People -> People
addPerson Person
birthedPerson People
currPeople)
                         , Identifier
newId)
                    Int
1 ->
                      ( AbsoluteTime
newEventTime
                      , AbsoluteTime -> Person -> EpidemicEvent
Removal AbsoluteTime
newEventTime Person
selectedPerson
                      , People -> BDSCODPopulation
BDSCODPopulation People
unselectedPeople
                      , Identifier
currId)
                    Int
2 ->
                      ( AbsoluteTime
newEventTime
                      , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
selectedPerson Bool
True
                      , People -> BDSCODPopulation
BDSCODPopulation People
unselectedPeople
                      , Identifier
currId)
                    Int
3 ->
                      ( AbsoluteTime
newEventTime
                      , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
selectedPerson Bool
False
                      , People -> BDSCODPopulation
BDSCODPopulation People
unselectedPeople
                      , Identifier
currId)
                    Int
_ ->
                      String
-> (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall a. HasCallStack => String -> a
error String
"no birth, death, sampling, occurrence event selected."
         else case Timed Double
-> Timed Double
-> AbsoluteTime
-> Maybe (AbsoluteTime, Either Double Double)
forall a b.
Timed a
-> Timed b -> AbsoluteTime -> Maybe (AbsoluteTime, Either a b)
maybeNextTimed Timed Double
catastInfo Timed Double
disastInfo AbsoluteTime
currTime of
                Just (AbsoluteTime
disastTime, Right Double
disastProb) ->
                 do (EpidemicEvent
disastEvent, BDSCODPopulation
postDisastPop) <-
                      (AbsoluteTime, Double)
-> BDSCODPopulation
-> GenIO
-> IO (EpidemicEvent, BDSCODPopulation)
randomDisasterEvent
                      (AbsoluteTime
disastTime, Double
disastProb)
                      BDSCODPopulation
currPop
                      GenIO
gen
                    (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
disastTime, EpidemicEvent
disastEvent, BDSCODPopulation
postDisastPop, Identifier
currId)
                Just (AbsoluteTime
catastTime, Left Double
catastProb) ->
                 do (EpidemicEvent
catastEvent, BDSCODPopulation
postCatastPop) <-
                      (AbsoluteTime, Double)
-> BDSCODPopulation
-> GenIO
-> IO (EpidemicEvent, BDSCODPopulation)
randomCatastropheEvent
                      (AbsoluteTime
catastTime, Double
catastProb)
                      BDSCODPopulation
currPop
                      GenIO
gen
                    (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
catastTime, EpidemicEvent
catastEvent, BDSCODPopulation
postCatastPop, Identifier
currId)
                Maybe (AbsoluteTime, Either Double Double)
Nothing -> String
-> IO (AbsoluteTime, EpidemicEvent, BDSCODPopulation, Identifier)
forall a. HasCallStack => String -> a
error String
"Missing a next scheduled event when there should be one."

-- | Return a randomly sampled Catastrophe event
randomCatastropheEvent ::
     (AbsoluteTime, Probability) -- ^ Time and probability of sampling in the catastrophe
  -> BDSCODPopulation -- ^ The state of the population prior to the catastrophe
  -> GenIO
  -> IO (EpidemicEvent, BDSCODPopulation)
randomCatastropheEvent :: (AbsoluteTime, Double)
-> BDSCODPopulation
-> GenIO
-> IO (EpidemicEvent, BDSCODPopulation)
randomCatastropheEvent (AbsoluteTime
catastTime, Double
rhoProb) (BDSCODPopulation (People Vector Person
currPeople)) GenIO
gen = 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) (Double -> Gen RealWorld -> IO Bool
forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Bool
bernoulli Double
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 :: 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
currPeople Vector Bool
rhoBernoullis
      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
currPeople Vector Bool
rhoBernoullis
   in (EpidemicEvent, BDSCODPopulation)
-> IO (EpidemicEvent, BDSCODPopulation)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( AbsoluteTime -> People -> Bool -> EpidemicEvent
PopulationSample AbsoluteTime
catastTime (Vector Person -> People
People Vector Person
sampledPeople) Bool
True
        , People -> BDSCODPopulation
BDSCODPopulation (Vector Person -> People
People Vector Person
unsampledPeople))

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