-
Notifications
You must be signed in to change notification settings - Fork 493
Expand file tree
/
Copy pathexample_tile_nbody.py
More file actions
187 lines (141 loc) · 5.93 KB
/
example_tile_nbody.py
File metadata and controls
187 lines (141 loc) · 5.93 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
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
###########################################################################
# Example N-Body
#
# Shows how to simulate an N-Body gravitational problem using an all-pairs
# approach with Warp tile primitives.
#
# References:
# L. Nyland, M. Harris, and J. Prins. "Fast N-Body Simulation with
# CUDA" in GPU Gems 3. H. Nguyen, Addison-Wesley Professional, 2007.
# https://developer.nvidia.com/gpugems/gpugems3/part-v-physics-simulation/chapter-31-fast-n-body-simulation-cuda
#
###########################################################################
import argparse
import numpy as np
import warp as wp
wp.init()
DT = wp.constant(0.016)
SOFTENING_SQ = wp.constant(0.1**2) # Softening factor for numerical stability
TILE_SIZE = wp.constant(64)
PARTICLE_MASS = wp.constant(1.0)
@wp.func
def body_body_interaction(p0: wp.vec3, pi: wp.vec3):
"""Return the acceleration of the particle at position `p0` due to the
particle at position `pi`."""
r = pi - p0
dist_sq = wp.length_sq(r) + SOFTENING_SQ
inv_dist = 1.0 / wp.sqrt(dist_sq)
inv_dist_cubed = inv_dist * inv_dist * inv_dist
acc = PARTICLE_MASS * inv_dist_cubed * r
return acc
@wp.kernel
def integrate_bodies_tiled(
old_position: wp.array[wp.vec3],
velocity: wp.array[wp.vec3],
new_position: wp.array[wp.vec3],
num_bodies: int,
):
i = wp.tid()
p0 = old_position[i]
accel = wp.vec3(0.0, 0.0, 0.0)
for k in range(num_bodies / TILE_SIZE):
k_tile = wp.tile_load(old_position, shape=TILE_SIZE, offset=k * TILE_SIZE)
for idx in range(TILE_SIZE):
pi = k_tile[idx]
accel += body_body_interaction(p0, pi)
# Advance the velocity one timestep (in-place)
velocity[i] = velocity[i] + accel * DT
# Advance the positions (using a second array)
new_position[i] = old_position[i] + DT * velocity[i]
class Example:
def __init__(self, headless=False, num_bodies=16384):
self.num_bodies = num_bodies
rng = np.random.default_rng(42)
# Sample the surface of a sphere
phi = np.arccos(1.0 - 2.0 * rng.uniform(low=0.0, high=1.0, size=self.num_bodies))
theta = rng.uniform(low=0.0, high=2.0 * np.pi, size=self.num_bodies)
x = np.cos(theta) * np.sin(phi)
y = np.sin(theta) * np.sin(phi)
z = np.cos(phi)
init_pos_np = np.stack((x, y, z), axis=1)
scale = (num_bodies / 1024) ** (1 / 2) # Scale factor to maintain a constant density
inner = 0.9625 * scale
outer = 1.54 * scale
radii = inner + (outer - inner) * rng.uniform(size=(self.num_bodies, 1))
init_pos_np = init_pos_np * radii
axis = np.array([0.0, 0.0, 1.0])
v_scale = scale * 3.08
init_vel_np = v_scale * np.cross(init_pos_np, axis)
self.graph_scale = np.max(radii) * 5.0
self.pos_array_0 = wp.array(init_pos_np, dtype=wp.vec3)
self.pos_array_1 = wp.empty_like(self.pos_array_0)
self.vel_array = wp.array(init_vel_np, dtype=wp.vec3)
if headless:
self.scatter_plot = None
else:
self.scatter_plot = self.create_plot()
def create_plot(self):
import matplotlib.pyplot as plt # noqa: PLC0415
# Create a figure and a 3D axis for the plot
self.fig = plt.figure()
ax = self.fig.add_subplot(111, projection="3d")
# Scatter plot of initial positions
point_size = 0.05 * self.graph_scale
init_pos_np = self.pos_array_0.numpy()
scatter_plot = ax.scatter(
init_pos_np[:, 0], init_pos_np[:, 1], init_pos_np[:, 2], s=point_size, c="#76b900", alpha=0.5
)
# Set axis limits
ax.set_xlim(-self.graph_scale, self.graph_scale)
ax.set_ylim(-self.graph_scale, self.graph_scale)
ax.set_zlim(-self.graph_scale, self.graph_scale)
return scatter_plot
def step(self):
wp.launch(
integrate_bodies_tiled,
dim=self.num_bodies,
inputs=[self.pos_array_0, self.vel_array, self.pos_array_1, self.num_bodies],
block_dim=TILE_SIZE,
)
# Swap arrays
(self.pos_array_0, self.pos_array_1) = (self.pos_array_1, self.pos_array_0)
def render(self):
positions_cpu = self.pos_array_0.numpy()
# Update scatter plot positions
self.scatter_plot._offsets3d = (
positions_cpu[:, 0],
positions_cpu[:, 1],
positions_cpu[:, 2],
)
# Function to update the scatter plot
def step_and_render(self, frame):
self.step()
self.render()
if __name__ == "__main__":
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
parser.add_argument("--num-frames", type=int, default=1000, help="Total number of frames.")
parser.add_argument("-N", help="Number of bodies. Should be a multiple of 64.", type=int, default=16384)
parser.add_argument(
"--headless",
action="store_true",
help="Run in headless mode, suppressing the opening of any graphical windows.",
)
args = parser.parse_known_args()[0]
if args.device == "cpu":
print("This example only runs on CUDA devices.")
exit()
with wp.ScopedDevice(args.device):
example = Example(headless=args.headless, num_bodies=args.N)
if not args.headless:
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
# Create the animation
ani = FuncAnimation(example.fig, example.step_and_render, frames=args.num_frames, interval=50, repeat=False)
# Display the animation
plt.show()
else:
for _ in range(args.num_frames):
example.step()