1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
module Main where
import Control.Monad
import Data.Fixed (mod')
import Data.Foldable (toList)
import Data.List hiding (group)
import qualified Data.Map as M
import GHC.Compact
import qualified Language.Haskell.TH.Syntax as TH
import Numeric.Log hiding (sum)
import Options.Applicative
import PPL hiding (Tree, binom)
import qualified Streaming as S
import Streaming.Prelude (Of, Stream, each, fold, yield)
import qualified Streaming.Prelude as S
import System.IO (BufferMode (NoBuffering), hSetBuffering, stdout)
import System.ProgressBar
import System.Random (StdGen, mkStdGen, random, setStdGen, split)
cumsum = scanl1 (+)
first f xs = snd . head . filter (f . fst) $ zip xs [0 ..]
stirling 0 = 1
stirling 1 = 1
stirling n = Exp $ n * (log n - 1) + log (sqrt (2 * pi)) + log n / 2
logFromInt = logFrom . fromIntegral
logFrom = Exp . log
binom :: Int -> Int -> Double -> Log Double
binom n@(logFromInt -> n') k@(logFromInt -> k') p =
n `choose` k * logFrom p ** k' * logFrom (1 - p) ** (n' - k')
where
choose (fromIntegral -> n) (fromIntegral -> k) = stirling n / stirling k / stirling (n - k)
-- (infinite) binary trees
data Tree a = Tree a (Tree a) (Tree a)
deriving (Show)
instance Foldable Tree where
foldMap f t = bftrav [t]
where
bftrav [] = mempty
bftrav ((Tree a l r) : ts) = f a <> bftrav (ts <> [l, r])
{-# INLINE group #-}
group (a : b : c : d : rs) = (a, b, c, d) : group rs
group [] = []
group s = error $ "unexpected number of columns, expecting 4:" <> show s
-- Infinite trees from infinite lists
-- NB: it's harder to partition a list so that it folds back to
-- equivalence. It doesn't really matter here since we're only
-- unfolding random uniforms anyway.
treeFromList (x : xs) = Tree x (treeFromList lpart) (treeFromList rpart)
where
(lpart, rpart) = unzip $ pairs xs
pairs (a : b : rs) = (a, b) : pairs rs
-- Constrain trees so leaves sum to node value
normTree :: Tree Double -> Tree Double
normTree (Tree x (Tree u l r) (Tree v l' r')) =
let s = x / (u + v)
in Tree x (normTree $ Tree (s * u) l r) (normTree $ Tree (s * v) l' r')
drawTreeProbs = toList . normTree . treeFromList <$> iid uniform
model :: [[Int]] -> Meas ([[Double]], [Int])
model xs = do
(ps, params, clusters) <- sample $ do
-- Sample hyperparameters
a <- bounded 1 10
-- CRP style
dir <- cumsum <$> dirichletProcess a
rs <- iid uniform
ps <- iid drawTreeProbs
let clusters = map (\r -> first (>= r) dir) rs
params = map (transpose ps !!) clusters
pure (ps, params, clusters)
mapM_ scoreLog $ zipWith likelihood params xs
let cls = take (length xs) clusters
k = maximum cls + 1
n = length (head xs) `div` 4
pure (map (take k) $ take n ps, cls)
where
likelihood ps cnts = product $ zipWith go ps (group cnts)
where
go p (c, d, e, f) = binom d c (min (1 - 1e-10) (p * fromIntegral (1 + e - f) / fromIntegral (1 + f)))
eps = $(TH.lift (2 * until ((== 1) . (1 +)) (/ 2) (1 :: Double)))
-- Tabulate list
tabulate xs = M.elems $ M.fromListWith (+) [(c, 1) | c <- xs]
-- Command line args
data Options = Options
{ seed :: Int,
nsamples :: Int,
mhfrac :: Double,
input :: FilePath,
outputPath :: FilePath
}
main = run =<< execParser opts
where
opts =
info
(parser <**> helper)
(fullDesc <> progDesc "Infer a phylogeny from SNV calls in multiple samples" <> header "phylogey - Bayesian phylogeny inference")
parser =
Options
<$> option auto (long "seed" <> short 's' <> help "random seed" <> showDefault <> value 42 <> metavar "INT")
<*> option auto (long "nsamples" <> short 'n' <> help "number of samples from posterior" <> value 100000 <> metavar "INT" <> showDefault)
<*> option probability (long "mhfrac" <> short 'm' <> help "Metropolis-Hastings mutation probability" <> value 0.01 <> metavar "[0,1]" <> showDefault)
<*> argument str (metavar "INPUT")
<*> argument str (metavar "OUTPUTDIR")
probability = eitherReader $ \arg -> case reads arg of
[(r, "")] -> if r <= 1 && r >= 0 then Right r else Left "mhfrac not a valid probability"
_ -> Left "mhfrac not a valid probability"
takeWithProgress :: (S.MonadIO m) => Int -> Stream (Of a) m r -> Stream (Of a) m ()
takeWithProgress n str = do
pb <- S.liftIO $ newProgressBar defStyle 10 (Progress 0 n ())
S.mapM (update pb) $ S.take n str
where
update pb x = do
S.liftIO $ incProgress pb 1
pure x
run opts = do
(hdr : lines) <- lines <$> readFile (input opts)
let dbl = round . read :: String -> Int
g = mkStdGen $ seed opts
parsed <- compact $ map (map dbl . tail . words) lines
hSetBuffering stdout NoBuffering
m@((ps, cl), _) <-
S.fold_ (\l r -> if mml l < mml r then l else r) (([[]], []), -1 / 0) id . takeWithProgress (nsamples opts) $
mh g (mhfrac opts) (model $ getCompact parsed)
writeFile (outputPath opts <> "/props") . unlines $ map (intercalate "," . map show) ps
writeFile (outputPath opts <> "/tree") . unlines $ map show cl
writeFile (outputPath opts <> "/mml") $ show (mml m)
where
mml ((ps, cl), lik) = sum' (sum' (log . (+ 1))) ps + sum' (log . (+ 1)) tab - sum' (ln . stirling) tab - ln lik
where
tab = tabulate cl
sum' f = sum . map f
|
