def computeGridIndex(x, y):

idx = np.floor(x / (2 * dh)).astype(int)

idy = np.floor(y / (2 * dh)).astype(int)

return idx, idy

def placeParticles(position_list, paticle_list, wall_mark, bound=0):

# position_list: [start_x, start_y, end_x, end_y]

start_x, start_y, end_x, end_y = position_list

vel_x, vel_y, p, rho = 0.0, 0.0, 0.0, 1000.0

for pos_x in np.arange(start_x, end_x, dx, dtype=np.float32):

for pos_y in np.arange(start_y, end_y, dx, dtype=np.float32):

paticle_list.append([pos_x, pos_y])

if bound:

wall_mark.append(0)

else:

wall_mark.append(1)

particle_list = []

wall_mark = []

#### Dam break #######

start_x = x_min + 0.5 * dx

start_y = y_min - 1.0 * dx

end_x = start_x + 0.5 * h

end_y = start_y + 0.6 * h

## Constrcut wall

# Bottom square

b_start_x = 0.0

b_start_y = 0.0

b_end_x = b_start_x + w

b_end_y = b_start_y + y_min - 2 * dx

bottom_bound = b_end_y

# Top square

t_start_x = 0.0

t_start_y = h - y_min + 2 * dx

t_end_x = t_start_x + w

t_end_y = h

top_bound = t_start_y

# left square

l_start_x = 0.0

l_start_y = y_min - 2 * dx

l_end_x = l_start_x + x_min - 2 * dx

l_end_y = l_start_y + h - 2 * y_min + 4 * dx

left_bound = l_end_x

# right square

r_start_x = w - x_min + 2 * dx

r_start_y = y_min - 2 * dx

r_end_x = w - dx

r_end_y = r_start_y + h - 2 * y_min + 4 * dx

right_bound = r_start_x

pos_list_fluid = [start_x, start_y, end_x, end_y]

placeParticles(pos_list_fluid, particle_list, wall_mark)

# These are boundaries

pos_list_bs = [b_start_x, b_start_y, b_end_x, b_end_y]

placeParticles(pos_list_bs, particle_list, wall_mark, bound=1)

pos_list_ts = [t_start_x, t_start_y, t_end_x, t_end_y]

placeParticles(pos_list_ts, particle_list, wall_mark, bound=1)

pos_list_ls = [l_start_x, l_start_y, l_end_x, l_end_y]

placeParticles(pos_list_ls, particle_list, wall_mark, bound=1)

pos_list_rs = [r_start_x, r_start_y, r_end_x, r_end_y]

placeParticles(pos_list_rs, particle_list, wall_mark, bound=1)

grid_size = 2*dh

num_x = np.ceil(w_bound / grid_size).astype(int)

num_y = np.ceil(h_bound / grid_size).astype(int)

grid_x = num_x

grid_y = num_y

def __init__(self, particle_list,

wall_mark,

grid,

bound, alpha=0.5, dx=0.2, max_time=10000, max_steps=1000,gui=None):

######## Solver parameters ##########

self.max_time = max_time

self.max_steps = max_steps

self.gui = gui

# Gravity

self.g = -9.80

# viscosity

self.alpha = alpha

# reference density

self.rho_0 = 1000.0

# CFL coefficient

self.CFL = 0.20

# Smooth kernel norm factor

self.kernel_norm = 1.0

# Pressure state function parameters

self.gamma = 7.0

self.c_0 = 20.0

###### Scene parameters ########

self.w = 20

self.h = 10

self.w_bound = 22

self.h_bound = 12

assert self.w_bound > self.w

assert self.h_bound > self.h

self.x_min = (self.w_bound - self.w) / 2.0

self.y_min = (self.h_bound - self.h) / 2.0

self.x_max = self.w_bound - (self.w_bound - self.w) / 2.0

self.y_max = self.h_bound - (self.h_bound - self.h) / 2.0

self.top_bound = bound[0] # top_bound

self.bottom_bound = bound[1] #bottom_bound

self.left_bound = bound[2] #left_bound

self.right_bound = bound[3] #right_bound

self.df_fac = 1.3

self.dx = 0.2

self.dh = self.dx * self.df_fac

self.kernel_norm = 10. / (7. * np.pi * self.dh ** 2)

###### Particles #######

self.dim = 2

self.particle_numbers = len(particle_list)

self.grid_x = grid[0]

self.grid_y = grid[1]

# Fluid particles

self.old_positions = ti.Vector(self.dim, dt=ti.f32)

self.particle_positions = ti.Vector(self.dim, dt=ti.f32)

self.particle_velocity = ti.Vector(self.dim, dt=ti.f32)

