def recast_depth_to_map(self, depthmap: ti.ext_arr(), texture: ti.ext_arr(), w: ti.i32, h: ti.i32, K:ti.ext_arr(), Kcolor:ti.ext_arr()): fx = K[0] fy = K[4] cx = K[2] cy = K[5] for j in range(h): for i in range(w): if depthmap[j, i] == 0 or depthmap[j, i]/1000 > ti.static(self.max_ray_length): continue dep = depthmap[j, i]/1000.0 pt = ti.Vector([ (i-cx)*dep/fx, (j-cy)*dep/fy, dep]) pt_map = self.input_R[None]@pt + self.input_T[None] if ti.static(self.TEXTURE_ENABLED): fx_c = Kcolor[0] fy_c = Kcolor[4] cx_c = Kcolor[2] cy_c = Kcolor[5] color_i = ti.cast((i-cx)/fx*fx_c+cx_c, ti.int32) color_j = ti.cast((j-cy)/fy*fy_c+cy_c, ti.int32) if color_i < 0 or color_i >= 640 or color_j < 0 or color_j >= 480: continue self.process_point(pt_map, [texture[color_j, color_i, 0], texture[color_j, color_i, 1], texture[color_j, color_i, 2]]) else: self.process_point(pt_map)
def initialize_particle(self, x: ti.ext_arr(), v: ti.ext_arr(), color: ti.ext_arr()): for i in range(self.num_particles[None]): for c in ti.static(range(3)): self.particle_x[i][c] = x[i, c] self.particle_v[i][c] = v[i, c] self.particle_color[i][c] = color[i, c]
def init_particles(p: ti.ext_arr(), v: ti.ext_arr()): for i in range(cur_num_particles[None]): velocity[i] = ti.Vector([v[i, 0], -5.0]) for c in ti.static(range(dim)): old_positions[i][c] = p[i, c] # velocity[i][c] = v[i, c] density[i] = rho0
def init(p: ti.ext_arr(), v: ti.ext_arr()): for i in range(num_particles): radius[i] = particle_radius_in_world for c in ti.static(range(dim)): positions[i][c] = p[i, c] velocities[i][c] = v[i, c] board_states[None] = ti.Vector([boundary[0] - epsilon, -0.0])
def initialize_particle(self, x: ti.ext_arr(), v: ti.ext_arr(), color: ti.ext_arr(), begin: ti.i32, end: ti.i32): for i in range(begin, end): for c in ti.static(range(3)): self.particle_x[i][c] = x[i - begin, c] if ti.static(self.enable_motion_blur): self.particle_v[i][c] = v[i - begin, c] self.particle_color[i][c] = color[i - begin, c]
def readframe(self, f: ti.i32, x: ti.ext_arr(), v: ti.ext_arr(), F: ti.ext_arr(), C: ti.ext_arr()): for i in range(self.n_particles): for j in ti.static(range(self.dim)): x[i, j] = self.x[f, i][j] v[i, j] = self.v[f, i][j] for k in ti.static(range(self.dim)): F[i, j, k] = self.F[f, i][j, k] C[i, j, k] = self.C[f, i][j, k]
def setframe(self, f: ti.i32, x: ti.ext_arr(), v: ti.ext_arr(), F: ti.ext_arr(), C: ti.ext_arr()): for i in range(self.n_particles): for j in ti.static(range(self.dim)): self.x[f, i][j] = x[i, j] self.v[f, i][j] = v[i, j] for k in ti.static(range(self.dim)): self.F[f, i][j, k] = F[i, j, k] self.C[f, i][j, k] = C[i, j, k]
def recast_pcl_to_map(self, xyz_array: ti.ext_arr(), rgb_array: ti.ext_arr(), n: ti.i32): for index in range(n): pt = ti.Vector([ xyz_array[index,0], xyz_array[index,1], xyz_array[index,2]]) pt = self.input_R[None]@pt + self.input_T[None] if ti.static(self.TEXTURE_ENABLED): self.process_point(pt, rgb_array[index]) else: self.process_point(pt)
def initialize_particle_x(x: ti.ext_arr(), v: ti.ext_arr(), color: ti.ext_arr()): for i in range(num_particles[None]): for c in ti.static(range(3)): particle_x[i][c] = x[i, c] particle_v[i][c] = v[i, c] particle_color[i][c] = color[i, c] for k in ti.static(range(27)): base_coord = (inv_dx * particle_x[i] - 0.5).cast(ti.i32) + ti.Vector( [k // 9, k // 3 % 3, k % 3]) grid_density[base_coord // grid_visualization_block_size] = 1
def init(self, p_list:ti.ext_arr(), w_list:ti.ext_arr()): for i in range(self.particle_numbers): for j in ti.static(range(self.dim)): self.particle_positions[i][j] = p_list[i,j] self.particle_velocity[i][j] = ti.cast(0.0, ti.f32) self.d_velocity[i][0] = ti.cast(0.0, ti.f32) self.d_velocity[i][1] = ti.cast(-9.8, ti.f32) self.wall_mark_list[i][0] = w_list[i] self.d_density[i][0] = ti.cast(0.0, ti.f32) self.particle_pressure[i][0] = ti.cast(0.0, ti.f32) self.particle_density[i][0] = ti.cast(1000.0, ti.f32)
def copyback_grid(np_idx: ti.ext_arr(), np_val: ti.ext_arr(), grid: ti.template(), solver: ti.template()): num_active_cell = np_idx.shape[0] for i in range(num_active_cell): idx = [] val = [] for j in ti.static(range(solver.dim)): idx.append(int(np_idx[i, j])) val.append(np_val[i, j]) ti_idx = ti.Vector(idx) ti_val = ti.Vector(val) grid[ti_idx] = ti_val
def add_obj(self, vn: ti.i32, en: ti.i32, node: ti.ext_arr(), element: ti.ext_arr()): for i in range(vn): self.node[self.vn_object_index[self.count] + i] = [node[i, 0], node[i, 1]] self.prev_node[self.vn_object_index[self.count] + i] = [node[i, 0], node[i, 1]] self.prev_t_node[self.vn_object_index[self.count] + i] = [node[i, 0], node[i, 1]] self.bar_node[self.vn_object_index[self.count] + i] = [node[i, 0], node[i, 1]] self.node_obj_idx[self.vn_object_index[self.count] + i] = self.count + 1 for i in range(en): # Mapping single object element id to system-wide self.element[self.en_object_index[self.count] + i] = \ [self.vn_object_index[self.count] + element[i, 0], self.vn_object_index[self.count] + element[i, 1], self.vn_object_index[self.count] + element[i, 2]] self.element_obj_idx[self.en_object_index[self.count] + i] = self.count + 1 # update vn_object_index and en_object_index self.vn_object_index[self.count + 1] = self.vn_object_index[self.count] + vn self.en_object_index[self.count + 1] = self.en_object_index[self.count] + en self.count += 1 for i in range(self.en_object_index[self.count - 1], self.en_object_index[self.count]): D = self.D(i) self.B[i] = D.inverse() a, b, c = self.element[i][0], self.element[i][1], self.element[i][ 2] self.element_volume[i] = abs(D.determinant()) / 2 # space in 2d self.element_mass[i] = self.element_volume[i] self.node_mass[a] += self.element_mass[i] self.node_mass[b] += self.element_mass[i] self.node_mass[c] += self.element_mass[i] self.neighbor_element_count[a] += 1 self.neighbor_element_count[b] += 1 self.neighbor_element_count[c] += 1 for i in range(self.vn_object_index[self.count - 1], self.vn_object_index[self.count]): self.node_mass[i] /= max(self.neighbor_element_count[i], 1)
def initialize_particle_x(x: ti.ext_arr(), v: ti.ext_arr(), color: ti.ext_arr()): for i in range(max_num_particles): if i < num_particles: for c in ti.static(range(3)): particle_x[i][c] = x[i, c] particle_v[i][c] = v[i, c] particle_color[i][c] = color[i, c] # reconstruct grid using particle position and MPM p2g. for k in ti.static(range(27)): base_coord = (inv_dx * particle_x[i] - 0.5).cast( ti.i32) + ti.Vector([k // 9, k // 3 % 3, k % 3]) grid_density[base_coord // grid_visualization_block_size] = 1
def vector_to_fast_image(img: ti.template(), out: ti.ext_arr()): # FIXME: Why is ``for i, j in img:`` slower than: for i, j in ti.ndrange(*img.shape): u, v, w = min(255, max(0, int(img[i, img.shape[1] - 1 - j] * 255))) # We use i32 for |out| since OpenGL and Metal doesn't support u8 types # TODO: treat Cocoa and Big-endian machines, with XOR logic out[j * img.shape[0] + i] = w + (v << 8) + (u << 16)
def apply_mouse_input_and_render(self, vf: ti.template(), dyef: ti.template(), imp_data: ti.ext_arr(), dt: ti.template()): for i, j in vf: mdir = ti.Vector([imp_data[0], imp_data[1]]) omx, omy = imp_data[2], imp_data[3] # move to cell center dx, dy = (i + 0.5 - omx), (j + 0.5 - omy) d2 = dx * dx + dy * dy # ref: https://developer.download.nvidia.cn/books/HTML/gpugems/gpugems_ch38.html # apply the force factor = ti.exp(-d2 * self.cfg.inv_force_radius) momentum = mdir * self.cfg.f_strength * dt * factor vf[i, j] += momentum # add dye dc = dyef[i, j] # TODO what the hell is this? if mdir.norm() > 0.5: dc += ti.exp(-d2 * self.cfg.inv_dye_denom) * ti.Vector( [imp_data[4], imp_data[5], imp_data[6]]) dc *= self.cfg.dye_decay dyef[i, j] = dc
def set_color_by_material(material_colors: ti.ext_arr()): 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 ])
def copy_back_and_clear(img: ti.ext_arr()): for i in range(res[0]): for j in range(res[1]): coord = ((res[1] - 1 - j) * res[0] + i) * 3 for c in ti.static(range(3)): img[coord + c] = screen[i, j][2 - c] screen[i, j][2 - c] = 0
def set_vertices(self, data: ti.ext_arr()): self.n_vertices[None] = min(ti.Expr(self.pos.shape[0]), data.shape[0]) for i in range(self.n_vertices[None]): for j in ti.static(range(3)): self.pos[i][j] = data[i, 0, j] self.tex[i][j] = data[i, 1, j] self.nrm[i][j] = data[i, 2, j]
def copy_grid(np_idx: ti.ext_arr(), np_val: ti.ext_arr(), grid: ti.template(), solver: ti.template()): """ save the sparse grid_v or grid_m :param np_idx: :param np_val: :param grid: :return: """ i = 0 for I in ti.grouped(grid): j = ti.atomic_add(i, 1) # print(j, I) for d in ti.static(range(solver.dim)): np_idx[j, d] = I[d] np_val[j, d] = grid[I][d]
def translate(self, vec3: ti.ext_arr()): _vec3 = ti.Vector([vec3[0], vec3[1], vec3[2]]) for i in ti.ndrange(self.ver.shape[0]): self.ver[i] = self.ver[i] + _vec3 # 更新bounding box的边界范围 self.minmax_coordinate()
def render(self, out: ti.ext_arr(), res: ti.template()): for i in range(res[0] * res[1]): r, g, b = self.image_at(i % res[0], res[1] - 1 - i // res[0]) if ti.static(ti.get_os_name() != 'osx'): out[i] = (r << 16) + (g << 8) + b else: alpha = -16777216 out[i] = (b << 16) + (g << 8) + r + alpha
def scale(self, vec3: ti.ext_arr()): _ti_mat = ti.Matrix([[vec3[0], 0, 0], [0, vec3[1], 0], [0, 0, vec3[2]]]) for i in ti.ndrange(self.ver.shape[0]): self.ver[i] = _ti_mat @ self.ver[i] # 更新bounding box的边界范围 self.minmax_coordinate()
def step_flowfield(z: ti.f32, ff: ti.ext_arr(), nsarr: ti.ext_arr()): t = z for x in range(nx): for y in range(ny): # [0:2]: normalized delta direction of flow # [2:4]: current mouse xy # [4:7]: color ns = nsarr[x, y] # ns = 1 radians = ns * math.tau c = ti.cos(radians) s = ti.sin(radians) v = ti.Vector([c, s]) v.normalized() c = v[0] s = v[1] ff[x, y, 0] = -c ff[x, y, 1] = -s cx = ff[x, y, 2] cy = ff[x, y, 3] ## these store the velocity speed = 1 # vel accel = 0.05 ax = (accel * c) ay = (accel * s) # x, y = np.linalg.norm((sx,sy)) ## velocity minS = -3 maxS = 3 ff[x, y, 8] = min(max(ff[x, y, 8] + ax, minS), maxS) ff[x, y, 9] = min(max(ff[x, y, 9] + ay, minS), maxS) # ff[x, y, 8] = ff[x, y, 8] * sx # ff[x, y, 9] = Nx = cx + ff[x, y, 8] Ny = cy + ff[x, y, 9] dx = Nx - cx dy = Ny - cy ff[x, y, 2] = Nx % res ff[x, y, 3] = Ny % res ### Colors ff[x, y, 4] = ns ff[x, y, 5] = ns * ns ff[x, y, 6] = 1 - ns ff[x, y, 7] = 0.0
def ext_arr_to_matrix(arr: ti.ext_arr(), mat: ti.template(), as_vector: ti.template()): for I in ti.grouped(mat): for p in ti.static(range(mat.n)): for q in ti.static(range(mat.m)): if ti.static(as_vector): mat[I][p] = arr[I, p] else: mat[I][p, q] = arr[I, p, q]
def reset_kernel(self, x: ti.ext_arr()): for i in range(self.n_particles): for j in ti.static(range(self.dim)): self.x[0, i][j] = x[i, j] self.v[0, i] = ti.Vector.zero(self.dtype, self.dim) self.F[0, i] = ti.Matrix.identity( self.dtype, self.dim) #ti.Matrix([[1, 0], [0, 1]]) self.C[0, i] = ti.Matrix.zero(self.dtype, self.dim, self.dim)
def set_color(ti_color: ti.template(), material_color: ti.ext_arr(), ti_material: ti.template()): for I in ti.grouped(ti_material): material_id = ti_material[I] color_4d = ti.Vector([0.0, 0.0, 0.0, 1.0]) for d in ti.static(range(3)): color_4d[d] = material_color[material_id, d] ti_color[I] = color_4d
def hub_get_image(imgout: ti.ext_arr()): for I in ti.grouped(img): if ti.static(isinstance(img, ti.Matrix)): for j in ti.static(range(img.n)): imgout[I, j] = cook_color(img[I][j]) else: val = cook_color(img[I]) for j in ti.static(range(3)): imgout[I, j] = val
def voxelize_triangles(self, num_triangles: ti.i32, triangles: ti.ext_arr()): for i in range(num_triangles): jitter_scale = ti.cast(0, self.precision) if ti.static(self.precision is ti.f32): jitter_scale = 1e-4 else: jitter_scale = 1e-8 # We jitter the vertices to prevent voxel samples from lying precicely at triangle edges jitter = ti.Vector([-0.057616723909439505, -0.25608986292614977, 0.06716309129743714]) * jitter_scale a = ti.Vector([triangles[i, 0], triangles[i, 1], triangles[i, 2]]) + jitter b = ti.Vector([triangles[i, 3], triangles[i, 4], triangles[i, 5]]) + jitter c = ti.Vector([triangles[i, 6], triangles[i, 7], triangles[i, 8]]) + jitter bound_min = ti.Vector.zero(self.precision, 3) bound_max = ti.Vector.zero(self.precision, 3) for k in ti.static(range(3)): bound_min[k] = min(a[k], b[k], c[k]) bound_max[k] = max(a[k], b[k], c[k]) p_min = int(ti.floor(bound_min[0] * self.inv_dx)) p_max = int(ti.floor(bound_max[0] * self.inv_dx)) + 1 p_min = max(self.padding, p_min) p_max = min(self.res[0] - self.padding, p_max) q_min = int(ti.floor(bound_min[1] * self.inv_dx)) q_max = int(ti.floor(bound_max[1] * self.inv_dx)) + 1 q_min = max(self.padding, q_min) q_max = min(self.res[1] - self.padding, q_max) normal = ti.normalized(ti.cross(b - a, c - a)) if abs(normal[2]) < 1e-10: continue a_proj = ti.Vector([a[0], a[1]]) b_proj = ti.Vector([b[0], b[1]]) c_proj = ti.Vector([c[0], c[1]]) for p in range(p_min, p_max): for q in range(q_min, q_max): pos2d = ti.Vector([(p + 0.5) * self.dx, (q + 0.5) * self.dx]) if inside_ccw(pos2d, a_proj, b_proj, c_proj) or inside_ccw(pos2d, a_proj, c_proj, b_proj): base_voxel = ti.Vector([pos2d[0], pos2d[1], 0]) height = int( -ti.dot(normal, base_voxel - a) / normal[2] * self.inv_dx + 0.5) height = min(height, self.res[1] - self.padding) inc = 0 if normal[2] > 0: inc = 1 else: inc = -1 self.fill(p, q, height, inc)
def add_texture_2d(self, offset_x: ti.f32, offset_y: ti.f32, texture: ti.ext_arr()): for i, j in ti.ndrange(texture.shape[0], texture.shape[1]): if texture[i, j] > 0.1: pid = ti.atomic_add(self.n_particles[None], 1) x = ti.Vector([ offset_x + i * self.dx * 0.5, offset_y + j * self.dx * 0.5 ]) self.seed_particle(pid, x, self.material_elastic, 0xFFFFFF, self.source_velocity[None])
def recast_depth_to_map_debug(self, depthmap: ti.ext_arr(), rgb_array: ti.ext_arr(), w: ti.i32, h: ti.i32, K:ti.ext_arr()): fx = K[0] fy = K[4] cx = K[2] cy = K[5] self.num_export_particles[None] = 0 for j in range(h): for i in range(w): if depthmap[j, i] == 0 or depthmap[j, i]/1000 > ti.static(self.max_ray_length): continue dep = depthmap[j, i]/1000.0 pt = ti.Vector([ (i-cx)*dep/fx, (j-cy)*dep/fy, dep]) pt = self.input_R@pt + self.input_T index = ti.atomic_add(self.num_export_particles[None], 1) self.export_x[index] = pt