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

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()

video

2 Likes