{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
module Epidemic.Types.Observations
( Observation(..)
, ReconstructedTree(..)
, maybeReconstructedTree
, PointProcessEvents(..)
, pointProcessEvents
, reconstructedTreeEvents
, observedEvents
) where
import Control.Monad (liftM)
import qualified Data.Aeson as Json
import qualified Data.ByteString.Builder as BBuilder
import qualified Data.List as List
import qualified Data.Vector as V
import Epidemic.Types.Events
( EpidemicEvent(..)
, EpidemicTree(..)
, maybeEpidemicTree
)
import Epidemic.Types.Time (TimeDelta(..), timeDelta)
import Epidemic.Types.Population (People(..), personByteString)
import GHC.Generics
newtype Observation =
Observation EpidemicEvent
deriving (Int -> Observation -> ShowS
[Observation] -> ShowS
Observation -> String
(Int -> Observation -> ShowS)
-> (Observation -> String)
-> ([Observation] -> ShowS)
-> Show Observation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Observation] -> ShowS
$cshowList :: [Observation] -> ShowS
show :: Observation -> String
$cshow :: Observation -> String
showsPrec :: Int -> Observation -> ShowS
$cshowsPrec :: Int -> Observation -> ShowS
Show, Eq Observation
Eq Observation
-> (Observation -> Observation -> Ordering)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Observation)
-> (Observation -> Observation -> Observation)
-> Ord Observation
Observation -> Observation -> Bool
Observation -> Observation -> Ordering
Observation -> Observation -> Observation
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Observation -> Observation -> Observation
$cmin :: Observation -> Observation -> Observation
max :: Observation -> Observation -> Observation
$cmax :: Observation -> Observation -> Observation
>= :: Observation -> Observation -> Bool
$c>= :: Observation -> Observation -> Bool
> :: Observation -> Observation -> Bool
$c> :: Observation -> Observation -> Bool
<= :: Observation -> Observation -> Bool
$c<= :: Observation -> Observation -> Bool
< :: Observation -> Observation -> Bool
$c< :: Observation -> Observation -> Bool
compare :: Observation -> Observation -> Ordering
$ccompare :: Observation -> Observation -> Ordering
$cp1Ord :: Eq Observation
Ord, Observation -> Observation -> Bool
(Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool) -> Eq Observation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Observation -> Observation -> Bool
$c/= :: Observation -> Observation -> Bool
== :: Observation -> Observation -> Bool
$c== :: Observation -> Observation -> Bool
Eq, (forall x. Observation -> Rep Observation x)
-> (forall x. Rep Observation x -> Observation)
-> Generic Observation
forall x. Rep Observation x -> Observation
forall x. Observation -> Rep Observation x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Observation x -> Observation
$cfrom :: forall x. Observation -> Rep Observation x
Generic)
instance Json.FromJSON Observation
instance Json.ToJSON Observation
newtype PointProcessEvents =
PointProcessEvents [Observation]
pointProcessEvents :: EpidemicTree -> PointProcessEvents
pointProcessEvents :: EpidemicTree -> PointProcessEvents
pointProcessEvents Shoot {} = [Observation] -> PointProcessEvents
PointProcessEvents []
pointProcessEvents (Leaf 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
..} -> [Observation] -> PointProcessEvents
PointProcessEvents ([Observation] -> PointProcessEvents)
-> [Observation] -> PointProcessEvents
forall a b. (a -> b) -> a -> b
$ if Bool -> Bool
not Bool
indSampSeq then [EpidemicEvent -> Observation
Observation EpidemicEvent
e] else []
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> [Observation] -> PointProcessEvents
PointProcessEvents ([Observation] -> PointProcessEvents)
-> [Observation] -> PointProcessEvents
forall a b. (a -> b) -> a -> b
$ if Bool -> Bool
not Bool
popSampSeq then [EpidemicEvent -> Observation
Observation EpidemicEvent
e] else []
EpidemicEvent
_ -> [Observation] -> PointProcessEvents
PointProcessEvents []
pointProcessEvents (Branch EpidemicEvent
_ EpidemicTree
lt EpidemicTree
rt) =
let (PointProcessEvents [Observation]
lEs) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
lt
(PointProcessEvents [Observation]
rEs) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
rt
allEs :: [Observation]
allEs = [Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation]
lEs [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ [Observation]
rEs
in [Observation] -> PointProcessEvents
PointProcessEvents [Observation]
allEs
data ReconstructedTree
= RBranch Observation ReconstructedTree ReconstructedTree
| RLeaf Observation
deriving (Int -> ReconstructedTree -> ShowS
[ReconstructedTree] -> ShowS
ReconstructedTree -> String
(Int -> ReconstructedTree -> ShowS)
-> (ReconstructedTree -> String)
-> ([ReconstructedTree] -> ShowS)
-> Show ReconstructedTree
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ReconstructedTree] -> ShowS
$cshowList :: [ReconstructedTree] -> ShowS
show :: ReconstructedTree -> String
$cshow :: ReconstructedTree -> String
showsPrec :: Int -> ReconstructedTree -> ShowS
$cshowsPrec :: Int -> ReconstructedTree -> ShowS
Show, ReconstructedTree -> ReconstructedTree -> Bool
(ReconstructedTree -> ReconstructedTree -> Bool)
-> (ReconstructedTree -> ReconstructedTree -> Bool)
-> Eq ReconstructedTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ReconstructedTree -> ReconstructedTree -> Bool
$c/= :: ReconstructedTree -> ReconstructedTree -> Bool
== :: ReconstructedTree -> ReconstructedTree -> Bool
$c== :: ReconstructedTree -> ReconstructedTree -> Bool
Eq)
maybeReconstructedTree :: EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree :: EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree Shoot {} = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"EpidemicTree is only a Shoot"
maybeReconstructedTree (Leaf EpidemicEvent
e) =
case EpidemicEvent
e of
IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
..} -> if Bool
indSampSeq
then ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation -> ReconstructedTree
RLeaf (EpidemicEvent -> Observation
Observation EpidemicEvent
e)
else String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Leaf with non-sequenced event individual sample"
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> if Bool
popSampSeq
then ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation -> ReconstructedTree
RLeaf (EpidemicEvent -> Observation
Observation EpidemicEvent
e)
else String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Leaf with non-sequenced event population sample"
EpidemicEvent
_ -> String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Bad leaf in the EpidemicTree"
maybeReconstructedTree (Branch e :: EpidemicEvent
e@Infection {} EpidemicTree
lt EpidemicTree
rt)
| EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt Bool -> Bool -> Bool
&& EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt = do
ReconstructedTree
rlt <- EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
lt
ReconstructedTree
rrt <- EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
rt
ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation
-> ReconstructedTree -> ReconstructedTree -> ReconstructedTree
RBranch (EpidemicEvent -> Observation
Observation EpidemicEvent
e) ReconstructedTree
rlt ReconstructedTree
rrt
| EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt = EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
lt
| EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt = EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
rt
| Bool
otherwise = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Neither subtree has a sequenced leaf"
maybeReconstructedTree Branch {} = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"EpidemicTree is a bad branch"
hasSequencedLeaf :: EpidemicTree -> Bool
hasSequencedLeaf :: EpidemicTree -> Bool
hasSequencedLeaf Shoot {} = Bool
False
hasSequencedLeaf (Leaf EpidemicEvent
e) =
case EpidemicEvent
e of
IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Bool
indSampSeq
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Bool
popSampSeq
EpidemicEvent
_ -> Bool
False
hasSequencedLeaf (Branch EpidemicEvent
_ EpidemicTree
lt EpidemicTree
rt) = EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt Bool -> Bool -> Bool
|| EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt
observedEvents :: [EpidemicEvent] -> Either String [Observation]
observedEvents :: [EpidemicEvent] -> Either String [Observation]
observedEvents [EpidemicEvent]
epiEvents = do
EpidemicTree
epiTree <- [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [EpidemicEvent]
epiEvents
let (PointProcessEvents [Observation]
unseqObss) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
epiTree
[Observation]
reconTreeEvents <-
if EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
epiTree
then ((ReconstructedTree -> [Observation])
-> Either String ReconstructedTree -> Either String [Observation]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ReconstructedTree -> [Observation]
reconstructedTreeEvents) (Either String ReconstructedTree -> Either String [Observation])
-> Either String ReconstructedTree -> Either String [Observation]
forall a b. (a -> b) -> a -> b
$ EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
epiTree
else [Observation] -> Either String [Observation]
forall a b. b -> Either a b
Right []
[Observation] -> Either String [Observation]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Observation] -> Either String [Observation])
-> [Observation] -> Either String [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation])
-> ([Observation] -> [Observation])
-> [Observation]
-> [Observation]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Observation] -> [Observation]
forall a. Eq a => [a] -> [a]
List.nub ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation]
unseqObss [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ [Observation]
reconTreeEvents
reconstructedTreeEvents :: ReconstructedTree -> [Observation]
reconstructedTreeEvents :: ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rt =
case ReconstructedTree
rt of
RBranch Observation
obs ReconstructedTree
rtl ReconstructedTree
rtr ->
[Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$
Observation
obs Observation -> [Observation] -> [Observation]
forall a. a -> [a] -> [a]
: (ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rtl [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rtr)
RLeaf Observation
obs -> [Observation
obs]