Parsing and plotting complex numbers

Question

I am trying to write a Haskell application that parses a bunch of complex numbers from a Mathematica output and then produces a PPM image.

The end product looks like this:

The code to create this is as follows:

module Main where

import PPM
import System.IO
import Control.Applicative as A
import Control.Monad.ST
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as LB
import qualified Data.Attoparsec.ByteString.Char8 as P
import qualified Data.Attoparsec.ByteString.Lazy as LP


apply :: a -> (a -> b) -> b
apply a = \f -> f a


---------- PARSING ----------

iT = B.pack "*I"

data Complex = Complex Double Double
    deriving (Show)

cAdd:: Complex -> [Complex] -> Complex
cAdd (Complex a b) [(Complex c d)] = Complex (a + c) (b + d)
cAdd c _ = c
cConv :: Complex -> Complex
cConv (Complex a b) = Complex b a

parseTerm :: P.Parser Complex
parseTerm = parseDouble <*> parseConv
    where 
        parseDouble = liftA (\x -> apply (Complex x 0)) $
            liftA2 (\x y -> read ((x:y) ++ "0")) P.anyChar $ many $ P.satisfy $ \x -> P.isDigit x || x == '.'
        parseConv = P.option id $ liftA (const cConv) $ P.string iT

test :: P.Parser String
test = liftA2 (\x y -> (x:y)) P.anyChar $ many $ P.satisfy $ \x -> P.isDigit x || x == '.'

parseComplex :: P.Parser Complex
parseComplex = liftA2 cAdd parseTerm $ many $ many (P.char '+') A.*> parseTerm


---------- IMAGE PROCESSING ----------

bounds = 2.0 :: Double
imgSize = 400 :: Int

intensity = 300.0 :: Double
falloff = 0.8 :: Double
(r, g, b) = (255.0, 76.5, 25.5) :: (Double, Double, Double)

computeIndex :: Double -> Int
computeIndex x = (+) 1 $ floor $ (x + bounds) / (2 * bounds / fromIntegral imgSize)

convCoord ::(Int, Int) -> Int
convCoord (h, v) = v * imgSize + h

computeCIndex :: Complex -> Int
computeCIndex (Complex r i) = convCoord (computeIndex r, computeIndex i)

genVec :: [Complex] -> V.Vector Int
genVec xs = runST $ do
    mv <- M.replicate (imgSize ^ 2) 0
    mapM_ (incr mv) $ map computeCIndex xs -- unsafeModify
    V.freeze mv
        where
            incr mv i = do
                pre <- M.unsafeRead mv i
                M.unsafeWrite mv i $ pre + 1

colorFunction :: Double -> (Int, Int, Int)
colorFunction d = (floor (y * r), floor(y * g), floor(y * b))
    where y = (intensity * d) ** falloff

colorFunction1 n = (n, n, n)

genImage :: V.Vector Int -> PPM
genImage v = PPM (V.map (\x -> colorFunction((fromIntegral x) / (fromIntegral mx))) v) imgSize imgSize 255
    where mx = V.maximum v

extr :: LP.Result a -> a
extr (LP.Done _ r) = r
extr (LP.Fail _ _ y) = error y

---------- MAIN ----------

main = do
    rawData <- liftA LB.words (LB.readFile "/mnt/hgfs/outputs/out_1-14.txt")
    let formatedData = map (extr.LP.parse parseComplex) rawData

    h <- openFile "test.ppm" WriteMode
    writeImage (genImage (genVec formatedData)) h
    hClose h

This uses a module:

module PPM where

import System.Environment
import System.IO
import GHC.IO.Device
import qualified Data.Vector as V

----- Data Type -----
---------------------
-- Unsafe Vector data type for PPM Images
data PPM = PPM {pixels :: V.Vector (Int, Int, Int) , height :: Int, width :: Int, cap :: Int}
    deriving (Show)



----- Helper Functions -----
----------------------------
writeTup :: (Int, Int, Int) -> Handle -> IO ()
writeTup (r,g,b) h = hPutStrLn h $ (show r) ++ " " ++ (show g) ++ " " ++ (show b)

writeRow :: V.Vector (Int, Int, Int) -> Handle -> IO ()
writeRow v h = helper v h 0
    where helper v h n
            | n < V.length v = do
                writeTup (v V.! n) h
                helper v h $ n + 1
            | otherwise = return ()

-- coordinates start at (0,0) and top left corner
-- returns 0 when out of range
getIndex :: Int -> Int -> Int -> Int -> Int
getIndex h w x y
    | x >= 0 && x < w && y < h && y >= 0 = y * w + x
    | otherwise = 0



----- Library Functions -----
-----------------------------
createBlank :: Int -> Int -> PPM
createBlank h w = PPM pix h w 255
    where pix = V.replicate (w * h) (0, 0, 0)

writePixel :: PPM -> (Int, Int, Int) -> Int -> Int -> PPM
writePixel ppm w x y = PPM nPixels (height ppm) (width ppm) (cap ppm)
    where
        n = getIndex (height ppm) (width ppm) x y
        nPixels = (pixels ppm) V.// [(n, w)]

readPixel :: PPM -> Int -> Int -> (Int, Int, Int)
readPixel ppm x y = (pixels ppm) V.! n
    where n = getIndex (height ppm) (width ppm) x y


