微妙にバグる ほぼhttp://github.com/githole/eduptの丸写し

peryapt.hs

import Debug.Trace
import Data.List
import Data.Bits

--main = putStrLn (imageToPpm (render 640 480 4 2))
main = putStrLn (imageToPpm (render 320 240 8 4))
-- it takes about one hour in my laptop

-- constants

inf = 1e128
eps = 1e-6

depthLowerBound = 5
depthUpperBound = 64

-- vectors

type Vec = (Double, Double, Double)
type Color = Vec
type Image = [[Color]]

infixl 6 %+
infixl 6 %-
infixl 7 %%*
infixl 7 %%/
infixl 7 %*
infixl 7 %<*>
infixl 7 %.

-- addition
(ax, ay, az) %+ (bx, by, bz) = (ax + bx, ay + by, az + bz)
-- subtraction
(ax, ay, az) %- (bx, by, bz) = (ax - bx, ay - by, az - bz)
-- scalar product
k %%* (x, y, z) = (k * x, k * y, k * z)
(x, y, z) %%/ k = (x / k, y / k, z / k)
-- interior product
(ax, ay, az) %. (bx, by, bz) = ax * bx + ay * by + az * bz
-- exterior product 
(ax, ay, az) %* (bx, by, bz) =
  (ay * bz - by * az,
   az * bx - bz * ax,
   ax * by - bx * ay)
-- element product
(ax, ay, az) %<*> (bx, by, bz) = (ax * bx, ay * by, az * bz)

-- generate orthonormal basis whose one component is the argument
genBasis :: Vec -> (Vec, Vec, Vec)
genBasis w@(wx, _, _) = (w, u, w %* u)
  where u = normalize $ (if wx > eps then (0.0, 1.0, 0.0) else (1.0, 0.0, 0.0)) %* w

-- make the normal vector oriented
orient :: Vec -> Vec -> Vec
orient normal dir = (if normal %. dir < 0.0 then 1.0 else (-1.0)) %%* normal

-- normalize the vector
normalize :: Vec -> Vec
normalize (x, y, z) = (x / k, y / k, z / k)
  where k = sqrt (x ** 2 + y ** 2 + z ** 2)

-- take the average of the vectors
average :: [Vec] -> Vec
average xs = foldl1' (\prev cur -> prev %+ cur) xs %%/ fromIntegral (length xs)

-- objects 

data Ray = Ray { rayOrg :: Vec, dir :: Vec }

data HitPoint = HitPoint { dist :: Double, normal :: Vec, hpPos :: Vec }
  deriving (Show, Eq)
data Intersection = Intersection HitPoint Sphere
  deriving (Show, Eq)

data ReflType = ReflTypeDiffuse | ReflTypeSpecular | ReflTypeRefraction
  deriving (Show, Eq)
data Sphere = Sphere { rad :: Double, spPos :: Vec, emission :: Color, color :: Color, reflType :: ReflType }
  deriving (Show, Eq)

-- check whether the ray intersects with the sphere
intersectSphere :: Ray -> Sphere -> Maybe HitPoint
intersectSphere (Ray org dir) s@(Sphere rad pos _ _ _)
  | d4 < 0.0             = Nothing
  | t1 < eps && t2 < eps = Nothing
  | otherwise = Just $ HitPoint { dist = dist',
                                  hpPos = hpPos',
                                  normal = normalize $ hpPos' %- pos }

  where hpPos' = org %+ (dist' %%* dir)
        po = pos %- org
        b  = po %. dir
        d4 = b ** 2 - po %. po + rad ** 2

        t1 = b - sqrt d4
        t2 = b + sqrt d4
        dist' = max t1 t2

-- scene data

bgColor = (0.0, 0.0, 0.0)