self.particle_pressure = ti.Vector(1, dt=ti.f32)

self.particle_density = ti.Vector(1, dt=ti.f32)

self.wall_mark_list = ti.Vector(1, dt=ti.f32)

self.d_velocity = ti.Vector(self.dim, dt=ti.f32)

self.d_density = ti.Vector(1, dt=ti.f32)

self.dx = dx

self.m = self.dx**2 * 1000.0

self.particle_list = np.array(particle_list, dtype=np.float32)

self.wall_mark = np.array(wall_mark, dtype=np.int32)

self.grid_num_particles = ti.var(ti.i32)

self.grid2particles = ti.var(ti.i32)

self.particle_num_neighbors = ti.var(ti.i32)

self.particle_neighbors = ti.var(ti.i32)

self.max_num_particles_per_cell = 100

self.max_num_neighbors = 100

ti.root.dense(ti.i, self.particle_numbers).place(self.old_positions, self.particle_positions,

self.particle_velocity, self.particle_pressure,

self.particle_density, self.d_velocity, self.d_density,

self.wall_mark_list)

grid_snode = ti.root.dense(ti.ij, (self.grid_x, self.grid_y))

grid_snode.place(self.grid_num_particles)

grid_snode.dense(ti.k, self.max_num_particles_per_cell).place(self.grid2particles)

nb_node = ti.root.dense(ti.i, self.particle_numbers)

nb_node.place(self.particle_num_neighbors)

nb_node.dense(ti.j, self.max_num_neighbors).place(self.particle_neighbors)

@ti.kernel

def init(self, p_list:ti.ext_arr(), w_list:ti.ext_arr()):

for i in range(self.particle_numbers):

for j in ti.static(range(self.dim)):

self.particle_positions[i][j] = p_list[i,j]

self.particle_velocity[i][j] = ti.cast(0.0, ti.f32)

self.d_velocity[i][0] = ti.cast(0.0, ti.f32)

self.d_velocity[i][1] = ti.cast(-9.8, ti.f32)

self.wall_mark_list[i][0] = w_list[i]

self.d_density[i][0] = ti.cast(0.0, ti.f32)

self.particle_pressure[i][0] = ti.cast(0.0, ti.f32)

self.particle_density[i][0] = ti.cast(1000.0, ti.f32)

@ti.func

def computeGridIndex(self, pos):

return (pos / (2 * dh)).cast(int)

@ti.kernel

def allocateParticles(self):

# Ref to pbf2d example from by Ye Kuang (k-ye)

# https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py

# allocate particles to grid

for p_i in self.particle_positions:

# Compute the grid index on the fly

cell = self.computeGridIndex(self.particle_positions[p_i])

offs = self.grid_num_particles[cell].atomic_add(1)

self.grid2particles[cell, offs] = p_i

@ti.func

def is_in_grid(self, c):

# Ref to pbf2d example from by Ye Kuang (k-ye)

# https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py

return 0 <= c[0] and c[0] < self.grid_x and 0 <= c[1] and c[1] < self.grid_y

@ti.func

def isFluid(self, p):

# check fluid particle or bound particle

return self.wall_mark_list[p][0]

@ti.kernel

def search_neighbors(self):

# Ref to pbf2d example from by Ye Kuang (k-ye)

# https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py

for p_i in self.particle_positions:

pos_i = self.particle_positions[p_i]

nb_i = 0

if self.isFluid(p_i) == 1:

# Compute the grid index on the fly

cell = self.computeGridIndex(self.particle_positions[p_i])

for offs in ti.static(ti.grouped(ti.ndrange((-1, 2), (-1, 2)))):

cell_to_check = cell + offs

if self.is_in_grid(cell_to_check):

for j in range(self.grid_num_particles[cell_to_check]):

p_j = self.grid2particles[cell_to_check, j]

if nb_i < self.max_num_neighbors and p_j != p_i and (

pos_i - self.particle_positions[p_j]).norm() < self.dh * 2.00:

self.particle_neighbors[p_i, nb_i] = p_j

nb_i += 1

self.particle_num_neighbors[p_i] = nb_i

@ti.func

def cubicKernel(self, r, h):

# value of cubic spline smoothing kernel

k = 10. / (7. * np.pi * h ** 2)

q = r / h

# assert q >= 0.0

res = ti.cast(0.0, ti.f32)

if q <= 1.0:

res = k * (1 - 1.5 * q ** 2 + 0.75 * q ** 3)

elif q < 2.0:

res = k * 0.25 * (2 - q) ** 3

return res

@ti.func

def cubicKernelDerivative(self, r, h):

