def calc_loss(arr: ti.template(), loss: ti.template()):
for i in arr:
loss[None] += arr[i]
def mark_valid_v(velocities_v: ti.template(), valid: ti.template()):
for i, j in velocities_v:
if is_fluid(i, j - 1) or is_fluid(i, j):
valid[i, j] = 1
else:
valid[i, j] = 0
def inc(a: ti.template(), b: ti.template()):
for i in a:
a[i] += b
def clear_loss(l: ti.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 apply_gravity(velocities_v: ti.template()):
for i, j in velocities_v:
velocities_v[i, j] += -9.8 * dt
def tensor_to_tensor(tensor: ti.template(), other: ti.template()):
for I in ti.grouped(tensor):
tensor[I] = other[I]
def fill_tensor(tensor: ti.template(), val: ti.template()):
for I in ti.grouped(tensor):
tensor[I] = val
def double(a: ti.template(), b: ti.template()):
for i in range(16):
b[i] = a[i] * 2 + 1
def fill(l: ti.template()):
for I in ti.grouped(a[l]):
a[l][I] = l
def init_kernel(self, cell_type : ti.template()):
for i, j in cell_type:
if (i - self.x0) ** 2 + (j - self.y0) ** 2 <= self.r ** 2:
cell_type[i, j] = utils.FLUID
def update(self, I, res: ti.template()):
for k, v in ti.static(res.items()):
self[k][I] = v
def calc_loss(input_field: ti.template(), loss: ti.template()):
for i in input_field:
loss[None] += input_field[i]
def ti_frac(input_field: ti.template(), output_field: ti.template()):
for i in input_field:
output_field[i] = frac(input_field[i])**2
def advect(vf: ti.template(), qf: ti.template(), new_qf: ti.template()):
for i, j in vf:
p = ti.Vector([i, j]) + 0.5
p = backtrace(vf, p, dt)
new_qf[i, j] = bilerp(qf, p) * dye_decay
def tensor_to_image(tensor: ti.template(), arr: ti.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 aTob(l: ti.template()):
for I in ti.grouped(b[l]):
b[l][I] = sample(a, l, I)
def vector_to_image(mat: ti.template(), arr: ti.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 rand_vector(n: ti.template()):
'''
Samples a n-dimensional random uniform vector.
'''
return ti.Vector([ti.random() for _ in range(n)])
def ext_arr_to_tensor(arr: ti.ext_arr(), tensor: ti.template()):
for I in ti.grouped(tensor):
tensor[I] = arr[I]
def backtrace(v: ti.template(), I, dt):
midI = I - 0.5 * ts.bilerp(v, I) * dt
finI = I - dt * ts.bilerp(v, midI)
return finI
def clear_gradients(vars: ti.template()):
for I in ti.grouped(ti.Expr(vars[0])):
for s in ti.static(vars):
ti.Expr(s)[I] = 0
def _vector_product(self: ti.template()):
ret = self[0]
for i in ti.static(range(1, self.n)):
ret *= self[i]
return ret
def fill_matrix(mat: ti.template(), vals: ti.template()):
for I in ti.grouped(mat):
for p in ti.static(range(mat.n)):
for q in ti.static(range(mat.m)):
mat[I][p, q] = vals[p][q]
def snode_deactivate(b: ti.template()):
for I in ti.grouped(b):
ti.deactivate(b, I)
def mark_valid_u(velocities_u: ti.template(), valid: ti.template()):
for i, j in velocities_u:
if is_fluid(i - 1, j) or is_fluid(i, j):
valid[i, j] = 1
else:
valid[i, j] = 0
def snode_deactivate_dynamic(b: ti.template()):
for I in ti.grouped(b.parent()):
ti.deactivate(b, I)
def sample(x: ti.template(), I):
return x[I]
def tensor_to_ext_arr(tensor: ti.template(), arr: ti.ext_arr()):
for I in ti.grouped(tensor):
arr[I] = tensor[I]
def inc2(z: ti.i32, a: ti.template(), b: ti.i32):
for i in a:
a[i] += b + z
def mul(arr: ti.template(), out: ti.template()):
for i in arr:
out[i] = arr[i] * 2.0
