def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: 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 matrix_to_ext_arr(mat: template(), arr: ext_arr(), as_vector: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
if static(as_vector):
arr[I, p] = mat[I][p]
else:
arr[I, p, q] = mat[I][p, q]
def ext_arr_to_matrix(arr: ndarray_type.ndarray(), mat: template(),
as_vector: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
if static(as_vector):
mat[I][p] = arr[I, p]
else:
mat[I][p, q] = arr[I, p, q]
def gen_normals_kernel(vertices: template(), normals: template()):
N = vertices.shape[0]
for i in range(N / 3):
a = vertices[i * 3]
b = vertices[i * 3 + 1]
c = vertices[i * 3 + 2]
n = (a - b).cross(a - c).normalized()
normals[i * 3] = n
normals[i * 3 + 1] = n
normals[i * 3 + 2] = n
def arr_vulkan_layout_to_field_normal_layout(vk_arr: ndarray_type.ndarray(),
normal_field: template()):
static_assert(len(normal_field.shape) == 2)
w = normal_field.shape[0]
h = normal_field.shape[1]
for i, j in ndrange(w, h):
normal_field[i, j] = vk_arr[(h - 1 - j) * w + i]
def ndarray_matrix_to_ext_arr(ndarray: ndarray_type.ndarray(),
arr: ndarray_type.ndarray(),
layout_is_aos: template(),
as_vector: template()):
for I in grouped(ndarray):
for p in static(range(ndarray[I].n)):
for q in static(range(ndarray[I].m)):
if static(as_vector):
if static(layout_is_aos):
arr[I, p] = ndarray[I][p]
else:
arr[p, I] = ndarray[I][p]
else:
if static(layout_is_aos):
arr[I, p, q] = ndarray[I][p, q]
else:
arr[p, q, I] = ndarray[I][p, q]
def sort_stage(keys: template(), use_values: int, values: template(), N: int,
p: int, k: int, invocations: int):
for inv in range(invocations):
j = k % p + inv * 2 * k
for i in range(0, ops.min(k, N - j - k)):
a = i + j
b = i + j + k
if int(a / (p * 2)) == int(b / (p * 2)):
key_a = keys[a]
key_b = keys[b]
if key_a > key_b:
keys[a] = key_b
keys[b] = key_a
if use_values != 0:
temp = values[a]
values[a] = values[b]
values[b] = temp
def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(),
as_vector: template()):
for I in grouped(ndarray):
for p in static(range(ndarray[I].n)):
for q in static(range(ndarray[I].m)):
if static(as_vector):
ndarray[I][p] = arr[I, p]
else:
ndarray[I][p, q] = arr[I, p, q]
def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(),
as_vector: template()):
for I in ti.grouped(ndarray):
for p in ti.static(range(ndarray[I].n)):
for q in ti.static(range(ndarray[I].m)):
if ti.static(as_vector):
arr[I, p] = ndarray[I][p]
else:
arr[I, p, q] = ndarray[I][p, q]
def gen_normals_kernel_indexed(vertices: template(), indices: template(),
normals: template(), weights: template()):
num_triangles = indices.shape[0] / 3
num_vertices = vertices.shape[0]
for i in range(num_vertices):
normals[i] = Vector([0.0, 0.0, 0.0])
weights[i] = 0.0
for i in range(num_triangles):
i_a = indices[i * 3]
i_b = indices[i * 3 + 1]
i_c = indices[i * 3 + 2]
a = vertices[i_a]
b = vertices[i_b]
c = vertices[i_c]
n = (a - b).cross(a - c).normalized()
atomic_add(normals[i_a], n)
atomic_add(normals[i_b], n)
atomic_add(normals[i_c], n)
atomic_add(weights[i_a], 1.0)
atomic_add(weights[i_b], 1.0)
atomic_add(weights[i_c], 1.0)
for i in range(num_vertices):
if weights[i] > 0.0:
normals[i] = normals[i] / weights[i]
def vector_to_fast_image(img: template(), out: ndarray_type.ndarray()):
# FIXME: Why is ``for i, j in img:`` slower than:
for i, j in ndrange(*img.shape):
r, g, b = 0, 0, 0
color = img[i, img.shape[1] - 1 - j]
if static(img.dtype in [f16, f32, f64]):
r, g, b = min(255, max(0, int(color * 255)))
else:
static_assert(img.dtype == u8)
r, g, b = color
idx = j * img.shape[0] + i
# We use i32 for |out| since OpenGL and Metal doesn't support u8 types
if static(get_os_name() != 'osx'):
out[idx] = (r << 16) + (g << 8) + b
else:
# What's -16777216?
#
# On Mac, we need to set the alpha channel to 0xff. Since Mac's GUI
# is big-endian, the color is stored in ABGR order, and we need to
# add 0xff000000, which is -16777216 in I32's legit range. (Albeit
# the clarity, adding 0xff000000 doesn't work.)
alpha = -16777216
out[idx] = (b << 16) + (g << 8) + r + alpha
def fill_ndarray(ndarray: ndarray_type.ndarray(), val: template()):
for I in grouped(ndarray):
ndarray[I] = val
def snode_deactivate_dynamic(b: template()):
for I in grouped(b.parent()):
deactivate(b, I)
def snode_deactivate(b: template()):
for I in grouped(b):
deactivate(b, I)
def fill_matrix(mat: template(), vals: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
mat[I][p, q] = vals[p][q]
def clear_loss(l: template()):
# Using SNode writers would result in a forced sync, therefore we wrap these
# writes into a kernel.
l[None] = 0
l.grad[None] = 1
def clear_gradients(_vars: template()):
for I in grouped(ScalarField(Expr(_vars[0]))):
for s in static(_vars):
ScalarField(Expr(s))[I] = 0
def tensor_to_image(tensor: template(), arr: ext_arr()):
for I in ti.grouped(tensor):
t = ti.cast(tensor[I], ti.f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
def tensor_to_ext_arr(tensor: template(), arr: ndarray_type.ndarray()):
for I in grouped(tensor):
arr[I] = tensor[I]
def vector_to_image(mat: template(), arr: ext_arr()):
for I in ti.grouped(mat):
for p in ti.static(range(mat.n)):
arr[I, p] = ti.cast(mat[I][p], ti.f32)
if ti.static(mat.n <= 2):
arr[I, 2] = 0
def vector_to_image(mat: template(), arr: ndarray_type.ndarray()):
for I in grouped(mat):
for p in static(range(mat.n)):
arr[I, p] = ops.cast(mat[I][p], f32)
if static(mat.n <= 2):
arr[I, 2] = 0
def fill_vbo(vbo: template(), value: template(), offset: template(),
num_components: template()):
for i in vbo:
for c in ti.static(range(num_components)):
vbo[i][offset + c] = value
def copy_to_vbo(vbo: template(), src: template(), offset: template(),
num_components: template()):
for i in src:
for c in ti.static(range(num_components)):
vbo[i][offset + c] = src[i][c]
def tensor_to_ext_arr(tensor: template(), arr: ext_arr()):
for I in ti.grouped(tensor):
arr[I] = tensor[I]
def fill_ndarray_matrix(ndarray: ndarray_type.ndarray(), val: template()):
for I in grouped(ndarray):
ndarray[I].fill(val)
def fill_tensor(tensor: template(), val: template()):
for I in grouped(tensor):
tensor[I] = val
def tensor_to_tensor(tensor: template(), other: template()):
for I in grouped(tensor):
tensor[I] = other[I]
def fill_ndarray_matrix(ndarray: any_arr(), val: template()):
for I in ti.grouped(ndarray):
ndarray[I].fill(val)
def tensor_to_image(tensor: template(), arr: ndarray_type.ndarray()):
for I in grouped(tensor):
t = ops.cast(tensor[I], f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
def ext_arr_to_tensor(arr: ndarray_type.ndarray(), tensor: template()):
for I in grouped(tensor):
tensor[I] = arr[I]
