biohazard-0.6.9: tools/redeye-dar.hs
{-# LANGUAGE FlexibleContexts #-}
-- Co-estimates aDNA damage with parameters for a simple genotype prior.
--
-- We want to estimate on only a subset of the genome. For the time
-- being, this is by definition a subset of the large blocks of the
-- mappability track for the human genome (so this doesn't work for
-- other genomes). To make this less crude, we need a differently
-- prepared input, but right now, input is one BAM file with read group
-- annotations.
--
-- We run the EM algorithm, repeatedly estimating one damage/error model
-- per read group and the global genotype parameters. Convergence is
-- achieved when the changes in the damage models are sufficiently
-- small. The damage/error model is one substitution matrix for each
-- position within a read near the ends, and one for what remains in the
-- middle.
import Bio.Adna
import Bio.Bam
import Bio.Bam.Pileup
import Bio.Genocall
import Bio.Genocall.Estimators
import Bio.Prelude
import Bio.Util.AD
import Data.Aeson.Encode.Pretty
import System.Console.GetOpt
import qualified Data.ByteString.Lazy.Char8 as L
import qualified Data.HashMap.Strict as H
import qualified Data.Sequence as Z
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M
data Conf = Conf {
conf_output :: LazyBytes -> IO (),
conf_report :: LazyBytes -> IO (),
conf_params :: Parameters,
conf_length :: Int,
conf_eps :: Double }
defaultConf :: Conf
defaultConf = Conf (L.hPutStrLn stdout) (\_ -> return ()) quietParameters 16 1.0E-6
options :: [OptDescr (Conf -> IO Conf)]
options = [
Option "o" ["output"] (ReqArg set_output "FILE") "Write output to FILE (stdout)",
Option "l" ["model-length"] (ReqArg set_len "NUM") "Set size of subst. model to NUM (16)",
Option "e" ["precision"] (ReqArg set_prec "NUM") "Set precision for fit to NUM (1E-6)",
Option "v" ["verbose"] (NoArg set_verbose) "Print progress reports",
Option "h?" ["help","usage"] (NoArg disp_usage) "Print this message and exit" ]
where
set_verbose c = return $ c { conf_report = L.hPutStrLn stderr, conf_params = debugParameters }
set_output f c = return $ c { conf_output = L.writeFile f }
set_len a c = readIO a >>= \x -> return $ c { conf_length = x }
set_prec a c = readIO a >>= \x -> return $ c { conf_eps = x }
disp_usage _ = do pn <- getProgName
let blah = "Usage: " ++ pn ++ " [OPTION...] [LIBRARY-NAME...]"
putStrLn $ usageInfo blah options
exitSuccess
main :: IO ()
main = do
(opts, files, errors) <- getOpt Permute options <$> getArgs
unless (null errors) $ mapM_ (hPutStrLn stderr) errors >> exitFailure
Conf{..} <- foldl (>>=) (return defaultConf) opts
-- For each iteration: read the input, decompose, pileup. The
-- "prior" damage model is the usual 'SubstModel', the "posterior"
-- damage model needs to be a mutable 'MSubstModel'. We feed
-- likelihoods into the div/het estimation (just as in
-- 'redeye-pileup'), but also into a caller that will estimate
-- damage.
let iter sp0 mod0 = do (((de1,de2),mod1),syms) <- emIter conf_length sp0 mod0 files
conf_report $ encodePretty de2
if diffSubstMod mod0 mod1 > conf_eps
then iter (case point_est de1 of [a,b] -> SinglePop a b) mod1
else return . ExtModel de1 (Just de2) . SubstModels
$ H.map ((mod1 V.!) . fromDmgToken) syms
final_model <- iter (SinglePop 0.001 0.002) V.empty
conf_output $ encodePretty (final_model :: ExtModel)
diffSubstMod :: V.Vector SubstModel -> V.Vector SubstModel -> Double
diffSubstMod v1 v2 =
V.foldl' max 0 (V.zipWith diff1 v1 v2) `max`
V.foldl' max 0 (V.map abs1 (V.drop (V.length v2) v1)) `max`
V.foldl' max 0 (V.map abs1 (V.drop (V.length v1) v2))
where
diff1 :: SubstModel -> SubstModel -> Double
diff1 sm1 sm2 = maximum $
[ V.maximum $ V.zipWith diff2 (left_substs_fwd sm1) (left_substs_fwd sm2)
, diff2 (middle_substs_fwd sm1) (middle_substs_fwd sm2)
, V.maximum $ V.zipWith diff2 (right_substs_fwd sm1) (right_substs_fwd sm2) ]
diff2 :: Mat44D -> Mat44D -> Double
diff2 (Mat44D u) (Mat44D v) = U.maximum $ U.map abs $ U.zipWith (-) u v
abs1 :: SubstModel -> Double
abs1 sm1 = V.maximum (V.map abs2 (left_substs_fwd sm1)) `max`
abs2 (middle_substs_fwd sm1) `max`
V.maximum (V.map abs2 (right_substs_fwd sm1))
abs2 :: Mat44D -> Double
abs2 (Mat44D v) = U.maximum v
-- One iteration of EM algorithm. We go in with a substitution model
-- and het/div, we come out with new estimates for same. We get het/div from
-- tabulation followed by numerical optimization. For damage, we have to
-- compute posterior probabilities using the old model, then update the
-- damage matrices with pseudo counts. (This amounts to a maximum
-- likelihood estimate for a weighted multinomial distribution.)
emIter :: Int -> SinglePop -> V.Vector SubstModel -> [FilePath] -> IO (((DivEst,DivEst), V.Vector SubstModel), HashMap Bytes DmgToken)
emIter msize divest mod0 infiles =
liftIO (newIORef V.empty) >>= \mmod ->
liftIO (newIORef H.empty) >>= \symtab ->
concatInputs infiles >=> run $ \hdr ->
concatMapStreamM (decompose_dmg_from symtab) =$
pileup =$
filterPilesWith (the_regions hdr) =$
mapStream ( id &&& calls mod0 ) =$
let div_estimation :: MonadIO m => Iteratee [(a, Calls)] m (DivEst,DivEst)
div_estimation = mapStream snd =$
tabulateSingle >>=
liftIO . estimateSingle
dmg_estimation :: MonadIO m => Iteratee [(Pile, Calls)] m (V.Vector SubstModel)
dmg_estimation = mapStreamM_ (\(p,c) ->
liftIO . updateSubstModel msize mod0 mmod p $
single_pop_posterior divest
(refix $ snp_refbase $ p_snp_pile c)
(snp_gls $ p_snp_pile c)) >>
liftIO (readIORef mmod) >>=
liftIO . V.mapM freezeSubstModel
in (,) <$> zipStreams div_estimation dmg_estimation
<*> liftIO (readIORef symtab)
where
the_regions hdr = sort [ Region (Refseq $ fromIntegral ri) p (p+l)
| (ch, p, l) <- good_regions
, let Just ri = Z.findIndexL ((==) ch . sq_name) (meta_refs hdr) ]
refix ref = U.fromListN 16 [0,0,2,0,5,0,0,0,9,0,0,0,0,0,0,0] U.! fromIntegral (unNs ref)
filterPilesWith :: Monad m => [Region] -> Enumeratee [Pile] [Pile] m b
filterPilesWith = unfoldConvStream go
where
go [ ] = skipToEof >> return ([],[])
go (r:rs) = do mp <- peekStream
case mp of
Just p | (p_refseq p, p_pos p) < (refseq r, start r) -> headStream >> go (r:rs)
| (p_refseq p, p_pos p) >= (refseq r, end r) -> go rs
| otherwise -> (\x -> (r:rs, [x])) `liftM` headStream
Nothing -> return ([],[])
-- Probabilistically count substitutions. We infer from posterior
-- genotype probabilities what the base must have been, then count
-- substitutions from that to the actual base.
updateSubstModel :: Int -> V.Vector SubstModel -> IORef (V.Vector MSubstModel) -> Pile -> U.Vector Prob -> IO ()
updateSubstModel msize mods0 vmods1 pile postp = case p_snp_pile pile of
(basesF, basesR) -> do mapM_ (count_base False) basesF
mapM_ (count_base True) basesR
where
-- Posterior probalities of the haploid base before damage
-- @P(H) = \sum_{G} P(H|G) P(G|D)@
pH_A = fromProb $ postp U.! 0 + 0.5 * ( postp U.! 1 + postp U.! 3 + postp U.! 6 )
pH_C = fromProb $ postp U.! 2 + 0.5 * ( postp U.! 1 + postp U.! 4 + postp U.! 7 )
pH_G = fromProb $ postp U.! 5 + 0.5 * ( postp U.! 3 + postp U.! 4 + postp U.! 8 )
pH_T = fromProb $ postp U.! 9 + 0.5 * ( postp U.! 6 + postp U.! 7 + postp U.! 8 )
-- P(H:->X) = P(H|X)
-- = P(X|H) P(H) / P(X)
-- = P(X|H) P(H) / \sum_H' P(X|H') P(H')
--
-- We get P(X|H) from the old substitution model.
count_base str (q,b) = do
let old_mat = case mods0 V.!? fromDmgToken (db_dmg_tk b) of
Nothing -> initmat
Just sm -> lookupSubstModel sm (db_dmg_pos b) str
new_mat <- do mods1 <- readIORef vmods1
sm <- case mods1 V.!? fromDmgToken (db_dmg_tk b) of
Nothing -> do m <- new_mmodel
writeIORef vmods1 (V.snoc mods1 m)
return m
Just m -> return m
return $ lookupSubstModel sm (db_dmg_pos b) str
-- Use map quality to de-emphasize badly mapped reads. If
-- everything is badly mapped, the estimation still works; if
-- some data is well mapped, that naturally dominates the
-- estimate.
let pHX = (1 - fromQual q) / (pHX_A + pHX_C + pHX_G + pHX_T)
pHX_A = (old_mat `bang` nucA :-> db_call b) * pH_A
pHX_C = (old_mat `bang` nucC :-> db_call b) * pH_C
pHX_G = (old_mat `bang` nucG :-> db_call b) * pH_G
pHX_T = (old_mat `bang` nucT :-> db_call b) * pH_T
nudge new_mat (nucA :-> db_call b) (pHX_A * pHX)
nudge new_mat (nucC :-> db_call b) (pHX_C * pHX)
nudge new_mat (nucG :-> db_call b) (pHX_G * pHX)
nudge new_mat (nucT :-> db_call b) (pHX_T * pHX)
new_mmodel = SubstModel <$> V.replicateM msize nullmat
<*> nullmat
<*> V.replicateM msize nullmat
<*> V.replicateM msize nullmat
<*> nullmat
<*> V.replicateM msize nullmat
nullmat = MMat44D <$> M.replicate 16 (0::Double)
calls :: V.Vector SubstModel -> Pile -> Calls
calls dmg pile = pile { p_snp_pile = s, p_indel_pile = i }
where
!s = simple_snp_call get_dmg $ p_snp_pile pile
!i = simple_indel_call get_dmg $ p_indel_pile pile
get_dmg :: DmgToken -> Int -> Bool -> Mat44D
get_dmg (DmgToken dt) di ds = case dmg V.!? dt of
Nothing -> initmat -- not found, happens in first round
Just sm -> lookupSubstModel sm di ds
initmat :: Mat44D
initmat = Mat44D $ U.fromListN 16 [ 0.91, 0.03, 0.03, 0.03
, 0.03, 0.91, 0.03, 0.03
, 0.03, 0.03, 0.91, 0.03
, 0.03, 0.03, 0.03, 0.91 ]
{-# INLINE decompose_dmg_from #-}
decompose_dmg_from :: IORef (HashMap Bytes DmgToken) -> BamRaw -> IO [PosPrimChunks]
decompose_dmg_from ref raw = do
hm <- readIORef ref
let rg = extAsString "RG" (unpackBam raw)
token <- case H.lookup rg hm of
Just tk -> return tk
Nothing -> do let tk = DmgToken $ H.size hm
writeIORef ref $! H.insert rg tk hm
return tk
return $ decompose token raw