module ELynx.Simulate.Coalescent
( simulate,
)
where
import Control.Monad.Primitive
import ELynx.Data.Tree.Measurable
import ELynx.Data.Tree.Phylogeny
import ELynx.Data.Tree.Rooted
import ELynx.Distribution.CoalescentContinuous
import Statistics.Distribution
import System.Random.MWC
simulate ::
(PrimMonad m) =>
Int ->
Gen (PrimState m) ->
m (Tree Length Int)
simulate n = simulate' n 0 trs
where
trs = [Node (Length 0) i [] | i <- [0 .. n - 1]]
simulate' ::
(PrimMonad m) =>
Int ->
Int ->
Forest Length Int ->
Gen (PrimState m) ->
m (Tree Length Int)
simulate' n a trs g
| n <= 0 = error "Cannot construct trees without leaves."
| n == 1 && length trs /= 1 = error "Too many trees provided."
| n == 1 && length trs == 1 = return $ head trs
| otherwise = do
i <- uniformR (1, n - 1) g
t <- genContVar (coalescentDistributionCont n) g
let trs' = map (applyStem (+ t)) trs
tl = trs' !! (i - 1)
tr = trs' !! i
tm = Node (Length 0) a [tl, tr]
trs'' = take (i - 1) trs' ++ [tm] ++ drop (i + 1) trs'
simulate' (n - 1) a trs'' g