tweak colours
[maximus:gmndl.git] / Calculate.hs
1 {-# LANGUAGE BangPatterns #-}
2
3 {-
4
5     gmndl -- Mandelbrot Set explorer
6     Copyright (C) 2010,2011  Claude Heiland-Allen <claudiusmaximus@goto10.org>
7
8     This program is free software; you can redistribute it and/or modify
9     it under the terms of the GNU General Public License as published by
10     the Free Software Foundation; either version 2 of the License, or
11     (at your option) any later version.
12
13     This program is distributed in the hope that it will be useful,
14     but WITHOUT ANY WARRANTY; without even the implied warranty of
15     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16     GNU General Public License for more details.
17
18     You should have received a copy of the GNU General Public License along
19     with this program; if not, write to the Free Software Foundation, Inc.,
20     51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
21
22 -}
23
24 module Calculate (convert, renderer) where
25
26 -- simple helpers
27 import Control.Monad (when)
28
29 -- concurrent renderer with capability-specific scheduling
30 import Control.Concurrent (MVar, newEmptyMVar, takeMVar, tryTakeMVar, tryPutMVar, threadDelay, forkIO, killThread)
31 import GHC.Conc (forkOnIO, numCapabilities)
32
33 -- each worker uses a mutable unboxed array of Bool to know which pixels
34 -- it has already started to render, to avoid pointless work duplication
35 import Data.Array.IO (IOUArray, newArray, readArray, writeArray, inRange)
36
37 -- each worker thread keeps a queue of pixels that it needs to render or
38 -- to continue rendering later
39 import Data.PriorityQueue (PriorityQueue, newPriorityQueue, enqueue, enqueueBatch, dequeue)
40
41 -- poking bytes into memory is dirty, but it's quick and allows use of
42 -- other fast functions like memset and easy integration with OpenGL
43 import Foreign (Word8)
44
45 -- higher precision arithmetic using libqd
46 import Numeric.QD.DoubleDouble (DoubleDouble())
47 import Numeric.QD.QuadDouble (QuadDouble())
48
49 import Data.Complex (Complex((:+)))
50 import Complex (Turbo(sqr, twice), convert)
51
52 -- some type aliases to shorten things
53 type B = Word8
54 type N = Int
55 type R = Double
56
57 {-
58 -- colour space conversion from HSV [0..1] to RGB [0..1]
59 -- HSV looks quite 'chemical' to my eyes, need to investigate something
60 -- better to make it feel more 'natural'
61 hsv2rgb :: R -> R -> R -> (R, R, R)
62 hsv2rgb !h !s !v
63   | s == 0 = (v, v, v)
64   | h == 1 = hsv2rgb 0 s v
65   | otherwise =
66       let !i = floor (h * 6) `mod` 6 :: N
67           !f = (h * 6) - fromIntegral i
68           !p = v * (1 - s)
69           !q = v * (1 - s * f)
70           !t = v * (1 - s * (1 - f))
71       in  case i of
72             0 -> (v, t, p)
73             1 -> (q, v, p)
74             2 -> (p, v, t)
75             3 -> (p, q, v)
76             4 -> (t, p, v)
77             5 -> (v, p, q)
78             _ -> (0, 0, 0)
79 -}
80
81 hsv2rgb' :: R -> R -> R -> (R, R, R)
82 hsv2rgb' !h !s !l =
83   let !a = 2 * pi * h
84       !ca = cos a
85       !sa = sin a
86       !y = l / 2
87       !u = s / 2 * ca
88       !v = s / 2 * sa
89       !r = y + 1.407 * u
90       !g = y - 0.677 * u - 0.236 * v
91       !b = y + 1.848 * v
92   in  (r, g, b)
93
94 -- compute RGB [0..255] bytes from the results of the complex iterations
95 -- don't need very high precision for this, as spatial aliasing will be
96 -- much more of a problem in intricate regions of the fractal
97 colour :: Complex Double -> Complex Double -> N -> (B, B, B)
98 colour !(zr:+zi) !(dzr:+dzi) !n =
99   let -- micro-optimization - there is no log2 function
100       !il2 = 1 / log 2
101       !zd2 = sqr zr + sqr zi
102       !dzd2 = sqr dzr + sqr dzi
103       -- normalized escape time
104       !d = (fromIntegral n :: R) - log (log zd2 / log escapeR2) * il2
105       !dwell = fromIntegral (floor d :: N)
106       -- final angle of the iterate
107       !finala = atan2 zi zr
108       -- distance estimate
109       !de = (log zd2 * il2) * sqrt (zd2 / dzd2)
110       !dscale = -log de * il2
111       -- HSV is based on escape time, distance estimate, and angle
112       !hue = log ((- log de * il2) `max` 1) * il2 / 16 + log d * il2
113       !saturation = 0 `max` (log d * il2 / 8) `min` 1
114       !value = 0 `max` (1 - dscale / 256) `min` 1
115       !h = hue - fromIntegral (floor hue :: N)
116       -- adjust saturation to give concentric striped pattern
117       !k = dwell / 2
118       !satf = if k - fromIntegral (floor k :: N) >= (0.5 :: R) then 0.9 else 1
119       -- adjust value to give tiled pattern
120       !valf = if finala < 0 then 0.9 else 1
121       -- convert to RGB
122       (!r, !g, !b) = hsv2rgb' h (satf * saturation) (valf * value)
123       -- convert to bytes
124       !rr = floor $ 0 `max` (255 * r) `min` 255
125       !gg = floor $ 0 `max` (255 * g) `min` 255
126       !bb = floor $ 0 `max` (255 * b) `min` 255
127   in  (rr, gg, bb)
128
129 -- a Job stores a pixel undergoing iterations
130 data Job c = Job !N !N !(Complex c) !(Complex c) !(Complex c) !N
131
132 -- the priority of a Job is how many iterations have been computed:
133 -- so 'fresher' pixels drop to the front of the queue in the hope of
134 -- avoiding too much work iterating pixels that will never escape
135 priority :: Job c -> N
136 priority !(Job _ _ _ _ _ n) = n
137
138 -- add a job to a work queue, taking care not to duplicate work
139 -- there is no race condition here as each worker has its own queue
140 addJob :: RealFloat c => N -> N -> Complex c -> c -> PriorityQueue IO (Job c) -> IOUArray (N,N) Bool -> N -> N -> IO ()
141 addJob !w !h !c !zradius' todo sync !i !j = do
142   already <- readArray sync (j, i)
143   when (not already) $ do
144     writeArray sync (j, i) True
145     enqueue todo $! Job i j (coords w h c zradius' i j) 0 0 0
146
147 -- spawns a new batch of workers to render an image
148 -- returns an action that stops the rendering
149 renderer' :: (Turbo c, RealFloat c) => MVar () -> ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex c -> c -> IO (IO ())
150 renderer' done rng output !c !zradius' = do
151   wdog <- newEmptyMVar
152   workerts <- mapM (\w -> forkOnIO w $ worker wdog rng c zradius' output w) [ 0 .. workers - 1 ]
153   watcher <- forkIO $ do
154     () <- takeMVar wdog
155     let loop = do
156           threadDelay 10000000 -- 10 seconds
157           m <- tryTakeMVar wdog
158           case m of
159             Nothing -> mapM_ killThread workerts >> tryPutMVar done () >> return ()
160             Just () -> loop
161     loop
162   return $ killThread watcher >> mapM_ killThread workerts
163
164 -- compute the Complex 'c' coordinate for a pixel in the image
165 coords :: RealFloat c => N -> N -> Complex c -> c -> N -> N -> Complex c
166 coords !w !h !c !zradius' !i !j = c + ( (fromIntegral (i - w`div`2) * k)
167                                       :+(fromIntegral (h`div`2 - j) * k))
168   where !k = zradius' / (fromIntegral $ (w `div` 2) `min` (h `div` 2))
169
170 -- the worker thread enqueues its border and starts computing iterations
171 worker :: (Turbo c, RealFloat c) => MVar () -> ((N,N),(N,N)) -> Complex c -> c -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO ()
172 worker wdog rng@((y0,x0),(y1,x1)) !c !zradius' output !me = do
173   sync <- newArray rng False
174   queue <- newPriorityQueue priority
175   let addJ = addJob w h c zradius' queue sync
176       js = filter mine (border w h)
177       w = x1 - x0 + 1
178       h = y1 - y0 + 1
179   mapM_ (flip (writeArray sync) True) js
180   enqueueBatch queue (map (\(j,i) -> Job i j (coords w h c zradius' i j) 0 0 0) js)
181   compute wdog rng addJ output queue
182   where mine (_, i) = i `mod` workers == me -- another dependency on spread
183
184 -- the compute engine pulls pixels from the queue until there are no
185 -- more, and calculates a batch of iterations for each
186 compute :: (Turbo c, RealFloat c) => MVar () -> ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job c) -> IO ()
187 compute wdog rng addJ output queue = do
188   mjob <- dequeue queue
189   case mjob of
190     Just (Job i j c z dz n) -> do
191       let -- called when the pixel escapes
192           done' !(zr:+zi) !(dzr:+dzi) !n' = {-# SCC "done" #-} do
193             _ <- tryPutMVar wdog ()
194             let (r, g, b) = colour (convert zr :+ convert zi) (convert dzr :+ convert dzi) n'
195             output i j r g b
196             -- a wavefront of computation spreads to neighbouring pixels
197             sequence_
198               [ addJ x y
199               | u <- spreadX
200               , v <- spreadY
201               , let x = i + u
202               , let y = j + v
203               , inRange rng (y, x)
204               ]
205           -- called when the pixel doesn't escape yet
206           todo' !z' !dz' !n' = {-# SCC "todo" #-} {- output i j 255 0 0 >> -} enqueue queue $! Job i j c z' dz' n'
207       calculate c limit z dz n done' todo'
208       compute wdog rng addJ output queue
209     Nothing -> return () -- no pixels left to render, so finish quietly
210
211 -- the raw z->z^2+c calculation engine
212 -- also computes the derivative for distance estimation calculations
213 -- this function is crucial for speed, too much allocation will slooow
214 -- everything down severely
215 calculate :: (Turbo c, RealFloat c) => Complex c -> N -> Complex c -> Complex c -> N -> (Complex c -> Complex c -> N -> IO ()) -> (Complex c -> Complex c -> N -> IO ()) -> IO ()
216 calculate !c !m0 !z0 !dz0 !n0 done todo = go m0 z0 dz0 n0
217   where
218     go !m !z@(zr:+zi) !dz !n
219       | not (sqr zr + sqr zi < er2) = done z dz n
220       | m <= 0 = todo z dz n
221       | otherwise = go (m - 1) (sqr z + c) (let !zdz = z * dz in twice zdz + 1) (n + 1)
222     !er2 = convert escapeR2
223
224 -- dispatch to different instances of renderer depending on required precision
225 -- if zoom is low, single precision Float is ok, but as soon as pixel spacing
226 -- gets really small, it's necessary to increase it
227 -- it's probably not even worth using Float - worth benchmarking this and
228 -- also the DD and QD types (which cause a massively noticeable slowdown)
229 renderer :: (RealFloat c) => MVar () -> ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex c -> c -> IO (IO ())
230 renderer done rng output !c !zradius'
231   | zoom' < 20  = {-# SCC "rF"  #-} renderer' done rng output (f c :: Complex Float       ) (g zradius')
232   | zoom' < 50  = {-# SCC "rD"  #-} renderer' done rng output (f c :: Complex Double      ) (g zradius')
233   | zoom' < 100 = {-# SCC "rDD" #-} renderer' done rng output (f c :: Complex DoubleDouble) (g zradius')
234   | otherwise   = {-# SCC "rQD" #-} renderer' done rng output (f c :: Complex QuadDouble  ) (g zradius')
235   where f !(cr :+ ci) = convert cr :+ convert ci
236         g !x = convert x
237         zoom' = - logBase 2 (zradius' / (fromIntegral $ w `min` h))
238         ((x0,y0), (x1, y1)) = rng
239         w = x1 - x0 + 1
240         h = y1 - y0 + 1
241
242 -- start rendering pixels from the edge of the image
243 -- the Mandelbrot Set and its complement are both simply-connected
244 -- discounting spatial aliasing any point inside the boundary that is
245 -- in the complement is 'reachable' from a point on the boundary that
246 -- is also in the complement - probably some heavy math involved to
247 -- prove this though
248 -- note: this implicitly depends on the spread values below - it's
249 -- necessary for each interlaced subimage (one per worker) to have
250 -- at least a one pixel deep border
251 border :: N -> N -> [(N, N)]
252 border !w !h = concat $
253   [ [ (j, i) | i <- [ 0 .. w - 1 ], j <- [ 0 ] ]
254   , [ (j, i) | j <- [ 0 .. h - 1 ], i <- [ 0 .. workers - 1 ] ]
255   , [ (j, i) | j <- [ 0 .. h - 1 ], i <- [ w - workers .. w - 1 ] ]
256   , [ (j, i) | i <- [ 0 .. w - 1 ], j <- [ h - 1 ] ]
257   ]
258
259 -- which neighbours to activate once a pixel has escaped
260 -- there are essentially two choices, with x<->y swapped
261 -- choose greater X spread because images are often wider than tall
262 -- other schemes wherein the spread is split in both directions
263 -- might benefit appearance with large worker count, but too complicated
264 spreadX, spreadY :: [ N ]
265 spreadX = [ -workers, 0, workers ]
266 spreadY = [ -1, 0, 1 ]
267
268 -- number of worker threads
269 -- use as many worker threads as capabilities, with the workers
270 -- distributed 1-1 onto capabilities to maximize CPU utilization
271 workers :: N
272 workers = numCapabilities
273
274 -- iteration limit per pixel
275 -- at most this many iterations are performed on each pixel before it
276 -- is shunted to the back of the work queue
277 -- this should be tuneable to balance display updates against overheads
278 limit :: N
279 limit = (2^(13::N)-1)
280
281 -- escape radius for fractal iteration calculations
282 -- once the complex iterate exceeds this, it's never coming back
283 -- theoretically escapeR = 2 would work
284 -- but higher values like this give a significantly smoother picture
285 escapeR, escapeR2 :: R
286 escapeR = 65536
287 escapeR2 = escapeR * escapeR