defaultScene =
  [Sphere 1e5  (1e5 + 1  , 40.8       , 81.6      ) (0.0 , 0.0 , 0.0 ) (0.75, 0.25, 0.25) ReflTypeDiffuse,
   Sphere 1e5  (-1e5 + 99, 40.8       , 81.6      ) (0.0 , 0.0 , 0.0 ) (0.25, 0.25, 0.75) ReflTypeDiffuse,
   Sphere 1e5  (50.0     , 40.8       , 1e5       ) (0.0 , 0.0 , 0.0 ) (0.75, 0.75, 0.75) ReflTypeDiffuse,
   Sphere 1e5  (50.0     , 40.8       , -1e5 + 250) (0.0 , 0.0 , 0.0 ) (0.0 , 0.0 , 0.0 ) ReflTypeDiffuse,
   Sphere 1e5  (50.0     , 1e5        , 81.6      ) (0.0 , 0.0 , 0.0 ) (0.75, 0.75, 0.75) ReflTypeDiffuse,
   Sphere 1e5  (50.0     , -1e5 + 81.6, 81.6      ) (0.0 , 0.0 , 0.0 ) (0.75, 0.75, 0.75) ReflTypeDiffuse,
   Sphere 20   (65.0     , 20.0       , 20.0      ) (0.0 , 0.0 , 0.0 ) (0.25, 0.75, 0.25) ReflTypeDiffuse,
   Sphere 16.5 (27.0     , 16.5       , 47.0      ) (0.0 , 0.0 , 0.0 ) (0.99, 0.99, 0.99) ReflTypeSpecular,
   Sphere 16.5 (77.0     , 16.5       , 78.0      ) (0.0 , 0.0 , 0.0 ) (0.99, 0.99, 0.99) ReflTypeRefraction,
   Sphere 15.0 (50.0     , 90.0       , 81.6      ) (36.0, 36.0, 36.0) (0.0 , 0.0 , 0.0 ) ReflTypeDiffuse]

