module ELynx.Tree.Simulate.Coalescent
( simulate,
)
where
import ELynx.Tree.Distribution.CoalescentContinuous
import ELynx.Tree.Length
import ELynx.Tree.Rooted
import Statistics.Distribution
import System.Random.Stateful
simulate ::
StatefulGen g m =>
Int ->
g ->
m (Tree Length Int)
simulate :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> g -> m (Tree Length Int)
simulate Int
n = forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' Int
n Int
0 Forest Length Int
trs
where
trs :: Forest Length Int
trs = [forall e a. e -> a -> Forest e a -> Tree e a
Node Length
0 Int
i [] | Int
i <- [Int
0 .. Int
n forall a. Num a => a -> a -> a
- Int
1]]
simulate' ::
StatefulGen g m =>
Int ->
Int ->
Forest Length Int ->
g ->
m (Tree Length Int)
simulate' :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' Int
n Int
a Forest Length Int
trs g
g
| Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot construct trees without leaves."
| Int
n forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Int
trs forall a. Eq a => a -> a -> Bool
/= Int
1 = forall a. HasCallStack => [Char] -> a
error [Char]
"Too many trees provided."
| Int
n forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Int
trs forall a. Eq a => a -> a -> Bool
== Int
1 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head Forest Length Int
trs
| Bool
otherwise = do
Int
i <- forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (Int
1, Int
n forall a. Num a => a -> a -> a
- Int
1) g
g
Length
t <- Double -> Length
toLengthUnsafe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall d g (m :: * -> *).
(ContGen d, StatefulGen g m) =>
d -> g -> m Double
genContVar (Int -> ExponentialDistribution
coalescentDistributionCont Int
n) g
g
let trs' :: Forest Length Int
trs' = forall a b. (a -> b) -> [a] -> [b]
map (forall e a. (e -> e) -> Tree e a -> Tree e a
modifyStem (forall a. Num a => a -> a -> a
+ Length
t)) Forest Length Int
trs
tl :: Tree Length Int
tl = Forest Length Int
trs' forall a. [a] -> Int -> a
!! (Int
i forall a. Num a => a -> a -> a
- Int
1)
tr :: Tree Length Int
tr = Forest Length Int
trs' forall a. [a] -> Int -> a
!! Int
i
tm :: Tree Length Int
tm = forall e a. e -> a -> Forest e a -> Tree e a
Node Length
0 Int
a [Tree Length Int
tl, Tree Length Int
tr]
trs'' :: Forest Length Int
trs'' = forall a. Int -> [a] -> [a]
take (Int
i forall a. Num a => a -> a -> a
- Int
1) Forest Length Int
trs' forall a. [a] -> [a] -> [a]
++ [Tree Length Int
tm] forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
drop (Int
i forall a. Num a => a -> a -> a
+ Int
1) Forest Length Int
trs'
forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
a Forest Length Int
trs'' g
g