-
Notifications
You must be signed in to change notification settings - Fork 20
/
material.py
126 lines (102 loc) · 3.86 KB
/
material.py
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
import taichi as ti
from vector import *
@ti.func
def reflectance(cosine, idx):
r0 = ((1.0 - idx) / (1.0 + idx))**2
return r0 + (1.0 - r0) * ((1.0 - cosine)**5)
@ti.func
def reflect(v, n):
return v - 2.0 * v.dot(n) * n
@ti.func
def refract(v, n, etai_over_etat):
cos_theta = min(-v.dot(n), 1.0)
r_out_perp = etai_over_etat * (v + cos_theta * n)
r_out_parallel = -ti.sqrt(abs(1.0 - r_out_perp.norm_sqr())) * n
return r_out_perp + r_out_parallel
class _material:
def scatter(self, in_direction, p, n):
pass
class Lambert(_material):
def __init__(self, color):
self.color = color
self.index = 0
self.roughness = 0.0
self.ior = 1.0
@staticmethod
@ti.func
def scatter(in_direction, p, n, color):
out_direction = n + random_in_hemisphere(n)
attenuation = color
return True, p, out_direction, attenuation
class Metal(_material):
def __init__(self, color, roughness):
self.color = color
self.index = 1
self.roughness = min(roughness, 1.0)
self.ior = 1.0
@staticmethod
@ti.func
def scatter(in_direction, p, n, color, roughness):
out_direction = reflect(in_direction.normalized(),
n) + roughness * random_in_unit_sphere()
attenuation = color
reflected = out_direction.dot(n) > 0.0
return reflected, p, out_direction, attenuation
class Dielectric(_material):
def __init__(self, ior):
self.color = Color(1.0, 1.0, 1.0)
self.index = 2
self.roughness = 0.0
self.ior = ior
@staticmethod
@ti.func
def scatter(in_direction, p, n, color, ior, front_facing):
refraction_ratio = 1.0 / ior if front_facing else ior
unit_dir = in_direction.normalized()
cos_theta = min(-unit_dir.dot(n), 1.0)
sin_theta = ti.sqrt(1.0 - cos_theta * cos_theta)
out_direction = Vector(0.0, 0.0, 0.0)
cannot_refract = refraction_ratio * sin_theta > 1.0
if cannot_refract or reflectance(cos_theta,
refraction_ratio) > ti.random():
out_direction = reflect(unit_dir, n)
else:
out_direction = refract(unit_dir, n, refraction_ratio)
attenuation = color
return True, p, out_direction, attenuation
@ti.data_oriented
class Materials:
''' List of materials for a scene.'''
def __init__(self, n):
self.roughness = ti.field(ti.f32)
self.colors = ti.Vector.field(3, dtype=ti.f32)
self.mat_index = ti.field(ti.u32)
self.ior = ti.field(ti.f32)
ti.root.dense(ti.i, n).place(self.roughness, self.colors,
self.mat_index, self.ior)
def set(self, i, material):
self.colors[i] = material.color
self.mat_index[i] = material.index
self.roughness[i] = material.roughness
self.ior[i] = material.ior
@ti.func
def scatter(self, i, ray_direction, p, n, front_facing):
''' Get the scattered ray that hits a material '''
mat_index = self.mat_index[i]
color = self.colors[i]
roughness = self.roughness[i]
ior = self.ior[i]
reflected = True
out_origin = Point(0.0, 0.0, 0.0)
out_direction = Vector(0.0, 0.0, 0.0)
attenuation = Color(0.0, 0.0, 0.0)
if mat_index == 0:
reflected, out_origin, out_direction, attenuation = Lambert.scatter(
ray_direction, p, n, color)
elif mat_index == 1:
reflected, out_origin, out_direction, attenuation = Metal.scatter(
ray_direction, p, n, color, roughness)
else:
reflected, out_origin, out_direction, attenuation = Dielectric.scatter(
ray_direction, p, n, color, ior, front_facing)
return reflected, out_origin, out_direction, attenuation