-- check whether there's a sphere which intersects with the ray
intersectScene :: Ray -> Maybe Intersection
intersectScene ray = fst (foldl' check (Nothing, inf) defaultScene)
  where check prev@(res, dist) curSphere =
          case (ray `intersectSphere` curSphere) of
            Just (curHP@(HitPoint curDist _ _))
              | curDist < dist -> (Just (Intersection curHP curSphere), curDist)
            _ -> prev

-- 玉が漆黒になるの、必ずradianceかそっから呼び出されてる関数に原因がある
-- 上下逆になるのの原因はrenderとどっちかは分からない

-- radiance function
radiance :: Ray -> [Double] -> Int -> (Color, [Double])
radiance ray (rnd:rnds) depth =
  case intersectScene ray of
    Nothing -> (bgColor, rnds)
    Just (inter@(Intersection _ object)) ->
      if rnd >= termProb then
        (emission object, rnds)
      else
        radiance' (reflType object) inter rnds termProb

      where termProb = getTermProb (color object)

            getTermProb :: Color -> Double
            getTermProb (x, y, z)
              | depth < depthLowerBound = 1.0
              | depth > depthUpperBound = xyzmax * (0.5 ** fromIntegral (depth - depthUpperBound))
              | otherwise = xyzmax
                where xyzmax = max x $ max y z

  where
    radiance' :: ReflType -> Intersection -> [Double] -> Double -> (Color, [Double])

    radiance' ReflTypeDiffuse (Intersection point object) (rnd:rnd':rnds) termProb =
      ((emission object %+ (weight %<*> inRadience)), rnds')
      -- 自然発光していない球に外側からレイを当てると黒になる問題

      where (w, u, v) = genBasis $ orient (normal point) (dir ray)

            r1 = 2 * pi * rnd
            r2 = rnd'

            dir' = normalize ((cos r1 * sqrt r2) %%* u %+
                              (sin r1 * sqrt r2) %%* v %+
                              (sqrt (1.0 - r2)) %%* w)

            (inRadience, rnds') = radiance (Ray (hpPos point) dir') rnds (depth + 1)
            weight = color object %%/ termProb

    radiance' ReflTypeSpecular (Intersection point object) rnds termProb =
      ((emission object %+ weight %<*> inRadience), rnds')

      where (inRadience, rnds') =
              radiance (Ray (hpPos point)
                            (dir ray %- 2.0 * (normal point %. dir ray) %%* normal point)) rnds (depth + 1)
            weight = color object %%/ termProb

    -- temporary
    radiance' ReflTypeRefraction j k l = radiance' ReflTypeDiffuse j k l


-- renderers

render :: Int -> Int -> Int -> Int -> Image
render width height sample supersample =
  do y <- [0 .. height - 1]
     return $ trace ("rendering y = " ++ show y) $ fst $ 
                foldr
                  (\x (prev, rnds) ->
                        let (cur, rnds') = renderSubpixels x y rnds in (cur : prev, rnds'))
                  ([], xorShiftRands y)
                  [0 .. width - 1]
  where
    subpixels = [(sx, sy) | sx <- [0 .. supersample - 1],
                            sy <- [0 .. supersample - 1],
                            _  <- [0 .. sample - 1]]

    renderSubpixels x y rnds =
      let (res, resRnds) = foldl' renderSubpixel ((0.0, 0.0, 0.0), rnds) subpixels
      in (res %%/ fromIntegral (length subpixels), resRnds)

      where

        renderSubpixel (prev, rnds) (sx, sy) =
          let (cur, rnds') = radiance (Ray camPos dir) rnds 0 in (prev %+ cur, rnds')

          where 
            rate = 1.0 / fromIntegral supersample
            r1 = fromIntegral sx * rate + rate / 2.0
            r2 = fromIntegral sy * rate + rate / 2.0

            scrPos = scrCenter %+
                     ((r1 + fromIntegral x) / fromIntegral width - 0.5) %%* scrX %+
                     ((r2 + fromIntegral y) / fromIntegral height - 0.5) %%* scrY

            dir = normalize $ scrPos %- camPos

    camPos = (50.0, 52.0, 220.0)
    camDir = normalize (0.0, -0.04, -1.0)
    camUp = (0.0, 1.0, 0.0)

    scrWidth = 30.0 * fromIntegral width / fromIntegral height
    scrHeight = 30.0
    scrDist = 40.0

    scrX = scrWidth %%* (normalize (camDir %* camUp))
    scrY = scrHeight %%* (normalize (scrX %* camDir))
    scrCenter = camPos %+ scrDist %%* camDir

--xorShiftRands :: Int -> Int -> Int -> Int -> [Double]
--xorShiftRands a b c d = xorShiftRands' 123456789 362436069 521288629 ((((a * p) + b) * p + c) * p + d)
--  where p = 100000007
xorShiftRands :: Int -> [Double]
xorShiftRands seed = xorShiftRands' 123456789 362436069 521288629 seed
--xorShiftRands a b c d = randomRs (0.0, 1.0) (mkStdGen ((((a * p) + b) * p + c) * p + d))

xorShiftRands' :: Int -> Int -> Int -> Int -> [Double]
xorShiftRands' x y z w = 
  let t = x `xor` (x `shiftL` 11)
      x' = y
      y' = z
      z' = w
      w' = (w `xor` (w `shiftR` 19)) `xor` (t `xor` (t `shiftR` 8))
  in (fromIntegral w' / fromIntegral (maxBound :: Int)) : xorShiftRands' x' y' z' w'

-- ppm exporters 

clamp :: Double -> Double
clamp = max 0.0 . min 1.0

hdrToLdr :: Double -> Int
hdrToLdr x = truncate (clamp x ** (1.0 / 2.2) * 255)

imageToPpm :: Image -> String
imageToPpm image = header ++ body
  where header = "P3\n" ++ show (length (head image)) ++ " " ++ show (length image) ++ "\n255\n"
        body = concat (map (\xs -> (concat (map pixel xs))) image)

        pixel (x, y, z) = show (hdrToLdr x) ++ " " ++ show (hdrToLdr y) ++ " " ++ show (hdrToLdr z) ++ " "