Created
September 10, 2021 17:14
-
-
Save noughtmare/4ae65499da7576bc54262e3bfe39953f to your computer and use it in GitHub Desktop.
Changing Haskell implementation of k-nucleotide to vector-hashtables.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{-# LANGUAGE ScopedTypeVariables #-} | |
-- The Computer Language Benchmarks Game | |
-- https://salsa.debian.org/benchmarksgame-team/benchmarksgame/ | |
-- | |
-- Contributed by cahu ette | |
module Main where | |
import Data.Bits | |
import Data.List | |
import Data.Word | |
import Data.Hashable | |
import Data.Traversable | |
import Text.Printf | |
import Data.Maybe | |
import Control.Monad | |
import Control.Monad.ST | |
import Control.Parallel.Strategies | |
import qualified Data.Map.Strict as M | |
import qualified Data.Vector.Hashtables.Internal as H | |
import qualified Data.Vector.Unboxed.Mutable as UV | |
import qualified Data.ByteString.Char8 as B | |
type HashTable s k v = H.Dictionary s UV.MVector k UV.MVector v | |
{- By using only 2 bits to encode keys, it's important to use a different table | |
- for different key sizes. Otherwise, if we encode 'A' as 0x00, "AT" and | |
- "AAT" would map to the same bucket in the table. | |
- | |
- We could use 3 bits per letter to avoid this problem if needed. | |
-} | |
bitsForChar :: Char -> Word64 | |
bitsForChar 'a' = 0 | |
bitsForChar 'A' = 0 | |
bitsForChar 'c' = 1 | |
bitsForChar 'C' = 1 | |
bitsForChar 'g' = 2 | |
bitsForChar 'G' = 2 | |
bitsForChar 't' = 3 | |
bitsForChar 'T' = 3 | |
bitsForChar _ = error "Ay, Caramba!" | |
charForBits :: Word64 -> Char | |
charForBits 0 = 'A' | |
charForBits 1 = 'C' | |
charForBits 2 = 'G' | |
charForBits 3 = 'T' | |
charForBits _ = error "Ay, Caramba!" | |
packKey :: B.ByteString -> Word64 | |
packKey = go zeroBits | |
where | |
go k bs = case B.uncons bs of | |
Nothing -> k | |
Just (c, cs) -> go (unsafeShiftL k 2 .|. bitsForChar c) cs | |
{-# INLINE packKey #-} | |
unpackKey :: Int -> Word64 -> B.ByteString | |
unpackKey = go [] | |
where | |
go s 0 _ = B.pack s | |
go s l i = go (charForBits (i .&. 3) : s) (l-1) (unsafeShiftR i 2) | |
{-# INLINE unpackKey #-} | |
countOccurrences :: Int -> Int -> B.ByteString -> ST s (HashTable s Word64 Int) | |
countOccurrences jumpSize frameSize input = do | |
t <- H.initialize 1024 | |
let go bs = unless (B.length bs < frameSize) $ do | |
let k = takeFrame bs | |
H.alter t (Just . maybe 1 (+1)) (packKey k) | |
go (nextFrame bs) | |
go input | |
return t | |
where | |
takeFrame = B.take frameSize | |
nextFrame = B.drop jumpSize | |
extractSequence :: String -> B.ByteString -> B.ByteString | |
extractSequence s = findSeq | |
where | |
prefix = B.pack ('>' : s) | |
skipSeq = | |
B.dropWhile (/= '>') | |
. B.drop 1 | |
takeSeq = | |
B.filter (/= '\n') | |
. B.takeWhile (/= '>') -- extract until next header | |
. B.dropWhile (/= '\n') -- skip header | |
findSeq str | |
| prefix `B.isPrefixOf` str = takeSeq str | |
| otherwise = findSeq (skipSeq str) | |
main :: IO () | |
main = do | |
s <- extractSequence "THREE" <$> B.getContents | |
let keys = ["GGT","GGTA","GGTATT","GGTATTTTAATT","GGTATTTTAATTTATAGT"] | |
let threads = [0 .. 63] | |
let threadWorkOcc key tid = runST $ do | |
t <- countOccurrences (length threads) (B.length key) (B.drop tid s) | |
fromMaybe 0 <$> H.lookup t (packKey key) | |
let calcOcc key = sum $ runEval $ | |
mapM (rpar . threadWorkOcc (B.pack key)) threads | |
let threadWorkFreq len tid = runST $ do | |
t <- countOccurrences (length threads) len (B.drop tid s) | |
vs <- H.toList t | |
return $ map (\(k, v) -> (B.unpack (unpackKey len k), freq v)) vs | |
where | |
freq v = 100 * fromIntegral v / fromIntegral (B.length s - len + 1) | |
let calcFreq len = | |
let l = concat $ runEval $ mapM (rpar . threadWorkFreq len) threads | |
m = foldr (uncurry $ M.insertWith (+)) M.empty l | |
in | |
M.toList m | |
let resultsOcc = map (\k -> (k, calcOcc k)) keys | |
printFreq (calcFreq 1) | |
putStrLn "" | |
printFreq (calcFreq 2) | |
putStrLn "" | |
forM_ resultsOcc $ \(k, r) -> printf "%d\t%s\n" r k | |
where | |
sortFreq = sortBy | |
(\ (k :: String, v :: Double) (k', v') -> | |
(compare v' v) `mappend` (compare k k')) | |
printFreq :: [(String, Double)] -> IO () | |
printFreq l = forM_ (sortFreq l) $ uncurry (printf "%s %.3f\n") |
Could you please share the input file on which program runs 24 sec?
Here: https://surfdrive.surf.nl/files/index.php/s/tohcryAeBL3dPHp (66.6MB)
On my machine it runs around 48s
.
With +RTS -N4
and compiled with -O2 -threaded
? It takes 81 seconds for me without multithreading.
Yes.
-O2 -threaded
compilation flags.- runtime:
time knucleotids +RTS -N4 -A64m -n4m -qb0 -K2048M -RTS < fasta25000000.txt
A 30.295
T 30.151
C 19.800
G 19.754
AA 9.177
TA 9.132
AT 9.131
TT 9.091
CA 6.002
AC 6.001
AG 5.987
GA 5.984
CT 5.971
TC 5.971
GT 5.957
TG 5.956
CC 3.917
GC 3.911
CG 3.909
GG 3.902
1471758 GGT
446535 GGTA
47336 GGTATT
893 GGTATTTTAATT
893 GGTATTTTAATTTATAGT
knucleotids +RTS -N4 -A64m -n4m -qb0 -K2048M -RTS < fasta25000000.txt 171.68s user 0.73s system 364% cpu 47.367 total
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This is what I mean by SAT: