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