{-# LANGUAGE ScopedTypeVariables #-}

-- |
--   Description :  Work with transition probability matrices on rooted trees
--   Copyright   :  (c) Dominik Schrempf 2021
--   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 Control.Monad.Primitive
import Data.Tree
import qualified Data.Vector as V
import Data.Word (Word32)
import ELynx.Data.MarkovProcess.RateMatrix
import ELynx.Simulate.MarkovProcess
import System.Random.MWC
import System.Random.MWC.Distributions (categorical)

-- 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 :: [[[a]]] -> [[a]]
horizontalConcat [[[a]]
xs] = [[a]]
xs
horizontalConcat [[[a]]]
xss = ([[a]] -> [[a]] -> [[a]]) -> [[[a]]] -> [[a]]
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (([a] -> [a] -> [a]) -> [[a]] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++)) [[[a]]]
xss

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 g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical StationaryDistribution
d Gen (PrimState m)
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 ::
  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

-- 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' ::
  (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]

-- | See 'simulateAndFlatten', parallel version.
simulateAndFlattenPar ::
  Int ->
  StationaryDistribution ->
  ExchangeabilityMatrix ->
  Tree Double ->
  GenIO ->
  IO [[State]]
simulateAndFlattenPar :: Int
-> StationaryDistribution
-> RateMatrix
-> Tree Double
-> GenIO
-> IO [[Int]]
simulateAndFlattenPar Int
n StationaryDistribution
d RateMatrix
e Tree Double
t GenIO
g = do
  Int
c <- IO Int
getNumCapabilities
  [Gen RealWorld]
gs <- Int -> GenIO -> IO [GenIO]
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
c GenIO
g
  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 <-
    ((Int, Gen RealWorld) -> IO [[Int]])
