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]))