Haskell: Julia Set function

I am trying to write a function that returns the length of a Julia Set where before it passes a certain threshold, 2, with a maximum number of iterations of 255.

i have function f z = z * z - (a :+ b) I iterate this function over a complex number until the magnitude of a + bi crosses 2 and then count the number of iterations it took.

my first attempt was

iter1 = genericLength . take 255 . takeWhile ((<=2) . magnitude) . iterate f

but it is painfully slow.

my second attempt is

iters2 x y = iters2' f ((<=2) . magnitude) (a :+ b)
iters2' f' p c = len c 0
  where
    len c acc = if p c && acc < 255 then len (f' c) (acc + 1) else acc

This is still slow, considering that I will have to do millions of iterations.

Can someone help with how to make this faster?

Also, is there list fusion built into the Prelude list functions?

Solution

For a loop like this handcrafted recursion usually is the best approach.

Here is a comparison of three implementations:

iter1 - original code
iter2 - handcrafted recursion
iter3 - @Alec's solution

The timings are (with -O2):

benchmarking julia/1
time                 17.30 μs   (17.06 μs .. 17.56 μs)
                     0.995 R²   (0.991 R² .. 0.998 R²)
mean                 17.79 μs   (17.37 μs .. 18.51 μs)
std dev              1.839 μs   (1.008 μs .. 3.285 μs)
variance introduced by outliers: 86% (severely inflated)

benchmarking julia/2
time                 1.456 μs   (1.448 μs .. 1.466 μs)
                     0.999 R²   (0.999 R² .. 1.000 R²)
mean                 1.457 μs   (1.444 μs .. 1.470 μs)
std dev              42.03 ns   (34.14 ns .. 54.37 ns)
variance introduced by outliers: 38% (moderately inflated)

benchmarking julia/3
time                 13.53 μs   (13.35 μs .. 13.69 μs)
                     0.999 R²   (0.998 R² .. 0.999 R²)
mean                 13.42 μs   (13.26 μs .. 13.58 μs)
std dev              550.8 ns   (445.3 ns .. 768.8 ns)
variance introduced by outliers: 50% (moderately inflated)

I think the main reason why iter2 is better is because it avoids allocating Complex values. In any case it does a lot fewer allocations.

Code follows:

import Data.List
import Data.Ratio
import Data.Complex
import Criterion
import Criterion.Main

iter1 :: Double -> Double -> Complex Double -> Int
iter1 a b = genericLength . take 255 . takeWhile ((<=2) . magnitude) . iterate f
  where
   f z = z * z - (a :+ b)

iter2 :: Double -> Double -> Complex Double -> Int
iter2 a b (x :+ y) =  loop 0 x y
  where
    loop n x y
      | n > 255 || x*x + y*y > 4 = n
    loop n x y = loop (n+1)  (x*x-y*y-a) (2*x*y-b)

iter3 :: Double -> Double -> Complex Double -> Int
iter3 a b z = fst $ until (\(n,z) -> magnitude z > 2 || n == 255)
                        (\(n,z) -> (n+1,f z))
                        (0,z)
  where
   f z = z * z - (a :+ b)

-- should give 101 iterations
z0 = 7.396147400000001e-3 :+ 0.6418972592000001

main = defaultMain [
  bgroup "julia"
    [
      bench "3"  $ whnf (iter3 (negate 0.40) 0.65) z0
    , bench "2"  $ whnf (iter2 (negate 0.40) 0.65) z0
    , bench "1"  $ whnf (iter1 (negate 0.40) 0.65) z0
    ]
  ]