-
Notifications
You must be signed in to change notification settings - Fork 493
Expand file tree
/
Copy pathexample_sph.py
More file actions
399 lines (315 loc) · 11.9 KB
/
example_sph.py
File metadata and controls
399 lines (315 loc) · 11.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
###########################################################################
# Example Smoothed Particle Hydrodynamics
#
# Shows how to implement a SPH fluid simulation.
#
# Neighbors are found using the wp.HashGrid class, and
# wp.hash_grid_query(), wp.hash_grid_query_next() kernel methods.
#
# Reference Publication
# Matthias Müller, David Charypar, and Markus H. Gross.
# "Particle-based fluid simulation for interactive applications."
# Symposium on Computer animation. Vol. 2. 2003.
#
###########################################################################
import numpy as np
import warp as wp
import warp.render
@wp.func
def square(x: float):
return x * x
@wp.func
def cube(x: float):
return x * x * x
@wp.func
def fifth(x: float):
return x * x * x * x * x
@wp.func
def density_kernel(xyz: wp.vec3, smoothing_length: float):
# calculate distance
distance = wp.dot(xyz, xyz)
return wp.max(cube(square(smoothing_length) - distance), 0.0)
@wp.func
def diff_pressure_kernel(
xyz: wp.vec3, pressure: float, neighbor_pressure: float, neighbor_rho: float, smoothing_length: float
):
# calculate distance
distance = wp.sqrt(wp.dot(xyz, xyz))
if distance < smoothing_length:
# calculate terms of kernel
term_1 = -xyz / distance
term_2 = (neighbor_pressure + pressure) / (2.0 * neighbor_rho)
term_3 = square(smoothing_length - distance)
return term_1 * term_2 * term_3
else:
return wp.vec3()
@wp.func
def diff_viscous_kernel(xyz: wp.vec3, v: wp.vec3, neighbor_v: wp.vec3, neighbor_rho: float, smoothing_length: float):
# calculate distance
distance = wp.sqrt(wp.dot(xyz, xyz))
# calculate terms of kernel
if distance < smoothing_length:
term_1 = (neighbor_v - v) / neighbor_rho
term_2 = smoothing_length - distance
return term_1 * term_2
else:
return wp.vec3()
@wp.kernel
def compute_density(
grid: wp.uint64,
particle_x: wp.array[wp.vec3],
particle_rho: wp.array[float],
density_normalization: float,
smoothing_length: float,
):
tid = wp.tid()
# order threads by cell
i = wp.hash_grid_point_id(grid, tid)
# get local particle variables
x = particle_x[i]
# store density
rho = float(0.0)
# particle contact
neighbors = wp.hash_grid_query(grid, x, smoothing_length)
# loop through neighbors to compute density
for index in neighbors:
# compute distance
distance = x - particle_x[index]
# compute kernel derivative
rho += density_kernel(distance, smoothing_length)
# add external potential
particle_rho[i] = density_normalization * rho
@wp.kernel
def get_acceleration(
grid: wp.uint64,
particle_x: wp.array[wp.vec3],
particle_v: wp.array[wp.vec3],
particle_rho: wp.array[float],
particle_a: wp.array[wp.vec3],
isotropic_exp: float,
base_density: float,
gravity: float,
pressure_normalization: float,
viscous_normalization: float,
smoothing_length: float,
):
tid = wp.tid()
# order threads by cell
i = wp.hash_grid_point_id(grid, tid)
# get local particle variables
x = particle_x[i]
v = particle_v[i]
rho = particle_rho[i]
pressure = isotropic_exp * (rho - base_density)
# store forces
pressure_force = wp.vec3()
viscous_force = wp.vec3()
# particle contact
neighbors = wp.hash_grid_query(grid, x, smoothing_length)
# loop through neighbors to compute acceleration
for index in neighbors:
if index != i:
# get neighbor velocity
neighbor_v = particle_v[index]
# get neighbor density and pressures
neighbor_rho = particle_rho[index]
neighbor_pressure = isotropic_exp * (neighbor_rho - base_density)
# compute relative position
relative_position = particle_x[index] - x
# calculate pressure force
pressure_force += diff_pressure_kernel(
relative_position, pressure, neighbor_pressure, neighbor_rho, smoothing_length
)
# compute kernel derivative
viscous_force += diff_viscous_kernel(relative_position, v, neighbor_v, neighbor_rho, smoothing_length)
# sum all forces
force = pressure_normalization * pressure_force + viscous_normalization * viscous_force
# add external potential
particle_a[i] = force / rho + wp.vec3(0.0, gravity, 0.0)
@wp.kernel
def apply_bounds(
particle_x: wp.array[wp.vec3],
particle_v: wp.array[wp.vec3],
damping_coef: float,
width: float,
height: float,
length: float,
):
tid = wp.tid()
# get pos and velocity
x = particle_x[tid]
v = particle_v[tid]
# clamp x left
if x[0] < 0.0:
x = wp.vec3(0.0, x[1], x[2])
v = wp.vec3(v[0] * damping_coef, v[1], v[2])
# clamp x right
if x[0] > width:
x = wp.vec3(width, x[1], x[2])
v = wp.vec3(v[0] * damping_coef, v[1], v[2])
# clamp y bot
if x[1] < 0.0:
x = wp.vec3(x[0], 0.0, x[2])
v = wp.vec3(v[0], v[1] * damping_coef, v[2])
# clamp z left
if x[2] < 0.0:
x = wp.vec3(x[0], x[1], 0.0)
v = wp.vec3(v[0], v[1], v[2] * damping_coef)
# clamp z right
if x[2] > length:
x = wp.vec3(x[0], x[1], length)
v = wp.vec3(v[0], v[1], v[2] * damping_coef)
# apply clamps
particle_x[tid] = x
particle_v[tid] = v
@wp.kernel
def kick(particle_v: wp.array[wp.vec3], particle_a: wp.array[wp.vec3], dt: float):
tid = wp.tid()
v = particle_v[tid]
particle_v[tid] = v + particle_a[tid] * dt
@wp.kernel
def drift(particle_x: wp.array[wp.vec3], particle_v: wp.array[wp.vec3], dt: float):
tid = wp.tid()
x = particle_x[tid]
particle_x[tid] = x + particle_v[tid] * dt
@wp.kernel
def initialize_particles(
particle_x: wp.array[wp.vec3], smoothing_length: float, width: float, height: float, length: float
):
tid = wp.tid()
# grid size
nr_x = int(width / 4.0 / smoothing_length)
nr_y = int(height / smoothing_length)
nr_z = int(length / 4.0 / smoothing_length)
# calculate particle position
z = float(tid % nr_z)
y = float((tid // nr_z) % nr_y)
x = float((tid // (nr_z * nr_y)) % nr_x)
pos = smoothing_length * wp.vec3(x, y, z)
# add small jitter
state = wp.rand_init(123, tid)
pos = pos + 0.001 * smoothing_length * wp.vec3(wp.randn(state), wp.randn(state), wp.randn(state))
# set position
particle_x[tid] = pos
class Example:
def __init__(self, stage_path="example_sph.usd", verbose=False):
self.verbose = verbose
# render params
fps = 60
self.frame_dt = 1.0 / fps
self.sim_time = 0.0
# simulation params
self.smoothing_length = 0.8 # NOTE change this to adjust number of particles
self.width = 80.0 # x
self.height = 80.0 # y
self.length = 80.0 # z
self.isotropic_exp = 20
self.base_density = 1.0
self.particle_mass = 0.01 * self.smoothing_length**3 # reduce according to smoothing length
self.dt = 0.01 * self.smoothing_length # decrease sim dt by smoothing length
self.dynamic_visc = 0.025
self.damping_coef = -0.95
self.gravity = -0.1
self.n = int(
self.height * (self.width / 4.0) * (self.height / 4.0) / (self.smoothing_length**3)
) # number particles (small box in corner)
self.sim_step_to_frame_ratio = int(32 / self.smoothing_length)
# constants
self.density_normalization = (315.0 * self.particle_mass) / (
64.0 * np.pi * self.smoothing_length**9
) # integrate density kernel
self.pressure_normalization = -(45.0 * self.particle_mass) / (np.pi * self.smoothing_length**6)
self.viscous_normalization = (45.0 * self.dynamic_visc * self.particle_mass) / (
np.pi * self.smoothing_length**6
)
# allocate arrays
self.x = wp.empty(self.n, dtype=wp.vec3)
self.v = wp.zeros(self.n, dtype=wp.vec3)
self.rho = wp.zeros(self.n, dtype=float)
self.a = wp.zeros(self.n, dtype=wp.vec3)
# set random positions
wp.launch(
kernel=initialize_particles,
dim=self.n,
inputs=[self.x, self.smoothing_length, self.width, self.height, self.length],
) # initialize in small area
# create hash array
grid_size = int(self.height / (4.0 * self.smoothing_length))
self.grid = wp.HashGrid(grid_size, grid_size, grid_size)
# renderer
self.renderer = None
if stage_path:
self.renderer = wp.render.UsdRenderer(stage_path)
def step(self):
with wp.ScopedTimer("step"):
for _ in range(self.sim_step_to_frame_ratio):
with wp.ScopedTimer("grid build", active=self.verbose):
# build grid
self.grid.build(self.x, self.smoothing_length)
with wp.ScopedTimer("forces", active=self.verbose):
# compute density of points
wp.launch(
kernel=compute_density,
dim=self.n,
inputs=[self.grid.id, self.x, self.rho, self.density_normalization, self.smoothing_length],
)
# get new acceleration
wp.launch(
kernel=get_acceleration,
dim=self.n,
inputs=[
self.grid.id,
self.x,
self.v,
self.rho,
self.a,
self.isotropic_exp,
self.base_density,
self.gravity,
self.pressure_normalization,
self.viscous_normalization,
self.smoothing_length,
],
)
# apply bounds
wp.launch(
kernel=apply_bounds,
dim=self.n,
inputs=[self.x, self.v, self.damping_coef, self.width, self.height, self.length],
)
# kick
wp.launch(kernel=kick, dim=self.n, inputs=[self.v, self.a, self.dt])
# drift
wp.launch(kernel=drift, dim=self.n, inputs=[self.x, self.v, self.dt])
self.sim_time += self.frame_dt
def render(self):
if self.renderer is None:
return
with wp.ScopedTimer("render"):
self.renderer.begin_frame(self.sim_time)
self.renderer.render_points(
points=self.x.numpy(), radius=self.smoothing_length, name="points", colors=(0.8, 0.3, 0.2)
)
self.renderer.end_frame()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
parser.add_argument(
"--stage-path",
type=lambda x: None if x == "None" else str(x),
default="example_sph.usd",
help="Path to the output USD file.",
)
parser.add_argument("--num-frames", type=int, default=480, help="Total number of frames.")
parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
args = parser.parse_known_args()[0]
with wp.ScopedDevice(args.device):
example = Example(stage_path=args.stage_path, verbose=args.verbose)
for _ in range(args.num_frames):
example.render()
example.step()
if example.renderer:
example.renderer.save()