{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Epidemic.Utility ( initialIdentifier
, inhomExponential
, randomPerson
, maybeToRight
, newPerson
, isReconTreeLeaf
, simulationWithSystem
, simulationWithFixedSeed
, simulationWithGenIO
) where
import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Vector as V
import Epidemic
import Epidemic.Types.Events
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import Epidemic.Types.Simulation
import Epidemic.Types.Time (AbsoluteTime (..),
TimeDelta (..), Timed (..),
cadlagValue, nextTime,
timeAfterDelta)
import System.Random.MWC
import System.Random.MWC.Distributions (exponential)
initialIdentifier :: Identifier
initialIdentifier :: Identifier
initialIdentifier = Integer -> Identifier
Identifier Integer
1
newPerson :: Identifier -> (Person, Identifier)
newPerson :: Identifier -> (Person, Identifier)
newPerson idntty :: Identifier
idntty@(Identifier Integer
idInt) = (Identifier -> Person
Person Identifier
idntty, Integer -> Identifier
Identifier (Integer
idInt Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1))
selectElem :: V.Vector a -> Int -> (a, V.Vector a)
selectElem :: Vector a -> Int -> (a, Vector a)
selectElem Vector a
v Int
n
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (Vector a -> a
forall a. Vector a -> a
V.head Vector a
v, Vector a -> Vector a
forall a. Vector a -> Vector a
V.tail Vector a
v)
| Bool
otherwise =
let (Vector a
foo, Vector a
bar) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
n Vector a
v
in (Vector a -> a
forall a. Vector a -> a
V.head Vector a
bar, Vector a
foo Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
V.++ (Vector a -> Vector a
forall a. Vector a -> Vector a
V.tail Vector a
bar))
randomPerson :: People -> GenIO -> IO (Person, People)
randomPerson :: People -> GenIO -> IO (Person, People)
randomPerson people :: People
people@(People Vector Person
persons) GenIO
gen = do
Double
u <- GenIO -> IO Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform GenIO
gen
let personIx :: Int
personIx = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ People -> Int
numPeople People
people :: Double))
(Person
person, Vector Person
remPeople) = Vector Person -> Int -> (Person, Vector Person)
forall a. Vector a -> Int -> (a, Vector a)
selectElem Vector Person
persons Int
personIx
in (Person, People) -> IO (Person, People)
forall (m :: * -> *) a. Monad m => a -> m a
return (Person
person, Vector Person -> People
People Vector Person
remPeople)
type NName = Maybe String
type NLength = Maybe Double
data NBranch =
NBranch NSubtree NLength
deriving (NBranch -> NBranch -> Bool
(NBranch -> NBranch -> Bool)
-> (NBranch -> NBranch -> Bool) -> Eq NBranch
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NBranch -> NBranch -> Bool
$c/= :: NBranch -> NBranch -> Bool
== :: NBranch -> NBranch -> Bool
$c== :: NBranch -> NBranch -> Bool
Eq)
instance Show NBranch where
show :: NBranch -> String
show (NBranch NSubtree
st (Just Double
l)) = NSubtree -> String
forall a. Show a => a -> String
show NSubtree
st String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
":" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Double -> String
forall a. Show a => a -> String
show Double
l
show (NBranch NSubtree
st Maybe Double
Nothing) = NSubtree -> String
forall a. Show a => a -> String
show NSubtree
st
data NBranchSet =
NBranchSet [NBranch]
deriving (NBranchSet -> NBranchSet -> Bool
(NBranchSet -> NBranchSet -> Bool)
-> (NBranchSet -> NBranchSet -> Bool) -> Eq NBranchSet
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NBranchSet -> NBranchSet -> Bool
$c/= :: NBranchSet -> NBranchSet -> Bool
== :: NBranchSet -> NBranchSet -> Bool
$c== :: NBranchSet -> NBranchSet -> Bool
Eq)
instance Show NBranchSet where
show :: NBranchSet -> String
show (NBranchSet [NBranch]
bs) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ (String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
List.intercalate String
"," ((NBranch -> String) -> [NBranch] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map NBranch -> String
forall a. Show a => a -> String
show [NBranch]
bs)) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
data NSubtree
= NLeaf NName
| NInternal NBranchSet
deriving (NSubtree -> NSubtree -> Bool
(NSubtree -> NSubtree -> Bool)
-> (NSubtree -> NSubtree -> Bool) -> Eq NSubtree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NSubtree -> NSubtree -> Bool
$c/= :: NSubtree -> NSubtree -> Bool
== :: NSubtree -> NSubtree -> Bool
$c== :: NSubtree -> NSubtree -> Bool
Eq)
instance Show NSubtree where
show :: NSubtree -> String
show (NLeaf (Just String
n)) = String
n
show (NLeaf Maybe String
Nothing) = String
""
show (NInternal NBranchSet
bs) = NBranchSet -> String
forall a. Show a => a -> String
show NBranchSet
bs
data NTree =
NTree [NBranch]
deriving (NTree -> NTree -> Bool
(NTree -> NTree -> Bool) -> (NTree -> NTree -> Bool) -> Eq NTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NTree -> NTree -> Bool
$c/= :: NTree -> NTree -> Bool
== :: NTree -> NTree -> Bool
$c== :: NTree -> NTree -> Bool
Eq)
instance Show NTree where
show :: NTree -> String
show (NTree [NBranch]
bs) = NBranchSet -> String
forall a. Show a => a -> String
show ([NBranch] -> NBranchSet
NBranchSet [NBranch]
bs) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
";"
count' :: (a -> Bool) -> [a] -> Int
count' :: (a -> Bool) -> [a] -> Int
count' a -> Bool
p [a]
xs = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [if a -> Bool
p a
x then Int
1 else Int
0 | a
x <- [a]
xs]
simulationWithGenIO ::
(ModelParameters a b, Population b)
=> SimulationConfiguration a b c
-> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithGenIO :: SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithGenIO config :: SimulationConfiguration a b c
config@SimulationConfiguration {a
b
Bool
Maybe (TerminationHandler b c)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: forall r p s. SimulationConfiguration r p s -> Bool
scTerminationHandler :: forall r p s.
SimulationConfiguration r p s -> Maybe (TerminationHandler p s)
scSimDuration :: forall r p s. SimulationConfiguration r p s -> TimeDelta
scStartTime :: forall r p s. SimulationConfiguration r p s -> AbsoluteTime
scNewIdentifier :: forall r p s. SimulationConfiguration r p s -> Identifier
scPopulation :: forall r p s. SimulationConfiguration r p s -> p
scRates :: forall r p s. SimulationConfiguration r p s -> r
scRequireCherry :: Bool
scTerminationHandler :: Maybe (TerminationHandler b c)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
..} a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc GenIO
gen =
if Bool
scRequireCherry
then
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b c.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationAtLeastCherry SimulationConfiguration a b c
config a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc GenIO
gen
else do
SimulationState b c
simState <-
a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc
a
scRates
(AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
Maybe (TerminationHandler b c)
scTerminationHandler
((AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState (AbsoluteTime
scStartTime, [], b
scPopulation, Identifier
scNewIdentifier))
GenIO
gen
Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ case SimulationState b c
simState of
SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) -> [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. b -> Either a b
Right ([EpidemicEvent] -> Either (Maybe c) [EpidemicEvent])
-> [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
List.sort [EpidemicEvent]
events
TerminatedSimulation Maybe c
maybeSummary -> Maybe c -> Either (Maybe c) [EpidemicEvent]
forall a b. a -> Either a b
Left Maybe c
maybeSummary
simulationWithFixedSeed ::
(ModelParameters a b, Population b)
=> SimulationConfiguration a b c
-> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithFixedSeed :: SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithFixedSeed SimulationConfiguration a b c
config a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc = do
Gen RealWorld
gen <- IO (Gen RealWorld)
IO GenIO
genIOFromFixed
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b c.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithGenIO SimulationConfiguration a b c
config a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc Gen RealWorld
GenIO
gen
simulationAtLeastCherry ::
(ModelParameters a b, Population b)
=> SimulationConfiguration a b c
-> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationAtLeastCherry :: SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationAtLeastCherry config :: SimulationConfiguration a b c
config@SimulationConfiguration {a
b
Bool
Maybe (TerminationHandler b c)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: Bool
scTerminationHandler :: Maybe (TerminationHandler b c)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
scRequireCherry :: forall r p s. SimulationConfiguration r p s -> Bool
scTerminationHandler :: forall r p s.
SimulationConfiguration r p s -> Maybe (TerminationHandler p s)
scSimDuration :: forall r p s. SimulationConfiguration r p s -> TimeDelta
scStartTime :: forall r p s. SimulationConfiguration r p s -> AbsoluteTime
scNewIdentifier :: forall r p s. SimulationConfiguration r p s -> Identifier
scPopulation :: forall r p s. SimulationConfiguration r p s -> p
scRates :: forall r p s. SimulationConfiguration r p s -> r
..} a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc GenIO
gen = do
SimulationState b c
simState <-
a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc
a
scRates
(AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
Maybe (TerminationHandler b c)
scTerminationHandler
((AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState (AbsoluteTime
scStartTime, [], b
scPopulation, Identifier
scNewIdentifier))
GenIO
gen
case SimulationState b c
simState of
SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) ->
if (EpidemicEvent -> Bool) -> [EpidemicEvent] -> Int
forall a. (a -> Bool) -> [a] -> Int
count' EpidemicEvent -> Bool
isReconTreeLeaf [EpidemicEvent]
events Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
then Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. b -> Either a b
Right ([EpidemicEvent] -> Either (Maybe c) [EpidemicEvent])
-> [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
List.sort [EpidemicEvent]
events
else SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b c.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> GenIO
-> IO (Either (Maybe c) [EpidemicEvent])
simulationAtLeastCherry SimulationConfiguration a b c
config a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc GenIO
gen
TerminatedSimulation Maybe c
maybeSummary -> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ Maybe c -> Either (Maybe c) [EpidemicEvent]
forall a b. a -> Either a b
Left Maybe c
maybeSummary
simulationWithSystem ::
(ModelParameters a b, Population b)
=> SimulationConfiguration a b c
-> (a -> AbsoluteTime -> Maybe (TerminationHandler b c) -> SimulationState b c -> GenIO -> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithSystem :: SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithSystem config :: SimulationConfiguration a b c
config@SimulationConfiguration {a
b
Bool
Maybe (TerminationHandler b c)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: Bool
scTerminationHandler :: Maybe (TerminationHandler b c)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
scRequireCherry :: forall r p s. SimulationConfiguration r p s -> Bool
scTerminationHandler :: forall r p s.
SimulationConfiguration r p s -> Maybe (TerminationHandler p s)
scSimDuration :: forall r p s. SimulationConfiguration r p s -> TimeDelta
scStartTime :: forall r p s. SimulationConfiguration r p s -> AbsoluteTime
scNewIdentifier :: forall r p s. SimulationConfiguration r p s -> Identifier
scPopulation :: forall r p s. SimulationConfiguration r p s -> p
scRates :: forall r p s. SimulationConfiguration r p s -> r
..} a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc = do
SimulationState b c
simState <-
(GenIO -> IO (SimulationState b c)) -> IO (SimulationState b c)
forall (m :: * -> *) a.
PrimBase m =>
(Gen (PrimState m) -> m a) -> IO a
withSystemRandom ((GenIO -> IO (SimulationState b c)) -> IO (SimulationState b c))
-> (GenIO -> IO (SimulationState b c)) -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$ \GenIO
g ->
a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc
a
scRates
(AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
Maybe (TerminationHandler b c)
scTerminationHandler
((AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState (AbsoluteTime
scStartTime, [], b
scPopulation, Identifier
scNewIdentifier))
GenIO
g
case SimulationState b c
simState of
SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) ->
if Bool
scRequireCherry
then (if (EpidemicEvent -> Bool) -> [EpidemicEvent] -> Int
forall a. (a -> Bool) -> [a] -> Int
count' EpidemicEvent -> Bool
isReconTreeLeaf [EpidemicEvent]
events Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
then Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. b -> Either a b
Right ([EpidemicEvent] -> Either (Maybe c) [EpidemicEvent])
-> [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
List.sort [EpidemicEvent]
events
else SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b c.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b c
-> (a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c))
-> IO (Either (Maybe c) [EpidemicEvent])
simulationWithSystem SimulationConfiguration a b c
config a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEventsFunc)
else Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. b -> Either a b
Right ([EpidemicEvent] -> Either (Maybe c) [EpidemicEvent])
-> [EpidemicEvent] -> Either (Maybe c) [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
List.sort [EpidemicEvent]
events
TerminatedSimulation Maybe c
maybeSummary -> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent]))
-> Either (Maybe c) [EpidemicEvent]
-> IO (Either (Maybe c) [EpidemicEvent])
forall a b. (a -> b) -> a -> b
$ Maybe c -> Either (Maybe c) [EpidemicEvent]
forall a b. a -> Either a b
Left Maybe c
maybeSummary
isReconTreeLeaf :: EpidemicEvent -> Bool
isReconTreeLeaf :: EpidemicEvent -> Bool
isReconTreeLeaf EpidemicEvent
e =
case EpidemicEvent
e of
IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
..} -> Bool
indSampSeq
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> Bool
popSampSeq Bool -> Bool -> Bool
&& Bool -> Bool
not (People -> Bool
nullPeople People
popSampPeople)
EpidemicEvent
_ -> Bool
False
finalSize ::
[EpidemicEvent]
-> Integer
finalSize :: [EpidemicEvent] -> Integer
finalSize = (Integer -> EpidemicEvent -> Integer)
-> Integer -> [EpidemicEvent] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Integer
x EpidemicEvent
y -> Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ EpidemicEvent -> Integer
eventPopDelta EpidemicEvent
y) Integer
1
inhomExponential ::
PrimMonad m
=> Timed Double
-> AbsoluteTime
-> Gen (PrimState m)
-> m (Maybe AbsoluteTime)
inhomExponential :: Timed Double
-> AbsoluteTime -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
inhomExponential Timed Double
stepFunc AbsoluteTime
t0 = AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp AbsoluteTime
t0 Timed Double
stepFunc
randInhomExp ::
PrimMonad m
=> AbsoluteTime
-> Timed Double
-> Gen (PrimState m)
-> m (Maybe AbsoluteTime)
randInhomExp :: AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp AbsoluteTime
crrT Timed Double
stepFunc Gen (PrimState m)
gen =
let crrR :: Maybe Double
crrR = Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
stepFunc AbsoluteTime
crrT
nxtT :: Maybe AbsoluteTime
nxtT = Timed Double -> AbsoluteTime -> Maybe AbsoluteTime
forall a. Timed a -> AbsoluteTime -> Maybe AbsoluteTime
nextTime Timed Double
stepFunc AbsoluteTime
crrT
in if Maybe Double -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Maybe Double
crrR Bool -> Bool -> Bool
&& Maybe AbsoluteTime -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Maybe AbsoluteTime
nxtT
then do
Double
crrD <- Double -> Gen (PrimState m) -> m Double
forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Double
exponential (Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe Double
crrR) Gen (PrimState m)
gen
let propT :: AbsoluteTime
propT = AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
crrT (Double -> TimeDelta
TimeDelta Double
crrD)
if AbsoluteTime
propT AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< Maybe AbsoluteTime -> AbsoluteTime
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe AbsoluteTime
nxtT
then Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe AbsoluteTime -> m (Maybe AbsoluteTime))
-> Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall a b. (a -> b) -> a -> b
$ AbsoluteTime -> Maybe AbsoluteTime
forall a. a -> Maybe a
Just AbsoluteTime
propT
else AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp (Maybe AbsoluteTime -> AbsoluteTime
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe AbsoluteTime
nxtT) Timed Double
stepFunc Gen (PrimState m)
gen
else Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe AbsoluteTime
forall a. Maybe a
Nothing
maybeToRight :: a -> Maybe b -> Either a b
maybeToRight :: a -> Maybe b -> Either a b
maybeToRight a
a Maybe b
maybeB =
case Maybe b
maybeB of
(Just b
b) -> b -> Either a b
forall a b. b -> Either a b
Right b
b
Maybe b
Nothing -> a -> Either a b
forall a b. a -> Either a b
Left a
a