{-# LANGUAGE ScopedTypeVariables #-}

-- |
--   Description :  Work with transition probability matrices on rooted trees
--   Copyright   :  2021 Dominik Schrempf
--   License     :  GPLv3
--
--   Maintainer  :  dominik.schrempf@gmail.com
--   Stability   :  unstable
--   Portability :  non-portable (not tested)
--
-- Calculate transition probability matrices, map rate matrices on trees, populate
-- a tree with states according to a stationary distribution, etc.
--
-- The implementation of the Markov process is more than basic and can be improved
-- in a lot of ways.
module ELynx.Simulate.MarkovProcessAlongTree
  ( -- * Single rate matrix.
    simulate,
    simulateAndFlatten,
    simulateAndFlattenPar,

    -- * Mixture models.
    simulateMixtureModel,
    simulateAndFlattenMixtureModel,
    simulateAndFlattenMixtureModelPar,
  )
where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad
import Data.Tree
import qualified Data.Vector as V
import ELynx.MarkovProcess.RateMatrix
import ELynx.Simulate.MarkovProcess
import System.Random.MWC.Distributions
import System.Random.Stateful

-- XXX @performace. The horizontal concatenation might be slow. If so,
-- 'concatenateSeqs' or 'concatenateAlignments' can be used, which directly
-- appends vectors.

-- A brain f***. As an example, let @xss@ be a list of alignments (i.e., a list
-- of a list of a list of alleles). This function horizontally concatenates the
-- sites. The number of species needs to be same in each alignment. No checks
-- are performed!
horizontalConcat :: [[[a]]] -> [[a]]
horizontalConcat :: forall a. [[[a]]] -> [[a]]
horizontalConcat [[[a]]
xs] = [[a]]
xs
horizontalConcat [[[a]]]
xss = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. [a] -> [a] -> [a]
(++)) [[[a]]]
xss

toProbTree :: RateMatrix -> Tree Double -> Tree ProbMatrix
toProbTree :: RateMatrix -> Tree Double -> Tree RateMatrix
toProbTree RateMatrix
q = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (RateMatrix -> Double -> RateMatrix
probMatrix RateMatrix
q)

getRootStates ::
  StatefulGen g m =>
  Int ->
  StationaryDistribution ->
  g ->
  m [State]
getRootStates :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
g = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical StationaryDistribution
d g
g

-- | Simulate a number of sites for a given substitution model. Only the states
-- at the leafs are retained. The states at internal nodes are removed. This has
-- a lower memory footprint.

-- XXX: Improve performance. Use vectors, not lists. I am actually not sure if
-- this improves performance...
simulateAndFlatten ::
  StatefulGen g m =>
  Int ->
  StationaryDistribution ->
  ExchangeabilityMatrix ->
  Tree Double ->
  g ->
  m [[State]]
simulateAndFlatten :: forall g (m :: * -> *).
StatefulGen g m =>
Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> g
-> m [[Int]]
simulateAndFlatten Int
n StationaryDistribution
d RateMatrix
e Tree Double
t g
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 <- forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
g
  forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt g
g

-- This is the heart of the simulation. Take a tree and a list of root states.
-- Recursively jump down the branches to the leafs. Forget states at internal
-- nodes.
simulateAndFlatten' ::
  StatefulGen g m =>
  [State] ->
  Tree ProbMatrix ->
  g ->
  m [[State]]
simulateAndFlatten' :: forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is (Node RateMatrix
p [Tree RateMatrix]
f) g
g = do
  [Int]
is' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
i -> forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i RateMatrix
p g
g) [Int]
is
  if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Tree RateMatrix]
f
    then forall (m :: * -> *) a. Monad m => a -> m a
return [[Int]
is']
    else forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is' Tree RateMatrix
t g
g | Tree RateMatrix
t <- [Tree RateMatrix]
f]

