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 paint():
for i, j in img:
ii = min(max(1, i * N // res), N - 2)
jj = min(max(1, j * N // res), N - 2)
if img_field == 0: # density
img[i, j] = Q[ii, jj][0]
elif img_field == 1: # numerical schlieren
img[i,
j] = ti.sqrt(((Q[ii + 1, jj][0] - Q[ii - 1, jj][0]) / h)**2 +
((Q[ii, jj + 1][0] - Q[ii, jj - 1][0]) / h)**2)
elif img_field == 2: # vorticity
img[i, j] = (Q[ii + 1, jj][2] - Q[ii - 1, jj][2]) / h - (
Q[ii, jj + 1][1] - Q[ii, jj - 1][1]) / h
elif img_field == 3: # velocity magnitude
img[i, j] = ti.sqrt(Q[ii, jj][1]**2 + Q[ii, jj][2]**2)
max = -1.0e10
min = 1.0e10
for i, j in img:
ti.atomic_max(max, img[i, j])
ti.atomic_min(min, img[i, j])
for i, j in img:
if use_fixed_caxis:
min = fixed_caxis[0]
max = fixed_caxis[1]
img[i, j] = (img[i, j] - min) / (max - min)
def checkStatic():
static[None] = 1
maxvn2[None] = 0.0
for i in ti.static(range(n_nodes)):
v2dn2 = v2d[i].dot(v2d[i])
ti.atomic_max(maxvn2[None], v2dn2)
if v2dn2 >= staticEpsilon * staticEpsilon:
static[None] = 0
def compute_max_grid_velocity(self, grid_v: ti.template()) -> ti.f32:
max_velocity = 0.0
for I in ti.grouped(grid_v):
v = grid_v[I]
v_max = 0.0
for i in ti.static(range(self.dim)):
v_max = max(v_max, abs(v[i]))
ti.atomic_max(max_velocity, v_max)
return max_velocity
def compute_max_velocity(self) -> ti.f32:
max_velocity = 0.0
for p in self.v:
v = self.v[p]
v_max = 0.0
for i in ti.static(range(self.dim)):
v_max = max(v_max, abs(v[i]))
ti.atomic_max(max_velocity, v_max)
return max_velocity
def foo():
# Parallel max
for i in range(1000):
ti.atomic_max(f[0], 1.12 * i)
# Serial max
for _ in range(1):
for i in range(1000):
f[1] = ti.max(1.12 * i, f[1])
def func():
for i in x:
x[i] = y[None]
x[0] = z[None]
x[1] += z[None]
x[2] -= z[None]
x[3] *= z[None]
x[4] //= z[None]
x[5] %= z[None]
x[6] &= z[None]
x[7] |= z[None]
x[8] ^= z[None]
ti.atomic_min(x[10], z[None])
ti.atomic_max(x[11], z[None])
def iou_kernel(self)->ti.float64:
ma = ti.cast(0., self.dtype)
mb = ti.cast(0., self.dtype)
I = ti.cast(0., self.dtype)
Ua = ti.cast(0., self.dtype)
Ub = ti.cast(0., self.dtype)
for i in ti.grouped(self.grid_mass):
ti.atomic_max(ma, self.grid_mass[i])
ti.atomic_max(mb, self.target_density[i])
I += self.grid_mass[i] * self.target_density[i]
Ua += self.grid_mass[i]
Ub += self.target_density[i]
I = I/ma/mb
U = Ua/ma + Ub/mb
return I/(U - I)
def func():
for i in x:
x[i] = y[None]
if ti.static(rhs_is_mat):
x[0] = z[None]
else:
x[0].fill(z[None])
x[1] += z[None]
x[2] -= z[None]
x[3] *= z[None]
x[4] /= z[None]
x[5] //= z[None]
x[6] %= z[None]
ti.atomic_min(x[7], z[None])
ti.atomic_max(x[8], z[None])
def updateCorrectedPressure(self):
# predict density, max density error, correction pressure
for pa_idx in self.predicted_position:
if self.particle_is_fluid[pa_idx] == 1:
pos_a = self.predicted_position[pa_idx] # position
vel_a = self.predicted_velocity[pa_idx] # velocity
# using two different method to predict density and
# choose the one which is closer to reference density
predicted_density = self.particle_density[pa_idx]
for nb_idx in range(self.particle_num_neighbors[pa_idx]):
pb_idx = self.particle_neighbors[pa_idx, nb_idx]
pos_b = self.predicted_position[pb_idx]
vel_b = self.predicted_velocity[pb_idx]
x_ab = pos_a - pos_b
v_ab = vel_a - vel_b
if self.particle_is_fluid[nb_idx] == 1:
predicted_density += self.rho_dt(v_ab, x_ab) * self.dt
# only handle positive error (thus incompressible)
density_error = max(predicted_density - self.rho_0, 0)
self.predicted_density[pa_idx] = predicted_density
self.corrected_pressure[
pa_idx] += self.delta * density_error # ref to PCISPH equation (9),(10)
err = abs(density_error)
# update max density error
self.max_density_error[None] = ti.atomic_max(
err, self.max_density_error[None])
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 compute_dist(t: ti.int32):
for bs, i in ti.ndrange(BATCH_SIZE, particle_num):
if material[i] == 0:
dist = 0.0
for j in ti.static(range(3)):
dist += (F_pos[bs, t, i][j] - target_centers[bs][j])**2
dist_sqr = ti.sqrt(dist)
ti.atomic_min(min_dist[bs], dist_sqr)
ti.atomic_max(max_height[bs], F_pos[bs, t, i][1])
ti.atomic_max(max_left[bs], F_pos[bs, t, i][0])
ti.atomic_max(max_right[bs], F_pos[bs, t, i][2])
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 render_line(camera, p1, p2):
scene = camera.scene
a = camera.untrans_pos(p1)
b = camera.untrans_pos(p2)
A = camera.uncook(a)
B = camera.uncook(b)
# dda algorithm
step = abs((B - A)).max()
delta = (B - A) / step
dz = (b.z - a.z) / step
# needs to do screen clipping
for i in range(step):
X = int(A + i * delta)
if X.x < 0 or X.x >= camera.res[0] or X.y < 0 or X.y >= camera.res[1]:
continue
zindex = 1 / (a.z + i * dz)
if zindex >= ti.atomic_max(camera.zbuf[X], zindex):
camera.img[X] = ts.vec3(0, 0.7, 0)
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 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 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 atomic_depth(self, X, depth):
idepth = self.idepth_fixp(1 / depth)
return idepth < ti.atomic_max(self['idepth'][X], idepth)
def test0():
for I in x:
x[I] = 0
x[1] = ti.cast(1, ti.u64) << 63
for I in x:
ti.atomic_max(x[0], x[I])
def render_triangle(model, camera, face):
scene = model.scene
L2W = model.L2W
posa, posb, posc = model.pos[face[0, 0]], model.pos[face[1, 0]], model.pos[
face[2, 0]]
texa, texb, texc = model.tex[face[0, 1]], model.tex[face[1, 1]], model.tex[
face[2, 1]]
nrma, nrmb, nrmc = model.nrm[face[0, 2]], model.nrm[face[1, 2]], model.nrm[
face[2, 2]]
posa = camera.untrans_pos(L2W @ posa)
posb = camera.untrans_pos(L2W @ posb)
posc = camera.untrans_pos(L2W @ posc)
nrma = camera.untrans_dir(L2W.matrix @ nrma)
nrmb = camera.untrans_dir(L2W.matrix @ nrmb)
nrmc = camera.untrans_dir(L2W.matrix @ nrmc)
pos_center = (posa + posb + posc) / 3
if ti.static(camera.type == camera.ORTHO):
pos_center = ts.vec3(0.0, 0.0, 1.0)
dpab = posa - posb
dpac = posa - posc
dtab = texa - texb
dtac = texa - texc
normal = ts.cross(dpab, dpac)
tan, bitan = compute_tangent(-dpab, -dpac, -dtab, -dtac)
# 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.
if ts.dot(pos_center, normal) <= 0:
clra = model.vertex_shader(posa, texa, nrma, tan, bitan)
clrb = model.vertex_shader(posb, texb, nrmb, tan, bitan)
clrc = model.vertex_shader(posc, texc, nrmc, tan, bitan)
A = camera.uncook(posa)
B = camera.uncook(posb)
C = camera.uncook(posc)
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
# screen space bounding box
M = int(ti.floor(min(A, B, C) - 1))
N = int(ti.ceil(max(A, B, C) + 1))
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))):
# 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
eps = ti.get_rel_eps() * 0.2
is_inside = w_A >= -eps and w_B >= -eps and w_C >= -eps
if not is_inside:
continue
zindex = 1 / (posa.z * w_A + posb.z * w_B + posc.z * w_C)
if zindex < ti.atomic_max(camera.fb['idepth'][X], zindex):
continue
clr = [
a * w_A + b * w_B + c * w_C
for a, b, c in zip(clra, clrb, clrc)
]
camera.fb.update(X, model.pixel_shader(*clr))
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 func():
a = -1000
for i in range(10):
ti.atomic_max(a, i)
A[None] = a
def func():
for i in range(n):
# this is an expr with side effect, make sure it's not optimized out.
ti.atomic_max(c[None], i * step)