Fractal program

gistfile1.hs

{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}

import Prelude hiding (writeFile)
import Control.Monad.ST (ST, runST)
import Data.Array (Array)
import Data.Array.IO (IOArray)
import Data.Array.Unboxed (UArray)
import Data.Array.IArray (IArray, bounds, (!))
import Data.Array.ST (STArray, runSTArray, runSTUArray)
import Data.Array.MArray (MArray, newArray, getBounds, freeze, readArray, writeArray)
import Data.Ix (Ix, inRange)
import Data.Monoid (Monoid, (<>), mempty, mappend, Sum)
import Data.Fixed (mod')
import Debug.Trace (trace)
import Data.Maybe (catMaybes)
import Data.Foldable (foldl')
import Codec.Picture (Pixel, Image, generateImage, Pixel8)
import Codec.Picture.Png (encodePng)
import Data.ByteString.Lazy (ByteString, writeFile)

data Fractal o = Fractal { output :: o , successors :: [Fractal o] }

data ColorPoint p c = ColorPoint {location :: p, color :: c}

data RGB = RGB {r :: !Double, g :: !Double, b :: !Double}

data Bounds i = Bounds i i

instance Monoid RGB where
	mempty = RGB 0 0 0
	mappend (RGB ra ga ba) (RGB rb gb bb) = RGB (ra+rb) (ga+gb) (ba+bb)

instance Monoid Double where
	mempty = 0
	mappend = (+)

data F1State = F1State Double (Int, Int) Double

fractal1gen :: F1State -> Fractal (ColorPoint (Int, Int) Double)
fractal1gen (F1State intensity (x,y) direction) = Fractal current nexts
	where
	current = ColorPoint (x,y) intensity
	nexts = map fractal1gen newStates
	newStates = catMaybes $ map toState deltas
	deltas = filter (/=(0,0)) [(dx,dy) | dx <- [-3..3], dy <- [-3..3]]
	toState (dx,dy) = if intensity' > 0.001
		then Just $ F1State intensity' (x+dx,y+dy) direction'
		else Nothing
		where
		intensity' = 0.99 * intensity * directionality * falloff
		direction' = (direction + 0.1) `mod'` (2 * pi)
		directionality = exp (negate ((3 * angleDifference)^2))
		angleDifference = angle - direction
		angle = atan2 (fromIntegral dy) (fromIntegral dx) `mod'` (2 * pi)
		falloff = 1.0 / sqrt ((fromIntegral dx)^2 + (fromIntegral dy)^2)

fractal1 :: Fractal (ColorPoint (Int, Int) Double)
fractal1 = fractal1gen (F1State 1.0 (0,0) 0.0)

nearbyes :: (Int, Int) -> [(Int, Int)]
nearbyes (x,y) = filter (/=(x,y)) [(x+dx, y+dy) | dx <- [-3..3], dy <- [-3..3]]

-- `c` is the "Color" type and `i` is the coordinate type.
-- Example color: `Sum Double` (See `Data.Monoid`'s `Sum` type.)
-- Example coordinate: (Int, Int)
render :: (Monoid c, Ix i) => Bounds i -> Fractal (ColorPoint i c) -> Array i c
render (Bounds a b) f = runSTArray $ do
	image <- newArray (a,b) mempty
	render' 100 image f
	return image

renderIO :: forall i c .(Monoid c, Ix i) => Bounds i -> Fractal (ColorPoint i c) -> IO (Array i c)
renderIO (Bounds a b) f = do
	image <- newArray (a,b) mempty :: IO (IOArray i c)
	render' 100 image f
	freeze image 

renderUnboxed :: (Ix i) => Bounds i -> Fractal (ColorPoint i Double) -> UArray i Double
renderUnboxed (Bounds a b) f = runSTUArray $ do
	image <- newArray (a,b) mempty
	render' 100 image f
	return image 

-- `c` is the "Color" type, `a` is the array type,
-- `m` is the monad we're calculating the mutable array in,
-- and `i` is the coordinate type.
-- Example color: `Sum Double` (See `Data.Monoid`'s `Sum` type.)
-- Example array: STArray, STUArray, IOArray
-- Example monad: ST, IO
-- Example coordinate: (Int, Int)
render' :: (Monoid c, MArray a c m, Ix i) => 
	Int -> a i c -> Fractal (ColorPoint i c) -> m ()
render' 0 _ _ = return () -- Depth exceeded
render' n array (Fractal (ColorPoint i c) successors) = do
	bounds <- getBounds array
	if inRange bounds i 
		then do
			old <- readArray array i
			writeArray array i $! (old <> c) -- trace "ayy" (old <> c)
			mapM_ (render' (n-1) array) successors
		else return ()

arrayToImage :: (IArray a c, Pixel p) => (c -> p) -> a (Int, Int) c -> Image p
arrayToImage f arr = generateImage getPixel dimx dimy
	where
	getPixel x y = f $ arr ! (x + minx, y + miny)
	((minx, miny), (maxx, maxy)) = bounds arr
	(dimx, dimy) = (maxx - minx, maxy - miny)

doubleArrayToPNG :: Array (Int, Int) Double -> ByteString
doubleArrayToPNG arr = encodePng $ arrayToImage toPixel8 arr
	where toPixel8 double = round (double * 255) :: Pixel8


--fsum = foldl' (\count val -> if val > 0 then count+1 else count) 0 
main = do
	let png = doubleArrayToPNG $  render (Bounds (-100, -100) (100, 100)) fractal1 
	writeFile "out.png" png
--main = do
--	result <- renderIO (Bounds (-100, -100) (100, 100)) fractal1 
--	print . fsum $ result