Môshe van der Sterre - Public notes

-- NOTE: Only Part 1 is done, the other parts have rough edges.

-- PART 0: Imports

-- We use complex numbers, haskell has excelent support for them.

import Data.Complex

-- And we use Either in a Monadic way, this instance comes from EitherT.

import Control.Monad.Either

-- For the matrix datatype, we use the hmatrix package.

import Data.Packed.Matrix

-- Also from the hmatrix package is the multiplication operator.

import Numeric.LinearAlgebra

-- We use Word8 for the binary output

import Data.Word

-- And putWord8

import qualified Data.ByteString as B

-- LiftM2

import Control.Monad (liftM2)

-- intersperse

import Data.List

-- PART 1: The mandelbrot set

-- Implementing the family of mandelbrot polynomials is a copy-paste from
-- Wikipedia: http://en.wikipedia.org/wiki/Mandelbrot_set#Formal_definition

p :: Complex Double -> Complex Double -> Complex Double
p c z = z ^ 2 + c

-- We need to track the moment of 'escape', otherwise we can not calculate the
-- gradient.

track :: (a -> Bool) -> b -> a -> Either b a
track cond nr val | cond val  = Left nr
                  | otherwise = Right val

-- Because we can not iterate infinitely in the real world, we do two things:
-- 1. we short-cut when the result is found.

canStop :: Complex Double -> Bool
canStop n = 2 < magnitude n

-- 2. we take only a finite amount of steps.

steps :: (Int -> a) -> Int -> [a]
steps step count = map step $ enumFromThenTo count (count - 1) 1

-- We also define the calculation for each iteration.

calc :: Complex Double -> Int -> Complex Double -> Either Int (Complex Double)
calc coord stepnr value = track canStop stepnr $ p coord value

-- Chaining the steps lazily can be done with foldr.
-- This function has some similarity with sequence and sequence_, only the
-- functions are chained to use the results.

sequenceChain :: Monad m => a -> [a -> m a] -> m a
sequenceChain initial = foldr (=<<) (return initial)

-- Now we define a function that checks if a number is part of the mandelbrot.
-- Right is inside the set, Left is outside the Set.
-- The value of Right is the result after stepcount iterations.
-- The value of Left is the iteration number where we determined the
-- coordinate is outside the mandelbrot set.

inMbrot :: Int -> Complex Double -> Either Int (Complex Double)
inMbrot stepcount coord = sequenceChain 0 $ steps (calc coord) stepcount

-- This inMbrot function takes a complex number, and returns the escape time,
-- This is all there is to the mandelbrot set.

-- All of the above fit in 2 lines of haskell, here are they:

inMbrot' a b = foldr (=<<) (return 0) $ map c [a,(a-1)..1] where
         c e f = if 2 < magnitude g then Left e else Right g where g = f*f+b



-- PART 2: The transformation matrix

-- I find the SVG specification a good reference to build the basic
-- transformation matrices, look at it to see what is going on here.
-- http://www.w3.org/TR/SVG/coords.html#TransformMatrixDefined

tmScale, tmTranslate :: Double -> Double -> Matrix Double
tmScale x y = (3><3) [x, 0, 0, 0, y, 0, 0, 0, 1]
tmTranslate x y = (3><3) [1, 0, x, 0, 1, y, 0, 0, 1]

-- We normalize the output coordinate system for clarity.
-- Note: we use Double instead of Integer for the width and height, this is
-- lazy coding to avoid calls to fromIntegral.

-- We normalize the x-axis to -2, 2.

normScale :: Double -> Double -> Matrix Double
normScale wdt hgt = tmScale s s where s = 4 / (wdt - 1)

-- And we center on 0, 0.

normTranslate :: Double -> Double -> Matrix Double
normTranslate wdt hgt = tmTranslate (-2) (-2 / (wdt - 1) * (hgt - 1))

-- Combining the normalization:

normOut :: Double -> Double -> Matrix Double
normOut wdt hgt = normTranslate wdt hgt <> normScale wdt hgt

-- Transformation to locate the wanted point on the complex plane, given a
-- Complex number and scale factor. (assumes the normalized coordinate system)

locate :: Complex Double -> Double -> Matrix Double
locate (real :+ imag) scale = tmScale scale scale <> tmTranslate real imag

-- Combine all the above.

getMatrix :: Double -> Double -> (Complex Double) -> Double -> Matrix Double
getMatrix width height pos scale = locate pos scale <> normOut width height


-- The following functions deal with the input and output to these
-- transformation matrices.

-- Each pixel is plugged in a Matrix Double for use as input to the transform.

pixel :: (Double, Double) -> Matrix Double
pixel (x, y) = (3><1) [x, y, 1]

-- After the transform, a Complex Double is pulled out of the Matrix Double.

coord :: Matrix Double -> Complex Double
coord m = m @@> (0, 0) :+ m @@> (1, 0)

-- With a transformation matrix and the position for the pixel, we get the
-- complex number that can be used to calculate the escape time.

pixelCoord :: Matrix Double -> Double -> Double -> Complex Double
pixelCoord transf y x = coord $ transf <> pixel (x, y)

-- Now, we have all tools to calculate the desired mandelbrot results.

-- This places all coordinates in a flat list (with the horizontal dimension
-- first). To get the results, simply map the mandelbrot on this list.

coords :: Double -> Double -> (Complex Double) -> Double -> [Complex Double]
coords wdt hgt pos scale = liftM2 (pixelCoord $ getMatrix wdt hgt pos scale)
                           (enumFromTo 0 (wdt - 1)) (enumFromTo 0 (hgt - 1))

-- The steps above can also be written a bit more dense, by pre-calculating the
-- matrix we can even ommit most of it. The following function includes a
-- rounded equivalent for 'coords 300 250 ((-0.5):+0) 1'

coords' = liftM2(\x y->(liftM2(:+)(@@>(0,0))(@@>(1,0)))(((3><3)[1.34e-2,0,-2.5,0,1.34e-2,-1.67,0,0,1])<>(3><1)[x,y,1]))[0..299][0..249]


-- PART 3: Outputting

-- We will output the mandelbrot to a PGM image, this format is easy, and we
-- can use the Word8 type to output it on standard output.
-- PGM images are graymaps, we simply use the escape time as pixel value.

-- We output a binary PGM (P5) file.
-- This function has a bug, ASCII assumed.

ppmHeader :: Int -> Int -> Int -> String
ppmHeader w h i = "P5 " ++ (concat $ intersperse " " $ map show [w, h, i])++" "

-- Convert the mandelbrot result to a usable byte value

color :: Either Int (Complex Double) -> Word8
color (Left x) = fromIntegral x
color (Right _) = 0

-- We start the calculations here.

ppmData :: Double -> Double -> Complex Double -> Double -> Int -> [IO ()]
ppmData wdt hgt pos scale iters = map (B.putStr . B.singleton . color . inMbrot iters) $ coords wdt hgt pos scale

-- The complete ppm.

ppm :: Double -> Double -> Complex Double -> Double -> Int -> IO ()
ppm wdt hgt pos scale iters = do
        putStr $ ppmHeader (floor wdt) (floor hgt) iters
        sequence $ ppmData wdt hgt pos scale iters
        return ()

main = ppm 900 750 ((-0.5) :+ 0) 1 255