def refract(I, N, eta):
R = ti.Vector([0.0, 0.0, 0.0])
k = 1.0 - eta * eta * (1.0 - ti.dot(N, I) * ti.dot(N, I))
if (k < 0.0):
R = ti.Vector([0.0, 0.0, 0.0])
else:
R = eta * I - (eta * ti.dot(N, I) + ti.sqrt(k)) * N
return R
def sdf_Capsule(p, a, b, r):
ab = b - a
ap = p - a
t = ti.dot(ab, ap) / ti.dot(ab, ab)
t_clamped = clamp(t)
c = a + t_clamped * ab
return length(p - c) - r
def ray_plane_intersect(pos, d, pt_on_plane, norm):
dist = inf
hit_pos = ti.Vector([0.0, 0.0, 0.0])
denom = ti.dot(d, norm)
if abs(denom) > eps:
dist = ti.dot((pt_on_plane - pos), norm) / denom
hit_pos = pos + d * dist
return dist, hit_pos
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 intersect_light(pos, d):
light_loc = ti.Vector(light_pos)
dot = -ti.dot(d, ti.Vector(light_normal))
dist = ti.dot(d, light_loc - pos)
dist_to_light = inf
if dot > 0 and dist > 0:
D = dist / dot
dist_to_center = (light_loc - (pos + D * d)).norm_sqr()
if dist_to_center < light_radius**2:
dist_to_light = D
return dist_to_light
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 render_func(self, pos, normal, viewdir, light, color):
lightdir = light.get_dir(pos)
costheta = max(0, ti.dot(normal, lightdir))
l_out = ts.vec3(0.0)
if costheta > 0:
l_out = self.brdf(normal, -viewdir, lightdir, color)\
* costheta * light.get_color(pos)
return l_out
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 GetLight(p, t, hit, nor, step, rd):
lightPos = ti.Vector([0, 37, 1.0])
l = normalize(lightPos - p)
n = ti.Vector([0.0, 0.0, 0.0])
if hit == PARTICLES: #particles
n = nor
else: #sphere or plane
n = GetNormal(p, t, hit)
# attenuating the light
atten = 1.0 / (1.0 + l * 0.2 + l * l * 0.1)
spec = pow(max(ti.dot(reflect(-l, n), -rd), 0.0), 8.0)
diff = clamp(ti.dot(n, l))
diff = (diff + 1.0) / 2.0
sceneCol = (getColor(hit) * (diff + 0.15) +
ti.Vector([0.8, 0.8, 0.2]) * spec * 0.5) * atten
return sceneCol, n
def voxelize_triangles(self, num_triangles: ti.i32,
triangles: ti.ext_arr()):
for i in range(num_triangles):
jitter_scale = ti.cast(0, self.precision)
if ti.static(self.precision is ti.f32):
jitter_scale = 1e-4
else:
jitter_scale = 1e-8
# We jitter the vertices to prevent voxel samples from lying precicely at triangle edges
jitter = ti.Vector([-0.057616723909439505, -0.25608986292614977, 0.06716309129743714]) * jitter_scale
a = ti.Vector([triangles[i, 0], triangles[i, 1], triangles[i, 2]]) + jitter
b = ti.Vector([triangles[i, 3], triangles[i, 4], triangles[i, 5]]) + jitter
c = ti.Vector([triangles[i, 6], triangles[i, 7], triangles[i, 8]]) + jitter
bound_min = ti.Vector.zero(self.precision, 3)
bound_max = ti.Vector.zero(self.precision, 3)
for k in ti.static(range(3)):
bound_min[k] = min(a[k], b[k], c[k])
bound_max[k] = max(a[k], b[k], c[k])
p_min = int(ti.floor(bound_min[0] * self.inv_dx))
p_max = int(ti.floor(bound_max[0] * self.inv_dx)) + 1
p_min = max(self.padding, p_min)
p_max = min(self.res[0] - self.padding, p_max)
q_min = int(ti.floor(bound_min[1] * self.inv_dx))
q_max = int(ti.floor(bound_max[1] * self.inv_dx)) + 1
q_min = max(self.padding, q_min)
q_max = min(self.res[1] - self.padding, q_max)
normal = ti.normalized(ti.cross(b - a, c - a))
if abs(normal[2]) < 1e-10:
continue
a_proj = ti.Vector([a[0], a[1]])
b_proj = ti.Vector([b[0], b[1]])
c_proj = ti.Vector([c[0], c[1]])
for p in range(p_min, p_max):
for q in range(q_min, q_max):
pos2d = ti.Vector([(p + 0.5) * self.dx, (q + 0.5) * self.dx])
if inside_ccw(pos2d, a_proj, b_proj, c_proj) or inside_ccw(pos2d, a_proj, c_proj, b_proj):
base_voxel = ti.Vector([pos2d[0], pos2d[1], 0])
height = int(
-ti.dot(normal, base_voxel - a) /
normal[2] * self.inv_dx + 0.5)
height = min(height, self.res[1] - self.padding)
inc = 0
if normal[2] > 0:
inc = 1
else:
inc = -1
self.fill(p, q, height, inc)
def collide(dt: ti.f32):
for I in ti.grouped(self.grid_m):
offset = I * self.dx - ti.Vector(point)
n = ti.Vector(normal)
if ti.dot(offset, n) < 0:
if ti.static(surface == self.surface_sticky):
self.grid_v[I] = ti.Vector.zero(ti.f32, self.dim)
else:
v = self.grid_v[I]
normal_component = ti.dot(n, v)
if ti.static(surface == self.surface_slip):
# Project out all normal component
v = v - n * normal_component
else:
# Project out only inward normal component
v = v - n * min(normal_component, 0)
self.grid_v[I] = v
def get_color(z, timef32):
zoo = 0.62 + 0.38 * ti.cos(0.02 * timef32)
coa = ti.cos(0.1 * (1.0 - zoo) * timef32)
sia = ti.sin(0.1 * (1.0 - zoo) * timef32)
zoo = ti.pow(zoo, 8.0)
xy = ti.Vector([z[0] * coa - z[1] * sia, z[0] * sia + z[1] * coa])
cc = (
ti.Vector([1.0, 0.0])
+ smoothstep(1.0, 0.5, zoo) * ti.Vector([0.24814, 0.7369])
+ xy * zoo * 2.0
)
col = ti.Vector([0.0, 0.0, 0.0])
sc = ti.Vector([ti.abs(cc[0] - 1.0) - 1.0, cc[1]])
if ti.dot(sc, sc) >= 1.0:
co = 0.0
w = ti.Vector([0.5, 0.0])
for _ in range(256):
if ti.dot(w, w) > 1024:
break
w = cmul(cc, cmul(w, cadd(1.0, -w)))
co += 1.0
sco = co + 1.0 - log2(0.5 * (log2(ti.dot(cc, cc)) + log2(ti.dot(w, w))))
col = 0.5 + 0.5 * ti.cos(3.0 + sco * 0.1 + ti.Vector([0.0, 0.5, 1.0]))
if co > 255.5:
col = ti.Vector([0.0, 0.0, 0.0])
if ti.abs(cc.x - 1.0) < 3.0:
al = smoothstep(17.0, 12.0, timef32)
col = clamp(col, 0.0, 1.0)
x = 0.5
for _ in range(200):
x = cc[0] * (1.0 - x) * x
for _ in range(200):
x = cc[0] * (1.0 - x) * x
col = col + mix(
col,
ti.Vector([1.0, 1.0, 0.0]),
0.15 + 0.85 * ti.pow(clamp(ti.abs(sc.x + 1.0) * 0.4, 0.0, 1.0), 4.0),
) * al * 0.06 * ti.exp(-15000.0 * (cc.y - x) * (cc.y - x))
return col
def form_factor_ball(n, dpos, r):
d = dpos.norm()
z = ti.dot(n, dpos) / d
f = ff_clamp if z > 0 else 0.0
if d > r and z > 0:
#f1 = (d ** 2 + r ** 2 - d * r) / (d - r)
#f2 = ti.sqrt((d - r) * (d + r))
#f = min(2.0 * z / (3.0 * d) * (f1 - f2), ff_clamp)
x = d / r
f = min(2.0 * z / (3.0 * x) * (1. / (x - 1.) + 1. / 2. / x), ff_clamp)
return f
def refract(d, n, ni_over_nt):
# Assuming |d| and |n| are normalized
has_r, rd = 0, d
dt = ti.dot(d, n)
discr = 1.0 - ni_over_nt * ni_over_nt * (1.0 - dt * dt)
if discr > 0.0:
has_r = 1
rd = ti.normalized(ni_over_nt * (d - n * dt) - n * ti.sqrt(discr))
else:
rd *= 0.0
return has_r, rd
def intensity(self, pos):
'''
Calculate the distance from the object to the light
as the distance to the normal plane of the light position,
i.e., x0x+y0y+z0z=0.
'''
dist = self.lightdist[None] - ti.dot(pos - self.viewtarget[None],
self.viewdir[None])
f = 0.0
if dist > 0:
f = 1. / (1. + self.c1 * dist + self.c2 * dist**2)
return f
def calc_gi(self, camera):
scene = camera.scene
for X in ti.grouped(camera.img):
if self.enable_gi[None] == 0:
continue
# use an additional normal buffer at this time
if camera.zbuf[X] > 0:
z = 1 / camera.zbuf[X]
pos = cook_coord(camera,
ti.cast(ti.Vector([X.x, X.y, z]), ti.f32))
pos_world = camera.trans_pos(pos)
norm_world = camera.trans_dir(camera.nbuf[X])
dy = pos_world[1] - floor_y
ftot = 0.0
if dy > 0:
ftot = form_factor_floor(norm_world, dy, self.radius)
gtot = ts.vec3(0.0)
if ti.static(self.grid is not None):
base = self.grid.grid_index(pos_world)
for dI in ti.grouped(ti.ndrange((-4, 5), (-4, 5),
(-4, 5))):
I = max(0, min(base + dI, self.grid.gridsize - 1))
for p in range(self.grid_n[I + dI]):
i = self.grid_list[I + dI, p]
if (self.particles[i] - pos_world).norm_sqr() \
> (self.gi_cutoff * self.boxlength) ** 2:
continue
f, gi = get_gi_factors(camera, scene,
self.particles[i],
pos_world, norm_world,
self.radius,
self.colors[self.type[i]])
ftot += f
gtot += gi
else:
for i in range(self.particles.shape[0]):
# don't compute for particles behind the normal
if ti.dot(self.particles[i] - pos_world,
norm_world) < 0:
continue
f, gi = get_gi_factors(camera, scene,
self.particles[i], pos_world,
norm_world, self.radius,
self.colors[self.type[i]])
ftot += f
gtot += gi
camera.img[X] = camera.img[X] * ts.vec3(max(1 - ftot,
0)) + gtot * 0.3
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 collide(dt: ti.f32):
for I in ti.grouped(self.grid_m):
offset = I * self.dx - ti.Vector(center)
if offset.norm_sqr() < radius * radius:
if ti.static(surface == self.surface_sticky):
self.grid_v[I] = ti.Vector.zero(ti.f32, self.dim)
else:
v = self.grid_v[I]
normal = ti.Vector.normalized(offset, eps=1e-5)
normal_component = ti.dot(normal, v)
if ti.static(surface == self.surface_slip):
# Project out all normal component
v = v - normal * normal_component
else:
# Project out only inward normal component
v = v - normal * min(normal_component, 0)
self.grid_v[I] = v
def grid_op():
for i, j in grid_m:
if grid_m[i, j] > 0:
inv_m = 1 / grid_m[i, j]
grid_v[i, j] = inv_m * grid_v[i, j]
grid_v(1)[i, j] -= dt * 9.8
# center sticky circle
dist = ti.Vector([i * dx - 0.5, j * dx - 0.5])
if dist.norm_sqr() < 0.005:
dist = ti.Vector.normalized(dist)
grid_v[i, j] -= dist * ti.dot(grid_v[i, j], dist)
# box
if i < bound and grid_v(0)[i, j] < 0:
grid_v(0)[i, j] = 0
if i > n_grid - bound and grid_v(0)[i, j] > 0:
grid_v(0)[i, j] = 0
if j < bound and grid_v(1)[i, j] < 0:
grid_v(1)[i, j] = 0
if j > n_grid - bound and grid_v(1)[i, j] > 0:
grid_v(1)[i, j] = 0
def grid_op():
for i, j in grid_m:
if grid_m[i, j] > 0:
inv_m = 1 / grid_m[i, j]
v_out = inv_m * grid_v_in[i, j] # momentum to velocity
v_out[1] -= dt * gravity # gravity
# center sticky circle
dist = ti.Vector([i * dx - 0.5, j * dx - 0.5])
if dist.norm_sqr() < 0.005:
dist = ti.normalized(dist)
v_out -= dist * ti.dot(v_out, dist)
# boundary conditions
if i < bound and v_out[0] < 0:
v_out[0] = 0
if i > n_grid - bound and v_out[0] > 0:
v_out[0] = 0
if j < bound and v_out[1] < 0:
v_out[1] = 0
if j > n_grid - bound and v_out[1] > 0:
v_out[1] = 0
grid_v_out[i, j] = v_out
def reflect(d, n):
# Assuming |d| and |n| are normalized
return d - 2.0 * ti.dot(d, n) * n
def capsule(grid_pos, a, b):
ab = a - b
ap = grid_pos - b
t = ti.dot(ab, ap) / ti.dot(ab, ab)
t_clamped = clamp(t)
return b + t_clamped*ab
def frensel(self, view, halfway, color):
f0 = ts.mix(ts.vec3(self.specular), color, self.metallic)
hdotv = min(1, max(ti.dot(halfway, view), 0))
return (f0 + (1.0 - f0) * (1.0 - hdotv)**5) * self.ks
def microfacet(self, normal, halfway):
alpha = self.roughness
ggx = alpha**2 / math.pi
ggx /= (ti.dot(normal, halfway)**2 * (alpha**2 - 1.0) + 1.0)**2
return ggx
def planeSDF(p, n):
nxyz = ti.Vector([n[0], n[1], n[2]])
return ti.dot(p, nxyz) + n[3]
def reflect(I, N):
return I - 2.0 * ti.dot(N, I) * N
