-
Notifications
You must be signed in to change notification settings - Fork 19
/
Raytracing.hs
190 lines (169 loc) · 5.6 KB
/
Raytracing.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE DataKinds #-}
module Raytracing
( Pos, Dir, Camera, mkCamera
, Sphere(..), sphereAABB
, Objs
, Colour, black, white
, render
)
where
import Data.Maybe (fromMaybe)
import BVH
import Image
import Vec3
import Data.Massiv.Array (S)
import Data.Massiv.Array.IO
type Pos = Vec3
type Dir = Vec3
type Colour = Vec3
black, white :: Colour
black = 0
white = 1
data Ray = Ray { rayOrigin :: Pos
, rayDir :: Dir }
deriving (Show)
pointAtParam :: Ray -> Float -> Pos
pointAtParam ray t = rayOrigin ray + (t `scale` rayDir ray)
data Hit = Hit { hitT :: Float
, hitP :: Pos
, hitNormal :: Dir
, hitColour :: Colour
}
data Sphere = Sphere { spherePos :: Pos
, sphereColour :: Colour
, sphereRadius :: Float }
deriving (Show, Read)
sphereAABB :: Sphere -> AABB
sphereAABB (Sphere centre _ radius) =
AABB
(centre - Vec3 radius radius radius)
(centre + Vec3 radius radius radius)
sphereHit :: Sphere -> Ray -> Float -> Float -> Maybe Hit
sphereHit (Sphere center colour radius) r t_min t_max =
let oc = rayOrigin r - center
a = dot (rayDir r) (rayDir r)
b = dot oc (rayDir r)
c = dot oc oc - radius*radius
discriminant = b*b - a*c
sqrtDisc = sqrt discriminant
tryHit temp =
if temp < t_max && temp > t_min
then Just $ Hit
{ hitT = temp
, hitP = pointAtParam r temp
, hitNormal = (1/radius) `scale` (pointAtParam r temp - center)
, hitColour = colour
}
else Nothing
in if discriminant <= 0
then Nothing
else case tryHit ((-b - sqrtDisc)/a) of
Just hit -> Just hit
Nothing -> tryHit ((-b + sqrtDisc)/a)
type Objs = BVH Sphere
aabbHit :: AABB -> Ray -> Float -> Float -> Bool
aabbHit aabb (Ray origin direction) tmin0 tmax0 =
let iter min' max' origin' dir' tmin' tmax' =
let invD = 1 / dir'
t0 = (min' - origin') * invD
t1 = (max' - origin') * invD
(t0', t1') = if invD < 0 then (t1, t0) else (t0, t1)
tmin'' = max t0' tmin'
tmax'' = min t1' tmax'
in (tmin'', tmax'')
(tmin1, tmax1) =
iter
(vecX (aabbMin aabb)) (vecX (aabbMax aabb))
(vecX origin) (vecX direction)
tmin0 tmax0
in not $
tmax1 <= tmin1 ||
let (tmin2, tmax2) =
iter (vecY (aabbMin aabb)) (vecY (aabbMax aabb))
(vecY origin) (vecY direction)
tmin1 tmax1
in tmax2 <= tmin2 ||
let (tmin3, tmax3) =
iter (vecZ (aabbMin aabb)) (vecZ (aabbMax aabb))
(vecZ origin) (vecZ direction)
tmin2 tmax2
in tmax3 <= tmin3
objsHit :: Objs -> Ray -> Float -> Float -> Maybe Hit
objsHit (BVHLeaf _ sphere) r t_min t_max =
sphereHit sphere r t_min t_max
objsHit (BVHSplit box left right) r t_min t_max
| not $ aabbHit box r t_min t_max =
Nothing
| otherwise =
case objsHit left r t_min t_max of
Nothing -> objsHit right r t_min t_max
Just h1 ->
Just $ fromMaybe h1 $ objsHit right r t_min (hitT h1)
data Camera = Camera { camOrigin :: Pos
, camLLC :: Pos
, camHorizontal :: Dir
, camVertical :: Dir
}
deriving (Show, Read)
mkCamera :: Pos -> Pos -> Dir -> Float -> Float -> Camera
mkCamera lookfrom lookat vup vfov aspect =
let theta = vfov * pi / 180
half_height = tan (theta / 2)
half_width = aspect * half_height
origin = lookfrom
w = normalise (lookfrom - lookat)
u = normalise (cross vup w)
v = cross w u
in Camera { camOrigin = lookfrom
, camLLC = origin -
(half_width `scale` u) -
(half_height `scale` v) -
w
, camHorizontal = (2*half_width) `scale` u
, camVertical = (2*half_height) `scale` v
}
getRay :: Camera -> Float -> Float -> Ray
getRay cam s t =
Ray
(camOrigin cam)
(camLLC cam +
(s `scale` camHorizontal cam) +
(t `scale` camVertical cam) -
camOrigin cam)
reflect :: Vec3 -> Vec3 -> Vec3
reflect v n = v - (2 * dot v n) `scale` n
scatter :: Ray -> Hit -> Maybe (Ray, Vec3)
scatter r hit =
let reflected =
reflect (normalise (rayDir r)) (hitNormal hit)
scattered = Ray (hitP hit) reflected
in if dot (rayDir scattered) (hitNormal hit) > 0
then Just (scattered, hitColour hit)
else Nothing
rayColour :: Objs -> Ray -> Int -> Colour
rayColour objs r depth =
case objsHit objs r 0.001 (1/0) of
Just hit ->
case scatter r hit of
Just (scattered, attenuation)
| depth < 50 ->
attenuation * (rayColour objs scattered (depth+1))
_ -> black
Nothing ->
let unit_dir = normalise (rayDir r)
t = 0.5 * (vecY unit_dir + 1)
in ((1-t) `scale` Vec3 1 1 1) + (t `scale` Vec3 0.5 0.7 1)
traceRay :: Objs -> Int -> Int -> Camera -> Int -> Int -> Colour
traceRay objs width height cam =
\j i -> let u = fromIntegral i / fromIntegral width
v = fromIntegral j / fromIntegral height
ray = getRay cam u v
in rayColour objs ray 0
colourToPixel :: Colour -> Pixel (SRGB 'NonLinear) Word8
colourToPixel (Vec3 r g b) = toPixel8 $ PixelSRGB r g b
render :: Objs -> Int -> Int -> Camera -> Image S (SRGB 'NonLinear) Word8
render objs height width cam =
mkImage height width $ \j i ->
colourToPixel $ traceRay objs width height cam j i