Ulam spirals

Colour.hs

module Colour where
 
data Colour = Colour {redPart, greenPart, bluePart :: Double} deriving (Eq, Show)
 
cmap :: (Double -> Double) -> Colour -> Colour
cmap f (Colour r g b) = Colour (f r) (f g) (f b)
 
czip :: (Double -> Double -> Double) -> Colour -> Colour -> Colour
czip f (Colour r1 g1 b1) (Colour r2 g2 b2) = Colour (f r1 r2) (f g1 g2) (f b1 b2)
 
cfold :: (Double -> Double -> Double) -> Colour -> Double
cfold f (Colour r g b) = r `f` g `f` b
 
cpromote :: Double -> Colour
cpromote x = Colour x x x
 
instance Num Colour where
  (+) = czip (+)
  (-) = czip (-)
  (*) = czip (*)
  negate = cmap negate
  abs    = cmap abs
  signum = cmap signum
  fromInteger x = cpromote (fromInteger x)
 
instance Fractional Colour where
  (/) = czip (/)
  recip = cmap recip
  fromRational x = cpromote (fromRational x)
 
clip :: (Num n, Ord n) => n -> n
clip n
  | n < 0 = 0
  | n > 1 = 1
  | otherwise = n
 
cclip :: Colour -> Colour
cclip = cmap clip



black = Colour 0 0 0
blue = Colour 0 0 1
green =Colour  0 0.5 0 
cyan = Colour 0 0.54296875 0.54296875
red = Colour 0.545098 0 0
magenta = Colour 0.54296875 0  0.54296875
yellow = Colour 1 1 0
white = Colour 1 1 1 

PPM6.hs

module PPM6 (Point,Image,module Colour,make_ppm, save_ppm, mapPixel, mapDouble) where
 
import Data.Word
import qualified Data.ByteString as BIN
import Colour
import Data.Array.IO
import System.IO(openBinaryFile,hClose,IOMode(..))
import Control.Monad(when)

type Point = (Double, Double)
type Image = Point -> Colour

----------------------------------------------------------
-- dealing with Word8 entities
 
quant8 :: RealFrac n => n -> Word8
quant8 x = floor $ x * 0xFF
 
cquant8 :: Colour -> [Word8]
cquant8 (Colour r g b) = [quant8 r, quant8 g, quant8 b]


----------------------------------------------------------
-- List based bit-mapped interface

string_to_bin :: String -> BIN.ByteString
string_to_bin = BIN.pack . map (fromIntegral . fromEnum)
 
header :: [[Colour]] -> BIN.ByteString
header pss =
  let nx = length $ head pss
      ny = length        pss
  in  string_to_bin $ "P6\n" ++ show nx ++ " " ++ show ny ++ " 255\n"
 
body :: [[Colour]] -> BIN.ByteString
body pss = BIN.pack $ concatMap (cquant8 . cclip) $ concat pss
 
make_ppm :: [[Colour]] -> BIN.ByteString
make_ppm pss = BIN.append (header pss) (body pss)

 
save_ppm :: FilePath -> [[Colour]] -> IO ()
save_ppm f pss = BIN.writeFile f (make_ppm pss)

-----------------------------------------------------------------
-- Pixel function based bit-mapped interface

storeList :: (a -> Word8) -> [a] -> IOUArray Int Word8 -> Int -> IO Int
storeList inject s arr next = do { mapM put (zip [next .. count - 1] s); return count }
  where put (i,c) = writeArray arr i (inject c)
        count = next + length s
        
charToWord8 c = (fromIntegral(fromEnum c)) 


storePixel:: (Int -> Int -> a -> IO a) -> Int -> Int -> a -> IO a
storePixel perform rowSize colSize next = go 1 1 next
  where go i j next
           | i > rowSize = return next
           | j > colSize = go (i+1) 1 next
           | otherwise = do { n <- perform i j next; go i (j+1) n }
           
new:: Int -> IO (IOUArray Int Word8)
new n  = newArray (0,n) (fromIntegral 0)

make :: (Int -> Int -> Colour) -> Int -> Int -> IO (IOUArray Int Word8,Int)
make f nx ny = 
  do { let prefix = "P6\n" ++ show nx ++ " " ++ show ny ++ " 255\n"
           size = length prefix + 3*(nx * ny)
     ; arr <- new size
     ; next0 <- storeList charToWord8 prefix arr 0
     ; let perform i j next = storeList id (cquant8(cclip(f i j))) arr next
     ; next1 <- storePixel perform nx ny next0
     ; return (arr,next1)
     }

mapPixel file f nx ny = 
  do { (arr,size) <- make f nx ny 
     ; handle <- openBinaryFile file WriteMode
     ; hPutArray handle arr size
     ; hClose handle
     ; return ()
     }
     

-----------------------------------------------------------------
-- Real function based bit-mapped interface

-- Lift a real funtion of color to an IO action that writes that color to an Array

lift :: IOUArray Int Word8 -> (Double -> Double -> Colour) -> Double -> Double -> Int -> IO Int
lift arr f x y next = 
    case cclip(f x y) of
      (Colour r g b) -> do writeArray arr next     (quant8 r)
                           writeArray arr (next+1) (quant8 g)
                           writeArray arr (next+2) (quant8 b)
                           return (next +3)
 
-- Iterate over all the real valued coordinates, in the order 
-- that corresponds to the bit-map order.

iterDouble:: (Ord n,Num n) => (n -> n -> a -> IO a) -> (n,n,Int) -> (n,n,Int) -> a -> IO a
iterDouble perform 
           x@(xStart,xDelta,xCount) 
           y@(yStart,yDelta,yCount) 
           next = go xStart 1 yStart 1 next
  where go x i y j next
           | i > xCount =  go xStart 1 (y+yDelta) (j+1) next
           | j > yCount = return next
           | otherwise = do { n <- perform x y next; go (x+xDelta) (i+1) y j n }  


-- do we generate the points in the correct order?
testIterDouble = 
       do a <- iterDouble perf (-2,0.5,5) (5,-0.5,9) []
          putStrLn (show(length a))
          print a
   where perf x y ans = return(ans ++ [(x,y)])
           
        
-- Build an array of colors, open a file, and write the array to bit-mapped graphics file
mapDouble
  :: FilePath   -- Image file to write output to
  -> Image      -- Image function
  -> Point      -- Lower left corner of image
  -> Point      -- Upper right corner of image
  -> (Int, Int) -- Number of pixels in x and y directions
  -> IO ()
