{-# 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 :: RateMatrix -> Tree Double -> Tree RateMatrix
toProbTree RateMatrix
q = (Double -> RateMatrix) -> Tree Double -> Tree RateMatrix
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (RateMatrix -> Double -> RateMatrix
probMatrix RateMatrix
q)
getRootStates ::
PrimMonad m =>
Int ->
StationaryDistribution ->
Gen (PrimState m) ->
m [State]
getRootStates :: Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
getRootStates Int
n StationaryDistribution
d Gen (PrimState m)
g = Int -> m Int -> m [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (m Int -> m [Int]) -> m Int -> m [Int]
forall a b. (a -> b) -> a -> b
$ StationaryDistribution -> Gen (PrimState m) -> m Int
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Double) =>
v Double -> Gen (PrimState m) -> m Int
categorical StationaryDistribution
d Gen (PrimState m)
g
simulateAndFlatten ::
PrimMonad m =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
Gen (PrimState m) ->
m [[State]]
simulateAndFlatten :: Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> Gen (PrimState m)
-> m [[Int]]
simulateAndFlatten Int
n StationaryDistribution
d RateMatrix
e Tree Double
t Gen (PrimState m)
g = do
let q :: RateMatrix
q = RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix RateMatrix
e StationaryDistribution
d
pt :: Tree RateMatrix
pt = RateMatrix -> Tree Double -> Tree RateMatrix
toProbTree RateMatrix
q Tree Double
t
[Int]
is <- Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
forall (m :: * -> *).
PrimMonad m =>
Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
getRootStates Int
n StationaryDistribution
d Gen (PrimState m)
g
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt Gen (PrimState m)
g
simulateAndFlatten' ::
(PrimMonad m) =>
[State] ->
Tree ProbMatrix ->
Gen (PrimState m) ->
m [[State]]
simulateAndFlatten' :: [Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
simulateAndFlatten' [Int]
is (Node RateMatrix
p Forest RateMatrix
f) Gen (PrimState m)
g = do
[Int]
is' <- (Int -> m Int) -> [Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
i -> Int -> RateMatrix -> Gen (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> RateMatrix -> Gen (PrimState m) -> m Int
jump Int
i RateMatrix
p Gen (PrimState m)
g) [Int]
is
if Forest RateMatrix -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Forest RateMatrix
f
then [[Int]] -> m [[Int]]
forall (m :: * -> *) a. Monad m => a -> m a
return [[Int]
is']
else [[[Int]]] -> [[Int]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[[Int]]] -> [[Int]]) -> m [[[Int]]] -> m [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [m [[Int]]] -> m [[[Int]]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
simulateAndFlatten' [Int]
is' Tree RateMatrix
t Gen (PrimState m)
g | Tree RateMatrix
t <- Forest RateMatrix
f]
simulate ::
PrimMonad m =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
Gen (PrimState m) ->
m (Tree [State])
simulate :: Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> Gen (PrimState m)
-> m (Tree [Int])
simulate Int
n StationaryDistribution
d RateMatrix
e Tree Double
t Gen (PrimState m)
g = do
let q :: RateMatrix
q = RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix RateMatrix
e StationaryDistribution
d
pt :: Tree RateMatrix
pt = RateMatrix -> Tree Double -> Tree RateMatrix
toProbTree RateMatrix
q Tree Double
t
[Int]
is <- Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
forall (m :: * -> *).
PrimMonad m =>
Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
getRootStates Int
n StationaryDistribution
d Gen (PrimState m)
g
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m (Tree [Int])
simulate' [Int]
is Tree RateMatrix
pt Gen (PrimState m)
g
simulate' ::
(PrimMonad m) =>
[State] ->
Tree ProbMatrix ->
Gen (PrimState m) ->
m (Tree [State])
simulate' :: [Int] -> Tree RateMatrix -> Gen (PrimState m) -> m (Tree [Int])
simulate' [Int]
is (Node RateMatrix
p Forest RateMatrix
f) Gen (PrimState m)
g = do
[Int]
is' <- (Int -> m Int) -> [Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
i -> Int -> RateMatrix -> Gen (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> RateMatrix -> Gen (PrimState m) -> m Int
jump Int
i RateMatrix
p Gen (PrimState m)
g) [Int]
is
[Tree [Int]]
f' <- [m (Tree [Int])] -> m [Tree [Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m (Tree [Int])
simulate' [Int]
is' Tree RateMatrix
t Gen (PrimState m)
g | Tree RateMatrix
t <- Forest RateMatrix
f]
Tree [Int] -> m (Tree [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return (Tree [Int] -> m (Tree [Int])) -> Tree [Int] -> m (Tree [Int])
forall a b. (a -> b) -> a -> b
$ [Int] -> [Tree [Int]] -> Tree [Int]
forall a. a -> Forest a -> Tree a
Node [Int]
is' [Tree [Int]]
f'
toProbTreeMixtureModel ::
[RateMatrix] -> Tree Double -> Tree [ProbMatrix]
toProbTreeMixtureModel :: [RateMatrix] -> Tree Double -> Tree [RateMatrix]
toProbTreeMixtureModel [RateMatrix]
qs =
(Double -> [RateMatrix]) -> Tree Double -> Tree [RateMatrix]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Double
a -> [RateMatrix -> Double -> RateMatrix
probMatrix RateMatrix
q Double
a | RateMatrix
q <- [RateMatrix]
qs] [RateMatrix] -> Strategy [RateMatrix] -> [RateMatrix]
forall a. a -> Strategy a -> a
`using` Strategy RateMatrix -> Strategy [RateMatrix]
forall a. Strategy a -> Strategy [a]
parList Strategy RateMatrix
forall a. Strategy a
rpar)
getComponentsAndRootStates ::
PrimMonad m =>
Int ->
Vector R ->
[StationaryDistribution] ->
Gen (PrimState m) ->
m ([Int], [State])
getComponentsAndRootStates :: Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n StationaryDistribution
ws [StationaryDistribution]
ds Gen (PrimState m)
g = do
[Int]
cs <- Int -> m Int -> m [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (m Int -> m [Int]) -> m Int -> m [Int]
forall a b. (a -> b) -> a -> b
$ StationaryDistribution -> Gen (PrimState m) -> m Int
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Double) =>
v Double -> Gen (PrimState m) -> m Int
categorical StationaryDistribution
ws Gen (PrimState m)
g
[Int]
is <- [m Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [StationaryDistribution -> Gen (PrimState m) -> m Int
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Double) =>
v Double -> Gen (PrimState m) -> m Int
categorical ([StationaryDistribution]
ds [StationaryDistribution] -> Int -> StationaryDistribution
forall a. [a] -> Int -> a
!! Int
c) Gen (PrimState m)
g | Int
c <- [Int]
cs]
([Int], [Int]) -> m ([Int], [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [Int]
is)
simulateAndFlattenMixtureModel ::
PrimMonad m =>
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
Gen (PrimState m) ->
m [[State]]
simulateAndFlattenMixtureModel :: Int
-> StationaryDistribution
-> [StationaryDistribution]
-> [RateMatrix]
-> Tree Double
-> Gen (PrimState m)
-> m [[Int]]
simulateAndFlattenMixtureModel Int
n StationaryDistribution
ws [StationaryDistribution]
ds [RateMatrix]
es Tree Double
t Gen (PrimState m)
g = do
let qs :: [RateMatrix]
qs = (RateMatrix -> StationaryDistribution -> RateMatrix)
-> [RateMatrix] -> [StationaryDistribution] -> [RateMatrix]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix [RateMatrix]
es [StationaryDistribution]
ds
pt :: Tree [RateMatrix]
pt = [RateMatrix] -> Tree Double -> Tree [RateMatrix]
toProbTreeMixtureModel [RateMatrix]
qs Tree Double
t
([Int]
cs, [Int]
is) <- Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
forall (m :: * -> *).
PrimMonad m =>
Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n StationaryDistribution
ws [StationaryDistribution]
ds Gen (PrimState m)
g
[Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs Tree [RateMatrix]
pt Gen (PrimState m)
g
simulateAndFlattenMixtureModel' ::
(PrimMonad m) =>
[State] ->
[Int] ->
Tree [ProbMatrix] ->
Gen (PrimState m) ->
m [[State]]
simulateAndFlattenMixtureModel' :: [Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs (Node [RateMatrix]
ps Forest [RateMatrix]
f) Gen (PrimState m)
g = do
[Int]
is' <- [m Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Int -> RateMatrix -> Gen (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> RateMatrix -> Gen (PrimState m) -> m Int
jump Int
i ([RateMatrix]
ps [RateMatrix] -> Int -> RateMatrix
forall a. [a] -> Int -> a
!! Int
c) Gen (PrimState m)
g | (Int
i, Int
c) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
is [Int]
cs]
if Forest [RateMatrix] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Forest [RateMatrix]
f
then [[Int]] -> m [[Int]]
forall (m :: * -> *) a. Monad m => a -> m a
return [[Int]
is']
else
[[[Int]]] -> [[Int]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
([[[Int]]] -> [[Int]]) -> m [[[Int]]] -> m [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [m [[Int]]] -> m [[[Int]]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [[Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is' [Int]
cs Tree [RateMatrix]
t Gen (PrimState m)
g | Tree [RateMatrix]
t <- Forest [RateMatrix]
f]
getChunks :: Int -> Int -> [Int]
getChunks :: Int -> Int -> [Int]
getChunks Int
c Int
n = [Int]
ns
where
n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
c
r :: Int
r = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
c
ns :: [Int]
ns = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
r (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r) Int
n'
splitGen :: PrimMonad m => Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen :: Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
n Gen (PrimState m)
gen
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Gen (PrimState m)] -> m [Gen (PrimState m)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
| Bool
otherwise = do
[Vector Word32]
seeds :: [V.Vector Word32] <- Int -> m (Vector Word32) -> m [Vector Word32]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (m (Vector Word32) -> m [Vector Word32])
-> m (Vector Word32) -> m [Vector Word32]
forall a b. (a -> b) -> a -> b
$ Gen (PrimState m) -> Int -> m (Vector Word32)
forall (m :: * -> *) a (v :: * -> *).
(PrimMonad m, Variate a, Vector v a) =>
Gen (PrimState m) -> Int -> m (v a)
uniformVector Gen (PrimState m)
gen Int
256
([Gen (PrimState m)] -> [Gen (PrimState m)])
-> m [Gen (PrimState m)] -> m [Gen (PrimState m)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Gen (PrimState m)
gen Gen (PrimState m) -> [Gen (PrimState m)] -> [Gen (PrimState m)]
forall a. a -> [a] -> [a]
:) ((Vector Word32 -> m (Gen (PrimState m)))
-> [Vector Word32] -> m [Gen (PrimState m)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Vector Word32 -> m (Gen (PrimState m))
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
initialize [Vector Word32]
seeds)
parComp :: Int -> (Int -> GenIO -> IO b) -> GenIO -> IO [b]
parComp :: Int -> (Int -> GenIO -> IO b) -> GenIO -> IO [b]
parComp Int
num Int -> GenIO -> IO b
fun GenIO
gen = do
Int
ncap <- IO Int
getNumCapabilities
let chunks :: [Int]
chunks = Int -> Int -> [Int]
getChunks Int
ncap Int
num
[Gen RealWorld]
gs <- Int -> GenIO -> IO [GenIO]
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
ncap GenIO
gen
((Int, Gen RealWorld) -> IO b) -> [(Int, Gen RealWorld)] -> IO [b]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently ((Int -> Gen RealWorld -> IO b) -> (Int, Gen RealWorld) -> IO b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> Gen RealWorld -> IO b
Int -> GenIO -> IO b
fun) ([Int] -> [Gen RealWorld] -> [(Int, Gen RealWorld)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [Gen RealWorld]
gs)
simulateAndFlattenMixtureModelPar ::
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
GenIO ->
IO [[[State]]]
simulateAndFlattenMixtureModelPar :: Int
-> StationaryDistribution
-> [StationaryDistribution]
-> [RateMatrix]
-> Tree Double
-> GenIO
-> IO [[[Int]]]
simulateAndFlattenMixtureModelPar Int
n StationaryDistribution
ws [StationaryDistribution]
ds [RateMatrix]
es Tree Double
t GenIO
g = do
let qs :: [RateMatrix]
qs = (RateMatrix -> StationaryDistribution -> RateMatrix)
-> [RateMatrix] -> [StationaryDistribution] -> [RateMatrix]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix [RateMatrix]
es [StationaryDistribution]
ds
pt :: Tree [RateMatrix]
pt = [RateMatrix] -> Tree Double -> Tree [RateMatrix]
toProbTreeMixtureModel [RateMatrix]
qs Tree Double
t
Int -> (Int -> GenIO -> IO [[Int]]) -> GenIO -> IO [[[Int]]]
forall b. Int -> (Int -> GenIO -> IO b) -> GenIO -> IO [b]
parComp
Int
n
( \Int
n' GenIO
g' -> do
([Int]
cs, [Int]
is) <- Int
-> StationaryDistribution
-> [StationaryDistribution]
-> GenIO
-> IO ([Int], [Int])
forall (m :: * -> *).
PrimMonad m =>
Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n' StationaryDistribution
ws [StationaryDistribution]
ds GenIO
g'
[Int] -> [Int] -> Tree [RateMatrix] -> GenIO -> IO [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int] -> Tree [RateMatrix] -> Gen (PrimState m) -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs Tree [RateMatrix]
pt GenIO
g'
)
GenIO
g
simulateMixtureModel ::
PrimMonad m =>
Int ->
Vector R ->
[StationaryDistribution] ->
[ExchangeabilityMatrix] ->
Tree Double ->
Gen (PrimState m) ->
m (Tree [State])
simulateMixtureModel :: Int
-> StationaryDistribution
-> [StationaryDistribution]
-> [RateMatrix]
-> Tree Double
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel Int
n StationaryDistribution
ws [StationaryDistribution]
ds [RateMatrix]
es Tree Double
t Gen (PrimState m)
g = do
let qs :: [RateMatrix]
qs = (RateMatrix -> StationaryDistribution -> RateMatrix)
-> [RateMatrix] -> [StationaryDistribution] -> [RateMatrix]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix [RateMatrix]
es [StationaryDistribution]
ds
pt :: Tree [RateMatrix]
pt = [RateMatrix] -> Tree Double -> Tree [RateMatrix]
toProbTreeMixtureModel [RateMatrix]
qs Tree Double
t
([Int]
cs, [Int]
is) <- Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
forall (m :: * -> *).
PrimMonad m =>
Int
-> StationaryDistribution
-> [StationaryDistribution]
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n StationaryDistribution
ws [StationaryDistribution]
ds Gen (PrimState m)
g
[Int]
-> [Int]
-> Tree [RateMatrix]
-> Gen (PrimState m)
-> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree [RateMatrix]
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs Tree [RateMatrix]
pt Gen (PrimState m)
g
simulateMixtureModel' ::
(PrimMonad m) =>
[State] ->
[Int] ->
Tree [ProbMatrix] ->
Gen (PrimState m) ->
m (Tree [State])
simulateMixtureModel' :: [Int]
-> [Int]
-> Tree [RateMatrix]
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs (Node [RateMatrix]
ps Forest [RateMatrix]
f) Gen (PrimState m)
g = do
[Int]
is' <- [m Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Int -> RateMatrix -> Gen (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> RateMatrix -> Gen (PrimState m) -> m Int
jump Int
i ([RateMatrix]
ps [RateMatrix] -> Int -> RateMatrix
forall a. [a] -> Int -> a
!! Int
c) Gen (PrimState m)
g | (Int
i, Int
c) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
is [Int]
cs]
[Tree [Int]]
f' <- [m (Tree [Int])] -> m [Tree [Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [[Int]
-> [Int]
-> Tree [RateMatrix]
-> Gen (PrimState m)
-> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree [RateMatrix]
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is' [Int]
cs Tree [RateMatrix]
t Gen (PrimState m)
g | Tree [RateMatrix]
t <- Forest [RateMatrix]
f]
Tree [Int] -> m (Tree [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return (Tree [Int] -> m (Tree [Int])) -> Tree [Int] -> m (Tree [Int])
forall a b. (a -> b) -> a -> b
$ [Int] -> [Tree [Int]] -> Tree [Int]
forall a. a -> Forest a -> Tree a
Node [Int]
is' [Tree [Int]]
f'