def advect_semilag(vf: ti.template(), qf: ti.template(),
new_qf: ti.template()):
ti.cache_read_only(qf, vf)
for i, j in vf:
p = ti.Vector([i, j]) + 0.5
p = backtrace(vf, p, dt)
new_qf[i, j] = bilerp(qf, p)
def advect_bfecc(vf: ti.template(), qf: ti.template(), new_qf: ti.template(),
intermedia_qf: ti.template()):
ti.cache_read_only(qf, vf)
for i, j in vf:
p = ti.Vector([i, j]) + 0.5
p = backtrace(vf, p, dt)
intermedia_qf[i, j] = bilerp(qf, p)
ti.cache_read_only(intermedia_qf, qf, vf)
for i, j in vf:
p = ti.Vector([i, j]) + 0.5
# star means the temp value after a back tracing (forward advection)
# two star means the temp value after a forward tracing (reverse advection)
p_two_star = backtrace(vf, p, -dt)
p_star = backtrace(vf, p, dt)
q_star = intermedia_qf[i, j]
new_qf[i, j] = bilerp(intermedia_qf, p_two_star)
new_qf[i, j] = q_star + 0.5 * (qf[i, j] - new_qf[i, j])
min_val, max_val = sample_minmax(qf, p_star)
cond = min_val < new_qf[i, j] < max_val
for k in ti.static(range(cond.n)):
if not cond[k]:
new_qf[i, j][k] = q_star[k]
def subtract_gradient(vf: ti.template(), pf: ti.template()):
ti.cache_read_only(pf)
for i, j in vf:
pl = sample(pf, i - 1, j)
pr = sample(pf, i + 1, j)
pb = sample(pf, i, j - 1)
pt = sample(pf, i, j + 1)
vf[i, j] -= 0.5 * ti.Vector([pr - pl, pt - pb])
def pressure_jacobi_single(pf: ti.template(), new_pf: ti.template()):
ti.cache_read_only(pf)
for i, j in pf:
pl = sample(pf, i - 1, j)
pr = sample(pf, i + 1, j)
pb = sample(pf, i, j - 1)
pt = sample(pf, i, j + 1)
div = velocity_divs[i, j]
new_pf[i, j] = (pl + pr + pb + pt - div) * 0.25
def vorticity(vf: ti.template()):
ti.cache_read_only(vf)
for i, j in vf:
vl = sample(vf, i - 1, j).y
vr = sample(vf, i + 1, j).y
vb = sample(vf, i, j - 1).x
vt = sample(vf, i, j + 1).x
vc = sample(vf, i, j)
velocity_curls[i, j] = (vr - vl - vt + vb) * 0.5
def enhance_vorticity(vf: ti.template(), cf: ti.template()):
# anti-physics visual enhancement...
ti.cache_read_only(cf)
for i, j in vf:
cl = sample(cf, i - 1, j)
cr = sample(cf, i + 1, j)
cb = sample(cf, i, j - 1)
ct = sample(cf, i, j + 1)
cc = sample(cf, i, j)
force = ti.Vector([abs(ct) - abs(cb),
abs(cl) - abs(cr)]).normalized(1e-3)
force *= curl_strength * cc
vf[i, j] = min(max(vf[i, j] + force * dt, -1e3), 1e3)
def advect_bfecc(vf: ti.template(), qf: ti.template(), new_qf: ti.template()):
ti.cache_read_only(qf, vf)
for i, j in vf:
p = ti.Vector([i, j]) + 0.5
p_mid = backtrace(vf, p, dt)
q_mid = bilerp(qf, p_mid)
p_fin = backtrace(vf, p_mid, -dt)
q_fin = bilerp(qf, p_fin)
new_qf[i, j] = q_mid + 0.5 * (q_fin - qf[i, j])
min_val, max_val = sample_minmax(qf, p_mid)
cond = min_val < new_qf[i, j] < max_val
for k in ti.static(range(cond.n)):
if not cond[k]:
new_qf[i, j][k] = q_mid[k]
def divergence(vf: ti.template()):
ti.cache_read_only(vf)
for i, j in vf:
vl = sample(vf, i - 1, j).x
vr = sample(vf, i + 1, j).x
vb = sample(vf, i, j - 1).y
vt = sample(vf, i, j + 1).y
vc = sample(vf, i, j)
if i == 0:
vl = -vc.x
if i == res - 1:
vr = -vc.x
if j == 0:
vb = -vc.y
if j == res - 1:
vt = -vc.y
velocity_divs[i, j] = (vr - vl + vt - vb) * 0.5
def pressure_jacobi_dual(pf: ti.template(), new_pf: ti.template()):
ti.cache_read_only(pf)
for i, j in pf:
pcc = sample(pf, i, j)
pll = sample(pf, i - 2, j)
prr = sample(pf, i + 2, j)
pbb = sample(pf, i, j - 2)
ptt = sample(pf, i, j + 2)
plb = sample(pf, i - 1, j - 1)
prb = sample(pf, i + 1, j - 1)
plt = sample(pf, i - 1, j + 1)
prt = sample(pf, i + 1, j + 1)
div = sample(velocity_divs, i, j)
divl = sample(velocity_divs, i - 1, j)
divr = sample(velocity_divs, i + 1, j)
divb = sample(velocity_divs, i, j - 1)
divt = sample(velocity_divs, i, j + 1)
new_pf[i,
j] = (pll + prr + pbb + ptt - divl - divr - divb - divt - div +
(plt + prt + prb + plb) * 2 + pcc * 4) * 0.0625
def test(data: ti.f32):
ti.cache_read_only(x)
assert x[None] == data