-- | See 'simulateAndFlatten', parallel version.
simulateAndFlattenPar ::
  RandomGen g =>
  Int ->
  StationaryDistribution ->
  ExchangeabilityMatrix ->
  Tree Double ->
  IOGenM g ->
  IO [[State]]
simulateAndFlattenPar :: forall g.
RandomGen g =>
Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> IOGenM g
-> IO [[Int]]
simulateAndFlattenPar Int
n StationaryDistribution
d RateMatrix
e Tree Double
t IOGenM g
g = do
  Int
c <- IO Int
getNumCapabilities
  [g]
rs <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
c forall a b. (a -> b) -> a -> b
$ forall g r (m :: * -> *). RandomGenM g r m => g -> m r
splitGenM IOGenM g
g
  [IOGenM g]
gs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
newIOGenM [g]
rs
  let chunks :: [Int]
chunks = Int -> Int -> [Int]
getChunks Int
c Int
n
      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
  -- The concurrent map returns a list of [[State]] objects. They have to be
  -- concatenated horizontally.
  [[[Int]]]
sss <-
    forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently
      ( \(Int
num, IOGenM g
gen) -> do
          [Int]
is <- forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
num StationaryDistribution
d IOGenM g
gen
          forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt IOGenM g
gen
      )
      (forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [IOGenM g]
gs)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. [[[a]]] -> [[a]]
horizontalConcat [[[Int]]]
sss

-- | Simulate a number of sites for a given substitution model. Keep states at
-- internal nodes. The result is a tree with the list of simulated states as
-- node labels.
simulate ::
  StatefulGen g m =>
  Int ->
  StationaryDistribution ->
  ExchangeabilityMatrix ->
  Tree Double ->
  g ->
  m (Tree [State])
simulate :: forall g (m :: * -> *).
StatefulGen g m =>
Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> g
-> m (Tree [Int])
simulate Int
n StationaryDistribution
d RateMatrix
e Tree Double
t g
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 <- forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
g
  forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is Tree RateMatrix
pt g
g

-- This is the heart of the simulation. Take a tree and a list of root states.
-- Recursively jump down the branches to the leafs.
simulate' ::
  StatefulGen g m =>
  [State] ->
  Tree ProbMatrix ->
  g ->
  m (Tree [State])
simulate' :: forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is (Node RateMatrix
p [Tree RateMatrix]
f) g
g = do
  [Int]
is' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
i -> forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i RateMatrix
p g
g) [Int]
is
  [Tree [Int]]
f' <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is' Tree RateMatrix
t g
g | Tree RateMatrix
t <- [Tree RateMatrix]
f]
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> [Tree a] -> Tree a
Node [Int]
is' [Tree [Int]]
f'

toProbTreeMixtureModel ::
  V.Vector RateMatrix -> Tree Double -> Tree (V.Vector ProbMatrix)
toProbTreeMixtureModel :: Vector RateMatrix -> Tree Double -> Tree (Vector RateMatrix)
toProbTreeMixtureModel Vector RateMatrix
qs =
  -- XXX: This function is slow. Parallelization?
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Double
a -> forall a b. (a -> b) -> Vector a -> Vector b
V.map (RateMatrix -> Double -> RateMatrix
`probMatrix` Double
a) Vector RateMatrix
qs)

getComponentsAndRootStates ::
  StatefulGen g m =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  g ->
  m ([Int], [State])
getComponentsAndRootStates :: forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> g
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector StationaryDistribution
ds g
g = do
  [Int]
cs <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical Vector Double
ws g
g
  [Int]
is <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical (Vector StationaryDistribution
ds forall a. Vector a -> Int -> a
V.! Int
c) g
g | Int
c <- [Int]
cs]
  forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [Int]
is)

-- | Simulate a number of sites for a given set of substitution models with
-- corresponding weights. Forget states at internal nodes. See also
-- 'simulateAndFlatten'.
simulateAndFlattenMixtureModel ::
  StatefulGen g m =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  V.Vector ExchangeabilityMatrix ->
  Tree Double ->
  g ->
  -- | (IndicesOfComponents, [SimulatedSequenceForEachTip])
  m ([Int], [[State]])
