{-# OPTIONS_GHC -Wno-partial-fields #-}
module Bio.TwoBit.Tool
( EncodeProgress(..)
, buildFasta
, faToTwoBit
, formatCdna
, parseAnno
, twoBitToFa
, vcfToTwoBit
)
where
import Bio.TwoBit
import Control.Applicative
import Control.Exception
import Control.Monad
import Data.Bits
import Data.Bool
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy.Char8 as L
import Data.ByteString.Short ( ShortByteString, toShort )
import qualified Data.ByteString.Short as H
import Data.Char ( isSpace, isUpper )
import Data.Foldable
import qualified Data.HashMap.Strict as M
import Data.Int ( Int64 )
import Data.Word ( Word8, Word32 )
import System.IO ( stdout )
type Bytes = B.ByteString
type LazyBytes = L.ByteString
-- | A cDNA or mRNA or transcript (these are all synonymous), with some
-- metainformation collected from the annotation. Whatever the input
-- was called, we call it 'cdna' in the transciptome.
data Cdna = Cdna
{ c_id :: !Bytes -- identifier (typically an ENST number)
, c_pos :: !Range -- genomic position
, c_gene_id :: !Bytes -- gene identifier (typically an ENSG number)
, c_gene_symbol :: !Bytes -- colloquial name, aka locus
, c_gene_biotype :: !Bytes -- whatever, just pass it on
, c_biotype :: !Bytes -- whatever, just pass it on
, c_description :: !Bytes -- unclear; always empty for now
, c_exons :: [Range] -- list of exon coordinates (sorted backwards)
}
deriving Show
data Range = Range
{ r_chrom :: !C.ByteString
, r_start :: !Int
, r_len :: !Int }
deriving Show
reverseRange :: Range -> Range
reverseRange (Range sq pos len) = Range sq (-pos-len) len
null_cdna :: Cdna
null_cdna = Cdna "" (Range "" 0 0) "" "" "" "" "" []
formatCdna :: TwoBitFile -> Cdna -> B.Builder
formatCdna tbf Cdna{..} = descr <> buildFasta 60 getExons
where
(_,tbf_fn) = C.breakEnd (=='/') $ tbf_path tbf
(tbf_base,_) = C.breakEnd (=='.') tbf_fn
genome_id = if C.null tbf_base then tbf_fn else C.init tbf_base
descr = B.char7 '>' <> B.byteString c_id <> " cdna chromosome:" <>
B.byteString genome_id <> B.char7 ':' <> formatRange c_pos <>
" gene:" <> B.byteString c_gene_id <>
maybeBS " gene_biotype:" c_gene_biotype <>
maybeBS " transcript_biotype:" c_biotype <>
maybeBS " gene_symbol:" c_gene_symbol <>
maybeBS " description:" c_description <> B.char7 '\n'
formatRange r | r_start r < 0 = formatRange1 (reverseRange r) <> ":-1"
| otherwise = formatRange1 r <> ":1"
formatRange1 r = B.byteString (r_chrom r) <> B.char7 ':' <>
B.intDec (r_start r) <> B.char7 ':' <>
B.intDec (r_start r + r_len r - 1)
maybeBS p s = if B.null s then mempty else p <> B.byteString s
getExons | r_start c_pos < 0 = concatMap getExon c_exons
| otherwise = concatMap getExon (reverse c_exons)
getExon :: Range -> [Word8]
getExon (Range ch start len) =
case findChrom ch tbf of
Just tbs | start >= 0 -> take len $ unpackRS $ tbc_fwd_seq tbs start
| otherwise -> take len $ unpackRS $ tbc_rev_seq tbs (-start-len)
Nothing -> error $ "unknown reference " ++ show ch
buildFasta :: Int -> [Word8] -> B.Builder
buildFasta n = go
where
go [ ] = mempty
go s = let (u,v) = splitAt n s
in foldMap B.word8 u <> B.char7 '\n' <> go v
{-# INLINE buildFasta #-}
twoBitToFa :: Int -> TwoBitSequence' dir -> IO ()
twoBitToFa ln = B.hPutBuilder stdout . buildFasta 60 . take ln . unpackRSMasked
data EncodeProgress
= EncodeProgress
{ ep_seqname :: !ShortByteString
, ep_position :: !Word32
, ep_hardmasked :: !Word32
, ep_softmasked :: !Word32
, ep_enclength :: !Int64
, ep_tail :: EncodeProgress }
| Encoded B.Builder
-- Strategy: We can only write the packedDNA after we wrote the nBlocks
-- and mBlocks. So packedDNA needs to be buffered. We have to do three
-- simultaneous strict folds of the input, all of which result in reasonably
-- compact structures (name table, mask table, encoded dna), which get
-- concatenated at the end.
--
-- We also have to buffer everything, since the header with the sequence
-- names must be written first. Oh joy.
--
-- We return a list of progress notifications terminated by the
-- 'B.Builder' for the whole 2bit file. The progress messages can be
-- printed or ignored; in either case, they should ensure enough
-- strictness to not waste more memory than necessary.
faToTwoBit :: L.ByteString -> EncodeProgress
faToTwoBit = get_each []
where
get_each acc inp = case L.uncons $ L.dropWhile (/= '>') inp of
Nothing -> Encoded $ seqs_to_twobit $ reverse acc
Just (_,s2) ->
let (nm, s') = L.break (<= ' ') s2
in get_one acc (toShort (L.toStrict nm)) 0 (GapList maxBound L2i_Nil)
(GapList maxBound L2i_Nil) (BaseAccu 0 0 emptyAccu)
(L.dropWhile (/= '\n') s')
get_one acc !nm !pos !ns !ms !bs inp = case L.uncons inp of
Nothing -> fin L.empty
Just (c,s')
| c <= ' ' -> get_one acc nm pos ns ms bs s'
| c == '>' -> fin (L.cons c s')
| otherwise -> get_one acc nm (succ pos)
(collect_Ns ns pos c)
(collect_ms ms pos c)
(collect_bases bs c) s'
where
fin k = let !r = encode_seq pos ns ms bs
in EncodeProgress nm pos (sum_L2i pos ns) (sum_L2i pos ms) (L.length r) $
get_each ((nm,r):acc) k
-- | Extracts the reference from a VCF. This assumes the presence of at
-- least one record per site. The VCF must be sorted by position. When
-- writing out, we try to match the order of the contigs as listed in
-- the header. Unlisted contigs follow at the end with their order
-- preserved; contigs without data are not written at all.
vcfToTwoBit :: [B.ByteString] -> EncodeProgress
vcfToTwoBit s0 = let (lns, s1) = read_header [] s0
in get_each lns [] $ filter (\s -> not (B.null s) && C.head s /= '#') s1
where
-- Collects the "contig" stanzas, parses their lengths. Returns the
-- length map and the remaining stream.
read_header acc [ ] = (reverse acc, [])
read_header acc (l:ls) | "##contig=" `C.isPrefixOf` l
, (Just !nm, Just !ln) <- parse_cline l = read_header ((nm,ln):acc) ls
| "#" `C.isPrefixOf` l = read_header acc ls
| otherwise = (reverse acc, l:ls)
parse_cline = p1 . C.filter (not . isSpace) . C.takeWhile (/='>') . C.drop 1 . C.dropWhile (/='<')
where
p1 s | "ID=" `C.isPrefixOf` s = let (nm,t) = C.break (==',') $ C.drop 3 s
(_,ln) = p1 $ C.drop 1 t
in (Just (toShort nm),ln)
| "length=" `C.isPrefixOf` s = case C.readInt $ C.drop 7 s of
Just (ln,u) -> let (nm,_) = p1 $ C.drop 1 $ C.dropWhile (/=',') u in (nm,Just (fromIntegral ln))
Nothing -> p1 $ C.drop 1 $ C.dropWhile (/=',') s
| C.null s = (Nothing,Nothing)
| otherwise = p1 $ C.drop 1 $ C.dropWhile (/=',') s
get_each :: [(ShortByteString,Word32)]
-> [(ShortByteString, LazyBytes)]
-> [B.ByteString]
-> EncodeProgress
get_each lns acc [ ] = Encoded $ seqs_to_twobit $ reorder (map fst lns) $ reverse acc
get_each lns acc (l:s2) = EncodeProgress nm' ln' (sum_L2i ln' ns') 0 (L.length r) $
get_each lns ((nm',r):acc) s3
where
nm = B.takeWhile (/= 9) l
(pos,ns,bs,s3) = get_one nm 0 (GapList maxBound L2i_Nil) (BaseAccu 0 0 emptyAccu) (l:s2)
!nm' = toShort nm
(ns',bs',ln') = case find ((==) nm' . fst) lns of
Just (_,ln) | ln > pos -> (extend_gap ns ln, pad_bases bs (fromIntegral $ ln-pos), ln)
_ -> (ns,bs,pos)
!r = encode_seq ln' ns' (GapList maxBound L2i_Nil) bs'
-- important: 1-based coordinates!
get_one !_nm !pos !ns !bs [ ] = (pos,ns,bs,[])
get_one !nm !pos !ns !bs (l:s')
| B.takeWhile (/=9) l /= nm = (pos,ns,bs,l:s')
| Just (pos',l3) <- C.readInt . C.drop 1 $ B.dropWhile (/=9) l
, ref <- B.takeWhile (/=9) . B.drop 1 . B.dropWhile (/=9) $ B.drop 1 l3
, fromIntegral pos' >= pos + 1
, not (C.null ref) =
if fromIntegral pos' == pos + 1
-- record in sequence
then get_one nm (succ pos) (collect_Ns ns pos $ C.head ref)
(collect_bases bs $ C.head ref) s'
-- gap: handle the gap, reprocess the record
else let gap_len = pos' - fromIntegral pos - 1
in get_one nm (fromIntegral pos' - 1) (extend_gap ns pos)
(pad_bases bs gap_len) (l:s')
-- anything else can be ignored (parse errors or additional records)
| otherwise = get_one nm pos ns bs s'
pad_bases bs n = foldl' collect_bases bs $ replicate n 'T'
-- Reorder a key-value list so it matches the order of a list of
-- keys. Missing keys are ignored, leftover pairs retain their
-- original order.
reorder :: Eq a => [a] -> [(a,b)] -> [(a,b)]
reorder [ ] vs = vs
reorder (k:ks) vs = go [] vs
where
go xs ((k1,v1):ys) | k == k1 = (k1,v1) : reorder ks (reverse xs ++ ys)
| otherwise = go ((k1,v1):xs) ys
go xs [ ] = reorder ks (reverse xs)
-- List of pairs of 'Word32's. Specialized and unpacked to conserve space. Probably overkill...
data L2i = L2i {-# UNPACK #-} !Word32 {-# UNPACK #-} !Word32 L2i | L2i_Nil
data GapList = GapList !Word32 !L2i
sum_L2i :: Word32 -> GapList -> Word32
sum_L2i p (GapList q xs) = go (if q == maxBound then 0 else p-q) xs
where
go !a (L2i x y z) = go (a+y-x) z
go !a L2i_Nil = a
encodeL2i :: L2i -> B.Builder
encodeL2i = go 0 mempty mempty
where
go !n ss ls L2i_Nil = B.word32LE n <> ss <> ls
go !n ss ls (L2i s e rs) = go (succ n) (B.word32LE s <> ss) (B.word32LE (e-s) <> ls) rs
seqs_to_twobit :: [(ShortByteString, LazyBytes)] -> B.Builder
seqs_to_twobit seqs = B.word32LE 0x1A412743 <> B.word32LE 0 <>
B.word32LE (fromIntegral $ length seqs) <> B.word32LE 0 <>
mconcat (zipWith (\nm off -> B.word8 (fromIntegral (H.length nm)) <>
B.shortByteString nm <>
B.word32LE (fromIntegral off))
(map fst seqs) offsets) <>
foldMap (B.lazyByteString . snd) seqs
where
offset0 = 16 + 5 * length seqs + sum (map (H.length . fst) seqs)
offsets = scanl (\a b -> a + fromIntegral (L.length b)) offset0 $ map snd seqs
-- | A way to accumulate bytes. If the accumulated bytes will hang
-- around in memory, this has much lower overhead than 'Builder'. If it
-- has short lifetime, 'Builder' is much more convenient.
newtype Accu = Accu [Bytes]
emptyAccu :: Accu
emptyAccu = Accu []
-- | Appends bytes to a collection of 'Bytes' in such a way that the
-- 'Bytes' keep doubling in size. This ensures O(n) time and space
-- complexity and fairly low overhead.
grow :: Word8 -> Accu -> Accu
grow w = go 1 [B.singleton w]
where
go l acc (Accu (s:ss))
| B.length s <= l = go (l+B.length s) (s:acc) (Accu ss)
| otherwise = let !s' = B.concat acc in Accu (s' : s : ss)
go _ acc (Accu [ ]) = let !s' = B.concat acc in Accu [s']
buildAccu :: Accu -> B.Builder
buildAccu (Accu ss) = foldMap B.byteString $ reverse ss
encode_seq :: Word32 -- ^ length
-> GapList -- ^ list of N stretches
-> GapList -- ^ list of mask stretches
-> BaseAccu -- ^ accumulated bases
-> LazyBytes
encode_seq pos ns ms bs = L.length r `seq` r
where
ss' = case bs of (BaseAccu 0 _ ss) -> ss
(BaseAccu n w ss) -> grow (w `shiftL` (8-2*n)) ss
r = B.toLazyByteString $
B.word32LE pos <>
encodeL2i (case ns of GapList p rs | p == maxBound -> rs ; GapList p rs -> L2i p pos rs) <>
encodeL2i (case ms of GapList p rs | p == maxBound -> rs ; GapList p rs -> L2i p pos rs) <>
B.word32LE 0 <>
buildAccu ss'
-- | Collects stretches of Ns by looking at one character at a time. In
-- reality, anything that isn't one of \"ACGT\" is treated as an N.
collect_Ns :: GapList -> Word32 -> Char -> GapList
collect_Ns (GapList spos rs) pos c
| spos == maxBound && c `C.elem` "ACGTacgt" = GapList maxBound rs
| spos == maxBound = GapList pos rs
| c `C.elem` "ACGTacgt" = GapList maxBound (L2i spos pos rs)
| otherwise = GapList spos rs
-- | Collects stretches of masked dna by looking at one letter at a
-- time. Anything lowercase is considered masked.
collect_ms :: GapList -> Word32 -> Char -> GapList
collect_ms (GapList spos rs) pos c
| spos == maxBound && isUpper c = GapList maxBound rs
| spos == maxBound = GapList pos rs
| isUpper c = GapList maxBound (L2i spos pos rs)
| otherwise = GapList spos rs
extend_gap :: GapList -> Word32 -> GapList
extend_gap (GapList spos rs) pos
| spos == maxBound = GapList pos rs
| otherwise = GapList spos rs
data BaseAccu = BaseAccu !Int !Word8 !Accu
-- | Collects bases in 2bit format. It accumulates 4 bases in one word,
-- then collects bytes in an 'Accu'. From the 2bit spec:
--
-- packedDna - the DNA packed to two bits per base, represented as
-- so: T - 00, C - 01, A - 10, G - 11. The first base is
-- in the most significant 2-bit byte; the last base is
-- in the least significant 2 bits. For example, the
-- sequence TCAG is represented as 00011011.
collect_bases :: BaseAccu -> Char -> BaseAccu
collect_bases (BaseAccu n w ss) c
= let code = case c of 'C'->1;'c'->1;'A'->2;'a'->2;'G'->3;'g'->3;_->0
w' = shiftL w 2 .|. code
in if n == 3 then BaseAccu 0 0 (grow w' ss) else BaseAccu (succ n) w' ss
data Gene = Gene { g_id :: Bytes, g_symbol :: Bytes, g_biotype :: Bytes }
null_gene :: Gene
null_gene = Gene "" "" ""
data GffError = GffError String Int GffErrorDetail deriving Show
data GffErrorDetail = GffParseError | GffIdMismatch | GffUnknownRef Bytes deriving Show
instance Exception GffError where
displayException (GffError fp ln dt) = displayDetail dt ++ " in line " ++ show ln ++ " of gff file " ++ fp
where
displayDetail GffParseError = "parse error"
displayDetail GffIdMismatch = "identifier does not match"
displayDetail (GffUnknownRef ch) = "unknown reference " ++ show ch
-- | Parses annotations in GFF format. We want to turn an annotation
-- and a 2bit file into a FastA of the transcriptome (one sequence per
-- annotated transcript), that looks like the stuff Lior Pachter feeds
-- into Kallisto. Annotations come in two dialects of GFF, either GFF3
-- or GTF. We autodetect and understand both.
parseAnno :: String -> L.ByteString -> [Either GffError Cdna]
parseAnno fp = filter (either (const True) (not . null . c_exons)) .
go null_gene null_cdna .
map (fmap (C.split '\t')) .
filter (\(_,s) -> not (B.null s) && C.head s /= '#') .
zip (enumFrom 1) .
map L.toStrict .
L.lines
where
go gene xscript ((ln, ch:_:tp:fro_:tho_:_:strand:_:stuff_:_) : strm)
| Just (fro,"") <- C.readInt fro_
, Just (tho,"") <- C.readInt tho_
, Just stuff <- parseStuff stuff_ =
let rng = bool id reverseRange (strand == "-") $ Range ch (fro-1) (tho-fro+1)
in go2 ln gene xscript strm rng (B.map (.|. 32) tp) stuff
go gene xscript ((ln, _) : strm) = Left (GffError fp ln GffParseError) : go gene xscript strm
go _gene xscript [ ] = Right xscript : []
go2 ln gene xscript strm rng tp (Left stuff)
| tp == "exon" = if M.lookup "Parent" stuff == Just (c_id xscript)
then go gene xscript { c_exons = rng : c_exons xscript } strm
else Left (GffError fp ln GffIdMismatch) : go gene xscript strm
| tp == "transcript" || tp == "cdna" || tp == "mrna" =
Right xscript :
case (M.lookup "ID" stuff, M.lookup "Parent" stuff) of
(Just tid, Just gid)
| gid == g_id gene -> let xscript' = Cdna { c_id = tid
, c_pos = rng
, c_gene_id = gid
, c_gene_symbol = g_symbol gene
, c_gene_biotype = g_biotype gene
, c_biotype = M.lookupDefault "" "biotype" stuff
, c_description = "" -- XXX
, c_exons = [] }
in go gene xscript' strm
| otherwise -> Left (GffError fp ln GffIdMismatch) : go gene null_cdna strm
_ -> Left (GffError fp ln GffParseError) : go gene null_cdna strm
| tp == "gene" =
Right xscript :
case M.lookup "ID" stuff of
Just gid -> let gene' = Gene { g_id = gid
, g_symbol = "" -- XXX
, g_biotype = M.lookupDefault "" "biotype" stuff }
in go gene' null_cdna strm
Nothing -> Left (GffError fp ln GffParseError) : go null_gene null_cdna strm
| otherwise = go gene xscript strm
go2 ln gene xscript strm rng tp (Right stuff)
| tp == "exon" =
case M.lookup "transcript_id" stuff of
Just tid
| tid == c_id xscript -> go gene xscript { c_exons = rng : c_exons xscript } strm
| otherwise -> Left (GffError fp ln GffIdMismatch) : go gene xscript strm
Nothing -> Left (GffError fp ln GffParseError) : go gene xscript strm
| tp == "transcript" || tp == "cdna" || tp == "mrna" =
Right xscript :
case (M.lookup "transcript_id" stuff, M.lookup "gene_id" stuff) of
(Just tid, Just gid) -> let xscript' = Cdna { c_id = tid
, c_pos = rng
, c_gene_id = gid
, c_gene_symbol = M.lookupDefault "" "gene_name" stuff
, c_gene_biotype = "" -- XXX
, c_biotype = "" -- XXX
, c_description = "" -- XXX
, c_exons = [] }
in go gene xscript' strm
_ -> Left (GffError fp ln GffParseError) : go gene null_cdna strm
| otherwise = go gene xscript strm
-- | Parses the random stuff in GFF into a hash table. Returns 'Just
-- (Left _)' if the file uses assignment style ("foo=bar;"), returns
-- 'Just (Right _)' if the file uses statement style ("foo \"bar\";"),
-- otherwise returns Nothing.
parseStuff :: Bytes -> Maybe (Either (M.HashMap Bytes Bytes) (M.HashMap Bytes Bytes))
parseStuff s = Left <$> parse_assignments M.empty s <|>
Right <$> parse_quoted M.empty s
where
parse_assignments !h s0
| C.null s0 = Just h
| otherwise = do let (k,s1) = C.break (=='=') s0
guard . not $ C.null k
guard . not $ C.null s1
let (v,s2) = C.break (==';') $ C.tail s1
parse_assignments (M.insert k v h) (C.drop 1 s2)
parse_quoted !h s0
| C.null s0 || C.head s0 == '#' = Just h
| otherwise = do let (k,s1) = C.break (==' ') s0
guard . not $ C.null k
guard $ C.isPrefixOf " \"" s1
let (v,s2) = C.break (=='"') $ C.drop 2 s1
guard . not $ C.null s2
let s3 = C.dropWhile (/=';') s2
parse_quoted (M.insert k v h) . C.dropWhile isSpace $ C.drop 1 s3