writeImage :: PPM -> Handle -> IO ()
writeImage img h = do
    hPutStrLn h "P3"
    hPutStrLn h $ (show (height img)) ++ " " ++ (show (width img))
    hPutStrLn h $ show $ cap img
    writeRow (pixels img) h 

Parsing is done with completely with Applicative Attoparsec parsers on Lazy Bytestrings. Input files are expected to be roughly 3Gbs-8Gbs in size, containing anywhere from 70-600 million roots. Each complex number is separated by a newline and looks something like this:

1.9238394+0.32423423*I

Mathematica outputs are a bit funky, retaining a decimal point, for whole numbers:

1.+123123125123*I

After parsing, I generate the PPM, represented by a mutable vector by first computing the index of the vector that corresponds to each Complex number, then mapping a Monad action over each of these indices. The end result is I have a mutable vector that kind of resembles a sort of histogram of points. Finally, the function colorFunction assigns a color to each point and I write the image to a file.

---

Questions

1. This program takes in a lot of data so I would like it to perform very well. However, I'm not experienced enough to really know what kind of experience I should expect. Is this program structured in a good way? Am I using the right data structures and functions? Are there any small refinements I can make?
2. Since data structures here are lazy, I'm a little bit unsure about how I can do benchmarks on this program. For instance, if I want to just measure the performance of the parser, how can I force ghc to evaluate that portion of the code without generating an image or doing a bunch of IO?
3. I have access to the Research Computing Center at the University of Chicago. I'm interested in whether I would be able to get better performance out of this if the code could be run in parallel. If so, where would I get started?

Answer 1

The program consists of a pipeline of operations:

• breaking up the input file into lines
• parsing each line as a complex number
• updating the histogram vector
• writing out the PPM file

The first step to optimizing the program is to measure how long each of these steps is taking.

criterion is a great tool, but because it runs the test function multiple times it isn't well suited for operations which take many seconds. Also, criterion is designed to get precise timings, and we just need ballpark figures. For our timings we're just going to use the shell's time command.

I first created a test file containing a million random complex numbers, and ran the entire pipeline to get a base line reading. Creating the the PPM from the million complex numbers took 33 secs. To understand if that was a reasonable number or not I wrote a quick perl script to just read in and add up the real parts of the complex numbers:

my $sum = 0;
while (<>) {
  if (m/^(.*?)[+]/) { $sum += $1; }
}
print $sum, "\n";

and it only took 1.7 seconds. Clearly there is a big problem somewhere with the Haskell program.

The next step it to write alternative main functions which cut off the pipeline at various stages, e.g.:

A main which just prints out the number of lines read:

mainLength path =  do
    rawData <- liftA LB.words (LB.readFile path)
    let formatedData = map (extr.LP.parse parseComplex) rawData
    print $ length formatedData

A main which just adds up the real parts of the parsed numbers:

mainSum path = do
    rawData <- liftA LB.words (LB.readFile path)
    let formatedData = map (extr.LP.parse parseComplex) rawData
        rpart (Complex r _) = r
    print $ sum $ map rpart $ formatedData

and the timings I got are as follows:

mainLength   0.18 secs
mainSum     17.37 secs

So what work is being done in mainSum which is not being done in mainLength? Due to laziness it is the conversion of the parsed bytestrings to Double values.

A quick scan of the code reveals that you are using read to perform this conversion. read is notoriously slow, and should be avoided for performance critical code.

A replacement for read :: ByteString -> Double

If you search Google you can find this Stackoverflow article:

https://stackoverflow.com/questions/4489518/efficient-number-reading-in-haskell

Mostly because the answer was submitted by Don Stewart I decided on using the bytestring-lexing module, and came up with this version of parseComplex:

parseComplex' = do
  r <- P.takeTill (== '+')
  P.char '+'
  i <- P.takeTill (== '*')
  P.string iT  -- or just skip this
  return $ Complex (toDouble r) (toDouble i)
  where toDouble s = case (readSigned readDecimal) s of
                       Nothing -> 0
                       Just (d, _) -> d

The timing for mainSum using parseComplex' is:

mainSum'  0.6 secs

So now the first two stages of the pipeline are very performant. The next step is to figure out why updating the histogram vector and writing out the PPM file is taking so long.

Using Ubboxed Vectors

I've found another important optimization - you want to use unboxed vectors instead of the regular vectors. Here is an alternate version of genVec:

import qualified Data.Vector.Unboxed as UnboxedV
import qualified Data.Vector.Unboxed.Mutable as UnboxedM

genVec :: [Complex] -> UnboxedV.Vector Int
genVec xs = runST $ do
  mv <- UnboxedM.replicate (imgSize*imgSize) (0::Int)
  forM_ xs $ \c -> do
    let x = computeCIndex c
    count <- UnboxedM.unsafeRead mv x
    UnboxedM.unsafeWrite mv x (count+1)
  UnboxedV.freeze mv

This will cut down the run time by another couple of seconds (which is now significant since the whole pipeline takes now only takes about 4 secs to run.)

Update:

I think you'll find that improving the Double parsing is about the best you can do for a single threaded program. To scale to 600M points you are going to have to use multiple threads / machines. Fortunately this is a classic map-reduce problem, so there's a lot of tools and libraries (not necessarily in Haskell) that you can draw upon.

If you just want to scale on a single machine using multiple threads, you can put together your own solution using a module like Control.Monad.Par. For a clustering solution you'll probably have to use a third-party framework like Hadoop in which case you might be interested in the hadron package - there is also a video describing it here: https://vimeo.com/90189610