diff --git a/.gitignore b/.gitignore index 958272e9a7b4d..95f26b57b141e 100644 --- a/.gitignore +++ b/.gitignore @@ -86,3 +86,4 @@ _build imgui.ini /venv/ /_skbuild/ +.cache diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0dc07c04fff9e..9764922063e38 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -33,4 +33,3 @@ repos: - id: pylint args: ['-rn', '-sn'] files: ^python/taichi/ - exclude: ^python/taichi/examples/.*.py diff --git a/python/taichi/examples/algorithm/marching_squares.py b/python/taichi/examples/algorithm/marching_squares.py index 29ef1e670c198..1de7c6d11ac12 100644 --- a/python/taichi/examples/algorithm/marching_squares.py +++ b/python/taichi/examples/algorithm/marching_squares.py @@ -74,10 +74,14 @@ def march(level: float) -> int: for i, j in ti.ndrange(N - 1, N - 1): case_id = 0 - if pixels[i, j] > level: case_id |= 1 - if pixels[i + 1, j] > level: case_id |= 2 - if pixels[i + 1, j + 1] > level: case_id |= 4 - if pixels[i, j + 1] > level: case_id |= 8 + if pixels[i, j] > level: + case_id |= 1 + if pixels[i + 1, j] > level: + case_id |= 2 + if pixels[i + 1, j + 1] > level: + case_id |= 4 + if pixels[i, j + 1] > level: + case_id |= 8 for k in range(2): if edge_table[case_id, k][0] == -1: @@ -92,16 +96,21 @@ def march(level: float) -> int: return n_edges -level = 0.2 +def main(): + level = 0.2 -gui = ti.GUI('Marching squares') -while gui.running and not gui.get_event(gui.ESCAPE): - touch(*gui.get_cursor_pos(), 0.05) - n_edges = march(level) - edge_coords_np = edge_coords.to_numpy()[:n_edges] / N - gui.set_image(ti.tools.imresize(pixels, *gui.res) / level) - gui.lines(edge_coords_np[:, 0], - edge_coords_np[:, 1], - color=0xff66cc, - radius=1.5) - gui.show() + gui = ti.GUI('Marching squares') + while gui.running and not gui.get_event(gui.ESCAPE): + touch(*gui.get_cursor_pos(), 0.05) + n_edges = march(level) + edge_coords_np = edge_coords.to_numpy()[:n_edges] / N + gui.set_image(ti.tools.imresize(pixels, *gui.res) / level) + gui.lines(edge_coords_np[:, 0], + edge_coords_np[:, 1], + color=0xff66cc, + radius=1.5) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/algorithm/mciso_advanced.py b/python/taichi/examples/algorithm/mciso_advanced.py index fd79563982751..ee556b07b9cf6 100644 --- a/python/taichi/examples/algorithm/mciso_advanced.py +++ b/python/taichi/examples/algorithm/mciso_advanced.py @@ -327,31 +327,43 @@ def march(self) -> int: r_n = 0 for I in ti.grouped(ti.ndrange(*[[1, self.N - 1]] * self.dim)): - id = 0 + idx = 0 if ti.static(self.dim == 2): i, j = I - if self.m[i, j] > 1: id |= 1 - if self.m[i + 1, j] > 1: id |= 2 - if self.m[i, j + 1] > 1: id |= 4 - if self.m[i + 1, j + 1] > 1: id |= 8 + if self.m[i, j] > 1: + idx |= 1 + if self.m[i + 1, j] > 1: + idx |= 2 + if self.m[i, j + 1] > 1: + idx |= 4 + if self.m[i + 1, j + 1] > 1: + idx |= 8 else: i, j, k = I - if self.m[i, j, k] > 1: id |= 1 - if self.m[i + 1, j, k] > 1: id |= 2 - if self.m[i + 1, j + 1, k] > 1: id |= 4 - if self.m[i, j + 1, k] > 1: id |= 8 - if self.m[i, j, k + 1] > 1: id |= 16 - if self.m[i + 1, j, k + 1] > 1: id |= 32 - if self.m[i + 1, j + 1, k + 1] > 1: id |= 64 - if self.m[i, j + 1, k + 1] > 1: id |= 128 + if self.m[i, j, k] > 1: + idx |= 1 + if self.m[i + 1, j, k] > 1: + idx |= 2 + if self.m[i + 1, j + 1, k] > 1: + idx |= 4 + if self.m[i, j + 1, k] > 1: + idx |= 8 + if self.m[i, j, k + 1] > 1: + idx |= 16 + if self.m[i + 1, j, k + 1] > 1: + idx |= 32 + if self.m[i + 1, j + 1, k + 1] > 1: + idx |= 64 + if self.m[i, j + 1, k + 1] > 1: + idx |= 128 for m in range(self.et.shape[1]): - if self.et[id, m][0] == -1: + if self.et[idx, m][0] == -1: break n = ti.atomic_add(r_n, 1) for l in ti.static(range(self.dim)): - e = self.et[id, m][l] + e = self.et[idx, m][l] R = float(I) if ti.static(self.dim == 2): @@ -435,8 +447,9 @@ def march(self) -> int: class MCISO_Example(MCISO): + @staticmethod @ti.func - def gauss(self, x): + def gauss(x): return ti.exp(-6 * x**2) @ti.kernel diff --git a/python/taichi/examples/algorithm/mgpcg.py b/python/taichi/examples/algorithm/mgpcg.py index 7e49057ebc89f..b8b50e38f701a 100644 --- a/python/taichi/examples/algorithm/mgpcg.py +++ b/python/taichi/examples/algorithm/mgpcg.py @@ -27,16 +27,17 @@ Ap = ti.field(dtype=real) # matrix-vector product alpha = ti.field(dtype=real) # step size beta = ti.field(dtype=real) # step size -sum = ti.field(dtype=real) # storage for reductions +sum_ = ti.field(dtype=real) # storage for reductions pixels = ti.field(dtype=real, shape=(N_gui, N_gui)) # image buffer ti.root.pointer(ti.ijk, [N_tot // 4]).dense(ti.ijk, 4).place(x, p, Ap) -for l in range(n_mg_levels): - ti.root.pointer(ti.ijk, [N_tot // (4 * 2**l)]).dense(ti.ijk, - 4).place(r[l], z[l]) +for lvl in range(n_mg_levels): + ti.root.pointer(ti.ijk, + [N_tot // (4 * 2**lvl)]).dense(ti.ijk, + 4).place(r[lvl], z[lvl]) -ti.root.place(alpha, beta, sum) +ti.root.place(alpha, beta, sum_) @ti.kernel @@ -65,9 +66,9 @@ def compute_Ap(): @ti.kernel -def reduce(p: ti.template(), q: ti.template()): - for I in ti.grouped(p): - sum[None] += p[I] * q[I] +def reduce(p_: ti.template(), q_: ti.template()): + for I in ti.grouped(p_): + sum_[None] += p_[I] * q_[I] @ti.kernel @@ -144,69 +145,74 @@ def paint(): pixels[i, j] = x[ii, jj, kk] / N_tot -gui = ti.GUI("mgpcg", res=(N_gui, N_gui)) - -init() - -sum[None] = 0.0 -reduce(r[0], r[0]) -initial_rTr = sum[None] - -# r = b - Ax = b since x = 0 -# p = r = r + 0 p -if use_multigrid: - apply_preconditioner() -else: - z[0].copy_from(r[0]) - -update_p() - -sum[None] = 0.0 -reduce(z[0], r[0]) -old_zTr = sum[None] +def main(): + gui = ti.GUI("mgpcg", res=(N_gui, N_gui)) -# CG -for i in range(400): - # alpha = rTr / pTAp - compute_Ap() - sum[None] = 0.0 - reduce(p, Ap) - pAp = sum[None] - alpha[None] = old_zTr / pAp + init() - # x = x + alpha p - update_x() - - # r = r - alpha Ap - update_r() - - # check for convergence - sum[None] = 0.0 + sum_[None] = 0.0 reduce(r[0], r[0]) - rTr = sum[None] - if rTr < initial_rTr * 1.0e-12: - break + initial_rTr = sum_[None] - # z = M^-1 r + # r = b - Ax = b since x = 0 + # p = r = r + 0 p if use_multigrid: apply_preconditioner() else: z[0].copy_from(r[0]) - # beta = new_rTr / old_rTr - sum[None] = 0.0 - reduce(z[0], r[0]) - new_zTr = sum[None] - beta[None] = new_zTr / old_zTr - - # p = z + beta p update_p() - old_zTr = new_zTr - - print(' ') - print(f'Iter = {i:4}, Residual = {rTr:e}') - paint() - gui.set_image(pixels) - gui.show() -ti.profiler.print_kernel_profiler_info() + sum_[None] = 0.0 + reduce(z[0], r[0]) + old_zTr = sum_[None] + + # CG + for i in range(400): + # alpha = rTr / pTAp + compute_Ap() + sum_[None] = 0.0 + reduce(p, Ap) + pAp = sum_[None] + alpha[None] = old_zTr / pAp + + # x = x + alpha p + update_x() + + # r = r - alpha Ap + update_r() + + # check for convergence + sum_[None] = 0.0 + reduce(r[0], r[0]) + rTr = sum_[None] + if rTr < initial_rTr * 1.0e-12: + break + + # z = M^-1 r + if use_multigrid: + apply_preconditioner() + else: + z[0].copy_from(r[0]) + + # beta = new_rTr / old_rTr + sum_[None] = 0.0 + reduce(z[0], r[0]) + new_zTr = sum_[None] + beta[None] = new_zTr / old_zTr + + # p = z + beta p + update_p() + old_zTr = new_zTr + + print(' ') + print(f'Iter = {i:4}, Residual = {rTr:e}') + paint() + gui.set_image(pixels) + gui.show() + + ti.profiler.print_kernel_profiler_info() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/algorithm/mgpcg_advanced.py b/python/taichi/examples/algorithm/mgpcg_advanced.py index 17f1a037f3120..8393b6a834c15 100644 --- a/python/taichi/examples/algorithm/mgpcg_advanced.py +++ b/python/taichi/examples/algorithm/mgpcg_advanced.py @@ -200,8 +200,8 @@ def solve(self, old_zTr = self.sum[None] # Conjugate gradients - iter = 0 - while max_iters == -1 or iter < max_iters: + it = 0 + while max_iters == -1 or it < max_iters: # self.alpha = rTr / pTAp self.compute_Ap() self.reduce(self.p, self.Ap) @@ -219,7 +219,7 @@ def solve(self, rTr = self.sum[None] if verbose: - print(f'iter {iter}, |residual|_2={math.sqrt(rTr)}') + print(f'iter {it}, |residual|_2={math.sqrt(rTr)}') if rTr < tol: break @@ -239,7 +239,7 @@ def solve(self, self.update_p() old_zTr = new_zTr - iter += 1 + it += 1 class MGPCG_Example(MGPCG): diff --git a/python/taichi/examples/autodiff/diff_sph/diff_sph.py b/python/taichi/examples/autodiff/diff_sph/diff_sph.py index 8be436a3b26ec..8ca102de77e43 100644 --- a/python/taichi/examples/autodiff/diff_sph/diff_sph.py +++ b/python/taichi/examples/autodiff/diff_sph/diff_sph.py @@ -141,10 +141,12 @@ def dump_weights(self, name="save.pkl"): for w in self.parameters(): w = w.to_numpy() w_val.append(w[0]) - pkl.dump(w_val, open(name, "wb")) + with open(name, "wb") as f: + pkl.dump(w_val, f) def load_weights(self, name="save.pkl", model_id=0): - w_val = pkl.load(open(name, 'rb')) + with open(name, 'rb') as f: + w_val = pkl.load(f) self.load_weights_from_value(w_val, model_id) def load_weights_from_value(self, w_val, model_id=0): @@ -153,95 +155,103 @@ def load_weights_from_value(self, w_val, model_id=0): val = val[0] self.copy_from_numpy(w, val, model_id) + @staticmethod @ti.kernel - def copy_from_numpy(self, dst: ti.template(), src: ti.ext_arr(), - model_id: ti.i32): + def copy_from_numpy( + dst: ti.template(), src: ti.ext_arr(), model_id: ti.i32): for I in ti.grouped(src): dst[model_id, I] = src[I] -# NN model -model_num = 1 -steps = 128 -n_input = 3 -n_hidden = 32 -n_output = 16 -n_output_act = 3 -learning_rate = 1e-3 -loss = ti.field(float, shape=(), needs_grad=True) - -if TRAIN: - batch_size = 16 - input_states = ti.field(float, - shape=(model_num, steps, batch_size, n_input), - needs_grad=True) - fc1 = Linear(n_models=model_num, - batch_size=batch_size, - n_steps=steps, - n_input=n_input, - n_hidden=n_hidden, - n_output=n_output, - needs_grad=True, - activation=False) - fc2 = Linear(n_models=model_num, - batch_size=batch_size, - n_steps=steps, - n_input=n_output, - n_hidden=n_hidden, - n_output=n_output_act, - needs_grad=True, - activation=True) - fc1.weights_init() - fc2.weights_init() - NNs = [fc1, fc2] - parameters = [] - for layer in NNs: - parameters.extend(layer.parameters()) - optimizer = SGD(params=parameters, lr=learning_rate) - - # Training data generation - sample_num = batch_size * 25 - x_range = (0.05, 0.45) - y_range = (0.4, 1.0) - z_range = (0.05, 0.45) - - def targets_generation(num, x_range, y_range, z_range): - low = np.array([x_range[0], y_range[0], z_range[0]]) - high = np.array([x_range[1], y_range[1], z_range[1]]) - return np.array( - [np.random.uniform(low=low, high=high) for _ in range(num)]) - - np.random.seed(0) - all_data = targets_generation(sample_num, x_range, y_range, z_range) - training_sample_num = batch_size * 4 - training_data = all_data[:training_sample_num, :] - test_data = all_data[training_sample_num:, :] - print("training data ", training_data.shape, "test data ", test_data.shape) -else: - batch_size = 1 - input_states = ti.field(float, - shape=(model_num, steps, batch_size, n_input), - needs_grad=False) - fc1 = Linear(n_models=model_num, - batch_size=batch_size, - n_steps=steps, - n_input=n_input, - n_hidden=n_hidden, - n_output=n_output, - needs_grad=False, - activation=False) - fc2 = Linear(n_models=model_num, - batch_size=batch_size, - n_steps=steps, - n_input=n_output, - n_hidden=n_hidden, - n_output=n_output_act, - needs_grad=False, - activation=True) - file_dir_path = os.path.dirname(os.path.realpath(__file__)) - fc1.load_weights(f"{file_dir_path}/fc1_pretrained.pkl", model_id=0) - fc2.load_weights(f"{file_dir_path}/fc2_pretrained.pkl", model_id=0) - print(f"Model at {file_dir_path} loaded. ") +def init_nn_model(): + global BATCH_SIZE, steps, input_states, fc1, fc2 + global training_sample_num, training_data, loss + # NN model + model_num = 1 + steps = 128 + n_input = 3 + n_hidden = 32 + n_output = 16 + n_output_act = 3 + learning_rate = 1e-3 + loss = ti.field(float, shape=(), needs_grad=True) + + if TRAIN: + BATCH_SIZE = 16 + input_states = ti.field(float, + shape=(model_num, steps, BATCH_SIZE, n_input), + needs_grad=True) + fc1 = Linear(n_models=model_num, + batch_size=BATCH_SIZE, + n_steps=steps, + n_input=n_input, + n_hidden=n_hidden, + n_output=n_output, + needs_grad=True, + activation=False) + fc2 = Linear(n_models=model_num, + batch_size=BATCH_SIZE, + n_steps=steps, + n_input=n_output, + n_hidden=n_hidden, + n_output=n_output_act, + needs_grad=True, + activation=True) + fc1.weights_init() + fc2.weights_init() + NNs = [fc1, fc2] + parameters = [] + for layer in NNs: + parameters.extend(layer.parameters()) + optimizer = SGD(params=parameters, lr=learning_rate) + + # Training data generation + sample_num = BATCH_SIZE * 25 + x_range = (0.05, 0.45) + y_range = (0.4, 1.0) + z_range = (0.05, 0.45) + + def targets_generation(num, x_range_, y_range_, z_range_): + low = np.array([x_range_[0], y_range_[0], z_range_[0]]) + high = np.array([x_range_[1], y_range_[1], z_range_[1]]) + return np.array( + [np.random.uniform(low=low, high=high) for _ in range(num)]) + + np.random.seed(0) + all_data = targets_generation(sample_num, x_range, y_range, z_range) + training_sample_num = BATCH_SIZE * 4 + training_data = all_data[:training_sample_num, :] + test_data = all_data[training_sample_num:, :] + print("training data ", training_data.shape, "test data ", + test_data.shape) + else: + BATCH_SIZE = 1 + input_states = ti.field(float, + shape=(model_num, steps, BATCH_SIZE, n_input), + needs_grad=False) + fc1 = Linear(n_models=model_num, + batch_size=BATCH_SIZE, + n_steps=steps, + n_input=n_input, + n_hidden=n_hidden, + n_output=n_output, + needs_grad=False, + activation=False) + fc2 = Linear(n_models=model_num, + batch_size=BATCH_SIZE, + n_steps=steps, + n_input=n_output, + n_hidden=n_hidden, + n_output=n_output_act, + needs_grad=False, + activation=True) + file_dir_path = os.path.dirname(os.path.realpath(__file__)) + fc1.load_weights(f"{file_dir_path}/fc1_pretrained.pkl", model_id=0) + fc2.load_weights(f"{file_dir_path}/fc2_pretrained.pkl", model_id=0) + print(f"Model at {file_dir_path} loaded. ") + + +init_nn_model() # Simulation configuration boundary_box_np = np.array([[0, 0, 0], [0.5, 1.5, 0.5]], dtype=dtype_f_np) @@ -249,12 +259,12 @@ def targets_generation(num, x_range, y_range, z_range): target_box_np = np.array([[0.15, 0.90, 0.15], [0.2, 0.95, 0.2]], dtype=dtype_f_np) -target_centers = ti.Vector.field(3, float, shape=batch_size, needs_grad=True) -min_dist = ti.field(float, shape=batch_size, needs_grad=True) -max_dist = ti.field(float, shape=batch_size, needs_grad=True) -max_height = ti.field(float, shape=batch_size, needs_grad=True) -max_left = ti.field(float, shape=batch_size, needs_grad=True) -max_right = ti.field(float, shape=batch_size, needs_grad=True) +target_centers = ti.Vector.field(3, float, shape=BATCH_SIZE, needs_grad=True) +min_dist = ti.field(float, shape=BATCH_SIZE, needs_grad=True) +max_dist = ti.field(float, shape=BATCH_SIZE, needs_grad=True) +max_height = ti.field(float, shape=BATCH_SIZE, needs_grad=True) +max_left = ti.field(float, shape=BATCH_SIZE, needs_grad=True) +max_right = ti.field(float, shape=BATCH_SIZE, needs_grad=True) jet_force_max = ti.Vector([9.81 * 3, 9.81 * 10, 9.81 * 3]) # Simulation parameters @@ -265,19 +275,19 @@ def targets_generation(num, x_range, y_range, z_range): N_target_np = ((target_box_np[1] - target_box_np[0]) / particle_diameter + 1).astype(int) -h = 4.0 * particle_radius +H = 4.0 * particle_radius fluid_particle_num = N_np[0] * N_np[1] * N_np[2] target_particle_num = N_target_np[0] * N_target_np[1] * N_target_np[2] particle_num = fluid_particle_num + target_particle_num print(f"Particle num: {particle_num}") -pos = ti.Vector.field(3, float) -vel = ti.Vector.field(3, float) -acc = ti.Vector.field(3, float) -jet_force = ti.Vector.field(3, - float, - shape=(steps, batch_size), - needs_grad=True) +F_pos = ti.Vector.field(3, float) +F_vel = ti.Vector.field(3, float) +F_acc = ti.Vector.field(3, float) +F_jet_force = ti.Vector.field(3, + float, + shape=(steps, BATCH_SIZE), + needs_grad=True) col = ti.Vector.field(3, float) material = ti.field(int) @@ -286,8 +296,8 @@ def targets_generation(num, x_range, y_range, z_range): pos_vis_buffer = ti.Vector.field(3, float, shape=particle_num) pos_output_buffer = ti.Vector.field(3, float, shape=(steps, particle_num)) -ti.root.dense(ti.ijk, (batch_size, steps, int(particle_num))).place( - pos, vel, acc, den, pre) +ti.root.dense(ti.ijk, (BATCH_SIZE, steps, int(particle_num))).place( + F_pos, F_vel, F_acc, den, pre) ti.root.dense(ti.i, int(particle_num)).place(material, col) ti.root.lazy_grad() @@ -353,25 +363,25 @@ def W_spiky_gradient(R, h): @ti.kernel def initialize_fluid_particle(t: ti.int32, pos: ti.template(), - N_fluid: ti.template()): + N_fluid_: ti.template()): # Allocate fluid - for bs, i in ti.ndrange(batch_size, fluid_particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, fluid_particle_num): pos[bs, t, i] = (ti.Vector([ - int(i % N_fluid[0]), - int(i / N_fluid[0]) % N_fluid[1], - int(i / N_fluid[0] / N_fluid[1] % N_fluid[2]) + int(i % N_fluid_[0]), + int(i / N_fluid_[0]) % N_fluid_[1], + int(i / N_fluid_[0] / N_fluid_[1] % N_fluid_[2]) ]) * particle_diameter + spawn_box[0]) - vel[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_vel[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) material[i] = 0 col[i] = ti.Vector([0.4, 0.7, 1.0]) - acc.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_acc.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) pos.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) - vel.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_vel.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) @ti.kernel def initialize_dists(): - for bs in range(batch_size): + for bs in range(BATCH_SIZE): min_dist[bs] = 1000.0 max_height[bs] = 0.0 max_left[bs] = 0.0 @@ -380,157 +390,156 @@ def initialize_dists(): @ti.kernel def initialize_target_particle(t: ti.int32, pos: ti.template(), - N_target: ti.template(), current_pos: ti.int32): + N_target_: ti.template(), + current_pos: ti.int32): # Allocate target cube for bs, i in ti.ndrange( - batch_size, + BATCH_SIZE, (fluid_particle_num, fluid_particle_num + target_particle_num)): pos[bs, t, i] = (ti.Vector([ - int(i % N_target[0]), - int(i / N_target[0]) % N_target[1], - int(i / N_target[0] / N_target[1] % N_target[2]) + int(i % N_target_[0]), + int(i / N_target_[0]) % N_target_[1], + int(i / N_target_[0] / N_target_[1] % N_target_[2]) ]) * particle_diameter + target_centers[current_pos]) - vel[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_vel[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) material[i] = 1 col[i] = ti.Vector([1.0, 0.65, 0.0]) - acc.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_acc.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) pos.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) - vel.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_vel.grad[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) @ti.kernel def initialize_density(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): den[bs, t, i] = 0.0 @ti.kernel def update_density(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): for j in range(particle_num): - R = pos[bs, t, i] - pos[bs, t, j] - den[bs, t, i] += mass * W(R, h) + R = F_pos[bs, t, i] - F_pos[bs, t, j] + den[bs, t, i] += mass * W(R, H) @ti.kernel def update_pressure(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): pre[bs, t, i] = pressure_scale * max( pow(den[bs, t, i] / rest_density, gamma) - 1, 0) @ti.kernel def controller_output(t: ti.int32): - for bs in range(batch_size): + for bs in range(BATCH_SIZE): for j in ti.static(range(3)): - jet_force[t, bs][j] = fc2.output[0, t, bs, j] * jet_force_max[j] + F_jet_force[t, bs][j] = fc2.output[0, t, bs, j] * jet_force_max[j] @ti.kernel def apply_force(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): if material[i] == 1: - acc[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) + F_acc[bs, t, i] = ti.Vector([0.0, 0.0, 0.0]) else: - if pos[bs, t, i][0] > 0.2 and pos[bs, t, i][0] < 0.3 and pos[ - bs, t, i][1] < 0.2 and pos[bs, t, - i][2] > 0.2 and pos[bs, t, - i][2] < 0.3: + if F_pos[bs, t, i][0] > 0.2 and F_pos[bs, t, i][0] < 0.3 and F_pos[ + bs, t, i][1] < 0.2 and F_pos[bs, t, i][2] > 0.2 and F_pos[ + bs, t, i][2] < 0.3: indicator = (steps - t) // (steps // 2) - acc[bs, t, - i] = jet_force[t, - bs] + gravity + indicator * (-gravity) * 0.1 + F_acc[bs, t, i] = F_jet_force[ + t, bs] + gravity + indicator * (-gravity) * 0.1 else: - acc[bs, t, i] = gravity + F_acc[bs, t, i] = gravity @ti.kernel def update_force(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): for j in range(particle_num): - R = pos[bs, t, i] - pos[bs, t, j] + R = F_pos[bs, t, i] - F_pos[bs, t, j] # Pressure forces - acc[bs, t, i] += -mass * ( + F_acc[bs, t, i] += -mass * ( pre[bs, t, i] / (den[bs, t, i] * den[bs, t, i]) + pre[bs, t, j] / - (den[bs, t, j] * den[bs, t, j])) * W_gradient(R, h) + (den[bs, t, j] * den[bs, t, j])) * W_gradient(R, H) # Viscosity forces - acc[bs, t, i] += viscosity_scale * mass \ - * (vel[bs, t, i] - vel[bs, t, j]).dot(R) / (R.norm(eps) + 0.01 * h * h) / den[bs, t, j] \ - * W_gradient(R, h) + F_acc[bs, t, i] += viscosity_scale * mass \ + * (F_vel[bs, t, i] - F_vel[bs, t, j]).dot(R) / (R.norm(eps) + 0.01 * H * H) / den[bs, t, j] \ + * W_gradient(R, H) @ti.kernel def advance(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): if material[i] == 0: - vel[bs, t, i] = vel[bs, t - 1, i] + acc[bs, t - 1, i] * dt - pos[bs, t, i] = pos[bs, t - 1, i] + vel[bs, t, i] * dt + F_vel[bs, t, i] = F_vel[bs, t - 1, i] + F_acc[bs, t - 1, i] * dt + F_pos[bs, t, i] = F_pos[bs, t - 1, i] + F_vel[bs, t, i] * dt @ti.kernel def boundary_handle(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): collision_normal = ti.Vector([0.0, 0.0, 0.0]) for j in ti.static(range(3)): - if pos[bs, t, i][j] < boundary_box[0][j]: - pos[bs, t, i][j] = boundary_box[0][j] + if F_pos[bs, t, i][j] < boundary_box[0][j]: + F_pos[bs, t, i][j] = boundary_box[0][j] collision_normal[j] += -1.0 for j in ti.static(range(3)): - if pos[bs, t, i][j] > boundary_box[1][j]: - pos[bs, t, i][j] = boundary_box[1][j] + if F_pos[bs, t, i][j] > boundary_box[1][j]: + F_pos[bs, t, i][j] = boundary_box[1][j] collision_normal[j] += 1.0 collision_normal_length = collision_normal.norm() if collision_normal_length > eps: collision_normal /= collision_normal_length - vel[bs, t, i] -= (1.0 + damping) * collision_normal.dot( - vel[bs, t, i]) * collision_normal + F_vel[bs, t, i] -= (1.0 + damping) * collision_normal.dot( + F_vel[bs, t, i]) * collision_normal @ti.kernel def compute_dist(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): + 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 += (pos[bs, t, i][j] - target_centers[bs][j])**2 + 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], pos[bs, t, i][1]) - ti.atomic_max(max_left[bs], pos[bs, t, i][0]) - ti.atomic_max(max_right[bs], pos[bs, t, i][2]) + 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]) @ti.kernel def compute_loss(t: ti.int32): - for bs in range(batch_size): + for bs in range(BATCH_SIZE): max_dist[bs] = ti.sqrt((max_left[bs] - target_centers[bs][0])**2 + (max_right[bs] - target_centers[bs][2])**2 + (max_height[bs] - target_centers[bs][1])**2) - loss[None] += (min_dist[bs] + 0.2 * max_dist[bs]) / batch_size + loss[None] += (min_dist[bs] + 0.2 * max_dist[bs]) / BATCH_SIZE @ti.kernel def copy_back(t: ti.int32): - for bs, i in ti.ndrange(batch_size, particle_num): - pos[bs, 0, i] = pos[bs, t, i] - vel[bs, 0, i] = vel[bs, t, i] - acc[bs, 0, i] = ti.Vector([0.0, 0.0, 0.0]) + for bs, i in ti.ndrange(BATCH_SIZE, particle_num): + F_pos[bs, 0, i] = F_pos[bs, t, i] + F_vel[bs, 0, i] = F_vel[bs, t, i] + F_acc[bs, 0, i] = ti.Vector([0.0, 0.0, 0.0]) @ti.kernel def copy_to_vis(t: ti.int32, bs: ti.int32): for i in range(particle_num): for j in ti.static(range(3)): - pos_vis_buffer[i][j] = pos[bs, t, i][j] + pos_vis_buffer[i][j] = F_pos[bs, t, i][j] @ti.kernel def copy_to_output_buffer(t: ti.int32, bs: ti.int32): for i in range(particle_num): for j in ti.static(range(3)): - pos_output_buffer[t, i][j] = pos[bs, t, i][j] + pos_output_buffer[t, i][j] = F_pos[bs, t, i][j] @ti.kernel @@ -542,7 +551,7 @@ def copy_from_output_to_vis(t: ti.int32): @ti.kernel def fill_target_centers(current_pos: ti.int32, data: ti.any_arr()): - for i in range(current_pos, current_pos + batch_size): + for i in range(current_pos, current_pos + BATCH_SIZE): for j in ti.static(range(3)): target_centers[i][j] = data[i, j] print('target_centers ', target_centers[current_pos]) @@ -550,12 +559,12 @@ def fill_target_centers(current_pos: ti.int32, data: ti.any_arr()): @ti.kernel def fill_input_states(current_pos: ti.int32): - for t, bs in ti.ndrange(steps, (current_pos, current_pos + batch_size)): + for t, bs in ti.ndrange(steps, (current_pos, current_pos + BATCH_SIZE)): for j in ti.static(range(3)): input_states[0, t, bs, j] = target_centers[bs][j] -if __name__ == "__main__": +def main(): show_window = True if TRAIN: show_window = False @@ -578,12 +587,12 @@ def fill_input_states(current_pos: ti.int32): loss_epoch = 0.0 cnt = 0 for current_data_offset in range(0, training_sample_num, - batch_size): + BATCH_SIZE): fill_target_centers(current_data_offset, training_data) fill_input_states(current_data_offset) - initialize_fluid_particle(0, pos, N_fluid) + initialize_fluid_particle(0, F_pos, N_fluid) initialize_dists() - initialize_target_particle(0, pos, N_target, + initialize_target_particle(0, F_pos, N_target, current_data_offset) fc1.clear() fc2.clear() @@ -605,7 +614,7 @@ def fill_input_states(current_pos: ti.int32): compute_loss(steps - 1) optimizer.step() print( - f"current opt progress: {current_data_offset + batch_size}/{training_sample_num}, loss: {loss[None]}" + f"current opt progress: {current_data_offset + BATCH_SIZE}/{training_sample_num}, loss: {loss[None]}" ) losses.append(loss[None]) loss_epoch += loss[None] @@ -644,7 +653,7 @@ def fill_input_states(current_pos: ti.int32): losses, label='loss per iteration') plt.plot([ - i * (training_sample_num // batch_size) + i * (training_sample_num // BATCH_SIZE) for i in range(len(losses_epoch_avg)) ], losses_epoch_avg, @@ -656,7 +665,7 @@ def fill_input_states(current_pos: ti.int32): plt.show() else: current_data_offset = 0 - initialize_fluid_particle(0, pos, N_fluid) + initialize_fluid_particle(0, F_pos, N_fluid) target_centers[current_data_offset][0] = 0.25 target_centers[current_data_offset][1] = 0.50 target_centers[current_data_offset][2] = 0.25 @@ -679,7 +688,7 @@ def fill_input_states(current_pos: ti.int32): if not paused[None]: fill_input_states(current_data_offset) - initialize_target_particle(0, pos, N_target, + initialize_target_particle(0, F_pos, N_target, current_data_offset) fc1.clear() fc2.clear() @@ -723,6 +732,10 @@ def fill_input_states(current_pos: ti.int32): canvas.scene(scene) if INFER_OUTPUT_IMG: if cnt % 2 == 0: - os.makedirs(f"demo_output_interactive/", exist_ok=True) + os.makedirs("demo_output_interactive/", exist_ok=True) window.write_image(f'demo_output_interactive/{cnt:04}.png') window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/autodiff/regression.py b/python/taichi/examples/autodiff/regression.py index c70273e973903..8bdafbbd88c51 100644 --- a/python/taichi/examples/autodiff/regression.py +++ b/python/taichi/examples/autodiff/regression.py @@ -69,8 +69,8 @@ def regress_raw(): regress.grad() print('Loss =', loss[None]) update() - for i in range(number_coeffs): - print(coeffs[i], end=', ') + for j in range(number_coeffs): + print(coeffs[j], end=', ') print() diff --git a/python/taichi/examples/features/gui/fullscreen.py b/python/taichi/examples/features/gui/fullscreen.py index 412747be1868b..77015fb5c1ee0 100644 --- a/python/taichi/examples/features/gui/fullscreen.py +++ b/python/taichi/examples/features/gui/fullscreen.py @@ -1,5 +1,3 @@ -import numpy as np - import taichi as ti ti.init(ti.gpu) diff --git a/python/taichi/examples/features/io/export_mesh.py b/python/taichi/examples/features/io/export_mesh.py index 21b2bfa495ac9..f3fbbefe5d887 100644 --- a/python/taichi/examples/features/io/export_mesh.py +++ b/python/taichi/examples/features/io/export_mesh.py @@ -1,5 +1,3 @@ -import os - import numpy as np import taichi as ti diff --git a/python/taichi/examples/features/io/export_ply.py b/python/taichi/examples/features/io/export_ply.py index d54208a43d688..4430bfc11777c 100644 --- a/python/taichi/examples/features/io/export_ply.py +++ b/python/taichi/examples/features/io/export_ply.py @@ -1,6 +1,3 @@ -import os -import random - import numpy as np import taichi as ti diff --git a/python/taichi/examples/features/io/export_videos.py b/python/taichi/examples/features/io/export_videos.py index 4061d45d6140e..87c4b4c4f9897 100644 --- a/python/taichi/examples/features/io/export_videos.py +++ b/python/taichi/examples/features/io/export_videos.py @@ -11,20 +11,25 @@ def paint(): pixels[i, j, k] = ti.random() * 255 -result_dir = "./results" -video_manager = ti.tools.VideoManager(output_dir=result_dir, - framerate=24, - automatic_build=False) - -for i in range(50): - paint() - - pixels_img = pixels.to_numpy() - video_manager.write_frame(pixels_img) - print(f'\rFrame {i+1}/50 is recorded', end='') - -print() -print('Exporting .mp4 and .gif videos...') -video_manager.make_video(gif=True, mp4=True) -print(f'MP4 video is saved to {video_manager.get_output_filename(".mp4")}') -print(f'GIF video is saved to {video_manager.get_output_filename(".gif")}') +def main(): + result_dir = "./results" + video_manager = ti.tools.VideoManager(output_dir=result_dir, + framerate=24, + automatic_build=False) + + for i in range(50): + paint() + + pixels_img = pixels.to_numpy() + video_manager.write_frame(pixels_img) + print(f'\rFrame {i+1}/50 is recorded', end='') + + print() + print('Exporting .mp4 and .gif videos...') + video_manager.make_video(gif=True, mp4=True) + print(f'MP4 video is saved to {video_manager.get_output_filename(".mp4")}') + print(f'GIF video is saved to {video_manager.get_output_filename(".gif")}') + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/features/sparse/explicit_activation.py b/python/taichi/examples/features/sparse/explicit_activation.py index d7b9af314d037..742c08f4be26a 100644 --- a/python/taichi/examples/features/sparse/explicit_activation.py +++ b/python/taichi/examples/features/sparse/explicit_activation.py @@ -15,7 +15,7 @@ def sparse_api_demo(): ti.activate(block2, [1, 2]) for i, j in x: - print('field x[{}, {}] = {}'.format(i, j, x[i, j])) + print(f'field x[{i}, {j}] = {x[i, j]}') # outputs: # field x[2, 4] = 0 # field x[2, 5] = 0 @@ -23,25 +23,24 @@ def sparse_api_demo(): # field x[3, 5] = 0 for i, j in block2: - print('Active block2: [{}, {}]'.format(i, j)) + print(f'Active block2: [{i}, {j}]') # output: Active block2: [1, 2] for i, j in block1: - print('Active block1: [{}, {}]'.format(i, j)) + print(f'Active block1: [{i}, {j}]') # output: Active block1: [0, 1] for j in range(4): - print('Activity of block2[2, {}] = {}'.format( - j, ti.is_active(block2, [1, j]))) + print(f'Activity of block2[2, {j}] = {ti.is_active(block2, [1, j])}') ti.deactivate(block2, [1, 2]) for i, j in block2: - print('Active block2: [{}, {}]'.format(i, j)) + print(f'Active block2: [{i}, {j}]') # output: nothing for i, j in block1: - print('Active block1: [{}, {}]'.format(i, j)) + print(f'Active block1: [{i}, {j}]') # output: Active block1: [0, 1] print(ti.rescale_index(x, block1, ti.Vector([9, 17]))) diff --git a/python/taichi/examples/features/sparse/taichi_bitmasked.py b/python/taichi/examples/features/sparse/taichi_bitmasked.py index 1437f29b82b5b..eb19b75e5f0f2 100644 --- a/python/taichi/examples/features/sparse/taichi_bitmasked.py +++ b/python/taichi/examples/features/sparse/taichi_bitmasked.py @@ -35,13 +35,18 @@ def paint_all_pixels(color: ti.f32): x[i, j] = color -ti.deactivate_all_snodes() -activate() - -gui = ti.GUI('bitmasked', (n, n)) -for frame in range(10000): - color = math.sin(frame * 0.05) * 0.5 + 0.5 - paint_active_pixels(color) - #paint_all_pixels(color) # try this and compare the difference! - gui.set_image(x) - gui.show() +def main(): + ti.deactivate_all_snodes() + activate() + + gui = ti.GUI('bitmasked', (n, n)) + for frame in range(10000): + color = math.sin(frame * 0.05) * 0.5 + 0.5 + paint_active_pixels(color) + #paint_all_pixels(color) # try this and compare the difference! + gui.set_image(x) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/features/sparse/taichi_dynamic.py b/python/taichi/examples/features/sparse/taichi_dynamic.py index 4dfbae62ad66a..7755f79097efb 100644 --- a/python/taichi/examples/features/sparse/taichi_dynamic.py +++ b/python/taichi/examples/features/sparse/taichi_dynamic.py @@ -18,9 +18,14 @@ def make_lists(): l[i] = ti.length(x.parent(), i) -make_lists() +def main(): + make_lists() -for i in range(n): - assert l[i] == i - for j in range(n): - assert x[i, j] == (j * j if j < i else 0) + for i in range(n): + assert l[i] == i + for j in range(n): + assert x[i, j] == (j * j if j < i else 0) + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/features/sparse/taichi_sparse.py b/python/taichi/examples/features/sparse/taichi_sparse.py index ebd8ea9091a6c..daec52366ecd0 100644 --- a/python/taichi/examples/features/sparse/taichi_sparse.py +++ b/python/taichi/examples/features/sparse/taichi_sparse.py @@ -43,13 +43,18 @@ def paint(): img[scatter(i), scatter(j)] = 1 - t / 4 -img.fill(0.05) +def main(): + img.fill(0.05) -gui = ti.GUI('Sparse Grids', (res, res)) + gui = ti.GUI('Sparse Grids', (res, res)) -for i in range(100000): - block1.deactivate_all() - activate(i * 0.05) - paint() - gui.set_image(img) - gui.show() + for i in range(100000): + block1.deactivate_all() + activate(i * 0.05) + paint() + gui.set_image(img) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/features/sparse/tutorial.py b/python/taichi/examples/features/sparse/tutorial.py index c268e2c7ac745..2b33e241cf00c 100644 --- a/python/taichi/examples/features/sparse/tutorial.py +++ b/python/taichi/examples/features/sparse/tutorial.py @@ -19,11 +19,11 @@ def sparse_struct_for(): x[5, 6] = 3 for i, j in x: - print('field x[{}, {}] = {}'.format(i, j, x[i, j])) + print(f'field {x[i, j]=}') for i, j in block: - print('Active block: [{}, {}]'.format(i, j)) + print(f'Active block: [{i}, {j}]') -print('use_bitmask = {}'.format(use_bitmask)) +print(f'{use_bitmask=}') sparse_struct_for() diff --git a/python/taichi/examples/ggui_examples/fem128_ggui.py b/python/taichi/examples/ggui_examples/fem128_ggui.py index 187befdcfa205..47b102181dcbc 100644 --- a/python/taichi/examples/ggui_examples/fem128_ggui.py +++ b/python/taichi/examples/ggui_examples/fem128_ggui.py @@ -60,11 +60,13 @@ def advance(): disp2 = disp.norm_sqr() if disp2 <= ball_radius**2: NoV = vel[i].dot(disp) - if NoV < 0: vel[i] -= NoV * disp / disp2 + if NoV < 0: + vel[i] -= NoV * disp / disp2 cond = pos[i] < 0 and vel[i] < 0 or pos[i] > 1 and vel[i] > 0 # rect boundary condition: for j in ti.static(range(pos.n)): - if cond[j]: vel[i][j] = 0 + if cond[j]: + vel[i][j] = 0 pos[i] += dt * vel[i] @@ -142,35 +144,40 @@ def render(): canvas.circles(vertexPositions, radius=0.003, color=(1, 0.6, 0.2)) -init_mesh() -init_pos() -gravity[None] = [0, -1] - -print( - "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." -) -while window.running: - for e in window.get_events(ti.ui.PRESS): - if e.key == ti.ui.ESCAPE: - window.running = False - elif e.key == 'r': - init_pos() - elif e.key in ('a', ti.ui.LEFT): - gravity[None] = [-1, 0] - elif e.key in ('d', ti.ui.RIGHT): - gravity[None] = [+1, 0] - elif e.key in ('s', ti.ui.DOWN): - gravity[None] = [0, -1] - elif e.key in ('w', ti.ui.UP): - gravity[None] = [0, +1] - - mouse_pos = window.get_cursor_pos() - attractor_pos[None] = mouse_pos - attractor_strength[None] = window.is_pressed( - ti.ui.LMB) - window.is_pressed(ti.ui.RMB) - for i in range(50): - with ti.Tape(loss=U): - update_U() - advance() - render() - window.show() +def main(): + init_mesh() + init_pos() + gravity[None] = [0, -1] + + print( + "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." + ) + while window.running: + for e in window.get_events(ti.ui.PRESS): + if e.key == ti.ui.ESCAPE: + window.running = False + elif e.key == 'r': + init_pos() + elif e.key in ('a', ti.ui.LEFT): + gravity[None] = [-1, 0] + elif e.key in ('d', ti.ui.RIGHT): + gravity[None] = [+1, 0] + elif e.key in ('s', ti.ui.DOWN): + gravity[None] = [0, -1] + elif e.key in ('w', ti.ui.UP): + gravity[None] = [0, +1] + + mouse_pos = window.get_cursor_pos() + attractor_pos[None] = mouse_pos + attractor_strength[None] = window.is_pressed( + ti.ui.LMB) - window.is_pressed(ti.ui.RMB) + for i in range(50): + with ti.Tape(loss=U): + update_U() + advance() + render() + window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/ggui_examples/fractal3d_ggui.py b/python/taichi/examples/ggui_examples/fractal3d_ggui.py index deedb922703f4..b4de653818d1a 100644 --- a/python/taichi/examples/ggui_examples/fractal3d_ggui.py +++ b/python/taichi/examples/ggui_examples/fractal3d_ggui.py @@ -69,12 +69,12 @@ def compute_sdf(z, c): md2 = 1.0 mz2 = dot(z, z) - for iter in range(iters): + for _ in range(iters): md2 *= max_norm * mz2 z = quat_mul(z, z) + c mz2 = z.dot(z) - if (mz2 > max_norm): + if mz2 > max_norm: break return 0.25 * ti.sqrt(mz2 / md2) * ti.log(mz2) @@ -129,6 +129,7 @@ def __init__(self): @ti.func def shade(self, pos, surface_color, normal, light_pos): + _ = self # make pylint happy light_color = ti.Vector([1, 1, 1]) light_dir = (light_pos - pos).normalized() @@ -184,16 +185,21 @@ def get_image(self, time): return self.image -julia = Julia() +def main(): + julia = Julia() -window = ti.ui.Window("Fractal 3D", image_res, vsync=True) -canvas = window.get_canvas() + window = ti.ui.Window("Fractal 3D", image_res, vsync=True) + canvas = window.get_canvas() -frame_id = 0 + frame_id = 0 -while window.running: - frame_id += 1 + while window.running: + frame_id += 1 - canvas.set_image(julia.get_image(frame_id / 60)) + canvas.set_image(julia.get_image(frame_id / 60)) - window.show() + window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py b/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py index 0f23939d12fb3..72a5d29ff2ecf 100644 --- a/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py +++ b/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py @@ -63,17 +63,23 @@ def initialize_mesh_indices(): initialize_mesh_indices() spring_offsets = [] -if bending_springs: - for i in range(-1, 2): - for j in range(-1, 2): - if (i, j) != (0, 0): - spring_offsets.append(ti.Vector([i, j])) -else: - for i in range(-2, 3): - for j in range(-2, 3): - if (i, j) != (0, 0) and abs(i) + abs(j) <= 2: - spring_offsets.append(ti.Vector([i, j])) + +def initialize_spring_offsets(): + if bending_springs: + for i in range(-1, 2): + for j in range(-1, 2): + if (i, j) != (0, 0): + spring_offsets.append(ti.Vector([i, j])) + + else: + for i in range(-2, 3): + for j in range(-2, 3): + if (i, j) != (0, 0) and abs(i) + abs(j) <= 2: + spring_offsets.append(ti.Vector([i, j])) + + +initialize_spring_offsets() @ti.kernel @@ -90,7 +96,7 @@ def substep(): v_ij = v[i] - v[j] d = x_ij.normalized() current_dist = x_ij.norm() - original_dist = quad_size * float(i - j).norm() + original_dist = quad_size * float(i - j).norm() # pylint: disable=no-member # Spring force force += -spring_Y * d * (current_dist / original_dist - 1) # Dashpot damping @@ -114,43 +120,48 @@ def update_vertices(): vertices[i * n + j] = x[i, j] -window = ti.ui.Window("Taichi Cloth Simulation on GGUI", (1024, 1024), - vsync=True) -canvas = window.get_canvas() -canvas.set_background_color((1, 1, 1)) -scene = ti.ui.Scene() -camera = ti.ui.make_camera() - -current_t = 0.0 -initialize_mass_points() - -while window.running: - if current_t > 1.5: - # Reset - initialize_mass_points() - current_t = 0 - - for i in range(substeps): - substep() - current_t += dt - update_vertices() - - camera.position(0.0, 0.0, 3) - camera.lookat(0.0, 0.0, 0) - scene.set_camera(camera) - - scene.point_light(pos=(0, 1, 2), color=(1, 1, 1)) - scene.ambient_light((0.5, 0.5, 0.5)) - scene.mesh(vertices, - indices=indices, - per_vertex_color=colors, - two_sided=True) - - # Draw a smaller ball to avoid visual penetration - scene.particles(ball_center, - radius=ball_radius * 0.95, - color=(0.5, 0.42, 0.8)) - canvas.scene(scene) - window.show() - -#TODO: include self-collision handling +def main(): + window = ti.ui.Window("Taichi Cloth Simulation on GGUI", (1024, 1024), + vsync=True) + canvas = window.get_canvas() + canvas.set_background_color((1, 1, 1)) + scene = ti.ui.Scene() + camera = ti.ui.make_camera() + + current_t = 0.0 + initialize_mass_points() + + while window.running: + if current_t > 1.5: + # Reset + initialize_mass_points() + current_t = 0 + + for i in range(substeps): + substep() + current_t += dt + update_vertices() + + camera.position(0.0, 0.0, 3) + camera.lookat(0.0, 0.0, 0) + scene.set_camera(camera) + + scene.point_light(pos=(0, 1, 2), color=(1, 1, 1)) + scene.ambient_light((0.5, 0.5, 0.5)) + scene.mesh(vertices, + indices=indices, + per_vertex_color=colors, + two_sided=True) + + # Draw a smaller ball to avoid visual penetration + scene.particles(ball_center, + radius=ball_radius * 0.95, + color=(0.5, 0.42, 0.8)) + canvas.scene(scene) + window.show() + + #TODO: include self-collision handling + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/ggui_examples/mpm128_ggui.py b/python/taichi/examples/ggui_examples/mpm128_ggui.py index 7c1eb5a6345f8..71c34ee38279b 100644 --- a/python/taichi/examples/ggui_examples/mpm128_ggui.py +++ b/python/taichi/examples/ggui_examples/mpm128_ggui.py @@ -1,5 +1,3 @@ -import numpy as np - import taichi as ti arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda @@ -47,10 +45,10 @@ def substep(): 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 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - F[p] = (ti.Matrix.identity(float, 2) + - dt * C[p]) @ F[p] # deformation gradient update - h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p]))) - ) # Hardening coefficient: snow gets harder when compressed + # deformation gradient update + F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p] + # Hardening coefficient: snow gets harder when compressed + h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p])))) if material[p] == 1: # jelly, make it softer h = 0.3 mu, la = mu_0 * h, lambda_0 * h @@ -66,18 +64,18 @@ def substep(): 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 + if material[p] == 0: + # Reset deformation gradient to avoid numerical instability F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J) elif material[p] == 2: - F[p] = U @ sig @ V.transpose( - ) # Reconstruct elastic deformation gradient after plasticity + # Reconstruct elastic deformation gradient after plasticity + F[p] = U @ sig @ V.transpose() stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose( ) + ti.Matrix.identity(float, 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 + # Loop over 3x3 grid node neighborhood + 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] @@ -94,17 +92,20 @@ def substep(): 0.01 + dist.norm()) * attractor_strength[None] * dt * 100 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 + 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 x: # grid to particle (G2P) base = (x[p] * inv_dx - 0.5).cast(int) fx = x[p] * inv_dx - base.cast(float) w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2] new_v = ti.Vector.zero(float, 2) new_C = ti.Matrix.zero(float, 2, 2) - for i, j in ti.static(ti.ndrange( - 3, 3)): # loop over 3x3 grid node neighborhood + # loop over 3x3 grid node neighborhood + for i, j in ti.static(ti.ndrange(3, 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] @@ -136,42 +137,54 @@ def render(): snow[i] = x[i + 2 * group_size] -print( - "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." -) - -res = (512, 512) -window = ti.ui.Window("Taichi MLS-MPM-128", res=res, vsync=True) -canvas = window.get_canvas() -radius = 0.003 - -reset() -gravity[None] = [0, -1] - -while window.running: - if window.get_event(ti.ui.PRESS): - if window.event.key == 'r': reset() - elif window.event.key in [ti.ui.ESCAPE]: break - if window.event is not None: gravity[None] = [0, 0] # if had any event - if window.is_pressed(ti.ui.LEFT, 'a'): gravity[None][0] = -1 - if window.is_pressed(ti.ui.RIGHT, 'd'): gravity[None][0] = 1 - if window.is_pressed(ti.ui.UP, 'w'): gravity[None][1] = 1 - if window.is_pressed(ti.ui.DOWN, 's'): gravity[None][1] = -1 - mouse = window.get_cursor_pos() - mouse_circle[0] = ti.Vector([mouse[0], mouse[1]]) - canvas.circles(mouse_circle, color=(0.2, 0.4, 0.6), radius=0.05) - attractor_pos[None] = [mouse[0], mouse[1]] - attractor_strength[None] = 0 - if window.is_pressed(ti.ui.LMB): - attractor_strength[None] = 1 - if window.is_pressed(ti.ui.RMB): - attractor_strength[None] = -1 - - for s in range(int(2e-3 // dt)): - substep() - render() - canvas.set_background_color((0.067, 0.184, 0.255)) - canvas.circles(water, radius=radius, color=(0, 0.5, 0.5)) - canvas.circles(jelly, radius=radius, color=(0.93, 0.33, 0.23)) - canvas.circles(snow, radius=radius, color=(1, 1, 1)) - window.show() +def main(): + print( + "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." + ) + + res = (512, 512) + window = ti.ui.Window("Taichi MLS-MPM-128", res=res, vsync=True) + canvas = window.get_canvas() + radius = 0.003 + + reset() + gravity[None] = [0, -1] + + while window.running: + if window.get_event(ti.ui.PRESS): + if window.event.key == 'r': + reset() + elif window.event.key in [ti.ui.ESCAPE]: + break + if window.event is not None: + gravity[None] = [0, 0] # if had any event + if window.is_pressed(ti.ui.LEFT, 'a'): + gravity[None][0] = -1 + if window.is_pressed(ti.ui.RIGHT, 'd'): + gravity[None][0] = 1 + if window.is_pressed(ti.ui.UP, 'w'): + gravity[None][1] = 1 + if window.is_pressed(ti.ui.DOWN, 's'): + gravity[None][1] = -1 + mouse = window.get_cursor_pos() + mouse_circle[0] = ti.Vector([mouse[0], mouse[1]]) + canvas.circles(mouse_circle, color=(0.2, 0.4, 0.6), radius=0.05) + attractor_pos[None] = [mouse[0], mouse[1]] + attractor_strength[None] = 0 + if window.is_pressed(ti.ui.LMB): + attractor_strength[None] = 1 + if window.is_pressed(ti.ui.RMB): + attractor_strength[None] = -1 + + for s in range(int(2e-3 // dt)): + substep() + render() + canvas.set_background_color((0.067, 0.184, 0.255)) + canvas.circles(water, radius=radius, color=(0, 0.5, 0.5)) + canvas.circles(jelly, radius=radius, color=(0.93, 0.33, 0.23)) + canvas.circles(snow, radius=radius, color=(1, 1, 1)) + window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/ggui_examples/mpm3d_ggui.py b/python/taichi/examples/ggui_examples/mpm3d_ggui.py index 3c4edab01f639..6b302f4b8df2b 100644 --- a/python/taichi/examples/ggui_examples/mpm3d_ggui.py +++ b/python/taichi/examples/ggui_examples/mpm3d_ggui.py @@ -20,28 +20,26 @@ p_rho = 1 p_vol = (dx * 0.5)**2 p_mass = p_vol * p_rho -g_x = 0 -g_y = -9.8 -g_z = 0 +GRAVITY = [0, -9.8, 0] bound = 3 E = 1000 # Young's modulus nu = 0.2 # Poisson's ratio mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ( (1 + nu) * (1 - 2 * nu)) # Lame parameters -x = ti.Vector.field(dim, float, n_particles) -v = ti.Vector.field(dim, float, n_particles) -C = ti.Matrix.field(dim, dim, float, n_particles) -F = ti.Matrix.field(3, 3, dtype=float, - shape=n_particles) # deformation gradient -Jp = ti.field(float, n_particles) +F_x = ti.Vector.field(dim, float, n_particles) +F_v = ti.Vector.field(dim, float, n_particles) +F_C = ti.Matrix.field(dim, dim, float, n_particles) +F_dg = ti.Matrix.field(3, 3, dtype=float, + shape=n_particles) # deformation gradient +F_Jp = ti.field(float, n_particles) -colors = ti.Vector.field(4, float, n_particles) -colors_random = ti.Vector.field(4, float, n_particles) -materials = ti.field(int, n_particles) -grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim) -grid_m = ti.field(float, (n_grid, ) * dim) -used = ti.field(int, n_particles) +F_colors = ti.Vector.field(4, float, n_particles) +F_colors_random = ti.Vector.field(4, float, n_particles) +F_materials = ti.field(int, n_particles) +F_grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim) +F_grid_m = ti.field(float, (n_grid, ) * dim) +F_used = ti.field(int, n_particles) neighbour = (3, ) * dim @@ -52,89 +50,87 @@ @ti.kernel def substep(g_x: float, g_y: float, g_z: float): - for I in ti.grouped(grid_m): - grid_v[I] = ti.zero(grid_v[I]) - grid_m[I] = 0 + for I in ti.grouped(F_grid_m): + F_grid_v[I] = ti.zero(F_grid_v[I]) + F_grid_m[I] = 0 ti.loop_config(block_dim=n_grid) - for p in x: - if used[p] == 0: + for p in F_x: + if F_used[p] == 0: continue - Xp = x[p] / dx + Xp = F_x[p] / dx base = int(Xp - 0.5) fx = Xp - base w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - F[p] = (ti.Matrix.identity(float, 3) + - dt * C[p]) @ F[p] # deformation gradient update - - h = ti.exp( - 10 * - (1.0 - - Jp[p])) # Hardening coefficient: snow gets harder when compressed - if materials[p] == JELLY: # jelly, make it softer + F_dg[p] = (ti.Matrix.identity(float, 3) + + dt * F_C[p]) @ F_dg[p] # deformation gradient update + # Hardening coefficient: snow gets harder when compressed + h = ti.exp(10 * (1.0 - F_Jp[p])) + if F_materials[p] == JELLY: # jelly, make it softer h = 0.3 mu, la = mu_0 * h, lambda_0 * h - if materials[p] == WATER: # liquid + if F_materials[p] == WATER: # liquid mu = 0.0 - U, sig, V = ti.svd(F[p]) + U, sig, V = ti.svd(F_dg[p]) J = 1.0 for d in ti.static(range(3)): new_sig = sig[d, d] - if materials[p] == SNOW: # Snow + if F_materials[p] == SNOW: # Snow new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3) # Plasticity - Jp[p] *= sig[d, d] / new_sig + F_Jp[p] *= sig[d, d] / new_sig sig[d, d] = new_sig J *= new_sig - if materials[ - p] == WATER: # Reset deformation gradient to avoid numerical instability + if F_materials[p] == WATER: + # Reset deformation gradient to avoid numerical instability new_F = ti.Matrix.identity(float, 3) new_F[0, 0] = J - F[p] = new_F - elif materials[p] == SNOW: - F[p] = U @ sig @ V.transpose( - ) # Reconstruct elastic deformation gradient after plasticity - stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose( + F_dg[p] = new_F + elif F_materials[p] == SNOW: + # Reconstruct elastic deformation gradient after plasticity + F_dg[p] = U @ sig @ V.transpose() + stress = 2 * mu * (F_dg[p] - U @ V.transpose()) @ F_dg[p].transpose( ) + ti.Matrix.identity(float, 3) * la * J * (J - 1) stress = (-dt * p_vol * 4) * stress / dx**2 - affine = stress + p_mass * C[p] + affine = stress + p_mass * F_C[p] for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))): dpos = (offset - fx) * dx weight = 1.0 for i in ti.static(range(dim)): weight *= w[offset[i]][i] - grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos) - grid_m[base + offset] += weight * p_mass - for I in ti.grouped(grid_m): - if grid_m[I] > 0: - grid_v[I] /= grid_m[I] - grid_v[I] += dt * ti.Vector([g_x, g_y, g_z]) - cond = (I < bound) & (grid_v[I] < 0) | \ - (I > n_grid - bound) & (grid_v[I] > 0) - grid_v[I] = 0 if cond else grid_v[I] + F_grid_v[base + + offset] += weight * (p_mass * F_v[p] + affine @ dpos) + F_grid_m[base + offset] += weight * p_mass + for I in ti.grouped(F_grid_m): + if F_grid_m[I] > 0: + F_grid_v[I] /= F_grid_m[I] + F_grid_v[I] += dt * ti.Vector([g_x, g_y, g_z]) + cond = (I < bound) & (F_grid_v[I] < 0) | \ + (I > n_grid - bound) & (F_grid_v[I] > 0) + F_grid_v[I] = 0 if cond else F_grid_v[I] ti.loop_config(block_dim=n_grid) - for p in x: - if used[p] == 0: + for p in F_x: + if F_used[p] == 0: continue - Xp = x[p] / dx + Xp = F_x[p] / dx base = int(Xp - 0.5) fx = Xp - base w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - new_v = ti.zero(v[p]) - new_C = ti.zero(C[p]) + new_v = ti.zero(F_v[p]) + new_C = ti.zero(F_C[p]) for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))): dpos = (offset - fx) * dx weight = 1.0 for i in ti.static(range(dim)): weight *= w[offset[i]][i] - g_v = grid_v[base + offset] + g_v = F_grid_v[base + offset] new_v += weight * g_v new_C += 4 * weight * g_v.outer_product(dpos) / dx**2 - v[p] = new_v - x[p] += dt * v[p] - C[p] = new_C + F_v[p] = new_v + F_x[p] += dt * F_v[p] + F_C[p] = new_C class CubeVolume: @@ -150,28 +146,28 @@ def init_cube_vol(first_par: int, last_par: int, x_begin: float, y_begin: float, z_begin: float, x_size: float, y_size: float, z_size: float, material: int): for i in range(first_par, last_par): - x[i] = ti.Vector([ti.random() for i in range(dim)]) * ti.Vector( + F_x[i] = ti.Vector([ti.random() for i in range(dim)]) * ti.Vector( [x_size, y_size, z_size]) + ti.Vector([x_begin, y_begin, z_begin]) - Jp[i] = 1 - F[i] = ti.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) - v[i] = ti.Vector([0.0, 0.0, 0.0]) - materials[i] = material - colors_random[i] = ti.Vector( + F_Jp[i] = 1 + F_dg[i] = ti.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + F_v[i] = ti.Vector([0.0, 0.0, 0.0]) + F_materials[i] = material + F_colors_random[i] = ti.Vector( [ti.random(), ti.random(), ti.random(), ti.random()]) - used[i] = 1 + F_used[i] = 1 @ti.kernel def set_all_unused(): - for p in used: - used[p] = 0 + for p in F_used: + F_used[p] = 0 # basically throw them away so they aren't rendered - x[p] = ti.Vector([533799.0, 533799.0, 533799.0]) - Jp[p] = 1 - F[p] = ti.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) - C[p] = ti.Matrix([[0, 0, 0], [0, 0, 0], [0, 0, 0]]) - v[p] = ti.Vector([0.0, 0.0, 0.0]) + F_x[p] = ti.Vector([533799.0, 533799.0, 533799.0]) + F_Jp[p] = 1 + F_dg[p] = ti.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + F_C[p] = ti.Matrix([[0, 0, 0], [0, 0, 0], [0, 0, 0]]) + F_v[p] = ti.Vector([0.0, 0.0, 0.0]) def init_vols(vols): @@ -181,7 +177,7 @@ def init_vols(vols): total_vol += v.volume next_p = 0 - for i in range(len(vols)): + for i, v in enumerate(vols): v = vols[i] if isinstance(v, CubeVolume): par_count = int(v.volume / total_vol * n_particles) @@ -197,13 +193,11 @@ def init_vols(vols): @ti.kernel -def set_color_by_material(material_colors: ti.types.ndarray()): +def set_color_by_material(mat_color: ti.types.ndarray()): for i in range(n_particles): - mat = materials[i] - colors[i] = ti.Vector([ - material_colors[mat, 0], material_colors[mat, 1], - material_colors[mat, 2], 1.0 - ]) + mat = F_materials[i] + F_colors[i] = ti.Vector( + [mat_color[mat, 0], mat_color[mat, 1], mat_color[mat, 2], 1.0]) print("Loading presets...this might take a minute") @@ -252,7 +246,6 @@ def init(): res = (1920, 1080) window = ti.ui.Window("Real MPM 3D", res, vsync=True) -frame_id = 0 canvas = window.get_canvas() scene = ti.ui.Scene() camera = ti.ui.make_camera() @@ -266,7 +259,6 @@ def show_options(): global paused global particles_radius global curr_preset_id - global g_x, g_y, g_z window.GUI.begin("Presets", 0.05, 0.1, 0.2, 0.15) old_preset = curr_preset_id @@ -279,9 +271,9 @@ def show_options(): window.GUI.end() window.GUI.begin("Gravity", 0.05, 0.3, 0.2, 0.1) - g_x = window.GUI.slider_float("x", g_x, -10, 10) - g_y = window.GUI.slider_float("y", g_y, -10, 10) - g_z = window.GUI.slider_float("z", g_z, -10, 10) + GRAVITY[0] = window.GUI.slider_float("x", GRAVITY[0], -10, 10) + GRAVITY[1] = window.GUI.slider_float("y", GRAVITY[1], -10, 10) + GRAVITY[2] = window.GUI.slider_float("z", GRAVITY[2], -10, 10) window.GUI.end() window.GUI.begin("Options", 0.05, 0.45, 0.2, 0.4) @@ -315,8 +307,8 @@ def render(): scene.ambient_light((0, 0, 0)) - colors_used = colors_random if use_random_colors else colors - scene.particles(x, per_vertex_color=colors_used, radius=particles_radius) + colors_used = F_colors_random if use_random_colors else F_colors + scene.particles(F_x, per_vertex_color=colors_used, radius=particles_radius) scene.point_light(pos=(0.5, 1.5, 0.5), color=(0.5, 0.5, 0.5)) scene.point_light(pos=(0.5, 1.5, 1.5), color=(0.5, 0.5, 0.5)) @@ -324,17 +316,22 @@ def render(): canvas.scene(scene) -while window.running: - #print("heyyy ",frame_id) - frame_id += 1 - frame_id = frame_id % 256 +def main(): + frame_id = 0 + + while window.running: + #print("heyyy ",frame_id) + frame_id += 1 + frame_id = frame_id % 256 - if not paused: - for s in range(steps): - substep(g_x, g_y, g_z) + if not paused: + for _ in range(steps): + substep(*GRAVITY) - render() + render() + show_options() + window.show() - show_options() - window.show() +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/ggui_examples/stable_fluid_ggui.py b/python/taichi/examples/ggui_examples/stable_fluid_ggui.py index 9ccd083c91264..5b6487e8b90dd 100644 --- a/python/taichi/examples/ggui_examples/stable_fluid_ggui.py +++ b/python/taichi/examples/ggui_examples/stable_fluid_ggui.py @@ -19,7 +19,6 @@ force_radius = res / 2.0 gravity = True debug = False -paused = False arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda ti.init(arch=arch) @@ -78,13 +77,13 @@ def bilerp(vf, p): # 3rd order Runge-Kutta @ti.func -def backtrace(vf: ti.template(), p, dt: ti.template()): +def backtrace(vf: ti.template(), p, dt_: ti.template()): v1 = bilerp(vf, p) - p1 = p - 0.5 * dt * v1 + p1 = p - 0.5 * dt_ * v1 v2 = bilerp(vf, p1) - p2 = p - 0.75 * dt * v2 + p2 = p - 0.75 * dt_ * v2 v3 = bilerp(vf, p2) - p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) + p -= dt_ * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) return p @@ -214,7 +213,7 @@ def step(mouse_data): print(f'divergence={div_s}') -class MouseDataGen(object): +class MouseDataGen: def __init__(self): self.prev_mouse = None self.prev_color = None @@ -249,33 +248,41 @@ def reset(): dyes_pair.cur.fill(0) -window = ti.ui.Window('Stable Fluid', (res, res), vsync=True) -canvas = window.get_canvas() -md_gen = MouseDataGen() - -while window.running: - if window.get_event(ti.ui.PRESS): - e = window.event - if e.key == ti.ui.ESCAPE: - break - elif e.key == 'r': - paused = False - reset() - elif e.key == 's': - if curl_strength: - curl_strength = 0 - else: - curl_strength = 7 - elif e.key == 'p': - paused = not paused - elif e.key == 'd': - debug = not debug - - # Debug divergence: - # print(max((abs(velocity_divs.to_numpy().reshape(-1))))) - - if not paused: - mouse_data = md_gen(window) - step(mouse_data) - canvas.set_image(dyes_pair.cur) - window.show() +def main(): + global curl_strength, debug + + paused = False + window = ti.ui.Window('Stable Fluid', (res, res), vsync=True) + canvas = window.get_canvas() + md_gen = MouseDataGen() + + while window.running: + if window.get_event(ti.ui.PRESS): + e = window.event + if e.key == ti.ui.ESCAPE: + break + elif e.key == 'r': + paused = False + reset() + elif e.key == 's': + if curl_strength: + curl_strength = 0 + else: + curl_strength = 7 + elif e.key == 'p': + paused = not paused + elif e.key == 'd': + debug = not debug + + # Debug divergence: + # print(max((abs(velocity_divs.to_numpy().reshape(-1))))) + + if not paused: + mouse_data = md_gen(window) + step(mouse_data) + canvas.set_image(dyes_pair.cur) + window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/graph/mpm88_graph.py b/python/taichi/examples/graph/mpm88_graph.py index c59263d7ffb71..a8e8913d8e3c0 100644 --- a/python/taichi/examples/graph/mpm88_graph.py +++ b/python/taichi/examples/graph/mpm88_graph.py @@ -1,7 +1,5 @@ import argparse -import numpy as np - import taichi as ti ti.init(arch=ti.vulkan) @@ -97,15 +95,16 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1), J[i] = 1 -x = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) -v = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) +F_x = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) +F_v = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) + +F_C = ti.Matrix.ndarray(2, 2, ti.f32, shape=(n_particles)) +F_J = ti.ndarray(ti.f32, shape=(n_particles)) +F_grid_v = ti.Vector.ndarray(2, ti.f32, shape=(n_grid, n_grid)) +F_grid_m = ti.ndarray(ti.f32, shape=(n_grid, n_grid)) -C = ti.Matrix.ndarray(2, 2, ti.f32, shape=(n_particles)) -J = ti.ndarray(ti.f32, shape=(n_particles)) -grid_v = ti.Vector.ndarray(2, ti.f32, shape=(n_grid, n_grid)) -grid_m = ti.ndarray(ti.f32, shape=(n_grid, n_grid)) -if __name__ == "__main__": +def main(): parser = argparse.ArgumentParser() parser.add_argument('--baseline', action='store_true') args, unknown = parser.parse_known_args() @@ -151,30 +150,34 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1), g_update = g_update_builder.compile() # Run - g_init.run({'x': x, 'v': v, 'J': J}) + g_init.run({'x': F_x, 'v': F_v, 'J': F_J}) gui = ti.GUI('MPM88') while gui.running: g_update.run({ - 'x': x, - 'v': v, - 'C': C, - 'J': J, - 'grid_v': grid_v, - 'grid_m': grid_m + 'x': F_x, + 'v': F_v, + 'C': F_C, + 'J': F_J, + 'grid_v': F_grid_v, + 'grid_m': F_grid_m }) gui.clear(0x112F41) - gui.circles(x.to_numpy(), radius=1.5, color=0x068587) + gui.circles(F_x.to_numpy(), radius=1.5, color=0x068587) # false+, pylint: disable=no-member gui.show() else: - init_particles(x, v, J) + init_particles(F_x, F_v, F_J) gui = ti.GUI('MPM88') while gui.running and not gui.get_event(gui.ESCAPE): for s in range(N_ITER): - substep_reset_grid(grid_v, grid_m) - substep_p2g(x, v, C, J, grid_v, grid_m) - substep_update_grid_v(grid_v, grid_m) - substep_g2p(x, v, C, J, grid_v) + substep_reset_grid(F_grid_v, F_grid_m) + substep_p2g(F_x, F_v, F_C, F_J, F_grid_v, F_grid_m) + substep_update_grid_v(F_grid_v, F_grid_m) + substep_g2p(F_x, F_v, F_C, F_J, F_grid_v) gui.clear(0x112F41) - gui.circles(x.to_numpy(), radius=1.5, color=0x068587) + gui.circles(F_x.to_numpy(), radius=1.5, color=0x068587) # false+, pylint: disable=no-member gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/graph/stable_fluid_graph.py b/python/taichi/examples/graph/stable_fluid_graph.py index 9a33ac5342fe2..ef896f8cc83f4 100644 --- a/python/taichi/examples/graph/stable_fluid_graph.py +++ b/python/taichi/examples/graph/stable_fluid_graph.py @@ -21,8 +21,6 @@ maxfps = 60 dye_decay = 1 - 1 / (maxfps * time_c) force_radius = res / 2.0 -gravity = True -paused = False class TexPair: @@ -64,13 +62,13 @@ def bilerp(vf: ti.template(), p): # 3rd order Runge-Kutta @ti.func -def backtrace(vf: ti.template(), p, dt: ti.template()): +def backtrace(vf: ti.template(), p, dt_: ti.template()): v1 = bilerp(vf, p) - p1 = p - 0.5 * dt * v1 + p1 = p - 0.5 * dt_ * v1 v2 = bilerp(vf, p1) - p2 = p - 0.75 * dt * v2 + p2 = p - 0.75 * dt_ * v2 v3 = bilerp(vf, p2) - p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) + p -= dt_ * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) return p @@ -180,7 +178,7 @@ def step_orig(mouse_data): mouse_data_ti = ti.ndarray(ti.f32, shape=(8, )) -class MouseDataGen(object): +class MouseDataGen: def __init__(self): self.prev_mouse = None self.prev_color = None @@ -206,7 +204,7 @@ def __call__(self, gui): else: self.prev_mouse = None self.prev_color = None - mouse_data_ti.from_numpy(mouse_data) + mouse_data_ti.from_numpy(mouse_data) # false+, pylint: disable=no-member return mouse_data_ti @@ -216,10 +214,14 @@ def reset(): dyes_pair.cur.fill(0) -if __name__ == "__main__": +def main(): + global velocities_pair, pressures_pair, dyes_pair, curl_strength + + paused = False + parser = argparse.ArgumentParser() parser.add_argument('--baseline', action='store_true') - args, unknown = parser.parse_known_args() + args, _ = parser.parse_known_args() gui = ti.GUI('Stable Fluid', (res, res)) md_gen = MouseDataGen() @@ -227,7 +229,6 @@ def reset(): _velocities = ti.Vector.ndarray(2, float, shape=(res, res)) _new_velocities = ti.Vector.ndarray(2, float, shape=(res, res)) _velocity_divs = ti.ndarray(float, shape=(res, res)) - velocity_curls = ti.ndarray(float, shape=(res, res)) _pressures = ti.ndarray(float, shape=(res, res)) _new_pressures = ti.ndarray(float, shape=(res, res)) _dye_buffer = ti.Vector.ndarray(3, float, shape=(res, res)) @@ -314,8 +315,6 @@ def reset(): curl_strength = 0 else: curl_strength = 7 - elif e.key == 'g': - gravity = not gravity elif e.key == 'p': paused = not paused @@ -337,10 +336,14 @@ def reset(): } if swap: g1.run(invoke_args) - gui.set_image(_dye_buffer.to_numpy()) + gui.set_image(_dye_buffer.to_numpy()) # false+, pylint: disable=no-member swap = False else: g2.run(invoke_args) - gui.set_image(_new_dye_buffer.to_numpy()) + gui.set_image(_new_dye_buffer.to_numpy()) # false+, pylint: disable=no-member swap = True gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/rendering/cornell_box.py b/python/taichi/examples/rendering/cornell_box.py index cc100a8af9b01..b7d2a08aa8a75 100644 --- a/python/taichi/examples/rendering/cornell_box.py +++ b/python/taichi/examples/rendering/cornell_box.py @@ -1,7 +1,6 @@ import time import numpy as np -from numpy.lib.function_base import average import taichi as ti @@ -65,9 +64,9 @@ def make_box_transform_matrices(): # left box -box_min = ti.Vector([0.0, 0.0, 0.0]) -box_max = ti.Vector([0.55, 1.1, 0.55]) -box_m_inv, box_m_inv_t = make_box_transform_matrices() +BOX_MIN = ti.Vector([0.0, 0.0, 0.0]) +BOX_MAX = ti.Vector([0.55, 1.1, 0.55]) +BOX_M_INV, BOX_M_INV_T = make_box_transform_matrices() @ti.func @@ -194,13 +193,13 @@ def intersect_aabb(box_min, box_max, o, d): @ti.func def intersect_aabb_transformed(box_min, box_max, o, d): # Transform the ray to the box's local space - obj_o = mat_mul_point(box_m_inv, o) - obj_d = mat_mul_vec(box_m_inv, d) + obj_o = mat_mul_point(BOX_M_INV, o) + obj_d = mat_mul_vec(BOX_M_INV, d) intersect, near_t, _, near_norm = intersect_aabb(box_min, box_max, obj_o, obj_d) if intersect and 0 < near_t: # Transform the normal in the box's local space to world space - near_norm = mat_mul_vec(box_m_inv_t, near_norm) + near_norm = mat_mul_vec(BOX_M_INV_T, near_norm) else: intersect = 0 return intersect, near_t, near_norm @@ -230,7 +229,7 @@ def intersect_scene(pos, ray_dir): normal = (hit_pos - sp1_center).normalized() c, mat = ti.Vector([1.0, 1.0, 1.0]), mat_glass # left box - hit, cur_dist, pnorm = intersect_aabb_transformed(box_min, box_max, pos, + hit, cur_dist, pnorm = intersect_aabb_transformed(BOX_MIN, BOX_MAX, pos, ray_dir) if hit and 0 < cur_dist < closest: closest = cur_dist @@ -501,8 +500,9 @@ def main(): interval = 10 if i % interval == 0: tonemap(i) - print("{:.2f} samples/s ({} iters)".format( - interval / (time.time() - last_t), i)) + print( + f"{interval / (time.time() - last_t):.2f} samples/s ({i} iters)" + ) last_t = time.time() gui.set_image(tonemapped_buffer) gui.show() diff --git a/python/taichi/examples/rendering/rasterizer.py b/python/taichi/examples/rendering/rasterizer.py index 4325bae93a121..2cab0c1fc5fc2 100644 --- a/python/taichi/examples/rendering/rasterizer.py +++ b/python/taichi/examples/rendering/rasterizer.py @@ -107,25 +107,30 @@ def rasterize(self): pixels[i, j] = samples_sum / num_samples_per_pixel -gui = ti.GUI("Rasterizer", res=(width, height)) - -triangles = TriangleRasterizer(num_triangles) - -i = 0 -while gui.running: - # Set a triangle to a new random triangle - triangles.set_triangle(i % num_triangles, - ti.Vector(np.random.rand(2) * [width, height]), - ti.Vector(np.random.rand(2) * [width, height]), - ti.Vector(np.random.rand(2) * [width, height]), - ti.Vector(np.random.rand(3)), - ti.Vector(np.random.rand(3)), - ti.Vector(np.random.rand(3))) - i = i + 1 - - samples.fill(ti.Vector([1.0, 1.0, 1.0])) - pixels.fill(ti.Vector([1.0, 1.0, 1.0])) - triangles.tile_culling() - triangles.rasterize() - gui.set_image(pixels) - gui.show() +def main(): + gui = ti.GUI("Rasterizer", res=(width, height)) + + triangles = TriangleRasterizer(num_triangles) + + i = 0 + while gui.running: + # Set a triangle to a new random triangle + triangles.set_triangle(i % num_triangles, + ti.Vector(np.random.rand(2) * [width, height]), + ti.Vector(np.random.rand(2) * [width, height]), + ti.Vector(np.random.rand(2) * [width, height]), + ti.Vector(np.random.rand(3)), + ti.Vector(np.random.rand(3)), + ti.Vector(np.random.rand(3))) + i = i + 1 + + samples.fill(ti.Vector([1.0, 1.0, 1.0])) + pixels.fill(ti.Vector([1.0, 1.0, 1.0])) + triangles.tile_culling() + triangles.rasterize() + gui.set_image(pixels) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/rendering/sdf_renderer.py b/python/taichi/examples/rendering/sdf_renderer.py index 284e00d4b586c..5ac787391909d 100644 --- a/python/taichi/examples/rendering/sdf_renderer.py +++ b/python/taichi/examples/rendering/sdf_renderer.py @@ -151,15 +151,20 @@ def render(): color_buffer[u, v] += throughput * hit_light -gui = ti.GUI('SDF Path Tracer', res) -last_t = 0 -for i in range(50000): - render() - interval = 10 - if i % interval == 0 and i > 0: - print("{:.2f} samples/s".format(interval / (time.time() - last_t))) - last_t = time.time() - img = color_buffer.to_numpy() * (1 / (i + 1)) - img = img / img.mean() * 0.24 - gui.set_image(np.sqrt(img)) - gui.show() +def main(): + gui = ti.GUI('SDF Path Tracer', res) + last_t = 0 + for i in range(50000): + render() + interval = 10 + if i % interval == 0 and i > 0: + print(f"{interval / (time.time() - last_t):.2f} samples/s") + last_t = time.time() + img = color_buffer.to_numpy() * (1 / (i + 1)) + img = img / img.mean() * 0.24 + gui.set_image(np.sqrt(img)) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/rendering/simple_texture.py b/python/taichi/examples/rendering/simple_texture.py index 0394c53d5cf93..5f42ea74ba4c3 100644 --- a/python/taichi/examples/rendering/simple_texture.py +++ b/python/taichi/examples/rendering/simple_texture.py @@ -8,7 +8,7 @@ pixels = ti.Vector.field(3, dtype=float, shape=res) tex_format = ti.u8 -tex = ti.Texture(tex_format, 1, (128, 128)) +texture = ti.Texture(tex_format, 1, (128, 128)) tex_ndarray = ti.ndarray(tex_format, shape=(128, 128)) @@ -21,7 +21,7 @@ def make_texture(arr: ti.types.ndarray()): make_texture(tex_ndarray) -tex.from_ndarray(tex_ndarray) +texture.from_ndarray(tex_ndarray) @ti.kernel @@ -39,12 +39,18 @@ def paint(t: ti.f32, tex: ti.types.texture()): pixels[i, j] = [c.r, c.r, c.r] -window = ti.ui.Window('UV', res) -canvas = window.get_canvas() +def main(): -t = 0.0 -while window.running: - paint(t, tex) - canvas.set_image(pixels) - window.show() - t += 0.03 + window = ti.ui.Window('UV', res) + canvas = window.get_canvas() + + t = 0.0 + while window.running: + paint(t, texture) + canvas.set_image(pixels) + window.show() + t += 0.03 + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/comet.py b/python/taichi/examples/simulation/comet.py index ea89ab6e55c08..40d9f509a64ed 100644 --- a/python/taichi/examples/simulation/comet.py +++ b/python/taichi/examples/simulation/comet.py @@ -81,19 +81,24 @@ def render(): img[p] += color[i] -inv_m[0] = 0 -x[0].x = +0.5 -x[0].y = -0.01 -v[0].x = +0.6 -v[0].y = +0.4 -color[0] = 1 - -gui = ti.GUI('Comet', res) -while gui.running: - gui.running = not gui.get_event(gui.ESCAPE) - generate() - for s in range(steps): - substep() - render() - gui.set_image(img) - gui.show() +def main(): + inv_m[0] = 0 + x[0].x = +0.5 + x[0].y = -0.01 + v[0].x = +0.6 + v[0].y = +0.4 + color[0] = 1 + + gui = ti.GUI('Comet', res) + while gui.running: + gui.running = not gui.get_event(gui.ESCAPE) + generate() + for s in range(steps): + substep() + render() + gui.set_image(img) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/euler.py b/python/taichi/examples/simulation/euler.py index f59b8da013c2f..421ebdd39e5ce 100644 --- a/python/taichi/examples/simulation/euler.py +++ b/python/taichi/examples/simulation/euler.py @@ -1,4 +1,4 @@ -import matplotlib.cm as cm +from matplotlib import cm import taichi as ti @@ -213,19 +213,19 @@ def HLLC_flux(qL, qR, n): # HLLC flux. HLLC = ti.Vector([0.0, 0.0, 0.0, 0.0]) - if (0 <= sL): + if 0 <= sL: HLLC = fL - elif (0 <= sM): + elif 0 <= sM: qsL = rL * (sL-vnL)/(sL-sM) \ * ti.Vector([1.0, sM*nx-vtL*ny,sM*ny+vtL*nx, \ qL[3]/rL + (sM-vnL)*(sM+pL/(rL*(sL-vnL)))]) HLLC = fL + sL * (qsL - qL) - elif (0 <= sR): + elif 0 <= sR: qsR = rR * (sR-vnR)/(sR-sM) \ * ti.Vector([1.0, sM*nx-vtR*ny,sM*ny+vtR*nx, \ qR[3]/rR + (sM-vnR)*(sM+pR/(rR*(sR-vnR)))]) HLLC = fR + sR * (qsR - qR) - elif (0 >= sR): + elif 0 >= sR: HLLC = fR return HLLC @@ -430,39 +430,44 @@ def paint(): 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 + 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]) + 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) - - -gui = ti.GUI('Euler Equations', (res, res)) -cmap = cm.get_cmap(cmap_name) -set_ic() -set_bc() - -n = 0 -while gui.running: - calc_dt() - copy_to_old() - for rk_step in range(2): - compute_W() - if method == 0: - compute_F_muscl() - else: - compute_F_thinc() - update_Q(rk_step) - set_bc() - - if n % 10 == 0: - paint() - gui.set_image(cmap(img.to_numpy())) - gui.show() - n += 1 + min_ = fixed_caxis[0] + max_ = fixed_caxis[1] + img[i, j] = (img[i, j] - min_) / (max_ - min_) + + +def main(): + gui = ti.GUI('Euler Equations', (res, res)) + cmap = cm.get_cmap(cmap_name) + set_ic() + set_bc() + + n = 0 + while gui.running: + calc_dt() + copy_to_old() + for rk_step in range(2): + compute_W() + if method == 0: + compute_F_muscl() + else: + compute_F_thinc() + update_Q(rk_step) + set_bc() + + if n % 10 == 0: + paint() + gui.set_image(cmap(img.to_numpy())) + gui.show() + n += 1 + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/fem128.py b/python/taichi/examples/simulation/fem128.py index 09062dff1cbac..f3d37145ee461 100644 --- a/python/taichi/examples/simulation/fem128.py +++ b/python/taichi/examples/simulation/fem128.py @@ -59,11 +59,13 @@ def advance(): disp2 = disp.norm_sqr() if disp2 <= ball_radius**2: NoV = vel[i].dot(disp) - if NoV < 0: vel[i] -= NoV * disp / disp2 + if NoV < 0: + vel[i] -= NoV * disp / disp2 cond = (pos[i] < 0) & (vel[i] < 0) | (pos[i] > 1) & (vel[i] > 0) # rect boundary condition: for j in ti.static(range(pos.n)): - if cond[j]: vel[i][j] = 0 + if cond[j]: + vel[i][j] = 0 pos[i] += dt * vel[i] @@ -102,38 +104,43 @@ def paint_phi(gui): gui.triangles(a, b, c, color=ti.rgb_to_hex([k + gb, gb, gb])) -init_mesh() -init_pos() -gravity[None] = [0, -1] - -gui = ti.GUI('FEM128') -print( - "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." -) -while gui.running: - for e in gui.get_events(gui.PRESS): - if e.key == gui.ESCAPE: - gui.running = False - elif e.key == 'r': - init_pos() - elif e.key in ('a', gui.LEFT): - gravity[None] = [-1, 0] - elif e.key in ('d', gui.RIGHT): - gravity[None] = [+1, 0] - elif e.key in ('s', gui.DOWN): - gravity[None] = [0, -1] - elif e.key in ('w', gui.UP): - gravity[None] = [0, +1] - mouse_pos = gui.get_cursor_pos() - attractor_pos[None] = mouse_pos - attractor_strength[None] = gui.is_pressed(gui.LMB) - gui.is_pressed( - gui.RMB) - for i in range(50): - with ti.Tape(loss=U): - update_U() - advance() - paint_phi(gui) - gui.circle(mouse_pos, radius=15, color=0x336699) - gui.circle(ball_pos, radius=ball_radius * 512, color=0x666666) - gui.circles(pos.to_numpy(), radius=2, color=0xffaa33) - gui.show() +def main(): + init_mesh() + init_pos() + gravity[None] = [0, -1] + + gui = ti.GUI('FEM128') + print( + "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." + ) + while gui.running: + for e in gui.get_events(gui.PRESS): + if e.key == gui.ESCAPE: + gui.running = False + elif e.key == 'r': + init_pos() + elif e.key in ('a', gui.LEFT): + gravity[None] = [-1, 0] + elif e.key in ('d', gui.RIGHT): + gravity[None] = [+1, 0] + elif e.key in ('s', gui.DOWN): + gravity[None] = [0, -1] + elif e.key in ('w', gui.UP): + gravity[None] = [0, +1] + mouse_pos = gui.get_cursor_pos() + attractor_pos[None] = mouse_pos + attractor_strength[None] = gui.is_pressed(gui.LMB) - gui.is_pressed( + gui.RMB) + for i in range(50): + with ti.Tape(loss=U): + update_U() + advance() + paint_phi(gui) + gui.circle(mouse_pos, radius=15, color=0x336699) + gui.circle(ball_pos, radius=ball_radius * 512, color=0x666666) + gui.circles(pos.to_numpy(), radius=2, color=0xffaa33) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/fem99.py b/python/taichi/examples/simulation/fem99.py index 4683096178288..64d25cf12c25d 100644 --- a/python/taichi/examples/simulation/fem99.py +++ b/python/taichi/examples/simulation/fem99.py @@ -54,11 +54,13 @@ def advance(): disp2 = disp.norm_sqr() if disp2 <= ball_radius**2: NoV = vel[i].dot(disp) - if NoV < 0: vel[i] -= NoV * disp / disp2 + if NoV < 0: + vel[i] -= NoV * disp / disp2 # rect boundary condition: cond = (pos[i] < 0) & (vel[i] < 0) | (pos[i] > 1) & (vel[i] > 0) for j in ti.static(range(pos.n)): - if cond[j]: vel[i][j] = 0 + if cond[j]: + vel[i][j] = 0 pos[i] += dt * vel[i] @@ -87,19 +89,24 @@ def init_mesh(): f2v[k + 1] = [c, d, a] -init_mesh() -init_pos() -gui = ti.GUI('FEM99') -while gui.running: - for e in gui.get_events(): - if e.key == gui.ESCAPE: - gui.running = False - elif e.key == 'r': - init_pos() - for i in range(30): - with ti.Tape(loss=U): - update_U() - advance() - gui.circles(pos.to_numpy(), radius=2, color=0xffaa33) - gui.circle(ball_pos, radius=ball_radius * 512, color=0x666666) - gui.show() +def main(): + init_mesh() + init_pos() + gui = ti.GUI('FEM99') + while gui.running: + for e in gui.get_events(): + if e.key == gui.ESCAPE: + gui.running = False + elif e.key == 'r': + init_pos() + for i in range(30): + with ti.Tape(loss=U): + update_U() + advance() + gui.circles(pos.to_numpy(), radius=2, color=0xffaa33) + gui.circle(ball_pos, radius=ball_radius * 512, color=0x666666) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/fractal.py b/python/taichi/examples/simulation/fractal.py index 0ebba64a2dce1..304b80cdd9240 100644 --- a/python/taichi/examples/simulation/fractal.py +++ b/python/taichi/examples/simulation/fractal.py @@ -23,9 +23,14 @@ def paint(t: float): pixels[i, j] = 1 - iterations * 0.02 -gui = ti.GUI("Julia Set", res=(n * 2, n)) +def main(): + gui = ti.GUI("Julia Set", res=(n * 2, n)) -for i in range(1000000): - paint(i * 0.03) - gui.set_image(pixels) - gui.show() + for i in range(1000000): + paint(i * 0.03) + gui.set_image(pixels) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/implicit_fem.py b/python/taichi/examples/simulation/implicit_fem.py index ad5c5c6c77637..8d8520dbb3b31 100644 --- a/python/taichi/examples/simulation/implicit_fem.py +++ b/python/taichi/examples/simulation/implicit_fem.py @@ -35,18 +35,18 @@ n_cells = 5 * np.product(n_cube - 1) dx = 1 / (n_cube.max() - 1) -vertices = ti.Vector.field(4, dtype=ti.i32, shape=n_cells) +F_vertices = ti.Vector.field(4, dtype=ti.i32, shape=n_cells) -x = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) -ox = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) -v = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) -f = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) -mul_ans = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) -m = ti.field(dtype=ti.f32, shape=n_verts) +F_x = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) +F_ox = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) +F_v = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) +F_f = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) +F_mul_ans = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts) +F_m = ti.field(dtype=ti.f32, shape=n_verts) n_cells = (n_cube - 1).prod() * 5 -B = ti.Matrix.field(args.dim, args.dim, dtype=ti.f32, shape=n_cells) -W = ti.field(dtype=ti.f32, shape=n_cells) +F_B = ti.Matrix.field(args.dim, args.dim, dtype=ti.f32, shape=n_cells) +F_W = ti.field(dtype=ti.f32, shape=n_cells) @ti.func @@ -57,7 +57,8 @@ def i2p(I): @ti.func def set_element(e, I, verts): for i in ti.static(range(args.dim + 1)): - vertices[e][i] = i2p(I + (([verts[i] >> k for k in range(3)] ^ I) & 1)) + F_vertices[e][i] = i2p(I + (([verts[i] >> k + for k in range(3)] ^ I) & 1)) @ti.kernel @@ -72,12 +73,12 @@ def get_vertices(): set_element(e + i, I, (j, j ^ 1, j ^ 2, j ^ 4)) set_element(e + 4, I, (1, 2, 4, 7)) for I in ti.grouped(ti.ndrange(*(n_cube))): - ox[i2p(I)] = I * dx + F_ox[i2p(I)] = I * dx @ti.func def Ds(verts): - return ti.Matrix.cols([x[verts[i]] - x[verts[3]] for i in range(3)]) + return ti.Matrix.cols([F_x[verts[i]] - F_x[verts[3]] for i in range(3)]) @ti.func @@ -96,33 +97,33 @@ def ssvd(F): @ti.func def get_force_func(c, verts): - F = Ds(verts) @ B[c] + F = Ds(verts) @ F_B[c] P = ti.Matrix.zero(ti.f32, 3, 3) U, sig, V = ssvd(F) P = 2 * mu * (F - U @ V.transpose()) - H = -W[c] * P @ B[c].transpose() + H = -F_W[c] * P @ F_B[c].transpose() for i in ti.static(range(3)): force = ti.Vector([H[j, i] for j in range(3)]) - f[verts[i]] += force - f[verts[3]] -= force + F_f[verts[i]] += force + F_f[verts[3]] -= force @ti.kernel def get_force(): - for c in vertices: - get_force_func(c, vertices[c]) - for u in f: - f[u].y -= 9.8 * m[u] + for c in F_vertices: + get_force_func(c, F_vertices[c]) + for u in F_f: + F_f[u].y -= 9.8 * F_m[u] @ti.kernel def matmul_cell(ret: ti.template(), vel: ti.template()): for i in ret: - ret[i] = vel[i] * m[i] - for c in vertices: - verts = vertices[c] - W_c = W[c] - B_c = B[c] + ret[i] = vel[i] * F_m[i] + for c in F_vertices: + verts = F_vertices[c] + W_c = F_W[c] + B_c = F_B[c] for u in range(4): for d in range(3): dD = ti.Matrix.zero(ti.f32, 3, 3) @@ -154,80 +155,81 @@ def dot(a: ti.template(), b: ti.template()) -> ti.f32: return ans -b = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) -r0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) -p0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) +F_b = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) +F_r0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) +F_p0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts) @ti.kernel def get_b(): - for i in b: - b[i] = m[i] * v[i] + dt * f[i] + for i in F_b: + F_b[i] = F_m[i] * F_v[i] + dt * F_f[i] def cg(): def mul(x): - matmul_cell(mul_ans, x) - return mul_ans + matmul_cell(F_mul_ans, x) + return F_mul_ans get_force() get_b() - mul(v) - add(r0, b, -1, mul(v)) + mul(F_v) + add(F_r0, F_b, -1, mul(F_v)) - d = p0 - d.copy_from(r0) - r_2 = dot(r0, r0) + d = F_p0 + d.copy_from(F_r0) + r_2 = dot(F_r0, F_r0) n_iter = 50 epsilon = 1e-6 r_2_init = r_2 r_2_new = r_2 - for iter in range(n_iter): + for _ in range(n_iter): q = mul(d) alpha = r_2_new / dot(d, q) - add(v, v, alpha, d) - add(r0, r0, -alpha, q) + add(F_v, F_v, alpha, d) + add(F_r0, F_r0, -alpha, q) r_2 = r_2_new - r_2_new = dot(r0, r0) - if r_2_new <= r_2_init * epsilon**2: break + r_2_new = dot(F_r0, F_r0) + if r_2_new <= r_2_init * epsilon**2: + break beta = r_2_new / r_2 - add(d, r0, beta, d) - f.fill(0) - add(x, x, dt, v) + add(d, F_r0, beta, d) + F_f.fill(0) + add(F_x, F_x, dt, F_v) @ti.kernel def advect(): - for p in x: - v[p] += dt * (f[p] / m[p]) - x[p] += dt * v[p] - f[p] = ti.Vector([0, 0, 0]) + for p in F_x: + F_v[p] += dt * (F_f[p] / F_m[p]) + F_x[p] += dt * F_v[p] + F_f[p] = ti.Vector([0, 0, 0]) @ti.kernel def init(): - for u in x: - x[u] = ox[u] - v[u] = [0.0] * 3 - f[u] = [0.0] * 3 - m[u] = 0.0 - for c in vertices: - F = Ds(vertices[c]) - B[c] = F.inverse() - W[c] = ti.abs(F.determinant()) / 6 + for u in F_x: + F_x[u] = F_ox[u] + F_v[u] = [0.0] * 3 + F_f[u] = [0.0] * 3 + F_m[u] = 0.0 + for c in F_vertices: + F = Ds(F_vertices[c]) + F_B[c] = F.inverse() + F_W[c] = ti.abs(F.determinant()) / 6 for i in range(4): - m[vertices[c][i]] += W[c] / 4 * density - for u in x: - x[u].y += 1.0 + F_m[F_vertices[c][i]] += F_W[c] / 4 * density + for u in F_x: + F_x[u].y += 1.0 @ti.kernel def floor_bound(): - for u in x: - if x[u].y < 0: - x[u].y = 0 - if v[u].y < 0: - v[u].y = 0 + for u in F_x: + if F_x[u].y < 0: + F_x[u].y = 0 + if F_v[u].y < 0: + F_v[u].y = 0 @ti.func @@ -237,30 +239,36 @@ def check(u): for i in ti.static(range(3)): k = rest % n_cube[2 - i] rest = rest // n_cube[2 - i] - if k == 0: ans |= (1 << (i * 2)) - if k == n_cube[2 - i] - 1: ans |= (1 << (i * 2 + 1)) + if k == 0: + ans |= (1 << (i * 2)) + if k == n_cube[2 - i] - 1: + ans |= (1 << (i * 2 + 1)) return ans -su = 0 -for i in range(3): - su += (n_cube[i] - 1) * (n_cube[(i + 1) % 3] - 1) -indices = ti.field(ti.i32, shape=2 * su * 2 * 3) +def gen_indices(): + su = 0 + for i in range(3): + su += (n_cube[i] - 1) * (n_cube[(i + 1) % 3] - 1) + return ti.field(ti.i32, shape=2 * su * 2 * 3) + + +indices = gen_indices() @ti.kernel def get_indices(): # calculate all the meshes on surface cnt = 0 - for c in vertices: + for c in F_vertices: if c % 5 != 4: for i in ti.static([0, 2, 3]): - verts = [vertices[c][(i + j) % 4] for j in range(3)] - sum = check(verts[0]) & check(verts[1]) & check(verts[2]) - if sum: + verts = [F_vertices[c][(i + j) % 4] for j in range(3)] + sum_ = check(verts[0]) & check(verts[1]) & check(verts[2]) + if sum_: m = ti.atomic_add(cnt, 1) det = ti.Matrix.rows([ - x[verts[i]] - [0.5, 1.5, 0.5] for i in range(3) + F_x[verts[i]] - [0.5, 1.5, 0.5] for i in range(3) ]).determinant() if det < 0: tmp = verts[1] @@ -282,7 +290,7 @@ def substep(): floor_bound() -if __name__ == '__main__': +def main(): get_vertices() init() get_indices() @@ -310,7 +318,7 @@ def render(): scene.point_light(pos=(0.5, 10.0, 0.5), color=(0.5, 0.5, 0.5)) scene.point_light(pos=(10.0, 10.0, 10.0), color=(0.5, 0.5, 0.5)) - scene.mesh(x, indices, color=(0.73, 0.33, 0.23)) + scene.mesh(F_x, indices, color=(0.73, 0.33, 0.23)) canvas.scene(scene) @@ -345,9 +353,14 @@ def T(a): while gui.running: substep() if gui.get_event(ti.GUI.PRESS): - if gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break + if gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: + break if gui.is_pressed('r'): init() gui.clear(0x000000) - gui.circles(T(x.to_numpy() / 3), radius=1.5, color=0xba543a) + gui.circles(T(F_x.to_numpy() / 3), radius=1.5, color=0xba543a) gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/implicit_mass_spring.py b/python/taichi/examples/simulation/implicit_mass_spring.py index 1ed7e3616dbd7..a9cfc743d1907 100644 --- a/python/taichi/examples/simulation/implicit_mass_spring.py +++ b/python/taichi/examples/simulation/implicit_mass_spring.py @@ -217,7 +217,7 @@ def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)): canvas.circles(self.pos, radius, color) -if __name__ == "__main__": +def main(): ti.init(arch=ti.cpu) h = 0.01 cloth = Cloth(N=5) @@ -261,3 +261,7 @@ def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)): canvas = window.get_canvas() cloth.displayGGUI(canvas) window.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/inital_value_problem.py b/python/taichi/examples/simulation/initial_value_problem.py similarity index 60% rename from python/taichi/examples/simulation/inital_value_problem.py rename to python/taichi/examples/simulation/initial_value_problem.py index 289c09c7a8e50..058c628ec465e 100644 --- a/python/taichi/examples/simulation/inital_value_problem.py +++ b/python/taichi/examples/simulation/initial_value_problem.py @@ -1,4 +1,3 @@ -import math import time import numpy as np @@ -31,17 +30,21 @@ def paint(t: float): y = locations[i, 1] dirs[i, 0] = ti.sin((t * x - y)) dirs[i, 1] = ti.cos(t * y - x) - len = (dirs[i, 0]**2 + dirs[i, 1]**2)**0.5 - dirs[i, 0] /= len * 40 - dirs[i, 1] /= len * 40 + l = (dirs[i, 0]**2 + dirs[i, 1]**2)**0.5 + dirs[i, 0] /= l * 40 + dirs[i, 1] /= l * 40 -gui = ti.GUI("Vector Field", res=(500, 500)) +def main(): + gui = ti.GUI("Vector Field", res=(500, 500)) -begining = time.time_ns() -for k in range(1000000): - start_time = time.time_ns() - paint((time.time_ns() - begining) * 0.00000001) - dirs_np = dirs.to_numpy() - gui.arrows(locations_np, dirs_np, radius=1) - gui.show() + begining = time.time_ns() + for k in range(1000000): + paint((time.time_ns() - begining) * 0.00000001) + dirs_np = dirs.to_numpy() + gui.arrows(locations_np, dirs_np, radius=1) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/mandelbrot_zoom.py b/python/taichi/examples/simulation/mandelbrot_zoom.py index 90ef7acb77083..c7f966fb78bbf 100644 --- a/python/taichi/examples/simulation/mandelbrot_zoom.py +++ b/python/taichi/examples/simulation/mandelbrot_zoom.py @@ -1,5 +1,5 @@ import taichi as ti -from taichi.math import * +from taichi.math import cmul, dot, log2, vec2, vec3 ti.init(arch=ti.gpu) @@ -43,8 +43,13 @@ def render(time: ti.f32): pixels[i, j] = setcolor(z, count) -gui = ti.GUI("Mandelbrot set zoom", res=(width, height)) -for i in range(100000): - render(i * 0.03) - gui.set_image(pixels) - gui.show() +def main(): + gui = ti.GUI("Mandelbrot set zoom", res=(width, height)) + for i in range(100000): + render(i * 0.03) + gui.set_image(pixels) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/mass_spring_game.py b/python/taichi/examples/simulation/mass_spring_game.py index 774dee46846ce..d4fdc9393ff37 100644 --- a/python/taichi/examples/simulation/mass_spring_game.py +++ b/python/taichi/examples/simulation/mass_spring_game.py @@ -163,10 +163,10 @@ def main(): gui.text( content= - f'Left click: add mass point (with shift to fix); Right click: attract', + 'Left click: add mass point (with shift to fix); Right click: attract', pos=(0, 0.99), color=0x0) - gui.text(content=f'C: clear all; Space: pause', + gui.text(content='C: clear all; Space: pause', pos=(0, 0.95), color=0x0) gui.text(content=f'Y: Spring Young\'s modulus {spring_Y[None]:.1f}', diff --git a/python/taichi/examples/simulation/mpm128.py b/python/taichi/examples/simulation/mpm128.py index 8dffdd7ffa526..dfa5172010c17 100644 --- a/python/taichi/examples/simulation/mpm128.py +++ b/python/taichi/examples/simulation/mpm128.py @@ -1,5 +1,3 @@ -import numpy as np - import taichi as ti ti.init(arch=ti.gpu) # Try to run on GPU @@ -40,10 +38,10 @@ def substep(): 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 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - F[p] = (ti.Matrix.identity(float, 2) + - dt * C[p]) @ F[p] # deformation gradient update - h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p]))) - ) # Hardening coefficient: snow gets harder when compressed + # deformation gradient update + F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p] + # Hardening coefficient: snow gets harder when compressed + h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p])))) if material[p] == 1: # jelly, make it softer h = 0.3 mu, la = mu_0 * h, lambda_0 * h @@ -59,18 +57,18 @@ def substep(): 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 + if material[p] == 0: + # Reset deformation gradient to avoid numerical instability F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J) elif material[p] == 2: - F[p] = U @ sig @ V.transpose( - ) # Reconstruct elastic deformation gradient after plasticity + # Reconstruct elastic deformation gradient after plasticity + F[p] = U @ sig @ V.transpose() stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose( ) + ti.Matrix.identity(float, 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 + 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] @@ -78,26 +76,28 @@ def substep(): grid_m[base + offset] += weight * p_mass for i, j in grid_m: 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 + # Momentum to velocity + grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] grid_v[i, j] += dt * gravity[None] * 30 # gravity dist = attractor_pos[None] - dx * ti.Vector([i, j]) - grid_v[i, j] += dist / ( - 0.01 + dist.norm()) * attractor_strength[None] * dt * 100 + grid_v[i, j] += \ + dist / (0.01 + dist.norm()) * attractor_strength[None] * dt * 100 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 + 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 x: # grid to particle (G2P) base = (x[p] * inv_dx - 0.5).cast(int) fx = x[p] * inv_dx - base.cast(float) w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2] new_v = ti.Vector.zero(float, 2) new_C = ti.Matrix.zero(float, 2, 2) - for i, j in ti.static(ti.ndrange( - 3, 3)): # loop over 3x3 grid node neighborhood + 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] @@ -131,13 +131,20 @@ def reset(): for frame in range(20000): if gui.get_event(ti.GUI.PRESS): - if gui.event.key == 'r': reset() - elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break - if gui.event is not None: gravity[None] = [0, 0] # if had any event - if gui.is_pressed(ti.GUI.LEFT, 'a'): gravity[None][0] = -1 - if gui.is_pressed(ti.GUI.RIGHT, 'd'): gravity[None][0] = 1 - if gui.is_pressed(ti.GUI.UP, 'w'): gravity[None][1] = 1 - if gui.is_pressed(ti.GUI.DOWN, 's'): gravity[None][1] = -1 + if gui.event.key == 'r': + reset() + elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: + break + if gui.event is not None: + gravity[None] = [0, 0] # if had any event + if gui.is_pressed(ti.GUI.LEFT, 'a'): + gravity[None][0] = -1 + if gui.is_pressed(ti.GUI.RIGHT, 'd'): + gravity[None][0] = 1 + if gui.is_pressed(ti.GUI.UP, 'w'): + gravity[None][1] = 1 + if gui.is_pressed(ti.GUI.DOWN, 's'): + gravity[None][1] = -1 mouse = gui.get_cursor_pos() gui.circle((mouse[0], mouse[1]), color=0x336699, radius=15) attractor_pos[None] = [mouse[0], mouse[1]] @@ -152,5 +159,6 @@ def reset(): radius=1.5, palette=[0x068587, 0xED553B, 0xEEEEF0], palette_indices=material) - gui.show( - ) # Change to gui.show(f'{frame:06d}.png') to write images to disk + + # Change to gui.show(f'{frame:06d}.png') to write images to disk + gui.show() diff --git a/python/taichi/examples/simulation/mpm3d.py b/python/taichi/examples/simulation/mpm3d.py index eb96f477b990e..939847ebc1417 100644 --- a/python/taichi/examples/simulation/mpm3d.py +++ b/python/taichi/examples/simulation/mpm3d.py @@ -22,71 +22,72 @@ bound = 3 E = 400 -x = ti.Vector.field(dim, float, n_particles) -v = ti.Vector.field(dim, float, n_particles) -C = ti.Matrix.field(dim, dim, float, n_particles) -J = ti.field(float, n_particles) +F_x = ti.Vector.field(dim, float, n_particles) +F_v = ti.Vector.field(dim, float, n_particles) +F_C = ti.Matrix.field(dim, dim, float, n_particles) +F_J = ti.field(float, n_particles) -grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim) -grid_m = ti.field(float, (n_grid, ) * dim) +F_grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim) +F_grid_m = ti.field(float, (n_grid, ) * dim) neighbour = (3, ) * dim @ti.kernel def substep(): - for I in ti.grouped(grid_m): - grid_v[I] = ti.zero(grid_v[I]) - grid_m[I] = 0 + for I in ti.grouped(F_grid_m): + F_grid_v[I] = ti.zero(F_grid_v[I]) + F_grid_m[I] = 0 ti.block_dim(n_grid) - for p in x: - Xp = x[p] / dx + for p in F_x: + Xp = F_x[p] / dx base = int(Xp - 0.5) fx = Xp - base w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - stress = -dt * 4 * E * p_vol * (J[p] - 1) / dx**2 - affine = ti.Matrix.identity(float, dim) * stress + p_mass * C[p] + stress = -dt * 4 * E * p_vol * (F_J[p] - 1) / dx**2 + affine = ti.Matrix.identity(float, dim) * stress + p_mass * F_C[p] for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))): dpos = (offset - fx) * dx weight = 1.0 for i in ti.static(range(dim)): weight *= w[offset[i]][i] - grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos) - grid_m[base + offset] += weight * p_mass - for I in ti.grouped(grid_m): - if grid_m[I] > 0: - grid_v[I] /= grid_m[I] - grid_v[I][1] -= dt * gravity - cond = (I < bound) & (grid_v[I] < 0) | \ - (I > n_grid - bound) & (grid_v[I] > 0) - grid_v[I] = 0 if cond else grid_v[I] + F_grid_v[base + + offset] += weight * (p_mass * F_v[p] + affine @ dpos) + F_grid_m[base + offset] += weight * p_mass + for I in ti.grouped(F_grid_m): + if F_grid_m[I] > 0: + F_grid_v[I] /= F_grid_m[I] + F_grid_v[I][1] -= dt * gravity + cond = (I < bound) & (F_grid_v[I] < 0) | \ + (I > n_grid - bound) & (F_grid_v[I] > 0) + F_grid_v[I] = 0 if cond else F_grid_v[I] ti.block_dim(n_grid) - for p in x: - Xp = x[p] / dx + for p in F_x: + Xp = F_x[p] / dx base = int(Xp - 0.5) fx = Xp - base w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - new_v = ti.zero(v[p]) - new_C = ti.zero(C[p]) + new_v = ti.zero(F_v[p]) + new_C = ti.zero(F_C[p]) for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))): dpos = (offset - fx) * dx weight = 1.0 for i in ti.static(range(dim)): weight *= w[offset[i]][i] - g_v = grid_v[base + offset] + g_v = F_grid_v[base + offset] new_v += weight * g_v new_C += 4 * weight * g_v.outer_product(dpos) / dx**2 - v[p] = new_v - x[p] += dt * v[p] - J[p] *= 1 + dt * new_C.trace() - C[p] = new_C + F_v[p] = new_v + F_x[p] += dt * F_v[p] + F_J[p] *= 1 + dt * new_C.trace() + F_C[p] = new_C @ti.kernel def init(): for i in range(n_particles): - x[i] = ti.Vector([ti.random() for i in range(dim)]) * 0.4 + 0.15 - J[i] = 1 + F_x[i] = ti.Vector([ti.random() for i in range(dim)]) * 0.4 + 0.15 + F_J[i] = 1 def T(a): @@ -97,22 +98,27 @@ def T(a): a = a - 0.5 x, y, z = a[:, 0], a[:, 1], a[:, 2] - c, s = np.cos(phi), np.sin(phi) - C, S = np.cos(theta), np.sin(theta) - x, z = x * c + z * s, z * c - x * s - u, v = x, y * C + z * S + cp, sp = np.cos(phi), np.sin(phi) + ct, st = np.cos(theta), np.sin(theta) + x, z = x * cp + z * sp, z * cp - x * sp + u, v = x, y * ct + z * st return np.array([u, v]).swapaxes(0, 1) + 0.5 -init() -gui = ti.GUI('MPM3D', background_color=0x112F41) -while gui.running and not gui.get_event(gui.ESCAPE): - for s in range(steps): - substep() - pos = x.to_numpy() - if export_file: - writer = ti.tools.PLYWriter(num_vertices=n_particles) - writer.add_vertex_pos(pos[:, 0], pos[:, 1], pos[:, 2]) - writer.export_frame(gui.frame, export_file) - gui.circles(T(pos), radius=1.5, color=0x66ccff) - gui.show() +def main(): + init() + gui = ti.GUI('MPM3D', background_color=0x112F41) + while gui.running and not gui.get_event(gui.ESCAPE): + for s in range(steps): + substep() + pos = F_x.to_numpy() + if export_file: + writer = ti.tools.PLYWriter(num_vertices=n_particles) + writer.add_vertex_pos(pos[:, 0], pos[:, 1], pos[:, 2]) + writer.export_frame(gui.frame, export_file) + gui.circles(T(pos), radius=1.5, color=0x66ccff) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/mpm99.py b/python/taichi/examples/simulation/mpm99.py index 05bdd608ab776..5f956c35af083 100644 --- a/python/taichi/examples/simulation/mpm99.py +++ b/python/taichi/examples/simulation/mpm99.py @@ -1,5 +1,3 @@ -import numpy as np - import taichi as ti ti.init(arch=ti.gpu) # Try to run on GPU @@ -35,12 +33,10 @@ def substep(): 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 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2] - F[p] = (ti.Matrix.identity(float, 2) + - dt * C[p]) @ F[p] # deformation gradient update - h = ti.exp( - 10 * - (1.0 - - Jp[p])) # Hardening coefficient: snow gets harder when compressed + # F[p]: deformation gradient update + F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p] + # h: Hardening coefficient: snow gets harder when compressed + h = ti.exp(10 * (1.0 - Jp[p])) if material[p] == 1: # jelly, make it softer h = 0.3 mu, la = mu_0 * h, lambda_0 * h @@ -56,18 +52,18 @@ def substep(): 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 + if material[p] == 0: + # Reset deformation gradient to avoid numerical instability F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J) elif material[p] == 2: - F[p] = U @ sig @ V.transpose( - ) # Reconstruct elastic deformation gradient after plasticity + # Reconstruct elastic deformation gradient after plasticity + F[p] = U @ sig @ V.transpose() stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose( ) + ti.Matrix.identity(float, 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 + # Loop over 3x3 grid node neighborhood + 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] @@ -75,23 +71,25 @@ def substep(): grid_m[base + offset] += weight * p_mass for i, j in grid_m: 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 / 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 + 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 x: # grid to particle (G2P) base = (x[p] * inv_dx - 0.5).cast(int) fx = x[p] * inv_dx - base.cast(float) w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2] new_v = ti.Vector.zero(float, 2) new_C = ti.Matrix.zero(float, 2, 2) - for i, j in ti.static(ti.ndrange( - 3, 3)): # loop over 3x3 grid node neighborhood + 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] @@ -127,8 +125,8 @@ def main(): radius=1.5, palette=[0x068587, 0xED553B, 0xEEEEF0], palette_indices=material) - gui.show( - ) # Change to gui.show(f'{frame:06d}.png') to write images to disk + # Change to gui.show(f'{frame:06d}.png') to write images to disk + gui.show() if __name__ == '__main__': diff --git a/python/taichi/examples/simulation/odop_solar.py b/python/taichi/examples/simulation/odop_solar.py index 72a72f4b53d36..e51c336277c8f 100644 --- a/python/taichi/examples/simulation/odop_solar.py +++ b/python/taichi/examples/simulation/odop_solar.py @@ -42,22 +42,31 @@ def integrate(self): # Semi-implicit Euler time integration self.v[i] += self.dt * self.gravity(self.x[i]) self.x[i] += self.dt * self.v[i] - def render(self, gui): # Render the scene on GUI + @staticmethod + def render(gui): # Render the scene on GUI gui.circle([0.5, 0.5], radius=10, color=0xffaa88) gui.circles(solar.x.to_numpy(), radius=3, color=0xffffff) -solar = SolarSystem(8, 0.0001) -solar.center[None] = [0.5, 0.5] -solar.initialize_particles() +def main(): + global solar + + solar = SolarSystem(8, 0.0001) + solar.center[None] = [0.5, 0.5] + solar.initialize_particles() + + gui = ti.GUI("Solar System", background_color=0x0071a) + while gui.running: + if gui.get_event() and gui.is_pressed(gui.SPACE): + solar.initialize_particles( + ) # reinitialize when space bar pressed. + + for _ in range(10): # Time integration + solar.integrate() -gui = ti.GUI("Solar System", background_color=0x0071a) -while gui.running: - if gui.get_event() and gui.is_pressed(gui.SPACE): - solar.initialize_particles() # reinitialize when space bar pressed. + solar.render(gui) + gui.show() - for i in range(10): # Time integration - solar.integrate() - solar.render(gui) - gui.show() +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/pbf2d.py b/python/taichi/examples/simulation/pbf2d.py index 3ea258b3a898b..a4346a931597e 100644 --- a/python/taichi/examples/simulation/pbf2d.py +++ b/python/taichi/examples/simulation/pbf2d.py @@ -37,7 +37,7 @@ def round_up(f, s): particle_radius_in_world = particle_radius / screen_to_world_ratio # PBF params -h = 1.1 +h_ = 1.1 mass = 1.0 rho0 = 1.0 lambda_epsilon = 100.0 @@ -46,7 +46,7 @@ def round_up(f, s): corrK = 0.001 # Need ti.pow() # corrN = 4.0 -neighbor_radius = h * 1.05 +neighbor_radius = h_ * 1.05 poly6_factor = 315.0 / 64.0 / math.pi spiky_grad_factor = -45.0 / math.pi @@ -97,7 +97,8 @@ def spiky_gradient(r, h): @ti.func def compute_scorr(pos_ji): # Eq (13) - x = poly6_value(pos_ji.norm(), h) / poly6_value(corr_deltaQ_coeff * h, h) + x = poly6_value(pos_ji.norm(), h_) / poly6_value(corr_deltaQ_coeff * h_, + h_) # pow(x, 4) x = x * x x = x * x @@ -202,11 +203,11 @@ def substep(): if p_j < 0: break pos_ji = pos_i - positions[p_j] - grad_j = spiky_gradient(pos_ji, h) + grad_j = spiky_gradient(pos_ji, h_) grad_i += grad_j sum_gradient_sqr += grad_j.dot(grad_j) # Eq(2) - density_constraint += poly6_value(pos_ji.norm(), h) + density_constraint += poly6_value(pos_ji.norm(), h_) # Eq(1) density_constraint = (mass * density_constraint / rho0) - 1.0 @@ -229,7 +230,7 @@ def substep(): pos_ji = pos_i - positions[p_j] scorr_ij = compute_scorr(pos_ji) pos_delta_i += (lambda_i + lambda_j + scorr_ij) * \ - spiky_gradient(pos_ji, h) + spiky_gradient(pos_ji, h_) pos_delta_i /= rho0 position_deltas[p_i] = pos_delta_i @@ -272,7 +273,7 @@ def render(gui): @ti.kernel def init_particles(): for i in range(num_particles): - delta = h * 0.8 + delta = h_ * 0.8 offs = ti.Vector([(boundary[0] - delta * num_particles_x) * 0.5, boundary[1] * 0.02]) positions[i] = ti.Vector([i % num_particles_x, i // num_particles_x @@ -285,11 +286,11 @@ def init_particles(): def print_stats(): print('PBF stats:') num = grid_num_particles.to_numpy() - avg, max = np.mean(num), np.max(num) - print(f' #particles per cell: avg={avg:.2f} max={max}') + avg, max_ = np.mean(num), np.max(num) + print(f' #particles per cell: avg={avg:.2f} max={max_}') num = particle_num_neighbors.to_numpy() - avg, max = np.mean(num), np.max(num) - print(f' #neighbors per particle: avg={avg:.2f} max={max}') + avg, max_ = np.mean(num), np.max(num) + print(f' #neighbors per particle: avg={avg:.2f} max={max_}') def main(): diff --git a/python/taichi/examples/simulation/physarum.py b/python/taichi/examples/simulation/physarum.py index 9eadf30211d0a..92dadaaf156a7 100644 --- a/python/taichi/examples/simulation/physarum.py +++ b/python/taichi/examples/simulation/physarum.py @@ -68,14 +68,19 @@ def step(phase: ti.i32): grid[1 - phase, i, j] = a -print("[Hint] Press A/Z to change the simulation speed.") -gui = ti.GUI('Physarum') -init() -i = 0 -step_per_frame = gui.slider('step_per_frame', 1, 100, 1) -while gui.running and not gui.get_event(gui.ESCAPE): - for _ in range(int(step_per_frame.value)): - step(i % 2) - i += 1 - gui.set_image(grid.to_numpy()[0]) - gui.show() +def main(): + print("[Hint] Use slider to change simulation speed.") + gui = ti.GUI('Physarum') + init() + i = 0 + step_per_frame = gui.slider('step_per_frame', 1, 100, 1) + while gui.running and not gui.get_event(gui.ESCAPE): + for _ in range(int(step_per_frame.value)): + step(i % 2) + i += 1 + gui.set_image(grid.to_numpy()[0]) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/stable_fluid.py b/python/taichi/examples/simulation/stable_fluid.py index 03cb98c91fa18..aac12c3799717 100644 --- a/python/taichi/examples/simulation/stable_fluid.py +++ b/python/taichi/examples/simulation/stable_fluid.py @@ -29,10 +29,7 @@ maxfps = 60 dye_decay = 1 - 1 / (maxfps * time_c) force_radius = res / 2.0 -gravity = True debug = False -paused = False -use_sparse_matrix = False use_sparse_matrix = args.use_sp_mat @@ -89,7 +86,7 @@ def fill_laplacian_matrix(A: ti.types.sparse_matrix_builder()): N = res * res K = ti.linalg.SparseMatrixBuilder(N, N, max_num_triplets=N * 6) - b = ti.field(ti.f32, shape=N) + F_b = ti.field(ti.f32, shape=N) fill_laplacian_matrix(K) L = K.build() @@ -128,13 +125,13 @@ def bilerp(vf, p): # 3rd order Runge-Kutta @ti.func -def backtrace(vf: ti.template(), p, dt: ti.template()): +def backtrace(vf: ti.template(), p, dt_: ti.template()): v1 = bilerp(vf, p) - p1 = p - 0.5 * dt * v1 + p1 = p - 0.5 * dt_ * v1 v2 = bilerp(vf, p1) - p2 = p - 0.75 * dt * v2 + p2 = p - 0.75 * dt_ * v2 v3 = bilerp(vf, p2) - p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) + p -= dt_ * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3) return p @@ -251,8 +248,8 @@ def apply_pressure(p_in: ti.types.ndarray(), p_out: ti.template()): def solve_pressure_sp_mat(): - copy_divergence(velocity_divs, b) - x = solver.solve(b) + copy_divergence(velocity_divs, F_b) + x = solver.solve(F_b) apply_pressure(x, pressures_pair.cur) @@ -289,7 +286,7 @@ def step(mouse_data): print(f'divergence={div_s}') -class MouseDataGen(object): +class MouseDataGen: def __init__(self): self.prev_mouse = None self.prev_color = None @@ -324,56 +321,62 @@ def reset(): dyes_pair.cur.fill(0) -visualize_d = True #visualize dye (default) -visualize_v = False #visualize velocity -visualize_c = False #visualize curl - -gui = ti.GUI('Stable Fluid', (res, res)) -md_gen = MouseDataGen() - -while gui.running: - if gui.get_event(ti.GUI.PRESS): - e = gui.event - if e.key == ti.GUI.ESCAPE: - break - elif e.key == 'r': - paused = False - reset() - elif e.key == 's': - if curl_strength: - curl_strength = 0 - else: - curl_strength = 7 - elif e.key == 'g': - gravity = not gravity - elif e.key == 'v': - visualize_v = True - visualize_c = False - visualize_d = False - elif e.key == 'd': - visualize_d = True - visualize_v = False - visualize_c = False - elif e.key == 'c': - visualize_c = True - visualize_d = False - visualize_v = False - elif e.key == 'p': - paused = not paused - elif e.key == 'd': - debug = not debug - - # Debug divergence: - # print(max((abs(velocity_divs.to_numpy().reshape(-1))))) - - if not paused: - mouse_data = md_gen(gui) - step(mouse_data) - if visualize_c: - vorticity(velocities_pair.cur) - gui.set_image(velocity_curls.to_numpy() * 0.03 + 0.5) - elif visualize_d: - gui.set_image(dyes_pair.cur) - elif visualize_v: - gui.set_image(velocities_pair.cur.to_numpy() * 0.01 + 0.5) - gui.show() +def main(): + global debug, curl_strength + visualize_d = True #visualize dye (default) + visualize_v = False #visualize velocity + visualize_c = False #visualize curl + + paused = False + + gui = ti.GUI('Stable Fluid', (res, res)) + md_gen = MouseDataGen() + + while gui.running: + if gui.get_event(ti.GUI.PRESS): + e = gui.event + if e.key == ti.GUI.ESCAPE: + break + elif e.key == 'r': + paused = False + reset() + elif e.key == 's': + if curl_strength: + curl_strength = 0 + else: + curl_strength = 7 + elif e.key == 'v': + visualize_v = True + visualize_c = False + visualize_d = False + elif e.key == 'd': + visualize_d = True + visualize_v = False + visualize_c = False + elif e.key == 'c': + visualize_c = True + visualize_d = False + visualize_v = False + elif e.key == 'p': + paused = not paused + elif e.key == 'd': + debug = not debug + + # Debug divergence: + # print(max((abs(velocity_divs.to_numpy().reshape(-1))))) + + if not paused: + mouse_data = md_gen(gui) + step(mouse_data) + if visualize_c: + vorticity(velocities_pair.cur) + gui.set_image(velocity_curls.to_numpy() * 0.03 + 0.5) + elif visualize_d: + gui.set_image(dyes_pair.cur) + elif visualize_v: + gui.set_image(velocities_pair.cur.to_numpy() * 0.01 + 0.5) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/vortex_rings.py b/python/taichi/examples/simulation/vortex_rings.py index ac8eaedb761c8..3d75655de18da 100644 --- a/python/taichi/examples/simulation/vortex_rings.py +++ b/python/taichi/examples/simulation/vortex_rings.py @@ -75,18 +75,22 @@ def init_tracers(): tracer[i] = [ti.random() - 0.5, ti.random() * 3 - 1.5] -init_tracers() +def main(): + init_tracers() + gui = ti.GUI("Vortex Rings", (1024, 512), background_color=0xFFFFFF) -gui = ti.GUI("Vortex Rings", (1024, 512), background_color=0xFFFFFF) + for T in range(1000): + for i in range(4): # substeps + advect() + integrate_vortex() -for T in range(1000): - for i in range(4): # substeps - advect() - integrate_vortex() + gui.circles(tracer.to_numpy() * np.array([[0.05, 0.1]]) + + np.array([[0.0, 0.5]]), + radius=0.5, + color=0x0) - gui.circles(tracer.to_numpy() * np.array([[0.05, 0.1]]) + - np.array([[0.0, 0.5]]), - radius=0.5, - color=0x0) + gui.show() - gui.show() + +if __name__ == '__main__': + main() diff --git a/python/taichi/examples/simulation/waterwave.py b/python/taichi/examples/simulation/waterwave.py index 717e3bc001b25..7487a52ffdb82 100644 --- a/python/taichi/examples/simulation/waterwave.py +++ b/python/taichi/examples/simulation/waterwave.py @@ -71,20 +71,25 @@ def visualize_wave(): pixels[i, j] = (1 - brightness) * color + brightness * light_color -print("[Hint] click on the window to create waves") - -reset() -gui = ti.GUI('Water Wave', shape) -while gui.running: - for e in gui.get_events(ti.GUI.PRESS): - if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: - gui.running = False - elif e.key == 'r': - reset() - elif e.key == ti.GUI.LMB: - x, y = e.pos - create_wave(3, x * shape[0], y * shape[1]) - update() - visualize_wave() - gui.set_image(pixels) - gui.show() +def main(): + print("[Hint] click on the window to create waves") + + reset() + gui = ti.GUI('Water Wave', shape) + while gui.running: + for e in gui.get_events(ti.GUI.PRESS): + if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: + gui.running = False + elif e.key == 'r': + reset() + elif e.key == ti.GUI.LMB: + x, y = e.pos + create_wave(3, x * shape[0], y * shape[1]) + update() + visualize_wave() + gui.set_image(pixels) + gui.show() + + +if __name__ == '__main__': + main() diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py index a00a56c7b2c4b..8337beaabacf0 100644 --- a/python/taichi/lang/impl.py +++ b/python/taichi/lang/impl.py @@ -23,7 +23,7 @@ from taichi.lang.struct import Struct, StructField, _IntermediateStruct from taichi.lang.util import (cook_dtype, get_traceback, is_taichi_class, python_scope, taichi_scope, warning) -from taichi.types.primitive_types import f16, f32, f64, i32, i64, types +from taichi.types.primitive_types import all_types, f16, f32, f64, i32, i64 @taichi_scope @@ -585,7 +585,7 @@ def ndarray(dtype, shape, layout=Layout.NULL): """ if isinstance(shape, numbers.Number): shape = (shape, ) - if dtype in types: + if dtype in all_types: assert layout == Layout.NULL return ScalarNdarray(dtype, shape) if isinstance(dtype, MatrixType): diff --git a/python/taichi/types/primitive_types.py b/python/taichi/types/primitive_types.py index 3149ffc450a8b..53fd84101040a 100644 --- a/python/taichi/types/primitive_types.py +++ b/python/taichi/types/primitive_types.py @@ -157,8 +157,8 @@ def ref(tp): integer_types = [i8, i16, i32, i64, u8, u16, u32, u64, int] integer_type_ids = [id(t) for t in integer_types] -types = real_types + integer_types -type_ids = [id(t) for t in types] +all_types = real_types + integer_types +type_ids = [id(t) for t in all_types] __all__ = [ 'float32',