{-# 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 :: 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
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
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]
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
[[[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 ::
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
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 =
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)
simulateAndFlattenMixtureModel ::
StatefulGen g m =>
Int ->
V.Vector Double ->
V.Vector StationaryDistribution ->
V.Vector ExchangeabilityMatrix ->
Tree Double ->
g ->
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'
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)
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
[([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)
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
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'