def refract(d, n, ni_over_nt):
# Assuming |d| and |n| are normalized
has_r, rd = 0, d
dt = d.dot(n)
discr = 1.0 - ni_over_nt * ni_over_nt * (1.0 - dt * dt)
if discr > 0.0:
has_r = 1
rd = (ni_over_nt * (d - n * dt) - n * ti.sqrt(discr)).normalized()
else:
rd *= 0.0
return has_r, rd
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 struct_cg(
A: ti.template(), b: ti.template(), x: ti.template(),
Ax: ti.template(), Ap: ti.template(), r: ti.template(),
p: ti.template(), nx: ti.i32, ny: ti.i32, eps: ti.f64, output: ti.i32):
n = (nx + 1) * ny
# dot(A,x)
for i in range(n):
Ax[i] = 0.0
for j in range(i - ny, i + ny + 1):
Ax[i] += A[i, j] * x[j]
# r = b - dot(A,x)
# p = r
for i in range(n):
r[i] = b[i] - Ax[i]
p[i] = r[i]
rsold = 0.0
for i in range(n):
rsold += r[i] * r[i]
for steps in range(100 * n):
# dot(A,p)
for i in range(n):
Ap[i] = 0.0
for j in range(i - ny, i + ny + 1):
Ap[i] += A[i, j] * p[j]
# dot(p, Ap) => pAp
pAp = 0.0
for i in range(n):
pAp += p[i] * Ap[i]
alpha = rsold / pAp
# x = x + dot(alpha,p)
# r = r - dot(alpha,Ap)
for i in range(n):
x[i] += alpha * p[i]
r[i] -= alpha * Ap[i]
rsnew = 0.0
for i in range(n):
rsnew += r[i] * r[i]
if ti.sqrt(rsnew) < eps:
if output:
print(" >> The solution has converged...")
break
for i in range(n):
p[i] = r[i] + (rsnew / rsold) * p[i]
rsold = rsnew
if output:
print(" >> Iteration ", steps, ", residual = ", rsold)
def sym_eig2x2(A, dt):
tr = A.trace()
det = A.determinant()
gap = tr**2 - 4 * det
lambda1 = (tr + ti.sqrt(gap)) * 0.5
lambda2 = (tr - ti.sqrt(gap)) * 0.5
eigenvalues = ti.Vector([lambda1, lambda2]).cast(dt)
A1 = A - lambda1 * ti.Matrix.identity(dt, 2)
A2 = A - lambda2 * ti.Matrix.identity(dt, 2)
v1 = ti.Vector.zero(dt, 2)
v2 = ti.Vector.zero(dt, 2)
if all(A1 == ti.Matrix.zero(dt, 2, 2)) and all(
A1 == ti.Matrix.zero(dt, 2, 2)):
v1 = ti.Vector([0.0, 1.0]).cast(dt)
v2 = ti.Vector([1.0, 0.0]).cast(dt)
else:
v1 = ti.Vector([A2[0, 0], A2[1, 0]]).cast(dt).normalized()
v2 = ti.Vector([A1[0, 0], A1[1, 0]]).cast(dt).normalized()
eigenvectors = ti.Matrix.cols([v1, v2])
return eigenvalues, eigenvectors
def rotmat2quaternion(rotmat):
m00, m01, m02 = rotmat[0, 0], rotmat[0, 1], rotmat[0, 2]
m10, m11, m12 = rotmat[1, 0], rotmat[1, 1], rotmat[1, 2]
m20, m21, m22 = rotmat[2, 0], rotmat[2, 1], rotmat[2, 2]
tr = m00 + m11 + m22 # ti.tr(rotmat)
qw = ti.sqrt(1.0 + tr) / 2.0
qx = (m21 - m12) / (4.0 * qw)
qy = (m02 - m20) / (4.0 * qw)
qz = (m10 - m01) / (4.0 * qw)
return ti.Vector([qw, qx, qy, qz])
def mask2sdf(self, to_rgb=True, output=True):
self.gen_udf_w_h()
if to_rgb: # grey value == 0.5 means sdf == 0, scale sdf proportionally
max_positive_dist = ti.sqrt(self.find_max(self.bit_pic_white))
min_negative_dist = ti.sqrt(self.find_max(
self.bit_pic_black)) # this value is positive
coefficient = 127.5 / max(max_positive_dist, min_negative_dist)
offset = 127.5
self.post_process_sdf(self.bit_pic_white, self.bit_pic_black,
self.num, coefficient, offset)
if output:
cv2.imwrite(self.output_filename('_sdf'),
self.output_pic.to_numpy())
else: # no normalization
if output:
pass
else:
self.post_process_sdf_linear_1channel(self.bit_pic_white,
self.bit_pic_black,
self.num)
def NeoHookeanElasticity(U, sig):
J = sig[0, 0] * sig[1, 1]
mu_J_1_2 = mu_0 * ti.sqrt(J)
J_prime = kappa_0 * 0.5 * (J * J - 1.0)
sqrS_1_2 = (sig[0, 0] * sig[0, 0] + sig[1, 1] * sig[1, 1]) / 2.0
stress = ti.Matrix.identity(float, 2)
stress[0, 0] = (sig[0, 0] * sig[0, 0] - sqrS_1_2) * mu_J_1_2
stress[1, 1] = (sig[1, 1] * sig[1, 1] - sqrS_1_2) * mu_J_1_2
stress = U @ stress @ U.transpose()
stress[0, 0] += J_prime
stress[1, 1] += J_prime
return stress
def svd2d(A, dt):
"""Perform singular value decomposition (A=USV^T) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U`, 'S' and `V`.
"""
R, S = polar_decompose2d(A, dt)
c, s = ti.cast(0.0, dt), ti.cast(0.0, dt)
s1, s2 = ti.cast(0.0, dt), ti.cast(0.0, dt)
if abs(S[0, 1]) < 1e-5:
c, s = 1, 0
s1, s2 = S[0, 0], S[1, 1]
else:
tao = ti.cast(0.5, dt) * (S[0, 0] - S[1, 1])
w = ti.sqrt(tao**2 + S[0, 1]**2)
t = ti.cast(0.0, dt)
if tao > 0:
t = S[0, 1] / (tao + w)
else:
t = S[0, 1] / (tao - w)
c = 1 / ti.sqrt(t**2 + 1)
s = -t * c
s1 = c**2 * S[0, 0] - 2 * c * s * S[0, 1] + s**2 * S[1, 1]
s2 = s**2 * S[0, 0] + 2 * c * s * S[0, 1] + c**2 * S[1, 1]
V = ti.Matrix.zero(dt, 2, 2)
if s1 < s2:
tmp = s1
s1 = s2
s2 = tmp
V = [[-s, c], [-c, -s]]
else:
V = [[c, s], [-s, c]]
U = R @ V
return U, ti.Matrix([[s1, ti.cast(0, dt)], [ti.cast(0, dt), s2]]), V
def range_conjgrad():
# dot(A,x)
for i in range(n):
Ax[i] = 0.0
for j in range(n):
Ax[i] += A[i, j] * x[j]
# r = b - dot(A,x)
# p = r
for i in range(n):
r[i] = b[i] - Ax[i]
p[i] = r[i]
rsold = 0.0
for i in range(n):
rsold += r[i] * r[i]
for steps in range(10000 * n):
# dot(A,p)
for i in range(n):
Ap[i] = 0.0
for j in range(n):
Ap[i] += A[i, j] * p[j]
# dot(p, Ap) => pAp
pAp = 0.0
for i in range(n):
pAp += p[i] * Ap[i]
alpha = rsold / pAp
# x = x + dot(alpha,p)
# r = r - dot(alpha,Ap)
for i in range(n):
x[i] += alpha * p[i]
r[i] -= alpha * Ap[i]
rsnew = 0.0
for i in range(n):
rsnew += r[i] * r[i]
if ti.sqrt(rsnew) < 1e-8:
print("The solution has converged...")
break
for i in range(n):
p[i] = r[i] + (rsnew / rsold) * p[i]
rsold = rsnew
print("Iteration ", steps, ", residual = ", rsold)
if steps == n - 1 and rsold > 1e-8:
print("The solution did NOT converge...")
return steps
def p2g(self, dt: ti.f32):
for p in self.x:
base = (self.x[p] * self.inv_dx - 0.5).cast(int)
fx = self.x[p] * self.inv_dx - base.cast(float)
# Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
# deformation gradient update
self.F[p] = (ti.Matrix.identity(ti.f32, self.dim) +
dt * self.C[p]) @ self.F[p]
# Hardening coefficient: snow gets harder when compressed
h = ti.exp(10 * (1.0 - self.Jp[p]))
if self.material[
p] == self.material_elastic: # jelly, make it softer
h = 0.3
mu, la = self.mu_0 * h, self.lambda_0 * h
if self.material[p] == self.material_water: # liquid
mu = 0.0
U, sig, V = ti.svd(self.F[p])
J = 1.0
for d in ti.static(range(self.dim)):
new_sig = sig[d, d]
if self.material[p] == self.material_snow: # Snow
new_sig = min(max(sig[d, d], 1 - 2.5e-2),
1 + 4.5e-3) # Plasticity
self.Jp[p] *= sig[d, d] / new_sig
sig[d, d] = new_sig
J *= new_sig
if self.material[p] == self.material_water:
# Reset deformation gradient to avoid numerical instability
self.F[p] = ti.Matrix.identity(ti.f32, self.dim) * ti.sqrt(J)
elif self.material[p] == self.material_snow:
# Reconstruct elastic deformation gradient after plasticity
self.F[p] = U @ sig @ V.T()
stress = 2 * mu * (self.F[p] - U @ V.T()) @ self.F[p].T(
) + ti.Matrix.identity(ti.f32, self.dim) * la * J * (J - 1)
stress = (-dt * self.p_vol * 4 * self.inv_dx**2) * stress
affine = stress + self.p_mass * self.C[p]
# Loop over 3x3 grid node neighborhood
for offset in ti.static(ti.grouped(self.stencil_range())):
dpos = (offset.cast(float) - fx) * self.dx
weight = 1.0
for d in ti.static(range(self.dim)):
weight *= w[offset[d]][d]
self.grid_v[base +
offset] += weight * (self.p_mass * self.v[p] +
affine @ dpos)
self.grid_m[base + offset] += weight * self.p_mass
def Dectect_Tearing_Part():
for p in x:
b = F[p] @ F[p].transpose()
Cpp = F[p].determinant()**(-2 /
3) * F[p].transpose() @ b.inverse() @ F[p]
tr = Cpp[0, 0] + Cpp[1, 1] + Cpp[2, 2]
dev = Cpp - (tr / 3) * ti.Matrix.identity(ti.f32, 3)
dev_sqr = dev @ dev
tr = ti.sqrt(dev_sqr[0, 0] + dev_sqr[1, 1] + dev_sqr[2, 2])
if tr > tear_stress:
weak[p] = 1
else:
weak[p] = 0
def copy(img: np.ndarray):
for i in range(res[0]):
for j in range(res[1]):
u = 1.0 * i / res[0]
v = 1.0 * j / res[1]
darken = 1.0 - vignette_strength * ti.max(
(ti.sqrt((u - vignette_center[0])**2 +
(v - vignette_center[1])) - vignette_radius), 0)**2
coord = ((res[1] - 1 - j) * res[0] + i) * 3
for c in ti.static(range(3)):
img[coord + c] = color_buffer[i, j][2 - c] * darken
def blueorange(rgb):
'''
Convert RGB value (vector) into blue-orange colormap (vector).
:parameter rgb: (3D vector)
The RGB value, X compoment is the R value, can be eitier float or int.
:return:
The return value is calculated as `sqrt(mix(vec(0, 0.01, 0.05), vec(1.19, 1.04, 0.98), color))`.
'''
blue = ts.vec(0.00, 0.01, 0.05)
orange = ts.vec(1.19, 1.04, 0.98)
return ti.sqrt(ts.mix(blue, orange, rgb))
def tonemap(accumulated: ti.f32) -> ti.f32:
sum = 0.0
sum_sq = 0.0
for i, j in color_buffer:
luma = color_buffer[i, j][0] * 0.2126 + color_buffer[
i, j][1] * 0.7152 + color_buffer[i, j][2] * 0.0722
sum += luma
sum_sq += ti.pow(luma / accumulated, 2.0)
mean = sum / (res[0] * res[1])
var = sum_sq / (res[0] * res[1]) - ti.pow(mean / accumulated, 2.0)
for i, j in tonemapped_buffer:
tonemapped_buffer[i, j] = ti.sqrt(color_buffer[i, j] / mean * 0.6)
return var
def random_point_in_unit_polygon(self, sides, angle):
point = ti.Vector.zero(ti.f32, 2)
central_angle = 2 * math.pi / sides
while True:
point = ti.Vector([ti.random(), ti.random()]) * 2 - 1
point_angle = ti.atan2(point.y, point.x)
theta = (point_angle -
angle) % central_angle # polygon angle is from +X axis
phi = central_angle / 2
dist = ti.sqrt((point**2).sum())
if dist < ti.cos(phi) / ti.cos(phi - theta):
break
return point
def make_tests():
assert isnan(nan) == isnan(-nan) == True
x = -1.0
assert isnan(ti.sqrt(x)) == True
assert isnan(inf) == isnan(1.0) == isnan(-1) == False
assert isinf(inf) == isinf(-inf) == True
assert isinf(nan) == isinf(1.0) == isinf(-1) == False
v = ti.math.vec4(inf, -inf, 1.0, nan)
assert all(isinf(v) == [1, 1, 0, 0])
v = ti.math.vec4(nan, -nan, 1, inf)
assert all(isnan(v) == [1, 1, 0, 0])
def sample_brdf(normal):
# cosine hemisphere sampling
# first, uniformly sample on a disk (r, theta)
r, theta = 0.0, 0.0
sx = ti.random() * 2.0 - 1.0
sy = ti.random() * 2.0 - 1.0
if sx >= -sy:
if sx > sy:
# first region
r = sx
div = abs(sy / r)
if sy > 0.0:
theta = div
else:
theta = 7.0 + div
else:
# second region
r = sy
div = abs(sx / r)
if sx > 0.0:
theta = 1.0 + sx / r
else:
theta = 2.0 + sx / r
else:
if sx <= sy:
# third region
r = -sx
div = abs(sy / r)
if sy > 0.0:
theta = 3.0 + div
else:
theta = 4.0 + div
else:
# fourth region
r = -sy
div = abs(sx / r)
if sx < 0.0:
theta = 5.0 + div
else:
theta = 6.0 + div
# Malley's method
u = ti.Vector([1.0, 0.0, 0.0])
if abs(normal[1]) < 1 - eps:
u = normal.cross(ti.Vector([0.0, 1.0, 0.0]))
v = normal.cross(u)
theta = theta * math.pi * 0.25
costt, sintt = ti.cos(theta), ti.sin(theta)
xy = (u * costt + v * sintt) * r
zlen = ti.sqrt(max(0.0, 1.0 - xy.dot(xy)))
return xy + zlen * normal
def intersect_prim_any(self, origin, direction, primitive_id):
prim_type = UF.get_prim_type(self.primitive, primitive_id)
hit_t = UF.INF_VALUE
if prim_type == SCD.PRIMITIVE_TRI:
hit_t, u, v = self.intersect_tri(origin, direction, primitive_id)
else:
sha_id = UF.get_prim_vindex(self.primitive, primitive_id)
sha_type = UF.get_shape_type(self.shape, sha_id)
if sha_type == SCD.SHPAE_SPHERE:
r = UF.get_shape_radius(self.shape, sha_id)
centre = UF.get_shape_pos(self.shape, sha_id)
oc = centre - origin
dis_oc_square = oc.dot(oc)
dis_op = direction.dot(oc)
dis_cp = ti.sqrt(dis_oc_square - dis_op * dis_op)
if (dis_cp < r):
a = direction.dot(direction)
b = -2.0 * dis_op
c = dis_oc_square - r * r
hit_t = (-b - ti.sqrt(b * b - 4.0 * a * c)) / 2.0 / a
return hit_t
def compute_sdf(z, c):
md2 = 1.0
mz2 = dot(z, z)
for _ in range(iters):
md2 *= max_norm * mz2
z = quat_mul(z, z) + c
mz2 = z.dot(z)
if mz2 > max_norm:
break
return 0.25 * ti.sqrt(mz2 / md2) * ti.log(mz2)
def inversesqrt(x):
'''
Return the inverse of the square root of the parameter.
`inversesqrt` returns the inverse of the square root of x; i.e.
the value `1 / sqrt(x)`. The result is undefined if x <= 0.
:parameter x:
Specify the value of which to take the inverse of the square root.
:return:
The return value can be calculated as `1 / sqrt(x)` or `pow(x, -0.5)`.
'''
return 1 / ti.sqrt(x)
def eig2x2(A, dt):
tr = A.trace()
det = A.determinant()
gap = tr**2 - 4 * det
lambda1 = ti.Vector.zero(dt, 2)
lambda2 = ti.Vector.zero(dt, 2)
v1 = ti.Vector.zero(dt, 4)
v2 = ti.Vector.zero(dt, 4)
if gap > 0:
lambda1 = ti.Vector([tr + ti.sqrt(gap), ti.cast(0.0, dt)]) * 0.5
lambda2 = ti.Vector([tr - ti.sqrt(gap), ti.cast(0.0, dt)]) * 0.5
A1 = A - lambda1[0] * ti.Matrix.identity(dt, 2)
A2 = A - lambda2[0] * ti.Matrix.identity(dt, 2)
if all(A1 == ti.Matrix.zero(dt, 2, 2)) and all(
A1 == ti.Matrix.zero(dt, 2, 2)):
v1 = ti.Vector([0.0, 0.0, 1.0, 0.0]).cast(dt)
v2 = ti.Vector([1.0, 0.0, 0.0, 0.0]).cast(dt)
else:
v1 = ti.Vector([A2[0, 0], 0.0, A2[1, 0],
0.0]).cast(dt).normalized()
v2 = ti.Vector([A1[0, 0], 0.0, A1[1, 0],
0.0]).cast(dt).normalized()
else:
lambda1 = ti.Vector([tr, ti.sqrt(-gap)]) * 0.5
lambda2 = ti.Vector([tr, -ti.sqrt(-gap)]) * 0.5
A1r = A - lambda1[0] * ti.Matrix.identity(dt, 2)
A1i = -lambda1[1] * ti.Matrix.identity(dt, 2)
A2r = A - lambda2[0] * ti.Matrix.identity(dt, 2)
A2i = -lambda2[1] * ti.Matrix.identity(dt, 2)
v1 = ti.Vector([A2r[0, 0], A2i[0, 0], A2r[1, 0],
A2i[1, 0]]).cast(dt).normalized()
v2 = ti.Vector([A1r[0, 0], A1i[0, 0], A1r[1, 0],
A1i[1, 0]]).cast(dt).normalized()
eigenvalues = ti.Matrix.rows([lambda1, lambda2])
eigenvectors = ti.Matrix.cols([v1, v2])
return eigenvalues, eigenvectors
def computeVisualLaplacian():
"""
Create the laplacian at any point using the method described here:
- paper: https://www.cs.cmu.edu/~kmcrane/Projects/MonteCarloGeometryProcessing/paper.pdf
- project page: https://www.cs.cmu.edu/~kmcrane/Projects/MonteCarloGeometryProcessing/index.html
Uses the SDF (SDF and Laplacian tensor must have the same shape).
"""
epsilon = ti.cast(1, ti.f32)
max_iter = 32
n_walks = ti.cast(128, ti.i32)
for i,j in sdf:
sum_u_x = ti.cast(0, ti.f32)
sum_u_y = ti.cast(0, ti.f32)
for walk in range(n_walks):
d = sdf[i,j][0]
n = 0
x_index = i
y_index = j
while d > epsilon and n < max_iter:
alpha = ti.random() * 2 * 3.14159
x = x_index + d * ti.cos(alpha)
y = y_index + d * ti.sin(alpha)
x_index = ti.cast(x, ti.i32)
y_index = ti.cast(y, ti.i32)
d = sdf[x_index, y_index][0]
n += 1
#sum_u += getBoundaryValue(x_index, y_index)
dx = sdf[x_index+1, y_index-1][0] - sdf[x_index-1, y_index-1][0] \
+ sdf[x_index+1, y_index][0] - sdf[x_index-1, y_index][0] \
+ sdf[x_index+1, y_index+1][0] - sdf[x_index-1, y_index+1][0]
dy = sdf[x_index-1, y_index+1][0] - sdf[x_index-1, y_index-1][0] \
+ sdf[x_index, y_index+1][0] - sdf[x_index, y_index-1][0] \
+ sdf[x_index+1, y_index+1][0] - sdf[x_index+1, y_index-1][0]
sum_u_x += ti.abs(dx)
sum_u_y += ti.abs(dy)
visual_laplacian[i,j][0] = sum_u_x / n_walks / ti.sqrt(sdf[i,j][0])
visual_laplacian[i,j][1] = sum_u_y / n_walks / ti.sqrt(sdf[i,j][0])
def p2g(f: ti.i32):
for p in range(0, n_particles):
base = ti.cast(x[f, p] * inv_dx - 0.5, ti.i32)
fx = x[f, p] * inv_dx - ti.cast(base, ti.i32)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
new_F = (ti.Matrix.diag(dim=dim, val=1) + dt * C[f, p]) @ F[f, p]
J = ti.determinant(new_F)
if particle_type[p] == 0: # fluid
sqrtJ = ti.sqrt(J)
# TODO: need pow(x, 1/3)
new_F = ti.Matrix([[sqrtJ, 0, 0], [0, sqrtJ, 0], [0, 0, 1]])
F[f + 1, p] = new_F
# r, s = ti.polar_decompose(new_F)
act_id = actuator_id[p]
act = actuation[f, ti.max(0, act_id)] * act_strength
if act_id == -1:
act = 0.0
# ti.print(act)
A = ti.Matrix([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]
]) * act
cauchy = ti.Matrix(zero_matrix())
mass = 0.0
ident = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
if particle_type[p] == 0:
mass = 4
cauchy = ti.Matrix(ident) * (J - 1) * E
else:
mass = 1
cauchy = mu * (new_F @ ti.transposed(new_F)) + ti.Matrix(ident) * (
la * ti.log(J) - mu)
cauchy += new_F @ A @ ti.transposed(new_F)
stress = -(dt * p_vol * 4 * inv_dx * inv_dx) * cauchy
affine = stress + mass * C[f, p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
for k in ti.static(range(3)):
offset = ti.Vector([i, j, k])
dpos = (ti.cast(ti.Vector([i, j, k]), real) - fx) * dx
weight = w[i](0) * w[j](1) * w[k](2)
grid_v_in[base + offset].atomic_add(
weight * (mass * v[f, p] + affine @ dpos))
grid_m_in[base + offset].atomic_add(weight * mass)
def pressure_jacobi_iter(pf: ti.template(), pf_new: ti.template(),
divf: ti.template()) -> ti.f32:
norm_new = 0
norm_diff = 0
for i, j in pf:
pf_new[i, j] = 0.25 * (p_with_boundary(pf, i + 1, j) + p_with_boundary(
pf, i - 1, j) + p_with_boundary(pf, i, j + 1) +
p_with_boundary(pf, i, j - 1) - divf[i, j])
pf_diff = ti.abs(pf_new[i, j] - p_with_boundary(pf, i, j))
norm_new += (pf_new[i, j] * pf_new[i, j])
norm_diff += (pf_diff * pf_diff)
residual = ti.sqrt(norm_diff / norm_new)
if norm_new == 0:
residual = 0.0
return residual
def update_velocity(self) -> ti.f64:
nx, ny, dx, dy = self.nx, self.ny, self.dx, self.dy
max_udiff = 0.0
max_vdiff = 0.0
for i, j in ti.ndrange((2, nx + 1), (1, ny + 1)):
if ti.abs(self.u_mid[i, j] - self.u[i, j]) > max_udiff:
max_udiff = ti.abs(self.u_mid[i, j] - self.u[i, j])
self.u[i, j] = self.alpha_m * self.u_mid[i, j] + (
1 - self.alpha_m) * self.u[i, j]
for i, j in ti.ndrange((1, nx + 1), (2, ny + 1)):
if ti.abs(self.v_mid[i, j] - self.v[i, j]) > max_vdiff:
max_vdiff = ti.abs(self.v_mid[i, j] - self.v[i, j])
self.v[i, j] = self.alpha_m * self.v_mid[i, j] + (
1 - self.alpha_m) * self.v[i, j]
return ti.sqrt(max_udiff**2 + max_vdiff**2)
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 ggx_sample(n, wo, roughness):
u = ti.Vector([1.0, 0.0, 0.0])
if abs(n[1]) < 1.0 - eps:
u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized()
v = n.cross(u)
r0 = ti.random()
r1 = ti.random()
a = roughness ** 2.0
a2 = a * a
theta = ti.acos(ti.sqrt((1.0 - r0) / ((a2-1.0)*r0 + 1.0)))
phi = 2.0 * math.pi * r1
wm = ti.cos(theta) * n + ti.sin(theta) * ti.cos(phi) * \
u + ti.sin(theta) * ti.sin(phi) * v
wi = reflect(-wo, wm)
return wi
def randSolid2D():
'''
Generate a 2-D random unit vector whose length is <= 1.0.
The return value is a 2-D vector, whose tip distributed evenly
**inside** a unit circle.
:return:
The return value is computed as::
a = rand() * math.tau
r = sqrt(rand())
return vec(cos(a), sin(a)) * r
'''
a = rand() * math.tau
r = ti.sqrt(rand())
return ti.Vector([ti.cos(a), ti.sin(a)]) * r
def create_torus_density():
for i in range(density_res):
for j in range(density_res):
for k in range(density_res):
# Convert to density coordinates
x = ti.cast(k, ti.f32) * inv_density_res - 0.5
y = ti.cast(j, ti.f32) * inv_density_res - 0.5
z = ti.cast(i, ti.f32) * inv_density_res - 0.5
l = ti.sqrt(x * x + y * y + z * z)
# Swap x, y to rotate the torus
if in_torus(y, x, z):
density[i, j, k] = inv_density_res
else:
density[i, j, k] = 0.0
def electric_field(t: ti.f64):
t_eff = t - t_h
a = parameters[0] / (ti.sqrt(2 * np.pi) * parameters[1])
w = ti.sqr(parameters[1])
e = a * ti.exp(-ti.sqr(t_eff) /
(2 * w)) * (ti.sin(parameters[2] * t_eff) +
parameters[3] * ti.sin(parameters[4] * t_eff))
e_field[None] = e
grad[0] = e / parameters[0]
grad[1] = -e / parameters[1] + e * ti.sqr(t_eff) / (parameters[1]**3)
grad[2] = a * ti.exp(-ti.sqr(t_eff) /
(2 * w)) * ti.cos(parameters[2] * t_eff) * t_eff
grad[3] = a * ti.exp(-ti.sqr(t_eff) /
(2 * w)) * ti.sin(parameters[4] * t_eff)
grad[4] = a * ti.exp(-ti.sqr(t_eff) / (2 * w)) * parameters[3] * ti.cos(
parameters[4] * t_eff) * t_eff
