├── .gitignore
├── DFSPH.py
├── IISPH.py
├── LICENSE.txt
├── README.md
├── WCSPH.py
├── config_builder.py
├── data
├── BoxOpenedHole.obj
├── gif
│ ├── armadillo_bath.gif
│ └── dragon_bath_large.gif
├── models
│ ├── Dragon_50k.obj
│ ├── armadillo_small.obj
│ ├── bunny.stl
│ └── bunny_sparse.obj
└── scenes
│ ├── armadillo_bath_dynamic.json
│ ├── armadillo_bath_dynamic_dfsph.json
│ ├── dragon_bath.json
│ ├── dragon_bath_dfsph.json
│ ├── dragon_bath_dynamic_dfsph.json
│ ├── high_fluid_dfsph.json
│ └── high_fluid_wcsph.json
├── demo_high_fluid.py
├── legacy
├── README.md
├── engine
│ ├── __init__.py
│ └── sph_solver.py
├── img
│ ├── DFSPH.gif
│ ├── PCISPH.gif
│ ├── WCSPH.gif
│ ├── sph_hv.gif
│ └── wcsph_alpha030.gif
├── scene.py
└── test_sample.py
├── particle_system.py
├── requirements.txt
├── run_simulation.py
├── scan_single_buffer.py
└── sph_base.py
/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 |
--------------------------------------------------------------------------------
/DFSPH.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 | from sph_base import SPHBase
3 |
4 |
5 | class DFSPHSolver(SPHBase):
6 | def __init__(self, particle_system):
7 | super().__init__(particle_system)
8 |
9 | self.surface_tension = 0.01
10 | self.dt[None] = self.ps.cfg.get_cfg("timeStepSize")
11 |
12 | self.enable_divergence_solver = True
13 |
14 | self.m_max_iterations_v = 100
15 | self.m_max_iterations = 100
16 |
17 | self.m_eps = 1e-5
18 |
19 | self.max_error_V = 0.1
20 | self.max_error = 0.05
21 |
22 |
23 | @ti.func
24 | def compute_densities_task(self, p_i, p_j, ret: ti.template()):
25 | x_i = self.ps.x[p_i]
26 | if self.ps.material[p_j] == self.ps.material_fluid:
27 | # Fluid neighbors
28 | x_j = self.ps.x[p_j]
29 | ret += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
30 | elif self.ps.material[p_j] == self.ps.material_solid:
31 | # Boundary neighbors
32 | ## Akinci2012
33 | x_j = self.ps.x[p_j]
34 | ret += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
35 |
36 |
37 | @ti.kernel
38 | def compute_densities(self):
39 | # for p_i in range(self.ps.particle_num[None]):
40 | for p_i in ti.grouped(self.ps.x):
41 | if self.ps.material[p_i] != self.ps.material_fluid:
42 | continue
43 | self.ps.density[p_i] = self.ps.m_V[p_i] * self.cubic_kernel(0.0)
44 | den = 0.0
45 | self.ps.for_all_neighbors(p_i, self.compute_densities_task, den)
46 | self.ps.density[p_i] += den
47 | self.ps.density[p_i] *= self.density_0
48 |
49 |
50 | @ti.func
51 | def compute_non_pressure_forces_task(self, p_i, p_j, ret: ti.template()):
52 | x_i = self.ps.x[p_i]
53 |
54 | ############## Surface Tension ###############
55 | if self.ps.material[p_j] == self.ps.material_fluid:
56 | # Fluid neighbors
57 | diameter2 = self.ps.particle_diameter * self.ps.particle_diameter
58 | x_j = self.ps.x[p_j]
59 | r = x_i - x_j
60 | r2 = r.dot(r)
61 | if r2 > diameter2:
62 | ret -= self.surface_tension / self.ps.m[p_i] * self.ps.m[p_j] * r * self.cubic_kernel(r.norm())
63 | else:
64 | ret -= self.surface_tension / self.ps.m[p_i] * self.ps.m[p_j] * r * self.cubic_kernel(ti.Vector([self.ps.particle_diameter, 0.0, 0.0]).norm())
65 |
66 |
67 | ############### Viscosoty Force ###############
68 | d = 2 * (self.ps.dim + 2)
69 | x_j = self.ps.x[p_j]
70 | # Compute the viscosity force contribution
71 | r = x_i - x_j
72 | v_xy = (self.ps.v[p_i] -
73 | self.ps.v[p_j]).dot(r)
74 |
75 | if self.ps.material[p_j] == self.ps.material_fluid:
76 | f_v = d * self.viscosity * (self.ps.m[p_j] / (self.ps.density[p_j])) * v_xy / (
77 | r.norm()**2 + 0.01 * self.ps.support_radius**2) * self.cubic_kernel_derivative(r)
78 | ret += f_v
79 | elif self.ps.material[p_j] == self.ps.material_solid:
80 | boundary_viscosity = 0.0
81 | # Boundary neighbors
82 | ## Akinci2012
83 | f_v = d * boundary_viscosity * (self.density_0 * self.ps.m_V[p_j] / (self.ps.density[p_i])) * v_xy / (
84 | r.norm()**2 + 0.01 * self.ps.support_radius**2) * self.cubic_kernel_derivative(r)
85 | ret += f_v
86 | if self.ps.is_dynamic_rigid_body(p_j):
87 | self.ps.acceleration[p_j] += -f_v * self.ps.density[p_i] / self.ps.density[p_j]
88 |
89 |
90 | @ti.kernel
91 | def compute_non_pressure_forces(self):
92 | for p_i in ti.grouped(self.ps.x):
93 | if self.ps.is_static_rigid_body(p_i):
94 | self.ps.acceleration[p_i].fill(0.0)
95 | continue
96 | ############## Body force ###############
97 | # Add body force
98 | d_v = ti.Vector(self.g)
99 | self.ps.acceleration[p_i] = d_v
100 | if self.ps.material[p_i] == self.ps.material_fluid:
101 | self.ps.for_all_neighbors(p_i, self.compute_non_pressure_forces_task, d_v)
102 | self.ps.acceleration[p_i] = d_v
103 |
104 |
105 | @ti.kernel
106 | def advect(self):
107 | # Update position
108 | for p_i in ti.grouped(self.ps.x):
109 | if self.ps.is_dynamic[p_i]:
110 | if self.ps.is_dynamic_rigid_body(p_i):
111 | self.ps.v[p_i] += self.dt[None] * self.ps.acceleration[p_i]
112 | self.ps.x[p_i] += self.dt[None] * self.ps.v[p_i]
113 |
114 |
115 | @ti.kernel
116 | def compute_DFSPH_factor(self):
117 | for p_i in ti.grouped(self.ps.x):
118 | if self.ps.material[p_i] != self.ps.material_fluid:
119 | continue
120 | sum_grad_p_k = 0.0
121 | grad_p_i = ti.Vector([0.0 for _ in range(self.ps.dim)])
122 |
123 | # `ret` concatenates `grad_p_i` and `sum_grad_p_k`
124 | ret = ti.Vector([0.0 for _ in range(self.ps.dim + 1)])
125 |
126 | self.ps.for_all_neighbors(p_i, self.compute_DFSPH_factor_task, ret)
127 |
128 | sum_grad_p_k = ret[3]
129 | for i in ti.static(range(3)):
130 | grad_p_i[i] = ret[i]
131 | sum_grad_p_k += grad_p_i.norm_sqr()
132 |
133 | # Compute pressure stiffness denominator
134 | factor = 0.0
135 | if sum_grad_p_k > 1e-6:
136 | factor = -1.0 / sum_grad_p_k
137 | else:
138 | factor = 0.0
139 | self.ps.dfsph_factor[p_i] = factor
140 |
141 |
142 | @ti.func
143 | def compute_DFSPH_factor_task(self, p_i, p_j, ret: ti.template()):
144 | if self.ps.material[p_j] == self.ps.material_fluid:
145 | # Fluid neighbors
146 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
147 | ret[3] += grad_p_j.norm_sqr() # sum_grad_p_k
148 | for i in ti.static(range(3)): # grad_p_i
149 | ret[i] -= grad_p_j[i]
150 | elif self.ps.material[p_j] == self.ps.material_solid:
151 | # Boundary neighbors
152 | ## Akinci2012
153 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
154 | for i in ti.static(range(3)): # grad_p_i
155 | ret[i] -= grad_p_j[i]
156 |
157 |
158 | @ti.kernel
159 | def compute_density_change(self):
160 | for p_i in ti.grouped(self.ps.x):
161 | if self.ps.material[p_i] != self.ps.material_fluid:
162 | continue
163 | ret = ti.Struct(density_adv=0.0, num_neighbors=0)
164 | self.ps.for_all_neighbors(p_i, self.compute_density_change_task, ret)
165 |
166 | # only correct positive divergence
167 | density_adv = ti.max(ret.density_adv, 0.0)
168 | num_neighbors = ret.num_neighbors
169 |
170 | # Do not perform divergence solve when paritlce deficiency happens
171 | if self.ps.dim == 3:
172 | if num_neighbors < 20:
173 | density_adv = 0.0
174 | else:
175 | if num_neighbors < 7:
176 | density_adv = 0.0
177 |
178 | self.ps.density_adv[p_i] = density_adv
179 |
180 |
181 | @ti.func
182 | def compute_density_change_task(self, p_i, p_j, ret: ti.template()):
183 | v_i = self.ps.v[p_i]
184 | v_j = self.ps.v[p_j]
185 | if self.ps.material[p_j] == self.ps.material_fluid:
186 | # Fluid neighbors
187 | ret.density_adv += self.ps.m_V[p_j] * (v_i - v_j).dot(self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j]))
188 | elif self.ps.material[p_j] == self.ps.material_solid:
189 | # Boundary neighbors
190 | ## Akinci2012
191 | ret.density_adv += self.ps.m_V[p_j] * (v_i - v_j).dot(self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j]))
192 |
193 | # Compute the number of neighbors
194 | ret.num_neighbors += 1
195 |
196 |
197 | @ti.kernel
198 | def compute_density_adv(self):
199 | for p_i in ti.grouped(self.ps.x):
200 | if self.ps.material[p_i] != self.ps.material_fluid:
201 | continue
202 | delta = 0.0
203 | self.ps.for_all_neighbors(p_i, self.compute_density_adv_task, delta)
204 | density_adv = self.ps.density[p_i] /self.density_0 + self.dt[None] * delta
205 | self.ps.density_adv[p_i] = ti.max(density_adv, 1.0)
206 |
207 |
208 | @ti.func
209 | def compute_density_adv_task(self, p_i, p_j, ret: ti.template()):
210 | v_i = self.ps.v[p_i]
211 | v_j = self.ps.v[p_j]
212 | if self.ps.material[p_j] == self.ps.material_fluid:
213 | # Fluid neighbors
214 | ret += self.ps.m_V[p_j] * (v_i - v_j).dot(self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j]))
215 | elif self.ps.material[p_j] == self.ps.material_solid:
216 | # Boundary neighbors
217 | ## Akinci2012
218 | ret += self.ps.m_V[p_j] * (v_i - v_j).dot(self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j]))
219 |
220 |
221 | @ti.kernel
222 | def compute_density_error(self, offset: float) -> float:
223 | density_error = 0.0
224 | for I in ti.grouped(self.ps.x):
225 | if self.ps.material[I] == self.ps.material_fluid:
226 | density_error += self.density_0 * self.ps.density_adv[I] - offset
227 | return density_error
228 |
229 | @ti.kernel
230 | def multiply_time_step(self, field: ti.template(), time_step: float):
231 | for I in ti.grouped(self.ps.x):
232 | if self.ps.material[I] == self.ps.material_fluid:
233 | field[I] *= time_step
234 |
235 |
236 | def divergence_solve(self):
237 | # TODO: warm start
238 | # Compute velocity of density change
239 | self.compute_density_change()
240 | inv_dt = 1 / self.dt[None]
241 | self.multiply_time_step(self.ps.dfsph_factor, inv_dt)
242 |
243 | m_iterations_v = 0
244 |
245 | # Start solver
246 | avg_density_err = 0.0
247 |
248 | while m_iterations_v < 1 or m_iterations_v < self.m_max_iterations_v:
249 |
250 | avg_density_err = self.divergence_solver_iteration()
251 | # Max allowed density fluctuation
252 | # use max density error divided by time step size
253 | eta = 1.0 / self.dt[None] * self.max_error_V * 0.01 * self.density_0
254 | # print("eta ", eta)
255 | if avg_density_err <= eta:
256 | break
257 | m_iterations_v += 1
258 | print(f"DFSPH - iteration V: {m_iterations_v} Avg density err: {avg_density_err}")
259 |
260 | # Multiply by h, the time step size has to be removed
261 | # to make the stiffness value independent
262 | # of the time step size
263 |
264 | # TODO: if warm start
265 | # also remove for kappa v
266 |
267 | self.multiply_time_step(self.ps.dfsph_factor, self.dt[None])
268 |
269 |
270 | def divergence_solver_iteration(self):
271 | self.divergence_solver_iteration_kernel()
272 | self.compute_density_change()
273 | density_err = self.compute_density_error(0.0)
274 | return density_err / self.ps.fluid_particle_num
275 |
276 |
277 | @ti.kernel
278 | def divergence_solver_iteration_kernel(self):
279 | # Perform Jacobi iteration
280 | for p_i in ti.grouped(self.ps.x):
281 | if self.ps.material[p_i] != self.ps.material_fluid:
282 | continue
283 | # evaluate rhs
284 | b_i = self.ps.density_adv[p_i]
285 | k_i = b_i*self.ps.dfsph_factor[p_i]
286 | ret = ti.Struct(dv=ti.Vector([0.0 for _ in range(self.ps.dim)]), k_i=k_i)
287 | # TODO: if warm start
288 | # get_kappa_V += k_i
289 | self.ps.for_all_neighbors(p_i, self.divergence_solver_iteration_task, ret)
290 | self.ps.v[p_i] += ret.dv
291 |
292 |
293 | @ti.func
294 | def divergence_solver_iteration_task(self, p_i, p_j, ret: ti.template()):
295 | if self.ps.material[p_j] == self.ps.material_fluid:
296 | # Fluid neighbors
297 | b_j = self.ps.density_adv[p_j]
298 | k_j = b_j * self.ps.dfsph_factor[p_j]
299 | k_sum = ret.k_i + self.density_0 / self.density_0 * k_j # TODO: make the neighbor density0 different for multiphase fluid
300 | if ti.abs(k_sum) > self.m_eps:
301 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
302 | ret.dv -= self.dt[None] * k_sum * grad_p_j
303 | elif self.ps.material[p_j] == self.ps.material_solid:
304 | # Boundary neighbors
305 | ## Akinci2012
306 | if ti.abs(ret.k_i) > self.m_eps:
307 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
308 | vel_change = -self.dt[None] * 1.0 * ret.k_i * grad_p_j
309 | ret.dv += vel_change
310 | if self.ps.is_dynamic_rigid_body(p_j):
311 | self.ps.acceleration[p_j] += -vel_change * (1 / self.dt[None]) * self.ps.density[p_i] / self.ps.density[p_j]
312 |
313 |
314 | def pressure_solve(self):
315 | inv_dt = 1 / self.dt[None]
316 | inv_dt2 = 1 / (self.dt[None] * self.dt[None])
317 |
318 | # TODO: warm start
319 |
320 | # Compute rho_adv
321 | self.compute_density_adv()
322 |
323 | self.multiply_time_step(self.ps.dfsph_factor, inv_dt2)
324 |
325 | m_iterations = 0
326 |
327 | # Start solver
328 | avg_density_err = 0.0
329 |
330 | while m_iterations < 1 or m_iterations < self.m_max_iterations:
331 |
332 | avg_density_err = self.pressure_solve_iteration()
333 | # Max allowed density fluctuation
334 | eta = self.max_error * 0.01 * self.density_0
335 | if avg_density_err <= eta:
336 | break
337 | m_iterations += 1
338 | print(f"DFSPH - iterations: {m_iterations} Avg density Err: {avg_density_err:.4f}")
339 | # Multiply by h, the time step size has to be removed
340 | # to make the stiffness value independent
341 | # of the time step size
342 |
343 | # TODO: if warm start
344 | # also remove for kappa v
345 |
346 | def pressure_solve_iteration(self):
347 | self.pressure_solve_iteration_kernel()
348 | self.compute_density_adv()
349 | density_err = self.compute_density_error(self.density_0)
350 | return density_err / self.ps.fluid_particle_num
351 |
352 |
353 | @ti.kernel
354 | def pressure_solve_iteration_kernel(self):
355 | # Compute pressure forces
356 | for p_i in ti.grouped(self.ps.x):
357 | if self.ps.material[p_i] != self.ps.material_fluid:
358 | continue
359 | # Evaluate rhs
360 | b_i = self.ps.density_adv[p_i] - 1.0
361 | k_i = b_i * self.ps.dfsph_factor[p_i]
362 |
363 | # TODO: if warmstart
364 | # get kappa V
365 | self.ps.for_all_neighbors(p_i, self.pressure_solve_iteration_task, k_i)
366 |
367 |
368 | @ti.func
369 | def pressure_solve_iteration_task(self, p_i, p_j, k_i: ti.template()):
370 | if self.ps.material[p_j] == self.ps.material_fluid:
371 | # Fluid neighbors
372 | b_j = self.ps.density_adv[p_j] - 1.0
373 | k_j = b_j * self.ps.dfsph_factor[p_j]
374 | k_sum = k_i + self.density_0 / self.density_0 * k_j # TODO: make the neighbor density0 different for multiphase fluid
375 | if ti.abs(k_sum) > self.m_eps:
376 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
377 | # Directly update velocities instead of storing pressure accelerations
378 | self.ps.v[p_i] -= self.dt[None] * k_sum * grad_p_j # ki, kj already contain inverse density
379 | elif self.ps.material[p_j] == self.ps.material_solid:
380 | # Boundary neighbors
381 | ## Akinci2012
382 | if ti.abs(k_i) > self.m_eps:
383 | grad_p_j = -self.ps.m_V[p_j] * self.cubic_kernel_derivative(self.ps.x[p_i] - self.ps.x[p_j])
384 |
385 | # Directly update velocities instead of storing pressure accelerations
386 | vel_change = - self.dt[None] * 1.0 * k_i * grad_p_j # kj already contains inverse density
387 | self.ps.v[p_i] += vel_change
388 | if self.ps.is_dynamic_rigid_body(p_j):
389 | self.ps.acceleration[p_j] += -vel_change * 1.0 / self.dt[None] * self.ps.density[p_i] / self.ps.density[p_j]
390 |
391 |
392 | @ti.kernel
393 | def predict_velocity(self):
394 | # compute new velocities only considering non-pressure forces
395 | for p_i in ti.grouped(self.ps.x):
396 | if self.ps.is_dynamic[p_i] and self.ps.material[p_i] == self.ps.material_fluid:
397 | self.ps.v[p_i] += self.dt[None] * self.ps.acceleration[p_i]
398 |
399 |
400 | def substep(self):
401 | self.compute_densities()
402 | self.compute_DFSPH_factor()
403 | if self.enable_divergence_solver:
404 | self.divergence_solve()
405 | self.compute_non_pressure_forces()
406 | self.predict_velocity()
407 | self.pressure_solve()
408 | self.advect()
409 |
--------------------------------------------------------------------------------
/IISPH.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 | from sph_base import SPHBase
3 |
4 |
5 | class IISPHSolver(SPHBase):
6 | def __init__(self, particle_system):
7 | super().__init__(particle_system)
8 |
9 | self.a_ii = ti.field(dtype=float, shape=self.ps.particle_max_num)
10 | self.density_deviation = ti.field(dtype=float, shape=self.ps.particle_max_num)
11 | self.last_pressure = ti.field(dtype=float, shape=self.ps.particle_max_num)
12 | self.avg_density_error = ti.field(dtype=float, shape=())
13 |
14 | self.ps.acceleration = ti.Vector.field(self.ps.dim, dtype=float)
15 | self.pressure_accel = ti.Vector.field(self.ps.dim, dtype=float)
16 | particle_node = ti.root.dense(ti.i, self.ps.particle_max_num)
17 | particle_node.place(self.ps.acceleration, self.pressure_accel)
18 | self.dt[None] = 2e-4
19 |
20 | @ti.kernel
21 | def predict_advection(self):
22 | # Compute a_ii
23 | for p_i in range(self.ps.particle_num[None]):
24 | x_i = self.ps.x[p_i]
25 | sum_neighbor = 0.0
26 | sum_neighbor_of_neighbor = 0.0
27 | m_Vi = self.ps.m_V[p_i]
28 | density_i = self.ps.density[p_i]
29 | density_i2 = density_i * density_i
30 | density_02 = self.density_0 * self.density_0
31 | self.a_ii[p_i] = 0.0
32 | # Fluid neighbors
33 | for j in range(self.ps.fluid_neighbors_num[p_i]):
34 | p_j = self.ps.fluid_neighbors[p_i, j]
35 | x_j = self.ps.x[p_j]
36 | sum_neighbor_inner = ti.Vector([0.0 for _ in range(self.ps.dim)])
37 | for k in range(self.ps.fluid_neighbors_num[p_i]):
38 | density_k = self.ps.density[k]
39 | density_k2 = density_k * density_k
40 | p_k = self.ps.fluid_neighbors[p_i, j]
41 | x_k = self.ps.x[p_k]
42 | sum_neighbor_inner += self.ps.m_V[p_k] * self.cubic_kernel_derivative(x_i - x_k) / density_k2
43 |
44 | kernel_grad_ij = self.cubic_kernel_derivative(x_i - x_j)
45 | sum_neighbor -= (self.ps.m_V[p_j] * sum_neighbor_inner).dot(kernel_grad_ij)
46 |
47 | sum_neighbor_of_neighbor -= (self.ps.m_V[p_j] * kernel_grad_ij).dot(kernel_grad_ij)
48 | sum_neighbor_of_neighbor *= m_Vi / density_i2
49 | self.a_ii[p_i] += (sum_neighbor + sum_neighbor_of_neighbor) * self.dt[None] * self.dt[None] * density_02
50 |
51 | # Boundary neighbors
52 | ## Akinci2012
53 | for j in range(self.ps.solid_neighbors_num[p_i]):
54 | p_j = self.ps.solid_neighbors[p_i, j]
55 | x_j = self.ps.x[p_j]
56 | sum_neighbor_inner = ti.Vector([0.0 for _ in range(self.ps.dim)])
57 | for k in range(self.ps.solid_neighbors_num[p_i]):
58 | density_k = self.ps.density[k]
59 | density_k2 = density_k * density_k
60 | p_k = self.ps.solid_neighbors[p_i, j]
61 | x_k = self.ps.x[p_k]
62 | sum_neighbor_inner += self.ps.m_V[p_k] * self.cubic_kernel_derivative(x_i - x_k) / density_k2
63 |
64 | kernel_grad_ij = self.cubic_kernel_derivative(x_i - x_j)
65 | sum_neighbor -= (self.ps.m_V[p_j] * sum_neighbor_inner).dot(kernel_grad_ij)
66 |
67 | sum_neighbor_of_neighbor -= (self.ps.m_V[p_j] * kernel_grad_ij).dot(kernel_grad_ij)
68 | sum_neighbor_of_neighbor *= m_Vi / density_i2
69 | self.a_ii[p_i] += (sum_neighbor + sum_neighbor_of_neighbor) * self.dt[None] * self.dt[None] * density_02
70 |
71 | # Compute source term (i.e., density deviation)
72 | # Compute the predicted v^star
73 | for p_i in range(self.ps.particle_num[None]):
74 | if self.ps.material[p_i] == self.ps.material_fluid:
75 | self.ps.v[p_i] += self.dt[None] * self.ps.acceleration[p_i]
76 |
77 | for p_i in range(self.ps.particle_num[None]):
78 | x_i = self.ps.x[p_i]
79 | density_i = self.ps.density[p_i]
80 | divergence = 0.0
81 | # Fluid neighbors
82 | for j in range(self.ps.fluid_neighbors_num[p_i]):
83 | p_j = self.ps.fluid_neighbors[p_i, j]
84 | x_j = self.ps.x[p_j]
85 | divergence += self.ps.m_V[p_j] * (self.ps.v[p_i] - self.ps.v[p_j]).dot(self.cubic_kernel_derivative(x_i - x_j))
86 |
87 | # Boundary neighbors
88 | ## Akinci2012
89 | for j in range(self.ps.solid_neighbors_num[p_i]):
90 | p_j = self.ps.solid_neighbors[p_i, j]
91 | x_j = self.ps.x[p_j]
92 | divergence += self.ps.m_V[p_j] * (self.ps.v[p_i] - self.ps.v[p_j]).dot(self.cubic_kernel_derivative(x_i - x_j))
93 |
94 | self.density_deviation[p_i] = self.density_0 - density_i - self.dt[None] * divergence * self.density_0
95 |
96 | # Clear all pressures
97 | for p_i in range(self.ps.particle_num[None]):
98 | # self.last_pressure[p_i] = 0.0
99 | # self.ps.pressure[p_i] = 0.0
100 | self.last_pressure[p_i] = 0.5 * self.ps.pressure[p_i]
101 |
102 | def pressure_solve(self):
103 | iteration = 0
104 | while iteration < 1000:
105 | self.avg_density_error[None] = 0.0
106 | self.pressure_solve_iteration()
107 | iteration += 1
108 | if iteration % 100 == 0:
109 | print(f'iter {iteration}, density err {self.avg_density_error[None]}')
110 | if self.avg_density_error[None] < 1e-3:
111 | # print(f'Stop criterion satisfied at iter {iteration}, density err {self.avg_density_error[None]}')
112 | break
113 |
114 | @ti.kernel
115 | def pressure_solve_iteration(self):
116 | omega = 0.5
117 | # Compute pressure acceleration
118 | for p_i in range(self.ps.particle_num[None]):
119 | # if self.ps.material[p_i] != self.ps.material_fluid:
120 | # self.pressure_accel[p_i].fill(0)
121 | # continue
122 | x_i = self.ps.x[p_i]
123 | d_v = ti.Vector([0.0 for _ in range(self.ps.dim)])
124 |
125 | dpi = self.last_pressure[p_i] / self.ps.density[p_i] ** 2
126 | # Fluid neighbors
127 | for j in range(self.ps.fluid_neighbors_num[p_i]):
128 | p_j = self.ps.fluid_neighbors[p_i, j]
129 | x_j = self.ps.x[p_j]
130 | dpj = self.last_pressure[p_j] / self.ps.density[p_j] ** 2
131 | # Compute the pressure force contribution, Symmetric Formula
132 | d_v += -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
133 | * self.cubic_kernel_derivative(x_i - x_j)
134 |
135 | # Boundary neighbors
136 | dpj = self.last_pressure[p_i] / self.density_0 ** 2
137 | ## Akinci2012
138 | for j in range(self.ps.solid_neighbors_num[p_i]):
139 | p_j = self.ps.solid_neighbors[p_i, j]
140 | x_j = self.ps.x[p_j]
141 | # Compute the pressure force contribution, Symmetric Formula
142 | d_v += -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
143 | * self.cubic_kernel_derivative(x_i - x_j)
144 | self.pressure_accel[p_i] += d_v
145 |
146 | # Compute Ap and compute new pressure
147 | for p_i in range(self.ps.particle_num[None]):
148 | x_i = self.ps.x[p_i]
149 | Ap = 0.0
150 | dt2 = self.dt[None] * self.dt[None]
151 | accel_p_i = self.pressure_accel[p_i]
152 | # Fluid neighbors
153 | for j in range(self.ps.fluid_neighbors_num[p_i]):
154 | p_j = self.ps.fluid_neighbors[p_i, j]
155 | x_j = self.ps.x[p_j]
156 | Ap += self.ps.m_V[p_j] * (accel_p_i - self.pressure_accel[p_j]).dot(self.cubic_kernel_derivative(x_i - x_j))
157 | # Boundary neighbors
158 | ## Akinci2012
159 | for j in range(self.ps.solid_neighbors_num[p_i]):
160 | p_j = self.ps.solid_neighbors[p_i, j]
161 | x_j = self.ps.x[p_j]
162 | Ap += self.ps.m_V[p_j] * (accel_p_i - self.pressure_accel[p_j]).dot(self.cubic_kernel_derivative(x_i - x_j))
163 | Ap *= dt2 * self.density_0
164 | # print(self.a_ii[1])
165 | if abs(self.a_ii[p_i]) > 1e-6:
166 | # Relaxed Jacobi
167 | self.ps.pressure[p_i] = ti.max(self.last_pressure[p_i] + omega * (self.density_deviation[p_i] - Ap) / self.a_ii[p_i], 0.0)
168 | else:
169 | self.ps.pressure[p_i] = 0.0
170 |
171 | if self.ps.pressure[p_i] != 0.0:
172 | # new_density = self.density_0
173 | # if p_i == 100:
174 | # print(" Ap ", Ap, " density deviation ", self.density_deviation[p_i], 'a_ii ', self.a_ii[p_i])
175 | self.avg_density_error[None] += abs(Ap - self.density_deviation[p_i]) / self.density_0
176 | self.avg_density_error[None] /= self.ps.particle_num[None]
177 | for p_i in range(self.ps.particle_num[None]):
178 | # Update the pressure
179 | self.last_pressure[p_i] = self.ps.pressure[p_i]
180 |
181 |
182 | @ti.kernel
183 | def compute_densities(self):
184 | for p_i in range(self.ps.particle_num[None]):
185 | if self.ps.material[p_i] != self.ps.material_fluid:
186 | continue
187 | x_i = self.ps.x[p_i]
188 | self.ps.density[p_i] = self.ps.m_V[p_i] * self.cubic_kernel(0.0)
189 | # Fluid neighbors
190 | for j in range(self.ps.fluid_neighbors_num[p_i]):
191 | p_j = self.ps.fluid_neighbors[p_i, j]
192 | x_j = self.ps.x[p_j]
193 | self.ps.density[p_i] += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
194 | # Boundary neighbors
195 | ## Akinci2012
196 | for j in range(self.ps.solid_neighbors_num[p_i]):
197 | p_j = self.ps.solid_neighbors[p_i, j]
198 | x_j = self.ps.x[p_j]
199 | self.ps.density[p_i] += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
200 | self.ps.density[p_i] *= self.density_0
201 |
202 | @ti.kernel
203 | def compute_pressure_forces(self):
204 | for p_i in range(self.ps.particle_num[None]):
205 | if self.ps.material[p_i] != self.ps.material_fluid:
206 | self.pressure_accel[p_i].fill(0)
207 | continue
208 | self.pressure_accel[p_i].fill(0)
209 | x_i = self.ps.x[p_i]
210 | d_v = ti.Vector([0.0 for _ in range(self.ps.dim)])
211 |
212 | dpi = self.ps.pressure[p_i] / self.ps.density[p_i] ** 2
213 | # Fluid neighbors
214 | for j in range(self.ps.fluid_neighbors_num[p_i]):
215 | p_j = self.ps.fluid_neighbors[p_i, j]
216 | x_j = self.ps.x[p_j]
217 | dpj = self.ps.pressure[p_j] / self.ps.density[p_j] ** 2
218 | # Compute the pressure force contribution, Symmetric Formula
219 | d_v += -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
220 | * self.cubic_kernel_derivative(x_i - x_j)
221 |
222 | # Boundary neighbors
223 | dpj = self.ps.pressure[p_i] / self.density_0 ** 2
224 | # dpj = 0.0
225 | ## Akinci2012
226 | for j in range(self.ps.solid_neighbors_num[p_i]):
227 | p_j = self.ps.solid_neighbors[p_i, j]
228 | x_j = self.ps.x[p_j]
229 | # Compute the pressure force contribution, Symmetric Formula
230 | d_v += -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
231 | * self.cubic_kernel_derivative(x_i - x_j)
232 |
233 | self.pressure_accel[p_i] = d_v
234 |
235 | @ti.kernel
236 | def compute_non_pressure_forces(self):
237 | for p_i in range(self.ps.particle_num[None]):
238 | # if self.ps.material[p_i] != self.ps.material_fluid:
239 | # self.ps.acceleration[p_i].fill(0)
240 | # continue
241 | x_i = self.ps.x[p_i]
242 | # Add body force
243 | d_v = ti.Vector([0.0 for _ in range(self.ps.dim)])
244 | d_v[1] = self.g
245 | for j in range(self.ps.fluid_neighbors_num[p_i]):
246 | p_j = self.ps.fluid_neighbors[p_i, j]
247 | x_j = self.ps.x[p_j]
248 | d_v += self.viscosity_force(p_i, p_j, x_i - x_j)
249 | self.ps.acceleration[p_i] = d_v
250 |
251 | @ti.kernel
252 | def advect(self):
253 | # Symplectic Euler
254 | for p_i in range(self.ps.particle_num[None]):
255 | if self.ps.material[p_i] == self.ps.material_fluid:
256 | self.ps.v[p_i] += self.dt[None] * self.pressure_accel[p_i]
257 | self.ps.x[p_i] += self.dt[None] * self.ps.v[p_i]
258 |
259 | def substep(self):
260 | self.compute_densities()
261 | self.compute_non_pressure_forces()
262 |
263 | self.predict_advection()
264 | self.pressure_solve()
265 |
266 | self.compute_pressure_forces()
267 | self.advect()
268 |
--------------------------------------------------------------------------------
/LICENSE.txt:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2022 Mingrui Zhang
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # SPH Taichi
2 |
3 | A high-performance implementation of Smooth Particle Hydrodynamics (SPH) simulator in [Taichi](https://github.com/taichi-dev/taichi). (working in progress)
4 |
5 | ## Examples
6 |
7 | - Dragon Bath (~420 K particles, ~280 FPS on RTX 3090 GPU, with timestep 4e-4)
8 |
9 |
10 | 
11 |
12 |
13 | - Armadillo Bath (~1.74 M particles, ~80 FPS on RTX 3090 GPU, with timestep 4e-4)
14 |
15 |
16 | 
17 |
18 |
19 | ## Features
20 |
21 | Currently, the following features have been implemented:
22 | - Cross-platform: Windows, Linux
23 | - Support massively parallel GPU computing
24 | - Weakly Compressible SPH (WCSPH)[1]
25 | - One-way/two-way fluid-solid coupling[2]
26 | - Shape-matching based rigid-body simulator[3]
27 | - Neighborhood search accelerated by GPU parallel prefix sum + counting sort
28 |
29 | ### Note
30 | The GPU parallel prefix sum is only supported by cuda/vulkan backend currently.
31 |
32 | ## Install
33 |
34 | ```
35 | python -m pip install -r requirements.txt
36 | ```
37 |
38 | To reproduce the demos show above:
39 |
40 | ```
41 | python run_simulation.py --scene_file ./data/scenes/dragon_bath.json
42 | ```
43 |
44 | ```
45 | python run_simulation.py --scene_file ./data/scenes/armadillo_bath_dynamic.json
46 | ```
47 |
48 |
49 | ## Reference
50 | 1. M. Becker and M. Teschner (2007). "Weakly compressible SPH for free surface flows". In:Proceedings of the 2007 ACM SIGGRAPH/Eurographics symposium on Computer animation. Eurographics Association, pp. 209–217.
51 | 2. N. Akinci, M. Ihmsen, G. Akinci, B. Solenthaler, and M. Teschner. 2012. Versatile
52 | rigid-fluid coupling for incompressible SPH. ACM Transactions on Graphics 31, 4 (2012), 62:1–62:8.
53 | 3. Miles Macklin, Matthias Müller, Nuttapong Chentanez, and Tae-Yong Kim. 2014. Unified particle physics for real-time applications. ACM Trans. Graph. 33, 4, Article 153 (July 2014), 12 pages.
54 |
55 |
56 | ## Acknowledgement
57 | Implementation is largely inspired by [SPlisHSPlasH](https://github.com/InteractiveComputerGraphics/SPlisHSPlasH).
58 |
--------------------------------------------------------------------------------
/WCSPH.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 | from sph_base import SPHBase
3 |
4 |
5 | class WCSPHSolver(SPHBase):
6 | def __init__(self, particle_system):
7 | super().__init__(particle_system)
8 | # Pressure state function parameters(WCSPH)
9 | self.exponent = 7.0
10 | self.exponent = self.ps.cfg.get_cfg("exponent")
11 |
12 | self.stiffness = 50000.0
13 | self.stiffness = self.ps.cfg.get_cfg("stiffness")
14 |
15 | self.surface_tension = 0.01
16 | self.dt[None] = self.ps.cfg.get_cfg("timeStepSize")
17 |
18 |
19 | @ti.func
20 | def compute_densities_task(self, p_i, p_j, ret: ti.template()):
21 | x_i = self.ps.x[p_i]
22 | if self.ps.material[p_j] == self.ps.material_fluid:
23 | # Fluid neighbors
24 | x_j = self.ps.x[p_j]
25 | ret += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
26 | elif self.ps.material[p_j] == self.ps.material_solid:
27 | # Boundary neighbors
28 | ## Akinci2012
29 | x_j = self.ps.x[p_j]
30 | ret += self.ps.m_V[p_j] * self.cubic_kernel((x_i - x_j).norm())
31 |
32 |
33 | @ti.kernel
34 | def compute_densities(self):
35 | # for p_i in range(self.ps.particle_num[None]):
36 | for p_i in ti.grouped(self.ps.x):
37 | if self.ps.material[p_i] != self.ps.material_fluid:
38 | continue
39 | self.ps.density[p_i] = self.ps.m_V[p_i] * self.cubic_kernel(0.0)
40 | den = 0.0
41 | self.ps.for_all_neighbors(p_i, self.compute_densities_task, den)
42 | self.ps.density[p_i] += den
43 | self.ps.density[p_i] *= self.density_0
44 |
45 |
46 | @ti.func
47 | def compute_pressure_forces_task(self, p_i, p_j, ret: ti.template()):
48 | x_i = self.ps.x[p_i]
49 | dpi = self.ps.pressure[p_i] / self.ps.density[p_i] ** 2
50 | # Fluid neighbors
51 | if self.ps.material[p_j] == self.ps.material_fluid:
52 | x_j = self.ps.x[p_j]
53 | density_j = self.ps.density[p_j] * self.density_0 / self.density_0 # TODO: The density_0 of the neighbor may be different when the fluid density is different
54 | dpj = self.ps.pressure[p_j] / (density_j * density_j)
55 | # Compute the pressure force contribution, Symmetric Formula
56 | ret += -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
57 | * self.cubic_kernel_derivative(x_i-x_j)
58 | elif self.ps.material[p_j] == self.ps.material_solid:
59 | # Boundary neighbors
60 | dpj = self.ps.pressure[p_i] / self.density_0 ** 2
61 | ## Akinci2012
62 | x_j = self.ps.x[p_j]
63 | # Compute the pressure force contribution, Symmetric Formula
64 | f_p = -self.density_0 * self.ps.m_V[p_j] * (dpi + dpj) \
65 | * self.cubic_kernel_derivative(x_i-x_j)
66 | ret += f_p
67 | if self.ps.is_dynamic_rigid_body(p_j):
68 | self.ps.acceleration[p_j] += -f_p * self.density_0 / self.ps.density[p_j]
69 |
70 | @ti.kernel
71 | def compute_pressure_forces(self):
72 | for p_i in ti.grouped(self.ps.x):
73 | if self.ps.material[p_i] != self.ps.material_fluid:
74 | continue
75 | self.ps.density[p_i] = ti.max(self.ps.density[p_i], self.density_0)
76 | self.ps.pressure[p_i] = self.stiffness * (ti.pow(self.ps.density[p_i] / self.density_0, self.exponent) - 1.0)
77 | for p_i in ti.grouped(self.ps.x):
78 | if self.ps.is_static_rigid_body(p_i):
79 | self.ps.acceleration[p_i].fill(0)
80 | continue
81 | elif self.ps.is_dynamic_rigid_body(p_i):
82 | continue
83 | dv = ti.Vector([0.0 for _ in range(self.ps.dim)])
84 | self.ps.for_all_neighbors(p_i, self.compute_pressure_forces_task, dv)
85 | self.ps.acceleration[p_i] += dv
86 |
87 |
88 | @ti.func
89 | def compute_non_pressure_forces_task(self, p_i, p_j, ret: ti.template()):
90 | x_i = self.ps.x[p_i]
91 |
92 | ############## Surface Tension ###############
93 | if self.ps.material[p_j] == self.ps.material_fluid:
94 | # Fluid neighbors
95 | diameter2 = self.ps.particle_diameter * self.ps.particle_diameter
96 | x_j = self.ps.x[p_j]
97 | r = x_i - x_j
98 | r2 = r.dot(r)
99 | if r2 > diameter2:
100 | ret -= self.surface_tension / self.ps.m[p_i] * self.ps.m[p_j] * r * self.cubic_kernel(r.norm())
101 | else:
102 | ret -= self.surface_tension / self.ps.m[p_i] * self.ps.m[p_j] * r * self.cubic_kernel(ti.Vector([self.ps.particle_diameter, 0.0, 0.0]).norm())
103 |
104 |
105 | ############### Viscosoty Force ###############
106 | d = 2 * (self.ps.dim + 2)
107 | x_j = self.ps.x[p_j]
108 | # Compute the viscosity force contribution
109 | r = x_i - x_j
110 | v_xy = (self.ps.v[p_i] -
111 | self.ps.v[p_j]).dot(r)
112 |
113 | if self.ps.material[p_j] == self.ps.material_fluid:
114 | f_v = d * self.viscosity * (self.ps.m[p_j] / (self.ps.density[p_j])) * v_xy / (
115 | r.norm()**2 + 0.01 * self.ps.support_radius**2) * self.cubic_kernel_derivative(r)
116 | ret += f_v
117 | elif self.ps.material[p_j] == self.ps.material_solid:
118 | boundary_viscosity = 0.0
119 | # Boundary neighbors
120 | ## Akinci2012
121 | f_v = d * boundary_viscosity * (self.density_0 * self.ps.m_V[p_j] / (self.ps.density[p_i])) * v_xy / (
122 | r.norm()**2 + 0.01 * self.ps.support_radius**2) * self.cubic_kernel_derivative(r)
123 | ret += f_v
124 | if self.ps.is_dynamic_rigid_body(p_j):
125 | self.ps.acceleration[p_j] += -f_v * self.density_0 / self.ps.density[p_j]
126 |
127 |
128 | @ti.kernel
129 | def compute_non_pressure_forces(self):
130 | for p_i in ti.grouped(self.ps.x):
131 | if self.ps.is_static_rigid_body(p_i):
132 | self.ps.acceleration[p_i].fill(0.0)
133 | continue
134 | ############## Body force ###############
135 | # Add body force
136 | d_v = ti.Vector(self.g)
137 | self.ps.acceleration[p_i] = d_v
138 | if self.ps.material[p_i] == self.ps.material_fluid:
139 | self.ps.for_all_neighbors(p_i, self.compute_non_pressure_forces_task, d_v)
140 | self.ps.acceleration[p_i] = d_v
141 |
142 |
143 | @ti.kernel
144 | def advect(self):
145 | # Symplectic Euler
146 | for p_i in ti.grouped(self.ps.x):
147 | if self.ps.is_dynamic[p_i]:
148 | self.ps.v[p_i] += self.dt[None] * self.ps.acceleration[p_i]
149 | self.ps.x[p_i] += self.dt[None] * self.ps.v[p_i]
150 |
151 |
152 | def substep(self):
153 | self.compute_densities()
154 | self.compute_non_pressure_forces()
155 | self.compute_pressure_forces()
156 | self.advect()
157 |
--------------------------------------------------------------------------------
/config_builder.py:
--------------------------------------------------------------------------------
1 | import json
2 |
3 |
4 | class SimConfig:
5 | def __init__(self, scene_file_path) -> None:
6 | self.config = None
7 | with open(scene_file_path, "r") as f:
8 | self.config = json.load(f)
9 | print(self.config)
10 |
11 | def get_cfg(self, name, enforce_exist=False):
12 | if enforce_exist:
13 | assert name in self.config["Configuration"]
14 | if name not in self.config["Configuration"]:
15 | if enforce_exist:
16 | assert name in self.config["Configuration"]
17 | else:
18 | return None
19 | return self.config["Configuration"][name]
20 |
21 | def get_rigid_bodies(self):
22 | if "RigidBodies" in self.config:
23 | return self.config["RigidBodies"]
24 | else:
25 | return []
26 |
27 | def get_rigid_blocks(self):
28 | if "RigidBlocks" in self.config:
29 | return self.config["RigidBlocks"]
30 | else:
31 | return []
32 |
33 | def get_fluid_blocks(self):
34 | if "FluidBlocks" in self.config:
35 | return self.config["FluidBlocks"]
36 | else:
37 | return []
38 |
--------------------------------------------------------------------------------
/data/BoxOpenedHole.obj:
--------------------------------------------------------------------------------
1 | # Blender v2.81 (sub 16) OBJ File: ''
2 | # www.blender.org
3 | v -0.500000 1.000000 0.500000
4 | v 0.500000 1.000000 0.500000
5 | v -0.500000 2.000000 0.500000
6 | v 0.500000 2.000000 0.500000
7 | v -0.500000 2.000000 -0.500000
8 | v 0.500000 2.000000 -0.500000
9 | v -0.500000 1.000000 -0.500000
10 | v 0.500000 1.000000 -0.500000
11 | v 0.000000 1.000000 0.098995
12 | v 0.098995 1.000000 0.000000
13 | v 0.000000 1.000000 -0.098995
14 | v -0.098995 1.000000 0.000000
15 | v -0.550000 0.905825 0.550000
16 | v -0.550000 2.005825 0.550000
17 | v 0.550000 0.905825 0.550000
18 | v 0.550000 2.005825 0.550000
19 | v -0.550000 2.005825 -0.550000
20 | v -0.550000 0.905825 -0.550000
21 | v 0.550000 2.005825 -0.550000
22 | v 0.550000 0.905825 -0.550000
23 | v -0.000000 0.905825 0.108895
24 | v 0.108894 0.905825 0.000000
25 | v -0.000000 0.905825 -0.108894
26 | v -0.108894 0.905825 0.000000
27 | vn 0.0000 0.0000 -1.0000
28 | vn 0.0000 -0.0000 1.0000
29 | vn -1.0000 0.0000 0.0000
30 | vn 1.0000 0.0000 0.0000
31 | vn 0.0000 1.0000 0.0000
32 | vn 0.0000 -1.0000 -0.0000
33 | vn -0.1157 0.9933 0.0000
34 | vn -0.7052 -0.0741 0.7052
35 | vn 0.0000 0.9933 0.1157
36 | vn 0.1157 0.9933 0.0000
37 | vn 0.7052 -0.0741 0.7052
38 | vn -0.7052 -0.0741 -0.7052
39 | vn 0.7052 -0.0741 -0.7052
40 | vn 0.0000 0.9933 -0.1157
41 | s off
42 | f 1//1 3//1 2//1
43 | f 3//1 4//1 2//1
44 | f 5//2 7//2 6//2
45 | f 7//2 8//2 6//2
46 | f 2//3 4//3 8//3
47 | f 4//3 6//3 8//3
48 | f 7//4 5//4 1//4
49 | f 5//4 3//4 1//4
50 | f 1//5 2//5 9//5
51 | f 9//5 2//5 10//5
52 | f 2//5 8//5 10//5
53 | f 10//5 8//5 11//5
54 | f 11//5 8//5 7//5
55 | f 11//5 7//5 12//5
56 | f 7//5 1//5 12//5
57 | f 12//5 1//5 9//5
58 | f 13//2 15//2 14//2
59 | f 14//2 15//2 16//2
60 | f 17//1 19//1 18//1
61 | f 18//1 19//1 20//1
62 | f 15//4 20//4 16//4
63 | f 16//4 20//4 19//4
64 | f 18//3 13//3 17//3
65 | f 17//3 13//3 14//3
66 | f 13//6 21//6 15//6
67 | f 21//6 22//6 15//6
68 | f 15//6 22//6 20//6
69 | f 22//6 23//6 20//6
70 | f 23//6 18//6 20//6
71 | f 23//6 24//6 18//6
72 | f 18//6 24//6 13//6
73 | f 24//6 21//6 13//6
74 | f 4//7 19//7 6//7
75 | f 11//8 22//8 10//8
76 | f 6//9 17//9 5//9
77 | f 3//10 17//10 14//10
78 | f 12//11 23//11 11//11
79 | f 9//12 22//12 21//12
80 | f 9//13 24//13 12//13
81 | f 3//14 16//14 4//14
82 | f 4//7 16//7 19//7
83 | f 11//8 23//8 22//8
84 | f 6//9 19//9 17//9
85 | f 3//10 5//10 17//10
86 | f 12//11 24//11 23//11
87 | f 9//12 10//12 22//12
88 | f 9//13 21//13 24//13
89 | f 3//14 14//14 16//14
90 |
--------------------------------------------------------------------------------
/data/gif/armadillo_bath.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/data/gif/armadillo_bath.gif
--------------------------------------------------------------------------------
/data/gif/dragon_bath_large.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/data/gif/dragon_bath_large.gif
--------------------------------------------------------------------------------
/data/models/bunny.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/data/models/bunny.stl
--------------------------------------------------------------------------------
/data/scenes/armadillo_bath_dynamic.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [5.0, 3.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 0,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.0004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "RigidBodies": [
20 | {
21 | "objectId": 1,
22 | "geometryFile": "./data/models/armadillo_small.obj",
23 | "translation": [4.0, 2.0, 1.2],
24 | "rotationAxis": [0, 1, 0],
25 | "rotationAngle": 180,
26 | "scale": [0.25, 0.25, 0.25],
27 | "velocity": [0.0, -5.0, 0.0],
28 | "density": 7874.0,
29 | "color": [255, 255, 255],
30 | "isDynamic": true
31 | },
32 | {
33 | "objectId": 2,
34 | "geometryFile": "./data/models/armadillo_small.obj",
35 | "translation": [2.5, 2.0, 1.2],
36 | "rotationAxis": [0, 1, 0],
37 | "rotationAngle": 180,
38 | "scale": [0.25, 0.25, 0.25],
39 | "velocity": [0.0, -5.0, 0.0],
40 | "density": 1700.0,
41 | "color": [255, 100, 50],
42 | "isDynamic": true
43 | },
44 | {
45 | "objectId": 3,
46 | "geometryFile": "./data/models/armadillo_small.obj",
47 | "translation": [1.0, 2.0, 1.2],
48 | "rotationAxis": [0, 1, 0],
49 | "rotationAngle": 180,
50 | "scale": [0.25, 0.25, 0.25],
51 | "velocity": [0.0, -5.0, 0.0],
52 | "density": 300.0,
53 | "color": [100, 100, 50],
54 | "isDynamic": true
55 | }
56 | ],
57 | "FluidBlocks": [
58 | {
59 | "objectId": 0,
60 | "start": [0.04, 0.04, 0.04],
61 | "end": [4.96, 1.50, 1.96],
62 | "translation": [0.0, 0.0, 0.0],
63 | "scale": [1, 1, 1],
64 | "velocity": [0.0, 0.0, 0.0],
65 | "density": 1000.0,
66 | "color": [50, 100, 200]
67 |
68 | }
69 | ]
70 | }
71 |
--------------------------------------------------------------------------------
/data/scenes/armadillo_bath_dynamic_dfsph.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [5.0, 3.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 4,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "RigidBodies": [
20 | {
21 | "objectId": 1,
22 | "geometryFile": "./data/models/armadillo_small.obj",
23 | "translation": [4.0, 2.0, 1.2],
24 | "rotationAxis": [0, 1, 0],
25 | "rotationAngle": 180,
26 | "scale": [0.25, 0.25, 0.25],
27 | "velocity": [0.0, -5.0, 0.0],
28 | "density": 7874.0,
29 | "color": [255, 255, 255],
30 | "isDynamic": true
31 | },
32 | {
33 | "objectId": 2,
34 | "geometryFile": "./data/models/armadillo_small.obj",
35 | "translation": [2.5, 2.0, 1.2],
36 | "rotationAxis": [0, 1, 0],
37 | "rotationAngle": 180,
38 | "scale": [0.25, 0.25, 0.25],
39 | "velocity": [0.0, -5.0, 0.0],
40 | "density": 1700.0,
41 | "color": [255, 100, 50],
42 | "isDynamic": true
43 | },
44 | {
45 | "objectId": 3,
46 | "geometryFile": "./data/models/armadillo_small.obj",
47 | "translation": [1.0, 2.0, 1.2],
48 | "rotationAxis": [0, 1, 0],
49 | "rotationAngle": 180,
50 | "scale": [0.25, 0.25, 0.25],
51 | "velocity": [0.0, -5.0, 0.0],
52 | "density": 300.0,
53 | "color": [100, 100, 50],
54 | "isDynamic": true
55 | }
56 | ],
57 | "FluidBlocks": [
58 | {
59 | "objectId": 0,
60 | "start": [0.04, 0.04, 0.04],
61 | "end": [4.96, 1.50, 1.96],
62 | "translation": [0.0, 0.0, 0.0],
63 | "scale": [1, 1, 1],
64 | "velocity": [0.0, 0.0, 0.0],
65 | "density": 1000.0,
66 | "color": [50, 100, 200]
67 |
68 | }
69 | ]
70 | }
71 |
--------------------------------------------------------------------------------
/data/scenes/dragon_bath.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [5.0, 3.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 0,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.0004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "RigidBodies": [
20 | {
21 | "objectId": 1,
22 | "geometryFile": "./data/models/Dragon_50k.obj",
23 | "translation": [3.5, 0.05, 1.0],
24 | "rotationAxis": [0, 1, 0],
25 | "rotationAngle": 0,
26 | "scale": [1, 1, 1],
27 | "velocity": [0.0, 0.0, 0.0],
28 | "density": 1000.0,
29 | "color": [255, 255, 255],
30 | "isDynamic": false
31 | }
32 | ],
33 | "FluidBlocks": [
34 | {
35 | "objectId": 0,
36 | "start": [0.1, 0.1, 0.5],
37 | "end": [1.2, 2.9, 1.6],
38 | "translation": [0.2, 0.0, 0.2],
39 | "scale": [1, 1, 1],
40 | "velocity": [0.0, -1.0, 0.0],
41 | "density": 1000.0,
42 | "color": [50, 100, 200]
43 | }
44 | ]
45 | }
46 |
--------------------------------------------------------------------------------
/data/scenes/dragon_bath_dfsph.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [5.0, 3.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 4,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "RigidBodies": [
20 | {
21 | "objectId": 1,
22 | "geometryFile": "./data/models/Dragon_50k.obj",
23 | "translation": [3.5, 0.05, 1.0],
24 | "rotationAxis": [0, 1, 0],
25 | "rotationAngle": 0,
26 | "scale": [1, 1, 1],
27 | "velocity": [0.0, 0.0, 0.0],
28 | "density": 1000.0,
29 | "color": [255, 255, 255],
30 | "isDynamic": false
31 | }
32 | ],
33 | "FluidBlocks": [
34 | {
35 | "objectId": 0,
36 | "start": [0.1, 0.1, 0.5],
37 | "end": [1.2, 2.9, 1.6],
38 | "translation": [0.2, 0.0, 0.2],
39 | "scale": [1, 1, 1],
40 | "velocity": [0.0, -1.0, 0.0],
41 | "density": 1000.0,
42 | "color": [50, 100, 200]
43 | }
44 | ]
45 | }
46 |
--------------------------------------------------------------------------------
/data/scenes/dragon_bath_dynamic_dfsph.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [5.0, 3.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 4,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "RigidBodies": [
20 | {
21 | "objectId": 1,
22 | "geometryFile": "./data/models/Dragon_50k.obj",
23 | "translation": [3.5, 0.05, 1.0],
24 | "rotationAxis": [0, 1, 0],
25 | "rotationAngle": 0,
26 | "scale": [1, 1, 1],
27 | "velocity": [0.0, 0.0, 0.0],
28 | "density": 1000.0,
29 | "color": [255, 255, 255],
30 | "isDynamic": true
31 | }
32 | ],
33 | "FluidBlocks": [
34 | {
35 | "objectId": 0,
36 | "start": [0.1, 0.1, 0.5],
37 | "end": [1.2, 2.9, 1.6],
38 | "translation": [0.2, 0.0, 0.2],
39 | "scale": [1, 1, 1],
40 | "velocity": [0.0, -1.0, 0.0],
41 | "density": 1000.0,
42 | "color": [50, 100, 200]
43 | }
44 | ]
45 | }
46 |
--------------------------------------------------------------------------------
/data/scenes/high_fluid_dfsph.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [2.0, 6.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 4,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "FluidBlocks": [
20 | {
21 | "objectId": 0,
22 | "start": [0.0, 0.0, 0.0],
23 | "end": [0.6, 5.4, 0.6],
24 | "translation": [0.1, 0.1, 0.1],
25 | "scale": [1, 1, 1],
26 | "velocity": [0.0, 0.0, 0.0],
27 | "density": 1000.0,
28 | "color": [50, 100, 200]
29 | }
30 | ]
31 | }
32 |
--------------------------------------------------------------------------------
/data/scenes/high_fluid_wcsph.json:
--------------------------------------------------------------------------------
1 | {
2 | "Configuration":
3 | {
4 | "domainStart": [0.0, 0.0, 0.0],
5 | "domainEnd": [2.0, 6.0, 2.0],
6 | "particleRadius": 0.01,
7 | "numberOfStepsPerRenderUpdate": 1,
8 | "density0": 1000,
9 | "simulationMethod": 0,
10 | "gravitation": [0.0, -9.81, 0.0],
11 | "timeStepSize": 0.0004,
12 | "stiffness": 50000,
13 | "exponent": 7,
14 | "boundaryHandlingMethod": 0,
15 | "exportFrame": false,
16 | "exportPly": false,
17 | "exportObj": false
18 | },
19 | "FluidBlocks": [
20 | {
21 | "objectId": 0,
22 | "start": [0.0, 0.0, 0.0],
23 | "end": [0.6, 5.4, 0.6],
24 | "translation": [0.1, 0.1, 0.1],
25 | "scale": [1, 1, 1],
26 | "velocity": [0.0, 0.0, 0.0],
27 | "density": 1000.0,
28 | "color": [50, 100, 200]
29 | }
30 | ]
31 | }
32 |
--------------------------------------------------------------------------------
/demo_high_fluid.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 | import numpy as np
3 | import trimesh as tm
4 | from particle_system import ParticleSystem
5 | from WCSPH import WCSPHSolver
6 | from IISPH import IISPHSolver
7 |
8 | # ti.init(arch=ti.cpu)
9 |
10 | # Use GPU for higher peformance if available
11 | ti.init(arch=ti.cuda, device_memory_GB=2, kernel_profiler=True)
12 |
13 |
14 | if __name__ == "__main__":
15 | x_max = 2.0
16 | y_max = 6.0
17 | z_max = 2.0
18 |
19 | domain_size = np.array([x_max, y_max, z_max])
20 |
21 | box_anchors = ti.Vector.field(3, dtype=ti.f32, shape = 8)
22 | box_anchors[0] = ti.Vector([0.0, 0.0, 0.0])
23 | box_anchors[1] = ti.Vector([0.0, y_max, 0.0])
24 | box_anchors[2] = ti.Vector([x_max, 0.0, 0.0])
25 | box_anchors[3] = ti.Vector([x_max, y_max, 0.0])
26 |
27 | box_anchors[4] = ti.Vector([0.0, 0.0, z_max])
28 | box_anchors[5] = ti.Vector([0.0, y_max, z_max])
29 | box_anchors[6] = ti.Vector([x_max, 0.0, z_max])
30 | box_anchors[7] = ti.Vector([x_max, y_max, z_max])
31 |
32 | box_lines_indices = ti.field(int, shape=(2 * 12))
33 |
34 | for i, val in enumerate([0, 1, 0, 2, 1, 3, 2, 3, 4, 5, 4, 6, 5, 7, 6, 7, 0, 4, 1, 5, 2, 6, 3, 7]):
35 | box_lines_indices[i] = val
36 |
37 | dim = 3
38 | substeps = 1
39 | output_frames = False
40 | output_ply = False
41 | solver_type = "WCSPH"
42 | # solver_type = "IISPH"
43 | ps = ParticleSystem(domain_size, GGUI=True)
44 |
45 | x_offset = 0.2
46 | y_offset = 0.2
47 | z_offset = 0.2
48 |
49 | # mesh = tm.load("./data/Dragon_50k.obj")
50 | # # mesh = tm.load("./data/bunny_sparse.obj")
51 | # # mesh = tm.load("./data/bunny.stl")
52 | # mesh_scale = 0.8
53 | # mesh.apply_scale(mesh_scale)
54 | # offset = np.array([1.5, 0.0 + y_offset, 1.0])
55 | # is_success = tm.repair.fill_holes(mesh)
56 | # print("Is the mesh successfully repaired? ", is_success)
57 | # voxelized_mesh = mesh.voxelized(pitch=ps.particle_diameter).fill()
58 | # # voxelized_mesh = mesh.voxelized(pitch=ps.particle_diameter).hollow()
59 | # # voxelized_mesh.show()
60 | # voxelized_points_np = voxelized_mesh.points + offset
61 | # num_particles_obj = voxelized_points_np.shape[0]
62 | # voxelized_points = ti.Vector.field(3, ti.f32, num_particles_obj)
63 | # voxelized_points.from_numpy(voxelized_points_np)
64 |
65 | # print("Rigid body, num of particles: ", num_particles_obj)
66 |
67 | # ps.add_particles(2,
68 | # num_particles_obj,
69 | # voxelized_points_np, # position
70 | # 0.0 * np.ones((num_particles_obj, 3)), # velocity
71 | # 10 * np.ones(num_particles_obj), # density
72 | # np.zeros(num_particles_obj), # pressure
73 | # np.array([0 for _ in range(num_particles_obj)], dtype=int), # material
74 | # 1 * np.ones(num_particles_obj), # is_dynamic
75 | # 255 * np.ones((num_particles_obj, 3))) # color
76 |
77 | # Fluid -1
78 | ps.add_cube(object_id=0,
79 | lower_corner=[0.1+x_offset, 0.1 + y_offset, 0.5+z_offset],
80 | cube_size=[0.6, 5.4, 0.6],
81 | velocity=[0.0, -0.0, 0.0],
82 | density=1000.0,
83 | is_dynamic=1,
84 | color=(50,100,200),
85 | material=1)
86 |
87 | # # Bottom boundary
88 | # ps.add_cube(object_id=1,
89 | # lower_corner=[0.0+x_offset, 0.0 + y_offset, 0.0+z_offset],
90 | # cube_size=[x_max-x_offset*2, ps.particle_diameter-0.001, z_max-z_offset*2],
91 | # velocity=[0.0, 0.0, 0.0],
92 | # density=1000.0,
93 | # is_dynamic=0,
94 | # color=(255,255,255),
95 | # material=0)
96 |
97 | # # left boundary
98 | # ps.add_cube(object_id=1,
99 | # lower_corner=[0.0+x_offset, 0.0 + y_offset, 0.0+z_offset],
100 | # cube_size=[ps.particle_diameter-0.001, y_max-y_offset*2, z_max-z_offset*2],
101 | # velocity=[0.0, 0.0, 0.0],
102 | # density=1000.0,
103 | # is_dynamic=0,
104 | # color=(255,255,255),
105 | # material=0)
106 |
107 | # # back
108 | # ps.add_cube(object_id=1,
109 | # lower_corner=[0.0+x_offset, 0.0 + y_offset, 0.0+z_offset],
110 | # cube_size=[x_max-x_offset*2, y_max-y_offset*2, ps.particle_diameter-0.001],
111 | # velocity=[0.0, 0.0, 0.0],
112 | # density=1000.0,
113 | # is_dynamic=0,
114 | # color=(255,255,255),
115 | # material=0)
116 |
117 | # # front
118 | # ps.add_cube(object_id=1,
119 | # lower_corner=[0.0+x_offset, 0.0 + y_offset, z_max - z_offset],
120 | # cube_size=[x_max-x_offset*2, y_max-y_offset*2, ps.particle_diameter-0.001],
121 | # velocity=[0.0, 0.0, 0.0],
122 | # density=1000.0,
123 | # is_dynamic=0,
124 | # color=(255,255,255),
125 | # material=0)
126 |
127 | # # right
128 | # ps.add_cube(object_id=1,
129 | # lower_corner=[x_max-x_offset, 0.0 + y_offset, 0.0+z_offset],
130 | # cube_size=[ps.particle_diameter-0.001, y_max-y_offset*2, z_max-z_offset*2],
131 | # velocity=[0.0, 0.0, 0.0],
132 | # density=1000.0,
133 | # is_dynamic=0,
134 | # color=(255,255,255),
135 | # material=0)
136 |
137 | if solver_type == "WCSPH":
138 | solver = WCSPHSolver(ps)
139 | elif solver_type == "IISPH":
140 | solver = IISPHSolver(ps)
141 |
142 |
143 | window = ti.ui.Window('SPH', (1024, 1024), show_window = True, vsync=False)
144 |
145 | scene = ti.ui.Scene()
146 | camera = ti.ui.make_camera()
147 | camera.position(0.0, 3.0, 5.0)
148 | camera.up(0.0, 1.0, 0.0)
149 | camera.lookat(0.0, 0.0, 0.0)
150 | camera.fov(60)
151 | scene.set_camera(camera)
152 |
153 | canvas = window.get_canvas()
154 | radius = 0.002
155 | movement_speed = 0.02
156 | background_color = (0, 0, 0) # 0xFFFFFF
157 | particle_color = (1, 1, 1)
158 |
159 | cnt = 0
160 | cnt_ply = 0
161 | solver.initialize_solver()
162 | series_prefix = "output/object_{}_demo_test.ply"
163 |
164 |
165 | # ti.profiler.clear_kernel_profiler_info()
166 | while window.running:
167 | for i in range(substeps):
168 | solver.step()
169 | # ps.copy_to_vis_buffer(invisible_objects=[1])
170 | if ps.dim == 2:
171 | canvas.set_background_color(background_color)
172 | canvas.circles(ps.x_vis_buffer, radius=ps.particle_radius / 5, color=particle_color)
173 | elif ps.dim == 3:
174 | # # user controlling of camera
175 | # position_change = ti.Vector([0.0, 0.0, 0.0])
176 | # up = ti.Vector([0.0, 1.0, 0.0])
177 | # # move camera up and down
178 | # if window.is_pressed("e"):
179 | # position_change = up * movement_speed
180 | # if window.is_pressed("q"):
181 | # position_change = -up * movement_speed
182 | # camera.position(*(camera.curr_position + position_change))
183 | # camera.lookat(*(camera.curr_lookat + position_change))
184 | camera.track_user_inputs(window, movement_speed=movement_speed, hold_key=ti.ui.LMB)
185 | scene.set_camera(camera)
186 |
187 | scene.point_light((2.0, 2.0, 2.0), color=(1.0, 1.0, 1.0))
188 | # scene.particles(ps.x_vis_buffer, radius=ps.particle_radius, per_vertex_color=ps.color_vis_buffer)
189 | scene.particles(ps.x, radius=ps.particle_radius, color=(50/255,100/255,200/255))
190 |
191 | scene.lines(box_anchors, indices=box_lines_indices, color = (0.99, 0.68, 0.28), width = 1.0)
192 | canvas.scene(scene)
193 |
194 | if output_frames:
195 | if cnt % 2 == 0:
196 | window.write_image(f"img_high_fluid_output/{cnt:04}.png")
197 | if output_ply:
198 | if cnt % 20 == 0:
199 | obj_id = 0
200 | obj_data = ps.dump(obj_id=obj_id)
201 | np_pos = obj_data["position"]
202 | writer = ti.tools.PLYWriter(num_vertices=ps.object_collection[obj_id])
203 | writer.add_vertex_pos(np_pos[:, 0], np_pos[:, 1], np_pos[:, 2])
204 | writer.export_frame_ascii(cnt_ply, series_prefix.format(0))
205 | cnt_ply += 1
206 | cnt += 1
207 | window.show()
208 | ti.profiler.print_kernel_profiler_info()
209 |
--------------------------------------------------------------------------------
/legacy/README.md:
--------------------------------------------------------------------------------
1 | # SPH_Taichi
2 | A [Taichi](https://github.com/taichi-dev/taichi) implementation of Smooth Particle Hydrodynamics (SPH) simulator. Taichi is a productive & portable programming language for high-performance, sparse & differentiable computing. The 2D case below can be run and rendered efficiently on a laptop thanks to Taichi.
3 |
4 | ## Features
5 | Currently, the following features have been implemented:
6 | - Weakly Compressible SPH (WCSPH)[1]
7 | - Predictive-Corrective Incompressible SPH (PCISPH)[2]
8 | - Divergence free SPH (DFSPH)[3]
9 |
10 | ### Note: Updates on November 6, 2021
11 | The code is now compatible with Taichi `v0.8.3`.
12 |
13 | ## Example
14 |
15 | - Demo (PCISPH, 4.5k particles)
16 |
17 | Run ```python scene.py --method PCISPH```
18 |
19 | 
20 |
21 |
22 | Demos for the other two methods: ```python scene.py --method WCSPH``` or ```python scene.py --method DFSPH```
23 |
24 | ## Reference
25 | 1. M. Becker and M. Teschner (2007). "Weakly compressible SPH for free surface flows". In:Proceedings of the 2007 ACM SIGGRAPH/Eurographics symposium on Computer animation. Eurographics Association, pp. 209–217.
26 | 2. B. Solenthaler and R. Pajarola (2009). “Predictive-corrective incompressible SPH”. In: ACM SIGGRAPH 2009 papers, pp. 1–6.
27 | 3. J. Bender, D. Koschier (2015) Divergence-free smoothed particle hydrodynamics[C]//Proceedings of the 14th ACM SIGGRAPH/Eurographics symposium on computer animation. ACM, 2015: 147-155.
28 |
29 |
--------------------------------------------------------------------------------
/legacy/engine/__init__.py:
--------------------------------------------------------------------------------
1 | from . import sph_solver
2 |
--------------------------------------------------------------------------------
/legacy/engine/sph_solver.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import time
3 | from functools import reduce
4 | import taichi as ti
5 |
6 |
7 | @ti.data_oriented
8 | class SPHSolver:
9 | method_WCSPH = 0
10 | method_PCISPH = 1
11 | method_DFSPH = 2
12 | methods = {
13 | 'WCSPH': method_WCSPH,
14 | 'PCISPH': method_PCISPH,
15 | 'DFSPH': method_DFSPH
16 | }
17 | material_fluid = 1
18 | material_bound = 0
19 | materials = {'fluid': material_fluid, 'bound': material_bound}
20 |
21 | def __init__(self,
22 | res,
23 | screen_to_world_ratio,
24 | bound,
25 | alpha=0.5,
26 | dx=0.2,
27 | max_num_particles=2**20,
28 | padding=12,
29 | max_time=5.0,
30 | max_steps=50000,
31 | dynamic_allocate=False,
32 | adaptive_time_step=True,
33 | method=0):
34 | self.method = method
35 | self.adaptive_time_step = adaptive_time_step
36 | self.dim = len(res)
37 | self.res = res
38 | self.screen_to_world_ratio = screen_to_world_ratio
39 | self.dynamic_allocate = dynamic_allocate
40 | # self.padding = padding / screen_to_world_ratio
41 | self.padding = 2 * dx
42 | self.max_time = max_time
43 | self.max_steps = max_steps
44 | self.max_num_particles = max_num_particles
45 |
46 | self.g = -9.80 # Gravity
47 | self.alpha = alpha # viscosity
48 | self.rho_0 = 1000.0 # reference density
49 | self.CFL_v = 0.25 # CFL coefficient for velocity
50 | self.CFL_a = 0.05 # CFL coefficient for acceleration
51 |
52 | self.df_fac = 1.3
53 | self.dx = dx # Particle radius
54 | self.dh = self.dx * self.df_fac # Smooth length
55 | # Declare dt as ti var for adaptive time step
56 | self.dt = ti.field(ti.f32, shape=())
57 | # Particle parameters
58 | self.m = self.dx**self.dim * self.rho_0
59 |
60 | self.grid_size = 2 * self.dh
61 | self.grid_pos = np.ceil(
62 | np.array(res) / self.screen_to_world_ratio /
63 | self.grid_size).astype(int)
64 |
65 | self.top_bound = bound[0] # top_bound
66 | self.bottom_bound = bound[1] # bottom_bound
67 | self.left_bound = bound[2] # left_bound
68 | self.right_bound = bound[3] # right_bound
69 |
70 | # ----------- WCSPH parameters-----------------
71 | # Pressure state function parameters(WCSPH)
72 | self.gamma = 7.0
73 | self.c_0 = 200.0
74 |
75 | # ----------- PCISPH parameters-----------------
76 | # Scaling factor for PCISPH
77 | self.s_f = ti.field(ti.f32, shape=())
78 | # Max iteration steps for pressure correction
79 | self.it = 0
80 | self.max_it = 0
81 | self.sub_max_iteration = 3
82 | self.rho_err = ti.Vector.field(1, dtype=ti.f32, shape=())
83 | self.max_rho_err = ti.Vector.field(1, dtype=ti.f32, shape=())
84 |
85 | # ----------- DFSPH parameters-----------------
86 | # Summing up the rho for all particles to compute the average rho
87 | self.sum_rho_err = ti.field(ti.f32, shape=())
88 | self.sum_drho = ti.field(ti.f32, shape=())
89 |
90 | # Dynamic Fill particles use
91 | self.source_bound = ti.Vector.field(self.dim, dtype=ti.f32, shape=2)
92 | self.source_velocity = ti.Vector.field(self.dim, dtype=ti.f32, shape=())
93 | self.source_pressure = ti.Vector.field(1, dtype=ti.f32, shape=())
94 | self.source_density = ti.Vector.field(1, dtype=ti.f32, shape=())
95 |
96 | self.particle_num = ti.field(ti.i32, shape=())
97 | self.particle_positions = ti.Vector.field(self.dim, dtype=ti.f32)
98 | self.particle_velocity = ti.Vector.field(self.dim, dtype=ti.f32)
99 | self.particle_positions_new = ti.Vector.field(
100 | self.dim, dtype=ti.f32) # Prediction values for P-C scheme use
101 | self.particle_velocity_new = ti.Vector.field(
102 | self.dim, dtype=ti.f32) # Prediction values for P-C scheme use
103 | self.particle_pressure = ti.Vector.field(1, dtype=ti.f32)
104 | self.particle_pressure_acc = ti.Vector.field(
105 | self.dim, dtype=ti.f32) # pressure force for PCISPH use
106 | self.particle_density = ti.Vector.field(1, dtype=ti.f32)
107 | self.particle_density_new = ti.Vector.field(
108 | 1, dtype=ti.f32) # Prediction values for P-C scheme use
109 | self.particle_alpha = ti.Vector.field(1, dtype=ti.f32) # For DFSPH use
110 | self.particle_stiff = ti.Vector.field(1, dtype=ti.f32) # For DFSPH use
111 |
112 | self.color = ti.field(dtype=ti.f32)
113 | self.material = ti.field(dtype=ti.f32)
114 |
115 | self.d_velocity = ti.Vector.field(self.dim, dtype=ti.f32)
116 | self.d_density = ti.Vector.field(1, dtype=ti.f32)
117 |
118 | self.grid_num_particles = ti.field(ti.i32)
119 | self.grid2particles = ti.field(ti.i32)
120 | self.particle_num_neighbors = ti.field(ti.i32)
121 | self.particle_neighbors = ti.field(ti.i32)
122 |
123 | self.max_num_particles_per_cell = 100
124 | self.max_num_neighbors = 100
125 |
126 | self.max_v = 0.0
127 | self.max_a = 0.0
128 | self.max_rho = 0.0
129 | self.max_pressure = 0.0
130 |
131 | if dynamic_allocate:
132 | ti.root.dynamic(ti.i, max_num_particles, 2**18).place(
133 | self.particle_positions, self.particle_velocity,
134 | self.particle_pressure, self.particle_density,
135 | self.particle_density_new, self.d_velocity, self.d_density,
136 | self.material, self.color, self.particle_positions_new,
137 | self.particle_velocity_new, self.particle_pressure_acc,
138 | self.particle_alpha, self.particle_stiff)
139 | else:
140 | # Allocate enough memory
141 | ti.root.dense(ti.i, 2**18).place(
142 | self.particle_positions, self.particle_velocity,
143 | self.particle_pressure, self.particle_density,
144 | self.particle_density_new, self.d_velocity, self.d_density,
145 | self.material, self.color, self.particle_positions_new,
146 | self.particle_velocity_new, self.particle_pressure_acc,
147 | self.particle_alpha, self.particle_stiff)
148 |
149 | if self.dim == 2:
150 | grid_snode = ti.root.dense(ti.ij, self.grid_pos)
151 | grid_snode.place(self.grid_num_particles)
152 | grid_snode.dense(ti.k, self.max_num_particles_per_cell).place(
153 | self.grid2particles)
154 | else:
155 | grid_snode = ti.root.dense(ti.ijk, self.grid_pos)
156 | grid_snode.place(self.grid_num_particles)
157 | grid_snode.dense(ti.l, self.max_num_particles_per_cell).place(
158 | self.grid2particles)
159 |
160 | nb_node = ti.root.dynamic(ti.i, max_num_particles)
161 | nb_node.place(self.particle_num_neighbors)
162 | nb_node.dense(ti.j,
163 | self.max_num_neighbors).place(self.particle_neighbors)
164 |
165 | # Initialize dt
166 | if method == SPHSolver.method_WCSPH:
167 | self.dt.from_numpy(
168 | np.array(0.1 * self.dh / self.c_0, dtype=np.float32))
169 | self.CFL_v = 0.20 # CFL coefficient for velocity
170 | self.CFL_a = 0.20 # CFL coefficient for acceleration
171 |
172 | if method == SPHSolver.method_PCISPH:
173 | self.s_f.from_numpy(np.array(1.0, dtype=np.float32))
174 | if self.adaptive_time_step:
175 | self.dt.from_numpy(np.array(0.0015, dtype=np.float32))
176 | else:
177 | self.dt.from_numpy(np.array(0.00015, dtype=np.float32))
178 |
179 | if method == SPHSolver.method_DFSPH:
180 | self.dt.from_numpy(
181 | np.array(1.0 * self.dh / self.c_0, dtype=np.float32))
182 | self.CFL_v = 0.30 # CFL coefficient for velocity
183 | self.CFL_a = 0.05 # CFL coefficient for acceleration
184 |
185 | @ti.func
186 | def compute_grid_index(self, pos):
187 | return (pos / (2 * self.dh)).cast(int)
188 |
189 | @ti.kernel
190 | def allocate_particles(self):
191 | # Ref to pbf2d example from by Ye Kuang (k-ye)
192 | # https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py
193 | # allocate particles to grid
194 | for p_i in range(self.particle_num[None]):
195 | # Compute the grid index
196 | cell = self.compute_grid_index(self.particle_positions[p_i])
197 | offs = self.grid_num_particles[cell].atomic_add(1)
198 | self.grid2particles[cell, offs] = p_i
199 |
200 | @ti.func
201 | def is_in_grid(self, c):
202 | res = 1
203 | for i in ti.static(range(self.dim)):
204 | res = ti.atomic_and(res, (0 <= c[i] < self.grid_pos[i]))
205 | return res
206 |
207 | @ti.func
208 | def is_fluid(self, p):
209 | # check fluid particle or bound particle
210 | return self.material[p]
211 |
212 | @ti.kernel
213 | def search_neighbors(self):
214 | # Ref to pbf2d example from by Ye Kuang (k-ye)
215 | # https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py
216 | for p_i in range(self.particle_num[None]):
217 | pos_i = self.particle_positions[p_i]
218 | nb_i = 0
219 | if self.is_fluid(p_i) == 1 or self.is_fluid(p_i) == 0:
220 | # Compute the grid index on the fly
221 | cell = self.compute_grid_index(self.particle_positions[p_i])
222 | for offs in ti.static(
223 | ti.grouped(ti.ndrange(*((-1, 2), ) * self.dim))):
224 | cell_to_check = cell + offs
225 | if self.is_in_grid(cell_to_check) == 1:
226 | for j in range(self.grid_num_particles[cell_to_check]):
227 | p_j = self.grid2particles[cell_to_check, j]
228 | if nb_i < self.max_num_neighbors and p_j != p_i and (
229 | pos_i - self.particle_positions[p_j]
230 | ).norm() < self.dh * 2.00:
231 | self.particle_neighbors[p_i, nb_i] = p_j
232 | nb_i.atomic_add(1)
233 | self.particle_num_neighbors[p_i] = nb_i
234 |
235 | @ti.func
236 | def cubic_kernel(self, r, h):
237 | # value of cubic spline smoothing kernel
238 | k = 10. / (7. * np.pi * h**self.dim)
239 | q = r / h
240 | # assert q >= 0.0 # Metal backend is not happy with assert
241 | res = ti.cast(0.0, ti.f32)
242 | if q <= 1.0:
243 | res = k * (1 - 1.5 * q**2 + 0.75 * q**3)
244 | elif q < 2.0:
245 | res = k * 0.25 * (2 - q)**3
246 | return res
247 |
248 | @ti.func
249 | def cubic_kernel_derivative(self, r, h):
250 | # derivative of cubic spline smoothing kernel
251 | k = 10. / (7. * np.pi * h**self.dim)
252 | q = r / h
253 | # assert q > 0.0
254 | res = ti.cast(0.0, ti.f32)
255 | if q < 1.0:
256 | res = (k / h) * (-3 * q + 2.25 * q**2)
257 | elif q < 2.0:
258 | res = -0.75 * (k / h) * (2 - q)**2
259 | return res
260 |
261 | @ti.func
262 | def rho_derivative(self, ptc_i, ptc_j, r, r_mod):
263 | # density delta, i.e. divergence
264 | return self.m * self.cubic_kernel_derivative(r_mod, self.dh) \
265 | * (self.particle_velocity[ptc_i] - self.particle_velocity[ptc_j]).dot(r / r_mod)
266 |
267 | @ti.func
268 | def p_update(self, rho, rho_0=1000.0, gamma=7.0, c_0=20.0):
269 | # Weakly compressible, tait function
270 | b = rho_0 * c_0**2 / gamma
271 | return b * ((rho / rho_0)**gamma - 1.0)
272 |
273 | @ti.func
274 | def pressure_force(self, ptc_i, ptc_j, r, r_mod):
275 | # Compute the pressure force contribution, Symmetric Formula
276 | res = ti.Vector([0.0 for _ in range(self.dim)])
277 | res = -self.m * (self.particle_pressure[ptc_i][0] / self.particle_density[ptc_i][0] ** 2
278 | + self.particle_pressure[ptc_j][0] / self.particle_density[ptc_j][0] ** 2) \
279 | * self.cubic_kernel_derivative(r_mod, self.dh) * r / r_mod
280 | return res
281 |
282 | @ti.func
283 | def viscosity_force(self, ptc_i, ptc_j, r, r_mod):
284 | # Compute the viscosity force contribution, artificial viscosity
285 | res = ti.Vector([0.0 for _ in range(self.dim)])
286 | v_xy = (self.particle_velocity[ptc_i] -
287 | self.particle_velocity[ptc_j]).dot(r)
288 | if v_xy < 0:
289 | # Artifical viscosity
290 | vmu = -2.0 * self.alpha * self.dx * self.c_0 / (
291 | self.particle_density[ptc_i][0] +
292 | self.particle_density[ptc_j][0])
293 | res = -self.m * vmu * v_xy / (
294 | r_mod**2 + 0.01 * self.dx**2) * self.cubic_kernel_derivative(
295 | r_mod, self.dh) * r / r_mod
296 | return res
297 |
298 | @ti.func
299 | def simulate_collisions(self, ptc_i, vec, d):
300 | # Collision factor, assume roughly (1-c_f)*velocity loss after collision
301 | c_f = 0.3
302 | self.particle_positions[ptc_i] += vec * d
303 | self.particle_velocity[ptc_i] -= (
304 | 1.0 + c_f) * self.particle_velocity[ptc_i].dot(vec) * vec
305 | if self.method == SPHSolver.method_DFSPH:
306 | self.particle_velocity_new[ptc_i] -= (
307 | 1.0 + c_f) * self.particle_velocity_new[ptc_i].dot(vec) * vec
308 |
309 | @ti.kernel
310 | def enforce_boundary(self):
311 | # TODO: only handle 2D case currently
312 | for p_i in range(self.particle_num[None]):
313 | if self.is_fluid(p_i) == 1:
314 | pos = self.particle_positions[p_i]
315 | if pos[0] < self.left_bound + 0.5 * self.padding:
316 | self.simulate_collisions(
317 | p_i, ti.Vector([1.0, 0.0]),
318 | self.left_bound + 0.5 * self.padding - pos[0])
319 | if pos[0] > self.right_bound - 0.5 * self.padding:
320 | self.simulate_collisions(
321 | p_i, ti.Vector([-1.0, 0.0]),
322 | pos[0] - self.right_bound + 0.5 * self.padding)
323 | if pos[1] > self.top_bound - self.padding:
324 | self.simulate_collisions(
325 | p_i, ti.Vector([0.0, -1.0]),
326 | pos[1] - self.top_bound + self.padding)
327 | if pos[1] < self.bottom_bound + self.padding:
328 | self.simulate_collisions(
329 | p_i, ti.Vector([0.0, 1.0]),
330 | self.bottom_bound + self.padding - pos[1])
331 |
332 | @ti.kernel
333 | def wc_compute_deltas(self):
334 | for p_i in range(self.particle_num[None]):
335 | pos_i = self.particle_positions[p_i]
336 | d_v = ti.Vector([0.0 for _ in range(self.dim)])
337 | d_rho = 0.0
338 | for j in range(self.particle_num_neighbors[p_i]):
339 | p_j = self.particle_neighbors[p_i, j]
340 | pos_j = self.particle_positions[p_j]
341 |
342 | # Compute distance and its mod
343 | r = pos_i - pos_j
344 | r_mod = ti.max(r.norm(), 1e-5)
345 |
346 | # Compute Density change
347 | d_rho += self.rho_derivative(p_i, p_j, r, r_mod)
348 |
349 | if self.is_fluid(p_i) == 1:
350 | # Compute Viscosity force contribution
351 | d_v += self.viscosity_force(p_i, p_j, r, r_mod)
352 |
353 | # Compute Pressure force contribution
354 | d_v += self.pressure_force(p_i, p_j, r, r_mod)
355 |
356 | # Add body force
357 | if self.is_fluid(p_i) == 1:
358 | val = [0.0 for _ in range(self.dim - 1)]
359 | val.extend([self.g])
360 | d_v += ti.Vector(val)
361 | self.d_velocity[p_i] = d_v
362 | self.d_density[p_i][0] = d_rho
363 |
364 | @ti.kernel
365 | def wc_update_time_step(self):
366 | # Simple Forward Euler currently
367 | for p_i in range(self.particle_num[None]):
368 | if self.is_fluid(p_i) == 1:
369 | self.particle_velocity[p_i] += self.dt[None] * self.d_velocity[p_i]
370 | self.particle_positions[
371 | p_i] += self.dt[None] * self.particle_velocity[p_i]
372 | self.particle_density[p_i][0] += self.dt[None] * self.d_density[p_i][0]
373 | self.particle_pressure[p_i][0] = self.p_update(
374 | self.particle_density[p_i][0], self.rho_0, self.gamma,
375 | self.c_0)
376 |
377 | @ti.kernel
378 | def pci_scaling_factor(self):
379 | grad_sum = ti.Vector([0.0 for _ in range(self.dim)])
380 | grad_dot_sum = 0.0
381 | range_num = ti.cast(self.dh * 2.0 / self.dx, ti.i32)
382 | half_range = ti.cast(0.5 * range_num, ti.i32)
383 | # TODO: only handle 2D case currently
384 | for x in range(-half_range, half_range):
385 | for y in range(-half_range, half_range):
386 | r = ti.Vector([-x * self.dx, -y * self.dx])
387 | r_mod = r.norm()
388 | if 2.0 * self.dh > r_mod > 1e-5:
389 | grad = self.cubic_kernel_derivative(r_mod,
390 | self.dh) * r / r_mod
391 | grad_sum += grad
392 | grad_dot_sum += grad.dot(grad)
393 |
394 | beta = 2 * (self.dt[None] * self.m / self.rho_0)**2
395 | self.s_f[None] = 1.0 / ti.max(
396 | beta * (grad_sum.dot(grad_sum) + grad_dot_sum), 1e-6)
397 |
398 | @ti.kernel
399 | def pci_pos_vel_prediction(self):
400 | for p_i in range(self.particle_num[None]):
401 | if self.is_fluid(p_i) == 1:
402 | self.particle_velocity_new[
403 | p_i] = self.particle_velocity[p_i] + self.dt[None] * (
404 | self.d_velocity[p_i] + self.particle_pressure_acc[p_i])
405 | self.particle_positions_new[p_i] = self.particle_positions[
406 | p_i] + self.dt[None] * self.particle_velocity_new[p_i]
407 | # Initialize the max_rho_err
408 | self.max_rho_err[None][0] = 0.0
409 |
410 | @ti.kernel
411 | def pci_update_pressure(self):
412 | for p_i in range(self.particle_num[None]):
413 | pos_i = self.particle_positions_new[p_i]
414 | d_rho = 0.0
415 | curr_rho = 0.0
416 | for j in range(self.particle_num_neighbors[p_i]):
417 | p_j = self.particle_neighbors[p_i, j]
418 | pos_j = self.particle_positions_new[p_j]
419 |
420 | # Compute distance and its mod
421 | r = pos_i - pos_j
422 | r_mod = r.norm()
423 | if r_mod > 1e-5:
424 | # Compute Density change
425 | d_rho += self.cubic_kernel_derivative(r_mod, self.dh) \
426 | * (self.particle_velocity_new[p_i] - self.particle_velocity_new[p_j]).dot(r / r_mod)
427 |
428 | self.d_density[p_i][0] = d_rho
429 | # Avoid negative density variation
430 | self.rho_err[None][0] = self.particle_density[p_i][
431 | 0] + self.dt[None] * d_rho - self.rho_0
432 | self.max_rho_err[None][0] = max(abs(self.rho_err[None][0]),
433 | self.max_rho_err[None][0])
434 | self.particle_pressure[p_i][
435 | 0] += self.s_f[None] * self.rho_err[None][0]
436 |
437 | @ti.kernel
438 | def pci_update_pressure_force(self):
439 | for p_i in range(self.particle_num[None]):
440 | pos_i = self.particle_positions_new[p_i]
441 | d_vp = ti.Vector([0.0 for _ in range(self.dim)])
442 | for j in range(self.particle_num_neighbors[p_i]):
443 | p_j = self.particle_neighbors[p_i, j]
444 | pos_j = self.particle_positions_new[p_j]
445 | # Compute distance and its mod
446 | r = pos_i - pos_j
447 | r_mod = r.norm()
448 | if r_mod > 1e-5:
449 | # Compute Pressure force contribution
450 | d_vp += self.pressure_force(p_i, p_j, r, r_mod)
451 | self.particle_pressure_acc[p_i] = d_vp
452 |
453 | def pci_pc_iteration(self):
454 | self.pci_pos_vel_prediction()
455 | self.pci_update_pressure()
456 | self.pci_update_pressure_force()
457 |
458 | @ti.kernel
459 | def pci_compute_deltas(self):
460 | for p_i in range(self.particle_num[None]):
461 | pos_i = self.particle_positions[p_i]
462 | d_v = ti.Vector([0.0 for _ in range(self.dim)])
463 |
464 | for j in range(self.particle_num_neighbors[p_i]):
465 | p_j = self.particle_neighbors[p_i, j]
466 | pos_j = self.particle_positions[p_j]
467 |
468 | # Compute distance and its mod
469 | r = pos_i - pos_j
470 | r_mod = r.norm()
471 |
472 | if r_mod > 1e-5 and self.is_fluid(p_i) == 1:
473 | # Compute Viscosity force contribution
474 | d_v += self.viscosity_force(p_i, p_j, r, r_mod)
475 |
476 | # Add body force
477 | if self.is_fluid(p_i) == 1:
478 | val = [0.0 for _ in range(self.dim - 1)]
479 | val.extend([self.g])
480 | d_v += ti.Vector(val)
481 | self.d_velocity[p_i] = d_v
482 | # Initialize the pressure
483 | self.particle_pressure[p_i][0] = 0.0
484 | self.particle_pressure_acc[p_i] = ti.Vector(
485 | [0.0 for _ in range(self.dim)])
486 |
487 | @ti.kernel
488 | def pci_update_time_step(self):
489 | # Final position and velocity update
490 | for p_i in range(self.particle_num[None]):
491 | if self.is_fluid(p_i) == 1:
492 | self.particle_velocity[p_i] += self.dt[None] * (
493 | self.d_velocity[p_i] + self.particle_pressure_acc[p_i])
494 | self.particle_positions[
495 | p_i] += self.dt[None] * self.particle_velocity[p_i]
496 | # Update density
497 | self.particle_density[p_i][0] += self.dt[None] * self.d_density[p_i][0]
498 |
499 | @ti.kernel
500 | def df_compute_deltas(self):
501 | for p_i in range(self.particle_num[None]):
502 | pos_i = self.particle_positions[p_i]
503 | d_v = ti.Vector([0.0 for _ in range(self.dim)])
504 | for j in range(self.particle_num_neighbors[p_i]):
505 | p_j = self.particle_neighbors[p_i, j]
506 | pos_j = self.particle_positions[p_j]
507 |
508 | # Compute distance and its mod
509 | r = pos_i - pos_j
510 | r_mod = r.norm()
511 |
512 | if r_mod > 1e-4 and self.is_fluid(p_i) == 1:
513 | # Compute Viscosity force contribution
514 | d_v += self.viscosity_force(p_i, p_j, r, r_mod)
515 |
516 | # Add body force
517 | if self.is_fluid(p_i) == 1:
518 | val = [0.0 for _ in range(self.dim - 1)]
519 | val.extend([self.g])
520 | d_v += ti.Vector(val)
521 | self.d_velocity[p_i] = d_v
522 |
523 | @ti.kernel
524 | def df_predict_velocities(self):
525 | for p_i in range(self.particle_num[None]):
526 | if self.is_fluid(p_i) == 1:
527 | self.particle_velocity_new[p_i] = self.particle_velocity[
528 | p_i] + self.dt[None] * self.d_velocity[p_i]
529 |
530 | @ti.kernel
531 | def df_correct_density_predict(self):
532 | for p_i in range(self.particle_num[None]):
533 | pos_i = self.particle_positions[p_i]
534 | d_rho = 0.0
535 | for j in range(self.particle_num_neighbors[p_i]):
536 | p_j = self.particle_neighbors[p_i, j]
537 | pos_j = self.particle_positions[p_j]
538 |
539 | # Compute distance and its mod
540 | r = pos_i - pos_j
541 | r_mod = ti.max(r.norm(), 1e-5)
542 |
543 | # Compute Density change
544 | if self.is_fluid(p_j) == 1:
545 | d_rho += self.m * self.cubic_kernel_derivative(r_mod, self.dh) \
546 | * (self.particle_velocity_new[p_i] - self.particle_velocity_new[p_j]).dot(r / r_mod)
547 | elif self.is_fluid(p_j) == 0:
548 | d_rho += self.m * self.cubic_kernel_derivative(r_mod, self.dh) \
549 | * self.particle_velocity_new[p_i].dot(r / r_mod)
550 |
551 | # Compute the predicted density rho star
552 | self.particle_density_new[p_i][
553 | 0] = self.particle_density[p_i][0] + self.dt[None] * d_rho
554 |
555 | # Only consider compressed
556 | err = ti.max(0.0, self.particle_density_new[p_i][0] - self.rho_0)
557 | self.particle_stiff[p_i][0] = err * self.particle_alpha[p_i][0]
558 |
559 | # Compute the density error sum for average use
560 | self.sum_rho_err[None] += err
561 |
562 | @ti.kernel
563 | def df_correct_density_adapt_vel(self):
564 | for p_i in range(self.particle_num[None]):
565 | pos_i = self.particle_positions[p_i]
566 | d_v = ti.Vector([0.0 for _ in range(self.dim)])
567 | for j in range(self.particle_num_neighbors[p_i]):
568 | p_j = self.particle_neighbors[p_i, j]
569 | pos_j = self.particle_positions[p_j]
570 | # Compute distance and its mod
571 | r = pos_i - pos_j
572 | r_mod = r.norm()
573 |
574 | if r_mod > 1e-4:
575 | if self.is_fluid(p_j) == 1:
576 | d_v += self.m * (self.particle_stiff[p_i][0] +
577 | self.particle_stiff[p_j][0]
578 | ) * self.cubic_kernel_derivative(
579 | r_mod, self.dh) * r / r_mod
580 | elif self.is_fluid(p_j) == 0:
581 | d_v += self.m * self.particle_stiff[
582 | p_i][0] * self.cubic_kernel_derivative(
583 | r_mod, self.dh) * r / r_mod
584 |
585 | # Predict velocity using pressure contribution
586 | self.particle_velocity_new[p_i] += d_v / ti.max(self.dt[None], 1e-5)
587 | # Store the pressure contribution to acceleration
588 | self.particle_pressure_acc[p_i] = d_v / ti.max(
589 | self.dt[None] * self.dt[None], 1e-8)
590 |
591 | @ti.kernel
592 | def df_update_positions(self):
593 | for p_i in range(self.particle_num[None]):
594 | # Update the positions
595 | if self.is_fluid(p_i) == 1:
596 | self.particle_positions[
597 | p_i] += self.dt[None] * self.particle_velocity_new[p_i]
598 |
599 | @ti.kernel
600 | def df_compute_density_alpha(self):
601 | for p_i in range(self.particle_num[None]):
602 | pos_i = self.particle_positions[p_i]
603 | grad_sum = ti.Vector([0.0 for _ in range(self.dim)])
604 | grad_square_sum = 0.0
605 | curr_rho = 0.0
606 | for j in range(self.particle_num_neighbors[p_i]):
607 | p_j = self.particle_neighbors[p_i, j]
608 | pos_j = self.particle_positions[p_j]
609 | # Compute distance and its mod
610 | r = pos_i - pos_j
611 | r_mod = r.norm()
612 | if r_mod > 1e-4:
613 | # Compute the grad sum and grad square sum for denominator alpha
614 | grad_val = self.m * self.cubic_kernel_derivative(
615 | r_mod, self.dh) * r / r_mod
616 | grad_sum += grad_val
617 |
618 | if self.is_fluid(p_j) == 1:
619 | grad_square_sum += grad_val.dot(grad_val)
620 | # Compute the density
621 | curr_rho += self.m * self.cubic_kernel(r_mod, self.dh)
622 | # Update the density
623 | self.particle_density[p_i][0] = curr_rho
624 | # Set a threshold of 10^-6 to avoid instability
625 | self.particle_alpha[p_i][0] = -1.0 / ti.max(
626 | grad_sum.dot(grad_sum) + grad_square_sum, 1e-6)
627 |
628 | @ti.kernel
629 | def df_correct_divergence_compute_drho(self):
630 | for p_i in range(self.particle_num[None]):
631 | pos_i = self.particle_positions[p_i]
632 | d_rho = 0.0
633 | for j in range(self.particle_num_neighbors[p_i]):
634 | p_j = self.particle_neighbors[p_i, j]
635 | pos_j = self.particle_positions[p_j]
636 | # Compute distance and its mod
637 | r = pos_i - pos_j
638 | r_mod = r.norm()
639 |
640 | if r_mod > 1e-4:
641 | if self.is_fluid(p_j) == 1:
642 | d_rho += self.m * (
643 | self.particle_velocity_new[p_i] -
644 | self.particle_velocity_new[p_j]).dot(
645 | r / r_mod) * self.cubic_kernel_derivative(
646 | r_mod, self.dh)
647 | # Boundary particles have no contributions to pressure force
648 | elif self.is_fluid(p_j) == 0:
649 | d_rho += self.m * self.particle_velocity_new[p_i].dot(
650 | r / r_mod) * self.cubic_kernel_derivative(
651 | r_mod, self.dh)
652 |
653 | self.d_density[p_i][0] = ti.max(d_rho, 0.0)
654 |
655 | # if the density is less than the rest density, skip update in this iteration
656 | if self.particle_density[p_i][0] + self.dt[None] * self.d_density[p_i][
657 | 0] < self.rho_0 and self.particle_density[p_i][
658 | 0] < self.rho_0:
659 | self.d_density[p_i][0] = 0.0
660 | self.particle_stiff[p_i][
661 | 0] = self.d_density[p_i][0] * self.particle_alpha[p_i][0]
662 |
663 | # Compute the predicted density rho star
664 | self.sum_drho[None] += self.d_density[p_i][0]
665 |
666 | @ti.kernel
667 | def df_correct_divergence_adapt_vel(self):
668 | for p_i in range(self.particle_num[None]):
669 | pos_i = self.particle_positions[p_i]
670 | d_v = ti.Vector([0.0 for _ in range(self.dim)])
671 |
672 | for j in range(self.particle_num_neighbors[p_i]):
673 | p_j = self.particle_neighbors[p_i, j]
674 | pos_j = self.particle_positions[p_j]
675 | # Compute distance and its mod
676 | r = pos_i - pos_j
677 | r_mod = r.norm()
678 |
679 | if r_mod > 1e-5:
680 | if self.is_fluid(p_j) == 1:
681 | d_v += self.m * (self.particle_stiff[p_i][0] +
682 | self.particle_stiff[p_j][0]
683 | ) * self.cubic_kernel_derivative(
684 | r_mod, self.dh) * r / r_mod
685 | elif self.is_fluid(p_j) == 0:
686 | d_v += self.m * self.particle_stiff[
687 | p_i][0] * self.cubic_kernel_derivative(
688 | r_mod, self.dh) * r / r_mod
689 |
690 | # Predict velocity using pressure contribution, dt has been cancelled
691 | self.particle_velocity_new[p_i] += d_v
692 | # Store the pressure contribution to acceleration
693 | self.particle_pressure_acc[p_i] = d_v / self.dt[None]
694 |
695 | @ti.kernel
696 | def df_update_velocities(self):
697 | for p_i in range(self.particle_num[None]):
698 | if self.is_fluid(p_i) == 1:
699 | # Update the velocities from prediction values to next step
700 | self.particle_velocity[p_i] = self.particle_velocity_new[p_i]
701 |
702 | def sim_info(self, output=False):
703 | print("Time step: ", self.dt[None])
704 | print(
705 | "Domain: (%s, %s, %s, %s)" %
706 | (self.x_min, self.x_max, self.y_min, self.y_max), )
707 | print("Fluid area: (%s, %s, %s, %s)" %
708 | (self.left_bound, self.right_bound, self.bottom_bound,
709 | self.top_bound))
710 | print("Grid: ", self.grid_pos)
711 |
712 | def sim_info_realtime(self, frame, t, curr_start, curr_end, total_start):
713 | print(
714 | "Step: %d, physics time: %s, progress: %s %%, time used: %s, total time used: %s"
715 | % (frame, t,
716 | 100 * np.max([t / self.max_time, frame / self.max_steps]),
717 | curr_end - curr_start, curr_end - total_start))
718 | print(
719 | "Max velocity: %s, Max acceleration: %s, Max density: %s, Max pressure: %s"
720 | % (self.max_v, self.max_a, self.max_rho, self.max_pressure))
721 | if self.method == SPHSolver.methods['PCISPH']:
722 | print("Max iter: %d, Max density variation: %s" %
723 | (self.max_it, self.max_rho_err[None][0]))
724 | if self.method == SPHSolver.methods['DFSPH']:
725 | print("Max iter: %d, Max density variation: %s" %
726 | (self.it, self.sum_rho_err[None] / self.particle_num[None]))
727 | print("Max iter: %d, Max divergence variation: %s" %
728 | (self.it, self.sum_drho[None] / self.particle_num[None]))
729 | print("Adaptive time step: ", self.dt[None])
730 |
731 | def adaptive_step(self):
732 | total_num = self.particle_num[None]
733 | self.max_v = np.max(
734 | np.linalg.norm(self.particle_velocity.to_numpy()[:total_num],
735 | 2,
736 | axis=1))
737 | # CFL analysis, constrained by v_max
738 | dt_cfl = self.CFL_v * self.dh / self.max_v
739 |
740 | self.max_a = np.max(
741 | np.linalg.norm((self.d_velocity.to_numpy() +
742 | self.particle_pressure_acc.to_numpy())[:total_num],
743 | 2,
744 | axis=1))
745 | # Constrained by a_max
746 | dt_f = self.CFL_a * np.sqrt(self.dh / self.max_a)
747 |
748 | if self.adaptive_time_step and self.method == SPHSolver.method_DFSPH:
749 | self.dt[None] = np.min([dt_cfl, dt_f])
750 | return
751 |
752 | self.max_rho = np.max(self.particle_density.to_numpy()[:total_num])
753 | self.max_pressure = np.max(
754 | self.particle_pressure.to_numpy()[:total_num])
755 | dt_a = 0.20 * self.dh / (self.c_0 * np.sqrt(
756 | (self.max_rho / self.rho_0)**self.gamma))
757 |
758 | if self.adaptive_time_step and self.method == SPHSolver.method_WCSPH:
759 | self.dt[None] = np.min([dt_cfl, dt_f, dt_a])
760 | if self.adaptive_time_step and self.method == SPHSolver.method_PCISPH:
761 | self.dt[None] = np.min([dt_cfl, dt_f])
762 |
763 | def step(self, frame, t, total_start):
764 | curr_start = time.process_time()
765 |
766 | self.grid_num_particles.fill(0)
767 | self.particle_neighbors.fill(-1)
768 | self.allocate_particles()
769 | self.search_neighbors()
770 |
771 | if self.method == SPHSolver.methods['WCSPH']:
772 | # Compute deltas
773 | self.wc_compute_deltas()
774 | # timestep Update
775 | self.wc_update_time_step()
776 | elif self.method == SPHSolver.methods['PCISPH']:
777 | # Compute viscosity and gravity force
778 | self.pci_compute_deltas()
779 | # Compute the scaling factor for pressure update
780 | self.pci_scaling_factor()
781 | # Start density prediction-correction iteration process
782 | self.it = 0
783 | self.max_it = 0
784 | # 1% rho_0
785 | while self.max_rho_err[None][
786 | 0] >= 0.01 * self.rho_0 or self.it < self.sub_max_iteration:
787 | self.pci_pc_iteration()
788 | self.it += 1
789 | self.max_it += 1
790 | self.max_it = max(self.it, self.max_it)
791 | if self.it > 1000:
792 | print(
793 | "Warning: PCISPH density does not converge, iterated %d steps"
794 | % self.it)
795 | break
796 | # Compute new velocity, position and density
797 | self.pci_update_time_step()
798 | elif self.method == SPHSolver.methods['DFSPH']:
799 | # The step order changed a bit to fit the step function
800 | # update rho and alpha
801 | self.df_compute_density_alpha()
802 | # Correct divergence error
803 | self.it = 0
804 | self.sum_drho[None] = 0.0
805 | # 1% rho_0
806 | while self.sum_drho[None] >= 0.01 * self.particle_num[
807 | None] * self.rho_0 or self.it < 1:
808 | self.sum_drho[None] = 0.0
809 | self.df_correct_divergence_compute_drho()
810 | self.df_correct_divergence_adapt_vel()
811 | self.it += 1
812 | if self.it > 1000:
813 | print(
814 | "Warning: DFSPH divergence does not converge, iterated %d steps"
815 | % self.it)
816 | break
817 | # Update velocities v
818 | self.df_update_velocities()
819 | # Compute non-pressure forces
820 | self.df_compute_deltas()
821 | # Adapt time step
822 | self.adaptive_step()
823 | # Predict velocities v_star
824 | self.df_predict_velocities()
825 | # Correct density error
826 | self.it = 0
827 | self.sum_rho_err[None] = 0.0
828 | # 1% rho_0
829 | while self.sum_rho_err[None] >= 0.01 * self.particle_num[
830 | None] * self.rho_0 or self.it < 2:
831 | self.sum_rho_err[None] = 0.0
832 | self.df_correct_density_predict()
833 | self.df_correct_density_adapt_vel()
834 | self.it += 1
835 | if self.it > 1000:
836 | print(
837 | "Warning: DFSPH density does not converge, iterated %d steps"
838 | % self.it)
839 | break
840 | self.df_update_positions()
841 |
842 | # Handle potential leak particles
843 | self.enforce_boundary()
844 | if self.method != SPHSolver.methods['DFSPH']:
845 | self.adaptive_step()
846 |
847 | curr_end = time.process_time()
848 |
849 | if frame % 10 == 0:
850 | self.sim_info_realtime(frame, t, curr_start, curr_end, total_start)
851 | return self.dt[None]
852 |
853 | @ti.func
854 | def fill_particle(self, i, x, material, color, velocity, pressure,
855 | density):
856 | self.particle_positions[i] = x
857 | self.particle_positions_new[i] = x
858 | self.particle_velocity[i] = velocity
859 | self.particle_velocity_new[i] = velocity
860 | self.d_velocity[i] = ti.Vector([0.0 for _ in range(self.dim)])
861 | self.particle_pressure[i] = pressure
862 | self.particle_pressure_acc[i] = ti.Vector([0.0 for _ in range(self.dim)])
863 | self.particle_density[i] = density
864 | self.particle_density_new[i] = density
865 | self.d_density[i][0] = 0.0
866 | self.particle_alpha[i][0] = 0.0
867 | self.particle_stiff[i][0] = 0.0
868 | self.color[i] = color
869 | self.material[i] = material
870 |
871 | @ti.kernel
872 | def fill(self, new_particles: ti.i32, new_positions: ti.types.ndarray(),
873 | new_material: ti.i32, color: ti.i32):
874 | for i in range(self.particle_num[None],
875 | self.particle_num[None] + new_particles):
876 | self.material[i] = new_material
877 | x = ti.Vector.zero(ti.f32, self.dim)
878 | for k in ti.static(range(self.dim)):
879 | x[k] = new_positions[k, i - self.particle_num[None]]
880 | self.fill_particle(i, x, new_material, color,
881 | self.source_velocity[None],
882 | self.source_pressure[None],
883 | self.source_density[None])
884 |
885 | def set_source_velocity(self, velocity):
886 | if velocity is not None:
887 | velocity = list(velocity)
888 | assert len(velocity) == self.dim
889 | self.source_velocity[None] = velocity
890 | else:
891 | for i in range(self.dim):
892 | self.source_velocity[None][i] = 0
893 |
894 | def set_source_pressure(self, pressure):
895 | if pressure is not None:
896 | self.source_pressure[None] = pressure
897 | else:
898 | self.source_pressure[None][0] = 0.0
899 |
900 | def set_source_density(self, density):
901 | if density is not None:
902 | self.source_density[None] = density
903 | else:
904 | self.source_density[None][0] = 0.0
905 |
906 | def add_cube(self,
907 | lower_corner,
908 | cube_size,
909 | material,
910 | color=0xFFFFFF,
911 | density=None,
912 | pressure=None,
913 | velocity=None):
914 |
915 | num_dim = []
916 | for i in range(self.dim):
917 | num_dim.append(
918 | np.arange(lower_corner[i], lower_corner[i] + cube_size[i],
919 | self.dx))
920 | num_new_particles = reduce(lambda x, y: x * y,
921 | [len(n) for n in num_dim])
922 | assert self.particle_num[
923 | None] + num_new_particles <= self.max_num_particles
924 |
925 | new_positions = np.array(np.meshgrid(*num_dim,
926 | sparse=False,
927 | indexing='ij'),
928 | dtype=np.float32)
929 | new_positions = new_positions.reshape(
930 | -1, reduce(lambda x, y: x * y, list(new_positions.shape[1:])))
931 | print(new_positions.shape)
932 |
933 | for i in range(self.dim):
934 | self.source_bound[0][i] = lower_corner[i]
935 | self.source_bound[1][i] = cube_size[i]
936 |
937 | self.set_source_velocity(velocity=velocity)
938 | self.set_source_pressure(pressure=pressure)
939 | self.set_source_density(density=density)
940 |
941 | self.fill(num_new_particles, new_positions, material, color)
942 | # Add to current particles count
943 | self.particle_num[None] += num_new_particles
944 |
945 | @ti.kernel
946 | def copy_dynamic_nd(self, np_x: ti.types.ndarray(), input_x: ti.template()):
947 | for i in range(self.particle_num[None]):
948 | for j in ti.static(range(self.dim)):
949 | np_x[i, j] = input_x[i][j]
950 |
951 | @ti.kernel
952 | def copy_dynamic(self, np_x: ti.types.ndarray(), input_x: ti.template()):
953 | for i in range(self.particle_num[None]):
954 | np_x[i] = input_x[i]
955 |
956 | def particle_info(self):
957 | np_x = np.ndarray((self.particle_num[None], self.dim),
958 | dtype=np.float32)
959 | self.copy_dynamic_nd(np_x, self.particle_positions)
960 | np_v = np.ndarray((self.particle_num[None], self.dim),
961 | dtype=np.float32)
962 | self.copy_dynamic_nd(np_v, self.particle_velocity)
963 | np_material = np.ndarray((self.particle_num[None], ), dtype=np.int32)
964 | self.copy_dynamic(np_material, self.material)
965 | np_color = np.ndarray((self.particle_num[None], ), dtype=np.int32)
966 | self.copy_dynamic(np_color, self.color)
967 | return {
968 | 'position': np_x,
969 | 'velocity': np_v,
970 | 'material': np_material,
971 | 'color': np_color
972 | }
973 |
--------------------------------------------------------------------------------
/legacy/img/DFSPH.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/legacy/img/DFSPH.gif
--------------------------------------------------------------------------------
/legacy/img/PCISPH.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/legacy/img/PCISPH.gif
--------------------------------------------------------------------------------
/legacy/img/WCSPH.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/legacy/img/WCSPH.gif
--------------------------------------------------------------------------------
/legacy/img/sph_hv.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/legacy/img/sph_hv.gif
--------------------------------------------------------------------------------
/legacy/img/wcsph_alpha030.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/erizmr/SPH_Taichi/b02e13a7ce4c4abda144d6d3aaf2a7697f0f72e2/legacy/img/wcsph_alpha030.gif
--------------------------------------------------------------------------------
/legacy/scene.py:
--------------------------------------------------------------------------------
1 | # SPH taichi implementation by mzhang
2 | import taichi as ti
3 | from engine.sph_solver import *
4 | import argparse
5 |
6 | # Default run on CPU
7 | # cuda performance has not been tested
8 | ti.init(arch=ti.cpu)
9 |
10 |
11 | def main(opt):
12 | dynamic_allocate = opt.dynamic_allocate
13 | save_frames = opt.save
14 | adaptive_time_step = opt.adaptive
15 | method_name = opt.method
16 | sim_physical_time = 5.0
17 | max_frame = 50000
18 |
19 | res = (400, 400)
20 | screen_to_world_ratio = 35
21 | dx = 0.1
22 | u, b, l, r = np.array([res[1], 0, 0, res[0]]) / screen_to_world_ratio
23 |
24 | gui = ti.GUI('SPH', res, background_color=0x112F41)
25 | sph = SPHSolver(res,
26 | screen_to_world_ratio, [u, b, l, r],
27 | alpha=0.30,
28 | dx=dx,
29 | max_time=sim_physical_time,
30 | max_steps=max_frame,
31 | dynamic_allocate=dynamic_allocate,
32 | adaptive_time_step=adaptive_time_step,
33 | method=SPHSolver.methods[method_name])
34 |
35 | print("Method use: %s" % method_name)
36 | # Add fluid particles
37 | sph.add_cube(lower_corner=[res[0] / 2 / screen_to_world_ratio - 3, 4 * dx],
38 | cube_size=[6, 6],
39 | velocity=[0.0, -5.0],
40 | density=[1000],
41 | color=0x068587,
42 | material=SPHSolver.material_fluid)
43 |
44 | colors = np.array([0xED553B, 0x068587, 0xEEEEF0, 0xFFFF00],
45 | dtype=np.uint32)
46 | add_cnt = 0.0
47 | add = True
48 | save_cnt = 0.0
49 | output_fps = 60
50 | save_point = 1.0 / output_fps
51 | t = 0.0
52 | frame = 0
53 | total_start = time.process_time()
54 | while frame < max_frame and t < sim_physical_time:
55 | dt = sph.step(frame, t, total_start)
56 | particles = sph.particle_info()
57 |
58 | # if frame == 1000:
59 | if add and add_cnt > 0.40:
60 | sph.add_cube(lower_corner=[6, 6],
61 | cube_size=[2.0, 2.0],
62 | velocity=[0.0, -5.0],
63 | density=[1000.0],
64 | color=0xED553B,
65 | material=SPHSolver.material_fluid)
66 |
67 | # if frame == 1000:
68 | if add and add_cnt > 0.40:
69 | sph.add_cube(lower_corner=[3, 8],
70 | cube_size=[1.0, 1.0],
71 | velocity=[0.0, -10.0],
72 | density=[1000.0],
73 | color=0xEEEEF0,
74 | material=SPHSolver.material_fluid)
75 | add = False
76 |
77 | for pos in particles['position']:
78 | for j in range(len(res)):
79 | pos[j] *= screen_to_world_ratio / res[j]
80 |
81 | gui.circles(particles['position'],
82 | radius=1.5,
83 | color=particles['color'])
84 |
85 | # Save in fixed frame interval, for fixed time step
86 | if not adaptive_time_step:
87 | if frame % 50 == 0:
88 | gui.show(f'{frame:06d}.png' if save_frames else None)
89 | else:
90 | gui.show()
91 |
92 | # Save in fixed frame per second, for adaptive time step
93 | if adaptive_time_step:
94 | if save_cnt >= save_point:
95 | gui.show(f'{frame:06d}.png' if save_frames else None)
96 | else:
97 | gui.show()
98 | save_cnt = 0.0
99 |
100 | frame += 1
101 | t += dt
102 | save_cnt += dt
103 | add_cnt += dt
104 |
105 | print('done')
106 |
107 |
108 | if __name__ == '__main__':
109 | parser = argparse.ArgumentParser()
110 | parser.add_argument("--method",
111 | type=str,
112 | default="PCISPH",
113 | help="SPH methods: WCSPH, PCISPH, DFSPH")
114 | parser.add_argument("--save",
115 | action='store_true',
116 | help="save frames")
117 | parser.add_argument("--adaptive",
118 | action='store_true',
119 | help="whether apply adaptive step size")
120 | parser.add_argument("--dynamic-allocate",
121 | action='store_true',
122 | help="whether apply dynamic allocation")
123 | opt = parser.parse_args()
124 | main(opt)
125 |
--------------------------------------------------------------------------------
/legacy/test_sample.py:
--------------------------------------------------------------------------------
1 | # WCSPH implementation by mzhang
2 | import numpy as np
3 | import matplotlib.pyplot as plt
4 | import time
5 | from itertools import count
6 | import taichi as ti
7 |
8 | ti.init(arch=ti.cpu)
9 |
10 | df_fac = 1.3
11 | dx = 0.2
12 | dh = dx * df_fac
13 |
14 | ###### Scene parameters ########
15 | w = 20
16 | h = 10
17 | w_bound = 22
18 | h_bound = 12
19 |
20 | bottom_bound = 0.0
21 | top_bound = 0.0
22 | left_bound = 0.0
23 | right_bound = 0.0
24 |
25 | assert w_bound > w
26 | assert h_bound > h
27 | x_min = (w_bound - w) / 2.0
28 | y_min = (h_bound - h) / 2.0
29 | x_max = w_bound - (w_bound - w) / 2.0
30 | y_max = h_bound - (h_bound - h) / 2.0
31 |
32 | screen_res = (800, 400)
33 | screen_to_world_ratio = 35.0
34 | bg_color = 0x112f41
35 | particle_color = 0x068587
36 | boundary_color = 0xebaca2
37 | particle_radius = 3.0
38 | particle_radius_in_world = particle_radius / screen_to_world_ratio
39 |
40 |
41 | def setup():
42 | def computeGridIndex(x, y):
43 | idx = np.floor(x / (2 * dh)).astype(int)
44 | idy = np.floor(y / (2 * dh)).astype(int)
45 | return idx, idy
46 |
47 | def placeParticles(position_list, paticle_list, wall_mark, bound=0):
48 | # position_list: [start_x, start_y, end_x, end_y]
49 | start_x, start_y, end_x, end_y = position_list
50 | vel_x, vel_y, p, rho = 0.0, 0.0, 0.0, 1000.0
51 | for pos_x in np.arange(start_x, end_x, dx, dtype=np.float32):
52 | for pos_y in np.arange(start_y, end_y, dx, dtype=np.float32):
53 | paticle_list.append([pos_x, pos_y])
54 | if bound:
55 | wall_mark.append(0)
56 | else:
57 | wall_mark.append(1)
58 | particle_list = []
59 | wall_mark = []
60 |
61 | #### Dam break #######
62 | start_x = x_min + 0.5 * dx
63 | start_y = y_min - 1.0 * dx
64 | end_x = start_x + 0.5 * h
65 | end_y = start_y + 0.6 * h
66 |
67 | ## Constrcut wall
68 | # Bottom square
69 | b_start_x = 0.0
70 | b_start_y = 0.0
71 | b_end_x = b_start_x + w
72 | b_end_y = b_start_y + y_min - 2 * dx
73 |
74 | bottom_bound = b_end_y
75 |
76 | # Top square
77 | t_start_x = 0.0
78 | t_start_y = h - y_min + 2 * dx
79 | t_end_x = t_start_x + w
80 | t_end_y = h
81 |
82 | top_bound = t_start_y
83 |
84 | # left square
85 | l_start_x = 0.0
86 | l_start_y = y_min - 2 * dx
87 | l_end_x = l_start_x + x_min - 2 * dx
88 | l_end_y = l_start_y + h - 2 * y_min + 4 * dx
89 |
90 | left_bound = l_end_x
91 |
92 | # right square
93 | r_start_x = w - x_min + 2 * dx
94 | r_start_y = y_min - 2 * dx
95 | r_end_x = w - dx
96 | r_end_y = r_start_y + h - 2 * y_min + 4 * dx
97 |
98 | right_bound = r_start_x
99 |
100 | pos_list_fluid = [start_x, start_y, end_x, end_y]
101 | placeParticles(pos_list_fluid, particle_list, wall_mark)
102 |
103 | # These are boundaries
104 | pos_list_bs = [b_start_x, b_start_y, b_end_x, b_end_y]
105 | placeParticles(pos_list_bs, particle_list, wall_mark, bound=1)
106 |
107 | pos_list_ts = [t_start_x, t_start_y, t_end_x, t_end_y]
108 | placeParticles(pos_list_ts, particle_list, wall_mark, bound=1)
109 |
110 | pos_list_ls = [l_start_x, l_start_y, l_end_x, l_end_y]
111 | placeParticles(pos_list_ls, particle_list, wall_mark, bound=1)
112 |
113 | pos_list_rs = [r_start_x, r_start_y, r_end_x, r_end_y]
114 | placeParticles(pos_list_rs, particle_list, wall_mark, bound=1)
115 |
116 | return particle_list,wall_mark, top_bound, bottom_bound, left_bound, right_bound
117 |
118 | def makeGrid():
119 | grid_size = 2*dh
120 | num_x = np.ceil(w_bound / grid_size).astype(int)
121 | num_y = np.ceil(h_bound / grid_size).astype(int)
122 |
123 | grid_x = num_x
124 | grid_y = num_y
125 | return grid_x, grid_y
126 |
127 | @ti.data_oriented
128 | class sph_solver:
129 | def __init__(self, particle_list,
130 | wall_mark,
131 | grid,
132 | bound, alpha=0.5, dx=0.2, max_time=10000, max_steps=1000,gui=None):
133 | ######## Solver parameters ##########
134 | self.max_time = max_time
135 | self.max_steps = max_steps
136 | self.gui = gui
137 | # Gravity
138 | self.g = -9.80
139 | # viscosity
140 | self.alpha = alpha
141 | # reference density
142 | self.rho_0 = 1000.0
143 | # CFL coefficient
144 | self.CFL = 0.20
145 |
146 | # Smooth kernel norm factor
147 | self.kernel_norm = 1.0
148 |
149 | # Pressure state function parameters
150 | self.gamma = 7.0
151 | self.c_0 = 20.0
152 |
153 | ###### Scene parameters ########
154 | self.w = 20
155 | self.h = 10
156 | self.w_bound = 22
157 | self.h_bound = 12
158 |
159 | assert self.w_bound > self.w
160 | assert self.h_bound > self.h
161 | self.x_min = (self.w_bound - self.w) / 2.0
162 | self.y_min = (self.h_bound - self.h) / 2.0
163 | self.x_max = self.w_bound - (self.w_bound - self.w) / 2.0
164 | self.y_max = self.h_bound - (self.h_bound - self.h) / 2.0
165 |
166 | self.top_bound = bound[0] # top_bound
167 | self.bottom_bound = bound[1] #bottom_bound
168 | self.left_bound = bound[2] #left_bound
169 | self.right_bound = bound[3] #right_bound
170 |
171 | self.df_fac = 1.3
172 | self.dx = 0.2
173 | self.dh = self.dx * self.df_fac
174 | self.kernel_norm = 10. / (7. * np.pi * self.dh ** 2)
175 |
176 | ###### Particles #######
177 | self.dim = 2
178 | self.particle_numbers = len(particle_list)
179 |
180 | self.grid_x = grid[0]
181 | self.grid_y = grid[1]
182 |
183 | # Fluid particles
184 | self.old_positions = ti.Vector(self.dim, dt=ti.f32)
185 | self.particle_positions = ti.Vector(self.dim, dt=ti.f32)
186 | self.particle_velocity = ti.Vector(self.dim, dt=ti.f32)
187 | self.particle_pressure = ti.Vector(1, dt=ti.f32)
188 | self.particle_density = ti.Vector(1, dt=ti.f32)
189 | self.wall_mark_list = ti.Vector(1, dt=ti.f32)
190 |
191 | self.d_velocity = ti.Vector(self.dim, dt=ti.f32)
192 | self.d_density = ti.Vector(1, dt=ti.f32)
193 |
194 | self.dx = dx
195 | self.m = self.dx**2 * 1000.0
196 |
197 | self.particle_list = np.array(particle_list, dtype=np.float32)
198 | self.wall_mark = np.array(wall_mark, dtype=np.int32)
199 |
200 | self.grid_num_particles = ti.var(ti.i32)
201 | self.grid2particles = ti.var(ti.i32)
202 | self.particle_num_neighbors = ti.var(ti.i32)
203 | self.particle_neighbors = ti.var(ti.i32)
204 |
205 | self.max_num_particles_per_cell = 100
206 | self.max_num_neighbors = 100
207 |
208 | ti.root.dense(ti.i, self.particle_numbers).place(self.old_positions, self.particle_positions,
209 | self.particle_velocity, self.particle_pressure,
210 | self.particle_density, self.d_velocity, self.d_density,
211 | self.wall_mark_list)
212 |
213 | grid_snode = ti.root.dense(ti.ij, (self.grid_x, self.grid_y))
214 | grid_snode.place(self.grid_num_particles)
215 | grid_snode.dense(ti.k, self.max_num_particles_per_cell).place(self.grid2particles)
216 |
217 | nb_node = ti.root.dense(ti.i, self.particle_numbers)
218 | nb_node.place(self.particle_num_neighbors)
219 | nb_node.dense(ti.j, self.max_num_neighbors).place(self.particle_neighbors)
220 |
221 | @ti.kernel
222 | def init(self, p_list:ti.types.ndarray(), w_list:ti.types.ndarray()):
223 | for i in range(self.particle_numbers):
224 | for j in ti.static(range(self.dim)):
225 | self.particle_positions[i][j] = p_list[i,j]
226 | self.particle_velocity[i][j] = ti.cast(0.0, ti.f32)
227 | self.d_velocity[i][0] = ti.cast(0.0, ti.f32)
228 | self.d_velocity[i][1] = ti.cast(-9.8, ti.f32)
229 |
230 | self.wall_mark_list[i][0] = w_list[i]
231 | self.d_density[i][0] = ti.cast(0.0, ti.f32)
232 | self.particle_pressure[i][0] = ti.cast(0.0, ti.f32)
233 | self.particle_density[i][0] = ti.cast(1000.0, ti.f32)
234 |
235 | @ti.func
236 | def computeGridIndex(self, pos):
237 | return (pos / (2 * dh)).cast(int)
238 |
239 | @ti.kernel
240 | def allocateParticles(self):
241 | # Ref to pbf2d example from by Ye Kuang (k-ye)
242 | # https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py
243 | # allocate particles to grid
244 | for p_i in self.particle_positions:
245 | # Compute the grid index on the fly
246 | cell = self.computeGridIndex(self.particle_positions[p_i])
247 | offs = self.grid_num_particles[cell].atomic_add(1)
248 | self.grid2particles[cell, offs] = p_i
249 |
250 | @ti.func
251 | def is_in_grid(self, c):
252 | # Ref to pbf2d example from by Ye Kuang (k-ye)
253 | # https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py
254 | return 0 <= c[0] and c[0] < self.grid_x and 0 <= c[1] and c[1] < self.grid_y
255 |
256 | @ti.func
257 | def isFluid(self, p):
258 | # check fluid particle or bound particle
259 | return self.wall_mark_list[p][0]
260 |
261 | @ti.kernel
262 | def search_neighbors(self):
263 | # Ref to pbf2d example from by Ye Kuang (k-ye)
264 | # https://github.com/taichi-dev/taichi/blob/master/examples/pbf2d.py
265 | for p_i in self.particle_positions:
266 | pos_i = self.particle_positions[p_i]
267 | nb_i = 0
268 | if self.isFluid(p_i) == 1:
269 | # Compute the grid index on the fly
270 | cell = self.computeGridIndex(self.particle_positions[p_i])
271 | for offs in ti.static(ti.grouped(ti.ndrange((-1, 2), (-1, 2)))):
272 | cell_to_check = cell + offs
273 | if self.is_in_grid(cell_to_check):
274 | for j in range(self.grid_num_particles[cell_to_check]):
275 | p_j = self.grid2particles[cell_to_check, j]
276 | if nb_i < self.max_num_neighbors and p_j != p_i and (
277 | pos_i - self.particle_positions[p_j]).norm() < self.dh * 2.00:
278 | self.particle_neighbors[p_i, nb_i] = p_j
279 | nb_i += 1
280 | self.particle_num_neighbors[p_i] = nb_i
281 |
282 | @ti.func
283 | def cubicKernel(self, r, h):
284 | # value of cubic spline smoothing kernel
285 | k = 10. / (7. * np.pi * h ** 2)
286 | q = r / h
287 | # assert q >= 0.0
288 | res = ti.cast(0.0, ti.f32)
289 | if q <= 1.0:
290 | res = k * (1 - 1.5 * q ** 2 + 0.75 * q ** 3)
291 | elif q < 2.0:
292 | res = k * 0.25 * (2 - q) ** 3
293 | return res
294 |
295 | @ti.func
296 | def cubicKernelDerivative(self, r, h):
297 | # derivative of cubcic spline smoothing kernel
298 | k = 10. / (7. * np.pi * h ** 2)
299 | q = r / h
300 | # assert q > 0.0
301 | res = ti.cast(0.0, ti.f32)
302 | if q < 1.0:
303 | res = (k / h) * (-3 * q + 2.25 * q ** 2)
304 | elif q < 2.0:
305 | res = -0.75 * (k / h) * (2 - q) ** 2
306 | return res
307 |
308 | @ti.func
309 | def rhoDerivative(self, ptc_i, ptc_j, r, r_mod):
310 | # density delta
311 | return self.m * self.cubicKernelDerivative(r_mod, self.dh) \
312 | * (self.particle_velocity[ptc_i] - self.particle_velocity[ptc_j]).dot(r / r_mod)
313 |
314 |
315 | @ti.func
316 | def pUpdate(self, rho, rho_0=1000, gamma=7.0, c_0=20.0):
317 | # Weakly compressible, tait function
318 | b = rho_0 * c_0 ** 2 / gamma
319 | return b * ((rho / rho_0) ** gamma - 1.0)
320 |
321 | @ti.func
322 | def pressureForce(self, ptc_i, ptc_j, r, r_mod, mirror_pressure=0):
323 | # Compute the pressure force contribution, Symmetric Formula
324 | res = ti.Vector([0.0, 0.0])
325 | # Disable the mirror force, use collision instead
326 | # Use pressure mirror method to handle boundary leak
327 | # if mirror_pressure == 1:
328 | # res = - self.m * (self.particle_pressure[ptc_i][0]/ self.particle_density[ptc_i][0] ** 2
329 | # + self.particle_pressure[ptc_i][0]/self.rho_0**2)* self.cubicKernelDerivative(r_mod, h) * r / r_mod
330 | # else:
331 | res = -self.m * (self.particle_pressure[ptc_i][0] / self.particle_density[ptc_i][0] ** 2
332 | + self.particle_pressure[ptc_j][0] / self.particle_density[ptc_j][0] ** 2) \
333 | * self.cubicKernelDerivative(r_mod, self.dh) * r / r_mod
334 | return res
335 |
336 | @ti.func
337 | def viscosityForce(self, ptc_i, ptc_j, r, r_mod):
338 | # Compute the viscosity force contribution, artificial viscosity
339 | res = ti.Vector([0.0, 0.0])
340 | v_xy= (self.particle_velocity[ptc_i]- self.particle_velocity[ptc_j]).dot(r)
341 | if v_xy < 0:
342 | # Artifical viscosity
343 | vmu = -2.0 * self.alpha * self.dx * self.c_0 / (self.particle_density[ptc_i][0] + self.particle_density[ptc_j][0])
344 | res = -self.m * vmu * v_xy/(r_mod**2 + 0.01*self.dx**2)* self.cubicKernelDerivative(r_mod, self.dh) * r / r_mod
345 | return res
346 |
347 | @ti.func
348 | def simualteCollisions(self, ptc_i, vec, d):
349 | # Collision factor, assume roughly 50% velocity loss after collision, i.e. m_f /(m_f + m_b)
350 | c_f = 0.5
351 | self.particle_positions[ptc_i] += vec * d
352 | self.particle_velocity[ptc_i] -= (1.0+c_f) * self.particle_velocity[ptc_i].dot(vec) * vec
353 |
354 | @ti.kernel
355 | def enforceBoundary(self):
356 | for p_i in self.particle_positions:
357 | if self.isFluid(p_i) == 1:
358 | pos = self.particle_positions[p_i]
359 | if pos[0] < self.left_bound:
360 | self.simualteCollisions(p_i, ti.Vector([1.0, 0.0]), self.left_bound - pos[0])
361 | if pos[0] > self.right_bound:
362 | self.simualteCollisions(p_i, ti.Vector([-1.0, 0.0]), pos[0] - self.right_bound)
363 | if pos[1] > self.top_bound:
364 | self.simualteCollisions(p_i, ti.Vector([0.0, -1.0]), pos[1] - self.top_bound)
365 | if pos[1] < self.bottom_bound:
366 | self.simualteCollisions(p_i, ti.Vector([0.0, 1.0]), self.bottom_bound - pos[1])
367 |
368 | @ti.kernel
369 | def computeDeltas(self):
370 | for p_i in self.particle_positions:
371 | pos_i = self.particle_positions[p_i]
372 | d_v = ti.Vector([0.0, 0.0], dt=ti.f32)
373 | d_rho = ti.cast(0.0, ti.f32)
374 | # if self.isFluid(p_i) == 1:
375 | # d_v = ti.Vector([0.0, -9.8])
376 | for j in range(self.particle_num_neighbors[p_i]):
377 | p_j = self.particle_neighbors[p_i, j]
378 | pos_j = self.particle_positions[p_j]
379 |
380 | # Disable mirror force
381 | # mirror_pressure = 0
382 | # if self.isFluid(p_j) == 0:
383 | # mirror_pressure = 1
384 |
385 | # Compute distance and its mod
386 | r = pos_i - pos_j
387 | r_mod = r.norm()
388 |
389 | # Compute Density change
390 | d_rho += self.rhoDerivative(p_i, p_j, r, r_mod)
391 |
392 | if self.isFluid(p_i) == 1:
393 | # Compute Viscosity force contribution
394 | d_v += self.viscosityForce(p_i, p_j, r, r_mod)
395 |
396 | # Compute Pressure force contribution
397 | d_v += self.pressureForce(p_i, p_j, r, r_mod)
398 |
399 | # Add body force
400 | if self.isFluid(p_i) == 1:
401 | d_v += ti.Vector([0.0, self.g], dt=ti.f32)
402 | self.d_velocity[p_i] = d_v
403 | self.d_density[p_i][0] = d_rho
404 |
405 | @ti.kernel
406 | def updateTimeStep(self):
407 | # Simple Forward Euler currently
408 | for p_i in self.particle_positions:
409 | if self.isFluid(p_i) == 1:
410 | self.particle_positions[p_i] += self.dt * self.particle_velocity[p_i]
411 | self.particle_velocity[p_i] += self.dt * self.d_velocity[p_i]
412 | self.particle_density[p_i][0] += self.dt * self.d_density[p_i][0]
413 | self.particle_pressure[p_i][0] = self.pUpdate(self.particle_density[p_i][0], self.rho_0, self.gamma, self.c_0)
414 |
415 | def solve(self, output=False):
416 | # Compute dt, a naive initial test value
417 | self.dt = 0.1 * self.dh / self.c_0
418 | print("Time step: ", self.dt)
419 | print("Domain: (%s, %s, %s, %s)" % (self.x_min, self.x_max, self.y_min, self.y_max), )
420 | print("Fluid area: (%s, %s, %s, %s)"%(self.left_bound, self.right_bound, self.bottom_bound, self.top_bound))
421 | print("Grid: (%d, %d)"%(self.grid_x, self.grid_y))
422 |
423 | step = 1
424 | t = 0.0
425 | total_start = time.process_time()
426 | while t < self.max_time and step < self.max_steps:
427 | curr_start = time.process_time()
428 | self.solveUpdate()
429 | max_v = np.max(np.linalg.norm(self.particle_velocity.to_numpy(),2, axis=1))
430 | max_a = np.max(np.linalg.norm(self.d_velocity.to_numpy(),2, axis=1))
431 | max_rho = np.max(self.particle_density.to_numpy())
432 | max_pressure = np.max(self.particle_pressure.to_numpy())
433 |
434 | curr_end = time.process_time()
435 | t += self.dt
436 | step += 1
437 |
438 | # CFL analysis, adaptive dt
439 | dt_cfl = self.dh / max_v
440 | dt_f = np.sqrt(self.dh / max_a)
441 | dt_a = self.dh / (self.c_0 * np.sqrt((max_rho / self.rho_0)**self.gamma))
442 | self.dt = self.CFL * np.min([dt_cfl, dt_f, dt_a])
443 | if step % 10 == 0:
444 | print("Step: %d, physics time: %s, progress: %s %%, time used: %s, total time used: %s"
445 | % (step, t, 100*np.max([t / self.max_time, step / self.max_steps]), curr_end-curr_start, curr_end-total_start))
446 | print("Max velocity: %s, Max acceleration: %s, Max density: %s, Max pressure: %s" % (max_v, max_a, max_rho, max_pressure))
447 | print("Adaptive time step: ", self.dt)
448 | self.render(step, self.gui, output)
449 | total_end = time.process_time()
450 | print("Total time used: %s " % (total_end - total_start))
451 |
452 | def solveUpdate(self):
453 | self.grid_num_particles.fill(0)
454 | self.particle_neighbors.fill(-1)
455 | self.allocateParticles()
456 | self.search_neighbors()
457 | # Compute deltas
458 | self.computeDeltas()
459 | # timestep Update
460 | self.updateTimeStep()
461 | # Handle potential leak particles
462 | self.enforceBoundary()
463 |
464 | def isFluidNP(self, p):
465 | # ti.func cannot be called in python scope
466 | # for render use
467 | return self.wall_mark[p]
468 |
469 | def render(self, step, gui, output=False):
470 | canvas = gui.canvas
471 | canvas.clear(bg_color)
472 | pos_np = self.particle_positions.to_numpy()
473 | fluid_p = []
474 | wall_p = []
475 | for i, pos in enumerate(pos_np):
476 | if self.isFluidNP(i) == 1:
477 | fluid_p.append(pos)
478 | else:
479 | wall_p.append(pos)
480 | fluid_p = np.array(fluid_p)
481 | wall_p = np.array(wall_p)
482 |
483 | for pos in fluid_p:
484 | for j in range(self.dim):
485 | pos[j] *= screen_to_world_ratio / screen_res[j]
486 |
487 | for pos in wall_p:
488 | for j in range(self.dim):
489 | pos[j] *= screen_to_world_ratio / screen_res[j]
490 |
491 | gui.circles(fluid_p, radius=particle_radius, color=particle_color)
492 | gui.circles(wall_p, radius=particle_radius, color=boundary_color)
493 | if output:
494 | if step%10 == 0:
495 | gui.show(f"{step:04d}.png")
496 | else:
497 | gui.show()
498 |
499 | def main():
500 | OUTPUT = False
501 | gui = ti.GUI('SPH2D', screen_res)
502 | grid_shape = makeGrid()
503 | particle_list, wall_mark, u, b, l, r = setup()
504 | sph = sph_solver(particle_list, wall_mark, grid_shape, [u,b,l,r],alpha=1.0, dx = dx, gui=gui, max_steps=10000)
505 | sph.init(sph.particle_list, sph.wall_mark)
506 | sph.solve(output=OUTPUT)
507 |
508 | print('done')
509 |
510 | if __name__ == '__main__':
511 | main()
512 |
513 |
514 |
515 |
516 |
517 |
518 |
519 |
520 |
521 |
522 |
523 |
--------------------------------------------------------------------------------
/particle_system.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 | import numpy as np
3 | import trimesh as tm
4 | from functools import reduce
5 | from config_builder import SimConfig
6 | from WCSPH import WCSPHSolver
7 | from DFSPH import DFSPHSolver
8 | from scan_single_buffer import parallel_prefix_sum_inclusive_inplace
9 |
10 | @ti.data_oriented
11 | class ParticleSystem:
12 | def __init__(self, config: SimConfig, GGUI=False):
13 | self.cfg = config
14 | self.GGUI = GGUI
15 |
16 | self.domain_start = np.array([0.0, 0.0, 0.0])
17 | self.domain_start = np.array(self.cfg.get_cfg("domainStart"))
18 |
19 | self.domain_end = np.array([1.0, 1.0, 1.0])
20 | self.domian_end = np.array(self.cfg.get_cfg("domainEnd"))
21 |
22 | self.domain_size = self.domian_end - self.domain_start
23 |
24 | self.dim = len(self.domain_size)
25 | assert self.dim > 1
26 | # Simulation method
27 | self.simulation_method = self.cfg.get_cfg("simulationMethod")
28 |
29 | # Material
30 | self.material_solid = 0
31 | self.material_fluid = 1
32 |
33 | self.particle_radius = 0.01 # particle radius
34 | self.particle_radius = self.cfg.get_cfg("particleRadius")
35 |
36 | self.particle_diameter = 2 * self.particle_radius
37 | self.support_radius = self.particle_radius * 4.0 # support radius
38 | self.m_V0 = 0.8 * self.particle_diameter ** self.dim
39 |
40 | self.particle_num = ti.field(int, shape=())
41 |
42 | # Grid related properties
43 | self.grid_size = self.support_radius
44 | self.grid_num = np.ceil(self.domain_size / self.grid_size).astype(int)
45 | print("grid size: ", self.grid_num)
46 | self.padding = self.grid_size
47 |
48 | # All objects id and its particle num
49 | self.object_collection = dict()
50 | self.object_id_rigid_body = set()
51 |
52 | #========== Compute number of particles ==========#
53 | #### Process Fluid Blocks ####
54 | fluid_blocks = self.cfg.get_fluid_blocks()
55 | fluid_particle_num = 0
56 | for fluid in fluid_blocks:
57 | particle_num = self.compute_cube_particle_num(fluid["start"], fluid["end"])
58 | fluid["particleNum"] = particle_num
59 | self.object_collection[fluid["objectId"]] = fluid
60 | fluid_particle_num += particle_num
61 |
62 | #### Process Rigid Blocks ####
63 | rigid_blocks = self.cfg.get_rigid_blocks()
64 | rigid_particle_num = 0
65 | for rigid in rigid_blocks:
66 | particle_num = self.compute_cube_particle_num(rigid["start"], rigid["end"])
67 | rigid["particleNum"] = particle_num
68 | self.object_collection[rigid["objectId"]] = rigid
69 | rigid_particle_num += particle_num
70 |
71 | #### Process Rigid Bodies ####
72 | rigid_bodies = self.cfg.get_rigid_bodies()
73 | for rigid_body in rigid_bodies:
74 | voxelized_points_np = self.load_rigid_body(rigid_body)
75 | rigid_body["particleNum"] = voxelized_points_np.shape[0]
76 | rigid_body["voxelizedPoints"] = voxelized_points_np
77 | self.object_collection[rigid_body["objectId"]] = rigid_body
78 | rigid_particle_num += voxelized_points_np.shape[0]
79 |
80 | self.fluid_particle_num = fluid_particle_num
81 | self.solid_particle_num = rigid_particle_num
82 | self.particle_max_num = fluid_particle_num + rigid_particle_num
83 | self.num_rigid_bodies = len(rigid_blocks)+len(rigid_bodies)
84 |
85 | #### TODO: Handle the Particle Emitter ####
86 | # self.particle_max_num += emitted particles
87 | print(f"Current particle num: {self.particle_num[None]}, Particle max num: {self.particle_max_num}")
88 |
89 | #========== Allocate memory ==========#
90 | # Rigid body properties
91 | if self.num_rigid_bodies > 0:
92 | # TODO: Here we actually only need to store rigid boides, however the object id of rigid may not start from 0, so allocate center of mass for all objects
93 | self.rigid_rest_cm = ti.Vector.field(self.dim, dtype=float, shape=self.num_rigid_bodies + len(fluid_blocks))
94 |
95 | # Particle num of each grid
96 | self.grid_particles_num = ti.field(int, shape=int(self.grid_num[0]*self.grid_num[1]*self.grid_num[2]))
97 | self.grid_particles_num_temp = ti.field(int, shape=int(self.grid_num[0]*self.grid_num[1]*self.grid_num[2]))
98 |
99 | self.prefix_sum_executor = ti.algorithms.PrefixSumExecutor(self.grid_particles_num.shape[0])
100 |
101 | # Particle related properties
102 | self.object_id = ti.field(dtype=int, shape=self.particle_max_num)
103 | self.x = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
104 | self.x_0 = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
105 | self.v = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
106 | self.acceleration = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
107 | self.m_V = ti.field(dtype=float, shape=self.particle_max_num)
108 | self.m = ti.field(dtype=float, shape=self.particle_max_num)
109 | self.density = ti.field(dtype=float, shape=self.particle_max_num)
110 | self.pressure = ti.field(dtype=float, shape=self.particle_max_num)
111 | self.material = ti.field(dtype=int, shape=self.particle_max_num)
112 | self.color = ti.Vector.field(3, dtype=int, shape=self.particle_max_num)
113 | self.is_dynamic = ti.field(dtype=int, shape=self.particle_max_num)
114 |
115 | if self.cfg.get_cfg("simulationMethod") == 4:
116 | self.dfsph_factor = ti.field(dtype=float, shape=self.particle_max_num)
117 | self.density_adv = ti.field(dtype=float, shape=self.particle_max_num)
118 |
119 | # Buffer for sort
120 | self.object_id_buffer = ti.field(dtype=int, shape=self.particle_max_num)
121 | self.x_buffer = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
122 | self.x_0_buffer = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
123 | self.v_buffer = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
124 | self.acceleration_buffer = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
125 | self.m_V_buffer = ti.field(dtype=float, shape=self.particle_max_num)
126 | self.m_buffer = ti.field(dtype=float, shape=self.particle_max_num)
127 | self.density_buffer = ti.field(dtype=float, shape=self.particle_max_num)
128 | self.pressure_buffer = ti.field(dtype=float, shape=self.particle_max_num)
129 | self.material_buffer = ti.field(dtype=int, shape=self.particle_max_num)
130 | self.color_buffer = ti.Vector.field(3, dtype=int, shape=self.particle_max_num)
131 | self.is_dynamic_buffer = ti.field(dtype=int, shape=self.particle_max_num)
132 |
133 | if self.cfg.get_cfg("simulationMethod") == 4:
134 | self.dfsph_factor_buffer = ti.field(dtype=float, shape=self.particle_max_num)
135 | self.density_adv_buffer = ti.field(dtype=float, shape=self.particle_max_num)
136 |
137 | # Grid id for each particle
138 | self.grid_ids = ti.field(int, shape=self.particle_max_num)
139 | self.grid_ids_buffer = ti.field(int, shape=self.particle_max_num)
140 | self.grid_ids_new = ti.field(int, shape=self.particle_max_num)
141 |
142 | self.x_vis_buffer = None
143 | if self.GGUI:
144 | self.x_vis_buffer = ti.Vector.field(self.dim, dtype=float, shape=self.particle_max_num)
145 | self.color_vis_buffer = ti.Vector.field(3, dtype=float, shape=self.particle_max_num)
146 |
147 |
148 | #========== Initialize particles ==========#
149 |
150 | # Fluid block
151 | for fluid in fluid_blocks:
152 | obj_id = fluid["objectId"]
153 | offset = np.array(fluid["translation"])
154 | start = np.array(fluid["start"]) + offset
155 | end = np.array(fluid["end"]) + offset
156 | scale = np.array(fluid["scale"])
157 | velocity = fluid["velocity"]
158 | density = fluid["density"]
159 | color = fluid["color"]
160 | self.add_cube(object_id=obj_id,
161 | lower_corner=start,
162 | cube_size=(end-start)*scale,
163 | velocity=velocity,
164 | density=density,
165 | is_dynamic=1, # enforce fluid dynamic
166 | color=color,
167 | material=1) # 1 indicates fluid
168 |
169 | # TODO: Handle rigid block
170 | # Rigid block
171 | for rigid in rigid_blocks:
172 | obj_id = rigid["objectId"]
173 | offset = np.array(rigid["translation"])
174 | start = np.array(rigid["start"]) + offset
175 | end = np.array(rigid["end"]) + offset
176 | scale = np.array(rigid["scale"])
177 | velocity = rigid["velocity"]
178 | density = rigid["density"]
179 | color = rigid["color"]
180 | is_dynamic = rigid["isDynamic"]
181 | self.add_cube(object_id=obj_id,
182 | lower_corner=start,
183 | cube_size=(end-start)*scale,
184 | velocity=velocity,
185 | density=density,
186 | is_dynamic=is_dynamic,
187 | color=color,
188 | material=0) # 1 indicates solid
189 |
190 | # Rigid bodies
191 | for rigid_body in rigid_bodies:
192 | obj_id = rigid_body["objectId"]
193 | self.object_id_rigid_body.add(obj_id)
194 | num_particles_obj = rigid_body["particleNum"]
195 | voxelized_points_np = rigid_body["voxelizedPoints"]
196 | is_dynamic = rigid_body["isDynamic"]
197 | if is_dynamic:
198 | velocity = np.array(rigid_body["velocity"], dtype=np.float32)
199 | else:
200 | velocity = np.array([0.0 for _ in range(self.dim)], dtype=np.float32)
201 | density = rigid_body["density"]
202 | color = np.array(rigid_body["color"], dtype=np.int32)
203 | self.add_particles(obj_id,
204 | num_particles_obj,
205 | np.array(voxelized_points_np, dtype=np.float32), # position
206 | np.stack([velocity for _ in range(num_particles_obj)]), # velocity
207 | density * np.ones(num_particles_obj, dtype=np.float32), # density
208 | np.zeros(num_particles_obj, dtype=np.float32), # pressure
209 | np.array([0 for _ in range(num_particles_obj)], dtype=np.int32), # material is solid
210 | is_dynamic * np.ones(num_particles_obj, dtype=np.int32), # is_dynamic
211 | np.stack([color for _ in range(num_particles_obj)])) # color
212 |
213 |
214 | def build_solver(self):
215 | solver_type = self.cfg.get_cfg("simulationMethod")
216 | if solver_type == 0:
217 | return WCSPHSolver(self)
218 | elif solver_type == 4:
219 | return DFSPHSolver(self)
220 | else:
221 | raise NotImplementedError(f"Solver type {solver_type} has not been implemented.")
222 |
223 | @ti.func
224 | def add_particle(self, p, obj_id, x, v, density, pressure, material, is_dynamic, color):
225 | self.object_id[p] = obj_id
226 | self.x[p] = x
227 | self.x_0[p] = x
228 | self.v[p] = v
229 | self.density[p] = density
230 | self.m_V[p] = self.m_V0
231 | self.m[p] = self.m_V0 * density
232 | self.pressure[p] = pressure
233 | self.material[p] = material
234 | self.is_dynamic[p] = is_dynamic
235 | self.color[p] = color
236 |
237 | def add_particles(self,
238 | object_id: int,
239 | new_particles_num: int,
240 | new_particles_positions: ti.types.ndarray(),
241 | new_particles_velocity: ti.types.ndarray(),
242 | new_particle_density: ti.types.ndarray(),
243 | new_particle_pressure: ti.types.ndarray(),
244 | new_particles_material: ti.types.ndarray(),
245 | new_particles_is_dynamic: ti.types.ndarray(),
246 | new_particles_color: ti.types.ndarray()
247 | ):
248 |
249 | self._add_particles(object_id,
250 | new_particles_num,
251 | new_particles_positions,
252 | new_particles_velocity,
253 | new_particle_density,
254 | new_particle_pressure,
255 | new_particles_material,
256 | new_particles_is_dynamic,
257 | new_particles_color
258 | )
259 |
260 | @ti.kernel
261 | def _add_particles(self,
262 | object_id: int,
263 | new_particles_num: int,
264 | new_particles_positions: ti.types.ndarray(),
265 | new_particles_velocity: ti.types.ndarray(),
266 | new_particle_density: ti.types.ndarray(),
267 | new_particle_pressure: ti.types.ndarray(),
268 | new_particles_material: ti.types.ndarray(),
269 | new_particles_is_dynamic: ti.types.ndarray(),
270 | new_particles_color: ti.types.ndarray()):
271 | for p in range(self.particle_num[None], self.particle_num[None] + new_particles_num):
272 | v = ti.Vector.zero(float, self.dim)
273 | x = ti.Vector.zero(float, self.dim)
274 | for d in ti.static(range(self.dim)):
275 | v[d] = new_particles_velocity[p - self.particle_num[None], d]
276 | x[d] = new_particles_positions[p - self.particle_num[None], d]
277 | self.add_particle(p, object_id, x, v,
278 | new_particle_density[p - self.particle_num[None]],
279 | new_particle_pressure[p - self.particle_num[None]],
280 | new_particles_material[p - self.particle_num[None]],
281 | new_particles_is_dynamic[p - self.particle_num[None]],
282 | ti.Vector([new_particles_color[p - self.particle_num[None], i] for i in range(3)])
283 | )
284 | self.particle_num[None] += new_particles_num
285 |
286 |
287 | @ti.func
288 | def pos_to_index(self, pos):
289 | return (pos / self.grid_size).cast(int)
290 |
291 |
292 | @ti.func
293 | def flatten_grid_index(self, grid_index):
294 | return grid_index[0] * self.grid_num[1] * self.grid_num[2] + grid_index[1] * self.grid_num[2] + grid_index[2]
295 |
296 | @ti.func
297 | def get_flatten_grid_index(self, pos):
298 | return self.flatten_grid_index(self.pos_to_index(pos))
299 |
300 |
301 | @ti.func
302 | def is_static_rigid_body(self, p):
303 | return self.material[p] == self.material_solid and (not self.is_dynamic[p])
304 |
305 |
306 | @ti.func
307 | def is_dynamic_rigid_body(self, p):
308 | return self.material[p] == self.material_solid and self.is_dynamic[p]
309 |
310 |
311 | @ti.kernel
312 | def update_grid_id(self):
313 | for I in ti.grouped(self.grid_particles_num):
314 | self.grid_particles_num[I] = 0
315 | for I in ti.grouped(self.x):
316 | grid_index = self.get_flatten_grid_index(self.x[I])
317 | self.grid_ids[I] = grid_index
318 | ti.atomic_add(self.grid_particles_num[grid_index], 1)
319 | for I in ti.grouped(self.grid_particles_num):
320 | self.grid_particles_num_temp[I] = self.grid_particles_num[I]
321 |
322 | @ti.kernel
323 | def counting_sort(self):
324 | # FIXME: make it the actual particle num
325 | for i in range(self.particle_max_num):
326 | I = self.particle_max_num - 1 - i
327 | base_offset = 0
328 | if self.grid_ids[I] - 1 >= 0:
329 | base_offset = self.grid_particles_num[self.grid_ids[I]-1]
330 | self.grid_ids_new[I] = ti.atomic_sub(self.grid_particles_num_temp[self.grid_ids[I]], 1) - 1 + base_offset
331 |
332 | for I in ti.grouped(self.grid_ids):
333 | new_index = self.grid_ids_new[I]
334 | self.grid_ids_buffer[new_index] = self.grid_ids[I]
335 | self.object_id_buffer[new_index] = self.object_id[I]
336 | self.x_0_buffer[new_index] = self.x_0[I]
337 | self.x_buffer[new_index] = self.x[I]
338 | self.v_buffer[new_index] = self.v[I]
339 | self.acceleration_buffer[new_index] = self.acceleration[I]
340 | self.m_V_buffer[new_index] = self.m_V[I]
341 | self.m_buffer[new_index] = self.m[I]
342 | self.density_buffer[new_index] = self.density[I]
343 | self.pressure_buffer[new_index] = self.pressure[I]
344 | self.material_buffer[new_index] = self.material[I]
345 | self.color_buffer[new_index] = self.color[I]
346 | self.is_dynamic_buffer[new_index] = self.is_dynamic[I]
347 |
348 | if ti.static(self.simulation_method == 4):
349 | self.dfsph_factor_buffer[new_index] = self.dfsph_factor[I]
350 | self.density_adv_buffer[new_index] = self.density_adv[I]
351 |
352 | for I in ti.grouped(self.x):
353 | self.grid_ids[I] = self.grid_ids_buffer[I]
354 | self.object_id[I] = self.object_id_buffer[I]
355 | self.x_0[I] = self.x_0_buffer[I]
356 | self.x[I] = self.x_buffer[I]
357 | self.v[I] = self.v_buffer[I]
358 | self.acceleration[I] = self.acceleration_buffer[I]
359 | self.m_V[I] = self.m_V_buffer[I]
360 | self.m[I] = self.m_buffer[I]
361 | self.density[I] = self.density_buffer[I]
362 | self.pressure[I] = self.pressure_buffer[I]
363 | self.material[I] = self.material_buffer[I]
364 | self.color[I] = self.color_buffer[I]
365 | self.is_dynamic[I] = self.is_dynamic_buffer[I]
366 |
367 | if ti.static(self.simulation_method == 4):
368 | self.dfsph_factor[I] = self.dfsph_factor_buffer[I]
369 | self.density_adv[I] = self.density_adv_buffer[I]
370 |
371 |
372 | def initialize_particle_system(self):
373 | self.update_grid_id()
374 | self.prefix_sum_executor.run(self.grid_particles_num)
375 | self.counting_sort()
376 |
377 |
378 | @ti.func
379 | def for_all_neighbors(self, p_i, task: ti.template(), ret: ti.template()):
380 | center_cell = self.pos_to_index(self.x[p_i])
381 | for offset in ti.grouped(ti.ndrange(*((-1, 2),) * self.dim)):
382 | grid_index = self.flatten_grid_index(center_cell + offset)
383 | for p_j in range(self.grid_particles_num[ti.max(0, grid_index-1)], self.grid_particles_num[grid_index]):
384 | if p_i[0] != p_j and (self.x[p_i] - self.x[p_j]).norm() < self.support_radius:
385 | task(p_i, p_j, ret)
386 |
387 | @ti.kernel
388 | def copy_to_numpy(self, np_arr: ti.types.ndarray(), src_arr: ti.template()):
389 | for i in range(self.particle_num[None]):
390 | np_arr[i] = src_arr[i]
391 |
392 | def copy_to_vis_buffer(self, invisible_objects=[]):
393 | if len(invisible_objects) != 0:
394 | self.x_vis_buffer.fill(0.0)
395 | self.color_vis_buffer.fill(0.0)
396 | for obj_id in self.object_collection:
397 | if obj_id not in invisible_objects:
398 | self._copy_to_vis_buffer(obj_id)
399 |
400 | @ti.kernel
401 | def _copy_to_vis_buffer(self, obj_id: int):
402 | assert self.GGUI
403 | # FIXME: make it equal to actual particle num
404 | for i in range(self.particle_max_num):
405 | if self.object_id[i] == obj_id:
406 | self.x_vis_buffer[i] = self.x[i]
407 | self.color_vis_buffer[i] = self.color[i] / 255.0
408 |
409 | def dump(self, obj_id):
410 | np_object_id = self.object_id.to_numpy()
411 | mask = (np_object_id == obj_id).nonzero()
412 | np_x = self.x.to_numpy()[mask]
413 | np_v = self.v.to_numpy()[mask]
414 |
415 | return {
416 | 'position': np_x,
417 | 'velocity': np_v
418 | }
419 |
420 |
421 | def load_rigid_body(self, rigid_body):
422 | obj_id = rigid_body["objectId"]
423 | mesh = tm.load(rigid_body["geometryFile"])
424 | mesh.apply_scale(rigid_body["scale"])
425 | offset = np.array(rigid_body["translation"])
426 |
427 | angle = rigid_body["rotationAngle"] / 360 * 2 * 3.1415926
428 | direction = rigid_body["rotationAxis"]
429 | rot_matrix = tm.transformations.rotation_matrix(angle, direction, mesh.vertices.mean(axis=0))
430 | mesh.apply_transform(rot_matrix)
431 | mesh.vertices += offset
432 |
433 | # Backup the original mesh for exporting obj
434 | mesh_backup = mesh.copy()
435 | rigid_body["mesh"] = mesh_backup
436 | rigid_body["restPosition"] = mesh_backup.vertices
437 | rigid_body["restCenterOfMass"] = mesh_backup.vertices.mean(axis=0)
438 | is_success = tm.repair.fill_holes(mesh)
439 | # print("Is the mesh successfully repaired? ", is_success)
440 | voxelized_mesh = mesh.voxelized(pitch=self.particle_diameter)
441 | voxelized_mesh = mesh.voxelized(pitch=self.particle_diameter).fill()
442 | # voxelized_mesh = mesh.voxelized(pitch=self.particle_diameter).hollow()
443 | # voxelized_mesh.show()
444 | voxelized_points_np = voxelized_mesh.points
445 | print(f"rigid body {obj_id} num: {voxelized_points_np.shape[0]}")
446 |
447 | return voxelized_points_np
448 |
449 |
450 | def compute_cube_particle_num(self, start, end):
451 | num_dim = []
452 | for i in range(self.dim):
453 | num_dim.append(
454 | np.arange(start[i], end[i], self.particle_diameter))
455 | return reduce(lambda x, y: x * y,
456 | [len(n) for n in num_dim])
457 |
458 | def add_cube(self,
459 | object_id,
460 | lower_corner,
461 | cube_size,
462 | material,
463 | is_dynamic,
464 | color=(0,0,0),
465 | density=None,
466 | pressure=None,
467 | velocity=None):
468 |
469 | num_dim = []
470 | for i in range(self.dim):
471 | num_dim.append(
472 | np.arange(lower_corner[i], lower_corner[i] + cube_size[i],
473 | self.particle_diameter))
474 | num_new_particles = reduce(lambda x, y: x * y,
475 | [len(n) for n in num_dim])
476 | print('particle num ', num_new_particles)
477 |
478 | new_positions = np.array(np.meshgrid(*num_dim,
479 | sparse=False,
480 | indexing='ij'),
481 | dtype=np.float32)
482 | new_positions = new_positions.reshape(-1,
483 | reduce(lambda x, y: x * y, list(new_positions.shape[1:]))).transpose()
484 | print("new position shape ", new_positions.shape)
485 | if velocity is None:
486 | velocity_arr = np.full_like(new_positions, 0, dtype=np.float32)
487 | else:
488 | velocity_arr = np.array([velocity for _ in range(num_new_particles)], dtype=np.float32)
489 |
490 | material_arr = np.full_like(np.zeros(num_new_particles, dtype=np.int32), material)
491 | is_dynamic_arr = np.full_like(np.zeros(num_new_particles, dtype=np.int32), is_dynamic)
492 | color_arr = np.stack([np.full_like(np.zeros(num_new_particles, dtype=np.int32), c) for c in color], axis=1)
493 | density_arr = np.full_like(np.zeros(num_new_particles, dtype=np.float32), density if density is not None else 1000.)
494 | pressure_arr = np.full_like(np.zeros(num_new_particles, dtype=np.float32), pressure if pressure is not None else 0.)
495 | self.add_particles(object_id, num_new_particles, new_positions, velocity_arr, density_arr, pressure_arr, material_arr, is_dynamic_arr, color_arr)
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | taichi>=1.2.0
2 | trimesh
--------------------------------------------------------------------------------
/run_simulation.py:
--------------------------------------------------------------------------------
1 | import os
2 | import argparse
3 | import taichi as ti
4 | import numpy as np
5 | from config_builder import SimConfig
6 | from particle_system import ParticleSystem
7 |
8 | ti.init(arch=ti.gpu, device_memory_fraction=0.5)
9 |
10 |
11 | if __name__ == "__main__":
12 | parser = argparse.ArgumentParser(description='SPH Taichi')
13 | parser.add_argument('--scene_file',
14 | default='',
15 | help='scene file')
16 | args = parser.parse_args()
17 | scene_path = args.scene_file
18 | config = SimConfig(scene_file_path=scene_path)
19 | scene_name = scene_path.split("/")[-1].split(".")[0]
20 |
21 | substeps = config.get_cfg("numberOfStepsPerRenderUpdate")
22 | output_frames = config.get_cfg("exportFrame")
23 | output_interval = int(0.016 / config.get_cfg("timeStepSize"))
24 | output_ply = config.get_cfg("exportPly")
25 | output_obj = config.get_cfg("exportObj")
26 | series_prefix = "{}_output/particle_object_{}.ply".format(scene_name, "{}")
27 | if output_frames:
28 | os.makedirs(f"{scene_name}_output_img", exist_ok=True)
29 | if output_ply:
30 | os.makedirs(f"{scene_name}_output", exist_ok=True)
31 |
32 |
33 | ps = ParticleSystem(config, GGUI=True)
34 | solver = ps.build_solver()
35 | solver.initialize()
36 |
37 | window = ti.ui.Window('SPH', (1024, 1024), show_window = True, vsync=False)
38 |
39 | scene = ti.ui.Scene()
40 | camera = ti.ui.Camera()
41 | camera.position(5.5, 2.5, 4.0)
42 | camera.up(0.0, 1.0, 0.0)
43 | camera.lookat(-1.0, 0.0, 0.0)
44 | camera.fov(70)
45 | scene.set_camera(camera)
46 |
47 | canvas = window.get_canvas()
48 | radius = 0.002
49 | movement_speed = 0.02
50 | background_color = (0, 0, 0) # 0xFFFFFF
51 | particle_color = (1, 1, 1)
52 |
53 | # Invisible objects
54 | invisible_objects = config.get_cfg("invisibleObjects")
55 | if not invisible_objects:
56 | invisible_objects = []
57 |
58 | # Draw the lines for domain
59 | x_max, y_max, z_max = config.get_cfg("domainEnd")
60 | box_anchors = ti.Vector.field(3, dtype=ti.f32, shape = 8)
61 | box_anchors[0] = ti.Vector([0.0, 0.0, 0.0])
62 | box_anchors[1] = ti.Vector([0.0, y_max, 0.0])
63 | box_anchors[2] = ti.Vector([x_max, 0.0, 0.0])
64 | box_anchors[3] = ti.Vector([x_max, y_max, 0.0])
65 |
66 | box_anchors[4] = ti.Vector([0.0, 0.0, z_max])
67 | box_anchors[5] = ti.Vector([0.0, y_max, z_max])
68 | box_anchors[6] = ti.Vector([x_max, 0.0, z_max])
69 | box_anchors[7] = ti.Vector([x_max, y_max, z_max])
70 |
71 | box_lines_indices = ti.field(int, shape=(2 * 12))
72 |
73 | for i, val in enumerate([0, 1, 0, 2, 1, 3, 2, 3, 4, 5, 4, 6, 5, 7, 6, 7, 0, 4, 1, 5, 2, 6, 3, 7]):
74 | box_lines_indices[i] = val
75 |
76 | cnt = 0
77 | cnt_ply = 0
78 |
79 | while window.running:
80 | for i in range(substeps):
81 | solver.step()
82 | ps.copy_to_vis_buffer(invisible_objects=invisible_objects)
83 | if ps.dim == 2:
84 | canvas.set_background_color(background_color)
85 | canvas.circles(ps.x_vis_buffer, radius=ps.particle_radius, color=particle_color)
86 | elif ps.dim == 3:
87 | camera.track_user_inputs(window, movement_speed=movement_speed, hold_key=ti.ui.LMB)
88 | scene.set_camera(camera)
89 |
90 | scene.point_light((2.0, 2.0, 2.0), color=(1.0, 1.0, 1.0))
91 | scene.particles(ps.x_vis_buffer, radius=ps.particle_radius, per_vertex_color=ps.color_vis_buffer)
92 |
93 | scene.lines(box_anchors, indices=box_lines_indices, color = (0.99, 0.68, 0.28), width = 1.0)
94 | canvas.scene(scene)
95 |
96 | if output_frames:
97 | if cnt % output_interval == 0:
98 | window.write_image(f"{scene_name}_output_img/{cnt:06}.png")
99 |
100 | if cnt % output_interval == 0:
101 | if output_ply:
102 | obj_id = 0
103 | obj_data = ps.dump(obj_id=obj_id)
104 | np_pos = obj_data["position"]
105 | writer = ti.tools.PLYWriter(num_vertices=ps.object_collection[obj_id]["particleNum"])
106 | writer.add_vertex_pos(np_pos[:, 0], np_pos[:, 1], np_pos[:, 2])
107 | writer.export_frame_ascii(cnt_ply, series_prefix.format(0))
108 | if output_obj:
109 | for r_body_id in ps.object_id_rigid_body:
110 | with open(f"{scene_name}_output/obj_{r_body_id}_{cnt_ply:06}.obj", "w") as f:
111 | e = ps.object_collection[r_body_id]["mesh"].export(file_type='obj')
112 | f.write(e)
113 | cnt_ply += 1
114 |
115 | cnt += 1
116 | # if cnt > 6000:
117 | # break
118 | window.show()
119 |
--------------------------------------------------------------------------------
/scan_single_buffer.py:
--------------------------------------------------------------------------------
1 | import taichi as ti
2 |
3 |
4 | @ti.func
5 | def warp_inclusive_add_cuda(val: ti.template()):
6 | global_tid = ti.global_thread_idx()
7 | lane_id = global_tid % 32
8 | # Intra-warp scan, manually unroll
9 | offset_j = 1
10 | n = ti.simt.warp.shfl_up_i32(ti.simt.warp.active_mask(), val, offset_j)
11 | if (lane_id >= offset_j):
12 | val += n
13 | offset_j = 2
14 | n = ti.simt.warp.shfl_up_i32(ti.simt.warp.active_mask(), val, offset_j)
15 | if (lane_id >= offset_j):
16 | val += n
17 | offset_j = 4
18 | n = ti.simt.warp.shfl_up_i32(ti.simt.warp.active_mask(), val, offset_j)
19 | if (lane_id >= offset_j):
20 | val += n
21 | offset_j = 8
22 | n = ti.simt.warp.shfl_up_i32(ti.simt.warp.active_mask(), val, offset_j)
23 | if (lane_id >= offset_j):
24 | val += n
25 | offset_j = 16
26 | n = ti.simt.warp.shfl_up_i32(ti.simt.warp.active_mask(), val, offset_j)
27 | if (lane_id >= offset_j):
28 | val += n
29 | return val
30 |
31 |
32 | # target = ti.cfg.cuda
33 | target = ti.cuda
34 | if target == ti.cuda:
35 | inclusive_add = warp_inclusive_add_cuda
36 | barrier = ti.simt.block.sync
37 | elif target == ti.vulkan:
38 | inclusive_add = ti.simt.subgroup.inclusive_add
39 | barrier = ti.simt.subgroup.barrier
40 | else:
41 | raise RuntimeError(f"Arch {target} not supported for parallel scan.")
42 |
43 |
44 | @ti.kernel
45 | def shfl_scan(arr_in: ti.template(), in_beg: ti.i32, in_end: ti.i32,
46 | sum_smem: ti.template(), single_block: ti.template()):
47 | ti.loop_config(block_dim=BLOCK_SZ)
48 | for i in range(in_beg, in_end):
49 | val = arr_in[i]
50 |
51 | thread_id = i % BLOCK_SZ
52 | block_id = int((i - in_beg) // BLOCK_SZ)
53 | lane_id = thread_id % WARP_SZ
54 | warp_id = thread_id // WARP_SZ
55 |
56 | val = inclusive_add(val)
57 | barrier()
58 |
59 | # Put warp scan results to smem
60 | if (thread_id % WARP_SZ == WARP_SZ - 1):
61 | sum_smem[block_id, warp_id] = val
62 | barrier()
63 |
64 | # Inter-warp scan, use the first thread in the first warp
65 | if (warp_id == 0 and lane_id == 0):
66 | for k in range(1, BLOCK_SZ / WARP_SZ):
67 | sum_smem[block_id, k] += sum_smem[block_id, k - 1]
68 | barrier()
69 |
70 | # Update data with warp sums
71 | warp_sum = 0
72 | if (warp_id > 0):
73 | warp_sum = sum_smem[block_id, warp_id - 1]
74 | val += warp_sum
75 | arr_in[i] = val
76 |
77 | # Update partial sums
78 | if not single_block and (thread_id == BLOCK_SZ - 1):
79 | arr_in[in_end + block_id] = val
80 |
81 |
82 | @ti.kernel
83 | def uniform_add(arr_in: ti.template(), in_beg: ti.i32, in_end: ti.i32):
84 | ti.loop_config(block_dim=BLOCK_SZ)
85 | for i in range(in_beg + BLOCK_SZ, in_end):
86 | block_id = int((i - in_beg) // BLOCK_SZ)
87 | arr_in[i] += arr_in[in_end + block_id - 1]
88 |
89 |
90 | WARP_SZ = 32
91 | BLOCK_SZ = 128
92 |
93 | @ti.kernel
94 | def blit_from_field_to_field(
95 | dst: ti.template(), src: ti.template(), offset: ti.i32, size: ti.i32
96 | ):
97 | for i in range(size):
98 | dst[i + offset] = src[i]
99 |
100 |
101 | smem = None
102 | arrs = [0]
103 | ele_nums = [0]
104 | ele_nums_pos = [0]
105 | large_arr = None
106 |
107 |
108 | def parallel_prefix_sum_inclusive_inplace(input_arr, length):
109 |
110 | global smem
111 | global arrs
112 | global ele_nums
113 | global large_arr
114 |
115 | GRID_SZ = int((length + BLOCK_SZ - 1) / BLOCK_SZ)
116 | if smem is None:
117 |
118 | # Declare input array and all partial sums
119 | ele_num = length
120 |
121 | # Get starting position and length
122 | ele_nums[0] = ele_num
123 | start_pos = 0
124 | ele_nums_pos[0] = start_pos
125 |
126 | while (ele_num > 1):
127 | ele_num = int((ele_num + BLOCK_SZ - 1) / BLOCK_SZ)
128 | ele_nums.append(ele_num)
129 | start_pos += BLOCK_SZ * ele_num
130 | ele_nums_pos.append(start_pos)
131 |
132 | large_arr = ti.field(ti.i32, shape = start_pos)
133 | smem = ti.field(ti.i32, shape=(int(GRID_SZ), 64))
134 |
135 | blit_from_field_to_field(large_arr, input_arr, 0, length)
136 |
137 | for i in range(len(ele_nums) - 1):
138 | if i == len(ele_nums) - 2:
139 | shfl_scan(large_arr, ele_nums_pos[i], ele_nums_pos[i + 1], smem, True)
140 | else:
141 | shfl_scan(large_arr, ele_nums_pos[i], ele_nums_pos[i + 1], smem, False)
142 |
143 | for i in range(len(ele_nums) - 3, -1, -1):
144 | uniform_add(large_arr, ele_nums_pos[i], ele_nums_pos[i + 1])
145 |
146 | blit_from_field_to_field(input_arr, large_arr, 0, length)
147 |
--------------------------------------------------------------------------------
/sph_base.py:
--------------------------------------------------------------------------------
1 | from matplotlib.pyplot import axis
2 | import taichi as ti
3 | import numpy as np
4 |
5 |
6 | @ti.data_oriented
7 | class SPHBase:
8 | def __init__(self, particle_system):
9 | self.ps = particle_system
10 | self.g = ti.Vector([0.0, -9.81, 0.0]) # Gravity
11 | if self.ps.dim == 2:
12 | self.g = ti.Vector([0.0, -9.81])
13 | self.g = np.array(self.ps.cfg.get_cfg("gravitation"))
14 |
15 | self.viscosity = 0.01 # viscosity
16 |
17 | self.density_0 = 1000.0 # reference density
18 | self.density_0 = self.ps.cfg.get_cfg("density0")
19 |
20 | self.dt = ti.field(float, shape=())
21 | self.dt[None] = 1e-4
22 |
23 | @ti.func
24 | def cubic_kernel(self, r_norm):
25 | res = ti.cast(0.0, ti.f32)
26 | h = self.ps.support_radius
27 | # value of cubic spline smoothing kernel
28 | k = 1.0
29 | if self.ps.dim == 1:
30 | k = 4 / 3
31 | elif self.ps.dim == 2:
32 | k = 40 / 7 / np.pi
33 | elif self.ps.dim == 3:
34 | k = 8 / np.pi
35 | k /= h ** self.ps.dim
36 | q = r_norm / h
37 | if q <= 1.0:
38 | if q <= 0.5:
39 | q2 = q * q
40 | q3 = q2 * q
41 | res = k * (6.0 * q3 - 6.0 * q2 + 1)
42 | else:
43 | res = k * 2 * ti.pow(1 - q, 3.0)
44 | return res
45 |
46 | @ti.func
47 | def cubic_kernel_derivative(self, r):
48 | h = self.ps.support_radius
49 | # derivative of cubic spline smoothing kernel
50 | k = 1.0
51 | if self.ps.dim == 1:
52 | k = 4 / 3
53 | elif self.ps.dim == 2:
54 | k = 40 / 7 / np.pi
55 | elif self.ps.dim == 3:
56 | k = 8 / np.pi
57 | k = 6. * k / h ** self.ps.dim
58 | r_norm = r.norm()
59 | q = r_norm / h
60 | res = ti.Vector([0.0 for _ in range(self.ps.dim)])
61 | if r_norm > 1e-5 and q <= 1.0:
62 | grad_q = r / (r_norm * h)
63 | if q <= 0.5:
64 | res = k * q * (3.0 * q - 2.0) * grad_q
65 | else:
66 | factor = 1.0 - q
67 | res = k * (-factor * factor) * grad_q
68 | return res
69 |
70 | @ti.func
71 | def viscosity_force(self, p_i, p_j, r):
72 | # Compute the viscosity force contribution
73 | v_xy = (self.ps.v[p_i] -
74 | self.ps.v[p_j]).dot(r)
75 | res = 2 * (self.ps.dim + 2) * self.viscosity * (self.ps.m[p_j] / (self.ps.density[p_j])) * v_xy / (
76 | r.norm()**2 + 0.01 * self.ps.support_radius**2) * self.cubic_kernel_derivative(
77 | r)
78 | return res
79 |
80 | def initialize(self):
81 | self.ps.initialize_particle_system()
82 | for r_obj_id in self.ps.object_id_rigid_body:
83 | self.compute_rigid_rest_cm(r_obj_id)
84 | self.compute_static_boundary_volume()
85 | self.compute_moving_boundary_volume()
86 |
87 | @ti.kernel
88 | def compute_rigid_rest_cm(self, object_id: int):
89 | self.ps.rigid_rest_cm[object_id] = self.compute_com(object_id)
90 |
91 | @ti.kernel
92 | def compute_static_boundary_volume(self):
93 | for p_i in ti.grouped(self.ps.x):
94 | if not self.ps.is_static_rigid_body(p_i):
95 | continue
96 | delta = self.cubic_kernel(0.0)
97 | self.ps.for_all_neighbors(p_i, self.compute_boundary_volume_task, delta)
98 | self.ps.m_V[p_i] = 1.0 / delta * 3.0 # TODO: the 3.0 here is a coefficient for missing particles by trail and error... need to figure out how to determine it sophisticatedly
99 |
100 | @ti.func
101 | def compute_boundary_volume_task(self, p_i, p_j, delta: ti.template()):
102 | if self.ps.material[p_j] == self.ps.material_solid:
103 | delta += self.cubic_kernel((self.ps.x[p_i] - self.ps.x[p_j]).norm())
104 |
105 |
106 | @ti.kernel
107 | def compute_moving_boundary_volume(self):
108 | for p_i in ti.grouped(self.ps.x):
109 | if not self.ps.is_dynamic_rigid_body(p_i):
110 | continue
111 | delta = self.cubic_kernel(0.0)
112 | self.ps.for_all_neighbors(p_i, self.compute_boundary_volume_task, delta)
113 | self.ps.m_V[p_i] = 1.0 / delta * 3.0 # TODO: the 3.0 here is a coefficient for missing particles by trail and error... need to figure out how to determine it sophisticatedly
114 |
115 | def substep(self):
116 | pass
117 |
118 | @ti.func
119 | def simulate_collisions(self, p_i, vec):
120 | # Collision factor, assume roughly (1-c_f)*velocity loss after collision
121 | c_f = 0.5
122 | self.ps.v[p_i] -= (
123 | 1.0 + c_f) * self.ps.v[p_i].dot(vec) * vec
124 |
125 | @ti.kernel
126 | def enforce_boundary_2D(self, particle_type:int):
127 | for p_i in ti.grouped(self.ps.x):
128 | if self.ps.material[p_i] == particle_type and self.ps.is_dynamic[p_i]:
129 | pos = self.ps.x[p_i]
130 | collision_normal = ti.Vector([0.0, 0.0])
131 | if pos[0] > self.ps.domain_size[0] - self.ps.padding:
132 | collision_normal[0] += 1.0
133 | self.ps.x[p_i][0] = self.ps.domain_size[0] - self.ps.padding
134 | if pos[0] <= self.ps.padding:
135 | collision_normal[0] += -1.0
136 | self.ps.x[p_i][0] = self.ps.padding
137 |
138 | if pos[1] > self.ps.domain_size[1] - self.ps.padding:
139 | collision_normal[1] += 1.0
140 | self.ps.x[p_i][1] = self.ps.domain_size[1] - self.ps.padding
141 | if pos[1] <= self.ps.padding:
142 | collision_normal[1] += -1.0
143 | self.ps.x[p_i][1] = self.ps.padding
144 | collision_normal_length = collision_normal.norm()
145 | if collision_normal_length > 1e-6:
146 | self.simulate_collisions(
147 | p_i, collision_normal / collision_normal_length)
148 |
149 | @ti.kernel
150 | def enforce_boundary_3D(self, particle_type:int):
151 | for p_i in ti.grouped(self.ps.x):
152 | if self.ps.material[p_i] == particle_type and self.ps.is_dynamic[p_i]:
153 | pos = self.ps.x[p_i]
154 | collision_normal = ti.Vector([0.0, 0.0, 0.0])
155 | if pos[0] > self.ps.domain_size[0] - self.ps.padding:
156 | collision_normal[0] += 1.0
157 | self.ps.x[p_i][0] = self.ps.domain_size[0] - self.ps.padding
158 | if pos[0] <= self.ps.padding:
159 | collision_normal[0] += -1.0
160 | self.ps.x[p_i][0] = self.ps.padding
161 |
162 | if pos[1] > self.ps.domain_size[1] - self.ps.padding:
163 | collision_normal[1] += 1.0
164 | self.ps.x[p_i][1] = self.ps.domain_size[1] - self.ps.padding
165 | if pos[1] <= self.ps.padding:
166 | collision_normal[1] += -1.0
167 | self.ps.x[p_i][1] = self.ps.padding
168 |
169 | if pos[2] > self.ps.domain_size[2] - self.ps.padding:
170 | collision_normal[2] += 1.0
171 | self.ps.x[p_i][2] = self.ps.domain_size[2] - self.ps.padding
172 | if pos[2] <= self.ps.padding:
173 | collision_normal[2] += -1.0
174 | self.ps.x[p_i][2] = self.ps.padding
175 |
176 | collision_normal_length = collision_normal.norm()
177 | if collision_normal_length > 1e-6:
178 | self.simulate_collisions(
179 | p_i, collision_normal / collision_normal_length)
180 |
181 |
182 | @ti.func
183 | def compute_com(self, object_id):
184 | sum_m = 0.0
185 | cm = ti.Vector([0.0, 0.0, 0.0])
186 | for p_i in range(self.ps.particle_num[None]):
187 | if self.ps.is_dynamic_rigid_body(p_i) and self.ps.object_id[p_i] == object_id:
188 | mass = self.ps.m_V0 * self.ps.density[p_i]
189 | cm += mass * self.ps.x[p_i]
190 | sum_m += mass
191 | cm /= sum_m
192 | return cm
193 |
194 |
195 | @ti.kernel
196 | def compute_com_kernel(self, object_id: int)->ti.types.vector(3, float):
197 | return self.compute_com(object_id)
198 |
199 |
200 | @ti.kernel
201 | def solve_constraints(self, object_id: int) -> ti.types.matrix(3, 3, float):
202 | # compute center of mass
203 | cm = self.compute_com(object_id)
204 | # A
205 | A = ti.Matrix([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
206 | for p_i in range(self.ps.particle_num[None]):
207 | if self.ps.is_dynamic_rigid_body(p_i) and self.ps.object_id[p_i] == object_id:
208 | q = self.ps.x_0[p_i] - self.ps.rigid_rest_cm[object_id]
209 | p = self.ps.x[p_i] - cm
210 | A += self.ps.m_V0 * self.ps.density[p_i] * p.outer_product(q)
211 |
212 | R, S = ti.polar_decompose(A)
213 |
214 | if all(abs(R) < 1e-6):
215 | R = ti.Matrix.identity(ti.f32, 3)
216 |
217 | for p_i in range(self.ps.particle_num[None]):
218 | if self.ps.is_dynamic_rigid_body(p_i) and self.ps.object_id[p_i] == object_id:
219 | goal = cm + R @ (self.ps.x_0[p_i] - self.ps.rigid_rest_cm[object_id])
220 | corr = (goal - self.ps.x[p_i]) * 1.0
221 | self.ps.x[p_i] += corr
222 | return R
223 |
224 |
225 | # @ti.kernel
226 | # def compute_rigid_collision(self):
227 | # # FIXME: This is a workaround, rigid collision failure in some cases is expected
228 | # for p_i in range(self.ps.particle_num[None]):
229 | # if not self.ps.is_dynamic_rigid_body(p_i):
230 | # continue
231 | # cnt = 0
232 | # x_delta = ti.Vector([0.0 for i in range(self.ps.dim)])
233 | # for j in range(self.ps.solid_neighbors_num[p_i]):
234 | # p_j = self.ps.solid_neighbors[p_i, j]
235 |
236 | # if self.ps.is_static_rigid_body(p_i):
237 | # cnt += 1
238 | # x_j = self.ps.x[p_j]
239 | # r = self.ps.x[p_i] - x_j
240 | # if r.norm() < self.ps.particle_diameter:
241 | # x_delta += (r.norm() - self.ps.particle_diameter) * r.normalized()
242 | # if cnt > 0:
243 | # self.ps.x[p_i] += 2.0 * x_delta # / cnt
244 |
245 |
246 |
247 | def solve_rigid_body(self):
248 | for i in range(1):
249 | for r_obj_id in self.ps.object_id_rigid_body:
250 | if self.ps.object_collection[r_obj_id]["isDynamic"]:
251 | R = self.solve_constraints(r_obj_id)
252 |
253 | if self.ps.cfg.get_cfg("exportObj"):
254 | # For output obj only: update the mesh
255 | cm = self.compute_com_kernel(r_obj_id)
256 | ret = R.to_numpy() @ (self.ps.object_collection[r_obj_id]["restPosition"] - self.ps.object_collection[r_obj_id]["restCenterOfMass"]).T
257 | self.ps.object_collection[r_obj_id]["mesh"].vertices = cm.to_numpy() + ret.T
258 |
259 | # self.compute_rigid_collision()
260 | self.enforce_boundary_3D(self.ps.material_solid)
261 |
262 |
263 | def step(self):
264 | self.ps.initialize_particle_system()
265 | self.compute_moving_boundary_volume()
266 | self.substep()
267 | self.solve_rigid_body()
268 | if self.ps.dim == 2:
269 | self.enforce_boundary_2D(self.ps.material_fluid)
270 | elif self.ps.dim == 3:
271 | self.enforce_boundary_3D(self.ps.material_fluid)
272 |
--------------------------------------------------------------------------------