{-# LANGUAGE ScopedTypeVariables #-}
module ELynx.Simulate.MarkovProcessAlongTree
(
simulate,
simulateAndFlatten,
simulateMixtureModel,
simulateAndFlattenMixtureModel,
simulateAndFlattenMixtureModelPar,
)
where
import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad
import Control.Monad.Primitive
import Control.Parallel.Strategies
import Data.Tree
import qualified Data.Vector.Storable as V
import Data.Word (Word32)
import ELynx.Data.MarkovProcess.RateMatrix
import ELynx.Simulate.MarkovProcess
import Numeric.LinearAlgebra
import System.Random.MWC
import System.Random.MWC.Distributions (categorical)
toProbTree :: RateMatrix -> Tree Double -> Tree ProbMatrix
toProbTree q = fmap (probMatrix q)
getRootStates ::
PrimMonad m =>
Int ->
StationaryDistribution ->
Gen (PrimState m) ->
m [State]
getRootStates n d g = replicateM n $ categorical d g
simulateAndFlatten ::
PrimMonad m =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
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 =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
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 ::
[RateMatrix] -> Tree Double -> Tree [ProbMatrix]
toProbTreeMixtureModel qs =
fmap (\a -> [probMatrix q 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 =>
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
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]
getChunks :: Int -> Int -> [Int]
getChunks c n = ns
where
n' = n `div` c
r = n `mod` c
ns = replicate r (n' + 1) ++ replicate (c - r) n'
splitGen :: PrimMonad m => Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen n gen
| n <= 0 = return []
| otherwise = do
seeds :: [V.Vector Word32] <- replicateM (n -1) $ uniformVector gen 256
fmap (gen :) (mapM initialize seeds)
parComp :: Int -> (Int -> GenIO -> IO b) -> GenIO -> IO [b]
parComp num fun gen = do
ncap <- getNumCapabilities
let chunks = getChunks ncap num
gs <- splitGen ncap gen
mapConcurrently (uncurry fun) (zip chunks gs)
simulateAndFlattenMixtureModelPar ::
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
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 =>
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
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'