{-# LANGUAGE ScopedTypeVariables #-}
module ELynx.Simulate.MarkovProcessAlongTree
(
simulate,
simulateAndFlatten,
simulateAndFlattenPar,
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
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 ::
StatefulGen g m =>
Int ->
StationaryDistribution ->
g ->
m [State]
getRootStates :: Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
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 -> g -> m Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical StationaryDistribution
d g
g
simulateAndFlatten ::
StatefulGen g m =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
g ->
m [[State]]
simulateAndFlatten :: 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 <- Int -> StationaryDistribution -> g -> m [Int]
forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
g
[Int] -> Tree RateMatrix -> g -> m [[Int]]
forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt g
g
simulateAndFlatten' ::
StatefulGen g m =>
[State] ->
Tree ProbMatrix ->
g ->
m [[State]]
simulateAndFlatten' :: [Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is (Node RateMatrix
p Forest RateMatrix
f) g
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 -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i RateMatrix
p g
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 -> g -> m [[Int]]
forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is' Tree RateMatrix
t g
g | Tree RateMatrix
t <- Forest RateMatrix
f]
simulateAndFlattenPar ::
RandomGen g =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
IOGenM g ->
IO [[State]]
simulateAndFlattenPar :: 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 <- Int -> IO g -> IO [g]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
c (IO g -> IO [g]) -> IO g -> IO [g]
forall a b. (a -> b) -> a -> b
$ IOGenM g -> IO g
forall g r (m :: * -> *). RandomGenM g r m => g -> m r
splitGenM IOGenM g
g
[IOGenM g]
gs <- (g -> IO (IOGenM g)) -> [g] -> IO [IOGenM g]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM g -> IO (IOGenM g)
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
[[[Int]]]
sss <-
((Int, IOGenM g) -> IO [[Int]])
-> [(Int, IOGenM g)] -> IO [[[Int]]]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently
( \(Int
num, IOGenM g
gen) -> do
[Int]
is <- Int -> StationaryDistribution -> IOGenM g -> IO [Int]
forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
num StationaryDistribution
d IOGenM g
gen
[Int] -> Tree RateMatrix -> IOGenM g -> IO [[Int]]
forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m [[Int]]
simulateAndFlatten' [Int]
is Tree RateMatrix
pt IOGenM g
gen
)
([Int] -> [IOGenM g] -> [(Int, IOGenM g)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [IOGenM g]
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 ::
StatefulGen g m =>
Int ->
StationaryDistribution ->
ExchangeabilityMatrix ->
Tree Double ->
g ->
m (Tree [State])
simulate :: 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 <- Int -> StationaryDistribution -> g -> m [Int]
forall g (m :: * -> *).
StatefulGen g m =>
Int -> StationaryDistribution -> g -> m [Int]
getRootStates Int
n StationaryDistribution
d g
g
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is Tree RateMatrix
pt g
g
simulate' ::
StatefulGen g m =>
[State] ->
Tree ProbMatrix ->
g ->
m (Tree [State])
simulate' :: [Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is (Node RateMatrix
p Forest RateMatrix
f) g
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 -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i RateMatrix
p g
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 -> g -> m (Tree [Int])
forall g (m :: * -> *).
StatefulGen g m =>
[Int] -> Tree RateMatrix -> g -> m (Tree [Int])
simulate' [Int]
is' Tree RateMatrix
t g
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 =
(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 ::
StatefulGen g m =>
Int ->
V.Vector Double ->
V.Vector StationaryDistribution ->
g ->
m ([Int], [State])
getComponentsAndRootStates :: Int
-> Vector Double
-> Vector StationaryDistribution
-> g
-> m ([Int], [Int])
getComponentsAndRootStates Int
n Vector Double
ws Vector StationaryDistribution
ds g
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 -> g -> m Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical Vector Double
ws g
g
[Int]
is <- [m Int] -> m [Int]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [StationaryDistribution -> g -> 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) g
g | Int
c <- [Int]
cs]
([Int], [Int]) -> m ([Int], [Int])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [Int]
is)
simulateAndFlattenMixtureModel ::
StatefulGen g m =>
Int ->
V.Vector Double ->
V.Vector StationaryDistribution ->
V.Vector ExchangeabilityMatrix ->
Tree Double ->
g ->
m ([Int], [[State]])
simulateAndFlattenMixtureModel :: 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 = (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
-> g
-> m ([Int], [Int])
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 <- [Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
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
([Int], [[Int]]) -> m ([Int], [[Int]])
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' :: [Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m [[Int]]
simulateAndFlattenMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps Forest (Vector RateMatrix)
f) g
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 -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i (Vector RateMatrix
ps Vector RateMatrix -> Int -> RateMatrix
forall a. Vector a -> Int -> a
V.! Int
c) g
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) -> g -> m [[Int]]
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 <- 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'
parComp :: RandomGen g => Int -> (Int -> IOGenM g -> IO b) -> IOGenM g -> IO [b]
parComp :: 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 <- Int -> IO g -> IO [g]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
ncap (IO g -> IO [g]) -> IO g -> IO [g]
forall a b. (a -> b) -> a -> b
$ IOGenM g -> IO g
forall g r (m :: * -> *). RandomGenM g r m => g -> m r
splitGenM IOGenM g
gen
[IOGenM g]
gs <- (g -> IO (IOGenM g)) -> [g] -> IO [IOGenM g]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM g -> IO (IOGenM g)
forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
newIOGenM [g]
rs
((Int, IOGenM g) -> IO b) -> [(Int, IOGenM g)] -> IO [b]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently ((Int -> IOGenM g -> IO b) -> (Int, IOGenM g) -> IO b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> IOGenM g -> IO b
fun) ([Int] -> [IOGenM g] -> [(Int, IOGenM g)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
chunks [IOGenM g]
gs)
simulateAndFlattenMixtureModelPar ::
RandomGen g =>
Int ->
V.Vector Double ->
V.Vector StationaryDistribution ->
V.Vector ExchangeabilityMatrix ->
Tree Double ->
IOGenM g ->
IO ([Int], [[State]])
simulateAndFlattenMixtureModelPar :: 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 = (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], [[Int]])]
csss <-
Int
-> (Int -> IOGenM g -> IO ([Int], [[Int]]))
-> IOGenM g
-> IO [([Int], [[Int]])]
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) <- Int
-> Vector Double
-> Vector StationaryDistribution
-> IOGenM g
-> IO ([Int], [Int])
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 <- [Int]
-> [Int] -> Tree (Vector RateMatrix) -> IOGenM g -> IO [[Int]]
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'
([Int], [[Int]]) -> IO ([Int], [[Int]])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cs, [[Int]]
ss)
)
IOGenM g
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)
simulateMixtureModel ::
StatefulGen g m =>
Int ->
V.Vector Double ->
V.Vector StationaryDistribution ->
V.Vector ExchangeabilityMatrix ->
Tree Double ->
g ->
m (Tree [State])
simulateMixtureModel :: 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 = (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
-> g
-> m ([Int], [Int])
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] -> [Int] -> Tree (Vector RateMatrix) -> g -> m (Tree [Int])
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
simulateMixtureModel' ::
StatefulGen g m =>
[State] ->
[Int] ->
Tree (V.Vector ProbMatrix) ->
g ->
m (Tree [State])
simulateMixtureModel' :: [Int] -> [Int] -> Tree (Vector RateMatrix) -> g -> m (Tree [Int])
simulateMixtureModel' [Int]
is [Int]
cs (Node Vector RateMatrix
ps Forest (Vector RateMatrix)
f) g
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 -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> RateMatrix -> g -> m Int
jump Int
i (Vector RateMatrix
ps Vector RateMatrix -> Int -> RateMatrix
forall a. Vector a -> Int -> a
V.! Int
c) g
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) -> g -> m (Tree [Int])
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 <- 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'