# derivative of cubcic spline smoothing kernel

k = 10. / (7. * np.pi * h ** 2)

q = r / h

# assert q > 0.0

res = ti.cast(0.0, ti.f32)

if q < 1.0:

res = (k / h) * (-3 * q + 2.25 * q ** 2)

elif q < 2.0:

res = -0.75 * (k / h) * (2 - q) ** 2

return res

@ti.func

def rhoDerivative(self, ptc_i, ptc_j, r, r_mod):

# density delta

return self.m * self.cubicKernelDerivative(r_mod, self.dh) \

* (self.particle_velocity[ptc_i] - self.particle_velocity[ptc_j]).dot(r / r_mod)

@ti.func

def pUpdate(self, rho, rho_0=1000, gamma=7.0, c_0=20.0):

# Weakly compressible, tait function

b = rho_0 * c_0 ** 2 / gamma

return b * ((rho / rho_0) ** gamma - 1.0)

@ti.func

def pressureForce(self, ptc_i, ptc_j, r, r_mod, mirror_pressure=0):

# Compute the pressure force contribution, Symmetric Formula

res = ti.Vector([0.0, 0.0])

# Disable the mirror force, use collision instead

# Use pressure mirror method to handle boundary leak

# if mirror_pressure == 1:

# res = - self.m * (self.particle_pressure[ptc_i][0]/ self.particle_density[ptc_i][0] ** 2

# + self.particle_pressure[ptc_i][0]/self.rho_0**2)* self.cubicKernelDerivative(r_mod, h) * r / r_mod

# else:

res = -self.m * (self.particle_pressure[ptc_i][0] / self.particle_density[ptc_i][0] ** 2

+ self.particle_pressure[ptc_j][0] / self.particle_density[ptc_j][0] ** 2) \

* self.cubicKernelDerivative(r_mod, self.dh) * r / r_mod

return res

@ti.func

def viscosityForce(self, ptc_i, ptc_j, r, r_mod):

# Compute the viscosity force contribution, artificial viscosity

res = ti.Vector([0.0, 0.0])

v_xy= (self.particle_velocity[ptc_i]- self.particle_velocity[ptc_j]).dot(r)

if v_xy < 0:

# Artifical viscosity

vmu = -2.0 * self.alpha * self.dx * self.c_0 / (self.particle_density[ptc_i][0] + self.particle_density[ptc_j][0])

res = -self.m * vmu * v_xy/(r_mod**2 + 0.01*self.dx**2)* self.cubicKernelDerivative(r_mod, self.dh) * r / r_mod

return res

@ti.func

def simualteCollisions(self, ptc_i, vec, d):

# Collision factor, assume roughly 50% velocity loss after collision, i.e. m_f /(m_f + m_b)

c_f = 0.5

self.particle_positions[ptc_i] += vec * d

self.particle_velocity[ptc_i] -= (1.0+c_f) * self.particle_velocity[ptc_i].dot(vec) * vec

@ti.kernel

def enforceBoundary(self):

for p_i in self.particle_positions:

if self.isFluid(p_i) == 1:

pos = self.particle_positions[p_i]

if pos[0] < self.left_bound:

self.simualteCollisions(p_i, ti.Vector([1.0, 0.0]), self.left_bound - pos[0])

if pos[0] > self.right_bound:

self.simualteCollisions(p_i, ti.Vector([-1.0, 0.0]), pos[0] - self.right_bound)

if pos[1] > self.top_bound:

self.simualteCollisions(p_i, ti.Vector([0.0, -1.0]), pos[1] - self.top_bound)

if pos[1] < self.bottom_bound:

self.simualteCollisions(p_i, ti.Vector([0.0, 1.0]), self.bottom_bound - pos[1])

@ti.kernel

def computeDeltas(self):

for p_i in self.particle_positions:

pos_i = self.particle_positions[p_i]

d_v = ti.Vector([0.0, 0.0], dt=ti.f32)

d_rho = ti.cast(0.0, ti.f32)

# if self.isFluid(p_i) == 1:

# d_v = ti.Vector([0.0, -9.8])

for j in range(self.particle_num_neighbors[p_i]):

p_j = self.particle_neighbors[p_i, j]

pos_j = self.particle_positions[p_j]

# Disable mirror force

# mirror_pressure = 0

# if self.isFluid(p_j) == 0:

# mirror_pressure = 1

# Compute distance and its mod

r = pos_i - pos_j

r_mod = r.norm()

# Compute Density change

d_rho += self.rhoDerivative(p_i, p_j, r, r_mod)

if self.isFluid(p_i) == 1:

# Compute Viscosity force contribution

