# Homework2: PBF-driven game (Water gun)

Sorry, my pc platform cann’t easily use a chinese input method.
Writter: Zeyuan Jin. Email: qqqqzeyuan@gmail.com
Main pbf code, I try to impletement cell sort. Still copy a lot from example

``````#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Pw @ 2020-08-07 05:44:34

import taichi as ti
import numpy as np
ti.init(arch=ti.gpu)

from sort import *

dim = 2
quality = 1  # Use a larger value for higher-res simulations
screen_res = (800, 400)
screen_to_world_ratio = 10.0
boundary = (screen_res[0] / screen_to_world_ratio,
screen_res[1] / screen_to_world_ratio)
cell_size = 2.51
cell_recpr = 1.0 / cell_size

pbf_num_iters = 5
def round(f, s):
return (math.floor(f * cell_recpr / s) + 1) * s

n_particles = 2000
#n_particles = 100 * quality**2
dt = 1 /20 / quality
inv_dt = 20 * quality
v = ti.Vector(2, dt=ti.f32, shape=n_particles)
old_positions = ti.Vector(dim, dt=ti.f32, shape = n_particles)
x = ti.Vector(2, dt=ti.f32, shape=n_particles)

g = ti.Vector(2, dt=ti.f32, shape= ())

#max_num_neighbors = 50
cell_num_x = round(boundary[0], 1)
cell_num_y = round( boundary[1], 1)
cell_total_num = cell_num_x * cell_num_y + 1#!/usr/bin/env python # -*- coding: utf-8 -*- # Pw @ 2020-08-08 11:55:32  import taichi as ti import math import time  #import numpy as np  #ti.init(arch=ti.cuda)  #array = ti.var(ti.i32, shape=n)   @ti.data_oriented class Sorter:      #self.data = ti.var(ti.i32, shape=n)     def sortStage(self,stage):         gap = 2**(stage + 1)         for i in range(stage + 1):             length = 2**(stage - i)             twoStage = 2 * length             self.sortSubStage(self.data, twoStage, gap, self.n)      def sortStageByKey(self,stage):         gap = 2**(stage + 1)         for i in range(stage + 1):             length = 2**(stage - i)             twoStage = 2 * length             self.sortSubStageByKey(self.data, self.values, twoStage, gap, self.n)      @ti.kernel     def sortSubStage(self, d: ti.template(), twoStage: ti.i32, gap: ti.i32, n: ti.i32):         for idx in range(n):             stage = twoStage // 2             direction = (idx // gap % 2 == 0)  # direction = True means decrease             #direction = False means increase             if direction:                 gapCount = idx // gap                 gapEnd = (gapCount + 1) * gap - 1                 if gapEnd >= n:                     i = (n - idx - 1) % twoStage                     if i < stage:                         partner = idx - stage                         if partner <= gapEnd - gap:                             partner = idx                         if d[idx] > d[partner]:                             tmp = d[idx]                             d[idx] = d[partner]                             d[partner] = tmp                 else:                     i = idx % twoStage                     if i < stage:                         partner = idx + stage                         if partner >= n:                             partner = idx                         if d[idx] < d[partner]:                             tmp = d[idx]                             d[idx] = d[partner]                             d[partner] = tmp              else:                 i = idx % twoStage                 if i < stage:                     partner = idx + stage                     if partner >= n:                         partner = idx                     if d[idx] > d[partner]:                         tmp = d[idx]                         d[idx] = d[partner]                         d[partner] = tmp      @ti.kernel     def sortSubStageByKey(self, d: ti.template(),v:ti.template(), twoStage: ti.i32, gap: ti.i32, n: ti.i32):         for idx in range(n):             stage = twoStage // 2             direction = (idx // gap % 2 == 0)  # direction = True means decrease             #direction = False means increase             if direction:                 gapCount = idx // gap                 gapEnd = (gapCount + 1) * gap - 1                 if gapEnd >= n:                     i = (n - idx - 1) % twoStage                     if i < stage:                         partner = idx - stage                         if partner <= gapEnd - gap:                             partner = idx                         if d[idx] > d[partner]:                             tmp = d[idx]                             d[idx] = d[partner]                             d[partner] = tmp                             tmp = v[idx]                             v[idx] = v[partner]                             v[partner] = tmp                 else:                     i = idx % twoStage                     if i < stage:                         partner = idx + stage                         if partner >= n:                             partner = idx                         if d[idx] < d[partner]:                             tmp = d[idx]                             d[idx] = d[partner]                             d[partner] = tmp                             tmp = v[idx]                             v[idx] = v[partner]                             v[partner] = tmp              else:                 i = idx % twoStage                 if i < stage:                     partner = idx + stage                     if partner >= n:                         partner = idx                     if d[idx] > d[partner]:                         tmp = d[idx]                         d[idx] = d[partner]                         d[partner] = tmp                         tmp = v[idx]                         v[idx] = v[partner]                         v[partner] = tmp      @ti.kernel     def reverse(self, n:ti.i32, d:ti.template(), v:ti.template()):         for i in range(n//2):             partner = n - i-1             tmp = d[i]             d[i] = d[partner]             d[partner] = tmp             tmp = v[i]             v[i] = v[partner]             v[partner] = tmp      #def __init__(self, num, t:ti.template()):         #self.n = num         #self.data = t      def sort(self, num, t:ti.template()):         self.n = num         self.data = t         for stage in range(math.ceil(math.log2(num))):             self.sortStage(stage)      def sortByKey(self, num, k:ti.template(), v:ti.template()):         self.n = num         self.data = k         self.values = v         for stage in range(math.ceil(math.log2(num))):             self.sortStageByKey(stage)      def reverseSelf(self):         self.reverse(self.n, self.data, self.values)      #@ti.kernel     #def check_ray():     #    #for i in range(n-1):     #    for i in range(n - 1):     #        #print(data[i])     #        if data[i] < data[i + 1]:     #            print(i, data[i], data[i + 1])       #def print_ray():     #    for i in range(n):     #        print(i, data[i])       #@ti.kernel     #def reset():     #    for i in range(n):     #        data[i] = int(ti.random() * 1000)  #check_ray()
cells = ti.var(dt=ti.i32, shape=cell_total_num)
cellEndIdx = ti.var(dt=ti.i32, shape=cell_total_num)

