def paintArrow(img: ti.template(),
orig,
dir,
color=1,
width=3,
max_size=12,
min_scale=0.4):
res = ts.vec(*img.shape)
I = orig * res
D = dir * res
J = I + D
DL = ts.length(D)
S = min(max_size, DL * min_scale)
DS = D / (DL + 1e-4) * S
SW = S + width
D1 = ti.Matrix.rotation2d(+math.pi * 3 / 4) @ DS
D2 = ti.Matrix.rotation2d(-math.pi * 3 / 4) @ DS
bmin, bmax = ti.floor(max(0,
min(I, J) - SW)), ti.ceil(
min(res - 1,
max(I, J) + SW))
for P in ti.grouped(ti.ndrange((bmin.x, bmax.x), (bmin.y, bmax.y))):
c0 = ts.smoothstep(abs(sdLine(I, J, P)), width, width / 2)
c1 = ts.smoothstep(abs(sdLine(J, J + D1, P)), width, width / 2)
c2 = ts.smoothstep(abs(sdLine(J, J + D2, P)), width, width / 2)
ti.atomic_max(img[P], max(c0, c1, c2) * color)
def render_line(model, camera, face, w0=0, w1=1):
posa, posb = face.pos # Position
texa, texb = face.tex # TexCoord
nrma, nrmb = face.nrm # Normal
clra = [posa, texa, nrma]
clrb = [posb, texb, nrmb]
A = camera.uncook(posa)
B = camera.uncook(posb)
M = int(ti.floor(min(A, B) - 1))
N = int(ti.ceil(max(A, B) + 1))
M = ts.clamp(M, 0, ti.Vector(camera.fb.res))
N = ts.clamp(N, 0, ti.Vector(camera.fb.res))
B_A = (B - A).normalized()
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
udf = abs((X - A).cross(B_A))
if udf >= w1:
continue
strength = ts.smoothstep(udf, w1, w0)
color = ts.vec3(strength)
camera.img[X] += color
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 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 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 do_render(self, scene):
W = 1
A = scene.uncook_coor(scene.camera.untrans_pos(self.a))
B = scene.uncook_coor(scene.camera.untrans_pos(self.b))
M, N = int(ti.floor(min(A, B) - W)), int(ti.ceil(max(A, B) + W))
for X in ti.grouped(ti.ndrange((M.x, N.x), (M.y, N.y))):
P = B - A
udf = (ts.cross(X, P) + ts.cross(B, A))**2 / P.norm_sqr()
XoP = ts.dot(X, P)
AoB = ts.dot(A, B)
if XoP > B.norm_sqr() - AoB:
udf = (B - X).norm_sqr()
elif XoP < AoB - A.norm_sqr():
udf = (A - X).norm_sqr()
if udf < 0:
scene.img[X] = ts.vec3(1.0)
elif udf < W**2:
t = ts.smoothstep(udf, 0, W**2)
ti.atomic_min(scene.img[X], ts.vec3(t))
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 imageGrid(p, n=20, tol=0.2):
return ts.smoothstep(abs(ts.fract(p * n) - 0.5).min(), 0.0, tol)