d_v += self.viscosityForce(p_i, p_j, r, r_mod)

# Compute Pressure force contribution

d_v += self.pressureForce(p_i, p_j, r, r_mod)

# Add body force

if self.isFluid(p_i) == 1:

d_v += ti.Vector([0.0, self.g], dt=ti.f32)

self.d_velocity[p_i] = d_v

self.d_density[p_i][0] = d_rho

@ti.kernel

def updateTimeStep(self):

# Simple Forward Euler currently

for p_i in self.particle_positions:

if self.isFluid(p_i) == 1:

self.particle_positions[p_i] += self.dt * self.particle_velocity[p_i]

self.particle_velocity[p_i] += self.dt * self.d_velocity[p_i]

self.particle_density[p_i][0] += self.dt * self.d_density[p_i][0]

self.particle_pressure[p_i][0] = self.pUpdate(self.particle_density[p_i][0], self.rho_0, self.gamma, self.c_0)

def solve(self, output=False):

# Compute dt, a naive initial test value

self.dt = 0.1 * self.dh / self.c_0

print("Time step: ", self.dt)

print("Domain: (%s, %s, %s, %s)" % (self.x_min, self.x_max, self.y_min, self.y_max), )

print("Fluid area: (%s, %s, %s, %s)"%(self.left_bound, self.right_bound, self.bottom_bound, self.top_bound))

print("Grid: (%d, %d)"%(self.grid_x, self.grid_y))

step = 1

t = 0.0

total_start = time.process_time()

while t < self.max_time and step < self.max_steps:

curr_start = time.process_time()

self.solveUpdate()

max_v = np.max(np.linalg.norm(self.particle_velocity.to_numpy(),2, axis=1))

max_a = np.max(np.linalg.norm(self.d_velocity.to_numpy(),2, axis=1))

max_rho = np.max(self.particle_density.to_numpy())

max_pressure = np.max(self.particle_pressure.to_numpy())

curr_end = time.process_time()

t += self.dt

step += 1

# CFL analysis, adaptive dt

dt_cfl = self.dh / max_v

dt_f = np.sqrt(self.dh / max_a)

dt_a = self.dh / (self.c_0 * np.sqrt((max_rho / self.rho_0)**self.gamma))

self.dt = self.CFL * np.min([dt_cfl, dt_f, dt_a])

if step % 10 == 0:

print("Step: %d, physics time: %s, progress: %s %%, time used: %s, total time used: %s"

% (step, t, 100*np.max([t / self.max_time, step / self.max_steps]), curr_end-curr_start, curr_end-total_start))

print("Max velocity: %s, Max acceleration: %s, Max density: %s, Max pressure: %s" % (max_v, max_a, max_rho, max_pressure))

print("Adaptive time step: ", self.dt)

self.render(step, self.gui, output)

total_end = time.process_time()

print("Total time used: %s " % (total_end - total_start))

def solveUpdate(self):

self.grid_num_particles.fill(0)

self.particle_neighbors.fill(-1)

self.allocateParticles()

self.search_neighbors()

# Compute deltas

self.computeDeltas()

# timestep Update

self.updateTimeStep()

# Handle potential leak particles

self.enforceBoundary()

def isFluidNP(self, p):

# ti.func cannot be called in python scope

# for render use

return self.wall_mark[p]

def render(self, step, gui, output=False):

canvas = gui.canvas

canvas.clear(bg_color)

pos_np = self.particle_positions.to_numpy()

fluid_p = []

wall_p = []

for i, pos in enumerate(pos_np):

if self.isFluidNP(i) == 1:

fluid_p.append(pos)

else:

wall_p.append(pos)

fluid_p = np.array(fluid_p)

wall_p = np.array(wall_p)

for pos in fluid_p:

for j in range(self.dim):

pos[j] *= screen_to_world_ratio / screen_res[j]

for pos in wall_p:

for j in range(self.dim):

pos[j] *= screen_to_world_ratio / screen_res[j]

gui.circles(fluid_p, radius=particle_radius, color=particle_color)

gui.circles(wall_p, radius=particle_radius, color=boundary_color)

if output:

if step%10 == 0:

gui.show(f"{step:04d}.png")

else:

OUTPUT = False

gui = ti.GUI('SPH2D', screen_res)

grid_shape = makeGrid()

particle_list, wall_mark, u, b, l, r = setup()

sph = sph_solver(particle_list, wall_mark, grid_shape, [u,b,l,r],alpha=1.0, dx = dx, gui=gui, max_steps=10000)

sph.init(sph.particle_list, sph.wall_mark)

sph.solve(output=OUTPUT)