particle2cell = ti.var(dt=ti.i32, shape = n_particles)
cell2particle = ti.var(dt=ti.i32, shape = n_particles)

is_wall = ti.var(dt = ti.i32, shape = (cell_num_x, cell_num_y))

sorter = Sorter()
h = 1.1
mass = 1.0
rho0 = 1.0
lambda_epsilon = 100.0
corr_deltaQ_coeff = 0.3
corrK = 0.001

bg_color = 0x112f41
particle_color = 0x068587
boundary_color = 0xebaca2
epsilon = 1e-5

lambdas = ti.var(ti.f32, shape = n_particles)
position_deltas = ti.Vector(dim, dt=ti.f32, shape = n_particles)
poly6_factor = 315.0 / 64.0 / np.pi

@ti.func
def poly6_value(s, h):
result = 0.0
if 0 < s and s < h:
x = (h * h - s * s) / (h * h * h)
result = poly6_factor * x * x * x
return result

@ti.func
result = ti.Vector([0.0, 0.0])
r_len = r.norm()
if 0 < r_len and r_len < h:
x = (h - r_len) / (h * h * h)
g_factor = spiky_grad_factor * x * x
result = r * g_factor / r_len
return result

@ti.kernel
def reset():
for i in range(n_particles):
v[i] = [0, 0]
g[None] = ti.Vector([0, -9.8])

@ti.kernel
for i in range(n_particles):
new_v = v[i] + g[None]*dt
new_x = new_v*dt + x[i]
new_x = confine_par_to_boundary(new_x, x[i])
v[i] = new_v
x[i] = new_x
#re

@ti.kernel
def confine_to_boundary():
for i in x:
pos = x[i]
x[i] = confine_par_to_boundary(pos, old_positions[i])

@ti.func
def confine_par_to_boundary(new_pos, old_pos):
cell = get_cell(old_pos)
for d in ti.static(range(dim)):
cell[d] -= 1
if is_wall[cell]==1:
#tmp = 1
tmp = (cell[d] +1)*cell_size
new_pos[d] = ti.max(new_pos[d], tmp+particle_radius_in_world + epsilon)
#print(new_pos, 114)
cell[d] += 2
if is_wall[cell]==1:
tmp = cell[d] *cell_size
new_pos[d] = ti.min(new_pos[d], tmp-particle_radius_in_world - epsilon)
#print(new_pos, 119)
cell[d] -= 1
new_cell = get_cell(new_pos)
if is_wall[ new_cell ]==1:
new_pos = old_pos
return new_pos

@ti.func
def get_cell(pos):
return (pos * cell_recpr).cast(int)

@ti.func
def hash_cell(cell):
return cell[0] + cell[1] * cell_num_x

@ti.func
def get_cell_particle(cell):
curCellEnd = cellEndIdx[hash_cell(cell)]
prevCell = (cell[0]-1, cell[1])
return (cellEndIdx[hash_cell(prevCell)], curCellEnd)
#ret = ti.Vector([-1,0])
#if cell[0] == 0 and cell[1] == 0:
#    ret[1] = curCellEnd
#elif cell[0] == 0:
#    prevCell = (cell_num_x-1, cell[1]-1)
#    ret[0] = cellEndIdx[hash_cell(prevCell)]
#    ret[1] = curCellEnd
#else:
#    prevCell = (cell[0]-1, cell[1])
#    ret[0] = cellEndIdx[hash_cell(prevCell)]
#    ret[1] = curCellEnd
#return ret

