From 2df041789e0371fa4be807f3f26affdea874d012 Mon Sep 17 00:00:00 2001 From: ailzhang Date: Sat, 11 Jun 2022 13:26:49 +0800 Subject: [PATCH] [ci] Enable yapf and isort on example files Note we explicitly exclude running pylint on them as it requires a bunch of manual fixes first. --- .pre-commit-config.yaml | 3 +- .../examples/algorithm/mciso_advanced.py | 3 +- .../examples/autodiff/diff_sph/diff_sph.py | 268 ++++++++++++------ .../examples/features/io/export_mesh.py | 4 +- .../ggui_examples/mass_spring_3d_ggui.py | 8 +- python/taichi/examples/graph/mpm88_graph.py | 30 +- .../examples/graph/stable_fluid_graph.py | 48 ++-- tests/python/examples/__init__.py | 1 - tests/python/examples/autodiff/__init__.py | 1 - 9 files changed, 239 insertions(+), 127 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7c0fdf9f995d9..0dc07c04fff9e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ ci: autoupdate_schedule: quarterly autoupdate_commit_msg: '[misc] Update pre-commit hooks' -exclude: ^((tests/python/test_exception|.*/examples/.*)\.py$|external/) +exclude: ^((tests/python/test_exception)\.py$|external/) repos: - repo: https://github.com/google/yapf rev: v0.32.0 @@ -33,3 +33,4 @@ repos: - id: pylint args: ['-rn', '-sn'] files: ^python/taichi/ + exclude: ^python/taichi/examples/.*.py diff --git a/python/taichi/examples/algorithm/mciso_advanced.py b/python/taichi/examples/algorithm/mciso_advanced.py index 991e56a0fbdc5..fd79563982751 100644 --- a/python/taichi/examples/algorithm/mciso_advanced.py +++ b/python/taichi/examples/algorithm/mciso_advanced.py @@ -489,7 +489,8 @@ def main(self): (0, 1)) if gui.is_pressed(gui.SPACE): num = ret.shape[0] - writer = ti.tools.PLYWriter(num_vertices=num * 3, num_faces=num) + writer = ti.tools.PLYWriter(num_vertices=num * 3, + num_faces=num) vertices = ret.reshape(num * 3, 3) * 2 - 1 writer.add_vertex_pos(vertices[:, 0], vertices[:, 1], vertices[:, 2]) diff --git a/python/taichi/examples/autodiff/diff_sph/diff_sph.py b/python/taichi/examples/autodiff/diff_sph/diff_sph.py index 3070b4aceaeea..8ee5495560e8f 100644 --- a/python/taichi/examples/autodiff/diff_sph/diff_sph.py +++ b/python/taichi/examples/autodiff/diff_sph/diff_sph.py @@ -1,23 +1,24 @@ # Smoothed-particle hydrodynamics (SPH) is a computational method used for simulating the mechanics of continuum media, such as solid mechanics and fluid flows. # Here we utilize SPH to simulate a fountain, who tries to hit a target given by the user. -# The SPH simulator here implemented using Taichi is differentiable. -# Therefore, it can be easily embedding into the training pipeline of a neural network modelled controller. +# The SPH simulator here implemented using Taichi is differentiable. +# Therefore, it can be easily embedding into the training pipeline of a neural network modelled controller. -import taichi as ti -import numpy as np -import matplotlib.pyplot as plt +import argparse import os import pickle as pkl -import argparse + +import matplotlib.pyplot as plt +import numpy as np + +import taichi as ti + parser = argparse.ArgumentParser() -parser.add_argument( - '--train', - action='store_true', - help='whether train model, default false') +parser.add_argument('--train', + action='store_true', + help='whether train model, default false') parser.add_argument('place_holder', nargs='*') args = parser.parse_args() - TRAIN = args.train TRAIN_OUTPUT_IMG = False TRAIN_VISUAL = False @@ -52,7 +53,15 @@ def zero_grad(self): @ti.data_oriented class Linear: - def __init__(self, n_models, batch_size, n_steps, n_input, n_hidden, n_output, needs_grad=False, activation=False): + def __init__(self, + n_models, + batch_size, + n_steps, + n_input, + n_hidden, + n_output, + needs_grad=False, + activation=False): self.n_models = n_models self.batch_size = batch_size self.n_steps = n_steps @@ -69,8 +78,12 @@ def __init__(self, n_models, batch_size, n_steps, n_input, n_hidden, n_output, n self.n_hidden_node = self.batch_node.dense(ti.j, self.n_hidden) self.weights1_node = self.n_hidden_node.dense(ti.k, self.n_input) - self.batch_node.dense(ti.axes(1, 2, 3), (self.n_steps, self.batch_size, self.n_hidden)).place(self.hidden) - self.batch_node.dense(ti.axes(1, 2, 3), (self.n_steps, self.batch_size, self.n_output)).place(self.output) + self.batch_node.dense( + ti.axes(1, 2, 3), + (self.n_steps, self.batch_size, self.n_hidden)).place(self.hidden) + self.batch_node.dense( + ti.axes(1, 2, 3), + (self.n_steps, self.batch_size, self.n_output)).place(self.output) self.weights1 = scalar() self.bias1 = scalar() @@ -87,19 +100,29 @@ def parameters(self): @ti.kernel def weights_init(self): q1 = ti.sqrt(6 / self.n_input) * 0.01 - for model_id, i, j in ti.ndrange(self.n_models, self.n_hidden, self.n_input): + for model_id, i, j in ti.ndrange(self.n_models, self.n_hidden, + self.n_input): self.weights1[model_id, i, j] = (ti.random() * 2 - 1) * q1 @ti.kernel def _forward(self, t: ti.i32, nn_input: ti.template()): - for model_id, k, i, j in ti.ndrange(self.n_models, self.batch_size, self.n_hidden, self.n_input): - self.hidden[model_id, t, k, i] += self.weights1[model_id, i, j] * nn_input[model_id, t, k, j] + for model_id, k, i, j in ti.ndrange(self.n_models, self.batch_size, + self.n_hidden, self.n_input): + self.hidden[model_id, t, k, + i] += self.weights1[model_id, i, + j] * nn_input[model_id, t, k, j] if ti.static(self.activation): - for model_id, k, i in ti.ndrange(self.n_models, self.batch_size, self.n_hidden): - self.output[model_id, t, k, i] = ti.tanh(self.hidden[model_id, t, k, i] + self.bias1[model_id, i]) + for model_id, k, i in ti.ndrange(self.n_models, self.batch_size, + self.n_hidden): + self.output[model_id, t, k, + i] = ti.tanh(self.hidden[model_id, t, k, i] + + self.bias1[model_id, i]) else: - for model_id, k, i in ti.ndrange(self.n_models, self.batch_size, self.n_hidden): - self.output[model_id, t, k, i] = self.hidden[model_id, t, k, i] + self.bias1[model_id, i] + for model_id, k, i in ti.ndrange(self.n_models, self.batch_size, + self.n_hidden): + self.output[model_id, t, k, + i] = self.hidden[model_id, t, k, + i] + self.bias1[model_id, i] @ti.kernel def clear(self): @@ -129,7 +152,8 @@ def load_weights_from_value(self, w_val, model_id=0): self.copy_from_numpy(w, val, model_id) @ti.kernel - def copy_from_numpy(self, dst: ti.template(), src: ti.ext_arr(), model_id: ti.i32): + def copy_from_numpy(self, dst: ti.template(), src: ti.ext_arr(), + model_id: ti.i32): for I in ti.grouped(src): dst[model_id, I] = src[I] @@ -144,14 +168,27 @@ def copy_from_numpy(self, dst: ti.template(), src: ti.ext_arr(), model_id: ti.i3 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) + 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] @@ -169,7 +206,8 @@ def copy_from_numpy(self, dst: ti.template(), src: ti.ext_arr(), model_id: ti.i3 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)]) + 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) @@ -179,11 +217,25 @@ def targets_generation(num, x_range, y_range, z_range): 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) + 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) @@ -200,13 +252,15 @@ def targets_generation(num, x_range, y_range, z_range): 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]) +jet_force_max = ti.Vector([9.81 * 3, 9.81 * 10, 9.81 * 3]) # Simulation parameters particle_radius = 0.01 particle_diameter = particle_radius * 2 -N_np = ((spawn_box_np[1] - spawn_box_np[0]) / particle_diameter + 1).astype(int) -N_target_np = ((target_box_np[1] - target_box_np[0]) / particle_diameter + 1).astype(int) +N_np = ((spawn_box_np[1] - spawn_box_np[0]) / particle_diameter + + 1).astype(int) +N_target_np = ((target_box_np[1] - target_box_np[0]) / particle_diameter + + 1).astype(int) h = 4.0 * particle_radius fluid_particle_num = N_np[0] * N_np[1] * N_np[2] @@ -217,8 +271,10 @@ def targets_generation(num, x_range, y_range, z_range): 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) - +jet_force = ti.Vector.field(3, + float, + shape=(steps, batch_size), + needs_grad=True) col = ti.Vector.field(3, float) material = ti.field(int) @@ -227,7 +283,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( + pos, vel, acc, den, pre) ti.root.dense(ti.i, int(particle_num)).place(material, col) ti.root.lazy_grad() @@ -292,16 +349,15 @@ def W_spiky_gradient(R, h): @ti.kernel -def initialize_fluid_particle(t: ti.int32, pos: ti.template(), N_fluid: ti.template()): +def initialize_fluid_particle(t: ti.int32, pos: ti.template(), + N_fluid: ti.template()): # Allocate fluid 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])] - ) - * particle_diameter - + spawn_box[0] - ) + 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]) + ]) * particle_diameter + spawn_box[0]) 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]) @@ -320,16 +376,17 @@ def initialize_dists(): @ti.kernel -def initialize_target_particle(t: ti.int32, pos: ti.template(), N_target:ti.template(), current_pos: ti.int32): +def initialize_target_particle(t: ti.int32, pos: ti.template(), + N_target: ti.template(), current_pos: ti.int32): # Allocate target cube - for bs, i in ti.ndrange(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])] - ) - * particle_diameter - + target_centers[current_pos] - ) + for bs, i in ti.ndrange( + 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]) + ]) * particle_diameter + target_centers[current_pos]) 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]) @@ -355,7 +412,8 @@ def update_density(t: ti.int32): @ti.kernel def update_pressure(t: ti.int32): 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) + pre[bs, t, i] = pressure_scale * max( + pow(den[bs, t, i] / rest_density, gamma) - 1, 0) @ti.kernel @@ -371,9 +429,14 @@ def apply_force(t: ti.int32): if material[i] == 1: 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 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: indicator = (steps - t) // (steps // 2) - acc[bs, t, i] = jet_force[t, bs] + gravity + indicator * (- gravity) * 0.1 + acc[bs, t, + i] = jet_force[t, + bs] + gravity + indicator * (-gravity) * 0.1 else: acc[bs, t, i] = gravity @@ -384,7 +447,10 @@ def update_force(t: ti.int32): for j in range(particle_num): R = pos[bs, t, i] - pos[bs, t, j] # Pressure forces - 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) + 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) # Viscosity forces acc[bs, t, i] += viscosity_scale * mass \ @@ -415,7 +481,8 @@ def boundary_handle(t: ti.int32): 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 + vel[bs, t, i] -= (1.0 + damping) * collision_normal.dot( + vel[bs, t, i]) * collision_normal @ti.kernel @@ -435,7 +502,9 @@ def compute_dist(t: ti.int32): @ti.kernel def compute_loss(t: ti.int32): 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) + 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 @@ -470,7 +539,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]) @@ -478,7 +547,7 @@ 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] @@ -505,12 +574,14 @@ def fill_input_states(current_pos: ti.int32): for opt_iter in range(opt_iters): loss_epoch = 0.0 cnt = 0 - for current_data_offset in range(0, training_sample_num, batch_size): + for current_data_offset in range(0, training_sample_num, + batch_size): fill_target_centers(current_data_offset, training_data) fill_input_states(current_data_offset) initialize_fluid_particle(0, pos, N_fluid) initialize_dists() - initialize_target_particle(0, pos, N_target, current_data_offset) + initialize_target_particle(0, pos, N_target, + current_data_offset) fc1.clear() fc2.clear() with ti.Tape(loss=loss): @@ -530,11 +601,14 @@ def fill_input_states(current_pos: ti.int32): compute_dist(i) compute_loss(steps - 1) optimizer.step() - print(f"current opt progress: {current_data_offset + batch_size}/{training_sample_num}, loss: {loss[None]}") + print( + f"current opt progress: {current_data_offset + batch_size}/{training_sample_num}, loss: {loss[None]}" + ) losses.append(loss[None]) loss_epoch += loss[None] cnt += 1 - print(f'opt iter {opt_iter} done. Average loss: {loss_epoch / cnt}') + print( + f'opt iter {opt_iter} done. Average loss: {loss_epoch / cnt}') losses_epoch_avg.append(loss_epoch / cnt) if TRAIN_VISUAL: @@ -544,21 +618,34 @@ def fill_input_states(current_pos: ti.int32): if i % substeps == 0: copy_from_output_to_vis(i) scene.set_camera(camera) - scene.point_light((2.0, 2.0, 2.0), color=(1.0, 1.0, 1.0)) - scene.particles(pos_vis_buffer, radius=particle_radius, per_vertex_color=col) + scene.point_light((2.0, 2.0, 2.0), + color=(1.0, 1.0, 1.0)) + scene.particles(pos_vis_buffer, + radius=particle_radius, + per_vertex_color=col) canvas.scene(scene) if TRAIN_OUTPUT_IMG: if i % substeps == 0: - window.write_image(f'output_img/{opt_iter}/{i:04}.png') + window.write_image( + f'output_img/{opt_iter}/{i:04}.png') if TRAIN_VISUAL_SHOW: window.show() if opt_iter % 2 == 0: os.makedirs(f"saved_models/{opt_iter}", exist_ok=True) - fc1.dump_weights(name=f"saved_models/{opt_iter}/fc1_{opt_iter:04}.pkl") - fc2.dump_weights(name=f"saved_models/{opt_iter}/fc2_{opt_iter:04}.pkl") - - plt.plot([i for i in range(len(losses))], losses, label='loss per iteration') - plt.plot([i * (training_sample_num // batch_size) for i in range(len(losses_epoch_avg))], losses_epoch_avg, label='loss epoch avg.') + fc1.dump_weights( + name=f"saved_models/{opt_iter}/fc1_{opt_iter:04}.pkl") + fc2.dump_weights( + name=f"saved_models/{opt_iter}/fc2_{opt_iter:04}.pkl") + + plt.plot([i for i in range(len(losses))], + losses, + label='loss per iteration') + plt.plot([ + i * (training_sample_num // batch_size) + for i in range(len(losses_epoch_avg)) + ], + losses_epoch_avg, + label='loss epoch avg.') plt.title("Training Loss") plt.xlabel("Training Iterations") plt.ylabel("Loss") @@ -579,20 +666,18 @@ def fill_input_states(current_pos: ti.int32): window.GUI.text("Space: pause") window.GUI.text("Set target positions:") - target_centers[current_data_offset][0] = window.GUI.slider_float("X", - target_centers[current_data_offset][0], - 0.05, 0.45) - target_centers[current_data_offset][1] = window.GUI.slider_float("Y", - target_centers[current_data_offset][1], - 0.4, 1.0) - target_centers[current_data_offset][2] = window.GUI.slider_float("Z", - target_centers[current_data_offset][2], - 0.05, 0.45) + target_centers[current_data_offset][0] = window.GUI.slider_float( + "X", target_centers[current_data_offset][0], 0.05, 0.45) + target_centers[current_data_offset][1] = window.GUI.slider_float( + "Y", target_centers[current_data_offset][1], 0.4, 1.0) + target_centers[current_data_offset][2] = window.GUI.slider_float( + "Z", target_centers[current_data_offset][2], 0.05, 0.45) window.GUI.end() if not paused[None]: fill_input_states(current_data_offset) - initialize_target_particle(0, pos, N_target, current_data_offset) + initialize_target_particle(0, pos, N_target, + current_data_offset) fc1.clear() fc2.clear() for i in range(1, substeps): @@ -623,15 +708,18 @@ def fill_input_states(current_pos: ti.int32): paused[None] = not paused[None] camera.position(*(camera.curr_position + position_change)) camera.lookat(*(camera.curr_lookat + position_change)) - camera.track_user_inputs(window, movement_speed=movement_speed, hold_key=ti.ui.RMB) + camera.track_user_inputs(window, + movement_speed=movement_speed, + hold_key=ti.ui.RMB) scene.set_camera(camera) scene.point_light((2.0, 2.0, 2.0), color=(1.0, 1.0, 1.0)) - scene.particles(pos_vis_buffer, radius=particle_radius, per_vertex_color=col) + scene.particles(pos_vis_buffer, + radius=particle_radius, + per_vertex_color=col) canvas.scene(scene) if INFER_OUTPUT_IMG: if cnt % 2 == 0: os.makedirs(f"demo_output_interactive/", exist_ok=True) window.write_image(f'demo_output_interactive/{cnt:04}.png') window.show() - diff --git a/python/taichi/examples/features/io/export_mesh.py b/python/taichi/examples/features/io/export_mesh.py index c2090b0e95b6f..21b2bfa495ac9 100644 --- a/python/taichi/examples/features/io/export_mesh.py +++ b/python/taichi/examples/features/io/export_mesh.py @@ -92,7 +92,9 @@ b = np.random.rand(20) alpha = np.random.rand(20) # re-fill - writer = ti.tools.PLYWriter(num_vertices=20, num_faces=12, face_type="quad") + writer = ti.tools.PLYWriter(num_vertices=20, + num_faces=12, + face_type="quad") writer.add_vertex_pos(x, y, z) writer.add_faces(indices) writer.add_vertex_channel("vdata1", "double", vdata) 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 4be661f9fbead..0f23939d12fb3 100644 --- a/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py +++ b/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py @@ -27,6 +27,7 @@ bending_springs = False + @ti.kernel def initialize_mass_points(): random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1 @@ -58,6 +59,7 @@ def initialize_mesh_indices(): else: colors[i * n + j] = (1, 0.334, 0.52) + initialize_mesh_indices() spring_offsets = [] @@ -73,6 +75,7 @@ def initialize_mesh_indices(): if (i, j) != (0, 0) and abs(i) + abs(j) <= 2: spring_offsets.append(ti.Vector([i, j])) + @ti.kernel def substep(): for i in ti.grouped(x): @@ -105,7 +108,6 @@ def substep(): x[i] += dt * v[i] - @ti.kernel def update_vertices(): for i, j in ti.ndrange(n, n): @@ -145,7 +147,9 @@ def update_vertices(): 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)) + scene.particles(ball_center, + radius=ball_radius * 0.95, + color=(0.5, 0.42, 0.8)) canvas.scene(scene) window.show() diff --git a/python/taichi/examples/graph/mpm88_graph.py b/python/taichi/examples/graph/mpm88_graph.py index 50ba7554fa7b8..c59263d7ffb71 100644 --- a/python/taichi/examples/graph/mpm88_graph.py +++ b/python/taichi/examples/graph/mpm88_graph.py @@ -1,5 +1,7 @@ import argparse + import numpy as np + import taichi as ti ti.init(arch=ti.vulkan) @@ -17,7 +19,6 @@ N_ITER = 500 # Use 500 to make speed diff more obvious - @ti.kernel def substep_reset_grid(grid_v: ti.any_arr(field_dim=2), grid_m: ti.any_arr(field_dim=2)): @@ -95,6 +96,7 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1), v[i] = [0, -1] J[i] = 1 + x = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) v = ti.Vector.ndarray(2, ti.f32, shape=(n_particles)) @@ -105,19 +107,29 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1), if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument( - '--baseline', - action='store_true') + parser.add_argument('--baseline', action='store_true') args, unknown = parser.parse_known_args() if not args.baseline: print('running in graph mode') # Build graph - sym_x = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'x', ti.f32, element_shape=(2, )) - sym_v = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'v', ti.f32, element_shape=(2, )) - sym_C = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'C', ti.f32, element_shape=(2, 2)) + sym_x = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, + 'x', + ti.f32, + element_shape=(2, )) + sym_v = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, + 'v', + ti.f32, + element_shape=(2, )) + sym_C = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, + 'C', + ti.f32, + element_shape=(2, 2)) sym_J = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'J', ti.f32) - sym_grid_v = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'grid_v', ti.f32, element_shape=(2, )) + sym_grid_v = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, + 'grid_v', + ti.f32, + element_shape=(2, )) sym_grid_m = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'grid_m', ti.f32) g_init_builder = ti.graph.GraphBuilder() g_init_builder.dispatch(init_particles, sym_x, sym_v, sym_J) @@ -127,7 +139,7 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1), substep.dispatch(substep_reset_grid, sym_grid_v, sym_grid_m) substep.dispatch(substep_p2g, sym_x, sym_v, sym_C, sym_J, sym_grid_v, - sym_grid_m) + sym_grid_m) substep.dispatch(substep_update_grid_v, sym_grid_v, sym_grid_m) substep.dispatch(substep_g2p, sym_x, sym_v, sym_C, sym_J, sym_grid_v) diff --git a/python/taichi/examples/graph/stable_fluid_graph.py b/python/taichi/examples/graph/stable_fluid_graph.py index c4b0826e953f6..9a33ac5342fe2 100644 --- a/python/taichi/examples/graph/stable_fluid_graph.py +++ b/python/taichi/examples/graph/stable_fluid_graph.py @@ -7,6 +7,7 @@ import argparse import numpy as np + import taichi as ti ti.init(arch=ti.vulkan) @@ -214,11 +215,10 @@ def reset(): pressures_pair.cur.fill(0) dyes_pair.cur.fill(0) + if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument( - '--baseline', - action='store_true') + parser.add_argument('--baseline', action='store_true') args, unknown = parser.parse_known_args() gui = ti.GUI('Stable Fluid', (res, res)) @@ -261,38 +261,44 @@ def reset(): 'pressures_pair_nxt', ti.f32) velocity_divs = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'velocity_divs', ti.f32) - mouse_data = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'mouse_data', ti.f32) + mouse_data = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'mouse_data', + ti.f32) g1_builder = ti.graph.GraphBuilder() g1_builder.dispatch(advect, velocities_pair_cur, velocities_pair_cur, - velocities_pair_nxt) - g1_builder.dispatch(advect, velocities_pair_cur, dyes_pair_cur, dyes_pair_nxt) - g1_builder.dispatch(apply_impulse, velocities_pair_nxt, dyes_pair_nxt, mouse_data) + velocities_pair_nxt) + g1_builder.dispatch(advect, velocities_pair_cur, dyes_pair_cur, + dyes_pair_nxt) + g1_builder.dispatch(apply_impulse, velocities_pair_nxt, dyes_pair_nxt, + mouse_data) g1_builder.dispatch(divergence, velocities_pair_nxt, velocity_divs) # swap is unrolled in the loop so we only need p_jacobi_iters // 2 iterations. for _ in range(p_jacobi_iters // 2): - g1_builder.dispatch(pressure_jacobi, pressures_pair_cur, pressures_pair_nxt, - velocity_divs) - g1_builder.dispatch(pressure_jacobi, pressures_pair_nxt, pressures_pair_cur, - velocity_divs) - g1_builder.dispatch(subtract_gradient, velocities_pair_nxt, pressures_pair_cur) + g1_builder.dispatch(pressure_jacobi, pressures_pair_cur, + pressures_pair_nxt, velocity_divs) + g1_builder.dispatch(pressure_jacobi, pressures_pair_nxt, + pressures_pair_cur, velocity_divs) + g1_builder.dispatch(subtract_gradient, velocities_pair_nxt, + pressures_pair_cur) g1 = g1_builder.compile() g2_builder = ti.graph.GraphBuilder() g2_builder.dispatch(advect, velocities_pair_nxt, velocities_pair_nxt, - velocities_pair_cur) - g2_builder.dispatch(advect, velocities_pair_nxt, dyes_pair_nxt, dyes_pair_cur) - g2_builder.dispatch(apply_impulse, velocities_pair_cur, dyes_pair_cur, mouse_data) + velocities_pair_cur) + g2_builder.dispatch(advect, velocities_pair_nxt, dyes_pair_nxt, + dyes_pair_cur) + g2_builder.dispatch(apply_impulse, velocities_pair_cur, dyes_pair_cur, + mouse_data) g2_builder.dispatch(divergence, velocities_pair_cur, velocity_divs) for _ in range(p_jacobi_iters // 2): - g2_builder.dispatch(pressure_jacobi, pressures_pair_cur, pressures_pair_nxt, - velocity_divs) - g2_builder.dispatch(pressure_jacobi, pressures_pair_nxt, pressures_pair_cur, - velocity_divs) - g2_builder.dispatch(subtract_gradient, velocities_pair_cur, pressures_pair_cur) + g2_builder.dispatch(pressure_jacobi, pressures_pair_cur, + pressures_pair_nxt, velocity_divs) + g2_builder.dispatch(pressure_jacobi, pressures_pair_nxt, + pressures_pair_cur, velocity_divs) + g2_builder.dispatch(subtract_gradient, velocities_pair_cur, + pressures_pair_cur) g2 = g2_builder.compile() - swap = True while gui.running: diff --git a/tests/python/examples/__init__.py b/tests/python/examples/__init__.py index 8b137891791fe..e69de29bb2d1d 100644 --- a/tests/python/examples/__init__.py +++ b/tests/python/examples/__init__.py @@ -1 +0,0 @@ - diff --git a/tests/python/examples/autodiff/__init__.py b/tests/python/examples/autodiff/__init__.py index 8b137891791fe..e69de29bb2d1d 100644 --- a/tests/python/examples/autodiff/__init__.py +++ b/tests/python/examples/autodiff/__init__.py @@ -1 +0,0 @@ -