MandelbrotSet practice --- sequence and basic parallel version

Mandelbrot.hs

{-# LANGUAGE BangPatterns #-}

-- Compile command: ghc -O2 Mandelbrot.hs -rtsopts -threaded -fllvm

import qualified Control.Monad.Par           as Par
import           Control.Parallel.Strategies (parList, rseq, using)
import           Data.Binary.IEEE754         (putFloat64le)
import           Data.Binary.Put
import qualified Data.ByteString.Lazy        as BL
import           Data.Time.Clock             (diffUTCTime, getCurrentTime)
import qualified Data.Vector                 as V
import           System.Environment          (getArgs)
import           Text.Printf

type Point = (Double, Double)
type Range = (Double, Double)
type Plane = (Range, Range)

maxIteration :: Integer
maxIteration = 3000

inMandel :: Point -> Bool
inMandel (!ox, !oy) = q maxIteration (ox, oy)
  where q 0 _ = True
        q t (!x, !y) = not (x^2 + y^2 >= 4.0) && q (t - 1) (x', y')
          where x' = x^2 - y^2 + ox
                y' = 2 * x * y + oy

planePoints :: Plane -> V.Vector Point
planePoints ((xb, xe), (yb, ye)) =
  do x <- V.enumFromStepN xb pointDist x_size
     y <- V.enumFromStepN yb pointDist y_size
     return (x, y)
  where pointDist = 0.001
        x_size = truncate $ (xe - xb) / pointDist
        y_size = truncate $ (ye - yb) / pointDist

planeToMandelPoints :: Plane -> V.Vector Point
planeToMandelPoints = V.filter inMandel . planePoints

splitPlane :: Integer -> Plane -> [Plane]
splitPlane size ((xb,xe), (yb,ye)) = [(rx,ry)| rx <- r xb xe, ry <- r yb ye]
  where side =  sqrt . fromIntegral $ size
        step b e = (e - b) / side
        r b e = zip rs (tail rs)
          where rs = [b, b + step b e .. e]

-- sequential version
mandelSet :: Plane -> V.Vector Point
mandelSet = planeToMandelPoints

-- basic parllel version with `parList`
mandelSetStart :: Integer -> Plane -> V.Vector Point
mandelSetStart size p = V.concat (map planeToMandelPoints (splitPlane size p) `using` parList rseq)

-- parallel with ParMonad using `parMap`
mandelSetPar :: Integer -> Plane -> V.Vector Point
mandelSetPar size p = V.concat $ Par.runPar $ Par.parMap planeToMandelPoints (splitPlane size p)

-- mandelSetPar' :: Integer -> Plane -> V.Vector Point
-- mandelSetPar' size p = V.concat $ Par.runPar $ v
--     where v = do p' <- mapM (Par.spawn . return . planeToMandelPoints) (splitPlane size p)
--                  mapM Par.get p'

-- serialize data
serDoublesData :: V.Vector Point -> Put
serDoublesData !ps = V.mapM_ putFloat64le $ uncurry (V.++) . V.unzip $ ps

main :: IO ()
main = do
  let fn = "mandelbrot.bin"
  let plane = ((-2.0, 1.0), (0.0, 1.0))

  args <- getArgs
  t0 <- getCurrentTime

  let !p = case args of
             ["seq"     ] -> mandelSet plane
             ["start", n] -> mandelSetStart (read n) plane
             ["par",   n] -> mandelSetPar (read n) plane
             _            -> error "Error argument(s)"
  BL.writeFile fn (runPut (serDoublesData p))
  t1 <- getCurrentTime
  printf "Total time: %.2f\n" (realToFrac (diffUTCTime t1 t0) :: Double)

plotMandelbrot.py

#!/usr/bin/env python3

import matplotlib.pyplot as plt
import struct


def read_point_list(fn):
    with open(fn, 'rb') as f:
        s = f.read()
        l = [struct.unpack('d', s[x:x+8])[0] for x in xrange(0, len(s), 8)]

        l_len_half = len(l)/2
        xs = l[0:l_len_half]
        ys = l[l_len_half:]
        return (xs, ys)

    return ([],[])

def main():
    xs, ys = read_point_list('./mandelbrot.bin')

    plt.plot(xs, ys, 'k,')
    plt.plot(xs, [-y for y in ys], 'k,')

    plt.axis('equal')
    plt.show()


if __name__ == '__main__':
    main()