module ELynx.Simulate.MarkovProcessAlongTree
(
simulate
, simulateAndFlatten
, simulateMixtureModel
, simulateAndFlattenMixtureModel
, simulateAndFlattenMixtureModelPar
)
where
import Control.Monad
import Control.Monad.Primitive
import Control.Parallel.Strategies
import Data.Tree
import Numeric.LinearAlgebra
import System.Random.MWC (Gen, GenIO)
import System.Random.MWC.Distributions (categorical)
import ELynx.Data.MarkovProcess.RateMatrix
import ELynx.Data.Tree.MeasurableTree
import ELynx.Simulate.MarkovProcess
import ELynx.Tools.Concurrent
toProbTree :: (Measurable a) => RateMatrix -> Tree a -> Tree ProbMatrix
toProbTree q = fmap (probMatrix q . getLen)
getRootStates :: PrimMonad m
=> Int -> StationaryDistribution -> Gen (PrimState m) -> m [State]
getRootStates n d g = replicateM n $ categorical d g
simulateAndFlatten :: (PrimMonad m, Measurable a)
=> Int -> StationaryDistribution -> ExchangeabilityMatrix -> Tree a -> Gen (PrimState m) -> m [[State]]
simulateAndFlatten n d e t g = do
let q = fromExchangeabilityMatrix e d
pt = toProbTree q t
is <- getRootStates n d g
simulateAndFlatten' is pt g
simulateAndFlatten' :: (PrimMonad m)
=> [State] -> Tree ProbMatrix -> Gen (PrimState m) -> m [[State]]
simulateAndFlatten' is (Node p f) g = do
is' <- mapM (\i -> jump i p g) is
if null f
then return [is']
else concat <$> sequence [simulateAndFlatten' is' t g | t <- f]
simulate :: (PrimMonad m, Measurable a)
=> Int -> StationaryDistribution -> ExchangeabilityMatrix -> Tree a -> Gen (PrimState m) -> m (Tree [State])
simulate n d e t g = do
let q = fromExchangeabilityMatrix e d
pt = toProbTree q t
is <- getRootStates n d g
simulate' is pt g
simulate' :: (PrimMonad m)
=> [State] -> Tree ProbMatrix -> Gen (PrimState m) -> m (Tree [State])
simulate' is (Node p f) g = do
is' <- mapM (\i -> jump i p g) is
f' <- sequence [simulate' is' t g | t <- f]
return $ Node is' f'
toProbTreeMixtureModel :: (Measurable a)
=> [RateMatrix] -> Tree a -> Tree [ProbMatrix]
toProbTreeMixtureModel qs =
fmap (\a -> [probMatrix q . getLen $ a | q <- qs] `using` parList rpar)
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)
simulateAndFlattenMixtureModel :: (PrimMonad m, Measurable a)
=> Int -> Vector R -> [StationaryDistribution] -> [ExchangeabilityMatrix] -> Tree a
-> Gen (PrimState m) -> m [[State]]
simulateAndFlattenMixtureModel n ws ds es t g = do
let qs = zipWith fromExchangeabilityMatrix es ds
pt = toProbTreeMixtureModel qs t
(cs, is) <- getComponentsAndRootStates n ws ds g
simulateAndFlattenMixtureModel' is cs pt g
simulateAndFlattenMixtureModel' :: (PrimMonad m)
=> [State] -> [Int] -> Tree [ProbMatrix] -> Gen (PrimState m) -> m [[State]]
simulateAndFlattenMixtureModel' 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 [ simulateAndFlattenMixtureModel' is' cs t g | t <- f ]
simulateAndFlattenMixtureModelPar
:: Measurable a
=> Int -> Vector R -> [StationaryDistribution] -> [ExchangeabilityMatrix] -> Tree a
-> GenIO -> IO [[[State]]]
simulateAndFlattenMixtureModelPar n ws ds es t g = do
let qs = zipWith fromExchangeabilityMatrix es ds
pt = toProbTreeMixtureModel qs t
parComp n (\n' g' -> do
(cs, is) <- getComponentsAndRootStates n' ws ds g'
simulateAndFlattenMixtureModel' is cs pt g') g
simulateMixtureModel :: (PrimMonad m, Measurable a)
=> Int -> Vector R -> [StationaryDistribution] -> [ExchangeabilityMatrix] -> Tree a
-> Gen (PrimState m) -> m (Tree [State])
simulateMixtureModel n ws ds es t g = do
let qs = zipWith fromExchangeabilityMatrix es ds
pt = toProbTreeMixtureModel qs t
(cs, is) <- getComponentsAndRootStates n ws ds g
simulateMixtureModel' is cs pt g
simulateMixtureModel' :: (PrimMonad m)
=> [State] -> [Int] -> Tree [ProbMatrix] -> Gen (PrimState m) -> m (Tree [State])
simulateMixtureModel' is cs (Node ps f) g = do
is' <- sequence [ jump i (ps !! c) g | (i, c) <- zip is cs ]
f' <- sequence [ simulateMixtureModel' is' cs t g | t <- f ]
return $ Node is' f'