def render_particle(model, camera, index):
scene = model.scene
a = (model.L2C[None] @ ts.vec4(model.pos[index], 1)).xyz
r = model.radius[index]
A = camera.uncook(a)
rad = camera.uncook(ts.vec3(r, r, a.z), False)
M = int(ti.floor(A - rad))
N = int(ti.ceil(A + rad))
M = ts.clamp(M, 0, ti.Vector(camera.res))
N = ts.clamp(N, 0, ti.Vector(camera.res))
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
pos = camera.cook(float(ts.vec3(X, a.z)))
dp = pos - a
dp2 = dp.norm_sqr()
if dp2 > r**2:
continue
dz = ti.sqrt(r**2 - dp2)
if camera.fb.atomic_depth(X, a.z - dz):
continue
n = ts.vec3(dp.xy, -dz)
normal = ts.normalize(n)
view = ts.normalize(a + n)
color = model.colorize(pos, normal)
camera.fb['img'][X] = color
camera.fb['normal'][X] = normal
def render_floor(self, camera):
scene = camera.scene
c_floor = ts.vec3(*self.floor_color)
c_sky = ts.vec3(*self.sky_color)
for X in ti.grouped(camera.img):
W = ti.cast(ti.Vector([X.x, X.y, 1]), ti.f32)
x = cook_coord(camera, W)
world_dir = camera.trans_dir(ts.normalize(x))
if world_dir[1] > 0 or camera.pos[None][1] < floor_y:
camera.img[X] = c_sky
else:
f = min(1, -world_dir[1] * 50)
n = camera.untrans_dir(ts.vec3(0, 1, 0))
pos = camera.pos[None] + world_dir * (
camera.pos[None][1] - floor_y) / abs(world_dir[1])
#print(pos)
pos = camera.untrans_pos(pos)
#print(pos)
color = ambient * c_sky * c_floor
for light in ti.static(scene.lights):
light_color = scene.opt.render_func(pos, n, \
ts.normalize(pos), light, c_floor)
color += light_color
camera.img[X] = c_sky * (1 - f) + color * f
camera.nbuf[X] = n
camera.zbuf[X] = 1 / pos.z
def get_gi_factors(camera, scene, pos, fragpos, normal, radius, color):
a = camera.untrans_pos(pos)
f = form_factor_ball(normal, (pos - fragpos), radius)
cosfactor = 0.0
for light in ti.static(scene.lights):
cosfactor += max(0.5, ti.dot(light.get_dir(a), ts.normalize(pos - a)))
gi = min(f, gi_clamp) / gi_clamp / 5 * color * cosfactor
return f / ff_clamp, gi
def render_cylinder(model, camera, v1, v2, radius, c1, c2):
scene = model.scene
a = camera.untrans_pos(v1)
b = camera.untrans_pos(v2)
A = camera.uncook(a)
B = camera.uncook(b)
bx = int(ti.ceil(camera.fx * radius / min(a.z, b.z)))
by = int(ti.ceil(camera.fy * radius / min(a.z, b.z)))
M, N = ti.floor(min(A, B)), ti.ceil(max(A, B))
M.x -= bx
N.x += bx
M.y -= by
N.y += by
M.x, N.x = min(max(M.x, 0),
camera.img.shape[0]), min(max(N.x, 0), camera.img.shape[1])
M.y, N.y = min(max(M.y, 0),
camera.img.shape[0]), min(max(N.y, 0), camera.img.shape[1])
if (M.x < N.x and M.y < N.y):
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
t = ti.dot(X - A, B - A) / (B - A).norm_sqr()
if t < 0 or t > 1:
continue
proj = a * (1 - t) + b * t
W = ti.cast(ti.Vector([X.x, X.y, proj.z]), ti.f32)
w = cook_coord(camera, W)
dw = w - proj
dw2 = dw.norm_sqr()
if dw2 > radius**2:
continue
dz = ti.sqrt(radius**2 - dw2)
n = ti.Vector([dw.x, dw.y, -dz])
zindex = 1 / (proj.z - dz)
if zindex >= ti.atomic_max(camera.zbuf[X], zindex):
basecolor = c1 if t < 0.5 else c2
normal = ts.normalize(n)
view = ts.normalize(a + n)
color = get_ambient(camera, normal, view) * basecolor
for light in ti.static(scene.lights):
light_color = scene.opt.render_func(a + n, normal, \
view, light, basecolor)
color += light_color
camera.img[X] = color
camera.nbuf[X] = normal
def interact(i, j):
disp = pos[i] - pos[j]
disv = vel[i] - vel[j]
if disp.norm_sqr() < (radius[i] + radius[j])**2 and disp.dot(disv) < 0:
mass_i = radius[i]**3
mass_j = radius[j]**3
disp = ts.normalize(disp)
vel[i], vel[j] = ts.momentumExchange(vel[i], vel[j], disp, mass_i,
mass_j, 0.8)
def ballBoundReflect(pos, vel, center, radius):
ret = vel
above = tl.distance(pos, center) - radius
if above <= 0:
normal = tl.normalize(pos - center)
NoV = tl.dot(vel, normal) - 6 * tl.smoothstep(above, 0, -0.1)
if NoV < 0:
ret -= NoV * normal
return ret
def brdf(self, normal, viewdir, lightdir, color):
halfway = ts.normalize(viewdir + lightdir)
ndotv = max(ti.dot(viewdir, normal), self.eps)
ndotl = max(ti.dot(lightdir, normal), self.eps)
diffuse = self.kd * color / math.pi
specular = self.microfacet(normal, halfway)\
* self.frensel(viewdir, halfway, color)\
* self.geometry(ndotv, ndotl)
specular /= 4 * ndotv * ndotl
return diffuse + specular
def trace(self, orig, dir):
pos = orig
color = ts.vec3(0.0)
normal = ts.vec3(0.0)
for s in range(100):
t = self.calc_sdf(pos)
if t <= 0:
normal = ts.normalize(self.calc_grad(pos) - t)
break
pos += dir * t
return pos, normal
def normal(pos):
eps = 0.005
v1 = ts.vec3(1.0, -1.0, -1.0)
v2 = ts.vec3(-1.0, -1.0, 1.0)
v3 = ts.vec3(-1.0, 1.0, -1.0)
v4 = ts.vec3(1.0, 1.0, 1.0)
return ts.normalize(v1 * scene(pos + v1 * eps) +
v2 * scene(pos + v2 * eps) +
v3 * scene(pos + v3 * eps) +
v4 * scene(pos + v4 * eps))
def fill_H(self, Hx: ti.template(), Hy: ti.template()):
for i, j in self.color_buffer:
x, y = self.ij_to_xy(i, j)
if ti.abs(Hx[x, y]) < 1e-4 and ti.abs(Hy[x, y]) < 1e-4:
self.angle_x[x, y] = 0
self.angle_y[x, y] = 0
else:
H = ts.normalize(ti.Vector([Hx[x, y], Hy[x, y]]))
self.angle_x[x, y] = H.x
self.angle_y[x, y] = H.y
def ballBoundReflect(pos, vel, center, radius, anti_fall=0, anti_depth=0.1):
ret = vel
above = tl.distance(pos, center) - radius
if above <= 0:
normal = tl.normalize(pos - center)
NoV = tl.dot(vel, normal)
if ti.static(anti_fall):
NoV -= anti_fall * tl.smoothstep(above, 0, -anti_depth)
if NoV < 0:
ret -= NoV * normal
return ret
def render(self, renderer):
width = 1
A = self.begin
B = self.end
A, B = min(A, B), max(A, B)
mold = tg.normalize(B - A)
for X in ti.grouped(
ti.ndrange((A.x - width, B.x + width),
(A.y - width, B.y + width))):
udf = abs(tg.cross(X - A, mold))
renderer.image[int(X)] = tg.vec3(tg.smoothstep(udf, width, 0))
def material(pos, camdir):
norm = normal(pos)
d1 = -ts.normalize(ts.vec(5.0, 10.0, -20.0))
d2 = -ts.normalize(ts.vec(-5, 10.0, 20.0))
d3 = -ts.normalize(ts.vec(20, 5.0, -5.0))
d4 = -ts.normalize(ts.vec(-20.0, 5.0, 5.0))
tex = ts.vec(0.2, 0.2, 0.2)
if pos[1] > -5.95:
tex = ts.vec3(0.32, 0.28, 0.0)
sha = 0.7 * softray(pos, -d1, 32.0) + 0.3 * softray(pos, -d4, 16.0)
ao = ambocc(pos, norm)
l1 = light(d1, ts.vec3(1.0, 0.9, 0.8), tex, norm, camdir)
l2 = light(d2, ts.vec3(0.8, 0.7, 0.6), tex, norm, camdir)
l3 = light(d3, ts.vec3(0.3, 0.3, 0.4), tex, norm, camdir)
l4 = light(d4, ts.vec3(0.5, 0.5, 0.5), tex, norm, camdir)
return 0.2 * ao + 0.8 * (l1 + l2 + l3 + l4) * sha
def get_ambient(camera, normal, view):
c_floor = ts.vec3(0.75)
c_sky = ts.vec3(0.8, 0.9, 0.95)
refl_dir = -ts.reflect(-view, normal)
refl_dir = ts.normalize(camera.trans_dir(refl_dir))
color = ts.vec3(ambient)
if refl_dir[1] > 0:
color *= c_sky
else:
f = min(1.0, -refl_dir[1] * 25.0)
color *= c_sky * (1 - f) + c_floor * f
return color
def render_particle(model, camera, vertex, radius, basecolor):
scene = model.scene
a = camera.untrans_pos(vertex)
A = camera.uncook(a)
bx = camera.fx * radius / a.z
by = camera.fy * radius / a.z
M = A
N = A
M.x -= bx
N.x += bx
M.y -= by
N.y += by
M.x, N.x = min(max(M.x, 0),
camera.img.shape[0]), min(max(N.x, 0), camera.img.shape[1])
M.y, N.y = min(max(M.y, 0),
camera.img.shape[0]), min(max(N.y, 0), camera.img.shape[1])
if (M.x < N.x and M.y < N.y):
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
W = ti.cast(ti.Vector([X.x, X.y, a.z]), ti.f32)
w = cook_coord(camera, W)
dw = w - a
dw2 = dw.norm_sqr()
if dw2 > radius**2:
continue
dz = ti.sqrt(radius**2 - dw2)
n = ti.Vector([dw.x, dw.y, -dz])
zindex = 1 / (a.z - dz)
if zindex >= ti.atomic_max(camera.zbuf[X], zindex):
normal = ts.normalize(n)
view = ts.normalize(a + n)
color = get_ambient(camera, normal, view) * basecolor
for light in ti.static(scene.lights):
light_color = scene.opt.render_func(
a + n, normal, view, light, basecolor)
color += light_color
camera.img[X] = color
camera.nbuf[X] = normal
def brdf(self, normal, lightdir, viewdir):
halfway = ts.normalize(lightdir + viewdir)
NoH = max(EPS, ts.dot(normal, halfway))
NoL = max(EPS, ts.dot(normal, lightdir))
NoV = max(EPS, ts.dot(normal, viewdir))
HoV = min(1 - EPS, max(EPS, ts.dot(halfway, viewdir)))
ndf = self.roughness**2 / (NoH**2 * (self.roughness**2 - 1) + 1)**2
vdf = 0.25 / (self.ischlick(NoL) * self.ischlick(NoV))
f0 = self.metallic * self.color + (1 - self.metallic) * self.specular
ks, kd = self.ks * f0, self.kd * (1 - f0) * (1 - self.metallic)
fdf = self.fresnel(f0, NoV)
strength = kd * self.color + ks * fdf * vdf * ndf / math.pi
return strength
def render_particle(model, camera, index):
scene = model.scene
L2W = model.L2W
a = model.pos[index]
r = model.radius[index]
a = camera.untrans_pos(L2W @ a)
A = camera.uncook(a)
rad = camera.uncook(ts.vec3(r, r, a.z), False)
M = int(ti.floor(A - rad))
N = int(ti.ceil(A + rad))
M = ts.clamp(M, 0, ti.Vector(camera.res))
N = ts.clamp(N, 0, ti.Vector(camera.res))
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
pos = camera.cook(float(ts.vec3(X, a.z)))
dp = pos - a
dp2 = dp.norm_sqr()
if dp2 > r**2:
continue
dz = ti.sqrt(r**2 - dp2)
zindex = 1 / (a.z - dz)
if zindex < ti.atomic_max(camera.fb['idepth'][X], zindex):
continue
n = ts.vec3(dp.xy, -dz)
normal = ts.normalize(n)
view = ts.normalize(a + n)
color = ts.vec3(1.0)
color = model.colorize(pos, normal, color)
camera.fb['img'][X] = color
camera.fb['normal'][X] = normal
def do_intersect(self, orig, dir):
op = self.pos - orig
b = op.dot(dir)
det = b ** 2 - op.norm_sqr() + self.radius ** 2
ret = INF
if det > 0.0:
det = ti.sqrt(det)
t = b - det
if t > EPS:
ret = t
else:
t = b + det
if t > EPS:
ret = t
return ret, ts.normalize(dir * ret - op)
def on_advance(self):
for i in self.pos:
acc = ts.vec(0.0, -1.0)
if any(self.iKeyDirection): # ASWD?
acc = self.iKeyDirection
if any(self.iMouseButton):
dir = ts.normalize(self.iMouse - self.pos[i]) * 2
if self.iMouseButton[0]: # LMB pressed?
acc += dir
if self.iMouseButton[1]: # RMB pressed?
acc -= dir
self.vel[i] += acc * self.dt
for i in self.pos:
self.vel[i] = ts.boundReflect(self.pos[i], self.vel[i], 0, 1, 0.8)
for i in self.pos:
self.pos[i] += self.vel[i] * self.dt
def generate(self, coor):
orig = ts.vec3(0.0)
dir = ts.vec3(0.0, 0.0, 1.0)
if ti.static(self.type == self.ORTHO):
orig = ts.vec3(coor, 0.0)
elif ti.static(self.type == self.TAN_FOV):
uv = coor * self.fov
dir = ts.normalize(ts.vec3(uv, 1))
elif ti.static(self.type == self.COS_FOV):
uv = coor * self.fov
dir = ts.vec3(ti.sin(uv), ti.cos(uv.norm()))
orig = (self.L2W[None] @ ts.vec4(orig, 1)).xyz
dir = (self.L2W[None] @ ts.vec4(dir, 0)).xyz
return orig, dir
def do_render(self):
model = self.model
scene = model.scene
L2W = model.L2W
a = scene.camera.untrans_pos(L2W @ self.vertex(0).pos)
b = scene.camera.untrans_pos(L2W @ self.vertex(1).pos)
c = scene.camera.untrans_pos(L2W @ self.vertex(2).pos)
A = scene.uncook_coor(a)
B = scene.uncook_coor(b)
C = scene.uncook_coor(c)
B_A = B - A
C_B = C - B
A_C = A - C
ilB_A = 1 / ts.length(B_A)
ilC_B = 1 / ts.length(C_B)
ilA_C = 1 / ts.length(A_C)
B_A *= ilB_A
C_B *= ilC_B
A_C *= ilA_C
BxA = ts.cross(B, A) * ilB_A
CxB = ts.cross(C, B) * ilC_B
AxC = ts.cross(A, C) * ilA_C
normal = ts.normalize(ts.cross(a - c, a - b))
light_dir = scene.camera.untrans_dir(scene.light_dir[None])
pos = (a + b + c) / 3
color = scene.opt.render_func(pos, normal, ts.vec3(0.0), light_dir)
color = scene.opt.pre_process(color)
Ak = 1 / (ts.cross(A, C_B) + CxB)
Bk = 1 / (ts.cross(B, A_C) + AxC)
Ck = 1 / (ts.cross(C, B_A) + BxA)
W = 1
M, N = int(ti.floor(min(A, B, C) - W)), int(ti.ceil(max(A, B, C) + W))
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
AB = ts.cross(X, B_A) + BxA
BC = ts.cross(X, C_B) + CxB
CA = ts.cross(X, A_C) + AxC
if AB <= 0 and BC <= 0 and CA <= 0:
zindex = pos.z #(Ak * a.z * BC + Bk * b.z * CA + Ck * c.z * AB)
zstep = zindex - ti.atomic_max(scene.zbuf[X], zindex)
if zstep >= 0:
scene.img[X] = color
def generate(self, coor):
fov = ti.static(math.radians(self.fov))
tan_fov = ti.static(math.tan(fov))
orig = ts.vec3(0.0)
dir = ts.vec3(0.0, 0.0, 1.0)
if ti.static(self.type == self.ORTHO):
orig = ts.vec3(coor, 0.0)
elif ti.static(self.type == self.TAN_FOV):
uv = coor * fov
dir = ts.normalize(ts.vec3(uv, 1))
elif ti.static(self.type == self.COS_FOV):
uv = coor * fov
dir = ts.vec3(ti.sin(uv), ti.cos(uv.norm()))
orig = self.trans_pos(orig)
dir = self.trans_dir(dir)
return orig, dir
def do_render(self, scene):
a = scene.camera.untrans_pos(self.a)
b = scene.camera.untrans_pos(self.b)
c = scene.camera.untrans_pos(self.c)
A = scene.uncook_coor(a)
B = scene.uncook_coor(b)
C = scene.uncook_coor(c)
B_A = B - A
C_B = C - B
A_C = A - C
ilB_A = 1 / ts.length(B_A)
ilC_B = 1 / ts.length(C_B)
ilA_C = 1 / ts.length(A_C)
B_A *= ilB_A
C_B *= ilC_B
A_C *= ilA_C
BxA = ts.cross(B, A) * ilB_A
CxB = ts.cross(C, B) * ilC_B
AxC = ts.cross(A, C) * ilA_C
normal = ts.normalize(ts.cross(a - c, a - b))
light_dir = scene.camera.untrans_dir(scene.light_dir[None])
pos = (a + b + c) * (1 / 3)
dir = ts.vec3(0.0)
color = scene.opt.render_func(pos, normal, dir, light_dir)
color = scene.opt.pre_process(color)
W = 0.4
M, N = int(ti.floor(min(A, B, C) - W)), int(ti.ceil(max(A, B, C) + W))
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
AB = ts.cross(X, B_A) + BxA
BC = ts.cross(X, C_B) + CxB
CA = ts.cross(X, A_C) + AxC
udf = max(AB, BC, CA)
if udf < 0:
scene.img[X] = color
elif udf < W:
t = ts.smoothstep(udf, W, 0)
ti.atomic_max(scene.img[X], t * color)
def get_dir(self, pos):
return ts.normalize(self.viewpos[None] - pos)
def brdf(self, normal, lightdir, viewdir):
NoH = max(0, ts.dot(normal, ts.normalize(lightdir + viewdir)))
ndf = (self.shineness + 8) / 8 * pow(NoH, self.shineness)
strength = self.color + ndf * self.specular
return strength
def render_triangle(model, camera, face):
scene = model.scene
L2W = model.L2W
_1 = ti.static(min(1, model.faces.m - 1))
_2 = ti.static(min(2, model.faces.m - 1))
ia, ib, ic = model.vi[face[0, 0]], model.vi[face[1, 0]], model.vi[face[2,
0]]
ta, tb, tc = model.vt[face[0, _1]], model.vt[face[1,
_1]], model.vt[face[2,
_1]]
na, nb, nc = model.vn[face[0, _2]], model.vn[face[1,
_2]], model.vn[face[2,
_2]]
a = camera.untrans_pos(L2W @ ia)
b = camera.untrans_pos(L2W @ ib)
c = camera.untrans_pos(L2W @ ic)
# NOTE: the normal computation indicates that # a front-facing face should
# be COUNTER-CLOCKWISE, i.e., glFrontFace(GL_CCW);
# this is to be compatible with obj model loading.
normal = ts.normalize(ts.cross(a - b, a - c))
pos = (a + b + c) / 3
view_pos = (a + b + c) / 3
if ti.static(camera.type == camera.ORTHO):
view_pos = ts.vec3(0.0, 0.0, 1.0)
if ts.dot(view_pos, normal) <= 0:
# shading
color = ts.vec3(0.0)
for light in ti.static(scene.lights):
color += scene.opt.render_func(pos, normal, ts.vec3(0.0), light)
color = scene.opt.pre_process(color)
A = camera.uncook(a)
B = camera.uncook(b)
C = camera.uncook(c)
scr_norm = 1 / ts.cross(A - C, B - A)
B_A = (B - A) * scr_norm
C_B = (C - B) * scr_norm
A_C = (A - C) * scr_norm
W = 1
# screen space bounding box
M, N = int(ti.floor(min(A, B, C) - W)), int(ti.ceil(max(A, B, C) + W))
M.x, N.x = min(max(M.x, 0),
camera.img.shape[0]), min(max(N.x, 0),
camera.img.shape[1])
M.y, N.y = min(max(M.y, 0),
camera.img.shape[0]), min(max(N.y, 0),
camera.img.shape[1])
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
# barycentric coordinates using the area method
X_A = X - A
w_C = ts.cross(B_A, X_A)
w_B = ts.cross(A_C, X_A)
w_A = 1 - w_C - w_B
# draw
in_screen = w_A >= 0 and w_B >= 0 and w_C >= 0 and 0 < X[
0] < camera.img.shape[0] and 0 < X[1] < camera.img.shape[1]
if not in_screen:
continue
zindex = 1 / (a.z * w_A + b.z * w_B + c.z * w_C)
if zindex < ti.atomic_max(camera.zbuf[X], zindex):
continue
coor = (ta * w_A + tb * w_B + tc * w_C)
camera.img[X] = color * model.texSample(coor)
def brdf(self, normal, lightdir, viewdir):
NoH = max(0, ts.dot(normal, ts.normalize(lightdir + viewdir)))
ndf = 5 / 2 * pow(NoH, 12)
strength = self.diffuse_color * self.diffuse + ndf * self.specular_color * self.specular
return strength
