{-# 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

-- | A new person constructed from the given identifier and a new identifier.
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))

-- | An element of a vector and the vector with that element removed.
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))

-- | A random person and the remaining group of people after they have been
-- sampled with removal.
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
";"

-- | The number of elements of the list that map to @True@ under the predicate.
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]

-- | Run a simulation described by a configuration object and the model's
-- @allEvents@ style function (see the example in
-- "Epidemic.Model.InhomogeneousBDSCOD") using the provided PRNG.
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

-- | Run a simulation using a fixed PRNG random seed.
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

-- | Simulation conditioned upon there being at least two sequenced samples.
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

-- | Run a simulation described by a configuration object but using a random
-- seed generated by the system rather than a seed
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

-- | Predicate for whether an epidemic event will appear as a leaf in the
-- reconstructed tree. For scheduled sequenced samples this will only return
-- true if there was at least one lineage observed.
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

-- | The number of lineages at the end of a simulation.
finalSize ::
     [EpidemicEvent] -- ^ The events from the simulation
  -> 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

-- | Generate exponentially distributed random variates with inhomogeneous rate
-- starting from a particular point in time.
--
-- Assuming the @stepFunc@ is the intensity of arrivals and @t0@ is the start
-- time this returns @t1@ the time of the next arrival.
inhomExponential ::
     PrimMonad m
  => Timed Double -- ^ Step function
  -> AbsoluteTime -- ^ Start time
  -> Gen (PrimState m) -- ^ Generator
  -> 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

-- | Generate exponentially distributed random variates with inhomogeneous rate.
--
-- __TODO__ The algorithm used here generates more variates than are needed. It
-- would be nice to use a more efficient implementation.
--
randInhomExp ::
     PrimMonad m
  => AbsoluteTime -- ^ Timer
  -> Timed Double -- ^ Step function
  -> Gen (PrimState m) -- ^ Generator.
  -> 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

-- | Helper function for converting between the Maybe monad and the Either
-- monad.
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