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
particle_radius = 3
particle_radius_in_world = particle_radius / screen_to_world_ratio
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=(boundary[0]/particle_radius, boundary[1]/particle_radius))
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
spiky_grad_factor = -45.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
def spiky_gradient(r, h):
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):
x[i] = [10+i%75*2*particle_radius_in_world, i//75*2*particle_radius_in_world+15]
v[i] = [0, 0]
g[None] = ti.Vector([0, -9.8])
@ti.kernel
def advector():
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]
grad_i = ti.Vector([0.0, 0.0])
sum_gradient_sqr = 0.0
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]
grad_j = spiky_gradient(pos_ji, h)
grad_i += grad_j
sum_gradient_sqr += grad_j.dot(grad_j)
# Eq(2)
density_constraint += poly6_value(pos_ji.norm(), h)
# Eq(1)
density_constraint = (mass * density_constraint / rho0) - 1.0
sum_gradient_sqr += grad_i.dot(grad_i)
lambdas[p_i] = (-density_constraint) / (sum_gradient_sqr +
lambda_epsilon)
#print(lambdas[p_i], density_constraint, sum_gradient_sqr)
@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) * \
spiky_gradient(pos_ji, h)
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)
advector()
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
#gui.circles(pos_np, radius=particle_radius, color=particle_color)
#for i in range(cell_num_x):
# for j in range(cell_num_y):
# val = is_wall[i,j]
#gui.rect(topLeft, rightBottom, radius=1, color=0xff00ff)
#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()
