{-# LANGUAGE BangPatterns #-}
module ELynx.Simulate.PointProcess
( PointProcess(..)
, TimeSpec
, simulate
, toReconstructedTree
, simulateReconstructedTree
, simulateNReconstructedTrees
) where
import Control.Monad
import Control.Monad.Primitive
import Data.List (mapAccumL)
import Data.Tree
import qualified Statistics.Distribution as D (genContVar)
import System.Random.MWC
import ELynx.Data.Tree.MeasurableTree
import ELynx.Data.Tree.PhyloTree
import ELynx.Data.Tree.Tree
import ELynx.Distribution.BirthDeath
import ELynx.Distribution.BirthDeathCritical
import ELynx.Distribution.BirthDeathCriticalNoTime
import ELynx.Distribution.BirthDeathNearlyCritical
import ELynx.Distribution.TimeOfOrigin
import ELynx.Distribution.TimeOfOriginNearCritical
import ELynx.Distribution.Types
import ELynx.Tools.Equality
import ELynx.Tools.List
epsNearCriticalPointProcess :: Double
epsNearCriticalPointProcess = 1e-5
epsNearCriticalTimeOfOrigin :: Double
epsNearCriticalTimeOfOrigin = 1e-8
data PointProcess a b = PointProcess
{ points :: ![a]
, values :: ![b]
, origin :: !b } deriving (Read, Show, Eq)
type TimeSpec = Maybe (Time, Bool)
simulate :: (PrimMonad m)
=> Int
-> TimeSpec
-> Rate
-> Rate
-> Gen (PrimState m)
-> m (PointProcess Int Double)
simulate n Nothing l m g
| m > l = error "Time of origin distribution formula not available when mu > lambda. Please specify height for the moment."
| m =~= l = do
!vs <- replicateM (n-1) (D.genContVar (BDCNTD l) g)
let t = maximum vs
return $ PointProcess [0..(n-1)] vs t
| abs (m-l) <= epsNearCriticalTimeOfOrigin = do
t <- D.genContVar (TONCD n l m) g
simulate n (Just (t, False)) l m g
| otherwise = do
t <- D.genContVar (TOD n l m) g
simulate n (Just (t, False)) l m g
simulate n (Just (t, c)) l m g
| n < 1 = error "Number of samples needs to be one or larger."
| t < 0.0 = error "Time of origin needs to be positive."
| l < 0.0 = error "Birth rate needs to be positive."
| (m =~= l) && not c = do
!vs <- replicateM (n-1) (D.genContVar (BDCD t l) g)
return $ PointProcess [0..(n-1)] vs t
| (abs (m - l) <= epsNearCriticalPointProcess) && not c = do
!vs <- replicateM (n-1) (D.genContVar (BDNCD t l m) g)
return $ PointProcess [0..(n-1)] vs t
| not c = do
!vs <- replicateM (n-1) (D.genContVar (BDD t l m) g)
return $ PointProcess [0..(n-1)] vs t
| (m =~= l) && c = do
!vs <- replicateM (n-2) (D.genContVar (BDCD t l) g)
vs' <- randomInsert t vs g
return $ PointProcess [0..(n-1)] vs' t
| (abs (m - l) <= epsNearCriticalPointProcess) && c = do
!vs <- replicateM (n-2) (D.genContVar (BDNCD t l m) g)
vs' <- randomInsert t vs g
return $ PointProcess [0..(n-1)] vs' t
| c = do
!vs <- replicateM (n-2) (D.genContVar (BDD t l m) g)
vs' <- randomInsert t vs g
return $ PointProcess [0..(n-1)] vs' t
| otherwise = error "simulate: Fell through guard, this should never happen."
sort :: (Ord b) => PointProcess a b -> ([b], [Int])
sort (PointProcess _ vs _) = (vsSorted, isSorted)
where vsIsSorted = sortWithIndices vs
vsSorted = map fst vsIsSorted
isSorted = flattenIndices $ map snd vsIsSorted
flattenIndices :: [Int] -> [Int]
flattenIndices is = snd $ mapAccumL fAcc [] is
fAcc :: [Int] -> Int -> ([Int], Int)
fAcc is i = (i:is, i')
where i' = i - length (filter (<i) is)
simulateNReconstructedTrees
:: (PrimMonad m)
=> Int
-> Int
-> TimeSpec
-> Rate
-> Rate
-> Gen (PrimState m)
-> m [Tree PhyloIntLabel]
simulateNReconstructedTrees nT nP t l m g
| nT <= 0 = return []
| otherwise = replicateM nT $ simulateReconstructedTree nP t l m g
simulateReconstructedTree
:: (PrimMonad m)
=> Int
-> TimeSpec
-> Rate
-> Rate
-> Gen (PrimState m)
-> m (Tree PhyloIntLabel)
simulateReconstructedTree n t l m g = toReconstructedTree <$> simulate n t l m g
toReconstructedTree :: PointProcess Int Double
-> Tree PhyloIntLabel
toReconstructedTree pp@(PointProcess ps vs o)
| length ps /= length vs + 1 = error "Too few or too many points."
| length vs <= 1 = error "Too few values."
| otherwise = treeOrigin
where (vsSorted, isSorted) = sort pp
!lvs = [ singleton (PhyloLabel p Nothing 0) | p <- ps ]
!heights = replicate (length ps) 0
!treeRoot = toReconstructedTree' isSorted vsSorted lvs heights
!h = last vsSorted
!treeOrigin = lengthenRoot (o-h) treeRoot
toReconstructedTree' :: [Int]
-> [Double]
-> [Tree PhyloIntLabel]
-> [Double]
-> Tree PhyloIntLabel
toReconstructedTree' [] [] trs _ = head trs
toReconstructedTree' is vs trs hs = toReconstructedTree' is' vs' trs'' hs'
where !i = head is
!is' = tail is
!v = head vs
!vs' = tail vs
!hl = hs !! i
!hr = hs !! (i+1)
!dvl = v - hl
!dvr = v - hr
!tl = lengthenRoot dvl $ trs !! i
!tr = lengthenRoot dvr $ trs !! (i+1)
!h' = hl + dvl
!tm = Node (PhyloLabel 0 Nothing 0) [tl, tr]
!trs'' = take i trs ++ [tm] ++ drop (i+2) trs
!hs' = take i hs ++ [h'] ++ drop (i+2) hs