simulateAndFlattenMixtureModel :: forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> Vector RateMatrix
-> Tree Double
-> g
-> m ([Int], [[Int]])
simulateAndFlattenMixtureModel Int
n Vector Double
ws Vector StationaryDistribution
ds Vector RateMatrix
es Tree Double
t g
g = do
  let qs :: Vector RateMatrix
qs = forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix Vector RateMatrix
es Vector StationaryDistribution
ds
      pt :: Tree (Vector RateMatrix)
pt = Vector RateMatrix -> Tree Double -> Tree (Vector RateMatrix)
toProbTreeMixtureModel Vector RateMatrix
qs Tree Double
t
  ([Int]
cs, [Int]
is) <- forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> g
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector StationaryDistribution
ds g
g
  [[Int]]
ss <- forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs Tree (Vector RateMatrix)
pt g
g
  forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [[Int]]
ss)

simulateAndFlattenMixtureModel' ::
  StatefulGen g m =>
  [State] ->
  [Int] ->
  Tree (V.Vector ProbMatrix) ->
  g ->
  m [[State]]
simulateAndFlattenMixtureModel' :: forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps [Tree (Vector RateMatrix)]
f) g
g = do
  [Int]
is' <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i (Vector RateMatrix
ps forall a. Vector a -> Int -> a
V.! Int
c) g
g | (Int
i, Int
c) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
is [Int]
cs]
  if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Tree (Vector RateMatrix)]
f
    then forall (m :: * -> *) a. Monad m => a -> m a
return [[Int]
is']
    else
      forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is' [Int]
cs Tree (Vector RateMatrix)
t g
g | Tree (Vector RateMatrix)
t <- [Tree (Vector RateMatrix)]
f]

getChunks :: Int -> Int -> [Int]
getChunks :: Int -> Int -> [Int]
getChunks Int
c Int
n = [Int]
ns
  where
    n' :: Int
n' = Int
n forall a. Integral a => a -> a -> a
`div` Int
c
    r :: Int
r = Int
n forall a. Integral a => a -> a -> a
`mod` Int
c
    ns :: [Int]
ns = forall a. Int -> a -> [a]
replicate Int
r (Int
n' forall a. Num a => a -> a -> a
+ Int
1) forall a. [a] -> [a] -> [a]
++ forall a. Int -> a -> [a]
replicate (Int
c forall a. Num a => a -> a -> a
- Int
r) Int
n'

-- NOTE: We could move away from IO here, but moving away from 'mapConcurrently'
-- requires benchmarks. I am just not sure if it makes sense to spend more time
-- on this since the parallelization itself is a bit weird. Like so, we walk
-- along separate trees in each process.
parComp :: RandomGen g => Int -> (Int -> IOGenM g -> IO b) -> IOGenM g -> IO [b]
parComp :: forall g b.
RandomGen g =>
Int -> (Int -> IOGenM g -> IO b) -> IOGenM g -> IO [b]
parComp Int
num Int -> IOGenM g -> IO b
fun IOGenM g
gen = do
  Int
ncap <- IO Int
getNumCapabilities
  let chunks :: [Int]
chunks = Int -> Int -> [Int]
getChunks Int
ncap Int
num
  [g]
rs <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
ncap forall a b. (a -> b) -> a -> b
$ forall g r (m :: * -> *). RandomGenM g r m => g -> m r
splitGenM IOGenM g
gen
  [IOGenM g]
gs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
newIOGenM [g]
rs
  forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> IOGenM g -> IO b
fun) (forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [IOGenM g]
gs)

-- | See 'simulateAndFlattenMixtureModel', parallel version.
simulateAndFlattenMixtureModelPar ::
  RandomGen g =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  V.Vector ExchangeabilityMatrix ->
  Tree Double ->
  IOGenM g ->
  IO ([Int], [[State]])
