diff --git a/Bio/Adna.hs b/Bio/Adna.hs new file mode 100644 --- /dev/null +++ b/Bio/Adna.hs @@ -0,0 +1,655 @@ +-- | Things specific to ancient DNA, e.g. damage models. +-- +-- For aDNA, we need a substitution probability. We have three options: +-- use an empirically determined PSSM, use an arithmetically defined +-- PSSM based on the /Johnson/ model, use a context sensitive PSSM based +-- on the /Johnson/ model and an alignment. Using /Dindel/, actual +-- substitutions relative to a called haplotype would be taken into +-- account. Since we're not going to do that, taking alignments into +-- account is difficult, somewhat approximate, and therefore not worth +-- the hassle. + +module Bio.Adna ( + DmgStats(..), + CompositionStats, + SubstitutionStats, + addFragType, + damagePatternsIter, + damagePatternsIterMD, + damagePatternsIter2Bit, + alnFromMd, + + DamageParameters(..), + NewDamageParameters(..), + GenDamageParameters(..), + DamageModel, + bang, nudge, + Alignment(..), + FragType(..), + Subst(..), + + NPair, + npair, + fst_np, + snd_np, + + noDamage, + univDamage, + empDamage, + Mat44D(..), + MMat44D(..), + scalarMat, + complMat, + freezeMats, + + bwa_cal_maxdiff + ) where + +import Bio.Bam +import Bio.Prelude +import Bio.TwoBit + +import qualified Data.Vector as V +import qualified Data.Vector.Generic as G +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import qualified Streaming.Prelude as Q + +-- | We represent substitution matrices by the type 'Mat44D'. Internally, +-- this is a vector of packed vectors. Conveniently, each of the packed +-- vectors represents all transitions /into/ the given nucleotide. + +newtype Mat44D = Mat44D (U.Vector Double) deriving (Show, Generic) +newtype MMat44D = MMat44D (UM.IOVector Double) + +-- | A 'DamageModel' is a function that gives substitution matrices for +-- each position in a read. The 'DamageModel' can depend on whether the +-- alignment is reversed, the length of the read and the position. (In +-- practice, we should probably memoize precomputed damage models +-- somehow.) + +type DamageModel = Bool -> Int -> Int -> Mat44D +data Subst = Nucleotide :-> Nucleotide deriving (Eq, Ord, Ix, Show) + +infix 9 :-> +infix 8 `bang` + +-- | Convenience function to access a substitution matrix that has a +-- mnemonic reading. +{-# INLINE bang #-} +bang :: Mat44D -> Subst -> Double +bang (Mat44D v) (N x :-> N y) + | U.length v == 16 = v U.! (fromIntegral x + 4 * fromIntegral y) + | otherwise = error $ "Huh? " ++ show (U.length v) + +{-# INLINE nudge #-} +nudge :: MMat44D -> Subst -> Double -> IO () +nudge (MMat44D v) (N x :-> N y) a = UM.read v i >>= UM.write v i . (+) a + where i = fromIntegral x + 4 * fromIntegral y + +scalarMat :: Double -> Mat44D +scalarMat s = Mat44D $ U.fromListN 16 [ s, 0, 0, 0 + , 0, s, 0, 0 + , 0, 0, s, 0 + , 0, 0, 0, s ] +complMat :: Mat44D -> Mat44D +complMat v = Mat44D $ U.fromListN 16 [ v `bang` compl x :-> compl y + | y <- range (nucA, nucT) + , x <- range (nucA, nucT) ] + +-- | Adds the two matrices of a mutable substitution model (one for each +-- strand) appropriately, normalizes the result (to make probabilities +-- from pseudo-counts), and freezes that into one immutable matrix. We +-- add a single count everywhere to avoid getting NaN from bizarre data. +freezeMats :: MMat44D -> MMat44D -> IO Mat44D +freezeMats (MMat44D vv) (MMat44D ww) = do + v <- Mat44D <$> U.freeze vv + w <- complMat . Mat44D <$> U.freeze ww + + let sums = U.generate 4 $ \x0 -> + let x = N $ fromIntegral x0 + in sum [ v `bang` x :-> z + w `bang` x :-> z + | z <- range (nucA, nucT) ] + 4 + + return . Mat44D $ U.fromListN 16 + [ (v `bang` x :-> y + w `bang` x :-> y + 1) / s + | y <- range (nucA, nucT) + , x <- range (nucA, nucT) + , let s = sums U.! fromIntegral (unN x) ] + + +-- | 'DamageModel' for undamaged DNA. The likelihoods follow directly +-- from the quality score. This needs elaboration to see what to do +-- with amibiguity codes (even though those haven't actually been +-- observed in the wild). +noDamage :: DamageModel +noDamage _ _ _ = one where !one = scalarMat 1 + + +-- | Parameters for the universal damage model. +-- +-- We assume the correct model is either no damage, or single strand +-- damage, or double strand damage. Each of them comes with a +-- probability. It turns out that blending them into one is simply +-- accomplished by multiplying these probabilities onto the deamination +-- probabilities. +-- +-- For single stranded library prep, only one kind of damage occurs (C +-- frequency ('ssd_sigma') in single stranded parts, and the overhang +-- length is distributed exponentially with parameter 'ssd_lambda' at +-- the 5' end and 'ssd_kappa' at the 3' end. (Without UDG treatment, +-- those will be equal. With UDG, those are much smaller and in fact +-- don't literally represent overhangs.) +-- +-- For double stranded library prep, we get C->T damage at the 5' end +-- and G->A at the 3' end with rate 'dsd_sigma' and both in the interior +-- with rate 'dsd_delta'. Everything is symmetric, and therefore the +-- orientation of the aligned read doesn't matter either. Both +-- overhangs follow a distribution with parameter 'dsd_lambda'. + +data DamageParameters float = DP { ssd_sigma :: !float -- deamination rate in ss DNA, SS model + , ssd_delta :: !float -- deamination rate in ds DNA, SS model + , ssd_lambda :: !float -- param for geom. distribution, 5' end, SS model + , ssd_kappa :: !float -- param for geom. distribution, 3' end, SS model + , dsd_sigma :: !float -- deamination rate in ss DNA, DS model + , dsd_delta :: !float -- deamination rate in ds DNA, DS model + , dsd_lambda :: !float } -- param for geom. distribution, DS model + deriving (Read, Show, Generic) + +data NewDamageParameters vec float = NDP { dp_gc_frac :: !float + , dp_mu :: !float + , dp_nu :: !float + , dp_alpha5 :: !(vec float) + , dp_beta5 :: !(vec float) + , dp_alpha :: !float + , dp_beta :: !float + , dp_alpha3 :: !(vec float) + , dp_beta3 :: !(vec float) } + deriving (Read, Show, Generic) + +data GenDamageParameters vec float + = UnknownDamage + | OldDamage (DamageParameters float) + | NewDamage (NewDamageParameters vec float) + deriving (Show, Generic, Read) + + + +-- | Generic substitution matrix, has C->T and G->A deamination as +-- parameters. Setting 'p' or 'q' to 0 as appropriate makes this apply +-- to the single stranded or undamaged case. + +{-# INLINE genSubstMat #-} +genSubstMat :: Double -> Double -> Mat44D +genSubstMat p q = Mat44D $ U.fromListN 16 [ 1, 0, q, 0 + , 0, 1-p, 0, 0 + , 0, 0, 1-q, 0 + , 0, p, 0, 1 ] + +univDamage :: DamageParameters Double -> DamageModel +univDamage DP{..} r l i = genSubstMat (p1+p2) (q1+q2) + where + (p1, q1) = if r then let lam5 = ssd_lambda ^ (l-i) + lam3 = ssd_kappa ^ (1+i) + lam = lam3 + lam5 - lam3 * lam5 + p = ssd_sigma * lam + ssd_delta * (1-lam) + in (0,p) + else let lam5 = ssd_lambda ^ (1+i) + lam3 = ssd_kappa ^ (l-i) + lam = lam3 + lam5 - lam3 * lam5 + p = ssd_sigma * lam + ssd_delta * (1-lam) + in (p,0) + + p2 = dsd_sigma * lam5_ds + dsd_delta * (1-lam5_ds) + q2 = dsd_sigma * lam3_ds + dsd_delta * (1-lam3_ds) + lam5_ds = dsd_lambda ^ (1+i) + lam3_ds = dsd_lambda ^ (l-i) + +empDamage :: NewDamageParameters U.Vector Double -> DamageModel +empDamage NDP{..} = + \r l i -> if i+i < l then + if r then fromMaybe middleRev (rev5 V.!? i) + else fromMaybe middle (fwd5 V.!? i) + else + if r then fromMaybe middleRev (rev3 V.!? (l-i-1)) + else fromMaybe middle (fwd3 V.!? (l-i-1)) + where + !middle = genSubstMat' dp_alpha dp_beta + !middleRev = genSubstMat' dp_beta dp_alpha + + !fwd5 = V.zipWith genSubstMat' (G.convert dp_alpha5) (G.convert dp_beta5) + !fwd3 = V.zipWith genSubstMat' (G.convert dp_alpha3) (G.convert dp_beta3) + + !rev5 = V.zipWith genSubstMat' (G.convert dp_beta5) (G.convert dp_alpha5) + !rev3 = V.zipWith genSubstMat' (G.convert dp_beta3) (G.convert dp_alpha3) + + genSubstMat' a b = genSubstMat (recip $ 1 + exp (-a)) (recip $ 1 + exp (-b)) + + +-- | Collected \"traditional\" statistics: +-- +-- * Base composition near 5' end and near 3' end. Each consists of +-- five vectors of counts of A,C,G,T, and everything else. +-- 'basecompo5' begins with 'context' bases to the left of the 5' end, +-- 'basecompo3' ends with 'context' bases to the right of the 3' end. +-- +-- * Substitutions. Counted from the reconstructed alignment, once +-- around the 5' end and once around the 3' end. For a total of 2*4*4 +-- different substitutions. Positions where the query has a gap are +-- skipped. +-- +-- * Substitutions at CpG motifs. Also counted from the reconstructed +-- alignment, and a CpG site is simply the sequence CG in the +-- reference. Gaps may confuse that definition, so that CpHpG still +-- counts as CpG, because the H is gapped. That might actually +-- be desirable. +-- +-- * Conditional substitutions. The 5' and 3' ends count as damaged if +-- the very last position has a C-to-T substitution. With that in +-- mind, 'substs5d5', 'substs5d3', 'substs5dd' are like 'substs5', but +-- counting only reads where the 5' end is damaged, where the 3' end +-- is damaged, and where both ends are damaged, respectively. + +data DmgStats = DmgStats { + basecompo5 :: CompositionStats, + basecompo3 :: CompositionStats, + substs5 :: SubstitutionStats, + substs3 :: SubstitutionStats, + substs5d5 :: SubstitutionStats, + substs3d5 :: SubstitutionStats, + substs5d3 :: SubstitutionStats, + substs3d3 :: SubstitutionStats, + substs5dd :: SubstitutionStats, + substs3dd :: SubstitutionStats, + substs5cpg :: SubstitutionStats, + substs3cpg :: SubstitutionStats } + deriving Show + +type CompositionStats = [( Maybe Nucleotide, U.Vector Int )] +type SubstitutionStats = [( Subst, U.Vector Int )] + + +data FragType = Complete | Leading | Trailing deriving (Show, Eq) + +-- | Compact storage of a pair of ambiguous 'Nucleotides'. Used to +-- represent alignments in a way that is accessible even to assembly +-- code. The first and sencond field are stored in the low and high +-- nybble, respectively. See 'fst_np', 'snd_np', 'npair'. +newtype NPair = NPair Word8 deriving (Eq, Ord) + +npair :: Nucleotides -> Nucleotides -> NPair +npair (Ns r) (Ns q) = NPair $ shiftL q 4 .|. r .&. 0xF + +fst_np, snd_np :: NPair -> Nucleotides +fst_np (NPair w) = Ns (w .&. 0xF) +snd_np (NPair w) = Ns (shiftR w 4) + +instance Storable NPair where + sizeOf _ = 1 + alignment _ = 1 + peek p = NPair <$> peek (castPtr p :: Ptr Word8) + poke p (NPair v) = poke (castPtr p :: Ptr Word8) v + +instance Show NPair where + showsPrec _ p = shows (fst_np p) . (:) '/' . shows (snd_np p) + +-- | Alignment record. The reference sequence is filled with Ns if +-- missing. +data Alignment = ALN + { a_sequence :: !(VS.Vector NPair) -- the alignment proper + , a_fragment_type :: !FragType } -- was the adapter trimmed? + +addFragType :: Monad m => BamMeta -> Stream (Of BamRaw) m b -> Stream (Of (BamRaw,FragType)) m b +addFragType meta = Q.map $ \br -> (br, case unpackBam br of + b | isFirstMate b && isPaired b -> Leading + | isSecondMate b && isPaired b -> Trailing + | not sane -> Complete -- leeHom fscked it up + | isFirstMate b || isSecondMate b -> Complete -- old style flagging + | isTrimmed b || isMerged b -> Complete -- new style flagging + | otherwise -> Leading) + where + sane = null [ () | ("PG",line) <- meta_other_shit meta + , ("PN","mergeTrimReadsBAM") <- line ] + +-- | Stream transformer that computes some statistics from plain BAM +-- (no MD field needed) and a 2bit file. The 'Alignment' is also +-- reconstructed and passed downstream. The final value of the source +-- stream ends up in the 'stats_more' field of the result. +-- +-- * Get the reference sequence including both contexts once. If this +-- includes invalid sequence (negative coordinate), pad suitably. +-- * Accumulate counts for the valid parts around 5' and 3' ends as +-- appropriate from flags and config. +-- * Combine the part that was aligned to (so no context) with the read +-- to reconstruct the alignment. +-- +-- Arguments are the table of reference names, the 2bit file with the +-- reference, the amount of context outside the alignment desired, and +-- the amount of context inside desired. +-- +-- For 'Complete' fragments, we cut the read in the middle, so the 5' +-- and 3' plots stay clean from each other's influence. 'Leading' and +-- 'Trailing' fragments count completely towards the appropriate end. + +damagePatternsIter2Bit :: MonadIO m + => Refs -> TwoBitFile -> Int -> Int + -> Stream (Of (BamRaw,FragType)) m x + -> Stream (Of Alignment) m DmgStats +damagePatternsIter2Bit refs tbf ctx rng = + damagePatternsIter ctx rng . + Q.mapMaybe (\(br,ft) -> do + let b@BamRec{..} = unpackBam br + guard (not $ isUnmapped b) + let ref_nm = sq_name $ getRef refs b_rname + ref = getFragment tbf ref_nm (b_pos - ctx) (alignedLength b_cigar + 2*ctx) + pps = aln_from_ref (U.drop ctx ref) b_seq b_cigar + return (b, ft, ref, pps)) + +-- | Stream transformer that computes some statistics from plain BAM +-- with a valid MD field. The 'Alignment' is also reconstructed and +-- passed downstream. The final value of the source stream becomes +-- the 'stats_more' field of the result. +-- +-- * Reconstruct the alignment from CIGAR, SEQ, and MD. +-- * Filter the alignment to get the reference sequence, accumulate it. +-- * Accumulate everything over the alignment. +-- +-- The argument is the amount of context inside desired. +-- +-- For 'Complete' fragments, we cut the read in the middle, so the 5' +-- and 3' plots stay clean from each other's influence. 'Leading' and +-- 'Trailing' fragments count completely towards the appropriate end. + +damagePatternsIterMD :: MonadIO m + => Int + -> Stream (Of (BamRaw,FragType)) m x + -> Stream (Of Alignment) m DmgStats +damagePatternsIterMD rng = + damagePatternsIter 0 rng . + Q.mapMaybe (\(br,ft) -> do + let b@BamRec{..} = unpackBam br + guard (not $ isUnmapped b) + md <- getMd b + let pps = alnFromMd b_seq b_cigar md + ref = U.convert $ VS.map fromN $ VS.filter ((/=) gap) $ VS.map fst_np pps + return (b, ft, ref, pps)) + where + fromN ns | ns == nucsA = 2 + | ns == nucsC = 1 + | ns == nucsG = 3 + | ns == nucsT = 0 + | otherwise = 4 + +-- | Common logic for statistics. The function 'get_ref_and_aln' +-- reconstructs reference sequence and alignment from a Bam record. It +-- is expected to construct the alignment with respect to the forwards +-- strand of the reference; we reverse-complement it if necessary. +damagePatternsIter :: MonadIO m + => Int -> Int + -> Stream (Of (BamRec, FragType, U.Vector Word8, VS.Vector NPair)) m x + -> Stream (Of Alignment) m DmgStats +damagePatternsIter ctx rng stream = do + let maxwidth = ctx + rng + acc_bc <- liftIO $ UM.replicate (2 * 5 * maxwidth) (0::Int) + acc_st <- liftIO $ UM.replicate (2 * 4 * 4 * 4 * rng) (0::Int) + acc_cg <- liftIO $ UM.replicate (2 * 2 * 4 * rng) (0::Int) + + void $ flip Q.mapM (Q.map revcom_both stream) $ + \(BamRec{..}, a_fragment_type, ref, a_sequence) -> liftIO $ do + -- basecompositon near 5' end, near 3' end + let (width5, width3) = case a_fragment_type of + Leading -> (full_width, 0) + Trailing -> (0, full_width) + Complete -> (half_width, half_width) + where full_width = min (U.length ref) $ ctx + min rng (alignedLength b_cigar) + half_width = min (U.length ref) $ ctx + min rng (alignedLength b_cigar `div` 2) + mapM_ (\i -> bump (fromIntegral (ref U.! i ) * maxwidth + i) acc_bc) [0 .. width5-1] + mapM_ (\i -> bump (fromIntegral (ref U.! (i + U.length ref) +6) * maxwidth + i) acc_bc) [-width3 .. -1] + + -- For substitutions, decide what damage class we're in: + -- 0 - no damage, 1 - damaged 5' end, 2 - damaged 3' end, 3 - both + let dmgbase = 2*4*4*rng * + ((if VS.null a_sequence || VS.head a_sequence /= npair nucsC nucsT then 1 else 0) + +(if VS.null a_sequence || VS.last a_sequence /= npair nucsC nucsT then 2 else 0)) + + -- substitutions near 5' end + let len_at_5 = case a_fragment_type of Leading -> min rng (G.length b_seq) + Complete -> min rng (G.length b_seq `div` 2) + Trailing -> 0 + flip G.imapM_ (VS.take len_at_5 a_sequence) $ + \i uv -> withPair uv $ \j -> bump (j * rng + i + dmgbase) acc_st + + -- substitutions at CpG sites near 5' end + G.izipWithM_ + (\i uv wz -> + when (fst_np uv == nucsC && fst_np wz == nucsG) $ do + withNs (snd_np uv) $ \y -> bump ( y * rng + i ) acc_cg + withNs (snd_np wz) $ \y -> bump ((y+4) * rng + i+1) acc_cg) + (VS.take len_at_5 a_sequence) (VS.drop 1 a_sequence) + + -- substitutions near 3' end + let len_at_3 = case a_fragment_type of Leading -> 0 + Complete -> min rng (G.length b_seq `div` 2) + Trailing -> min rng (G.length b_seq) + flip G.imapM_ (VS.take len_at_3 (VS.reverse a_sequence)) $ + \i uv -> withPair uv $ \j -> bump ((17+j) * rng -i -1 + dmgbase) acc_st + + -- substitutions at CpG sites near 3' end + G.izipWithM_ + (\i wz uv -> + when (fst_np uv == nucsC && fst_np wz == nucsG) $ do + withNs (snd_np uv) $ \y -> bump ((y+ 9) * rng - i-2) acc_cg + withNs (snd_np wz) $ \y -> bump ((y+13) * rng - i-1) acc_cg) + (VS.take len_at_3 (VS.reverse a_sequence)) + (VS.drop 1 (VS.reverse a_sequence)) + + return ALN{..} + + + let nsubsts = 4*4*rng + mk_substs off = sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((4*i+j)*rng + off*nsubsts) rng acc_st) + | (i,n1) <- zip [0..] [nucA..nucT] + , (j,n2) <- zip [0..] [nucA..nucT] ] + + accs <- liftIO $ DmgStats <$> sequence [ (,) nuc <$> U.unsafeFreeze (UM.slice (i*maxwidth) maxwidth acc_bc) + | (i,nuc) <- zip [2,1,3,0,4] [Just nucA,Just nucC,Just nucG,Just nucT,Nothing] ] + <*> sequence [ (,) nuc <$> U.unsafeFreeze (UM.slice (i*maxwidth) maxwidth acc_bc) + | (i,nuc) <- zip [7,6,8,5,9] [Just nucA,Just nucC,Just nucG,Just nucT,Nothing] ] + + <*> mk_substs 0 + <*> mk_substs 1 + <*> mk_substs 2 + <*> mk_substs 3 + <*> mk_substs 4 + <*> mk_substs 5 + <*> mk_substs 6 + <*> mk_substs 7 + + <*> sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((i+j)*rng) rng acc_cg) + | (i,n1) <- [(0,nucC), (4,nucG)] + , (j,n2) <- zip [0..] [nucA..nucT] ] + + <*> sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((i+j)*rng) rng acc_cg) + | (i,n2) <- [(8,nucC), (12,nucG)] + , (j,n1) <- zip [0..] [nucA..nucT] ] + + return $ accs { substs5 = mconcat [ substs5 accs, substs5d5 accs, substs5d3 accs, substs5dd accs ] + , substs3 = mconcat [ substs3 accs, substs3d5 accs, substs3d3 accs, substs3dd accs ] + , substs5d5 = mconcat [ substs5d5 accs, substs5dd accs] + , substs3d5 = mconcat [ substs3d5 accs, substs3dd accs] + , substs5d3 = mconcat [ substs5d3 accs, substs5dd accs] + , substs3d3 = mconcat [ substs3d3 accs, substs3dd accs] } + where + {-# INLINE withPair #-} + withPair (NPair i) k = + case pairTab `U.unsafeIndex` fromIntegral i of + j -> when (j >= 0) (k j) + + !pairTab = U.replicate 256 (-1) U.// + [ (fromIntegral i, x*4+y) | (u,x) <- zip [nucsA, nucsC, nucsG, nucsT] [0,1,2,3] + , (v,y) <- zip [nucsA, nucsC, nucsG, nucsT] [0,1,2,3] + , let NPair i = npair u v ] + {-# INLINE bump #-} + bump i v = UM.unsafeRead v i >>= UM.unsafeWrite v i . succ + + {-# INLINE withNs #-} + withNs ns k | ns == nucsA = k 0 + | ns == nucsC = k 1 + | ns == nucsG = k 2 + | ns == nucsT = k 3 + | otherwise = return () + + +instance Monoid DmgStats where + mappend = (<>) + mempty = DmgStats { basecompo5 = empty_compo + , basecompo3 = empty_compo + , substs5 = empty_subst + , substs3 = empty_subst + , substs5d5 = empty_subst + , substs3d5 = empty_subst + , substs5d3 = empty_subst + , substs3d3 = empty_subst + , substs5dd = empty_subst + , substs3dd = empty_subst + , substs5cpg = empty_subst + , substs3cpg = empty_subst } + where + empty_compo = [ (nuc, U.empty) | nuc <- [Just nucA, Just nucC, Just nucG, Just nucT, Nothing] ] + empty_subst = [ (n1 :-> n2, U.empty) | n1 <- [nucA..nucT], n2 <- [nucA..nucT] ] + +instance Semigroup DmgStats where + (<>) = merge_dmg_stats + +merge_dmg_stats :: DmgStats -> DmgStats -> DmgStats +merge_dmg_stats a b = + DmgStats { basecompo5 = zipWith s1 (basecompo5 a) (basecompo5 b) + , basecompo3 = zipWith s1 (basecompo3 a) (basecompo3 b) + , substs5 = zipWith s2 (substs5 a) (substs5 b) + , substs3 = zipWith s2 (substs3 a) (substs3 b) + , substs5d5 = zipWith s2 (substs5d5 a) (substs5d5 b) + , substs3d5 = zipWith s2 (substs3d5 a) (substs3d5 b) + , substs5d3 = zipWith s2 (substs5d3 a) (substs5d3 b) + , substs3d3 = zipWith s2 (substs3d3 a) (substs3d3 b) + , substs5dd = zipWith s2 (substs5dd a) (substs5dd b) + , substs3dd = zipWith s2 (substs3dd a) (substs3dd b) + , substs5cpg = zipWith s2 (substs5cpg a) (substs5cpg b) + , substs3cpg = zipWith s2 (substs3cpg a) (substs3cpg b) } + where + s1 (x, u) (z, v) | x /= z = error "Mismatch in zip. This is a bug." + | U.null u = (x, v) + | U.null v = (x, u) + | otherwise = (x, U.zipWith (+) u v) + + s2 (x :-> y, u) (z :-> w, v) | x /= z || y /= w = error "Mismatch in zip. This is a bug." + | U.null u = (x :-> y, v) + | U.null v = (x :-> y, u) + | otherwise = (x :-> y, U.zipWith (+) u v) + + +revcom_both :: ( BamRec, FragType, U.Vector Word8, VS.Vector NPair ) + -> ( BamRec, FragType, U.Vector Word8, VS.Vector NPair ) +revcom_both (b, ft, ref, pps) + | isReversed b = ( b, ft, revcom_ref ref, revcom_pairs pps ) + | otherwise = ( b, ft, ref, pps ) + where + revcom_ref = U.reverse . U.map (\c -> if c > 3 then c else xor c 2) + revcom_pairs = VS.reverse . VS.map (\p -> npair (compls $ fst_np p) (compls $ snd_np p)) + + +-- | Reconstructs the alignment from reference, query, and cigar. Only +-- positions where the query is not gapped are produced. +aln_from_ref :: U.Vector Word8 -> Vector_Nucs_half Nucleotides -> VS.Vector Cigar -> VS.Vector NPair +aln_from_ref ref0 qry0 cig0 = VS.fromList $ step ref0 qry0 cig0 + where + step ref qry cig1 + | U.null ref || G.null qry || G.null cig1 = [] + | otherwise = case G.unsafeHead cig1 of { op :* n -> + case G.unsafeTail cig1 of { cig -> + case op of { + + Mat -> zipWith (npair . nn) (G.toList (G.take n ref)) + (G.toList (G.take n qry)) ++ step (G.drop n ref) (G.drop n qry) cig ; + Del -> step (G.drop n ref) qry cig ; + Ins -> map (npair gap) (G.toList (G.take n qry)) ++ step ref (G.drop n qry) cig ; + SMa -> map (npair gap) (G.toList (G.take n qry)) ++ step ref (G.drop n qry) cig ; + HMa -> replicate n (npair gap nucsN) ++ step ref qry cig ; + Nop -> step ref qry cig ; + Pad -> step ref qry cig }}} + + nn 0 = nucsT + nn 1 = nucsC + nn 2 = nucsA + nn 3 = nucsG + nn _ = nucsN + + +-- | Reconstructs the alignment from query, cigar, and md. Only +-- positions where the query is not gapped are produced. +alnFromMd :: Vector_Nucs_half Nucleotides -> VS.Vector Cigar -> [MdOp] -> VS.Vector NPair +alnFromMd qry0 cig0 md0 = VS.fromList $ step qry0 cig0 md0 + where + step qry cig1 md + | G.null qry || G.null cig1 || null md = [] + | otherwise = case G.unsafeHead cig1 of op :* n -> step' qry op n (G.unsafeTail cig1) md + + step' qry _ 0 cig md = step qry cig md + step' qry op n cig (MdNum 0 : md) = step' qry op n cig md + step' qry op n cig (MdDel [] : md) = step' qry op n cig md + + step' qry Mat n cig (MdNum m : md) + | n < m = map twin (G.toList (G.take n qry)) ++ step (G.drop n qry) cig (MdNum (m-n) : md) + | n > m = map twin (G.toList (G.take m qry)) ++ step' (G.drop m qry) Mat (n-m) cig md + | n == m = map twin (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md + step' qry Mat n cig (MdRep c : md) = npair c (G.head qry) : step' (G.tail qry) Mat (n-1) cig md + step' _ Mat _ _ _ = [] + + step' qry Del n cig (MdDel (_:ss) : md) = step' qry Del (n-1) cig (MdDel ss : md) + step' _ Del _ _ _ = [] + + step' qry Ins n cig md = map (npair gap) (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md + step' qry SMa n cig md = map (npair gap) (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md + step' qry HMa n cig md = replicate n (npair gap nucsN) ++ step qry cig md + step' qry Nop _ cig md = step qry cig md + step' qry Pad _ cig md = step qry cig md + + twin q = npair q q + +-- | Number of mismatches allowed by BWA. +-- @bwa_cal_maxdiff thresh len@ returns the number of mismatches +-- @bwa aln -n $tresh@ would allow in a read of length @len@. For +-- reference, here is the code from BWA that computes it (we assume @err +-- = 0.02@, just like BWA): +-- +-- @ +-- int bwa_cal_maxdiff(int l, double err, double thres) +-- { +-- double elambda = exp(-l * err); +-- double sum, y = 1.0; +-- int k, x = 1; +-- for (k = 1, sum = elambda; k < 1000; ++k) { +-- y *= l * err; +-- x *= k; +-- sum += elambda * y / x; +-- if (1.0 - sum < thres) return k; +-- } +-- return 2; +-- } +-- @ + +bwa_cal_maxdiff :: Double -> Int -> Int +bwa_cal_maxdiff thresh len = k_fin-1 + where + (k_fin, _, _, _) = head $ dropWhile bad $ iterate step (1,elambda,1,1) + + err = 0.02 + elambda = exp . negate $ fromIntegral len * err + + step (k, s, x, y) = (k+1, s', x', y') + where y' = y * fromIntegral len * err + x' = x * fromIntegral k + s' = s + elambda * y' / x' + + bad (_, s, _, _) = 1-s >= thresh + diff --git a/Bio/Align.hs b/Bio/Align.hs new file mode 100644 --- /dev/null +++ b/Bio/Align.hs @@ -0,0 +1,85 @@ +module Bio.Align ( + Mode(..), + myersAlign, + showAligned + ) where + +import Bio.Prelude hiding ( lefts, rights ) +import Foreign.C.String ( CString ) +import Foreign.C.Types ( CInt(..) ) +import Foreign.Marshal.Alloc ( allocaBytes ) + +import qualified Data.ByteString.Char8 as S +import qualified Data.ByteString.Unsafe as S +import qualified Data.ByteString.Lazy.Char8 as L + +foreign import ccall unsafe "myers_align.h myers_diff" myers_diff :: + CString -> CInt -> -- sequence A and length A + CInt -> -- mode (an enum) + CString -> CInt -> -- sequence B and length B + CInt -> -- max distance + CString -> -- backtracing space A + CString -> -- backtracing space B + IO CInt -- returns distance + +-- | Mode argument for 'myersAlign', determines where free gaps are +-- allowed. +data Mode = Globally -- ^ align globally, without gaps at either end + | HasPrefix -- ^ align so that the second sequence is a prefix of the first + | IsPrefix -- ^ align so that the first sequence is a prefix of the second + deriving Enum + +-- | Align two strings. @myersAlign maxd seqA mode seqB@ tries to align +-- @seqA@ to @seqB@, which will work as long as no more than @maxd@ gaps +-- or mismatches are incurred. The @mode@ argument determines if either +-- of the sequences is allowed to have an overhanging tail. +-- +-- The result is the triple of the actual distance (gaps + mismatches) +-- and the two padded sequences. These sequences are the original +-- sequences with dashes inserted for gaps. +-- +-- The algorithm is the O(nd) algorithm by Myers, implemented in C. A +-- gap and a mismatch score the same. The strings are supposed to code +-- for DNA, the code understands IUPAC-IUB ambiguity codes. Two +-- characters match iff there is at least one nucleotide both can code +-- for. Note that N is a wildcard, while X matches nothing. + +myersAlign :: Int -> Bytes -> Mode -> Bytes -> (Int, Bytes, Bytes) +myersAlign maxd seqA mode seqB = + unsafePerformIO $ + S.unsafeUseAsCStringLen seqA $ \(seq_a, len_a) -> + S.unsafeUseAsCStringLen seqB $ \(seq_b, len_b) -> + + -- size of output buffers derives from this: + -- char *out_a = bt_a + len_a + maxd +2 ; + -- char *out_b = bt_b + len_b + maxd +2 ; + allocaBytes (len_a + maxd + 2) $ \bt_a -> + allocaBytes (len_b + maxd + 2) $ \bt_b -> + + myers_diff seq_a (fromIntegral len_a) + (fromIntegral $ fromEnum mode) + seq_b (fromIntegral len_b) + (fromIntegral maxd) bt_a bt_b >>= \dist -> + if dist < 0 + then return (maxBound, S.empty, S.empty) + else (,,) (fromIntegral dist) <$> + S.packCString bt_a <*> + S.packCString bt_b + + +-- | Nicely print an alignment. An alignment is simply a list of +-- strings with inserted gaps to make them align. We split them into +-- manageable chunks, stack them vertically and add a line showing +-- asterisks in every column where all aligned strings agree. The +-- result is /almost/ the Clustal format. +showAligned :: Int -> [Bytes] -> [L.ByteString] +showAligned w ss | all S.null ss = [] + | otherwise = map (L.fromChunks . (:[])) lefts ++ + L.pack agreement : + L.empty : + showAligned w rights + where + (lefts, rights) = unzip $ map (S.splitAt w) ss + agreement = map star $ S.transpose lefts + star str = if S.null str || S.all (== S.head str) str then '*' else ' ' + diff --git a/Bio/Bam.hs b/Bio/Bam.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam.hs @@ -0,0 +1,24 @@ +-- | Umbrella module for most of what's under 'Bio.Bam'. + +module Bio.Bam ( + module Bio.Bam.Fastq, + module Bio.Bam.Filter, + module Bio.Bam.Header, + module Bio.Bam.Index, + module Bio.Bam.Reader, + module Bio.Bam.Rec, + module Bio.Bam.Trim, + module Bio.Bam.Writer, + module Bio.Streaming + ) where + +import Bio.Bam.Fastq +import Bio.Bam.Filter +import Bio.Bam.Header +import Bio.Bam.Index +import Bio.Bam.Reader +import Bio.Bam.Rec +import Bio.Bam.Trim +import Bio.Bam.Writer +import Bio.Streaming + diff --git a/Bio/Bam/Evan.hs b/Bio/Bam/Evan.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Evan.hs @@ -0,0 +1,100 @@ +-- | This module contains stuff relating to conventions local to MPI +-- EVAN. The code is needed regularly, but it can be harmful when +-- applied to BAM files that follow different conventions. Most +-- importantly, no program should call these functions by default. + +module Bio.Bam.Evan + ( fixupFlagAbuse + , fixupBwaFlags + , removeWarts + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude + +import qualified Data.ByteString.Char8 as S + +-- | Fixes abuse of flags valued 0x800 and 0x1000. We used them for +-- low quality and low complexity, but they have since been redefined. +-- If set, we clear them and store them into the ZQ field. Also fixes +-- abuse of the combination of the paired, 1st mate and 2nd mate flags +-- used to indicate merging or trimming. These are canonicalized and +-- stored into the FF field. This function is unsafe on BAM files of +-- unclear origin! +fixupFlagAbuse :: BamRec -> BamRec +fixupFlagAbuse b = + (if b_flag b .&. flag_low_quality /= 0 then setQualFlag 'Q' else id) $ -- low qual, new convention + (if b_flag b .&. flag_low_complexity /= 0 then setQualFlag 'C' else id) $ -- low complexity, new convention + b { b_flag = cleaned_flags, b_exts = cleaned_exts } + where + -- removes old flag abuse + flags' = b_flag b .&. complement (flag_low_quality .|. flag_low_complexity) + cleaned_flags | flags' .&. flagPaired == 0 = flags' .&. complement (flagFirstMate .|. flagSecondMate) + | otherwise = flags' + + flag_low_quality = 0x800 + flag_low_complexity = 0x1000 + + -- merged & trimmed from old flag abuse + is_merged = flags' .&. (flagPaired .|. flagFirstMate .|. flagSecondMate) == flagFirstMate .|. flagSecondMate + is_trimmed = flags' .&. (flagPaired .|. flagFirstMate .|. flagSecondMate) == flagSecondMate + newflags = (if is_merged then eflagMerged else 0) .|. (if is_trimmed then eflagTrimmed else 0) + + -- Extended flags, renamed to avoid collision with BWA. Goes like this: if FF is there, use + -- it. Else check if XF is there __and is numeric__. If so, use it, remove it, and set FF + -- instead. Else use 0 and leave it alone. Note that this resolves the collision with BWA, + -- since BWA puts a character there, not an int. + cleaned_exts = case (lookup "FF" (b_exts b), lookup "XF" (b_exts b)) of + ( Just (Int i), _ ) -> updateE "FF" (Int (i .|. newflags)) (b_exts b) + ( _, Just (Int i) ) -> updateE "FF" (Int (i .|. newflags)) $ deleteE "XF" (b_exts b) + _ | newflags /= 0 -> updateE "FF" (Int newflags ) (b_exts b) + | otherwise -> b_exts b + + +-- | Fixes typical inconsistencies produced by Bwa: sometimes, 'mate unmapped' should be set, and we +-- can see it, because we match the mate's coordinates. Sometimes 'properly paired' should not be +-- set, because one mate is unmapped. This function is generally safe, but needs to be called only +-- on the output of affected (older?) versions of Bwa. +fixupBwaFlags :: BamRec -> BamRec +fixupBwaFlags b = b { b_flag = fixPP $ b_flag b .|. if mu then flagMateUnmapped else 0 } + where + -- Set "mate unmapped" if self coordinates and mate coordinates are equal, but self is + -- paired and mapped. (BWA forgets this flag for invalid mate alignments) + mu = and [ isPaired b, not (isUnmapped b) + , isReversed b == isMateReversed b + , b_rname b == b_mrnm b, b_pos b == b_mpos b ] + + -- If either mate is unmapped, remove "properly paired". + fixPP f | f .&. (flagUnmapped .|. flagMateUnmapped) == 0 = f + | otherwise = f .&. complement flagProperlyPaired + + +-- | Removes syntactic warts from old read names or the read names used +-- in FastQ files. +removeWarts :: BamRec -> BamRec +removeWarts br = br { b_qname = name, b_flag = flags, b_exts = tags } + where + (name, flags, tags) = checkFR $ checkC $ checkSharp (b_qname br, b_flag br, b_exts br) + + checkFR (n,f,t) | "F_" `S.isPrefixOf` n = checkC (S.drop 2 n, f .|. flagFirstMate .|. flagPaired, t) + | "R_" `S.isPrefixOf` n = checkC (S.drop 2 n, f .|. flagSecondMate .|. flagPaired, t) + | "M_" `S.isPrefixOf` n = checkC (S.drop 2 n, f, insertE "FF" (Int eflagMerged) t) + | "T_" `S.isPrefixOf` n = checkC (S.drop 2 n, f, insertE "FF" (Int eflagTrimmed) t) + | "/1" `S.isSuffixOf` n = ( rdrop 2 n, f .|. flagFirstMate .|. flagPaired, t) + | "/2" `S.isSuffixOf` n = ( rdrop 2 n, f .|. flagSecondMate .|. flagPaired, t) + | otherwise = ( n, f, t) + + checkC (n,f,t) | "C_" `S.isPrefixOf` n = (S.drop 2 n, f, insertE "XP" (Int (-1)) t) + | otherwise = ( n, f, t) + + rdrop n s = S.take (S.length s - n) s + + checkSharp (n,f,t) = case S.split '#' n of [n',ts] -> (n', f, insertTags ts t) + _ -> ( n, f, t) + + insertTags ts t | S.null y = insertE "XI" (Text ts) t + | otherwise = insertE "XI" (Text x) $ insertE "XJ" (Text $ S.tail y) t + where (x,y) = S.break (== ',') ts + + diff --git a/Bio/Bam/Fastq.hs b/Bio/Bam/Fastq.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Fastq.hs @@ -0,0 +1,121 @@ +-- | Parser for @FastA/FastQ@, 'ByteStream' style, based on +-- "Data.Attoparsec", and written such that it is compatible with module +-- "Bio.Bam". This gives import of @FastA/FastQ@ while respecting some +-- local (to MPI EVAN) conventions. + +module Bio.Bam.Fastq ( parseFastq, parseFastqWith, parseFastqCassava ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude hiding ( isSpace ) +import Bio.Streaming +import Bio.Streaming.Parse ( parse, atto ) +import Data.Attoparsec.ByteString.Char8 + ( char, skipSpace, satisfy, inClass, skipWhile, takeTill + , scan, isSpace, isSpace_w8, () ) + +import qualified Data.Attoparsec.ByteString.Char8 as P +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as C +import qualified Data.Vector.Generic as V +import qualified Streaming.Prelude as Q +import qualified Bio.Streaming.Bytes as S + +-- | Reader for DNA (not protein) sequences in FastA and FastQ. We read +-- everything vaguely looking like FastA or FastQ, then shoehorn it into +-- a BAM record. We strive to extract information following more or +-- less established conventions from the header, but don't aim for +-- completeness. The recognized syntactical warts are converted into +-- appropriate flags and removed. Only the canonical variant of FastQ +-- is supported (qualities stored as raw bytes with offset 33). +-- +-- Supported additional conventions: +-- +-- * A name suffix of @/1@ or @/2@ is turned into the first mate or second +-- mate flag and the read is flagged as paired. +-- +-- * Same for name prefixes of @F_@ or @R_@, respectively. +-- +-- * A name prefix of @M_@ flags the sequence as unpaired and merged +-- +-- * A name prefix of @T_@ flags the sequence as unpaired and trimmed +-- +-- * A name prefix of @C_@, optionally before or after any of the other +-- prefixes, is turned into the extra flag @XP:i:-1@ (result of +-- duplicate removal with unknown duplicate count). +-- +-- * A collection of tags separated from the name by an octothorpe is +-- removed and put into the fields @XI@ and @XJ@ as text. +-- +-- Everything before the first sequence header is ignored. Headers can +-- start with @\>@ or @\@@, we treat both equally. The first word of +-- the header becomes the read name, the remainder of the header is +-- ignored. The sequence can be split across multiple lines; +-- whitespace, dashes and dots are ignored, IUPAC-IUB ambiguity codes +-- are accepted as bases, anything else causes an error. The sequence +-- ends at a line that is either a header or starts with @\+@, in the +-- latter case, that line is ignored and must be followed by quality +-- scores. There must be exactly as many Q-scores as there are bases, +-- followed immediately by a header or end-of-file. Whitespace is +-- ignored. + +parseFastq :: Monad m => ByteStream m r -> Stream (Of BamRec) m (Either SomeException r) +parseFastq = parseFastqWith (const id) + +-- | Like 'parseFastq', but also +-- +-- * If the first word of the description has at least four colon +-- separated subfields, the first is used to flag first/second mate, +-- the second is the \"QC failed\" flag, and the fourth is the index +-- sequence. + +parseFastqCassava :: Monad m => ByteStream m r -> Stream (Of BamRec) m (Either SomeException r) +parseFastqCassava = parseFastqWith (pdesc . C.split ':' . C.takeWhile (' ' /=)) + where + pdesc (num:flg:_:idx:_) br = br { b_flag = sum [ if num == "1" then flagFirstMate .|. flagPaired else 0 + , if num == "2" then flagSecondMate .|. flagPaired else 0 + , if flg == "Y" then flagFailsQC else 0 + , b_flag br .&. complement (flagFailsQC .|. flagSecondMate .|. flagPaired) ] + , b_exts = if C.all (`C.elem` "ACGTN") idx then insertE "XI" (Text idx) (b_exts br) else b_exts br } + pdesc _ br = br + +-- | Same as 'parseFastq', but a custom function can be applied to the +-- description string (the part of the header after the sequence name), +-- which can modify the parsed record. Note that the quality field can +-- end up empty. + +parseFastqWith :: Monad m => (Bytes -> BamRec -> BamRec) -> ByteStream m r -> Stream (Of BamRec) m (Either SomeException r) +parseFastqWith descr = Q.unfoldr (liftM twiddle . parse (const $ atto pRec)) . skipJunk + where + twiddle (Left e) = Left (Left e) + twiddle (Right (Left r)) = Left (Right r) + twiddle (Right (Right a)) = Right a + + isCBase = inClass "ACGTUBDHVSWMKRYNacgtubdhvswmkryn" + canSkip c = isSpace c || c == '.' || c == '-' + isHdr c = c == '@' || c == '>' + + pRec = (satisfy isHdr "start marker") *> (makeRecord <$> pName <*> (descr <$> P.takeWhile ('\n' /=)) <*> (pSeq >>= pQual)) + pName = takeTill isSpace <* skipWhile (\c -> c /= '\n' && isSpace c) "read name" + pSeq = (:) <$> satisfy isCBase <*> pSeq + <|> satisfy canSkip *> pSeq + <|> pure [] "sequence" + + pQual sq = (,) sq <$> (char '+' *> skipWhile ('\n' /=) *> pQual' (length sq) <* skipSpace <|> return C.empty) "qualities" + pQual' n = B.filter (not . isSpace_w8) <$> scan n step + step 0 _ = Nothing + step i c | isSpace c = Just i + | otherwise = Just (i-1) + +skipJunk :: Monad m => ByteStream m r -> ByteStream m r +skipJunk = lift . S.nextByte >=> check + where + check (Right (c,s)) | bad c = skipJunk . S.drop 1 . S.dropWhile (c2w '\n' /=) $ s + | otherwise = S.cons c s + check (Left r) = pure r + bad c = c /= c2w '>' && c /= c2w '@' + +makeRecord :: Bytes -> (BamRec->BamRec) -> (String, Bytes) -> BamRec +makeRecord name extra (sq,qual) = extra $ nullBamRec + { b_qname = name, b_seq = V.fromList $ read sq, b_qual = V.fromList $ map (Q . subtract 33) $ B.unpack qual } + diff --git a/Bio/Bam/Filter.hs b/Bio/Bam/Filter.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Filter.hs @@ -0,0 +1,121 @@ +-- | Quality filters adapted from prehistoric pipeline. + +module Bio.Bam.Filter ( + filterPairs, QualFilter, + complexSimple, complexEntropy, + qualityAverage, qualityMinimum, + qualityFromOldIllumina, qualityFromNewIllumina + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude +import Bio.Streaming + +import qualified Data.Vector.Generic as V + +-- | A filter/transformation applied to pairs of reads. We supply a +-- predicate to be applied to single reads and one to be applied to +-- pairs, the latter can get incomplete pairs, too, if mates have been +-- separated or filtered asymmetrically. This fails spectacularly if +-- the input isn't grouped by name. + +filterPairs :: Monad m => (BamRec -> [BamRec]) + -> (Maybe BamRec -> Maybe BamRec -> [BamRec]) + -> Stream (Of BamRec) m r -> Stream (Of BamRec) m r +filterPairs ps pp = step + where + step = lift . inspect >=> either pure step' + step' (b :> s) + | isPaired b = lift (inspect s) >>= step'' b + | otherwise = each (ps b) >> step s + + step'' b (Left r) = each (pp (Just b) Nothing) >> pure r + + step'' b (Right (c :> s)) + | b_rname b /= b_rname c || not (isPaired c) = + let b' = if isSecondMate b then pp Nothing (Just b) else pp (Just b) Nothing + in each b' >> step' (c :> s) + + | isFirstMate c && isSecondMate b = step''' c b s + | otherwise = step''' b c s + + step''' b c s = each (pp (Just b) (Just c)) >> step s + + +-- | A quality filter is simply a transformation on @BamRec@s. By +-- convention, quality filters should set @flagFailsQC@, a further step +-- can then remove the failed reads. Filtering of individual reads +-- tends to result in mate pairs with inconsistent flags, which in turn +-- will result in lone mates and all sort of troubles with programs that +-- expect non-broken BAM files. It is therefore recommended to use +-- @pairFilter@ with suitable predicates to do the post processing. + +type QualFilter = BamRec -> BamRec + +{-# INLINE count #-} +count :: (V.Vector v a, Eq a) => a -> v a -> Int +count x = V.foldl' (\acc y -> if x == y then acc+1 else acc) 0 + +-- | Simple complexity filter aka "Nancy Filter". A read is considered +-- not-sufficiently-complex if the most common base accounts for greater +-- than the @cutoff@ fraction of all non-N bases. +{-# INLINE complexSimple #-} +complexSimple :: Double -> QualFilter +complexSimple r b = if p then b else b' + where + b' = setQualFlag 'C' $ b { b_flag = b_flag b .|. flagFailsQC } + p = let counts = [ count x $ b_seq b | x <- properBases ] + lim = floor $ r * fromIntegral (sum counts) + in all (<= lim) counts + +-- | Filter on order zero empirical entropy. Entropy per base must be +-- greater than cutoff. +{-# INLINE complexEntropy #-} +complexEntropy :: Double -> QualFilter +complexEntropy r b = if p then b else b' + where + b' = setQualFlag 'C' $ b { b_flag = b_flag b .|. flagFailsQC } + p = ent >= r * total + + counts = [ count x $ b_seq b | x <- properBases ] + total = fromIntegral $ V.length $ b_seq b + ent = sum [ fromIntegral c * log (total / fromIntegral c) | c <- counts, c /= 0 ] / log 2 + +-- | Filter on average quality. Reads without quality string pass. +{-# INLINE qualityAverage #-} +qualityAverage :: Int -> QualFilter +qualityAverage q b = if p then b else b' + where + b' = setQualFlag 'Q' $ b { b_flag = b_flag b .|. flagFailsQC } + p = let total = V.foldl' (\a x -> a + fromIntegral (unQ x)) 0 $ b_qual b + in total >= q * V.length (b_qual b) + +-- | Filter on minimum quality. In @qualityMinimum n q@, a read passes +-- if it has no more than @n@ bases with quality less than @q@. Reads +-- without quality string pass. +{-# INLINE qualityMinimum #-} +qualityMinimum :: Int -> Qual -> QualFilter +qualityMinimum n (Q q) b = if p then b else b' + where + b' = setQualFlag 'Q' $ b { b_flag = b_flag b .|. flagFailsQC } + p = V.length (V.filter (< Q q) (b_qual b)) <= n + + +-- | Convert quality scores from old Illumina scale (different formula +-- and offset 64 in FastQ). +qualityFromOldIllumina :: BamRec -> BamRec +qualityFromOldIllumina b = b { b_qual = V.map conv $ b_qual b } + where + conv (Q s) = let s' :: Double + s' = exp $ log 10 * (fromIntegral s - 31) / (-10) + p = s' / (1+s') + q = - 10 * log p / log 10 + in Q (round q) + +-- | Convert quality scores from new Illumina scale (standard formula +-- but offset 64 in FastQ). +qualityFromNewIllumina :: BamRec -> BamRec +qualityFromNewIllumina b = b { b_qual = V.map (Q . subtract 31 . unQ) $ b_qual b } + + diff --git a/Bio/Bam/Header.hs b/Bio/Bam/Header.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Header.hs @@ -0,0 +1,570 @@ +{-# LANGUAGE UndecidableInstances #-} +module Bio.Bam.Header ( + BamMeta(..), + parseBamMeta, + showBamMeta, + addPG, + + BamKey(..), + BamHeader(..), + BamSQ(..), + BamSorting(..), + BamOtherShit, + + Refseq(..), + invalidRefseq, + isValidRefseq, + invalidPos, + isValidPos, + unknownMapq, + isKnownMapq, + + Refs(..), + getRef, + + compareNames, + + flagPaired, + flagProperlyPaired, + flagUnmapped, + flagMateUnmapped, + flagReversed, + flagMateReversed, + flagFirstMate, + flagSecondMate, + flagAuxillary, + flagSecondary, + flagFailsQC, + flagDuplicate, + flagSupplementary, + eflagTrimmed, + eflagMerged, + eflagAlternative, + eflagExactIndex, + + distinctBin, + + MdOp(..), + readMd, + showMd + ) where + +import Bio.Prelude hiding ( uncons ) +import Bio.Util.Nub +import Control.Monad.Trans.RWS +import Data.ByteString ( uncons ) +import Data.ByteString.Builder ( Builder, byteString, char7, intDec, word16LE ) + +import qualified Data.Attoparsec.ByteString.Char8 as P +import qualified Data.ByteString.Char8 as S +import qualified Data.HashMap.Strict as H +import qualified Data.Vector as V + +data BamMeta = BamMeta { + meta_hdr :: !BamHeader, + meta_refs :: !Refs, + meta_pgs :: [Fix BamPG], + meta_other_shit :: [(BamKey, BamOtherShit)], + meta_comment :: [Bytes] + } deriving ( Show, Generic ) + +-- | Exactly two characters, for the \"named\" fields in bam. +newtype BamKey = BamKey Word16 + deriving ( Eq, Ord, Hashable, Generic ) + +instance IsString BamKey where + {-# INLINE fromString #-} + fromString [a,b] + | ord a < 256 && ord b < 256 + = BamKey . fromIntegral $ ord a .|. shiftL (ord b) 8 + + fromString s + = error $ "Not a legal BAM key: " ++ show s + +instance Show BamKey where + show (BamKey a) = [ chr (fromIntegral a .&. 0xff), chr (shiftR (fromIntegral a) 8 .&. 0xff) ] + +-- | Adds a new program line to a header. The new entry is +-- (arbitrarily) prepended to the first existing chain, or forms a new +-- singleton chain if none exists. + +addPG :: Maybe Version -> IO (BamMeta -> BamMeta) +addPG vn = do + args <- getArgs + pn <- getProgName + + let more = ("PN", S.pack pn) : + ("CL", S.pack $ unwords args) : + maybe [] (\v -> [("VN",S.pack (showVersion v))]) vn + + return $ \bm -> case meta_pgs bm of + [ ] -> bm { meta_pgs = Fix (BamPG (S.pack pn) Nothing more) : [ ] } + pg:pgs -> bm { meta_pgs = Fix (BamPG (S.pack pn) (Just pg) more) : pgs } + + +instance Semigroup BamMeta where (<>) = combineBamMeta +instance Monoid BamMeta where mempty = BamMeta mempty mempty mempty [] [] + mappend = (<>) + +{- | Combines two bam headers into one. + +The overarching goal is to combine headers in such a way that no +information is lost, but redundant information is removed. In +particular, we sometimes \"merge\" headers with the same references, at +other times we \"meld\" headers with entirely different references. In +the former case, we must concatenate the reference lists, in the latter +case we want to keep it as is. + +* If both headers have a version number, the result is the smaller of + the two. + +* The resulting sort order is the most specific one compatible with both + input sort orders. The stupid 'Unknown' state is compatible with + everything. + +* Reference sequences are appended and run through 'nub'. The numbering + of reference may thus change, which has to be dealt with in an + appropriate way, see 'concatInputs', 'mergeInputsOn', and \"bam-meld\" + for details. (It is also possible that different sequences are left + with the same name. We cannot solve this right here, and there is no + reliable way to do it in general.) + +* Comments are appended and run through 'nub'. This should work in + most case, and if it doesn't, someone needs to \"samtools reheader\" + the file anyway. + +* Program chains are just collected, but when formatting, they are + (effectively) run through 'nub' and are potentially assigned new + unique identifiers. +-} +combineBamMeta :: BamMeta -> BamMeta -> BamMeta +combineBamMeta a b = BamMeta + { meta_hdr = meta_hdr a <> meta_hdr b + , meta_refs = meta_refs a `mappend` meta_refs b + , meta_pgs = meta_pgs a <> meta_pgs b + , meta_other_shit = nubHash $ meta_other_shit a ++ meta_other_shit b + , meta_comment = nubHash $ meta_comment a ++ meta_comment b } + +data BamHeader = BamHeader { + hdr_version :: (Int, Int), + hdr_sorting :: BamSorting, + hdr_other_shit :: BamOtherShit + } deriving (Show, Eq) + +instance Monoid BamHeader where + mempty = BamHeader (1,0) Unknown [] + mappend = (<>) + +instance Semigroup BamHeader where + a <> b = BamHeader { hdr_version = max (hdr_version a) (hdr_version b) + , hdr_sorting = hdr_sorting a <> hdr_sorting b + , hdr_other_shit = nubHashBy fst $ hdr_other_shit a ++ hdr_other_shit b } + +data BamSQ = BamSQ { + sq_name :: Bytes, + sq_length :: Int, + sq_other_shit :: BamOtherShit + } deriving (Show, Eq, Generic) + +instance Hashable BamSQ + +bad_seq :: BamSQ +bad_seq = BamSQ (error "no SN field") (error "no LN field") [] + +data BamPG pp = BamPG { + pg_pref_name :: Bytes, + pg_prev_pg :: Maybe pp, + pg_other_shit :: BamOtherShit + } deriving (Show, Eq, Generic1) + +newtype Fix f = Fix (f (Fix f)) + +instance Eq (f (Fix f)) => Eq (Fix f) where + Fix f == Fix g = f == g + +instance Show (f (Fix f)) => Show (Fix f) where + showsPrec p (Fix f) = showsPrec p f + +instance Hashable1 f => Hashable (Fix f) where + hashWithSalt salt (Fix f) = liftHashWithSalt hashWithSalt salt f + +instance Hashable1 BamPG + +bad_pg :: BamPG pp +bad_pg = BamPG (error "no ID field") Nothing [] + + +-- | Possible sorting orders from bam header. Thanks to samtools, which +-- doesn't declare sorted files properly, we have to have the stupid +-- 'Unknown' state, too. +data BamSorting = Unknown -- ^ undeclared sort order + | Unsorted -- ^ definitely not sorted + | Grouped -- ^ grouped by query name + | Queryname -- ^ sorted by query name + | Coordinate -- ^ sorted by coordinate + deriving (Show, Eq) + +instance Semigroup BamSorting where + Unknown <> b = b + a <> Unknown = a + Grouped <> Grouped = Grouped + Grouped <> Queryname = Grouped + Queryname <> Grouped = Grouped + Queryname <> Queryname = Queryname + Coordinate <> Coordinate = Coordinate + _ <> _ = Unsorted + + +type BamOtherShit = [(BamKey, Bytes)] + +parseBamMeta :: P.Parser BamMeta +parseBamMeta = fixupMeta . foldl' (flip ($)) emptyHeader + <$> many (parseBamMetaLine <* P.skipWhile (=='\t') <* P.char '\n') + +-- Bam header in the process of being parsed. Better suited for +-- collecting lines than 'BamMeta'. +data PreBamMeta = PreBamMeta { + pmeta_hdr :: BamHeader, + pmeta_refs :: [BamSQ], + pmeta_pgs :: HashMap Bytes (BamPG Bytes), + pmeta_other_shit :: [(BamKey, BamOtherShit)], + pmeta_comment :: [Bytes] } + +emptyHeader :: PreBamMeta +emptyHeader = PreBamMeta mempty [] H.empty [] [] + + +-- | Fixes a bam header after parsing. It turns accumulated lists in to +-- vectors, and it handles the program lines. Program lines come in as +-- an arbitrary graph. It chould be a linear chain, but this +-- isn't guaranteed in practice. We decompose the graph into chains by +-- tracing from nodes with no predecessor, or from an arbitrary node if +-- all nodes have predecessors. Tracing stops once it would form a +-- cycle. +fixupMeta :: PreBamMeta -> BamMeta +fixupMeta PreBamMeta{..} = BamMeta + { meta_hdr = pmeta_hdr + , meta_refs = Refs . V.fromList . reverse $ pmeta_refs + , meta_pgs = snd $ evalRWS trace_pgs () pmeta_pgs + , meta_other_shit = reverse pmeta_other_shit + , meta_comment = reverse pmeta_comment } + where + -- keep tracing from roots until no nodes are left + trace_pgs :: RWS () [Fix BamPG] (HashMap Bytes (BamPG Bytes)) () + trace_pgs = do + gg <- get + case foldl' (flip H.delete) gg + [ pp | p <- H.elems gg + , pp <- maybe [] pure (pg_prev_pg p) ] of + orphans + -- the empty graph has no roots: + | H.null gg -> return () + -- an arbitrary node is picked as root: + | H.null orphans -> trace_pg H.empty (head $ H.keys gg) >> trace_pgs + -- nodes without parents are roots: + | otherwise -> mapM_ (trace_pg H.empty) (H.keys orphans) >> trace_pgs + + -- Trace one PG line. Do not trace into nodes in the 'closed' set, + -- remove reached nodes from the 'open' set (the state) and add them + -- to the 'closed' set. + trace_pg :: HashMap Bytes () -> Bytes -> RWS () [Fix BamPG] (HashMap Bytes x) (Maybe (Fix BamPG)) + trace_pg closed name = + case H.lookup name pmeta_pgs of + _ | H.member name closed -> return Nothing + Nothing -> return Nothing + Just pg -> do + modify $ H.delete name + pp <- mapM (trace_pg (H.insert name () closed)) (pg_prev_pg pg) + let self = Fix $ pg { pg_prev_pg = join pp } + tell [ self ] + return $ Just self + + +parseBamMetaLine :: P.Parser (PreBamMeta -> PreBamMeta) +parseBamMetaLine = P.char '@' >> P.choice [hdLine, sqLine, pgLine, coLine, otherLine] + where + hdLine = P.string "HD\t" >> + (\fns meta -> meta { pmeta_hdr = foldr ($) (pmeta_hdr meta) fns }) + <$> P.sepBy1 (P.choice [hdvn, hdso, hdother]) tabs + + sqLine = P.string "SQ\t" >> + (\fns meta -> meta { pmeta_refs = foldr ($) bad_seq fns : pmeta_refs meta }) + <$> P.sepBy1 (P.choice [sqnm, sqln, sqother]) tabs + + pgLine = P.string "PG\t" >> + (\fns meta -> let pg = foldr ($) bad_pg fns + in meta { pmeta_pgs = H.insert (pg_pref_name pg) pg (pmeta_pgs meta) }) + <$> P.sepBy1 (P.choice [pgid, pgpp, pgother]) tabs + + hdvn = P.string "VN:" >> + (\a b hdr -> hdr { hdr_version = (a,b) }) + <$> P.decimal <*> ((P.char '.' <|> P.char ':') >> P.decimal) + + hdso = P.string "SO:" >> + (\s hdr -> hdr { hdr_sorting = s }) + <$> P.choice [ Grouped <$ P.string "grouped" + , Queryname <$ P.string "queryname" + , Coordinate <$ P.string "coordinate" + , Unsorted <$ P.string "unsorted" + , Unknown <$ P.skipWhile (\c -> c/='\t' && c/='\n') ] + + sqnm = P.string "SN:" >> (\s sq -> sq { sq_name = s }) <$> pall + sqln = P.string "LN:" >> (\i sq -> sq { sq_length = i }) <$> P.decimal + + pgid = P.string "ID:" >> (\s pg -> pg { pg_pref_name = s }) <$> pall + pgpp = P.string "PP:" >> (\s pg -> pg { pg_prev_pg = Just s }) <$> pall + + hdother = (\t hdr -> hdr { hdr_other_shit = t : hdr_other_shit hdr }) <$> tagother + sqother = (\t sq -> sq { sq_other_shit = t : sq_other_shit sq }) <$> tagother + pgother = (\t p -> p { pg_other_shit = t : pg_other_shit p }) <$> tagother + + coLine = P.string "CO\t" >> + (\s meta -> s `seq` meta { pmeta_comment = s : pmeta_comment meta }) + <$> P.takeWhile (/= 'n') + + otherLine = (\k ts meta -> meta { pmeta_other_shit = (k,ts) : pmeta_other_shit meta }) + <$> bamkey <*> (tabs >> P.sepBy1 tagother tabs) + + tagother :: P.Parser (BamKey,Bytes) + tagother = (,) <$> bamkey <*> (P.char ':' >> pall) + + tabs = P.char '\t' >> P.skipWhile (== '\t') + + pall :: P.Parser Bytes + pall = P.takeWhile (\c -> c/='\t' && c/='\n') + + bamkey :: P.Parser BamKey + bamkey = (\a b -> fromString [a,b]) <$> P.anyChar <*> P.anyChar + + +-- | Creates the textual form of Bam meta data. +-- +-- Formatting is straight forward, only program lines are a bit +-- involved. Our multiple chains may lead to common nodes, and we do +-- not want to print multiple identical lines. At the same time, we may +-- need to print multiple different lines that carry the same id. The +-- solution is to memoize printed lines, and to reuse their identity if +-- an identical line is needed. When printing a line, it gets its +-- preferred identifier, but if it's already taken, a new identifier is +-- made up by first removing any trailing number and then by appending +-- numeric suffixes. + +showBamMeta :: BamMeta -> Builder +showBamMeta (BamMeta h (Refs ss) pgs os cs) = + show_bam_meta_hdr h <> + foldMap show_bam_meta_seq ss <> + show_bam_pgs <> + foldMap show_bam_meta_other os <> + foldMap show_bam_meta_comment cs + where + show_bam_meta_hdr (BamHeader (major,minor) so os') = + "@HD\tVN:" <> + intDec major <> char7 '.' <> intDec minor <> + byteString (case so of Unsorted -> "\tSO:unsorted" + Grouped -> "\tSO:grouped" + Queryname -> "\tSO:queryname" + Coordinate -> "\tSO:coordinate" + Unknown -> mempty) <> + show_bam_others os' + + show_bam_meta_seq (BamSQ nm ln ts) = + byteString "@SQ\tSN:" <> byteString nm <> + byteString "\tLN:" <> intDec ln <> show_bam_others ts + + show_bam_meta_comment cm = byteString "@CO\t" <> byteString cm <> char7 '\n' + + show_bam_meta_other (BamKey k,ts) = + char7 '@' <> word16LE k <> show_bam_others ts + + show_bam_others ts = + foldMap show_bam_other ts <> char7 '\n' + + show_bam_other (BamKey k,v) = + char7 '\t' <> word16LE k <> char7 ':' <> byteString v + + show_bam_pgs = snd $ evalRWS (mapM_ show_bam_pg pgs) () (H.empty, H.empty) + + show_bam_pg p@(Fix (BamPG pn pp po)) = do + ppid <- case pp of Nothing -> return Nothing + Just p' -> Just <$> show_bam_pg p' + + gets (H.lookup p . fst) >>= \case + Just pid -> return pid + Nothing -> do + -- preferred name without a trailing dash-and-number + let pn' = case dropWhile isDigit . reverse $ S.unpack pn of + '-':xs -> reverse xs + _ -> S.unpack pn + + -- find unused preferable PG:ID: try prefered name, + -- preferred name without number, preferred name + -- without number and increasing numbers attached + pid <- gets $ \(_,hs) -> + head . filter (not . flip H.member hs) $ + pn : S.pack pn' : [ S.pack $ pn' ++ '-' : (show i) | i <- [2::Int ..] ] + + modify . first $ H.insert p pid + modify . second $ H.insert pid () + + tell $ byteString "@PG\tID:" <> byteString pid <> + maybe mempty (\x -> byteString "\tPP:" <> byteString x) ppid <> + show_bam_others po + return pid + + +-- | Reference sequence in Bam +-- Bam enumerates the reference sequences and then sorts by index. We +-- need to track that index if we want to reproduce the sorting order. +newtype Refseq = Refseq { unRefseq :: Word32 } deriving (Eq, Ord, Ix, Bounded) + +instance Show Refseq where + showsPrec p (Refseq r) = showsPrec p r + +instance Enum Refseq where + succ = Refseq . succ . unRefseq + pred = Refseq . pred . unRefseq + toEnum = Refseq . fromIntegral + fromEnum = fromIntegral . unRefseq + enumFrom = map Refseq . enumFrom . unRefseq + enumFromThen (Refseq a) (Refseq b) = map Refseq $ enumFromThen a b + enumFromTo (Refseq a) (Refseq b) = map Refseq $ enumFromTo a b + enumFromThenTo (Refseq a) (Refseq b) (Refseq c) = map Refseq $ enumFromThenTo a b c + + +-- | Tests whether a reference sequence is valid. +-- Returns true unless the the argument equals @invalidRefseq@. +isValidRefseq :: Refseq -> Bool +isValidRefseq = (/=) invalidRefseq + +-- | The invalid Refseq. +-- Bam uses this value to encode a missing reference sequence. +invalidRefseq :: Refseq +invalidRefseq = Refseq 0xffffffff + +-- | The invalid position. +-- Bam uses this value to encode a missing position. +{-# INLINE invalidPos #-} +invalidPos :: Int +invalidPos = -1 + +-- | Tests whether a position is valid. +-- Returns true unless the the argument equals @invalidPos@. +{-# INLINE isValidPos #-} +isValidPos :: Int -> Bool +isValidPos = (/=) invalidPos + +{-# INLINE unknownMapq #-} +unknownMapq :: Int +unknownMapq = 255 + +isKnownMapq :: Int -> Bool +isKnownMapq = (/=) unknownMapq + +-- | A list of reference sequences. +newtype Refs = Refs { unRefs :: V.Vector BamSQ } deriving Show + +instance Monoid Refs where + mempty = Refs V.empty + mappend = (<>) + +instance Semigroup Refs where + Refs a <> Refs b = Refs . V.fromList . nubHash $ V.toList a ++ V.toList b + +getRef :: Refs -> Refseq -> BamSQ +getRef (Refs refs) (Refseq i) = fromMaybe (BamSQ "*" 0 []) $ refs V.!? fromIntegral i + +flagPaired, flagProperlyPaired, flagUnmapped, flagMateUnmapped, + flagReversed, flagMateReversed, flagFirstMate, flagSecondMate, + flagAuxillary, flagSecondary, flagFailsQC, flagDuplicate, + flagSupplementary :: Int + +flagPaired = 0x1 +flagProperlyPaired = 0x2 +flagUnmapped = 0x4 +flagMateUnmapped = 0x8 +flagReversed = 0x10 +flagMateReversed = 0x20 +flagFirstMate = 0x40 +flagSecondMate = 0x80 +flagAuxillary = 0x100 +flagSecondary = 0x100 +flagFailsQC = 0x200 +flagDuplicate = 0x400 +flagSupplementary = 0x800 + +eflagTrimmed, eflagMerged, eflagAlternative, eflagExactIndex :: Int +eflagTrimmed = 0x1 +eflagMerged = 0x2 +eflagAlternative = 0x4 +eflagExactIndex = 0x8 + + +-- | Compares two sequence names the way samtools does. +-- samtools sorts by \"strnum_cmp\": +-- +-- * if both strings start with a digit, parse the initial +-- sequence of digits and compare numerically, if equal, +-- continue behind the numbers +-- * else compare the first characters (possibly NUL), if equal +-- continue behind them +-- * else both strings ended and the shorter one counts as +-- smaller (and that part is stupid) + +compareNames :: Bytes -> Bytes -> Ordering +compareNames n m = case (uncons n, uncons m) of + ( Nothing, Nothing ) -> EQ + ( Just _, Nothing ) -> GT + ( Nothing, Just _ ) -> LT + ( Just (c,n'), Just (d,m') ) + | is_digit c || is_digit d + , Just (u,n'') <- S.readInt n + , Just (v,m'') <- S.readInt m -> + case u `compare` v of + LT -> LT + GT -> GT + EQ -> n'' `compareNames` m'' + | otherwise -> + case c `compare` d of + LT -> LT + GT -> GT + EQ -> n' `compareNames` m' + where + is_digit c = c2w '0' <= c && c <= c2w '9' + + +data MdOp = MdNum Int | MdRep Nucleotides | MdDel [Nucleotides] deriving Show + +readMd :: Bytes -> Maybe [MdOp] +readMd s | S.null s = return [] + | isDigit (S.head s) = do (n,t) <- S.readInt s + (MdNum n :) <$> readMd t + | S.head s == '^' = let (a,b) = S.break isDigit (S.tail s) + in (MdDel (map toNucleotides $ S.unpack a) :) <$> readMd b + | otherwise = (MdRep (toNucleotides $ S.head s) :) <$> readMd (S.tail s) + +-- | Normalizes a series of 'MdOp's and encodes them in the way BAM and +-- SAM expect it. +showMd :: [MdOp] -> Bytes +showMd = S.pack . flip s1 [] + where + s1 (MdNum i : MdNum j : ms) = s1 (MdNum (i+j) : ms) + s1 (MdNum 0 : ms) = s1 ms + s1 (MdNum i : ms) = shows i . s1 ms + + s1 (MdRep r : ms) = shows r . s1 ms + + s1 (MdDel d1 : MdDel d2 : ms) = s1 (MdDel (d1++d2) : ms) + s1 (MdDel [] : ms) = s1 ms + s1 (MdDel ns : MdRep r : ms) = (:) '^' . shows ns . (:) '0' . shows r . s1 ms + s1 (MdDel ns : ms) = (:) '^' . shows ns . s1 ms + s1 [ ] = id + + +-- | Computes the "distinct bin" according to the BAM binning scheme. If +-- an alignment starts at @pos@ and its CIGAR implies a length of @len@ +-- on the reference, then it goes into bin @distinctBin pos len@. +distinctBin :: Int -> Int -> Int +distinctBin beg len = mkbin 14 $ mkbin 17 $ mkbin 20 $ mkbin 23 $ mkbin 26 0 + where end = beg + len - 1 + mkbin n x = if beg `shiftR` n /= end `shiftR` n then x + else ((1 `shiftL` (29-n))-1) `div` 7 + (beg `shiftR` n) diff --git a/Bio/Bam/Index.hs b/Bio/Bam/Index.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Index.hs @@ -0,0 +1,384 @@ +module Bio.Bam.Index ( + BamIndex(..), + withIndexedBam, + readBamIndex, + readBaiIndex, + readTabix, + + Region(..), + Subsequence(..), + streamBamRefseq, + streamBamRegions, + streamBamSubseq, + streamBamUnaligned +) where + +import Bio.Bam.Header +import Bio.Bam.Reader +import Bio.Bam.Rec +import Bio.Bam.Regions ( Region(..), Subsequence(..) ) +import Bio.Prelude +import Bio.Streaming +import Bio.Streaming.Bgzf ( bgunzip ) +import System.Directory ( doesFileExist ) + +import qualified Bio.Bam.Regions as R +import qualified Bio.Streaming.Bytes as S +import qualified Bio.Streaming.Parse as P +import qualified Data.IntMap.Strict as M +import qualified Data.ByteString as B +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as W +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as N +import qualified Data.Vector.Algorithms.Intro as N +import qualified Streaming.Prelude as Q + +-- | Full index, unifying BAI and CSI style. In both cases, we have the +-- binning scheme, parameters are fixed in BAI, but variable in CSI. +-- Checkpoints are created from the linear index in BAI or from the +-- `loffset' field in CSI. + +data BamIndex a = BamIndex { + -- | Minshift parameter from CSI + minshift :: {-# UNPACK #-} !Int, + + -- | Depth parameter from CSI + depth :: {-# UNPACK #-} !Int, + + -- | Best guess at where the unaligned records start + unaln_off :: {-# UNPACK #-} !Int64, + + -- | Room for stuff (needed for tabix) + extensions :: a, + + -- | Records for the binning index, where each bin has a list of + -- segments belonging to it. + refseq_bins :: {-# UNPACK #-} !(V.Vector Bins), + + -- | Known checkpoints of the form (pos,off) where off is the + -- virtual offset of the first record crossing pos. + refseq_ckpoints :: {-# UNPACK #-} !(V.Vector Ckpoints) } + + deriving Show + +-- | Mapping from bin number to vector of clusters. +type Bins = IntMap Segments +type Segments = U.Vector (Int64,Int64) + + +-- | Checkpoints. Each checkpoint is a position with the virtual offset +-- where the first alignment crossing the position is found. In BAI, we +-- get this from the 'ioffset' vector, in CSI we get it from the +-- 'loffset' field: "Given a region [beg,end), we only need to visit +-- chunks whose end file offset is larger than 'ioffset' of the 16kB +-- window containing 'beg'." (Sounds like a marginal gain, though.) + +type Ckpoints = IntMap Int64 + + +-- | Decode only those reads that fall into one of several regions. +-- Strategy: We will scan the file mostly linearly, but only those +-- regions that are actually needed. We filter the decoded stuff so +-- that it actually overlaps our regions. +-- +-- From the binning index, we get a list of segments per requested +-- region. Using the checkpoints, we prune them: if we have a +-- checkpoint to the left of the beginning of the interesting region, we +-- can move the start of each segment forward to the checkpoint. If +-- that makes the segment empty, it can be droppped. +-- +-- The resulting segment lists are merged, then traversed. We seek to +-- the beginning of the earliest segment and start decoding. Once the +-- virtual file position leaves the segment or the alignment position +-- moves past the end of the requested region, we move to the next. +-- Moving is a seek if it spans a sufficiently large gap or points +-- backwards, else we just keep going. + +-- | A 'Segment' has a start and an end offset, and an "end coordinate" +-- from the originating region. +data Segment = Segment {-# UNPACK #-} !Int64 {-# UNPACK #-} !Int64 {-# UNPACK #-} !Int deriving Show + +segmentLists :: BamIndex a -> Refseq -> R.Subsequence -> [[Segment]] +segmentLists bi@BamIndex{..} (Refseq ref) (R.Subsequence imap) + | Just bins <- refseq_bins V.!? fromIntegral ref, + Just cpts <- refseq_ckpoints V.!? fromIntegral ref + = [ rgnToSegments bi beg end bins cpts | (beg,end) <- M.toList imap ] +segmentLists _ _ _ = [] + +-- from region to list of bins, then to list of segments +rgnToSegments :: BamIndex a -> Int -> Int -> Bins -> Ckpoints -> [Segment] +rgnToSegments bi@BamIndex{..} beg end bins cpts = + [ Segment boff' eoff end + | bin <- binList bi beg end + , (boff,eoff) <- maybe [] U.toList $ M.lookup bin bins + , let boff' = max boff cpt + , boff' < eoff ] + where + !cpt = maybe 0 snd $ M.lookupLE beg cpts + +-- list of bins for given range of coordinates, from Heng's horrible code +binList :: BamIndex a -> Int -> Int -> [Int] +binList BamIndex{..} beg end = binlist' 0 (minshift + 3*depth) 0 + where + binlist' l s t = if l > depth then [] else [b..e] ++ go + where + b = t + beg `shiftR` s + e = t + (end-1) `shiftR` s + go = binlist' (l+1) (s-3) (t + 1 `shiftL` (3*l)) + + +-- | Merges two lists of segments. Lists must be sorted, the merge sort +-- merges overlapping segments into one. +infix 4 ~~ +(~~) :: [Segment] -> [Segment] -> [Segment] +Segment a b e : xs ~~ Segment u v f : ys + | b < u = Segment a b e : (xs ~~ Segment u v f : ys) -- no overlap + | a < u && b < v = Segment a v (max e f) : (xs ~~ ys) -- some overlap + | b < v = Segment u v (max e f) : (xs ~~ ys) -- contained + | v < a = Segment u v f : (xs ~~ Segment a b e : ys) -- no overlap + | u < a = Segment u b (max e f) : (xs ~~ ys) -- some overlap + | otherwise = Segment a b (max e f) : (xs ~~ ys) -- contained +[] ~~ ys = ys +xs ~~ [] = xs + + +data IndexFormatError = IndexFormatError Bytes deriving Typeable + +instance Exception IndexFormatError +instance Show IndexFormatError where + show (IndexFormatError m) = "index signature " ++ show m ++ " not recognized" + +{- | Reads any index we can find for a file. + +If the file name has a .bai or .csi extension, optionally followed by +.gz, we read it. Else we look for the index by adding such an extension +and by replacing the extension with these two, and finally try the file +itself. The first file that exists is used. +-} +readBamIndex :: FilePath -> IO (BamIndex ()) +readBamIndex fp1 | any (`isSuffixOf` fp1) exts = streamFile fp1 readBaiIndex + | otherwise = tryAll exts + where + exts = words ".bai .bai.gz .csi .csi.gz" + + fp2 = reverse $ case dropWhile (/='.') f of [] -> f++d ; _:b -> b++d + (f,d) = break (=='/') $ reverse fp1 + + tryAll [ ] = streamFile fp1 readBaiIndex + tryAll (e:es) = do x1 <- liftIO $ doesFileExist (fp1 ++ e) + x2 <- liftIO $ doesFileExist (fp2 ++ e) + case () of + _ | x1 -> streamFile (fp1 ++ e) readBaiIndex + | x2 -> streamFile (fp2 ++ e) readBaiIndex + _ -> tryAll es + +-- | Reads an index in BAI or CSI format, recognized automatically. The +-- index can be compressed, even though this isn't standard. +readBaiIndex :: MonadIO m => ByteStream m r -> m (BamIndex ()) +readBaiIndex = either (const . liftIO $ throwM P.EofException) (return . fst) <=< + P.parseIO (const $ P.getString 4 >>= switch) . S.gunzip + where + switch "BAI\1" = do nref <- fromIntegral `liftM` P.getWord32 + getIndexArrays nref 14 5 (const return) getIntervals + + switch "CSI\1" = do minshift <- fromIntegral `liftM` P.getWord32 + depth <- fromIntegral `liftM` P.getWord32 + P.getWord32 >>= P.drop . fromIntegral -- aux data + nref <- fromIntegral `liftM` P.getWord32 + getIndexArrays nref minshift depth (addOneCheckpoint minshift depth) return + + switch magic = throwM $ IndexFormatError magic + + + -- Insert one checkpoint. If we already have an entry (can happen + -- if it comes from a different bin), we conservatively take the min + addOneCheckpoint minshift depth bin cp = do + loffset <- fromIntegral `liftM` P.getWord64 + let key = llim (fromIntegral bin) (3*depth) minshift + return $! M.insertWith min key loffset cp + + -- compute left limit of bin + llim bin dp sf | dp == 0 = 0 + | bin >= ix = (bin - ix) `shiftL` sf + | otherwise = llim bin (dp-3) (sf+3) + where ix = (1 `shiftL` dp - 1) `div` 7 + +withIndexedBam :: (MonadIO m, MonadMask m) => FilePath -> (BamMeta -> BamIndex () -> Handle -> m r) -> m r +withIndexedBam f k = do + idx <- liftIO $ readBamIndex f + bracket (liftIO $ openBinaryFile f ReadMode) (liftIO . hClose) $ \hdl -> do + (hdr,_) <- decodeBam $ streamHandle hdl + k hdr idx hdl + + +type TabIndex = BamIndex TabMeta + +data TabMeta = TabMeta { format :: TabFormat + , col_seq :: Int -- Column for the sequence name + , col_beg :: Int -- Column for the start of a region + , col_end :: Int -- Column for the end of a region + , comment_char :: Char + , skip_lines :: Int + , names :: V.Vector Bytes } + deriving Show + +data TabFormat = Generic | SamFormat | VcfFormat | ZeroBased deriving Show + +-- | Reads a Tabix index. Note that tabix indices are compressed, this +-- is taken care of automatically. +readTabix :: MonadIO m => ByteStream m r -> m TabIndex +readTabix = either (const . liftIO $ throwM P.EofException) (return . fst) <=< + P.parseIO (const $ P.getString 4 >>= switch) . S.gunzip + where + switch "TBI\1" = do nref <- fromIntegral `liftM` P.getWord32 + format <- liftM toFormat P.getWord32 + col_seq <- liftM fromIntegral P.getWord32 + col_beg <- liftM fromIntegral P.getWord32 + col_end <- liftM fromIntegral P.getWord32 + comment_char <- liftM (chr . fromIntegral) P.getWord32 + skip_lines <- liftM fromIntegral P.getWord32 + names <- liftM (V.fromList . B.split 0) . P.getString . fromIntegral =<< P.getWord32 + + ix <- getIndexArrays nref 14 5 (const return) getIntervals + fin <- P.isFinished + if fin then return $! ix { extensions = TabMeta{..} } + else do unaln <- fromIntegral `liftM` P.getWord64 + return $! ix { unaln_off = unaln, extensions = TabMeta{..} } + + switch magic = throwM $ IndexFormatError magic + + toFormat 1 = SamFormat + toFormat 2 = VcfFormat + toFormat x = if testBit x 16 then ZeroBased else Generic + +-- Read the intervals. Each one becomes a checkpoint. +getIntervals :: Monad m => (IntMap Int64, Int64) -> P.Parser r m (IntMap Int64, Int64) +getIntervals (cp,mx0) = do + nintv <- fromIntegral `liftM` P.getWord32 + reduceM 0 nintv (cp,mx0) $ \(!im,!mx) int -> do + oo <- fromIntegral `liftM` P.getWord64 + return (if oo == 0 then im else M.insert (int * 0x4000) oo im, max mx oo) + + +getIndexArrays :: MonadIO m => Int -> Int -> Int + -> (Word32 -> Ckpoints -> P.Parser r m Ckpoints) + -> ((Ckpoints, Int64) -> P.Parser r m (Ckpoints, Int64)) + -> P.Parser r m (BamIndex ()) +getIndexArrays nref minshift depth addOneCheckpoint addManyCheckpoints + | nref < 1 = return $ BamIndex minshift depth 0 () V.empty V.empty + | otherwise = do + rbins <- liftIO $ W.new nref + rckpts <- liftIO $ W.new nref + mxR <- reduceM 0 nref 0 $ \mx0 r -> do + nbins <- P.getWord32 + (!bins,!cpts,!mx1) <- reduceM 0 nbins (M.empty,M.empty,mx0) $ \(!im,!cp,!mx) _ -> do + bin <- P.getWord32 -- the "distinct bin" + cp' <- addOneCheckpoint bin cp + segsarr <- getSegmentArray + let !mx' = if U.null segsarr then mx else max mx (snd (U.last segsarr)) + return (M.insert (fromIntegral bin) segsarr im, cp', mx') + (!cpts',!mx2) <- addManyCheckpoints (cpts,mx1) + liftIO $ W.write rbins r bins >> W.write rckpts r cpts' + return mx2 + liftM2 (BamIndex minshift depth mxR ()) (liftIO $ V.unsafeFreeze rbins) (liftIO $ V.unsafeFreeze rckpts) + +-- | Reads the list of segments from an index file and makes sure +-- it is sorted. +getSegmentArray :: MonadIO m => P.Parser r m Segments +getSegmentArray = do + nsegs <- fromIntegral `liftM` P.getWord32 + segsarr <- liftIO $ N.new nsegs + loopM 0 nsegs $ \i -> do beg <- fromIntegral `liftM` P.getWord64 + end <- fromIntegral `liftM` P.getWord64 + liftIO $ N.write segsarr i (beg,end) + liftIO $ N.sort segsarr >> U.unsafeFreeze segsarr + +{-# INLINE reduceM #-} +reduceM :: (Monad m, Enum ix, Eq ix) => ix -> ix -> a -> (a -> ix -> m a) -> m a +reduceM beg end acc cons = if beg /= end then cons acc beg >>= \n -> reduceM (succ beg) end n cons else return acc + +{-# INLINE loopM #-} +loopM :: (Monad m, Enum ix, Eq ix) => ix -> ix -> (ix -> m ()) -> m () +loopM beg end k = if beg /= end then k beg >> loopM (succ beg) end k else return () + +{-| Seeks to a virtual offset in a BGZF file and streams from there. + +If the optional end offset is supplied, streaming stops when it is +reached. Else, streaming goes on to the end of file. +-} +streamBgzf :: MonadIO m => Handle -> Int64 -> Maybe Int64 -> ByteStream m () +streamBgzf hdl off eoff = + S.drop (off .&. 0xffff) . bgunzip $ do + when (off /= 0) (liftIO $ hSeek hdl AbsoluteSeek $ fromIntegral $ shiftR off 16) + maybe id S.trim eoff $ streamHandle hdl +{-# INLINE streamBgzf #-} + +{- | Streams one reference from a bam file. + +Seeks to a given sequence in a Bam file and enumerates only those +records aligning to that reference. We use the first checkpoint +available for the sequence, which an appropriate index. Streams the +'BamRaw' records of the correct reference sequence only, and produces an +empty stream if the sequence isn't found. +-} +streamBamRefseq :: MonadIO m => BamIndex b -> Handle -> Refseq -> Stream (Of BamRaw) m () +streamBamRefseq BamIndex{..} hdl (Refseq r) + | Just ckpts <- refseq_ckpoints V.!? fromIntegral r + , Just (voff, _) <- M.minView ckpts + , voff /= 0 = void $ + Q.takeWhile ((Refseq r ==) . b_rname . unpackBam) $ + Q.unfoldr (P.parseIO getBamRaw) $ + streamBgzf hdl voff Nothing + | otherwise = pure () + +{- | Reads from a Bam file the part with unaligned reads. + +Sort of the dual to 'streamBamRefseq'. Since the index does not +actually point to the unaligned part at the end, we use a best guess at +where the unaligned stuff might start, then skip over any aligned +records. Our \"fallback guess\" is to decode from the current position; +this only works if something else already consumed the Bam header. +-} +streamBamUnaligned :: MonadIO m => BamIndex b -> Handle -> Stream (Of BamRaw) m () +streamBamUnaligned BamIndex{..} hdl = + Q.filter (not . isValidRefseq . b_rname . unpackBam) $ + Q.unfoldr (P.parseIO getBamRaw) $ + streamBgzf hdl unaln_off Nothing + +{- | Streams one 'Segment'. + +Takes a 'Handle', a 'Segment' and a 'Stream' coming from that handle. +If skipping ahead in the stream looks cheap enough, that is done. Else +we seek the handle to the start offset and stream from it. Either way, +the part of the stream before it crosses either the end offset or the +max position is returned, and the remaining stream after it is returned +in its functorial value so it can be passed to another invocation of +e.g. 'streamBamSegment'. Note that the stream passed in becomes +unusable. +-} +streamBamSegment :: MonadIO m + => Handle -> Segment -> Stream (Of BamRaw) m () + -> Stream (Of BamRaw) m (Stream (Of BamRaw) m ()) +streamBamSegment hdl (Segment beg end mpos) = + lift . Q.uncons >=> \case + -- don't seek if it's a forwards seek of less than 512k + Just (br,brs) | near (virt_offset br) + -> Q.span in_seg $ Q.cons br brs + _ -> Q.span in_seg $ Q.unfoldr (P.parseIO getBamRaw) $ streamBgzf hdl beg (Just end) + where + near o = beg <= fromIntegral o && beg + 0x800000000 > fromIntegral o + in_seg br = virt_offset br <= end && b_pos (unpackBam br) <= mpos + +streamBamSubseq :: MonadIO m + => BamIndex b -> Handle -> Refseq -> R.Subsequence -> Stream (Of BamRaw) m () + -> Stream (Of BamRaw) m (Stream (Of BamRaw) m ()) +streamBamSubseq bi hdl ref subs str = Q.filter olap $ foldM (flip $ streamBamSegment hdl) str segs + where + segs = foldr (~~) [] $ segmentLists bi ref subs + olap br = case unpackBam br of + BamRec{..} -> b_rname == ref && R.overlaps b_pos (b_pos + alignedLength b_cigar) subs + +streamBamRegions :: MonadIO m => BamIndex b -> Handle -> [R.Region] -> Stream (Of BamRaw) m () +streamBamRegions bi hdl = void . foldM (\s (r,is) -> streamBamSubseq bi hdl r is s) (pure ()) . R.toList . R.fromList + diff --git a/Bio/Bam/Pileup.hs b/Bio/Bam/Pileup.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Pileup.hs @@ -0,0 +1,475 @@ +{-# LANGUAGE Rank2Types #-} + +-- | Pileup, similar to Samtools +-- +-- Pileup turns a sorted sequence of reads into a sequence of \"piles\", +-- one for each site where a genetic variant might be called. We will +-- scan each read's CIGAR line and MD field in concert with the sequence +-- and effective quality. Effective quality is the lowest available +-- quality score of QUAL, MAPQ, and BQ. For aDNA calling, a base is +-- represented as four probabilities, derived from a position dependent +-- damage model. + +module Bio.Bam.Pileup + ( PrimChunks(..) + , PrimBase(..) + , PosPrimChunks + , DamagedBase(..) + , DmgToken(..) + , dissect + , CallStats(..) + , V_Nuc(..) + , V_Nucs(..) + , IndelVariant(..) + , BasePile + , IndelPile + , Pile'(..) + , Pile + , pileup + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude +import Bio.Streaming + +import qualified Data.ByteString as B +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Storable as U +import qualified Streaming.Prelude as Q + +-- | The primitive pieces for genotype calling: A position, a base +-- represented as four likelihoods, an inserted sequence, and the +-- length of a deleted sequence. The logic is that we look at a base +-- followed by some indel, and all those indels are combined into a +-- single insertion and a single deletion. +data PrimChunks = Seek {-# UNPACK #-} !Int PrimBase -- ^ skip to position (at start or after N operation) + | Indel [Nucleotides] [DamagedBase] PrimBase -- ^ observed deletion and insertion between two bases + | EndOfRead -- ^ nothing anymore + deriving Show + +data PrimBase = Base { _pb_wait :: {-# UNPACK #-} !Int -- ^ number of bases to wait due to a deletion + , _pb_base :: {-# UNPACK #-} !DamagedBase + , _pb_mapq :: {-# UNPACK #-} !Qual -- ^ map quality + , _pb_chunks :: PrimChunks } -- ^ more chunks + deriving Show + +type PosPrimChunks = (Refseq, Int, Bool, PrimChunks) + +-- | Represents our knowledge about a certain base, which consists of +-- the base itself (A,C,G,T, encoded as 0..3; no Ns), the quality score +-- (anything that isn't A,C,G,T becomes A with quality 0), and a +-- substitution matrix representing post-mortem but pre-sequencing +-- substitutions. +-- +-- Unfortunately, none of this can be rolled into something more simple, +-- because damage and sequencing error behave so differently. +-- +-- Damage information is polymorphic. We might run with a simple +-- version (a matrix) for calling, but we need more (a matrix and a +-- mutable matrix, I think) for estimation. + +data DamagedBase = DB { db_call :: {-# UNPACK #-} !Nucleotide -- ^ called base + , db_qual :: {-# UNPACK #-} !Qual -- ^ quality of called base + , db_dmg_tk :: {-# UNPACK #-} !DmgToken -- ^ damage information + , db_dmg_pos :: {-# UNPACK #-} !Int -- ^ damage information + , db_ref :: {-# UNPACK #-} !Nucleotides } -- ^ reference base from MD field + +newtype DmgToken = DmgToken { fromDmgToken :: Int } + +instance Show DamagedBase where + showsPrec _ (DB n q _ _ r) + | nucToNucs n == r = shows n . (:) '@' . shows q + | otherwise = shows n . (:) '/' . shows r . (:) '@' . shows q + + +-- | Decomposes a BAM record into chunks suitable for piling up. We +-- pick apart the CIGAR and MD fields, and combine them with sequence +-- and quality as appropriate. Clipped bases are removed/skipped as +-- needed. We also apply a substitution matrix to each base, which must +-- be supplied along with the read. +{-# INLINE dissect #-} +dissect :: DmgToken -> BamRaw -> [PosPrimChunks] +dissect dtok br = + if isUnmapped b || isDuplicate b || not (isValidRefseq b_rname) + then [] else [(b_rname, b_pos, isReversed b, pchunks)] + where + b@BamRec{..} = unpackBam br + pchunks = firstBase b_pos 0 0 (fromMaybe [] $ getMd b) + + !max_cig = V.length b_cigar + !max_seq = V.length b_seq + !baq = extAsString "BQ" b + + -- This will compute the effective quality. As far as I can see + -- from the BAM spec V1.4, the qualities that matter are QUAL, MAPQ, + -- and BAQ. If QUAL is invalid, we replace it (arbitrarily) with + -- 23 (assuming a rather conservative error rate of ~0.5%), BAQ is + -- added to QUAL, and MAPQ is an upper limit for effective quality. + + get_seq :: Int -> (Nucleotides -> Nucleotides) -> DamagedBase + get_seq i f = case b_seq `V.unsafeIndex` i of -- nucleotide + n | n == nucsA -> DB nucA qe dtok dmg (f n) + | n == nucsC -> DB nucC qe dtok dmg (f n) + | n == nucsG -> DB nucG qe dtok dmg (f n) + | n == nucsT -> DB nucT qe dtok dmg (f n) + | otherwise -> DB nucA (Q 0) dtok dmg (f n) + where + !q = case b_qual `V.unsafeIndex` i of Q 0xff -> Q 30 ; x -> x -- quality; invalid (0xff) becomes 30 + !q' | i >= B.length baq = q -- no BAQ available + | otherwise = Q (unQ q + (B.index baq i - 64)) -- else correct for BAQ + !qe = min q' b_mapq -- use MAPQ as upper limit + !dmg = if i+i > max_seq then i-max_seq else i + + -- Look for first base following the read's start or a gap (CIGAR + -- code N). Indels are skipped, since these are either bugs in the + -- aligner or the aligner getting rid of essentially unalignable + -- bases. + firstBase :: Int -> Int -> Int -> [MdOp] -> PrimChunks + firstBase !pos !is !ic mds + | is >= max_seq || ic >= max_cig = EndOfRead + | otherwise = case b_cigar `V.unsafeIndex` ic of + Ins :* cl -> firstBase pos (cl+is) (ic+1) mds + SMa :* cl -> firstBase pos (cl+is) (ic+1) mds + Del :* cl -> firstBase (pos+cl) is (ic+1) (drop_del cl mds) + Nop :* cl -> firstBase (pos+cl) is (ic+1) mds + HMa :* _ -> firstBase pos is (ic+1) mds + Pad :* _ -> firstBase pos is (ic+1) mds + Mat :* 0 -> firstBase pos is (ic+1) mds + Mat :* _ -> Seek pos $ nextBase 0 pos is ic 0 mds + where + -- We have to treat (MdNum 0), because samtools actually + -- generates(!) it all over the place and if not handled as a + -- special case, it looks like an inconsistent MD field. + drop_del n (MdDel ns : mds') + | n < length ns = MdDel (drop n ns) : mds' + | n > length ns = drop_del (n - length ns) mds' + | otherwise = mds' + drop_del n (MdNum 0 : mds') = drop_del n mds' + drop_del _ mds' = mds' + + -- Generate likelihoods for the next base. When this gets called, + -- we are looking at an M CIGAR operation and all the subindices are + -- valid. + -- I don't think we can ever get (MdDel []), but then again, who + -- knows what crazy shit samtools decides to generate. There is + -- little harm in special-casing it. + nextBase :: Int -> Int -> Int -> Int -> Int -> [MdOp] -> PrimBase + nextBase !wt !pos !is !ic !io mds = case mds of + MdNum 0 : mds' -> nextBase wt pos is ic io mds' + MdDel [] : mds' -> nextBase wt pos is ic io mds' + MdNum 1 : mds' -> nextBase' (get_seq is id ) mds' + MdNum n : mds' -> nextBase' (get_seq is id ) (MdNum (n-1) : mds') + MdRep ref : mds' -> nextBase' (get_seq is $ const ref ) mds' + MdDel _ : _ -> nextBase' (get_seq is $ const nucsN) mds + [ ] -> nextBase' (get_seq is $ const nucsN) [ ] + where + nextBase' ref mds' = Base wt ref b_mapq $ nextIndel [] [] (pos+1) (is+1) ic (io+1) mds' + + -- Look for the next indel after a base. We collect all indels (I + -- and D codes) into one combined operation. If we hit N or the + -- read's end, we drop all of it (indels next to a gap indicate + -- trouble). Other stuff is skipped: we could check for stuff that + -- isn't valid in the middle of a read (H and S), but then what + -- would we do about it anyway? Just ignoring it is much easier and + -- arguably at least as correct. + nextIndel :: [[DamagedBase]] -> [Nucleotides] -> Int -> Int -> Int -> Int -> [MdOp] -> PrimChunks + nextIndel ins del !pos !is !ic !io mds + | is >= max_seq || ic >= max_cig = EndOfRead + | otherwise = case b_cigar `V.unsafeIndex` ic of + Ins :* cl -> nextIndel (isq cl) del pos (cl+is) (ic+1) 0 mds + SMa :* cl -> nextIndel ins del pos (cl+is) (ic+1) 0 mds + Del :* cl -> nextIndel ins (del++dsq) (pos+cl) is (ic+1) 0 mds' + where (dsq,mds') = split_del cl mds + Pad :* _ -> nextIndel ins del pos is (ic+1) 0 mds + HMa :* _ -> nextIndel ins del pos is (ic+1) 0 mds + Nop :* cl -> firstBase (pos+cl) is (ic+1) mds -- ends up generating a 'Seek' + Mat :* cl | io == cl -> nextIndel ins del pos is (ic+1) 0 mds + | otherwise -> indel del out $ nextBase (length del) pos is ic io mds -- ends up generating a 'Base' + where + indel d o k = foldr seq (Indel d o k) o + out = concat $ reverse ins + isq cl = [ get_seq i $ const gap | i <- [is..is+cl-1] ] : ins + + -- We have to treat (MdNum 0), because samtools actually + -- generates(!) it all over the place and if not handled as a + -- special case, it looks like an incinsistend MD field. + split_del n (MdDel ns : mds') + | n < length ns = (take n ns, MdDel (drop n ns) : mds') + | n > length ns = let (ns', mds'') = split_del (n - length ns) mds' in (ns++ns', mds'') + | otherwise = (ns, mds') + split_del n (MdNum 0 : mds') = split_del n mds' + split_del n mds' = (replicate n nucsN, mds') + +-- | Statistics about a genotype call. Probably only useful for +-- filtering (so not very useful), but we keep them because it's easy to +-- track them. + +data CallStats = CallStats + { read_depth :: {-# UNPACK #-} !Int -- number of contributing reads + , reads_mapq0 :: {-# UNPACK #-} !Int -- number of (non-)contributing reads with MAPQ==0 + , sum_mapq :: {-# UNPACK #-} !Int -- sum of map qualities of contributing reads + , sum_mapq_squared :: {-# UNPACK #-} !Int } -- sum of squared map qualities of contributing reads + deriving (Show, Eq, Generic) + +instance Monoid CallStats where + mempty = CallStats { read_depth = 0 + , reads_mapq0 = 0 + , sum_mapq = 0 + , sum_mapq_squared = 0 } + mappend = (<>) + +instance Semigroup CallStats where + x <> y = CallStats { read_depth = read_depth x + read_depth y + , reads_mapq0 = reads_mapq0 x + reads_mapq0 y + , sum_mapq = sum_mapq x + sum_mapq y + , sum_mapq_squared = sum_mapq_squared x + sum_mapq_squared y } + +newtype V_Nuc = V_Nuc (U.Vector Nucleotide) deriving (Eq, Ord, Show) +newtype V_Nucs = V_Nucs (U.Vector Nucleotides) deriving (Eq, Ord, Show) + +data IndelVariant = IndelVariant { deleted_bases :: !V_Nucs, inserted_bases :: !V_Nuc } + deriving (Eq, Ord, Show, Generic) + + +-- | Map quality and a list of encountered bases, with damage +-- information and reference base if known. +type BasePile = [DamagedBase] + +-- | Map quality and a list of encountered indel variants. The deletion +-- has the reference sequence, if known, an insertion has the inserted +-- sequence with damage information. +type IndelPile = [( Qual, ([Nucleotides], [DamagedBase]) )] -- a list of indel variants + +-- | Running pileup results in a series of piles. A 'Pile' has the +-- basic statistics of a 'VarCall', but no likelihood values and a +-- pristine list of variants instead of a proper call. We emit one pile +-- with two 'BasePile's (one for each strand) and one 'IndelPile' (the +-- one immediately following) at a time. + +data Pile' a b = Pile { p_refseq :: {-# UNPACK #-} !Refseq + , p_pos :: {-# UNPACK #-} !Int + , p_snp_stat :: {-# UNPACK #-} !CallStats + , p_snp_pile :: a + , p_indel_stat :: {-# UNPACK #-} !CallStats + , p_indel_pile :: b } + deriving Show + +-- | Raw pile. Bases and indels are piled separately on forward and +-- backward strands. +type Pile = Pile' (BasePile, BasePile) (IndelPile, IndelPile) + +-- | The pileup enumeratee takes 'BamRaw's, dissects them, interleaves +-- the pieces appropriately, and generates 'Pile's. The output will +-- contain at most one 'BasePile' and one 'IndelPile' for each position, +-- piles are sorted by position. +-- +-- This top level driver receives 'BamRaw's. Unaligned reads and +-- duplicates are skipped (but not those merely failing quality checks). +-- Processing stops when the first read with invalid 'br_rname' is +-- encountered or a t end of file. + +{-# INLINE pileup #-} +pileup :: Stream (Of PosPrimChunks) IO b -> Stream (Of Pile) IO b +pileup = runPileM pileup' finish (Refseq 0) 0 ([],[]) (Empty,Empty) + where + finish () _ _ ([],[]) (Empty,Empty) inp = lift (Q.effects inp) + finish () _ _ _ _ _ = error "logic error: leftovers after pileup" + + +-- | The pileup logic keeps a current coordinate (just two integers) and +-- two running queues: one of /active/ 'PrimBase's that contribute to +-- current genotype calling and on of /waiting/ 'PrimBase's that will +-- contribute at a later point. +-- +-- This is the CPS version of multiple state and environment monads. It +-- is somewhat faster than direct style and gives more control over when +-- evaluation happens. + +newtype PileM m a = PileM { runPileM :: forall r . (a -> PileF m r) -> PileF m r } + +-- | The things we drag along in 'PileM'. Notes: +-- * The /active/ queue is a simple stack. We add at the front when we +-- encounter reads, which reverses them. When traversing it, we traverse +-- reads backwards, but since we accumulate the 'BasePile', it gets reversed +-- back. The new /active/ queue, however, is no longer reversed (as it should +-- be). So after the traversal, we reverse it again. (Yes, it is harder to +-- understand than using a proper deque type, but it is cheaper. +-- There may not be much point in the reversing, though.) + +type PileF m r = Refseq -> Int -> -- current position + ( [PrimBase], [PrimBase] ) -> -- active queues + ( Heap, Heap ) -> -- waiting queues + Stream (Of PosPrimChunks) m r -> -- input stream + Stream (Of Pile) m r -- output stream + +instance Functor (PileM m) where + {-# INLINE fmap #-} + fmap f (PileM m) = PileM $ \k -> m (k . f) + +instance Applicative (PileM m) where + {-# INLINE pure #-} + pure a = PileM $ \k -> k a + {-# INLINE (<*>) #-} + u <*> v = PileM $ \k -> runPileM u (\a -> runPileM v (k . a)) + +instance Monad (PileM m) where + {-# INLINE return #-} + return a = PileM $ \k -> k a + {-# INLINE (>>=) #-} + m >>= k = PileM $ \k' -> runPileM m (\a -> runPileM (k a) k') + +-- -- XXX +-- {-# INLINE get_refseq #-} +-- get_refseq :: PileM m Refseq +-- get_refseq = PileM $ \k r -> k r r + +-- {-# INLINE get_pos #-} +-- get_pos :: PileM m Int +-- get_pos = PileM $ \k r p -> k p r p + +{-# INLINE upd_pos #-} +upd_pos :: (Int -> Int) -> PileM m () +upd_pos f = PileM $ \k r p -> k () r $! f p + +{-# INLINE yieldPile #-} +yieldPile :: Monad m => CallStats -> BasePile -> BasePile -> CallStats -> IndelPile -> IndelPile -> PileM m () +yieldPile x1 x2a x2b x3 x4a x4b = PileM $ \ !kont !r !p !a !w !inp -> + let pile = Pile r p x1 (x2a,x2b) x3 (x4a,x4b) + in Q.cons pile $ kont () r p a w inp + +-- | The actual pileup algorithm. If /active/ contains something, +-- continue here. Else find the coordinate to continue from, which is +-- the minimum of the next /waiting/ coordinate and the next coordinate +-- in input; if found, continue there, else we're all done. +pileup' :: Monad m => PileM m () +pileup' = PileM $ \ !k !refseq !pos !active !waiting inp0 -> do + + inp <- lift $ inspect inp0 + let inp1 = either pure wrap inp + cont2 rs po = runPileM pileup'' k rs po active waiting inp1 + leave = k () refseq pos active waiting inp1 + + case (active, getMinKeysH waiting, inp) of + ( (_:_,_), _, _ ) -> cont2 refseq pos + ( (_,_:_), _, _ ) -> cont2 refseq pos + ( _, Just nw, Left _ ) -> cont2 refseq nw + ( _, Nothing, Left _ ) -> leave + ( _, Nothing, Right ((r,p,_,_):>_) ) -> cont2 r p + ( _, Just nw, Right ((r,p,_,_):>_) ) + | (refseq,nw) <= (r,p) -> cont2 refseq nw + | otherwise -> cont2 r p + where + getMinKeysH :: (Heap, Heap) -> Maybe Int + getMinKeysH (a,b) = case (getMinKeyH a, getMinKeyH b) of + ( Nothing, Nothing ) -> Nothing + ( Just x, Nothing ) -> Just x + ( Nothing, Just y ) -> Just y + ( Just x, Just y ) -> Just (min x y) + + +pileup'' :: Monad m => PileM m () +pileup'' = do + -- Input is still 'BamRaw', since these can be relied on to be + -- sorted. First see if there is any input at the current location, + -- if so, dissect it and add it to the appropriate queue. + p'feed_input + p'check_waiting + ((fin_bsL, fin_bpL), (fin_bsR, fin_bpR), (fin_isL, fin_ipL), (fin_isR, fin_ipR)) <- p'scan_active + + -- Output, but don't bother emitting empty piles. Note that a plain + -- basecall still yields an entry in the 'IndelPile'. This is necessary, + -- because actual indel calling will want to know how many reads /did not/ + -- show the variant. However, if no reads show any variant, and here is the + -- first place where we notice that, the pile is useless. + let uninteresting (_,(d,i)) = null d && null i + unless (null fin_bpL && null fin_bpR && all uninteresting fin_ipL && all uninteresting fin_ipR) $ + yieldPile (fin_bsL <> fin_bsR) fin_bpL fin_bpR + (fin_isL <> fin_isR) fin_ipL fin_ipR + + -- Bump coordinate and loop. (Note that the bump to the next + -- reference /sequence/ is done implicitly, because we will run out of + -- reads and restart in 'pileup''.) + upd_pos succ + pileup' + +-- | Feeds input as long as it starts at the current position +p'feed_input :: Monad m => PileM m () +p'feed_input = PileM $ \kont rs po ac@(af,ar) wt@(wf,wr) -> + lift . inspect >=> \case + Right ((rs', po', str, prim) :> bs) + | rs == rs' && po == po' -> + case prim of + EndOfRead -> runPileM p'feed_input kont rs po ac wt bs + Indel _ _ !pb -> runPileM p'feed_input kont rs po (if str then (af,pb:ar) else (pb:af,ar)) wt bs + Seek !p !pb -> runPileM p'feed_input kont rs po ac (if str then (wf,wr') else (wf',wr)) bs + where wf' = Node p pb Empty Empty `unionH` wf + wr' = Node p pb Empty Empty `unionH` wr + inp -> kont () rs po ac wt $ either pure wrap inp + +-- | Checks /waiting/ queue. If there is anything waiting for the +-- current position, moves it to /active/ queue. +p'check_waiting :: PileM m () +p'check_waiting = PileM $ \kont rs po (af0,ar0) (wf0,wr0) -> + let go1 af wf = case viewMinH wf of + Just (!mk, !pb, !wf') | mk == po -> go1 (pb:af) wf' + _ -> go2 af wf ar0 wr0 + + go2 af wf ar wr = case viewMinH wr of + Just (!mk, !pb, !wr') | mk == po -> go2 af wf (pb:ar) wr' + _ -> kont () rs po (af,ar) (wf,wr) + + in go1 af0 wf0 + + +-- | Separately scans the two /active/ queues and makes one 'BasePile' +-- from each. Also sees what's next in the 'PrimChunks': 'Indel's +-- contribute to two separate 'IndelPile's, 'Seek's are pushed back to +-- the /waiting/ queue, 'EndOfRead's are removed, and everything else is +-- added to two fresh /active/ queues. +p'scan_active :: PileM m (( CallStats, BasePile ), ( CallStats, BasePile ), + ( CallStats, IndelPile ), ( CallStats, IndelPile )) +p'scan_active = do + (bpf,ipf) <- PileM $ \kont rs pos (af,ar) (wf,wr) -> go (\r af' wf' -> kont r rs pos (af',ar) (wf',wr)) [] wf mempty mempty af + (bpr,ipr) <- PileM $ \kont rs pos (af,ar) (wf,wr) -> go (\r ar' wr' -> kont r rs pos (af,ar') (wf,wr')) [] wr mempty mempty ar + return (bpf,bpr,ipf,ipr) + where + go k !ac !wt !bpile !ipile [ ] = k (bpile, ipile) (reverse ac) wt + go k !ac !wt !bpile !ipile (Base nwt qs mq pchunks : bs) = + case pchunks of + _ | nwt > 0 -> b' `seq` go k (b':ac) wt bpile ipile bs + Seek p' pb' -> go k ac (ins p' pb' wt) (z bpile) ipile bs + Indel nd ni pb' -> go k (pb':ac) wt (z bpile) (y ipile) bs where y = put (,) mq (nd,ni) + EndOfRead -> go k ac wt (z bpile) ipile bs + where + b' = Base (nwt-1) qs mq pchunks + z = put (\q x -> x { db_qual = min q (db_qual x) }) mq qs + + ins q v w = Node q v Empty Empty `unionH` w + + put f (Q !q) !x (!st,!vs) = ( st { read_depth = read_depth st + 1 + , reads_mapq0 = reads_mapq0 st + (if q == 0 then 1 else 0) + , sum_mapq = sum_mapq st + fromIntegral q + , sum_mapq_squared = sum_mapq_squared st + fromIntegral q * fromIntegral q } + , f (Q q) x : vs ) + + +-- | We need a simple priority queue. Here's a skew heap (specialized +-- to strict 'Int' priorities and 'PrimBase' values). +data Heap = Empty | Node {-# UNPACK #-} !Int PrimBase Heap Heap + +unionH :: Heap -> Heap -> Heap +Empty `unionH` t2 = t2 +t1 `unionH` Empty = t1 +t1@(Node k1 x1 l1 r1) `unionH` t2@(Node k2 x2 l2 r2) + | k1 <= k2 = Node k1 x1 (t2 `unionH` r1) l1 + | otherwise = Node k2 x2 (t1 `unionH` r2) l2 + +getMinKeyH :: Heap -> Maybe Int +getMinKeyH Empty = Nothing +getMinKeyH (Node x _ _ _) = Just x + +viewMinH :: Heap -> Maybe (Int, PrimBase, Heap) +viewMinH Empty = Nothing +viewMinH (Node k v l r) = Just (k, v, l `unionH` r) + diff --git a/Bio/Bam/Reader.hs b/Bio/Bam/Reader.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Reader.hs @@ -0,0 +1,261 @@ +-- | Parsers for BAM and SAM. + +module Bio.Bam.Reader ( + decodeBam, + decodeBamFile, + decodeBamFiles, + + decodePlainBam, + decodePlainSam, + getBamMeta, + getBamRaw, + getSamRec, + + concatInputs, + mergeInputsOn, + guardRefCompat, + coordinates, + qnames + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Bam.Writer ( packBam ) +import Bio.Streaming +import Bio.Streaming.Bgzf ( getBgzfHdr, bgunzip ) +import Bio.Prelude +import Data.Attoparsec.ByteString ( anyWord8 ) + +import qualified Data.Attoparsec.ByteString.Char8 as P +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as C +import qualified Data.HashMap.Strict as M +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Storable as W +import qualified Data.Vector.Unboxed as U +import qualified Bio.Streaming.Bytes as S +import qualified Bio.Streaming.Parse as S +import qualified Streaming.Prelude as Q + +{- | Decodes either BAM or SAM. + +The input can be plain, gzip'ed or bgzf'd and either BAM or SAM. BAM +is reliably recognized, anything else is treated as SAM. The offsets +stored in BAM records make sense only for uncompressed or bgzf'd BAM. +-} +decodeBam :: MonadIO m + => S.ByteStream m r + -> m (BamMeta, Stream (Of BamRaw) m r) +decodeBam = getBgzfHdr >=> S.splitAt' 4 . pgunzip >=> unbam + where + unbam ("BAM\SOH" :> s) = decodePlainBam s + unbam (magic :> s) = decodePlainSam (S.consChunk magic s) + + pgunzip (Nothing, hdr, s) = S.gunzip (S.consChunk hdr s) + pgunzip (Just _, hdr, s) = bgunzip (S.consChunk hdr s) +{-# INLINE decodeBam #-} + +decodeBamFile :: (MonadIO m, MonadMask m) => FilePath -> (BamMeta -> Stream (Of BamRaw) m () -> m r) -> m r +decodeBamFile f k = streamFile f $ decodeBam >=> uncurry k +{-# INLINE decodeBamFile #-} + +{- | Reads multiple bam files. + +A continuation is run on the list of headers and streams. Since no +attempt is made to unify the headers, this will work for completely +unrelated bam files. All files are opened at the same time, which might +run into the file descriptor limit given some ridiculous workflows. +-} +decodeBamFiles :: (MonadMask m, MonadIO m) => [FilePath] -> ([(BamMeta, Stream (Of BamRaw) m ())] -> m r) -> m r +decodeBamFiles [ ] k = k [] +decodeBamFiles ("-":fs) k = decodeBam (streamHandle stdin) >>= \b -> decodeBamFiles fs $ \bs -> k (b:bs) +decodeBamFiles ( f :fs) k = streamFile f $ \s -> decodeBam s >>= \b -> decodeBamFiles fs $ \bs -> k (b:bs) +{-# INLINE decodeBamFiles #-} + +decodePlainBam :: MonadIO m => S.ByteStream m r -> m (BamMeta, Stream (Of BamRaw) m r) +decodePlainBam = + S.parseIO (const getBamMeta) >=> + either (const . liftIO $ throwM S.EofException) + (return . fmap (Q.unfoldr (S.parseIO getBamRaw))) + +getBamMeta :: Monad m => S.Parser r m BamMeta +getBamMeta = liftM2 mmerge get_bam_header get_ref_array + where + get_bam_header = do hdr_len <- S.getWord32 + S.isolate (fromIntegral hdr_len) (S.atto parseBamMeta) + + get_ref_array = do nref <- S.getWord32 + V.fromList `liftM` replicateM (fromIntegral nref) + (do nm <- S.getWord32 >>= S.getString . fromIntegral + ln <- S.getWord32 + return $! BamSQ (C.init nm) (fromIntegral ln) []) + + -- Need to merge information from header into actual reference list. + -- The latter is the authoritative source for the *order* of the + -- sequences, so leftovers from the header are discarded. Merging + -- is by name. So we merge information from the header into the + -- list, then replace the header information. + mmerge meta refs = + let tbl = M.fromList [ (sq_name sq, sq) | sq <- V.toList (unRefs (meta_refs meta)) ] + in meta { meta_refs = Refs $ fmap (\s -> maybe s (mmerge' s) (M.lookup (sq_name s) tbl)) refs } + + mmerge' l r | sq_length l == sq_length r = l { sq_other_shit = sq_other_shit l ++ sq_other_shit r } + | otherwise = l -- contradiction in header, but we'll just ignore it +{-# INLINABLE getBamMeta #-} + + +data ShortRecord = ShortRecord deriving Typeable + +instance Exception ShortRecord +instance Show ShortRecord where show _ = "short BAM record" + +getBamRaw :: Monad m => Int64 -> S.Parser r m BamRaw +getBamRaw o = do + bsize <- fromIntegral `liftM` S.getWord32 + when (bsize < 32) $ throwM ShortRecord + s <- S.getString bsize + unless (B.length s == bsize) S.abortParse + return $ bamRaw o s +{-# INLINABLE getBamRaw #-} + +{- | Streaming parser for SAM files. + +It parses plain uncompressed SAM +and returns a result compatible with 'decodePlainBam'. Since it is +supposed to work the same way as the BAM parser, it requires a +symbol table for the reference names. This is extracted from the @SQ +lines in the header. Note that reading SAM tends to be inefficient; +if you care about performance at all, use BAM. +-} +decodePlainSam :: MonadIO m => S.ByteStream m r -> m (BamMeta, Stream (Of BamRaw) m r) +decodePlainSam s = do + (hdr,rest) <- either (\r -> (mempty, pure r)) id `liftM` S.parseIO (const $ S.atto parseBamMeta) s + let !refs = M.fromList $ zip [ nm | BamSQ { sq_name = nm } <- V.toList (unRefs (meta_refs hdr))] [toEnum 0..] + ref x = M.lookupDefault invalidRefseq x refs + return (hdr, Q.mapM (liftIO . either throwM packBam . getSamRec ref) (S.lines' rest)) + + +getSamRec :: (Bytes -> Refseq) -> Bytes -> Either S.ParseError BamRec +getSamRec ref s = case P.parseOnly record s of + Left e -> Left $ S.ParseError [unpack s] e + Right b -> Right b + where + record = do b_qname <- word + b_flag <- num + b_rname <- ref <$> word + b_pos <- subtract 1 <$> num + b_mapq <- Q <$> num' + b_cigar <- W.fromList <$> cigar + b_mrnm <- rnext <*> pure b_rname + b_mpos <- subtract 1 <$> num + b_isize <- snum + b_seq <- sequ + b_qual <- quals <*> pure b_seq + b_exts <- exts + let b_virtual_offset = 0 + return BamRec{..} + + sep = P.endOfInput <|> () <$ P.char '\t' + word = P.takeTill ('\t' ==) <* sep + num = P.decimal <* sep + num' = P.decimal <* sep + snum = P.signed P.decimal <* sep + + rnext = id <$ P.char '=' <* sep <|> const . ref <$> word + sequ = (V.empty <$ P.char '*' <|> + V.fromList . map toNucleotides . C.unpack <$> P.takeWhile is_nuc) <* sep + + quals = defaultQs <$ P.char '*' <* sep <|> bsToVec <$> word + where + defaultQs sq = W.replicate (V.length sq) (Q 0xff) + bsToVec qs _ = W.fromList . map (Q . subtract 33) $ B.unpack qs + + cigar = [] <$ P.char '*' <* sep <|> + P.manyTill (flip (:*) <$> P.decimal <*> cigop) sep + + cigop = P.choice $ zipWith (\c r -> r <$ P.char c) "MIDNSHP" [Mat,Ins,Del,Nop,SMa,HMa,Pad] + exts = ext `P.sepBy` sep + ext = (\a b v -> (fromString [a,b],v)) <$> P.anyChar <*> P.anyChar <*> (P.char ':' *> value) + + value = P.char 'A' *> P.char ':' *> (Char <$> anyWord8) <|> + P.char 'i' *> P.char ':' *> (Int <$> P.signed P.decimal) <|> + P.char 'Z' *> P.char ':' *> (Text <$> P.takeTill ('\t' ==)) <|> + P.char 'H' *> P.char ':' *> (Bin <$> hexarray) <|> + P.char 'f' *> P.char ':' *> (Float . realToFrac <$> P.double) <|> + P.char 'B' *> P.char ':' *> ( + P.satisfy (P.inClass "cCsSiI") *> (intArr <$> many (P.char ',' *> P.signed P.decimal)) <|> + P.char 'f' *> (floatArr <$> many (P.char ',' *> P.double))) + + intArr is = IntArr $ U.fromList is + floatArr fs = FloatArr $ U.fromList $ map realToFrac fs + hexarray = B.pack . repack . C.unpack <$> P.takeWhile (P.inClass "0-9A-Fa-f") + repack (a:b:cs) = fromIntegral (digitToInt a * 16 + digitToInt b) : repack cs ; repack _ = [] + is_nuc = P.inClass "acgtswkmrybdhvnACGTSWKMRYBDHVN" + + +data IncompatibleRefs = IncompatibleRefs FilePath FilePath deriving Typeable + +instance Exception IncompatibleRefs +instance Show IncompatibleRefs where + show (IncompatibleRefs a b) = "references in " ++ a ++ " and " ++ b ++ " are incompatible" + +guardRefCompat :: MonadThrow m => (FilePath,BamMeta) -> (FilePath,BamMeta) -> m () +guardRefCompat (f0,hdr0) (f1,hdr1) = + unless (p hdr1 `isPrefixOf` p hdr0) $ throwM $ IncompatibleRefs f0 f1 + where + p = V.toList . unRefs . meta_refs + + +{- | Reads multiple bam inputs in sequence. + +Only one file is opened at a time, so they must also be consumed in +sequence. If you can afford to open all inputs simultaneously, you +probably want to use 'mergeInputsOn' instead. The filename \"-\" refers +to stdin, if no filenames are given, stdin is read. Since we can't look +ahead into further files, the header of the first input is used +for the result, and an exception is thrown if one of the subsequent +headers is incompatible with the first one. +-} +concatInputs :: (MonadIO m, MonadMask m) => [FilePath] -> (BamMeta -> Stream (Of BamRaw) m () -> m r) -> m r +concatInputs fs0 k = streamInputs fs0 (go1 $ fs0 ++ repeat "-") + where + go1 fs = inspect >=> \case + Left () -> k mempty (pure ()) + Right s -> do (hdr,bs) <- decodeBam s + k hdr (bs >>= go (head fs) hdr (tail fs)) + + go f0 hdr0 fs = lift . inspect >=> \case + Left () -> pure () + Right s -> do (hdr,bs) <- lift $ decodeBam s + lift $ guardRefCompat (f0,hdr0) (head fs,hdr) + bs >>= go f0 hdr0 (tail fs) +{-# INLINABLE concatInputs #-} + +{- | Reads multiple bam files and merges them. + +If the inputs are all sorted by the thing being merged on, the output +will be sorted, too. The headers are all merged sensibly, even if their +reference lists differ. However, for performance reasons, we don't want +to change the rname and mrnm fields in potentially all records. So +instead of allowing arbitrary reference lists to be merged, we throw an +exception unless every input is compatible with the effective reference +list. +-} +mergeInputsOn :: (Ord x, MonadIO m, MonadMask m) + => (BamRaw -> x) -> [FilePath] + -> (BamMeta -> Stream (Of BamRaw) m () -> m r) -> m r +mergeInputsOn _ [] k = decodeBam (streamHandle stdin) >>= uncurry k +mergeInputsOn p fs k = decodeBamFiles fs $ \bs -> do + let hdr = foldMap fst bs + sequence_ $ zipWith (\f (h,_) -> guardRefCompat ("*",hdr) (f,h)) fs bs + k hdr (foldr (\a b -> void $ mergeStreamsOn p (snd a) b) (pure ()) bs) +{-# INLINABLE mergeInputsOn #-} + +coordinates :: BamRaw -> (Refseq, Int) +coordinates = (b_rname &&& b_pos) . unpackBam +{-# INLINE coordinates #-} + +qnames :: BamRaw -> Bytes +qnames = b_qname . unpackBam +{-# INLINE qnames #-} + diff --git a/Bio/Bam/Rec.hs b/Bio/Bam/Rec.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Rec.hs @@ -0,0 +1,378 @@ +{-# LANGUAGE TypeFamilies #-} +-- | Parsers and Printers for BAM and SAM. + +module Bio.Bam.Rec ( + BamRaw, + bamRaw, + virt_offset, + raw_data, + + BamRec(..), + unpackBam, + nullBamRec, + getMd, + + Cigar(..), + CigOp(..), + alignedLength, + + Nucleotides(..), Vector_Nucs_half, + Extensions, Ext(..), + extAsInt, extAsString, setQualFlag, + deleteE, insertE, updateE, adjustE, + + isPaired, + isProperlyPaired, + isUnmapped, + isMateUnmapped, + isReversed, + isMateReversed, + isFirstMate, + isSecondMate, + isAuxillary, + isSecondary, + isFailsQC, + isDuplicate, + isSupplementary, + isTrimmed, + isMerged, + isAlternative, + isExactIndex, + type_mask +) where + +import Bio.Bam.Header +import Bio.Prelude +import Bio.Util.Storable + +import Control.Monad.Primitive ( unsafePrimToPrim, unsafeInlineIO ) +import Foreign.C.Types ( CInt(..), CSize(..) ) +import Foreign.Marshal.Alloc ( alloca ) + +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as S +import qualified Data.ByteString.Internal as B +import qualified Data.ByteString.Unsafe as B +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Generic.Mutable as VM +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as U + + +-- | Cigar line in BAM coding +-- Bam encodes an operation and a length into a single integer, we keep +-- those integers in an array. +data Cigar = !CigOp :* !Int deriving (Eq, Ord) +infix 9 :* + +data CigOp = Mat | Ins | Del | Nop | SMa | HMa | Pad + deriving ( Eq, Ord, Enum, Show, Bounded, Ix ) + +instance Show Cigar where + showsPrec _ (op :* num) = shows num . (:) (S.index "MIDNSHP" (fromEnum op)) + +instance Storable Cigar where + sizeOf _ = 4 + alignment _ = 1 + + peek p = do w <- fromIntegral <$> peekUnalnWord32LE p + return $ toEnum (w .&. 0xf) :* shiftR w 4 + + poke p (op :* num) = pokeUnalnWord32LE p . fromIntegral $ fromEnum op .|. shiftL num 4 + + +-- | Extracts the aligned length from a cigar line. +-- This gives the length of an alignment as measured on the reference, +-- which is different from the length on the query or the length of the +-- alignment. +{-# INLINE alignedLength #-} +alignedLength :: V.Vector v Cigar => v Cigar -> Int +alignedLength = V.foldl' (\a -> (a +) . l) 0 + where l (op :* n) = if op == Mat || op == Del || op == Nop then n else 0 + + +-- | internal representation of a BAM record +data BamRec = BamRec { + b_qname :: Bytes, + b_flag :: Int, + b_rname :: Refseq, + b_pos :: Int, + b_mapq :: Qual, + b_cigar :: VS.Vector Cigar, + b_mrnm :: Refseq, + b_mpos :: Int, + b_isize :: Int, + b_seq :: Vector_Nucs_half Nucleotides, + b_qual :: VS.Vector Qual, + b_exts :: Extensions, + b_virtual_offset :: Int64 -- ^ virtual offset for indexing purposes + } deriving Show + +nullBamRec :: BamRec +nullBamRec = BamRec { + b_qname = S.empty, + b_flag = flagUnmapped, + b_rname = invalidRefseq, + b_pos = invalidPos, + b_mapq = Q 0, + b_cigar = VS.empty, + b_mrnm = invalidRefseq, + b_mpos = invalidPos, + b_isize = 0, + b_seq = V.empty, + b_qual = VS.empty, + b_exts = [], + b_virtual_offset = 0 + } + +getMd :: BamRec -> Maybe [MdOp] +getMd r = case lookup "MD" $ b_exts r of + Just (Text mdfield) -> readMd mdfield + Just (Char mdfield) -> readMd $ B.singleton mdfield + _ -> Nothing + +-- | A vector that packs two 'Nucleotides' into one byte, just like Bam does. +data Vector_Nucs_half a = Vector_Nucs_half !Int !Int !(ForeignPtr Word8) + +-- | A mutable vector that packs two 'Nucleotides' into one byte, just like Bam does. +data MVector_Nucs_half s a = MVector_Nucs_half !Int !Int !(ForeignPtr Word8) + +type instance V.Mutable Vector_Nucs_half = MVector_Nucs_half + +instance V.Vector Vector_Nucs_half Nucleotides where + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeFreeze (MVector_Nucs_half o l fp) = return $ Vector_Nucs_half o l fp + {-# INLINE basicUnsafeThaw #-} + basicUnsafeThaw (Vector_Nucs_half o l fp) = return $ MVector_Nucs_half o l fp + + {-# INLINE basicLength #-} + basicLength (Vector_Nucs_half _ l _) = l + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice s l (Vector_Nucs_half o _ fp) = Vector_Nucs_half (o + s) l fp + + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeIndexM (Vector_Nucs_half o _ fp) i + | even (o+i) = return . Ns $! (b `shiftR` 4) .&. 0xF + | otherwise = return . Ns $! b .&. 0xF + where !b = unsafeInlineIO $ withForeignPtr fp $ \p -> peekByteOff p ((o+i) `shiftR` 1) + +instance VM.MVector MVector_Nucs_half Nucleotides where + {-# INLINE basicLength #-} + basicLength (MVector_Nucs_half _ l _) = l + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice s l (MVector_Nucs_half o _ fp) = MVector_Nucs_half (o + s) l fp + + {-# INLINE basicOverlaps #-} + basicOverlaps (MVector_Nucs_half _ _ fp1) (MVector_Nucs_half _ _ fp2) = fp1 == fp2 + {-# INLINE basicUnsafeNew #-} + basicUnsafeNew l = unsafePrimToPrim $ MVector_Nucs_half 0 l <$> mallocForeignPtrBytes ((l+1) `shiftR` 1) + + {-# INLINE basicInitialize #-} + basicInitialize v@(MVector_Nucs_half o l fp) + + | even o = do unsafePrimToPrim $ withForeignPtr fp $ \p -> + memset (plusPtr p (o `shiftR` 1)) 0 (fromIntegral $ l `shiftR` 1) + when (odd l) $ VM.basicUnsafeWrite v (l-1) (Ns 0) + + | otherwise = do when (odd o) $ VM.basicUnsafeWrite v 0 (Ns 0) + unsafePrimToPrim $ withForeignPtr fp $ \p -> + memset (plusPtr p ((o+1) `shiftR` 1)) 0 (fromIntegral $ (l-1) `shiftR` 1) + when (even l) $ VM.basicUnsafeWrite v (l-1) (Ns 0) + + + {-# INLINE basicUnsafeRead #-} + basicUnsafeRead (MVector_Nucs_half o _ fp) i + | even (o+i) = liftM (Ns . (.&.) 0xF . (`shiftR` 4)) b + | otherwise = liftM (Ns . (.&.) 0xF ) b + where b = unsafePrimToPrim $ withForeignPtr fp $ \p -> peekByteOff p ((o+i) `shiftR` 1) + + {-# INLINE basicUnsafeWrite #-} + basicUnsafeWrite (MVector_Nucs_half o _ fp) i (Ns x) = + unsafePrimToPrim $ withForeignPtr fp $ \p -> do + y <- peekByteOff p ((o+i) `shiftR` 1) + let y' | even (o+i) = x `shiftL` 4 .|. y .&. 0x0F + | otherwise = x .|. y .&. 0xF0 + pokeByteOff p ((o+i) `shiftR` 1) y' + +foreign import ccall unsafe "string.h memset" memset + :: Ptr Word8 -> CInt -> CSize -> IO () + +instance Show (Vector_Nucs_half Nucleotides) where + show = show . V.toList + +-- | Bam record in its native encoding along with virtual address. +data BamRaw = BamRaw { virt_offset :: {-# UNPACK #-} !Int64 + , raw_data :: {-# UNPACK #-} !Bytes } + +-- | Smart constructor. Makes sure we got a at least a full record. +{-# INLINE bamRaw #-} +bamRaw :: Int64 -> Bytes -> BamRaw +bamRaw o s = if good then BamRaw o s else error $ "broken BAM record " ++ shows (S.length s, m) " " ++ show (S.unpack (S.take 10 s)) + where + good | S.length s < 32 = False + | otherwise = S.length s >= sum m + m = [ 32, l_rnm, l_seq, (l_seq+1) `div` 2, l_cig * 4 ] + l_rnm = fromIntegral (B.unsafeIndex s 8) - 1 + l_cig = fromIntegral (B.unsafeIndex s 12) .|. fromIntegral (B.unsafeIndex s 13) `shiftL` 8 + l_seq = fromIntegral (B.unsafeIndex s 16) .|. fromIntegral (B.unsafeIndex s 17) `shiftL` 8 .|. + fromIntegral (B.unsafeIndex s 18) `shiftL` 16 .|. fromIntegral (B.unsafeIndex s 19) `shiftL` 24 + +{-# INLINE[1] unpackBam #-} +unpackBam :: BamRaw -> BamRec +unpackBam br = BamRec { + b_rname = Refseq $ getWord32 0, + b_pos = getInt32 4, + b_mapq = Q $ getInt8 9, + b_flag = getInt16 14, + b_mrnm = Refseq $ getWord32 20, + b_mpos = getInt32 24, + b_isize = getInt32 28, + + b_qname = B.unsafeTake l_read_name $ B.unsafeDrop 32 $ raw_data br, + b_cigar = VS.unsafeCast $ VS.unsafeFromForeignPtr fp (off0+off_c) (4*l_cigar), + b_seq = Vector_Nucs_half (2 * (off_s+off0)) l_seq fp, + b_qual = VS.unsafeCast $ VS.unsafeFromForeignPtr fp (off0+off_q) l_seq, + + b_exts = unpackExtensions $ S.drop off_e $ raw_data br, + b_virtual_offset = virt_offset br } + where + (fp, off0, _) = B.toForeignPtr $ raw_data br + off_c = 33 + l_read_name + off_s = off_c + 4 * l_cigar + off_q = off_s + (l_seq + 1) `div` 2 + off_e = off_q + l_seq + + l_read_name = getInt8 8 - 1 + l_seq = getWord32 16 + l_cigar = getInt16 12 + + getInt8 :: Num a => Int -> a + getInt8 o = fromIntegral (B.unsafeIndex (raw_data br) o) + + getInt16 :: Num a => Int -> a + getInt16 o = unsafeDupablePerformIO $ B.unsafeUseAsCString (raw_data br) $ + fmap fromIntegral . peekUnalnWord16LE . flip plusPtr o + + getWord32 :: Num a => Int -> a + getWord32 o = unsafeDupablePerformIO $ B.unsafeUseAsCString (raw_data br) $ + fmap fromIntegral . peekUnalnWord32LE . flip plusPtr o + + -- ensures proper sign extension + getInt32 :: Num a => Int -> a + getInt32 o = fromIntegral (getWord32 o :: Int32) + +-- | A collection of extension fields. A 'BamKey' is actually two ASCII +-- characters. +type Extensions = [( BamKey, Ext )] + +-- | Deletes all occurences of some extension field. +deleteE :: BamKey -> Extensions -> Extensions +deleteE k = filter ((/=) k . fst) + +-- | Blindly inserts an extension field. This can create duplicates +-- (and there is no telling how other tools react to that). +insertE :: BamKey -> Ext -> Extensions -> Extensions +insertE k v = (:) (k,v) + +-- | Deletes all occurences of an extension field, then inserts it with +-- a new value. This is safer than 'insertE', but also more expensive. +updateE :: BamKey -> Ext -> Extensions -> Extensions +updateE k v = insertE k v . deleteE k + +-- | Adjusts a named extension by applying a function. +adjustE :: (Ext -> Ext) -> BamKey -> Extensions -> Extensions +adjustE _ _ [ ] = [] +adjustE f k ((k',v):es) | k == k' = (k', f v) : es + | otherwise = (k', v) : adjustE f k es + +data Ext = Int Int | Float Float | Text Bytes | Bin Bytes | Char Word8 + | IntArr (U.Vector Int) | FloatArr (U.Vector Float) + deriving (Show, Eq, Ord) + +{-# INLINE unpackExtensions #-} +unpackExtensions :: Bytes -> Extensions +unpackExtensions = go + where + go s | S.length s < 4 = [] + | otherwise = let key = fromString [ S.index s 0, S.index s 1 ] + in case S.index s 2 of + 'Z' -> case S.break (== '\0') (S.drop 3 s) of (l,r) -> (key, Text l) : go (S.drop 1 r) + 'H' -> case S.break (== '\0') (S.drop 3 s) of (l,r) -> (key, Bin l) : go (S.drop 1 r) + 'A' -> (key, Char (B.index s 3)) : go (S.drop 4 s) + 'B' -> let tp = S.index s 3 + n = getInt 'I' (S.drop 4 s) + in case tp of + 'f' -> (key, FloatArr (U.fromListN (n+1) [ getFloat (S.drop i s) | i <- [8, 12 ..] ])) + : go (S.drop (12+4*n) s) + _ -> (key, IntArr (U.fromListN (n+1) [ getInt tp (S.drop i s) | i <- [8, 8 + size tp ..] ])) + : go (S.drop (8 + size tp * (n+1)) s) + 'f' -> (key, Float (getFloat (S.drop 3 s))) : go (S.drop 7 s) + tp -> (key, Int (getInt tp (S.drop 3 s))) : go (S.drop (3 + size tp) s) + + size 'C' = 1 + size 'c' = 1 + size 'S' = 2 + size 's' = 2 + size 'I' = 4 + size 'i' = 4 + size 'f' = 4 + size _ = 0 + + getInt 'C' s | S.length s >= 1 = fromIntegral ((B.index s 0) :: Word8) + getInt 'c' s | S.length s >= 1 = fromIntegral (fromIntegral (B.index s 0) :: Int8) + getInt 'S' s | S.length s >= 2 = fromIntegral (i :: Word16) + where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord16LE + getInt 's' s | S.length s >= 2 = fromIntegral (fromIntegral i :: Int16) + where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord16LE + getInt 'I' s | S.length s >= 4 = fromIntegral (i :: Word32) + where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord32LE + getInt 'i' s | S.length s >= 4 = fromIntegral (fromIntegral i :: Int32) + where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord32LE + getInt _ _ = 0 + + getFloat s = unsafeDupablePerformIO $ alloca $ \buf -> + pokeByteOff buf 0 (getInt 'I' s :: Word32) >> peek buf + + +isPaired, isProperlyPaired, isUnmapped, isMateUnmapped, isReversed, + isMateReversed, isFirstMate, isSecondMate, isAuxillary, isSecondary, + isFailsQC, isDuplicate, isSupplementary, + isTrimmed, isMerged, isAlternative, isExactIndex :: BamRec -> Bool + +isPaired = flip testBit 0 . b_flag +isProperlyPaired = flip testBit 1 . b_flag +isUnmapped = flip testBit 2 . b_flag +isMateUnmapped = flip testBit 3 . b_flag +isReversed = flip testBit 4 . b_flag +isMateReversed = flip testBit 5 . b_flag +isFirstMate = flip testBit 6 . b_flag +isSecondMate = flip testBit 7 . b_flag +isAuxillary = flip testBit 8 . b_flag +isSecondary = flip testBit 8 . b_flag +isFailsQC = flip testBit 9 . b_flag +isDuplicate = flip testBit 10 . b_flag +isSupplementary = flip testBit 11 . b_flag + +isTrimmed = flip testBit 0 . extAsInt 0 "FF" +isMerged = flip testBit 1 . extAsInt 0 "FF" +isAlternative = flip testBit 2 . extAsInt 0 "FF" +isExactIndex = flip testBit 3 . extAsInt 0 "FF" + +type_mask :: Int +type_mask = flagFirstMate .|. flagSecondMate .|. flagPaired + +extAsInt :: Int -> BamKey -> BamRec -> Int +extAsInt d nm br = case lookup nm (b_exts br) of Just (Int i) -> i ; _ -> d + +extAsString :: BamKey -> BamRec -> Bytes +extAsString nm br = case lookup nm (b_exts br) of + Just (Char c) -> B.singleton c + Just (Text s) -> s + _ -> B.empty + +setQualFlag :: Char -> BamRec -> BamRec +setQualFlag c br = br { b_exts = updateE "ZQ" (Text s') $ b_exts br } + where + s = extAsString "ZQ" br + s' = if c `S.elem` s then s else c `S.cons` s + diff --git a/Bio/Bam/Regions.hs b/Bio/Bam/Regions.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Regions.hs @@ -0,0 +1,53 @@ +module Bio.Bam.Regions + ( Region(..) + , Regions(..) + , Subsequence(..) + + , toList + , fromList + , overlaps + ) where + +import Bio.Bam.Header ( Refseq(..) ) +import Bio.Prelude hiding ( toList ) + +import qualified Data.IntMap.Strict as IM + +data Region = Region { refseq :: !Refseq, start :: !Int, end :: !Int } + deriving (Eq, Ord, Show) + +-- | A subset of a genome. The idea is to map the reference sequence +-- (represented by its number) to a 'Subseqeunce'. +newtype Regions = Regions (IntMap Subsequence) deriving Show + +-- | A mostly contiguous subset of a sequence, stored as a set of +-- non-overlapping intervals in an 'IntMap' from start position to end +-- position (half-open intervals, naturally). +newtype Subsequence = Subsequence (IntMap Int) deriving Show + +toList :: Regions -> [(Refseq, Subsequence)] +toList (Regions m) = [ (Refseq $ fromIntegral k, v) | (k,v) <- IM.toList m ] + +fromList :: [Region] -> Regions +fromList = foldl' (flip add) (Regions IM.empty) + +add :: Region -> Regions -> Regions +add (Region (Refseq r) b e) (Regions m) = + let single = Just . Subsequence $ IM.singleton b e + in Regions $ IM.alter (maybe single (Just . addInt b e)) (fromIntegral r) m + + +addInt :: Int -> Int -> Subsequence -> Subsequence +addInt b e (Subsequence m0) = Subsequence $ merge_into b e m0 + where + merge_into x y m = case IM.lookupLT y m of + Just (u,v) | x < u && y <= v -> merge_into x v $ IM.delete u m -- extend to the left + | x < u -> merge_into x y $ IM.delete u m -- subsume + | y <= v -> m -- subsumed + | x <= v -> merge_into u y $ IM.delete u m -- extend to the right + _ -> IM.insert x y m -- no overlap + +overlaps :: Int -> Int -> Subsequence -> Bool +overlaps b e (Subsequence m) = case IM.lookupLT e m of + Just (_,v) -> b < v + Nothing -> False diff --git a/Bio/Bam/Rmdup.hs b/Bio/Bam/Rmdup.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Rmdup.hs @@ -0,0 +1,691 @@ +module Bio.Bam.Rmdup( + rmdup, Collapse, cons_collapse, cheap_collapse, + cons_collapse_keep, cheap_collapse_keep, + check_sort, normalizeTo, wrapTo, + ECig(..), toECig, setMD, toCigar + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude hiding ( left, right ) +import Bio.Streaming + +import qualified Data.ByteString as B +import qualified Data.ByteString.Char8 as T +import qualified Data.Map.Strict as M +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Storable as W +import qualified Data.Vector.Unboxed as U +import qualified Streaming.Prelude as Q + +data Collapse = Collapse { + collapse :: [BamRec] -> (Decision,[BamRec]), -- cluster to consensus and stuff or representative and stuff + originals :: [BamRec] -> [BamRec] } -- treatment of the redundant original reads + +data Decision = Consensus { fromDecision :: BamRec } + | Representative { fromDecision :: BamRec } + +cons_collapse :: Qual -> Collapse +cons_collapse maxq = Collapse (do_collapse maxq) (const []) + +cons_collapse_keep :: Qual -> Collapse +cons_collapse_keep maxq = Collapse (do_collapse maxq) (map (\b -> b { b_flag = b_flag b .|. flagDuplicate })) + +cheap_collapse :: Collapse +cheap_collapse = Collapse do_cheap_collapse (const []) + +cheap_collapse_keep :: Collapse +cheap_collapse_keep = Collapse do_cheap_collapse (map (\b -> b { b_flag = b_flag b .|. flagDuplicate })) + + +-- | Removes duplicates from an aligned, sorted BAM stream. +-- +-- The incoming stream must be sorted by coordinate, and we check for +-- violations of that assumption. We cannot assume that length was +-- taken into account when sorting (samtools doesn't do so), so +-- duplicates may be separated by reads that start at the same position +-- but have different length or different strand. +-- +-- We are looking at three different kinds of reads: paired reads, true +-- single ended reads, merged or trimmed reads. They are somewhat +-- different, but here's the situation if we wanted to treat them +-- separately. These conditions define a set of duplicates: +-- +-- Merged or trimmed: We compare the leftmost coordinates and the +-- aligned length. If the library prep is strand-preserving, we also +-- compare the strand. +-- +-- Paired: We compare both left-most coordinates (b_pos and b_mpos). If +-- the library prep is strand-preserving, only first-mates can be +-- duplicates of first-mates. Else a first-mate can be the duplicate of +-- a second-mate. There may be pairs with one unmapped mate. This is +-- not a problem as they get assigned synthetic coordinates and will be +-- handled smoothly. +-- +-- True singles: We compare only the leftmost coordinate. It does not +-- matter if the library prep is strand-preserving, the strand always +-- matters. +-- +-- Across these classes, we can see more duplicates: +-- +-- Merged/trimmed and paired: these can be duplicates if the merging +-- failed for the pair. We would need to compare the outer coordinates +-- of the merged reads to the two 5' coordinates of the pair. However, +-- since we don't have access to the mate, we cannot actually do +-- anything right here. This case should be solved externally by +-- merging those pairs that overlap in coordinate space. +-- +-- Single and paired: in the single case, we only have one coordinate +-- to compare. This will inevitably lead to trouble, as we could find +-- that the single might be the duplicate of two pairs, but those two +-- pairs are definitely not duplicates of each other. We solve it by +-- removing the single read(s). +-- +-- Single and merged/trimmed: same trouble as in the single+paired +-- case. We remove the single to solve it. +-- +-- +-- In principle, we might want to allow some wiggle room in the +-- coordinates. So far, this has not been implemented. It adds the +-- complication that groups of separated reads can turn into a set of +-- duplicates because of the appearance of a new reads. Needs some +-- thinking about... or maybe it's not too important. +-- +-- Once a set of duplicates is collected, we perform a majority vote on +-- the correct CIGAR line. Of all those reads that agree on this CIGAR +-- line, a consensus is called, quality scores are adjusted and clamped +-- to a maximum, the MD field is updated and the XP field is assigned +-- the number of reads in the original cluster. The new MAPQ becomes +-- the RMSQ of the map qualities of all reads. +-- +-- Treatment of Read Groups: We generalize by providing a "label" +-- function; only reads that have the same label are considered +-- duplicates of each other. The typical label function would extract +-- read groups, libraries or samples. + +rmdup :: (Monad m, Ord l) => (BamRec -> l) -> Bool -> Collapse + -> Q.Stream (Of BamRec) m r -> Q.Stream (Of (Int,BamRec)) m r +rmdup label strand_preserved collapse_cfg = + -- Easiest way to go about this: We simply collect everything that + -- starts at some specific coordinate and group it appropriately. + -- Treat the groups separately, output, go on. + check_sort (b_cpos . snd) "internal error, output isn't sorted anymore" . + mapGroups ( sortBy (comparing $ V.length . b_seq . snd) + . concatMap (do_rmdup strand_preserved collapse_cfg) + . M.elems . accumMap label id ) . + check_sort b_cpos "input must be sorted for rmdup to work" + where + b_cpos = b_rname &&& b_pos + + mapGroups f = lift . inspect >=> either pure (\(a :> s) -> mg1 f [] a s) + + mg1 f acc a s = + lift (inspect s) >>= \case + Left r -> Q.each (f (a:acc)) >> pure r + Right (b :> s') + | b_cpos a == b_cpos b -> mg1 f (b:acc) a s' + | otherwise -> Q.each (f (a:acc)) >> mg1 f [] b s' + +check_sort :: (Monad m, Ord b) => (a -> b) -> String -> Stream (Of a) m r -> Stream (Of a) m r +check_sort f msg = lift . inspect >=> either pure step + where + step (a :> s) = lift (inspect s) >>= either (Q.cons a . pure) (step' a) + step' a (b :> s) + | f a > f b = fail $ "rmdup: " ++ msg + | otherwise = Q.cons a $ step (b :> s) +{-# INLINE check_sort #-} + +{- | Workhorse for duplicate removal. + + - Unmapped fragments should not be considered to be duplicates of + mapped fragments. The /unmapped/ flag can serve for that: while + there are two classes of /unmapped/ reads (those that are not mapped + and those that are mapped to an invalid position), the two sets will + always have different coordinates. (Unfortunately, correct duplicate + removal now relies on correct /unmapped/ and /mate unmapped/ flags, + and we don't get them from unmodified BWA. So correct operation + requires patched BWA or a run of @bam-fixpair@.) + + (1) Other definitions (e.g. lack of CIGAR) don't work, because that + information won't be available for the mate. + + (2) This would amount to making the /unmapped/ flag part of the + coordinate, but samtools is not going to take it into account + when sorting. + + (3) Instead, both flags become part of the /mate pos/ grouping + criterion. + + - First Mates should (probably) not be considered duplicates of Second + Mates. This is unconditionally true for libraries with A\/B-style + adapters (definitely 454, probably Mathias' ds protocol) and the ss + protocol, it is not true for fork adapter protocols (vanilla Illumina + protocol). So it has to be an option, which would ideally be derived + from header information. + + - This code ignores read groups, but it will do a majority vote on the + @RG@ field and call consensi for the index sequences. If you believe + that duplicates across read groups are impossible, you must call it + with an appropriately filtered stream. + + - Half-Aligned Pairs (meaning one known coordinate, while the validity + of the alignments is immaterial) are rather complicated: + + (1) Given that only one coordinate is known (5' of the aligned mate), + we want to treat them like true singles. But the unaligned mate + should be kept if possible, though it should not contribute to a + consensus sequence. We assume nothing about the unaligned mate, + not even that it /shouldn't/ be aligned, never mind the fact that + it /couldn't/ be. (The difference is in the finite abilities of + real world aligners, naturally.) + + (2) Therefore, aligned reads with unaligned mates go to the same + potential duplicate set as true singletons. If at least one pair + exists that might be a duplicate of those, all singletons and + half-aligned mates are discarded. Else a consensus is computed + and replaces the aligned mates. + + (3) The unaligned mates end up in the same place in a BAM stream as + the aligned mates (therefore we see them and can treat them + locally). We cannot call a consensus, since these molecules may + well have different length, so we select one. It doesn't really + matter which one is selected, and since we're treating both mates + at the same time, it doesn't even need to be reproducible without + local information. This is made to be the mate of the consensus. + + (4) See 'merge_singles' for how it's actually done. +-} + +do_rmdup :: Bool -> Collapse -> [BamRec] -> [(Int,BamRec)] +do_rmdup strand_preserved Collapse{..} rds = + results ++ map ((,) 0) (originals (leftovers ++ r1 ++ r2 ++ r3)) + where + (results, leftovers) = merge_singles singles' unaligned' $ + [ (str, second fromDecision b) | ((_,str ),b) <- M.toList merged' ] ++ + [ (str, second fromDecision b) | ((_,str,_),b) <- M.toList pairs' ] + + (raw_pairs, raw_singles) = partition isPaired rds + (merged, true_singles) = partition (liftA2 (||) isMerged isTrimmed) raw_singles + + (pairs, raw_half_pairs) = partition b_totally_aligned raw_pairs + (half_unaligned, half_aligned) = partition isUnmapped raw_half_pairs + + mkMap :: Ord a => (BamRec -> a) -> [BamRec] -> (M.Map a (Int,Decision), [BamRec]) + mkMap f x = let m1 = M.map (length &&& collapse) $ accumMap f id x + in (M.map (second fst) m1, concatMap (snd.snd) $ M.elems m1) + + (pairs',r1) = mkMap (\b -> (b_mate_pos b, b_strand b, b_mate b)) pairs + (merged',r2) = mkMap (\b -> (alignedLength (b_cigar b), b_strand b)) merged + (singles',r3) = mkMap b_strand (true_singles++half_aligned) + unaligned' = accumMap b_strand id half_unaligned + + b_strand b = strand_preserved && isReversed b + b_mate b = strand_preserved && isFirstMate b + + +-- | Merging information about true singles, merged singles, +-- half-aligned pairs, actually aligned pairs. +-- +-- We collected aligned reads with unaligned mates together with aligned +-- true singles (@singles@). We collected the unaligned mates, which +-- necessarily have the exact same alignment coordinates, separately +-- (@unaligned@). If we don't find a matching true pair (that case is +-- already handled smoothly), we keep the highest quality unaligned +-- mate, pair it with the consensus of the aligned mates and aligned +-- singletons, and give it the lexically smallest name of the +-- half-aligned pairs. + +-- NOTE: I need to decide when to run 'make_singleton'. Basically, +-- when we call a consensus for half-aligned pairs and keep +-- everything(?). Then we don't have a mate for the consensus... though +-- we could decide to duplicate one mate read to get it. + +merge_singles :: M.Map Bool (Int,Decision) -- strand --> true singles & half aligned + -> M.Map Bool [BamRec] -- strand --> half unaligned + -> [ (Bool, (Int, BamRec)) ] -- strand --> paireds & mergeds + -> ([(Int,BamRec)],[BamRec]) -- results, leftovers + +merge_singles singles unaligneds = go + where + -- Say we generated a consensus or passed something through. If + -- there is a singleton consensus with the same strand, we should + -- add in its XP field and discard it. If there is a singleton + -- representative, we add in its XP field and put it into the + -- leftovers. If there is unaligned stuff here that has the same + -- strand, it goes to the leftovers. + go ( (str, (m,v)) : paireds) = + let (r,l) = merge_singles (M.delete str singles) (M.delete str unaligneds) paireds + unal = M.findWithDefault [] str unaligneds ++ l + + in case M.lookup str singles of + Nothing -> ( (m,v) : r, unal ) + Just (n, Consensus w) -> ( (n, add_xp_of w v) : r, unal ) + Just (n, Representative w) -> ( (n, add_xp_of w v) : r, w : unal ) + + -- No more pairs, delegate the problem + go [] = merge_halves unaligneds (M.toList singles) + + add_xp_of w v = v { b_exts = updateE "XP" (Int $ extAsInt 1 "XP" w `oplus` extAsInt 1 "XP" v) (b_exts v) } + +-- | Merging of half-aligned reads. The first argument is a map of +-- unaligned reads (their mates are aligned to the current position), +-- the second is a list of reads that are aligned (their mates are not +-- aligned). +-- +-- So, suppose we're looking at a 'Representative' that was passed +-- through. We need to emit it along with its mate, which may be hidden +-- inside a list. (Alternatively, we could force it to single, but that +-- fails if we're passing everything along somehow.) +-- +-- Suppose we're looking at a 'Consensus'. We could pair it with some +-- mate (which we'd need to duplicate), or we could turn it into a +-- singleton. Duplication is ugly, so in this case, we force it to +-- singleton. + +merge_halves :: M.Map Bool [BamRec] -- strand --> half unaligned + -> [(Bool, (Int,Decision))] -- strand --> true singles & half aligned + -> ([(Int,BamRec)],[BamRec]) -- results, leftovers + +-- Emitting a consensus: make it a single. Nothing goes to leftovers; +-- we may still need it for something else to be emitted. (While that +-- would be strange, making sure the BAM file stays completely valid is +-- probably better.) +merge_halves unaligneds ((_, (n, Consensus v)) : singles) = ( (n, v { b_flag = b_flag v .&. complement pflags }) : r, l ) + where + (r,l) = merge_halves unaligneds singles + pflags = flagPaired .|. flagProperlyPaired .|. flagMateUnmapped .|. flagMateReversed .|. flagFirstMate .|. flagSecondMate + + +-- Emitting a representative: find the mate in the list of unaligned +-- reads (take up to one match to be robust), and emit that, too, as a +-- result. Everything else goes to leftovers. If the representative +-- happens to be unpaired, no mate is found and that case therefore is +-- handled smoothly. +merge_halves unaligneds ((str, (n, Representative v)) : singles) = ((n,v) : map ((,)1) (take 1 same) ++ r, drop 1 same ++ diff ++ l) + where + (r,l) = merge_halves (M.delete str unaligneds) singles + (same,diff) = partition (is_mate_of v) $ M.findWithDefault [] str unaligneds + is_mate_of a b = b_qname a == b_qname b && isPaired a && isPaired b && isFirstMate a == isSecondMate b + +-- No more singles, all unaligneds are leftovers. +merge_halves unaligneds [] = ( [], concat $ M.elems unaligneds ) + + + + +type MPos = (Refseq, Int, Bool, Bool) + +b_mate_pos :: BamRec -> MPos +b_mate_pos b = (b_mrnm b, b_mpos b, isUnmapped b, isMateUnmapped b) + +b_totally_aligned :: BamRec -> Bool +b_totally_aligned b = not (isUnmapped b || isMateUnmapped b) + + +accumMap :: Ord k => (a -> k) -> (a -> v) -> [a] -> M.Map k [v] +accumMap f g = go M.empty + where + go m [ ] = m + go m (a:as) = let ws = M.findWithDefault [] (f a) m ; g' = g a + in g' `seq` go (M.insert (f a) (g':ws) m) as + + +{- We need to deal sensibly with each field, but different fields have + different needs. We can take the value from the first read to + preserve determinism or because all reads should be equal anyway, + aggregate over all reads computing either RMSQ or the most common + value, delete a field because it wouldn't make sense anymore or + because doing something sensible would be hard and we're going to + ignore it anyway, or we calculate some special value; see below. + Unknown fields will be taken from the first read, which seems to be a + safe default. + + QNAME and most fields taken from first + FLAG qc fail majority vote + dup deleted + MAPQ rmsq + CIGAR, SEQ, QUAL, MD, NM, XP generated + XA concatenate all + + BQ, CM, FZ, Q2, R2, XM, XO, XG, YQ, EN + deleted because they would become wrong + + CQ, CS, E2, FS, OQ, OP, OC, U2, H0, H1, H2, HI, NH, IH, ZQ + delete because they will be ignored anyway + + AM, AS, MQ, PQ, SM, UQ + compute rmsq + + X0, X1, XT, XS, XF, XE, BC, LB, RG, XI, YI, XJ, YJ + majority vote -} + +do_collapse :: Qual -> [BamRec] -> (Decision, [BamRec]) +do_collapse maxq [br] | V.all (<= maxq) (b_qual br) = ( Representative br, [ ] ) -- no modifcation, pass through + | otherwise = ( Consensus lq_br, [br] ) -- qualities reduced, must keep original + where + lq_br = br { b_qual = V.map (min maxq) $ b_qual br + , b_virtual_offset = 0 + , b_qname = b_qname br `B.snoc` c2w 'c' } + +do_collapse maxq brs = ( Consensus b0 { b_exts = modify_extensions $ b_exts b0 + , b_flag = failflag .&. complement flagDuplicate + , b_mapq = Q $ rmsq $ map (unQ . b_mapq) $ good brs + , b_cigar = cigar' + , b_seq = V.fromList $ map fst cons_seq_qual + , b_qual = V.fromList $ map snd cons_seq_qual + , b_qname = b_qname b0 `B.snoc` 99 + , b_virtual_offset = 0 }, brs ) -- many modifications, must keep everything + where + !b0 = minimumBy (comparing b_qname) brs + !most_fail = 2 * length (filter isFailsQC brs) > length brs + !failflag | most_fail = b_flag b0 .|. flagFailsQC + | otherwise = b_flag b0 .&. complement flagFailsQC + + rmsq xs = case foldl' (\(!n,!d) x -> (n + fromIntegral x * fromIntegral x, d + 1)) (0,0) xs of + (!n,!d) -> round $ sqrt $ (n::Double) / fromIntegral (d::Int) + + maj xs = head . maximumBy (comparing length) . group . sort $ xs + nub' = concatMap head . group . sort + + -- majority vote on the cigar lines, then filter + !cigar' = maj $ map b_cigar brs + good = filter ((==) cigar' . b_cigar) + + cons_seq_qual = [ consensus maxq [ (V.unsafeIndex (b_seq b) i, q) + | b <- good brs, let q = if V.null (b_qual b) then Q 23 else b_qual b V.! i ] + | i <- [0 .. len - 1] ] + where !len = V.length . b_seq . head $ good brs + + md' = case [ (b_seq b,md,b) | b <- good brs, Just md <- [ getMd b ] ] of + [ ] -> [] + (seq1, md1,b) : _ -> case mk_new_md' [] (V.toList cigar') md1 (V.toList seq1) (map fst cons_seq_qual) of + Right x -> x + Left (MdFail cigs ms osq nsq) -> error $ unlines + [ "Broken MD field when trying to construct new MD!" + , "QNAME: " ++ show (b_qname b) + , "POS: " ++ shows (unRefseq (b_rname b)) ":" ++ show (b_pos b) + , "CIGAR: " ++ show cigs + , "MD: " ++ show ms + , "refseq: " ++ show osq + , "readseq: " ++ show nsq ] + + + nm' = sum $ [ n | Ins :* n <- W.toList cigar' ] ++ [ n | Del :* n <- W.toList cigar' ] ++ [ 1 | MdRep _ <- md' ] + xa' = nub' [ T.split ';' xas | Just (Text xas) <- map (lookup "XA" . b_exts) brs ] + + modify_extensions es = foldr ($!) es $ + [ let vs = mapMaybe (lookup k . b_exts) brs + in if null vs then id else updateE k $! maj vs | k <- do_maj ] ++ + [ let vs = [ v | Just (Int v) <- map (lookup k . b_exts) brs ] + in if null vs then id else updateE k $! Int (rmsq vs) | k <- do_rmsq ] ++ + map deleteE useless ++ + [ updateE "NM" $! Int nm' + , updateE "XP" $! Int (foldl' (\a b -> a `oplus` extAsInt 1 "XP" b) 0 brs) + , if null xa' then id else updateE "XA" $! (Text $ T.intercalate (T.singleton ';') xa') + , if null md' then id else updateE "MD" $! (Text $ showMd md') ] + + useless = map fromString $ words "BQ CM FZ Q2 R2 XM XO XG YQ EN CQ CS E2 FS OQ OP OC U2 H0 H1 H2 HI NH IH ZQ" + do_rmsq = map fromString $ words "AM AS MQ PQ SM UQ" + do_maj = map fromString $ words "X0 X1 XT XS XF XE BC LB RG XI XJ YI YJ" + +minViewBy :: (a -> a -> Ordering) -> [a] -> (a,[a]) +minViewBy _ [ ] = error "minViewBy on empty list" +minViewBy cmp (x:xs) = go x [] xs + where + go m acc [ ] = (m,acc) + go m acc (a:as) = case m `cmp` a of GT -> go a (m:acc) as + _ -> go m (a:acc) as + +data MdFail = MdFail [Cigar] [MdOp] [Nucleotides] [Nucleotides] + +mk_new_md' :: [MdOp] -> [Cigar] -> [MdOp] -> [Nucleotides] -> [Nucleotides] -> Either MdFail [MdOp] +mk_new_md' acc [] [] [] [] = Right $ normalize [] acc + where + normalize a (MdNum 0:os) = normalize a os + normalize (MdNum n:a) (MdNum m:os) = normalize (MdNum (n+m):a) os + normalize a (MdDel []:os) = normalize a os + normalize (MdDel u:a) (MdDel v:os) = normalize (MdDel (v++u):a) os + normalize a ( o:os) = normalize (o:a) os + normalize a [ ] = a + +mk_new_md' acc ( _ :* 0 : cigs) mds osq nsq = mk_new_md' acc cigs mds osq nsq +mk_new_md' acc cigs (MdNum 0 : mds) osq nsq = mk_new_md' acc cigs mds osq nsq +mk_new_md' acc cigs (MdDel [] : mds) osq nsq = mk_new_md' acc cigs mds osq nsq + +mk_new_md' acc (Mat :* u : cigs) (MdRep b : mds) (_:osq) (n:nsq) + | b == n = mk_new_md' (MdNum 1 : acc) (Mat :* (u-1):cigs) mds osq nsq + | otherwise = mk_new_md' (MdRep b : acc) (Mat :* (u-1):cigs) mds osq nsq + +mk_new_md' acc (Mat :* u : cigs) (MdNum v : mds) (o:osq) (n:nsq) + | o == n = mk_new_md' (MdNum 1 : acc) (Mat :* (u-1):cigs) (MdNum (v-1) : mds) osq nsq + | otherwise = mk_new_md' (MdRep o : acc) (Mat :* (u-1):cigs) (MdNum (v-1) : mds) osq nsq + +mk_new_md' acc (Del :* n : cigs) (MdDel bs : mds) osq nsq | n == length bs = mk_new_md' (MdDel bs : acc) cigs mds osq nsq +mk_new_md' acc (Del :* n : cigs) (MdDel (b:bs) : mds) osq nsq = mk_new_md' (MdDel [b] : acc) (Del :* (n-1) : cigs) (MdDel bs:mds) osq nsq +mk_new_md' acc (Del :* n : cigs) (MdRep b : mds) osq nsq = mk_new_md' (MdDel [b] : acc) (Del :* (n-1) : cigs) mds osq nsq +mk_new_md' acc (Del :* n : cigs) (MdNum m : mds) osq nsq = mk_new_md' (MdDel [nucsN] : acc) (Del :* (n-1) : cigs) (MdNum (m-1):mds) osq nsq + +mk_new_md' acc (Ins :* n : cigs) md osq nsq = mk_new_md' acc cigs md (drop n osq) (drop n nsq) +mk_new_md' acc (SMa :* n : cigs) md osq nsq = mk_new_md' acc cigs md (drop n osq) (drop n nsq) +mk_new_md' acc (HMa :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq +mk_new_md' acc (Pad :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq +mk_new_md' acc (Nop :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq + +mk_new_md' _acc cigs ms osq nsq = Left $ MdFail cigs ms osq nsq + +consensus :: Qual -> [ (Nucleotides, Qual) ] -> (Nucleotides, Qual) +consensus (Q maxq) nqs = + let accs :: U.Vector Int + accs = U.accum (+) (U.replicate 16 0) [ (fromIntegral n, fromIntegral q) | (Ns n,Q q) <- nqs ] + in case sortBy (flip $ comparing snd) $ zip [Ns 0 ..] $ U.toList accs of + (n0,q0) : (_,q1) : _ -> + let qr = fromIntegral $ (q0-q1) `min` fromIntegral maxq + in if qr > 3 then (n0, Q qr) else (nucsN, Q 0) + _ -> error "can't happen" + + +-- Cheap version: simply takes the lexically first record, adds XP field +do_cheap_collapse :: [BamRec] -> ( Decision, [BamRec] ) +do_cheap_collapse [b] = ( Representative b, [] ) +do_cheap_collapse bs = ( Representative $ replaceXP new_xp b0, bx ) + where + (b0, bx) = minViewBy (comparing b_qname) bs + new_xp = foldl' (\a b -> a `oplus` extAsInt 1 "XP" b) 0 bs + +replaceXP :: Int -> BamRec -> BamRec +replaceXP new_xp b0 = b0 { b_exts = updateE "XP" (Int new_xp) $ b_exts b0 } + +oplus :: Int -> Int -> Int +_ `oplus` (-1) = -1 +(-1) `oplus` _ = -1 +a `oplus` b = a + b + +-- | Normalize a read's alignment to fall into the canonical region +-- of [0..l]. Takes the name of the reference sequence and its length. +-- Returns @Left x@ if the coordinate decreased so the result is out of +-- order now, @Right x@ if the coordinate is unchanged. +normalizeTo :: Bytes -> Int -> BamRec -> Either BamRec BamRec +normalizeTo nm l b = lr $ b { b_pos = b_pos b `mod` l + , b_mpos = b_mpos b `mod` l + , b_mapq = if dups_are_fine then Q 37 else b_mapq b + , b_exts = if dups_are_fine then deleteE "XA" (b_exts b) else b_exts b } + where + lr = if b_pos b >= l then Left else Right + dups_are_fine = all_match_XA (extAsString "XA" b) + all_match_XA s = case T.split ';' s of [xa1, xa2] | T.null xa2 -> one_match_XA xa1 + [xa1] -> one_match_XA xa1 + _ -> False + one_match_XA s = case T.split ',' s of (sq:pos:_) | sq == nm -> pos_match_XA pos ; _ -> False + pos_match_XA s = case T.readInt s of Just (p,z) | T.null z -> int_match_XA p ; _ -> False + int_match_XA p | p >= 0 = (p-1) `mod` l == b_pos b `mod` l && not (isReversed b) + | otherwise = (-p-1) `mod` l == b_pos b `mod` l && isReversed b + + +-- | Wraps a read to be fully contained in the canonical interval +-- [0..l]. If the read overhangs, it is duplicated and both copies are +-- suitably masked. A piece with changed coordinate that is now out of +-- order is returned as @Left x@, if the order is fine, it is returned +-- as @Right x@. +wrapTo :: Int -> BamRec -> [Either BamRec BamRec] +wrapTo l b = if overhangs then do_wrap else [Right b] + where + overhangs = not (isUnmapped b) && b_pos b < l && l < b_pos b + alignedLength (b_cigar b) + + do_wrap = case split_ecig (l - b_pos b) $ toECig (b_cigar b) (fromMaybe [] $ getMd b) of + (left,right) -> [ Right $ b { b_cigar = toCigar left } `setMD` left + , Left $ b { b_cigar = toCigar right, b_pos = 0 } `setMD` right ] + +-- | Split an 'ECig' into two at some position. The position is counted +-- in terms of the reference (therefore, deletions count, insertions +-- don't). The parts that would be skipped if we were splitting lists +-- are replaced by soft masks. +split_ecig :: Int -> ECig -> (ECig, ECig) +split_ecig _ WithMD = (WithMD, WithMD) +split_ecig _ WithoutMD = (WithoutMD, WithoutMD) +split_ecig 0 ecs = (mask_all ecs, ecs) + +split_ecig i (Ins' n ecs) = case split_ecig i ecs of (u,v) -> (Ins' n u, SMa' n v) +split_ecig i (SMa' n ecs) = case split_ecig i ecs of (u,v) -> (SMa' n u, SMa' n v) +split_ecig i (HMa' n ecs) = case split_ecig i ecs of (u,v) -> (HMa' n u, HMa' n v) +split_ecig i (Pad' n ecs) = case split_ecig i ecs of (u,v) -> (Pad' n u, v) + +split_ecig i (Mat' n ecs) + | i >= n = case split_ecig (i-n) ecs of (u,v) -> (Mat' n u, SMa' n v) + | otherwise = (Mat' i $ SMa' (n-i) $ mask_all ecs, SMa' i $ Mat' (n-i) ecs) + +split_ecig i (Rep' x ecs) = case split_ecig (i-1) ecs of (u,v) -> (Rep' x u, SMa' 1 v) +split_ecig i (Del' x ecs) = case split_ecig (i-1) ecs of (u,v) -> (Del' x u, v) + +split_ecig i (Nop' n ecs) + | i >= n = case split_ecig (i-n) ecs of (u,v) -> (Nop' n u, v) + | otherwise = (Nop' i $ mask_all ecs, Nop' (n-i) ecs) + +mask_all :: ECig -> ECig +mask_all WithMD = WithMD +mask_all WithoutMD = WithoutMD +mask_all (Nop' _ ec) = mask_all ec +mask_all (HMa' _ ec) = mask_all ec +mask_all (Pad' _ ec) = mask_all ec +mask_all (Del' _ ec) = mask_all ec +mask_all (Rep' _ ec) = SMa' 1 $ mask_all ec +mask_all (Mat' n ec) = SMa' n $ mask_all ec +mask_all (Ins' n ec) = SMa' n $ mask_all ec +mask_all (SMa' n ec) = SMa' n $ mask_all ec + +-- | Extended CIGAR. This subsumes both the CIGAR string and the +-- optional MD field. If we have MD on input, we generate it on output, +-- too. And in between, we break everything into /very small/ +-- operations. + +data ECig = WithMD -- terminate, do generate MD field + | WithoutMD -- terminate, don't bother with MD + | Mat' Int ECig + | Rep' Nucleotides ECig + | Ins' Int ECig + | Del' Nucleotides ECig + | Nop' Int ECig + | SMa' Int ECig + | HMa' Int ECig + | Pad' Int ECig + + +toECig :: W.Vector Cigar -> [MdOp] -> ECig +toECig = go . W.toList + where + go cs (MdNum 0:mds) = go cs mds + go cs (MdDel []:mds) = go cs mds + go (_:*0 :cs) mds = go cs mds + go [ ] [ ] = WithMD -- all was fine to the very end + go [ ] _ = WithoutMD -- here it wasn't fine + + go (Mat :* n : cs) (MdRep x:mds) = Rep' x $ go (Mat :* (n-1) : cs) mds + go (Mat :* n : cs) (MdNum m:mds) + | n < m = Mat' n $ go cs (MdNum (m-n):mds) + | n > m = Mat' m $ go (Mat :* (n-m) : cs) mds + | otherwise = Mat' n $ go cs mds + go (Mat :* n : cs) _ = Mat' n $ go' cs + + go (Ins :* n : cs) mds = Ins' n $ go cs mds + go (Del :* n : cs) (MdDel (x:xs):mds) = Del' x $ go (Del :* (n-1) : cs) (MdDel xs:mds) + go (Del :* n : cs) _ = Del' nucsN $ go' (Del :* (n-1) : cs) + + go (Nop :* n : cs) mds = Nop' n $ go cs mds + go (SMa :* n : cs) mds = SMa' n $ go cs mds + go (HMa :* n : cs) mds = HMa' n $ go cs mds + go (Pad :* n : cs) mds = Pad' n $ go cs mds + + -- We jump here once the MD fiels ran out early or was messed up. + -- We no longer bother with it (this also happens if the MD isn't + -- present to begin with). + go' (_ :* 0 : cs) = go' cs + go' [ ] = WithoutMD -- we didn't have MD or it was broken + + go' (Mat :* n : cs) = Mat' n $ go' cs + go' (Ins :* n : cs) = Ins' n $ go' cs + go' (Del :* n : cs) = Del' nucsN $ go' (Del :* (n-1) : cs) + + go' (Nop :* n : cs) = Nop' n $ go' cs + go' (SMa :* n : cs) = SMa' n $ go' cs + go' (HMa :* n : cs) = HMa' n $ go' cs + go' (Pad :* n : cs) = Pad' n $ go' cs + + +-- We normalize matches, deletions and soft masks, because these are the +-- operations we generate. Everything else is either already normalized +-- or nobody really cares anyway. +toCigar :: ECig -> W.Vector Cigar +toCigar = V.fromList . go + where + go WithMD = [] + go WithoutMD = [] + + go (Ins' n ecs) = Ins :* n : go ecs + go (Nop' n ecs) = Nop :* n : go ecs + go (HMa' n ecs) = HMa :* n : go ecs + go (Pad' n ecs) = Pad :* n : go ecs + go (SMa' n ecs) = go_sma n ecs + go (Mat' n ecs) = go_mat n ecs + go (Rep' _ ecs) = go_mat 1 ecs + go (Del' _ ecs) = go_del 1 ecs + + go_sma !n (SMa' m ecs) = go_sma (n+m) ecs + go_sma !n ecs = SMa :* n : go ecs + + go_mat !n (Mat' m ecs) = go_mat (n+m) ecs + go_mat !n (Rep' _ ecs) = go_mat (n+1) ecs + go_mat !n ecs = Mat :* n : go ecs + + go_del !n (Del' _ ecs) = go_del (n+1) ecs + go_del !n ecs = Del :* n : go ecs + + + +-- | Create an MD field from an extended CIGAR and place it in a record. +-- We build it piecemeal (in 'go'), call out to 'addNum', 'addRep', +-- 'addDel' to make sure the operations are not generated in a +-- degenerate manner, and finally check if we're even supposed to create +-- an MD field. +setMD :: BamRec -> ECig -> BamRec +setMD b ec = case go ec of + Just md -> b { b_exts = updateE "MD" (Text $ showMd md) (b_exts b) } + Nothing -> b { b_exts = deleteE "MD" (b_exts b) } + where + go WithMD = Just [] + go WithoutMD = Nothing + + go (Ins' _ ecs) = go ecs + go (Nop' _ ecs) = go ecs + go (SMa' _ ecs) = go ecs + go (HMa' _ ecs) = go ecs + go (Pad' _ ecs) = go ecs + go (Mat' n ecs) = (if n == 0 then id else fmap (addNum n)) $ go ecs + go (Rep' x ecs) = (if isGap x then id else fmap (addRep x)) $ go ecs + go (Del' x ecs) = (if isGap x then id else fmap (addDel x)) $ go ecs + + addNum n (MdNum m : mds) = MdNum (n+m) : mds + addNum n mds = MdNum n : mds + + addRep x mds = MdRep x : mds + + addDel x (MdDel y : mds) = MdDel (x:y) : mds + addDel x mds = MdDel [x] : mds diff --git a/Bio/Bam/Trim.hs b/Bio/Bam/Trim.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Trim.hs @@ -0,0 +1,442 @@ +-- | Trimming of reads as found in BAM files. Implements trimming low +-- quality sequence from the 3' end. + +module Bio.Bam.Trim + ( trim_3 + , trim_3' + , trim_low_quality + , default_fwd_adapters + , default_rev_adapters + , find_merge + , mergeBam + , find_trim + , trimBam + , mergeTrimBam + , twoMins + , merged_seq + , merged_qual + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Bam.Rmdup ( ECig(..), setMD, toECig ) +import Bio.Prelude +import Bio.Streaming +import Foreign.C.Types ( CInt(..) ) + +import qualified Data.ByteString as B +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Storable as W + +-- | Trims from the 3' end of a sequence. +-- @trim_3' p b@ trims the 3' end of the sequence in @b@ at the +-- earliest position such that @p@ evaluates to true on every suffix +-- that was trimmed off. Note that the 3' end may be the beginning of +-- the sequence if it happens to be stored in reverse-complemented form. +-- Also note that trimming from the 3' end may not make sense for reads +-- that were constructed by merging paired end data (but we cannot take +-- care of that here). Further note that trimming may break dependent +-- information, notably the "mate" information of the mate and many +-- optional fields. + +trim_3' :: ([Nucleotides] -> [Qual] -> Bool) -> BamRec -> BamRec +trim_3' p b | b_flag b `testBit` 4 = trim_rev + | otherwise = trim_fwd + where + trim_fwd = let l = subtract 1 . length . takeWhile (uncurry p) $ + zip (inits . reverse . V.toList $ b_seq b) + (inits . reverse . V.toList $ b_qual b) + in trim_3 l b + + trim_rev = let l = subtract 1 . length . takeWhile (uncurry p) $ + zip (inits . V.toList $ b_seq b) + (inits . V.toList $ b_qual b) + in trim_3 l b + +trim_3 :: Int -> BamRec -> BamRec +trim_3 l b | b_flag b `testBit` 4 = trim_rev + | otherwise = trim_fwd + where + trim_fwd = let (_, cigar') = trim_back_cigar (b_cigar b) l + c = modMd (takeECig (V.length (b_seq b) - l)) b + in c { b_seq = V.take (V.length (b_seq c) - l) (b_seq c) + , b_qual = V.take (V.length (b_qual c) - l) (b_qual c) + , b_cigar = cigar' + , b_exts = map (\(k,e) -> case e of + Text t | k `elem` trim_set + -> (k, Text (B.take (B.length t - l) t)) + _ -> (k,e) + ) (b_exts c) } + + trim_rev = let (off, cigar') = trim_fwd_cigar (b_cigar b) l + c = modMd (dropECig l) b + in c { b_seq = V.drop l (b_seq c) + , b_qual = V.drop l (b_qual c) + , b_pos = b_pos c + off + , b_cigar = cigar' + , b_exts = map (\(k,e) -> case e of + Text t | k `elem` trim_set + -> (k, Text (B.drop l t)) + _ -> (k,e) + ) (b_exts c) } + + trim_set = ["BQ","CQ","CS","E2","OQ","U2"] + + modMd :: (ECig -> ECig) -> BamRec -> BamRec + modMd f br = maybe br (setMD br . f . toECig (b_cigar br)) (getMd br) + + endOf :: ECig -> ECig + endOf WithMD = WithMD + endOf WithoutMD = WithoutMD + endOf (Mat' _ es) = endOf es + endOf (Ins' _ es) = endOf es + endOf (SMa' _ es) = endOf es + endOf (Rep' _ es) = endOf es + endOf (Del' _ es) = endOf es + endOf (Nop' _ es) = endOf es + endOf (HMa' _ es) = endOf es + endOf (Pad' _ es) = endOf es + + takeECig :: Int -> ECig -> ECig + takeECig 0 es = endOf es + takeECig _ WithMD = WithMD + takeECig _ WithoutMD = WithoutMD + takeECig n (Mat' m es) = Mat' n $ if n > m then takeECig (n-m) es else WithMD + takeECig n (Ins' m es) = Ins' n $ if n > m then takeECig (n-m) es else WithMD + takeECig n (SMa' m es) = SMa' n $ if n > m then takeECig (n-m) es else WithMD + takeECig n (Rep' ns es) = Rep' ns $ takeECig (n-1) es + takeECig n (Del' ns es) = Del' ns $ takeECig n es + takeECig n (Nop' m es) = Nop' m $ takeECig n es + takeECig n (HMa' m es) = HMa' m $ takeECig n es + takeECig n (Pad' m es) = Pad' m $ takeECig n es + + dropECig :: Int -> ECig -> ECig + dropECig 0 es = es + dropECig _ WithMD = WithMD + dropECig _ WithoutMD = WithoutMD + dropECig n (Mat' m es) = if n > m then dropECig (n-m) es else Mat' n WithMD + dropECig n (Ins' m es) = if n > m then dropECig (n-m) es else Ins' n WithMD + dropECig n (SMa' m es) = if n > m then dropECig (n-m) es else SMa' n WithMD + dropECig n (Rep' _ es) = dropECig (n-1) es + dropECig n (Del' _ es) = dropECig n es + dropECig n (Nop' _ es) = dropECig n es + dropECig n (HMa' _ es) = dropECig n es + dropECig n (Pad' _ es) = dropECig n es + + +trim_back_cigar, trim_fwd_cigar :: V.Vector v Cigar => v Cigar -> Int -> ( Int, v Cigar ) +trim_back_cigar c l = (o, V.fromList $ reverse c') where (o,c') = sanitize_cigar . trim_cigar l $ reverse $ V.toList c +trim_fwd_cigar c l = (o, V.fromList c') where (o,c') = sanitize_cigar $ trim_cigar l $ V.toList c + +sanitize_cigar :: (Int, [Cigar]) -> (Int, [Cigar]) +sanitize_cigar (o, [ ]) = (o, []) +sanitize_cigar (o, (op:*l):xs) | op == Pad = sanitize_cigar (o,xs) -- del P + | op == Del || op == Nop = sanitize_cigar (o + l, xs) -- adjust D,N + | op == Ins = (o, (SMa :* l):xs) -- I --> S + | otherwise = (o, (op :* l):xs) -- rest is fine + +trim_cigar :: Int -> [Cigar] -> (Int, [Cigar]) +trim_cigar 0 cs = (0, cs) +trim_cigar _ [] = (0, []) +trim_cigar l ((op:*ll):cs) | bad_op op = let (o,cs') = trim_cigar l cs in (o + reflen op ll, cs') + | otherwise = case l `compare` ll of + LT -> (reflen op l, (op :* (ll-l)):cs) + EQ -> (reflen op ll, cs) + GT -> let (o,cs') = trim_cigar (l - ll) cs in (o + reflen op ll, cs') + + where + reflen op' = if ref_op op' then id else const 0 + bad_op o = o /= Mat && o /= Ins && o /= SMa + ref_op o = o == Mat || o == Del + + +-- | Trim predicate to get rid of low quality sequence. +-- @trim_low_quality q ns qs@ evaluates to true if all qualities in @qs@ +-- are smaller (i.e. worse) than @q@. +trim_low_quality :: Qual -> a -> [Qual] -> Bool +trim_low_quality q = const $ all (< q) + + +-- | Finds the merge point. Input is list of forward adapters, list of +-- reverse adapters, sequence1, quality1, sequence2, quality2; output is +-- merge point and two qualities (YM, YN). +find_merge :: [W.Vector Nucleotides] -> [W.Vector Nucleotides] + -> W.Vector Nucleotides -> W.Vector Qual + -> W.Vector Nucleotides -> W.Vector Qual + -> (Int, Int, Int) +find_merge ads1 ads2 r1 q1 r2 q2 = (mlen, score2 - score1, plain_score - score1) + where + plain_score = 6 * (V.length r1 + V.length r2) + (score1, mlen, score2) = twoMins plain_score (V.length r1 + V.length r2) $ + merge_score ads1 ads2 r1 q1 r2 q2 + +-- | Overlap-merging of read pairs. We shall compute the likelihood +-- for every possible overlap, then select the most likely one (unless it +-- looks completely random), compute a quality from the second best +-- merge, then merge and clamp the quality accordingly. +-- (We could try looking for chimaera after completing the merge, if +-- only we knew which ones to expect?) +-- +-- Two reads go in, with two adapter lists. We return 'Nothing' if all +-- merges looked mostly random. Else we return the two original reads, +-- flagged as 'eflagVestigial' *and* the merged version, flagged as +-- 'eflagMerged' and optionally 'eflagTrimmed'. All reads contain the +-- computed qualities (in YM and YN), which we also return. +-- +-- The merging automatically limits quality scores some of the time. We +-- additionally impose a hard limit of 63 to avoid difficulties +-- representing the result, and even that is ridiculous. Sane people +-- would further limit the returned quality! (In practice, map quality +-- later imposes a limit anyway, so no worries...) + +mergeBam :: Int -> Int -> [W.Vector Nucleotides] -> [W.Vector Nucleotides] -> BamRec -> BamRec -> [BamRec] +mergeBam lowq highq ads1 ads2 r1 r2 + | V.null (b_seq r1) && V.null (b_seq r2) = [ ] + | qual1 < lowq || mlen < 0 = [ r1', r2' ] + | qual1 >= highq && mlen == 0 = [ ] + | qual1 >= highq = [ rm ] + | mlen < len_r1-20 || mlen < len_r2-20 = [ rm ] + | otherwise = map flag_alternative [ r1', r2', rm ] + where + len_r1 = V.length $ b_seq r1 + len_r2 = V.length $ b_seq r2 + + b_seq_r1 = V.convert $ b_seq r1 + b_seq_r2 = V.convert $ b_seq r2 + b_qual_r1 = V.convert $ b_qual r1 + b_qual_r2 = V.convert $ b_qual r2 + + (mlen, qual1, qual2) = find_merge ads1 ads2 b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 + + flag_alternative br = br { b_exts = updateE "FF" (Int $ extAsInt 0 "FF" br .|. eflagAlternative) $ b_exts br } + store_quals br = br { b_exts = updateE "YM" (Int qual1) $ updateE "YN" (Int qual2) $ b_exts br } + pair_flags = flagPaired.|.flagProperlyPaired.|.flagMateUnmapped.|.flagMateReversed.|.flagFirstMate.|.flagSecondMate + + r1' = store_quals r1 + r2' = store_quals r2 + rm = store_quals $ merged_read mlen (fromIntegral $ min 63 qual1) + + merged_read l qmax = nullBamRec { + b_qname = b_qname r1, + b_flag = flagUnmapped .|. complement pair_flags .&. b_flag r1, + b_seq = V.convert $ merged_seq l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2, + b_qual = V.convert $ merged_qual qmax l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2, + b_exts = let ff = if l < len_r1 then eflagTrimmed else 0 + in updateE "FF" (Int $ extAsInt 0 "FF" r1 .|. eflagMerged .|. ff) $ b_exts r1 } + +{-# INLINE merged_seq #-} +merged_seq :: (V.Vector v Nucleotides, V.Vector v Qual) + => Int -> v Nucleotides -> v Qual -> v Nucleotides -> v Qual -> v Nucleotides +merged_seq l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 = V.concat + [ V.take (l - len_r2) b_seq_r1 + , V.zipWith4 zz (V.take l $ V.drop (l - len_r2) b_seq_r1) + (V.take l $ V.drop (l - len_r2) b_qual_r1) + (V.reverse $ V.take l $ V.drop (l - len_r1) b_seq_r2) + (V.reverse $ V.take l $ V.drop (l - len_r1) b_qual_r2) + , V.reverse $ V.take (l - len_r1) b_seq_r2 ] + where + len_r1 = V.length b_qual_r1 + len_r2 = V.length b_qual_r2 + zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = n1 + | q1 > q2 = n1 + | otherwise = compls n2 + +{-# INLINE merged_qual #-} +merged_qual :: (V.Vector v Nucleotides, V.Vector v Qual) + => Word8 -> Int -> v Nucleotides -> v Qual -> v Nucleotides -> v Qual -> v Qual +merged_qual qmax l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 = V.concat + [ V.take (l - len_r2) b_qual_r1 + , V.zipWith4 zz (V.take l $ V.drop (l - len_r2) b_seq_r1) + (V.take l $ V.drop (l - len_r2) b_qual_r1) + (V.reverse $ V.take l $ V.drop (l - len_r1) b_seq_r2) + (V.reverse $ V.take l $ V.drop (l - len_r1) b_qual_r2) + , V.reverse $ V.take (l - len_r1) b_qual_r2 ] + where + len_r1 = V.length b_qual_r1 + len_r2 = V.length b_qual_r2 + zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = Q $ min qmax (q1 + q2) + | q1 > q2 = Q $ q1 - q2 + | otherwise = Q $ q2 - q1 + + + +-- | Finds the trimming point. Input is list of forward adapters, +-- sequence, quality; output is trim point and two qualities (YM, YN). +find_trim :: [W.Vector Nucleotides] + -> W.Vector Nucleotides -> W.Vector Qual + -> (Int, Int, Int) +find_trim ads1 r1 q1 = (mlen, score2 - score1, plain_score - score1) + where + plain_score = 6 * V.length r1 + (score1, mlen, score2) = twoMins plain_score (V.length r1) $ + merge_score ads1 [V.empty] r1 q1 V.empty V.empty + +-- | Trimming for a single read: we need one adapter only (the one coming +-- /after/ the read), here provided as a list of options, and then we +-- merge with an empty second read. Results in up to two reads (the +-- original, possibly flagged, and the trimmed one, definitely flagged, +-- and two qualities). +trimBam :: Int -> Int -> [W.Vector Nucleotides] -> BamRec -> [BamRec] +trimBam lowq highq ads1 r1 + | V.null (b_seq r1) = [ ] + | mlen == 0 && qual1 >= highq = [ ] + | qual1 < lowq || mlen < 0 = [ r1' ] + | qual1 >= highq = [ r1t ] + | otherwise = map flag_alternative [ r1', r1t ] + where + -- the "merge" score if there is no trimming + + b_seq_r1 = V.convert $ b_seq r1 + b_qual_r1 = V.convert $ b_qual r1 + + (mlen, qual1, qual2) = find_trim ads1 b_seq_r1 b_qual_r1 + + flag_alternative br = br { b_exts = updateE "FF" (Int $ extAsInt 0 "FF" br .|. eflagAlternative) $ b_exts br } + store_quals br = br { b_exts = updateE "YM" (Int qual1) $ updateE "YN" (Int qual2) $ b_exts br } + + r1' = store_quals r1 + r1t = store_quals $ trimmed_read mlen + + trimmed_read l = nullBamRec { + b_qname = b_qname r1, + b_flag = flagUnmapped .|. b_flag r1, + b_seq = V.take l $ b_seq r1, + b_qual = V.take l $ b_qual r1, + b_exts = updateE "FF" (Int $ extAsInt 0 "FF" r1 .|. eflagTrimmed) $ b_exts r1 } + + +-- | For merging, we don't need the complete adapters (length around 70!), +-- only a sufficient prefix. Taking only the more-or-less constant +-- part (length around 30), there aren't all that many different +-- adapters in the world. To deal with pretty much every library, we +-- only need the following forward adapters, which will be the default +-- (defined here in the direction they would be sequenced in): Genomic +-- R2, Multiplex R2, Fraft P7. + +default_fwd_adapters :: [ W.Vector Nucleotides ] +default_fwd_adapters = map (W.fromList. map toNucleotides) + [ {- Genomic R2 -} "AGATCGGAAGAGCGGTTCAG" + , {- Multiplex R2 -} "AGATCGGAAGAGCACACGTC" + , {- Graft P7 -} "AGATCGGAAGAGCTCGTATG" ] + +-- | Like 'default_rev_adapters', these are the few adapters needed for +-- the reverse read (defined in the direction they would be sequenced in +-- as part of the second read): Genomic R1, CL 72. + +default_rev_adapters :: [ W.Vector Nucleotides ] +default_rev_adapters = map (W.fromList. map toNucleotides) + [ {- Genomic_R1 -} "AGATCGGAAGAGCGTCGTGT" + , {- CL72 -} "GGAAGAGCGTCGTGTAGGGA" ] + +-- We need to compute the likelihood of a read pair given an assumed +-- insert length. The likelihood of the first read is the likelihood of +-- a match with the adapter where it overlaps the 3' adapter, elsewhere +-- it's 1/4 per position. The likelihood of the second read is the +-- likelihood of a match with the adapter where it overlaps the adapter, +-- the likehood of a read-read match where it overlaps read one, 1/4 per +-- position elsewhere. (Yes, this ignores base composition. It doesn't +-- matter enough.) + +merge_score + :: [ W.Vector Nucleotides ] -- 3' adapters as they appear in the first read + -> [ W.Vector Nucleotides ] -- 5' adapters as they appear in the second read + -> W.Vector Nucleotides -> W.Vector Qual -- first read, qual + -> W.Vector Nucleotides -> W.Vector Qual -- second read, qual + -> Int -- assumed insert length + -> Int -- score (roughly sum of qualities at mismatches) +merge_score fwd_adapters rev_adapters !read1 !qual1 !read2 !qual2 !l + = 6 * (l `min` V.length read1) -- read1, part before adapter + + 6 * (max 0 (l - V.length read1)) -- read2, part before overlap + + + foldl' (\acc fwd_ad -> min acc + (match_adapter l read1 qual1 fwd_ad + -- read1, match with forward adapter + 6 * (max 0 (V.length read1 - V.length fwd_ad - l))) -- read1, part after (known) adapter + ) maxBound fwd_adapters + + + foldl' (\acc rev_ad -> min acc + (match_adapter l read2 qual2 rev_ad + -- read2, match with reverse adapter + 6 * (max 0 (V.length read2 - V.length rev_ad - l))) -- read2, part after (known) adapter + ) maxBound rev_adapters + + + match_reads l read1 qual1 read2 qual2 + +{-# INLINE match_adapter #-} +match_adapter :: Int -> W.Vector Nucleotides -> W.Vector Qual -> W.Vector Nucleotides -> Int +match_adapter !off !rd !qs !ad + | V.length rd /= V.length qs = error "read/qual length mismatch" + | efflength <= 0 = 0 + | otherwise + = fromIntegral . unsafePerformIO $ + W.unsafeWith rd $ \p_rd -> + W.unsafeWith qs $ \p_qs -> + W.unsafeWith ad $ \p_ad -> + prim_match_ad (fromIntegral off) + (fromIntegral efflength) + p_rd p_qs p_ad + where + !efflength = (V.length rd - off) `min` V.length ad + +foreign import ccall unsafe "prim_match_ad" + prim_match_ad :: CInt -> CInt + -> Ptr Nucleotides -> Ptr Qual + -> Ptr Nucleotides -> IO CInt + + +-- | Computes overlap score for two reads (with qualities) assuming an +-- insert length. +{-# INLINE match_reads #-} +match_reads :: Int -> W.Vector Nucleotides -> W.Vector Qual -> W.Vector Nucleotides -> W.Vector Qual -> Int +match_reads !l !rd1 !qs1 !rd2 !qs2 + | V.length rd1 /= V.length qs1 || V.length rd2 /= V.length qs2 = error "read/qual length mismatch" + | efflength <= 0 = 0 + | otherwise + = fromIntegral . unsafePerformIO $ + W.unsafeWith rd1 $ \p_rd1 -> + W.unsafeWith qs1 $ \p_qs1 -> + W.unsafeWith rd2 $ \p_rd2 -> + W.unsafeWith qs2 $ \p_qs2 -> + prim_match_reads (fromIntegral minidx1) + (fromIntegral maxidx2) + (fromIntegral efflength) + p_rd1 p_qs1 p_rd2 p_qs2 + where + -- vec1, forward + !minidx1 = (l - V.length rd2) `max` 0 + -- vec2, backward + !maxidx2 = l `min` V.length rd2 + -- effective length + !efflength = ((V.length rd1 + V.length rd2 - l) `min` l) `max` 0 + + +foreign import ccall unsafe "prim_match_reads" + prim_match_reads :: CInt -> CInt -> CInt + -> Ptr Nucleotides -> Ptr Qual + -> Ptr Nucleotides -> Ptr Qual -> IO CInt + + +{-# INLINE twoMins #-} +twoMins :: (Bounded a, Ord a) => a -> Int -> (Int -> a) -> (a,Int,a) +twoMins a0 imax f = go a0 (-1) maxBound 0 0 + where + go !m1 !i1 !m2 !i2 !i + | i == imax = (m1,i1,m2) + | otherwise = + case f i of + x | x < m1 -> go x i m1 i1 (i+1) + | x < m2 -> go m1 i1 x i (i+1) + | otherwise -> go m1 i1 m2 i2 (i+1) + + +mergeTrimBam :: Monad m + => Int -> Int -> [W.Vector Nucleotides] -> [W.Vector Nucleotides] + -> Stream (Of BamRec) m r -> Stream (Of BamRec) m r +mergeTrimBam lowq highq fwd_ads rev_ads = go + where + go = lift . inspect >=> either pure go1 + + go1 (r1 :> s) | isPaired r1 = lift (inspect s) >>= go2 r1 + | otherwise = each (trimBam lowq highq fwd_ads r1) >> go s + + go2 r1 (Left _) = error $ "Lone mate found: " ++ show (b_qname r1) + go2 r1 (Right (r2 :> s)) = each (mergeBam lowq highq fwd_ads rev_ads r1 r2) >> go s + diff --git a/Bio/Bam/Writer.hs b/Bio/Bam/Writer.hs new file mode 100644 --- /dev/null +++ b/Bio/Bam/Writer.hs @@ -0,0 +1,247 @@ +-- | Printers for BAM and SAM. BAM is properly supported, SAM can be +-- piped to standard output. + +module Bio.Bam.Writer ( + IsBamRec(..), + encodeBamWith, + + packBam, + writeBamFile, + writeBamHandle, + pipeBamOutput, + pipeSamOutput + ) where + +import Bio.Bam.Header +import Bio.Bam.Rec +import Bio.Prelude +import Bio.Streaming +import Bio.Streaming.Bgzf + +import Data.ByteString.Builder.Prim ( (>*<) ) +import Data.ByteString.Internal ( ByteString(..) ) +import Data.ByteString.Lazy ( foldrChunks ) +import Foreign.Marshal.Alloc ( alloca ) + +import qualified Bio.Streaming.Bytes as S +import qualified Data.ByteString as B +import qualified Data.ByteString.Builder as B +import qualified Data.ByteString.Builder.Extra as B +import qualified Data.ByteString.Builder.Prim as E +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Storable as W +import qualified Data.Vector.Unboxed as U +import qualified Streaming.Prelude as Q + +{- | write in SAM format to stdout + +This is useful for piping to other tools (say, AWK scripts) or for +debugging. No convenience functions to send SAM to a file or to +compress it exist, because these are stupid ideas. +-} +pipeSamOutput :: (IsBamRec a, MonadIO m) => BamMeta -> Stream (Of a) m r -> m r +pipeSamOutput meta s = do + liftIO . B.hPutBuilder stdout $ showBamMeta meta + Q.mapM_ (liftIO . B.hPutBuilder stdout . encodeSamEntry (meta_refs meta) . unpackBamRec) s +{-# INLINE pipeSamOutput #-} + +encodeSamEntry :: Refs -> BamRec -> B.Builder +encodeSamEntry refs b = + B.byteStringCopy (b_qname b) <> B.char7 '\t' <> + B.intDec (b_flag b .&. 0xffff) <> B.char7 '\t' <> + B.byteStringCopy (sq_name $ getRef refs $ b_rname b) <> B.char7 '\t' <> + B.intDec (b_pos b + 1) <> B.char7 '\t' <> + B.word8Dec (unQ $ b_mapq b) <> B.char7 '\t' <> + buildCigar (b_cigar b) <> B.char7 '\t' <> + buildMrnm (b_mrnm b) (b_rname b) <> B.char7 '\t' <> + B.intDec (b_mpos b + 1) <> B.char7 '\t' <> + B.intDec (b_isize b) <> B.char7 '\t' <> + buildSeq (b_seq b) <> B.char7 '\t' <> + buildQual (b_qual b) <> + foldMap buildExt (b_exts b) <> B.char7 '\n' + where + buildCigar = E.primUnfoldrBounded + (E.intDec >*< E.liftFixedToBounded E.word8) + (vuncons $ \(op :* num) -> (num, B.index "MIDNSHP" (fromEnum op))) + + buildMrnm mrnm rname + | isValidRefseq mrnm && mrnm == rname = B.char7 '=' + | otherwise = B.byteString (sq_name $ getRef refs mrnm) + + buildSeq = E.primUnfoldrFixed E.word8 (vuncons $ \(Ns x) -> B.index "-ACMGRSVTWYHKDBN" $ fromIntegral $ x .&. 15) + buildQual = E.primUnfoldrFixed E.word8 (vuncons $ \(Q q) -> q + 33) + + buildExt (BamKey k,v) = B.char7 '\t' <> + B.word8 (fromIntegral k .&. 0xff) <> + B.word8 (shiftR (fromIntegral k) 8 .&. 0xff) <> + B.char7 ':' <> + buildExtVal v + + buildExtVal (Int i) = B.char7 'i' <> B.char7 ':' <> B.intDec i + buildExtVal (Float f) = B.char7 'f' <> B.char7 ':' <> B.floatDec f + buildExtVal (Text t) = B.char7 'Z' <> B.char7 ':' <> B.byteStringCopy t + buildExtVal (Bin x) = B.char7 'H' <> B.char7 ':' <> B.byteStringHex x + buildExtVal (Char c) = B.char7 'A' <> B.char7 ':' <> B.word8 c + buildExtVal (IntArr a) = B.char7 'B' <> B.char7 ':' <> B.char7 'i' <> buildArr B.intDec a + buildExtVal (FloatArr a) = B.char7 'B' <> B.char7 ':' <> B.char7 'f' <> buildArr B.floatDec a + + buildArr p = U.foldr (\x k -> B.char7 ',' <> p x <> k) mempty + + vuncons f v | V.null v = Nothing + | otherwise = Just (f (V.unsafeHead v), V.unsafeTail v) + + +class IsBamRec a where + pushBam :: a -> BgzfTokens -> BgzfTokens + unpackBamRec :: a -> BamRec + +instance IsBamRec BamRaw where + {-# INLINE pushBam #-} + pushBam = pushBamRaw + {-# INLINE unpackBamRec #-} + unpackBamRec = unpackBam + +instance IsBamRec BamRec where + {-# INLINE pushBam #-} + pushBam = pushBamRec + {-# INLINE unpackBamRec #-} + unpackBamRec = id + +instance (IsBamRec a, IsBamRec b) => IsBamRec (Either a b) where + {-# INLINE pushBam #-} + pushBam = either pushBam pushBam + {-# INLINE unpackBamRec #-} + unpackBamRec = either unpackBamRec unpackBamRec + +-- | Encodes BAM records straight into a dynamic buffer, then BGZF's it. +-- Should be fairly direct and perform well. +encodeBamWith :: (IsBamRec a, MonadIO m) => Int -> BamMeta -> Stream (Of a) m r -> ByteStream m r +encodeBamWith lv meta = encodeBgzf lv . enc_bam + where + enc_bam bs = Q.cons pushHeader $ Q.map (Endo . pushBam) bs + + pushHeader :: Endo BgzfTokens + pushHeader = Endo $ TkString "BAM\1" + . TkSetMark -- the length byte + . pushBuilder (showBamMeta meta) + . TkEndRecord -- fills the length in + . TkWord32 (fromIntegral . V.length . unRefs $ meta_refs meta) + . appEndo (foldMap (Endo . pushRef) (unRefs $ meta_refs meta)) + + pushRef :: BamSQ -> BgzfTokens -> BgzfTokens + pushRef bs = TkWord32 (fromIntegral $ B.length (sq_name bs) + 1) + . TkString (sq_name bs) + . TkWord8 0 + . TkWord32 (fromIntegral $ sq_length bs) + + pushBuilder :: B.Builder -> BgzfTokens -> BgzfTokens + pushBuilder b tk = foldrChunks TkString tk (B.toLazyByteString b) +{-# INLINE encodeBamWith #-} + +pushBamRaw :: BamRaw -> BgzfTokens -> BgzfTokens +pushBamRaw r = TkWord32 (fromIntegral $ B.length $ raw_data r) . + TkString (raw_data r) +{-# INLINE pushBamRaw #-} + +-- | Writes BAM encoded stuff to a file. +-- In reality, it cleverly writes to a temporary file and renames it +-- when done. +writeBamFile :: (IsBamRec a, MonadIO m, MonadMask m) => FilePath -> BamMeta -> Stream (Of a) m r -> m r +writeBamFile fp meta = S.writeFile fp . encodeBamWith 6 meta + +-- | Write BAM encoded stuff to stdout. +-- This sends uncompressed(!) BAM to stdout. Useful for piping to other +-- tools. The output is still wrapped in a BGZF stream, because that's +-- what all tools expect; but the individuals blocks are not compressed. +pipeBamOutput :: (IsBamRec a, MonadIO m) => BamMeta -> Stream (Of a) m r -> m r +pipeBamOutput meta = S.hPut stdout . encodeBamWith 0 meta +{-# INLINE pipeBamOutput #-} + +-- | Writes BAM encoded stuff to a 'Handle'. +writeBamHandle :: (IsBamRec a, MonadIO m) => Handle -> BamMeta -> Stream (Of a) m r -> m r +writeBamHandle hdl meta = S.hPut hdl . encodeBamWith 6 meta + +{-# RULES + "pushBam/unpackBam" forall b . pushBamRec (unpackBam b) = pushBamRaw b + #-} + +{-# INLINE[1] pushBamRec #-} +pushBamRec :: BamRec -> BgzfTokens -> BgzfTokens +pushBamRec BamRec{..} = + TkSetMark + . TkWord32 (unRefseq b_rname) + . TkWord32 (fromIntegral b_pos) + . TkWord8 (fromIntegral $ B.length b_qname + 1) + . TkWord8 (unQ b_mapq) + . TkWord16 (fromIntegral bin) + . TkWord16 (fromIntegral $ W.length b_cigar) + . TkWord16 (fromIntegral b_flag) + . TkWord32 (fromIntegral $ V.length b_seq) + . TkWord32 (unRefseq b_mrnm) + . TkWord32 (fromIntegral b_mpos) + . TkWord32 (fromIntegral b_isize) + . TkString b_qname + . TkWord8 0 + . W.foldr ((.) . TkWord8) id (W.unsafeCast b_cigar :: W.Vector Word8) + . pushSeq b_seq + . W.foldr ((.) . TkWord8 . unQ) id b_qual + . foldr ((.) . pushExt) id b_exts + . TkEndRecord + where + bin = distinctBin b_pos (alignedLength b_cigar) + + pushSeq :: V.Vector vec Nucleotides => vec Nucleotides -> BgzfTokens -> BgzfTokens + pushSeq v = case v V.!? 0 of + Nothing -> id + Just a -> case v V.!? 1 of + Nothing -> TkWord8 (unNs a `shiftL` 4) + Just b -> TkWord8 (unNs a `shiftL` 4 .|. unNs b) . pushSeq (V.drop 2 v) + + pushExt :: (BamKey, Ext) -> BgzfTokens -> BgzfTokens + pushExt (BamKey k, e) = case e of + Text t -> common 'Z' . TkString t . TkWord8 0 + Bin t -> common 'H' . TkString t . TkWord8 0 + Char c -> common 'A' . TkWord8 c + Float f -> common 'f' . TkWord32 (fromFloat f) + + Int i -> case put_some_int (U.singleton i) of + (c,op) -> common c . op i + + IntArr ia -> case put_some_int ia of + (c,op) -> common 'B' . TkWord8 (fromIntegral $ ord c) + . TkWord32 (fromIntegral $ U.length ia-1) + . U.foldr ((.) . op) id ia + + FloatArr fa -> common 'B' . TkWord8 (fromIntegral $ ord 'f') + . TkWord32 (fromIntegral $ U.length fa-1) + . U.foldr ((.) . TkWord32 . fromFloat) id fa + where + common :: Char -> BgzfTokens -> BgzfTokens + common z = TkWord16 k . TkWord8 (fromIntegral $ ord z) + + put_some_int :: U.Vector Int -> (Char, Int -> BgzfTokens -> BgzfTokens) + put_some_int is + | U.all (between 0 0xff) is = ('C', TkWord8 . fromIntegral) + | U.all (between (-0x80) 0x7f) is = ('c', TkWord8 . fromIntegral) + | U.all (between 0 0xffff) is = ('S', TkWord16 . fromIntegral) + | U.all (between (-0x8000) 0x7fff) is = ('s', TkWord16 . fromIntegral) + | U.all (> 0) is = ('I', TkWord32 . fromIntegral) + | otherwise = ('i', TkWord32 . fromIntegral) + + between :: Int -> Int -> Int -> Bool + between l r x = l <= x && x <= r + + fromFloat :: Float -> Word32 + fromFloat float = unsafeDupablePerformIO $ alloca $ \buf -> + pokeByteOff buf 0 float >> peek buf + +packBam :: BamRec -> IO BamRaw +packBam br = do bb <- newBuffer 1000 + (bb', TkEnd) <- store_loop bb (pushBamRec br TkEnd) + return . bamRaw 0 $ PS (buffer bb') 4 (used bb' - 4) + where + store_loop bb tk = do (bb',tk') <- fillBuffer bb tk + case tk' of TkEnd -> return (bb',tk') + _ -> do bb'' <- expandBuffer (128*1024) bb' + store_loop bb'' tk' + diff --git a/Bio/Base.hs b/Bio/Base.hs new file mode 100644 --- /dev/null +++ b/Bio/Base.hs @@ -0,0 +1,339 @@ +{-# LANGUAGE CPP #-} +-- | Common data types used everywhere. This module is a collection of +-- very basic "bioinformatics" data types that are simple, but don't +-- make sense to define over and over. + +module Bio.Base( + Nucleotide(..), Nucleotides(..), + Qual(..), toQual, fromQual, fromQualRaised, probToQual, + Prob'(..), Prob, toProb, fromProb, qualToProb, pow, + + Pair(..), + Word8, + nucA, nucC, nucG, nucT, + nucsA, nucsC, nucsG, nucsT, nucsN, gap, + toNucleotide, toNucleotides, nucToNucs, + showNucleotide, showNucleotides, + isGap, + isBase, + isProperBase, + properBases, + compl, compls, + + Position(..), + shiftPosition, + p_is_reverse, + + Range(..), + shiftRange, + reverseRange, + extendRange, + insideRange, + wrapRange +) where + +import BasePrelude +#if MIN_VERSION_base(4,9,0) + hiding ( log1pexp, log1mexp ) +#endif +import Bio.Util.Numeric ( log1pexp, log1mexp ) + +import qualified Data.ByteString.Char8 as C +import qualified Data.Vector.Unboxed as U + +-- | A nucleotide base. We only represent A,C,G,T. The contained +-- 'Word8' ist guaranteed to be 0..3. +newtype Nucleotide = N { unN :: Word8 } deriving ( Eq, Ord, Enum, Ix, Storable ) + +instance Bounded Nucleotide where + minBound = N 0 + maxBound = N 3 + +-- | A nucleotide base in an alignment. +-- Experience says we're dealing with Ns and gaps all the type, so +-- purity be damned, they are included as if they were real bases. +-- +-- To allow @Nucleotides@s to be unpacked and incorporated into +-- containers, we choose to represent them the same way as the BAM file +-- format: as a 4 bit wide field. Gaps are encoded as 0 where they +-- make sense, N is 15. The contained 'Word8' is guaranteed to be +-- 0..15. + +newtype Nucleotides = Ns { unNs :: Word8 } deriving ( Eq, Ord, Enum, Ix, Storable ) + +instance Bounded Nucleotides where + minBound = Ns 0 + maxBound = Ns 15 + +nucToNucs :: Nucleotide -> Nucleotides +nucToNucs (N x) = Ns $ 1 `shiftL` fromIntegral (x .&. 3) + +-- | Qualities are stored in deciban, also known as the Phred scale. To +-- represent a value @p@, we store @-10 * log_10 p@. Operations work +-- directly on the \"Phred\" value, as the name suggests. The same goes +-- for the 'Ord' instance: greater quality means higher \"Phred\" +-- score, meand lower error probability. + +newtype Qual = Q { unQ :: Word8 } deriving ( Eq, Ord, Storable, Bounded ) + +instance Show Qual where + showsPrec p (Q q) = (:) 'q' . showsPrec p q + +toQual :: (Floating a, RealFrac a) => a -> Qual +toQual a = Q $ round (-10 * log a / log 10) + +fromQual :: Qual -> Double +fromQual (Q q) = 10 ** (- fromIntegral q / 10) + +fromQualRaised :: Double -> Qual -> Double +fromQualRaised k (Q q) = 10 ** (- k * fromIntegral q / 10) + +-- | A positive floating point value stored in log domain. We store the +-- natural logarithm (makes computation easier), but allow conversions +-- to the familiar \"Phred\" scale used for 'Qual' values. +newtype Prob' a = Pr { unPr :: a } deriving ( Eq, Ord, Storable ) + +-- | Common way of using 'Prob''. +type Prob = Prob' Double + +instance RealFloat a => Show (Prob' a) where + showsPrec _ (Pr p) = (:) 'q' . showFFloat (Just 1) q + where q = - 10 * p / log 10 + +instance (Floating a, Ord a) => Num (Prob' a) where + {-# INLINE fromInteger #-} + fromInteger a = Pr (log (fromInteger a)) + {-# INLINE (+) #-} + Pr x + Pr y = Pr $ if x >= y then x + log1pexp (y-x) else y + log1pexp (x-y) + {-# INLINE (-) #-} + Pr x - Pr y = Pr $ if x >= y then x + log1mexp (y-x) else error "no negative error probabilities" + {-# INLINE (*) #-} + Pr a * Pr b = Pr $ a + b + {-# INLINE negate #-} + negate _ = Pr $ error "no negative error probabilities" + {-# INLINE abs #-} + abs x = x + {-# INLINE signum #-} + signum _ = Pr 0 + +instance (Floating a, Fractional a, Ord a) => Fractional (Prob' a) where + fromRational a = Pr (log (fromRational a)) + Pr a / Pr b = Pr (a - b) + recip (Pr a) = Pr (negate a) + +infixr 8 `pow` +pow :: Num a => Prob' a -> a -> Prob' a +pow (Pr a) e = Pr $ a * e + + +toProb :: Floating a => a -> Prob' a +toProb p = Pr (log p) + +fromProb :: Floating a => Prob' a -> a +fromProb (Pr q) = exp q + +qualToProb :: Floating a => Qual -> Prob' a +qualToProb (Q q) = Pr (- log 10 * fromIntegral q / 10) + +probToQual :: (Floating a, RealFrac a) => Prob' a -> Qual +probToQual (Pr p) = Q (round (- 10 * p / log 10)) + +nucA, nucC, nucG, nucT :: Nucleotide +nucA = N 0 +nucC = N 1 +nucG = N 2 +nucT = N 3 + +gap, nucsA, nucsC, nucsG, nucsT, nucsN :: Nucleotides +gap = Ns 0 +nucsA = Ns 1 +nucsC = Ns 2 +nucsG = Ns 4 +nucsT = Ns 8 +nucsN = Ns 15 + + +-- | Coordinates in a genome. +-- The position is zero-based, no questions about it. Think of the +-- position as pointing to the crack between two bases: looking forward +-- you see the next base to the right, looking in the reverse direction +-- you see the complement of the first base to the left. +-- +-- To encode the strand, we (virtually) reverse-complement any sequence +-- and prepend it to the normal one. That way, reversed coordinates +-- have a negative sign and automatically make sense. Position 0 could +-- either be the beginning of the sequence or the end on the reverse +-- strand... that ambiguity shouldn't really matter. + +data Position = Pos { + p_seq :: {-# UNPACK #-} !C.ByteString, -- ^ sequence (e.g. some chromosome) + p_start :: {-# UNPACK #-} !Int -- ^ offset, zero-based + } deriving (Show, Eq, Ord) + +p_is_reverse :: Position -> Bool +p_is_reverse = (< 0) . p_start + +-- | Ranges in genomes +-- We combine a position with a length. In 'Range pos len', 'pos' is +-- always the start of a stretch of length 'len'. Positions therefore +-- move in the opposite direction on the reverse strand. To get the +-- same stretch on the reverse strand, shift r_pos by r_length, then +-- reverse direction (or call reverseRange). +data Range = Range { + r_pos :: {-# UNPACK #-} !Position, + r_length :: {-# UNPACK #-} !Int + } deriving (Show, Eq, Ord) + + +-- | Converts a character into a 'Nucleotides'. +-- The usual codes for A,C,G,T and U are understood, '-' and '.' become +-- gaps and everything else is an N. +toNucleotide :: Char -> Nucleotide +toNucleotide c = if ord c < 128 then N (ar `U.unsafeIndex` ord c) else N 0 + where + ar = U.replicate 128 0 U.// + ( [ (ord x, n) | (x, N n) <- pairs ] ++ + [ (ord (toUpper x), n) | (x, N n) <- pairs ] ) + + pairs = [ ('a', nucA), ('c', nucC), ('g', nucG), ('t', nucT) ] + + +-- | Converts a character into a 'Nucleotides'. +-- The usual codes for A,C,G,T and U are understood, '-' and '.' become +-- gaps and everything else is an N. +toNucleotides :: Char -> Nucleotides +toNucleotides c = if ord c < 128 then Ns (ar `U.unsafeIndex` ord c) else nucsN + where + ar = U.replicate 128 (unNs nucsN) U.// + ( [ (ord x, n) | (x, Ns n) <- pairs ] ++ + [ (ord (toUpper x), n) | (x, Ns n) <- pairs ] ) + + Ns a `plus` Ns b = Ns (a .|. b) + + pairs = [ ('a', nucsA), ('c', nucsC), ('g', nucsG), ('t', nucsT), + ('u', nucsT), ('-', gap), ('.', gap), + ('b', nucsC `plus` nucsG `plus` nucsT), + ('d', nucsA `plus` nucsG `plus` nucsT), + ('h', nucsA `plus` nucsC `plus` nucsT), + ('v', nucsA `plus` nucsC `plus` nucsG), + ('k', nucsG `plus` nucsT), + ('m', nucsA `plus` nucsC), + ('s', nucsC `plus` nucsG), + ('w', nucsA `plus` nucsT), + ('r', nucsA `plus` nucsG), + ('y', nucsC `plus` nucsT) ] + +-- | Tests if a 'Nucleotides' is a base. +-- Returns 'True' for everything but gaps. +isBase :: Nucleotides -> Bool +isBase (Ns n) = n /= 0 + +-- | Tests if a 'Nucleotides' is a proper base. +-- Returns 'True' for A,C,G,T only. +isProperBase :: Nucleotides -> Bool +isProperBase x = x == nucsA || x == nucsC || x == nucsG || x == nucsT + +properBases :: [ Nucleotides ] +properBases = [ nucsA, nucsC, nucsG, nucsT ] + +-- | Tests if a 'Nucleotides' is a gap. +-- Returns true only for the gap. +isGap :: Nucleotides -> Bool +isGap x = x == gap + + +{-# INLINE showNucleotide #-} +showNucleotide :: Nucleotide -> Char +showNucleotide (N x) = C.index "ACGT" $ fromIntegral $ x .&. 3 + +{-# INLINE showNucleotides #-} +showNucleotides :: Nucleotides -> Char +showNucleotides (Ns x) = C.index "-ACMGRSVTWYHKDBN" $ fromIntegral $ x .&. 15 + +instance Show Nucleotide where + show x = [ showNucleotide x ] + showList l = (map showNucleotide l ++) + +instance Read Nucleotide where + readsPrec _ ('a':cs) = [(nucA, cs)] + readsPrec _ ('A':cs) = [(nucA, cs)] + readsPrec _ ('c':cs) = [(nucC, cs)] + readsPrec _ ('C':cs) = [(nucC, cs)] + readsPrec _ ('g':cs) = [(nucG, cs)] + readsPrec _ ('G':cs) = [(nucG, cs)] + readsPrec _ ('t':cs) = [(nucT, cs)] + readsPrec _ ('T':cs) = [(nucT, cs)] + readsPrec _ ('u':cs) = [(nucT, cs)] + readsPrec _ ('U':cs) = [(nucT, cs)] + readsPrec _ _ = [ ] + + readList ('-':cs) = readList cs + readList (c:cs) | isSpace c = readList cs + | otherwise = case reads (c:cs) of + [] -> [ ([],c:cs) ] + xs -> [ (n:ns,r2) | (n,r1) <- xs, (ns,r2) <- readList r1 ] + readList [] = [([],[])] + +instance Show Nucleotides where + show x = [ showNucleotides x ] + showList l = (map showNucleotides l ++) + +instance Read Nucleotides where + readsPrec _ (c:cs) = [(toNucleotides c, cs)] + readsPrec _ [ ] = [] + readList s = let (hd,tl) = span (\c -> isAlpha c || isSpace c || '-' == c) s + in [(map toNucleotides $ filter (not . isSpace) hd, tl)] + +-- | Complements a Nucleotides. +{-# INLINE compl #-} +compl :: Nucleotide -> Nucleotide +compl (N n) = N $ n `xor` 3 + +-- | Complements a Nucleotides. +{-# INLINE compls #-} +compls :: Nucleotides -> Nucleotides +compls (Ns x) = Ns $ ar `U.unsafeIndex` fromIntegral (x .&. 15) + where + !ar = U.fromListN 16 [ 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15 ] + + +-- | Moves a @Position@. The position is moved forward according to the +-- strand, negative indexes move backward accordingly. +shiftPosition :: Int -> Position -> Position +shiftPosition a p = p { p_start = p_start p + a } + +-- | Moves a @Range@. This is just @shiftPosition@ lifted. +shiftRange :: Int -> Range -> Range +shiftRange a r = r { r_pos = shiftPosition a (r_pos r) } + +-- | Reverses a 'Range' to give the same @Range@ on the opposite strand. +reverseRange :: Range -> Range +reverseRange (Range (Pos sq pos) len) = Range (Pos sq (-pos-len)) len + +-- | Extends a range. The length of the range is simply increased. +extendRange :: Int -> Range -> Range +extendRange a r = r { r_length = r_length r + a } + +-- | Expands a subrange. +-- @(range1 `insideRange` range2)@ interprets @range1@ as a subrange of +-- @range2@ and computes its absolute coordinates. The sequence name of +-- @range1@ is ignored. +insideRange :: Range -> Range -> Range +insideRange r1@(Range (Pos _ start1) length1) r2@(Range (Pos sq start2) length2) + | start2 < 0 = reverseRange (insideRange r1 (reverseRange r2)) + | start1 <= length2 = Range (Pos sq (start2 + start1)) (min length1 (length2 - start1)) + | otherwise = Range (Pos sq (start2 + length2)) 0 + + +-- | Wraps a range to a region. This simply normalizes the start +-- position to be in the interval '[0,n)', which only makes sense if the +-- @Range@ is to be mapped onto a circular genome. This works on both +-- strands and the strand information is retained. +wrapRange :: Int -> Range -> Range +wrapRange n (Range (Pos sq s) l) = Range (Pos sq (s `mod` n)) l + + +-- | A strict pair. +data Pair a b = !a :!: !b deriving(Eq, Ord, Show, Read, Bounded, Ix) +infixl 2 :!: + diff --git a/Bio/Prelude.hs b/Bio/Prelude.hs new file mode 100644 --- /dev/null +++ b/Bio/Prelude.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE CPP #-} +module Bio.Prelude ( + module BasePrelude, + module Bio.Base, + module Bio.Util.Text, + module Control.Monad.Catch, + module Control.Monad.IO.Class, + module Control.Monad.Trans.Class, + module Data.Bifunctor, + module Data.List.NonEmpty, + module Data.Semigroup, + module System.IO, + + Bytes, LazyBytes, + Generic1(..), + Hashable(..), + Hashable1(..), + Hashable2(..), + HashMap, + HashSet, + IntMap, + IntSet, + NonEmpty(..), + Semigroup(..), + Text, LazyText + ) where + +import BasePrelude hiding ( (<>), EOF + , bracket, bracket_, bracketOnError + , catch, catches, catchIOError, catchJust + , Handler, handle, handleJust + , finally, try, tryJust, onException + , mask, mask_, uninterruptibleMask, uninterruptibleMask_ +#if MIN_VERSION_base(4,9,0) + , log1p, log1pexp, log1mexp, expm1 +#endif + ) + +import Bio.Base +import Bio.Util.Text +import Control.Monad.Catch +import Control.Monad.IO.Class +import Control.Monad.Trans.Class +import Data.Bifunctor +import Data.ByteString ( ByteString ) +import Data.List.NonEmpty ( NonEmpty(..) ) +import Data.Semigroup ( Semigroup(..) ) +import Data.Text ( Text ) +import Data.Hashable ( Hashable(..) ) +import Data.Hashable.Lifted ( Hashable1(..), Hashable2(..) ) +import Data.HashMap.Strict ( HashMap ) +import Data.HashSet ( HashSet ) +import Data.IntMap ( IntMap ) +import Data.IntSet ( IntSet ) +import GHC.Generics ( Generic1(..) ) +import System.IO ( hPrint, hPutStr, hPutStrLn, stderr, stdout, stdin + , openBinaryFile, withBinaryFile, IOMode(..) + , hFlush, hSeek, hClose, SeekMode(..) ) + +import qualified Data.ByteString.Lazy as BL +import qualified Data.Text.Lazy as TL + +type Bytes = ByteString +type LazyBytes = BL.ByteString +type LazyText = TL.Text diff --git a/Bio/Streaming.hs b/Bio/Streaming.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming.hs @@ -0,0 +1,195 @@ +module Bio.Streaming + ( MonadIO(..) + , MonadMask + , ByteStream + + , streamFile + , streamHandle + , streamInput + , streamInputs + , withOutputFile + + , protectTerm + , psequence + , progressGen + , progressNum + , progressPos + + , mergeStreams + , mergeStreamsBy + , mergeStreamsOn + + , module Streaming + , module Streaming.Prelude ) + where + +import Bio.Bam.Header +import Bio.Prelude +import Bio.Streaming.Bytes +import Bio.Util.Numeric ( showNum ) +import Streaming hiding ( (<>) ) +import Streaming.Internal ( Stream(..) ) +import Streaming.Prelude ( each ) +import System.IO + +import qualified Streaming.Prelude as Q + +{- | Default buffer size in elements. + +Since we often want to merge many files, a read should take more time +than a seek. Assuming a rotating hard drive, this sets the sensible +buffer size to somewhat more than one MB. A smaller buffer size would +surely work on SSDs, but the large buffer doesn't hurt either. +-} +defaultBufSize :: Int +defaultBufSize = 2*1024*1024 + +streamFile :: (MonadIO m, MonadMask m) => FilePath -> (ByteStream m () -> m r) -> m r +streamFile f k = bracket (liftIO $ openBinaryFile f ReadMode) (liftIO . hClose) (k . streamHandle) +{-# INLINE streamFile #-} + +streamHandle :: MonadIO m => Handle -> ByteStream m () +streamHandle = hGetContentsN defaultBufSize +{-# INLINE streamHandle #-} + +-- | Reads 'stdin' if the filename is \"-\", else reads the named file. +streamInput :: (MonadIO m, MonadMask m) => FilePath -> (ByteStream m () -> m r) -> m r +streamInput "-" k = k (streamHandle stdin) +streamInput f k = streamFile f k +{-# INLINE streamInput #-} + +{- | Reads multiple inputs in sequence. + +Only one file is opened at a time, so they must also be consumed in +sequence. The filename \"-\" refers to stdin, if no filenames are +given, stdin is read. +-} +streamInputs :: MonadIO m => [FilePath] -> (Stream (ByteStream m) m () -> r) -> r +streamInputs [] k = k $ yields (streamHandle stdin) +streamInputs fs k = k $ mapM_ go fs + where + go "-" = yields (streamHandle stdin) + go f = yields $ do h <- liftIO $ openBinaryFile f ReadMode + streamHandle h + liftIO $ hClose h +{-# INLINE streamInputs #-} + +{- | Protects the terminal from binary junk. + +If @s@ is a 'Stream', then @protectTerm s@ throws an error if 'stdout' +is a terminal device, followed by the same 'Stream'. This is most +usefully composed with functions that might otherwise write binary data +to an interactive terminal. +-} +protectTerm :: (Functor f, MonadIO m) => Stream f m r -> Stream f m r +protectTerm str = do + t <- liftIO $ hIsTerminalDevice stdout + when t . liftIO . throwM $ ErrorCall "cowardly refusing to write binary data to terminal" + str +{-# INLINE protectTerm #-} + +{- Like 'Streaming.sequence', but parallel. + +This runs each element of a stream of actions. A configurable number of +actions are buffered and run asynchronously. +-} +psequence :: MonadIO m => Int -> Stream (Of (IO a)) m b -> Stream (Of a) m b +psequence np = go emptyQ + where + -- if the queue is full, wait for the head element to complete + go !qq s = case popQ qq of + Just (a,qq') | lengthQ qq == np -> reap a >>= wrap . (:> go qq' s) + _ -> lift (inspect s) >>= go' qq + + -- if we have room for input, we get input + go' !qq (Right (k :> s)) = liftIO (spawn k) >>= \a -> go (pushQ a qq) s + go' !qq (Left r) = goE r qq + + -- input ended, empty the queue + goE r !qq = case popQ qq of + Nothing -> pure r + Just (a,qq') -> reap a >>= wrap . (:> goE r qq') + + spawn :: IO a -> IO (MVar (Either SomeException a)) + spawn k = newEmptyMVar >>= \mv -> + forkIO (try k >>= putMVar mv) >> + return mv + + reap mv = liftIO (takeMVar mv) >>= either (liftIO . throwM) return + + +-- A very simple queue data type. +-- Invariants: q = QQ l f b --> l == length f + length b +-- --> l == 0 ==> null f + +data QQ a = QQ !Int [a] [a] + +emptyQ :: QQ a +emptyQ = QQ 0 [] [] + +lengthQ :: QQ a -> Int +lengthQ (QQ l _ _) = l + +pushQ :: a -> QQ a -> QQ a +pushQ a (QQ l [] b) = QQ (l+1) (reverse (a:b)) [] +pushQ a (QQ l f b) = QQ (l+1) f (a:b) + +popQ :: QQ a -> Maybe (a, QQ a) +popQ (QQ _ [ ] _) = Nothing +popQ (QQ l [ a ] b) = Just (a, QQ (l-1) (reverse b) []) +popQ (QQ l (a:fs) b) = Just (a, QQ (l-1) fs b) + + +mergeStreams :: (Monad m, Ord a) + => Stream (Of a) m r -> Stream (Of a) m s -> Stream (Of a) m (r, s) +mergeStreams = mergeStreamsBy compare +{-# INLINE mergeStreams #-} + +mergeStreamsOn :: (Monad m, Ord b) + => (a -> b) -> Stream (Of a) m r -> Stream (Of a) m s -> Stream (Of a) m (r, s) +mergeStreamsOn f = mergeStreamsBy (comparing f) +{-# INLINE mergeStreamsOn #-} + +mergeStreamsBy :: Monad m + => (a -> a -> Ordering) + -> Stream (Of a) m r -> Stream (Of a) m s -> Stream (Of a) m (r, s) +mergeStreamsBy cmp = go + where + go str0 str1 = case str0 of + Return r0 -> (\r1 -> (r0, r1)) <$> str1 + Effect m -> Effect $ liftM (\str -> go str str1) m + Step (a :> rest0) -> case str1 of + Return r1 -> (\r0 -> (r0, r1)) <$> str0 + Effect m -> Effect $ liftM (go str0) m + Step (b :> rest1) -> case cmp a b of + LT -> Step (a :> go rest0 str1) + EQ -> Step (a :> go rest0 str1) -- left-biased + GT -> Step (b :> go str0 rest1) +{-# INLINABLE mergeStreamsBy #-} + +-- | A general progress indicator that prints some message after a set +-- number of records have passed through. +progressGen :: MonadIO m + => (Int -> a -> String) -> Int -> (String -> IO ()) + -> Q.Stream (Q.Of a) m r -> Q.Stream (Q.Of a) m r +progressGen msg sz put = go 0 + where + go !n = lift . Q.next >=> either pure (step $ succ n) + step !n (a,s) = do when (n `mod` sz == 0) . liftIO . put $ "\27[K" ++ msg n a ++ "\r" + Q.cons a (go n s) + +-- | A simple progress indicator that prints the number of records. +progressNum :: MonadIO m + => String -> Int -> (String -> IO ()) + -> Q.Stream (Q.Of a) m r -> Q.Stream (Q.Of a) m r +progressNum msg = progressGen (\n _ -> msg ++ " " ++ showNum n) + +-- | A simple progress indicator that prints a position every set number +-- of passed records. +progressPos :: MonadIO m + => (a -> (Refseq, Int)) -> String -> Refs -> Int -> (String -> IO ()) + -> Q.Stream (Q.Of a) m r -> Q.Stream (Q.Of a) m r +progressPos f msg refs = + progressGen $ \_ a -> let (!rs1, !po1) = f a + !nm = unpack . sq_name $ getRef refs rs1 + in msg ++ " " ++ nm ++ ":" ++ showNum po1 diff --git a/Bio/Streaming/Bgzf.hs b/Bio/Streaming/Bgzf.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming/Bgzf.hs @@ -0,0 +1,331 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +-- | Buffer builder to assemble Bgzf blocks. The idea is to serialize +-- stuff (BAM and BCF) into a buffer, then bgzf chunks from the buffer. +-- We use a large buffer, and we always make sure there is plenty of +-- space in it (to avoid redundant checks). + +module Bio.Streaming.Bgzf ( + bgunzip, + getBgzfHdr, + BB(..), + newBuffer, + fillBuffer, + expandBuffer, + encodeBgzf, + BgzfTokens(..), + BclArgs(..), + BclSpecialType(..), + loop_dec_int, + loop_bcl_special + ) where + +import Bio.Prelude +import Bio.Streaming +import Foreign.C.Types ( CInt(..) ) +import Foreign.Marshal.Utils ( copyBytes, with ) + +import qualified Bio.Streaming.Bytes as S +import qualified Data.ByteString as B +import qualified Data.ByteString.Internal as B +import qualified Data.ByteString.Unsafe as B +import qualified Data.Vector.Storable as V + +{-| Decompresses a bgzip stream. Individual chunks are decompressed in + parallel. Leftovers are discarded (some compressed HETFA files + appear to have junk at the end). -} + +bgunzip :: MonadIO m => ByteStream m r -> ByteStream m r +bgunzip s = do + np <- liftIO $ getNumCapabilities + S.fromChunks $ psequence (2*np) $ lift (getBgzfHdr s) >>= go + where + go (Nothing, _hdr, s1) = lift (S.effects s1) + go (Just bsize, _hdr, s1) = do + blk :> s2 <- lift $ S.splitAt' bsize s1 + wrap (decompressChunk blk :> (lift (getBgzfHdr s2) >>= go)) +{-# INLINABLE bgunzip #-} + +getBgzfHdr :: Monad m => ByteStream m r -> m (Maybe Int, B.ByteString, ByteStream m r) +getBgzfHdr s0 = do + hdr :> s1 <- S.splitAt' 18 s0 + if or [ B.length hdr < 18 + , B.index hdr 0 /= 139 + , B.index hdr 1 /= 31 + , B.index hdr 3 .&. 4 /= 4 + , B.index hdr 10 /= 6 + , B.index hdr 11 /= 0 + , B.index hdr 12 /= 66 + , B.index hdr 13 /= 67 ] + then return (Nothing, hdr, s1) + else do + let bsize = fromIntegral (B.index hdr 16) + fromIntegral (B.index hdr 17) `shiftL` 8 - 16 + return (Just bsize, hdr, s1) +{-# INLINE getBgzfHdr #-} + +-- | We manage a large buffer (multiple megabytes), of which we fill an +-- initial portion. We remember the size, the used part, and two marks +-- where we later fill in sizes for the length prefixed BAM or BCF +-- records. We move the buffer down when we yield a piece downstream, +-- and when we run out of space, we simply move to a new buffer. +-- Garbage collection should take care of the rest. Unused 'mark' must +-- be set to (maxBound::Int) so it doesn't interfere with flushing. + +data BB = BB { buffer :: {-# UNPACK #-} !(ForeignPtr Word8) + , size :: {-# UNPACK #-} !Int -- total size of buffer + , off :: {-# UNPACK #-} !Int -- offset of active portion + , used :: {-# UNPACK #-} !Int -- used portion (inactive & active) + , mark :: {-# UNPACK #-} !Int -- offset of mark + , mark2 :: {-# UNPACK #-} !Int } -- offset of mark2 + +-- | Things we are able to encode. Taking inspiration from +-- binary-serialise-cbor, we define these as a lazy list-like thing and +-- consume it in a interpreter. + +data BgzfTokens = TkWord32 {-# UNPACK #-} !Word32 BgzfTokens -- a 4-byte int + | TkWord16 {-# UNPACK #-} !Word16 BgzfTokens -- a 2-byte int + | TkWord8 {-# UNPACK #-} !Word8 BgzfTokens -- a byte + | TkFloat {-# UNPACK #-} !Float BgzfTokens -- a float + | TkDouble {-# UNPACK #-} !Double BgzfTokens -- a double + | TkString {-# UNPACK #-} !B.ByteString BgzfTokens -- a raw string + | TkDecimal {-# UNPACK #-} !Int BgzfTokens -- roughly ':%d' + + | TkSetMark BgzfTokens -- sets the first mark + | TkEndRecord BgzfTokens -- completes a BAM record + | TkEndRecordPart1 BgzfTokens -- completes part 1 of a BCF record + | TkEndRecordPart2 BgzfTokens -- completes part 2 of a BCF record + | TkEnd -- nothing more, for now + + | TkBclSpecial !BclArgs BgzfTokens + | TkLowLevel {-# UNPACK #-} !Int (BB -> IO BB) BgzfTokens + +data BclSpecialType = BclNucsBin | BclNucsAsc | BclNucsAscRev | BclNucsWide + | BclQualsBin | BclQualsAsc | BclQualsAscRev + +data BclArgs = BclArgs BclSpecialType + {-# UNPACK #-} !(V.Vector Word8) -- bcl matrix + {-# UNPACK #-} !Int -- stride + {-# UNPACK #-} !Int -- first cycle + {-# UNPACK #-} !Int -- last cycle + {-# UNPACK #-} !Int -- cluster index + +-- | Creates a buffer. +newBuffer :: Int -> IO BB +newBuffer sz = mallocForeignPtrBytes sz >>= \ar -> return $ BB ar sz 0 0 maxBound maxBound + +-- | Creates a new buffer, copying the active content from an old one, +-- with higher capacity. The size of the new buffer is twice the free +-- space in the old buffer, but at least @minsz@. +expandBuffer :: Int -> BB -> IO BB +expandBuffer minsz b = do + let sz' = max (2 * (size b - used b)) minsz + arr1 <- mallocForeignPtrBytes sz' + withForeignPtr arr1 $ \d -> + withForeignPtr (buffer b) $ \s -> + copyBytes d (plusPtr s (off b)) (used b - off b) + return BB{ buffer = arr1 + , size = sz' + , off = 0 + , used = used b - off b + , mark = if mark b == maxBound then maxBound else mark b - off b + , mark2 = if mark2 b == maxBound then maxBound else mark2 b - off b } + +compressChunk :: Int -> ForeignPtr Word8 -> Int -> Int -> IO B.ByteString +compressChunk lv fptr off slen = + withForeignPtr fptr $ \ptr -> + B.createAndTrim 65536 $ \buf -> + with 65536 $ \p_len -> do + rc <- compress_chunk buf p_len (plusPtr ptr off) (fromIntegral slen) (fromIntegral lv) + when (rc /= 0 && rc /= 1) . error $ "compress_chunk failed: " ++ show rc + fromIntegral <$> peek p_len + +decompressChunk :: B.ByteString -> IO B.ByteString +decompressChunk ck = + B.unsafeUseAsCString ck $ \psrc -> + peekByteOff psrc (B.length ck - 4) >>= \dlen -> + B.create (fromIntegral (dlen::Word32)) $ \pdest -> do + rc <- decompress_chunk pdest (fromIntegral dlen) (castPtr psrc) (fromIntegral $ B.length ck) + when (rc /= 0) . error $ "decompress_chunk failed: " ++ show rc + + +-- | Expand a chain of tokens into a buffer, sending finished pieces +-- downstream as soon as possible. +encodeBgzf :: MonadIO m => Int -> Stream (Of (Endo BgzfTokens)) m b -> S.ByteStream m b +encodeBgzf lv str = do + np <- liftIO $ getNumCapabilities + bb <- liftIO $ newBuffer (1024*1024) + S.fromChunks $ psequence (2*np) $ lift (inspect str) >>= go bb + where + go :: MonadIO m + => BB + -> Either b (Of (Endo BgzfTokens) (Stream (Of (Endo BgzfTokens)) m b)) + -> Stream (Of (IO Bytes)) m b + go bb0 (Left r) = final_flush bb0 r + go bb0 (Right (f :> s)) + -- initially, we make sure we have reasonable space. this may not be enough. + | size bb0 - used bb0 < 1024 = liftIO (expandBuffer (1024*1024) bb0) >>= \bb' -> go' bb' (appEndo f TkEnd) s + | otherwise = go' bb0 (appEndo f TkEnd) s + + go' bb0 tk s = liftIO (fillBuffer bb0 tk) >>= \(bb',tk') -> flush_blocks tk' bb' s + + -- We can flush anything that is between 'off' and the lower of 'mark' + -- and 'used'. When done, we bump 'off'. + flush_blocks tk bb s + | min (mark bb) (used bb) - off bb < maxBlockSize = + case tk of TkEnd -> lift (inspect s) >>= go bb + _ -> -- we arrive here because we ran out of buffer space, so we expand it. + liftIO (expandBuffer (1024*1024) bb) >>= \bb' -> go' bb' tk s + | otherwise = + wrap $ compressChunk lv (buffer bb) (off bb) maxBlockSize + :> flush_blocks tk bb { off = off bb + maxBlockSize } s + + final_flush bb r + | used bb > off bb = + wrap $ compressChunk lv (buffer bb) (off bb) (used bb - off bb) + :> wrap (return bgzfEofMarker :> pure r) + | otherwise = + wrap (return bgzfEofMarker :> pure r) + + -- maximum block size for Bgzf: 64k with some room for + -- headers and uncompressible stuff + maxBlockSize = 65478 + + +fillBuffer :: BB -> BgzfTokens -> IO (BB, BgzfTokens) +fillBuffer bb0 tk = withForeignPtr (buffer bb0) (\p -> go_slowish p bb0 tk) + where + go_slowish p bb = go_fast p bb (used bb) + + go_fast p bb use tk1 = case tk1 of + -- no space? not our job. + _ | size bb - use < 1024 -> return (bb { used = use },tk1) + + -- the actual end. + TkEnd -> return (bb { used = use },tk1) + + -- I'm cheating. This stuff works only if the platform allows + -- unaligned accesses, is little-endian and uses IEEE floats. + -- It's true on i386 and ix86_64. + TkWord32 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 4) tk' + + TkWord16 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 2) tk' + + TkWord8 x tk' -> do pokeByteOff p use x + go_fast p bb (use + 1) tk' + + TkFloat x tk' -> do pokeByteOff p use x + go_fast p bb (use + 4) tk' + + TkDouble x tk' -> do pokeByteOff p use x + go_fast p bb (use + 8) tk' + + TkString s tk' + -- Too big, can't handle. By returning with unfinished + -- business, we will get progressively bigger buffers and + -- eventually handle it. + | B.length s > size bb - use -> return (bb { used = use },tk1) + + | otherwise -> do let ln = B.length s + B.unsafeUseAsCString s $ \q -> + copyBytes (p `plusPtr` use) q ln + go_fast p bb (use + ln) tk' + + TkDecimal x tk' -> do ln <- int_loop (p `plusPtr` use) (fromIntegral x) + go_fast p bb (use + fromIntegral ln) tk' + + TkSetMark tk' -> go_slowish p bb { used = use + 4, mark = use } tk' + + TkEndRecord tk' -> do let !l = use - mark bb - 4 + pokeByteOff p (mark bb) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark = maxBound } tk' + + TkEndRecordPart1 tk' -> do let !l = use - mark bb - 4 + pokeByteOff p (mark bb - 4) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark2 = use } tk' + + TkEndRecordPart2 tk' -> do let !l = use - mark2 bb + pokeByteOff p (mark bb) (fromIntegral l :: Word32) + go_slowish p bb { used = use, mark = maxBound } tk' + + + TkBclSpecial special_args tk' -> do + l <- loop_bcl_special (p `plusPtr` use) special_args + go_fast p bb (use + l) tk' + + TkLowLevel minsize proc tk' + | size bb - use < minsize -> return (bb { used = use },tk1) + | otherwise -> do bb' <- proc bb { used = use } + go_slowish p bb' tk' + +-- | The EOF marker for BGZF files. +-- This is just an empty string compressed as BGZF. Appended to BAM +-- files to indicate their end. +bgzfEofMarker :: Bytes +bgzfEofMarker = "\x1f\x8b\x8\x4\0\0\0\0\0\xff\x6\0\x42\x43\x2\0\x1b\0\x3\0\0\0\0\0\0\0\0\0" + +loop_dec_int :: Ptr Word8 -> Int -> IO Int +loop_dec_int p i = fromIntegral <$> int_loop p (fromIntegral i) + +loop_bcl_special :: Ptr Word8 -> BclArgs -> IO Int +loop_bcl_special p (BclArgs tp vec stride u v i) = + + V.unsafeWith vec $ \q -> case tp of + BclNucsBin -> do + nuc_loop p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ (v - u + 2) `div` 2 + + BclNucsWide -> do + nuc_loop_wide p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + + BclNucsAsc -> do + nuc_loop_asc p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + + BclNucsAscRev -> do + nuc_loop_asc_rev p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + + BclQualsBin -> do + qual_loop p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + + BclQualsAsc -> do + qual_loop_asc p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + + BclQualsAscRev -> do + qual_loop_asc_rev p (fromIntegral stride) (plusPtr q i) (fromIntegral u) (fromIntegral v) + return $ v - u + 1 + +foreign import ccall unsafe "nuc_loop" + nuc_loop :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "nuc_loop_wide" + nuc_loop_wide :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "nuc_loop_asc" + nuc_loop_asc :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "nuc_loop_asc_rev" + nuc_loop_asc_rev :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "qual_loop" + qual_loop :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "qual_loop_asc" + qual_loop_asc :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "qual_loop_asc_rev" + qual_loop_asc_rev :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> CInt -> IO () + +foreign import ccall unsafe "int_loop" + int_loop :: Ptr Word8 -> CInt -> IO CInt + +foreign import ccall unsafe "compress_chunk" + compress_chunk :: Ptr Word8 -> Ptr CInt -> Ptr Word8 -> CInt -> CInt -> IO CInt + +foreign import ccall unsafe "decompress_chunk" + decompress_chunk :: Ptr Word8 -> CInt -> Ptr Word8 -> CInt -> IO CInt + diff --git a/Bio/Streaming/Bytes.hs b/Bio/Streaming/Bytes.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming/Bytes.hs @@ -0,0 +1,780 @@ +{-# LANGUAGE Rank2Types, TypeFamilies #-} + +-- This library emulates Data.ByteStream.Lazy but includes a monadic element +-- and thus at certain points uses a `Stream`/`FreeT` type in place of lists. + +-- | +-- Module : ByteStream +-- Copyright : (c) Don Stewart 2006 +-- (c) Duncan Coutts 2006-2011 +-- (c) Michael Thompson 2015 +-- (c) Udo Stenzel 2018 +-- License : BSD-style +-- +-- Maintainer : u.stenzel@web.de +-- Stability : experimental +-- Portability : portable +-- +-- See the simple examples of use . +-- We begin with a slight modification of the documentation to "Data.ByteStream.Lazy": +-- +-- A time and space-efficient implementation of effectful byte streams +-- using a stream of packed 'Word8' arrays, suitable for high performance +-- use, both in terms of large data quantities, or high speed +-- requirements. ByteStreams are encoded as streams of strict chunks +-- of bytes. +-- +-- A key feature of ByteStreams is the means to manipulate large or +-- unbounded streams of data without requiring the entire sequence to be +-- resident in memory. To take advantage of this you have to write your +-- functions in a streaming style, e.g. classic pipeline composition. The +-- default I\/O chunk size is 32k, which should be good in most circumstances. +-- +-- Some operations, such as 'concat', 'append', 'reverse' and 'cons', have +-- better complexity than their "Data.ByteStream" equivalents, due to +-- optimisations resulting from the list spine structure. For other +-- operations streaming, like lazy, ByteStreams are usually within a few percent of +-- strict ones. +-- +-- This module is intended to be imported @qualified@, to avoid name +-- clashes with "Prelude" functions. eg. +-- +-- > import qualified Bio.Streaming.Bytes as B +-- +-- Original GHC implementation by Bryan O\'Sullivan. +-- Rewritten to use 'Data.Array.Unboxed.UArray' by Simon Marlow. +-- Rewritten to support slices and use 'Foreign.ForeignPtr.ForeignPtr' +-- by David Roundy. +-- Rewritten again and extended by Don Stewart and Duncan Coutts. +-- Lazy variant by Duncan Coutts and Don Stewart. +-- Streaming variant by Michael Thompson, following the ideas of Gabriel Gonzales' pipes-bytestring +-- Adapted for use in biohazard by Udo Stenzel. +-- +module Bio.Streaming.Bytes ( + -- * The @ByteStream@ type + ByteStream(..) + + -- * Introducing and eliminating 'ByteStream's + , empty -- empty :: ByteStream m () + , singleton -- singleton :: Monad m => Word8 -> ByteStream m () + , fromLazy -- fromLazy :: Monad m => ByteStream -> ByteStream m () + , fromChunks -- fromChunks :: Monad m => Stream (Of Bytes) m r -> ByteStream m r + , toLazy -- toLazy :: Monad m => ByteStream m () -> m ByteStream + , toStrict -- toStrict :: Monad m => ByteStream m () -> m ByteStream + , effects + , mwrap + + -- * Basic interface + , cons -- cons :: Monad m => Word8 -> ByteStream m r -> ByteStream m r + , nextByte -- nextByte :: Monad m => ByteStream m r -> m (Either r (Word8, ByteStream m r)) + , nextByteOff -- nextByteOff :: Monad m => ByteStream m r -> m (Either r (Word8, Int64, ByteStream m r)) + + -- * Substrings + + -- ** Breaking strings + , break -- break :: Monad m => (Word8 -> Bool) -> ByteStream m r -> ByteStream m (ByteStream m r) + , drop -- drop :: Monad m => GHC.Int.Int64 -> ByteStream m r -> ByteStream m r + , dropWhile + , splitAt -- splitAt :: Monad m => GHC.Int.Int64 -> ByteStream m r -> ByteStream m (ByteStream m r) + , splitAt' -- splitAt' :: Monad m => Int -> ByteStream m r -> m (Of Bytes (ByteStream m r)) + , trim + + -- ** Breaking into many substrings + , lines + , lines' + + -- ** Special folds + , concat -- concat :: Monad m => Stream (ByteStream m) m r -> ByteStream m r + + -- * Builders + , toByteStream + , toByteStreamWith + , concatBuilders + + -- * I\/O with 'ByteStream's + + -- ** Files + , withOutputFile + , writeFile -- writeFile :: FilePath -> ByteStream IO r -> IO r + + -- ** I\/O with Handles + , hGetContents -- hGetContents :: Handle -> ByteStream IO () + , hGetContentsN -- hGetContentsN :: Int -> Handle -> ByteStream IO () + , hPut -- hPut :: Handle -> ByteStream IO r -> IO r + + -- * Simple chunkwise operations + , nextChunk + , nextChunkOff + , consChunk -- :: Bytes -> ByteStream m r -> ByteStream m r + , consChunkOff -- :: Bytes -> Int64 -> ByteStream m r -> ByteStream m r + , chunk + , copy + , mapChunksM_ + + -- * compression support + , gzip + , gunzip + , gunzipWith + + ) where + +import Bio.Prelude hiding (break,concat,drop,dropWhile,lines,splitAt,writeFile,empty,loop) +import Data.ByteString.Builder.Internal + (Builder,builder,runBuilder,runBuilderWith,bufferSize + ,AllocationStrategy,ChunkIOStream(..),buildStepToCIOS + ,byteStringFromBuffer,safeStrategy,defaultChunkSize) +import GHC.Exts (SpecConstrAnnotation(..)) +import Streaming (Of(..),Identity(..),destroy) +import Streaming.Internal (Stream (..)) +import System.Directory (renameFile) + +import qualified Codec.Compression.Zlib.Internal as Z +import qualified Data.ByteString as B +import qualified Data.ByteString.Internal as B +import qualified Data.ByteString.Lazy.Internal as L (foldrChunks,ByteString(..),smallChunkSize,defaultChunkSize) +import qualified Data.ByteString.Unsafe as B +import qualified Streaming.Prelude as Q + +-- | A space-efficient representation of a succession of 'Word8' vectors, supporting many +-- efficient operations. +-- +-- An effectful 'ByteStream' contains 8-bit bytes, or by using certain +-- operations can be interpreted as containing 8-bit characters. It +-- also contains an offset, which will be needed to track the virtual +-- offsets in the BGZF decode. + +data ByteStream m r = + Empty r + | Chunk {-# UNPACK #-} !Bytes {-# UNPACK #-} !Int64 (ByteStream m r) + | Go (m (ByteStream m r)) + +instance Monad m => Functor (ByteStream m) where + fmap f x = case x of + Empty a -> Empty (f a) + Chunk bs o bss -> Chunk bs o (fmap f bss) + Go mbss -> Go (liftM (fmap f) mbss) + +instance Monad m => Applicative (ByteStream m) where + pure = Empty + {-# INLINE pure #-} + bf <*> bx = do {f <- bf; x <- bx; Empty (f x)} + {-# INLINE (<*>) #-} + (*>) = (>>) + {-# INLINE (*>) #-} + +instance Monad m => Monad (ByteStream m) where + return = Empty + {-# INLINE return #-} + x0 >> y = loop SPEC x0 where + loop !_ x = case x of -- this seems to be insanely effective + Empty _ -> y + Chunk a o b -> Chunk a o (loop SPEC b) + Go m -> Go (liftM (loop SPEC) m) + {-# INLINEABLE (>>) #-} + x >>= f = + loop SPEC2 x where -- unlike >> this SPEC seems pointless + loop !_ y = case y of + Empty a -> f a + Chunk bs o bss -> Chunk bs o (loop SPEC bss) + Go mbss -> Go (liftM (loop SPEC) mbss) + {-# INLINEABLE (>>=) #-} + +instance MonadIO m => MonadIO (ByteStream m) where + liftIO io = Go (liftM Empty (liftIO io)) + {-# INLINE liftIO #-} + +instance MonadTrans ByteStream where + lift ma = Go $ liftM Empty ma + {-# INLINE lift #-} + +instance (r ~ ()) => IsString (ByteStream m r) where + fromString = chunk . fromString + {-# INLINE fromString #-} + +instance (m ~ Identity, Show r) => Show (ByteStream m r) where + show bs0 = case bs0 of -- the implementation this instance deserves ... + Empty r -> "Empty (" ++ show r ++ ")" + Go (Identity bs') -> "Go (Identity (" ++ show bs' ++ "))" + Chunk bs'' o bs -> "Chunk " ++ show bs'' ++ " " ++ show o ++ " (" ++ show bs ++ ")" + +instance (Semigroup r, Monad m) => Semigroup (ByteStream m r) where + (<>) = liftM2 (<>) + {-# INLINE (<>) #-} + +instance (Semigroup r, Monoid r, Monad m) => Monoid (ByteStream m r) where + mempty = Empty mempty + {-# INLINE mempty #-} + mappend = (<>) + {-# INLINE mappend #-} + + +data SPEC = SPEC | SPEC2 +{-# ANN type SPEC ForceSpecConstr #-} + +-- -------------------------------------------------------------------------- + +-- | Smart constructor for 'Chunk'. +consChunk :: Bytes -> ByteStream m r -> ByteStream m r +consChunk c = consChunkOff c 0 +{-# INLINE consChunk #-} + +consChunkOff :: Bytes -> Int64 -> ByteStream m r -> ByteStream m r +consChunkOff c@(B.PS _ _ len) off cs + | len == 0 = cs + | otherwise = Chunk c off cs +{-# INLINE consChunkOff #-} + +-- | Yield-style smart constructor for 'Chunk'. +chunk :: Bytes -> ByteStream m () +chunk bs = consChunk bs empty +{-# INLINE chunk #-} + + +{- | Reconceive an effect that results in an effectful bytestring as an effectful bytestring. + Compare Streaming.mwrap. The closes equivalent of + +>>> Streaming.wrap :: f (Stream f m r) -> Stream f m r + + is here @consChunk@. @mwrap@ is the smart constructor for the internal @Go@ constructor. +-} +mwrap :: m (ByteStream m r) -> ByteStream m r +mwrap = Go +{-# INLINE mwrap #-} + +-- | Construct a succession of chunks from its Church encoding (compare @GHC.Exts.build@) +materialize :: (forall x . (r -> x) -> (Bytes -> Int64 -> x -> x) -> (m x -> x) -> x) + -> ByteStream m r +materialize phi = phi Empty Chunk Go +{-# INLINE[0] materialize #-} + +-- | Resolve a succession of chunks into its Church encoding; this is +-- not a safe operation; it is equivalent to exposing the constructors +dematerialize :: Monad m + => ByteStream m r + -> (forall x . (r -> x) -> (Bytes -> Int64 -> x -> x) -> (m x -> x) -> x) +dematerialize x0 nil con fin = loop SPEC x0 + where + loop !_ x = case x of + Empty r -> nil r + Chunk b o bs -> con b o (loop SPEC bs ) + Go ms -> fin (liftM (loop SPEC) ms) +{-# INLINE [1] dematerialize #-} + +{-# RULES + "dematerialize/materialize" forall (phi :: forall b . (r -> b) -> (Bytes -> Int64 -> b -> b) -> (m b -> b) -> b) . dematerialize (materialize phi) = phi ; + #-} +------------------------------------------------------------------------ + +copy :: Monad m => ByteStream m r -> ByteStream (ByteStream m) r +copy = loop where + loop str = case str of + Empty r -> Empty r + Go m -> Go (liftM loop (lift m)) + Chunk bs o rest -> Chunk bs o (Go (Chunk bs o (Empty (loop rest)))) +{-# INLINABLE copy #-} + +-- | /O(n)/ Concatenate a stream of byte streams. +concat :: Monad m => Stream (ByteStream m) m r -> ByteStream m r +concat x = destroy x join Go Empty +{-# INLINE concat #-} + +-- | Perform the effects contained in an effectful bytestring, ignoring the bytes. +effects :: Monad m => ByteStream m r -> m r +effects bs = case bs of + Empty r -> return r + Go m -> m >>= effects + Chunk _ _ rest -> effects rest +{-# INLINABLE effects #-} + + +-- ----------------------------------------------------------------------------- +-- Introducing and eliminating 'ByteStream's + +{-| /O(1)/ The empty 'ByteStream' -- i.e. @return ()@ Note that @ByteStream m w@ is + generally a monoid for monoidal values of @w@, like @()@ +-} +empty :: ByteStream m () +empty = Empty () +{-# INLINE empty #-} + +{-| /O(1)/ Yield a 'Word8' as a minimal 'ByteStream' +-} +singleton :: Word8 -> ByteStream m () +singleton w = Chunk (B.singleton w) 0 (Empty ()) +{-# INLINE singleton #-} + +-- | /O(c)/ Converts a byte stream into a stream of individual strict bytestrings. +-- This of course exposes the internal chunk structure. +toChunks :: Monad m => ByteStream m r -> Stream (Of Bytes) m r +toChunks bs = dematerialize bs return (\b _ mx -> Step (b :> mx)) Effect +{-# INLINE toChunks #-} + +mapChunksM_ :: Monad m => (Bytes -> m ()) -> ByteStream m r -> m r +mapChunksM_ f bs = dematerialize bs return (\c _ k -> f c >> k) join +{-# INLINE mapChunksM_ #-} + + +-- | /O(c)/ Converts a stream of strict bytestrings into a byte stream. +fromChunks :: Monad m => Stream (Of Bytes) m r -> ByteStream m r +fromChunks bs = destroy bs + (\(b :> mx) !i -> Chunk b i (mx (i + fromIntegral (B.length b)))) + (\k !i -> Go (k >>= \f -> return (f i))) + (\r !_ -> return r) 0 +{-# INLINE fromChunks #-} + +{-| /O(n)/ Convert a monadic byte stream into a single strict 'ByteStream', + retaining the return value of the original pair. This operation is + for use with 'mapped'. + +> mapped R.toStrict :: Monad m => Stream (ByteStream m) m r -> Stream (Of ByteStream) m r + + It is subject to all the objections one makes to Data.ByteStream.Lazy 'toStrict'; + all of these are devastating. +-} +toStrict :: Monad m => ByteStream m r -> m (Of Bytes r) +toStrict bs = do + (bss :> r) <- Q.toList (toChunks bs) + return $ (B.concat bss :> r) +{-# INLINE toStrict #-} + +{-| /O(c)/ Transmute a pseudo-pure lazy bytestring to its representation + as a monadic stream of chunks. + +>>> Q.putStrLn $ Q.fromLazy "hi" +hi +>>> Q.fromLazy "hi" +Chunk "hi" (Empty (())) -- note: a 'show' instance works in the identity monad +>>> Q.fromLazy $ BL.fromChunks ["here", "are", "some", "chunks"] +Chunk "here" (Chunk "are" (Chunk "some" (Chunk "chunks" (Empty (()))))) + +-} +fromLazy :: LazyBytes -> ByteStream m () +fromLazy = L.foldrChunks consChunk empty +{-# INLINE fromLazy #-} + +{-| /O(n)/ Convert an effectful byte stream into a single lazy 'ByteStream' + with the same internal chunk structure, retaining the original + return value. + + This is the canonical way of breaking streaming (@toStrict@ and the + like are far more demonic). Essentially one is dividing the interleaved + layers of effects and bytes into one immense layer of effects, + followed by the memory of the succession of bytes. + + Because one preserves the return value, @toLazy@ is a suitable argument + for 'Streaming.mapped' + +> B.mapped Q.toLazy :: Stream (ByteStream m) m r -> Stream (Of LazyBytes) m r + +>>> Q.toLazy "hello" +"hello" :> () +>>> B.toListM $ traverses Q.toLazy $ Q.lines "one\ntwo\nthree\nfour\nfive\n" +["one","two","three","four","five",""] -- [LazyBytes] + +-} +toLazy :: Monad m => ByteStream m r -> m (Of LazyBytes r) +toLazy bs0 = dematerialize bs0 + (\r -> return (L.Empty :> r)) + (\b _ mx -> do + (bs :> x) <- mx + return $ L.Chunk b bs :> x + ) + join +{-# INLINE toLazy #-} + +-- | /O(1)/ 'cons' is analogous to '(:)' for lists. +cons :: Word8 -> ByteStream m r -> ByteStream m r +cons c cs = Chunk (B.singleton c) 0 cs +{-# INLINE cons #-} + +-- | /O(1)/ Extract the head and tail of a 'ByteStream', or its return value +-- if it is empty. This is the \'natural\' uncons for an effectful byte stream. +nextByte :: Monad m => ByteStream m r -> m (Either r (Word8, ByteStream m r)) +nextByte = liftM (either Left (\(a,_,b) -> Right (a,b))) . nextByteOff +{-# INLINE nextByte #-} + +nextByteOff :: Monad m => ByteStream m r -> m (Either r (Word8, Int64, ByteStream m r)) +nextByteOff (Empty r) = return (Left r) +nextByteOff (Chunk c o cs) + = if B.null c + then nextByteOff cs + else return $ Right (B.unsafeHead c, o + , if B.length c == 1 + then cs + else Chunk (B.unsafeTail c) (o+1) cs) +nextByteOff (Go m) = m >>= nextByteOff +{-# INLINABLE nextByteOff #-} + +nextChunk :: Monad m => ByteStream m r -> m (Either r (Bytes, ByteStream m r)) +nextChunk = liftM (either Left (\(a,_,b) -> Right (a,b))) . nextChunkOff +{-# INLINE nextChunk #-} + +nextChunkOff :: Monad m => ByteStream m r -> m (Either r (Bytes, Int64, ByteStream m r)) +nextChunkOff (Empty r) = return (Left r) +nextChunkOff (Go m) = m >>= nextChunkOff +nextChunkOff (Chunk c o cs) + | B.null c = nextChunkOff cs + | otherwise = return (Right (c,o,cs)) +{-# INLINABLE nextChunkOff #-} + +{-| /O(n\/c)/ 'drop' @n xs@ returns the suffix of @xs@ after the first @n@ + elements, or @[]@ if @n > 'length' xs@. + +>>> Q.putStrLn $ Q.drop 6 "Wisconsin" +sin +>>> Q.putStrLn $ Q.drop 16 "Wisconsin" + +>>> +-} +drop :: Monad m => Int64 -> ByteStream m r -> ByteStream m r +drop i p | i <= 0 = p +drop i cs0 = drop' i cs0 + where drop' 0 cs = cs + drop' _ (Empty r) = Empty r + drop' n (Chunk c o cs) = + if n < fromIntegral (B.length c) + then Chunk (B.drop (fromIntegral n) c) (o+n) cs + else drop' (n - fromIntegral (B.length c)) cs + drop' n (Go m) = Go (liftM (drop' n) m) +{-# INLINABLE drop #-} + + +{-| /O(n\/c)/ 'splitAt' @n xs@ is equivalent to @('take' n xs, 'drop' n xs)@. + +>>> rest <- Q.putStrLn $ Q.splitAt 3 "therapist is a danger to good hyphenation, as Knuth notes" +the +>>> Q.putStrLn $ Q.splitAt 19 rest +rapist is a danger + +-} +splitAt :: Monad m => Int64 -> ByteStream m r -> ByteStream m (ByteStream m r) +splitAt i cs0 | i <= 0 = Empty cs0 +splitAt i cs0 = go i cs0 + where go 0 cs = Empty cs + go _ (Empty r) = Empty (Empty r) + go n (Chunk c o cs) = + if n < fromIntegral (B.length c) + then Chunk (B.take (fromIntegral n) c) o $ + Empty (Chunk (B.drop (fromIntegral n) c) (o+n) cs) + else Chunk c o (go (n - fromIntegral (B.length c)) cs) + go n (Go m) = Go (liftM (go n) m) +{-# INLINABLE splitAt #-} + +-- | Strictly splits off a piece. This breaks streaming, so reserve its +-- use for small strings or when conversion to strict 'Bytes' is needed +-- anyway. +splitAt' :: Monad m => Int -> ByteStream m r -> m (Of Bytes (ByteStream m r)) +splitAt' i cs0 | i <= 0 = return $! B.empty :> cs0 +splitAt' i cs0 = go i [] cs0 + where go 0 acc cs = return $! B.concat (reverse acc) :> cs + go _ acc (Empty r) = return $! B.concat (reverse acc) :> Empty r + go n acc (Chunk c o cs) = + if n < B.length c + then return $! B.concat (reverse (B.take n c : acc)) + :> Chunk (B.drop n c) (o + fromIntegral n) cs + else go (n - B.length c) (c:acc) cs + go n acc (Go m) = m >>= go n acc +{-# INLINABLE splitAt' #-} + +trim :: Monad m => Int64 -> ByteStream m () -> ByteStream m () +trim eoff = go + where + go (Empty _) = Empty () + go (Go m) = lift m >>= go + go (Chunk c o s) | o < eoff = Chunk c o (go s) + | otherwise = Empty () +{-# INLINABLE trim #-} + +-- | 'dropWhile' @p xs@ returns the suffix remaining after 'takeWhile' @p xs@. +dropWhile :: Monad m => (Word8 -> Bool) -> ByteStream m r -> ByteStream m r +dropWhile pr = drop' where + drop' bs = case bs of + Empty r -> Empty r + Go m -> Go (liftM drop' m) + Chunk c o cs -> case findIndexOrEnd (not.pr) c of + 0 -> Chunk c o cs + n | n < B.length c -> Chunk (B.drop n c) (o + fromIntegral n) cs + | otherwise -> drop' cs +{-# INLINABLE dropWhile #-} + +-- | 'break' @p@ is equivalent to @'span' ('not' . p)@. +break :: Monad m => (Word8 -> Bool) -> ByteStream m r -> ByteStream m (ByteStream m r) +break f cs0 = break' cs0 + where break' (Empty r) = Empty (Empty r) + break' (Chunk c o cs) = + case findIndexOrEnd f c of + 0 -> Empty (Chunk c o cs) + n | n < B.length c -> Chunk (B.take n c) o $ + Empty (Chunk (B.drop n c) (o + fromIntegral n) cs) + | otherwise -> Chunk c o (break' cs) + break' (Go m) = Go (liftM break' m) +{-# INLINABLE break #-} + +{- | Read entire handle contents /lazily/ into a 'ByteStream'. Chunks + are read on demand, in at most @k@-sized chunks. It does not block + waiting for a whole @k@-sized chunk, so if less than @k@ bytes are + available then they will be returned immediately as a smaller chunk. + + The handle is closed on EOF. + + Note: the 'Handle' should be placed in binary mode with + 'System.IO.hSetBinaryMode' for 'hGetContentsN' to + work correctly. +-} +hGetContentsN :: MonadIO m => Int -> Handle -> ByteStream m () +hGetContentsN k h = loop 0 + where + loop !o = do + c <- liftIO (B.hGetSome h k) + -- only blocks if there is no data available + if B.null c + then Empty () + else Chunk c o (loop (o + fromIntegral (B.length c))) +{-# INLINABLE hGetContentsN #-} -- very effective inline pragma + +{-| Read entire handle contents /lazily/ into a 'ByteStream'. Chunks + are read on demand, using the default chunk size. + + Note: the 'Handle' should be placed in binary mode with + 'System.IO.hSetBinaryMode' for 'hGetContents' to + work correctly. +-} +hGetContents :: MonadIO m => Handle -> ByteStream m () +hGetContents = hGetContentsN defaultChunkSize +{-# INLINE hGetContents #-} + + +withOutputFile :: (MonadIO m, MonadMask m) => FilePath -> (Handle -> m a) -> m a +withOutputFile "-" k = k stdout +withOutputFile f k = bracket (liftIO $ openBinaryFile (f++".#~#") WriteMode) (liftIO . hClose) $ \hdl -> + k hdl >>= \r -> liftIO (renameFile (f++".#~#") f) >> return r +{-# INLINE withOutputFile #-} + +{-| Writes a 'ByteStream' to a file. Actually writes to a temporary + file and renames it on successful completion. The filename \"-\" + causes it to write to stdout instead. + -} +writeFile :: (MonadIO m, MonadMask m) => FilePath -> ByteStream m r -> m r +writeFile f str = withOutputFile f $ \hdl -> hPut hdl str +{-# INLINE writeFile #-} + +-- | Outputs a 'ByteStream' to the specified 'Handle'. +hPut :: MonadIO m => Handle -> ByteStream m r -> m r +hPut h cs = dematerialize cs return (\x _ y -> liftIO (B.hPut h x) >> y) (>>= id) +{-# INLINE hPut #-} + +-- -- --------------------------------------------------------------------- +-- -- Internal utilities + +-- | 'findIndexOrEnd' is a variant of findIndex, that returns the length +-- of the string if no element is found, rather than Nothing. +findIndexOrEnd :: (Word8 -> Bool) -> Bytes -> Int +findIndexOrEnd k (B.PS x s l) = + unsafeDupablePerformIO $ + withForeignPtr x $ \f -> go (f `plusPtr` s) 0 + where + go !ptr !n | n >= l = return l + | otherwise = do w <- peek ptr + if k w + then return n + else go (ptr `plusPtr` 1) (n+1) +{-# INLINABLE findIndexOrEnd #-} + +{- Take a builder constructed otherwise and convert it to a genuine + streaming bytestring. + +>>> Q.putStrLn $ Q.toByteStream $ stringUtf8 "哈斯克尔" <> stringUtf8 " " <> integerDec 98 +哈斯克尔 98 + + shows its + indistinguishable performance is indistinguishable from @toLazyByteStream@ +-} + +toByteStream :: MonadIO m => Builder -> ByteStream m () +toByteStream = toByteStreamWith (safeStrategy L.smallChunkSize L.defaultChunkSize) +{-# INLINE toByteStream #-} + +{-| Take a builder and convert it to a genuine + streaming bytestring, using a specific allocation strategy. +-} +toByteStreamWith :: MonadIO m => AllocationStrategy -> Builder -> ByteStream m () +toByteStreamWith strategy builder0 = do + cios <- liftIO (buildStepToCIOS strategy (runBuilder builder0)) + let loop !o cios0 = case cios0 of + Yield1 bs io -> Chunk bs o $ do + cios1 <- liftIO io + loop (o + fromIntegral (B.length bs)) cios1 + Finished buf r -> trimmedChunkFromBuffer o buf (Empty r) + trimmedChunkFromBuffer o buffer k + | B.null bs = k + | 2 * B.length bs < bufferSize buffer = Chunk (B.copy bs) o k + | otherwise = Chunk bs o k + where + bs = byteStringFromBuffer buffer + loop 0 cios +{-# INLINABLE toByteStreamWith #-} +{-# SPECIALIZE toByteStreamWith :: AllocationStrategy -> Builder -> ByteStream IO () #-} + + +{- Concatenate a stream of builders (not a streaming bytestring!) into a single builder. + +>>> let aa = yield (integerDec 10000) >> yield (string8 " is a number.") >> yield (char8 '\n') +>>> hPutBuilder IO.stdout $ concatBuilders aa +10000 is a number. + +-} +concatBuilders :: Stream (Of Builder) IO () -> Builder +concatBuilders p = builder $ \bstep r -> do + case p of + Return _ -> runBuilderWith mempty bstep r + Step (b :> rest) -> runBuilderWith (b `mappend` concatBuilders rest) bstep r + Effect m -> m >>= \p' -> runBuilderWith (concatBuilders p') bstep r +{-# INLINABLE concatBuilders #-} + +{- | Turns a ByteStream into a connected stream of ByteStreams that + divide at newline characters. The resulting strings do not contain + newlines. This is the genuinely streaming 'lines' which only + breaks chunks, and thus never increases the use of memory. + + Because 'ByteStream's are usually read in binary mode, with no line + ending conversion, this function recognizes both @\\n@ and @\\r\\n@ + endings (regardless of the current platform). -} + +lines :: Monad m => ByteStream m r -> Stream (ByteStream m) m r +lines text0 = loop1 text0 + where + loop1 :: Monad m => ByteStream m r -> Stream (ByteStream m) m r + loop1 text = + case text of + Empty r -> Return r + Go m -> Effect $ liftM loop1 m + Chunk c _ cs + | B.null c -> loop1 cs + | otherwise -> Step (loop2 Nothing text) + loop2 :: Monad m => Maybe Int64 -> ByteStream m r -> ByteStream m (Stream (ByteStream m) m r) + loop2 prevCr text = + case text of + Empty r -> case prevCr of + Just o -> Chunk (B.singleton 13) o (Empty (Return r)) + Nothing -> Empty (Return r) + Go m -> Go $ liftM (loop2 prevCr) m + Chunk c o cs -> + case B.elemIndex 10 c of + Nothing -> case B.length c of + 0 -> loop2 prevCr cs + l -> if B.unsafeLast c == 13 + then Chunk (B.unsafeInit c) o (loop2 (Just (o-1 + fromIntegral l)) cs) + else Chunk c o (loop2 Nothing cs) + Just i -> do + let prefixLength = + if i >= 1 && B.unsafeIndex c (i-1) == 13 -- \r\n (dos) + then i-1 + else i + rest = + if B.length c > i+1 + then Chunk (B.drop (i+1) c) (o+1 + fromIntegral i) cs + else cs + result = Chunk (B.unsafeTake prefixLength c) o (Empty (loop1 rest)) + case prevCr of + Just oo | i > 0 -> Chunk (B.singleton 13) oo result + _ -> result +{-# INLINABLE lines #-} + +{- | Turns a 'ByteStream' into a stream of strict 'Bytes' that divide at + newline characters. The resulting strings do not contain newlines. + This will cost memory if the lines are very long, and it does not + recognize DOS line endings. -} + +lines' :: Monad m => ByteStream m r -> Stream (Of Bytes) m r +lines' = loop1 [] + where + loop1 :: Monad m => [Bytes] -> ByteStream m r -> Stream (Of Bytes) m r + loop1 acc text = + case text of + Empty r -> Return r + Go m -> Effect $ liftM (loop1 acc) m + Chunk c o cs + | B.null c -> loop1 acc cs + | otherwise -> + case B.elemIndex 10 c of + Just i -> Q.cons (if null acc then B.take i c else B.concat (reverse (B.take i c : acc))) + (loop1 [] (Chunk (B.drop (i+1) c) (o+1 + fromIntegral i) cs)) + Nothing -> loop1 (c:acc) cs +{-# INLINABLE lines' #-} + +-- -------------------------------------------------------------------------- + +{-| Decompresses GZip if present. If any GZip stream is found, all + such streams are decompressed and any remaining data is discarded. + Else, the input is returned unchanged. If the input is BGZF, the + result will contain meaningful virtual offsets. If the input + contains exactly one GZip stream, the result will have meaningfull + offsets into the uncompressed data. Else, the offsets will be + bogus. -} + +gunzip :: MonadIO m => ByteStream m r -> ByteStream m r +gunzip = gunzipWith id +{-# INLINABLE gunzip #-} + +{-| Checks if the input is GZip at all, and runs gunzip if it is. If + it isn't, it runs 'k' on the input. -} + +gunzipWith :: MonadIO m => (ByteStream m r -> ByteStream m r) + -> ByteStream m r -> ByteStream m r +gunzipWith k s0 = lift (nextByteOff s0) >>= \case + Right (31, o, s') -> lift (nextByte s') >>= \case + Right (139,s'') -> gunzipLoop o $ Chunk (B.pack [31,139]) o s'' + Right ( c, s'') -> k $ Chunk (B.pack [31,c]) o s'' + Left r -> k $ Chunk (B.singleton 31) o (pure r) + Right ( c, o, s') -> k $ Chunk (B.singleton c) o s' + Left r -> k $ pure r +{-# INLINABLE gunzipWith #-} + +{-| Decompresses a gzip stream. If the leftovers look like another + gzip stream, it recurses (some files, notably those produced by + bgzip, contain multiple streams). Otherwise, the leftovers are + discarded (some compressed HETFA files appear to have junk at the + end). -} + +gunzipLoop :: MonadIO m => Int64 -> ByteStream m r -> ByteStream m r +gunzipLoop o = go o (shiftL o 16) $ Z.decompressIO Z.gzipOrZlibFormat Z.defaultDecompressParams + where + -- get next chunk, make sure it is empty iff the input ended + go inoff outoff (Z.DecompressInputRequired next) inp = + lift (nextChunk inp) >>= \case + Left r -> do z <- liftIO (next B.empty) + go inoff outoff z (pure r) + Right (ck,inp') + | B.null ck -> go inoff outoff (Z.DecompressInputRequired next) inp' + | otherwise -> do z <- liftIO (next ck) + go (inoff + fromIntegral (B.length ck)) outoff z inp' + + go inoff outoff (Z.DecompressOutputAvailable outchunk next) inp = do + z <- Chunk outchunk outoff (liftIO next) + go inoff (outoff + fromIntegral (B.length outchunk)) z inp + + go inoff _outoff (Z.DecompressStreamEnd inchunk) inp = + -- decompress leftovers if possible, else return + gunzipWith (lift . effects) (Chunk inchunk (inoff - fromIntegral (B.length inchunk)) inp) + + go _inoff _outoff (Z.DecompressStreamError derr) _inp = + liftIO $ throwIO derr + +-- | Compresses a byte stream using GZip with default parameters. +gzip :: MonadIO m => ByteStream m r -> ByteStream m r +gzip = go $ Z.compressIO Z.gzipFormat Z.defaultCompressParams + where + -- get next chunk, make sure it is empty iff the input ended + go (Z.CompressInputRequired next) inp = + lift (nextChunk inp) >>= \case + Left r -> liftIO (next B.empty) >>= flip go (pure r) + Right (ck,inp') + | B.null ck -> go (Z.CompressInputRequired next) inp' + | otherwise -> liftIO (next ck) >>= flip go inp' + + go (Z.CompressOutputAvailable outchunk next) inp = + chunk outchunk >> liftIO next >>= flip go inp + + go Z.CompressStreamEnd inp = lift (effects inp) + + diff --git a/Bio/Streaming/Furrow.hs b/Bio/Streaming/Furrow.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming/Furrow.hs @@ -0,0 +1,45 @@ +module Bio.Streaming.Furrow + ( Furrow(..) + , evertStream + , afford + , drain + ) where + +import Bio.Prelude +import Bio.Streaming + +{- | A tiny stream that can be afforded to incrementally. + +The streaming abstraction works fine if multiple sources feed into a +small constant number of functions, but fails if there is an +unpredictable number of such consumers. In that case, 'evertStream' +should be used to turn each consumer into a 'Furrow'. It's then +possible to incrementally 'afford' stuff to each 'Furrow' in a +collection in a simple loop. To get the final value, 'drain' each +'Furrow'. +-} +newtype Furrow a m r = Furrow (Stream ((->) (Maybe a)) m r) deriving + (Functor, Applicative, Monad, MonadTrans, MonadIO, MFunctor, MMonad) + +instance MonadThrow m => MonadThrow (Furrow a m) where + throwM = Furrow . lift . throwM + +afford :: Monad m => Furrow a m b -> a -> m (Furrow a m b) +afford (Furrow s) a = inspect s >>= \case + Left b -> return (Furrow (pure b)) + Right f -> return (Furrow (f (Just a))) + +drain :: Monad m => Furrow a m b -> m b +drain (Furrow s) = inspect s >>= \case + Left b -> return b + Right f -> inspect (f Nothing) >>= \case + Left b -> return b + Right _ -> error "continuedAfterEOF" + +-- | Turns a function that consumes a stream into a furrow. Idea and +-- some code stolen from \"streaming-eversion\". +evertStream :: Monad m => (Stream (Of a) (Furrow a m) () -> Furrow a m b) -> Furrow a m b +evertStream consumer = consumer cat + where + cat = lift (Furrow (yields id)) >>= maybe (pure ()) (\a -> wrap (a :> cat)) + diff --git a/Bio/Streaming/Parse.hs b/Bio/Streaming/Parse.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming/Parse.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE Rank2Types #-} +module Bio.Streaming.Parse + ( Parser + , ParseError(..) + , EofException(..) + , parse + , parseIO + , parseM + , abortParse + , isFinished + , drop + , dropLine + , getByte + , getString + , getWord32 + , getWord64 + , isolate + , atto + ) where + +-- ^ Parsers for use with 'ByteStream's. + +import Bio.Prelude hiding ( drop ) +import Bio.Streaming.Bytes ( ByteStream ) + +import qualified Bio.Streaming.Bytes as S +import qualified Data.Attoparsec.ByteString as A +import qualified Data.ByteString as B +import qualified Streaming.Prelude as Q + +newtype Parser r m a = P { + runP :: forall x . + (a -> ByteStream m r -> m x) + -> (r -> m x) + -> (SomeException -> m x) + -> ByteStream m r -> m x } + +instance Functor (Parser r m) where + fmap f p = P $ \sk -> runP p (sk . f) + +instance Applicative (Parser r m) where + pure a = P $ \sk _rk _ek -> sk a + a <*> b = P $ \sk rk ek -> runP a (\f -> runP b (\x -> sk (f x)) rk ek) rk ek + +instance Monad (Parser r m) where + return = pure + m >>= k = P $ \sk rk ek -> runP m (\a -> runP (k a) sk rk ek) rk ek + +instance MonadIO m => MonadIO (Parser r m) where + liftIO m = P $ \sk _rk _ek s -> liftIO m >>= \a -> sk a s + +instance MonadTrans (Parser r) where + lift m = P $ \sk _rk _ek s -> m >>= \a -> sk a s + +instance MonadThrow (Parser r m) where + throwM e = P $ \_sk _rk ek _s -> ek (toException e) + +modify :: (ByteStream m r -> ByteStream m r) -> Parser r m () +modify f = P $ \sk _rk _ek -> sk () . f + +parse :: Monad m => (Int64 -> Parser r m a) -> ByteStream m r -> m (Either SomeException (Either r (a, ByteStream m r))) +parse p = go + where + go (S.Empty r) = return $ Right $ Left r + go (S.Go k) = k >>= go + go ck@(S.Chunk c o s) | B.null c = go s + | otherwise = runP (p o) (\a t -> return . Right $ Right (a,t)) + (return . Right . Left) + (return . Left) + ck + +parseIO :: MonadIO m => (Int64 -> Parser r m a) -> ByteStream m r -> m (Either r (a, ByteStream m r)) +parseIO p = parse p >=> either (liftIO . throwM) return + +parseM :: MonadThrow m => (Int64 -> Parser r m a) -> ByteStream m r -> m (Either r (a, ByteStream m r)) +parseM p = parse p >=> either throwM return + +abortParse :: Monad m => Parser r m a +abortParse = P $ \_sk rk _ek -> S.effects >=> rk + +liftFun :: Monad m => (ByteStream m r -> m (a, ByteStream m r)) -> Parser r m a +liftFun f = P $ \sk _rk _ek -> f >=> uncurry sk + +isFinished :: Monad m => Parser r m Bool +isFinished = liftFun go + where + go (S.Empty r) = return (True, S.Empty r) + go (S.Go k) = k >>= go + go ck@(S.Chunk c _ s) | B.null c = go s + | otherwise = return (False, ck) + +drop :: Monad m => Int -> Parser r m () +drop l = modify $ S.drop (fromIntegral l) + +dropLine :: Monad m => Parser r m () +dropLine = modify $ S.drop 1 . S.dropWhile (/= 10) + +getByte :: Monad m => Parser r m Word8 +getByte = P $ \sk _rk ek -> S.nextByte >=> either (const $ ek (toException EofException)) (uncurry sk) + +getString :: Monad m => Int -> Parser r m B.ByteString +getString l = liftFun $ liftM Q.lazily . S.splitAt' l + +getWord32 :: Monad m => Parser r m Word32 +getWord32 = liftM (fst . B.foldl (\(a,i) w -> (a + shiftL (fromIntegral w) i, i + 8)) (0,0)) (getString 4) + +getWord64 :: Monad m => Parser r m Word64 +getWord64 = liftM (fst . B.foldl (\(a,i) w -> (a + shiftL (fromIntegral w) i, i + 8)) (0,0)) (getString 8) + +isolate :: Monad m => Int -> Parser (ByteStream m r) m a -> Parser r m a +isolate l p = P $ \sk rk ek -> runP p (\a -> S.effects >=> sk a) + (S.effects >=> rk) + ek . S.splitAt (fromIntegral l) + +data ParseError = ParseError {errorContexts :: [String], errorMessage :: String} + deriving (Show, Typeable) + +data EofException = EofException + deriving (Show, Typeable) + +instance Exception ParseError +instance Exception EofException + +atto :: Monad m => A.Parser a -> Parser r m a +atto = go . A.parse + where + go k = P $ \sk rk ek -> + S.nextChunk >=> \case + Left r -> case k B.empty of + A.Fail _ err dsc -> ek $ toException (ParseError err dsc) + A.Partial _ -> ek $ toException EofException + A.Done rest v -> sk v (S.consChunk rest (pure r)) + Right (c,s') + | B.null c -> runP (go k) sk rk ek s' + | otherwise -> case k c of + A.Fail _ err dsc -> ek $ toException (ParseError err dsc) + A.Partial k' -> runP (go k') sk rk ek s' + A.Done rest v -> sk v (S.consChunk rest s') + + diff --git a/Bio/Streaming/Vector.hs b/Bio/Streaming/Vector.hs new file mode 100644 --- /dev/null +++ b/Bio/Streaming/Vector.hs @@ -0,0 +1,35 @@ +module Bio.Streaming.Vector ( stream2vector, stream2vectorN ) where + +import Bio.Prelude +import Streaming + +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Mutable as VM + +-- | Equivalent to @stream2vector . Streaming.Prelude.take n@, but +-- terminates early and is thereby more efficient. +stream2vectorN :: (MonadIO m, VG.Vector v a) => Int -> Stream (Of a) m () -> m (v a) +stream2vectorN n s0 = do + mv <- liftIO $ VM.new n + go mv 0 s0 + where + go !mv !i s + | i == n = liftIO $ VG.unsafeFreeze mv + | otherwise = + inspect s >>= \case + Left () -> liftIO $ VG.unsafeFreeze $ VM.take i mv + Right (a :> s') -> liftIO (VM.write mv i a) >> go mv (i+1) s' + +-- | Reads the whole stream into a 'VG.Vector'. +stream2vector :: (MonadIO m, VG.Vector v a) => Stream (Of a) m r -> m (Of (v a) r) +stream2vector s0 = do + mv <- liftIO $ VM.new 1024 + go mv 0 s0 + where + go !mv !i = + inspect >=> \case + Left r -> liftM (:> r) $ liftIO $ VG.unsafeFreeze $ VM.take i mv + Right (a :> s) -> do mv' <- if VM.length mv == i then liftIO (VM.grow mv (VM.length mv)) else return mv + liftIO $ VM.write mv' i a + go mv' (i+1) s + diff --git a/Bio/TwoBit.hs b/Bio/TwoBit.hs new file mode 100644 --- /dev/null +++ b/Bio/TwoBit.hs @@ -0,0 +1,296 @@ +-- | Would you believe it? The 2bit format stores blocks of Ns in a table at +-- the beginning of a sequence, then packs four bases into a byte. So it +-- is neither possible nor necessary to store Ns in the main sequence, and +-- you would think they aren't stored there, right? And they aren't. +-- Instead Ts are stored which the reader has to replace with Ns. +-- +-- The sensible way to treat these is probably to just say there are two +-- kinds of implied annotation (repeats and large gaps for a typical +-- genome), which can be interpreted in whatever way fits. And that's why +-- we have 'Mask' and 'getSubseqWith'. + +module Bio.TwoBit ( + TwoBitFile(..), + TwoBitSequence(..), + openTwoBit, + + getFwdSubseqWith, + getSubseq, + getSubseqWith, + getSubseqAscii, + getSubseqMasked, + getLazySubseq, + getFragment, + getFwdSubseqV, + getSeqnames, + lookupSequence, + getSeqLength, + clampPosition, + getRandomSeq, + + takeOverlap, + mergeBlocks, + Mask(..) + ) where + +import Bio.Prelude hiding ( left, right, chr ) +import Bio.Util.MMap +import Bio.Util.Storable +import Control.Monad.Trans.State ( StateT(..), get, evalStateT ) +import qualified Data.ByteString as B +import qualified Data.ByteString.Unsafe as B +import qualified Data.IntMap.Strict as I +import qualified Data.HashMap.Lazy as M +import qualified Data.Vector.Unboxed as U +import Foreign.C.Types ( CChar ) + +data TwoBitFile = TBF { + tbf_raw :: B.ByteString, + -- This map is intentionally lazy. May or may not be important. + tbf_seqs :: !(M.HashMap Bytes TwoBitSequence) +} + +data TwoBitSequence = TBS { tbs_n_blocks :: !(I.IntMap Int) + , tbs_m_blocks :: !(I.IntMap Int) + , tbs_dna_offset :: {-# UNPACK #-} !Int + , tbs_dna_size :: {-# UNPACK #-} !Int } + +-- | Brings a 2bit file into memory. The file is mmap'ed, so it will +-- not work on streams that are not actual files. It's also unsafe if +-- the file is modified in any way. +openTwoBit :: FilePath -> IO TwoBitFile +openTwoBit fp = do + raw <- unsafeMMapFile fp + B.unsafeUseAsCString raw $ \praw -> + -- return $ flip runGet (L.fromChunks [raw]) $ do + flip evalStateT praw $ do + sig <- getWord32be + getWord32 <- case sig :: Word32 of + 0x1A412743 -> return getWord32be + 0x4327411A -> return getWord32le + _ -> fail $ "invalid .2bit signature " ++ showHex sig [] + + version <- getWord32 + unless (version == 0) $ fail $ "wrong .2bit version " ++ show version + + nseqs <- getWord32 + _reserved <- getWord32 + + TBF raw <$> foldM (\ix _ -> do !key <- getWord8 >>= getByteString + !off <- getWord32 + return $! M.insert key (mkBlockIndex raw getWord32 off) ix + ) M.empty [1..nseqs] + +type Get = StateT (Ptr CChar) IO + +getWord8, getWord32be, getWord32le :: Num a => Get a +getWord8 = StateT $ \p -> peekUnalnWord32BE p >>= \w -> return (fromIntegral w, plusPtr p 1) +getWord32be = StateT $ \p -> peekUnalnWord32BE p >>= \w -> return (fromIntegral w, plusPtr p 4) +getWord32le = StateT $ \p -> peekUnalnWord32LE p >>= \w -> return (fromIntegral w, plusPtr p 4) + +getByteString :: Int -> Get Bytes +getByteString l = StateT $ \p -> B.packCStringLen (p,l) >>= \s -> return (s, plusPtr p l) + +mkBlockIndex :: B.ByteString -> Get Int -> Int -> TwoBitSequence +mkBlockIndex raw getWord32 ofs = + unsafePerformIO $ + B.unsafeUseAsCString raw $ \praw -> + evalStateT getBlock (plusPtr praw ofs) + where + getBlock = do p0 <- get + ds <- getWord32 + nb <- readBlockList + mb <- readBlockList + _ <- getWord32 + p1 <- get + return $! TBS (I.fromList nb) (I.fromList mb) (ofs + minusPtr p1 p0) ds + + readBlockList = getWord32 >>= \n -> liftM2 zip (repM n getWord32) (repM n getWord32) + +-- | Repeat monadic action @n@ times. Returns result in reverse(!) +-- order, but doesn't build a huge list of thunks in memory. +repM :: Monad m => Int -> m a -> m [a] +repM n0 m = go [] n0 + where + go acc 0 = return acc + go acc n = m >>= \x -> x `seq` go (x:acc) (n-1) + +takeOverlap :: Int -> I.IntMap Int -> [(Int,Int)] +takeOverlap k m = dropWhile far_left $ + maybe id (\(kv,_) -> (:) kv) (I.maxViewWithKey left) $ + maybe id (\v -> (:) (k,v)) middle $ + I.toAscList right + where + (left, middle, right) = I.splitLookup k m + far_left (s,l) = s+l <= k + +data Mask = None | Soft | Hard | Both deriving (Eq, Ord, Enum, Show) + +getFwdSubseqWith :: TwoBitFile -> TwoBitSequence -- raw data, sequence + -> (Word8 -> Mask -> a) -- mask function + -> Int -> [a] -- start, lazy result +getFwdSubseqWith TBF{..} TBS{..} nt start = + do_mask (takeOverlap start tbs_n_blocks `mergeBlocks` takeOverlap start tbs_m_blocks) start . + drop (start .&. 3) . + B.foldr toDNA [] . + B.drop (tbs_dna_offset + (start `shiftR` 2)) $ tbf_raw + where + toDNA b = (++) [ 3 .&. (b `shiftR` x) | x <- [6,4,2,0] ] + + do_mask _ _ [] = [] + do_mask [ ] _ ws = map (`nt` None) ws + do_mask ((s,l,m):is) p ws + | p < s = map (`nt` None) (take (s-p) ws) ++ do_mask ((s,l,m):is) s (drop (s-p) ws) + | otherwise = map (`nt` m) (take (s+l-p) ws) ++ do_mask is (s+l) (drop (s+l-p) ws) + +-- | Merge blocks of Ns and blocks of Ms into single list of blocks with +-- masking annotation. Gaps remain. Used internally only. +mergeBlocks :: [(Int,Int)] -> [(Int,Int)] -> [(Int,Int,Mask)] +mergeBlocks ((_,0):nbs) mbs = mergeBlocks nbs mbs +mergeBlocks nbs ((_,0):mbs) = mergeBlocks nbs mbs + +mergeBlocks ((ns,nl):nbs) ((ms,ml):mbs) + | ns < ms = let l = min (ms-ns) nl in (ns,l, Hard) : mergeBlocks ((ns+l,nl-l):nbs) ((ms,ml):mbs) + | ms < ns = let l = min (ns-ms) ml in (ms,l, Soft) : mergeBlocks ((ns,nl):nbs) ((ms+l,ml-l):mbs) + | otherwise = let l = min nl ml in (ns,l, Both) : mergeBlocks ((ns+l,nl-l):nbs) ((ms+l,ml-l):mbs) + +mergeBlocks ((ns,nl):nbs) [] = (ns,nl, Hard) : mergeBlocks nbs [] +mergeBlocks [] ((ms,ml):mbs) = (ms,ml, Soft) : mergeBlocks [] mbs + +mergeBlocks [ ] [ ] = [] + + +-- | Extract a subsequence and apply masking. TwoBit file can represent +-- two kinds of masking (hard and soft), where hard masking is usually +-- realized by replacing everything by Ns and soft masking is done by +-- lowercasing. Here, we take a user supplied function to apply +-- masking. +getSubseqWith :: (Nucleotide -> Mask -> a) -> TwoBitFile -> Range -> [a] +getSubseqWith maskf tbf Range{ r_pos = Pos { p_seq = chr, p_start = start }, r_length = len } = do + let sq1 = fromMaybe (error $ unpack chr ++ " doesn't exist") $ M.lookup chr (tbf_seqs tbf) + let go = getFwdSubseqWith tbf sq1 + if start < 0 + then reverse $ take len $ go (maskf . cmp_nt) (-start-len) + else take len $ go (maskf . fwd_nt) start + where + fwd_nt = (!!) [nucT, nucC, nucA, nucG] . fromIntegral + cmp_nt = (!!) [nucA, nucG, nucT, nucC] . fromIntegral + +-- | Works only in forward direction. +getLazySubseq :: TwoBitFile -> Position -> [Nucleotide] +getLazySubseq tbf Pos{ p_seq = chr, p_start = start } = do + let sq1 = fromMaybe (error $ unpack chr ++ " doesn't exist") $ M.lookup chr (tbf_seqs tbf) + let go = getFwdSubseqWith tbf sq1 + if start < 0 + then error "sorry, can't go backwards" + -- then reverse $ take len $ go (maskf . cmp_nt) (-start-len) + else go fwd_nt start + where + fwd_nt n _ = [nucT, nucC, nucA, nucG] !! fromIntegral n + + +-- | Extract a subsequence without masking. +getSubseq :: TwoBitFile -> Range -> [Nucleotide] +getSubseq = getSubseqWith const + +-- | Extract a subsequence with typical masking: soft masking is +-- ignored, hard masked regions are replaced with Ns. +getSubseqMasked :: TwoBitFile -> Range -> [Nucleotides] +getSubseqMasked = getSubseqWith mymask + where + mymask n None = nucToNucs n + mymask n Soft = nucToNucs n + mymask _ Hard = nucsN + mymask _ Both = nucsN + +-- | Extract a subsequence with masking for biologists: soft masking is +-- done by lowercasing, hard masking by printing an N. +getSubseqAscii :: TwoBitFile -> Range -> String +getSubseqAscii = getSubseqWith mymask + where + mymask n None = showNucleotide n + mymask n Soft = toLower (showNucleotide n) + mymask _ Hard = 'N' + mymask _ Both = 'N' + + +getSeqnames :: TwoBitFile -> [Bytes] +getSeqnames = M.keys . tbf_seqs + +lookupSequence :: TwoBitFile -> Bytes -> Maybe TwoBitSequence +lookupSequence tbf sq = M.lookup sq . tbf_seqs $ tbf + +getSeqLength :: TwoBitFile -> Bytes -> Int +getSeqLength tbf chr = + maybe (error $ shows chr " doesn't exist") tbs_dna_size $ + M.lookup chr (tbf_seqs tbf) + +-- | limits a range to a position within the actual sequence +clampPosition :: TwoBitFile -> Range -> Range +clampPosition tbf (Range (Pos n start) len) = Range (Pos n start') (end' - start') + where + size = getSeqLength tbf n + start' = if start < 0 then max start (-size) else start + end' = min (start + len) $ if start < 0 then 0 else size + + +-- | Sample a piece of random sequence uniformly from the genome. +-- Only pieces that are not hard masked are sampled, soft masking is +-- allowed, but not reported. +-- On a 32bit platform, this will fail for genomes larger than 1G bases. +-- However, if you're running this code on a 32bit platform, you have +-- bigger problems to worry about. +getRandomSeq :: TwoBitFile -- ^ 2bit file + -> Int -- ^ desired length + -> (Int -> g -> (Int, g)) -- ^ draw random int below limit + -> g -- ^ RNG + -> ((Range, [Nucleotide]), g) -- ^ position, sequence, new RNG +getRandomSeq tbf len rndInt = draw + where + names = getSeqnames tbf + lengths = map (getSeqLength tbf) names + total = sum lengths + frags = I.fromList $ zip (scanl (+) 0 lengths) names + + draw g0 | good = ((r', sq), gn) + | otherwise = draw gn + where + (p0, gn) = rndInt (2*total) g0 + p = p0 `shiftR` 1 + (o,s) = fst $ fromJust $ I.maxViewWithKey $ fst $ I.split (p+1) frags + r' = (if odd p0 then id else reverseRange) $ clampPosition tbf $ Range (Pos s (p-o)) len + sq = catMaybes $ getSubseqWith mask2maybe tbf r' + good = r_length r' == len && length sq == len + + mask2maybe n None = Just n + mask2maybe n Soft = Just n + mask2maybe _ Hard = Nothing + mask2maybe _ Both = Nothing + +-- | Gets a fragment from a 2bit file. The result always has the +-- desired length; if necessary, it is padded with Ns. Be careful about +-- the unconventional encoding: 0..4 == TCAGN +getFragment :: TwoBitFile -> Bytes -> Int -> Int -> U.Vector Word8 +getFragment tbf chr p l = + case lookupSequence tbf chr of + Nothing -> U.replicate l 4 + Just tbs -> getFwdSubseqV tbf tbs p l + +-- Careful about weird encoding: 0..4 == TCAGN +getFwdSubseqV :: TwoBitFile -> TwoBitSequence -> Int -> Int -> U.Vector Word8 +getFwdSubseqV TBF{..} TBS{..} start len = U.unfoldrN len step ini + where + ini = (start, takeOverlap start tbs_n_blocks) + + step (off, nbs) + | off < 0 = Just (4, (succ off, nbs)) + | off >= tbs_dna_size = Just (4, (succ off, nbs)) + | otherwise = case nbs of + [ ] -> Just (y, (succ off, [ ])) + (s,l):nbs' | off < s -> Just (y, (succ off, nbs)) + | off < s+l -> Just (4, (succ off, nbs)) + | otherwise -> Just (y, (succ off, nbs')) + where + x = B.index tbf_raw (tbs_dna_offset + off `shiftR` 2) + y = x `shiftR` (6 - 2 * (off .&. 3)) .&. 3 -- T,C,A,G + diff --git a/Bio/Util/MMap.hs b/Bio/Util/MMap.hs new file mode 100644 --- /dev/null +++ b/Bio/Util/MMap.hs @@ -0,0 +1,28 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +module Bio.Util.MMap ( unsafeMMapFile ) where + +import BasePrelude +import Data.ByteString.Internal ( fromForeignPtr, ByteString ) +import Foreign.C.Types +import System.Posix.Files +import System.Posix.IO + +unsafeMMapFile :: FilePath -> IO ByteString +unsafeMMapFile fp = + bracket (openFd fp ReadOnly Nothing defaultFileFlags) closeFd $ \fd -> do + stat <- getFdStatus fd + let size = fromIntegral (fileSize stat) + if size <= 0 + then return mempty + else do + ptr <- c_mmap size (fromIntegral fd) + if ptr == nullPtr + then error "unable to mmap file" + else do + fptr <- newForeignPtrEnv c_munmap (intPtrToPtr $ fromIntegral size) ptr + return $ fromForeignPtr fptr 0 (fromIntegral size) + +foreign import ccall unsafe "my_mmap" c_mmap :: CSize -> CInt -> IO (Ptr Word8) +foreign import ccall unsafe "&my_munmap" c_munmap :: FunPtr (Ptr () -> Ptr Word8 -> IO ()) + + diff --git a/Bio/Util/Nub.hs b/Bio/Util/Nub.hs new file mode 100644 --- /dev/null +++ b/Bio/Util/Nub.hs @@ -0,0 +1,16 @@ +module Bio.Util.Nub ( nubHash, nubHashBy )where + +import BasePrelude +import Data.Hashable ( Hashable ) +import qualified Data.HashSet as H + +nubHash :: (Hashable a, Eq a) => [a] -> [a] +nubHash = nubHashBy id + +nubHashBy :: (Hashable b, Eq b) => (a -> b) -> [a] -> [a] +nubHashBy f = go H.empty + where + go _ [] = [] + go h (x:xs) | f x `H.member` h = go h xs + | otherwise = x : go (H.insert (f x) h) xs + diff --git a/Bio/Util/Numeric.hs b/Bio/Util/Numeric.hs new file mode 100644 --- /dev/null +++ b/Bio/Util/Numeric.hs @@ -0,0 +1,213 @@ +-- | Random useful stuff I didn't know where to put. + +module Bio.Util.Numeric ( + wilson, invnormcdf, choose, + estimateComplexity, showNum, showOOM, + log1p, expm1, (<#>), + log1mexp, log1pexp, + lsum, llerp + ) where + +import Prelude +import Data.Char ( intToDigit ) +import Data.List ( foldl1' ) + +-- | Calculates the Wilson Score interval. +-- If @(l,m,h) = wilson c x n@, then @m@ is the binary proportion and +-- @(l,h)@ it's @c@-confidence interval for @x@ positive examples out of +-- @n@ observations. @c@ is typically something like 0.05. + +wilson :: Double -> Int -> Int -> (Double, Double, Double) +wilson c x n = ( (m - h) / d, p, (m + h) / d ) + where + nn = fromIntegral n + p = fromIntegral x / nn + + z = invnormcdf (1-c*0.5) + h = z * sqrt (( p * (1-p) + 0.25*z*z / nn ) / nn) + m = p + 0.5 * z * z / nn + d = 1 + z * z / nn + +showNum :: Show a => a -> String +showNum = triplets [] . reverse . show + where + triplets acc [] = acc + triplets acc [a] = a:acc + triplets acc [a,b] = b:a:acc + triplets acc [a,b,c] = c:b:a:acc + triplets acc (a:b:c:s) = triplets (',':c:b:a:acc) s + +showOOM :: Double -> String +showOOM x | x < 0 = '-' : showOOM (negate x) + | otherwise = findSuffix (x*10) ".kMGTPEZY" + where + findSuffix _ [] = "many" + findSuffix y (s:ss) | y < 100 = intToDigit (round y `div` 10) : case (round y `mod` 10, s) of + (0,'.') -> [] ; (0,_) -> [s] ; (d,_) -> [s, intToDigit d] + | y < 1000 = intToDigit (round y `div` 100) : intToDigit ((round y `mod` 100) `div` 10) : + if s == '.' then [] else [s] + | y < 10000 = intToDigit (round y `div` 1000) : intToDigit ((round y `mod` 1000) `div` 100) : + '0' : if s == '.' then [] else [s] + | otherwise = findSuffix (y*0.001) ss + +-- Stolen from Lennart Augustsson's erf package, who in turn took it from +-- Accurate to about 1e-9. +invnormcdf :: (Ord a, Floating a) => a -> a +invnormcdf p = + let a1 = -3.969683028665376e+01 + a2 = 2.209460984245205e+02 + a3 = -2.759285104469687e+02 + a4 = 1.383577518672690e+02 + a5 = -3.066479806614716e+01 + a6 = 2.506628277459239e+00 + + b1 = -5.447609879822406e+01 + b2 = 1.615858368580409e+02 + b3 = -1.556989798598866e+02 + b4 = 6.680131188771972e+01 + b5 = -1.328068155288572e+01 + + c1 = -7.784894002430293e-03 + c2 = -3.223964580411365e-01 + c3 = -2.400758277161838e+00 + c4 = -2.549732539343734e+00 + c5 = 4.374664141464968e+00 + c6 = 2.938163982698783e+00 + + d1 = 7.784695709041462e-03 + d2 = 3.224671290700398e-01 + d3 = 2.445134137142996e+00 + d4 = 3.754408661907416e+00 + + pLow = 0.02425 + + nan = 0/0 + + in if p < 0 then + nan + else if p == 0 then + -1/0 + else if p < pLow then + let q = sqrt(-2 * log p) + in (((((c1*q+c2)*q+c3)*q+c4)*q+c5)*q+c6) / + ((((d1*q+d2)*q+d3)*q+d4)*q+1) + else if p < 1 - pLow then + let q = p - 0.5 + r = q*q + in (((((a1*r+a2)*r+a3)*r+a4)*r+a5)*r+a6)*q / + (((((b1*r+b2)*r+b3)*r+b4)*r+b5)*r+1) + else if p <= 1 then + - invnormcdf (1 - p) + else + nan + + +-- | Try to estimate complexity of a whole from a sample. Suppose we +-- sampled @total@ things and among those @singles@ occured only once. +-- How many different things are there? +-- +-- Let the total number be @m@. The copy number follows a Poisson +-- distribution with paramter @\lambda@. Let \( z := e^{\lambda} \), then +-- we have: +-- +-- \[ +-- P( 0 ) = e^{-\lambda} = \frac{1}{z} \\ +-- P( 1 ) = \lambda e^{-\lambda} = \frac{\ln z}{z} \\ +-- P(\ge 1) = 1 - e^{-\lambda} = 1 - \frac{1}{z} \\ +-- \] +-- \[ +-- \mbox{singles} = m \frac{\ln z}{z} \\ +-- \mbox{total} = m \left( 1 - \frac{1}{z} \right) \\ +-- \] +-- \[ +-- D := \frac{\mbox{total}}{\mbox{singles}} = (1 - \frac{1}{z}) * \frac{z}{\ln z} \\ +-- f := z - 1 - D \ln z = 0 +-- \] +-- +-- To get @z@, we solve using Newton iteration and then substitute to +-- get @m@: +-- +-- \[ +-- df/dz = 1 - D/z \\ +-- z' = z - \frac{ z (z - 1 - D \ln z) }{ z - D } \\ +-- m = \mbox{singles} * \frac{z}{\ln z} +-- \] +-- +-- It converges as long as the initial @z@ is large enough, and @10D@ +-- (in the line for @zz@ below) appears to work well. + +estimateComplexity :: (Integral a, Floating b, Ord b) => a -> a -> Maybe b +estimateComplexity total singles | total <= singles = Nothing + | singles <= 0 = Nothing + | otherwise = Just m + where + d = fromIntegral total / fromIntegral singles + step z = z * (z - 1 - d * log z) / (z - d) + iter z = case step z of zd | abs zd < 1e-12 -> z + | otherwise -> iter $! z-zd + zz = iter $! 10*d + m = fromIntegral singles * zz / log zz + + +-- | Computes \( \ln \left( e^x + e^y \right) \) without leaving the log domain and +-- hence without losing precision. +infixl 5 <#> +{-# INLINE (<#>) #-} +(<#>) :: (Floating a, Ord a) => a -> a -> a +x <#> y = if x >= y then x + log1pexp (y-x) else y + log1pexp (x-y) + +-- | Computes @log (1+x)@ to a relative precision of @10^-8@ even for +-- very small @x@. Stolen from +{-# INLINE log1p #-} +log1p :: (Floating a, Ord a) => a -> a +log1p x | x < -1 = error "log1p: argument must be greater than -1" + -- x is large enough that the obvious evaluation is OK: + | x > 0.0001 || x < -0.0001 = log $ 1 + x + -- Use Taylor approx. log(1 + x) = x - x^2/2 with error roughly x^3/3 + -- Since |x| < 10^-4, |x|^3 < 10^-12, relative error less than 10^-8: + | otherwise = (1 - 0.5*x) * x + + +-- | Computes \( e^x - 1 \) to a relative precision of @10^-10@ even for +-- very small @x@. Stolen from +{-# INLINE expm1 #-} +expm1 :: (Floating a, Ord a) => a -> a +expm1 x | x > -0.00001 && x < 0.00001 = (1 + 0.5 * x) * x -- Taylor approx + | otherwise = exp x - 1 -- direct eval + +-- | Computes \( \ln (1 - e^x) \), following Martin Mächler. +{-# INLINE log1mexp #-} +log1mexp :: (Floating a, Ord a) => a -> a +log1mexp x | x > - log 2 = log (- expm1 x) + | otherwise = log1p (- exp x) + +-- | Computes \( \ln (1 + e^x) \), following Martin Mächler. +{-# INLINE log1pexp #-} +log1pexp :: (Floating a, Ord a) => a -> a +log1pexp x | x <= -37 = exp x + | x <= 18 = log1p $ exp x + | x <= 33.3 = x + exp (-x) + | otherwise = x + + +-- | Computes \( \ln ( \sum_i e^{x_i} ) \) sensibly. The list must be +-- sorted in descending(!) order. +{-# INLINE lsum #-} +lsum :: (Floating a, Ord a) => [a] -> a +lsum = foldl1' (\x y -> if x >= y then x + log1pexp (y-x) else err) + where err = error "lsum: argument list must be in descending order" + +-- | Computes \( \ln \left( c e^x + (1-c) e^y \right) \). +{-# INLINE llerp #-} +llerp :: (Floating a, Ord a) => a -> a -> a -> a +llerp c x y | c <= 0.0 = y + | c >= 1.0 = x + | x >= y = log c + x + log1p ( (1-c)/c * exp (y-x) ) -- Hmm. + | otherwise = log1p (-c) + y + log1p ( c/(1-c) * exp (x-y) ) -- Hmm. + +-- | Binomial coefficient: \( \mbox{choose n k} = \frac{n!}{(n-k)! k!} \) +{-# INLINE choose #-} +choose :: Integral a => a -> a -> a +choose n k = product [n-k+1 .. n] `div` product [2..k] + + diff --git a/Bio/Util/Storable.hs b/Bio/Util/Storable.hs new file mode 100644 --- /dev/null +++ b/Bio/Util/Storable.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE CPP #-} +-- | Utilities to read multibyte quantities from arbitrary positions. +module Bio.Util.Storable + ( peekWord8 + , peekUnalnWord16LE + , peekUnalnWord16BE + , peekUnalnWord32LE + , peekUnalnWord32BE + , pokeUnalnWord32LE + ) where + +#if __GLASGOW_HASKELL__ >= 710 +#define HAVE_BYTESWAP_PRIMOPS +#endif + +#if i386_HOST_ARCH || x86_64_HOST_ARCH +#define MEM_UNALIGNED_OPS +#endif + +import BasePrelude + +peekWord8 :: Ptr a -> IO Word8 +peekWord8 = peek . castPtr + +#if defined(MEM_UNALIGNED_OPS) && defined(WORDS_BIGENDIAN) && defined(HAVE_BYTESWAP_PRIMOPS) +peekUnalnWord16LE :: Ptr a -> IO Word16 +peekUnalnWord16LE = fmap byteSwap16 . peek . castPtr + +peekUnalnWord32LE :: Ptr a -> IO Word32 +peekUnalnWord32LE = fmap byteSwap32 . peek . castPtr + +pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () +pokeUnalnWord32LE p w = poke (castPtr p) (byteSwap32 w) + +#elif defined(MEM_UNALIGNED_OPS) && !defined(WORDS_BIGENDIAN) +peekUnalnWord16LE :: Ptr a -> IO Word16 +peekUnalnWord16LE = peek . castPtr + +peekUnalnWord32LE :: Ptr a -> IO Word32 +peekUnalnWord32LE = peek . castPtr + +pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () +pokeUnalnWord32LE p w = poke (castPtr p) w + +#else +peekUnalnWord16LE :: Ptr a -> IO Word16 +peekUnalnWord16LE p = do + x <- fromIntegral <$> peekWord8 (plusPtr p 0) + y <- fromIntegral <$> peekWord8 (plusPtr p 1) + return $! x .|. unsafeShiftL y 8 + +peekUnalnWord32LE :: Ptr a -> IO Word32 +peekUnalnWord32LE p = do + x <- fromIntegral <$> peekWord8 (plusPtr p 0) + y <- fromIntegral <$> peekWord8 (plusPtr p 1) + z <- fromIntegral <$> peekWord8 (plusPtr p 2) + w <- fromIntegral <$> peekWord8 (plusPtr p 3) + return $! x .|. unsafeShiftL y 8 .|. unsafeShiftL z 16 .|. unsafeShiftL w 24 + +pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () +pokeUnalnWord32LE p w = do pokeByteOff p 0 (fromIntegral $ shiftR w 0 :: Word8) + pokeByteOff p 1 (fromIntegral $ shiftR w 8 :: Word8) + pokeByteOff p 2 (fromIntegral $ shiftR w 16 :: Word8) + pokeByteOff p 3 (fromIntegral $ shiftR w 24 :: Word8) +#endif + + +#if defined(MEM_UNALIGNED_OPS) && !defined(WORDS_BIGENDIAN) && defined(HAVE_BYTESWAP_PRIMOPS) +peekUnalnWord16BE :: Ptr a -> IO Word16 +peekUnalnWord16BE = fmap byteSwap16 . peek . castPtr + +peekUnalnWord32BE :: Ptr a -> IO Word32 +peekUnalnWord32BE = fmap byteSwap32 . peek . castPtr + +#elif defined(MEM_UNALIGNED_OPS) && defined(WORDS_BIGENDIAN) +peekUnalnWord16BE :: Ptr a -> IO Word16 +peekUnalnWord16BE = peek . castPtr + +peekUnalnWord32BE :: Ptr a -> IO Word32 +peekUnalnWord32BE = peek . castPtr + +#else +peekUnalnWord16BE :: Ptr a -> IO Word16 +peekUnalnWord16BE p = do + x <- fromIntegral <$> peekWord8 (plusPtr p 0) + y <- fromIntegral <$> peekWord8 (plusPtr p 1) + return $! y .|. unsafeShiftL x 8 + +peekUnalnWord32BE :: Ptr a -> IO Word32 +peekUnalnWord32BE p = do + x <- fromIntegral <$> peekWord8 (plusPtr p 0) + y <- fromIntegral <$> peekWord8 (plusPtr p 1) + z <- fromIntegral <$> peekWord8 (plusPtr p 2) + w <- fromIntegral <$> peekWord8 (plusPtr p 3) + return $! w .|. unsafeShiftL z 8 .|. unsafeShiftL y 16 .|. unsafeShiftL x 24 +#endif + diff --git a/Bio/Util/Text.hs b/Bio/Util/Text.hs new file mode 100644 --- /dev/null +++ b/Bio/Util/Text.hs @@ -0,0 +1,51 @@ +module Bio.Util.Text + ( Unpack(..) + , w2c, c2w + , decodeBytes + , encodeBytes + , decompressGzip + ) where + +import BasePrelude +import Data.ByteString.Internal ( c2w, w2c ) +import Data.Text.Encoding ( encodeUtf8, decodeUtf8With ) + +import qualified Codec.Compression.Zlib.Internal as Z +import qualified Data.ByteString.Char8 as S +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString.Lazy.Internal as L ( ByteString(..) ) +import qualified Data.Text as T + +-- | Class of things that can be unpacked into 'String's. Kind of the +-- opposite of 'IsString'. +class Unpack s where unpack :: s -> String + +instance Unpack S.ByteString where unpack = S.unpack +instance Unpack T.Text where unpack = T.unpack +instance Unpack String where unpack = id + + +-- | Converts 'Bytes' into 'Text'. This uses UTF8, but if there is an +-- error, it pretends it was Latin1. Evil as this is, it tends to Just +-- Work on files where nobody ever wasted a thought on encodings. +decodeBytes :: S.ByteString -> T.Text +decodeBytes = decodeUtf8With (const $ fmap w2c) + +-- | Converts 'Text' into 'Bytes'. This uses UTF8. +encodeBytes :: T.Text -> S.ByteString +encodeBytes = encodeUtf8 + + +-- | Decompresses Gzip or Bgzf and passes everything else on. In +-- reality, it simply decompresses Gzip, and when done, looks for +-- another Gzip stream. Since there is a small chance to attempt +-- decompression of an uncompressed stream, the original data is +-- returned in case of an error. +decompressGzip :: L.ByteString -> L.ByteString +decompressGzip s = case L.uncons s of + Just (31, s') -> case L.uncons s' of + Just (139,_) -> Z.foldDecompressStreamWithInput L.Chunk decompressGzip (const s) + (Z.decompressST Z.gzipOrZlibFormat Z.defaultDecompressParams) s + _ -> s + _ -> s + diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,10 @@ +# 2.0.0 + + * Switched from "iteratee" to "streaming". The code looks much cleaner + and is easier to understand. + + * Repaired mistreatment of bam headers when merging files. + # 1.1.0 (2018-10-15) * Fix for previous workaround. Repairs writeBamFile. diff --git a/biohazard.cabal b/biohazard.cabal --- a/biohazard.cabal +++ b/biohazard.cabal @@ -1,5 +1,5 @@ Name: biohazard -Version: 1.1.1 +Version: 2.0 Synopsis: bioinformatics support library Description: This is a collection of modules I separated from various bioinformatics tools. @@ -16,7 +16,7 @@ Cabal-version: >= 1.10 Build-type: Simple -Tested-with: GHC == 7.8.4, GHC == 7.10.3, GHC == 8.0.2, GHC == 8.2.1, GHC == 8.4.3, GHC == 8.6.1 +Tested-with: GHC == 7.8.4, GHC == 7.10.3, GHC == 8.0.2, GHC == 8.2.2, GHC == 8.4.3, GHC == 8.6.1 source-repository head type: git @@ -39,53 +39,59 @@ Bio.Bam.Trim, Bio.Bam.Writer, Bio.Base, - Bio.Iteratee, - Bio.Iteratee.Base, - Bio.Iteratee.Bgzf, - Bio.Iteratee.Builder, - Bio.Iteratee.Bytes, - Bio.Iteratee.Exception, - Bio.Iteratee.IO, - Bio.Iteratee.Iteratee, - Bio.Iteratee.List, - Bio.Iteratee.ZLib, Bio.Prelude, + Bio.Streaming, + Bio.Streaming.Bgzf, + Bio.Streaming.Bytes, + Bio.Streaming.Furrow, + Bio.Streaming.Parse, + Bio.Streaming.Vector, Bio.TwoBit, Bio.Util.MMap, + Bio.Util.Nub, Bio.Util.Numeric, Bio.Util.Storable, - Bio.Util.Zlib + Bio.Util.Text - Build-depends: async >= 2.0 && < 2.3, - attoparsec >= 0.10 && < 0.14, + Build-depends: attoparsec >= 0.10 && < 0.14, base >= 4.7 && < 4.13, base-prelude == 1.2.0.*, bytestring >= 0.10.2 && < 0.11, containers >= 0.5 && < 0.7, + directory >= 1.2.1 && < 1.4, exceptions >= 0.6 && < 0.11, hashable >= 1.0 && < 1.3, primitive >= 0.5 && < 0.7, stm >= 2.4 && < 2.6, + streaming >= 0.1.4.2 && < 0.3, text >= 1.0 && < 1.3, transformers >= 0.4.1 && < 0.6, unix >= 2.5 && < 2.8, unordered-containers >= 0.2.3 && < 0.3, vector >= 0.11 && < 0.13, vector-algorithms >= 0.3 && < 0.8, - vector-th-unbox == 0.2.*, zlib == 0.6.* + if !impl(ghc >= 7.10) + build-depends: bifunctors == 5.* if !impl(ghc >= 8.0) build-depends: semigroups == 0.18.* Ghc-options: -Wall + if impl(ghc >= 8.0) + Ghc-options: -Wincomplete-uni-patterns -Wredundant-constraints + -Wcompat -Wincomplete-record-updates -Widentities + if impl(ghc >= 8.4) + Ghc-options: -Wmissing-export-lists -Wpartial-fields Default-Language: Haskell2010 Default-Extensions: BangPatterns, DeriveDataTypeable, + DeriveGeneric, FlexibleContexts, FlexibleInstances, + GeneralizedNewtypeDeriving, LambdaCase, MultiParamTypeClasses, NoImplicitPrelude, @@ -94,21 +100,19 @@ TypeSynonymInstances Other-Extensions: CPP, - DeriveGeneric, ExistentialQuantification, ForeignFunctionInterface, - GeneralizedNewtypeDeriving, Rank2Types, - TemplateHaskell, TypeFamilies - Hs-source-dirs: src - Include-dirs: src/cbits + Hs-source-dirs: . + Include-dirs: cbits Install-Includes: myers_align.h - C-sources: src/cbits/loops.c, - src/cbits/mmap.c, - src/cbits/myers_align.c, - src/cbits/trim.c + C-sources: cbits/loops.c, + cbits/mmap.c, + cbits/myers_align.c, + cbits/trim.c + cbits/zlib.c CC-options: -fPIC -- :vim:tw=132: diff --git a/cbits/loops.c b/cbits/loops.c new file mode 100644 --- /dev/null +++ b/cbits/loops.c @@ -0,0 +1,113 @@ +void nuc_loop( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u < v ) { + char a = q[ u ] ; + char b = q[ u + stride ] ; + char a1 = a ? 0x10 << (a&3) : 0xf0 ; + char b1 = b ? 0x1 << (b&3) : 0xf ; + *p++ = a1 | b1 ; + u += stride+stride ; + } + if( u == v ) { + char a = q[ u ] ; + char a1 = a ? 0x10 << (a&3) : 0xf0 ; + *p = a1 ; + } +} + +void nuc_loop_wide( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u < v ) { + char a = q[ u ] ; + *p++ = a ? 0x1 << (a&3) : 0xf ; + u += stride ; + } +} + +void nuc_loop_asc( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u <= v ) { + char a = q[ u ] ; + *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'A' : (a&3) == 1 ? 'C' : (a&3) == 2 ? 'G' : 'T' ; + u += stride ; + } + *p = 0 ; +} + +void nuc_loop_asc_rev( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + + while( u <= v ) { + char a = q[ v ] ; + *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'T' : (a&3) == 1 ? 'G' : (a&3) == 2 ? 'C' : 'A' ; + v -= stride ; + } + *p = 0 ; +} + +void qual_loop( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = (q[u] >> 2) & 0x3f ; + u += stride ; + } +} + +void qual_loop_asc( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = 33 + ((q[u] >> 2) & 0x3f) ; + u += stride ; + } + *p = 0 ; +} + +void qual_loop_asc_rev( char* p, int stride, char* q, int u, int v ) +{ + u *= stride ; + v *= stride ; + while( u <= v ) { + *p++ = 33 + ((q[v] >> 2) & 0x3f) ; + v -= stride ; + } + *p = 0 ; +} + +int int_loop( char* p, int x ) +{ + *p++ = ':' ; + if( x == 0 ) { + *p = '0' ; + return 2 ; + } + char *q = p ; + while( x > 0 ) { + *q++ = '0' + x % 10 ; + x /= 10 ; + } + int r = q-p ; + --q ; + while( p < q ) { + char c = *p ; + *p++ = *q ; + *q-- = c ; + } + return r+1 ; +} + + diff --git a/cbits/mmap.c b/cbits/mmap.c new file mode 100644 --- /dev/null +++ b/cbits/mmap.c @@ -0,0 +1,10 @@ +#include + +unsigned char *my_mmap(size_t len, int fd) { + void *result = mmap(0, len, PROT_READ, MAP_SHARED, fd, 0); + return (unsigned char*)( result == MAP_FAILED ? 0 : result ); +} + +void my_munmap(void *len, unsigned char *p) { + munmap( p, (size_t)len ) ; +} diff --git a/cbits/myers_align.c b/cbits/myers_align.c new file mode 100644 --- /dev/null +++ b/cbits/myers_align.c @@ -0,0 +1,102 @@ +#include "myers_align.h" + +#include +#include +#include + +// [*blech*, this looks and feels like FORTRAN.] +unsigned myers_diff( + const char *seq_a, int len_a, enum myers_align_mode mode, + const char* seq_b, int len_b, int maxd, + char *bt_a, char *bt_b ) +{ + // int len_a = strlen( seq_a ), len_b = strlen( seq_b ) ; + if( maxd > len_a + len_b ) maxd = len_a + len_b ; + + // in vee[d][k], d runs from 0 to maxd; k runs from -d to +d + int **vee = calloc( maxd, sizeof(int*) ) ; + + int d, dd, k, x, y, r = UINT_MAX ; + int *v_d_1 = 0, *v_d = 0 ; // "array slice" vee[.][d-1] + for( d = 0 ; d != maxd ; ++d, v_d_1 = v_d ) // D-paths in order of increasing D + { + v_d = d + (vee[d] = malloc( (2 * d + 1) * sizeof( int ) )) ; // "array slice" vee[.][d] + + for( k = max(-d,-len_a) ; k <= min(d,len_b) ; ++k ) // diagonals + { + if( d == 0 ) x = 0 ; + else if(d==1&&k==0) x = v_d_1[ k ]+1 ; + else if( k == -d ) x = v_d_1[ k+1 ] ; + else if( k == d ) x = v_d_1[ k-1 ]+1 ; // argh, need to check for d first, b/c -d+2 could be equal to d + else if( k == -d+1 ) x = max( v_d_1[ k ]+1, v_d_1[ k+1 ] ) ; + else if( k == d-1 ) x = max( v_d_1[ k-1 ]+1, v_d_1[ k ]+1 ) ; + else x = max3( v_d_1[ k-1 ]+1, v_d_1[ k ]+1, v_d_1[ k+1 ] ) ; + + y = x-k ; + while( x < len_b && y < len_a && match( seq_b[x], seq_a[y] ) ) ++x, ++y ; + v_d[ k ] = x ; + + if( + bt_a && bt_b && + (mode == myers_align_is_prefix || y == len_a) && + (mode == myers_align_has_prefix || x == len_b) ) + { + char *out_a = bt_a + len_a + d +2 ; + char *out_b = bt_b + len_b + d +2 ; + *--out_a = 0 ; + *--out_b = 0 ; + for( dd = d ; dd != 0 ; ) + { + if( k != -dd && k != dd && x == vee[ dd-1 ][ k + dd-1 ]+1 ) + { + --dd ; + --x ; + --y ; + *--out_b = seq_b[x] ; + *--out_a = seq_a[y] ; + } + else if( k > -dd+1 && x == vee[ dd-1 ][ k-1 + dd-1 ]+1 ) + { + --x ; + --k ; + --dd ; + *--out_b = seq_b[x] ; + *--out_a = '-' ; + } + else if( k < dd-1 && x == vee[ dd-1 ][ k+1 + dd-1 ] ) + { + ++k ; + --y ; + --dd ; + *--out_b = '-' ; + *--out_a = seq_a[y] ; + } + else // this better had been a match... + { + --x ; + --y ; + *--out_b = seq_b[x] ; + *--out_a = seq_a[y] ; + } + } + while( x > 0 ) + { + --x ; + *--out_b = seq_b[x] ; + *--out_a = seq_a[x] ; + } + memmove( bt_a, out_a, bt_a + len_a + d + 2 - out_a ) ; + memmove( bt_b, out_b, bt_b + len_b + d + 2 - out_b ) ; + r = d ; + goto cleanup ; + } + } + } + +cleanup: + for( dd = maxd ; dd != 0 ; --dd ) + free( vee[dd-1] ) ; + free( vee ) ; + return r ; +} + diff --git a/cbits/myers_align.h b/cbits/myers_align.h new file mode 100644 --- /dev/null +++ b/cbits/myers_align.h @@ -0,0 +1,76 @@ +#ifndef INCLUDED_MYERS_ALIGN +#define INCLUDED_MYERS_ALIGN + +enum myers_align_mode { + myers_align_globally = 0, + myers_align_is_prefix = 1, + myers_align_has_prefix = 2 } ; + +//! \brief aligns two sequences in O(nd) time +//! This alignment algorithm following Eugene W. Myers: "An O(ND) +//! Difference Algorithm and Its Variations". +//! Both input sequences are ASCIIZ-encoded with IUPAC-IUB ambiguity +//! codes. By definition, if ambiguity codes overlap, that's a match, +//! else a mismatch. Mismatches and gaps count a unit penalty. If mode +//! is myers_align_globally, both sequences must align completely. If +//! mode is myers_align_is_prefix, seq_a must align completely as prefix +//! of seq_b. If mode is myers_align_has_prefix, seq_b must align +//! completely as prefix of seq_a. +//! +//! Note that the calculation time is O(nd) where n is the length of the +//! best alignment and d the number of differences in it, memory +//! consumption is O(maxd^2). +//! +//! \param seq_a First input sequence. +//! \param mode How to align (i.e. what gaps to count). +//! \param seq_b Second input sequence. +//! \param maxd Maximum penalty to consider. +//! \param bt_a Space to backtrace seq_a into, must have room for +//! (strlen(seq_a)+maxd+1) characters. +//! \param bt_b Space to backtrace seq_b into, must have room for +//! (strlen(seq_b)+maxd+1) characters. +//! \return The actual edit distance or UINT_MAX if the edit distance +//! would be greater than maxd. +//! +unsigned myers_diff( + const char *seq_a, int len_a, enum myers_align_mode mode, + const char* seq_b, int len_b, int maxd, + char *bt_a, char *bt_b ) ; + +//! \brief converts an IUPAC-IUB ambiguity code to a bitmap Each base is +//! represented by a bit, makes checking for matches easier. +static inline int char_to_bitmap( char x ) +{ + switch( x & ~32 ) + { + case 'A': return 1 ; + case 'C': return 2 ; + case 'G': return 4 ; + case 'T': return 8 ; + case 'U': return 8 ; + + case 'S': return 6 ; + case 'W': return 9 ; + case 'R': return 5 ; + case 'Y': return 10 ; + case 'K': return 12 ; + case 'M': return 3 ; + + case 'B': return 14 ; + case 'D': return 13 ; + case 'H': return 11 ; + case 'V': return 7 ; + + case 'N': return 15 ; + default: return 0 ; + } +} + +static inline int compatible( char x, char y ) { return (char_to_bitmap(x) & char_to_bitmap(y)) != 0 ; } +static inline int match( char a, char b ) { return (char_to_bitmap(a) & char_to_bitmap(b)) != 0 ; } + +static inline int min( int a, int b ) { return a < b ? a : b ; } +static inline int max( int a, int b ) { return a < b ? b : a ; } +static inline int max3( int a, int b, int c ) { return a < b ? max( b, c ) : max( a, c ) ; } + +#endif diff --git a/cbits/trim.c b/cbits/trim.c new file mode 100644 --- /dev/null +++ b/cbits/trim.c @@ -0,0 +1,46 @@ +#include + +static const uint8_t compls[] = + { 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15 } ; + +int prim_match_reads( int i1 + , int i2 + , int r + , const uint8_t *rd1 + , const uint8_t *qs1 + , const uint8_t *rd2 + , const uint8_t *qs2 ) +{ + int acc = 0 ; + while( r != 0 ) + { + --i2 ; + uint8_t n1 = rd1[ i1 ] ; + uint8_t n2 = rd2[ i2 ] ; + uint8_t q1 = qs1[ i1 ] ; + uint8_t q2 = qs2[ i2 ] ; + + acc += (n1 & 0xF) == compls[ n2 & 0xF ] ? 0 : 5 + (q1 < q2 ? q1 : q2) ; + + ++i1 ; + --r ; + } + return acc ; +} + +int prim_match_ad( int off + , int i + , const uint8_t *rd + , const uint8_t *qs + , const uint8_t *ad ) +{ + int acc = 0 ; + while( i > 0 ) + { + --i; + acc += rd[ i+off ] == ad[ i ] ? 0 : 5 + + (qs[ i+off ] < 25 ? qs[ i+off ] : 25) ; + } + return acc ; +} + diff --git a/cbits/zlib.c b/cbits/zlib.c new file mode 100644 --- /dev/null +++ b/cbits/zlib.c @@ -0,0 +1,64 @@ +#include +#include +#include + +// header stolen from the EOF marker +static const char header[18] = "\x1f\x8b\x8\x4\0\0\0\0\0\xff\x6\0\x42\x43\x2\0" ; + +// Compresses one bgzf chunk. We assemble the header ourselves and call +// deflate in raw mode, because some tools seems to be extremely picky +// about irrelevant details. +int compress_chunk ( char *dest, int *dest_len, char *source, int source_len, int level ) +{ + z_stream stream = {0} ; + stream.next_in = source ; + stream.next_out = dest + 18 ; + stream.avail_in = source_len ; + stream.avail_out = *dest_len ; + + memmove( dest, header, 16 ) ; + + int rc = deflateInit2( &stream, level, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY ) ; + if( rc != Z_OK ) return rc ; + + rc = deflate( &stream, Z_FINISH ) ; + int rc2 = deflateEnd( &stream ) ; + if( rc != Z_STREAM_END ) return rc ; + if( rc2 != Z_OK ) return rc2 ; + + long crc = crc32( crc32( 0, 0, 0 ), source, source_len ) ; + + int compressed_length = 18 + 8 + stream.total_out ; + if (compressed_length > 65536) return Z_BUF_ERROR ; + + dest[16] = (compressed_length-1) & 0xff ; + dest[17] = (compressed_length-1) >> 8 ; + + *(uint32_t*)(dest+compressed_length-8) = crc ; + *(uint32_t*)(dest+compressed_length-4) = source_len ; + + *dest_len = compressed_length ; + return Z_OK ; +} + +// Decompresses one bgzf chunk. We receive no header, so we call +// inflate in raw mode and check the crc32 ourselves. +int decompress_chunk ( char *dest, int dest_len, char *source, int source_len ) +{ + z_stream stream = {0} ; + stream.next_in = source ; + stream.next_out = dest ; + stream.avail_in = source_len - 8 ; + stream.avail_out = dest_len ; + + int rc = inflateInit2( &stream, -15 ) ; + if( rc != Z_OK ) return rc ; + + rc = inflate( &stream, Z_FINISH ) ; + int rc2 = inflateEnd( &stream ) ; + if( rc != Z_STREAM_END ) return rc ; + if( rc2 != Z_OK ) return rc2 ; + + long crc = crc32( crc32( 0, 0, 0 ), dest, dest_len ) ; + return *(uint32_t*)(dest+dest_len-8) == crc && stream.avail_out == 0 ? Z_OK : Z_STREAM_ERROR ; +} diff --git a/src/Bio/Adna.hs b/src/Bio/Adna.hs deleted file mode 100644 --- a/src/Bio/Adna.hs +++ /dev/null @@ -1,661 +0,0 @@ -{-# LANGUAGE DeriveGeneric #-} - --- | Things specific to ancient DNA, e.g. damage models. --- --- For aDNA, we need a substitution probability. We have three options: --- use an empirically determined PSSM, use an arithmetically defined --- PSSM based on the /Johnson/ model, use a context sensitive PSSM based --- on the /Johnson/ model and an alignment. Using /Dindel/, actual --- substitutions relative to a called haplotype would be taken into --- account. Since we're not going to do that, taking alignments into --- account is difficult, somewhat approximate, and therefore not worth --- the hassle. - -module Bio.Adna ( - DmgStats(..), - CompositionStats, - SubstitutionStats, - addFragType, - damagePatternsIter, - damagePatternsIterMD, - damagePatternsIter2Bit, - alnFromMd, - - DamageParameters(..), - NewDamageParameters(..), - GenDamageParameters(..), - DamageModel, - bang, nudge, - Alignment(..), - FragType(..), - Subst(..), - - NPair, - npair, - fst_np, - snd_np, - - noDamage, - univDamage, - empDamage, - Mat44D(..), - MMat44D(..), - scalarMat, - complMat, - freezeMats, - - bwa_cal_maxdiff - ) where - -import Bio.Bam -import Bio.Prelude -import Bio.TwoBit - -import qualified Data.Vector as V -import qualified Data.Vector.Generic as G -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as U -import qualified Data.Vector.Unboxed.Mutable as UM - --- | We represent substitution matrices by the type 'Mat44D'. Internally, --- this is a vector of packed vectors. Conveniently, each of the packed --- vectors represents all transitions /into/ the given nucleotide. - -newtype Mat44D = Mat44D (U.Vector Double) deriving (Show, Generic) -newtype MMat44D = MMat44D (UM.IOVector Double) - --- | A 'DamageModel' is a function that gives substitution matrices for --- each position in a read. The 'DamageModel' can depend on whether the --- alignment is reversed, the length of the read and the position. (In --- practice, we should probably memoize precomputed damage models --- somehow.) - -type DamageModel = Bool -> Int -> Int -> Mat44D -data Subst = Nucleotide :-> Nucleotide deriving (Eq, Ord, Ix, Show) - -infix 9 :-> -infix 8 `bang` - --- | Convenience function to access a substitution matrix that has a --- mnemonic reading. -{-# INLINE bang #-} -bang :: Mat44D -> Subst -> Double -bang (Mat44D v) (N x :-> N y) - | U.length v == 16 = v U.! (fromIntegral x + 4 * fromIntegral y) - | otherwise = error $ "Huh? " ++ show (U.length v) - -{-# INLINE nudge #-} -nudge :: MMat44D -> Subst -> Double -> IO () -nudge (MMat44D v) (N x :-> N y) a = UM.read v i >>= UM.write v i . (+) a - where i = fromIntegral x + 4 * fromIntegral y - -scalarMat :: Double -> Mat44D -scalarMat s = Mat44D $ U.fromListN 16 [ s, 0, 0, 0 - , 0, s, 0, 0 - , 0, 0, s, 0 - , 0, 0, 0, s ] - -complMat :: Mat44D -> Mat44D -complMat v = Mat44D $ U.fromListN 16 [ v `bang` compl x :-> compl y - | y <- range (nucA, nucT) - , x <- range (nucA, nucT) ] - --- | Adds the two matrices of a mutable substitution model (one for each --- strand) appropriately, normalizes the result (to make probabilities --- from pseudo-counts), and freezes that into one immutable matrix. We --- add a single count everywhere to avoid getting NaNs from bizarre --- data. -freezeMats :: MMat44D -> MMat44D -> IO Mat44D -freezeMats (MMat44D vv) (MMat44D ww) = do - v <- Mat44D <$> U.freeze vv - w <- complMat . Mat44D <$> U.freeze ww - - let sums = U.generate 4 $ \x0 -> - let x = N $ fromIntegral x0 - in sum [ v `bang` x :-> z + w `bang` x :-> z - | z <- range (nucA, nucT) ] + 4 - - return . Mat44D $ U.fromListN 16 - [ (v `bang` x :-> y + w `bang` x :-> y + 1) / s - | y <- range (nucA, nucT) - , x <- range (nucA, nucT) - , let s = sums U.! fromIntegral (unN x) ] - - --- | 'DamageModel' for undamaged DNA. The likelihoods follow directly --- from the quality score. This needs elaboration to see what to do --- with amibiguity codes (even though those haven't actually been --- observed in the wild). - -noDamage :: DamageModel -noDamage _ _ _ = one - where !one = scalarMat 1 - - --- | Parameters for the universal damage model. --- --- We assume the correct model is either no damage, or single strand --- damage, or double strand damage. Each of them comes with a --- probability. It turns out that blending them into one is simply --- accomplished by multiplying these probabilities onto the deamination --- probabilities. --- --- For single stranded library prep, only one kind of damage occurs (C --- frequency ('ssd_sigma') in single stranded parts, and the overhang --- length is distributed exponentially with parameter 'ssd_lambda' at --- the 5' end and 'ssd_kappa' at the 3' end. (Without UDG treatment, --- those will be equal. With UDG, those are much smaller and in fact --- don't literally represent overhangs.) --- --- For double stranded library prep, we get C->T damage at the 5' end --- and G->A at the 3' end with rate 'dsd_sigma' and both in the interior --- with rate 'dsd_delta'. Everything is symmetric, and therefore the --- orientation of the aligned read doesn't matter either. Both --- overhangs follow a distribution with parameter 'dsd_lambda'. - -data DamageParameters float = DP { ssd_sigma :: !float -- deamination rate in ss DNA, SS model - , ssd_delta :: !float -- deamination rate in ds DNA, SS model - , ssd_lambda :: !float -- param for geom. distribution, 5' end, SS model - , ssd_kappa :: !float -- param for geom. distribution, 3' end, SS model - , dsd_sigma :: !float -- deamination rate in ss DNA, DS model - , dsd_delta :: !float -- deamination rate in ds DNA, DS model - , dsd_lambda :: !float } -- param for geom. distribution, DS model - deriving (Read, Show, Generic) - -data NewDamageParameters vec float = NDP { dp_gc_frac :: !float - , dp_mu :: !float - , dp_nu :: !float - , dp_alpha5 :: !(vec float) - , dp_beta5 :: !(vec float) - , dp_alpha :: !float - , dp_beta :: !float - , dp_alpha3 :: !(vec float) - , dp_beta3 :: !(vec float) } - deriving (Read, Show, Generic) - -data GenDamageParameters vec float - = UnknownDamage - | OldDamage (DamageParameters float) - | NewDamage (NewDamageParameters vec float) - deriving (Show, Generic, Read) - - - --- | Generic substitution matrix, has C->T and G->A deamination as --- parameters. Setting 'p' or 'q' to 0 as appropriate makes this apply --- to the single stranded or undamaged case. - -{-# INLINE genSubstMat #-} -genSubstMat :: Double -> Double -> Mat44D -genSubstMat p q = Mat44D $ U.fromListN 16 [ 1, 0, q, 0 - , 0, 1-p, 0, 0 - , 0, 0, 1-q, 0 - , 0, p, 0, 1 ] - -univDamage :: DamageParameters Double -> DamageModel -univDamage DP{..} r l i = genSubstMat (p1+p2) (q1+q2) - where - (p1, q1) = if r then let lam5 = ssd_lambda ^ (l-i) - lam3 = ssd_kappa ^ (1+i) - lam = lam3 + lam5 - lam3 * lam5 - p = ssd_sigma * lam + ssd_delta * (1-lam) - in (0,p) - else let lam5 = ssd_lambda ^ (1+i) - lam3 = ssd_kappa ^ (l-i) - lam = lam3 + lam5 - lam3 * lam5 - p = ssd_sigma * lam + ssd_delta * (1-lam) - in (p,0) - - p2 = dsd_sigma * lam5_ds + dsd_delta * (1-lam5_ds) - q2 = dsd_sigma * lam3_ds + dsd_delta * (1-lam3_ds) - lam5_ds = dsd_lambda ^ (1+i) - lam3_ds = dsd_lambda ^ (l-i) - -empDamage :: NewDamageParameters U.Vector Double -> DamageModel -empDamage NDP{..} = - \r l i -> if i+i < l then - if r then fromMaybe middleRev (rev5 V.!? i) - else fromMaybe middle (fwd5 V.!? i) - else - if r then fromMaybe middleRev (rev3 V.!? (l-i-1)) - else fromMaybe middle (fwd3 V.!? (l-i-1)) - where - !middle = genSubstMat' dp_alpha dp_beta - !middleRev = genSubstMat' dp_beta dp_alpha - - !fwd5 = V.zipWith genSubstMat' (G.convert dp_alpha5) (G.convert dp_beta5) - !fwd3 = V.zipWith genSubstMat' (G.convert dp_alpha3) (G.convert dp_beta3) - - !rev5 = V.zipWith genSubstMat' (G.convert dp_beta5) (G.convert dp_alpha5) - !rev3 = V.zipWith genSubstMat' (G.convert dp_beta3) (G.convert dp_alpha3) - - genSubstMat' a b = genSubstMat (recip $ 1 + exp (-a)) (recip $ 1 + exp (-b)) - - --- | Collected \"traditional\" statistics: --- --- * Base composition near 5' end and near 3' end. Each consists of --- five vectors of counts of A,C,G,T, and everything else. --- 'basecompo5' begins with 'context' bases to the left of the 5' end, --- 'basecompo3' ends with 'context' bases to the right of the 3' end. --- --- * Substitutions. Counted from the reconstructed alignment, once --- around the 5' end and once around the 3' end. For a total of 2*4*4 --- different substitutions. Positions where the query has a gap are --- skipped. --- --- * Substitutions at CpG motifs. Also counted from the reconstructed --- alignment, and a CpG site is simply the sequence CG in the --- reference. Gaps may confuse that definition, so that CpHpG still --- counts as CpG, because the H is gapped. That might actually --- be desirable. --- --- * Conditional substitutions. The 5' and 3' ends count as damaged if --- the very last position has a C-to-T substitution. With that in --- mind, 'substs5d5', 'substs5d3', 'substs5dd' are like 'substs5', but --- counting only reads where the 5' end is damaged, where the 3' end --- is damaged, and where both ends are damaged, respectively. - -data DmgStats a = DmgStats { - basecompo5 :: CompositionStats, - basecompo3 :: CompositionStats, - substs5 :: SubstitutionStats, - substs3 :: SubstitutionStats, - substs5d5 :: SubstitutionStats, - substs3d5 :: SubstitutionStats, - substs5d3 :: SubstitutionStats, - substs3d3 :: SubstitutionStats, - substs5dd :: SubstitutionStats, - substs3dd :: SubstitutionStats, - substs5cpg :: SubstitutionStats, - substs3cpg :: SubstitutionStats, - stats_more :: a } - deriving Show - -type CompositionStats = [( Maybe Nucleotide, U.Vector Int )] -type SubstitutionStats = [( Subst, U.Vector Int )] - - -data FragType = Complete | Leading | Trailing deriving (Show, Eq) - --- | Compact storage of a pair of ambiguous 'Nucleotides'. Used to --- represent alignments in a way that is accessible even to assembly --- code. The first and sencond field are stored in the low and high --- nybble, respectively. See 'fst_np', 'snd_np', 'npair'. -newtype NPair = NPair Word8 deriving (Eq, Ord) - -npair :: Nucleotides -> Nucleotides -> NPair -npair (Ns r) (Ns q) = NPair $ shiftL q 4 .|. r .&. 0xF - -fst_np, snd_np :: NPair -> Nucleotides -fst_np (NPair w) = Ns (w .&. 0xF) -snd_np (NPair w) = Ns (shiftR w 4) - -instance Storable NPair where - sizeOf _ = 1 - alignment _ = 1 - peek p = NPair <$> peek (castPtr p :: Ptr Word8) - poke p (NPair v) = poke (castPtr p :: Ptr Word8) v - -instance Show NPair where - showsPrec _ p = shows (fst_np p) . (:) '/' . shows (snd_np p) - --- | Alignment record. The reference sequence is filled with Ns if --- missing. -data Alignment = ALN - { a_sequence :: !(VS.Vector NPair) -- the alignment proper - , a_fragment_type :: !FragType } -- was the adapter trimmed? - -addFragType :: BamMeta -> Enumeratee [BamRaw] [(BamRaw,FragType)] m b -addFragType meta = mapStream $ \br -> (br, case unpackBam br of - b | isFirstMate b && isPaired b -> Leading - | isSecondMate b && isPaired b -> Trailing - | not sane -> Complete -- leeHom fscked it up - | isFirstMate b || isSecondMate b -> Complete -- old style flagging - | isTrimmed b || isMerged b -> Complete -- new style flagging - | otherwise -> Leading) - where - sane = null [ () | ("PG",line) <- meta_other_shit meta - , ("PN","mergeTrimReadsBAM") <- line ] - --- | Enumeratee (almost) that computes some statistics from plain BAM --- (no MD field needed) and a 2bit file. The 'Alignment' is also --- reconstructed and passed downstream. The result of any downstream --- processing is available in the 'stats_more' field of the result. --- --- * Get the reference sequence including both contexts once. If this --- includes invalid sequence (negative coordinate), pad suitably. --- * Accumulate counts for the valid parts around 5' and 3' ends as --- appropriate from flags and config. --- * Combine the part that was aligned to (so no context) with the read --- to reconstruct the alignment. --- --- Arguments are the table of reference names, the 2bit file with the --- reference, the amount of context outside the alignment desired, and --- the amount of context inside desired. --- --- For 'Complete' fragments, we cut the read in the middle, so the 5' --- and 3' plots stay clean from each other's influence. 'Leading' and --- 'Trailing' fragments count completely towards the appropriate end. - -damagePatternsIter2Bit :: MonadIO m - => Refs -> TwoBitFile -> Int -> Int - -> Iteratee [Alignment] m b - -> Iteratee [(BamRaw,FragType)] m (DmgStats b) -damagePatternsIter2Bit refs tbf ctx rng it = - mapMaybeStream (\(br,ft) -> do - let b@BamRec{..} = unpackBam br - guard (not $ isUnmapped b) - let ref_nm = sq_name $ getRef refs b_rname - ref = getFragment tbf ref_nm (b_pos - ctx) (alignedLength b_cigar + 2*ctx) - pps = aln_from_ref (U.drop ctx ref) b_seq b_cigar - return (b, ft, ref, pps)) =$ - damagePatternsIter ctx rng it - --- | Enumeratee (almost) that computes some statistics from plain BAM --- with a valid MD field. The 'Alignment' is also reconstructed and --- passed downstream. The result of any downstream processing is --- available in the 'stats_more' field of the result. --- --- * Reconstruct the alignment from CIGAR, SEQ, and MD. --- * Filter the alignment to get the reference sequence, accumulate it. --- * Accumulate everything over the alignment. --- --- The argument is the amount of context inside desired. --- --- For 'Complete' fragments, we cut the read in the middle, so the 5' --- and 3' plots stay clean from each other's influence. 'Leading' and --- 'Trailing' fragments count completely towards the appropriate end. - -damagePatternsIterMD :: MonadIO m - => Int -> Iteratee [Alignment] m b - -> Iteratee [(BamRaw,FragType)] m (DmgStats b) -damagePatternsIterMD rng it = - mapMaybeStream (\(br,ft) -> do - let b@BamRec{..} = unpackBam br - guard (not $ isUnmapped b) - md <- getMd b - let pps = alnFromMd b_seq b_cigar md - ref = U.convert $ VS.map fromN $ VS.filter ((/=) gap) $ VS.map fst_np pps - return (b, ft, ref, pps)) =$ - damagePatternsIter 0 rng it - where - fromN ns | ns == nucsA = 2 - | ns == nucsC = 1 - | ns == nucsG = 3 - | ns == nucsT = 0 - | otherwise = 4 - --- | Common logic for statistics. The function 'get_ref_and_aln' --- reconstructs reference sequence and alignment from a Bam record. It --- is expected to construct the alignment with respect to the forwards --- strand of the reference; we reverse-complement it if necessary. -damagePatternsIter :: MonadIO m - => Int -> Int - -> Iteratee [Alignment] m b - -> Iteratee [(BamRec, FragType, U.Vector Word8, VS.Vector NPair)] m (DmgStats b) -damagePatternsIter ctx rng it = mapStream revcom_both =$ do - let maxwidth = ctx + rng - acc_bc <- liftIO $ UM.replicate (2 * 5 * maxwidth) (0::Int) - acc_st <- liftIO $ UM.replicate (2 * 4 * 4 * 4 * rng) (0::Int) - acc_cg <- liftIO $ UM.replicate (2 * 2 * 4 * rng) (0::Int) - - it' <- flip mapStreamM it $ \(BamRec{..}, a_fragment_type, ref, a_sequence) -> liftIO $ do - -- basecompositon near 5' end, near 3' end - let (width5, width3) = case a_fragment_type of - Leading -> (full_width, 0) - Trailing -> (0, full_width) - Complete -> (half_width, half_width) - where full_width = min (U.length ref) $ ctx + min rng (alignedLength b_cigar) - half_width = min (U.length ref) $ ctx + min rng (alignedLength b_cigar `div` 2) - mapM_ (\i -> bump (fromIntegral (ref U.! i ) * maxwidth + i) acc_bc) [0 .. width5-1] - mapM_ (\i -> bump (fromIntegral (ref U.! (i + U.length ref) +6) * maxwidth + i) acc_bc) [-width3 .. -1] - - -- For substitutions, decide what damage class we're in: - -- 0 - no damage, 1 - damaged 5' end, 2 - damaged 3' end, 3 - both - let dmgbase = 2*4*4*rng * ( (if VS.null a_sequence || VS.head a_sequence /= npair nucsC nucsT then 1 else 0) - + (if VS.null a_sequence || VS.last a_sequence /= npair nucsC nucsT then 2 else 0) ) - - -- substitutions near 5' end - let len_at_5 = case a_fragment_type of Leading -> min rng (G.length b_seq) - Complete -> min rng (G.length b_seq `div` 2) - Trailing -> 0 - flip G.imapM_ (VS.take len_at_5 a_sequence) $ - \i uv -> withPair uv $ \j -> bump (j * rng + i + dmgbase) acc_st - - -- substitutions at CpG sites near 5' end - G.izipWithM_ - (\i uv wz -> - when (fst_np uv == nucsC && fst_np wz == nucsG) $ do - withNs (snd_np uv) $ \y -> bump ( y * rng + i ) acc_cg - withNs (snd_np wz) $ \y -> bump ((y+4) * rng + i+1) acc_cg) - (VS.take len_at_5 a_sequence) (VS.drop 1 a_sequence) - - -- substitutions near 3' end - let len_at_3 = case a_fragment_type of Leading -> 0 - Complete -> min rng (G.length b_seq `div` 2) - Trailing -> min rng (G.length b_seq) - flip G.imapM_ (VS.take len_at_3 (VS.reverse a_sequence)) $ - \i uv -> withPair uv $ \j -> bump ((17+j) * rng -i -1 + dmgbase) acc_st - - -- substitutions at CpG sites near 3' end - G.izipWithM_ - (\i wz uv -> - when (fst_np uv == nucsC && fst_np wz == nucsG) $ do - withNs (snd_np uv) $ \y -> bump ((y+ 9) * rng - i-2) acc_cg - withNs (snd_np wz) $ \y -> bump ((y+13) * rng - i-1) acc_cg) - (VS.take len_at_3 (VS.reverse a_sequence)) - (VS.drop 1 (VS.reverse a_sequence)) - - return ALN{..} - - - let nsubsts = 4*4*rng - mk_substs off = sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((4*i+j)*rng + off*nsubsts) rng acc_st) - | (i,n1) <- zip [0..] [nucA..nucT] - , (j,n2) <- zip [0..] [nucA..nucT] ] - - accs <- liftIO $ DmgStats <$> sequence [ (,) nuc <$> U.unsafeFreeze (UM.slice (i*maxwidth) maxwidth acc_bc) - | (i,nuc) <- zip [2,1,3,0,4] [Just nucA,Just nucC,Just nucG,Just nucT,Nothing] ] - <*> sequence [ (,) nuc <$> U.unsafeFreeze (UM.slice (i*maxwidth) maxwidth acc_bc) - | (i,nuc) <- zip [7,6,8,5,9] [Just nucA,Just nucC,Just nucG,Just nucT,Nothing] ] - - <*> mk_substs 0 - <*> mk_substs 1 - <*> mk_substs 2 - <*> mk_substs 3 - <*> mk_substs 4 - <*> mk_substs 5 - <*> mk_substs 6 - <*> mk_substs 7 - - <*> sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((i+j)*rng) rng acc_cg) - | (i,n1) <- [(0,nucC), (4,nucG)] - , (j,n2) <- zip [0..] [nucA..nucT] ] - - <*> sequence [ (,) (n1 :-> n2) <$> U.unsafeFreeze (UM.slice ((i+j)*rng) rng acc_cg) - | (i,n2) <- [(8,nucC), (12,nucG)] - , (j,n1) <- zip [0..] [nucA..nucT] ] - - accs' <- accs `liftM` lift (run it') - return $ accs' { substs5 = mconcat [ substs5 accs', substs5d5 accs', substs5d3 accs', substs5dd accs' ] - , substs3 = mconcat [ substs3 accs', substs3d5 accs', substs3d3 accs', substs3dd accs' ] - , substs5d5 = mconcat [ substs5d5 accs', substs5dd accs'] - , substs3d5 = mconcat [ substs3d5 accs', substs3dd accs'] - , substs5d3 = mconcat [ substs5d3 accs', substs5dd accs'] - , substs3d3 = mconcat [ substs3d3 accs', substs3dd accs'] } - where - {-# INLINE withPair #-} - withPair (NPair i) k = - case pairTab `U.unsafeIndex` fromIntegral i of - j -> when (j >= 0) (k j) - - !pairTab = U.replicate 256 (-1) U.// - [ (fromIntegral i, x*4+y) | (u,x) <- zip [nucsA, nucsC, nucsG, nucsT] [0,1,2,3] - , (v,y) <- zip [nucsA, nucsC, nucsG, nucsT] [0,1,2,3] - , let NPair i = npair u v ] - {-# INLINE bump #-} - bump i v = UM.unsafeRead v i >>= UM.unsafeWrite v i . succ - - {-# INLINE withNs #-} - withNs ns k | ns == nucsA = k 0 - | ns == nucsC = k 1 - | ns == nucsG = k 2 - | ns == nucsT = k 3 - | otherwise = return () - - -instance Monoid a => Monoid (DmgStats a) where - mappend = merge_dmg_stats mappend - mempty = DmgStats { basecompo5 = empty_compo - , basecompo3 = empty_compo - , substs5 = empty_subst - , substs3 = empty_subst - , substs5d5 = empty_subst - , substs3d5 = empty_subst - , substs5d3 = empty_subst - , substs3d3 = empty_subst - , substs5dd = empty_subst - , substs3dd = empty_subst - , substs5cpg = empty_subst - , substs3cpg = empty_subst - , stats_more = mempty } - where - empty_compo = [ (nuc, U.empty) | nuc <- [Just nucA, Just nucC, Just nucG, Just nucT, Nothing] ] - empty_subst = [ (n1 :-> n2, U.empty) | n1 <- [nucA..nucT], n2 <- [nucA..nucT] ] - -instance Semigroup a => Semigroup (DmgStats a) where - (<>) = merge_dmg_stats (<>) - -merge_dmg_stats :: (a -> a -> a) -> DmgStats a -> DmgStats a -> DmgStats a -merge_dmg_stats plus a b = - DmgStats { basecompo5 = zipWith s1 (basecompo5 a) (basecompo5 b) - , basecompo3 = zipWith s1 (basecompo3 a) (basecompo3 b) - , substs5 = zipWith s2 (substs5 a) (substs5 b) - , substs3 = zipWith s2 (substs3 a) (substs3 b) - , substs5d5 = zipWith s2 (substs5d5 a) (substs5d5 b) - , substs3d5 = zipWith s2 (substs3d5 a) (substs3d5 b) - , substs5d3 = zipWith s2 (substs5d3 a) (substs5d3 b) - , substs3d3 = zipWith s2 (substs3d3 a) (substs3d3 b) - , substs5dd = zipWith s2 (substs5dd a) (substs5dd b) - , substs3dd = zipWith s2 (substs3dd a) (substs3dd b) - , substs5cpg = zipWith s2 (substs5cpg a) (substs5cpg b) - , substs3cpg = zipWith s2 (substs3cpg a) (substs3cpg b) - , stats_more = plus (stats_more a) (stats_more b) } - where - s1 (x, u) (z, v) | x /= z = error "Mismatch in zip. This is a bug." - | U.null u = (x, v) - | U.null v = (x, u) - | otherwise = (x, U.zipWith (+) u v) - - s2 (x :-> y, u) (z :-> w, v) | x /= z || y /= w = error "Mismatch in zip. This is a bug." - | U.null u = (x :-> y, v) - | U.null v = (x :-> y, u) - | otherwise = (x :-> y, U.zipWith (+) u v) - - -revcom_both :: ( BamRec, FragType, U.Vector Word8, VS.Vector NPair ) - -> ( BamRec, FragType, U.Vector Word8, VS.Vector NPair ) -revcom_both (b, ft, ref, pps) - | isReversed b = ( b, ft, revcom_ref ref, revcom_pairs pps ) - | otherwise = ( b, ft, ref, pps ) - where - revcom_ref = U.reverse . U.map (\c -> if c > 3 then c else xor c 2) - revcom_pairs = VS.reverse . VS.map (\p -> npair (compls $ fst_np p) (compls $ snd_np p)) - - --- | Reconstructs the alignment from reference, query, and cigar. Only --- positions where the query is not gapped are produced. -aln_from_ref :: U.Vector Word8 -> Vector_Nucs_half Nucleotides -> VS.Vector Cigar -> VS.Vector NPair -aln_from_ref ref0 qry0 cig0 = VS.fromList $ step ref0 qry0 cig0 - where - step ref qry cig1 - | U.null ref || G.null qry || G.null cig1 = [] - | otherwise = case G.unsafeHead cig1 of { op :* n -> - case G.unsafeTail cig1 of { cig -> - case op of { - - Mat -> zipWith (npair . nn) (G.toList (G.take n ref)) - (G.toList (G.take n qry)) ++ step (G.drop n ref) (G.drop n qry) cig ; - Del -> step (G.drop n ref) qry cig ; - Ins -> map (npair gap) (G.toList (G.take n qry)) ++ step ref (G.drop n qry) cig ; - SMa -> map (npair gap) (G.toList (G.take n qry)) ++ step ref (G.drop n qry) cig ; - HMa -> replicate n (npair gap nucsN) ++ step ref qry cig ; - Nop -> step ref qry cig ; - Pad -> step ref qry cig }}} - - nn 0 = nucsT - nn 1 = nucsC - nn 2 = nucsA - nn 3 = nucsG - nn _ = nucsN - - --- | Reconstructs the alignment from query, cigar, and md. Only --- positions where the query is not gapped are produced. -alnFromMd :: Vector_Nucs_half Nucleotides -> VS.Vector Cigar -> [MdOp] -> VS.Vector NPair -alnFromMd qry0 cig0 md0 = VS.fromList $ step qry0 cig0 md0 - where - step qry cig1 md - | G.null qry || G.null cig1 || null md = [] - | otherwise = case G.unsafeHead cig1 of op :* n -> step' qry op n (G.unsafeTail cig1) md - - step' qry _ 0 cig md = step qry cig md - step' qry op n cig (MdNum 0 : md) = step' qry op n cig md - step' qry op n cig (MdDel [] : md) = step' qry op n cig md - - step' qry Mat n cig (MdNum m : md) - | n < m = map twin (G.toList (G.take n qry)) ++ step (G.drop n qry) cig (MdNum (m-n) : md) - | n > m = map twin (G.toList (G.take m qry)) ++ step' (G.drop m qry) Mat (n-m) cig md - | n == m = map twin (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md - step' qry Mat n cig (MdRep c : md) = npair c (G.head qry) : step' (G.tail qry) Mat (n-1) cig md - step' _ Mat _ _ _ = [] - - step' qry Del n cig (MdDel (_:ss) : md) = step' qry Del (n-1) cig (MdDel ss : md) - step' _ Del _ _ _ = [] - - step' qry Ins n cig md = map (npair gap) (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md - step' qry SMa n cig md = map (npair gap) (G.toList (G.take n qry)) ++ step (G.drop n qry) cig md - step' qry HMa n cig md = replicate n (npair gap nucsN) ++ step qry cig md - step' qry Nop _ cig md = step qry cig md - step' qry Pad _ cig md = step qry cig md - - twin q = npair q q - --- | Number of mismatches allowed by BWA. --- @bwa_cal_maxdiff thresh len@ returns the number of mismatches --- @bwa aln -n $tresh@ would allow in a read of length @len@. For --- reference, here is the code from BWA that computes it (we assume @err --- = 0.02@, just like BWA): --- --- @ --- int bwa_cal_maxdiff(int l, double err, double thres) --- { --- double elambda = exp(-l * err); --- double sum, y = 1.0; --- int k, x = 1; --- for (k = 1, sum = elambda; k < 1000; ++k) { --- y *= l * err; --- x *= k; --- sum += elambda * y / x; --- if (1.0 - sum < thres) return k; --- } --- return 2; --- } --- @ - -bwa_cal_maxdiff :: Double -> Int -> Int -bwa_cal_maxdiff thresh len = k_fin-1 - where - (k_fin, _, _, _) : _ = dropWhile bad $ iterate step (1,elambda,1,1) - - err = 0.02 - elambda = exp . negate $ fromIntegral len * err - - step (k, s, x, y) = (k+1, s', x', y') - where y' = y * fromIntegral len * err - x' = x * fromIntegral k - s' = s + elambda * y' / x' - - bad (_, s, _, _) = 1-s >= thresh - diff --git a/src/Bio/Align.hs b/src/Bio/Align.hs deleted file mode 100644 --- a/src/Bio/Align.hs +++ /dev/null @@ -1,85 +0,0 @@ -module Bio.Align ( - Mode(..), - myersAlign, - showAligned - ) where - -import Bio.Prelude hiding ( lefts, rights ) -import Foreign.C.String ( CString ) -import Foreign.C.Types ( CInt(..) ) -import Foreign.Marshal.Alloc ( allocaBytes ) - -import qualified Data.ByteString.Char8 as S -import qualified Data.ByteString.Unsafe as S -import qualified Data.ByteString.Lazy.Char8 as L - -foreign import ccall unsafe "myers_align.h myers_diff" myers_diff :: - CString -> CInt -> -- sequence A and length A - CInt -> -- mode (an enum) - CString -> CInt -> -- sequence B and length B - CInt -> -- max distance - CString -> -- backtracing space A - CString -> -- backtracing space B - IO CInt -- returns distance - --- | Mode argument for 'myersAlign', determines where free gaps are --- allowed. -data Mode = Globally -- ^ align globally, without gaps at either end - | HasPrefix -- ^ align so that the second sequence is a prefix of the first - | IsPrefix -- ^ align so that the first sequence is a prefix of the second - deriving Enum - --- | Align two strings. @myersAlign maxd seqA mode seqB@ tries to align --- @seqA@ to @seqB@, which will work as long as no more than @maxd@ gaps --- or mismatches are incurred. The @mode@ argument determines if either --- of the sequences is allowed to have an overhanging tail. --- --- The result is the triple of the actual distance (gaps + mismatches) --- and the two padded sequences. These sequences are the original --- sequences with dashes inserted for gaps. --- --- The algorithm is the O(nd) algorithm by Myers, implemented in C. A --- gap and a mismatch score the same. The strings are supposed to code --- for DNA, the code understands IUPAC-IUB ambiguity codes. Two --- characters match iff there is at least one nucleotide both can code --- for. Note that N is a wildcard, while X matches nothing. - -myersAlign :: Int -> Bytes -> Mode -> Bytes -> (Int, Bytes, Bytes) -myersAlign maxd seqA mode seqB = - unsafePerformIO $ - S.unsafeUseAsCStringLen seqA $ \(seq_a, len_a) -> - S.unsafeUseAsCStringLen seqB $ \(seq_b, len_b) -> - - -- size of output buffers derives from this: - -- char *out_a = bt_a + len_a + maxd +2 ; - -- char *out_b = bt_b + len_b + maxd +2 ; - allocaBytes (len_a + maxd + 2) $ \bt_a -> - allocaBytes (len_b + maxd + 2) $ \bt_b -> - - myers_diff seq_a (fromIntegral len_a) - (fromIntegral $ fromEnum mode) - seq_b (fromIntegral len_b) - (fromIntegral maxd) bt_a bt_b >>= \dist -> - if dist < 0 - then return (maxBound, S.empty, S.empty) - else (,,) (fromIntegral dist) <$> - S.packCString bt_a <*> - S.packCString bt_b - - --- | Nicely print an alignment. An alignment is simply a list of --- strings with inserted gaps to make them align. We split them into --- manageable chunks, stack them vertically and add a line showing --- asterisks in every column where all aligned strings agree. The --- result is /almost/ the Clustal format. -showAligned :: Int -> [Bytes] -> [L.ByteString] -showAligned w ss | all S.null ss = [] - | otherwise = map (L.fromChunks . (:[])) lefts ++ - L.pack agreement : - L.empty : - showAligned w rights - where - (lefts, rights) = unzip $ map (S.splitAt w) ss - agreement = map star $ S.transpose lefts - star str = if S.null str || S.all (== S.head str) str then '*' else ' ' - diff --git a/src/Bio/Bam.hs b/src/Bio/Bam.hs deleted file mode 100644 --- a/src/Bio/Bam.hs +++ /dev/null @@ -1,24 +0,0 @@ --- | Umbrella module for most of what's under 'Bio.Bam'. - -module Bio.Bam ( - module Bio.Bam.Fastq, - module Bio.Bam.Filter, - module Bio.Bam.Header, - module Bio.Bam.Index, - module Bio.Bam.Reader, - module Bio.Bam.Rec, - module Bio.Bam.Trim, - module Bio.Bam.Writer, - module Bio.Iteratee - ) where - -import Bio.Bam.Fastq -import Bio.Bam.Filter -import Bio.Bam.Header -import Bio.Bam.Index -import Bio.Bam.Reader -import Bio.Bam.Rec -import Bio.Bam.Trim -import Bio.Bam.Writer -import Bio.Iteratee - diff --git a/src/Bio/Bam/Evan.hs b/src/Bio/Bam/Evan.hs deleted file mode 100644 --- a/src/Bio/Bam/Evan.hs +++ /dev/null @@ -1,97 +0,0 @@ --- | This module contains stuff relating to conventions local to MPI --- EVAN. The code is needed regularly, but it can be harmful when --- applied to BAM files that follow different conventions. Most --- importantly, no program should call these functions by default. - -module Bio.Bam.Evan where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Data.Bits -import Prelude - -import qualified Data.ByteString.Char8 as S - --- | Fixes abuse of flags valued 0x800 and 0x1000. We used them for --- low quality and low complexity, but they have since been redefined. --- If set, we clear them and store them into the ZQ field. Also fixes --- abuse of the combination of the paired, 1st mate and 2nd mate flags --- used to indicate merging or trimming. These are canonicalized and --- stored into the FF field. This function is unsafe on BAM files of --- unclear origin! -fixupFlagAbuse :: BamRec -> BamRec -fixupFlagAbuse b = - (if b_flag b .&. flag_low_quality /= 0 then setQualFlag 'Q' else id) $ -- low qual, new convention - (if b_flag b .&. flag_low_complexity /= 0 then setQualFlag 'C' else id) $ -- low complexity, new convention - b { b_flag = cleaned_flags, b_exts = cleaned_exts } - where - -- removes old flag abuse - flags' = b_flag b .&. complement (flag_low_quality .|. flag_low_complexity) - cleaned_flags | flags' .&. flagPaired == 0 = flags' .&. complement (flagFirstMate .|. flagSecondMate) - | otherwise = flags' - - flag_low_quality = 0x800 - flag_low_complexity = 0x1000 - - -- merged & trimmed from old flag abuse - is_merged = flags' .&. (flagPaired .|. flagFirstMate .|. flagSecondMate) == flagFirstMate .|. flagSecondMate - is_trimmed = flags' .&. (flagPaired .|. flagFirstMate .|. flagSecondMate) == flagSecondMate - newflags = (if is_merged then eflagMerged else 0) .|. (if is_trimmed then eflagTrimmed else 0) - - -- Extended flags, renamed to avoid collision with BWA. Goes like this: if FF is there, use - -- it. Else check if XF is there __and is numeric__. If so, use it, remove it, and set FF - -- instead. Else use 0 and leave it alone. Note that this resolves the collision with BWA, - -- since BWA puts a character there, not an int. - cleaned_exts = case (lookup "FF" (b_exts b), lookup "XF" (b_exts b)) of - ( Just (Int i), _ ) -> updateE "FF" (Int (i .|. newflags)) (b_exts b) - ( _, Just (Int i) ) -> updateE "FF" (Int (i .|. newflags)) $ deleteE "XF" (b_exts b) - _ | newflags /= 0 -> updateE "FF" (Int newflags ) (b_exts b) - | otherwise -> b_exts b - - --- | Fixes typical inconsistencies produced by Bwa: sometimes, 'mate unmapped' should be set, and we --- can see it, because we match the mate's coordinates. Sometimes 'properly paired' should not be --- set, because one mate is unmapped. This function is generally safe, but needs to be called only --- on the output of affected (older?) versions of Bwa. -fixupBwaFlags :: BamRec -> BamRec -fixupBwaFlags b = b { b_flag = fixPP $ b_flag b .|. if mu then flagMateUnmapped else 0 } - where - -- Set "mate unmapped" if self coordinates and mate coordinates are equal, but self is - -- paired and mapped. (BWA forgets this flag for invalid mate alignments) - mu = and [ isPaired b, not (isUnmapped b) - , isReversed b == isMateReversed b - , b_rname b == b_mrnm b, b_pos b == b_mpos b ] - - -- If either mate is unmapped, remove "properly paired". - fixPP f | f .&. (flagUnmapped .|. flagMateUnmapped) == 0 = f - | otherwise = f .&. complement flagProperlyPaired - - --- | Removes syntactic warts from old read names or the read names used --- in FastQ files. -removeWarts :: BamRec -> BamRec -removeWarts br = br { b_qname = name, b_flag = flags, b_exts = tags } - where - (name, flags, tags) = checkFR $ checkC $ checkSharp (b_qname br, b_flag br, b_exts br) - - checkFR (n,f,t) | "F_" `S.isPrefixOf` n = checkC (S.drop 2 n, f .|. flagFirstMate .|. flagPaired, t) - | "R_" `S.isPrefixOf` n = checkC (S.drop 2 n, f .|. flagSecondMate .|. flagPaired, t) - | "M_" `S.isPrefixOf` n = checkC (S.drop 2 n, f, insertE "FF" (Int eflagMerged) t) - | "T_" `S.isPrefixOf` n = checkC (S.drop 2 n, f, insertE "FF" (Int eflagTrimmed) t) - | "/1" `S.isSuffixOf` n = ( rdrop 2 n, f .|. flagFirstMate .|. flagPaired, t) - | "/2" `S.isSuffixOf` n = ( rdrop 2 n, f .|. flagSecondMate .|. flagPaired, t) - | otherwise = ( n, f, t) - - checkC (n,f,t) | "C_" `S.isPrefixOf` n = (S.drop 2 n, f, insertE "XP" (Int (-1)) t) - | otherwise = ( n, f, t) - - rdrop n s = S.take (S.length s - n) s - - checkSharp (n,f,t) = case S.split '#' n of [n',ts] -> (n', f, insertTags ts t) - _ -> ( n, f, t) - - insertTags ts t | S.null y = insertE "XI" (Text ts) t - | otherwise = insertE "XI" (Text x) $ insertE "XJ" (Text $ S.tail y) t - where (x,y) = S.break (== ',') ts - - diff --git a/src/Bio/Bam/Fastq.hs b/src/Bio/Bam/Fastq.hs deleted file mode 100644 --- a/src/Bio/Bam/Fastq.hs +++ /dev/null @@ -1,113 +0,0 @@ --- | Parser for @FastA/FastQ@, 'Iteratee' style, based on --- "Data.Attoparsec", and written such that it is compatible with module --- 'Bio.Bam'. This gives import of @FastA/FastQ@ while respecting some --- local (to MPI EVAN) conventions. - -module Bio.Bam.Fastq ( parseFastq, parseFastq', parseFastqCassava ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Prelude hiding ( isSpace ) -import Bio.Iteratee -import Data.Attoparsec.ByteString.Char8 - ( char, skipSpace, satisfy, inClass, skipWhile, takeTill - , scan, isSpace, isSpace_w8, () ) - -import qualified Data.Attoparsec.ByteString.Char8 as P -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as S -import qualified Data.Vector.Generic as V - --- | Reader for DNA (not protein) sequences in FastA and FastQ. We read --- everything vaguely looking like FastA or FastQ, then shoehorn it into --- a BAM record. We strive to extract information following more or --- less established conventions from the header, but don't aim for --- completeness. The recognized syntactical warts are converted into --- appropriate flags and removed. Only the canonical variant of FastQ --- is supported (qualities stored as raw bytes with offset 33). --- --- Supported additional conventions: --- --- * A name suffix of @/1@ or @/2@ is turned into the first mate or second --- mate flag and the read is flagged as paired. --- --- * Same for name prefixes of @F_@ or @R_@, respectively. --- --- * A name prefix of @M_@ flags the sequence as unpaired and merged --- --- * A name prefix of @T_@ flags the sequence as unpaired and trimmed --- --- * A name prefix of @C_@, optionally before or after any of the other --- prefixes, is turned into the extra flag @XP:i:-1@ (result of --- duplicate removal with unknown duplicate count). --- --- * A collection of tags separated from the name by an octothorpe is --- removed and put into the fields @XI@ and @XJ@ as text. --- --- Everything before the first sequence header is ignored. Headers can --- start with @\>@ or @\@@, we treat both equally. The first word of --- the header becomes the read name, the remainder of the header is --- ignored. The sequence can be split across multiple lines; --- whitespace, dashes and dots are ignored, IUPAC-IUB ambiguity codes --- are accepted as bases, anything else causes an error. The sequence --- ends at a line that is either a header or starts with @\+@, in the --- latter case, that line is ignored and must be followed by quality --- scores. There must be exactly as many Q-scores as there are bases, --- followed immediately by a header or end-of-file. Whitespace is --- ignored. - -parseFastq :: Monad m => Enumeratee Bytes [ BamRec ] m a -parseFastq = parseFastq' (const id) - --- | Like 'parseFastq', but also --- --- * If the first word of the description has at least four colon --- separated subfields, the first is used to flag first/second mate, --- the second is the \"QC failed\" flag, and the fourth is the index --- sequence. - -parseFastqCassava :: Monad m => Enumeratee Bytes [ BamRec ] m a -parseFastqCassava = parseFastq' (pdesc . S.split ':' . S.takeWhile (' ' /=)) - where - pdesc (num:flg:_:idx:_) br = br { b_flag = sum [ if num == "1" then flagFirstMate .|. flagPaired else 0 - , if num == "2" then flagSecondMate .|. flagPaired else 0 - , if flg == "Y" then flagFailsQC else 0 - , b_flag br .&. complement (flagFailsQC .|. flagSecondMate .|. flagPaired) ] - , b_exts = if S.all (`S.elem` "ACGTN") idx then insertE "XI" (Text idx) (b_exts br) else b_exts br } - pdesc _ br = br - --- | Same as 'parseFastq', but a custom function can be applied to the --- description string (the part of the header after the sequence name), --- which can modify the parsed record. Note that the quality field can --- end up empty. - -parseFastq' :: Monad m => ( Bytes -> BamRec -> BamRec ) -> Enumeratee Bytes [ BamRec ] m a -parseFastq' descr it = do skipJunk ; convStream (parserToIteratee $ (:[]) <$> pRec) it - where - isCBase = inClass "ACGTUBDHVSWMKRYNacgtubdhvswmkryn" - canSkip c = isSpace c || c == '.' || c == '-' - isHdr c = c == '@' || c == '>' - - pRec = (satisfy isHdr "start marker") *> (makeRecord <$> pName <*> (descr <$> P.takeWhile ('\n' /=)) <*> (pSeq >>= pQual)) - pName = takeTill isSpace <* skipWhile (\c -> c /= '\n' && isSpace c) "read name" - pSeq = (:) <$> satisfy isCBase <*> pSeq - <|> satisfy canSkip *> pSeq - <|> pure [] "sequence" - - pQual sq = (,) sq <$> (char '+' *> skipWhile ('\n' /=) *> pQual' (length sq) <* skipSpace <|> return S.empty) "qualities" - pQual' n = B.filter (not . isSpace_w8) <$> scan n step - step 0 _ = Nothing - step i c | isSpace c = Just i - | otherwise = Just (i-1) - -skipJunk :: Monad m => Iteratee Bytes m () -skipJunk = peekStreamBS >>= check - where - check (Just c) | bad c = dropWhileStreamBS (c2w '\n' /=) >> dropStreamBS 1 >> skipJunk - check _ = return () - bad c = c /= c2w '>' && c /= c2w '@' - -makeRecord :: Seqid -> (BamRec->BamRec) -> (String, Bytes) -> BamRec -makeRecord name extra (sq,qual) = extra $ nullBamRec - { b_qname = name, b_seq = V.fromList $ read sq, b_qual = V.fromList $ map (Q . subtract 33) $ B.unpack qual } - diff --git a/src/Bio/Bam/Filter.hs b/src/Bio/Bam/Filter.hs deleted file mode 100644 --- a/src/Bio/Bam/Filter.hs +++ /dev/null @@ -1,127 +0,0 @@ --- | Quality filters adapted from prehistoric pipeline. - -module Bio.Bam.Filter ( - filterPairs, QualFilter, - complexSimple, complexEntropy, - qualityAverage, qualityMinimum, - qualityFromOldIllumina, qualityFromNewIllumina - ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Base -import Bio.Iteratee -import Data.Bits -import Prelude - -import qualified Data.Vector.Generic as V - --- | A filter/transformation applied to pairs of reads. We supply a --- predicate to be applied to single reads and one to be applied to --- pairs, tha latter can get incomplete pairs, too, if mates have been --- separated or filtered asymmetrically. - -filterPairs :: Monad m => (BamRec -> [BamRec]) - -> (Maybe BamRec -> Maybe BamRec -> [BamRec]) - -> Enumeratee [BamRec] [BamRec] m a -filterPairs ps pp = eneeCheckIfDone step - where - step k = tryHead >>= step' k - step' k Nothing = return $ liftI k - step' k (Just b) - | isPaired b = tryHead >>= step'' k b - | otherwise = case ps b of [] -> step k ; b' -> eneeCheckIfDone step . k $ Chunk b' - - step'' k b Nothing = case pp (Just b) Nothing of - [] -> return $ liftI k - b' -> return $ k $ Chunk b' - - step'' k b (Just c) - | b_rname b /= b_rname c || not (isPaired c) = - let b' = if isSecondMate b then pp Nothing (Just b) else pp (Just b) Nothing - in case b' of [] -> step' k (Just c) - _ -> eneeCheckIfDone (\k' -> step' k' (Just c)) . k $ Chunk b' - - | isFirstMate c && isSecondMate b = step''' k c b - | otherwise = step''' k b c - - step''' k b c = case pp (Just b) (Just c) of [] -> step k - b' -> eneeCheckIfDone step . k $ Chunk b' - - --- | A quality filter is simply a transformation on @BamRec@s. By --- convention, quality filters should set @flagFailsQC@, a further step --- can then remove the failed reads. Filtering of individual reads --- tends to result in mate pairs with inconsistent flags, which in turn --- will result in lone mates and all sort of troubles with programs that --- expect non-broken BAM files. It is therefore recommended to use --- @pairFilter@ with suitable predicates to do the post processing. - -type QualFilter = BamRec -> BamRec - -{-# INLINE count #-} -count :: (V.Vector v a, Eq a) => a -> v a -> Int -count x = V.foldl' (\acc y -> if x == y then acc+1 else acc) 0 - --- | Simple complexity filter aka "Nancy Filter". A read is considered --- not-sufficiently-complex if the most common base accounts for greater --- than the @cutoff@ fraction of all non-N bases. -{-# INLINE complexSimple #-} -complexSimple :: Double -> QualFilter -complexSimple r b = if p then b else b' - where - b' = setQualFlag 'C' $ b { b_flag = b_flag b .|. flagFailsQC } - p = let counts = [ count x $ b_seq b | x <- properBases ] - lim = floor $ r * fromIntegral (sum counts) - in all (<= lim) counts - --- | Filter on order zero empirical entropy. Entropy per base must be --- greater than cutoff. -{-# INLINE complexEntropy #-} -complexEntropy :: Double -> QualFilter -complexEntropy r b = if p then b else b' - where - b' = setQualFlag 'C' $ b { b_flag = b_flag b .|. flagFailsQC } - p = ent >= r * total - - counts = [ count x $ b_seq b | x <- properBases ] - total = fromIntegral $ V.length $ b_seq b - ent = sum [ fromIntegral c * log (total / fromIntegral c) | c <- counts, c /= 0 ] / log 2 - --- | Filter on average quality. Reads without quality string pass. -{-# INLINE qualityAverage #-} -qualityAverage :: Int -> QualFilter -qualityAverage q b = if p then b else b' - where - b' = setQualFlag 'Q' $ b { b_flag = b_flag b .|. flagFailsQC } - p = let total = V.foldl' (\a x -> a + fromIntegral (unQ x)) 0 $ b_qual b - in total >= q * V.length (b_qual b) - --- | Filter on minimum quality. In @qualityMinimum n q@, a read passes --- if it has no more than @n@ bases with quality less than @q@. Reads --- without quality string pass. -{-# INLINE qualityMinimum #-} -qualityMinimum :: Int -> Qual -> QualFilter -qualityMinimum n (Q q) b = if p then b else b' - where - b' = setQualFlag 'Q' $ b { b_flag = b_flag b .|. flagFailsQC } - p = V.length (V.filter (< Q q) (b_qual b)) <= n - - --- | Convert quality scores from old Illumina scale (different formula --- and offset 64 in FastQ). -qualityFromOldIllumina :: BamRec -> BamRec -qualityFromOldIllumina b = b { b_qual = V.map conv $ b_qual b } - where - conv (Q s) = let s' :: Double - s' = exp $ log 10 * (fromIntegral s - 31) / (-10) - p = s' / (1+s') - q = - 10 * log p / log 10 - in Q (round q) - --- | Convert quality scores from new Illumina scale (standard formula --- but offset 64 in FastQ). -qualityFromNewIllumina :: BamRec -> BamRec -qualityFromNewIllumina b = b { b_qual = V.map (Q . subtract 31 . unQ) $ b_qual b } - - diff --git a/src/Bio/Bam/Header.hs b/src/Bio/Bam/Header.hs deleted file mode 100644 --- a/src/Bio/Bam/Header.hs +++ /dev/null @@ -1,397 +0,0 @@ -module Bio.Bam.Header ( - BamMeta(..), - parseBamMeta, - parseBamMetaLine, - showBamMeta, - addPG, - - BamKey(..), - BamHeader(..), - BamSQ(..), - BamSorting(..), - BamOtherShit, - - Refseq(..), - invalidRefseq, - isValidRefseq, - invalidPos, - isValidPos, - unknownMapq, - isKnownMapq, - - Refs, - noRefs, - getRef, - - compareNames, - - flagPaired, - flagProperlyPaired, - flagUnmapped, - flagMateUnmapped, - flagReversed, - flagMateReversed, - flagFirstMate, - flagSecondMate, - flagAuxillary, - flagSecondary, - flagFailsQC, - flagDuplicate, - flagSupplementary, - eflagTrimmed, - eflagMerged, - eflagAlternative, - eflagExactIndex, - - distinctBin, - - MdOp(..), - readMd, - showMd - ) where - -import Bio.Prelude hiding ( uncons ) -import Data.ByteString ( uncons ) -import Data.ByteString.Builder ( Builder, byteString, char7, intDec, word16LE ) -import Data.Sequence ( (><), (|>) ) - -import qualified Data.Attoparsec.ByteString.Char8 as P -import qualified Data.ByteString.Char8 as S -import qualified Data.Sequence as Z - -data BamMeta = BamMeta { - meta_hdr :: !BamHeader, - meta_refs :: !Refs, - meta_other_shit :: [(BamKey, BamOtherShit)], - meta_comment :: [Bytes] - } deriving Show - --- | Exactly two characters, for the \"named\" fields in bam. -newtype BamKey = BamKey Word16 - deriving ( Eq, Ord ) - -instance IsString BamKey where - {-# INLINE fromString #-} - fromString [a,b] - | ord a < 256 && ord b < 256 - = BamKey . fromIntegral $ ord a .|. shiftL (ord b) 8 - - fromString s - = error $ "Not a legal BAM key: " ++ show s - -instance Show BamKey where - show (BamKey a) = [ chr (fromIntegral a .&. 0xff), chr (shiftR (fromIntegral a) 8 .&. 0xff) ] - -addPG :: Maybe Version -> IO (BamMeta -> BamMeta) -addPG vn = do - args <- getArgs - pn <- getProgName - return $ go args pn - where - go args pn bm = bm { meta_other_shit = ("PG",pg_line) : meta_other_shit bm } - where - pg_line = concat [ [ ("ID", pg_id) ] - , [ ("PN", S.pack pn) ] - , [ ("CL", S.pack $ unwords args) ] - , maybe [] (\v -> [("VN",S.pack (showVersion v))]) vn - , map (\p -> ("PP",p)) (take 1 pg_pp) - , map (\p -> ("pp",p)) (drop 1 pg_pp) ] - - pg_id : _ = filter (not . flip elem pg_ids) . map S.pack $ - pn : [ pn ++ '-' : show i | i <- [(1::Int)..] ] - - pg_ids = [ pgid | ("PG",fs) <- meta_other_shit bm, ("ID",pgid) <- fs ] - pg_pps = [ pgid | ("PG",fs) <- meta_other_shit bm, ("PP",pgid) <- fs ] - - pg_pp = pg_ids \\ pg_pps - - -instance Monoid BamMeta where - mempty = BamMeta mempty noRefs [] [] - mappend = (<>) - -instance Semigroup BamMeta where - a <> b = BamMeta { meta_hdr = meta_hdr a `mappend` meta_hdr b - , meta_refs = meta_refs a >< meta_refs b - , meta_other_shit = nub $ meta_other_shit a ++ meta_other_shit b - , meta_comment = nub $ meta_comment a ++ meta_comment b } - -data BamHeader = BamHeader { - hdr_version :: (Int, Int), - hdr_sorting :: !BamSorting, - hdr_other_shit :: BamOtherShit - } deriving (Show, Eq) - -instance Monoid BamHeader where - mempty = BamHeader (1,0) Unknown [] - mappend = (<>) - -instance Semigroup BamHeader where - a <> b = BamHeader { hdr_version = hdr_version a `min` hdr_version b - , hdr_sorting = let u = hdr_sorting a ; v = hdr_sorting b in if u == v then u else Unknown - , hdr_other_shit = hdr_other_shit a ++ hdr_other_shit b } - -data BamSQ = BamSQ { - sq_name :: Seqid, - sq_length :: Int, - sq_other_shit :: BamOtherShit - } deriving (Show, Eq) - -bad_seq :: BamSQ -bad_seq = BamSQ (error "no SN field") (error "no LN field") [] - --- | Possible sorting orders from bam header. Thanks to samtools, which --- doesn't declare sorted files properly, we have to have the stupid --- 'Unknown' state, too. -data BamSorting = Unknown | Unsorted | Grouped | Queryname | Coordinate | GroupSorted - deriving (Show, Eq) - -type BamOtherShit = [(BamKey, Bytes)] - -parseBamMeta :: P.Parser BamMeta -parseBamMeta = fixup . foldl' (flip ($)) mempty <$> P.sepBy parseBamMetaLine (P.skipWhile (=='\t') >> P.char '\n') - where - fixup meta = meta { meta_other_shit = reverse (meta_other_shit meta) - , meta_comment = reverse (meta_comment meta) } - -parseBamMetaLine :: P.Parser (BamMeta -> BamMeta) -parseBamMetaLine = P.char '@' >> P.choice [hdLine, sqLine, coLine, otherLine] - where - hdLine = P.string "HD\t" >> - (\fns meta -> let fixup hdr = hdr { hdr_other_shit = reverse (hdr_other_shit hdr) } - in meta { meta_hdr = fixup $! foldl' (flip ($)) (meta_hdr meta) fns }) - <$> P.sepBy1 (P.choice [hdvn, hdso, hdother]) tabs - - sqLine = P.string "SQ\t" >> - (\fns meta -> let fixup sq = sq { sq_other_shit = reverse (sq_other_shit sq) } - !s = fixup $ foldl' (flip ($)) bad_seq fns - in meta { meta_refs = meta_refs meta |> s }) - <$> P.sepBy1 (P.choice [sqnm, sqln, sqother]) tabs - - hdvn = P.string "VN:" >> - (\a b hdr -> hdr { hdr_version = (a,b) }) - <$> P.decimal <*> ((P.char '.' <|> P.char ':') >> P.decimal) - - hdso = P.string "SO:" >> - (\s hdr -> hdr { hdr_sorting = s }) - <$> P.choice [ Grouped <$ P.string "grouped" - , Queryname <$ P.string "queryname" - , Coordinate <$ P.string "coordinate" - , GroupSorted <$ P.string "groupsort" - , Unsorted <$ P.string "unsorted" - , Unknown <$ P.skipWhile (\c -> c/='\t' && c/='\n') ] - - sqnm = P.string "SN:" >> (\s sq -> sq { sq_name = s }) <$> pall - sqln = P.string "LN:" >> (\i sq -> sq { sq_length = i }) <$> P.decimal - - hdother = (\t hdr -> t `seq` hdr { hdr_other_shit = t : hdr_other_shit hdr }) <$> tagother - sqother = (\t sq -> t `seq` sq { sq_other_shit = t : sq_other_shit sq }) <$> tagother - - coLine = P.string "CO\t" >> - (\s meta -> s `seq` meta { meta_comment = s : meta_comment meta }) - <$> P.takeWhile (/= 'n') - - otherLine = (\k ts meta -> meta { meta_other_shit = (k,ts) : meta_other_shit meta }) - <$> bamkey <*> (tabs >> P.sepBy1 tagother tabs) - - tagother :: P.Parser (BamKey,Bytes) - tagother = (,) <$> bamkey <*> (P.char ':' >> pall) - - tabs = P.char '\t' >> P.skipWhile (== '\t') - - pall :: P.Parser Bytes - pall = P.takeWhile (\c -> c/='\t' && c/='\n') - - bamkey :: P.Parser BamKey - bamkey = (\a b -> fromString [a,b]) <$> P.anyChar <*> P.anyChar - -showBamMeta :: BamMeta -> Builder -showBamMeta (BamMeta h ss os cs) = - show_bam_meta_hdr h <> - foldMap show_bam_meta_seq ss <> - foldMap show_bam_meta_other os <> - foldMap show_bam_meta_comment cs - where - show_bam_meta_hdr (BamHeader (major,minor) so os') = - byteString "@HD\tVN:" <> - intDec major <> char7 '.' <> intDec minor <> - byteString (case so of Unknown -> mempty - Unsorted -> "\tSO:unsorted" - Grouped -> "\tSO:grouped" - Queryname -> "\tSO:queryname" - Coordinate -> "\tSO:coordinate" - GroupSorted -> "\tSO:groupsort") <> - show_bam_others os' - - show_bam_meta_seq (BamSQ _ _ []) = mempty - show_bam_meta_seq (BamSQ nm ln ts) = - byteString "@SQ\tSN:" <> byteString nm <> - byteString "\tLN:" <> intDec ln <> show_bam_others ts - - show_bam_meta_comment cm = byteString "@CO\t" <> byteString cm <> char7 '\n' - - show_bam_meta_other (BamKey k,ts) = - char7 '@' <> word16LE k <> show_bam_others ts - - show_bam_others ts = - foldMap show_bam_other ts <> char7 '\n' - - show_bam_other (BamKey k,v) = - char7 '\t' <> word16LE k <> char7 ':' <> byteString v - - --- | Reference sequence in Bam --- Bam enumerates the reference sequences and then sorts by index. We --- need to track that index if we want to reproduce the sorting order. -newtype Refseq = Refseq { unRefseq :: Word32 } deriving (Eq, Ord, Ix) - -instance Show Refseq where - showsPrec p (Refseq r) = showsPrec p r - -instance Enum Refseq where - succ = Refseq . succ . unRefseq - pred = Refseq . pred . unRefseq - toEnum = Refseq . fromIntegral - fromEnum = fromIntegral . unRefseq - enumFrom = map Refseq . enumFrom . unRefseq - enumFromThen (Refseq a) (Refseq b) = map Refseq $ enumFromThen a b - enumFromTo (Refseq a) (Refseq b) = map Refseq $ enumFromTo a b - enumFromThenTo (Refseq a) (Refseq b) (Refseq c) = map Refseq $ enumFromThenTo a b c - - --- | Tests whether a reference sequence is valid. --- Returns true unless the the argument equals @invalidRefseq@. -isValidRefseq :: Refseq -> Bool -isValidRefseq = (/=) invalidRefseq - --- | The invalid Refseq. --- Bam uses this value to encode a missing reference sequence. -invalidRefseq :: Refseq -invalidRefseq = Refseq 0xffffffff - --- | The invalid position. --- Bam uses this value to encode a missing position. -{-# INLINE invalidPos #-} -invalidPos :: Int -invalidPos = -1 - --- | Tests whether a position is valid. --- Returns true unless the the argument equals @invalidPos@. -{-# INLINE isValidPos #-} -isValidPos :: Int -> Bool -isValidPos = (/=) invalidPos - -{-# INLINE unknownMapq #-} -unknownMapq :: Int -unknownMapq = 255 - -isKnownMapq :: Int -> Bool -isKnownMapq = (/=) unknownMapq - --- | A list of reference sequences. -type Refs = Z.Seq BamSQ - --- | The empty list of references. Needed for BAM files that don't really store alignments. -noRefs :: Refs -noRefs = Z.empty - -getRef :: Refs -> Refseq -> BamSQ -getRef refs (Refseq i) - | i < fromIntegral (Z.length refs) = Z.index refs (fromIntegral i) - | otherwise = BamSQ "*" 0 [] - - -flagPaired, flagProperlyPaired, flagUnmapped, flagMateUnmapped, - flagReversed, flagMateReversed, flagFirstMate, flagSecondMate, - flagAuxillary, flagSecondary, flagFailsQC, flagDuplicate, - flagSupplementary :: Int - -flagPaired = 0x1 -flagProperlyPaired = 0x2 -flagUnmapped = 0x4 -flagMateUnmapped = 0x8 -flagReversed = 0x10 -flagMateReversed = 0x20 -flagFirstMate = 0x40 -flagSecondMate = 0x80 -flagAuxillary = 0x100 -flagSecondary = 0x100 -flagFailsQC = 0x200 -flagDuplicate = 0x400 -flagSupplementary = 0x800 - -eflagTrimmed, eflagMerged, eflagAlternative, eflagExactIndex :: Int -eflagTrimmed = 0x1 -eflagMerged = 0x2 -eflagAlternative = 0x4 -eflagExactIndex = 0x8 - - --- | Compares two sequence names the way samtools does. --- samtools sorts by \"strnum_cmp\": --- --- * if both strings start with a digit, parse the initial --- sequence of digits and compare numerically, if equal, --- continue behind the numbers --- * else compare the first characters (possibly NUL), if equal --- continue behind them --- * else both strings ended and the shorter one counts as --- smaller (and that part is stupid) - -compareNames :: Seqid -> Seqid -> Ordering -compareNames n m = case (uncons n, uncons m) of - ( Nothing, Nothing ) -> EQ - ( Just _, Nothing ) -> GT - ( Nothing, Just _ ) -> LT - ( Just (c,n'), Just (d,m') ) - | is_digit c && is_digit d -> - let Just (u,n'') = S.readInt n - Just (v,m'') = S.readInt m - in case u `compare` v of - LT -> LT - GT -> GT - EQ -> n'' `compareNames` m'' - | otherwise -> case c `compare` d of - LT -> LT - GT -> GT - EQ -> n' `compareNames` m' - where - is_digit c = 48 <= c && c < 58 - - -data MdOp = MdNum Int | MdRep Nucleotides | MdDel [Nucleotides] deriving Show - -readMd :: Bytes -> Maybe [MdOp] -readMd s | S.null s = return [] - | isDigit (S.head s) = do (n,t) <- S.readInt s - (MdNum n :) <$> readMd t - | S.head s == '^' = let (a,b) = S.break isDigit (S.tail s) - in (MdDel (map toNucleotides $ S.unpack a) :) <$> readMd b - | otherwise = (MdRep (toNucleotides $ S.head s) :) <$> readMd (S.tail s) - --- | Normalizes a series of 'MdOp's and encodes them in the way BAM and --- SAM expect it. -showMd :: [MdOp] -> Bytes -showMd = S.pack . flip s1 [] - where - s1 (MdNum i : MdNum j : ms) = s1 (MdNum (i+j) : ms) - s1 (MdNum 0 : ms) = s1 ms - s1 (MdNum i : ms) = shows i . s1 ms - - s1 (MdRep r : ms) = shows r . s1 ms - - s1 (MdDel d1 : MdDel d2 : ms) = s1 (MdDel (d1++d2) : ms) - s1 (MdDel [] : ms) = s1 ms - s1 (MdDel ns : MdRep r : ms) = (:) '^' . shows ns . (:) '0' . shows r . s1 ms - s1 (MdDel ns : ms) = (:) '^' . shows ns . s1 ms - s1 [ ] = id - - --- | Computes the "distinct bin" according to the BAM binning scheme. If --- an alignment starts at @pos@ and its CIGAR implies a length of @len@ --- on the reference, then it goes into bin @distinctBin pos len@. -distinctBin :: Int -> Int -> Int -distinctBin beg len = mkbin 14 $ mkbin 17 $ mkbin 20 $ mkbin 23 $ mkbin 26 0 - where end = beg + len - 1 - mkbin n x = if beg `shiftR` n /= end `shiftR` n then x - else ((1 `shiftL` (29-n))-1) `div` 7 + (beg `shiftR` n) diff --git a/src/Bio/Bam/Index.hs b/src/Bio/Bam/Index.hs deleted file mode 100644 --- a/src/Bio/Bam/Index.hs +++ /dev/null @@ -1,334 +0,0 @@ -module Bio.Bam.Index ( - BamIndex(..), - readBamIndex, - readBaiIndex, - readTabix, - - Region(..), - Subsequence(..), - eneeBamRefseq, - eneeBamSubseq, - eneeBamRegions, - eneeBamUnaligned -) where - -import Bio.Bam.Header -import Bio.Bam.Reader -import Bio.Bam.Rec -import Bio.Bam.Regions ( Region(..), Subsequence(..) ) -import Bio.Iteratee -import Bio.Prelude - -import qualified Bio.Bam.Regions as R -import qualified Data.IntMap.Strict as M -import qualified Data.ByteString as B -import qualified Data.Vector as V -import qualified Data.Vector.Mutable as W -import qualified Data.Vector.Unboxed as U -import qualified Data.Vector.Unboxed.Mutable as N -import qualified Data.Vector.Algorithms.Intro as N - --- | Full index, unifying BAI and CSI style. In both cases, we have the --- binning scheme, parameters are fixed in BAI, but variable in CSI. --- Checkpoints are created from the linear index in BAI or from the --- `loffset' field in CSI. - -data BamIndex a = BamIndex { - -- | Minshift parameter from CSI - minshift :: {-# UNPACK #-} !Int, - - -- | Depth parameter from CSI - depth :: {-# UNPACK #-} !Int, - - -- | Best guess at where the unaligned records start - unaln_off :: {-# UNPACK #-} !Int64, - - -- | Room for stuff (needed for tabix) - extensions :: a, - - -- | Records for the binning index, where each bin has a list of - -- segments belonging to it. - refseq_bins :: {-# UNPACK #-} !(V.Vector Bins), - - -- | Known checkpoints of the form (pos,off) where off is the - -- virtual offset of the first record crossing pos. - refseq_ckpoints :: {-# UNPACK #-} !(V.Vector Ckpoints) } - - deriving Show - --- | Mapping from bin number to vector of clusters. -type Bins = IntMap Segments -type Segments = U.Vector (Int64,Int64) - - --- | Checkpoints. Each checkpoint is a position with the virtual offset --- where the first alignment crossing the position is found. In BAI, we --- get this from the 'ioffset' vector, in CSI we get it from the --- 'loffset' field: "Given a region [beg,end), we only need to visit --- chunks whose end file offset is larger than 'ioffset' of the 16kB --- window containing 'beg'." (Sounds like a marginal gain, though.) - -type Ckpoints = IntMap Int64 - - --- | Decode only those reads that fall into one of several regions. --- Strategy: We will scan the file mostly linearly, but only those --- regions that are actually needed. We filter the decoded stuff so --- that it actually overlaps our regions. --- --- From the binning index, we get a list of segments per requested --- region. Using the checkpoints, we prune them: if we have a --- checkpoint to the left of the beginning of the interesting region, we --- can move the start of each segment forward to the checkpoint. If --- that makes the segment empty, it can be droppped. --- --- The resulting segment lists are merged, then traversed. We seek to --- the beginning of the earliest segment and start decoding. Once the --- virtual file position leaves the segment or the alignment position --- moves past the end of the requested region, we move to the next. --- Moving is a seek if it spans a sufficiently large gap or points --- backwards, else we just keep going. - --- | A 'Segment' has a start and an end offset, and an "end coordinate" --- from the originating region. -data Segment = Segment {-# UNPACK #-} !Int64 {-# UNPACK #-} !Int64 {-# UNPACK #-} !Int deriving Show - -segmentLists :: BamIndex a -> Refseq -> R.Subsequence -> [[Segment]] -segmentLists bi@BamIndex{..} (Refseq ref) (R.Subsequence imap) - | Just bins <- refseq_bins V.!? fromIntegral ref, - Just cpts <- refseq_ckpoints V.!? fromIntegral ref - = [ rgnToSegments bi beg end bins cpts | (beg,end) <- M.toList imap ] -segmentLists _ _ _ = [] - --- from region to list of bins, then to list of segments -rgnToSegments :: BamIndex a -> Int -> Int -> Bins -> Ckpoints -> [Segment] -rgnToSegments bi@BamIndex{..} beg end bins cpts = - [ Segment boff' eoff end - | bin <- binList bi beg end - , (boff,eoff) <- maybe [] U.toList $ M.lookup bin bins - , let boff' = max boff cpt - , boff' < eoff ] - where - !cpt = maybe 0 snd $ M.lookupLE beg cpts - --- list of bins for given range of coordinates, from Heng's horrible code -binList :: BamIndex a -> Int -> Int -> [Int] -binList BamIndex{..} beg end = binlist' 0 (minshift + 3*depth) 0 - where - binlist' l s t = if l > depth then [] else [b..e] ++ go - where - b = t + beg `shiftR` s - e = t + (end-1) `shiftR` s - go = binlist' (l+1) (s-3) (t + 1 `shiftL` (3*l)) - - --- | Merges two lists of segments. Lists must be sorted, the merge sort --- merges overlapping segments into one. -infix 4 ~~ -(~~) :: [Segment] -> [Segment] -> [Segment] -Segment a b e : xs ~~ Segment u v f : ys - | b < u = Segment a b e : (xs ~~ Segment u v f : ys) -- no overlap - | a < u && b < v = Segment a v (max e f) : (xs ~~ ys) -- some overlap - | b < v = Segment u v (max e f) : (xs ~~ ys) -- contained - | v < a = Segment u v f : (xs ~~ Segment a b e : ys) -- no overlap - | u < a = Segment u b (max e f) : (xs ~~ ys) -- some overlap - | otherwise = Segment a b (max e f) : (xs ~~ ys) -- contained -[] ~~ ys = ys -xs ~~ [] = xs - - --- | Reads any index we can find for a file. If the file name has a --- .bai or .csi extension, we read it. Else we look for the index by --- adding such an extension and by replacing the extension with these --- two, and finally in the file itself. The first file that exists and --- can actually be parsed, is used. -readBamIndex :: FilePath -> IO (BamIndex ()) -readBamIndex fp | ".bai" `isSuffixOf` fp = enumFile defaultBufSize fp readBaiIndex >>= run - | ".csi" `isSuffixOf` fp = enumFile defaultBufSize fp readBaiIndex >>= run - | otherwise = tryIx (fp ++ ".bai") $ - tryIx (dropExtension fp ++ "bai") $ - tryIx (fp ++ ".csi") $ - tryIx (dropExtension fp ++ "csi") $ - enumFile defaultBufSize fp readBaiIndex >>= run - where - dropExtension p = reverse $ (if null b then "." ++ f else b) ++ d - where - (f,d) = break (=='/') $ reverse p - b = dropWhile (/='.') f - - tryIx f k = do e <- fileExist f - if e then do r <- enumFile defaultBufSize f readBaiIndex >>= tryRun - case r of Right ix -> return ix - Left (IterStringException _) -> k - else k - --- | Read an index in BAI or CSI format, recognized automatically. --- Note that TBI is supposed to be compressed using bgzip; it must be --- decompressed before being passed to 'readBaiIndex'. - -readBaiIndex :: MonadIO m => Iteratee Bytes m (BamIndex ()) -readBaiIndex = iGetString 4 >>= switch - where - switch "BAI\1" = do nref <- fromIntegral `liftM` endianRead4 LSB - getIndexArrays nref 14 5 (const return) getIntervals - - switch "CSI\1" = do minshift <- fromIntegral `liftM` endianRead4 LSB - depth <- fromIntegral `liftM` endianRead4 LSB - endianRead4 LSB >>= dropStreamBS . fromIntegral -- aux data - nref <- fromIntegral `liftM` endianRead4 LSB - getIndexArrays nref minshift depth (addOneCheckpoint minshift depth) return - - switch magic = throwErr . iterStrExc $ "index signature " ++ show magic ++ " not recognized" - - - -- Insert one checkpoint. If we already have an entry (can happen - -- if it comes from a different bin), we conservatively take the min - addOneCheckpoint minshift depth bin cp = do - loffset <- fromIntegral `liftM` endianRead8 LSB - let key = llim (fromIntegral bin) (3*depth) minshift - return $! M.insertWith min key loffset cp - - -- compute left limit of bin - llim bin dp sf | dp == 0 = 0 - | bin >= ix = (bin - ix) `shiftL` sf - | otherwise = llim bin (dp-3) (sf+3) - where ix = (1 `shiftL` dp - 1) `div` 7 - -type TabIndex = BamIndex TabMeta - -data TabMeta = TabMeta { format :: TabFormat - , col_seq :: Int -- Column for the sequence name - , col_beg :: Int -- Column for the start of a region - , col_end :: Int -- Column for the end of a region - , comment_char :: Char - , skip_lines :: Int - , names :: V.Vector Bytes } - deriving Show - -data TabFormat = Generic | SamFormat | VcfFormat | ZeroBased deriving Show - --- | Reads a Tabix index. Note that tabix indices are compressed, this --- is taken care of. -readTabix :: MonadIO m => Iteratee Bytes m TabIndex -readTabix = joinI $ decompressBgzf $ iGetString 4 >>= switch - where - switch "TBI\1" = do nref <- fromIntegral `liftM` endianRead4 LSB - format <- liftM toFormat (endianRead4 LSB) - col_seq <- liftM fromIntegral (endianRead4 LSB) - col_beg <- liftM fromIntegral (endianRead4 LSB) - col_end <- liftM fromIntegral (endianRead4 LSB) - comment_char <- liftM (chr . fromIntegral) (endianRead4 LSB) - skip_lines <- liftM fromIntegral (endianRead4 LSB) - names <- liftM (V.fromList . B.split 0) . iGetString . fromIntegral =<< endianRead4 LSB - - ix <- getIndexArrays nref 14 5 (const return) getIntervals - fin <- isFinished - if fin then return $! ix { extensions = TabMeta{..} } - else do unaln <- fromIntegral `liftM` endianRead8 LSB - return $! ix { unaln_off = unaln, extensions = TabMeta{..} } - - switch magic = throwErr . iterStrExc $ "index signature " ++ show magic ++ " not recognized" - - toFormat 1 = SamFormat - toFormat 2 = VcfFormat - toFormat x = if testBit x 16 then ZeroBased else Generic - --- Read the intervals. Each one becomes a checkpoint. -getIntervals :: Monad m => (IntMap Int64, Int64) -> Iteratee Bytes m (IntMap Int64, Int64) -getIntervals (cp,mx0) = do - nintv <- fromIntegral `liftM` endianRead4 LSB - reduceM 0 nintv (cp,mx0) $ \(!im,!mx) int -> do - oo <- fromIntegral `liftM` endianRead8 LSB - return (if oo == 0 then im else M.insert (int * 0x4000) oo im, max mx oo) - - -getIndexArrays :: MonadIO m => Int -> Int -> Int - -> (Word32 -> Ckpoints -> Iteratee Bytes m Ckpoints) - -> ((Ckpoints, Int64) -> Iteratee Bytes m (Ckpoints, Int64)) - -> Iteratee Bytes m (BamIndex ()) -getIndexArrays nref minshift depth addOneCheckpoint addManyCheckpoints - | nref < 1 = return $ BamIndex minshift depth 0 () V.empty V.empty - | otherwise = do - rbins <- liftIO $ W.new nref - rckpts <- liftIO $ W.new nref - mxR <- reduceM 0 nref 0 $ \mx0 r -> do - nbins <- endianRead4 LSB - (!bins,!cpts,!mx1) <- reduceM 0 nbins (M.empty,M.empty,mx0) $ \(!im,!cp,!mx) _ -> do - bin <- endianRead4 LSB -- the "distinct bin" - cp' <- addOneCheckpoint bin cp - segsarr <- getSegmentArray - let !mx' = if U.null segsarr then mx else max mx (snd (U.last segsarr)) - return (M.insert (fromIntegral bin) segsarr im, cp', mx') - (!cpts',!mx2) <- addManyCheckpoints (cpts,mx1) - liftIO $ W.write rbins r bins >> W.write rckpts r cpts' - return mx2 - liftM2 (BamIndex minshift depth mxR ()) (liftIO $ V.unsafeFreeze rbins) (liftIO $ V.unsafeFreeze rckpts) - --- | Reads the list of segments from an index file and makes sure --- it is sorted. -getSegmentArray :: MonadIO m => Iteratee Bytes m Segments -getSegmentArray = do - nsegs <- fromIntegral `liftM` endianRead4 LSB - segsarr <- liftIO $ N.new nsegs - loopM 0 nsegs $ \i -> do beg <- fromIntegral `liftM` endianRead8 LSB - end <- fromIntegral `liftM` endianRead8 LSB - liftIO $ N.write segsarr i (beg,end) - liftIO $ N.sort segsarr >> U.unsafeFreeze segsarr - -{-# INLINE reduceM #-} -reduceM :: (Monad m, Enum ix, Eq ix) => ix -> ix -> a -> (a -> ix -> m a) -> m a -reduceM beg end acc cons = if beg /= end then cons acc beg >>= \n -> reduceM (succ beg) end n cons else return acc - -{-# INLINE loopM #-} -loopM :: (Monad m, Enum ix, Eq ix) => ix -> ix -> (ix -> m ()) -> m () -loopM beg end k = if beg /= end then k beg >> loopM (succ beg) end k else return () - - --- | Seeks to a given sequence in a Bam file and enumerates only those --- records aligning to that reference. We use the first checkpoint --- available for the sequence. This requires an appropriate index, and --- the file must have been opened in such a way as to allow seeking. --- Enumerates over the @BamRaw@ records of the correct sequence only, --- doesn't enumerate at all if the sequence isn't found. - -eneeBamRefseq :: Monad m => BamIndex b -> Refseq -> Enumeratee [BamRaw] [BamRaw] m a -eneeBamRefseq BamIndex{..} (Refseq r) iter - | Just ckpts <- refseq_ckpoints V.!? fromIntegral r - , Just (voff, _) <- M.minView ckpts - , voff /= 0 = do seek $ fromIntegral voff - breakE ((Refseq r /=) . b_rname . unpackBam) iter - | otherwise = return iter - --- | Seeks to the part of a Bam file that contains unaligned reads and --- enumerates those. Sort of the dual to 'eneeBamRefseq'. We use the --- best guess at where the unaligned stuff starts. If no such guess is --- available, we decode everything. - -eneeBamUnaligned :: Monad m => BamIndex b -> Enumeratee [BamRaw] [BamRaw] m a -eneeBamUnaligned BamIndex{..} iter = do when (unaln_off /= 0) $ seek $ fromIntegral unaln_off - filterStream (not . isValidRefseq . b_rname . unpackBam) iter - --- | Enumerates one 'Segment'. Seeks to the start offset, unless --- reading over the skipped part looks cheaper. Enumerates until we --- either cross the end offset or the max position. -eneeBamSegment :: Monad m => Segment -> Enumeratee [BamRaw] [BamRaw] m r -eneeBamSegment (Segment beg end mpos) out = do - -- seek if it's a backwards seek or more than 512k forwards - peekStream >>= \x -> case x of - Just br | beg <= o && beg + 0x8000 > o -> return () - where o = fromIntegral $ virt_offset br - _ -> seek $ fromIntegral beg - - let in_segment br = virt_offset br <= fromIntegral end && b_pos (unpackBam br) <= mpos - takeWhileE in_segment out - -eneeBamSubseq :: Monad m => BamIndex b -> Refseq -> R.Subsequence -> Enumeratee [BamRaw] [BamRaw] m a -eneeBamSubseq bi ref subs = foldr ((>=>) . eneeBamSegment) return segs ><> filterStream olap - where - segs = foldr (~~) [] $ segmentLists bi ref subs - olap br = b_rname == ref && R.overlaps b_pos (b_pos + alignedLength b_cigar) subs - where BamRec{..} = unpackBam br - -eneeBamRegions :: Monad m => BamIndex b -> [R.Region] -> Enumeratee [BamRaw] [BamRaw] m a -eneeBamRegions bi = foldr ((>=>) . uncurry (eneeBamSubseq bi)) return . R.toList . R.fromList - diff --git a/src/Bio/Bam/Pileup.hs b/src/Bio/Bam/Pileup.hs deleted file mode 100644 --- a/src/Bio/Bam/Pileup.hs +++ /dev/null @@ -1,469 +0,0 @@ -{-# LANGUAGE Rank2Types, DeriveGeneric #-} - --- | Pileup, similar to Samtools --- --- Pileup turns a sorted sequence of reads into a sequence of \"piles\", --- one for each site where a genetic variant might be called. We will --- scan each read's CIGAR line and MD field in concert with the sequence --- and effective quality. Effective quality is the lowest available --- quality score of QUAL, MAPQ, and BQ. For aDNA calling, a base is --- represented as four probabilities, derived from a position dependent --- damage model. - -module Bio.Bam.Pileup where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Iteratee -import Bio.Prelude - -import qualified Data.ByteString as B -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Unboxed as U - --- | The primitive pieces for genotype calling: A position, a base --- represented as four likelihoods, an inserted sequence, and the --- length of a deleted sequence. The logic is that we look at a base --- followed by some indel, and all those indels are combined into a --- single insertion and a single deletion. -data PrimChunks = Seek {-# UNPACK #-} !Int PrimBase -- ^ skip to position (at start or after N operation) - | Indel [Nucleotides] [DamagedBase] PrimBase -- ^ observed deletion and insertion between two bases - | EndOfRead -- ^ nothing anymore - deriving Show - -data PrimBase = Base { _pb_wait :: {-# UNPACK #-} !Int -- ^ number of bases to wait due to a deletion - , _pb_base :: {-# UNPACK #-} !DamagedBase - , _pb_mapq :: {-# UNPACK #-} !Qual -- ^ map quality - , _pb_chunks :: PrimChunks } -- ^ more chunks - deriving Show - -type PosPrimChunks = (Refseq, Int, Bool, PrimChunks) - --- | Represents our knowledge about a certain base, which consists of --- the base itself (A,C,G,T, encoded as 0..3; no Ns), the quality score --- (anything that isn't A,C,G,T becomes A with quality 0), and a --- substitution matrix representing post-mortem but pre-sequencing --- substitutions. --- --- Unfortunately, none of this can be rolled into something more simple, --- because damage and sequencing error behave so differently. --- --- Damage information is polymorphic. We might run with a simple --- version (a matrix) for calling, but we need more (a matrix and a --- mutable matrix, I think) for estimation. - -data DamagedBase = DB { db_call :: {-# UNPACK #-} !Nucleotide -- ^ called base - , db_qual :: {-# UNPACK #-} !Qual -- ^ quality of called base - , db_dmg_tk :: {-# UNPACK #-} !DmgToken -- ^ damage information - , db_dmg_pos :: {-# UNPACK #-} !Int -- ^ damage information - , db_ref :: {-# UNPACK #-} !Nucleotides } -- ^ reference base from MD field - -newtype DmgToken = DmgToken { fromDmgToken :: Int } - -instance Show DamagedBase where - showsPrec _ (DB n q _ _ r) - | nucToNucs n == r = shows n . (:) '@' . shows q - | otherwise = shows n . (:) '/' . shows r . (:) '@' . shows q - - --- | Decomposes a BAM record into chunks suitable for piling up. We --- pick apart the CIGAR and MD fields, and combine them with sequence --- and quality as appropriate. Clipped bases are removed/skipped as --- needed. We also apply a substitution matrix to each base, which must --- be supplied along with the read. -{-# INLINE decompose #-} -decompose :: DmgToken -> BamRaw -> [PosPrimChunks] -decompose dtok br = - if isUnmapped b || isDuplicate b || not (isValidRefseq b_rname) - then [] else [(b_rname, b_pos, isReversed b, pchunks)] - where - b@BamRec{..} = unpackBam br - pchunks = firstBase b_pos 0 0 (fromMaybe [] $ getMd b) - - !max_cig = V.length b_cigar - !max_seq = V.length b_seq - !baq = extAsString "BQ" b - - -- This will compute the effective quality. As far as I can see - -- from the BAM spec V1.4, the qualities that matter are QUAL, MAPQ, - -- and BAQ. If QUAL is invalid, we replace it (arbitrarily) with - -- 23 (assuming a rather conservative error rate of ~0.5%), BAQ is - -- added to QUAL, and MAPQ is an upper limit for effective quality. - - get_seq :: Int -> (Nucleotides -> Nucleotides) -> DamagedBase - get_seq i f = case b_seq `V.unsafeIndex` i of -- nucleotide - n | n == nucsA -> DB nucA qe dtok dmg (f n) - | n == nucsC -> DB nucC qe dtok dmg (f n) - | n == nucsG -> DB nucG qe dtok dmg (f n) - | n == nucsT -> DB nucT qe dtok dmg (f n) - | otherwise -> DB nucA (Q 0) dtok dmg (f n) - where - !q = case b_qual `V.unsafeIndex` i of Q 0xff -> Q 30 ; x -> x -- quality; invalid (0xff) becomes 30 - !q' | i >= B.length baq = q -- no BAQ available - | otherwise = Q (unQ q + (B.index baq i - 64)) -- else correct for BAQ - !qe = min q' b_mapq -- use MAPQ as upper limit - !dmg = if i+i > max_seq then i-max_seq else i - - -- Look for first base following the read's start or a gap (CIGAR - -- code N). Indels are skipped, since these are either bugs in the - -- aligner or the aligner getting rid of essentially unalignable - -- bases. - firstBase :: Int -> Int -> Int -> [MdOp] -> PrimChunks - firstBase !pos !is !ic mds - | is >= max_seq || ic >= max_cig = EndOfRead - | otherwise = case b_cigar `V.unsafeIndex` ic of - Ins :* cl -> firstBase pos (cl+is) (ic+1) mds - SMa :* cl -> firstBase pos (cl+is) (ic+1) mds - Del :* cl -> firstBase (pos+cl) is (ic+1) (drop_del cl mds) - Nop :* cl -> firstBase (pos+cl) is (ic+1) mds - HMa :* _ -> firstBase pos is (ic+1) mds - Pad :* _ -> firstBase pos is (ic+1) mds - Mat :* 0 -> firstBase pos is (ic+1) mds - Mat :* _ -> Seek pos $ nextBase 0 pos is ic 0 mds - where - -- We have to treat (MdNum 0), because samtools actually - -- generates(!) it all over the place and if not handled as a - -- special case, it looks like an inconsistent MD field. - drop_del n (MdDel ns : mds') - | n < length ns = MdDel (drop n ns) : mds' - | n > length ns = drop_del (n - length ns) mds' - | otherwise = mds' - drop_del n (MdNum 0 : mds') = drop_del n mds' - drop_del _ mds' = mds' - - -- Generate likelihoods for the next base. When this gets called, - -- we are looking at an M CIGAR operation and all the subindices are - -- valid. - -- I don't think we can ever get (MdDel []), but then again, who - -- knows what crazy shit samtools decides to generate. There is - -- little harm in special-casing it. - nextBase :: Int -> Int -> Int -> Int -> Int -> [MdOp] -> PrimBase - nextBase !wt !pos !is !ic !io mds = case mds of - MdNum 0 : mds' -> nextBase wt pos is ic io mds' - MdDel [] : mds' -> nextBase wt pos is ic io mds' - MdNum 1 : mds' -> nextBase' (get_seq is id ) mds' - MdNum n : mds' -> nextBase' (get_seq is id ) (MdNum (n-1) : mds') - MdRep ref : mds' -> nextBase' (get_seq is $ const ref ) mds' - MdDel _ : _ -> nextBase' (get_seq is $ const nucsN) mds - [ ] -> nextBase' (get_seq is $ const nucsN) [ ] - where - nextBase' ref mds' = Base wt ref b_mapq $ nextIndel [] [] (pos+1) (is+1) ic (io+1) mds' - - -- Look for the next indel after a base. We collect all indels (I - -- and D codes) into one combined operation. If we hit N or the - -- read's end, we drop all of it (indels next to a gap indicate - -- trouble). Other stuff is skipped: we could check for stuff that - -- isn't valid in the middle of a read (H and S), but then what - -- would we do about it anyway? Just ignoring it is much easier and - -- arguably at least as correct. - nextIndel :: [[DamagedBase]] -> [Nucleotides] -> Int -> Int -> Int -> Int -> [MdOp] -> PrimChunks - nextIndel ins del !pos !is !ic !io mds - | is >= max_seq || ic >= max_cig = EndOfRead - | otherwise = case b_cigar `V.unsafeIndex` ic of - Ins :* cl -> nextIndel (isq cl) del pos (cl+is) (ic+1) 0 mds - SMa :* cl -> nextIndel ins del pos (cl+is) (ic+1) 0 mds - Del :* cl -> nextIndel ins (del++dsq) (pos+cl) is (ic+1) 0 mds' - where (dsq,mds') = split_del cl mds - Pad :* _ -> nextIndel ins del pos is (ic+1) 0 mds - HMa :* _ -> nextIndel ins del pos is (ic+1) 0 mds - Nop :* cl -> firstBase (pos+cl) is (ic+1) mds -- ends up generating a 'Seek' - Mat :* cl | io == cl -> nextIndel ins del pos is (ic+1) 0 mds - | otherwise -> indel del out $ nextBase (length del) pos is ic io mds -- ends up generating a 'Base' - where - indel d o k = foldr seq (Indel d o k) o - out = concat $ reverse ins - isq cl = [ get_seq i $ const gap | i <- [is..is+cl-1] ] : ins - - -- We have to treat (MdNum 0), because samtools actually - -- generates(!) it all over the place and if not handled as a - -- special case, it looks like an incinsistend MD field. - split_del n (MdDel ns : mds') - | n < length ns = (take n ns, MdDel (drop n ns) : mds') - | n > length ns = let (ns', mds'') = split_del (n - length ns) mds' in (ns++ns', mds'') - | otherwise = (ns, mds') - split_del n (MdNum 0 : mds') = split_del n mds' - split_del n mds' = (replicate n nucsN, mds') - --- | Statistics about a genotype call. Probably only useful for --- fitlering (so not very useful), but we keep them because it's easy to --- track them. - -data CallStats = CallStats - { read_depth :: {-# UNPACK #-} !Int -- number of contributing reads - , reads_mapq0 :: {-# UNPACK #-} !Int -- number of (non-)contributing reads with MAPQ==0 - , sum_mapq :: {-# UNPACK #-} !Int -- sum of map qualities of contributing reads - , sum_mapq_squared :: {-# UNPACK #-} !Int } -- sum of squared map qualities of contributing reads - deriving (Show, Eq, Generic) - -instance Monoid CallStats where - mempty = CallStats { read_depth = 0 - , reads_mapq0 = 0 - , sum_mapq = 0 - , sum_mapq_squared = 0 } - mappend = (<>) - -instance Semigroup CallStats where - x <> y = CallStats { read_depth = read_depth x + read_depth y - , reads_mapq0 = reads_mapq0 x + reads_mapq0 y - , sum_mapq = sum_mapq x + sum_mapq y - , sum_mapq_squared = sum_mapq_squared x + sum_mapq_squared y } - -newtype V_Nuc = V_Nuc (U.Vector Nucleotide) deriving (Eq, Ord, Show) -newtype V_Nucs = V_Nucs (U.Vector Nucleotides) deriving (Eq, Ord, Show) - -data IndelVariant = IndelVariant { deleted_bases :: !V_Nucs, inserted_bases :: !V_Nuc } - deriving (Eq, Ord, Show, Generic) - - --- | Map quality and a list of encountered bases, with damage --- information and reference base if known. -type BasePile = [DamagedBase] - --- | Map quality and a list of encountered indel variants. The deletion --- has the reference sequence, if known, an insertion has the inserted --- sequence with damage information. -type IndelPile = [( Qual, ([Nucleotides], [DamagedBase]) )] -- a list of indel variants - --- | Running pileup results in a series of piles. A 'Pile' has the --- basic statistics of a 'VarCall', but no likelihood values and a --- pristine list of variants instead of a proper call. We emit one pile --- with two 'BasePile's (one for each strand) and one 'IndelPile' (the --- one immediately following) at a time. - -data Pile' a b = Pile { p_refseq :: {-# UNPACK #-} !Refseq - , p_pos :: {-# UNPACK #-} !Int - , p_snp_stat :: {-# UNPACK #-} !CallStats - , p_snp_pile :: a - , p_indel_stat :: {-# UNPACK #-} !CallStats - , p_indel_pile :: b } - deriving Show - --- | Raw pile. Bases and indels are piled separately on forward and --- backward strands. -type Pile = Pile' (BasePile, BasePile) (IndelPile, IndelPile) - --- | The pileup enumeratee takes 'BamRaw's, decomposes them, interleaves --- the pieces appropriately, and generates 'Pile's. The output will --- contain at most one 'BasePile' and one 'IndelPile' for each position, --- piles are sorted by position. --- --- This top level driver receives 'BamRaw's. Unaligned reads and --- duplicates are skipped (but not those merely failing quality checks). --- Processing stops when the first read with invalid 'br_rname' is --- encountered or a t end of file. - -{-# INLINE pileup #-} -pileup :: Enumeratee [PosPrimChunks] [Pile] IO b -pileup = eneeCheckIfDonePass (icont . runPileM pileup' finish (Refseq 0) 0 ([],[]) (Empty,Empty)) - where - finish () _r _p ([],[]) (Empty,Empty) out inp = idone (liftI out) inp - finish () _ _ _ _ _ _ = error "logic error: leftovers after pileup" - - --- | The pileup logic keeps a current coordinate (just two integers) and --- two running queues: one of /active/ 'PrimBase's that contribute to --- current genotype calling and on of /waiting/ 'PrimBase's that will --- contribute at a later point. --- --- Oppan continuation passing style! Not only is the CPS version of the --- state monad (we have five distinct pieces of state) somewhat faster, --- we also need CPS to interact with the mechanisms of 'Iteratee'. It --- makes implementing 'yield', 'peek', and 'bump' straight forward. - -newtype PileM m a = PileM { runPileM :: forall r . (a -> PileF m r) -> PileF m r } - --- | The things we drag along in 'PileM'. Notes: --- * The /active/ queue is a simple stack. We add at the front when we --- encounter reads, which reverses them. When traversing it, we traverse --- reads backwards, but since we accumulate the 'BasePile', it gets reversed --- back. The new /active/ queue, however, is no longer reversed (as it should --- be). So after the traversal, we reverse it again. (Yes, it is harder to --- understand than using a proper deque type, but it is cheaper. --- There may not be much point in the reversing, though.) - -type PileF m r = Refseq -> Int -> -- current position - ( [PrimBase], [PrimBase] ) -> -- active queues - ( Heap, Heap ) -> -- waiting queues - (Stream [Pile] -> Iteratee [Pile] m r) -> -- output function - Stream [PosPrimChunks] -> -- pending input - Iteratee [PosPrimChunks] m (Iteratee [Pile] m r) - -instance Functor (PileM m) where - {-# INLINE fmap #-} - fmap f (PileM m) = PileM $ \k -> m (k . f) - -instance Applicative (PileM m) where - {-# INLINE pure #-} - pure a = PileM $ \k -> k a - {-# INLINE (<*>) #-} - u <*> v = PileM $ \k -> runPileM u (\a -> runPileM v (k . a)) - -instance Monad (PileM m) where - {-# INLINE return #-} - return a = PileM $ \k -> k a - {-# INLINE (>>=) #-} - m >>= k = PileM $ \k' -> runPileM m (\a -> runPileM (k a) k') - -{-# INLINE get_refseq #-} -get_refseq :: PileM m Refseq -get_refseq = PileM $ \k r -> k r r - -{-# INLINE get_pos #-} -get_pos :: PileM m Int -get_pos = PileM $ \k r p -> k p r p - -{-# INLINE upd_pos #-} -upd_pos :: (Int -> Int) -> PileM m () -upd_pos f = PileM $ \k r p -> k () r $! f p - --- | Sends one piece of output downstream. You are not expected to --- understand how this works, but inlining 'eneeCheckIfDone' plugged an --- annoying memory leak. -{-# INLINE yieldPile #-} -yieldPile :: CallStats -> BasePile -> BasePile -> CallStats -> IndelPile -> IndelPile -> PileM m () -yieldPile x1 x2a x2b x3 x4a x4b = PileM $ \ !kont !r !p !a !w !out !inp -> Iteratee $ \od oc -> - let recurse = kont () r p a w - onDone y s = od (idone y s) inp - onCont k Nothing = runIter (recurse k inp) od oc - onCont k (Just e) = runIter (throwRecoverableErr e (recurse k . (<>) inp)) od oc - pile = Pile r p x1 (x2a,x2b) x3 (x4a,x4b) - in runIter (out (Chunk [pile])) onDone onCont - --- | The actual pileup algorithm. If /active/ contains something, --- continue here. Else find the coordinate to continue from, which is --- the minimum of the next /waiting/ coordinate and the next coordinate --- in input; if found, continue there, else we're all done. -pileup' :: PileM m () -pileup' = PileM $ \ !k !refseq !pos !active !waiting !out !inp -> - - let recurse = runPileM pileup' k refseq pos active waiting out - cont2 rs po = runPileM pileup'' k rs po active waiting out inp - leave = k () refseq pos active waiting out inp - - in case (active, getMinKeysH waiting, inp) of - ( (_:_,_), _, _ ) -> cont2 refseq pos - ( (_,_:_), _, _ ) -> cont2 refseq pos - ( _, Just nw, EOF _ ) -> cont2 refseq nw - ( _, Nothing, EOF _ ) -> leave - ( _, _, Chunk [ ] ) -> liftI recurse - ( _, Nothing, Chunk ((r,p,_,_):_) ) -> cont2 r p - ( _, Just nw, Chunk ((r,p,_,_):_) ) - | (refseq,nw) <= (r,p) -> cont2 refseq nw - | otherwise -> cont2 r p - where - getMinKeysH :: (Heap, Heap) -> Maybe Int - getMinKeysH (a,b) = case (getMinKeyH a, getMinKeyH b) of - ( Nothing, Nothing ) -> Nothing - ( Just x, Nothing ) -> Just x - ( Nothing, Just y ) -> Just y - ( Just x, Just y ) -> Just (min x y) - - -pileup'' :: PileM m () -pileup'' = do - -- Input is still 'BamRaw', since these can be relied on to be - -- sorted. First see if there is any input at the current location, - -- if so, decompose it and add it to the appropriate queue. - p'feed_input - p'check_waiting - ((fin_bsL, fin_bpL), (fin_bsR, fin_bpR), (fin_isL, fin_ipL), (fin_isR, fin_ipR)) <- p'scan_active - - -- Output, but don't bother emitting empty piles. Note that a plain - -- basecall still yields an entry in the 'IndelPile'. This is necessary, - -- because actual indel calling will want to know how many reads /did not/ - -- show the variant. However, if no reads show any variant, and here is the - -- first place where we notice that, the pile is useless. - let uninteresting (_,(d,i)) = null d && null i - unless (null fin_bpL && null fin_bpR && all uninteresting fin_ipL && all uninteresting fin_ipR) $ - yieldPile (fin_bsL <> fin_bsR) fin_bpL fin_bpR - (fin_isL <> fin_isR) fin_ipL fin_ipR - - -- Bump coordinate and loop. (Note that the bump to the next - -- reference /sequence/ is done implicitly, because we will run out of - -- reads and restart in 'pileup''.) - upd_pos succ - pileup' - --- | Feeds input as long as it starts at the current position -p'feed_input :: PileM m () -p'feed_input = PileM $ \kont rs po ac@(af,ar) wt@(wf,wr) out inp -> case inp of - Chunk [ ] -> liftI $ runPileM p'feed_input kont rs po ac wt out - Chunk ((rs', po', str, prim):bs) - | rs == rs' && po == po' -> - case prim of - Seek !p !pb -> let wf' = Node p pb Empty Empty `unionH` wf - wr' = Node p pb Empty Empty `unionH` wr - in runPileM p'feed_input kont rs po ac (if str then (wf,wr') else (wf',wr)) out (Chunk bs) - - Indel _ _ !pb -> runPileM p'feed_input kont rs po (if str then (af,pb:ar) else (pb:af,ar)) wt out (Chunk bs) - - EndOfRead -> runPileM p'feed_input kont rs po ac wt out (Chunk bs) - - _ -> kont () rs po ac wt out inp - --- | Checks /waiting/ queue. If there is anything waiting for the --- current position, moves it to /active/ queue. -p'check_waiting :: PileM m () -p'check_waiting = PileM $ \kont rs po (af0,ar0) (wf0,wr0) -> - let go1 af wf = case viewMinH wf of - Just (!mk, !pb, !wf') | mk == po -> go1 (pb:af) wf' - _ -> go2 af wf ar0 wr0 - - go2 af wf ar wr = case viewMinH wr of - Just (!mk, !pb, !wr') | mk == po -> go2 af wf (pb:ar) wr' - _ -> kont () rs po (af,ar) (wf,wr) - - in go1 af0 wf0 - - --- | Separately scans the two /active/ queues and makes one 'BasePile' --- from each. Also sees what's next in the 'PrimChunks': 'Indel's --- contribute to two separate 'IndelPile's, 'Seek's are pushed back to --- the /waiting/ queue, 'EndOfRead's are removed, and everything else is --- added to two fresh /active/ queues. -p'scan_active :: PileM m (( CallStats, BasePile ), ( CallStats, BasePile ), - ( CallStats, IndelPile ), ( CallStats, IndelPile )) -p'scan_active = do - (bpf,ipf) <- PileM $ \kont rs pos (af,ar) (wf,wr) -> go (\r af' wf' -> kont r rs pos (af',ar) (wf',wr)) [] wf mempty mempty af - (bpr,ipr) <- PileM $ \kont rs pos (af,ar) (wf,wr) -> go (\r ar' wr' -> kont r rs pos (af,ar') (wf,wr')) [] wr mempty mempty ar - return (bpf,bpr,ipf,ipr) - where - go k !ac !wt !bpile !ipile [ ] = k (bpile, ipile) (reverse ac) wt - go k !ac !wt !bpile !ipile (Base nwt qs mq pchunks : bs) = - case pchunks of - _ | nwt > 0 -> b' `seq` go k (b':ac) wt bpile ipile bs - Seek p' pb' -> go k ac (ins p' pb' wt) (z bpile) ipile bs - Indel nd ni pb' -> go k (pb':ac) wt (z bpile) (y ipile) bs where y = put (,) mq (nd,ni) - EndOfRead -> go k ac wt (z bpile) ipile bs - where - b' = Base (nwt-1) qs mq pchunks - z = put (\q x -> x { db_qual = min q (db_qual x) }) mq qs - - ins q v w = Node q v Empty Empty `unionH` w - - put f (Q !q) !x (!st,!vs) = ( st { read_depth = read_depth st + 1 - , reads_mapq0 = reads_mapq0 st + (if q == 0 then 1 else 0) - , sum_mapq = sum_mapq st + fromIntegral q - , sum_mapq_squared = sum_mapq_squared st + fromIntegral q * fromIntegral q } - , f (Q q) x : vs ) - - --- | We need a simple priority queue. Here's a skew heap (specialized --- to strict 'Int' priorities and 'PrimBase' values). -data Heap = Empty | Node {-# UNPACK #-} !Int PrimBase Heap Heap - -unionH :: Heap -> Heap -> Heap -Empty `unionH` t2 = t2 -t1 `unionH` Empty = t1 -t1@(Node k1 x1 l1 r1) `unionH` t2@(Node k2 x2 l2 r2) - | k1 <= k2 = Node k1 x1 (t2 `unionH` r1) l1 - | otherwise = Node k2 x2 (t1 `unionH` r2) l2 - -getMinKeyH :: Heap -> Maybe Int -getMinKeyH Empty = Nothing -getMinKeyH (Node x _ _ _) = Just x - -viewMinH :: Heap -> Maybe (Int, PrimBase, Heap) -viewMinH Empty = Nothing -viewMinH (Node k v l r) = Just (k, v, l `unionH` r) - diff --git a/src/Bio/Bam/Reader.hs b/src/Bio/Bam/Reader.hs deleted file mode 100644 --- a/src/Bio/Bam/Reader.hs +++ /dev/null @@ -1,295 +0,0 @@ --- | Parsers for BAM and SAM. --- --- TONOTDO: --- --- * Reader for gzipped\/bzipped\/bgzf'ed SAM. Storing SAM is a bad idea, --- so why would anyone ever want to compress, much less index it? --- * Proper support for the "=" symbol. It's completely alien to the --- ususal representation of sequences. - -module Bio.Bam.Reader ( - Block(..), - decompressBgzfBlocks, - decompressBgzf, - compressBgzf, - - decodeBam, - getBamRaw, - decodeAnyBam, - decodeAnyBamFile, - - BamrawEnumeratee, - BamEnumeratee, - isBamOrSam, - - isBam, - isPlainBam, - isGzipBam, - isBgzfBam, - - decodeSam, - decodeSam', - - decodeAnyBamOrSam, - decodeAnyBamOrSamFile, - - concatInputs, - concatDefaultInputs, - mergeInputs, - mergeDefaultInputs, - combineCoordinates, - combineNames, - ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Iteratee -import Bio.Iteratee.Bgzf -import Bio.Iteratee.ZLib hiding ( CompressionLevel ) -import Bio.Prelude - -import Data.Attoparsec.ByteString ( anyWord8 ) -import Data.Sequence ( (|>) ) - -import qualified Data.Attoparsec.ByteString.Char8 as P -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as S -import qualified Data.Foldable as F -import qualified Data.HashMap.Strict as M -import qualified Data.Sequence as Z -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as U - -type BamrawEnumeratee m b = Enumeratee' BamMeta Bytes [BamRaw] m b -type BamEnumeratee m b = Enumeratee' BamMeta Bytes [BamRec] m b - -isBamOrSam :: MonadIO m => Iteratee Bytes m (BamEnumeratee m a) -isBamOrSam = maybe decodeSam wrap `liftM` isBam - where - wrap enee it' = enee (mapStream unpackBam . it') >>= lift . run - - --- | Checks if a file contains BAM in any of the common forms, --- then decompresses it appropriately. If the stream doesn't contain --- BAM at all, it is instead decoded as SAM. Since SAM is next to --- impossible to recognize reliably, we don't even try. Any old junk is --- decoded as SAM and will fail later. -decodeAnyBamOrSam :: MonadIO m => BamEnumeratee m a -decodeAnyBamOrSam it = isBamOrSam >>= \k -> k it - -decodeAnyBamOrSamFile :: MonadBracketIO m - => FilePath -> (BamMeta -> Iteratee [BamRec] m a) -> m (Iteratee [BamRec] m a) -decodeAnyBamOrSamFile fn k = enumFileRandom defaultBufSize fn (decodeAnyBamOrSam k) >>= run - --- | Iteratee-style parser for SAM files, designed to be compatible with --- the BAM parsers. Parses plain uncompressed SAM, nothing else. Since --- it is supposed to work the same way as the BAM parser, it requires --- the presense of the SQ header lines. These are stripped from the --- header text and turned into the symbol table. -decodeSam :: Monad m => (BamMeta -> Iteratee [BamRec] m a) -> Iteratee Bytes m (Iteratee [BamRec] m a) -decodeSam inner = joinI $ enumLinesBS $ do - let pHeaderLine acc str = case P.parseOnly parseBamMetaLine str of Right f -> return $ f : acc - Left e -> fail $ e ++ ", " ++ show str - meta <- liftM (foldr ($) mempty . reverse) (joinI $ breakE (not . S.isPrefixOf "@") $ foldStreamM pHeaderLine []) - decodeSamLoop (meta_refs meta) (inner meta) - -decodeSamLoop :: Monad m => Refs -> Enumeratee [Bytes] [BamRec] m a -decodeSamLoop refs = convStream (liftI parse_record) - where - !refs' = M.fromList $ zip [ nm | BamSQ { sq_name = nm } <- F.toList refs ] [toEnum 0..] - ref x = M.lookupDefault invalidRefseq x refs' - - parse_record (EOF x) = icont parse_record x - parse_record (Chunk []) = liftI parse_record - parse_record (Chunk (l:ls)) | "@" `S.isPrefixOf` l = parse_record (Chunk ls) - parse_record (Chunk (l:ls)) = case P.parseOnly (parseSamRec ref) l of - Right r -> idone [r] (Chunk ls) - Left err -> icont parse_record (Just $ iterStrExc $ err ++ ", " ++ show l) - --- | Parser for SAM that doesn't look for a header. Has the advantage --- that it doesn't stall on a pipe that never delivers data. Has the --- disadvantage that it never reads the header and therefore needs a --- list of allowed RNAMEs. -decodeSam' :: Monad m => Refs -> Enumeratee Bytes [BamRec] m a -decodeSam' refs inner = joinI $ enumLinesBS $ decodeSamLoop refs inner - -parseSamRec :: (Bytes -> Refseq) -> P.Parser BamRec -parseSamRec ref = mkBamRec - <$> word <*> num <*> (ref <$> word) <*> (subtract 1 <$> num) - <*> (Q <$> num') <*> (VS.fromList <$> cigar) <*> rnext <*> (subtract 1 <$> num) - <*> snum <*> sequ <*> quals <*> exts <*> pure 0 - where - sep = P.endOfInput <|> () <$ P.char '\t' - word = P.takeTill ('\t' ==) <* sep - num = P.decimal <* sep - num' = P.decimal <* sep - snum = P.signed P.decimal <* sep - - rnext = id <$ P.char '=' <* sep <|> const . ref <$> word - sequ = {-# SCC "parseSamRec/sequ" #-} - (V.empty <$ P.char '*' <|> - V.fromList . map toNucleotides . S.unpack <$> P.takeWhile is_nuc) <* sep - - quals = {-# SCC "parseSamRec/quals" #-} defaultQs <$ P.char '*' <* sep <|> bsToVec <$> word - where - defaultQs sq = VS.replicate (V.length sq) (Q 0xff) - bsToVec qs _ = VS.fromList . map (Q . subtract 33) $ B.unpack qs - - cigar = [] <$ P.char '*' <* sep <|> - P.manyTill (flip (:*) <$> P.decimal <*> cigop) sep - - cigop = P.choice $ zipWith (\c r -> r <$ P.char c) "MIDNSHP" [Mat,Ins,Del,Nop,SMa,HMa,Pad] - exts = ext `P.sepBy` sep - ext = (\a b v -> (fromString [a,b],v)) <$> P.anyChar <*> P.anyChar <*> (P.char ':' *> value) - - value = P.char 'A' *> P.char ':' *> (Char <$> anyWord8) <|> - P.char 'i' *> P.char ':' *> (Int <$> P.signed P.decimal) <|> - P.char 'Z' *> P.char ':' *> (Text <$> P.takeTill ('\t' ==)) <|> - P.char 'H' *> P.char ':' *> (Bin <$> hexarray) <|> - P.char 'f' *> P.char ':' *> (Float . realToFrac <$> P.double) <|> - P.char 'B' *> P.char ':' *> ( - P.satisfy (P.inClass "cCsSiI") *> (intArr <$> many (P.char ',' *> P.signed P.decimal)) <|> - P.char 'f' *> (floatArr <$> many (P.char ',' *> P.double))) - - intArr is = IntArr $ U.fromList is - floatArr fs = FloatArr $ U.fromList $ map realToFrac fs - hexarray = B.pack . repack . S.unpack <$> P.takeWhile (P.inClass "0-9A-Fa-f") - repack (a:b:cs) = fromIntegral (digitToInt a * 16 + digitToInt b) : repack cs ; repack _ = [] - is_nuc = P.inClass "acgtswkmrybdhvnACGTSWKMRYBDHVN" - - mkBamRec nm fl rn po mq cg rn' mp is sq qs' = - BamRec nm fl rn po mq cg (rn' rn) mp is sq (qs' sq) - --- | Tests if a data stream is a Bam file. --- Recognizes plain Bam, gzipped Bam and bgzf'd Bam. If a file is --- recognized as Bam, a decoder (suitable Enumeratee) for it is --- returned. This uses 'iLookAhead' internally, so it shouldn't consume --- anything from the stream. -isBam, isEmptyBam, isPlainBam, isBgzfBam, isGzipBam :: MonadIO m - => Iteratee Bytes m (Maybe (BamrawEnumeratee m a)) -isBam = firstOf [ isEmptyBam, isPlainBam, isBgzfBam, isGzipBam ] - where - firstOf [] = return Nothing - firstOf (k:ks) = iLookAhead k >>= maybe (firstOf ks) (return . Just) - -isEmptyBam = (\e -> if e then Just (\k -> return $ k mempty) else Nothing) `liftM` isFinished - -isPlainBam = (\n -> if n == "BAM\SOH" then Just (joinI . decompressPlain . decodeBam) else Nothing) - `liftM` iGetString 4 - --- Interesting... iLookAhead interacts badly with the parallel --- decompression of BGZF. (The chosen interface doesn't allow the EOF --- signal to be passed on.) One workaround would be to run sequential --- BGZF decompression to check if the content is BAM, but since BGZF is --- actually GZip in disguise, the easier workaround if to use the --- ordinary GZip decompressor. --- (A clean workaround would be an @Alternative@ instance for --- @Iteratee@.) -isBgzfBam = do b <- isBgzf - k <- if b then joinI $ enumInflate GZip defaultDecompressParams isPlainBam else return Nothing - return $ const (joinI . decompressBgzfBlocks . decodeBam) `fmap` k - -isGzipBam = do b <- isGzip - k <- if b then joinI $ enumInflate GZip defaultDecompressParams isPlainBam else return Nothing - return $ ((joinI . enumInflate GZip defaultDecompressParams) .) `fmap` k - --- | Checks if a file contains BAM in any of the common forms, then --- decompresses it appropriately. We support plain BAM, Bgzf'd BAM, --- and Gzip'ed BAM. --- --- The recommendation for these functions is to use @decodeAnyBam@ (or --- @decodeAnyBamFile@) for any code that can handle @BamRaw@ input, but --- @decodeAnyBamOrSam@ (or @decodeAnyBamOrSamFile@) for code that needs --- @BamRec@. That way, SAM is supported automatically, and seeking will --- be supported if possible. -decodeAnyBam :: MonadIO m => BamrawEnumeratee m a -decodeAnyBam it = do mk <- isBam ; case mk of Just k -> k it - Nothing -> fail "this isn't BAM." - -decodeAnyBamFile :: MonadBracketIO m => FilePath -> (BamMeta -> Iteratee [BamRaw] m a) -> m (Iteratee [BamRaw] m a) -decodeAnyBamFile fn k = enumFileRandom defaultBufSize fn (decodeAnyBam k) >>= run - -concatDefaultInputs :: MonadBracketIO m => Enumerator' BamMeta [BamRaw] m a -concatDefaultInputs it0 = liftIO getArgs >>= \fs -> concatInputs fs it0 - -concatInputs :: MonadBracketIO m => [FilePath] -> Enumerator' BamMeta [BamRaw] m a -concatInputs [ ] = enumFd defaultBufSize stdInput . decodeAnyBam >=> run -concatInputs (fp0:fps) = enum1 fp0 >=> go fps - where - go = foldr (\fp1 -> (>=>) (enum1 fp1 . const)) return - enum1 "-" = enumFd defaultBufSize stdInput . decodeAnyBam >=> run - enum1 fp = enumFile defaultBufSize fp . decodeAnyBam >=> run - -mergeDefaultInputs :: MonadBracketIO m - => (BamMeta -> Enumeratee [BamRaw] [BamRaw] (Iteratee [BamRaw] m) a) - -> Enumerator' BamMeta [BamRaw] m a -mergeDefaultInputs (?) it0 = liftIO getArgs >>= \fs -> mergeInputs (?) fs it0 - -mergeInputs :: MonadBracketIO m - => (BamMeta -> Enumeratee [BamRaw] [BamRaw] (Iteratee [BamRaw] m) a) - -> [FilePath] -> Enumerator' BamMeta [BamRaw] m a -mergeInputs _ [ ] = \k -> enumFd defaultBufSize stdInput (decodeAnyBam k) >>= run -mergeInputs (?) (fp0:fps0) = go fp0 fps0 - where - enum1 "-" k1 = enumFd defaultBufSize stdInput (decodeAnyBam k1) >>= run - enum1 fp k1 = enumFile defaultBufSize fp (decodeAnyBam k1) >>= run - - go fp [ ] = enum1 fp - go fp (fp1:fps) = mergeEnums' (go fp1 fps) (enum1 fp) (?) - -{-# INLINE combineCoordinates #-} -combineCoordinates :: Monad m => BamMeta -> Enumeratee [BamRaw] [BamRaw] (Iteratee [BamRaw] m) a -combineCoordinates _ = mergeSortStreams (?) - where u ? v = if (b_rname &&& b_pos) (unpackBam u) < (b_rname &&& b_pos) (unpackBam v) then Less else NotLess - -{-# INLINE combineNames #-} -combineNames :: Monad m => BamMeta -> Enumeratee [BamRaw] [BamRaw] (Iteratee [BamRaw] m) a -combineNames _ = mergeSortStreams (?) - where u ? v = case b_qname (unpackBam u) `compareNames` b_qname (unpackBam v) of LT -> Less ; _ -> NotLess - --- | Decode a BAM stream into raw entries. Note that the entries can be --- unpacked using @decodeBamEntry@. Also note that this is an --- Enumeratee in spirit, only the @BamMeta@ and @Refs@ need to get --- passed separately. -{-# INLINE decodeBam #-} -decodeBam :: Monad m => (BamMeta -> Iteratee [BamRaw] m a) -> Iteratee Block m (Iteratee [BamRaw] m a) -decodeBam inner = do meta <- liftBlock get_bam_header - refs <- liftBlock get_ref_array - convStream getBamRaw $ inner $! mmerge meta refs - where - get_bam_header = do magic <- iGetString 4 - when (magic /= "BAM\SOH") $ do - s <- iGetString 10 - fail $ "BAM signature not found: " ++ show magic ++ " " ++ show s - hdr_len <- endianRead4 LSB - joinI $ takeStreamBS (fromIntegral hdr_len) $ parserToIteratee parseBamMeta - - get_ref_array = do nref <- endianRead4 LSB - foldM (\acc _ -> do - nm <- endianRead4 LSB >>= iGetString . fromIntegral - ln <- endianRead4 LSB - return $! acc |> BamSQ (S.init nm) (fromIntegral ln) [] - ) Z.empty [1..nref] - - -- Need to merge information from header into actual reference list. - -- The latter is the authoritative source for the *order* of the - -- sequences, so leftovers from the header are discarded. Merging - -- is by name. So we merge information from the header into the - -- list, then replace the header information. - mmerge meta refs = - let tbl = M.fromList [ (sq_name sq, sq) | sq <- F.toList (meta_refs meta) ] - in meta { meta_refs = fmap (\s -> maybe s (mmerge' s) (M.lookup (sq_name s) tbl)) refs } - - mmerge' l r | sq_length l == sq_length r = l { sq_other_shit = sq_other_shit l ++ sq_other_shit r } - | otherwise = l -- contradiction in header, but we'll just ignore it - - -{-# INLINE getBamRaw #-} -getBamRaw :: Monad m => Iteratee Block m [BamRaw] -getBamRaw = do off <- getOffset - raw <- liftBlock $ do - bsize <- endianRead4 LSB - when (bsize < 32) $ fail "short BAM record" - iGetString (fromIntegral bsize) - return [bamRaw off raw] diff --git a/src/Bio/Bam/Rec.hs b/src/Bio/Bam/Rec.hs deleted file mode 100644 --- a/src/Bio/Bam/Rec.hs +++ /dev/null @@ -1,391 +0,0 @@ -{-# LANGUAGE TypeFamilies #-} - --- | Parsers and Printers for BAM and SAM. We employ an @Iteratee@ --- interface, and we strive to support everything possible in BAM. So --- far, the implementation of the nucleotides is somewhat lacking: we --- do not have support for ambiguity codes, and the "=" symbol is not --- understood. - -module Bio.Bam.Rec ( - BamRaw, - bamRaw, - virt_offset, - raw_data, - - BamRec(..), - unpackBam, - nullBamRec, - getMd, - - Cigar(..), - CigOp(..), - alignedLength, - - Nucleotides(..), Vector_Nucs_half, - Extensions, Ext(..), - extAsInt, extAsString, setQualFlag, - deleteE, insertE, updateE, adjustE, - - isPaired, - isProperlyPaired, - isUnmapped, - isMateUnmapped, - isReversed, - isMateReversed, - isFirstMate, - isSecondMate, - isAuxillary, - isSecondary, - isFailsQC, - isDuplicate, - isSupplementary, - isTrimmed, - isMerged, - isAlternative, - isExactIndex, - type_mask, - - progressBam, - Word32 -) where - -import Bio.Bam.Header -import Bio.Iteratee -import Bio.Prelude -import Bio.Util.Storable - -import Control.Monad.Primitive ( unsafePrimToPrim, unsafeInlineIO ) -import Foreign.C.Types ( CInt(..), CSize(..) ) -import Foreign.Marshal.Alloc ( alloca ) - -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as S -import qualified Data.ByteString.Internal as B -import qualified Data.ByteString.Unsafe as B -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Generic.Mutable as VM -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as U - - --- | Cigar line in BAM coding --- Bam encodes an operation and a length into a single integer, we keep --- those integers in an array. -data Cigar = !CigOp :* !Int deriving (Eq, Ord) -infix 9 :* - -data CigOp = Mat | Ins | Del | Nop | SMa | HMa | Pad - deriving ( Eq, Ord, Enum, Show, Bounded, Ix ) - -instance Show Cigar where - showsPrec _ (op :* num) = shows num . (:) (S.index "MIDNSHP" (fromEnum op)) - -instance Storable Cigar where - sizeOf _ = 4 - alignment _ = 1 - - peek p = do w <- fromIntegral <$> peekUnalnWord32LE p - return $ toEnum (w .&. 0xf) :* shiftR w 4 - - poke p (op :* num) = pokeUnalnWord32LE p . fromIntegral $ fromEnum op .|. shiftL num 4 - - --- | Extracts the aligned length from a cigar line. --- This gives the length of an alignment as measured on the reference, --- which is different from the length on the query or the length of the --- alignment. -{-# INLINE alignedLength #-} -alignedLength :: V.Vector v Cigar => v Cigar -> Int -alignedLength = V.foldl' (\a -> (a +) . l) 0 - where l (op :* n) = if op == Mat || op == Del || op == Nop then n else 0 - - --- | internal representation of a BAM record -data BamRec = BamRec { - b_qname :: Seqid, - b_flag :: Int, - b_rname :: Refseq, - b_pos :: Int, - b_mapq :: Qual, - b_cigar :: VS.Vector Cigar, - b_mrnm :: Refseq, - b_mpos :: Int, - b_isize :: Int, - b_seq :: Vector_Nucs_half Nucleotides, - b_qual :: VS.Vector Qual, - b_exts :: Extensions, - b_virtual_offset :: FileOffset -- ^ virtual offset for indexing purposes - } deriving Show - -nullBamRec :: BamRec -nullBamRec = BamRec { - b_qname = S.empty, - b_flag = flagUnmapped, - b_rname = invalidRefseq, - b_pos = invalidPos, - b_mapq = Q 0, - b_cigar = VS.empty, - b_mrnm = invalidRefseq, - b_mpos = invalidPos, - b_isize = 0, - b_seq = V.empty, - b_qual = VS.empty, - b_exts = [], - b_virtual_offset = 0 - } - -getMd :: BamRec -> Maybe [MdOp] -getMd r = case lookup "MD" $ b_exts r of - Just (Text mdfield) -> readMd mdfield - Just (Char mdfield) -> readMd $ B.singleton mdfield - _ -> Nothing - --- | A vector that packs two 'Nucleotides' into one byte, just like Bam does. -data Vector_Nucs_half a = Vector_Nucs_half !Int !Int !(ForeignPtr Word8) - --- | A mutable vector that packs two 'Nucleotides' into one byte, just like Bam does. -data MVector_Nucs_half s a = MVector_Nucs_half !Int !Int !(ForeignPtr Word8) - -type instance V.Mutable Vector_Nucs_half = MVector_Nucs_half - -instance V.Vector Vector_Nucs_half Nucleotides where - {-# INLINE basicUnsafeFreeze #-} - basicUnsafeFreeze (MVector_Nucs_half o l fp) = return $ Vector_Nucs_half o l fp - {-# INLINE basicUnsafeThaw #-} - basicUnsafeThaw (Vector_Nucs_half o l fp) = return $ MVector_Nucs_half o l fp - - {-# INLINE basicLength #-} - basicLength (Vector_Nucs_half _ l _) = l - {-# INLINE basicUnsafeSlice #-} - basicUnsafeSlice s l (Vector_Nucs_half o _ fp) = Vector_Nucs_half (o + s) l fp - - {-# INLINE basicUnsafeIndexM #-} - basicUnsafeIndexM (Vector_Nucs_half o _ fp) i - | even (o+i) = return . Ns $! (b `shiftR` 4) .&. 0xF - | otherwise = return . Ns $! b .&. 0xF - where !b = unsafeInlineIO $ withForeignPtr fp $ \p -> peekByteOff p ((o+i) `shiftR` 1) - -instance VM.MVector MVector_Nucs_half Nucleotides where - {-# INLINE basicLength #-} - basicLength (MVector_Nucs_half _ l _) = l - {-# INLINE basicUnsafeSlice #-} - basicUnsafeSlice s l (MVector_Nucs_half o _ fp) = MVector_Nucs_half (o + s) l fp - - {-# INLINE basicOverlaps #-} - basicOverlaps (MVector_Nucs_half _ _ fp1) (MVector_Nucs_half _ _ fp2) = fp1 == fp2 - {-# INLINE basicUnsafeNew #-} - basicUnsafeNew l = unsafePrimToPrim $ MVector_Nucs_half 0 l <$> mallocForeignPtrBytes ((l+1) `shiftR` 1) - - {-# INLINE basicInitialize #-} - basicInitialize v@(MVector_Nucs_half o l fp) - - | even o = do unsafePrimToPrim $ withForeignPtr fp $ \p -> - memset (plusPtr p (o `shiftR` 1)) 0 (fromIntegral $ l `shiftR` 1) - when (odd l) $ VM.basicUnsafeWrite v (l-1) (Ns 0) - - | otherwise = do when (odd o) $ VM.basicUnsafeWrite v 0 (Ns 0) - unsafePrimToPrim $ withForeignPtr fp $ \p -> - memset (plusPtr p ((o+1) `shiftR` 1)) 0 (fromIntegral $ (l-1) `shiftR` 1) - when (even l) $ VM.basicUnsafeWrite v (l-1) (Ns 0) - - - {-# INLINE basicUnsafeRead #-} - basicUnsafeRead (MVector_Nucs_half o _ fp) i - | even (o+i) = liftM (Ns . (.&.) 0xF . (`shiftR` 4)) b - | otherwise = liftM (Ns . (.&.) 0xF ) b - where b = unsafePrimToPrim $ withForeignPtr fp $ \p -> peekByteOff p ((o+i) `shiftR` 1) - - {-# INLINE basicUnsafeWrite #-} - basicUnsafeWrite (MVector_Nucs_half o _ fp) i (Ns x) = - unsafePrimToPrim $ withForeignPtr fp $ \p -> do - y <- peekByteOff p ((o+i) `shiftR` 1) - let y' | even (o+i) = x `shiftL` 4 .|. y .&. 0x0F - | otherwise = x .|. y .&. 0xF0 - pokeByteOff p ((o+i) `shiftR` 1) y' - -foreign import ccall unsafe "string.h memset" memset - :: Ptr Word8 -> CInt -> CSize -> IO () - -instance Show (Vector_Nucs_half Nucleotides) where - show = show . V.toList - --- | Bam record in its native encoding along with virtual address. -data BamRaw = BamRaw { virt_offset :: {-# UNPACK #-} !FileOffset - , raw_data :: {-# UNPACK #-} !Bytes } - --- | Smart constructor. Makes sure we got a at least a full record. -{-# INLINE bamRaw #-} -bamRaw :: FileOffset -> Bytes -> BamRaw -bamRaw o s = if good then BamRaw o s else error $ "broken BAM record " ++ shows (S.length s, m) " " ++ show (S.unpack (S.take 10 s)) - where - good | S.length s < 32 = False - | otherwise = S.length s >= sum m - m = [ 32, l_rnm, l_seq, (l_seq+1) `div` 2, l_cig * 4 ] - l_rnm = fromIntegral (B.unsafeIndex s 8) - 1 - l_cig = fromIntegral (B.unsafeIndex s 12) .|. fromIntegral (B.unsafeIndex s 13) `shiftL` 8 - l_seq = fromIntegral (B.unsafeIndex s 16) .|. fromIntegral (B.unsafeIndex s 17) `shiftL` 8 .|. - fromIntegral (B.unsafeIndex s 18) `shiftL` 16 .|. fromIntegral (B.unsafeIndex s 19) `shiftL` 24 - -{-# INLINE[1] unpackBam #-} -unpackBam :: BamRaw -> BamRec -unpackBam br = BamRec { - b_rname = Refseq $ getWord32 0, - b_pos = getInt32 4, - b_mapq = Q $ getInt8 9, - b_flag = getInt16 14, - b_mrnm = Refseq $ getWord32 20, - b_mpos = getInt32 24, - b_isize = getInt32 28, - - b_qname = B.unsafeTake l_read_name $ B.unsafeDrop 32 $ raw_data br, - b_cigar = VS.unsafeCast $ VS.unsafeFromForeignPtr fp (off0+off_c) (4*l_cigar), - b_seq = Vector_Nucs_half (2 * (off_s+off0)) l_seq fp, - b_qual = VS.unsafeCast $ VS.unsafeFromForeignPtr fp (off0+off_q) l_seq, - - b_exts = unpackExtensions $ S.drop off_e $ raw_data br, - b_virtual_offset = virt_offset br } - where - (fp, off0, _) = B.toForeignPtr $ raw_data br - off_c = 33 + l_read_name - off_s = off_c + 4 * l_cigar - off_q = off_s + (l_seq + 1) `div` 2 - off_e = off_q + l_seq - - l_read_name = getInt8 8 - 1 - l_seq = getWord32 16 - l_cigar = getInt16 12 - - getInt8 :: Num a => Int -> a - getInt8 o = fromIntegral (B.unsafeIndex (raw_data br) o) - - getInt16 :: Num a => Int -> a - getInt16 o = unsafeDupablePerformIO $ B.unsafeUseAsCString (raw_data br) $ - fmap fromIntegral . peekUnalnWord16LE . flip plusPtr o - - getWord32 :: Num a => Int -> a - getWord32 o = unsafeDupablePerformIO $ B.unsafeUseAsCString (raw_data br) $ - fmap fromIntegral . peekUnalnWord32LE . flip plusPtr o - - -- ensures proper sign extension - getInt32 :: Num a => Int -> a - getInt32 o = fromIntegral (getWord32 o :: Int32) - --- | A collection of extension fields. A 'BamKey' is actually two ASCII --- characters. -type Extensions = [( BamKey, Ext )] - --- | Deletes all occurences of some extension field. -deleteE :: BamKey -> Extensions -> Extensions -deleteE k = filter ((/=) k . fst) - --- | Blindly inserts an extension field. This can create duplicates --- (and there is no telling how other tools react to that). -insertE :: BamKey -> Ext -> Extensions -> Extensions -insertE k v = (:) (k,v) - --- | Deletes all occurences of an extension field, then inserts it with --- a new value. This is safer than 'insertE', but also more expensive. -updateE :: BamKey -> Ext -> Extensions -> Extensions -updateE k v = insertE k v . deleteE k - --- | Adjusts a named extension by applying a function. -adjustE :: (Ext -> Ext) -> BamKey -> Extensions -> Extensions -adjustE _ _ [ ] = [] -adjustE f k ((k',v):es) | k == k' = (k', f v) : es - | otherwise = (k', v) : adjustE f k es - -data Ext = Int Int | Float Float | Text Bytes | Bin Bytes | Char Word8 - | IntArr (U.Vector Int) | FloatArr (U.Vector Float) - deriving (Show, Eq, Ord) - -{-# INLINE unpackExtensions #-} -unpackExtensions :: Bytes -> Extensions -unpackExtensions = go - where - go s | S.length s < 4 = [] - | otherwise = let key = fromString [ S.index s 0, S.index s 1 ] - in case S.index s 2 of - 'Z' -> case S.break (== '\0') (S.drop 3 s) of (l,r) -> (key, Text l) : go (S.drop 1 r) - 'H' -> case S.break (== '\0') (S.drop 3 s) of (l,r) -> (key, Bin l) : go (S.drop 1 r) - 'A' -> (key, Char (B.index s 3)) : go (S.drop 4 s) - 'B' -> let tp = S.index s 3 - n = getInt 'I' (S.drop 4 s) - in case tp of - 'f' -> (key, FloatArr (U.fromListN (n+1) [ getFloat (S.drop i s) | i <- [8, 12 ..] ])) - : go (S.drop (12+4*n) s) - _ -> (key, IntArr (U.fromListN (n+1) [ getInt tp (S.drop i s) | i <- [8, 8 + size tp ..] ])) - : go (S.drop (8 + size tp * (n+1)) s) - 'f' -> (key, Float (getFloat (S.drop 3 s))) : go (S.drop 7 s) - tp -> (key, Int (getInt tp (S.drop 3 s))) : go (S.drop (3 + size tp) s) - - size 'C' = 1 - size 'c' = 1 - size 'S' = 2 - size 's' = 2 - size 'I' = 4 - size 'i' = 4 - size 'f' = 4 - size _ = 0 - - getInt 'C' s | S.length s >= 1 = fromIntegral (fromIntegral (B.index s 0) :: Word8) - getInt 'c' s | S.length s >= 1 = fromIntegral (fromIntegral (B.index s 0) :: Int8) - getInt 'S' s | S.length s >= 2 = fromIntegral (i :: Word16) - where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord16LE - getInt 's' s | S.length s >= 2 = fromIntegral (fromIntegral i :: Int16) - where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord16LE - getInt 'I' s | S.length s >= 4 = fromIntegral (i :: Word32) - where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord32LE - getInt 'i' s | S.length s >= 4 = fromIntegral (fromIntegral i :: Int32) - where i = unsafeDupablePerformIO $ B.unsafeUseAsCString s $ peekUnalnWord32LE - getInt _ _ = 0 - - getFloat s = unsafeDupablePerformIO $ alloca $ \buf -> - pokeByteOff buf 0 (getInt 'I' s :: Word32) >> peek buf - - -isPaired, isProperlyPaired, isUnmapped, isMateUnmapped, isReversed, - isMateReversed, isFirstMate, isSecondMate, isAuxillary, isSecondary, - isFailsQC, isDuplicate, isSupplementary, - isTrimmed, isMerged, isAlternative, isExactIndex :: BamRec -> Bool - -isPaired = flip testBit 0 . b_flag -isProperlyPaired = flip testBit 1 . b_flag -isUnmapped = flip testBit 2 . b_flag -isMateUnmapped = flip testBit 3 . b_flag -isReversed = flip testBit 4 . b_flag -isMateReversed = flip testBit 5 . b_flag -isFirstMate = flip testBit 6 . b_flag -isSecondMate = flip testBit 7 . b_flag -isAuxillary = flip testBit 8 . b_flag -isSecondary = flip testBit 8 . b_flag -isFailsQC = flip testBit 9 . b_flag -isDuplicate = flip testBit 10 . b_flag -isSupplementary = flip testBit 11 . b_flag - -isTrimmed = flip testBit 0 . extAsInt 0 "FF" -isMerged = flip testBit 1 . extAsInt 0 "FF" -isAlternative = flip testBit 2 . extAsInt 0 "FF" -isExactIndex = flip testBit 3 . extAsInt 0 "FF" - -type_mask :: Int -type_mask = flagFirstMate .|. flagSecondMate .|. flagPaired - -extAsInt :: Int -> BamKey -> BamRec -> Int -extAsInt d nm br = case lookup nm (b_exts br) of Just (Int i) -> i ; _ -> d - -extAsString :: BamKey -> BamRec -> Bytes -extAsString nm br = case lookup nm (b_exts br) of - Just (Char c) -> B.singleton c - Just (Text s) -> s - _ -> B.empty - -setQualFlag :: Char -> BamRec -> BamRec -setQualFlag c br = br { b_exts = updateE "ZQ" (Text s') $ b_exts br } - where - s = extAsString "ZQ" br - s' = if c `S.elem` s then s else c `S.cons` s - --- | A simple progress indicator that prints sequence id and position. -progressBam :: MonadIO m => String -> Refs -> Int -> (String -> IO ()) -> Enumeratee [BamRaw] [BamRaw] m a -progressBam = progressPos (\br -> case unpackBam br of b -> (b_rname b, b_pos b)) - diff --git a/src/Bio/Bam/Regions.hs b/src/Bio/Bam/Regions.hs deleted file mode 100644 --- a/src/Bio/Bam/Regions.hs +++ /dev/null @@ -1,45 +0,0 @@ -module Bio.Bam.Regions where - -import Bio.Bam.Header ( Refseq(..) ) -import Bio.Prelude - -import qualified Data.IntMap.Strict as IM - -data Region = Region { refseq :: !Refseq, start :: !Int, end :: !Int } - deriving (Eq, Ord, Show) - --- | A subset of a genome. The idea is to map the reference sequence --- (represented by its number) to a 'Subseqeunce'. -newtype Regions = Regions (IM.IntMap Subsequence) deriving Show - --- | A mostly contiguous subset of a sequence, stored as a set of --- non-overlapping intervals in an 'IntMap' from start position to end --- position (half-open intervals, naturally). -newtype Subsequence = Subsequence (IM.IntMap Int) deriving Show - -toList :: Regions -> [(Refseq, Subsequence)] -toList (Regions m) = [ (Refseq $ fromIntegral k, v) | (k,v) <- IM.toList m ] - -fromList :: [Region] -> Regions -fromList = foldl' (flip add) (Regions IM.empty) - -add :: Region -> Regions -> Regions -add (Region (Refseq r) b e) (Regions m) = - let single = Just . Subsequence $ IM.singleton b e - in Regions $ IM.alter (maybe single (Just . addInt b e)) (fromIntegral r) m - - -addInt :: Int -> Int -> Subsequence -> Subsequence -addInt b e (Subsequence m0) = Subsequence $ merge_into b e m0 - where - merge_into x y m = case IM.lookupLT y m of - Just (u,v) | x < u && y <= v -> merge_into x v $ IM.delete u m -- extend to the left - | x < u -> merge_into x y $ IM.delete u m -- subsume - | y <= v -> m -- subsumed - | x <= v -> merge_into u y $ IM.delete u m -- extend to the right - _ -> IM.insert x y m -- no overlap - -overlaps :: Int -> Int -> Subsequence -> Bool -overlaps b e (Subsequence m) = case IM.lookupLT e m of - Just (_,v) -> b < v - Nothing -> False diff --git a/src/Bio/Bam/Rmdup.hs b/src/Bio/Bam/Rmdup.hs deleted file mode 100644 --- a/src/Bio/Bam/Rmdup.hs +++ /dev/null @@ -1,688 +0,0 @@ -module Bio.Bam.Rmdup( - rmdup, Collapse, cons_collapse, cheap_collapse, - cons_collapse_keep, cheap_collapse_keep, - check_sort, normalizeTo, wrapTo, - ECig(..), toECig, setMD, toCigar - ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Iteratee -import Bio.Prelude hiding ( left, right ) - -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as T -import qualified Data.Map.Strict as M -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as U - -data Collapse = Collapse { - collapse :: [BamRec] -> (Decision,[BamRec]), -- cluster to consensus and stuff or representative and stuff - originals :: [BamRec] -> [BamRec] } -- treatment of the redundant original reads - -data Decision = Consensus { fromDecision :: BamRec } - | Representative { fromDecision :: BamRec } - -cons_collapse :: Qual -> Collapse -cons_collapse maxq = Collapse (do_collapse maxq) (const []) - -cons_collapse_keep :: Qual -> Collapse -cons_collapse_keep maxq = Collapse (do_collapse maxq) (map (\b -> b { b_flag = b_flag b .|. flagDuplicate })) - -cheap_collapse :: Collapse -cheap_collapse = Collapse do_cheap_collapse (const []) - -cheap_collapse_keep :: Collapse -cheap_collapse_keep = Collapse do_cheap_collapse (map (\b -> b { b_flag = b_flag b .|. flagDuplicate })) - - --- | Removes duplicates from an aligned, sorted BAM stream. --- --- The incoming stream must be sorted by coordinate, and we check for --- violations of that assumption. We cannot assume that length was --- taken into account when sorting (samtools doesn't do so), so --- duplicates may be separated by reads that start at the same position --- but have different length or different strand. --- --- We are looking at three different kinds of reads: paired reads, true --- single ended reads, merged or trimmed reads. They are somewhat --- different, but here's the situation if we wanted to treat them --- separately. These conditions define a set of duplicates: --- --- Merged or trimmed: We compare the leftmost coordinates and the --- aligned length. If the library prep is strand-preserving, we also --- compare the strand. --- --- Paired: We compare both left-most coordinates (b_pos and b_mpos). If --- the library prep is strand-preserving, only first-mates can be --- duplicates of first-mates. Else a first-mate can be the duplicate of --- a second-mate. There may be pairs with one unmapped mate. This is --- not a problem as they get assigned synthetic coordinates and will be --- handled smoothly. --- --- True singles: We compare only the leftmost coordinate. It does not --- matter if the library prep is strand-preserving, the strand always --- matters. --- --- Across these classes, we can see more duplicates: --- --- Merged/trimmed and paired: these can be duplicates if the merging --- failed for the pair. We would need to compare the outer coordinates --- of the merged reads to the two 5' coordinates of the pair. However, --- since we don't have access to the mate, we cannot actually do --- anything right here. This case should be solved externally by --- merging those pairs that overlap in coordinate space. --- --- Single and paired: in the single case, we only have one coordinate --- to compare. This will inevitably lead to trouble, as we could find --- that the single might be the duplicate of two pairs, but those two --- pairs are definitely not duplicates of each other. We solve it by --- removing the single read(s). --- --- Single and merged/trimmed: same trouble as in the single+paired --- case. We remove the single to solve it. --- --- --- In principle, we might want to allow some wiggle room in the --- coordinates. So far, this has not been implemented. It adds the --- complication that groups of separated reads can turn into a set of --- duplicates because of the appearance of a new reads. Needs some --- thinking about... or maybe it's not too important. --- --- Once a set of duplicates is collected, we perform a majority vote on --- the correct CIGAR line. Of all those reads that agree on this CIGAR --- line, a consensus is called, quality scores are adjusted and clamped --- to a maximum, the MD field is updated and the XP field is assigned --- the number of reads in the original cluster. The new MAPQ becomes --- the RMSQ of the map qualities of all reads. --- --- Treatment of Read Groups: We generalize by providing a "label" --- function; only reads that have the same label are considered --- duplicates of each other. The typical label function would extract --- read groups, libraries or samples. - -rmdup :: (Monad m, Ord l) => (BamRec -> l) -> Bool -> Collapse -> Enumeratee [BamRec] [(Int,BamRec)] m r -rmdup label strand_preserved collapse_cfg = - -- Easiest way to go about this: We simply collect everything that - -- starts at some specific coordinate and group it appropriately. - -- Treat the groups separately, output, go on. - check_sort id "input must be sorted for rmdup to work" ><> - mapGroups rmdup_group ><> - check_sort snd "internal error, output isn't sorted anymore" - where - rmdup_group = nice_sort . do_rmdup label strand_preserved collapse_cfg - same_pos u v = b_cpos u == b_cpos v - b_cpos u = (b_rname u, b_pos u) - - nice_sort = sortBy $ comparing (V.length . b_seq . snd) - - mapGroups f o = tryHead >>= maybe (return o) (\a -> eneeCheckIfDone (mg1 f a []) o) - mg1 f a acc k = tryHead >>= \case - Nothing -> return . k . Chunk . f $ a:acc - Just b | same_pos a b -> mg1 f a (b:acc) k - | otherwise -> eneeCheckIfDone (mg1 f b []) . k . Chunk . f $ a:acc - -check_sort :: Monad m => (a -> BamRec) -> String -> Enumeratee [a] [a] m b -check_sort f msg out = tryHead >>= maybe (return out) (\a -> eneeCheckIfDone (step a) out) - where - step a k = tryHead >>= maybe (return . k $ Chunk [a]) (step' a k) - step' a k b | (b_rname (f a), b_pos (f a)) > (b_rname (f b), b_pos (f b)) = fail $ "rmdup: " ++ msg - | otherwise = eneeCheckIfDone (step b) . k $ Chunk [a] - - -{- | Workhorse for duplicate removal. - - - Unmapped fragments should not be considered to be duplicates of - mapped fragments. The /unmapped/ flag can serve for that: while - there are two classes of /unmapped/ reads (those that are not mapped - and those that are mapped to an invalid position), the two sets will - always have different coordinates. (Unfortunately, correct duplicate - removal now relies on correct /unmapped/ and /mate unmapped/ flags, - and we don't get them from unmodified BWA. So correct operation - requires patched BWA or a run of @bam-fixpair@.) - - (1) Other definitions (e.g. lack of CIGAR) don't work, because that - information won't be available for the mate. - - (2) This would amount to making the /unmapped/ flag part of the - coordinate, but samtools is not going to take it into account - when sorting. - - (3) Instead, both flags become part of the /mate pos/ grouping - criterion. - - - First Mates should (probably) not be considered duplicates of Second - Mates. This is unconditionally true for libraries with A\/B-style - adapters (definitely 454, probably Mathias' ds protocol) and the ss - protocol, it is not true for fork adapter protocols (vanilla Illumina - protocol). So it has to be an option, which would ideally be derived - from header information. - - - This code ignores read groups, but it will do a majority vote on the - @RG@ field and call consensi for the index sequences. If you believe - that duplicates across read groups are impossible, you must call it - with an appropriately filtered stream. - - - Half-Aligned Pairs (meaning one known coordinate, while the validity - of the alignments is immaterial) are rather complicated: - - (1) Given that only one coordinate is known (5' of the aligned mate), - we want to treat them like true singles. But the unaligned mate - should be kept if possible, though it should not contribute to a - consensus sequence. We assume nothing about the unaligned mate, - not even that it /shouldn't/ be aligned, never mind the fact that - it /couldn't/ be. (The difference is in the finite abilities of - real world aligners, naturally.) - - (2) Therefore, aligned reads with unaligned mates go to the same - potential duplicate set as true singletons. If at least one pair - exists that might be a duplicate of those, all singletons and - half-aligned mates are discarded. Else a consensus is computed - and replaces the aligned mates. - - (3) The unaligned mates end up in the same place in a BAM stream as - the aligned mates (therefore we see them and can treat them - locally). We cannot call a consensus, since these molecules may - well have different length, so we select one. It doesn't really - matter which one is selected, and since we're treating both mates - at the same time, it doesn't even need to be reproducible without - local information. This is made to be the mate of the consensus. - - (4) See 'merge_singles' for how it's actually done. --} - -do_rmdup :: Ord l => (BamRec -> l) -> Bool -> Collapse -> [BamRec] -> [(Int,BamRec)] -do_rmdup label strand_preserved Collapse{..} = - concatMap do_rmdup1 . M.elems . accumMap label id - where - do_rmdup1 rds = results ++ map ((,) 0) (originals (leftovers ++ r1 ++ r2 ++ r3)) - where - (results, leftovers) = merge_singles singles' unaligned' $ - [ (str, second fromDecision b) | ((_,str ),b) <- M.toList merged' ] ++ - [ (str, second fromDecision b) | ((_,str,_),b) <- M.toList pairs' ] - - (raw_pairs, raw_singles) = partition isPaired rds - (merged, true_singles) = partition (liftA2 (||) isMerged isTrimmed) raw_singles - - (pairs, raw_half_pairs) = partition b_totally_aligned raw_pairs - (half_unaligned, half_aligned) = partition isUnmapped raw_half_pairs - - mkMap :: Ord a => (BamRec -> a) -> [BamRec] -> (M.Map a (Int,Decision), [BamRec]) - mkMap f x = let m1 = M.map (length &&& collapse) $ accumMap f id x - in (M.map (second fst) m1, concatMap (snd.snd) $ M.elems m1) - - (pairs',r1) = mkMap (\b -> (b_mate_pos b, b_strand b, b_mate b)) pairs - (merged',r2) = mkMap (\b -> (alignedLength (b_cigar b), b_strand b)) merged - (singles',r3) = mkMap b_strand (true_singles++half_aligned) - unaligned' = accumMap b_strand id half_unaligned - - b_strand b = strand_preserved && isReversed b - b_mate b = strand_preserved && isFirstMate b - - --- | Merging information about true singles, merged singles, --- half-aligned pairs, actually aligned pairs. --- --- We collected aligned reads with unaligned mates together with aligned --- true singles (@singles@). We collected the unaligned mates, which --- necessarily have the exact same alignment coordinates, separately --- (@unaligned@). If we don't find a matching true pair (that case is --- already handled smoothly), we keep the highest quality unaligned --- mate, pair it with the consensus of the aligned mates and aligned --- singletons, and give it the lexically smallest name of the --- half-aligned pairs. - --- NOTE: I need to decide when to run 'make_singleton'. Basically, --- when we call a consensus for half-aligned pairs and keep --- everything(?). Then we don't have a mate for the consensus... though --- we could decide to duplicate one mate read to get it. - -merge_singles :: M.Map Bool (Int,Decision) -- strand --> true singles & half aligned - -> M.Map Bool [BamRec] -- strand --> half unaligned - -> [ (Bool, (Int, BamRec)) ] -- strand --> paireds & mergeds - -> ([(Int,BamRec)],[BamRec]) -- results, leftovers - -merge_singles singles unaligneds = go - where - -- Say we generated a consensus or passed something through. If - -- there is a singleton consensus with the same strand, we should - -- add in its XP field and discard it. If there is a singleton - -- representative, we add in its XP field and put it into the - -- leftovers. If there is unaligned stuff here that has the same - -- strand, it goes to the leftovers. - go ( (str, (m,v)) : paireds) = - let (r,l) = merge_singles (M.delete str singles) (M.delete str unaligneds) paireds - unal = M.findWithDefault [] str unaligneds ++ l - - in case M.lookup str singles of - Nothing -> ( (m,v) : r, unal ) - Just (n, Consensus w) -> ( (n, add_xp_of w v) : r, unal ) - Just (n, Representative w) -> ( (n, add_xp_of w v) : r, w : unal ) - - -- No more pairs, delegate the problem - go [] = merge_halves unaligneds (M.toList singles) - - add_xp_of w v = v { b_exts = updateE "XP" (Int $ extAsInt 1 "XP" w `oplus` extAsInt 1 "XP" v) (b_exts v) } - --- | Merging of half-aligned reads. The first argument is a map of --- unaligned reads (their mates are aligned to the current position), --- the second is a list of reads that are aligned (their mates are not --- aligned). --- --- So, suppose we're looking at a 'Representative' that was passed --- through. We need to emit it along with its mate, which may be hidden --- inside a list. (Alternatively, we could force it to single, but that --- fails if we're passing everything along somehow.) --- --- Suppose we're looking at a 'Consensus'. We could pair it with some --- mate (which we'd need to duplicate), or we could turn it into a --- singleton. Duplication is ugly, so in this case, we force it to --- singleton. - -merge_halves :: M.Map Bool [BamRec] -- strand --> half unaligned - -> [(Bool, (Int,Decision))] -- strand --> true singles & half aligned - -> ([(Int,BamRec)],[BamRec]) -- results, leftovers - --- Emitting a consensus: make it a single. Nothing goes to leftovers; --- we may still need it for something else to be emitted. (While that --- would be strange, making sure the BAM file stays completely valid is --- probably better.) -merge_halves unaligneds ((_, (n, Consensus v)) : singles) = ( (n, v { b_flag = b_flag v .&. complement pflags }) : r, l ) - where - (r,l) = merge_halves unaligneds singles - pflags = flagPaired .|. flagProperlyPaired .|. flagMateUnmapped .|. flagMateReversed .|. flagFirstMate .|. flagSecondMate - - --- Emitting a representative: find the mate in the list of unaligned --- reads (take up to one match to be robust), and emit that, too, as a --- result. Everything else goes to leftovers. If the representative --- happens to be unpaired, no mate is found and that case therefore is --- handled smoothly. -merge_halves unaligneds ((str, (n, Representative v)) : singles) = ((n,v) : map ((,)1) (take 1 same) ++ r, drop 1 same ++ diff ++ l) - where - (r,l) = merge_halves (M.delete str unaligneds) singles - (same,diff) = partition (is_mate_of v) $ M.findWithDefault [] str unaligneds - is_mate_of a b = b_qname a == b_qname b && isPaired a && isPaired b && isFirstMate a == isSecondMate b - --- No more singles, all unaligneds are leftovers. -merge_halves unaligneds [] = ( [], concat $ M.elems unaligneds ) - - - - -type MPos = (Refseq, Int, Bool, Bool) - -b_mate_pos :: BamRec -> MPos -b_mate_pos b = (b_mrnm b, b_mpos b, isUnmapped b, isMateUnmapped b) - -b_totally_aligned :: BamRec -> Bool -b_totally_aligned b = not (isUnmapped b || isMateUnmapped b) - - -accumMap :: Ord k => (a -> k) -> (a -> v) -> [a] -> M.Map k [v] -accumMap f g = go M.empty - where - go m [ ] = m - go m (a:as) = let ws = M.findWithDefault [] (f a) m ; g' = g a - in g' `seq` go (M.insert (f a) (g':ws) m) as - - -{- We need to deal sensibly with each field, but different fields have - different needs. We can take the value from the first read to - preserve determinism or because all reads should be equal anyway, - aggregate over all reads computing either RMSQ or the most common - value, delete a field because it wouldn't make sense anymore or - because doing something sensible would be hard and we're going to - ignore it anyway, or we calculate some special value; see below. - Unknown fields will be taken from the first read, which seems to be a - safe default. - - QNAME and most fields taken from first - FLAG qc fail majority vote - dup deleted - MAPQ rmsq - CIGAR, SEQ, QUAL, MD, NM, XP generated - XA concatenate all - - BQ, CM, FZ, Q2, R2, XM, XO, XG, YQ, EN - deleted because they would become wrong - - CQ, CS, E2, FS, OQ, OP, OC, U2, H0, H1, H2, HI, NH, IH, ZQ - delete because they will be ignored anyway - - AM, AS, MQ, PQ, SM, UQ - compute rmsq - - X0, X1, XT, XS, XF, XE, BC, LB, RG, XI, YI, XJ, YJ - majority vote -} - -do_collapse :: Qual -> [BamRec] -> (Decision, [BamRec]) -do_collapse maxq [br] | V.all (<= maxq) (b_qual br) = ( Representative br, [ ] ) -- no modifcation, pass through - | otherwise = ( Consensus lq_br, [br] ) -- qualities reduced, must keep original - where - lq_br = br { b_qual = V.map (min maxq) $ b_qual br - , b_virtual_offset = 0 - , b_qname = b_qname br `B.snoc` c2w 'c' } - -do_collapse maxq brs = ( Consensus b0 { b_exts = modify_extensions $ b_exts b0 - , b_flag = failflag .&. complement flagDuplicate - , b_mapq = Q $ rmsq $ map (unQ . b_mapq) $ good brs - , b_cigar = cigar' - , b_seq = V.fromList $ map fst cons_seq_qual - , b_qual = V.fromList $ map snd cons_seq_qual - , b_qname = b_qname b0 `B.snoc` 99 - , b_virtual_offset = 0 }, brs ) -- many modifications, must keep everything - where - !b0 = minimumBy (comparing b_qname) brs - !most_fail = 2 * length (filter isFailsQC brs) > length brs - !failflag | most_fail = b_flag b0 .|. flagFailsQC - | otherwise = b_flag b0 .&. complement flagFailsQC - - rmsq xs = case foldl' (\(!n,!d) x -> (n + fromIntegral x * fromIntegral x, d + 1)) (0,0) xs of - (!n,!d) -> round $ sqrt $ (n::Double) / fromIntegral (d::Int) - - maj xs = head . maximumBy (comparing length) . group . sort $ xs - nub' = concatMap head . group . sort - - -- majority vote on the cigar lines, then filter - !cigar' = maj $ map b_cigar brs - good = filter ((==) cigar' . b_cigar) - - cons_seq_qual = [ consensus maxq [ (V.unsafeIndex (b_seq b) i, q) - | b <- good brs, let q = if V.null (b_qual b) then Q 23 else b_qual b V.! i ] - | i <- [0 .. len - 1] ] - where !len = V.length . b_seq . head $ good brs - - md' = case [ (b_seq b,md,b) | b <- good brs, Just md <- [ getMd b ] ] of - [ ] -> [] - (seq1, md1,b) : _ -> case mk_new_md' [] (V.toList cigar') md1 (V.toList seq1) (map fst cons_seq_qual) of - Right x -> x - Left (MdFail cigs ms osq nsq) -> error $ unlines - [ "Broken MD field when trying to construct new MD!" - , "QNAME: " ++ show (b_qname b) - , "POS: " ++ shows (unRefseq (b_rname b)) ":" ++ show (b_pos b) - , "CIGAR: " ++ show cigs - , "MD: " ++ show ms - , "refseq: " ++ show osq - , "readseq: " ++ show nsq ] - - - nm' = sum $ [ n | Ins :* n <- VS.toList cigar' ] ++ [ n | Del :* n <- VS.toList cigar' ] ++ [ 1 | MdRep _ <- md' ] - xa' = nub' [ T.split ';' xas | Just (Text xas) <- map (lookup "XA" . b_exts) brs ] - - modify_extensions es = foldr ($!) es $ - [ let vs = mapMaybe (lookup k . b_exts) brs - in if null vs then id else updateE k $! maj vs | k <- do_maj ] ++ - [ let vs = [ v | Just (Int v) <- map (lookup k . b_exts) brs ] - in if null vs then id else updateE k $! Int (rmsq vs) | k <- do_rmsq ] ++ - map deleteE useless ++ - [ updateE "NM" $! Int nm' - , updateE "XP" $! Int (foldl' (\a b -> a `oplus` extAsInt 1 "XP" b) 0 brs) - , if null xa' then id else updateE "XA" $! (Text $ T.intercalate (T.singleton ';') xa') - , if null md' then id else updateE "MD" $! (Text $ showMd md') ] - - useless = map fromString $ words "BQ CM FZ Q2 R2 XM XO XG YQ EN CQ CS E2 FS OQ OP OC U2 H0 H1 H2 HI NH IH ZQ" - do_rmsq = map fromString $ words "AM AS MQ PQ SM UQ" - do_maj = map fromString $ words "X0 X1 XT XS XF XE BC LB RG XI XJ YI YJ" - -minViewBy :: (a -> a -> Ordering) -> [a] -> (a,[a]) -minViewBy _ [ ] = error "minViewBy on empty list" -minViewBy cmp (x:xs) = go x [] xs - where - go m acc [ ] = (m,acc) - go m acc (a:as) = case m `cmp` a of GT -> go a (m:acc) as - _ -> go m (a:acc) as - -data MdFail = MdFail [Cigar] [MdOp] [Nucleotides] [Nucleotides] - -mk_new_md' :: [MdOp] -> [Cigar] -> [MdOp] -> [Nucleotides] -> [Nucleotides] -> Either MdFail [MdOp] -mk_new_md' acc [] [] [] [] = Right $ normalize [] acc - where - normalize a (MdNum 0:os) = normalize a os - normalize (MdNum n:a) (MdNum m:os) = normalize (MdNum (n+m):a) os - normalize a (MdDel []:os) = normalize a os - normalize (MdDel u:a) (MdDel v:os) = normalize (MdDel (v++u):a) os - normalize a ( o:os) = normalize (o:a) os - normalize a [ ] = a - -mk_new_md' acc ( _ :* 0 : cigs) mds osq nsq = mk_new_md' acc cigs mds osq nsq -mk_new_md' acc cigs (MdNum 0 : mds) osq nsq = mk_new_md' acc cigs mds osq nsq -mk_new_md' acc cigs (MdDel [] : mds) osq nsq = mk_new_md' acc cigs mds osq nsq - -mk_new_md' acc (Mat :* u : cigs) (MdRep b : mds) (_:osq) (n:nsq) - | b == n = mk_new_md' (MdNum 1 : acc) (Mat :* (u-1):cigs) mds osq nsq - | otherwise = mk_new_md' (MdRep b : acc) (Mat :* (u-1):cigs) mds osq nsq - -mk_new_md' acc (Mat :* u : cigs) (MdNum v : mds) (o:osq) (n:nsq) - | o == n = mk_new_md' (MdNum 1 : acc) (Mat :* (u-1):cigs) (MdNum (v-1) : mds) osq nsq - | otherwise = mk_new_md' (MdRep o : acc) (Mat :* (u-1):cigs) (MdNum (v-1) : mds) osq nsq - -mk_new_md' acc (Del :* n : cigs) (MdDel bs : mds) osq nsq | n == length bs = mk_new_md' (MdDel bs : acc) cigs mds osq nsq -mk_new_md' acc (Del :* n : cigs) (MdDel (b:bs) : mds) osq nsq = mk_new_md' (MdDel [b] : acc) (Del :* (n-1) : cigs) (MdDel bs:mds) osq nsq -mk_new_md' acc (Del :* n : cigs) (MdRep b : mds) osq nsq = mk_new_md' (MdDel [b] : acc) (Del :* (n-1) : cigs) mds osq nsq -mk_new_md' acc (Del :* n : cigs) (MdNum m : mds) osq nsq = mk_new_md' (MdDel [nucsN] : acc) (Del :* (n-1) : cigs) (MdNum (m-1):mds) osq nsq - -mk_new_md' acc (Ins :* n : cigs) md osq nsq = mk_new_md' acc cigs md (drop n osq) (drop n nsq) -mk_new_md' acc (SMa :* n : cigs) md osq nsq = mk_new_md' acc cigs md (drop n osq) (drop n nsq) -mk_new_md' acc (HMa :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq -mk_new_md' acc (Pad :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq -mk_new_md' acc (Nop :* _ : cigs) md osq nsq = mk_new_md' acc cigs md osq nsq - -mk_new_md' _acc cigs ms osq nsq = Left $ MdFail cigs ms osq nsq - -consensus :: Qual -> [ (Nucleotides, Qual) ] -> (Nucleotides, Qual) -consensus (Q maxq) nqs = if qr > 3 then (n0, Q qr) else (nucsN, Q 0) - where - accs :: U.Vector Int - accs = U.accum (+) (U.replicate 16 0) [ (fromIntegral n, fromIntegral q) | (Ns n,Q q) <- nqs ] - - (n0,q0) : (_,q1) : _ = sortBy (flip $ comparing snd) $ zip [Ns 0 ..] $ U.toList accs - qr = fromIntegral $ (q0-q1) `min` fromIntegral maxq - - --- Cheap version: simply takes the lexically first record, adds XP field -do_cheap_collapse :: [BamRec] -> ( Decision, [BamRec] ) -do_cheap_collapse [b] = ( Representative b, [] ) -do_cheap_collapse bs = ( Representative $ replaceXP new_xp b0, bx ) - where - (b0, bx) = minViewBy (comparing b_qname) bs - new_xp = foldl' (\a b -> a `oplus` extAsInt 1 "XP" b) 0 bs - -replaceXP :: Int -> BamRec -> BamRec -replaceXP new_xp b0 = b0 { b_exts = updateE "XP" (Int new_xp) $ b_exts b0 } - -oplus :: Int -> Int -> Int -_ `oplus` (-1) = -1 -(-1) `oplus` _ = -1 -a `oplus` b = a + b - --- | Normalize a read's alignment to fall into the canonical region --- of [0..l]. Takes the name of the reference sequence and its length. --- Returns @Left x@ if the coordinate decreased so the result is out of --- order now, @Right x@ if the coordinate is unchanged. -normalizeTo :: Seqid -> Int -> BamRec -> Either BamRec BamRec -normalizeTo nm l b = lr $ b { b_pos = b_pos b `mod` l - , b_mpos = b_mpos b `mod` l - , b_mapq = if dups_are_fine then Q 37 else b_mapq b - , b_exts = if dups_are_fine then deleteE "XA" (b_exts b) else b_exts b } - where - lr = if b_pos b >= l then Left else Right - dups_are_fine = all_match_XA (extAsString "XA" b) - all_match_XA s = case T.split ';' s of [xa1, xa2] | T.null xa2 -> one_match_XA xa1 - [xa1] -> one_match_XA xa1 - _ -> False - one_match_XA s = case T.split ',' s of (sq:pos:_) | sq == nm -> pos_match_XA pos ; _ -> False - pos_match_XA s = case T.readInt s of Just (p,z) | T.null z -> int_match_XA p ; _ -> False - int_match_XA p | p >= 0 = (p-1) `mod` l == b_pos b `mod` l && not (isReversed b) - | otherwise = (-p-1) `mod` l == b_pos b `mod` l && isReversed b - - --- | Wraps a read to be fully contained in the canonical interval --- [0..l]. If the read overhangs, it is duplicated and both copies are --- suitably masked. A piece with changed coordinate that is now out of --- order is returned as @Left x@, if the order is fine, it is returned --- as @Right x@. -wrapTo :: Int -> BamRec -> [Either BamRec BamRec] -wrapTo l b = if overhangs then do_wrap else [Right b] - where - overhangs = not (isUnmapped b) && b_pos b < l && l < b_pos b + alignedLength (b_cigar b) - - do_wrap = case split_ecig (l - b_pos b) $ toECig (b_cigar b) (fromMaybe [] $ getMd b) of - (left,right) -> [ Right $ b { b_cigar = toCigar left } `setMD` left - , Left $ b { b_cigar = toCigar right, b_pos = 0 } `setMD` right ] - --- | Split an 'ECig' into two at some position. The position is counted --- in terms of the reference (therefore, deletions count, insertions --- don't). The parts that would be skipped if we were splitting lists --- are replaced by soft masks. -split_ecig :: Int -> ECig -> (ECig, ECig) -split_ecig _ WithMD = (WithMD, WithMD) -split_ecig _ WithoutMD = (WithoutMD, WithoutMD) -split_ecig 0 ecs = (mask_all ecs, ecs) - -split_ecig i (Ins' n ecs) = case split_ecig i ecs of (u,v) -> (Ins' n u, SMa' n v) -split_ecig i (SMa' n ecs) = case split_ecig i ecs of (u,v) -> (SMa' n u, SMa' n v) -split_ecig i (HMa' n ecs) = case split_ecig i ecs of (u,v) -> (HMa' n u, HMa' n v) -split_ecig i (Pad' n ecs) = case split_ecig i ecs of (u,v) -> (Pad' n u, v) - -split_ecig i (Mat' n ecs) - | i >= n = case split_ecig (i-n) ecs of (u,v) -> (Mat' n u, SMa' n v) - | otherwise = (Mat' i $ SMa' (n-i) $ mask_all ecs, SMa' i $ Mat' (n-i) ecs) - -split_ecig i (Rep' x ecs) = case split_ecig (i-1) ecs of (u,v) -> (Rep' x u, SMa' 1 v) -split_ecig i (Del' x ecs) = case split_ecig (i-1) ecs of (u,v) -> (Del' x u, v) - -split_ecig i (Nop' n ecs) - | i >= n = case split_ecig (i-n) ecs of (u,v) -> (Nop' n u, v) - | otherwise = (Nop' i $ mask_all ecs, Nop' (n-i) ecs) - -mask_all :: ECig -> ECig -mask_all WithMD = WithMD -mask_all WithoutMD = WithoutMD -mask_all (Nop' _ ec) = mask_all ec -mask_all (HMa' _ ec) = mask_all ec -mask_all (Pad' _ ec) = mask_all ec -mask_all (Del' _ ec) = mask_all ec -mask_all (Rep' _ ec) = SMa' 1 $ mask_all ec -mask_all (Mat' n ec) = SMa' n $ mask_all ec -mask_all (Ins' n ec) = SMa' n $ mask_all ec -mask_all (SMa' n ec) = SMa' n $ mask_all ec - --- | Extended CIGAR. This subsumes both the CIGAR string and the --- optional MD field. If we have MD on input, we generate it on output, --- too. And in between, we break everything into /very small/ --- operations. - -data ECig = WithMD -- terminate, do generate MD field - | WithoutMD -- terminate, don't bother with MD - | Mat' Int ECig - | Rep' Nucleotides ECig - | Ins' Int ECig - | Del' Nucleotides ECig - | Nop' Int ECig - | SMa' Int ECig - | HMa' Int ECig - | Pad' Int ECig - - -toECig :: VS.Vector Cigar -> [MdOp] -> ECig -toECig = go . VS.toList - where - go cs (MdNum 0:mds) = go cs mds - go cs (MdDel []:mds) = go cs mds - go (_:*0 :cs) mds = go cs mds - go [ ] [ ] = WithMD -- all was fine to the very end - go [ ] _ = WithoutMD -- here it wasn't fine - - go (Mat :* n : cs) (MdRep x:mds) = Rep' x $ go (Mat :* (n-1) : cs) mds - go (Mat :* n : cs) (MdNum m:mds) - | n < m = Mat' n $ go cs (MdNum (m-n):mds) - | n > m = Mat' m $ go (Mat :* (n-m) : cs) mds - | otherwise = Mat' n $ go cs mds - go (Mat :* n : cs) _ = Mat' n $ go' cs - - go (Ins :* n : cs) mds = Ins' n $ go cs mds - go (Del :* n : cs) (MdDel (x:xs):mds) = Del' x $ go (Del :* (n-1) : cs) (MdDel xs:mds) - go (Del :* n : cs) _ = Del' nucsN $ go' (Del :* (n-1) : cs) - - go (Nop :* n : cs) mds = Nop' n $ go cs mds - go (SMa :* n : cs) mds = SMa' n $ go cs mds - go (HMa :* n : cs) mds = HMa' n $ go cs mds - go (Pad :* n : cs) mds = Pad' n $ go cs mds - - -- We jump here once the MD fiels ran out early or was messed up. - -- We no longer bother with it (this also happens if the MD isn't - -- present to begin with). - go' (_ :* 0 : cs) = go' cs - go' [ ] = WithoutMD -- we didn't have MD or it was broken - - go' (Mat :* n : cs) = Mat' n $ go' cs - go' (Ins :* n : cs) = Ins' n $ go' cs - go' (Del :* n : cs) = Del' nucsN $ go' (Del :* (n-1) : cs) - - go' (Nop :* n : cs) = Nop' n $ go' cs - go' (SMa :* n : cs) = SMa' n $ go' cs - go' (HMa :* n : cs) = HMa' n $ go' cs - go' (Pad :* n : cs) = Pad' n $ go' cs - - --- We normalize matches, deletions and soft masks, because these are the --- operations we generate. Everything else is either already normalized --- or nobody really cares anyway. -toCigar :: ECig -> VS.Vector Cigar -toCigar = V.fromList . go - where - go WithMD = [] - go WithoutMD = [] - - go (Ins' n ecs) = Ins :* n : go ecs - go (Nop' n ecs) = Nop :* n : go ecs - go (HMa' n ecs) = HMa :* n : go ecs - go (Pad' n ecs) = Pad :* n : go ecs - go (SMa' n ecs) = go_sma n ecs - go (Mat' n ecs) = go_mat n ecs - go (Rep' _ ecs) = go_mat 1 ecs - go (Del' _ ecs) = go_del 1 ecs - - go_sma !n (SMa' m ecs) = go_sma (n+m) ecs - go_sma !n ecs = SMa :* n : go ecs - - go_mat !n (Mat' m ecs) = go_mat (n+m) ecs - go_mat !n (Rep' _ ecs) = go_mat (n+1) ecs - go_mat !n ecs = Mat :* n : go ecs - - go_del !n (Del' _ ecs) = go_del (n+1) ecs - go_del !n ecs = Del :* n : go ecs - - - --- | Create an MD field from an extended CIGAR and place it in a record. --- We build it piecemeal (in 'go'), call out to 'addNum', 'addRep', --- 'addDel' to make sure the operations are not generated in a --- degenerate manner, and finally check if we're even supposed to create --- an MD field. -setMD :: BamRec -> ECig -> BamRec -setMD b ec = case go ec of - Just md -> b { b_exts = updateE "MD" (Text $ showMd md) (b_exts b) } - Nothing -> b { b_exts = deleteE "MD" (b_exts b) } - where - go WithMD = Just [] - go WithoutMD = Nothing - - go (Ins' _ ecs) = go ecs - go (Nop' _ ecs) = go ecs - go (SMa' _ ecs) = go ecs - go (HMa' _ ecs) = go ecs - go (Pad' _ ecs) = go ecs - go (Mat' n ecs) = (if n == 0 then id else fmap (addNum n)) $ go ecs - go (Rep' x ecs) = (if isGap x then id else fmap (addRep x)) $ go ecs - go (Del' x ecs) = (if isGap x then id else fmap (addDel x)) $ go ecs - - addNum n (MdNum m : mds) = MdNum (n+m) : mds - addNum n mds = MdNum n : mds - - addRep x mds = MdRep x : mds - - addDel x (MdDel y : mds) = MdDel (x:y) : mds - addDel x mds = MdDel [x] : mds diff --git a/src/Bio/Bam/Trim.hs b/src/Bio/Bam/Trim.hs deleted file mode 100644 --- a/src/Bio/Bam/Trim.hs +++ /dev/null @@ -1,441 +0,0 @@ --- | Trimming of reads as found in BAM files. Implements trimming low --- quality sequence from the 3' end. - -module Bio.Bam.Trim - ( trim_3 - , trim_3' - , trim_low_quality - , default_fwd_adapters - , default_rev_adapters - , find_merge - , mergeBam - , find_trim - , trimBam - , mergeTrimBam - , twoMins - , merged_seq - , merged_qual - ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Bam.Rmdup ( ECig(..), setMD, toECig ) -import Bio.Iteratee -import Bio.Prelude - -import Foreign.C.Types ( CInt(..) ) - -import qualified Data.ByteString as B -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Storable as W - --- | Trims from the 3' end of a sequence. --- @trim_3' p b@ trims the 3' end of the sequence in @b@ at the --- earliest position such that @p@ evaluates to true on every suffix --- that was trimmed off. Note that the 3' end may be the beginning of --- the sequence if it happens to be stored in reverse-complemented form. --- Also note that trimming from the 3' end may not make sense for reads --- that were constructed by merging paired end data (but we cannot take --- care of that here). Further note that trimming may break dependent --- information, notably the "mate" information of the mate and many --- optional fields. - -trim_3' :: ([Nucleotides] -> [Qual] -> Bool) -> BamRec -> BamRec -trim_3' p b | b_flag b `testBit` 4 = trim_rev - | otherwise = trim_fwd - where - trim_fwd = let l = subtract 1 . fromIntegral . length . takeWhile (uncurry p) $ - zip (inits . reverse . V.toList $ b_seq b) - (inits . reverse . V.toList $ b_qual b) - in trim_3 l b - - trim_rev = let l = subtract 1 . fromIntegral . length . takeWhile (uncurry p) $ - zip (inits . V.toList $ b_seq b) - (inits . V.toList $ b_qual b) - in trim_3 l b - -trim_3 :: Int -> BamRec -> BamRec -trim_3 l b | b_flag b `testBit` 4 = trim_rev - | otherwise = trim_fwd - where - trim_fwd = let (_, cigar') = trim_back_cigar (b_cigar b) l - c = modMd (takeECig (V.length (b_seq b) - l)) b - in c { b_seq = V.take (V.length (b_seq c) - l) (b_seq c) - , b_qual = V.take (V.length (b_qual c) - l) (b_qual c) - , b_cigar = cigar' - , b_exts = map (\(k,e) -> case e of - Text t | k `elem` trim_set - -> (k, Text (B.take (B.length t - l) t)) - _ -> (k,e) - ) (b_exts c) } - - trim_rev = let (off, cigar') = trim_fwd_cigar (b_cigar b) l - c = modMd (dropECig l) b - in c { b_seq = V.drop l (b_seq c) - , b_qual = V.drop l (b_qual c) - , b_pos = b_pos c + off - , b_cigar = cigar' - , b_exts = map (\(k,e) -> case e of - Text t | k `elem` trim_set - -> (k, Text (B.drop l t)) - _ -> (k,e) - ) (b_exts c) } - - trim_set = ["BQ","CQ","CS","E2","OQ","U2"] - - modMd :: (ECig -> ECig) -> BamRec -> BamRec - modMd f br = maybe br (setMD br . f . toECig (b_cigar br)) (getMd br) - - endOf :: ECig -> ECig - endOf WithMD = WithMD - endOf WithoutMD = WithoutMD - endOf (Mat' _ es) = endOf es - endOf (Ins' _ es) = endOf es - endOf (SMa' _ es) = endOf es - endOf (Rep' _ es) = endOf es - endOf (Del' _ es) = endOf es - endOf (Nop' _ es) = endOf es - endOf (HMa' _ es) = endOf es - endOf (Pad' _ es) = endOf es - - takeECig :: Int -> ECig -> ECig - takeECig 0 es = endOf es - takeECig _ WithMD = WithMD - takeECig _ WithoutMD = WithoutMD - takeECig n (Mat' m es) = Mat' n $ if n > m then takeECig (n-m) es else WithMD - takeECig n (Ins' m es) = Ins' n $ if n > m then takeECig (n-m) es else WithMD - takeECig n (SMa' m es) = SMa' n $ if n > m then takeECig (n-m) es else WithMD - takeECig n (Rep' ns es) = Rep' ns $ takeECig (n-1) es - takeECig n (Del' ns es) = Del' ns $ takeECig n es - takeECig n (Nop' m es) = Nop' m $ takeECig n es - takeECig n (HMa' m es) = HMa' m $ takeECig n es - takeECig n (Pad' m es) = Pad' m $ takeECig n es - - dropECig :: Int -> ECig -> ECig - dropECig 0 es = es - dropECig _ WithMD = WithMD - dropECig _ WithoutMD = WithoutMD - dropECig n (Mat' m es) = if n > m then dropECig (n-m) es else Mat' n WithMD - dropECig n (Ins' m es) = if n > m then dropECig (n-m) es else Ins' n WithMD - dropECig n (SMa' m es) = if n > m then dropECig (n-m) es else SMa' n WithMD - dropECig n (Rep' _ es) = dropECig (n-1) es - dropECig n (Del' _ es) = dropECig n es - dropECig n (Nop' _ es) = dropECig n es - dropECig n (HMa' _ es) = dropECig n es - dropECig n (Pad' _ es) = dropECig n es - - -trim_back_cigar, trim_fwd_cigar :: V.Vector v Cigar => v Cigar -> Int -> ( Int, v Cigar ) -trim_back_cigar c l = (o, V.fromList $ reverse c') where (o,c') = sanitize_cigar . trim_cigar l $ reverse $ V.toList c -trim_fwd_cigar c l = (o, V.fromList c') where (o,c') = sanitize_cigar $ trim_cigar l $ V.toList c - -sanitize_cigar :: (Int, [Cigar]) -> (Int, [Cigar]) -sanitize_cigar (o, [ ]) = (o, []) -sanitize_cigar (o, (op:*l):xs) | op == Pad = sanitize_cigar (o,xs) -- del P - | op == Del || op == Nop = sanitize_cigar (o + l, xs) -- adjust D,N - | op == Ins = (o, (SMa :* l):xs) -- I --> S - | otherwise = (o, (op :* l):xs) -- rest is fine - -trim_cigar :: Int -> [Cigar] -> (Int, [Cigar]) -trim_cigar 0 cs = (0, cs) -trim_cigar _ [] = (0, []) -trim_cigar l ((op:*ll):cs) | bad_op op = let (o,cs') = trim_cigar l cs in (o + reflen op ll, cs') - | otherwise = case l `compare` ll of - LT -> (reflen op l, (op :* (ll-l)):cs) - EQ -> (reflen op ll, cs) - GT -> let (o,cs') = trim_cigar (l - ll) cs in (o + reflen op ll, cs') - - where - reflen op' = if ref_op op' then id else const 0 - bad_op o = o /= Mat && o /= Ins && o /= SMa - ref_op o = o == Mat || o == Del - - --- | Trim predicate to get rid of low quality sequence. --- @trim_low_quality q ns qs@ evaluates to true if all qualities in @qs@ --- are smaller (i.e. worse) than @q@. -trim_low_quality :: Qual -> a -> [Qual] -> Bool -trim_low_quality q = const $ all (< q) - - --- | Finds the merge point. Input is list of forward adapters, list of --- reverse adapters, sequence1, quality1, sequence2, quality2; output is --- merge point and two qualities (YM, YN). -find_merge :: [W.Vector Nucleotides] -> [W.Vector Nucleotides] - -> W.Vector Nucleotides -> W.Vector Qual - -> W.Vector Nucleotides -> W.Vector Qual - -> (Int, Int, Int) -find_merge ads1 ads2 r1 q1 r2 q2 = (mlen, score2 - score1, plain_score - score1) - where - plain_score = 6 * fromIntegral (V.length r1 + V.length r2) - (score1, mlen, score2) = twoMins plain_score (V.length r1 + V.length r2) $ - merge_score ads1 ads2 r1 q1 r2 q2 - --- | Overlap-merging of read pairs. We shall compute the likelihood --- for every possible overlap, then select the most likely one (unless it --- looks completely random), compute a quality from the second best --- merge, then merge and clamp the quality accordingly. --- (We could try looking for chimaera after completing the merge, if --- only we knew which ones to expect?) --- --- Two reads go in, with two adapter lists. We return 'Nothing' if all --- merges looked mostly random. Else we return the two original reads, --- flagged as 'eflagVestigial' *and* the merged version, flagged as --- 'eflagMerged' and optionally 'eflagTrimmed'. All reads contain the --- computed qualities (in YM and YN), which we also return. --- --- The merging automatically limits quality scores some of the time. We --- additionally impose a hard limit of 63 to avoid difficulties --- representing the result, and even that is ridiculous. Sane people --- would further limit the returned quality! (In practice, map quality --- later imposes a limit anyway, so no worries...) - -mergeBam :: Int -> Int -> [W.Vector Nucleotides] -> [W.Vector Nucleotides] -> BamRec -> BamRec -> [BamRec] -mergeBam lowq highq ads1 ads2 r1 r2 - | V.null (b_seq r1) && V.null (b_seq r2) = [ ] - | qual1 < lowq || mlen < 0 = [ r1', r2' ] - | qual1 >= highq && mlen == 0 = [ ] - | qual1 >= highq = [ rm ] - | mlen < len_r1-20 || mlen < len_r2-20 = [ rm ] - | otherwise = map flag_alternative [ r1', r2', rm ] - where - len_r1 = V.length $ b_seq r1 - len_r2 = V.length $ b_seq r2 - - b_seq_r1 = V.convert $ b_seq r1 - b_seq_r2 = V.convert $ b_seq r2 - b_qual_r1 = V.convert $ b_qual r1 - b_qual_r2 = V.convert $ b_qual r2 - - (mlen, qual1, qual2) = find_merge ads1 ads2 b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 - - flag_alternative br = br { b_exts = updateE "FF" (Int $ extAsInt 0 "FF" br .|. eflagAlternative) $ b_exts br } - store_quals br = br { b_exts = updateE "YM" (Int qual1) $ updateE "YN" (Int qual2) $ b_exts br } - pair_flags = flagPaired.|.flagProperlyPaired.|.flagMateUnmapped.|.flagMateReversed.|.flagFirstMate.|.flagSecondMate - - r1' = store_quals r1 - r2' = store_quals r2 - rm = store_quals $ merged_read mlen (fromIntegral $ min 63 qual1) - - merged_read l qmax = nullBamRec { - b_qname = b_qname r1, - b_flag = flagUnmapped .|. complement pair_flags .&. b_flag r1, - b_seq = V.convert $ merged_seq l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2, - b_qual = V.convert $ merged_qual qmax l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2, - b_exts = let ff = if l < len_r1 then eflagTrimmed else 0 - in updateE "FF" (Int $ extAsInt 0 "FF" r1 .|. eflagMerged .|. ff) $ b_exts r1 } - -{-# INLINE merged_seq #-} -merged_seq :: (V.Vector v Nucleotides, V.Vector v Qual) - => Int -> v Nucleotides -> v Qual -> v Nucleotides -> v Qual -> v Nucleotides -merged_seq l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 = V.concat - [ V.take (l - len_r2) b_seq_r1 - , V.zipWith4 zz (V.take l $ V.drop (l - len_r2) b_seq_r1) - (V.take l $ V.drop (l - len_r2) b_qual_r1) - (V.reverse $ V.take l $ V.drop (l - len_r1) b_seq_r2) - (V.reverse $ V.take l $ V.drop (l - len_r1) b_qual_r2) - , V.reverse $ V.take (l - len_r1) b_seq_r2 ] - where - len_r1 = V.length b_qual_r1 - len_r2 = V.length b_qual_r2 - zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = n1 - | q1 > q2 = n1 - | otherwise = compls n2 - -{-# INLINE merged_qual #-} -merged_qual :: (V.Vector v Nucleotides, V.Vector v Qual) - => Word8 -> Int -> v Nucleotides -> v Qual -> v Nucleotides -> v Qual -> v Qual -merged_qual qmax l b_seq_r1 b_qual_r1 b_seq_r2 b_qual_r2 = V.concat - [ V.take (l - len_r2) b_qual_r1 - , V.zipWith4 zz (V.take l $ V.drop (l - len_r2) b_seq_r1) - (V.take l $ V.drop (l - len_r2) b_qual_r1) - (V.reverse $ V.take l $ V.drop (l - len_r1) b_seq_r2) - (V.reverse $ V.take l $ V.drop (l - len_r1) b_qual_r2) - , V.reverse $ V.take (l - len_r1) b_qual_r2 ] - where - len_r1 = V.length b_qual_r1 - len_r2 = V.length b_qual_r2 - zz !n1 (Q !q1) !n2 (Q !q2) | n1 == compls n2 = Q $ min qmax (q1 + q2) - | q1 > q2 = Q $ q1 - q2 - | otherwise = Q $ q2 - q1 - - - --- | Finds the trimming point. Input is list of forward adapters, --- sequence, quality; output is trim point and two qualities (YM, YN). -find_trim :: [W.Vector Nucleotides] - -> W.Vector Nucleotides -> W.Vector Qual - -> (Int, Int, Int) -find_trim ads1 r1 q1 = (mlen, score2 - score1, plain_score - score1) - where - plain_score = 6 * fromIntegral (V.length r1) - (score1, mlen, score2) = twoMins plain_score (V.length r1) $ - merge_score ads1 [V.empty] r1 q1 V.empty V.empty - --- | Trimming for a single read: we need one adapter only (the one coming --- /after/ the read), here provided as a list of options, and then we --- merge with an empty second read. Results in up to two reads (the --- original, possibly flagged, and the trimmed one, definitely flagged, --- and two qualities). -trimBam :: Int -> Int -> [W.Vector Nucleotides] -> BamRec -> [BamRec] -trimBam lowq highq ads1 r1 - | V.null (b_seq r1) = [ ] - | mlen == 0 && qual1 >= highq = [ ] - | qual1 < lowq || mlen < 0 = [ r1' ] - | qual1 >= highq = [ r1t ] - | otherwise = map flag_alternative [ r1', r1t ] - where - -- the "merge" score if there is no trimming - - b_seq_r1 = V.convert $ b_seq r1 - b_qual_r1 = V.convert $ b_qual r1 - - (mlen, qual1, qual2) = find_trim ads1 b_seq_r1 b_qual_r1 - - flag_alternative br = br { b_exts = updateE "FF" (Int $ extAsInt 0 "FF" br .|. eflagAlternative) $ b_exts br } - store_quals br = br { b_exts = updateE "YM" (Int qual1) $ updateE "YN" (Int qual2) $ b_exts br } - - r1' = store_quals r1 - r1t = store_quals $ trimmed_read mlen - - trimmed_read l = nullBamRec { - b_qname = b_qname r1, - b_flag = flagUnmapped .|. b_flag r1, - b_seq = V.take l $ b_seq r1, - b_qual = V.take l $ b_qual r1, - b_exts = updateE "FF" (Int $ extAsInt 0 "FF" r1 .|. eflagTrimmed) $ b_exts r1 } - - --- | For merging, we don't need the complete adapters (length around 70!), --- only a sufficient prefix. Taking only the more-or-less constant --- part (length around 30), there aren't all that many different --- adapters in the world. To deal with pretty much every library, we --- only need the following forward adapters, which will be the default --- (defined here in the direction they would be sequenced in): Genomic --- R2, Multiplex R2, Fraft P7. - -default_fwd_adapters :: [ W.Vector Nucleotides ] -default_fwd_adapters = map (W.fromList. map toNucleotides) - [ {- Genomic R2 -} "AGATCGGAAGAGCGGTTCAG" - , {- Multiplex R2 -} "AGATCGGAAGAGCACACGTC" - , {- Graft P7 -} "AGATCGGAAGAGCTCGTATG" ] - --- | Like 'default_rev_adapters', these are the few adapters needed for --- the reverse read (defined in the direction they would be sequenced in --- as part of the second read): Genomic R1, CL 72. - -default_rev_adapters :: [ W.Vector Nucleotides ] -default_rev_adapters = map (W.fromList. map toNucleotides) - [ {- Genomic_R1 -} "AGATCGGAAGAGCGTCGTGT" - , {- CL72 -} "GGAAGAGCGTCGTGTAGGGA" ] - --- We need to compute the likelihood of a read pair given an assumed --- insert length. The likelihood of the first read is the likelihood of --- a match with the adapter where it overlaps the 3' adapter, elsewhere --- it's 1/4 per position. The likelihood of the second read is the --- likelihood of a match with the adapter where it overlaps the adapter, --- the likehood of a read-read match where it overlaps read one, 1/4 per --- position elsewhere. (Yes, this ignores base composition. It doesn't --- matter enough.) - -merge_score - :: [ W.Vector Nucleotides ] -- 3' adapters as they appear in the first read - -> [ W.Vector Nucleotides ] -- 5' adapters as they appear in the second read - -> W.Vector Nucleotides -> W.Vector Qual -- first read, qual - -> W.Vector Nucleotides -> W.Vector Qual -- second read, qual - -> Int -- assumed insert length - -> Int -- score (roughly sum of qualities at mismatches) -merge_score fwd_adapters rev_adapters !read1 !qual1 !read2 !qual2 !l - = 6 * fromIntegral (l `min` V.length read1) -- read1, part before adapter - + 6 * fromIntegral (max 0 (l - V.length read1)) -- read2, part before overlap - - + foldl' (\acc fwd_ad -> min acc - (match_adapter l read1 qual1 fwd_ad + -- read1, match with forward adapter - 6 * fromIntegral (max 0 (V.length read1 - V.length fwd_ad - l))) -- read1, part after (known) adapter - ) maxBound fwd_adapters - - + foldl' (\acc rev_ad -> min acc - (match_adapter l read2 qual2 rev_ad + -- read2, match with reverse adapter - 6 * fromIntegral (max 0 (V.length read2 - V.length rev_ad - l))) -- read2, part after (known) adapter - ) maxBound rev_adapters - - + match_reads l read1 qual1 read2 qual2 - -{-# INLINE match_adapter #-} -match_adapter :: Int -> W.Vector Nucleotides -> W.Vector Qual -> W.Vector Nucleotides -> Int -match_adapter !off !rd !qs !ad - | V.length rd /= V.length qs = error "read/qual length mismatch" - | efflength <= 0 = 0 - | otherwise - = fromIntegral . unsafePerformIO $ - W.unsafeWith rd $ \p_rd -> - W.unsafeWith qs $ \p_qs -> - W.unsafeWith ad $ \p_ad -> - prim_match_ad (fromIntegral off) - (fromIntegral efflength) - p_rd p_qs p_ad - where - !efflength = (V.length rd - off) `min` V.length ad - -foreign import ccall unsafe "prim_match_ad" - prim_match_ad :: CInt -> CInt - -> Ptr Nucleotides -> Ptr Qual - -> Ptr Nucleotides -> IO CInt - - --- | Computes overlap score for two reads (with qualities) assuming an --- insert length. -{-# INLINE match_reads #-} -match_reads :: Int -> W.Vector Nucleotides -> W.Vector Qual -> W.Vector Nucleotides -> W.Vector Qual -> Int -match_reads !l !rd1 !qs1 !rd2 !qs2 - | V.length rd1 /= V.length qs1 || V.length rd2 /= V.length qs2 = error "read/qual length mismatch" - | efflength <= 0 = 0 - | otherwise - = fromIntegral . unsafePerformIO $ - W.unsafeWith rd1 $ \p_rd1 -> - W.unsafeWith qs1 $ \p_qs1 -> - W.unsafeWith rd2 $ \p_rd2 -> - W.unsafeWith qs2 $ \p_qs2 -> - prim_match_reads (fromIntegral minidx1) - (fromIntegral maxidx2) - (fromIntegral efflength) - p_rd1 p_qs1 p_rd2 p_qs2 - where - -- vec1, forward - !minidx1 = (l - V.length rd2) `max` 0 - -- vec2, backward - !maxidx2 = l `min` V.length rd2 - -- effective length - !efflength = ((V.length rd1 + V.length rd2 - l) `min` l) `max` 0 - - -foreign import ccall unsafe "prim_match_reads" - prim_match_reads :: CInt -> CInt -> CInt - -> Ptr Nucleotides -> Ptr Qual - -> Ptr Nucleotides -> Ptr Qual -> IO CInt - - -{-# INLINE twoMins #-} -twoMins :: (Bounded a, Ord a) => a -> Int -> (Int -> a) -> (a,Int,a) -twoMins a0 imax f = go a0 (-1) maxBound 0 0 - where - go !m1 !i1 !m2 !i2 !i - | i == imax = (m1,i1,m2) - | otherwise = - case f i of - x | x < m1 -> go x i m1 i1 (i+1) - | x < m2 -> go m1 i1 x i (i+1) - | otherwise -> go m1 i1 m2 i2 (i+1) - - -mergeTrimBam :: Monad m => Int -> Int -> [W.Vector Nucleotides] -> [W.Vector Nucleotides] -> Enumeratee [BamRec] [BamRec] m a -mergeTrimBam lowq highq fwd_ads rev_ads = convStream go - where - go = do r1 <- headStream - if isPaired r1 - then tryHead >>= go2 r1 - else return $ trimBam lowq highq fwd_ads r1 - - go2 r1 Nothing = error $ "Lone mate found: " ++ show (b_qname r1) - go2 r1 (Just r2) = return $ mergeBam lowq highq fwd_ads rev_ads r1 r2 - diff --git a/src/Bio/Bam/Writer.hs b/src/Bio/Bam/Writer.hs deleted file mode 100644 --- a/src/Bio/Bam/Writer.hs +++ /dev/null @@ -1,218 +0,0 @@ --- | Printers for BAM and SAM. BAM is properly supported, SAM can be --- piped to standard output. - -module Bio.Bam.Writer ( - IsBamRec(..), - encodeBamWith, - - packBam, - writeBamFile, - writeBamHandle, - pipeBamOutput, - pipeSamOutput - ) where - -import Bio.Bam.Header -import Bio.Bam.Rec -import Bio.Iteratee -import Bio.Iteratee.Builder -import Bio.Prelude - -import Data.ByteString.Builder ( hPutBuilder, Builder, toLazyByteString ) -import Data.ByteString.Internal ( ByteString(..) ) -import Data.ByteString.Lazy ( foldrChunks ) -import Foreign.Marshal.Alloc ( alloca ) -import System.IO ( openBinaryFile, IOMode(..) ) - -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as S -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as U -import qualified Data.Sequence as Z - --- | write in SAM format to stdout --- This is useful for piping to other tools (say, AWK scripts) or for --- debugging. No convenience function to send SAM to a file exists, --- because that's a stupid idea. -pipeSamOutput :: MonadIO m => BamMeta -> Iteratee [BamRec] m () -pipeSamOutput meta = do liftIO . hPutBuilder stdout $ showBamMeta meta - mapStreamM_ $ \b -> liftIO . putStr $ encodeSamEntry (meta_refs meta) b "\n" - -encodeSamEntry :: Refs -> BamRec -> String -> String -encodeSamEntry refs b = conjoin '\t' [ - unpck (b_qname b), - shows (b_flag b .&. 0xffff), - unpck (sq_name $ getRef refs $ b_rname b), - shows (b_pos b + 1), - shows (unQ $ b_mapq b), - V.foldr ((.) . shows) id (b_cigar b), - if isValidRefseq (b_mrnm b) && b_mrnm b == b_rname b - then (:) '=' else unpck (sq_name $ getRef refs $ b_mrnm b), - shows (b_mpos b + 1), - shows (b_isize b), - shows (V.toList $ b_seq b), - (++) (V.toList . V.map (chr . (+33) . fromIntegral . unQ) $ b_qual b) ] . - foldr (\(k,v) f -> (:) '\t' . shows k . (:) ':' . extToSam v . f) id (b_exts b) - where - unpck = (++) . S.unpack - conjoin c = foldr1 (\a f -> a . (:) c . f) - - extToSam (Int i) = (:) 'i' . (:) ':' . shows i - extToSam (Float f) = (:) 'f' . (:) ':' . shows f - extToSam (Text t) = (:) 'Z' . (:) ':' . unpck t - extToSam (Bin x) = (:) 'H' . (:) ':' . tohex x - extToSam (Char c) = (:) 'A' . (:) ':' . (:) (w2c c) - extToSam (IntArr a) = (:) 'B' . (:) ':' . (:) 'i' . sarr a - extToSam (FloatArr a) = (:) 'B' . (:) ':' . (:) 'f' . sarr a - - tohex = B.foldr (\c f -> w2d (c `shiftR` 4) . w2d (c .&. 0xf) . f) id - w2d = (:) . S.index "0123456789ABCDEF" . fromIntegral - sarr v = conjoin ',' . map shows $ U.toList v - -class IsBamRec a where - pushBam :: a -> BgzfTokens -> BgzfTokens - -instance IsBamRec BamRaw where - {-# INLINE pushBam #-} - pushBam = pushBamRaw - -instance IsBamRec BamRec where - {-# INLINE pushBam #-} - pushBam = pushBamRec - -instance (IsBamRec a, IsBamRec b) => IsBamRec (Either a b) where - {-# INLINE pushBam #-} - pushBam = either pushBam pushBam - --- | Encodes BAM records straight into a dynamic buffer, the BGZF's it. --- Should be fairly direct and perform well. -{-# INLINE encodeBamWith #-} -encodeBamWith :: (MonadIO m, IsBamRec r) => Int -> BamMeta -> Enumeratee [r] S.ByteString m () -encodeBamWith lv meta = eneeBam ><> encodeBgzf lv - where - eneeBam = eneeCheckIfDone (\k -> mapChunks (foldMap (Endo . pushBam)) . k $ Chunk pushHeader) - - pushHeader :: Endo BgzfTokens - pushHeader = Endo $ TkString "BAM\1" - . TkSetMark -- the length byte - . pushBuilder (showBamMeta meta) - . TkEndRecord -- fills the length in - . TkWord32 (fromIntegral . Z.length $ meta_refs meta) - . appEndo (foldMap (Endo . pushRef) (meta_refs meta)) - - pushRef :: BamSQ -> BgzfTokens -> BgzfTokens - pushRef bs = TkWord32 (fromIntegral $ B.length (sq_name bs) + 1) - . TkString (sq_name bs) - . TkWord8 0 - . TkWord32 (fromIntegral $ sq_length bs) - - pushBuilder :: Builder -> BgzfTokens -> BgzfTokens - pushBuilder b tk = foldrChunks TkString tk (toLazyByteString b) - -{-# INLINE pushBamRaw #-} -pushBamRaw :: BamRaw -> BgzfTokens -> BgzfTokens -pushBamRaw r = TkWord32 (fromIntegral $ B.length $ raw_data r) . - TkString (raw_data r) - --- | Writes BAM encoded stuff to a file. -writeBamFile :: IsBamRec r => FilePath -> BamMeta -> Iteratee [r] IO () -writeBamFile fp meta = - bracketIO (openBinaryFile fp WriteMode) - (hClose) - (flip writeBamHandle meta) - --- | write BAM encoded stuff to stdout --- This send uncompressed BAM to stdout. Useful for piping to other --- tools. -pipeBamOutput :: IsBamRec r => BamMeta -> Iteratee [r] IO () -pipeBamOutput meta = encodeBamWith 0 meta =$ mapChunksM_ (liftIO . S.hPut stdout) - --- | Writes BAM encoded stuff to a 'Handle'. -writeBamHandle :: (MonadIO m, IsBamRec r) => Handle -> BamMeta -> Iteratee [r] m () -writeBamHandle hdl meta = encodeBamWith 6 meta =$ mapChunksM_ (liftIO . S.hPut hdl) - -{-# RULES - "pushBam/unpackBam" forall b . pushBamRec (unpackBam b) = pushBamRaw b - #-} - -{-# INLINE[1] pushBamRec #-} -pushBamRec :: BamRec -> BgzfTokens -> BgzfTokens -pushBamRec BamRec{..} = - TkSetMark - . TkWord32 (unRefseq b_rname) - . TkWord32 (fromIntegral b_pos) - . TkWord8 (fromIntegral $ B.length b_qname + 1) - . TkWord8 (unQ b_mapq) - . TkWord16 (fromIntegral bin) - . TkWord16 (fromIntegral $ VS.length b_cigar) - . TkWord16 (fromIntegral b_flag) - . TkWord32 (fromIntegral $ V.length b_seq) - . TkWord32 (unRefseq b_mrnm) - . TkWord32 (fromIntegral b_mpos) - . TkWord32 (fromIntegral b_isize) - . TkString b_qname - . TkWord8 0 - . VS.foldr ((.) . TkWord8) id (VS.unsafeCast b_cigar :: VS.Vector Word8) - . pushSeq b_seq - . VS.foldr ((.) . TkWord8 . unQ) id b_qual - . foldr ((.) . pushExt) id b_exts - . TkEndRecord - where - bin = distinctBin b_pos (alignedLength b_cigar) - - pushSeq :: V.Vector vec Nucleotides => vec Nucleotides -> BgzfTokens -> BgzfTokens - pushSeq v = case v V.!? 0 of - Nothing -> id - Just a -> case v V.!? 1 of - Nothing -> TkWord8 (unNs a `shiftL` 4) - Just b -> TkWord8 (unNs a `shiftL` 4 .|. unNs b) . pushSeq (V.drop 2 v) - - pushExt :: (BamKey, Ext) -> BgzfTokens -> BgzfTokens - pushExt (BamKey k, e) = case e of - Text t -> common 'Z' . TkString t . TkWord8 0 - Bin t -> common 'H' . TkString t . TkWord8 0 - Char c -> common 'A' . TkWord8 c - Float f -> common 'f' . TkWord32 (fromIntegral $ fromFloat f) - - Int i -> case put_some_int (U.singleton i) of - (c,op) -> common c . op i - - IntArr ia -> case put_some_int ia of - (c,op) -> common 'B' . TkWord8 (fromIntegral $ ord c) - . TkWord32 (fromIntegral $ U.length ia-1) - . U.foldr ((.) . op) id ia - - FloatArr fa -> common 'B' . TkWord8 (fromIntegral $ ord 'f') - . TkWord32 (fromIntegral $ U.length fa-1) - . U.foldr ((.) . TkWord32 . fromFloat) id fa - where - common :: Char -> BgzfTokens -> BgzfTokens - common z = TkWord16 k . TkWord8 (fromIntegral $ ord z) - - put_some_int :: U.Vector Int -> (Char, Int -> BgzfTokens -> BgzfTokens) - put_some_int is - | U.all (between 0 0xff) is = ('C', TkWord8 . fromIntegral) - | U.all (between (-0x80) 0x7f) is = ('c', TkWord8 . fromIntegral) - | U.all (between 0 0xffff) is = ('S', TkWord16 . fromIntegral) - | U.all (between (-0x8000) 0x7fff) is = ('s', TkWord16 . fromIntegral) - | U.all (> 0) is = ('I', TkWord32 . fromIntegral) - | otherwise = ('i', TkWord32 . fromIntegral) - - between :: Int -> Int -> Int -> Bool - between l r x = l <= x && x <= r - - fromFloat :: Float -> Word32 - fromFloat float = unsafeDupablePerformIO $ alloca $ \buf -> - pokeByteOff buf 0 float >> peek buf - -packBam :: BamRec -> IO BamRaw -packBam br = do bb <- newBuffer 1000 - (bb', TkEnd) <- store_loop bb (pushBamRec br TkEnd) - return . bamRaw 0 $ PS (buffer bb') 4 (used bb' - 4) - where - store_loop bb tk = do (bb',tk') <- fillBuffer bb tk - case tk' of TkEnd -> return (bb',tk') - _ -> do bb'' <- expandBuffer (128*1024) bb' - store_loop bb'' tk' - diff --git a/src/Bio/Base.hs b/src/Bio/Base.hs deleted file mode 100644 --- a/src/Bio/Base.hs +++ /dev/null @@ -1,369 +0,0 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, TypeFamilies, CPP #-} -{-# LANGUAGE ExistentialQuantification, TemplateHaskell #-} --- | Common data types used everywhere. This module is a collection of --- very basic "bioinformatics" data types that are simple, but don't --- make sense to define over and over. - -module Bio.Base( - Nucleotide(..), Nucleotides(..), - Qual(..), toQual, fromQual, fromQualRaised, probToQual, - Prob'(..), Prob, toProb, fromProb, qualToProb, pow, - - Word8, - nucA, nucC, nucG, nucT, - nucsA, nucsC, nucsG, nucsT, nucsN, gap, - toNucleotide, toNucleotides, nucToNucs, - showNucleotide, showNucleotides, - isGap, - isBase, - isProperBase, - properBases, - compl, compls, - - Seqid, - - Position(..), - shiftPosition, - p_is_reverse, - - Range(..), - shiftRange, - reverseRange, - extendRange, - insideRange, - wrapRange, - - w2c, - c2w, - - findAuxFile -) where - -import BasePrelude -#if MIN_VERSION_base(4,9,0) - hiding ( log1pexp, log1mexp ) -#endif -import Bio.Util.Numeric ( log1pexp, log1mexp ) -import Data.ByteString.Internal ( c2w, w2c ) -import Data.Vector.Unboxed.Deriving ( derivingUnbox ) -import System.Posix.Files ( fileExist ) - -import qualified Data.ByteString.Char8 as S -import qualified Data.Vector.Unboxed as U - --- | A nucleotide base. We only represent A,C,G,T. The contained --- 'Word8' ist guaranteed to be 0..3. -newtype Nucleotide = N { unN :: Word8 } deriving ( Eq, Ord, Enum, Ix, Storable ) - -derivingUnbox "Nucleotide" [t| Nucleotide -> Word8 |] [| unN |] [| N |] - -instance Bounded Nucleotide where - minBound = N 0 - maxBound = N 3 - --- | A nucleotide base in an alignment. --- Experience says we're dealing with Ns and gaps all the type, so --- purity be damned, they are included as if they were real bases. --- --- To allow @Nucleotides@s to be unpacked and incorporated into --- containers, we choose to represent them the same way as the BAM file --- format: as a 4 bit wide field. Gaps are encoded as 0 where they --- make sense, N is 15. The contained 'Word8' is guaranteed to be --- 0..15. - -newtype Nucleotides = Ns { unNs :: Word8 } deriving ( Eq, Ord, Enum, Ix, Storable ) - -derivingUnbox "Nucleotides" [t| Nucleotides -> Word8 |] [| unNs |] [| Ns |] - -instance Bounded Nucleotides where - minBound = Ns 0 - maxBound = Ns 15 - -nucToNucs :: Nucleotide -> Nucleotides -nucToNucs (N x) = Ns $ 1 `shiftL` fromIntegral (x .&. 3) - --- | Qualities are stored in deciban, also known as the Phred scale. To --- represent a value @p@, we store @-10 * log_10 p@. Operations work --- directly on the \"Phred\" value, as the name suggests. The same goes --- for the 'Ord' instance: greater quality means higher \"Phred\" --- score, meand lower error probability. - -newtype Qual = Q { unQ :: Word8 } deriving ( Eq, Ord, Storable, Bounded ) - -derivingUnbox "Qual" [t| Qual -> Word8 |] [| unQ |] [| Q |] - -instance Show Qual where - showsPrec p (Q q) = (:) 'q' . showsPrec p q - -toQual :: (Floating a, RealFrac a) => a -> Qual -toQual a = Q $ round (-10 * log a / log 10) - -fromQual :: Qual -> Double -fromQual (Q q) = 10 ** (- fromIntegral q / 10) - -fromQualRaised :: Double -> Qual -> Double -fromQualRaised k (Q q) = 10 ** (- k * fromIntegral q / 10) - --- | A positive floating point value stored in log domain. We store the --- natural logarithm (makes computation easier), but allow conversions --- to the familiar \"Phred\" scale used for 'Qual' values. -newtype Prob' a = Pr { unPr :: a } deriving ( Eq, Ord, Storable ) - --- | Common way of using 'Prob''. -type Prob = Prob' Double - -derivingUnbox "Prob'" [t| forall a . U.Unbox a => Prob' a -> a |] [| unPr |] [| Pr |] - -instance RealFloat a => Show (Prob' a) where - showsPrec _ (Pr p) = (:) 'q' . showFFloat (Just 1) q - where q = - 10 * p / log 10 - -instance (Floating a, Ord a) => Num (Prob' a) where - {-# INLINE fromInteger #-} - fromInteger a = Pr (log (fromInteger a)) - {-# INLINE (+) #-} - Pr x + Pr y = Pr $ if x >= y then x + log1pexp (y-x) else y + log1pexp (x-y) - {-# INLINE (-) #-} - Pr x - Pr y = Pr $ if x >= y then x + log1mexp (y-x) else error "no negative error probabilities" - {-# INLINE (*) #-} - Pr a * Pr b = Pr $ a + b - {-# INLINE negate #-} - negate _ = Pr $ error "no negative error probabilities" - {-# INLINE abs #-} - abs x = x - {-# INLINE signum #-} - signum _ = Pr 0 - -instance (Floating a, Fractional a, Ord a) => Fractional (Prob' a) where - fromRational a = Pr (log (fromRational a)) - Pr a / Pr b = Pr (a - b) - recip (Pr a) = Pr (negate a) - -infixr 8 `pow` -pow :: Num a => Prob' a -> a -> Prob' a -pow (Pr a) e = Pr $ a * e - - -toProb :: Floating a => a -> Prob' a -toProb p = Pr (log p) - -fromProb :: Floating a => Prob' a -> a -fromProb (Pr q) = exp q - -qualToProb :: Floating a => Qual -> Prob' a -qualToProb (Q q) = Pr (- log 10 * fromIntegral q / 10) - -probToQual :: (Floating a, RealFrac a) => Prob' a -> Qual -probToQual (Pr p) = Q (round (- 10 * p / log 10)) - -nucA, nucC, nucG, nucT :: Nucleotide -nucA = N 0 -nucC = N 1 -nucG = N 2 -nucT = N 3 - -gap, nucsA, nucsC, nucsG, nucsT, nucsN :: Nucleotides -gap = Ns 0 -nucsA = Ns 1 -nucsC = Ns 2 -nucsG = Ns 4 -nucsT = Ns 8 -nucsN = Ns 15 - - --- | Sequence identifiers are ASCII strings --- Since we tend to store them for a while, we use strict byte strings. -type Seqid = S.ByteString - --- | Coordinates in a genome. --- The position is zero-based, no questions about it. Think of the --- position as pointing to the crack between two bases: looking forward --- you see the next base to the right, looking in the reverse direction --- you see the complement of the first base to the left. --- --- To encode the strand, we (virtually) reverse-complement any sequence --- and prepend it to the normal one. That way, reversed coordinates --- have a negative sign and automatically make sense. Position 0 could --- either be the beginning of the sequence or the end on the reverse --- strand... that ambiguity shouldn't really matter. - -data Position = Pos { - p_seq :: {-# UNPACK #-} !Seqid, -- ^ sequence (e.g. some chromosome) - p_start :: {-# UNPACK #-} !Int -- ^ offset, zero-based - } deriving (Show, Eq, Ord) - -p_is_reverse :: Position -> Bool -p_is_reverse = (< 0) . p_start - --- | Ranges in genomes --- We combine a position with a length. In 'Range pos len', 'pos' is --- always the start of a stretch of length 'len'. Positions therefore --- move in the opposite direction on the reverse strand. To get the --- same stretch on the reverse strand, shift r_pos by r_length, then --- reverse direction (or call reverseRange). -data Range = Range { - r_pos :: {-# UNPACK #-} !Position, - r_length :: {-# UNPACK #-} !Int - } deriving (Show, Eq, Ord) - - --- | Converts a character into a 'Nucleotides'. --- The usual codes for A,C,G,T and U are understood, '-' and '.' become --- gaps and everything else is an N. -toNucleotide :: Char -> Nucleotide -toNucleotide c = if ord c < 128 then N (ar `U.unsafeIndex` ord c) else N 0 - where - ar = U.replicate 128 0 U.// - ( [ (ord x, n) | (x, N n) <- pairs ] ++ - [ (ord (toUpper x), n) | (x, N n) <- pairs ] ) - - pairs = [ ('a', nucA), ('c', nucC), ('g', nucG), ('t', nucT) ] - - --- | Converts a character into a 'Nucleotides'. --- The usual codes for A,C,G,T and U are understood, '-' and '.' become --- gaps and everything else is an N. -toNucleotides :: Char -> Nucleotides -toNucleotides c = if ord c < 128 then Ns (ar `U.unsafeIndex` ord c) else nucsN - where - ar = U.replicate 128 (unNs nucsN) U.// - ( [ (ord x, n) | (x, Ns n) <- pairs ] ++ - [ (ord (toUpper x), n) | (x, Ns n) <- pairs ] ) - - Ns a `plus` Ns b = Ns (a .|. b) - - pairs = [ ('a', nucsA), ('c', nucsC), ('g', nucsG), ('t', nucsT), - ('u', nucsT), ('-', gap), ('.', gap), - ('b', nucsC `plus` nucsG `plus` nucsT), - ('d', nucsA `plus` nucsG `plus` nucsT), - ('h', nucsA `plus` nucsC `plus` nucsT), - ('v', nucsA `plus` nucsC `plus` nucsG), - ('k', nucsG `plus` nucsT), - ('m', nucsA `plus` nucsC), - ('s', nucsC `plus` nucsG), - ('w', nucsA `plus` nucsT), - ('r', nucsA `plus` nucsG), - ('y', nucsC `plus` nucsT) ] - --- | Tests if a 'Nucleotides' is a base. --- Returns 'True' for everything but gaps. -isBase :: Nucleotides -> Bool -isBase (Ns n) = n /= 0 - --- | Tests if a 'Nucleotides' is a proper base. --- Returns 'True' for A,C,G,T only. -isProperBase :: Nucleotides -> Bool -isProperBase x = x == nucsA || x == nucsC || x == nucsG || x == nucsT - -properBases :: [ Nucleotides ] -properBases = [ nucsA, nucsC, nucsG, nucsT ] - --- | Tests if a 'Nucleotides' is a gap. --- Returns true only for the gap. -isGap :: Nucleotides -> Bool -isGap x = x == gap - - -{-# INLINE showNucleotide #-} -showNucleotide :: Nucleotide -> Char -showNucleotide (N x) = S.index str $ fromIntegral $ x .&. 3 - where str = S.pack "ACGT" - -{-# INLINE showNucleotides #-} -showNucleotides :: Nucleotides -> Char -showNucleotides (Ns x) = S.index str $ fromIntegral $ x .&. 15 - where str = S.pack "-ACMGRSVTWYHKDBN" - -instance Show Nucleotide where - show x = [ showNucleotide x ] - showList l = (map showNucleotide l ++) - -instance Read Nucleotide where - readsPrec _ ('a':cs) = [(nucA, cs)] - readsPrec _ ('A':cs) = [(nucA, cs)] - readsPrec _ ('c':cs) = [(nucC, cs)] - readsPrec _ ('C':cs) = [(nucC, cs)] - readsPrec _ ('g':cs) = [(nucG, cs)] - readsPrec _ ('G':cs) = [(nucG, cs)] - readsPrec _ ('t':cs) = [(nucT, cs)] - readsPrec _ ('T':cs) = [(nucT, cs)] - readsPrec _ ('u':cs) = [(nucT, cs)] - readsPrec _ ('U':cs) = [(nucT, cs)] - readsPrec _ _ = [ ] - - readList ('-':cs) = readList cs - readList (c:cs) | isSpace c = readList cs - | otherwise = case reads (c:cs) of - [] -> [ ([],c:cs) ] - xs -> [ (n:ns,r2) | (n,r1) <- xs, (ns,r2) <- readList r1 ] - readList [] = [([],[])] - -instance Show Nucleotides where - show x = [ showNucleotides x ] - showList l = (map showNucleotides l ++) - -instance Read Nucleotides where - readsPrec _ (c:cs) = [(toNucleotides c, cs)] - readsPrec _ [ ] = [] - readList s = let (hd,tl) = span (\c -> isAlpha c || isSpace c || '-' == c) s - in [(map toNucleotides $ filter (not . isSpace) hd, tl)] - --- | Complements a Nucleotides. -{-# INLINE compl #-} -compl :: Nucleotide -> Nucleotide -compl (N n) = N $ n `xor` 3 - --- | Complements a Nucleotides. -{-# INLINE compls #-} -compls :: Nucleotides -> Nucleotides -compls (Ns x) = Ns $ ar `U.unsafeIndex` fromIntegral (x .&. 15) - where - !ar = U.fromListN 16 [ 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15 ] - - --- | Moves a @Position@. The position is moved forward according to the --- strand, negative indexes move backward accordingly. -shiftPosition :: Int -> Position -> Position -shiftPosition a p = p { p_start = p_start p + a } - --- | Moves a @Range@. This is just @shiftPosition@ lifted. -shiftRange :: Int -> Range -> Range -shiftRange a r = r { r_pos = shiftPosition a (r_pos r) } - --- | Reverses a 'Range' to give the same @Range@ on the opposite strand. -reverseRange :: Range -> Range -reverseRange (Range (Pos sq pos) len) = Range (Pos sq (-pos-len)) len - --- | Extends a range. The length of the range is simply increased. -extendRange :: Int -> Range -> Range -extendRange a r = r { r_length = r_length r + a } - --- | Expands a subrange. --- @(range1 `insideRange` range2)@ interprets @range1@ as a subrange of --- @range2@ and computes its absolute coordinates. The sequence name of --- @range1@ is ignored. -insideRange :: Range -> Range -> Range -insideRange r1@(Range (Pos _ start1) length1) r2@(Range (Pos sq start2) length2) - | start2 < 0 = reverseRange (insideRange r1 (reverseRange r2)) - | start1 <= length2 = Range (Pos sq (start2 + start1)) (min length1 (length2 - start1)) - | otherwise = Range (Pos sq (start2 + length2)) 0 - - --- | Wraps a range to a region. This simply normalizes the start --- position to be in the interval '[0,n)', which only makes sense if the --- @Range@ is to be mapped onto a circular genome. This works on both --- strands and the strand information is retained. -wrapRange :: Int -> Range -> Range -wrapRange n (Range (Pos sq s) l) = Range (Pos sq (s `mod` n)) l - --- | Finds a file by searching the environment variable BIOHAZARD like a --- PATH. -findAuxFile :: FilePath -> IO FilePath -findAuxFile fn | "/" `isPrefixOf` fn = return fn - | otherwise = go . fromMaybe "." . lookup "BIOHAZARD" =<< getEnvironment - where - go "" = return fn - go pp = let (p,ps) = break (==':') pp - in fileExist (p ++ "/" ++ fn) >>= - bool (return $ p ++ "/" ++ fn) (go $ drop 1 ps) - diff --git a/src/Bio/Iteratee.hs b/src/Bio/Iteratee.hs deleted file mode 100644 --- a/src/Bio/Iteratee.hs +++ /dev/null @@ -1,318 +0,0 @@ -module Bio.Iteratee ( - iGetString, - iterLoop, - iLookAhead, - - protectTerm, - parMapChunksIO, - parRunIO, - progressGen, - progressNum, - progressPos, - - ($==), - MonadIO, MonadMask, - lift, liftIO, - stdin, stdout, stderr, - - enumAuxFile, - enumInputs, - enumDefaultInputs, - - Ordering'(..), - mergeSortStreams, - - Enumerator', - Enumeratee', - mergeEnums', - - QQ(..), - emptyQ, - lengthQ, - pushQ, - popQ, - cancelAll, - - ParseError(..), - parserToIteratee, - stream2vector, - stream2vectorN, - - Fd, - withFileFd, - - module Bio.Iteratee.Bytes, - module Bio.Iteratee.IO, - module Bio.Iteratee.Iteratee, - module Bio.Iteratee.List - ) where - -import Bio.Bam.Header -import Bio.Iteratee.Base -import Bio.Iteratee.Bytes -import Bio.Iteratee.IO -import Bio.Iteratee.Iteratee -import Bio.Iteratee.List -import Bio.Prelude -import Bio.Util.Numeric ( showNum ) -import Control.Concurrent.Async ( Async, async, wait, cancel ) -import Control.Monad.Catch ( MonadMask(..) ) -import Control.Monad.IO.Class ( MonadIO(..) ) -import Control.Monad.Trans.Class ( MonadTrans(..) ) -import System.IO ( hIsTerminalDevice ) - -import qualified Control.Monad.Catch as CMC -import qualified Data.Attoparsec.ByteString as A -import qualified Data.ByteString.Char8 as S -import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Generic.Mutable as VM - --- | Run an Iteratee, collect the input. When it finishes, return the --- result along with *all* input. Effectively allows lookahead. Be --- careful, this will eat memory if the @Iteratee@ doesn't return --- speedily. -iLookAhead :: Monoid s => Iteratee s m a -> Iteratee s m a -iLookAhead = go mempty - where - go acc it = Iteratee $ \od oc -> runIter it (\x _ -> od x (Chunk acc)) (oc . step acc) - - step acc k c@(Chunk str) = go (acc `mappend` str) (k c) - step acc k c@(EOF _) = Iteratee $ \od1 -> runIter (k c) (\x _ -> od1 x (Chunk acc)) - - --- | Collects a string of a given length. Don't use this for long --- strings, use 'takeStream' instead. -iGetString :: Int -> Iteratee S.ByteString m S.ByteString -iGetString 0 = idone S.empty (Chunk S.empty) -iGetString n = liftI $ step [] 0 - where - step acc l c@(EOF _) = icont (step acc l) (Just $ setEOF c) - step acc l (Chunk c) | l + S.length c >= n = let r = S.concat . reverse $ S.take (n-l) c : acc - in idone r (Chunk $ S.drop (n-l) c) - | otherwise = liftI $ step (c:acc) (l + S.length c) - --- | Repeatedly apply an 'Iteratee' to a value until end of stream. --- Returns the final value. -iterLoop :: (Nullable s, Monad m) => (a -> Iteratee s m a) -> a -> Iteratee s m a -iterLoop it a = do e <- isFinished - if e then return a - else it a >>= iterLoop it -infixl 1 $== -{-# INLINE ($==) #-} --- | Compose an 'Enumerator'' with an 'Enumeratee', giving a new --- 'Enumerator''. -($==) :: Monad m => Enumerator' hdr input m (Iteratee output m result) - -> Enumeratee input output m result - -> Enumerator' hdr output m result -($==) enum enee iter = run =<< enum (enee . iter) - --- | Merge two 'Enumerator''s into one. The header provided by the --- inner 'Enumerator'' is passed to the output iterator, the header --- provided by the outer 'Enumerator'' is passed to the merging iteratee - -{-# INLINE mergeEnums' #-} -mergeEnums' :: (Nullable s2, Nullable s1, Monad m) - => Enumerator' hi s1 m a -- ^ inner enumerator - -> Enumerator' ho s2 (Iteratee s1 m) a -- ^ outer enumerator - -> (ho -> Enumeratee s2 s1 (Iteratee s1 m) a) -- ^ merging enumeratee - -> Enumerator' hi s1 m a -mergeEnums' e1 e2 etee i = e1 $ \hi -> e2 (\ho -> joinI . etee ho $ ilift lift (i hi)) >>= run - -type Enumerator' h eo m b = (h -> Iteratee eo m b) -> m (Iteratee eo m b) -type Enumeratee' h ei eo m b = (h -> Iteratee eo m b) -> Iteratee ei m (Iteratee eo m b) - -enumAuxFile :: MonadBracketIO m => FilePath -> Iteratee S.ByteString m a -> m a -enumAuxFile fp it = liftIO (findAuxFile fp) >>= \f -> enumFile defaultBufSize f it >>= run - -enumDefaultInputs :: MonadBracketIO m => Enumerator S.ByteString m a -enumDefaultInputs it0 = liftIO getArgs >>= flip enumInputs it0 - -enumInputs :: MonadBracketIO m => [FilePath] -> Enumerator S.ByteString m a -enumInputs [] = enumFd defaultBufSize stdInput -enumInputs xs = go xs - where go ("-":fs) = enumFd defaultBufSize stdInput >=> go fs - go ( f :fs) = enumFile defaultBufSize f >=> go fs - go [ ] = return - -data Ordering' a = Less | Equal a | NotLess - -mergeSortStreams :: Monad m => (a -> a -> Ordering' a) -> Enumeratee [a] [a] (Iteratee [a] m) b -mergeSortStreams comp = eneeCheckIfDone step - where - step out = peekStream >>= \mx -> lift peekStream >>= \my -> case (mx, my) of - (Just x, Just y) -> case x `comp` y of - Less -> do dropStream 1 ; eneeCheckIfDone step . out $ Chunk [x] - NotLess -> do lift (dropStream 1) ; eneeCheckIfDone step . out $ Chunk [y] - Equal z -> do dropStream 1 ; lift (dropStream 1) ; eneeCheckIfDone step . out $ Chunk [z] - - (Just x, Nothing) -> do dropStream 1 ; eneeCheckIfDone step . out $ Chunk [x] - (Nothing, Just y) -> do lift (dropStream 1) ; eneeCheckIfDone step . out $ Chunk [y] - (Nothing, Nothing) -> idone (liftI out) $ EOF Nothing - - --- | Parallel map of an IO action over the elements of a stream --- --- This 'Enumeratee' applies an 'IO' action to every chunk of the input --- stream. These 'IO' actions are run asynchronously in a limited --- parallel way. Don't forget to `evaluate` - -parMapChunksIO :: (MonadIO m, Nullable s) => Int -> (s -> IO t) -> Enumeratee s t m a -parMapChunksIO np f = eneeCheckIfDonePass (go emptyQ) - where - -- check if the queue is full - go !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) - go !qq k Nothing = case popQ qq of - Just (a,qq') | lengthQ qq == np -> liftIO (wait a) >>= eneeCheckIfDonePass (go qq') . k . Chunk - _ -> liftI $ go' qq k - - -- we have room for input - go' !qq k (EOF mx) = do a <- liftIO (async (f emptyP)) - goE mx (pushQ a qq) k Nothing - go' !qq k (Chunk c) = do a <- liftIO (async (f c)) - go (pushQ a qq) k Nothing - - -- input ended, empty the queue - goE _ !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) - goE mx !qq k Nothing = case popQ qq of - Nothing -> idone (liftI k) (EOF mx) - Just (a,qq') -> liftIO (wait a) >>= eneeCheckIfDonePass (goE mx qq') . k . Chunk - -parRunIO :: MonadIO m => Int -> Enumeratee [IO a] a m b -parRunIO np = eneeCheckIfDonePass (go emptyQ) - where - -- check if the queue is full - go !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) - go !qq k Nothing = case popQ qq of - Just (a,qq') | lengthQ qq == np -> liftIO (wait a) >>= eneeCheckIfDonePass (go qq') . k . Chunk - _ -> liftI $ go' qq k - - -- we have room for input - go' !qq k (Chunk (c:cs)) = liftIO (async c) >>= \a -> go' (pushQ a qq) k (Chunk cs) - go' !qq k (Chunk [ ]) = go qq k Nothing - go' !qq k (EOF mx) = goE mx qq k Nothing - - -- input ended, empty the queue - goE _ !qq k (Just e) = cancelAll qq >> icont (go' emptyQ k) (Just e) - goE mx !qq k Nothing = case popQ qq of - Nothing -> idone (liftI k) (EOF mx) - Just (a,qq') -> liftIO (wait a) >>= eneeCheckIfDonePass (goE mx qq') . k . Chunk - --- | Protects the terminal from binary junk. If @i@ is an 'Iteratee' --- that might write binary to 'stdout', then @protectTerm i@ is the same --- 'Iteratee', but it will abort if 'stdout' is a terminal device. -protectTerm :: (Nullable s, MonadIO m) => Iteratee s m a -> Iteratee s m a -protectTerm itr = do - t <- liftIO $ hIsTerminalDevice stdout - if t then err else itr - where - err = error "cowardly refusing to write binary data to terminal" - --- | A general progress indicator that prints some message after a set --- number of records have passed through. -progressGen :: MonadIO m - => (Int -> a -> String) -> Int -> (String -> IO ()) -> Enumeratee [a] [a] m b -progressGen msg sz put = eneeCheckIfDonePass (icont . go 0) - where - go !_ k (EOF mx) = idone (liftI k) (EOF mx) - go !n k (Chunk as) - | null as = liftI $ go n k - | otherwise = let !n' = n + length as - in when (n' `div` sz /= n `div` sz) (liftIO . put $ - "\27[K" ++ msg n' (head as) ++ "\r") - `ioBind_` eneeCheckIfDonePass (icont . go n') (k $ Chunk as) - --- | A simple progress indicator that prints the number of records. -progressNum :: MonadIO m - => String -> Int -> (String -> IO ()) -> Enumeratee [a] [a] m b -progressNum msg = progressGen (\n _ -> msg ++ " " ++ showNum n) - --- | A simple progress indicator that prints a position every set number --- of passed records. -progressPos :: MonadIO m - => (a -> (Refseq, Int)) -> String -> Refs - -> Int -> (String -> IO ()) -> Enumeratee [a] [a] m b -progressPos f msg refs = - progressGen $ \_ a -> let (!rs1, !po1) = f a - !nm = unpack . sq_name $ getRef refs rs1 - in msg ++ " " ++ nm ++ ":" ++ showNum po1 - --- A very simple queue data type. --- Invariants: q = QQ l f b --> l == length f + length b --- --> l == 0 || not (null f) - -data QQ a = QQ !Int [a] [a] - -emptyQ :: QQ a -emptyQ = QQ 0 [] [] - -lengthQ :: QQ a -> Int -lengthQ (QQ l _ _) = l - -pushQ :: a -> QQ a -> QQ a -pushQ a (QQ l [] b) = QQ (l+1) (reverse (a:b)) [] -pushQ a (QQ l f b) = QQ (l+1) f (a:b) - -popQ :: QQ a -> Maybe (a, QQ a) -popQ (QQ _ [ ] _) = Nothing -popQ (QQ l [ a ] b) = Just (a, QQ (l-1) (reverse b) []) -popQ (QQ l (a:fs) b) = Just (a, QQ (l-1) fs b) - -cancelAll :: MonadIO m => QQ (Async a) -> m () -cancelAll (QQ _ ff bb) = liftIO $ mapM_ cancel (ff ++ bb) - -data ParseError = ParseError {errorContexts :: [String], errorMessage :: String} - deriving (Show, Typeable) - -instance Exception ParseError - --- | A function to convert attoparsec 'Parser's into 'Iteratee's. -parserToIteratee :: A.Parser a -> Iteratee S.ByteString m a -parserToIteratee p = icont (f (A.parse p)) Nothing - where - f k (EOF Nothing) = - case A.feed (k S.empty) S.empty of - A.Fail _ err dsc -> throwErr (toException $ ParseError err dsc) - A.Partial _ -> throwErr (toException EofException) - A.Done rest v | S.null rest -> idone v (EOF Nothing) - | otherwise -> idone v (Chunk rest) - f _ (EOF (Just e)) = throwErr e - f k (Chunk s) - | S.null s = icont (f k) Nothing - | otherwise = - case k s of - A.Fail _ err dsc -> throwErr (toException $ ParseError err dsc) - A.Partial k' -> icont (f k') Nothing - A.Done rest v -> idone v (Chunk rest) - - --- | Equivalent to @joinI $ takeStream n $ stream2vector@, but more --- efficient. -stream2vectorN :: (MonadIO m, VG.Vector v a) => Int -> Iteratee [a] m (v a) -stream2vectorN n = do - mv <- liftIO $ VM.new n - l <- go mv 0 - liftIO $ VG.unsafeFreeze $ VM.take l mv - where - go mv i - | i == n = return n - | otherwise = - tryHead >>= \case - Nothing -> return i - Just a -> liftIO (VM.write mv i a) >> go mv (i+1) - --- | Reads the whole stream into a 'VG.Vector'. -stream2vector :: (MonadIO m, VG.Vector v a) => Iteratee [a] m (v a) -stream2vector = liftIO (VM.new 1024) >>= go 0 - where - go !i !mv = tryHead >>= \case - Nothing -> liftIO $ VG.unsafeFreeze $ VM.take i mv - Just a -> do mv' <- if VM.length mv == i then liftIO (VM.grow mv (VM.length mv)) else return mv - when (i `rem` 0x10000 == 0) $ liftIO performGC - liftIO $ VM.write mv' i a - go (i+1) mv' - -withFileFd :: (MonadIO m, MonadMask m) => FilePath -> (Fd -> m a) -> m a -withFileFd filepath = CMC.bracket - (liftIO $ openFd filepath ReadOnly Nothing defaultFileFlags) - (liftIO . closeFd) - diff --git a/src/Bio/Iteratee/Base.hs b/src/Bio/Iteratee/Base.hs deleted file mode 100644 --- a/src/Bio/Iteratee/Base.hs +++ /dev/null @@ -1,242 +0,0 @@ -{-# LANGUAGE Rank2Types #-} - --- |Monadic Iteratees: --- incremental input parsers, processors and transformers - -module Bio.Iteratee.Base ( - -- * Types - Stream (..) - ,StreamStatus (..) - -- ** Exception types - ,module Bio.Iteratee.Exception - -- ** Iteratees - ,Iteratee (..) - -- * Functions - -- ** Control functions - ,run - ,tryRun - ,ilift - ,ifold - -- ** Creating Iteratees - ,idone - ,icont - ,liftI - ,idoneM - ,icontM - -- ** Stream Functions - ,setEOF - -- * Classes - ,NullPoint(..) - ,Nullable(..) -) -where - -import Bio.Iteratee.Exception -import Bio.Prelude - -import Control.Monad.Catch as CIO -import Control.Monad.IO.Class ( MonadIO(..) ) -import Control.Monad.Trans.Class ( MonadTrans(..) ) - -import qualified Control.Exception as E -import qualified Data.ByteString as B -import qualified Data.ByteString.Lazy as L - --- | NullPoint class. Containers that have a null representation, --- corresponding to Data.Monoid.mempty. -class NullPoint c where emptyP :: c - -instance NullPoint (Endo a) where emptyP = Endo id -instance NullPoint [a] where emptyP = [] -instance NullPoint B.ByteString where emptyP = B.empty -instance NullPoint L.ByteString where emptyP = L.empty - --- | Nullable container class -class NullPoint c => Nullable c where nullC :: c -> Bool - -instance Nullable [a] where nullC [] = True ; nullC _ = False -instance Nullable B.ByteString where nullC = B.null -instance Nullable L.ByteString where nullC = L.null - --- |A stream is a (continuing) sequence of elements bundled in Chunks. --- The first variant indicates termination of the stream. --- Chunk a gives the currently available part of the stream. --- The stream is not terminated yet. --- The case (null Chunk) signifies a stream with no currently available --- data but which is still continuing. A stream processor should, --- informally speaking, ``suspend itself'' and wait for more data --- to arrive. - -data Stream c = EOF (Maybe SomeException) | Chunk c - deriving (Show, Typeable) - -instance (Eq c) => Eq (Stream c) where - (Chunk c1) == (Chunk c2) = c1 == c2 - (EOF Nothing) == (EOF Nothing) = True - (EOF (Just e1)) == (EOF (Just e2)) = typeOf e1 == typeOf e2 - _ == _ = False - -instance Semigroup c => Semigroup (Stream c) where - EOF mErr <> _ = EOF mErr - _ <> EOF mErr = EOF mErr - Chunk s1 <> Chunk s2 = Chunk (s1 <> s2) - -instance Monoid c => Monoid (Stream c) where - mempty = Chunk mempty - EOF mErr `mappend` _ = EOF mErr - _ `mappend` EOF mErr = EOF mErr - Chunk s1 `mappend` Chunk s2 = Chunk (mappend s1 s2) - --- |Map a function over a stream. -instance Functor Stream where - fmap f (Chunk xs) = Chunk $ f xs - fmap _ (EOF mErr) = EOF mErr - --- |Describe the status of a stream of data. -data StreamStatus = - DataRemaining - | EofNoError - | EofError SomeException - deriving (Show, Typeable) - --- ---------------------------------------------- --- create exception type hierarchy - --- |Produce the 'EOF' error message. If the stream was terminated because --- of an error, keep the error message. -setEOF :: Stream c -> SomeException -setEOF (EOF (Just e)) = e -setEOF _ = toException EofException - --- ---------------------------------------------- --- | Monadic iteratee -newtype Iteratee s m a = Iteratee{ runIter :: forall r. - (a -> Stream s -> m r) -> - ((Stream s -> Iteratee s m a) -> Maybe SomeException -> m r) -> - m r} - --- ---------------------------------------------- - -idone :: a -> Stream s -> Iteratee s m a -idone a s = Iteratee $ \onDone _ -> onDone a s - -icont :: (Stream s -> Iteratee s m a) -> Maybe SomeException -> Iteratee s m a -icont k e = Iteratee $ \_ onCont -> onCont k e - -liftI :: (Stream s -> Iteratee s m a) -> Iteratee s m a -liftI k = Iteratee $ \_ onCont -> onCont k Nothing - --- Monadic versions, frequently used by enumerators -idoneM :: Monad m => a -> Stream s -> m (Iteratee s m a) -idoneM x str = return $ Iteratee $ \onDone _ -> onDone x str - -icontM - :: Monad m => - (Stream s -> Iteratee s m a) - -> Maybe SomeException - -> m (Iteratee s m a) -icontM k e = return $ Iteratee $ \_ onCont -> onCont k e - -instance (Functor m) => Functor (Iteratee s m) where - fmap f m = Iteratee $ \onDone onCont -> - let od = onDone . f - oc = onCont . (fmap f .) - in runIter m od oc - -instance (Functor m, Monad m, Nullable s) => Applicative (Iteratee s m) where - pure x = idone x (Chunk emptyP) - {-# INLINE (<*>) #-} - m <*> a = m >>= flip fmap a - -instance (Monad m, Nullable s) => Monad (Iteratee s m) where - {-# INLINE return #-} - return x = Iteratee $ \onDone _ -> onDone x (Chunk emptyP) - {-# INLINE (>>=) #-} - (>>=) = bindIteratee - -{-# INLINE bindIteratee #-} -bindIteratee :: Nullable s - => Iteratee s m a - -> (a -> Iteratee s m b) - -> Iteratee s m b -bindIteratee m f = Iteratee $ \onDone onCont -> - let m_done a (Chunk s) - | nullC s = runIter (f a) onDone onCont - m_done a stream = runIter (f a) (const . flip onDone stream) f_cont - where f_cont k Nothing = runIter (k stream) onDone onCont - f_cont k e = onCont k e - in runIter m m_done (onCont . (flip bindIteratee f .)) - - -instance NullPoint s => MonadTrans (Iteratee s) where - lift m = Iteratee $ \onDone _ -> m >>= flip onDone (Chunk emptyP) - -instance (MonadIO m, Nullable s, NullPoint s) => MonadIO (Iteratee s m) where - liftIO = lift . liftIO - -instance (MonadThrow m, Nullable s, NullPoint s) => MonadThrow (Iteratee s m) where - throwM e = lift $ CIO.throwM e - -instance (MonadCatch m, Nullable s, NullPoint s) => MonadCatch (Iteratee s m) where - m `catch` f = Iteratee $ \od oc -> runIter m od oc `CIO.catch` (\e -> runIter (f e) od oc) - --- |Send 'EOF' to the @Iteratee@ and disregard the unconsumed part of the --- stream. If the iteratee is in an exception state, that exception is --- thrown with 'Control.Exception.throw'. Iteratees that do not terminate --- on @EOF@ will throw 'EofException'. -run :: Monad m => Iteratee s m a -> m a -run iter = runIter iter onDone onCont - where - onDone x _ = return x - onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' - onCont _ (Just e) = E.throw e - onCont' _ Nothing = E.throw EofException - onCont' _ (Just e) = E.throw e - --- |Run an iteratee, returning either the result or the iteratee exception. --- Note that only internal iteratee exceptions will be returned; exceptions --- thrown with @Control.Exception.throw@ or @Control.Monad.CatchIO.throw@ will --- not be returned. --- --- See 'Data.Iteratee.Exception.IFException' for details. -tryRun :: (Exception e, Monad m) => Iteratee s m a -> m (Either e a) -tryRun iter = runIter iter onDone onCont - where - onDone x _ = return $ Right x - onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' - onCont _ (Just e) = return $ maybeExc e - onCont' _ Nothing = return $ maybeExc (toException EofException) - onCont' _ (Just e) = return $ maybeExc e - maybeExc e = maybe (Left (E.throw e)) Left (fromException e) - --- | Lift a computation in the inner monad of an iteratee. --- --- A simple use would be to lift a logger iteratee to a monad stack. --- --- > logger :: Iteratee String IO () --- > logger = mapChunksM_ putStrLn --- > --- > loggerG :: MonadIO m => Iteratee String m () --- > loggerG = ilift liftIO logger --- --- A more complex example would involve lifting an iteratee to work with --- interleaved streams. See the example at 'Data.Iteratee.ListLike.merge'. -ilift :: - (Monad m, Monad n) - => (forall r. m r -> n r) - -> Iteratee s m a - -> Iteratee s n a -ilift f i = Iteratee $ \od oc -> - let onDone a str = return $ Left (a,str) - onCont k mErr = return $ Right (ilift f . k, mErr) - in f (runIter i onDone onCont) >>= either (uncurry od) (uncurry oc) - --- | Lift a computation in the inner monad of an iteratee, while threading --- through an accumulator. -ifold :: (Monad m, Monad n) => (forall r. m r -> acc -> n (r, acc)) - -> acc -> Iteratee s m a -> Iteratee s n (a, acc) -ifold f acc i = Iteratee $ \ od oc -> do - (r, acc') <- flip f acc $ - runIter i (curry $ return . Left) (curry $ return . Right) - either (uncurry (od . flip (,) acc')) - (uncurry (oc . (ifold f acc .))) r diff --git a/src/Bio/Iteratee/Bgzf.hsc b/src/Bio/Iteratee/Bgzf.hsc deleted file mode 100644 --- a/src/Bio/Iteratee/Bgzf.hsc +++ /dev/null @@ -1,491 +0,0 @@ --- | Handling of BGZF files. Right now, we have an Enumeratee each for --- input and output. The input iteratee can optionally supply virtual --- file offsets, so that seeking is possible. - -module Bio.Iteratee.Bgzf ( - Block(..), decompressBgzfBlocks', decompressBgzfBlocks, - decompressBgzf, decompressPlain, - maxBlockSize, bgzfEofMarker, liftBlock, getOffset, - BgzfChunk(..), isBgzf, isGzip, parMapChunksIO, - compressBgzf, compressBgzfLv, compressBgzf', CompressParams(..), - compressChunk - ) where - -import Bio.Iteratee -import Bio.Prelude -import Control.Concurrent.Async ( async, wait ) -import Foreign.C.String ( withCAString ) -import Foreign.C.Types ( CInt(..), CChar(..), CUInt(..), CULong(..) ) -import Foreign.Marshal.Alloc ( mallocBytes, free, allocaBytes ) - -import qualified Data.ByteString as S -import qualified Data.ByteString.Unsafe as S - -#include - --- | One BGZF block: virtual offset and contents. Could also be a block --- of an uncompressed file, if we want to support indexing of --- uncompressed BAM or some silliness like that. -data Block = Block { block_offset :: {-# UNPACK #-} !FileOffset - , block_contents :: {-# UNPACK #-} !Bytes } - -instance NullPoint Block where emptyP = Block 0 S.empty -instance Nullable Block where nullC = S.null . block_contents - -instance Monoid Block where - mempty = Block 0 S.empty - mappend = (<>) - mconcat [] = Block 0 S.empty - mconcat bs@(Block x _:_) = Block x $ S.concat [s|Block _ s <- bs] - -instance Semigroup Block where - Block x s <> Block _ t = Block x (s `S.append` t) - sconcat (Block x b :| bs) = Block x . sconcat $ b :| map block_contents bs - - --- | "Decompresses" a plain file. What's actually happening is that the --- offset in the input stream is tracked and added to the @Bytes@s --- giving @Block@s. This results in the same interface as decompressing --- actual Bgzf. -decompressPlain :: MonadIO m => Enumeratee Bytes Block m a -decompressPlain = eneeCheckIfDone (liftI . step 0) - where - step !o it (Chunk s) = eneeCheckIfDone (liftI . step (o + fromIntegral (S.length s))) . it $ Chunk (Block o s) - step _ it (EOF mx) = idone (liftI it) (EOF mx) - --- | Decompress a BGZF stream into a stream of 'Bytes's. -decompressBgzf :: MonadIO m => Enumeratee Bytes Bytes m a -decompressBgzf = decompressBgzfBlocks ><> mapChunks block_contents - -decompressBgzfBlocks :: MonadIO m => Enumeratee Bytes Block m a -decompressBgzfBlocks out = do - np <- liftIO $ getNumCapabilities - decompressBgzfBlocks' np out - --- | Decompress a BGZF stream into a stream of 'Block's, 'np' fold parallel. -decompressBgzfBlocks' :: MonadIO m => Int -> Enumeratee Bytes Block m a -decompressBgzfBlocks' np = eneeCheckIfDonePass (go 0 emptyQ) - where - -- check if the queue is full - go !off !qq k (Just e) = handleSeek off qq k e - go !off !qq k Nothing = case popQ qq of - Just (a, qq') | lengthQ qq == np -> liftIO (wait a) >>= eneeCheckIfDonePass (go off qq') . k . Chunk - _ -> liftI $ go' off qq k - - -- we have room for input, so try and get a compressed block - go' !_ !qq k (EOF mx) = goE mx qq k Nothing - go' !off !qq k (Chunk c) - | S.null c = liftI $ go' off qq k - | otherwise = joinIM $ enumPure1Chunk c $ do - (off', op) <- get_bgzf_block off - a <- liftIO (async op) - go off' (pushQ a qq) k Nothing - - -- input ended, empty the queue - goE _ !qq k (Just e) = handleSeek 0 qq k e - goE mx !qq k Nothing = case popQ qq of - Nothing -> idone (liftI k) (EOF mx) - Just (a,qq') -> liftIO (wait a) >>= eneeCheckIfDonePass (goE mx qq') . k . Chunk - - handleSeek !off !qq k e = case fromException e of - Nothing -> throwRecoverableErr e $ go' off qq k - Just (SeekException o) -> do - cancelAll qq - seek $ o `shiftR` 16 - eneeCheckIfDonePass (go (o `shiftR` 16) emptyQ) $ do - block'drop . fromIntegral $ o .&. 0xffff - liftI k - - block'drop sz = liftI $ \s -> case s of - EOF _ -> throwErr $ setEOF s - Chunk (Block p c) - | S.length c < sz -> block'drop (sz - S.length c) - | otherwise -> let b' = Block (p + fromIntegral sz) (S.drop sz c) - in idone () (Chunk b') - -get_bgzf_block :: MonadIO m => FileOffset -> Iteratee Bytes m (FileOffset, IO Block) -get_bgzf_block off = do !(csize,xlen) <- get_bgzf_header - !comp <- get_block . fromIntegral $ csize - xlen - 19 - !crc <- endianRead4 LSB - !isize <- endianRead4 LSB - - let !off' = off + fromIntegral csize + 1 - op = decompress1 (off `shiftL` 16) comp crc (fromIntegral isize) - return (off',op) - where - -- Get a block of a prescribed size. Comes back as a list of chunks. - get_block sz = liftI $ \s -> case s of - EOF _ -> throwErr $ setEOF s - Chunk c | S.length c < sz -> (:) c `liftM` get_block (sz - S.length c) - | otherwise -> idone [S.take sz c] (Chunk (S.drop sz c)) - - --- | Decodes a BGZF block header and returns the block size if --- successful. -get_bgzf_header :: Monad m => Iteratee Bytes m (Word16, Word16) -get_bgzf_header = do x <- headStreamBS - y <- headStreamBS - _cm <- headStreamBS - flg <- headStreamBS - if flg `testBit` 2 then do - dropStreamBS 6 - xlen <- endianRead2 LSB - it <- takeStreamBS (fromIntegral xlen) get_bsize >>= lift . tryRun - case it of Left e -> throwErr e - Right s | x == 31 && y == 139 -> return (s,xlen) - _ -> throwErr $ iterStrExc "No BGZF" - else throwErr $ iterStrExc "No BGZF" - where - get_bsize = do i1 <- headStreamBS - i2 <- headStreamBS - len <- endianRead2 LSB - if i1 == 66 && i2 == 67 && len == 2 - then endianRead2 LSB - else dropStreamBS (fromIntegral len) >> get_bsize - --- | Tests whether a stream is in BGZF format. Does not consume any --- input. -isBgzf :: Monad m => Iteratee Bytes m Bool -isBgzf = liftM isRight $ checkErr $ iLookAhead $ get_bgzf_header - --- | Tests whether a stream is in GZip format. Also returns @True@ on a --- Bgzf stream, which is technically a special case of GZip. -isGzip :: Monad m => Iteratee Bytes m Bool -isGzip = liftM (either (const False) id) $ checkErr $ iLookAhead $ test - where - test = do x <- headStreamBS - y <- headStreamBS - dropStreamBS 24 - b <- isFinished - return $ not b && x == 31 && y == 139 - --- ------------------------------------------------------------------------- Output - --- | Maximum block size for Bgzf: 64k with some room for headers and --- uncompressible stuff -maxBlockSize :: Int -maxBlockSize = 65450 - - --- | The EOF marker for BGZF files. --- This is just an empty string compressed as BGZF. Appended to BAM --- files to indicate their end. -bgzfEofMarker :: Bytes -bgzfEofMarker = "\x1f\x8b\x8\x4\0\0\0\0\0\xff\x6\0\x42\x43\x2\0\x1b\0\x3\0\0\0\0\0\0\0\0\0" - --- | Decompress a collection of strings into a single BGZF block. --- --- Ideally, we receive one decode chunk from a BGZF file, decompress it, --- and return it, in the process attaching the virtual address. But we --- might actually get more than one chunk, depending on the internals of --- the @Iteratee@s used. If so, we concatenate them; the first gets to --- assign the address. --- --- Now allocate space for uncompressed data, decompress the chunks we --- got, compute crc for each and check it, finally convert to 'Bytes' --- and emit. --- --- We could probably get away with @unsafePerformIO@'ing everything in --- here, but then again, we only do this when we're writing output --- anyway. Hence, run in IO. - - -decompress1 :: FileOffset -> [Bytes] -> Word32 -> Int -> IO Block -decompress1 off ss crc usize = - allocaBytes (#{const sizeof(z_stream)}) $ \stream -> do - buf <- mallocBytes usize - - #{poke z_stream, msg} stream nullPtr - #{poke z_stream, zalloc} stream nullPtr - #{poke z_stream, zfree} stream nullPtr - #{poke z_stream, opaque} stream nullPtr - #{poke z_stream, next_in} stream nullPtr - #{poke z_stream, next_out} stream buf - #{poke z_stream, avail_in} stream (0 :: CUInt) - #{poke z_stream, avail_out} stream (fromIntegral usize :: CUInt) - - z_check "inflateInit2" =<< c_inflateInit2 stream (-15) - - -- loop over the fragments, forward order - forM_ ss $ \s -> case fromIntegral $ S.length s of - l | l > 0 -> S.unsafeUseAsCString s $ \p -> do - #{poke z_stream, next_in} stream p - #{poke z_stream, avail_in} stream (l :: CUInt) - z_check "inflate" =<< c_inflate stream #{const Z_NO_FLUSH} - _ -> return () - - z_check "inflate" =<< c_inflate stream #{const Z_FINISH} - z_check "inflateEnd" =<< c_inflateEnd stream - - pe <- #{peek z_stream, next_out} stream - when (pe `minusPtr` buf /= usize) $ error "size mismatch after deflate()" - - crc0 <- c_crc32 0 nullPtr 0 - crc' <- c_crc32 crc0 buf (fromIntegral usize) - when (fromIntegral crc /= crc') $ error "CRC error after deflate()" - - Block off `liftM` S.unsafePackCStringFinalizer (castPtr buf) usize (free buf) - - --- | Compress a collection of strings into a single BGZF block. --- --- Okay, performance was lacking... let's do it again, in a more direct --- style. We build our block manually. First check if the compressed --- data is going to fit---if not, that's a bug. Then alloc a buffer, --- fill with a dummy header, alloc a ZStream, compress the pieces we --- were handed one at a time. Calculate CRC32, finalize header, --- construct a byte string, return it. --- --- We could probably get away with @unsafePerformIO@'ing everything in --- here, but then again, we only do this when we're writing output --- anyway. Hence, run in IO. - -compress1 :: Int -> [Bytes] -> IO Bytes -compress1 _lv [] = return bgzfEofMarker -compress1 lv ss0 = - allocaBytes (#{const sizeof(z_stream)}) $ \stream -> do - - let input_length = sum (map S.length ss0) - when (input_length > maxBlockSize) $ error "Trying to create too big a BGZF block; this is a bug." - buf <- mallocBytes 65536 - - -- steal header from the EOF marker (length is wrong for now) - S.unsafeUseAsCString bgzfEofMarker $ \eof -> - forM_ [0,4..16] $ \o -> do x <- peekByteOff eof o - pokeByteOff buf o (x::Word32) - - #{poke z_stream, msg} stream nullPtr - #{poke z_stream, zalloc} stream nullPtr - #{poke z_stream, zfree} stream nullPtr - #{poke z_stream, opaque} stream nullPtr - #{poke z_stream, next_in} stream nullPtr - #{poke z_stream, next_out} stream (buf `plusPtr` 18) - #{poke z_stream, avail_in} stream (0 :: CUInt) - #{poke z_stream, avail_out} stream (65536-18-8 :: CUInt) - - z_check "deflateInit2" =<< c_deflateInit2 stream (fromIntegral lv) #{const Z_DEFLATED} - (-15) 8 #{const Z_DEFAULT_STRATEGY} - - -- loop over the fragments. In reverse order! - let go (s:ss) = do - crc <- go ss - S.unsafeUseAsCString s $ \p -> - case fromIntegral $ S.length s of - l | l > 0 -> do - #{poke z_stream, next_in} stream p - #{poke z_stream, avail_in} stream (l :: CUInt) - z_check "deflate" =<< c_deflate stream #{const Z_NO_FLUSH} - c_crc32 crc p l - _ -> return crc - go [] = c_crc32 0 nullPtr 0 - crc <- go ss0 - - z_check "deflate" =<< c_deflate stream #{const Z_FINISH} - z_check "deflateEnd" =<< c_deflateEnd stream - - compressed_length <- (+) (18+8) `fmap` #{peek z_stream, total_out} stream - when (compressed_length > 65536) $ error "produced too big a block" - - -- set length in header - pokeByteOff buf 16 (fromIntegral $ (compressed_length-1) .&. 0xff :: Word8) - pokeByteOff buf 17 (fromIntegral $ (compressed_length-1) `shiftR` 8 :: Word8) - - pokeByteOff buf (compressed_length-8) (fromIntegral crc :: Word32) - pokeByteOff buf (compressed_length-4) (fromIntegral input_length :: Word32) - - S.unsafePackCStringFinalizer buf compressed_length (free buf) - - -data ZStream - -{-# INLINE z_check #-} -z_check :: String -> CInt -> IO () -z_check msg c = when (c /= #{const Z_OK} && c /= #{const Z_STREAM_END}) $ - error $ msg ++ " failed: " ++ show c - - -c_deflateInit2 :: Ptr ZStream -> CInt -> CInt -> CInt -> CInt -> CInt -> IO CInt -c_deflateInit2 z a b c d e = withCAString #{const_str ZLIB_VERSION} $ \versionStr -> - c_deflateInit2_ z a b c d e versionStr (#{const sizeof(z_stream)} :: CInt) - -foreign import ccall unsafe "zlib.h deflateInit2_" c_deflateInit2_ :: - Ptr ZStream -> CInt -> CInt -> CInt -> CInt -> CInt - -> Ptr CChar -> CInt -> IO CInt - -c_inflateInit2 :: Ptr ZStream -> CInt -> IO CInt -c_inflateInit2 z a = withCAString #{const_str ZLIB_VERSION} $ \versionStr -> - c_inflateInit2_ z a versionStr (#{const sizeof(z_stream)} :: CInt) - -foreign import ccall unsafe "zlib.h inflateInit2_" c_inflateInit2_ :: - Ptr ZStream -> CInt -> Ptr CChar -> CInt -> IO CInt - -foreign import ccall unsafe "zlib.h deflate" c_deflate :: - Ptr ZStream -> CInt -> IO CInt - -foreign import ccall unsafe "zlib.h inflate" c_inflate :: - Ptr ZStream -> CInt -> IO CInt - -foreign import ccall unsafe "zlib.h deflateEnd" c_deflateEnd :: - Ptr ZStream -> IO CInt - -foreign import ccall unsafe "zlib.h inflateEnd" c_inflateEnd :: - Ptr ZStream -> IO CInt - -foreign import ccall unsafe "zlib.h crc32" c_crc32 :: - CULong -> Ptr CChar -> CUInt -> IO CULong - --- ------------------------------------------------------------------------------------------------- utils - --- | Get the current virtual offset. The virtual address in a BGZF --- stream contains the offset of the current block in the upper 48 bits --- and the current offset into that block in the lower 16 bits. This --- scheme is compatible with the way BAM files are indexed. -getOffset :: Iteratee Block m FileOffset -getOffset = liftI step - where - step s@(EOF _) = icont step (Just (setEOF s)) - step s@(Chunk (Block o _)) = idone o s - --- | Runs an @Iteratee@ for @Bytes@s when decompressing BGZF. Adds --- internal bookkeeping. -liftBlock :: Monad m => Iteratee Bytes m a -> Iteratee Block m a -liftBlock = liftI . step - where - step it (EOF ex) = joinI $ lift $ enumChunk (EOF ex) it - - step it (Chunk (Block !l !s)) = Iteratee $ \od oc -> - enumPure1Chunk s it >>= \it' -> runIter it' (onDone od) (oc . step . liftI) - where - !sl = S.length s - onDone od hdr (Chunk !rest) = od hdr . Chunk $! Block (l + fromIntegral (sl - S.length rest)) rest - onDone od hdr (EOF ex) = od hdr (EOF ex) - - --- | Compresses a stream of @Bytes@s into a stream of BGZF blocks, --- in parallel - --- We accumulate an uncompressed block as long as adding a new chunk to --- it doesn't exceed the max. block size. If we receive an empty chunk --- (used as a flush signal), or if we would exceed the block size, we --- write out a block. Then we continue writing until we're below block --- size. On EOF, we flush and write the end marker. - -compressBgzf' :: MonadIO m => CompressParams -> Enumeratee BgzfChunk Bytes m a -compressBgzf' (CompressParams lv np) = bgzfBlocks ><> parMapChunksIO np (compress1 lv) - -data BgzfChunk = SpecialChunk !Bytes BgzfChunk - | RecordChunk !Bytes BgzfChunk - | LeftoverChunk !Bytes BgzfChunk - | NoChunk - -instance NullPoint BgzfChunk where emptyP = NoChunk -instance Nullable BgzfChunk where - nullC NoChunk = True - nullC (SpecialChunk s c) = S.null s && nullC c - nullC (RecordChunk s c) = S.null s && nullC c - nullC (LeftoverChunk s c) = S.null s && nullC c - --- | Breaks a stream into chunks suitable to be compressed individually. --- Each chunk on output is represented as a list of 'Bytes', --- each list must be reversed and concatenated to be compressed. --- ('compress1' does that.) - -bgzfBlocks :: Monad m => Enumeratee BgzfChunk [Bytes] m a -bgzfBlocks = eneeCheckIfDone (liftI . to_blocks 0 []) - where - -- terminate by sending the last block and then an empty block, - -- which becomes the EOF marker - to_blocks _alen acc k (EOF mx) = - lift (enumPure1Chunk [S.empty] (k $ Chunk acc)) >>= flip idone (EOF mx) - - -- \'Empty list\', in a sense. - to_blocks alen acc k (Chunk NoChunk) = liftI $ to_blocks alen acc k - - to_blocks alen acc k (Chunk (SpecialChunk c cs)) -- special chunk, encode then flush - -- If it fits, flush. - | alen + S.length c < maxBlockSize = eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk cs)) . k $ Chunk (c:acc) - -- If nothing is pending, flush the biggest thing that does fit. - | null acc = let (l,r) = S.splitAt maxBlockSize c - in eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (SpecialChunk r cs))) . k $ Chunk [l] - -- Otherwise, flush what's pending and think again. - | otherwise = eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (SpecialChunk c cs))) . k $ Chunk acc - - to_blocks alen acc k (Chunk (RecordChunk c cs)) - -- if it fits, we accumulate, (needs to consider the length prefix!) - | alen + S.length c + 4 < maxBlockSize = to_blocks (alen + S.length c + 4) (c:encLength c:acc) k (Chunk cs) - -- else if nothing's pending, we break the chunk, (needs to consider the length prefix!) - | null acc = let (l,r) = S.splitAt (maxBlockSize-4) c - in eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (LeftoverChunk r cs))) . k $ - Chunk [l, encLength l] - -- else we flush the accumulator and think again. - | otherwise = eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (RecordChunk c cs))) . k $ Chunk acc - where - encLength s = let !l = S.length s in S.pack [ fromIntegral (l `shiftR` 0 .&. 0xff) - , fromIntegral (l `shiftR` 8 .&. 0xff) - , fromIntegral (l `shiftR` 16 .&. 0xff) - , fromIntegral (l `shiftR` 24 .&. 0xff) ] - - to_blocks alen acc k (Chunk (LeftoverChunk c cs)) - -- if it fits, we accumulate, - | alen + S.length c < maxBlockSize = to_blocks (alen + S.length c) (c:acc) k (Chunk cs) - -- else if nothing's pending, we break the chunk, - | null acc = let (l,r) = S.splitAt maxBlockSize c - in eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (LeftoverChunk r cs))) . k $ Chunk [l] - -- else we flush the accumulator and think again. - | otherwise = eneeCheckIfDone (\k' -> to_blocks 0 [] k' (Chunk (LeftoverChunk c cs))) . k $ Chunk acc - --- | Like 'compressBgzf'', with sensible defaults. -compressBgzf :: MonadIO m => Enumeratee BgzfChunk Bytes m a -compressBgzf = compressBgzfLv 6 - -compressBgzfLv :: MonadIO m => Int -> Enumeratee BgzfChunk Bytes m a -compressBgzfLv lv out = do - np <- liftIO $ getNumCapabilities - compressBgzf' (CompressParams lv (np+2)) out - -data CompressParams = CompressParams { - compression_level :: Int, - queue_depth :: Int } - deriving Show - -compressChunk :: Int -> Ptr Word8 -> CUInt -> IO Bytes -compressChunk lv ptr len = - allocaBytes (#{const sizeof(z_stream)}) $ \stream -> do - buf <- mallocBytes 65536 - - -- steal header from the EOF marker (length is wrong for now) - S.unsafeUseAsCString bgzfEofMarker $ \eof -> - forM_ [0,4..16] $ \o -> do x <- peekByteOff eof o - pokeByteOff buf o (x::Word32) - - -- set up ZStream - #{poke z_stream, msg} stream nullPtr - #{poke z_stream, zalloc} stream nullPtr - #{poke z_stream, zfree} stream nullPtr - #{poke z_stream, opaque} stream nullPtr - #{poke z_stream, next_in} stream ptr - #{poke z_stream, next_out} stream (buf `plusPtr` 18) - #{poke z_stream, avail_in} stream len - #{poke z_stream, avail_out} stream (65536-18-8 :: CUInt) - - z_check "deflateInit2" =<< c_deflateInit2 stream (fromIntegral lv) #{const Z_DEFLATED} - (-15) 8 #{const Z_DEFAULT_STRATEGY} - -- z_check "deflate" =<< c_deflate stream #{const Z_NO_FLUSH} - z_check "deflate" =<< c_deflate stream #{const Z_FINISH} - z_check "deflateEnd" =<< c_deflateEnd stream - - crc0 <- c_crc32 0 nullPtr 0 - crc <- c_crc32 crc0 (castPtr ptr) len - - compressed_length <- (+) (18+8) `fmap` #{peek z_stream, total_out} stream - when (compressed_length > 65536) $ error "produced too big a block" - - -- set length in header - pokeByteOff buf 16 (fromIntegral $ (compressed_length-1) .&. 0xff :: Word8) - pokeByteOff buf 17 (fromIntegral $ (compressed_length-1) `shiftR` 8 :: Word8) - - pokeByteOff buf (compressed_length-8) (fromIntegral crc :: Word32) - pokeByteOff buf (compressed_length-4) (fromIntegral len :: Word32) - - S.unsafePackCStringFinalizer buf compressed_length (free buf) - diff --git a/src/Bio/Iteratee/Builder.hs b/src/Bio/Iteratee/Builder.hs deleted file mode 100644 --- a/src/Bio/Iteratee/Builder.hs +++ /dev/null @@ -1,262 +0,0 @@ -{-# LANGUAGE ForeignFunctionInterface #-} --- | Buffer builder to assemble Bgzf blocks. The idea is to serialize --- stuff (BAM and BCF) into a buffer, then bgzf chunks from the buffer. --- We use a large buffer, and we always make sure there is plenty of --- space in it (to avoid redundant checks). - -module Bio.Iteratee.Builder ( - BB(..), - newBuffer, - fillBuffer, - expandBuffer, - encodeBgzf, - BgzfTokens(..), - BclArgs(..), - BclSpecialType(..), - int_loop, - loop_bcl_special - ) where - -import Bio.Iteratee -import Bio.Iteratee.Bgzf ( compressChunk, maxBlockSize, bgzfEofMarker ) -import Bio.Prelude -import Foreign.Marshal.Utils ( copyBytes ) - -import qualified Data.ByteString as B -import qualified Data.ByteString.Unsafe as B -import qualified Data.Vector.Storable as V - --- | We manage a large buffer (multiple megabytes), of which we fill an --- initial portion. We remember the size, the used part, and two marks --- where we later fill in sizes for the length prefixed BAM or BCF --- records. We move the buffer down when we yield a piece downstream, --- and when we run out of space, we simply move to a new buffer. --- Garbage collection should take care of the rest. Unused 'mark' must --- be set to (maxBound::Int) so it doesn't interfere with flushing. - -data BB = BB { buffer :: {-# UNPACK #-} !(ForeignPtr Word8) - , size :: {-# UNPACK #-} !Int -- total size of buffer - , off :: {-# UNPACK #-} !Int -- offset of active portion - , used :: {-# UNPACK #-} !Int -- used portion (inactive & active) - , mark :: {-# UNPACK #-} !Int -- offset of mark - , mark2 :: {-# UNPACK #-} !Int } -- offset of mark2 - --- | Things we are able to encode. Taking inspiration from --- binary-serialise-cbor, we define these as a lazy list-like thing and --- consume it in a interpreter. - -data BgzfTokens = TkWord32 {-# UNPACK #-} !Word32 BgzfTokens -- a 4-byte int - | TkWord16 {-# UNPACK #-} !Word16 BgzfTokens -- a 2-byte int - | TkWord8 {-# UNPACK #-} !Word8 BgzfTokens -- a byte - | TkFloat {-# UNPACK #-} !Float BgzfTokens -- a float - | TkDouble {-# UNPACK #-} !Double BgzfTokens -- a double - | TkString {-# UNPACK #-} !B.ByteString BgzfTokens -- a raw string - | TkDecimal {-# UNPACK #-} !Int BgzfTokens -- roughly ':%d' - -- lotsa stuff might be missing here - - | TkSetMark BgzfTokens -- sets the first mark - | TkEndRecord BgzfTokens -- completes a BAM record - | TkEndRecordPart1 BgzfTokens -- completes part 1 of a BCF record - | TkEndRecordPart2 BgzfTokens -- completes part 2 of a BCF record - | TkEnd -- nothing more, for now - - -- specialties - | TkBclSpecial !BclArgs BgzfTokens - | TkLowLevel {-# UNPACK #-} !Int (BB -> IO BB) BgzfTokens - -data BclSpecialType = BclNucsBin | BclNucsWide | BclNucsAsc | BclNucsAscRev | BclQualsBin | BclQualsAsc | BclQualsAscRev - -data BclArgs = BclArgs BclSpecialType - {-# UNPACK #-} !(V.Vector Word8) -- bcl matrix - {-# UNPACK #-} !Int -- stride - {-# UNPACK #-} !Int -- first cycle - {-# UNPACK #-} !Int -- last cycle - {-# UNPACK #-} !Int -- cluster index - --- | Creates a buffer. -newBuffer :: Int -> IO BB -newBuffer sz = mallocForeignPtrBytes sz >>= \ar -> return $ BB ar sz 0 0 maxBound maxBound - --- | Creates a new buffer, copying the active content from an old one, --- with higher capacity. The size of the new buffer is twice the free --- space in the old buffer, but at least @minsz@. -expandBuffer :: Int -> BB -> IO BB -expandBuffer minsz b = do - let sz' = max (2 * (size b - used b)) minsz - arr1 <- mallocForeignPtrBytes sz' - withForeignPtr arr1 $ \d -> - withForeignPtr (buffer b) $ \s -> - copyBytes d (plusPtr s (off b)) (used b - off b) - return BB{ buffer = arr1 - , size = sz' - , off = 0 - , used = used b - off b - , mark = if mark b == maxBound then maxBound else mark b - off b - , mark2 = if mark2 b == maxBound then maxBound else mark2 b - off b } - -compressChunk' :: Int -> ForeignPtr Word8 -> Int -> Int -> IO B.ByteString -compressChunk' lv fptr off len = - withForeignPtr fptr $ \ptr -> - compressChunk lv (plusPtr ptr off) (fromIntegral len) - -instance Nullable (Endo BgzfTokens) where - nullC f = case appEndo f TkEnd of TkEnd -> True ; _ -> False - --- | Expand a chain of tokens into a buffer, sending finished pieces --- downstream as soon as possible. -encodeBgzf :: MonadIO m => Int -> Enumeratee (Endo BgzfTokens) B.ByteString m b -encodeBgzf lv = (\out -> newBuffer (1024*1024) `ioBind` \bb -> eneeCheckIfDone (liftI . go bb) out) - ><> parRunIO (2*numCapabilities) - where - go bb0 k (EOF mx) = final_flush bb0 mx k - go bb0 k (Chunk f) - -- initially, we make sure we have reasonable space. this may not be enough. - | size bb0 - used bb0 < 1024 = expandBuffer (1024*1024) bb0 `ioBind` \bb' -> go' bb' k (appEndo f TkEnd) - | otherwise = go' bb0 k (appEndo f TkEnd) - - go' bb0 k tk = fillBuffer bb0 tk `ioBind` \(bb',tk') -> flush_blocks tk' bb' k - - -- We can flush anything that is between 'off' and the lower of 'mark' - -- and 'used'. When done, we bump 'off'. - flush_blocks tk bb k - | min (mark bb) (used bb) - off bb < maxBlockSize = - case tk of TkEnd -> liftI $ go bb k - _ -> -- we arrive here because we ran out of buffer space, so we expand it. - expandBuffer (1024*1024) bb `ioBind` \bb' -> go' bb' k tk - | otherwise = - eneeCheckIfDone (flush_blocks tk bb { off = off bb + maxBlockSize }) $ - k $ Chunk [compressChunk' lv (buffer bb) (off bb) maxBlockSize] - - final_flush bb mx k - | used bb > off bb = - idone (k $ Chunk [ compressChunk' lv (buffer bb) (off bb) (used bb - off bb) - , return bgzfEofMarker ]) (EOF mx) - | otherwise = - idone (k $ Chunk [ return bgzfEofMarker ]) (EOF mx) - - -fillBuffer :: BB -> BgzfTokens -> IO (BB, BgzfTokens) -fillBuffer bb0 tk = withForeignPtr (buffer bb0) (\p -> go_slowish p bb0 tk) - where - go_slowish p bb = go_fast p bb (used bb) - - go_fast p bb use tk1 = case tk1 of - -- no space? not our job. - _ | size bb - use < 1024 -> return (bb { used = use },tk1) - - -- the actual end. - TkEnd -> return (bb { used = use },tk1) - - -- I'm cheating. This stuff works only if the platform allows - -- unaligned accesses, is little-endian and uses IEEE floats. - -- It's true on i386 and ix86_64. - TkWord32 x tk' -> do pokeByteOff p use x - go_fast p bb (use + 4) tk' - - TkWord16 x tk' -> do pokeByteOff p use x - go_fast p bb (use + 2) tk' - - TkWord8 x tk' -> do pokeByteOff p use x - go_fast p bb (use + 1) tk' - - TkFloat x tk' -> do pokeByteOff p use x - go_fast p bb (use + 4) tk' - - TkDouble x tk' -> do pokeByteOff p use x - go_fast p bb (use + 8) tk' - - TkString s tk' - -- Too big, can't handle. By returning with unfinished - -- business, we will get progressively bigger buffers and - -- eventually handle it. - | B.length s > size bb - use -> return (bb { used = use },tk1) - - | otherwise -> do let ln = B.length s - B.unsafeUseAsCString s $ \q -> - copyBytes (p `plusPtr` use) q ln - go_fast p bb (use + ln) tk' - - TkDecimal x tk' -> do ln <- int_loop (p `plusPtr` use) x - go_fast p bb (use + ln) tk' - - TkSetMark tk' -> go_slowish p bb { used = use + 4, mark = use } tk' - - TkEndRecord tk' -> do let !l = use - mark bb - 4 - pokeByteOff p (mark bb) (fromIntegral l :: Word32) - go_slowish p bb { used = use, mark = maxBound } tk' - - TkEndRecordPart1 tk' -> do let !l = use - mark bb - 4 - pokeByteOff p (mark bb - 4) (fromIntegral l :: Word32) - go_slowish p bb { used = use, mark2 = use } tk' - - TkEndRecordPart2 tk' -> do let !l = use - mark2 bb - pokeByteOff p (mark bb) (fromIntegral l :: Word32) - go_slowish p bb { used = use, mark = maxBound } tk' - - - TkBclSpecial special_args tk' -> do - l <- loop_bcl_special (p `plusPtr` use) special_args - go_fast p bb (use + l) tk' - - TkLowLevel minsize proc tk' - | size bb - use < minsize -> return (bb { used = use },tk1) - | otherwise -> do bb' <- proc bb { used = use } - go_slowish p bb' tk' - - -loop_bcl_special :: Ptr Word8 -> BclArgs -> IO Int -loop_bcl_special p (BclArgs tp vec stride u v i) = - - V.unsafeWith vec $ \q -> case tp of - BclNucsBin -> do - nuc_loop p stride (plusPtr q i) u v - return $ (v - u + 2) `div` 2 - - BclNucsWide -> do - nuc_loop_wide p stride (plusPtr q i) u v - return $ v - u + 1 - - BclNucsAsc -> do - nuc_loop_asc p stride (plusPtr q i) u v - return $ v - u + 1 - - BclNucsAscRev -> do - nuc_loop_asc_rev p stride (plusPtr q i) u v - return $ v - u + 1 - - BclQualsBin -> do - qual_loop p stride (plusPtr q i) u v - return $ v - u + 1 - - BclQualsAsc -> do - qual_loop_asc p stride (plusPtr q i) u v - return $ v - u + 1 - - BclQualsAscRev -> do - qual_loop_asc_rev p stride (plusPtr q i) u v - return $ v - u + 1 - -foreign import ccall unsafe "nuc_loop" - nuc_loop :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "nuc_loop_wide" - nuc_loop_wide :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "nuc_loop_asc" - nuc_loop_asc :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "nuc_loop_asc_rev" - nuc_loop_asc_rev :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "qual_loop" - qual_loop :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "qual_loop_asc" - qual_loop_asc :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "qual_loop_asc_rev" - qual_loop_asc_rev :: Ptr Word8 -> Int -> Ptr Word8 -> Int -> Int -> IO () - -foreign import ccall unsafe "int_loop" - int_loop :: Ptr Word8 -> Int -> IO Int - diff --git a/src/Bio/Iteratee/Bytes.hs b/src/Bio/Iteratee/Bytes.hs deleted file mode 100644 --- a/src/Bio/Iteratee/Bytes.hs +++ /dev/null @@ -1,269 +0,0 @@ -{-# LANGUAGE FlexibleContexts, BangPatterns #-} - --- |Monadic Iteratees: --- incremental input parsers, processors, and transformers --- --- Iteratees for parsing binary data. - -module Bio.Iteratee.Bytes ( - -- * Types - Endian (..) - -- * Endian multi-byte iteratees - ,endianRead2 - ,endianRead3 - ,endianRead3i - ,endianRead4 - ,endianRead8 - - -- * Iteratees treating Bytes as list of Word8 - ,headStreamBS - ,tryHeadBS - ,peekStreamBS - ,takeStreamBS - ,dropStreamBS - ,dropWhileStreamBS - - -- * Iteratees treating Bytes as list of Char - ,enumLinesBS - ,enumWordsBS -) -where - -import Bio.Iteratee.Base -import Bio.Iteratee.Iteratee -import Bio.Prelude - -import qualified Data.ByteString as B -import qualified Data.ByteString.Char8 as C -import qualified Data.ByteString.Unsafe as B - --- ------------------------------------------------------------------------ --- Binary Random IO Iteratees - --- Iteratees to read unsigned integers written in Big- or Little-endian ways - --- | Indicate endian-ness. -data Endian = MSB -- ^ Most Significant Byte is first (big-endian) - | LSB -- ^ Least Significan Byte is first (little-endian) - deriving (Eq, Ord, Show, Enum) - -endianRead2 :: Endian -> Iteratee Bytes m Word16 -endianRead2 e = endianReadN e 2 word16' -{-# INLINE endianRead2 #-} - -endianRead3 :: Endian -> Iteratee Bytes m Word32 -endianRead3 e = endianReadN e 3 (word32' . (0:)) -{-# INLINE endianRead3 #-} - --- |Read 3 bytes in an endian manner. If the first bit is set (negative), --- set the entire first byte so the Int32 will be negative as --- well. -endianRead3i :: Monad m => Endian -> Iteratee Bytes m Int32 -endianRead3i e = do - c1 <- headStreamBS - c2 <- headStreamBS - c3 <- headStreamBS - case e of - MSB -> return $ (((fromIntegral c1 - `shiftL` 8) .|. fromIntegral c2) - `shiftL` 8) .|. fromIntegral c3 - LSB -> - let m :: Int32 - m = shiftR (shiftL (fromIntegral c3) 24) 8 - in return $ (((fromIntegral c3 - `shiftL` 8) .|. fromIntegral c2) - `shiftL` 8) .|. fromIntegral m -{-# INLINE endianRead3i #-} - -endianRead4 :: Endian -> Iteratee Bytes m Word32 -endianRead4 e = endianReadN e 4 word32' -{-# INLINE endianRead4 #-} - -endianRead8 :: Endian -> Iteratee Bytes m Word64 -endianRead8 e = endianReadN e 8 word64' -{-# INLINE endianRead8 #-} - --- This function does all the parsing work, depending upon provided arguments -endianReadN :: - Endian - -> Int - -> ([Word8] -> b) - -> Iteratee Bytes m b -endianReadN MSB n0 cnct = liftI (step n0 []) - where - step !n acc (Chunk c) - | B.null c = liftI (step n acc) - | B.length c >= n = let (this,next) = B.splitAt n c - !result = cnct $ acc ++ B.unpack this - in idone result (Chunk next) - | otherwise = liftI (step (n - B.length c) (acc ++ B.unpack c)) - step !n acc (EOF Nothing) = icont (step n acc) (Just $ toException EofException) - step !n acc (EOF (Just e)) = icont (step n acc) (Just e) -endianReadN LSB n0 cnct = liftI (step n0 []) - where - step !n acc (Chunk c) - | B.null c = liftI (step n acc) - | B.length c >= n = let (this,next) = B.splitAt n c - !result = cnct $ B.unpack (B.reverse this) ++ acc - in idone result (Chunk next) - | otherwise = liftI (step (n - B.length c) - (B.unpack (B.reverse c) ++ acc)) - step !n acc (EOF Nothing) = icont (step n acc) - (Just $ toException EofException) - step !n acc (EOF (Just e)) = icont (step n acc) (Just e) -{-# INLINE endianReadN #-} - - -word16' :: [Word8] -> Word16 -word16' [c1,c2] = word16 c1 c2 -word16' _ = error "iteratee: internal error in word16'" - -word16 :: Word8 -> Word8 -> Word16 -word16 c1 c2 = (fromIntegral c1 `shiftL` 8) .|. fromIntegral c2 -{-# INLINE word16 #-} - -word32' :: [Word8] -> Word32 -word32' [c1,c2,c3,c4] = word32 c1 c2 c3 c4 -word32' _ = error "iteratee: internal error in word32'" - -word32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32 -word32 c1 c2 c3 c4 = - (fromIntegral c1 `shiftL` 24) .|. - (fromIntegral c2 `shiftL` 16) .|. - (fromIntegral c3 `shiftL` 8) .|. - fromIntegral c4 -{-# INLINE word32 #-} - -word64' :: [Word8] -> Word64 -word64' [c1,c2,c3,c4,c5,c6,c7,c8] = word64 c1 c2 c3 c4 c5 c6 c7 c8 -word64' _ = error "iteratee: internal error in word64'" -{-# INLINE word64' #-} - -word64 - :: Word8 -> Word8 -> Word8 -> Word8 - -> Word8 -> Word8 -> Word8 -> Word8 - -> Word64 -word64 c1 c2 c3 c4 c5 c6 c7 c8 = - (fromIntegral c1 `shiftL` 56) .|. - (fromIntegral c2 `shiftL` 48) .|. - (fromIntegral c3 `shiftL` 40) .|. - (fromIntegral c4 `shiftL` 32) .|. - (fromIntegral c5 `shiftL` 24) .|. - (fromIntegral c6 `shiftL` 16) .|. - (fromIntegral c7 `shiftL` 8) .|. - fromIntegral c8 -{-# INLINE word64 #-} - -headStreamBS :: Iteratee Bytes m Word8 -headStreamBS = liftI step - where - step (Chunk c) - | B.null c = icont step Nothing - | otherwise = idone (B.unsafeHead c) (Chunk (B.unsafeTail c)) - step stream = icont step (Just (setEOF stream)) -{-# INLINE headStreamBS #-} - -peekStreamBS :: Iteratee Bytes m (Maybe Word8) -peekStreamBS = liftI step - where - step s@(Chunk vec) - | B.null vec = liftI step - | otherwise = idone (Just $ B.unsafeHead vec) s - step stream = idone Nothing stream -{-# INLINE peekStreamBS #-} - -tryHeadBS :: Iteratee Bytes m (Maybe Word8) -tryHeadBS = liftI step - where - step (Chunk vec) - | B.null vec = liftI step - | otherwise = idone (Just (B.unsafeHead vec)) (Chunk (B.unsafeTail vec)) - step stream = idone Nothing stream -{-# INLINE tryHeadBS #-} - -dropStreamBS :: Int -> Iteratee Bytes m () -dropStreamBS 0 = idone () (Chunk emptyP) -dropStreamBS n' = liftI (step n') - where - step n (Chunk str) - | B.length str < n = liftI (step (n - B.length str)) - | otherwise = idone () (Chunk (B.drop n str)) - step _ stream = idone () stream -{-# INLINE dropStreamBS #-} - -dropWhileStreamBS :: (Word8 -> Bool) -> Iteratee Bytes m () -dropWhileStreamBS p = liftI step - where - step (Chunk str) - | B.null rest = liftI step - | otherwise = idone () (Chunk rest) - where - rest = B.dropWhile p str - step stream = idone () stream -{-# INLINE dropWhileStreamBS #-} - -takeStreamBS :: - Monad m - => Int -- ^ number of elements to consume - -> Enumeratee Bytes Bytes m a -takeStreamBS n' iter - | n' <= 0 = return iter - | otherwise = Iteratee $ \od oc -> runIter iter (on_done od oc) (on_cont od oc) - where - on_done od oc x _ = runIter (dropStreamBS n' >> return (return x)) od oc - on_cont od oc k Nothing = if n' == 0 then od (liftI k) (Chunk mempty) - else runIter (liftI (step n' k)) od oc - on_cont od oc _ (Just e) = runIter (dropStreamBS n' >> throwErr e) od oc - step n k (Chunk str) - | B.null str = liftI (step n k) - | B.length str <= n = takeStreamBS (n - B.length str) $ k (Chunk str) - | otherwise = idone (k (Chunk s1)) (Chunk s2) - where (s1, s2) = B.splitAt n str - step _n k stream = idone (liftI k) stream -{-# INLINE takeStreamBS #-} - --- Like enumWords, but operates on ByteStrings. --- This is provided as a higher-performance alternative to enumWords, and --- is equivalent to treating the stream as a Data.ByteString.Char8.ByteString. -enumWordsBS :: Monad m => Enumeratee Bytes [Bytes] m a -enumWordsBS = convStream getter - where - getter = liftI step - lChar = isSpace . C.last - step (Chunk xs) - | C.null xs = getter - | lChar xs = idone (C.words xs) (Chunk C.empty) - | otherwise = icont (step' xs) Nothing - step str = idone mempty str - step' xs (Chunk ys) - | C.null ys = icont (step' xs) Nothing - | lChar ys = idone (C.words . C.append xs $ ys) mempty - | otherwise = let w' = C.words . C.append xs $ ys - ws = init w' - ck = last w' - in idone ws (Chunk ck) - step' xs str = idone (C.words xs) str -{-# INLINE enumWordsBS #-} - --- Like enumLines, but operates on ByteStrings. --- This is provided as a higher-performance alternative to enumLines, and --- is equivalent to treating the stream as a Data.ByteString.Char8.ByteString. -enumLinesBS :: Monad m => Enumeratee Bytes [Bytes] m a -enumLinesBS = convStream getter - where - getter = icont step Nothing - lChar = (== '\n') . C.last - step (Chunk xs) - | C.null xs = getter - | lChar xs = idone (C.lines xs) (Chunk C.empty) - | otherwise = icont (step' xs) Nothing - step str = idone mempty str - step' xs (Chunk ys) - | C.null ys = icont (step' xs) Nothing - | lChar ys = idone (C.lines . C.append xs $ ys) mempty - | otherwise = let w' = C.lines $ C.append xs ys - ws = init w' - ck = last w' - in idone ws (Chunk ck) - step' xs str = idone (C.lines xs) str -{-# INLINE enumLinesBS #-} diff --git a/src/Bio/Iteratee/Exception.hs b/src/Bio/Iteratee/Exception.hs deleted file mode 100644 --- a/src/Bio/Iteratee/Exception.hs +++ /dev/null @@ -1,212 +0,0 @@ -{-# LANGUAGE ExistentialQuantification #-} - --- |Monadic and General Iteratees: --- Messaging and exception handling. --- --- Iteratees use an internal exception handling mechanism that is parallel to --- that provided by 'Control.Exception'. This allows the iteratee framework --- to handle its own exceptions outside @IO@. --- --- Iteratee exceptions are divided into two categories, 'IterException' and --- 'EnumException'. @IterExceptions@ are exceptions within an iteratee, and --- @EnumExceptions@ are exceptions within an enumerator. --- --- Enumerators can be constructed to handle an 'IterException' with --- @Data.Iteratee.Iteratee.enumFromCallbackCatch@. If the enumerator detects --- an @iteratee exception@, the enumerator calls the provided exception handler. --- The enumerator is then able to continue feeding data to the iteratee, --- provided the exception was successfully handled. If the handler could --- not handle the exception, the 'IterException' is converted to an --- 'EnumException' and processing aborts. --- --- Exceptions can also be cleared by @Data.Iteratee.Iteratee.checkErr@, --- although in this case the iteratee continuation cannot be recovered. --- --- When viewed as Resumable Exceptions, iteratee exceptions provide a means --- for iteratees to send control messages to enumerators. The @seek@ --- implementation provides an example. @Data.Iteratee.Iteratee.seek@ stores --- the current iteratee continuation and throws a 'SeekException', which --- inherits from 'IterException'. @Data.Iteratee.IO.enumHandleRandom@ is --- constructed with @enumFromCallbackCatch@ and a handler that performs --- an @hSeek@. Upon receiving the 'SeekException', @enumHandleRandom@ calls --- the handler, checks that it executed properly, and then continues with --- the stored continuation. --- --- As the exception hierarchy is open, users can extend it with custom --- exceptions and exception handlers to implement sophisticated messaging --- systems based upon resumable exceptions. - - -module Bio.Iteratee.Exception ( - -- * Exception types - IFException (..) - ,Exception (..) -- from Control.Exception - ,FileOffset -- from System.Posix.Types - -- ** Enumerator exceptions - ,EnumException (..) - ,DivergentException (..) - ,EnumStringException (..) - ,EnumUnhandledIterException (..) - -- ** Iteratee exceptions - ,IException (..) - ,IterException (..) - ,SeekException (..) - ,EofException (..) - ,IterStringException (..) - -- * Functions - ,enStrExc - ,iterStrExc - ,wrapIterExc - ,iterExceptionToException - ,iterExceptionFromException -) -where - -import Control.Exception -import Data.Data -import Prelude -import System.Posix.Types ( FileOffset ) - - --- ---------------------------------------------- --- create exception type hierarchy - --- |Root of the Iteratee exception hierarchy. @IFException@ derives from --- @Control.Exception.SomeException@. 'EnumException', 'IterException', --- and all inheritants are descendents of 'IFException'. -data IFException = forall e . Exception e => IFException e - deriving Typeable - -instance Show IFException where - show (IFException e) = show e - -instance Exception IFException - -ifExceptionToException :: Exception e => e -> SomeException -ifExceptionToException = toException . IFException - -ifExceptionFromException :: Exception e => SomeException -> Maybe e -ifExceptionFromException x = do - IFException a <- fromException x - cast a - --- Root of enumerator exceptions. -data EnumException = forall e . Exception e => EnumException e - deriving Typeable - -instance Show EnumException where - show (EnumException e) = show e - -instance Exception EnumException where - toException = ifExceptionToException - fromException = ifExceptionFromException - -enumExceptionToException :: Exception e => e -> SomeException -enumExceptionToException = toException . IterException - -enumExceptionFromException :: Exception e => SomeException -> Maybe e -enumExceptionFromException x = do - IterException a <- fromException x - cast a - --- |The @iteratee@ diverged upon receiving 'EOF'. -data DivergentException = DivergentException - deriving (Show, Typeable) - -instance Exception DivergentException where - toException = enumExceptionToException - fromException = enumExceptionFromException - --- |Create an enumerator exception from a @String@. -newtype EnumStringException = EnumStringException String - deriving (Show, Typeable) - -instance Exception EnumStringException where - toException = enumExceptionToException - fromException = enumExceptionFromException - --- |Create an 'EnumException' from a string. -enStrExc :: String -> EnumException -enStrExc = EnumException . EnumStringException - --- |The enumerator received an 'IterException' it could not handle. -newtype EnumUnhandledIterException = EnumUnhandledIterException IterException - deriving (Show, Typeable) - -instance Exception EnumUnhandledIterException where - toException = enumExceptionToException - fromException = enumExceptionFromException - --- |Convert an 'IterException' to an 'EnumException'. Meant to be used --- within an @Enumerator@ to signify that it could not handle the --- @IterException@. -wrapIterExc :: IterException -> EnumException -wrapIterExc = EnumException . EnumUnhandledIterException - --- iteratee exceptions - --- |A class for @iteratee exceptions@. Only inheritants of @IterException@ --- should be instances of this class. -class Exception e => IException e where - toIterException :: e -> IterException - toIterException = IterException - fromIterException :: IterException -> Maybe e - fromIterException = fromException . toException - --- |Root of iteratee exceptions. -data IterException = forall e . Exception e => IterException e - deriving Typeable - -instance Show IterException where - show (IterException e) = show e - -instance Exception IterException where - toException = ifExceptionToException - fromException = ifExceptionFromException - -iterExceptionToException :: Exception e => e -> SomeException -iterExceptionToException = toException . IterException - -iterExceptionFromException :: Exception e => SomeException -> Maybe e -iterExceptionFromException x = do - IterException a <- fromException x - cast a - -instance IException IterException where - toIterException = id - fromIterException = Just - --- |A seek request within an @Iteratee@. -newtype SeekException = SeekException FileOffset - deriving (Typeable, Show) - -instance Exception SeekException where - toException = iterExceptionToException - fromException = iterExceptionFromException - -instance IException SeekException where - --- |The @Iteratee@ needs more data but received @EOF@. -data EofException = EofException - deriving (Typeable, Show) - -instance Exception EofException where - toException = iterExceptionToException - fromException = iterExceptionFromException - -instance IException EofException where - --- |An @Iteratee exception@ specified by a @String@. -newtype IterStringException = IterStringException String deriving (Typeable, Show) - -instance Exception IterStringException where - toException = iterExceptionToException - fromException = iterExceptionFromException - -instance IException IterStringException where - --- |Create an @iteratee exception@ from a string. --- This convenience function wraps 'IterStringException' and 'toException'. -iterStrExc :: String -> SomeException -iterStrExc= toException . IterStringException - diff --git a/src/Bio/Iteratee/IO.hs b/src/Bio/Iteratee/IO.hs deleted file mode 100644 --- a/src/Bio/Iteratee/IO.hs +++ /dev/null @@ -1,91 +0,0 @@ --- |Random and Binary IO with generic Iteratees, using File Descriptors for IO. --- when available, these are the preferred functions for performing IO as they --- run in constant space and function properly with sockets, pipes, etc. - -module Bio.Iteratee.IO( - -- * Data - defaultBufSize - -- * File enumerators - ,enumFile - ,enumFileRandom - -- * FileDescriptor based enumerators for monadic iteratees - ,enumFd - ,enumFdRandom -) -where - -import Bio.Iteratee.Iteratee -import Bio.Prelude hiding ( bracket, loop ) -import Control.Monad.IO.Class ( MonadIO(..) ) -import System.IO ( SeekMode(..) ) - -import qualified Data.ByteString as B - --- | Default buffer size in elements. This was 1024 in "Data.Iteratee", --- which is obviously too small. Since we often want to merge many --- files, a read should take more time than a seek. This sets the --- sensible buffer size to somewhat more than one MB. -defaultBufSize :: Int -defaultBufSize = 2*1024*1024 - --- |The enumerator of a POSIX File Descriptor. This version enumerates --- over the entire contents of a file, in order, unless stopped by --- the iteratee. In particular, seeking is not supported. -enumFd :: MonadIO m => Int -> Fd -> Enumerator Bytes m a -enumFd bufsize fd = loop - where - loop iter = runIter iter idoneM onCont - - onCont k j@(Just _) = return (icont k j) - onCont k Nothing = do - s <- liftIO $ fdGet bufsize fd - if B.null s then return $ liftI k - else loop . k $ Chunk s - - --- |The enumerator of a POSIX File Descriptor: a variation of @enumFd@ that --- supports RandomIO (seek requests). -enumFdRandom :: MonadIO m => Int -> Fd -> Enumerator Bytes m a -enumFdRandom bs fd = loop - where - loop iter = runIter iter idoneM onCont - - onCont k Nothing = do - s <- liftIO $ fdGet bs fd - if B.null s then return $ liftI k - else loop . k $ Chunk s - - onCont k j@(Just e) = case fromException e of - Just (SeekException off) -> do - liftIO . void $ fdSeek fd AbsoluteSeek (fromIntegral off) - onCont k Nothing - - Nothing -> return (icont k j) - - - -enumFile' :: MonadBracketIO m => - (Int -> Fd -> Enumerator s m a) - -> Int -- ^Buffer size - -> FilePath - -> Enumerator s m a -enumFile' enumf bufsize filepath iter = bracketIO - (openFd filepath ReadOnly Nothing defaultFileFlags) - (closeFd) - (flip (enumf bufsize) iter) - -enumFile :: - MonadBracketIO m - => Int -- ^Buffer size - -> FilePath - -> Enumerator Bytes m a -enumFile = enumFile' enumFd - -enumFileRandom :: - MonadBracketIO m - => Int -- ^Buffer size - -> FilePath - -> Enumerator Bytes m a -enumFileRandom = enumFile' enumFdRandom - - diff --git a/src/Bio/Iteratee/Iteratee.hs b/src/Bio/Iteratee/Iteratee.hs deleted file mode 100644 --- a/src/Bio/Iteratee/Iteratee.hs +++ /dev/null @@ -1,601 +0,0 @@ -{-# LANGUAGE Rank2Types, FlexibleContexts #-} - --- |Monadic and General Iteratees: --- incremental input parsers, processors and transformers - -module Bio.Iteratee.Iteratee ( - -- * Types - EnumerateeHandler - -- ** Error handling - ,throwErr - ,throwRecoverableErr - ,checkErr - -- ** Basic Iteratees - ,skipToEof - ,isStreamFinished - -- ** Iteratee composition - ,mBind - ,mBind_ - ,ioBind - ,ioBind_ - ,MonadBracketIO(..) - -- ** Chunkwise Iteratees - ,mapChunksM_ - ,foldChunksM - ,getChunk - ,getChunks - -- ** Nested iteratee combinators - ,mapChunks - ,mapChunksM - ,convStream - ,unfoldConvStream - ,unfoldConvStreamCheck - ,joinI - ,joinIM - -- * Enumerators - ,Enumerator - ,Enumeratee - -- ** Basic enumerators - ,enumChunk - ,enumEof - ,enumErr - ,enumPure1Chunk - ,enumList - ,enumCheckIfDone - ,enumFromCallback - ,enumFromCallbackCatch - -- ** Enumerator Combinators - ,eneeCheckIfDone - ,eneeCheckIfDoneHandle - ,eneeCheckIfDoneIgnore - ,eneeCheckIfDonePass - ,mergeEnums - -- ** Enumeratee Combinators - ,($=) - ,(=$) - ,(><>) - ,(<><) - -- * Misc. - ,seek - -- * Classes - ,module Bio.Iteratee.Base -) -where - -import Bio.Iteratee.Base -import Bio.Prelude hiding (loop) -import Control.Monad.Catch as CIO -import Control.Monad.IO.Class (MonadIO(..)) -import Control.Monad.Trans.Class (MonadTrans(..)) - --- exception helpers -excDivergent :: SomeException -excDivergent = toException DivergentException - --- ------------------------------------------------------------------------ --- Primitive iteratees - --- |Report and propagate an unrecoverable error. --- Disregard the input first and then propagate the error. This error --- cannot be handled by 'enumFromCallbackCatch', although it can be cleared --- by 'checkErr'. -throwErr :: SomeException -> Iteratee s m a -throwErr e = icont (const (throwErr e)) (Just e) - --- |Report and propagate a recoverable error. This error can be handled by --- both 'enumFromCallbackCatch' and 'checkErr'. -throwRecoverableErr :: - SomeException - -> (Stream s -> Iteratee s m a) - -> Iteratee s m a -throwRecoverableErr e i = icont i (Just e) - - --- |Check if an iteratee produces an error. --- Returns @Right a@ if it completes without errors, otherwise --- @Left SomeException@. 'checkErr' is useful for iteratees that may not --- terminate, such as @Data.Iteratee.head@ with an empty stream. -checkErr :: - (NullPoint s) => - Iteratee s m a - -> Iteratee s m (Either SomeException a) -checkErr iter = Iteratee $ \onDone onCont -> - let od = onDone . Right - oc k Nothing = onCont (checkErr . k) Nothing - oc _ (Just e) = onDone (Left e) (Chunk emptyP) - in runIter iter od oc - --- ------------------------------------------------------------------------ --- Parser combinators - --- |Get the stream status of an iteratee. -isStreamFinished :: (Nullable s) => Iteratee s m (Maybe SomeException) -isStreamFinished = liftI check - where - check s@(Chunk xs) - | nullC xs = isStreamFinished - | otherwise = idone Nothing s - check s@(EOF e) = idone (Just $ fromMaybe (toException EofException) e) s -{-# INLINE isStreamFinished #-} - - --- |Skip the rest of the stream -skipToEof :: Iteratee s m () -skipToEof = icont check Nothing - where - check (Chunk _) = skipToEof - check s = idone () s - - --- |Seek to a position in the stream -seek :: Nullable s => FileOffset -> Iteratee s m () -seek o = throwRecoverableErr (toException $ SeekException o) (idone ()) - - --- | Map a monadic function over the chunks of the stream and ignore the --- result. Useful for creating efficient monadic iteratee consumers, e.g. --- --- > logger = mapChunksM_ (liftIO . putStrLn) --- --- these can be efficiently run in parallel with other iteratees via --- @Data.Iteratee.ListLike.zip@. -mapChunksM_ :: (Monad m, Nullable s) => (s -> m b) -> Iteratee s m () -mapChunksM_ f = liftI step - where - step (Chunk xs) - | nullC xs = liftI step - | otherwise = lift (f xs) >> liftI step - step s@(EOF _) = idone () s -{-# INLINE mapChunksM_ #-} - --- | A fold over chunks -foldChunksM :: (Monad m, Nullable s) => (a -> s -> m a) -> a -> Iteratee s m a -foldChunksM f = liftI . go - where - go a (Chunk c) = lift (f a c) >>= liftI . go - go a e = idone a e -{-# INLINE foldChunksM #-} - --- | Get the current chunk from the stream. -getChunk :: Nullable s => Iteratee s m s -getChunk = liftI step - where - step (Chunk xs) - | nullC xs = liftI step - | otherwise = idone xs $ Chunk emptyP - step (EOF Nothing) = throwErr $ toException EofException - step (EOF (Just e)) = throwErr e -{-# INLINE getChunk #-} - --- | Get a list of all chunks from the stream. -getChunks :: (Nullable s) => Iteratee s m [s] -getChunks = liftI (step id) - where - step acc (Chunk xs) - | nullC xs = liftI (step acc) - | otherwise = liftI (step $ acc . (xs:)) - step acc stream = idone (acc []) stream -{-# INLINE getChunks #-} - --- --------------------------------------------------- --- The converters show a different way of composing two iteratees: --- `vertical' rather than `horizontal' - -type Enumeratee sFrom sTo m a = - Iteratee sTo m a - -> Iteratee sFrom m (Iteratee sTo m a) - --- The following pattern appears often in Enumeratee code -{-# INLINE eneeCheckIfDone #-} - --- | Utility function for creating enumeratees. Typical usage is demonstrated --- by the @breakE@ definition. --- --- > breakE --- > :: (Monad m, LL.ListLike s el, NullPoint s) --- > => (el -> Bool) --- > -> Enumeratee s s m a --- > breakE cpred = eneeCheckIfDone (liftI . step) --- > where --- > step k (Chunk s) --- > | LL.null s = liftI (step k) --- > | otherwise = case LL.break cpred s of --- > (str', tail') --- > | LL.null tail' -> eneeCheckIfDone (liftI . step) . k $ Chunk str' --- > | otherwise -> idone (k $ Chunk str') (Chunk tail') --- > step k stream = idone (k stream) stream --- -eneeCheckIfDone :: - (Monad m, NullPoint elo) => - ((Stream eli -> Iteratee eli m a) -> Iteratee elo m (Iteratee eli m a)) - -> Enumeratee elo eli m a -eneeCheckIfDone f = eneeCheckIfDonePass f' - where - f' k Nothing = f k - f' k (Just e) = throwRecoverableErr e (\s -> joinIM $ enumChunk s $ eneeCheckIfDone f (liftI k)) - -type EnumerateeHandler eli elo m a = - (Stream eli -> Iteratee eli m a) - -> SomeException - -> Iteratee elo m (Iteratee eli m a) - --- | The same as eneeCheckIfDonePass, with one extra argument: --- a handler which is used --- to process any exceptions in a separate method. -eneeCheckIfDoneHandle - :: (NullPoint elo) - => EnumerateeHandler eli elo m a - -> ((Stream eli -> Iteratee eli m a) - -> Maybe SomeException - -> Iteratee elo m (Iteratee eli m a) - ) - -> Enumeratee elo eli m a -eneeCheckIfDoneHandle h f inner = Iteratee $ \od oc -> - let onDone x s = od (idone x s) (Chunk emptyP) - onCont k Nothing = runIter (f k Nothing) od oc - onCont k (Just e) = runIter (h k e) od oc - in runIter inner onDone onCont -{-# INLINABLE eneeCheckIfDoneHandle #-} - -eneeCheckIfDonePass - :: (NullPoint elo) - => ((Stream eli -> Iteratee eli m a) - -> Maybe SomeException - -> Iteratee elo m (Iteratee eli m a) - ) - -> Enumeratee elo eli m a -eneeCheckIfDonePass f = eneeCheckIfDoneHandle (\k e -> f k (Just e)) f -{-# INLINABLE eneeCheckIfDonePass #-} - -eneeCheckIfDoneIgnore - :: (NullPoint elo) - => ((Stream eli -> Iteratee eli m a) - -> Maybe SomeException - -> Iteratee elo m (Iteratee eli m a) - ) - -> Enumeratee elo eli m a -eneeCheckIfDoneIgnore f = eneeCheckIfDoneHandle (\k _ -> f k Nothing) f - -{-# INLINE mBind #-} --- | Lifts a monadic action and combines it with a continuation. --- @mBind m f@ is the same as @lift m >>= f@, but does not require a --- 'Nullable' constraint on the stream type. -infixl 1 `mBind` -mBind :: Monad m => m a -> (a -> Iteratee s m b) -> Iteratee s m b -mBind m f = Iteratee $ \onDone onCont -> m >>= \a -> runIter (f a) onDone onCont - -{-# INLINE mBind_ #-} --- | Lifts a monadic action, ignored the result and combines it with a --- continuation. @mBind_ m f@ is the same as @lift m >>= f@, but does --- not require a 'Nullable' constraint on the stream type. -infixl 1 `mBind_` -mBind_ :: Monad m => m a -> Iteratee s m b -> Iteratee s m b -mBind_ m b = Iteratee $ \onDone onCont -> m >> runIter b onDone onCont - -{-# INLINE ioBind #-} --- | Lifts an IO action and combines it with a continuation. --- @ioBind m f@ is the same as @liftIO m >>= f@, but does not require a --- 'Nullable' constraint on the stream type. -infixl 1 `ioBind` -ioBind :: MonadIO m => IO a -> (a -> Iteratee s m b) -> Iteratee s m b -ioBind m f = Iteratee $ \onDone onCont -> liftIO m >>= \a -> runIter (f a) onDone onCont - -{-# INLINE ioBind_ #-} --- | Lifts an IO action, ignores its result, and combines it with a --- continuation. @ioBind_ m f@ is the same as @liftIO m >> f@, but does --- not require a 'Nullable' constraint on the stream type. -infixl 1 `ioBind_` -ioBind_ :: MonadIO m => IO a -> Iteratee s m b -> Iteratee s m b -ioBind_ m b = Iteratee $ \onDone onCont -> liftIO m >> runIter b onDone onCont - -class (MonadCatch m, MonadIO m) => MonadBracketIO m where - altmask :: ((forall a. m a -> m a) -> m b) -> m b - - -- | Runs an 'Iteratee' in between an 'IO' action to acquire a - -- resource and one to release it. 'Iteratee' can't be an instance - -- of 'CIO.MonadMask', so 'CIO.bracket' isn't defined for it. - -- However, if we restrict the acquire/release actions to 'IO', - -- which is the most important use case anyway, we can directly - -- implement this weaker version. - bracketIO :: IO a -> (a -> IO b) -> (a -> m c) -> m c - -instance MonadBracketIO IO where - {-# INLINE altmask #-} - altmask = CIO.mask - - {-# INLINE bracketIO #-} - bracketIO = CIO.bracket - -instance (MonadBracketIO m, Nullable s) => MonadBracketIO (Iteratee s m) where - {-# INLINE altmask #-} - altmask q = Iteratee $ \od oc -> altmask $ \u -> runIter (q $ ilift u) od oc - - {-# INLINE bracketIO #-} - bracketIO acquire release use = - Iteratee $ \od oc -> altmask $ \u -> - runIter (acquire `ioBind` \resource -> - ilift u (use resource) `CIO.onException` liftIO (release resource) >>= \result -> - release resource `ioBind_` return result) od oc - - --- | Convert one stream into another with the supplied mapping function. --- This function operates on whole chunks at a time, contrasting to --- @mapStream@ which operates on single elements. --- --- > unpacker :: Enumeratee B.ByteString [Word8] m a --- > unpacker = mapChunks B.unpack --- -mapChunks :: (NullPoint s) => (s -> s') -> Enumeratee s s' m a -mapChunks f = eneeCheckIfDonePass (icont . step) - where - step k (Chunk xs) = eneeCheckIfDonePass (icont . step) . k . Chunk $ f xs - step k str@(EOF mErr) = idone (k $ EOF mErr) str -{-# INLINE mapChunks #-} - --- | Convert a stream of @s@ to a stream of @s'@ using the supplied function. -mapChunksM :: (Monad m, NullPoint s) => (s -> m s') -> Enumeratee s s' m a -mapChunksM f = eneeCheckIfDonePass (icont . step) - where - step k (Chunk xs) = f xs `mBind` eneeCheckIfDonePass (icont . step) . k . Chunk - step k str = idone (liftI k) str -{-# INLINE mapChunksM #-} - --- |Convert one stream into another, not necessarily in lockstep. --- --- The transformer mapStream maps one element of the outer stream --- to one element of the nested stream. The transformer below is more --- general: it may take several elements of the outer stream to produce --- one element of the inner stream, or the other way around. --- The transformation from one stream to the other is specified as --- Iteratee s m s'. -convStream :: - (Monad m, Nullable s) => - Iteratee s m s' - -> Enumeratee s s' m a -convStream fi = eneeCheckIfDonePass check - where - check k (Just e) = throwRecoverableErr e (idone ()) >> check k Nothing - check k _ = isStreamFinished >>= maybe (step k) (idone (liftI k) . EOF . Just) - step k = fi >>= eneeCheckIfDonePass check . k . Chunk -{-# INLINABLE convStream #-} - --- |The most general stream converter. Given a function to produce iteratee --- transformers and an initial state, convert the stream using iteratees --- generated by the function while continually updating the internal state. -unfoldConvStream :: - (Monad m, Nullable s) => - (acc -> Iteratee s m (acc, s')) - -> acc - -> Enumeratee s s' m a -unfoldConvStream f acc0 = eneeCheckIfDonePass (check acc0) - where - check acc k (Just e) = throwRecoverableErr e (idone ()) >> check acc k Nothing - check acc k _ = isStreamFinished >>= - maybe (step acc k) (idone (liftI k) . EOF . Just) - step acc k = f acc >>= \(acc', s') -> - eneeCheckIfDonePass (check acc') . k . Chunk $ s' -{-# INLINABLE unfoldConvStream #-} - -unfoldConvStreamCheck - :: (Monad m, Nullable elo) - => (((Stream eli -> Iteratee eli m a) - -> Maybe SomeException - -> Iteratee elo m (Iteratee eli m a) - ) - -> Enumeratee elo eli m a - ) - -> (acc -> Iteratee elo m (acc, eli)) - -> acc - -> Enumeratee elo eli m a -unfoldConvStreamCheck checkDone f acc0 = checkDone (check acc0) - where - check acc k mX = isStreamFinished >>= - maybe (step acc k mX) (idone (icont k mX) . EOF . Just) - step acc k Nothing = f acc >>= \(acc', s') -> - (checkDone (check acc') . k $ Chunk s') - step acc k (Just ex) = throwRecoverableErr ex $ \str' -> - let i = f acc >>= \(acc', s') -> - (checkDone (check acc') . k $ Chunk s') - in joinIM $ enumChunk str' i -{-# INLINABLE unfoldConvStreamCheck #-} - --- | Collapse a nested iteratee. The inner iteratee is terminated by @EOF@. --- Errors are propagated through the result. --- --- The stream resumes from the point of the outer iteratee; any remaining --- input in the inner iteratee will be lost. --- Differs from 'Control.Monad.join' in that the inner iteratee is terminated, --- and may have a different stream type than the result. -joinI :: - (Monad m, Nullable s) => - Iteratee s m (Iteratee s' m a) - -> Iteratee s m a -joinI = (>>= - \inner -> Iteratee $ \od oc -> - let onDone x _ = od x (Chunk emptyP) - onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' - onCont _ (Just e) = runIter (throwErr e) od oc - onCont' _ e = runIter (throwErr (fromMaybe excDivergent e)) od oc - in runIter inner onDone onCont) -{-# INLINE joinI #-} - --- | Lift an iteratee inside a monad to an iteratee. -joinIM :: (Monad m) => m (Iteratee s m a) -> Iteratee s m a -joinIM mIter = Iteratee $ \od oc -> mIter >>= \iter -> runIter iter od oc - - --- ------------------------------------------------------------------------ --- Enumerators --- | Each enumerator takes an iteratee and returns an iteratee --- --- an Enumerator is an iteratee transformer. --- The enumerator normally stops when the stream is terminated --- or when the iteratee moves to the done state, whichever comes first. --- When to stop is of course up to the enumerator... - -type Enumerator s m a = Iteratee s m a -> m (Iteratee s m a) - --- |Applies the iteratee to the given stream. This wraps 'enumEof', --- 'enumErr', and 'enumPure1Chunk', calling the appropriate enumerator --- based upon 'Stream'. -enumChunk :: (Monad m) => Stream s -> Enumerator s m a -enumChunk (Chunk xs) = enumPure1Chunk xs -enumChunk (EOF Nothing) = enumEof -enumChunk (EOF (Just e)) = enumErr e - --- |The most primitive enumerator: applies the iteratee to the terminated --- stream. The result is the iteratee in the Done state. It is an error --- if the iteratee does not terminate on EOF. -enumEof :: (Monad m) => Enumerator s m a -enumEof iter = runIter iter onDone onCont - where - onDone x _str = return $ idone x (EOF Nothing) - onCont k Nothing = runIter (k (EOF Nothing)) onDone onCont' - onCont k e = return $ icont k e - onCont' _ Nothing = return $ throwErr excDivergent - onCont' k e = return $ icont k e - --- |Another primitive enumerator: tell the Iteratee the stream terminated --- with an error. -enumErr :: (Exception e, Monad m) => e -> Enumerator s m a -enumErr e iter = runIter iter onDone onCont - where - onDone x _ = return $ idone x (EOF . Just $ toException e) - onCont k Nothing = runIter (k (EOF (Just (toException e)))) onDone onCont' - onCont k e' = return $ icont k e' - onCont' _ Nothing = return $ throwErr excDivergent - onCont' k e' = return $ icont k e' - - -infixr 0 =$ - --- | Combines an Enumeratee from @s@ to @s'@ and an Iteratee that --- consumes @s'@ into an Iteratee which consumes @s@ -(=$) - :: (Nullable s, Monad m) - => Enumeratee s s' m a - -> Iteratee s' m a - -> Iteratee s m a -(=$) = (.) joinI - -infixl 1 $= - --- | Combines Enumerator which produces stream of @s@ and @Enumeratee@ --- which transforms stream of @s@ to stream --- of @s'@ to into Enumerator which produces stream of @s'@ -($=) - :: Monad m - => (forall a. Enumerator s m a) - -> Enumeratee s s' m b - -> Enumerator s' m b -($=) enum enee iter = enum (enee iter) >>= run - - --- | Enumeratee composition --- Run the second enumeratee within the first. In this example, stream2list --- is run within the 'takeStream 10', which is itself run within 'takeStream 15', resulting --- in 15 elements being consumed --- --- >>> run =<< enumPure1Chunk [1..1000 :: Int] (joinI $ (I.takeStream 15 ><> I.takeStream 10) I.stream2list) --- [1,2,3,4,5,6,7,8,9,10] --- -(><>) :: - (Nullable s1, Monad m) - => (forall x . Enumeratee s1 s2 m x) - -> Enumeratee s2 s3 m a - -> Enumeratee s1 s3 m a -f ><> g = joinI . f . g - --- | enumeratee composition with the arguments flipped, see '><>' -(<><) :: - (Nullable s1, Monad m) - => Enumeratee s2 s3 m a - -> (forall x. Enumeratee s1 s2 m x) - -> Enumeratee s1 s3 m a -f <>< g = joinI . g . f - --- | Combine enumeration over two streams. The merging enumeratee would --- typically be the result of 'Data.Iteratee.ListLike.merge' or --- 'Data.Iteratee.ListLike.mergeByChunks' (see @merge@ for example). -mergeEnums :: - (Nullable s2, Nullable s1, Monad m) - => Enumerator s1 m a -- ^ inner enumerator - -> Enumerator s2 (Iteratee s1 m) a -- ^ outer enumerator - -> Enumeratee s2 s1 (Iteratee s1 m) a -- ^ merging enumeratee - -> Enumerator s1 m a -mergeEnums e1 e2 etee i = e1 $ e2 (joinI . etee $ ilift lift i) >>= run -{-# INLINE mergeEnums #-} - --- | The pure 1-chunk enumerator --- --- It passes a given list of elements to the iteratee in one chunk --- This enumerator does no IO and is useful for testing of base parsing -enumPure1Chunk :: (Monad m) => s -> Enumerator s m a -enumPure1Chunk str iter = runIter iter idoneM onCont - where - onCont k Nothing = return $ k $ Chunk str - onCont k e = return $ icont k e - --- | Enumerate chunks from a list --- -enumList :: (Monad m) => [s] -> Enumerator s m a -enumList = go - where - go [] i = return i - go xs' i = runIter i idoneM (onCont xs') - where - onCont (x:xs) k Nothing = go xs . k $ Chunk x - onCont _ _ (Just e) = return $ throwErr e - onCont _ k Nothing = return $ icont k Nothing -{-# INLINABLE enumList #-} - --- | Checks if an iteratee has finished. --- --- This enumerator runs the iteratee, performing any monadic actions. --- If the result is True, the returned iteratee is done. -enumCheckIfDone :: (Monad m) => Iteratee s m a -> m (Bool, Iteratee s m a) -enumCheckIfDone iter = runIter iter onDone onCont - where - onDone x str = return (True, idone x str) - onCont k e = return (False, icont k e) -{-# INLINE enumCheckIfDone #-} - - --- |Create an enumerator from a callback function -enumFromCallback :: - (Monad m, NullPoint s) => - (st -> m (Either SomeException ((Bool, st), s))) - -> st - -> Enumerator s m a -enumFromCallback c = loop - where - loop st iter = runIter iter idoneM (onCont st) - check k (True, st') = loop st' . k . Chunk - check k (False,_st') = return . k . Chunk - onCont st k Nothing = c st >>= - either (return . k . EOF . Just) (uncurry (check k)) - onCont _st k j = return (icont k j) - --- |Create an enumerator from a callback function with an exception handler. --- The exception handler is called if an iteratee reports an exception. -enumFromCallbackCatch - :: (IException e, Monad m, NullPoint s) - => (st -> m (Either SomeException ((Bool, st), s))) - -> (e -> m (Maybe EnumException)) - -> st - -> Enumerator s m a -enumFromCallbackCatch c handler = loop - where - loop st iter = runIter iter idoneM (onCont st) - check k (True, st') = loop st' . k . Chunk - check k (False,_st') = return . k . Chunk - onCont st k Nothing = c st >>= - either (return . k . EOF . Just) (uncurry (check k)) - onCont st k j@(Just e) = case fromException e of - Just e' -> handler e' >>= - maybe (loop st . k $ Chunk emptyP) - (return . icont k . Just) . fmap toException - Nothing -> return (icont k j) -{-# INLINE enumFromCallbackCatch #-} - - diff --git a/src/Bio/Iteratee/List.hs b/src/Bio/Iteratee/List.hs deleted file mode 100644 --- a/src/Bio/Iteratee/List.hs +++ /dev/null @@ -1,810 +0,0 @@ --- |Monadic Iteratees: --- incremental input parsers, processors and transformers --- --- This module provides many basic iteratees from which more complicated --- iteratees can be built. In general these iteratees parallel those in --- @Data.List@, with some additions. - -module Bio.Iteratee.List ( - -- * Iteratees - -- ** Iteratee Utilities - isFinished - ,stream2list - ,stream2stream - -- ** Basic Iteratees - ,dropWhileStream - ,dropStream - ,headStream - ,tryHead - ,lastStream - ,heads - ,peekStream - ,roll - ,lengthStream - ,chunkLength - ,takeFromChunk - -- ** Nested iteratee combinators - ,breakStream - ,breakE - ,takeStream - ,takeUpTo - ,takeWhileE - ,mapStream - ,concatMapStream - ,concatMapStreamM - ,mapMaybeStream - ,filterStream - ,filterStreamM - ,groupStreamBy - ,groupStreamOn - ,mergeStreams - ,mergeByChunks - -- ** Folds - ,foldStream - -- * Enumerators - -- ** Basic enumerators - ,enumPureNChunk - -- ** Enumerator Combinators - ,enumWith - ,zipStreams - ,zipStreams3 - ,zipStreams4 - ,zipStreams5 - ,sequenceStreams_ - ,countConsumed - -- ** Monadic functions - ,mapStreamM - ,mapStreamM_ - ,foldStreamM - -- * Re-exported modules - ,module Bio.Iteratee.Iteratee -) -where - -import Bio.Iteratee.Iteratee -import Bio.Prelude -import Control.Monad.Trans.Class (MonadTrans(..)) - --- import qualified Data.ByteString as B - - --- Useful combinators for implementing iteratees and enumerators - --- | Check if a stream has received 'EOF'. -isFinished :: Nullable s => Iteratee s m Bool -isFinished = liftI check - where - check c@(Chunk xs) - | nullC xs = liftI check - | otherwise = idone False c - check s@(EOF _) = idone True s -{-# INLINE isFinished #-} - --- ------------------------------------------------------------------------ --- Primitive iteratees - --- |Read a stream to the end and return all of its elements as a list. --- This iteratee returns all data from the stream *strictly*. -stream2list :: Monad m => Iteratee [el] m [el] -stream2list = liftM concat getChunks -{-# INLINE stream2list #-} - --- |Read a stream to the end and return all of its elements as a stream. --- This iteratee returns all data from the stream *strictly*. -stream2stream :: (Monad m, Nullable s, Monoid s) => Iteratee s m s -stream2stream = liftM mconcat getChunks -{-# INLINE stream2stream #-} - - --- ------------------------------------------------------------------------ --- Parser combinators - --- |Attempt to read the next element of the stream and return it --- Raise a (recoverable) error if the stream is terminated. --- --- The analogue of @List.head@ --- --- Because @head@ can raise an error, it shouldn't be used when constructing --- iteratees for @convStream@. Use @tryHead@ instead. -headStream :: Iteratee [el] m el -headStream = liftI step - where - step (Chunk [ ]) = icont step Nothing - step (Chunk (hd:tl)) = idone hd (Chunk tl) - step stream = icont step (Just (setEOF stream)) -{-# INLINE headStream #-} - --- | Similar to @headStream@, except it returns @Nothing@ if the stream --- is terminated. -tryHead :: Iteratee [el] m (Maybe el) -tryHead = liftI step - where - step (Chunk [ ]) = liftI step - step (Chunk (hd:tl)) = idone (Just hd) (Chunk tl) - step stream = idone Nothing stream -{-# INLINE tryHead #-} - --- |Attempt to read the last element of the stream and return it --- Raise a (recoverable) error if the stream is terminated --- --- The analogue of @List.last@ -lastStream :: Iteratee [el] m el -lastStream = liftI (step Nothing) - where - step l (Chunk xs) - | nullC xs = liftI (step l) - | otherwise = liftI $ step (Just $ last xs) - step l s@(EOF _) = case l of - Nothing -> icont (step l) . Just . setEOF $ s - Just x -> idone x s -{-# INLINE lastStream #-} - - --- |Given a sequence of characters, attempt to match them against --- the characters on the stream. Return the count of how many --- characters matched. The matched characters are removed from the --- stream. --- For example, if the stream contains 'abd', then (heads 'abc') --- will remove the characters 'ab' and return 2. -heads :: (Monad m, Eq el) => [el] -> Iteratee [el] m Int -heads st | nullC st = return 0 -heads st = loopE 0 st - where - loopE cnt xs - | nullC xs = return cnt - | otherwise = liftI (step cnt xs) - step cnt str (Chunk []) = liftI (step cnt str) - step cnt [ ] stream = idone cnt stream - step cnt (y:ys) s@(Chunk (x:xs)) - | y == x = step (succ cnt) ys (Chunk xs) - | otherwise = idone cnt s - step cnt _ stream = idone cnt stream -{-# INLINE heads #-} - - --- |Look ahead at the next element of the stream, without removing --- it from the stream. --- Return @Just c@ if successful, return @Nothing@ if the stream is --- terminated by 'EOF'. -peekStream :: Iteratee [el] m (Maybe el) -peekStream = liftI step - where - step (Chunk [ ]) = liftI step - step s@(Chunk (x:_)) = idone (Just x) s - step stream = idone Nothing stream -{-# INLINE peekStream #-} - --- | Return a chunk of @t@ elements length while consuming @d@ elements --- from the stream. Useful for creating a 'rolling average' with --- 'convStream'. -roll - :: Monad m - => Int -- ^ length of chunk (t) - -> Int -- ^ amount to consume (d) - -> Iteratee [el] m [[el]] -roll t d | t > d = liftI step - where - step (Chunk vec) - | length vec >= t = - idone [take t vec] (Chunk $ drop d vec) - | null vec = liftI step - | otherwise = liftI (step' vec) - step stream = idone empty stream - step' v1 (Chunk vec) = step . Chunk $ v1 `mappend` vec - step' v1 stream = idone [v1] stream -roll t d = do r <- joinI (takeStream t stream2stream) - dropStream (d-t) - return [r] - -- d is >= t, so this version works -{-# INLINE roll #-} - - --- |Drop n elements of the stream, if there are that many. --- --- The analogue of @List.drop@ -dropStream :: Int -> Iteratee [el] m () -dropStream 0 = idone () (Chunk emptyP) -dropStream n' = liftI (step n') - where - step n (Chunk str) - | length str < n = liftI (step (n - length str)) - | otherwise = idone () (Chunk (drop n str)) - step _ stream = idone () stream -{-# INLINE dropStream #-} - --- |Skip all elements while the predicate is true. --- --- The analogue of @List.dropWhile@ -dropWhileStream :: (el -> Bool) -> Iteratee [el] m () -dropWhileStream p = liftI step - where - step (Chunk str) - | null rest = liftI step - | otherwise = idone () (Chunk rest) - where - rest = dropWhile p str - step stream = idone () stream -{-# INLINE dropWhileStream #-} - - --- | Return the total length of the remaining part of the stream. --- --- This forces evaluation of the entire stream. --- --- The analogue of @List.length@ -lengthStream :: Num a => Iteratee [el] m a -lengthStream = liftI (step 0) - where - step !i (Chunk xs) = liftI (step $ i + fromIntegral (length xs)) - step !i stream = idone i stream -{-# INLINE lengthStream #-} - --- | Get the length of the current chunk, or @Nothing@ if 'EOF'. --- --- This function consumes no input. -chunkLength :: Iteratee [el] m (Maybe Int) -chunkLength = liftI step - where - step s@(Chunk xs) = idone (Just $ length xs) s - step stream = idone Nothing stream -{-# INLINE chunkLength #-} - --- | Take @n@ elements from the current chunk, or the whole chunk if --- @n@ is greater. -takeFromChunk :: Int -> Iteratee [el] m [el] -takeFromChunk n | n <= 0 = idone emptyP (Chunk emptyP) -takeFromChunk n = liftI step - where - step (Chunk xs) = let (h,t) = splitAt n xs in idone h $ Chunk t - step stream = idone emptyP stream -{-# INLINE takeFromChunk #-} - --- |Takes an element predicate and returns the (possibly empty) prefix of --- the stream. None of the characters in the string satisfy the character --- predicate. --- If the stream is not terminated, the first character of the remaining stream --- satisfies the predicate. --- --- N.B. 'breakE' should be used in preference to @breakStream@. --- @breakStream@ will retain all data until the predicate is met, which may --- result in a space leak. --- --- The analogue of @List.break@ - -breakStream :: (el -> Bool) -> Iteratee [el] m [el] -breakStream cpred = icont (step mempty) Nothing - where - step bfr (Chunk str) - | null str = icont (step bfr) Nothing - | otherwise = case break cpred str of - (str', tail') - | null tail' -> icont (step (bfr `mappend` str)) Nothing - | otherwise -> idone (bfr `mappend` str') (Chunk tail') - step bfr stream = idone bfr stream -{-# INLINE breakStream #-} - --- --------------------------------------------------- --- The converters show a different way of composing two iteratees: --- `vertical' rather than `horizontal' - --- |Takes an element predicate and an iteratee, running the iteratee --- on all elements of the stream until the predicate is met. --- --- the following rule relates @break@ to @breakE@ --- @break@ pred === @joinI@ (@breakE@ pred stream2stream) --- --- @breakE@ should be used in preference to @break@ whenever possible. -breakE :: (el -> Bool) -> Enumeratee [el] [el] m a -breakE cpred = eneeCheckIfDonePass (icont . step) - where - step k (Chunk s) - | null s = liftI (step k) - | otherwise = case break cpred s of - (str', tail') - | null tail' -> eneeCheckIfDonePass (icont . step) . k $ Chunk str' - | otherwise -> idone (k $ Chunk str') (Chunk tail') - step k stream = idone (liftI k) stream -{-# INLINE breakE #-} - --- |Read n elements from a stream and apply the given iteratee to the --- stream of the read elements. Unless the stream is terminated early, we --- read exactly n elements, even if the iteratee has accepted fewer. --- --- The analogue of @List.take@ -takeStream :: - Monad m - => Int -- ^ number of elements to consume - -> Enumeratee [el] [el] m a -takeStream n' iter - | n' <= 0 = return iter - | otherwise = Iteratee $ \od oc -> runIter iter (on_done od oc) (on_cont od oc) - where - on_done od oc x _ = runIter (dropStream n' >> return (return x)) od oc - on_cont od oc k Nothing = if n' == 0 then od (liftI k) (Chunk mempty) - else runIter (liftI (step n' k)) od oc - on_cont od oc _ (Just e) = runIter (dropStream n' >> throwErr e) od oc - step n k (Chunk str) - | null str = liftI (step n k) - | length str <= n = takeStream (n - length str) $ k (Chunk str) - | otherwise = idone (k (Chunk s1)) (Chunk s2) - where (s1, s2) = splitAt n str - step _n k stream = idone (liftI k) stream -{-# INLINE takeStream #-} - --- |Read n elements from a stream and apply the given iteratee to the --- stream of the read elements. If the given iteratee accepted fewer --- elements, we stop. --- This is the variation of 'takeStream' with the early termination --- of processing of the outer stream once the processing of the inner stream --- finished early. --- --- Iteratees composed with 'takeUpTo' will consume only enough elements to --- reach a done state. Any remaining data will be available in the outer --- stream. --- --- > > let iter = do --- > h <- joinI $ takeUpTo 5 I.head --- > t <- stream2list --- > return (h,t) --- > --- > > enumPureNChunk [1..10::Int] 3 iter >>= run >>= print --- > (1,[2,3,4,5,6,7,8,9,10]) --- > --- > > enumPureNChunk [1..10::Int] 7 iter >>= run >>= print --- > (1,[2,3,4,5,6,7,8,9,10]) --- --- in each case, @I.head@ consumes only one element, returning the remaining --- 4 elements to the outer stream -takeUpTo :: Monad m => Int -> Enumeratee [el] [el] m a -takeUpTo i iter - | i <= 0 = idone iter (Chunk emptyP) - | otherwise = Iteratee $ \od oc -> - runIter iter (onDone od oc) (onCont od oc) - where - onDone od oc x str = runIter (idone (return x) str) od oc - onCont od oc k Nothing = if i == 0 then od (liftI k) (Chunk mempty) - else runIter (liftI (step i k)) od oc - onCont od oc _ (Just e) = runIter (throwErr e) od oc - step n k (Chunk str) - | null str = liftI (step n k) - | length str < n = takeUpTo (n - length str) $ k (Chunk str) - | otherwise = - -- check to see if the inner iteratee has completed, and if so, - -- grab any remaining stream to put it in the outer iteratee. - -- the outer iteratee is always complete at this stage, although - -- the inner may not be. - let (s1, s2) = splitAt n str - in Iteratee $ \od' _ -> do - res <- runIter (k (Chunk s1)) (\a s -> return $ Left (a, s)) - (\k' e -> return $ Right (k',e)) - case res of - Left (a,Chunk s1') -> od' (return a) - (Chunk $ s1' ++ s2) - Left (a,s') -> od' (idone a s') (Chunk s2) - Right (k',e) -> od' (icont k' e) (Chunk s2) - step _ k stream = idone (liftI k) stream -{-# INLINE takeUpTo #-} - - --- |Takes an element predicate and an iteratee, running the iteratee --- on all elements of the stream while the predicate is met. --- --- This is preferred to @takeWhile@. -takeWhileE :: (el -> Bool) -> Enumeratee [el] [el] m a -takeWhileE = breakE . (not .) -{-# INLINEABLE takeWhileE #-} - --- | Map a function over an 'Iteratee'. --- This one is reimplemented and differs from the the one in --- "Data.Iteratee.ListLike" in so far that it doesn't pass on an 'EOF' --- received in the input, which is the expected behavior. -mapStream :: (el -> el') -> Enumeratee [el] [el'] m a -mapStream = mapChunks . map -{-# INLINE mapStream #-} - --- | Apply a function to the elements of a stream, concatenate the --- results into a stream. No giant intermediate list is produced. -concatMapStream :: Monoid t => (a -> t) -> Enumeratee [a] t m r -concatMapStream = mapChunks . foldMap -{-# INLINE concatMapStream #-} - --- | Apply a monadic function to the elements of a stream, concatenate --- the results into a stream. No giant intermediate list is produced. -concatMapStreamM :: Monad m => (a -> m t) -> Enumeratee [a] t m r -concatMapStreamM f = eneeCheckIfDone (liftI . go) - where - go k (EOF mx) = idone (liftI k) (EOF mx) - go k (Chunk xs) | null xs = liftI (go k) - | otherwise = f (head xs) `mBind` - eneeCheckIfDone (flip go (Chunk (tail xs))) . k . Chunk -{-# INLINE concatMapStreamM #-} - -mapMaybeStream :: (a -> Maybe b) -> Enumeratee [a] [b] m r -mapMaybeStream = mapChunks . mapMaybe -{-# INLINE mapMaybeStream #-} - --- |Creates an 'enumeratee' with only elements from the stream that --- satisfy the predicate function. The outer stream is completely consumed. --- --- The analogue of @List.filter@ -filterStream :: (el -> Bool) -> Enumeratee [el] [el] m a -filterStream p = mapChunks (filter p) -{-# INLINE filterStream #-} - --- | Apply a monadic filter predicate to an 'Iteratee'. -filterStreamM :: Monad m => (a -> m Bool) -> Enumeratee [a] [a] m r -filterStreamM k = mapChunksM (go id) - where - go acc [ ] = return $! acc empty - go acc (h:t) = do p <- k h - let acc' = if p then (:) h . acc else acc - go acc' t -{-# INLINE filterStreamM #-} - --- | Grouping on 'Iteratee's. @groupStreamOn proj inner outer@ executes --- @inner (proj e)@, where @e@ is the first input element, to obtain an --- 'Iteratee' @i@, then passes elements @e@ to @i@ as long as @proj e@ --- produces the same result. If @proj e@ changes or the input ends, the --- pair of @proj e@ and the result of @run i@ is passed to @outer@. At --- end of input, the resulting @outer@ is returned. -groupStreamOn :: (Monad m, Eq t1) - => (e -> t1) - -> (t1 -> m (Iteratee [e] m t2)) - -> Enumeratee [e] [(t1, t2)] m a -groupStreamOn proj inner = eneeCheckIfDonePass (icont . step) - where - step outer (EOF mx) = idone (liftI outer) $ EOF mx - step outer (Chunk [ ]) = liftI $ step outer - step outer c@(Chunk (h:_)) = let x = proj h - in lift (inner x) >>= \i -> step' x i outer c - - -- We want to feed a 'Chunk' to the inner 'Iteratee', which might be - -- finished. In that case, we would want to abort, but we cannot, - -- since the outer iteration is still going on. So instead we - -- discard data we would have fed to the inner 'Iteratee'. (Use of - -- 'enumPure1Chunk' is not appropriate, it would accumulate the - -- data, just to have it discarded by the 'run' that eventually - -- happens. - - step' c it outer (Chunk as) - | null as = liftI $ step' c it outer - | (l,r) <- span ((==) c . proj) as, not (null l) = - let od a _str = idoneM a $ EOF Nothing - oc k Nothing = return $ k (Chunk l) - oc k m = icontM k m - in lift (runIter it od oc) >>= \it' -> step' c it' outer (Chunk r) - - step' c it outer str = - lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [(c,b)] - - --- | Grouping on 'Iteratee's. @groupStreamBy cmp inner outer@ executes --- @inner@ to obtain an 'Iteratee' @i@, then passes elements @e@ to @i@ --- as long as @cmp e0 e@, where @e0@ is some preceeding element, is --- true. Else, the result of @run i@ is passed to @outer@ and --- 'groupStreamBy' restarts. At end of input, the resulting @outer@ is --- returned. -groupStreamBy :: Monad m - => (t -> t -> Bool) - -> m (Iteratee [t] m t2) - -> Enumeratee [t] [t2] m a -groupStreamBy cmp inner = eneeCheckIfDonePass (icont . step) - where - step outer (EOF mx) = idone (liftI outer) $ EOF mx - step outer (Chunk [ ]) = liftI $ step outer - step outer c@(Chunk (h:_)) = lift inner >>= \i -> step' h i outer c - - step' c it outer (Chunk as) - | null as = liftI $ step' c it outer - | (l,r) <- span (cmp c) as, not (null l) = - let od a _str = idoneM a $ EOF Nothing - oc k Nothing = return $ k (Chunk l) - oc k m = icontM k m - in lift (runIter it od oc) >>= \it' -> step' (head l) it' outer (Chunk r) - - step' _ it outer str = - lift (run it) >>= \b -> eneeCheckIfDone (`step` str) . outer $ Chunk [b] - - --- | @mergeStreams@ offers another way to nest iteratees: as a monad stack. --- This allows for the possibility of interleaving data from multiple --- streams. --- --- > -- print each element from a stream of lines. --- > logger :: (MonadIO m) => Iteratee [ByteString] m () --- > logger = mapStreamM_ (liftIO . putStrLn . B.unpack) --- > --- > -- combine alternating lines from two sources --- > -- To see how this was derived, follow the types from --- > -- 'ileaveLines logger' and work outwards. --- > run =<< enumFile 10 "file1" (joinI $ enumLinesBS $ --- > ( enumFile 10 "file2" . joinI . enumLinesBS $ joinI --- > (ileaveLines logger)) >>= run) --- > --- > ileaveLines :: (Functor m, Monad m) --- > => Enumeratee [ByteString] [ByteString] (Iteratee [ByteString] m) --- > [ByteString] --- > ileaveLines = mergeStreams (\l1 l2 -> --- > [B.pack "f1:\n\t" ,l1 ,B.pack "f2:\n\t" ,l2 ] --- > --- > --- -mergeStreams :: Monad m => (el1 -> el2 -> b) -> Enumeratee [el2] b (Iteratee [el1] m) a -mergeStreams f = convStream $ liftM2 f (lift headStream) headStream -{-# INLINE mergeStreams #-} - --- | A version of mergeStreams which operates on chunks instead of --- elements. --- --- mergeByChunks offers more control than 'mergeStreams'. --- 'mergeStreams' terminates when the first stream terminates, however --- mergeByChunks will continue until both streams are exhausted. --- --- 'mergeByChunks' guarantees that both chunks passed to the merge --- function will have the same number of elements, although that number --- may vary between calls. -mergeByChunks :: - Monad m - => ([el1] -> [el2] -> c3) -- ^ merge function - -> ([el1] -> c3) - -> ([el2] -> c3) - -> Enumeratee [el2] c3 (Iteratee [el1] m) a -mergeByChunks f f1 f2 = unfoldConvStream iter (0 :: Int) - where - iter 1 = (\x -> (1,f1 x)) `liftM` lift getChunk - iter 2 = (\x -> (2,f2 x)) `liftM` getChunk - iter _ = do - ml1 <- lift chunkLength - ml2 <- chunkLength - case (ml1, ml2) of - (Just l1, Just l2) -> do - let tval = min l1 l2 - c1 <- lift $ takeFromChunk tval - c2 <- takeFromChunk tval - return (0, f c1 c2) - (Just _, Nothing) -> iter 1 - (Nothing, _) -> iter 2 -{-# INLINE mergeByChunks #-} - --- ------------------------------------------------------------------------ --- Folds - --- | Left-associative fold that is strict in the accumulator. --- This function should be used in preference to 'foldl' whenever possible. --- --- The analogue of @List.foldl'@. -foldStream :: (a -> el -> a) -> a -> Iteratee [el] m a -foldStream f i = liftI (step i) - where - step acc (Chunk xs) - | null xs = liftI (step acc) - | otherwise = liftI (step $! foldl' f acc xs) - step acc stream = idone acc stream -{-# INLINE foldStream #-} - --- ------------------------------------------------------------------------ --- Zips - --- |Enumerate two iteratees over a single stream simultaneously. --- --- Compare to @List.zip@. -zipStreams - :: Monad m - => Iteratee [el] m a - -> Iteratee [el] m b - -> Iteratee [el] m (a, b) -zipStreams x0 y0 = do - -- need to check if both iteratees are initially finished. If so, - -- we don't want to push a chunk which will be dropped - (a', x') <- lift $ runIter x0 od oc - (b', y') <- lift $ runIter y0 od oc - case checkDone a' b' of - Just (Right (a,b,s)) -> idone (a,b) s -- 's' may be EOF, needs to stay - Just (Left (Left a)) -> liftM ((,) a) y' - Just (Left (Right b)) -> liftM (flip (,) b) x' - Nothing -> liftI (step x' y') - where - step x y (Chunk xs) | nullC xs = liftI (step x y) - step x y (Chunk xs) = do - (a', x') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs x - (b', y') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs y - case checkDone a' b' of - Just (Right (a,b,s)) -> idone (a,b) s - Just (Left (Left a)) -> liftM ((,) a) y' - Just (Left (Right b)) -> liftM (flip (,) b) x' - Nothing -> liftI (step x' y') - step x y (EOF err) = joinIM $ case err of - Nothing -> (liftM2.liftM2) (,) (enumEof x) (enumEof y) - Just e -> (liftM2.liftM2) (,) (enumErr e x) (enumErr e y) - - od a s = return (Just (a, s), idone a s) - oc k e = return (Nothing , icont k e) - - checkDone r1 r2 = case (r1, r2) of - (Just (a, s1), Just (b,s2)) -> Just $ Right (a, b, shorter s1 s2) - (Just (a, _), Nothing) -> Just . Left $ Left a - (Nothing, Just (b, _)) -> Just . Left $ Right b - (Nothing, Nothing) -> Nothing - - shorter c1@(Chunk xs) c2@(Chunk ys) - | length xs < length ys = c1 - | otherwise = c2 - shorter e@(EOF _) _ = e - shorter _ e@(EOF _) = e -{-# INLINE zipStreams #-} - -zipStreams3 - :: Monad m - => Iteratee [el] m a -> Iteratee [el] m b - -> Iteratee [el] m c -> Iteratee [el] m (a, b, c) -zipStreams3 a b c = zipStreams a (zipStreams b c) >>= - \(r1, (r2, r3)) -> return (r1, r2, r3) -{-# INLINE zipStreams3 #-} - -zipStreams4 - :: Monad m - => Iteratee [el] m a -> Iteratee [el] m b - -> Iteratee [el] m c -> Iteratee [el] m d - -> Iteratee [el] m (a, b, c, d) -zipStreams4 a b c d = zipStreams a (zipStreams3 b c d) >>= - \(r1, (r2, r3, r4)) -> return (r1, r2, r3, r4) -{-# INLINE zipStreams4 #-} - -zipStreams5 - :: Monad m - => Iteratee [el] m a -> Iteratee [el] m b - -> Iteratee [el] m c -> Iteratee [el] m d - -> Iteratee [el] m e -> Iteratee [el] m (a, b, c, d, e) -zipStreams5 a b c d e = zipStreams a (zipStreams4 b c d e) >>= - \(r1, (r2, r3, r4, r5)) -> return (r1, r2, r3, r4, r5) -{-# INLINE zipStreams5 #-} - --- | Enumerate over two iteratees in parallel as long as the first iteratee --- is still consuming input. The second iteratee will be terminated with EOF --- when the first iteratee has completed. An example use is to determine --- how many elements an iteratee has consumed: --- --- > snd <$> enumWith (dropWhile (<5)) length --- --- Compare to @zipStreams@ -enumWith - :: Monad m - => Iteratee [el] m a - -> Iteratee [el] m b - -> Iteratee [el] m (a, b) -enumWith i1 i2 = do - -- as with zipStreams, first check to see if the initial iteratee is complete, - -- otherwise data would be dropped. - -- running the second iteratee as well to prevent a monadic effect mismatch - -- although I think that would be highly unlikely to happen in common - -- code - (a', x') <- lift $ runIter i1 od oc - (_, y') <- lift $ runIter i2 od oc - case a' of - Just (a, s) -> flip idone s =<< lift (liftM ((,) a) $ run i2) - Nothing -> go x' y' - where - od a s = return (Just (a, s), idone a s) - oc k e = return (Nothing , icont k e) - - getUsed xs (Chunk ys) = take (length xs - length ys) xs - getUsed xs (EOF _) = xs - - go x y = liftI step - where - step (Chunk xs) | nullC xs = liftI step - step (Chunk xs) = do - (a', x') <- lift $ (\i -> runIter i od oc) =<< enumPure1Chunk xs x - case a' of - Just (a, s) -> do - b <- lift $ run =<< enumPure1Chunk (getUsed xs s) y - idone (a, b) s - Nothing -> lift (enumPure1Chunk xs y) >>= go x' - step (EOF err) = joinIM $ case err of - Nothing -> (liftM2.liftM2) (,) (enumEof x) (enumEof y) - Just e -> (liftM2.liftM2) (,) (enumErr e x) (enumErr e y) -{-# INLINE enumWith #-} - --- |Enumerate a list of iteratees over a single stream simultaneously --- and discard the results. This is a different behavior than Prelude's --- sequence_ which runs iteratees in the list one after the other. --- --- Compare to @Prelude.sequence_@. -sequenceStreams_ - :: Monad m - => [Iteratee [el] m a] - -> Iteratee [el] m () -sequenceStreams_ = self - where - self is = liftI step - where - step (Chunk xs) | null xs = liftI step - step s@(Chunk _) = do - -- give a chunk to each iteratee - is' <- lift $ mapM (enumChunk s) is - -- filter done iteratees - (done, notDone) <- lift $ partition fst `liftM` mapM enumCheckIfDone is' - if null notDone - then idone () <=< remainingStream $ map snd done - else self $ map snd notDone - step s@(EOF _) = do - s' <- remainingStream <=< lift $ mapM (enumChunk s) is - case s' of - EOF (Just e) -> throwErr e - _ -> idone () s' - - -- returns the unconsumed part of the stream; "sequenceStreams_ is" consumes as - -- much of the stream as the iteratee in is that consumes the most; e.g. - -- sequenceStreams_ [I.head, I.last] consumes whole stream - remainingStream :: Monad m => [Iteratee [el] m a] -> Iteratee [el] m (Stream [el]) - remainingStream is = lift $ - return . foldl1 shorter <=< mapM (\i -> runIter i od oc) $ is - where - od _ s = return s - oc _ e = return $ case e of - Nothing -> mempty - _ -> EOF e - - -- return the shorter one of two streams; errors are propagated with the - -- priority given to the "left" - shorter c1@(Chunk xs) c2@(Chunk ys) - | length xs < length ys = c1 - | otherwise = c2 - shorter (EOF e1 ) (EOF e2 ) = EOF (e1 `mplus` e2) - shorter e@(EOF _) _ = e - shorter _ e@(EOF _) = e - --- |Transform an iteratee into one that keeps track of how much data it --- consumes. -countConsumed :: (Monad m, Integral n) => Iteratee [el] m a -> Iteratee [el] m (a, n) -countConsumed i = go 0 (const i) (Chunk emptyP) - where - go !n f str@(EOF _) = flip (,) n `liftM` f str - go !n f str@(Chunk c) = Iteratee rI - where - newLen = n + fromIntegral (length c) - rI od oc = runIter (f str) onDone onCont - where - onDone a str'@(Chunk c') = - od (a, newLen - fromIntegral (length c')) str' - onDone a str'@(EOF _) = od (a, n) str' - onCont f' = oc (go newLen f') -{-# INLINE countConsumed #-} - --- ------------------------------------------------------------------------ --- Enumerators - --- |The pure n-chunk enumerator --- It passes a given stream of elements to the iteratee in @n@-sized chunks. -enumPureNChunk :: Monad m => [el] -> Int -> Enumerator [el] m a -enumPureNChunk str n iter - | null str = return iter - | n > 0 = enum' str iter - | otherwise = error $ "enumPureNChunk called with n==" ++ show n - where - enum' str' iter' - | null str' = return iter' - | otherwise = let (s1, s2) = splitAt n str' - on_cont k Nothing = enum' s2 . k $ Chunk s1 - on_cont k e = return $ icont k e - in runIter iter' idoneM on_cont -{-# INLINE enumPureNChunk #-} - --- ------------------------------------------------------------------------ --- Monadic functions - --- | Maps a monadic function over the elements of the stream and ignores --- the result. -mapStreamM_ :: Monad m => (el -> m b) -> Iteratee [el] m () -mapStreamM_ = mapChunksM_ . mapM_ -{-# INLINE mapStreamM_ #-} - --- | Maps a monadic function over an 'Iteratee'. -mapStreamM :: Monad m => (el -> m el') -> Enumeratee [el] [el'] m a -mapStreamM = mapChunksM . mapM -{-# INLINE mapStreamM #-} - --- | Folds a monadic function over an 'Iteratee'. -foldStreamM :: Monad m => (b -> a -> m b) -> b -> Iteratee [a] m b -foldStreamM = foldChunksM . foldM -{-# INLINE foldStreamM #-} diff --git a/src/Bio/Iteratee/ZLib.hsc b/src/Bio/Iteratee/ZLib.hsc deleted file mode 100644 --- a/src/Bio/Iteratee/ZLib.hsc +++ /dev/null @@ -1,741 +0,0 @@ --- Stolen from iteratee-compress module, which doesn't work due to --- dependency problems. Modified for proper early-out behaviour. -module Bio.Iteratee.ZLib - ( - -- * Enumeratees - enumInflate, - enumInflateAny, - enumDeflate, - -- * Exceptions - ZLibParamsException(..), - ZLibException(..), - -- * Parameters - CompressParams(..), - defaultCompressParams, - DecompressParams(..), - defaultDecompressParams, - Format(..), - CompressionLevel(..), - Method(..), - WindowBits(..), - MemoryLevel(..), - CompressionStrategy(..), - enumSyncFlush, - enumFullFlush, - enumBlockFlush, - ) -where -#include - -import Bio.Iteratee -import Control.Applicative -import Control.Exception -import Control.Monad ( liftM, liftM2 ) -import Data.ByteString as BS -import Data.ByteString.Internal -import Data.Foldable -import Data.Typeable -import Foreign -import Foreign.C -import Prelude -#ifdef DEBUG -import qualified Foreign.Concurrent as C -import System.IO (stderr) -import qualified System.IO as IO -#endif - --- | Denotes error is user-supplied parameter -data ZLibParamsException - = IncorrectCompressionLevel !Int - -- ^ Incorrect compression level was chosen - | IncorrectWindowBits !Int - -- ^ Incorrect number of window bits was chosen - | IncorrectMemoryLevel !Int - -- ^ Incorrect memory level was chosen - deriving (Eq,Typeable) - --- | Denotes error in compression and decompression -data ZLibException - = NeedDictionary - -- ^ Decompression requires user-supplied dictionary (not supported) - | BufferError - -- ^ Buffer error - denotes a library error --- | File Error - | StreamError - -- ^ State of steam inconsistent - | DataError - -- ^ Input data corrupted - | MemoryError - -- ^ Not enough memory - | VersionError - -- ^ Version error - | Unexpected !CInt - -- ^ Unexpected or unknown error - please report as bug - | IncorrectState - -- ^ Incorrect state - denotes error in library - deriving (Eq,Typeable) - --- | Denotes the flush that can be sent to stream -data ZlibFlush - = SyncFlush - -- ^ All pending output is flushed and all input that is available is sent - -- to inner Iteratee. - | FullFlush - -- ^ Flush all pending output and reset the compression state. It allows to - -- restart from this point if compression was damaged but it can seriously - -- affect the compression rate. - -- - -- It may be only used during compression. - | Block - -- ^ If the iteratee is compressing it requests to stop when next block is - -- emmited. On the beginning it skips only header if and only if it exists. - deriving (Eq,Typeable) - -instance Show ZlibFlush where - show SyncFlush = "zlib: flush requested" - show FullFlush = "zlib: full flush requested" - show Block = "zlib: block flush requested" - -instance Exception ZlibFlush - -fromFlush :: ZlibFlush -> CInt -fromFlush SyncFlush = #{const Z_SYNC_FLUSH} -fromFlush FullFlush = #{const Z_FULL_FLUSH} -fromFlush Block = #{const Z_BLOCK} - -instance Show ZLibParamsException where - show (IncorrectCompressionLevel lvl) - = "zlib: incorrect compression level " ++ show lvl - show (IncorrectWindowBits lvl) - = "zlib: incorrect window bits " ++ show lvl - show (IncorrectMemoryLevel lvl) - = "zlib: incorrect memory level " ++ show lvl - -instance Show ZLibException where - show NeedDictionary = "zlib: needs dictionary" - show BufferError = "zlib: no progress is possible (internal error)" --- show FileError = "zlib: file I/O error" - show StreamError = "zlib: stream error" - show DataError = "zlib: data error" - show MemoryError = "zlib: memory error" - show VersionError = "zlib: version error" - show (Unexpected lvl) = "zlib: unknown error " ++ show lvl - show IncorrectState = "zlib: incorrect state" - -instance Exception ZLibParamsException -instance Exception ZLibException - -newtype ZStream = ZStream (ForeignPtr ZStream) -withZStream :: ZStream -> (Ptr ZStream -> IO a) -> IO a -withZStream (ZStream fptr) = withForeignPtr fptr - - --- Following code is copied from Duncan Coutts zlib haskell library version --- 0.5.2.0 ((c) 2006-2008 Duncan Coutts, published on BSD licence) and adapted - --- | Set of parameters for compression. For sane defaults use --- 'defaultCompressParams' -data CompressParams = CompressParams { - compressLevel :: !CompressionLevel, - compressMethod :: !Method, - compressWindowBits :: !WindowBits, - compressMemoryLevel :: !MemoryLevel, - compressStrategy :: !CompressionStrategy, - -- | The size of output buffer. That is the size of 'Chunk's that will be - -- emitted to inner iterator (except the last 'Chunk'). - compressBufferSize :: !Int, - compressDictionary :: !(Maybe ByteString) - } - -defaultCompressParams :: CompressParams -defaultCompressParams - = CompressParams DefaultCompression Deflated DefaultWindowBits - DefaultMemoryLevel DefaultStrategy (8*1024) Nothing - --- | Set of parameters for decompression. For sane defaults see --- 'defaultDecompressParams'. -data DecompressParams = DecompressParams { - -- | Window size - it have to be at least the size of - -- 'compressWindowBits' the stream was compressed with. - -- - -- Default in 'defaultDecompressParams' is the maximum window size - - -- please do not touch it unless you know what you are doing. - decompressWindowBits :: !WindowBits, - -- | The size of output buffer. That is the size of 'Chunk's that will be - -- emitted to inner iterator (except the last 'Chunk'). - decompressBufferSize :: !Int, - decompressDictionary :: !(Maybe ByteString) - } - -defaultDecompressParams :: DecompressParams -defaultDecompressParams = DecompressParams DefaultWindowBits (8*1024) Nothing - --- | Specify the format for compression and decompression -data Format - = GZip - -- ^ The gzip format is widely used and uses a header with checksum and - -- some optional metadata about the compress file. - -- - -- It is intended primarily for compressing individual files but is also - -- used for network protocols such as HTTP. - -- - -- The format is described in RFC 1952 - -- . - | Zlib - -- ^ The zlib format uses a minimal header with a checksum but no other - -- metadata. It is designed for use in network protocols. - -- - -- The format is described in RFC 1950 - -- - | Raw - -- ^ The \'raw\' format is just the DEFLATE compressed data stream without - -- and additionl headers. - -- - -- Thr format is described in RFC 1951 - -- - | GZipOrZlib - -- ^ "Format" for decompressing a 'Zlib' or 'GZip' stream. - deriving (Eq) - --- | The compression level specify the tradeoff between speed and compression. -data CompressionLevel - = DefaultCompression - -- ^ Default compression level set at 6. - | NoCompression - -- ^ No compression, just a block copy. - | BestSpeed - -- ^ The fastest compression method (however less compression) - | BestCompression - -- ^ The best compression method (however slowest) - | CompressionLevel Int - -- ^ Compression level set by number from 1 to 9 - --- | Specify the compression method. -data Method - = Deflated - -- ^ \'Deflate\' is so far the only method supported. - --- | This specify the size of compression level. Larger values result in better --- compression at the expense of highier memory usage. --- --- The compression window size is 2 to the power of the value of the window --- bits. --- --- The total memory used depends on windows bits and 'MemoryLevel'. -data WindowBits - = WindowBits Int - -- ^ The size of window bits. It have to be between @8@ (which corresponds - -- to 256b i.e. 32B) and @15@ (which corresponds to 32 kib i.e. 4kiB). - | DefaultWindowBits - -- ^ The default window size which is 4kiB - --- | The 'MemoryLevel' specifies how much memory should be allocated for the --- internal state. It is a tradeoff between memory usage, speed and --- compression. --- Using more memory allows faster and better compression. --- --- The memory used for interal state, excluding 'WindowBits', is 512 bits times --- 2 to power of memory level. --- --- The total amount of memory use depends on the 'WindowBits' and --- 'MemoryLevel'. -data MemoryLevel - = DefaultMemoryLevel - -- ^ Default memory level set to 8. - | MinMemoryLevel - -- ^ Use the small amount of memory (equivalent to memory level 1) - i.e. - -- 1024b or 256 B. - -- It slow and reduces the compresion ratio. - | MaxMemoryLevel - -- ^ Maximum memory level for optimal compression speed (equivalent to - -- memory level 9). - -- The internal state is 256kib or 32kiB. - | MemoryLevel Int - -- ^ A specific level. It have to be between 1 and 9. - --- | Tunes the compress algorithm but does not affact the correctness. -data CompressionStrategy - = DefaultStrategy - -- ^ Default strategy - | Filtered - -- ^ Use the filtered compression strategy for data produced by a filter - -- (or predictor). Filtered data consists mostly of small values with a - -- somewhat random distribution. In this case, the compression algorithm - -- is tuned to compress them better. The effect of this strategy is to - -- force more Huffman coding and less string matching; it is somewhat - -- intermediate between 'DefaultStrategy' and 'HuffmanOnly'. - | HuffmanOnly - -- ^ Use the Huffman-only compression strategy to force Huffman encoding - -- only (no string match). - -fromMethod :: Method -> CInt -fromMethod Deflated = #{const Z_DEFLATED} - -fromCompressionLevel :: CompressionLevel -> Either ZLibParamsException CInt -fromCompressionLevel DefaultCompression = Right $! -1 -fromCompressionLevel NoCompression = Right $! 0 -fromCompressionLevel BestSpeed = Right $! 1 -fromCompressionLevel BestCompression = Right $! 9 -fromCompressionLevel (CompressionLevel n) - | n >= 0 && n <= 9 = Right $! fromIntegral $! n - | otherwise = Left $! IncorrectCompressionLevel n - -fromWindowBits :: Format -> WindowBits -> Either ZLibParamsException CInt -fromWindowBits format bits - = formatModifier format <$> checkWindowBits bits - where checkWindowBits DefaultWindowBits = Right $! 15 - checkWindowBits (WindowBits n) - | n >= 8 && n <= 15 = Right $! fromIntegral $! n - | otherwise = Left $! IncorrectWindowBits $! n - formatModifier Zlib = id - formatModifier GZip = (+16) - formatModifier GZipOrZlib = (+32) - formatModifier Raw = negate - -fromMemoryLevel :: MemoryLevel -> Either ZLibParamsException CInt -fromMemoryLevel DefaultMemoryLevel = Right $! 8 -fromMemoryLevel MinMemoryLevel = Right $! 1 -fromMemoryLevel MaxMemoryLevel = Right $! 9 -fromMemoryLevel (MemoryLevel n) - | n >= 1 && n <= 9 = Right $! fromIntegral n - | otherwise = Left $! IncorrectMemoryLevel $! fromIntegral n - -fromCompressionStrategy :: CompressionStrategy -> CInt -fromCompressionStrategy DefaultStrategy = #{const Z_DEFAULT_STRATEGY} -fromCompressionStrategy Filtered = #{const Z_FILTERED} -fromCompressionStrategy HuffmanOnly = #{const Z_HUFFMAN_ONLY} - -fromErrno :: CInt -> Either ZLibException Bool -fromErrno (#{const Z_OK}) = Right $! True -fromErrno (#{const Z_STREAM_END}) = Right $! False -fromErrno (#{const Z_NEED_DICT}) = Left $! NeedDictionary -fromErrno (#{const Z_BUF_ERROR}) = Left $! BufferError ---fromErrno (#{const Z_ERRNO}) = Left $! FileError -fromErrno (#{const Z_STREAM_ERROR}) = Left $! StreamError -fromErrno (#{const Z_DATA_ERROR}) = Left $! DataError -fromErrno (#{const Z_MEM_ERROR}) = Left $! MemoryError -fromErrno (#{const Z_VERSION_ERROR}) = Left $! VersionError -fromErrno n = Left $! Unexpected n - --- Helper function -convParam :: Format - -> CompressParams - -> Either ZLibParamsException (CInt, CInt, CInt, CInt, CInt) -convParam f (CompressParams c m w l s _ _) - = let c' = fromCompressionLevel c - m' = fromMethod m - b' = fromWindowBits f w - l' = fromMemoryLevel l - s' = fromCompressionStrategy s - eit = either Left - r = Right - in eit (\c_ -> eit (\b_ -> eit (\l_ -> r (c_, m', b_, l_, s')) l') b') c' - --- In following code we go through 7 states. Some of the operations are --- 'deterministic' like 'insertOut' and some of them depends on input ('fill') --- or library call. --- --- (Finished) --- ^ --- | --- | --- | finish --- | --- insertOut fill[1] | ---- (Initial) -------------> (EmptyIn) -----------> (Finishing) --- ^ ^ | ^ | --- | run[2] | | | \------------------\ --- | | | | | --- | | | \------------------\ | --- | run[1] | | flush[0] | | --- \------------------\ | | fill[0] | | fill[3] --- | | | | | --- | | | | | --- swapOut | | v flush[1] | v --- (FullOut) -------------> (Invalid) <----------- (Flushing) --- --- Initial: Initial state, both buffers are empty --- EmptyIn: Empty in buffer, out waits untill filled --- FullOut: Out was filled and sent. In was not entirely read --- Invalid[1]: Both buffers non-empty --- Finishing: There is no more in data and in buffer is empty. Waits till --- all outs was sent. --- Finished: Operation finished --- Flushing: Flush requested --- --- Please note that the decompressing can finish also on flush and finish. --- --- [1] Named for 'historical' reasons - -newtype Initial = Initial ZStream -data EmptyIn = EmptyIn !ZStream !ByteString -data FullOut = FullOut !ZStream !ByteString -data Invalid = Invalid !ZStream !ByteString !ByteString -data Finishing = Finishing !ZStream !ByteString -data Flushing = Flushing !ZStream !ZlibFlush !ByteString - -withByteString :: ByteString -> (Ptr Word8 -> Int -> IO a) -> IO a -withByteString (PS ptr off len) f - = withForeignPtr ptr (\ptr' -> f (ptr' `plusPtr` off) len) - -#ifdef DEBUG -mkByteString :: MonadIO m => Int -> m ByteString -mkByteString s = liftIO $ do - base <- mallocForeignPtrArray s - withForeignPtr base $ \ptr -> C.addForeignPtrFinalizer base $ do - IO.hPutStrLn stderr $ "Freed buffer " ++ show ptr - IO.hPutStrLn stderr $ "Allocated buffer " ++ show base - return $! PS base 0 s - -dumpZStream :: ZStream -> IO () -dumpZStream zstr = withZStream zstr $ \zptr -> do - IO.hPutStr stderr $ "<>" -#else -mkByteString :: MonadIO m => Int -> m ByteString -mkByteString s = liftIO $ create s (\_ -> return ()) -#endif - -putOutBuffer :: Int -> ZStream -> IO ByteString -putOutBuffer size zstr = do - _out <- mkByteString size - withByteString _out $ \ptr len -> withZStream zstr $ \zptr -> do - #{poke z_stream, next_out} zptr ptr - #{poke z_stream, avail_out} zptr len - return _out - -putInBuffer :: ZStream -> ByteString -> IO () -putInBuffer zstr _in - = withByteString _in $ \ptr len -> withZStream zstr $ \zptr -> do - #{poke z_stream, next_in} zptr ptr - #{poke z_stream, avail_in} zptr len - -pullOutBuffer :: ZStream -> ByteString -> IO ByteString -pullOutBuffer zstr _out = withByteString _out $ \ptr _ -> do - next_out <- withZStream zstr $ \zptr -> #{peek z_stream, next_out} zptr - return $! BS.take (next_out `minusPtr` ptr) _out - -pullInBuffer :: ZStream -> ByteString -> IO ByteString -pullInBuffer zstr _in = withByteString _in $ \ptr _ -> do - next_in <- withZStream zstr $ \zptr -> #{peek z_stream, next_in} zptr - return $! BS.drop (next_in `minusPtr` ptr) _in - -type EnumerateeS eli elo m a = (Stream eli -> Iteratee eli m a) -> Iteratee elo m (Iteratee eli m a) - -eneeErr :: (Monad m, Exception err, Nullable elo) - => (Stream eli -> Iteratee eli m a) -> err -> Iteratee elo m () -eneeErr iter = liftM (const ()) . lift . run . iter . EOF . Just . toException - -insertOut :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> Initial - -> Enumeratee ByteString ByteString m a -insertOut size runf (Initial zstr) iter = do - _out <- liftIO $ putOutBuffer size zstr -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ "Inserted out buffer of size " ++ show size -#endif - eneeCheckIfDone (fill size runf (EmptyIn zstr _out)) iter - -fill :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> EmptyIn - -> EnumerateeS ByteString ByteString m a -fill size run' (EmptyIn zstr _out) iter - = let fill' (Chunk _in) - | not (BS.null _in) = do - liftIO $ putInBuffer zstr _in -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ - "Inserted in buffer of size " ++ show (BS.length _in) -#endif - doRun size run' (Invalid zstr _in _out) iter - | otherwise = fillI - fill' (EOF Nothing) = do - out <- liftIO $ pullOutBuffer zstr _out - eneeCheckIfDone (finish size run' (Finishing zstr BS.empty)) $ iter (Chunk out) - fill' (EOF (Just err)) - = case fromException err of - Just err' -> flush size run' (Flushing zstr err' _out) iter - Nothing -> throwRecoverableErr err fill' -#ifdef DEBUG - fillI = do - liftIO $ IO.hPutStrLn stderr $ "About to insert in buffer" - liftI fill' -#else - fillI = liftI fill' -#endif - in fillI - -swapOut :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> FullOut - -> Enumeratee ByteString ByteString m a -swapOut size run' (FullOut zstr _in) iter = do - _out <- liftIO $ putOutBuffer size zstr -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ "Swapped out buffer of size " ++ show size -#endif - eneeCheckIfDone (doRun size run' (Invalid zstr _in _out)) iter - -doRun :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> Invalid - -> EnumerateeS ByteString ByteString m a -doRun size run' (Invalid zstr _in _out) iter = do -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ "About to run" - liftIO $ dumpZStream zstr -#endif - status <- liftIO $ run' zstr #{const Z_NO_FLUSH} -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ "Runned" -#endif - case fromErrno status of - Left err -> do - eneeErr iter err - throwErr (toException err) - Right False -> do -- End of stream - remaining <- liftIO $ pullInBuffer zstr _in - out <- liftIO $ pullOutBuffer zstr _out - idone (iter (Chunk out)) (Chunk remaining) - Right True -> do -- Continue - (avail_in, avail_out) <- liftIO $ withZStream zstr $ \zptr -> do - avail_in <- liftIO $ #{peek z_stream, avail_in} zptr - avail_out <- liftIO $ #{peek z_stream, avail_out} zptr - return (avail_in, avail_out) :: IO (CInt, CInt) - case avail_out of - 0 -> do - out <- liftIO $ pullOutBuffer zstr _out - case avail_in of - 0 -> insertOut size run' (Initial zstr) $ iter (Chunk out) - _ -> swapOut size run' (FullOut zstr _in) $ iter (Chunk out) - _ -> case avail_in of - 0 -> fill size run' (EmptyIn zstr _out) iter - _ -> do - eneeErr iter IncorrectState - throwErr (toException IncorrectState) - -flush :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> Flushing - -> EnumerateeS ByteString ByteString m a -flush size run' (Flushing zstr _flush _out) iter = do - status <- liftIO $ run' zstr (fromFlush _flush) - case fromErrno status of - Left err -> do - eneeErr iter err - throwErr (toException err) - Right False -> do -- Finished - out <- liftIO $ pullOutBuffer zstr _out - idone (iter (Chunk out)) (Chunk BS.empty) - Right True -> do - avail_out <- liftIO $ withZStream zstr #{peek z_stream, avail_out} - case avail_out :: CInt of - 0 -> do - out <- liftIO $ pullOutBuffer zstr _out - out' <- liftIO $ putOutBuffer size zstr - eneeCheckIfDone (flush size run' (Flushing zstr _flush out')) $ iter (Chunk out) - _ -> insertOut size run' (Initial zstr) (liftI iter) - -finish :: MonadIO m - => Int - -> (ZStream -> CInt -> IO CInt) - -> Finishing - -> EnumerateeS ByteString ByteString m a -finish size run' fin@(Finishing zstr _in) iter = do -#ifdef DEBUG - liftIO $ IO.hPutStrLn stderr $ - "Finishing with out buffer of size " ++ show size -#endif - _out <- liftIO $ putOutBuffer size zstr - status <- liftIO $ run' zstr #{const Z_FINISH} - case fromErrno status of - Left err -> do - eneeErr iter err - throwErr (toException err) - Right False -> do -- Finished - remaining <- liftIO $ pullInBuffer zstr _in - out <- liftIO $ pullOutBuffer zstr _out - idone (iter (Chunk out)) (Chunk remaining) - Right True -> do - avail_out <- liftIO $ withZStream zstr #{peek z_stream, avail_out} - case avail_out :: CInt of - 0 -> do - out <- liftIO $ pullOutBuffer zstr _out - eneeCheckIfDone (finish size run' fin) $ iter (Chunk out) - _ -> do - eneeErr iter IncorrectState - throwErr $! toException IncorrectState - -foreign import ccall unsafe deflateInit2_ :: Ptr ZStream -> CInt -> CInt - -> CInt -> CInt -> CInt - -> CString -> CInt -> IO CInt -foreign import ccall unsafe inflateInit2_ :: Ptr ZStream -> CInt - -> CString -> CInt -> IO CInt -foreign import ccall unsafe inflate :: Ptr ZStream -> CInt -> IO CInt -foreign import ccall unsafe deflate :: Ptr ZStream -> CInt -> IO CInt -foreign import ccall unsafe "&deflateEnd" - deflateEnd :: FunPtr (Ptr ZStream -> IO ()) -foreign import ccall unsafe "&inflateEnd" - inflateEnd :: FunPtr (Ptr ZStream -> IO ()) -foreign import ccall unsafe deflateSetDictionary :: Ptr ZStream -> Ptr Word8 - -> CUInt -> IO CInt -foreign import ccall unsafe inflateSetDictionary :: Ptr ZStream -> Ptr Word8 - -> CUInt -> IO CInt - -deflateInit2 :: Ptr ZStream -> CInt -> CInt -> CInt -> CInt -> CInt -> IO CInt -deflateInit2 s l m wB mL s' - = withCString #{const_str ZLIB_VERSION} $ \v -> - deflateInit2_ s l m wB mL s' v #{size z_stream} - -inflateInit2 :: Ptr ZStream -> CInt -> IO CInt -inflateInit2 s wB - = withCString #{const_str ZLIB_VERSION} $ \v -> - inflateInit2_ s wB v #{size z_stream} - -#ifdef DEBUG -deflate' :: ZStream -> CInt -> IO CInt -deflate' z f = withZStream z $ \p -> do - IO.hPutStrLn stderr "About to run deflate" - deflate p f - -inflate' :: ZStream -> CInt -> IO CInt -inflate' z f = withZStream z $ \p -> do - IO.hPutStrLn stderr "About to run inflate" - inflate p f -#else -deflate' :: ZStream -> CInt -> IO CInt -deflate' z f = withZStream z $ \p -> deflate p f - -inflate' :: ZStream -> CInt -> IO CInt -inflate' z f = withZStream z $ \p -> inflate p f -#endif - -mkCompress :: Format -> CompressParams - -> IO (Either ZLibParamsException Initial) -mkCompress frm cp - = case convParam frm cp of - Left err -> return $! Left err - Right (c, m, b, l, s) -> do - zstr <- mallocForeignPtrBytes #{size z_stream} - withForeignPtr zstr $ \zptr -> do - _ <- memset (castPtr zptr) 0 #{size z_stream} - _ <- deflateInit2 zptr c m b l s `finally` - addForeignPtrFinalizer deflateEnd zstr - for_ (compressDictionary cp) $ \(PS fp off len) -> - withForeignPtr fp $ \ptr -> - deflateSetDictionary zptr (ptr `plusPtr` off) - (fromIntegral len) - return $! Right $! Initial $ ZStream zstr - -mkDecompress :: Format -> DecompressParams - -> IO (Either ZLibParamsException (Initial, Maybe ByteString)) -mkDecompress frm (DecompressParams w _ md) - = case fromWindowBits frm w of - Left err -> return $! Left err - Right wB' -> do - zstr <- mallocForeignPtrBytes #{size z_stream} - v <- withForeignPtr zstr $ \zptr -> do - _ <- memset (castPtr zptr) 0 #{size z_stream} - _ <- inflateInit2 zptr wB' `finally` - addForeignPtrFinalizer inflateEnd zstr - case (md, frm) of - (Just (PS fp off len), Raw) -> do - _ <- withForeignPtr fp $ \ptr -> - inflateSetDictionary zptr (ptr `plusPtr` off) - (fromIntegral len) - return Nothing - (Nothing, _) -> return $! Nothing - (Just bs, _) -> return $! (Just bs) - return $! Right $! (Initial $ ZStream zstr, v) - --- User-related code - --- | Compress the input and send to inner iteratee. -enumDeflate :: MonadIO m - => Format -- ^ Format of input - -> CompressParams -- ^ Parameters of compression - -> Enumeratee ByteString ByteString m a -enumDeflate f cp@(CompressParams _ _ _ _ _ size _) iter = do - cmp <- liftIO $ mkCompress f cp - case cmp of - Left err -> do - _ <- lift $ enumErr err iter - throwErr (toException err) - Right init' -> insertOut size deflate' init' iter - --- | Decompress the input and send to inner iteratee. If there is data --- after the end of zlib stream, it is left unprocessed. -enumInflate :: MonadIO m - => Format - -> DecompressParams - -> Enumeratee ByteString ByteString m a -enumInflate f dp@(DecompressParams _ size _md) iter = do - dcmp <- liftIO $ mkDecompress f dp - case dcmp of - Left err -> do - _ <- lift $ enumErr err iter - throwErr (toException err) - Right (init', Nothing) -> insertOut size inflate' init' iter - Right (init', (Just (PS fp off len))) -> - let inflate'' zstr param = do - ret <- inflate' zstr param - case fromErrno ret of - Left NeedDictionary -> do - _ <- withForeignPtr fp $ \ptr -> - withZStream zstr $ \zptr -> - inflateSetDictionary zptr (ptr `plusPtr` off) - (fromIntegral len) - inflate' zstr param - _ -> return ret - in insertOut size inflate'' init' iter - --- | Inflate if Gzip format is recognized, otherwise pass through. -enumInflateAny :: MonadIO m => Enumeratee ByteString ByteString m a -enumInflateAny it = do magic <- iLookAhead $ liftM2 (,) tryHeadBS tryHeadBS - case magic of - (Just 0x1f, Just 0x8b) -> - enumInflate GZip defaultDecompressParams it - >>= enumInflateAny - _ -> mapChunks id it - -enumSyncFlush :: Monad m => Enumerator ByteString m a --- ^ Enumerate synchronise flush. It cause the all pending output to be flushed --- and all available input is sent to inner Iteratee. -enumSyncFlush = enumErr SyncFlush - -enumFullFlush :: Monad m => Enumerator ByteString m a --- ^ Enumerate full flush. It flushes all pending output and reset the --- compression. It allows to restart from this point if compressed data was --- corrupted but it can affect the compression rate. --- --- It may be only used during compression. -enumFullFlush = enumErr FullFlush - -enumBlockFlush :: Monad m => Enumerator ByteString m a --- ^ Enumerate block flush. If the enumerator is compressing it allows to --- finish current block. If the enumerator is decompressing it forces to stop --- on next block boundary. -enumBlockFlush = enumErr Block - diff --git a/src/Bio/Prelude.hs b/src/Bio/Prelude.hs deleted file mode 100644 --- a/src/Bio/Prelude.hs +++ /dev/null @@ -1,117 +0,0 @@ -{-# LANGUAGE CPP #-} -module Bio.Prelude ( - module Bio.Base, - module BasePrelude, - module Data.List.NonEmpty, - module Data.Semigroup, - module System.IO, - module System.Posix.Files, - module System.Posix.IO, - module System.Posix.Types, - - Bytes, LazyBytes, - HashMap, - HashSet, - IntMap, - IntSet, - Text, LazyText, - Pair(..), - -#if !MIN_VERSION_base(4,8,0) - first, - second, -#endif - - decodeBytes, - encodeBytes, - - Hashable(..), - Unpack(..), - fdGet, - fdPut, - fdPutLazy, - withFd - ) where - -import BasePrelude -#if MIN_VERSION_base(4,9,0) - hiding ( (<>), EOF, log1p, log1pexp, log1mexp, expm1 ) -#else - hiding ( (<>), EOF ) -#endif - -#if !MIN_VERSION_base(4,8,0) --- Not as nice as Data.Bifunctor, but still useful. -import Control.Arrow ( first, second ) -#endif - -import Bio.Base -import Data.ByteString ( ByteString ) -import Data.ByteString.Internal ( createAndTrim ) -import Data.List.NonEmpty ( NonEmpty(..) ) -import Data.Semigroup ( Semigroup(..) ) -import Data.Text ( Text ) -import Data.Hashable ( Hashable(..) ) -import Data.HashMap.Strict ( HashMap ) -import Data.HashSet ( HashSet ) -import Data.IntMap ( IntMap ) -import Data.IntSet ( IntSet ) -import Data.Text.Encoding ( encodeUtf8, decodeUtf8With ) -import Foreign.C.Error ( throwErrnoIf_ ) -import System.IO ( hPrint, hPutStr, hPutStrLn, stderr, stdout, stdin ) -import System.Posix.Files -import System.Posix.IO -import System.Posix.Types - -import qualified Data.ByteString.Unsafe as B -import qualified Data.ByteString.Lazy as BL -import qualified Data.ByteString.Char8 as S -import qualified Data.Text as T -import qualified Data.Text.Lazy as TL - -type Bytes = ByteString -type LazyBytes = BL.ByteString -type LazyText = TL.Text - -infixl 2 :!: - --- | A strict pair. -data Pair a b = !a :!: !b deriving(Eq, Ord, Show, Read, Bounded, Ix) - --- | Class of things that can be unpacked into 'String's. Kind of the --- opposite of 'IsString'. -class Unpack s where unpack :: s -> String - -instance Unpack ByteString where unpack = S.unpack -instance Unpack Text where unpack = T.unpack -instance Unpack String where unpack = id - --- | @fdGet bs fd@ reads up to @bs@ 'Bytes' from file descriptor @Fd@. --- Returns an empty 'Bytes' at end of file. -fdGet :: Int -> Fd -> IO Bytes -fdGet bs fd = - createAndTrim bs $ \p -> - fromIntegral <$> fdReadBuf fd (castPtr p) (fromIntegral bs) - -fdPut :: Fd -> Bytes -> IO () -fdPut fd s = B.unsafeUseAsCStringLen s $ \(p,l) -> - throwErrnoIf_ (/= fromIntegral l) "fdPut" $ - fdWriteBuf fd (castPtr p) (fromIntegral l) - -fdPutLazy :: Fd -> LazyBytes -> IO () -fdPutLazy fd = mapM_ (fdPut fd) . BL.toChunks - -withFd :: FilePath -> OpenMode -> Maybe FileMode -> OpenFileFlags - -> (Fd -> IO a) -> IO a -withFd fp om fm ff = bracket (openFd fp om fm ff) closeFd - --- | Converts 'Bytes' into 'Text'. This uses UTF8, but if there is an --- error, it pretends it was Latin1. Evil as this is, it tends to Just --- Work on files where nobody ever wasted a thought on encodings. -decodeBytes :: Bytes -> Text -decodeBytes = decodeUtf8With (const $ fmap w2c) - --- | Converts 'Text' into 'Bytes'. This uses UTF8. -encodeBytes :: Text -> Bytes -encodeBytes = encodeUtf8 - diff --git a/src/Bio/TwoBit.hs b/src/Bio/TwoBit.hs deleted file mode 100644 --- a/src/Bio/TwoBit.hs +++ /dev/null @@ -1,296 +0,0 @@ --- | Would you believe it? The 2bit format stores blocks of Ns in a table at --- the beginning of a sequence, then packs four bases into a byte. So it --- is neither possible nor necessary to store Ns in the main sequence, and --- you would think they aren't stored there, right? And they aren't. --- Instead Ts are stored which the reader has to replace with Ns. --- --- The sensible way to treat these is probably to just say there are two --- kinds of implied annotation (repeats and large gaps for a typical --- genome), which can be interpreted in whatever way fits. And that's why --- we have 'Mask' and 'getSubseqWith'. - -module Bio.TwoBit ( - TwoBitFile(..), - TwoBitSequence(..), - openTwoBit, - - getFwdSubseqWith, - getSubseq, - getSubseqWith, - getSubseqAscii, - getSubseqMasked, - getLazySubseq, - getFragment, - getFwdSubseqV, - getSeqnames, - lookupSequence, - getSeqLength, - clampPosition, - getRandomSeq, - - takeOverlap, - mergeBlocks, - Mask(..) - ) where - -import Bio.Prelude hiding ( left, right, chr ) -import Bio.Util.MMap -import Bio.Util.Storable -import Control.Monad.Trans.State ( StateT(..), get, evalStateT ) -import qualified Data.ByteString as B -import qualified Data.ByteString.Unsafe as B -import qualified Data.IntMap.Strict as I -import qualified Data.HashMap.Lazy as M -import qualified Data.Vector.Unboxed as U -import Foreign.C.Types ( CChar ) - -data TwoBitFile = TBF { - tbf_raw :: B.ByteString, - -- This map is intentionally lazy. May or may not be important. - tbf_seqs :: !(M.HashMap Seqid TwoBitSequence) -} - -data TwoBitSequence = TBS { tbs_n_blocks :: !(I.IntMap Int) - , tbs_m_blocks :: !(I.IntMap Int) - , tbs_dna_offset :: {-# UNPACK #-} !Int - , tbs_dna_size :: {-# UNPACK #-} !Int } - --- | Brings a 2bit file into memory. The file is mmap'ed, so it will --- not work on streams that are not actual files. It's also unsafe if --- the file is modified in any way. -openTwoBit :: FilePath -> IO TwoBitFile -openTwoBit fp = do - raw <- unsafeMMapFile fp - B.unsafeUseAsCString raw $ \praw -> - -- return $ flip runGet (L.fromChunks [raw]) $ do - flip evalStateT praw $ do - sig <- getWord32be - getWord32 <- case sig :: Word32 of - 0x1A412743 -> return getWord32be - 0x4327411A -> return getWord32le - _ -> fail $ "invalid .2bit signature " ++ showHex sig [] - - version <- getWord32 - unless (version == 0) $ fail $ "wrong .2bit version " ++ show version - - nseqs <- getWord32 - _reserved <- getWord32 - - TBF raw <$> foldM (\ix _ -> do !key <- getWord8 >>= getByteString - !off <- getWord32 - return $! M.insert key (mkBlockIndex raw getWord32 off) ix - ) M.empty [1..nseqs] - -type Get = StateT (Ptr CChar) IO - -getWord8, getWord32be, getWord32le :: Num a => Get a -getWord8 = StateT $ \p -> peekUnalnWord32BE p >>= \w -> return (fromIntegral w, plusPtr p 1) -getWord32be = StateT $ \p -> peekUnalnWord32BE p >>= \w -> return (fromIntegral w, plusPtr p 4) -getWord32le = StateT $ \p -> peekUnalnWord32LE p >>= \w -> return (fromIntegral w, plusPtr p 4) - -getByteString :: Int -> Get Bytes -getByteString l = StateT $ \p -> B.packCStringLen (p,l) >>= \s -> return (s, plusPtr p l) - -mkBlockIndex :: B.ByteString -> Get Int -> Int -> TwoBitSequence -mkBlockIndex raw getWord32 ofs = - unsafePerformIO $ - B.unsafeUseAsCString raw $ \praw -> - evalStateT getBlock (plusPtr praw ofs) - where - getBlock = do p0 <- get - ds <- getWord32 - nb <- readBlockList - mb <- readBlockList - _ <- getWord32 - p1 <- get - return $! TBS (I.fromList nb) (I.fromList mb) (ofs + minusPtr p1 p0) ds - - readBlockList = getWord32 >>= \n -> liftM2 zip (repM n getWord32) (repM n getWord32) - --- | Repeat monadic action @n@ times. Returns result in reverse(!) --- order, but doesn't build a huge list of thunks in memory. -repM :: Monad m => Int -> m a -> m [a] -repM n0 m = go [] n0 - where - go acc 0 = return acc - go acc n = m >>= \x -> x `seq` go (x:acc) (n-1) - -takeOverlap :: Int -> I.IntMap Int -> [(Int,Int)] -takeOverlap k m = dropWhile far_left $ - maybe id (\(kv,_) -> (:) kv) (I.maxViewWithKey left) $ - maybe id (\v -> (:) (k,v)) middle $ - I.toAscList right - where - (left, middle, right) = I.splitLookup k m - far_left (s,l) = s+l <= k - -data Mask = None | Soft | Hard | Both deriving (Eq, Ord, Enum, Show) - -getFwdSubseqWith :: TwoBitFile -> TwoBitSequence -- raw data, sequence - -> (Word8 -> Mask -> a) -- mask function - -> Int -> [a] -- start, lazy result -getFwdSubseqWith TBF{..} TBS{..} nt start = - do_mask (takeOverlap start tbs_n_blocks `mergeBlocks` takeOverlap start tbs_m_blocks) start . - drop (start .&. 3) . - B.foldr toDNA [] . - B.drop (fromIntegral $ tbs_dna_offset + (start `shiftR` 2)) $ tbf_raw - where - toDNA b = (++) [ 3 .&. (b `shiftR` x) | x <- [6,4,2,0] ] - - do_mask _ _ [] = [] - do_mask [ ] _ ws = map (`nt` None) ws - do_mask ((s,l,m):is) p ws - | p < s = map (`nt` None) (take (s-p) ws) ++ do_mask ((s,l,m):is) s (drop (s-p) ws) - | otherwise = map (`nt` m) (take (s+l-p) ws) ++ do_mask is (s+l) (drop (s+l-p) ws) - --- | Merge blocks of Ns and blocks of Ms into single list of blocks with --- masking annotation. Gaps remain. Used internally only. -mergeBlocks :: [(Int,Int)] -> [(Int,Int)] -> [(Int,Int,Mask)] -mergeBlocks ((_,0):nbs) mbs = mergeBlocks nbs mbs -mergeBlocks nbs ((_,0):mbs) = mergeBlocks nbs mbs - -mergeBlocks ((ns,nl):nbs) ((ms,ml):mbs) - | ns < ms = let l = min (ms-ns) nl in (ns,l, Hard) : mergeBlocks ((ns+l,nl-l):nbs) ((ms,ml):mbs) - | ms < ns = let l = min (ns-ms) ml in (ms,l, Soft) : mergeBlocks ((ns,nl):nbs) ((ms+l,ml-l):mbs) - | otherwise = let l = min nl ml in (ns,l, Both) : mergeBlocks ((ns+l,nl-l):nbs) ((ms+l,ml-l):mbs) - -mergeBlocks ((ns,nl):nbs) [] = (ns,nl, Hard) : mergeBlocks nbs [] -mergeBlocks [] ((ms,ml):mbs) = (ms,ml, Soft) : mergeBlocks [] mbs - -mergeBlocks [ ] [ ] = [] - - --- | Extract a subsequence and apply masking. TwoBit file can represent --- two kinds of masking (hard and soft), where hard masking is usually --- realized by replacing everything by Ns and soft masking is done by --- lowercasing. Here, we take a user supplied function to apply --- masking. -getSubseqWith :: (Nucleotide -> Mask -> a) -> TwoBitFile -> Range -> [a] -getSubseqWith maskf tbf Range{ r_pos = Pos { p_seq = chr, p_start = start }, r_length = len } = do - let sq1 = fromMaybe (error $ unpack chr ++ " doesn't exist") $ M.lookup chr (tbf_seqs tbf) - let go = getFwdSubseqWith tbf sq1 - if start < 0 - then reverse $ take len $ go (maskf . cmp_nt) (-start-len) - else take len $ go (maskf . fwd_nt) start - where - fwd_nt = (!!) [nucT, nucC, nucA, nucG] . fromIntegral - cmp_nt = (!!) [nucA, nucG, nucT, nucC] . fromIntegral - --- | Works only in forward direction. -getLazySubseq :: TwoBitFile -> Position -> [Nucleotide] -getLazySubseq tbf Pos{ p_seq = chr, p_start = start } = do - let sq1 = fromMaybe (error $ unpack chr ++ " doesn't exist") $ M.lookup chr (tbf_seqs tbf) - let go = getFwdSubseqWith tbf sq1 - if start < 0 - then error "sorry, can't go backwards" - -- then reverse $ take len $ go (maskf . cmp_nt) (-start-len) - else go fwd_nt start - where - fwd_nt n _ = [nucT, nucC, nucA, nucG] !! fromIntegral n - - --- | Extract a subsequence without masking. -getSubseq :: TwoBitFile -> Range -> [Nucleotide] -getSubseq = getSubseqWith const - --- | Extract a subsequence with typical masking: soft masking is --- ignored, hard masked regions are replaced with Ns. -getSubseqMasked :: TwoBitFile -> Range -> [Nucleotides] -getSubseqMasked = getSubseqWith mymask - where - mymask n None = nucToNucs n - mymask n Soft = nucToNucs n - mymask _ Hard = nucsN - mymask _ Both = nucsN - --- | Extract a subsequence with masking for biologists: soft masking is --- done by lowercasing, hard masking by printing an N. -getSubseqAscii :: TwoBitFile -> Range -> String -getSubseqAscii = getSubseqWith mymask - where - mymask n None = showNucleotide n - mymask n Soft = toLower (showNucleotide n) - mymask _ Hard = 'N' - mymask _ Both = 'N' - - -getSeqnames :: TwoBitFile -> [Seqid] -getSeqnames = M.keys . tbf_seqs - -lookupSequence :: TwoBitFile -> Seqid -> Maybe TwoBitSequence -lookupSequence tbf sq = M.lookup sq . tbf_seqs $ tbf - -getSeqLength :: TwoBitFile -> Seqid -> Int -getSeqLength tbf chr = - maybe (error $ shows chr " doesn't exist") tbs_dna_size $ - M.lookup chr (tbf_seqs tbf) - --- | limits a range to a position within the actual sequence -clampPosition :: TwoBitFile -> Range -> Range -clampPosition tbf (Range (Pos n start) len) = Range (Pos n start') (end' - start') - where - size = getSeqLength tbf n - start' = if start < 0 then max start (-size) else start - end' = min (start + len) $ if start < 0 then 0 else size - - --- | Sample a piece of random sequence uniformly from the genome. --- Only pieces that are not hard masked are sampled, soft masking is --- allowed, but not reported. --- On a 32bit platform, this will fail for genomes larger than 1G bases. --- However, if you're running this code on a 32bit platform, you have --- bigger problems to worry about. -getRandomSeq :: TwoBitFile -- ^ 2bit file - -> Int -- ^ desired length - -> (Int -> g -> (Int, g)) -- ^ draw random int below limit - -> g -- ^ RNG - -> ((Range, [Nucleotide]), g) -- ^ position, sequence, new RNG -getRandomSeq tbf len rndInt = draw - where - names = getSeqnames tbf - lengths = map (getSeqLength tbf) names - total = sum lengths - frags = I.fromList $ zip (scanl (+) 0 lengths) names - - draw g0 | good = ((r', sq), gn) - | otherwise = draw gn - where - (p0, gn) = rndInt (2*total) g0 - p = p0 `shiftR` 1 - Just ((o,s),_) = I.maxViewWithKey $ fst $ I.split (p+1) frags - r' = (if odd p0 then id else reverseRange) $ clampPosition tbf $ Range (Pos s (p-o)) len - sq = catMaybes $ getSubseqWith mask2maybe tbf r' - good = r_length r' == len && length sq == len - - mask2maybe n None = Just n - mask2maybe n Soft = Just n - mask2maybe _ Hard = Nothing - mask2maybe _ Both = Nothing - --- | Gets a fragment from a 2bit file. The result always has the --- desired length; if necessary, it is padded with Ns. Be careful about --- the unconventional encoding: 0..4 == TCAGN -getFragment :: TwoBitFile -> Seqid -> Int -> Int -> U.Vector Word8 -getFragment tbf chr p l = - case lookupSequence tbf chr of - Nothing -> U.replicate l 4 - Just tbs -> getFwdSubseqV tbf tbs p l - --- Careful about weird encoding: 0..4 == TCAGN -getFwdSubseqV :: TwoBitFile -> TwoBitSequence -> Int -> Int -> U.Vector Word8 -getFwdSubseqV TBF{..} TBS{..} start len = U.unfoldrN len step ini - where - ini = (start, takeOverlap start tbs_n_blocks) - - step (off, nbs) - | off < 0 = Just (4, (succ off, nbs)) - | off >= tbs_dna_size = Just (4, (succ off, nbs)) - | otherwise = case nbs of - [ ] -> Just (y, (succ off, [ ])) - (s,l):nbs' | off < s -> Just (y, (succ off, nbs)) - | off < s+l -> Just (4, (succ off, nbs)) - | otherwise -> Just (y, (succ off, nbs')) - where - x = B.index tbf_raw (tbs_dna_offset + off `shiftR` 2) - y = x `shiftR` (6 - 2 * (off .&. 3)) .&. 3 -- T,C,A,G - diff --git a/src/Bio/Util/MMap.hs b/src/Bio/Util/MMap.hs deleted file mode 100644 --- a/src/Bio/Util/MMap.hs +++ /dev/null @@ -1,26 +0,0 @@ -{-# LANGUAGE ForeignFunctionInterface #-} -module Bio.Util.MMap where - -import Bio.Prelude -import Data.ByteString.Internal ( fromForeignPtr ) -import Foreign.C.Types - -unsafeMMapFile :: FilePath -> IO Bytes -unsafeMMapFile fp = - bracket (openFd fp ReadOnly Nothing defaultFileFlags) closeFd $ \fd -> do - stat <- getFdStatus fd - let size = fromIntegral (fileSize stat) - if size <= 0 - then return mempty - else do - ptr <- c_mmap size (fromIntegral fd) - if ptr == nullPtr - then error "unable to mmap file" - else do - fptr <- newForeignPtrEnv c_munmap (intPtrToPtr $ fromIntegral size) ptr - return $ fromForeignPtr fptr 0 (fromIntegral size) - -foreign import ccall unsafe "my_mmap" c_mmap :: CSize -> CInt -> IO (Ptr Word8) -foreign import ccall unsafe "&my_munmap" c_munmap :: FunPtr (Ptr () -> Ptr Word8 -> IO ()) - - diff --git a/src/Bio/Util/Numeric.hs b/src/Bio/Util/Numeric.hs deleted file mode 100644 --- a/src/Bio/Util/Numeric.hs +++ /dev/null @@ -1,213 +0,0 @@ --- | Random useful stuff I didn't know where to put. - -module Bio.Util.Numeric ( - wilson, invnormcdf, choose, - estimateComplexity, showNum, showOOM, - log1p, expm1, (<#>), - log1mexp, log1pexp, - lsum, llerp - ) where - -import Data.List ( foldl1' ) -import Data.Char ( intToDigit ) -import Prelude - --- | Calculates the Wilson Score interval. --- If @(l,m,h) = wilson c x n@, then @m@ is the binary proportion and --- @(l,h)@ it's @c@-confidence interval for @x@ positive examples out of --- @n@ observations. @c@ is typically something like 0.05. - -wilson :: Double -> Int -> Int -> (Double, Double, Double) -wilson c x n = ( (m - h) / d, p, (m + h) / d ) - where - nn = fromIntegral n - p = fromIntegral x / nn - - z = invnormcdf (1-c*0.5) - h = z * sqrt (( p * (1-p) + 0.25*z*z / nn ) / nn) - m = p + 0.5 * z * z / nn - d = 1 + z * z / nn - -showNum :: Show a => a -> String -showNum = triplets [] . reverse . show - where - triplets acc [] = acc - triplets acc [a] = a:acc - triplets acc [a,b] = b:a:acc - triplets acc [a,b,c] = c:b:a:acc - triplets acc (a:b:c:s) = triplets (',':c:b:a:acc) s - -showOOM :: Double -> String -showOOM x | x < 0 = '-' : showOOM (negate x) - | otherwise = findSuffix (x*10) ".kMGTPEZY" - where - findSuffix _ [] = "many" - findSuffix y (s:ss) | y < 100 = intToDigit (round y `div` 10) : case (round y `mod` 10, s) of - (0,'.') -> [] ; (0,_) -> [s] ; (d,_) -> [s, intToDigit d] - | y < 1000 = intToDigit (round y `div` 100) : intToDigit ((round y `mod` 100) `div` 10) : - if s == '.' then [] else [s] - | y < 10000 = intToDigit (round y `div` 1000) : intToDigit ((round y `mod` 1000) `div` 100) : - '0' : if s == '.' then [] else [s] - | otherwise = findSuffix (y*0.001) ss - --- Stolen from Lennart Augustsson's erf package, who in turn took it from --- Accurate to about 1e-9. -invnormcdf :: (Ord a, Floating a) => a -> a -invnormcdf p = - let a1 = -3.969683028665376e+01 - a2 = 2.209460984245205e+02 - a3 = -2.759285104469687e+02 - a4 = 1.383577518672690e+02 - a5 = -3.066479806614716e+01 - a6 = 2.506628277459239e+00 - - b1 = -5.447609879822406e+01 - b2 = 1.615858368580409e+02 - b3 = -1.556989798598866e+02 - b4 = 6.680131188771972e+01 - b5 = -1.328068155288572e+01 - - c1 = -7.784894002430293e-03 - c2 = -3.223964580411365e-01 - c3 = -2.400758277161838e+00 - c4 = -2.549732539343734e+00 - c5 = 4.374664141464968e+00 - c6 = 2.938163982698783e+00 - - d1 = 7.784695709041462e-03 - d2 = 3.224671290700398e-01 - d3 = 2.445134137142996e+00 - d4 = 3.754408661907416e+00 - - pLow = 0.02425 - - nan = 0/0 - - in if p < 0 then - nan - else if p == 0 then - -1/0 - else if p < pLow then - let q = sqrt(-2 * log p) - in (((((c1*q+c2)*q+c3)*q+c4)*q+c5)*q+c6) / - ((((d1*q+d2)*q+d3)*q+d4)*q+1) - else if p < 1 - pLow then - let q = p - 0.5 - r = q*q - in (((((a1*r+a2)*r+a3)*r+a4)*r+a5)*r+a6)*q / - (((((b1*r+b2)*r+b3)*r+b4)*r+b5)*r+1) - else if p <= 1 then - - invnormcdf (1 - p) - else - nan - - --- | Try to estimate complexity of a whole from a sample. Suppose we --- sampled @total@ things and among those @singles@ occured only once. --- How many different things are there? --- --- Let the total number be @m@. The copy number follows a Poisson --- distribution with paramter @\lambda@. Let \( z := e^{\lambda} \), then --- we have: --- --- \[ --- P( 0 ) = e^{-\lambda} = \frac{1}{z} \\ --- P( 1 ) = \lambda e^{-\lambda} = \frac{\ln z}{z} \\ --- P(\ge 1) = 1 - e^{-\lambda} = 1 - \frac{1}{z} \\ --- \] --- \[ --- \mbox{singles} = m \frac{\ln z}{z} \\ --- \mbox{total} = m \left( 1 - \frac{1}{z} \right) \\ --- \] --- \[ --- D := \frac{\mbox{total}}{\mbox{singles}} = (1 - \frac{1}{z}) * \frac{z}{\ln z} \\ --- f := z - 1 - D \ln z = 0 --- \] --- --- To get @z@, we solve using Newton iteration and then substitute to --- get @m@: --- --- \[ --- df/dz = 1 - D/z \\ --- z' = z - \frac{ z (z - 1 - D \ln z) }{ z - D } \\ --- m = \mbox{singles} * \frac{z}{\ln z} --- \] --- --- It converges as long as the initial @z@ is large enough, and @10D@ --- (in the line for @zz@ below) appears to work well. - -estimateComplexity :: (Integral a, Floating b, Ord b) => a -> a -> Maybe b -estimateComplexity total singles | total <= singles = Nothing - | singles <= 0 = Nothing - | otherwise = Just m - where - d = fromIntegral total / fromIntegral singles - step z = z * (z - 1 - d * log z) / (z - d) - iter z = case step z of zd | abs zd < 1e-12 -> z - | otherwise -> iter $! z-zd - zz = iter $! 10*d - m = fromIntegral singles * zz / log zz - - --- | Computes \( \ln \left( e^x + e^y \right) \) without leaving the log domain and --- hence without losing precision. -infixl 5 <#> -{-# INLINE (<#>) #-} -(<#>) :: (Floating a, Ord a) => a -> a -> a -x <#> y = if x >= y then x + log1pexp (y-x) else y + log1pexp (x-y) - --- | Computes @log (1+x)@ to a relative precision of @10^-8@ even for --- very small @x@. Stolen from -{-# INLINE log1p #-} -log1p :: (Floating a, Ord a) => a -> a -log1p x | x < -1 = error "log1p: argument must be greater than -1" - -- x is large enough that the obvious evaluation is OK: - | x > 0.0001 || x < -0.0001 = log $ 1 + x - -- Use Taylor approx. log(1 + x) = x - x^2/2 with error roughly x^3/3 - -- Since |x| < 10^-4, |x|^3 < 10^-12, relative error less than 10^-8: - | otherwise = (1 - 0.5*x) * x - - --- | Computes \( e^x - 1 \) to a relative precision of @10^-10@ even for --- very small @x@. Stolen from -{-# INLINE expm1 #-} -expm1 :: (Floating a, Ord a) => a -> a -expm1 x | x > -0.00001 && x < 0.00001 = (1 + 0.5 * x) * x -- Taylor approx - | otherwise = exp x - 1 -- direct eval - --- | Computes \( \ln (1 - e^x) \), following Martin Mächler. -{-# INLINE log1mexp #-} -log1mexp :: (Floating a, Ord a) => a -> a -log1mexp x | x > - log 2 = log (- expm1 x) - | otherwise = log1p (- exp x) - --- | Computes \( \ln (1 + e^x) \), following Martin Mächler. -{-# INLINE log1pexp #-} -log1pexp :: (Floating a, Ord a) => a -> a -log1pexp x | x <= -37 = exp x - | x <= 18 = log1p $ exp x - | x <= 33.3 = x + exp (-x) - | otherwise = x - - --- | Computes \( \ln ( \sum_i e^{x_i} ) \) sensibly. The list must be --- sorted in descending(!) order. -{-# INLINE lsum #-} -lsum :: (Floating a, Ord a) => [a] -> a -lsum = foldl1' (\x y -> if x >= y then x + log1pexp (y-x) else err) - where err = error "lsum: argument list must be in descending order" - --- | Computes \( \ln \left( c e^x + (1-c) e^y \right) \). -{-# INLINE llerp #-} -llerp :: (Floating a, Ord a) => a -> a -> a -> a -llerp c x y | c <= 0.0 = y - | c >= 1.0 = x - | x >= y = log c + x + log1p ( (1-c)/c * exp (y-x) ) -- Hmm. - | otherwise = log1p (-c) + y + log1p ( c/(1-c) * exp (x-y) ) -- Hmm. - --- | Binomial coefficient: \( \mbox{choose n k} = \frac{n!}{(n-k)! k!} \) -{-# INLINE choose #-} -choose :: Integral a => a -> a -> a -choose n k = product [n-k+1 .. n] `div` product [2..k] - - diff --git a/src/Bio/Util/Storable.hs b/src/Bio/Util/Storable.hs deleted file mode 100644 --- a/src/Bio/Util/Storable.hs +++ /dev/null @@ -1,97 +0,0 @@ -{-# LANGUAGE CPP #-} --- | Utilities to read multibyte quantities from arbitrary positions. -module Bio.Util.Storable - ( peekWord8 - , peekUnalnWord16LE - , peekUnalnWord16BE - , peekUnalnWord32LE - , peekUnalnWord32BE - , pokeUnalnWord32LE - ) where - -#if __GLASGOW_HASKELL__ >= 710 -#define HAVE_BYTESWAP_PRIMOPS -#endif - -#if i386_HOST_ARCH || x86_64_HOST_ARCH -#define MEM_UNALIGNED_OPS -#endif - -import Bio.Prelude - -peekWord8 :: Ptr a -> IO Word8 -peekWord8 = peek . castPtr - -#if defined(MEM_UNALIGNED_OPS) && defined(WORDS_BIGENDIAN) && defined(HAVE_BYTESWAP_PRIMOPS) -peekUnalnWord16LE :: Ptr a -> IO Word16 -peekUnalnWord16LE = fmap byteSwap16 . peek . castPtr - -peekUnalnWord32LE :: Ptr a -> IO Word32 -peekUnalnWord32LE = fmap byteSwap32 . peek . castPtr - -pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () -pokeUnalnWord32LE p w = poke (castPtr p) (byteSwap32 w) - -#elif defined(MEM_UNALIGNED_OPS) && !defined(WORDS_BIGENDIAN) -peekUnalnWord16LE :: Ptr a -> IO Word16 -peekUnalnWord16LE = peek . castPtr - -peekUnalnWord32LE :: Ptr a -> IO Word32 -peekUnalnWord32LE = peek . castPtr - -pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () -pokeUnalnWord32LE p w = poke (castPtr p) w - -#else -peekUnalnWord16LE :: Ptr a -> IO Word16 -peekUnalnWord16LE p = do - x <- fromIntegral <$> peekWord8 (plusPtr p 0) - y <- fromIntegral <$> peekWord8 (plusPtr p 1) - return $! x .|. unsafeShiftL y 8 - -peekUnalnWord32LE :: Ptr a -> IO Word32 -peekUnalnWord32LE p = do - x <- fromIntegral <$> peekWord8 (plusPtr p 0) - y <- fromIntegral <$> peekWord8 (plusPtr p 1) - z <- fromIntegral <$> peekWord8 (plusPtr p 2) - w <- fromIntegral <$> peekWord8 (plusPtr p 3) - return $! x .|. unsafeShiftL y 8 .|. unsafeShiftL z 16 .|. unsafeShiftL w 24 - -pokeUnalnWord32LE :: Ptr a -> Word32 -> IO () -pokeUnalnWord32LE p w = do pokeByteOff p 0 (fromIntegral $ shiftR w 0 :: Word8) - pokeByteOff p 1 (fromIntegral $ shiftR w 8 :: Word8) - pokeByteOff p 2 (fromIntegral $ shiftR w 16 :: Word8) - pokeByteOff p 3 (fromIntegral $ shiftR w 24 :: Word8) -#endif - - -#if defined(MEM_UNALIGNED_OPS) && !defined(WORDS_BIGENDIAN) && defined(HAVE_BYTESWAP_PRIMOPS) -peekUnalnWord16BE :: Ptr a -> IO Word16 -peekUnalnWord16BE = fmap byteSwap16 . peek . castPtr - -peekUnalnWord32BE :: Ptr a -> IO Word32 -peekUnalnWord32BE = fmap byteSwap32 . peek . castPtr - -#elif defined(MEM_UNALIGNED_OPS) && defined(WORDS_BIGENDIAN) -peekUnalnWord16BE :: Ptr a -> IO Word16 -peekUnalnWord16BE = peek . castPtr - -peekUnalnWord32BE :: Ptr a -> IO Word32 -peekUnalnWord32BE = peek . castPtr - -#else -peekUnalnWord16BE :: Ptr a -> IO Word16 -peekUnalnWord16BE p = do - x <- fromIntegral <$> peekWord8 (plusPtr p 0) - y <- fromIntegral <$> peekWord8 (plusPtr p 1) - return $! y .|. unsafeShiftL x 8 - -peekUnalnWord32BE :: Ptr a -> IO Word32 -peekUnalnWord32BE p = do - x <- fromIntegral <$> peekWord8 (plusPtr p 0) - y <- fromIntegral <$> peekWord8 (plusPtr p 1) - z <- fromIntegral <$> peekWord8 (plusPtr p 2) - w <- fromIntegral <$> peekWord8 (plusPtr p 3) - return $! w .|. unsafeShiftL z 8 .|. unsafeShiftL y 16 .|. unsafeShiftL x 24 -#endif - diff --git a/src/Bio/Util/Zlib.hs b/src/Bio/Util/Zlib.hs deleted file mode 100644 --- a/src/Bio/Util/Zlib.hs +++ /dev/null @@ -1,23 +0,0 @@ -{-# LANGUAGE CPP #-} -module Bio.Util.Zlib ( decompressGzip ) where - -import Prelude - -import qualified Data.ByteString.Lazy as L -import qualified Data.ByteString.Lazy.Internal as L ( ByteString(..) ) -import qualified Codec.Compression.Zlib.Internal as Z - - --- | Decompresses Gzip or Bgzf and passes everything else on. In --- reality, it simply decompresses Gzip, and when done, looks for --- another Gzip stream. Since there is a small chance to attempt --- decompression of an uncompressed stream, the original data is --- returned in case of an error. -decompressGzip :: L.ByteString -> L.ByteString -decompressGzip s = case L.uncons s of - Just (31, s') -> case L.uncons s' of - Just (139,_) -> Z.foldDecompressStreamWithInput L.Chunk decompressGzip (const s) - (Z.decompressST Z.gzipOrZlibFormat Z.defaultDecompressParams) s - _ -> s - _ -> s - diff --git a/src/cbits/loops.c b/src/cbits/loops.c deleted file mode 100644 --- a/src/cbits/loops.c +++ /dev/null @@ -1,113 +0,0 @@ -void nuc_loop( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - - while( u < v ) { - char a = q[ u ] ; - char b = q[ u + stride ] ; - char a1 = a ? 0x10 << (a&3) : 0xf0 ; - char b1 = b ? 0x1 << (b&3) : 0xf ; - *p++ = a1 | b1 ; - u += stride+stride ; - } - if( u == v ) { - char a = q[ u ] ; - char a1 = a ? 0x10 << (a&3) : 0xf0 ; - *p = a1 ; - } -} - -void nuc_loop_wide( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - - while( u < v ) { - char a = q[ u ] ; - *p++ = a ? 0x1 << (a&3) : 0xf ; - u += stride ; - } -} - -void nuc_loop_asc( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - - while( u <= v ) { - char a = q[ u ] ; - *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'A' : (a&3) == 1 ? 'C' : (a&3) == 2 ? 'G' : 'T' ; - u += stride ; - } - *p = 0 ; -} - -void nuc_loop_asc_rev( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - - while( u <= v ) { - char a = q[ v ] ; - *p++ = a == 0 ? 'N' : (a&3) == 0 ? 'T' : (a&3) == 1 ? 'G' : (a&3) == 2 ? 'C' : 'A' ; - v -= stride ; - } - *p = 0 ; -} - -void qual_loop( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - while( u <= v ) { - *p++ = (q[u] >> 2) & 0x3f ; - u += stride ; - } -} - -void qual_loop_asc( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - while( u <= v ) { - *p++ = 33 + ((q[u] >> 2) & 0x3f) ; - u += stride ; - } - *p = 0 ; -} - -void qual_loop_asc_rev( char* p, int stride, char* q, int u, int v ) -{ - u *= stride ; - v *= stride ; - while( u <= v ) { - *p++ = 33 + ((q[v] >> 2) & 0x3f) ; - v -= stride ; - } - *p = 0 ; -} - -int int_loop( char* p, int x ) -{ - *p++ = ':' ; - if( x == 0 ) { - *p = '0' ; - return 2 ; - } - char *q = p ; - while( x > 0 ) { - *q++ = '0' + x % 10 ; - x /= 10 ; - } - int r = q-p ; - --q ; - while( p < q ) { - char c = *p ; - *p++ = *q ; - *q-- = c ; - } - return r+1 ; -} - - diff --git a/src/cbits/mmap.c b/src/cbits/mmap.c deleted file mode 100644 --- a/src/cbits/mmap.c +++ /dev/null @@ -1,10 +0,0 @@ -#include - -unsigned char *my_mmap(size_t len, int fd) { - void *result = mmap(0, len, PROT_READ, MAP_SHARED, fd, 0); - return (unsigned char*)( result == MAP_FAILED ? 0 : result ); -} - -void my_munmap(void *len, unsigned char *p) { - munmap( p, (size_t)len ) ; -} diff --git a/src/cbits/myers_align.c b/src/cbits/myers_align.c deleted file mode 100644 --- a/src/cbits/myers_align.c +++ /dev/null @@ -1,102 +0,0 @@ -#include "myers_align.h" - -#include -#include -#include - -// [*blech*, this looks and feels like FORTRAN.] -unsigned myers_diff( - const char *seq_a, int len_a, enum myers_align_mode mode, - const char* seq_b, int len_b, int maxd, - char *bt_a, char *bt_b ) -{ - // int len_a = strlen( seq_a ), len_b = strlen( seq_b ) ; - if( maxd > len_a + len_b ) maxd = len_a + len_b ; - - // in vee[d][k], d runs from 0 to maxd; k runs from -d to +d - int **vee = calloc( maxd, sizeof(int*) ) ; - - int d, dd, k, x, y, r = UINT_MAX ; - int *v_d_1 = 0, *v_d = 0 ; // "array slice" vee[.][d-1] - for( d = 0 ; d != maxd ; ++d, v_d_1 = v_d ) // D-paths in order of increasing D - { - v_d = d + (vee[d] = malloc( (2 * d + 1) * sizeof( int ) )) ; // "array slice" vee[.][d] - - for( k = max(-d,-len_a) ; k <= min(d,len_b) ; ++k ) // diagonals - { - if( d == 0 ) x = 0 ; - else if(d==1&&k==0) x = v_d_1[ k ]+1 ; - else if( k == -d ) x = v_d_1[ k+1 ] ; - else if( k == d ) x = v_d_1[ k-1 ]+1 ; // argh, need to check for d first, b/c -d+2 could be equal to d - else if( k == -d+1 ) x = max( v_d_1[ k ]+1, v_d_1[ k+1 ] ) ; - else if( k == d-1 ) x = max( v_d_1[ k-1 ]+1, v_d_1[ k ]+1 ) ; - else x = max3( v_d_1[ k-1 ]+1, v_d_1[ k ]+1, v_d_1[ k+1 ] ) ; - - y = x-k ; - while( x < len_b && y < len_a && match( seq_b[x], seq_a[y] ) ) ++x, ++y ; - v_d[ k ] = x ; - - if( - bt_a && bt_b && - (mode == myers_align_is_prefix || y == len_a) && - (mode == myers_align_has_prefix || x == len_b) ) - { - char *out_a = bt_a + len_a + d +2 ; - char *out_b = bt_b + len_b + d +2 ; - *--out_a = 0 ; - *--out_b = 0 ; - for( dd = d ; dd != 0 ; ) - { - if( k != -dd && k != dd && x == vee[ dd-1 ][ k + dd-1 ]+1 ) - { - --dd ; - --x ; - --y ; - *--out_b = seq_b[x] ; - *--out_a = seq_a[y] ; - } - else if( k > -dd+1 && x == vee[ dd-1 ][ k-1 + dd-1 ]+1 ) - { - --x ; - --k ; - --dd ; - *--out_b = seq_b[x] ; - *--out_a = '-' ; - } - else if( k < dd-1 && x == vee[ dd-1 ][ k+1 + dd-1 ] ) - { - ++k ; - --y ; - --dd ; - *--out_b = '-' ; - *--out_a = seq_a[y] ; - } - else // this better had been a match... - { - --x ; - --y ; - *--out_b = seq_b[x] ; - *--out_a = seq_a[y] ; - } - } - while( x > 0 ) - { - --x ; - *--out_b = seq_b[x] ; - *--out_a = seq_a[x] ; - } - memmove( bt_a, out_a, bt_a + len_a + d + 2 - out_a ) ; - memmove( bt_b, out_b, bt_b + len_b + d + 2 - out_b ) ; - r = d ; - goto cleanup ; - } - } - } - -cleanup: - for( dd = maxd ; dd != 0 ; --dd ) - free( vee[dd-1] ) ; - free( vee ) ; - return r ; -} - diff --git a/src/cbits/myers_align.h b/src/cbits/myers_align.h deleted file mode 100644 --- a/src/cbits/myers_align.h +++ /dev/null @@ -1,76 +0,0 @@ -#ifndef INCLUDED_MYERS_ALIGN -#define INCLUDED_MYERS_ALIGN - -enum myers_align_mode { - myers_align_globally = 0, - myers_align_is_prefix = 1, - myers_align_has_prefix = 2 } ; - -//! \brief aligns two sequences in O(nd) time -//! This alignment algorithm following Eugene W. Myers: "An O(ND) -//! Difference Algorithm and Its Variations". -//! Both input sequences are ASCIIZ-encoded with IUPAC-IUB ambiguity -//! codes. By definition, if ambiguity codes overlap, that's a match, -//! else a mismatch. Mismatches and gaps count a unit penalty. If mode -//! is myers_align_globally, both sequences must align completely. If -//! mode is myers_align_is_prefix, seq_a must align completely as prefix -//! of seq_b. If mode is myers_align_has_prefix, seq_b must align -//! completely as prefix of seq_a. -//! -//! Note that the calculation time is O(nd) where n is the length of the -//! best alignment and d the number of differences in it, memory -//! consumption is O(maxd^2). -//! -//! \param seq_a First input sequence. -//! \param mode How to align (i.e. what gaps to count). -//! \param seq_b Second input sequence. -//! \param maxd Maximum penalty to consider. -//! \param bt_a Space to backtrace seq_a into, must have room for -//! (strlen(seq_a)+maxd+1) characters. -//! \param bt_b Space to backtrace seq_b into, must have room for -//! (strlen(seq_b)+maxd+1) characters. -//! \return The actual edit distance or UINT_MAX if the edit distance -//! would be greater than maxd. -//! -unsigned myers_diff( - const char *seq_a, int len_a, enum myers_align_mode mode, - const char* seq_b, int len_b, int maxd, - char *bt_a, char *bt_b ) ; - -//! \brief converts an IUPAC-IUB ambiguity code to a bitmap Each base is -//! represented by a bit, makes checking for matches easier. -static inline int char_to_bitmap( char x ) -{ - switch( x & ~32 ) - { - case 'A': return 1 ; - case 'C': return 2 ; - case 'G': return 4 ; - case 'T': return 8 ; - case 'U': return 8 ; - - case 'S': return 6 ; - case 'W': return 9 ; - case 'R': return 5 ; - case 'Y': return 10 ; - case 'K': return 12 ; - case 'M': return 3 ; - - case 'B': return 14 ; - case 'D': return 13 ; - case 'H': return 11 ; - case 'V': return 7 ; - - case 'N': return 15 ; - default: return 0 ; - } -} - -static inline int compatible( char x, char y ) { return (char_to_bitmap(x) & char_to_bitmap(y)) != 0 ; } -static inline int match( char a, char b ) { return (char_to_bitmap(a) & char_to_bitmap(b)) != 0 ; } - -static inline int min( int a, int b ) { return a < b ? a : b ; } -static inline int max( int a, int b ) { return a < b ? b : a ; } -static inline int max3( int a, int b, int c ) { return a < b ? max( b, c ) : max( a, c ) ; } - -#endif diff --git a/src/cbits/trim.c b/src/cbits/trim.c deleted file mode 100644 --- a/src/cbits/trim.c +++ /dev/null @@ -1,46 +0,0 @@ -#include - -static const uint8_t compls[] = - { 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15 } ; - -int prim_match_reads( int i1 - , int i2 - , int r - , const uint8_t *rd1 - , const uint8_t *qs1 - , const uint8_t *rd2 - , const uint8_t *qs2 ) -{ - int acc = 0 ; - while( r != 0 ) - { - --i2 ; - uint8_t n1 = rd1[ i1 ] ; - uint8_t n2 = rd2[ i2 ] ; - uint8_t q1 = qs1[ i1 ] ; - uint8_t q2 = qs2[ i2 ] ; - - acc += (n1 & 0xF) == compls[ n2 & 0xF ] ? 0 : 5 + (q1 < q2 ? q1 : q2) ; - - ++i1 ; - --r ; - } - return acc ; -} - -int prim_match_ad( int off - , int i - , const uint8_t *rd - , const uint8_t *qs - , const uint8_t *ad ) -{ - int acc = 0 ; - while( i > 0 ) - { - --i; - acc += rd[ i+off ] == ad[ i ] ? 0 : 5 + - (qs[ i+off ] < 25 ? qs[ i+off ] : 25) ; - } - return acc ; -} -