@ti.kernel
def genP2C():
for i in range(n_particles):
particle2cell[i] = hash_cell(get_cell(x[i]))

@ti.kernel
def genC2P():
for i in range(n_particles):
cell2particle[i] = i

def gen_cell_idx():
for i in range(cell_total_num):
cellEndIdx[i] = 0
for i in range(n_particles):
if i+1 == n_particles:
lastCell = particle2cell[n_particles-1]
cellEndIdx[lastCell] = n_particles-1
elif particle2cell[i] != particle2cell[i+1]:
cellEndIdx[particle2cell[i]] = i
for i in range(cell_total_num-1):
if cellEndIdx[i+1] == 0:
cellEndIdx[i+1] = cellEndIdx[i]

#for i in range(n_particles):
#print(i, particle2cell[i])

@ti.kernel
def compute_lambdas():
# Eq (8) ~ (11)
for p_i in x:
pos_i = x[p_i]

density_constraint = 0.0
cell = get_cell(pos_i)

for offs in ti.static(ti.grouped(ti.ndrange((-1, 2), (-1, 2)))):
cell_to_check = cell + offs
particleRange = get_cell_particle(cell_to_check)
for par in range(particleRange[1] - particleRange[0]):
p_j = cell2particle[par+ particleRange[0]+1]
# TODO: does taichi supports break?
if p_j >= 0:
pos_ji = pos_i - x[p_j]
# Eq(2)
density_constraint += poly6_value(pos_ji.norm(), h)

# Eq(1)
density_constraint = (mass * density_constraint / rho0) - 1.0

lambdas[p_i] = (-density_constraint) / (sum_gradient_sqr +
lambda_epsilon)

@ti.kernel
def blit_buffers(f: ti.template(), t: ti.template()):
for i in f:
t[i] = f[i]

@ti.kernel
def compute_position_deltas():
# Eq(12), (14)
for p_i in x:
pos_i = x[p_i]
lambda_i = lambdas[p_i]

pos_delta_i = ti.Vector([0.0, 0.0])
cell = get_cell(pos_i)

for offs in ti.static(ti.grouped(ti.ndrange((-1, 2), (-1, 2)))):
cell_to_check = cell + offs
particleRange = get_cell_particle(cell_to_check)
for par in range(particleRange[1] - particleRange[0]):
p_j = cell2particle[par+ particleRange[0]+1]
if p_j >= 0:
lambda_j = lambdas[p_j]
pos_ji = pos_i - x[p_j]
scorr_ij = compute_scorr(pos_ji)
pos_delta_i += (lambda_i + lambda_j + scorr_ij) * \

pos_delta_i /= rho0
position_deltas[p_i] = pos_delta_i

@ti.func
def compute_scorr(pos_ji):
# Eq (13)
x = poly6_value(pos_ji.norm(), h) / poly6_value(corr_deltaQ_coeff * h, h)
# pow(x, 4)
x = x * x
x = x * x
return (-corrK) * x

@ti.kernel
def apply_position_deltas():
for i in x:
x[i] += position_deltas[i]

@ti.kernel
def update_velocities():
for i in x:
v[i] = (x[i] - old_positions[i]) * inv_dt

#@ti.kernel
def get_neibors():
genP2C()
genC2P()
sorter.sortByKey(n_particles, particle2cell, cell2particle)
sorter.reverseSelf()
gen_cell_idx()

