def sample_ray_dir(indir, normal, hit_pos, mat):
u = ti.Vector([0.0, 0.0, 0.0])
pdf = 1.0
if mat == mat_lambertian:
u = sample_brdf(normal)
pdf = max(eps, compute_brdf_pdf(normal, u))
elif mat == mat_specular:
u = reflect(indir, normal)
elif mat == mat_glass:
cos = indir.dot(normal)
ni_over_nt = refr_idx
outn = normal
if cos > 0.0:
outn = -normal
cos = refr_idx * cos
else:
ni_over_nt = 1.0 / refr_idx
cos = -cos
has_refr, refr_dir = refract(indir, outn, ni_over_nt)
refl_prob = 1.0
if has_refr:
refl_prob = schlick(cos, refr_idx)
if ti.random() < refl_prob:
u = reflect(indir, normal)
else:
u = refr_dir
return u.normalized(), pdf
def out_dir(indir, n, mat):
u = ti.Vector([1.0, 0.0, 0.0])
if mat == mat_lambertian:
if abs(n[1]) < 1 - eps:
u = ti.normalized(ti.cross(n, ti.Vector([0.0, 1.0, 0.0])))
v = ti.cross(n, u)
phi = 2 * math.pi * ti.random()
ay = ti.sqrt(ti.random())
ax = ti.sqrt(1 - ay**2)
u = ax * (ti.cos(phi) * u + ti.sin(phi) * v) + ay * n
elif mat == mat_metal:
u = reflect(indir, n)
else:
# glass
cos = ti.dot(indir, n)
ni_over_nt = refr_idx
outn = n
if cos > 0.0:
outn = -n
cos = refr_idx * cos
else:
ni_over_nt = 1.0 / refr_idx
cos = -cos
has_refr, refr_dir = refract(indir, outn, ni_over_nt)
refl_prob = 1.0
if has_refr:
refl_prob = schlick(cos, refr_idx)
if ti.random() < refl_prob:
u = reflect(indir, n)
else:
u = refr_dir
return ti.normalized(u)
def sample(p):
a = ti.random(ti.f32) * tau
d = vec2(ti.cos(a), ti.sin(a))
ret = 0.0
depth = 0
steps = 0
sign = 1.0
f = sdf(p)
while depth < 5 and steps < 1e3:
if f[0] < 0.0:
sign = -1.0
steps += 1
f = sdf(p)
p += d * sign * f[0]
if sign * f[0] < 1e-6:
ret += f[1] # EMI
if random_in(f[2]): # RFL
depth += 1
n = sign * gradient(p)
d = reflect(d, n)
p += n * 1e-3
elif random_in(f[3]): # RFR
depth += 1
n = sign * gradient(p)
eta = 1.50
if sign > 0.0:
eta = 1 / eta
has, d = refract(d, n, eta)
if not has:
d = reflect(d, n)
else:
p += -n * 1e-3
else:
break
elif abs(f[0]) > 1e1:
break
if f[0] < 0:
sign = -1.0
else:
sign = 1.0
return ret