Ticket #2983: RayT.hs

File RayT.hs, 3.3 KB (added by thoughtpolice, 5 years ago)

ray tracer code

Line 
1import System
2infinity = 1/0
3delta = sqrt e where e = encodeFloat (floatRadix e) (-floatDigits e)
4infixl 7 .*, *|
5data Vector = V !Double !Double !Double deriving (Show, Eq)
6s *| V x y z = V (s * x) (s * y) (s * z)
7instance Num Vector where
8    V x y z + V x' y' z' = V (x + x') (y + y') (z + z')
9    V x y z - V x' y' z' = V (x - x') (y - y') (z - z')
10    fromInteger i = V x x x where x = fromInteger i
11V x y z .* V x' y' z' = x * x' + y * y' + z * z'
12vlength r = sqrt (r .* r)
13unitise r = 1 / vlength r *| r
14
15data Scene
16    = Sphere !Vector !Double
17    | Group !Vector !Double Scene Scene Scene Scene Scene
18    deriving (Show)
19
20ray_sphere (V dx dy dz) (V vx vy vz) r =
21  let disc = vx * vx + vy * vy + vz * vz - r * r
22  in  if disc < 0 then infinity else
23      let b = vx * dx + vy * dy + vz * dz
24          b2 = b * b
25      in  if b2 < disc then infinity else
26          let disk = sqrt(b2 - disc)
27              t1 = b - disk
28          in  if t1 > 0 then t1 else b + disk
29
30ray_sphere' (V ox oy oz) (V dx dy dz) (V cx cy cz) r =
31  let vx = cx - ox; vy = cy - oy; vz = cz - oz
32      vv = vx * vx + vy * vy + vz * vz
33      b = vx * dx + vy * dy + vz * dz
34      disc = b * b - vv + r * r
35  in  disc >= 0 && b + sqrt disc >= 0
36
37data Hit = H {l :: !Double, nv :: Vector }
38
39intersect dir@(V dx dy dz) hit s = case s of
40    Sphere center@(V cx cy cz) radius ->
41      let l' = ray_sphere dir center radius in
42      if l' >= l hit then hit else
43        let x = l' * dx - cx
44            y = l' * dy - cy
45            z = l' * dz - cz
46            il = 1 / sqrt(x * x + y * y + z * z)
47        in  H {l = l', nv = V (il * x) (il * y) (il * z) }
48    Group center radius a b c d e ->
49      let l' = ray_sphere dir center radius in
50      if l' >= l hit then hit else
51        let f h s = intersect dir h s in
52        f (f (f (f (f hit a) b) c) d) e
53
54intersect' orig dir s = case s of
55    Sphere center radius -> ray_sphere' orig dir center radius
56    Group center radius a b c d e ->
57      let f s = intersect' orig dir s in
58      ray_sphere' orig dir center radius && (f a || f b || f c || f d || f e)
59
60neg_light = unitise (V 1 3 (-2))
61
62ray_trace dir scene =
63  let hit = intersect dir (H infinity 0) scene in
64  if l hit == infinity then 0 else
65    let n = nv hit in
66    let g = n .* neg_light in
67    if g < 0 then 0 else
68      if intersect' (l hit *| dir + delta *| n) neg_light scene then 0 else g
69
70fold5 f x a b c d e = f (f (f (f (f x a) b) c) d) e
71
72create level c r =
73  let obj = Sphere c r in
74  if level == 1 then obj else
75    let a = 3 * r / sqrt 12 in
76    let bound (c, r) s = case s of
77         Sphere c' r' -> (c, max r (vlength (c - c') + r'))
78         Group _ _ v w x y z -> fold5 bound (c, r) v w x y z in
79    let aux x' z' = create (level - 1 :: Int) (c + V x' a z') (0.5 * r) in
80    let w = aux (-a) (-a); x = aux a (-a) in
81    let y = aux (-a) a; z = aux a a in
82    let (c1, r1) = fold5 bound (c + V 0 r 0, 0) obj w x y z in
83    Group c1 r1 obj w x y z
84
85ss = 4
86pixel_vals n scene y x = sum
87  [ let f a da = a - n / 2 + da / ss; d = unitise (V (f x dx) (f y dy) n)
88    in  ray_trace d scene | dx <- [0..ss-1], dy <- [0..ss-1] ]
89main = do 
90    [level,ni] <- fmap (map read) getArgs
91    let n = fromIntegral ni
92        scene = create level (V 0 (-1) 4) 1 
93        scale x = 0.5 + 255 * x / (ss*ss)
94        picture = [ toEnum $ truncate $ scale $ pixel_vals n scene y x | y <- [n-1,n-2..0], x <- [0..n-1]]
95    putStr $ "P5\n" ++ show ni ++ " " ++ show ni ++ "\n255\n" ++ picture