def pbf_simulate():
blit_buffers(x, old_positions)
get_neibors()
for _ in range(pbf_num_iters):
compute_lambdas()
compute_position_deltas()
apply_position_deltas()
`
confine_to_boundary()
update_velocities()

def gen_wall():
is_wall.fill(0)
for i in range(cell_num_x):
is_wall[i, 0] = 1
is_wall[i, cell_num_y-1] = 1
for i in range(cell_num_y):
is_wall[0,i] = 1
is_wall[cell_num_x-1,i] = 1
is_wall[13,1] = 1
is_wall[13,2] = 1

pixels = ti.field(ti.u8, shape =(screen_res[0], screen_res[1], 3))
@ti.kernel
def render():
for i, j, k in pixels:
cell = (ti.floor(i*0.1* cell_recpr), ti.floor(j*0.1* cell_recpr))
if is_wall[cell]:
pixels[i, j, 0] = 255
else:
pixels[i, j, k] = 0
for p in range(n_particles):
t = ti.Vector([0,0])
for j in ti.static(range(dim)):
#pos[j] *= screen_to_world_ratio / screen_res[j]
t[j] = ti.floor(x[p][j] * screen_to_world_ratio)
for offs in ti.static(ti.grouped(ti.ndrange((-1, 2), (-1, 2)))):
for c in ti.static(range(3)):
pixels[t[0] + offs[0],t[1]+offs[1],c] = 128

#for i in range(cell_num_x):
#    for j in range(cell_num_y):
#        val = is_wall[i,j]
#gui.show()

reset()
result_dir = "./results"
video_manager = ti.VideoManager(output_dir=result_dir, framerate=24, automatic_build=False)
gen_wall()
for frame in range(300):
pbf_simulate()
render()
pixels_img = pixels.to_numpy()
video_manager.write_frame(pixels_img)

video_manager.make_video(gif=True, mp4=True)
``````

sorter code part, it can sort any number below 2**28, bigger num is untest, but just decreasing order, maybe need reverse. Use Bitonic sort sequence.

``````#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Pw @ 2020-08-08 11:55:32

import taichi as ti
import math
import time

#import numpy as np

#ti.init(arch=ti.cuda)

#array = ti.var(ti.i32, shape=n)

@ti.data_oriented
class Sorter:

#self.data = ti.var(ti.i32, shape=n)
def sortStage(self,stage):
gap = 2**(stage + 1)
for i in range(stage + 1):
length = 2**(stage - i)
twoStage = 2 * length
self.sortSubStage(self.data, twoStage, gap, self.n)

def sortStageByKey(self,stage):
gap = 2**(stage + 1)
for i in range(stage + 1):
length = 2**(stage - i)
twoStage = 2 * length
self.sortSubStageByKey(self.data, self.values, twoStage, gap, self.n)

@ti.kernel
def sortSubStage(self, d: ti.template(), twoStage: ti.i32, gap: ti.i32, n: ti.i32):
for idx in range(n):
stage = twoStage // 2
direction = (idx // gap % 2 == 0)  # direction = True means decrease
#direction = False means increase
if direction:
gapCount = idx // gap
gapEnd = (gapCount + 1) * gap - 1
if gapEnd >= n:
i = (n - idx - 1) % twoStage
if i < stage:
partner = idx - stage
if partner <= gapEnd - gap:
partner = idx
if d[idx] > d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp
else:
i = idx % twoStage
if i < stage:
partner = idx + stage
if partner >= n:
partner = idx
if d[idx] < d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp

else:
i = idx % twoStage
if i < stage:
partner = idx + stage
if partner >= n:
partner = idx
if d[idx] > d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp

@ti.kernel
def sortSubStageByKey(self, d: ti.template(),v:ti.template(), twoStage: ti.i32, gap: ti.i32, n: ti.i32):
for idx in range(n):
stage = twoStage // 2
direction = (idx // gap % 2 == 0)  # direction = True means decrease
#direction = False means increase
if direction:
gapCount = idx // gap
gapEnd = (gapCount + 1) * gap - 1
if gapEnd >= n:
i = (n - idx - 1) % twoStage
if i < stage:
partner = idx - stage
if partner <= gapEnd - gap:
partner = idx
if d[idx] > d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp
tmp = v[idx]
v[idx] = v[partner]
v[partner] = tmp
else:
i = idx % twoStage
if i < stage:
partner = idx + stage
if partner >= n:
partner = idx
if d[idx] < d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp
tmp = v[idx]
v[idx] = v[partner]
v[partner] = tmp

else:
i = idx % twoStage
if i < stage:
partner = idx + stage
if partner >= n:
partner = idx
if d[idx] > d[partner]:
tmp = d[idx]
d[idx] = d[partner]
d[partner] = tmp
tmp = v[idx]
v[idx] = v[partner]
v[partner] = tmp

@ti.kernel
def reverse(self, n:ti.i32, d:ti.template(), v:ti.template()):
for i in range(n//2):
partner = n - i-1
tmp = d[i]
d[i] = d[partner]
d[partner] = tmp
tmp = v[i]
v[i] = v[partner]
v[partner] = tmp

#def __init__(self, num, t:ti.template()):
#self.n = num
#self.data = t

def sort(self, num, t:ti.template()):
self.n = num
self.data = t
for stage in range(math.ceil(math.log2(num))):
self.sortStage(stage)

def sortByKey(self, num, k:ti.template(), v:ti.template()):
self.n = num
self.data = k
self.values = v
for stage in range(math.ceil(math.log2(num))):
self.sortStageByKey(stage)

def reverseSelf(self):
self.reverse(self.n, self.data, self.values)

#@ti.kernel
#def check_ray():
#    #for i in range(n-1):
#    for i in range(n - 1):
#        #print(data[i])
#        if data[i] < data[i + 1]:
#            print(i, data[i], data[i + 1])

#def print_ray():
#    for i in range(n):
#        print(i, data[i])

#@ti.kernel
#def reset():
#    for i in range(n):
#        data[i] = int(ti.random() * 1000)

#check_ray()
``````

2 Likes