forked from erizmr/taichi_sph
/
wcsph.py
70 lines (63 loc) · 2.84 KB
/
wcsph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import taichi as ti
from sph_base import SPHBase
class WCSPHSolver(SPHBase):
def __init__(self, particle_system):
super().__init__(particle_system)
# Pressure state function parameters(WCSPH)
self.exponent = 7.0
self.stiffness = 50.0
self.d_velocity = ti.Vector.field(self.ps.dim, dtype=float)
particle_node = ti.root.dense(ti.i, self.ps.particle_max_num)
particle_node.place(self.d_velocity)
@ti.kernel
def compute_densities(self):
for p_i in range(self.ps.particle_num[None]):
x_i = self.ps.x[p_i]
self.ps.density[p_i] = 0.0
for j in range(self.ps.particle_neighbors_num[p_i]):
p_j = self.ps.particle_neighbors[p_i, j]
x_j = self.ps.x[p_j]
self.ps.density[p_i] += self.ps.m_V * self.cubic_kernel((x_i - x_j).norm())
self.ps.density[p_i] *= self.density_0
@ti.kernel
def compute_pressure_forces(self):
for p_i in range(self.ps.particle_num[None]):
self.ps.density[p_i] = ti.max(self.ps.density[p_i], self.density_0)
self.ps.pressure[p_i] = self.stiffness * (ti.pow(self.ps.density[p_i] / self.density_0, self.exponent) - 1.0)
for p_i in range(self.ps.particle_num[None]):
if self.ps.material[p_i] != self.ps.material_fluid:
continue
x_i = self.ps.x[p_i]
d_v = ti.Vector([0.0 for _ in range(self.ps.dim)])
for j in range(self.ps.particle_neighbors_num[p_i]):
p_j = self.ps.particle_neighbors[p_i, j]
x_j = self.ps.x[p_j]
# Compute Pressure force contribution
d_v += self.pressure_force(p_i, p_j, x_i-x_j)
self.d_velocity[p_i] += d_v
@ti.kernel
def compute_non_pressure_forces(self):
for p_i in range(self.ps.particle_num[None]):
if self.ps.material[p_i] != self.ps.material_fluid:
continue
x_i = self.ps.x[p_i]
# Add body force
d_v = ti.Vector([0.0 for _ in range(self.ps.dim)])
d_v[self.ps.dim-1] = self.g
for j in range(self.ps.particle_neighbors_num[p_i]):
p_j = self.ps.particle_neighbors[p_i, j]
x_j = self.ps.x[p_j]
d_v += self.viscosity_force(p_i, p_j, x_i - x_j)
self.d_velocity[p_i] = d_v
@ti.kernel
def advect(self):
# Symplectic Euler
for p_i in range(self.ps.particle_num[None]):
if self.ps.material[p_i] == self.ps.material_fluid:
self.ps.v[p_i] += self.dt[None] * self.d_velocity[p_i]
self.ps.x[p_i] += self.dt[None] * self.ps.v[p_i]
def substep(self):
self.compute_densities()
self.compute_non_pressure_forces()
self.compute_pressure_forces()
self.advect()