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
