module ELynx.Simulate.MarkovProcessAlongTree
(
simulateNSitesAlongTree
, simulateAndFlattenNSitesAlongTree
, simulateNSitesAlongTreeMixtureModel
, simulateAndFlattenNSitesAlongTreeMixtureModel
)
where
import Control.Monad
import Control.Monad.Primitive
import Data.Tree
import Numeric.LinearAlgebra
import System.Random.MWC
import System.Random.MWC.Distributions
import ELynx.Data.MarkovProcess.RateMatrix
import ELynx.Data.Tree.MeasurableTree
import ELynx.Simulate.MarkovProcess
measureableTreeToProbTree :: (Measurable a) => RateMatrix -> Tree a -> Tree ProbMatrix
measureableTreeToProbTree q = fmap (probMatrix q . getLen)
getRootStates :: PrimMonad m
=> Int -> StationaryDistribution -> Gen (PrimState m) -> m [State]
getRootStates n d g = replicateM n $ categorical d g
simulateAndFlattenNSitesAlongTree :: (PrimMonad m, Measurable a)
=> Int -> StationaryDistribution -> ExchangeabilityMatrix -> Tree a -> Gen (PrimState m) -> m [[State]]
simulateAndFlattenNSitesAlongTree n d e t g = do
let q = fromExchangeabilityMatrix e d
pt = measureableTreeToProbTree q t
is <- getRootStates n d g
simulateAndFlattenAlongProbTree is pt g
simulateAndFlattenAlongProbTree :: (PrimMonad m)
=> [State] -> Tree ProbMatrix -> Gen (PrimState m) -> m [[State]]
simulateAndFlattenAlongProbTree is (Node p f) g = do
is' <- mapM (\i -> jump i p g) is
if null f
then return [is']
else concat <$> sequence [simulateAndFlattenAlongProbTree is' t g | t <- f]
simulateNSitesAlongTree :: (PrimMonad m, Measurable a)
=> Int -> StationaryDistribution -> ExchangeabilityMatrix -> Tree a -> Gen (PrimState m) -> m (Tree [State])
simulateNSitesAlongTree n d e t g = do
let q = fromExchangeabilityMatrix e d
pt = measureableTreeToProbTree q t
is <- getRootStates n d g
simulateAlongProbTree is pt g
simulateAlongProbTree :: (PrimMonad m)
=> [State] -> Tree ProbMatrix -> Gen (PrimState m) -> m (Tree [State])
simulateAlongProbTree is (Node p f) g = do
is' <- mapM (\i -> jump i p g) is
f' <- sequence [simulateAlongProbTree is' t g | t <- f]
return $ Node is' f'
measureableTreeToProbTreeMixtureModel :: (Measurable a)
=> [RateMatrix] -> Tree a -> Tree [ProbMatrix]
measureableTreeToProbTreeMixtureModel qs =
fmap (\a -> [probMatrix q . getLen $ a | q <- qs])
getComponentsAndRootStates :: PrimMonad m
=> Int -> Vector R -> [StationaryDistribution] -> Gen (PrimState m) -> m ([Int], [State])
getComponentsAndRootStates n ws ds g = do
cs <- replicateM n $ categorical ws g
is <- sequence [ categorical (ds !! c) g | c <- cs ]
return (cs, is)
simulateAndFlattenNSitesAlongTreeMixtureModel :: (PrimMonad m, Measurable a)
=> Int -> Vector R -> [StationaryDistribution] -> [ExchangeabilityMatrix] -> Tree a
-> Gen (PrimState m) -> m [[State]]
simulateAndFlattenNSitesAlongTreeMixtureModel n ws ds es t g = do
let qs = zipWith fromExchangeabilityMatrix es ds
pt = measureableTreeToProbTreeMixtureModel qs t
(cs, is) <- getComponentsAndRootStates n ws ds g
simulateAndFlattenAlongProbTreeMixtureModel is cs pt g
simulateAndFlattenAlongProbTreeMixtureModel :: (PrimMonad m)
=> [State] -> [Int] -> Tree [ProbMatrix] -> Gen (PrimState m) -> m [[State]]
simulateAndFlattenAlongProbTreeMixtureModel is cs (Node ps f) g
= do is' <- sequence [ jump i (ps !! c) g | (i, c) <- zip is cs ]
if null f
then return [is']
else concat <$> sequence [ simulateAndFlattenAlongProbTreeMixtureModel is' cs t g | t <- f ]
simulateNSitesAlongTreeMixtureModel :: (PrimMonad m, Measurable a)
=> Int -> Vector R -> [StationaryDistribution] -> [ExchangeabilityMatrix] -> Tree a
-> Gen (PrimState m) -> m (Tree [State])
simulateNSitesAlongTreeMixtureModel n ws ds es t g = do
let qs = zipWith fromExchangeabilityMatrix es ds
pt = measureableTreeToProbTreeMixtureModel qs t
(cs, is) <- getComponentsAndRootStates n ws ds g
simulateAlongProbTreeMixtureModel is cs pt g
simulateAlongProbTreeMixtureModel :: (PrimMonad m)
=> [State] -> [Int] -> Tree [ProbMatrix] -> Gen (PrimState m) -> m (Tree [State])
simulateAlongProbTreeMixtureModel is cs (Node ps f) g = do
is' <- sequence [ jump i (ps !! c) g | (i, c) <- zip is cs ]
f' <- sequence [ simulateAlongProbTreeMixtureModel is' cs t g | t <- f ]
return $ Node is' f'