def complex_diff(z1, z2):
result = ti.Vector([0., 0.], dt=ti.f64, needs_grad=True)
result[0] = (z1[0] * z2[0] + z1[1] * z2[1]) / (ti.sqr(z2[0]) +
ti.sqr(z2[1]))
result[1] = (z1[1] * z2[0] - z1[0] * z2[1]) / (ti.sqr(z2[0]) +
ti.sqr(z2[1]))
return result
def apply_spring_force(t: ti.i32):
for i in range(n_springs):
a = spring_anchor_a[i]
b = spring_anchor_b[i]
pos_a, vel_a, rela_a = to_world(t, a, spring_offset_a[i])
pos_b, vel_b, rela_b = to_world(t, b, spring_offset_b[i])
dist = pos_a - pos_b
length = dist.norm() + 1e-4
act = actuation[t, i]
is_joint = spring_length[i] == -1
target_length = spring_length[i] * (1.0 + spring_actuation[i] * act)
if is_joint:
target_length = 0.0
impulse = dt * (length -
target_length) * spring_stiffness[i] / length * dist
if is_joint:
rela_vel = vel_a - vel_b
rela_vel_norm = rela_vel.norm() + 1e-1
impulse_dir = rela_vel / rela_vel_norm
impulse_contribution = inverse_mass[a] + ti.sqr(
cross(impulse_dir, rela_a)) * inverse_inertia[
a] + inverse_mass[b] + ti.sqr(cross(impulse_dir,
rela_b)) * \
inverse_inertia[
b]
# project relative velocity
impulse += rela_vel_norm / impulse_contribution * impulse_dir
apply_impulse(t, a, -impulse, pos_a, 0.0)
apply_impulse(t, b, impulse, pos_b, 0.0)
def p2g(f: ti.i32):
for p in range(0, n_particles):
new_Jp = Jp[f, p]
base = ti.cast(x[f, p] * inv_dx - 0.5, ti.i32)
fx = x[f, p] * inv_dx - ti.cast(base, ti.f32)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
] # quadratic kernels
new_F = (ti.Matrix.diag(dim=dim, val=1) +
dt * C[f, p]) @ F[f, p] # deformation gradient update
h = max(0.1, min(5, ti.exp(10 * (1.0 - new_Jp))))
if particle_type[p] == 1: # jelly, make it softer
h = 0.3
mu, la = mu_0 * h, lambda_0 * h
if particle_type[p] == 0: # liquid
mu = 0.0
U, sig, V = ti.svd(new_F)
J = 1.0
for d in ti.static(range(2)):
new_sig = sig[d, d]
if particle_type[p] == 2: # Snow
new_sig = min(max(sig[d, d], 1 - 2.5e-2),
1 + 4.5e-3) # Plasticity
new_Jp *= sig[d, d] / new_sig
sig[d, d] = new_sig
J *= new_sig
if particle_type[
p] == 0: # Reset deformation gradient to avoid numerical instability
new_F = ti.Matrix.diag(dim=dim, val=1) * ti.sqrt(J)
elif particle_type[p] == 2:
new_F = U @ sig @ ti.transposed(
V) # Reconstruct elastic deformation gradient after plasticity
F[f + 1, p] = new_F
Jp[f + 1, p] = new_Jp
stress = 2 * mu * (new_F - U @ ti.transposed(V)) @ ti.transposed(
new_F) + ti.Matrix.diag(dim=dim, val=1) * la * J * (J - 1)
stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
affine = stress + p_mass * C[f, p]
# loop over 3x3 node neighborhood
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (ti.cast(ti.Vector([i, j]), ti.f32) - fx) * dx
weight = w[i](0) * w[j](1)
ti.atomic_add(grid_v_in[base + offset],
weight * (p_mass * v[f, p] + affine @ dpos))
ti.atomic_add(grid_m[base + offset], weight * p_mass)
def fdtd(t: ti.i32):
"""
The update formula comes from the paper DiffTaichi
"""
for i in range(n):
for j in range(n):
u[t, i, j] = 2 * u[t - 1, i, j] \
+ (ti.sqr(c) * ti.sqr(dt) + c * alpha * dt) * laplace(t - 1, i, j) \
- u[t - 2, i, j] \
- c * alpha * dt * laplace(t - 2, i, j)
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]
#section 7 of snow paper ... dt being too big
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
new_F = ti.Matrix.identity(ti.f32, self.dim)
new_F[0, 0] = J
self.F[p] = new_F
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 copy(img: ti.ext_arr(), samples:ti.i32):
for i, j in color_buffer:
u = 1.0 * i / res[0]
v = 1.0 * j / res[1]
darken = 1.0 - vignette_strength * max((ti.sqrt(
ti.sqr(u - vignette_center[0]) + ti.sqr(v - vignette_center[1])) -
vignette_radius), 0)
for c in ti.static(range(3)):
img[i, j, c] = ti.sqrt(color_buffer[i, j][c] * darken * exposure / samples)
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(
ti.sqr(u - vignette_center[0]) + ti.sqr(v - vignette_center[1])) -
vignette_radius), 0)
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 p2g():
for p in x:
base = ti.cast(x[p] * inv_dx - 0.5, ti.i32)
fx = x[p] * inv_dx - ti.cast(base, ti.f32)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)]
affine = p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (ti.cast(ti.Vector([i, j]), ti.f32) - fx) * dx
weight = w[i](0) * w[j](1)
grid_v[base + offset].atomic_add(weight * (p_mass * v[p] - x.grad[p] + affine @ dpos))
grid_m[base + offset].atomic_add(weight * p_mass)
def collide(t: ti.i32):
for i in range(n_objects):
hs = halfsize[i]
for k in ti.static(range(4)):
# the corner for collision detection
offset_scale = ti.Vector([k % 2 * 2 - 1, k // 2 % 2 * 2 - 1])
corner_x, corner_v, rela_pos = to_world(t, i, offset_scale * hs)
corner_v = corner_v + dt * gravity * ti.Vector([0.0, 1.0])
# Apply impulse so that there's no sinking
normal = ti.Vector([0.0, 1.0])
tao = ti.Vector([1.0, 0.0])
rn = cross(rela_pos, normal)
rt = cross(rela_pos, tao)
impulse_contribution = inverse_mass[i] + ti.sqr(rn) * \
inverse_inertia[i]
timpulse_contribution = inverse_mass[i] + ti.sqr(rt) * \
inverse_inertia[i]
rela_v_ground = normal.dot(corner_v)
impulse = 0.0
timpulse = 0.0
new_corner_x = corner_x + dt * corner_v
toi = 0.0
if rela_v_ground < 0 and new_corner_x[1] < ground_height:
impulse = -(1 +
elasticity) * rela_v_ground / impulse_contribution
if impulse > 0:
# friction
timpulse = -corner_v.dot(tao) / timpulse_contribution
timpulse = ti.min(friction * impulse,
ti.max(-friction * impulse, timpulse))
if corner_x[1] > ground_height:
toi = -(corner_x[1] - ground_height) / ti.min(
corner_v[1], 1e-3)
apply_impulse(t, i, impulse * normal + timpulse * tao,
new_corner_x, toi)
penalty = 0.0
if new_corner_x[1] < ground_height:
# apply penalty
penalty = -dt * penalty * (
new_corner_x[1] - ground_height) / impulse_contribution
apply_impulse(t, i, penalty * normal, new_corner_x, 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 p2g():
for p in x:
base = ti.cast(x[p] * inv_dx - 0.5, ti.i32)
fx = x[p] * inv_dx - ti.cast(base, ti.f32)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)]
stress = -dt * p_vol * (J[p] - 1) * 4 * inv_dx * inv_dx * E
affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (ti.cast(ti.Vector([i, j]), ti.f32) - fx) * dx
weight = w[i](0) * w[j](1)
grid_v[base + offset].atomic_add(weight * (p_mass * v[p] + affine @ dpos))
grid_m[base + offset].atomic_add(weight * p_mass)
def grid_op():
for i, j in grid_m_in:
inv_m = 1 / (grid_m_in[i, j] + 1e-10)
v_out = inv_m * grid_v_in[i, j]
v_out[1] -= dt * gravity
if i < bound and v_out[0] < 0:
v_out[0] = 0
v_out[1] = 0
if i > n_grid - bound and v_out[0] > 0:
v_out[0] = 0
v_out[1] = 0
if j < bound and v_out[1] < 0:
v_out[0] = 0
v_out[1] = 0
normal = ti.Vector([0.0, 1.0])
lsq = ti.sqr(normal).sum()
if lsq > 0.5:
if ti.static(coeff < 0):
v_out(0).val = 0
v_out(1).val = 0
else:
lin = (ti.transposed(v_out) @ normal)(0)
if lin < 0:
vit = v_out - lin * normal
lit = vit.norm() + 1e-10
if lit + coeff * lin <= 0:
v_out(0).val = 0
v_out(1).val = 0
else:
v_out = (1 + coeff * lin / lit) * vit
if j > n_grid - bound and v_out[1] > 0:
v_out[0] = 0
v_out[1] = 0
grid_v_out[i, j] = v_out
def compute_loss():
for i in range(n_grid):
for j in range(n_grid):
ti.atomic_add(
loss,
ti.sqr(target[i, j] - smoke[steps - 1, i, j]) *
(1 / n_grid**2))
def regress():
for i in x:
v = x[i]
est = 0.0
for j in ti.static(range(number_coeffs)):
est += coeffs[j] * ti.pow(v, j)
loss.atomic_add(0.5 * ti.sqr(y[i] - est))
def compute_loss(view_id: ti.i32):
for i in range(res):
for j in range(res):
ti.atomic_add(
loss,
ti.sqr(images[view_id, i, j] - target_images[view_id, i, j]) *
(1.0 / (res * res)))
def compute_loss(t: ti.i32):
x01 = x[t, 0] - x[t, 1]
x02 = x[t, 0] - x[t, 2]
area = ti.abs(
0.5 * (x01[0] * x02[1] - x01[1] * x02[0])) # area from cross product
target_area = 0.1
loss[None] = ti.sqr(area - target_area)
def compute_loss():
for i in range(n_grid):
for j in range(n_grid):
for k in range(1):
# for t in range(150, steps):
ti.atomic_add(
loss,
ti.sqr(target_img[i, j, k] - refracted_image[i, j, k, steps-1]) * (1 / n_grid**2))
def test_poly(): grad_test(lambda x: x) grad_test(lambda x: -x) grad_test(lambda x: x * x) grad_test(lambda x: ti.sqr(x)) grad_test(lambda x: x * x * x) grad_test(lambda x: x * x * x * x) grad_test(lambda x: 0.4 * x * x - 3) grad_test(lambda x: (x - 3) * (x - 1)) grad_test(lambda x: (x - 3) * (x - 1) + x * x)
def g2p():
for p in x:
base = ti.cast(x[p] * inv_dx - 0.5, ti.i32)
fx = x[p] * inv_dx - ti.cast(base, ti.f32)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)]
new_v = ti.Vector([0.0, 0.0])
new_C = ti.Matrix([[0.0, 0.0], [0.0, 0.0]])
for i in ti.static(range(3)):
for j in ti.static(range(3)):
dpos = ti.cast(ti.Vector([i, j]), ti.f32) - fx
g_v = grid_v[base(0) + i, base(1) + j]
weight = w[i](0) * w[j](1)
new_v += weight * g_v
new_C += 4 * weight * ti.outer_product(g_v, dpos) * inv_dx
v[p] = new_v
x[p] += dt * v[p]
C[p] = new_C
def g2p(f: ti.i32):
for p in range(n_particles):
base = ti.cast(x[f, p] * inv_dx - 0.5, ti.i32)
fx = x[f, p] * inv_dx - ti.cast(base, real)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)]
new_v = ti.Vector([0.0, 0.0])
new_C = ti.Matrix([[0.0, 0.0], [0.0, 0.0]])
for i in ti.static(range(3)):
for j in ti.static(range(3)):
dpos = ti.cast(ti.Vector([i, j]), real) - fx
g_v = grid_v_out[base(0) + i, base(1) + j]
weight = w[i](0) * w[j](1)
new_v += weight * g_v
new_C += 4 * weight * ti.outer_product(g_v, dpos) * inv_dx
v[f + 1, p] = new_v
x[f + 1, p] = x[f, p] + dt * v[f + 1, p]
C[f + 1, p] = new_C
def p2g(f: ti.i32):
for p in range(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=2, val=1) + dt * C[f, p]) @ F[f, p]
J = ti.determinant(new_F)
if particle_type[p] == 0: # fluid
sqrtJ = ti.sqrt(J)
new_F = ti.Matrix([[sqrtJ, 0], [0, sqrtJ]])
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, 1.0]]) * act
cauchy = ti.Matrix([[0.0, 0.0], [0.0, 0.0]])
mass = 0.0
if particle_type[p] == 0:
mass = 4
cauchy = ti.Matrix([[1.0, 0.0], [0.0, 0.1]]) * (J - 1) * E
else:
mass = 1
cauchy = 2 * mu * (new_F - r) @ ti.transposed(new_F) + \
ti.Matrix.diag(2, la * (J - 1) * J)
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)):
offset = ti.Vector([i, j])
dpos = (ti.cast(ti.Vector([i, j]), real) - fx) * dx
weight = w[i](0) * w[j](1)
grid_v_in[base + offset] += weight * (mass * v[f, p] + affine @ dpos)
grid_m_in[base + offset] += weight * mass
def g2p(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)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
]
new_v = ti.Vector.zero(ti.f32, self.dim)
new_C = ti.Matrix.zero(ti.f32, self.dim, self.dim)
# loop over 3x3 grid node neighborhood
for I in ti.static(ti.grouped(self.stencil_range())):
dpos = I.cast(float) - fx
g_v = self.grid_v[base + I]
weight = 1.0
for d in ti.static(range(self.dim)):
weight *= w[I[d]][d]
new_v += weight * g_v
new_C += 4 * self.inv_dx * weight * ti.outer_product(g_v, dpos)
self.v[p], self.C[p] = new_v, new_C
self.x[p] += dt * self.v[p] # advection
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=2, val=1) + dt * C[f, p]) @ F[f, p]
F[f + 1, p] = new_F
J = ti.determinant(new_F)
r, s = ti.polar_decompose(new_F)
cauchy = 2 * mu * (new_F - r) @ ti.transposed(new_F) + \
ti.Matrix.diag(2, la * (J - 1) * J)
stress = -(dt * p_vol * 4 * inv_dx * inv_dx) * cauchy
affine = stress + p_mass * C[f, p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (ti.cast(ti.Vector([i, j]), real) - fx) * dx
weight = w[i](0) * w[j](1)
grid_v_in[f, base + offset] += weight * (
p_mass * v[f, p] + affine @ dpos)
grid_m_in[f, base + offset] += weight * p_mass
def substep():
for i, j in ti.ndrange(n_grid, n_grid):
grid_v[i, j] = [0, 0]
grid_m[i, j] = 0
for p in range(n_particles): # Particle state update and scatter to grid (P2G)
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * 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)]
F[p] = (ti.Matrix.identity(ti.f32, 2) + dt * C[p]) @ F[p] # deformation gradient update
h = ti.exp(10 * (1.0 - Jp[p])) # Hardening coefficient: snow gets harder when compressed
if material[p] == 1: # jelly, make it softer
h = 0.3
mu, la = mu_0 * h, lambda_0 * h
if material[p] == 0: # liquid
mu = 0.0
U, sig, V = ti.svd(F[p])
J = 1.0
for d in ti.static(range(2)):
new_sig = sig[d, d]
if material[p] == 2: # Snow
new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3) # Plasticity
Jp[p] *= sig[d, d] / new_sig
sig[d, d] = new_sig
J *= new_sig
if material[p] == 0: # Reset deformation gradient to avoid numerical instability
F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)
elif material[p] == 2:
F[p] = U @ sig @ V.T() # Reconstruct elastic deformation gradient after plasticity
stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T() + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)
stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
affine = stress + p_mass * C[p]
for i, j in ti.static(ti.ndrange(3, 3)): # Loop over 3x3 grid node neighborhood
offset = ti.Vector([i, j])
dpos = (offset.cast(float) - fx) * dx
weight = w[i][0] * w[j][1]
grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
grid_m[base + offset] += weight * p_mass
for i, j in ti.ndrange(n_grid, n_grid):
if grid_m[i, j] > 0: # No need for epsilon here
grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
grid_v[i, j][1] -= dt * 50 # gravity
if i < 3 and grid_v[i, j][0] < 0: grid_v[i, j][0] = 0 # Boundary conditions
if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
for p in range(n_particles): # grid to particle (G2P)
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0), 0.5 * ti.sqr(fx - 0.5)]
new_v = ti.Vector.zero(ti.f32, 2)
new_C = ti.Matrix.zero(ti.f32, 2, 2)
for i, j in ti.static(ti.ndrange(3, 3)): # loop over 3x3 grid node neighborhood
dpos = ti.Vector([i, j]).cast(float) - fx
g_v = grid_v[base + ti.Vector([i, j])]
weight = w[i][0] * w[j][1]
new_v += weight * g_v
new_C += 4 * inv_dx * weight * ti.outer_product(g_v, dpos)
v[p], C[p] = new_v, new_C
x[p] += dt * v[p] # advection
def substep():
for K in ti.static(range(16)):
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
F[p] = (ti.Matrix.identity(ti.f32, 2) + dt * C[p]) @ F[p]
h = ti.exp(10 * (1.0 - Jp[p]))
if material[p] == 1:
h = 0.3
mu, la = mu_0 * h, lambda_0 * h
if material[p] == 0: # liquid
mu = 0.0
U, sig, V = ti.svd(F[p])
J = 1.0
for d in ti.static(range(2)):
new_sig = sig[d, d]
if material[p] == 2: # Snow
new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)
Jp[p] *= sig[d, d] / new_sig
sig[d, d] = new_sig
J *= new_sig
if material[p] == 0:
F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)
elif material[p] == 2:
F[p] = U @ sig @ V.T()
stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T(
) + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)
stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
affine = stress + p_mass * C[p]
for i, j in ti.static(ti.ndrange(3, 3)):
offset = ti.Vector([i, j])
dpos = (offset.cast(float) - fx) * dx
weight = w[i][0] * w[j][1]
grid_v[base +
offset] += weight * (p_mass * v[p] + affine @ dpos)
grid_m[base + offset] += weight * p_mass
def test_poly():
import time
t = time.time()
grad_test(lambda x: x)
grad_test(lambda x: -x)
grad_test(lambda x: x * x)
grad_test(lambda x: ti.sqr(x))
grad_test(lambda x: x * x * x)
grad_test(lambda x: x * x * x * x)
grad_test(lambda x: 0.4 * x * x - 3)
grad_test(lambda x: (x - 3) * (x - 1))
grad_test(lambda x: (x - 3) * (x - 1) + x * x)
ti.core.print_profile_info()
print('total_time', time.time() - t)
def substep():
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
stress = -dt * p_vol * (J[p] - 1) * 4 * inv_dx * inv_dx * E
affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (offset.cast(float) - fx) * dx
weight = w[i][0] * w[j][1]
grid_v[base + offset].atomic_add(
weight * (p_mass * v[p] + affine @ dpos))
grid_m[base + offset].atomic_add(weight * p_mass)
for i, j in grid_m:
if grid_m[i, j] > 0:
bound = 3
inv_m = 1 / grid_m[i, j]
grid_v[i, j] = inv_m * grid_v[i, j]
grid_v[i, j][1] -= dt * 9.8
if i < bound and grid_v[i, j][0] < 0:
grid_v[i, j][0] = 0
if i > n_grid - bound and grid_v[i, j][0] > 0:
grid_v[i, j][0] = 0
if j < bound and grid_v[i, j][1] < 0:
grid_v[i, j][1] = 0
if j > n_grid - bound and grid_v[i, j][1] > 0:
grid_v[i, j][1] = 0
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
]
new_v = ti.Vector.zero(ti.f32, 2)
new_C = ti.Matrix.zero(ti.f32, 2, 2)
for i in ti.static(range(3)):
for j in ti.static(range(3)):
dpos = ti.Vector([i, j]).cast(float) - fx
g_v = grid_v[base + ti.Vector([i, j])]
weight = w[i][0] * w[j][1]
new_v = new_v + weight * g_v
new_C = new_C + 4 * weight * ti.outer_product(g_v,
dpos) * inv_dx
v[p] = new_v
x[p] += dt * v[p]
J[p] *= 1 + dt * new_C.trace()
C[p] = new_C
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
def compute_loss():
dist = ti.sqr(x_avg - ti.Vector(target))
loss[None] = 0.5 * (dist(0) + dist(1))
def compute_loss(t: ti.i32):
ti.atomic_add(loss[None], dt * ti.sqr(target_v[t][0] - v[t, head_id][0]))
