import taichi as ti ti.init(arch=ti.cpu) RES = 1024 K = 2 R = 7 N = K**R Broot = ti.root B = ti.root for r in range(R): B = B.bitmasked(ti.ij, (K, K)) qt = ti.field(ti.f32) B.place(qt) img = ti.Vector.field(3, dtype=ti.f32, shape=(RES, RES)) print('The quad tree layout is:\n', qt.snode) @ti.kernel def action(p: ti.ext_arr()): a = ti.cast(p[0] * N, ti.i32) b = ti.cast(p[1] * N, ti.i32) qt[a, b] = 1 @ti.func
res = 600 dt = 0.03 p_jacobi_iters = 30 f_strength = 10000.0 curl_strength = 0 dye_decay = 0.99 force_radius = res / 3.0 debug = False paused = False ti.init(arch=ti.gpu) _velocities = ti.Vector.field(2, ti.f32, shape=(res, res)) _new_velocities = ti.Vector.field(2, ti.f32, shape=(res, res)) velocity_divs = ti.field(ti.f32, shape=(res, res)) velocity_curls = ti.field(ti.f32, shape=(res, res)) _pressures = ti.field(ti.f32, shape=(res, res)) _new_pressures = ti.field(ti.f32, shape=(res, res)) _dye_buffer = ti.Vector.field(3, ti.f32, shape=(res, res)) _new_dye_buffer = ti.Vector.field(3, ti.f32, shape=(res, res)) class TexPair: def __init__(self, cur, nxt): self.cur = cur self.nxt = nxt def swap(self): self.cur, self.nxt = self.nxt, self.cur
import numpy as np import os import matplotlib.pyplot as plt real = ti.f32 ti.init(default_fp=real) max_steps = 2048 vis_interval = 8 output_vis_interval = 8 steps = 512 seg_size = 256 vis_resolution = 1024 scalar = lambda: ti.field(dtype=real) vec = lambda: ti.Vector.field(2, dtype=real) loss = scalar() hidden = scalar() damping = 0.2 x = vec() v = vec() n_gravitation = 8 goal = vec() goal_v = vec() gravitation = scalar()
N = 80 n_grid = 120 dx = 1 / n_grid inv_dx = 1 / dx dt = 3e-4 p_mass = 1 p_vol = 1 E = 100 mu = E la = E max_steps = 1024 steps = 1024 gravity = 9.8 target = [0.3, 0.6] scalar = lambda: ti.field(dtype=real) vec = lambda: ti.Vector.field(dim, dtype=real) mat = lambda: ti.Matrix.field(dim, dim, dtype=real) x = ti.Vector.field(dim, dtype=real, shape=(max_steps, n_particles), needs_grad=True) x_avg = ti.Vector.field(dim, dtype=real, shape=(), needs_grad=True) v = ti.Vector.field(dim, dtype=real, shape=(max_steps, n_particles), needs_grad=True) grid_v_in = ti.Vector.field(dim, dtype=real, shape=(n_grid, n_grid),
if not is_nan and (word_value <= lens_min or word_value >= lens_max): summed_score += word_value words_in_filter += 1 book_happiness_avg_ti[idx] = summed_score / words_in_filter moby_dick_book = words_of_book(MOBY_DICK_URL, force_lower_case=True) pride_prejucide_book = words_of_book(PRIDE_AND_PREJUDICE_URL, force_lower_case=True) happiness_dict = load_happiness_dict() max_words = max(len(moby_dick_book), len(pride_prejucide_book)) book_happiness = np.zeros(max_words, dtype=np.float32) book_happiness_ti = ti.field(ti.f32, (max_words)) book_happiness_avg_ti = ti.field(ti.f32, (max_words)) for book, name in [(moby_dick_book, "Moby Dick"), (pride_prejucide_book, "Pride and Prejudice")]: book_happiness[:] = 0 curr_book_len = len(book) for idx in range(curr_book_len): word = book[idx] book_happiness[ idx] = happiness_dict[word] if word in happiness_dict else np.nan book_happiness_ti.from_numpy(book_happiness) gfig, gax = plt.subplots() for N in [1000, 3200, 10000]:
def __init__( self, res, size=1, max_num_particles=2**27, # Max 128 MB particles padding=3, unbounded=False, dt_scale=1, E_scale=1, voxelizer_super_sample=2): self.dim = len(res) assert self.dim in ( 2, 3), "MPM solver supports only 2D and 3D simulations." self.t = 0.0 self.res = res self.n_particles = ti.field(ti.i32, shape=()) self.dx = size / res[0] self.inv_dx = 1.0 / self.dx self.default_dt = 2e-2 * self.dx / size * dt_scale self.p_vol = self.dx**self.dim self.p_rho = 1000 self.p_mass = self.p_vol * self.p_rho self.max_num_particles = max_num_particles self.gravity = ti.Vector.field(self.dim, dtype=ti.f32, shape=()) self.source_bound = ti.Vector.field(self.dim, dtype=ti.f32, shape=2) self.source_velocity = ti.Vector.field(self.dim, dtype=ti.f32, shape=()) self.pid = ti.field(ti.i32) # position self.x = ti.Vector.field(self.dim, dtype=ti.f32) # velocity self.v = ti.Vector.field(self.dim, dtype=ti.f32) # affine velocity field self.C = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32) # deformation gradient self.F = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32) # material id self.material = ti.field(dtype=ti.i32) self.color = ti.field(dtype=ti.i32) # plastic deformation volume ratio self.Jp = ti.field(dtype=ti.f32) if self.dim == 2: indices = ti.ij else: indices = ti.ijk offset = tuple(-self.grid_size // 2 for _ in range(self.dim)) self.offset = offset # grid node momentum/velocity self.grid_v = ti.Vector.field(self.dim, dtype=ti.f32) # grid node mass self.grid_m = ti.field(dtype=ti.f32) grid_block_size = 128 self.grid = ti.root.pointer(indices, self.grid_size // grid_block_size) if self.dim == 2: self.leaf_block_size = 16 else: self.leaf_block_size = 8 block = self.grid.pointer(indices, grid_block_size // self.leaf_block_size) def block_component(c): block.dense(indices, self.leaf_block_size).place(c, offset=offset) block_component(self.grid_m) for v in self.grid_v.entries: block_component(v) block.dynamic(ti.indices(self.dim), 1024 * 1024, chunk_size=self.leaf_block_size**self.dim * 8).place( self.pid, offset=offset + (0, )) self.padding = padding # Young's modulus and Poisson's ratio self.E, self.nu = 1e6 * size * E_scale, 0.2 # Lame parameters self.mu_0, self.lambda_0 = self.E / ( 2 * (1 + self.nu)), self.E * self.nu / ((1 + self.nu) * (1 - 2 * self.nu)) # Sand parameters friction_angle = math.radians(45) sin_phi = math.sin(friction_angle) self.alpha = math.sqrt(2 / 3) * 2 * sin_phi / (3 - sin_phi) self.particle = ti.root.dynamic(ti.i, max_num_particles, 2**20) self.particle.place(self.x, self.v, self.C, self.F, self.material, self.color, self.Jp) self.total_substeps = 0 self.unbounded = unbounded if self.dim == 2: self.voxelizer = None self.set_gravity((0, -9.8)) else: if USE_IN_BLENDER: from .voxelizer import Voxelizer else: from engine.voxelizer import Voxelizer self.voxelizer = Voxelizer(res=self.res, dx=self.dx, padding=self.padding, super_sample=voxelizer_super_sample) self.set_gravity((0, -9.8, 0)) self.voxelizer_super_sample = voxelizer_super_sample self.grid_postprocess = [] self.add_bounding_box(self.unbounded) self.writers = []
n_particles, n_grid = 9000 * quality**2, 128 * quality dx, inv_dx = 1 / n_grid, float(n_grid) dt = 1e-4 / quality p_vol, p_rho = (dx * 0.5)**2, 1 p_mass = p_vol * p_rho E, nu = 5e3, 0.2 # Young's modulus and Poisson's ratio mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ( (1 + nu) * (1 - 2 * nu)) # Lame parameters x = ti.Vector.field(2, dtype=float, shape=n_particles) # position v = ti.Vector.field(2, dtype=float, shape=n_particles) # velocity C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # affine velocity field F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # deformation gradient material = ti.field(dtype=int, shape=n_particles) # material id Jp = ti.field(dtype=float, shape=n_particles) # plastic deformation grid_v = ti.Vector.field(2, dtype=float, shape=(n_grid, n_grid)) # grid node momentum/velocity grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) # grid node mass gravity = ti.Vector.field(2, dtype=float, shape=()) attractor_strength = ti.field(dtype=float, shape=()) attractor_pos = ti.Vector.field(2, dtype=float, shape=()) group_size = n_particles // 3 water = ti.Vector.field(2, dtype=float, shape=group_size) # position jelly = ti.Vector.field(2, dtype=float, shape=group_size) # position snow = ti.Vector.field(2, dtype=float, shape=group_size) # position mouse_circle = ti.Vector.field(2, dtype=float, shape=(1, ))
def test_index_mismatch(): with pytest.raises(AssertionError): val = ti.field(ti.i32, shape=(1, 2, 3)) val[0, 0] = 1
import taichi as ti ti.init(arch=ti.cpu) n = 16 a = ti.field(dtype=ti.f32) # ti.root.pointer(ti.i, n // 4).dense(ti.i, 4).place(a) ti.root.dense(ti.i, n).place(a) @ti.kernel def init(): for i in ti.ndrange( (3,7) ): print(i) a[i] = i @ti.kernel def printa(): for i in a: print(f'a[{i}] = {a[i]}') init() printa()
rho0 = 1.0 lambda_epsilon = 100.0 pbf_num_iters = 5 corr_deltaQ_coeff = 0.3 corrK = 0.001 # Need ti.pow() # corrN = 4.0 neighbor_radius = h * 1.05 poly6_factor = 315.0 / 64.0 / math.pi spiky_grad_factor = -45.0 / math.pi old_positions = ti.Vector.field(dim, float) positions = ti.Vector.field(dim, float) velocities = ti.Vector.field(dim, float) grid_num_particles = ti.field(int) grid2particles = ti.field(int) particle_num_neighbors = ti.field(int) particle_neighbors = ti.field(int) lambdas = ti.field(float) position_deltas = ti.Vector.field(dim, float) # 0: x-pos, 1: timestep in sin() board_states = ti.Vector.field(2, float) ti.root.dense(ti.i, num_particles).place(old_positions, positions, velocities) grid_snode = ti.root.dense(ti.ij, grid_size) grid_snode.place(grid_num_particles) grid_snode.dense(ti.k, max_num_particles_per_cell).place(grid2particles) nb_node = ti.root.dense(ti.i, num_particles) nb_node.place(particle_num_neighbors) nb_node.dense(ti.j, max_num_neighbors).place(particle_neighbors)
res = 1024 # gui resolution cmap_name = 'magma_r' # python colormap use_fixed_caxis = 0 # 1: use fixed caxis limits, 0: automatic caxis limits fixed_caxis = [0.0, 5.0] # fixed caxis limits Q = ti.Vector.field(4, dtype=real, shape=(N, N)) # [rho, rho*u, rho*v, rho*e] consv vars Q_old = ti.Vector.field(4, dtype=real, shape=(N, N)) W = ti.Vector.field(4, dtype=real, shape=(N, N)) # [rho, u, v, p] cell avg W_xl = ti.Vector.field(4, dtype=real, shape=(N, N, 3)) # left side of x-face W_xr = ti.Vector.field(4, dtype=real, shape=(N, N, 3)) # right side of x-face W_yl = ti.Vector.field(4, dtype=real, shape=(N, N, 3)) # left side of y-face W_yr = ti.Vector.field(4, dtype=real, shape=(N, N, 3)) # right side of y-face F_x = ti.Vector.field(4, dtype=real, shape=(N, N)) # x-face flux F_y = ti.Vector.field(4, dtype=real, shape=(N, N)) # y-face flux dt = ti.field(dtype=real, shape=()) img = ti.field(dtype=ti.f32, shape=(res, res)) beta_smooth = 1.2 beta_sharp = 2.0 gamma = 1.4 # ratio of specific heats h = 1.0 / (N - 2) # cell size vol = h * h # cell volume @ti.func def is_interior_cell(i, j): return 0 < i < N - 1 and 0 < j < N - 1 @ti.func
import taichi as ti ti.init() n = 512 x = ti.field(dtype=ti.f32, shape=(n, n)) @ti.kernel def paint(): for i, j in ti.ndrange(n * 4, n * 4): # 4x4 super sampling: ret = ti.taichi_logo(ti.Vector([i, j]) / (n * 4)) x[i // 4, j // 4] += ret / 16 def main(): paint() gui = ti.GUI('Logo', (n, n)) while gui.running: gui.set_image(x) gui.show() if __name__ == '__main__': main()
sum = 0.0 for j in range(n): if i != j: sum += Anp[i, j] Anp[i, i] = -sum + 0.1 print('Solving in numpy...') xnp = np.linalg.solve(Anp, bnp) print('Checking the solution...') print(np.allclose(np.dot(Anp, xnp), bnp)) # # -- Solving in Taichi -- # print('Generating Matrix in Taichi...') A = ti.field(dtype=ti.f64) M = ti.field(dtype=ti.f64) ti.root.dense(ti.ij, (n, n)).place(A, M) b = ti.field(dtype=ti.f64) x = ti.field(dtype=ti.f64) x_old = ti.field(dtype=ti.f64) x_new = ti.field(dtype=ti.f64) Ax = ti.field(dtype=ti.f64) Ap = ti.field(dtype=ti.f64) Ap_tld = ti.field(dtype=ti.f64) r = ti.field(dtype=ti.f64) p = ti.field(dtype=ti.f64) z = ti.field(dtype=ti.f64)
import taichi as ti import os ti.init(arch=ti.gpu) max_num_particles = 256 curr_num_particles = 6 dt = 3e-3 num_particles = ti.field(ti.i32, shape=()) x = ti.Vector(2, dt=ti.f32, shape=max_num_particles) inv_m = ti.field(ti.f32, shape=max_num_particles) x_proposed = ti.Vector(2, dt=ti.f32, shape=max_num_particles) v = ti.Vector(2, dt=ti.f32, shape=max_num_particles) # There are two constraint types in this simulation: 0 -- point constraint; 1 -- distance constraint; # constraint_vector[c_idx][0] = c_type; # constraint_vector[c_idx][1] = p1_idx; # constraint_vector[c_idx][2] = p2_idx; c_num = 6 constraint_vector = ti.Vector(3, dt=ti.int32, shape=max_num_particles) gravity = ti.Vector([0, -9.8]) itr_num = 10 @ti.kernel def init(): # Init pos: for i in range(curr_num_particles): # p0 is the position constraint point; if i == 0:
import taichi as ti ti.init(random_seed=42, arch=ti.cpu, debug=True, advanced_optimization=False) import numpy as np # Matplotlib settings. import matplotlib import matplotlib.pyplot as plt import matplotlib.animation as anim # %matplotlib notebook # %matplotlib inline # ------ ti helper funcs dim = 2 real = ti.f32 discrete = lambda: ti.field(dtype=real, shape=(1)) scalar = lambda **kwargs: ti.field(dtype=real, **kwargs) vector = lambda **kwargs: ti.Vector.field(dim, dtype=real, **kwargs) # ------ init ti vars # ti.reset() # reset ti kernel to be allowed to init variables dt = scalar(needs_grad=False) steps = 100 height = 225 width_to_height_ratio = 16 / 9. width = int(width_to_height_ratio * height) dy = 1 / height # difference is smaller of height, width
import matplotlib.pyplot as plt import taichi as ti import math import numpy as np import os real = ti.f32 ti.init(default_fp=real) max_steps = 4096 vis_interval = 256 output_vis_interval = 8 steps = 2048 // 2 assert steps * 2 <= max_steps scalar = lambda: ti.field(dtype=real) vec = lambda: ti.Vector.field(2, dtype=real) loss = scalar() x = vec() v = vec() v_inc = vec() head_id = 10 goal = vec() n_objects = 0 # target_ball = 0 elasticity = 0.0 ground_height = 0.1
import time import numpy as np from renderer_utils import (eps, inf, inside_taichi, intersect_sphere, out_dir, ray_aabb_intersection, sphere_aabb_intersect_motion) import taichi as ti ti.init(arch=ti.cuda, device_memory_GB=4) res = 1280, 720 num_spheres = 1024 color_buffer = ti.Vector.field(3, dtype=ti.f32) bbox = ti.Vector.field(3, dtype=ti.f32, shape=2) grid_density = ti.field(dtype=ti.i32) voxel_has_particle = ti.field(dtype=ti.i32) max_ray_depth = 4 use_directional_light = True particle_x = ti.Vector.field(3, dtype=ti.f32) particle_v = ti.Vector.field(3, dtype=ti.f32) particle_color = ti.Vector.field(3, dtype=ti.f32) pid = ti.field(ti.i32) num_particles = ti.field(ti.i32, shape=()) fov = 0.23 dist_limit = 100 exposure = 1.5 camera_pos = ti.Vector([0.5, 0.32, 2.7])
import taichi as ti ti.init() pixels = ti.field(ti.u8, shape=(512, 512, 3)) @ti.kernel def paint(): for i, j, k in pixels: 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")}')
n_elements = (n_particle_x - 1) * (n_particle_y - 1) * 2 n_grid = 64 * quality dx = 1 / n_grid inv_dx = 1 / dx dt = 1e-4 / quality E = 25000 p_mass = 1 p_vol = 1 mu = 1 la = 1 x = ti.Vector.field(dim, dtype=float, shape=n_particles, needs_grad=True) v = ti.Vector.field(dim, dtype=float, shape=n_particles) C = ti.Matrix.field(dim, dim, dtype=float, shape=n_particles) grid_v = ti.Vector.field(dim, dtype=float, shape=(n_grid, n_grid)) grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) restT = ti.Matrix.field(dim, dim, dtype=float, shape=n_particles) total_energy = ti.field(dtype=float, shape=(), needs_grad=True) vertices = ti.field(dtype=ti.i32, shape=(n_elements, 3)) @ti.func def mesh(i, j): return i * n_particle_y + j @ti.func def compute_T(i): a = vertices[i, 0] b = vertices[i, 1] c = vertices[i, 2]
import matplotlib.pyplot as plt import time from matplotlib.pyplot import cm real = ti.f32 ti.init(default_fp=real) max_steps = 4096 vis_interval = 4 output_vis_interval = 16 steps = 204 assert steps * 2 <= max_steps vis_resolution = 1024 scalar = lambda: ti.field(dtype=real) vec = lambda: ti.Vector.field(2, dtype=real) loss = scalar() x = vec() v = vec() goal = [0.9, 0.15] n_objects = 1 ground_height = 0.1 tmp_x_field = ti.field(shape=(max_steps, n_objects, 2), dtype=real, needs_grad=False)
def test_cpu_debug_snode_reader_out_of_bound_negative(): ti.set_gdb_trigger(False) x = ti.field(ti.f32, shape=3) with pytest.raises(RuntimeError): a = x[-1]
def bls_particle_grid(N, ppc=8, block_size=16, scatter=True, benchmark=0, pointer_level=1, sort_points=True, use_offset=True): M = N * N * ppc m1 = ti.field(ti.f32) m2 = ti.field(ti.f32) m3 = ti.field(ti.f32) pid = ti.field(ti.i32) err = ti.field(ti.i32, shape=()) max_num_particles_per_block = block_size**2 * 4096 x = ti.Vector.field(2, dtype=ti.f32) s1 = ti.field(dtype=ti.f32) s2 = ti.field(dtype=ti.f32) s3 = ti.field(dtype=ti.f32) ti.root.dense(ti.i, M).place(x) ti.root.dense(ti.i, M).place(s1, s2, s3) if pointer_level == 1: block = ti.root.pointer(ti.ij, N // block_size) elif pointer_level == 2: block = ti.root.pointer(ti.ij, N // block_size // 4).pointer(ti.ij, 4) else: raise ValueError('pointer_level must be 1 or 2') if use_offset: grid_offset = (-N // 2, -N // 2) grid_offset_block = (-N // 2 // block_size, -N // 2 // block_size) world_offset = -0.5 else: grid_offset = (0, 0) grid_offset_block = (0, 0) world_offset = 0 block.dense(ti.ij, block_size).place(m1, offset=grid_offset) block.dense(ti.ij, block_size).place(m2, offset=grid_offset) block.dense(ti.ij, block_size).place(m3, offset=grid_offset) block.dynamic(ti.l, max_num_particles_per_block, chunk_size=block_size**2 * ppc * 4).place( pid, offset=grid_offset_block + (0, )) bound = 0.1 extend = 4 x_ = [(random.random() * (1 - 2 * bound) + bound + world_offset, random.random() * (1 - 2 * bound) + bound + world_offset) for _ in range(M)] if sort_points: x_.sort(key=lambda q: int(q[0] * N) // block_size * N + int(q[1] * N) // block_size) x.from_numpy(np.array(x_, dtype=np.float32)) @ti.kernel def insert(): ti.block_dim(256) for i in x: # It is important to ensure insert and p2g uses the exact same way to compute the base # coordinates. Otherwise there might be coordinate mismatch due to float-point errors. base = ti.Vector([ int(ti.floor(x[i][0] * N) - grid_offset[0]), int(ti.floor(x[i][1] * N) - grid_offset[1]) ]) base_p = ti.rescale_index(m1, pid, base) ti.append(pid.parent(), base_p, i) scatter_weight = (N * N / M) * 0.01 @ti.kernel def p2g(use_shared: ti.template(), m: ti.template()): ti.block_dim(256) if ti.static(use_shared): ti.block_local(m) for I in ti.grouped(pid): p = pid[I] u_ = ti.floor(x[p] * N).cast(ti.i32) Im = ti.rescale_index(pid, m, I) u0 = ti.assume_in_range(u_[0], Im[0], 0, 1) u1 = ti.assume_in_range(u_[1], Im[1], 0, 1) u = ti.Vector([u0, u1]) for offset in ti.static(ti.grouped(ti.ndrange(extend, extend))): m[u + offset] += scatter_weight @ti.kernel def p2g_naive(): ti.block_dim(256) for p in x: u = ti.floor(x[p] * N).cast(ti.i32) for offset in ti.static(ti.grouped(ti.ndrange(extend, extend))): m3[u + offset] += scatter_weight @ti.kernel def fill_m1(): for i, j in ti.ndrange(N, N): m1[i, j] = ti.random() @ti.kernel def g2p(use_shared: ti.template(), s: ti.template()): ti.block_dim(256) if ti.static(use_shared): ti.block_local(m1) for I in ti.grouped(pid): p = pid[I] u_ = ti.floor(x[p] * N).cast(ti.i32) Im = ti.rescale_index(pid, m1, I) u0 = ti.assume_in_range(u_[0], Im[0], 0, 1) u1 = ti.assume_in_range(u_[1], Im[1], 0, 1) u = ti.Vector([u0, u1]) tot = 0.0 for offset in ti.static(ti.grouped(ti.ndrange(extend, extend))): tot += m1[u + offset] s[p] = tot @ti.kernel def g2p_naive(s: ti.template()): ti.block_dim(256) for p in x: u = ti.floor(x[p] * N).cast(ti.i32) tot = 0.0 for offset in ti.static(ti.grouped(ti.ndrange(extend, extend))): tot += m1[u + offset] s[p] = tot insert() for i in range(benchmark): pid.parent(2).snode.deactivate_all() insert() @ti.kernel def check_m(): for i in range(grid_offset[0], grid_offset[0] + N): for j in range(grid_offset[1], grid_offset[1] + N): if abs(m1[i, j] - m3[i, j]) > 1e-4: err[None] = 1 if abs(m2[i, j] - m3[i, j]) > 1e-4: err[None] = 1 @ti.kernel def check_s(): for i in range(M): if abs(s1[i] - s2[i]) > 1e-4: err[None] = 1 if abs(s1[i] - s3[i]) > 1e-4: err[None] = 1 if scatter: for i in range(max(benchmark, 1)): p2g(True, m1) p2g(False, m2) p2g_naive() check_m() else: for i in range(max(benchmark, 1)): g2p(True, s1) g2p(False, s2) g2p_naive(s3) check_s() assert not err[None]
def __init__(self, m, n, u, v, cell_type, multigrid_level=4, pre_and_post_smoothing=2, bottom_smoothing=10): self.m = m self.n = n self.u = u self.v = v self.cell_type = cell_type self.multigrid_level = multigrid_level self.pre_and_post_smoothing = pre_and_post_smoothing self.bottom_smoothing = bottom_smoothing # rhs of linear system self.b = ti.field(dtype=ti.f32, shape=(self.m, self.n)) def grid_shape(l): return (self.m // 2**l, self.n // 2**l) # lhs of linear system and its corresponding form in coarse grids self.Adiag = [ ti.field(dtype=ti.f32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] self.Ax = [ ti.field(dtype=ti.f32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] self.Ay = [ ti.field(dtype=ti.f32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] # grid type self.grid_type = [ ti.field(dtype=ti.i32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] # pcg var self.r = [ ti.field(dtype=ti.f32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] self.z = [ ti.field(dtype=ti.f32, shape=grid_shape(l)) for l in range(self.multigrid_level) ] self.p = ti.field(dtype=ti.f32, shape=(self.m, self.n)) self.s = ti.field(dtype=ti.f32, shape=(self.m, self.n)) self.As = ti.field(dtype=ti.f32, shape=(self.m, self.n)) self.sum = ti.field(dtype=ti.f32, shape=()) self.alpha = ti.field(dtype=ti.f32, shape=()) self.beta = ti.field(dtype=ti.f32, shape=())
def bls_test_template(dim, N, bs, stencil, block_dim=None, scatter=False, benchmark=0, dense=False): x, y, y2 = ti.field(ti.i32), ti.field(ti.i32), ti.field(ti.i32) index = ti.indices(*range(dim)) mismatch = ti.field(ti.i32, shape=()) if not isinstance(bs, (tuple, list)): bs = [bs for _ in range(dim)] grid_size = [N // bs[i] for i in range(dim)] if dense: create_block = lambda: ti.root.dense(index, grid_size) else: create_block = lambda: ti.root.pointer(index, grid_size) if scatter: block = create_block() block.dense(index, bs).place(x) block.dense(index, bs).place(y) block.dense(index, bs).place(y2) else: create_block().dense(index, bs).place(x) create_block().dense(index, bs).place(y) create_block().dense(index, bs).place(y2) ndrange = ((bs[i], N - bs[i]) for i in range(dim)) if block_dim is None: block_dim = 1 for i in range(dim): block_dim *= bs[i] @ti.kernel def populate(): for I in ti.grouped(ti.ndrange(*ndrange)): s = 0 for i in ti.static(range(dim)): s += I[i]**(i + 1) x[I] = s @ti.kernel def apply(use_bls: ti.template(), y: ti.template()): if ti.static(use_bls and not scatter): ti.block_local(x) if ti.static(use_bls and scatter): ti.block_local(y) ti.block_dim(block_dim) for I in ti.grouped(x): if ti.static(scatter): for offset in ti.static(stencil): y[I + ti.Vector(offset)] += x[I] else: # gather s = 0 for offset in ti.static(stencil): s = s + x[I + ti.Vector(offset)] y[I] = s populate() if benchmark: for i in range(benchmark): x.snode.parent().deactivate_all() if not scatter: populate() y.snode.parent().deactivate_all() y2.snode.parent().deactivate_all() apply(False, y2) apply(True, y) else: # Simply test apply(False, y2) apply(True, y) @ti.kernel def check(): for I in ti.grouped(y2): if y[I] != y2[I]: print('check failed', I, y[I], y2[I]) mismatch[None] = 1 check() ti.kernel_profiler_print() assert mismatch[None] == 0
import taichi as ti import time ti.init() nx, ny = 10000,10000 p = ti.field(dtype=ti.f32, shape=(nx,ny)) @ti.kernel def main(): for i,j in p: p[i,j] = p[i,j] + 0.1 start = time.time() main() print(time.time()-start)
ti.init(arch=ti.gpu) n_particles, n_grid, n_lines = 9000, 128, 0 dx, inv_dx = 1 / n_grid, float(n_grid) dt = 1e-4 p_vol, p_rho = (dx * 0.5)**2, 1 p_mass = p_vol * p_rho # particle mass E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters x = ti.Vector.field(2, dtype=float, shape=n_particles) # particle position v = ti.Vector.field(2, dtype=float, shape=n_particles) # particle velocity C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # affine velocity field F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # deformation gradient material = ti.field(dtype=int, shape=n_particles) # material id: 0,1,2: fluid Jp = ti.field(dtype=float, shape=n_particles) # plastic deformation grid_v = ti.Vector.field(2, dtype=float, shape=(n_grid, n_grid)) # grid node momentum/velocity grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) # grid node mass @ti.func def P2G(): # Particle to grid (P2G) for p in x: base = (x[p] * inv_dx - 0.5).cast(int) fx = x[p] * inv_dx - base.cast(float) # Bspline 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 mu, la = mu_0 * h, lambda_0 * h mu = 0.0
from taichi.lang import impl import taichi as ti ti.init(arch=ti.cuda, async_mode=True) a = ti.field(dtype=ti.i32, shape=()) b = ti.field(dtype=ti.i32, shape=()) c = ti.field(dtype=ti.i32, shape=()) d = ti.field(dtype=ti.i32, shape=()) e = ti.field(dtype=ti.i32, shape=()) f = ti.field(dtype=ti.i32, shape=()) @ti.kernel def foo(): a[None] += 1 b[None] += 1 c[None] += 1 d[None] += 1 e[None] += 1 f[None] += 1 for i in range(1000): foo() impl.get_runtime().prog.benchmark_rebuild_graph()
import taichi as ti import os import ctypes ti.init() N = 1024 x = ti.field(ti.i32, shape=N) y = ti.field(ti.i32, shape=N) z = ti.field(ti.i32, shape=N) source = ''' extern "C" { void add_and_mul(float a, float b, float *c, float *d, int *e) { *c = a + b; *d = a * b; *e = int(a * b + a); } void pow_int(int a, int b, int *c) { int ret = 1; for (int i = 0; i < b; i++) ret = ret * a; *c = ret; } } ''' with open('a.cpp', 'w') as f: f.write(source) os.system("g++ a.cpp -o a.so -fPIC -shared")
import taichi as ti ti.init() N = 8 dt = 1e-5 x = ti.Vector.field(2, dtype=ti.f32, shape=N, needs_grad=True) # particle positions v = ti.Vector.field(2, dtype=ti.f32, shape=N) # particle velocities U = ti.field(dtype=ti.f32, shape=(), needs_grad=True) # potential energy @ti.kernel def compute_U(): for i, j in ti.ndrange(N, N): r = x[i] - x[j] # r.norm(1e-3) is equivalent to ti.sqrt(r.norm()**2 + 1e-3) # This is to prevent 1/0 error which can cause wrong derivative U[None] += -1 / r.norm(1e-3) # U += -1 / |r| @ti.kernel def advance(): for i in x: v[i] += dt * -x.grad[i] # dv/dt = -dU/dx for i in x: x[i] += dt * v[i] # dx/dt = v def substep():
import pickle import random import numpy as np import taichi as ti # ti.runtime.print_preprocessed = True # ti.cfg.print_ir = True input = ti.field(ti.f32) weight1 = ti.field(ti.f32) output1 = ti.field(ti.f32) output1_nonlinear = ti.field(ti.f32) weight2 = ti.field(ti.f32) output = ti.field(ti.f32) output_exp = ti.field(ti.f32) output_softmax = ti.field(ti.f32) softmax_sum = ti.field(ti.f32) gt = ti.field(ti.f32) loss = ti.field(ti.f32) learning_rate = ti.field(ti.f32) n_input = 28**2 n_hidden = 500 n_output = 10 ti.root.dense(ti.i, n_input).place(input) ti.root.dense(ti.ij, (n_input, n_hidden)).place(weight1)