From 6a15da856e412dc69484cb72d43519fef41a5910 Mon Sep 17 00:00:00 2001 From: Mingrui Zhang <33411325+erizmr@users.noreply.github.com> Date: Tue, 12 Apr 2022 19:59:25 +0800 Subject: [PATCH] [example] Add diff sph demo (#4769) * [example] add diff sph demo * add pretrained models * add comments and rename the file --- MANIFEST.in | 2 + .../examples/autodiff/diff_sph/diff_sph.py | 637 ++++++++++++++++++ .../autodiff/diff_sph/fc1_pretrained.pkl | Bin 0 -> 711 bytes .../autodiff/diff_sph/fc2_pretrained.pkl | Bin 0 -> 2375 bytes 4 files changed, 639 insertions(+) create mode 100644 python/taichi/examples/autodiff/diff_sph/diff_sph.py create mode 100644 python/taichi/examples/autodiff/diff_sph/fc1_pretrained.pkl create mode 100644 python/taichi/examples/autodiff/diff_sph/fc2_pretrained.pkl diff --git a/MANIFEST.in b/MANIFEST.in index d9c29fe03e974..eff0f4f8f88c7 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -6,6 +6,8 @@ include *.cfg include python/taichi/*.md recursive-include python/taichi/assets * recursive-include python/taichi/examples *.py +include python/taichi/examples/autodiff/diff_sph/fc1_pretrained.pkl +include python/taichi/examples/autodiff/diff_sph/fc2_pretrained.pkl include python/taichi/_lib/core/*.so include python/taichi/_lib/core/*.pyd include python/taichi/_lib/runtime/*.bc diff --git a/python/taichi/examples/autodiff/diff_sph/diff_sph.py b/python/taichi/examples/autodiff/diff_sph/diff_sph.py new file mode 100644 index 0000000000000..3070b4aceaeea --- /dev/null +++ b/python/taichi/examples/autodiff/diff_sph/diff_sph.py @@ -0,0 +1,637 @@ +# 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. + +import taichi as ti +import numpy as np +import matplotlib.pyplot as plt +import os +import pickle as pkl +import argparse +parser = argparse.ArgumentParser() +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 +TRAIN_VISUAL_SHOW = False +INFER_OUTPUT_IMG = False +ti.init(arch=ti.gpu, device_memory_fraction=0.5, random_seed=5) +screen_res = (1000, 1000) + +real = ti.f32 +scalar = lambda: ti.field(dtype=real) + + +@ti.data_oriented +class SGD: + def __init__(self, params, lr): + self.params = params + self.lr = lr + + def step(self): + for w in self.params: + self._step(w) + + @ti.kernel + def _step(self, w: ti.template()): + for I in ti.grouped(w): + w[I] -= min(max(w.grad[I], -20.0), 20.0) * self.lr + + def zero_grad(self): + for w in self.params: + w.grad.fill(0.0) + + +@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): + self.n_models = n_models + self.batch_size = batch_size + self.n_steps = n_steps + self.n_input = n_input + self.n_hidden = n_hidden + self.n_output = n_output + self.activation = activation + + self.hidden = scalar() + self.output = scalar() + + # array of structs + self.batch_node = ti.root.dense(ti.i, self.n_models) + 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.weights1 = scalar() + self.bias1 = scalar() + + self.weights1_node.place(self.weights1) + self.n_hidden_node.place(self.bias1) + + if needs_grad: + ti.root.lazy_grad() + + def parameters(self): + return [self.weights1, self.bias1] + + @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): + 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] + 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]) + 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] + + @ti.kernel + def clear(self): + for I in ti.grouped(self.hidden): + self.hidden[I] = 0. + for I in ti.grouped(self.output): + self.output[I] = 0. + + def forward(self, t, nn_input): + self._forward(t, nn_input) + + def dump_weights(self, name="save.pkl"): + w_val = [] + for w in self.parameters(): + w = w.to_numpy() + w_val.append(w[0]) + pkl.dump(w_val, open(name, "wb")) + + def load_weights(self, name="save.pkl", model_id=0): + w_val = pkl.load(open(name, 'rb')) + self.load_weights_from_value(w_val, model_id) + + def load_weights_from_value(self, w_val, model_id=0): + for w, val in zip(self.parameters(), w_val): + if val.shape[0] == 1: + val = val[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): + 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. ") + +# Simulation configuration +boundary_box_np = np.array([[0, 0, 0], [0.5, 1.5, 0.5]]) +spawn_box_np = np.array([[0.0, 0.0, 0.0], [0.5, 0.05, 0.5]]) +target_box_np = np.array([[0.15, 0.90, 0.15], [0.2, 0.95, 0.2]]) + +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 +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) + +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) + + +col = ti.Vector.field(3, float) +material = ti.field(int) +den = ti.field(float) +pre = ti.field(float) + +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.i, int(particle_num)).place(material, col) +ti.root.lazy_grad() + +boundary_box = ti.Vector.field(3, float, shape=2) +spawn_box = ti.Vector.field(3, float, shape=2) +target_box = ti.Vector.field(3, float, shape=2) +N_fluid = ti.Vector([N_np[0], N_np[1], N_np[2]]) +N_target = ti.Vector([N_target_np[0], N_target_np[1], N_target_np[2]]) + +gravity = ti.Vector([0.0, -9.8, 0.0]) + +boundary_box.from_numpy(boundary_box_np) +spawn_box.from_numpy(spawn_box_np) +target_box.from_numpy(target_box_np) + +rest_density = 1000.0 +mass = rest_density * particle_diameter * particle_diameter * particle_diameter * 0.8 +pressure_scale = 10000.0 +viscosity_scale = 0.1 * 3 +tension_scale = 0.005 +gamma = 1.0 +substeps = 5 +dt = 0.016 / substeps +eps = 1e-6 +damping = 0.5 +pi = 3.1415926535 + + +@ti.func +def W_poly6(R, h): + r = R.norm(eps) + res = 0.0 + if r <= h: + h2 = h * h + h4 = h2 * h2 + h9 = h4 * h4 * h + h2_r2 = h2 - r * r + res = 315.0 / (64 * pi * h9) * h2_r2 * h2_r2 * h2_r2 + else: + res = 0.0 + return res + + +@ti.func +def W_spiky_gradient(R, h): + r = R.norm(eps) + res = ti.Vector([0.0, 0.0, 0.0]) + if r == 0.0: + res = ti.Vector([0.0, 0.0, 0.0]) + elif r <= h: + h3 = h * h * h + h6 = h3 * h3 + h_r = h - r + res = -45.0 / (pi * h6) * h_r * h_r * (R / r) + else: + res = ti.Vector([0.0, 0.0, 0.0]) + return res + + +W = W_poly6 +W_gradient = W_spiky_gradient + + +@ti.kernel +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] + ) + 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]) + 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]) + + +@ti.kernel +def initialize_dists(): + for bs in range(batch_size): + min_dist[bs] = 1000.0 + max_height[bs] = 0.0 + max_left[bs] = 0.0 + max_right[bs] = 0.0 + + +@ti.kernel +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] + ) + 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]) + 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]) + + +@ti.kernel +def initialize_density(t: ti.int32): + 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 j in range(particle_num): + R = pos[bs, t, i] - 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): + 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 j in ti.static(range(3)): + 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): + 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: + indicator = (steps - t) // (steps // 2) + acc[bs, t, i] = jet_force[t, bs] + gravity + indicator * (- gravity) * 0.1 + else: + acc[bs, t, i] = gravity + + +@ti.kernel +def update_force(t: ti.int32): + 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] + # 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) + + # 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) + + +@ti.kernel +def advance(t: ti.int32): + 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 + + +@ti.kernel +def boundary_handle(t: ti.int32): + 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] + 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] + 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 + + +@ti.kernel +def compute_dist(t: ti.int32): + for bs, i in ti.ndrange(batch_size, particle_num): + if material[i] == 0: + dist = 0.0 + for j in ti.static(range(3)): + dist += (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.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) + 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]) + + +@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] + + +@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] + + +@ti.kernel +def copy_from_output_to_vis(t: ti.int32): + for i in range(particle_num): + for j in ti.static(range(3)): + pos_vis_buffer[i][j] = pos_output_buffer[t, i][j] + + +@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 j in ti.static(range(3)): + target_centers[i][j] = data[i, j] + print('target_centers ', target_centers[current_pos]) + + +@ti.kernel +def fill_input_states(current_pos: ti.int32): + 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__": + show_window = True + if TRAIN: + show_window = False + window = ti.ui.Window("Diff SPH", screen_res, show_window=show_window) + scene = ti.ui.Scene() + camera = ti.ui.make_camera() + camera.position(0.5, 1.0, 2.0) + camera.up(0.0, 1.0, 0.0) + camera.lookat(0.5, 0.5, 0.5) + camera.fov(70) + scene.set_camera(camera) + canvas = window.get_canvas() + movement_speed = 0.02 + + if TRAIN: + losses = [] + losses_epoch_avg = [] + opt_iters = 7 + 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): + 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) + fc1.clear() + fc2.clear() + with ti.Tape(loss=loss): + for i in range(1, steps): + initialize_density(i - 1) + update_density(i - 1) + update_pressure(i - 1) + fc1.forward(i - 1, input_states) + fc2.forward(i - 1, fc1.output) + controller_output(i - 1) + apply_force(i - 1) + update_force(i - 1) + advance(i) + boundary_handle(i) + if i % substeps == 0: + copy_to_output_buffer(i, 0) + 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]}") + losses.append(loss[None]) + loss_epoch += loss[None] + cnt += 1 + print(f'opt iter {opt_iter} done. Average loss: {loss_epoch / cnt}') + losses_epoch_avg.append(loss_epoch / cnt) + + if TRAIN_VISUAL: + if opt_iter % 1 == 0: + os.makedirs(f"output_img/{opt_iter}", exist_ok=True) + for i in range(1, steps): + 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) + canvas.scene(scene) + if TRAIN_OUTPUT_IMG: + if i % substeps == 0: + 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.') + plt.title("Training Loss") + plt.xlabel("Training Iterations") + plt.ylabel("Loss") + plt.legend() + plt.show() + else: + current_data_offset = 0 + initialize_fluid_particle(0, 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 + + print("Start... ") + cnt = 0 + paused = ti.field(int, shape=()) + while window.running: + window.GUI.begin("Diff SPH", 0.05, 0.05, 0.2, 0.2) + 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) + window.GUI.end() + + if not paused[None]: + fill_input_states(current_data_offset) + initialize_target_particle(0, pos, N_target, current_data_offset) + fc1.clear() + fc2.clear() + for i in range(1, substeps): + initialize_density(i - 1) + update_density(i - 1) + update_pressure(i - 1) + fc1.forward(i - 1, input_states) + fc2.forward(i - 1, fc1.output) + controller_output(i - 1) + apply_force(i - 1) + update_force(i - 1) + advance(i) + boundary_handle(i) + copy_to_vis(substeps - 1, 0) + copy_back(substeps - 1) + cnt += 1 + # user controlling of camera + position_change = ti.Vector([0.0, 0.0, 0.0]) + up = ti.Vector([0.0, 1.0, 0.0]) + # move camera up and down + if window.is_pressed("e"): + position_change = up * movement_speed + if window.is_pressed("q"): + position_change = -up * movement_speed + + for e in window.get_events(ti.ui.PRESS): + if e.key == ti.ui.SPACE: + 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) + + 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) + 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/autodiff/diff_sph/fc1_pretrained.pkl b/python/taichi/examples/autodiff/diff_sph/fc1_pretrained.pkl new file mode 100644 index 0000000000000000000000000000000000000000..c6791e6b3c9432057830171e4ecb0029abc54bf1 GIT binary patch literal 711 zcmZo*jxA)+NX{$GEvVE>&M!*U%Pq|*$xJLNO049HFG@|$&nqq|DorloDrAJH;>t^b z$QCkrGqe^mJ2NH~va}bn1{Jbtcr$t{cr&*ZvO|@olvEa^ausq!FflMNq?r_Qc6PKE zas?G~Yj`t9FapJF3VHne{Jj1H0hsV+C@JJk>U3%VO7cvaY?J-Yes7;|`o4=!pRBtb z=GjHNFSgzFt6-01&HR0?2|atCew}P{Yt8~Y+l`ZK{Tl4vz5;_vGdk%*?ay-{l4pemhRP`IDcQE$@0Ag`4{Z;<_KHQR-Ly` zoO$V9rPr(N+!jvQ$JeZC7x-!N9^*|}``#X(wTDM!iCz9toqeaWX4&?wGqzcizi{71 zrO$g$C-m*BD&B9KyhUy=(}|4TA!_sX$^Kokx2Wc#&G&Pc_AK0(y3gFGefON}%j|A{ z?y}`y(qR2KdDXs~Wi$6GGAytT__EyYB4dxOvkBKePwDP`*@s^5ogC}7|E}J@Juh>6 z_sxql+*|&!$hNgX+in?`>fY~ejr%4x%-?G_Ve!6ojmP&M2)$ufZhY5TAa3zK&D+!W zD)Ozcd-<_zA8(DiT_Mw!J?FP3?W=8_yeIR>Ji8f6virmywcGwE7TA+lzGUCYhp+e6 zX-?iZVds3?4NKJb?zQx=xp;m4K5m!AdkgtW3i*>V7&E}}ERX?=WWk_9A#fzO77Ax@ zcRDw4U01f{f#AI$n&Z0iUiIno?6|Hk-fNM{xtAL%2U5py`JfGuzc(nnc`uL;QV*gr U^eYn3?}ntmP^6?#G$~aN07ry42><{9 literal 0 HcmV?d00001 diff --git a/python/taichi/examples/autodiff/diff_sph/fc2_pretrained.pkl b/python/taichi/examples/autodiff/diff_sph/fc2_pretrained.pkl new file mode 100644 index 0000000000000000000000000000000000000000..6c73c5c75b6a47ffd2d67bdefcc7b3b4d28cc415 GIT binary patch literal 2375 zcmchZ*;`as7R3utQ6>dsZ~!6-RRSV~Ak;l|&n<&W4O9swA|SRh89_h`3PnW4Fv;Km zI6wm;4snbTB|z1==UgWswgb*2ih(GmZL|?Z#38hyP5%wOAJ+c%-s@$3`{kEyv|Pgz z#HOStuiNMuyEZl6GdVp;n-HCv8okkEMQVKP+LSbHYI>~JL}T!m%OoZ4udBvT!OPPa z1sKF=MippAt2JW;3IoMgiV67|VsOn`p*{8}fL(RX&<=h1&}>CTfkTKw%Va zFtQVC%v36s{PReC5d}}HF^?(ynm3lmn>%U&EA%WtcgzexskQ};n}9G+)50YV*D=HL zF?eNI4|j5QGT67ph|#V}#%%?!#KI^2U|g$;YoPO}HP7;}cjzK$b4pG=SA@V@z3sT= zo{&7~3WJ{>xR5Ay6|S^WkVunaj=#y7lw8x1N&BQk+ZhT^X#RzfR>dB zu&H7PtGMaG(c2+2ePIDueOxXHOn0CQUo}H2a4srIehVdeW^~6-q4=8wH@eCqPbYuf z3VLHI(eQc~aHCg>wD&Hfbw+XMhN>L7-pm1EN^kIPXDzTxsYb?kEWr;$m-NDI;rhbV zd;0Bz#q60G5x~MP4F&$~DCpnI;tyGE+^ji=m`zvA@Y%*5F6QA3aD2HDQ)93UZ#glV z_88`i2R(v0`{EgTwL=k>Jy*cyDN^EmRt_(<-oke{KcWMnaG`iQ*|qy1US%sMZt`{SIw-n6x!7yiyh%BGdMaH_O<%!6abVxw4@ zq$GM8ouNgr=(RsG`Y;4n9QCK$cC5j}FCpFK)32*~e-M-y??+|VzDFhT)xi2@CgqT| z130d9K%Wj=L`#NS^s@y4{u8I{NAH;3pvr0vie_^_k-ZCQ5}siBc9oznaubRgETobS zU%*$lwsOyPbqsy6lZ}e);o`Kp)V5ut7_>MA=Ky!w>Xi_L_62h>DbsRO8227n#!i;Ia&vnb#^doY^+_s|j9JO2eMhlS`q4zsHoSqOiUf3h zP&%&VFQnf`eh<7Kcmeo_0`&ESq31^$QFnvA-q`jW8nkmo>9aFY){7D#3`2T;|9-Hi zJsPw}+k(u#W;D?(cqF%gt$%q69kC|VnfNHsnZ)ALoh=-SKg2XP^ssHJ9?sd+5F}kO zWIEc__X|WO>JTE5}dqUuSbQ@P3l8~nB!O$5v5oK;J zR_BF~%3q5(eh{DRPSBC7jEv0vHUcgjQj(bOjLB^C1=!^qK$ll|;k7lDtWT#4cYhhf zBn{M4@0ZFXOvilMznH+>sCW7k=q-G)b`D*)RD+k-`_cD(n^1Y38YrEnfzr_zK>3F0 z=tyb{lKrNjrnT7UYc~1-_g+4zE4&0|2bCev)gj8b+eH~QDM9JeCegtIN>){q3M?|l zqPX^bdcW3IT<6rr^@Y?ilZ*%0#dCYO>vIS6874-|&C}6%ORf`b)jJ=&J)z=qj4P1P zBO6;Dm%(I9842hNgp=<-z^Bw+BrE+4J z6b7${h7i}rvE)C^0_;M9X|6|wD>v_D-6{m!jsgjjUojJ0IVY2(z4M}%C^1~w@wf z8@uEwcY1(5d{=SU%~62 z(8SUr6h6G=M8Y3d>Dh!>;s%}C)+ASiHcSE7Kql^sQuprd~GjaT5e|k99 z3%{1FQaFoV@<0b|w&tJ<0q@|)ovyUqNhLnNY91|cXa;uT znTRT_(mOp2M@PHLkTk9wC8qW1Cq25Xdu{HEu03t1SlgG>GtXlvx3w1ZSh;}AYXzwG zunjU76|vvCtUxiD^MR*t1c(f6!2B&O+|U22W324Qj|m%y+=KG|QNjWqd%Jy*l&DLl)qpI-rQ_Y5Ab90A98m(WwRCdnxRoqLg0I=o1z;^^p@QxMGr%AL6sv$!8a0FJLPe3$Iu8|v^;<; zs~Ygh6|LMhwvK5uv%@=Uo^riDJh0!-fa#O0!5vF}7oU6NE`F80i2LA?4T9NX{GO+R z2NX(@*|88V_WuQIdPL;E|GFW-^EH1=?VZv*Cslq*SiD11W&9Bl!hr5?ir^x`!?=Yj8 zvx~4=Ih9_}<_b&`s?pR^2^yM{h{AQ_^*e_Q0B?&HtYyo