{-# LANGUAGE RecordWildCards, BangPatterns, FlexibleContexts, OverloadedStrings #-}
import Bio.Bam
import Bio.Bam.Rmdup
import Bio.Base
import Bio.Util.Numeric ( showNum, showOOM, estimateComplexity )
import Control.Monad
import Control.Monad.ST ( runST )
import Data.Bits
import Data.Foldable ( toList )
import Data.List ( intercalate )
import Data.Maybe
import Data.Ord ( comparing )
import Data.Vector.Algorithms.Intro ( sortBy )
import Data.Version ( showVersion )
import Numeric ( showFFloat )
import Paths_biohazard ( version )
import System.Console.GetOpt
import System.Environment ( getArgs, getProgName )
import System.Exit
import System.IO
import qualified Data.ByteString.Char8 as S
import qualified Data.HashMap.Strict as M
import qualified Data.IntMap as IM
import qualified Data.Iteratee as I
import qualified Data.Sequence as Z
import qualified Data.Vector as VV
import qualified Data.Vector.Generic as V
data Conf = Conf {
output :: Maybe ((BamRec -> Seqid) -> BamMeta -> Iteratee [BamRec] IO ()),
strand_preserved :: Bool,
collapse :: Bool -> Collapse,
clean_multimap :: BamRec -> IO (Maybe BamRec),
keep_all :: Bool,
keep_unaligned :: Bool,
keep_improper :: Bool,
transform :: BamRec -> Maybe BamRec,
min_len :: Int,
min_qual :: Qual,
get_label :: M.HashMap Seqid Seqid -> BamRec -> Seqid,
putResult :: String -> IO (),
debug :: String -> IO (),
which :: Which,
circulars :: Refs -> IO (IM.IntMap (Seqid,Int), Refs) }
-- | Which reference sequences to scan
data Which = Allrefs | Some Refseq Refseq | Unaln deriving Show
defaults :: Conf
defaults = Conf { output = Nothing
, strand_preserved = True
, collapse = cons_collapse' (Q 60)
, clean_multimap = check_flags
, keep_all = False
, keep_unaligned = False
, keep_improper = False
, transform = Just
, min_len = 0
, min_qual = Q 0
, get_label = get_library
, putResult = putStr
, debug = \_ -> return ()
, which = Allrefs
, circulars = \rs -> return (IM.empty, rs) }
options :: [OptDescr (Conf -> IO Conf)]
options = [
Option "o" ["output"] (ReqArg set_output "FILE") "Write to FILE (default: no output, count only)",
Option "O" ["output-lib"] (ReqArg set_lib_out "PAT") "Write each lib to file named following PAT",
Option [ ] ["debug"] (NoArg set_debug_out) "Write textual debugging output",
Option "z" ["circular"] (ReqArg add_circular "CHR:LEN") "Refseq CHR is circular with length LEN",
Option "R" ["refseq"] (ReqArg set_range "RANGE") "Read only range of reference sequences",
Option "p" ["improper-pairs"] (NoArg set_improper) "Include improper pairs",
Option "u" ["unaligned"] (NoArg set_unaligned) "Include unaligned reads and pairs",
Option "1" ["single-read"] (NoArg set_single) "Pretend there is no second mate",
Option "m" ["multimappers"] (NoArg set_multi) "Process multi-mappers (by dropping secondary alignments)",
Option "c" ["cheap"] (NoArg set_cheap) "Cheap computation: skip the consensus calling",
Option "k" ["keep","mark-only"](NoArg set_keep) "Mark duplicates, but include them in output",
Option "Q" ["max-qual"] (ReqArg set_qual "QUAL") "Set maximum quality after consensus call to QUAL",
Option "l" ["min-length"] (ReqArg set_len "LEN") "Discard reads shorter than LEN",
Option "q" ["min-mapq"] (ReqArg set_mapq "QUAL") "Discard reads with map quality lower than QUAL",
Option "s" ["no-strand"] (NoArg set_no_strand) "Strand of alignments is uninformative",
Option "r" ["ignore-rg"] (NoArg set_no_rg) "Ignore read groups when looking for duplicates",
Option "v" ["verbose"] (NoArg set_verbose) "Print more diagnostics",
Option "h?" ["help","usage"] (NoArg (const usage)) "Display this message",
Option "V" ["version"] (NoArg (const vrsn)) "Display version number and exit" ]
where
set_output "-" c = return $ c { output = Just $ \_ -> pipeBamOutput, putResult = hPutStr stderr }
set_output f c = return $ c { output = Just $ \_ -> writeBamFile f }
set_lib_out f c = return $ c { output = Just $ writeLibBamFiles f }
set_debug_out c = return $ c { output = Just $ \_ -> pipeSamOutput, putResult = hPutStr stderr }
set_qual n c = readIO n >>= \q -> return $ c { collapse = cons_collapse' (Q q) }
set_no_strand c = return $ c { strand_preserved = False }
set_verbose c = return $ c { debug = hPutStr stderr }
set_improper c = return $ c { keep_improper = True }
set_single c = return $ c { transform = make_single }
set_cheap c = return $ c { collapse = cheap_collapse' }
set_keep c = return $ c { keep_all = True }
set_unaligned c = return $ c { keep_unaligned = True }
set_len n c = readIO n >>= \l -> return $ c { min_len = l }
set_mapq n c = readIO n >>= \q -> return $ c { min_qual = Q q }
set_no_rg c = return $ c { get_label = get_no_library }
set_multi c = return $ c { clean_multimap = clean_multi_flags }
set_range a c
| a == "A" || a == "a" = return $ c { which = Allrefs }
| a == "U" || a == "u" = return $ c { which = Unaln }
| otherwise = case reads a of
[ (x,"") ] -> return $ c { which = Some (Refseq $ x-1) (Refseq $ x-1) }
[ (x,'-':b) ] -> readIO b >>= \y ->
return $ c { which = Some (Refseq $ x-1) (Refseq $ y-1) }
_ -> fail $ "parse error in " ++ show a
add_circular a c = case break ((==) ':') a of
(nm,':':r) -> case reads r of
[(l,[])] | l > 0 -> return $ c { circulars = add_circular' (S.pack nm) l (circulars c) }
_ -> fail $ "couldn't parse length " ++ show r ++ " for " ++ show nm
_ -> fail $ "couldn't parse \"circular\" argument " ++ show a
add_circular' nm l io refs = do
(m1, refs') <- io refs
case filter (S.isPrefixOf nm . sq_name . snd) $ zip [0..] $ toList refs' of
[(k,a)] | sq_length a >= l -> let m2 = IM.insert k (sq_name a,l) m1
refs'' = Z.update k (a { sq_length = l }) refs'
in return (m2, refs'')
| otherwise -> fail $ "cannot wrap " ++ show nm ++ " to " ++ show l
++ ", which is more than the original " ++ show (sq_length a)
[] -> fail $ "no match for target sequence " ++ show nm
_ -> fail $ "target sequence " ++ show nm ++ " is ambiguous"
vrsn :: IO a
vrsn = do pn <- getProgName
hPutStrLn stderr $ pn ++ ", version " ++ showVersion version
exitSuccess
usage :: IO a
usage = do p <- getProgName
hPutStrLn stderr $ "Usage: " ++ usageInfo (p ++ info) options
exitSuccess
where
info = " [option...] [bam-file...]\n\
\Removes PCR duplicates from BAM files and calls a consensus for each duplicate set. \n\
\Input files must be sorted by coordinate and are merged on the fly. Options are:"
cons_collapse' :: Qual -> Bool -> Collapse
cons_collapse' m False = cons_collapse m
cons_collapse' m True = cons_collapse_keep m
cheap_collapse' :: Bool -> Collapse
cheap_collapse' False = cheap_collapse
cheap_collapse' True = cheap_collapse_keep
-- | Get library from BAM record.
-- This gets the read group from a bam record, then the library for read
-- group. This will work correctly if and only if the RG-LB field is
-- the name of the "Ur-Library", the common one before the first
-- amplification.
--
-- If no RG-LB field is present, RG-SM is used instead. This will work
-- if and only if no libraries were aliquotted and then pooled again.
--
-- Else the RG-ID field is used. This will work if and only if read
-- groups correspond directly to libraries.
--
-- If no RG is present, the empty string is returned. This serves as
-- fall-back.
get_library, get_no_library :: M.HashMap Seqid Seqid -> BamRec -> Seqid
get_library tbl br = M.lookupDefault rg rg tbl where rg = extAsString "RG" br
get_no_library _ _ = S.empty
mk_rg_tbl :: BamMeta -> M.HashMap Seqid Seqid
mk_rg_tbl hdr = M.fromList
[ (rg_id, rg_lb)
| ("RG",fields) <- meta_other_shit hdr
, rg_id <- take 1 [ i | ("ID",i) <- fields ]
, rg_lb <- take 1 $ [ l | ("LB",l) <- fields ]
++ [ s | ("SM",s) <- fields ]
++ [ rg_id ] ]
data Counts = Counts { tin :: !Int
, tout :: !Int
, good_singles :: !Int
, good_total :: !Int }
main :: IO ()
main = do
args <- getArgs
when (null args) usage
let (opts, files, errors) = getOpt Permute options args
unless (null errors) $ mapM_ (hPutStrLn stderr) errors >> exitFailure
Conf{..} <- foldr (>=>) return opts defaults
add_pg <- addPG $ Just version
(counts, ()) <- mergeInputRanges which files >=> run $ \hdr -> do
(circtable, refs') <- liftIO $ circulars (meta_refs hdr)
let tbl = mk_rg_tbl hdr
unless (M.null tbl) $ liftIO $ do
debug "mapping of read groups to libraries:\n"
mapM_ debug [ unpackSeqid k ++ " --> " ++ unpackSeqid v ++ "\n" | (k,v) <- M.toList tbl ]
let filters = progressBam "Rmdup at " debug refs' ><>
mapChunks (mapMaybe (transform . unpackBam)) ><>
mapChunksM (mapMM clean_multimap) ><>
filterStream (\br -> (keep_unaligned || is_aligned br) &&
(keep_improper || is_proper br) &&
eff_len br >= min_len)
let (co, ou) = case output of Nothing -> (cheap_collapse', skipToEof)
Just o -> (collapse, joinI $ wrapSortWith circtable $
o (get_label tbl) (add_pg hdr { meta_refs = refs' }))
ou' <- takeWhileE is_halfway_aligned ><> filters ><>
normalizeSortWith circtable ><>
filterStream (\b -> b_mapq b >= min_qual) ><>
rmdup (get_label tbl) strand_preserved (co keep_all) $
count_all (get_label tbl) `I.zip` ou
let do_copy = do liftIO $ debug "\27[Krmdup done; copying junk\n" ; joinI (filters ou')
do_bail = do liftIO $ debug "\27[Krmdup done\n" ; lift (run ou')
case which of
Unaln -> do_copy
_ | keep_unaligned -> do_copy
_ -> do_bail
putResult . unlines $
"\27[K#RG\tin\tout\tin@MQ20\tsingle@MQ20\tunseen\ttotal\t%unique\t%exhausted"
: map (uncurry do_report) (M.toList counts)
do_report :: Seqid -> Counts -> String
do_report lbl Counts{..} = intercalate "\t" fs
where
fs = label : showNum tin : showNum tout : showNum good_total : showNum good_singles :
report_estimate (estimateComplexity good_total good_singles)
label = if S.null lbl then "--" else unpackSeqid lbl
report_estimate Nothing = [ "N/A" ]
report_estimate (Just good_grand_total) =
[ showOOM (grand_total - fromIntegral tout)
, showOOM grand_total
, showFFloat (Just 1) rate []
, showFFloat (Just 1) exhaustion [] ]
where
grand_total = good_grand_total * fromIntegral tout / fromIntegral good_total
exhaustion = 100 * fromIntegral good_total / good_grand_total
rate = 100 * fromIntegral tout / fromIntegral tin :: Double
-- | Counting reads: we count total read in (ti), total reads out (to),
-- good (confidently mapped) singletons out (gs), total good
-- (confidently mapped) reads out (gt). Paired reads count 1, unpaired
-- reads count 2, and at the end we divide by 2. This ensures that we
-- don't double count mate pairs, while still working mostly sensibly in
-- the presence of broken BAM files.
count_all :: Functor m => (BamRec -> Seqid) -> Iteratee [BamRec] m (M.HashMap Seqid Counts)
count_all lbl = M.map fixup `fmap` I.foldl' plus M.empty
where
plus m b = M.insert (lbl b) cs m
where
!cs = plus1 (M.lookupDefault (Counts 0 0 0 0) (lbl b) m) b
plus1 (Counts ti to gs gt) b = Counts ti' to' gs' gt'
where
!w = if isPaired b then 1 else 2
!ti' = ti + w * extAsInt 1 "XP" b
!to' = to + w
!gs' = if b_mapq b >= Q 20 && extAsInt 1 "XP" b == 1 then gs + w else gs
!gt' = if b_mapq b >= Q 20 then gt + w else gt
fixup (Counts ti to gs gt) = Counts (div ti 2) (div to 2) (div gs 2) (div gt 2)
eff_len :: BamRec -> Int
eff_len b | isProperlyPaired b = abs $ b_isize b
| otherwise = V.length $ b_seq b
is_halfway_aligned :: BamRaw -> Bool
is_halfway_aligned = isValidRefseq . b_rname . unpackBam
is_aligned :: BamRec -> Bool
is_aligned b = not (isUnmapped b && isMateUnmapped b) && isValidRefseq (b_rname b)
is_proper :: BamRec -> Bool
is_proper b = not (isPaired b) || (isMateUnmapped b == isUnmapped b && isProperlyPaired b)
make_single :: BamRec -> Maybe BamRec
make_single b | isPaired b && isSecondMate b = Nothing
| isUnmapped b = Nothing
| not (isPaired b) = Just b
| otherwise = Just $ b { b_flag = b_flag b .&. complement pair_flags
, b_mrnm = invalidRefseq
, b_mpos = invalidPos
, b_isize = 0 }
where
pair_flags = flagPaired .|. flagProperlyPaired .|.
flagFirstMate .|. flagSecondMate .|.
flagMateUnmapped
mergeInputRanges :: (MonadIO m, MonadMask m)
=> Which -> [FilePath] -> Enumerator' BamMeta [BamRaw] m a
mergeInputRanges Allrefs fps = mergeInputs combineCoordinates fps
mergeInputRanges _ [ ] = \k -> return $ k mempty
mergeInputRanges rng (fp0:fps0) = go fp0 fps0
where
enum1 fp k1 = case rng of Allrefs -> decodeAnyBamFile fp k1
Some x y -> decodeBamFileRange x y fp k1
Unaln -> decodeWithIndex eneeBamUnaligned fp k1
go fp [ ] = enum1 fp
go fp (fp1:fps) = mergeEnums' (go fp1 fps) (enum1 fp) combineCoordinates
decodeBamFileRange x y = decodeWithIndex $
\idx -> foldr ((>=>) . eneeBamRefseq idx) return [x..y]
decodeWithIndex :: (MonadIO m, MonadMask m)
=> (BamIndex () -> Enumeratee [BamRaw] [BamRaw] m a)
-> FilePath -> (BamMeta -> Iteratee [BamRaw] m a)
-> m (Iteratee [BamRaw] m a)
decodeWithIndex enum fp k0 = do
idx <- liftIO $ readBamIndex fp
decodeAnyBamFile fp >=> run $ enum idx . k0
writeLibBamFiles :: (MonadIO m, MonadMask m)
=> FilePath -> (BamRec -> Seqid) -> BamMeta -> Iteratee [BamRec] m ()
writeLibBamFiles fp lbl hdr = tryHead >>= loop M.empty
where
loop m Nothing = liftIO . mapM_ run $ M.elems m
loop m (Just br) = do
let !l = lbl br
let !it = M.lookupDefault (writeBamFile (fp `subst` l) hdr) l m
it' <- liftIO $ enumPure1Chunk [br] it
let !m' = M.insert l it' m
tryHead >>= loop m'
subst [ ] _ = []
subst ('%':'s':rest) l = unpackSeqid l ++ subst rest l
subst ('%':'%':rest) l = '%' : subst rest l
subst ( c : rest) l = c : subst rest l
mapMM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMM f = go []
where
go acc [ ] = return $ reverse acc
go acc (a:as) = do b <- f a ; go (maybe acc (:acc) b) as
check_flags :: Monad m => BamRec -> m (Maybe BamRec)
check_flags b | extAsInt 1 "HI" b /= 1 = fail "cannot deal with HI /= 1"
| extAsInt 1 "IH" b /= 1 = fail "cannot deal with IH /= 1"
| extAsInt 1 "NH" b /= 1 = fail "cannot deal with NH /= 1"
| otherwise = return $ Just b
clean_multi_flags :: Monad m => BamRec -> m (Maybe BamRec)
clean_multi_flags b = return $ if extAsInt 1 "HI" b /= 1 then Nothing else Just b'
where
b' = b { b_exts = deleteE "HI" $ deleteE "IH" $ deleteE "NH" $ b_exts b }
-- Given a map from reference sequences to arguments, extract those
-- groups as list, apply a function to the argument and the list, pass
-- the result on. Absent groups are passed on as they are. Note that
-- ordering within groups is messed up (it doesn't matter here).
mapAtGroups :: Monad m => IM.IntMap a -> (a -> [BamRec] -> [BamRec]) -> Enumeratee [BamRec] [BamRec] m b
mapAtGroups m f = eneeCheckIfDonePass no_group
where
no_group k (Just e) = idone (liftI k) $ EOF (Just e)
no_group k Nothing = tryHead >>= maybe (idone (liftI k) $ EOF Nothing) (\a -> no_group_1 a k Nothing)
no_group_1 _ k (Just e) = idone (liftI k) $ EOF (Just e)
no_group_1 a k Nothing = case IM.lookup (b_rname_int a) m of
Nothing -> eneeCheckIfDonePass no_group . k $ Chunk [a]
Just arg -> cont_group (b_rname a) arg [a] k Nothing
cont_group _rn _arg _acc k (Just e) = idone (liftI k) $ EOF (Just e)
cont_group rn arg acc k Nothing = tryHead >>= maybe flush_eof check1
where
flush_eof = idone (k $ Chunk $ f arg acc) (EOF Nothing)
flush_go a = eneeCheckIfDonePass (no_group_1 a) . k . Chunk $ f arg acc
check1 a | b_rname a == rn = cont_group rn arg (a:acc) k Nothing
| otherwise = flush_go a
b_rname_int = fromIntegral . unRefseq . b_rname
normalizeSortWith :: Monad m => IM.IntMap (Seqid, Int) -> Enumeratee [BamRec] [BamRec] m a
normalizeSortWith m = mapAtGroups m $ \(nm,l) -> sortPos . map (normalizeTo nm l)
wrapSortWith :: Monad m => IM.IntMap (Seqid, Int) -> Enumeratee [BamRec] [BamRec] m a
wrapSortWith m = mapAtGroups m $ \(_,l) -> sortPos . concatMap (wrapTo l)
sortPos :: [BamRec] -> [BamRec]
sortPos l = VV.toList $ runST (VV.unsafeThaw (VV.fromList l) >>= \vm -> sortBy (comparing b_pos) vm >> VV.unsafeFreeze vm)