mapDouble file perform (lowerleft@(x1,y1)) (upperright@(x2,y2)) (nx,ny) = 
   do { let prefix = "P6\n" ++ show nx ++ " " ++ show ny ++ " 255\n"
            size = length prefix + 3*(nx * ny)
      ; when (x1>=x2 || y1>=y2) 
           (error ("In a call to mapDouble, the Lower left corner "++show lowerleft++"\nis not below and to the left of the upper right corner "++show upperright))
     ; putStrLn "Creating prefix header"
     ; arr <- new size
        -- Store the PPM style 6 prefix
      ; next0 <- storeList charToWord8 prefix arr 0 
      ; putStrLn "Mapping the function"
      ; iterDouble (lift arr (curry perform')) (x1,rowDelta,nx) (y2,colDelta,ny) next0
      ; handle <- openBinaryFile file WriteMode
      ; putStrLn $ "Writing the ppm file: "++file
      ; hPutArray handle arr size
      ; hClose handle
      ; return () }
  where rowDelta = (x2 - x1) / fromIntegral nx
        colDelta = (y1 - y2) / fromIntegral ny
        perform' = perform
        -- perform' = antiAlias' 0 0 rowDelta colDelta perform

-- Using a type class here is probably overkill ...
class Average t where
  average :: t -> t -> t
instance Average Double where
  average d1 d2 = (d1 + d2) / 2
instance Average Colour where
  average c1 c2 = cZipWith average c1 c2
   where
    -- This is 'Colour.czip'
    cZipWith f (Colour r1 g1 b1) (Colour r2 g2 b2) =
      Colour (r1 `f` r2) (g1 `f` g2) (b1 `f` b2)

-- Anti-alias using 'n' extra samples.
--
-- So, no anti-aliasing when 'n == 0'.
--
-- ??? Here it seems 'n == 1' works best ???
antiAlias :: Int -> Double -> Double -> Image -> Image
antiAlias n dx dy f (x,y) = averageAll $ map f grid
 where
  averageAll = foldl1 average
  grid       = points n dx dy (x,y)

-- Like 'antiAlias', but also sample 'points m' points in each
-- neighboring pixel.
antiAlias' :: Int -> Int -> Double -> Double -> Image -> Image
antiAlias' m n dx dy f (x,y) = averageAll $ map f grid
 where
  averageAll = foldl1 average
  grid       = points n dx dy (x,y) 
               ++ concatMap (shift $ points m dx dy (x,y)) shifts
  shifts     = [ (0, dy)   -- North
               , (0,-dy)   -- South
               , ( dx,0)   -- East
               , (-dx,0) ] -- West

points :: Int -> Double -> Double -> Point -> [Point]
-- E.g., for n = 1, we sample the points marked X:
--
--         |<---   dx  --->|
--       
--         +-------+-------+  ---
--         |       |       |   ^
--         |   X   |   X   |   |
--         |       |       |
--   ---   +-------+-------+   dy
--    ^    |       |       |
--   ddy   |   X   |   X   |   |
--    v    |       |       |   v
--   ---   +-------+-------+  ---
--       
--         |< ddx >|
--
-- where the lower left corner is '(x,y)'.
points n dx dy (x,y) =
  [ ( x + (i + 0.5)*ddx
    , y + (j + 0.5)*ddy )
  | i <- ns
  , j <- ns ]
 where
  -- Both endpoints are *included*, e.g. for 'n == 2' get '[0,1,2]'.
  ns  = map fromIntegral [0..n]
  ddx = dx / fromIntegral (n+1)
  ddy = dy / fromIntegral (n+1)

shift :: [Point] -> Point -> [Point]
shift ps (dx,dy) = [ (x+dx,y+dy) | (x,y) <- ps ]

--------------------------------
-- tests  


f 2 3 = white
f 5 1 = black
f _ 1 = red
f _ _ = blue

go = mapPixel "tim.ppm" f 10 10

Ulam.hs

-- Ulam.hs
--
-- This program generates graphics of Ulam spirals using the primes
-- package and the PPM6 library.

module Ulam where

import Data.Numbers.Primes -- requires the primes package from hackage,
                           -- used for efficient lazy prime generation
import PPM6 -- requires PPM6.hs and Colour.hs

-- This is a list whose i-th element is True if i+1 is a prime and False
-- if i+1 is composite
primalities :: [Bool]
primalities = primalities' primes 1
  where
    primalities' remainingPrimes currentNum
      | head remainingPrimes == currentNum = True : primalities'
                                             (tail remainingPrimes)
                                             (currentNum + 1)
      | otherwise = False : primalities' remainingPrimes (currentNum + 1)


-- Hopefully this will give some speedup since list lookup is O(n). I
-- expect that memory usage will not inccrease because of tail-sharing
-- (?)
primalitiesGraded = primalities:[drop (n * n) primalities | n <- [1,3..]]

-- for testing
numbersGraded = [1..]:[drop (n * n) [1..] | n <- [1,3..]]

-- given a Cartesian coordinate pair, gives a shell number and an offset
-- in that shell
ulamIndex :: Int -> Int -> (Int, Int)
ulamIndex x y = let
  grade = max (abs x) (abs y)
  right  = 0
  top    = 2 * grade
  left   = 4 * grade
  bottom = 6 * grade
  offset' -- how far around the shell from the positive x axis, in a
          -- counterclockwise spiral that starts going outwards to the
          -- right
    | grade == y = top - x
    | grade == x = right + y
    | grade == (-y) = bottom + x
    | grade == (-x) = left - y
  offset -- how far around the shell from the "beginning" of the shell
         -- in the lower rightish
    | grade == 0  = 0
    | otherwise   = mod (offset' + grade - 1) (8 * grade)
  in
   (grade, offset)

-- given a Cartesian coordinate pair, return whether or not there is a
-- prime in the Ulam spiral at that point
ulam :: Int -> Int -> Bool
ulam x y = let
  (grade, offset) = ulamIndex x y
  shell = primalitiesGraded !! grade
  in
   shell !! offset

-- An image showing the Ulam spiral with primes in white and composite
-- numbers in black
ulamFunc :: Image
ulamFunc (x, y) = case ulam (floor x) (floor y) of
  True  -> white
  False -> black

-- main
main :: IO ()
main = mapDouble "Ulam.ppm"
                 ulamFunc (-100, -100) (100, 100) (400, 400)