{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
module Bio.Data.Bed.Utils
( fetchSeq
, clipBed
, CutoffMotif(..)
, mkCutoffMotif
, scanMotif
, monoColonalize
, BaseMap(..)
, baseMap
, queryBaseMap
, rpkmBed
, rpkmSortedBed
, countTagsBed
, countTagsBinBed
, countTagsBinBed'
, tagCountDistr
, peakCluster
) where
import Conduit
import Lens.Micro
import Control.Monad.State.Strict
import qualified Data.ByteString.Char8 as B
import qualified Data.Foldable as F
import Data.Function (on)
import qualified Data.HashMap.Strict as M
import qualified Data.IntervalMap.Strict as IM
import Data.Maybe (fromJust, fromMaybe)
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as I
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import System.IO
import Bio.Data.Bed
import Bio.Data.Bed.Types
import Bio.Motif (Bkgd (..), Motif (..), CDF, PWM)
import qualified Bio.Motif as Motif
import qualified Bio.Motif.Search as Motif
import Bio.Seq hiding (length)
import Bio.Seq.IO
import qualified Bio.Utils.BitVector as BV
clipBed :: (BEDLike b, Monad m)
=> [(B.ByteString, Int)] -- ^ Chromosome sizes
-> ConduitT b b m ()
clipBed chrsize = concatMapC f
where
f x = case M.lookup (x^.chrom) chrsize' of
Nothing -> Nothing
Just n -> if x^.chromStart >= n
then Nothing
else Just $ chromStart %~ max 0 $ chromEnd %~ min n $ x
chrsize' = M.fromListWith (error "redundant chromosomes") chrsize
{-# INLINE clipBed #-}
-- | retreive sequences
fetchSeq :: BioSeq DNA a
=> Genome
-> BED
-> IO (Either String (DNA a))
fetchSeq g bed = do
dna <- getSeq g (bed^.chrom, bed^.chromStart, bed^.chromEnd)
return $ case bed^.strand of
Just False -> rc <$> dna
_ -> dna
{-# INLINE fetchSeq #-}
-- | Motif with predefined cutoff score. All necessary intermediate data
-- structure for motif scanning are stored.
data CutoffMotif = CutoffMotif
{ _motif_name :: B.ByteString
, _motif_pwm :: PWM
, _motif_sigma :: U.Vector Double
, _motif_pwm_rc :: PWM
, _motif_sigma_rc :: U.Vector Double
, _background :: Bkgd
, _cutoff :: Double
, _cdf :: CDF }
mkCutoffMotif :: Bkgd
-> Double -- ^ p-value
-> Motif -> CutoffMotif
mkCutoffMotif bg p motif = CutoffMotif (_name motif) (_pwm motif) sigma pwm'
sigma' bg sc $ Motif.truncateCDF (1 - p * 10) cdf
where
cdf = Motif.scoreCDF bg $ _pwm motif
sc = Motif.cdf' cdf $ 1 - p
sigma = Motif.optimalScoresSuffix bg $ _pwm motif
pwm' = Motif.rcPWM $ _pwm motif
sigma' = Motif.optimalScoresSuffix bg pwm'
-- | Motif score is in [0, 1000]: ( 1 / (1 + exp (-(-logP - 5))) ) * 1000.
scanMotif :: (BEDLike b, MonadIO m)
=> Genome -> [CutoffMotif] -> ConduitT b BED m ()
scanMotif g motifs = awaitForever $ \bed -> do
let (chr, start, end) = (bed^.chrom, bed^.chromStart, bed^.chromEnd)
liftIO (getSeq g (chr, start, end)) >>= \case
Left _ -> liftIO $ hPutStrLn stderr $
"Warning: no sequence for region: " ++ show (chr, start, end)
Right dna -> forM_ motifs $ \CutoffMotif{..} -> do
let mkBed str (i, sc) = BED chr (start + i) (start + i + n)
(Just $ _motif_name) (Just $ toAffinity $ 1 - Motif.cdf _cdf sc)
(Just str)
n = Motif.size _motif_pwm
-- Scan forward strand
Motif.findTFBSWith _motif_sigma _background _motif_pwm
(dna :: DNA IUPAC) _cutoff True .| mapC (mkBed True)
-- Scan reverse strand
Motif.findTFBSWith _motif_sigma_rc _background _motif_pwm_rc
dna _cutoff True .| mapC (mkBed False)
where
toAffinity x' = round $ sc * 1000
where
sc = 1 / (1 + exp (-(x - 5)))
x = negate $ logBase 10 $ max 1e-20 x'
{-# INLINE scanMotif #-}
-- | process a sorted BED stream, keep only mono-colonal tags
monoColonalize :: Monad m => ConduitT BED BED m ()
monoColonalize = do
x <- headC
case x of
Just b -> yield b >> concatMapAccumC f b
Nothing -> return ()
where
f cur prev = case compareBed prev cur of
GT -> error $
"Input is not sorted: " ++ show prev ++ " > " ++ show cur
LT -> (cur, [cur])
_ -> if prev^.strand == cur^.strand then (cur, []) else (cur, [cur])
{-# INLINE monoColonalize #-}
newtype BaseMap = BaseMap (M.HashMap B.ByteString BV.BitVector)
-- | Count the tags (starting positions) at each position in the genome.
baseMap :: PrimMonad m
=> [(B.ByteString, Int)] -- ^ chromosomes and their sizes
-> ConduitT BED o m BaseMap
baseMap chrs = do
bvs <- lift $ fmap M.fromList $ forM chrs $ \(chr, n) -> do
bv <- BV.zeros n
return (chr, bv)
mapM_C $ \bed -> case M.lookup (bed^.chrom) bvs of
Nothing -> return ()
Just bv -> if fromMaybe True $ bed^.strand
then BV.set bv $ bed^.chromStart
else BV.set bv $ bed^.chromEnd - 1
lift $ fmap BaseMap $ sequence $ fmap BV.unsafeFreeze bvs
queryBaseMap :: BEDLike b => b -> BaseMap -> Maybe [Bool]
queryBaseMap bed (BaseMap bm) = case M.lookup (bed^.chrom) bm of
Nothing -> Nothing
Just bv ->
let res = map (bv BV.!) [bed^.chromStart .. bed^.chromEnd - 1]
in case bed^.strand of
Just False -> Just $ reverse res
_ -> Just res
-- | calculate RPKM on a set of unique regions. Regions (in bed format) would be kept in
-- memory but not tag file.
-- RPKM: Readcounts per kilobase per million reads. Only counts the starts of tags
rpkmBed :: (PrimMonad m, BEDLike b, G.Vector v Double)
=> [b] -> ConduitT BED o m (v Double)
rpkmBed regions = do
v <- lift $ do v' <- V.unsafeThaw . V.fromList . zip [0..] $ regions
I.sortBy (compareBed `on` snd) v'
V.unsafeFreeze v'
let (idx, sortedRegions) = V.unzip v
n = G.length idx
readCount <- rpkmSortedBed $ Sorted sortedRegions
lift $ do
result <- GM.new n
G.sequence_ . G.imap (\x i -> GM.unsafeWrite result i (readCount U.! x)) $ idx
G.unsafeFreeze result
{-# INLINE rpkmBed #-}
-- | calculate RPKM on a set of regions. Regions must be sorted. The Sorted data
-- type is used to remind users to sort their data.
rpkmSortedBed :: (PrimMonad m, BEDLike b, G.Vector v Double)
=> Sorted (V.Vector b) -> ConduitT BED o m (v Double)
rpkmSortedBed (Sorted regions) = do
vec <- lift $ GM.replicate l 0
n <- foldMC (count vec) (0 :: Int)
let factor = fromIntegral n / 1e9
lift $ liftM (G.imap (\i x -> x / factor / (fromIntegral . size) (regions V.! i)))
$ G.unsafeFreeze vec
where
count v nTags tag = do
let p | tag^.strand == Just True = tag^.chromStart
| tag^.strand == Just False = tag^.chromEnd - 1
| otherwise = error "Unkown strand"
xs = concat $ IM.elems $
IM.containing (M.lookupDefault IM.empty (tag^.chrom) intervalMap) p
addOne v xs
return $ succ nTags
intervalMap = sortedBedToTree (++) . Sorted . G.toList . G.zip regions .
G.map return . G.enumFromN 0 $ l
addOne v' = mapM_ $ \x -> GM.unsafeRead v' x >>= GM.unsafeWrite v' x . (+1)
l = G.length regions
{-# INLINE rpkmSortedBed #-}
countTagsBed :: (PrimMonad m, BEDLike b, G.Vector v Int)
=> [b] -> ConduitT BED o m (v Int, Int)
countTagsBed regions = do
vec <- lift $ GM.replicate l 0
n <- foldMC (count vec) (0 :: Int)
vec' <- lift $ G.unsafeFreeze vec
return (vec', n)
where
count v nTags tag = do
let p | tag^.strand == Just True = tag^.chromStart
| tag^.strand == Just False = tag^.chromEnd - 1
| otherwise = error "Unkown strand"
xs = concat $ IM.elems $
IM.containing (M.lookupDefault IM.empty (tag^.chrom) intervalMap) p
addOne v xs
return $ succ nTags
intervalMap = bedToTree (++) $ zip regions $ map return [0..]
addOne v' = mapM_ $ \x -> GM.unsafeRead v' x >>= GM.unsafeWrite v' x . (+1)
l = length regions
{-# INLINE countTagsBed #-}
-- | divide each region into consecutive bins, and count tags for each bin and
-- return the number of all tags. Note: a tag is considered to be overlapped
-- with a region only if the starting position of the tag is in the region. For
-- the common sense overlapping, use countTagsBinBed'.
countTagsBinBed :: (Integral a, PrimMonad m, G.Vector v a, BEDLike b)
=> Int -- ^ bin size
-> [b] -- ^ regions
-> ConduitT BED o m ([v a], Int)
countTagsBinBed k beds = do
vs <- lift $ fmap V.fromList $ forM beds $ \bed -> do
let start = bed^.chromStart
num = ((bed^.chromEnd) - start) `div` k
index i = (i - start) `div` k
v <- GM.replicate num 0
return (v, index)
nTags <- foldMC (f vs) 0
readCount <- lift $ mapM (G.unsafeFreeze . fst) $ G.toList vs
return (readCount, nTags)
where
f vs n bed = do
let pos | bed^.strand == Just True = bed^.chromStart
| bed^.strand == Just False = bed^.chromEnd - 1
| otherwise = error "unkown strand."
overlaps = concat $ IM.elems $ IM.containing
(M.lookupDefault IM.empty (bed^.chrom) intervalMap) pos
forM_ overlaps $ \x -> do
let (v, idxFn) = vs `G.unsafeIndex` x
i = let i' = idxFn pos
l = GM.length v
in if i' >= l then l - 1 else i'
GM.unsafeModify v (+1) i
return $ n + 1
intervalMap = bedToTree (++) $ zip beds $ map return [0..]
{-# INLINE countTagsBinBed #-}
-- | Same as countTagsBinBed, except that tags are treated as complete intervals
-- instead of single points.
countTagsBinBed' :: (Integral a, PrimMonad m, G.Vector v a, BEDLike b1, BEDLike b2)
=> Int -- ^ bin size
-> [b1] -- ^ regions
-> ConduitT b2 o m ([v a], Int)
countTagsBinBed' k beds = do
initRC <- lift $ forM beds $ \bed -> do
let start = bed^.chromStart
end = bed^.chromEnd
num = (end - start) `div` k
index i = (i - start) `div` k
v <- GM.replicate num 0
return (v, index)
sink 0 $ V.fromList initRC
where
sink !nTags vs = do
tag <- await
case tag of
Just bed -> do
let chr = bed^.chrom
start = bed^.chromStart
end = bed^.chromEnd
overlaps = concat $ IM.elems $ IM.intersecting
(M.lookupDefault IM.empty chr intervalMap) $ IM.IntervalCO start end
lift $ forM_ overlaps $ \x -> do
let (v, idxFn) = vs `G.unsafeIndex` x
lo = let i = idxFn start
in if i < 0 then 0 else i
hi = let i = idxFn end
l = GM.length v
in if i >= l then l - 1 else i
forM_ [lo..hi] $ \i ->
GM.unsafeRead v i >>= GM.unsafeWrite v i . (+1)
sink (nTags+1) vs
_ -> do readCount <- lift $ mapM (G.unsafeFreeze . fst) $ G.toList vs
return (readCount, nTags)
intervalMap = bedToTree (++) $ zip beds $ map return [0..]
{-# INLINE countTagsBinBed' #-}
tagCountDistr :: PrimMonad m => G.Vector v Int => ConduitT BED o m (v Int)
tagCountDistr = loop M.empty
where
loop m = do
x <- await
case x of
Just bed -> do
let p | fromMaybe True (bed^.strand) = bed^.chromStart
| otherwise = 1 - bed^.chromEnd
case M.lookup (bed^.chrom) m of
Just table -> loop $ M.insert (bed^.chrom) (M.insertWith (+) p 1 table) m
_ -> loop $ M.insert (bed^.chrom) (M.fromList [(p,1)]) m
_ -> lift $ do
vec <- GM.replicate 100 0
F.forM_ m $ \table ->
F.forM_ table $ \v -> do
let i = min 99 v
GM.unsafeRead vec i >>= GM.unsafeWrite vec i . (+1)
G.unsafeFreeze vec
{-# INLINE tagCountDistr #-}
-- | cluster peaks
peakCluster :: (BEDLike b, Monad m)
=> [b] -- ^ peaks
-> Int -- ^ radius
-> Int -- ^ cutoff
-> ConduitT o BED m ()
peakCluster peaks r th = mergeBedWith mergeFn peaks' .| filterC g
where
peaks' = map f peaks
f b = let c = (b^.chromStart + b^.chromEnd) `div` 2
in asBed (b^.chrom) (c-r) (c+r) :: BED3
mergeFn xs = asBed (head xs ^. chrom) lo hi & score .~ Just (fromIntegral $ length xs)
where
lo = minimum $ map (^.chromStart) xs
hi = maximum $ map (^.chromEnd) xs
g b = fromJust (b^.score) >= fromIntegral th
{-# INLINE peakCluster #-}