-> [(Int, Gen RealWorld)] -> IO [[[Int]]]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently
      ( \(Int
num, Gen RealWorld
gen) -> do
          [Int]
is <- Int -> StationaryDistribution -> GenIO -> IO [Int]
forall (m :: * -> *).
PrimMonad m =>
Int -> StationaryDistribution -> Gen (PrimState m) -> m [Int]
getRootStates Int
num StationaryDistribution
d Gen RealWorld
GenIO
gen
          [Int] -> Tree RateMatrix -> GenIO -> IO [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int] -> Tree RateMatrix -> Gen (PrimState m) -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt Gen RealWorld
GenIO
gen
      )
      ([Int] -> [Gen RealWorld] -> [(Int, Gen RealWorld)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [Gen RealWorld]
gs)
  [[Int]] -> IO [[Int]]
forall (m :: * -> *) a. Monad m => a -> m a
return ([[Int]] -> IO [[Int]]) -> [[Int]] -> IO [[Int]]
forall a b. (a -> b) -> a -> b
$ [[[Int]]] -> [[Int]]
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 ::
  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

-- 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' ::
  (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 ::
  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?
  (Double -> Vector RateMatrix)
-> Tree Double -> Tree (Vector RateMatrix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Double
a -> (RateMatrix -> RateMatrix)
-> Vector RateMatrix -> Vector RateMatrix
forall a b. (a -> b) -> Vector a -> Vector b
V.map (RateMatrix -> Double -> RateMatrix
`probMatrix` Double
a) Vector RateMatrix
qs)

getComponentsAndRootStates ::
  PrimMonad m =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  Gen (PrimState m) ->
  m ([Int], [State])
getComponentsAndRootStates :: Int
-> Vector Double
-> Vector StationaryDistribution
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector 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
$ Vector Double -> Gen (PrimState m) -> m Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical Vector Double
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 g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical (Vector StationaryDistribution
ds Vector StationaryDistribution -> Int -> StationaryDistribution
forall a. Vector a -> Int -> a
V.! 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)

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

simulateAndFlattenMixtureModel' ::
  (PrimMonad m) =>
  [State] ->
  [Int] ->
  Tree (V.Vector ProbMatrix) ->
  Gen (PrimState m) ->
  m [[State]]
simulateAndFlattenMixtureModel' :: [Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps Forest (Vector 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 (Vector RateMatrix
ps Vector RateMatrix -> Int -> RateMatrix
forall a. Vector a -> Int -> a
V.! 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 (Vector RateMatrix) -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Forest (Vector 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 (Vector RateMatrix)
-> Gen (PrimState m)
-> m [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is' [Int]
cs Tree (Vector RateMatrix)
t Gen (PrimState m)
g | Tree (Vector RateMatrix)
t <- Forest (Vector 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 :: * -> *) g a (v :: * -> *).
(PrimMonad m, StatefulGen g m, Uniform a, Vector v a) =>
g -> 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)

-- | See 'simulateAndFlattenMixtureModel', parallel version.
simulateAndFlattenMixtureModelPar ::
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  V.Vector ExchangeabilityMatrix ->
  Tree Double ->
  GenIO ->
  IO ([Int], [[State]])
simulateAndFlattenMixtureModelPar :: Int
-> Vector Double
-> Vector StationaryDistribution
-> Vector RateMatrix
-> Tree Double
-> GenIO
-> IO ([Int], [[Int]])
simulateAndFlattenMixtureModelPar Int
n Vector Double
ws Vector StationaryDistribution
ds Vector RateMatrix
es Tree Double
t GenIO
g = do
  let qs :: Vector RateMatrix
qs = (RateMatrix -> StationaryDistribution -> RateMatrix)
-> Vector RateMatrix
-> Vector StationaryDistribution
-> Vector RateMatrix
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 <-
    Int
-> (Int -> GenIO -> IO ([Int], [[Int]]))
-> GenIO
-> IO [([Int], [[Int]])]
forall b. Int -> (Int -> GenIO -> IO b) -> GenIO -> IO [b]
parComp
      Int
n
      ( \Int
n' GenIO
g' ->
          do
            ([Int]
cs, [Int]
is) <- Int
-> Vector Double
-> Vector StationaryDistribution
-> GenIO
-> IO ([Int], [Int])
forall (m :: * -> *).
PrimMonad m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n' Vector Double
ws Vector StationaryDistribution
ds GenIO
g'
            [[Int]]
ss <- [Int] -> [Int] -> Tree (Vector RateMatrix) -> GenIO -> IO [[Int]]
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs Tree (Vector RateMatrix)
pt GenIO
g'
            ([Int], [[Int]]) -> IO ([Int], [[Int]])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [[Int]]
ss)
      )
      GenIO
g
  ([Int], [[Int]]) -> IO ([Int], [[Int]])
forall (m :: * -> *) a. Monad m => a -> m a
return ((([Int], [[Int]]) -> [Int]) -> [([Int], [[Int]])] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [[Int]]) -> [Int]
forall a b. (a, b) -> a
fst [([Int], [[Int]])]
csss, [[[Int]]] -> [[Int]]
forall a. [[[a]]] -> [[a]]
horizontalConcat ([[[Int]]] -> [[Int]]) -> [[[Int]]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ (([Int], [[Int]]) -> [[Int]]) -> [([Int], [[Int]])] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map ([Int], [[Int]]) -> [[Int]]
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 ::
  PrimMonad m =>
  Int ->
  V.Vector Double ->
  V.Vector StationaryDistribution ->
  V.Vector ExchangeabilityMatrix ->
  Tree Double ->
  Gen (PrimState m) ->
  m (Tree [State])
simulateMixtureModel :: Int
-> Vector Double
-> Vector StationaryDistribution
-> Vector RateMatrix
-> Tree Double
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel Int
n Vector Double
ws Vector StationaryDistribution
ds Vector RateMatrix
es Tree Double
t Gen (PrimState m)
g = do
  let qs :: Vector RateMatrix
qs = (RateMatrix -> StationaryDistribution -> RateMatrix)
-> Vector RateMatrix
-> Vector StationaryDistribution
-> Vector RateMatrix
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) <- Int
-> Vector Double
-> Vector StationaryDistribution
-> Gen (PrimState m)
-> m ([Int], [Int])
forall (m :: * -> *).
PrimMonad m =>
Int
-> Vector Double
-> Vector StationaryDistribution
-> Gen (PrimState m)
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector StationaryDistribution
ds Gen (PrimState m)
g
  [Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs Tree (Vector RateMatrix)
pt Gen (PrimState m)
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' ::
  (PrimMonad m) =>
  [State] ->
  [Int] ->
  Tree (V.Vector ProbMatrix) ->
  Gen (PrimState m) ->
  m (Tree [State])
simulateMixtureModel' :: [Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps Forest (Vector 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 (Vector RateMatrix
ps Vector RateMatrix -> Int -> RateMatrix
forall a. Vector a -> Int -> a
V.! 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 (Vector RateMatrix)
-> Gen (PrimState m)
-> m (Tree [Int])
forall (m :: * -> *).
PrimMonad m =>
[Int]
-> [Int]
-> Tree (Vector RateMatrix)
-> Gen (PrimState m)
-> m (Tree [Int])
simulateMixtureModel' [Int]
is' [Int]
cs Tree (Vector RateMatrix)
t Gen (PrimState m)
g | Tree (Vector RateMatrix)
t <- Forest (Vector 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'