simulateAndFlattenMixtureModelPar :: forall g.
RandomGen g =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> Vector RateMatrix
-> Tree Double
-> IOGenM g
-> IO ([Int], [[Int]])
simulateAndFlattenMixtureModelPar Int
n Vector Double
ws Vector StationaryDistribution
ds Vector RateMatrix
es Tree Double
t IOGenM g
g = do
  let qs :: Vector RateMatrix
qs = forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix Vector RateMatrix
es Vector StationaryDistribution
ds
      pt :: Tree (Vector RateMatrix)
pt = Vector RateMatrix -> Tree Double -> Tree (Vector RateMatrix)
toProbTreeMixtureModel Vector RateMatrix
qs Tree Double
t
  -- The concurrent computation returns a list of ([Int], [[State]]) objects.
  -- They have to be concatenated horizontally.
  [([Int], [[Int]])]
csss <-
    forall g b.
RandomGen g =>
Int -> (Int -> IOGenM g -> IO b) -> IOGenM g -> IO [b]
parComp
      Int
n
      ( \Int
n' IOGenM g
g' ->
          do
            ([Int]
cs, [Int]
is) <- forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> g
-> m ([Int], [Int])
getComponentsAndRootStates Int
n' Vector Double
ws Vector StationaryDistribution
ds IOGenM g
g'
            [[Int]]
ss <- forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs Tree (Vector RateMatrix)
pt IOGenM g
g'
            forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [[Int]]
ss)
      )
      IOGenM g
g
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> a
fst [([Int], [[Int]])]
csss, forall a. [[[a]]] -> [[a]]
horizontalConcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [([Int], [[Int]])]
csss)

-- | Simulate a number of sites for a given set of substitution models with
-- corresponding weights. Keep states at internal nodes. See also
-- 'simulate'.
simulateMixtureModel ::
  StatefulGen g m =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  V.Vector ExchangeabilityMatrix ->
  Tree Double ->
  g ->
  m (Tree [State])
simulateMixtureModel :: forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> Vector RateMatrix
-> Tree Double
-> g
-> m (Tree [Int])
simulateMixtureModel Int
n Vector Double
ws Vector StationaryDistribution
ds Vector RateMatrix
es Tree Double
t g
g = do
  let qs :: Vector RateMatrix
qs = forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith RateMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix Vector RateMatrix
es Vector StationaryDistribution
ds
      pt :: Tree (Vector RateMatrix)
pt = Vector RateMatrix -> Tree Double -> Tree (Vector RateMatrix)
toProbTreeMixtureModel Vector RateMatrix
qs Tree Double
t
  ([Int]
cs, [Int]
is) <- forall g (m :: * -> *).
StatefulGen g m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> g
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector StationaryDistribution
ds g
g
  forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs Tree (Vector RateMatrix)
pt g
g

-- See 'simulateAlongProbTree', only we have a number of mixture components. The
-- starting states and the components for each site have to be provided.
simulateMixtureModel' ::
  StatefulGen g m =>
  [State] ->
  [Int] ->
  Tree (V.Vector ProbMatrix) ->
  g ->
  m (Tree [State])
simulateMixtureModel' :: forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps [Tree (Vector RateMatrix)]
f) g
g = do
  [Int]
is' <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i (Vector RateMatrix
ps forall a. Vector a -> Int -> a
V.! Int
c) g
g | (Int
i, Int
c) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
is [Int]
cs]
  [Tree [Int]]
f' <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m (Tree [Int])
simulateMixtureModel' [Int]
is' [Int]
cs Tree (Vector RateMatrix)
t g
g | Tree (Vector RateMatrix)
t <- [Tree (Vector RateMatrix)]
f]
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> [Tree a] -> Tree a
Node [Int]
is' [Tree [Int]]
f'