Skip to content

Commit 878111d

Browse files
committed
Add a new quadtree nbody simulation using the Barnes Hut algorithm
Signed-off-by: Pablo Galindo <[email protected]>
1 parent 9e29e3f commit 878111d

File tree

2 files changed

+393
-0
lines changed

2 files changed

+393
-0
lines changed
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
[project]
2+
name = "pyperformance_bm_barnes_hut"
3+
requires-python = ">=3.8"
4+
dependencies = ["pyperf"]
5+
urls = {repository = "https://github.com/python/pyperformance"}
6+
dynamic = ["version"]
7+
8+
[tool.pyperformance]
9+
name = "barnes_hut"
10+
tags = "math"
Lines changed: 383 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,383 @@
1+
"""
2+
Quadtree N-body benchmark using the pyperf framework.
3+
4+
This benchmark simulates gravitational interactions between particles using
5+
a quadtree for spatial partitioning and the Barnes-Hut approximation algorithm.
6+
No visualization, pure Python implementation without dependencies.
7+
"""
8+
9+
import pyperf
10+
import math
11+
12+
DEFAULT_ITERATIONS = 100
13+
DEFAULT_PARTICLES = 200
14+
DEFAULT_THETA = 0.5
15+
16+
# Constants
17+
G = 6.67430e-11 # Gravitational constant
18+
SOFTENING = 5.0 # Softening parameter to avoid singularities
19+
TIME_STEP = 0.1 # Time step for simulation
20+
21+
class Point:
22+
def __init__(self, x, y):
23+
self.x = x
24+
self.y = y
25+
26+
class Particle:
27+
def __init__(self, x, y, mass=1.0):
28+
self.position = Point(x, y)
29+
self.velocity = Point(0, 0)
30+
self.acceleration = Point(0, 0)
31+
self.mass = mass
32+
33+
def update(self, time_step):
34+
# Update velocity using current acceleration
35+
self.velocity.x += self.acceleration.x * time_step
36+
self.velocity.y += self.acceleration.y * time_step
37+
38+
# Update position using updated velocity
39+
self.position.x += self.velocity.x * time_step
40+
self.position.y += self.velocity.y * time_step
41+
42+
# Reset acceleration for next frame
43+
self.acceleration.x = 0
44+
self.acceleration.y = 0
45+
46+
def apply_force(self, fx, fy):
47+
# F = ma -> a = F/m
48+
self.acceleration.x += fx / self.mass
49+
self.acceleration.y += fy / self.mass
50+
51+
class Rectangle:
52+
def __init__(self, x, y, w, h):
53+
self.x = x
54+
self.y = y
55+
self.w = w
56+
self.h = h
57+
58+
def contains(self, point):
59+
return (
60+
point.x >= self.x - self.w and
61+
point.x < self.x + self.w and
62+
point.y >= self.y - self.h and
63+
point.y < self.y + self.h
64+
)
65+
66+
def intersects(self, range_rect):
67+
return not (
68+
range_rect.x - range_rect.w > self.x + self.w or
69+
range_rect.x + range_rect.w < self.x - self.w or
70+
range_rect.y - range_rect.h > self.y + self.h or
71+
range_rect.y + range_rect.h < self.y - self.h
72+
)
73+
74+
class QuadTree:
75+
def __init__(self, boundary, capacity=4):
76+
self.boundary = boundary
77+
self.capacity = capacity
78+
self.particles = []
79+
self.divided = False
80+
self.center_of_mass = Point(0, 0)
81+
self.total_mass = 0
82+
self.northeast = None
83+
self.northwest = None
84+
self.southeast = None
85+
self.southwest = None
86+
87+
def insert(self, particle):
88+
if not self.boundary.contains(particle.position):
89+
return False
90+
91+
if len(self.particles) < self.capacity and not self.divided:
92+
self.particles.append(particle)
93+
self.update_mass_distribution(particle)
94+
return True
95+
96+
if not self.divided:
97+
self.subdivide()
98+
99+
if self.northeast.insert(particle): return True
100+
if self.northwest.insert(particle): return True
101+
if self.southeast.insert(particle): return True
102+
if self.southwest.insert(particle): return True
103+
104+
# This should never happen if the boundary check is correct
105+
return False
106+
107+
def update_mass_distribution(self, particle):
108+
# Update center of mass and total mass when adding a particle
109+
total_mass_new = self.total_mass + particle.mass
110+
111+
# Calculate new center of mass
112+
if total_mass_new > 0:
113+
self.center_of_mass.x = (self.center_of_mass.x * self.total_mass +
114+
particle.position.x * particle.mass) / total_mass_new
115+
self.center_of_mass.y = (self.center_of_mass.y * self.total_mass +
116+
particle.position.y * particle.mass) / total_mass_new
117+
118+
self.total_mass = total_mass_new
119+
120+
def subdivide(self):
121+
x = self.boundary.x
122+
y = self.boundary.y
123+
w = self.boundary.w / 2
124+
h = self.boundary.h / 2
125+
126+
ne = Rectangle(x + w, y - h, w, h)
127+
self.northeast = QuadTree(ne, self.capacity)
128+
129+
nw = Rectangle(x - w, y - h, w, h)
130+
self.northwest = QuadTree(nw, self.capacity)
131+
132+
se = Rectangle(x + w, y + h, w, h)
133+
self.southeast = QuadTree(se, self.capacity)
134+
135+
sw = Rectangle(x - w, y + h, w, h)
136+
self.southwest = QuadTree(sw, self.capacity)
137+
138+
self.divided = True
139+
140+
# Redistribute particles to children
141+
for particle in self.particles:
142+
self.northeast.insert(particle)
143+
self.northwest.insert(particle)
144+
self.southeast.insert(particle)
145+
self.southwest.insert(particle)
146+
147+
# Clear the particles at this node
148+
self.particles = []
149+
150+
def calculate_force(self, particle, theta):
151+
if self.total_mass == 0:
152+
return 0, 0
153+
154+
# If this is an external node (leaf with one particle) and it's the same particle, skip
155+
if len(self.particles) == 1 and self.particles[0] == particle:
156+
return 0, 0
157+
158+
# Calculate distance between particle and center of mass
159+
dx = self.center_of_mass.x - particle.position.x
160+
dy = self.center_of_mass.y - particle.position.y
161+
distance = math.sqrt(dx*dx + dy*dy)
162+
163+
# If this is a leaf node or the distance is sufficient for approximation
164+
if not self.divided or (self.boundary.w * 2) / distance < theta:
165+
# Avoid division by zero and extreme forces at small distances
166+
if distance < SOFTENING:
167+
distance = SOFTENING
168+
169+
# Calculate gravitational force
170+
f = G * particle.mass * self.total_mass / (distance * distance)
171+
172+
# Resolve force into x and y components
173+
fx = f * dx / distance
174+
fy = f * dy / distance
175+
176+
return fx, fy
177+
178+
# Otherwise, recursively calculate forces from children
179+
fx, fy = 0, 0
180+
181+
if self.northeast:
182+
fx_ne, fy_ne = self.northeast.calculate_force(particle, theta)
183+
fx += fx_ne
184+
fy += fy_ne
185+
186+
if self.northwest:
187+
fx_nw, fy_nw = self.northwest.calculate_force(particle, theta)
188+
fx += fx_nw
189+
fy += fy_nw
190+
191+
if self.southeast:
192+
fx_se, fy_se = self.southeast.calculate_force(particle, theta)
193+
fx += fx_se
194+
fy += fy_se
195+
196+
if self.southwest:
197+
fx_sw, fy_sw = self.southwest.calculate_force(particle, theta)
198+
fx += fx_sw
199+
fy += fy_sw
200+
201+
return fx, fy
202+
203+
def create_deterministic_galaxy(num_particles, center_x, center_y, radius=300, spiral_factor=0.1):
204+
"""Create a deterministic galaxy-like distribution of particles"""
205+
particles = []
206+
207+
# Create central bulge (30% of particles)
208+
bulge_count = int(num_particles * 0.3)
209+
for i in range(bulge_count):
210+
# Use deterministic distribution for the bulge
211+
# Distribute in a spiral pattern from center
212+
distance_factor = i / bulge_count
213+
r = radius / 5 * distance_factor
214+
angle = 2 * math.pi * i / bulge_count * 7 # 7 rotations for central bulge
215+
216+
x = center_x + r * math.cos(angle)
217+
y = center_y + r * math.sin(angle)
218+
219+
# Deterministic mass values
220+
mass = 50 + (50 * (1 - distance_factor))
221+
222+
particle = Particle(x, y, mass)
223+
particles.append(particle)
224+
225+
# Create spiral arms (70% of particles)
226+
spiral_count = num_particles - bulge_count
227+
arms = 2 # Number of spiral arms
228+
particles_per_arm = spiral_count // arms
229+
230+
for arm in range(arms):
231+
base_angle = 2 * math.pi * arm / arms
232+
233+
for i in range(particles_per_arm):
234+
# Distance from center (linearly distributed from inner to outer radius)
235+
distance_factor = i / particles_per_arm
236+
distance = radius * (0.2 + 0.8 * distance_factor)
237+
238+
# Spiral angle based on distance from center
239+
angle = base_angle + spiral_factor * distance
240+
241+
# Deterministic variation
242+
variation = 0.2 * math.sin(i * 0.1)
243+
angle += variation
244+
245+
x = center_x + distance * math.cos(angle)
246+
y = center_y + distance * math.sin(angle)
247+
248+
# Deterministic mass values
249+
mass = 1 + 9 * (1 - distance_factor)
250+
251+
particle = Particle(x, y, mass)
252+
particles.append(particle)
253+
254+
# Remaining particles (if odd number doesn't divide evenly into arms)
255+
remaining = spiral_count - (particles_per_arm * arms)
256+
for i in range(remaining):
257+
angle = 2 * math.pi * i / remaining
258+
distance = radius * 0.5
259+
260+
x = center_x + distance * math.cos(angle)
261+
y = center_y + distance * math.sin(angle)
262+
263+
particle = Particle(x, y, 5.0) # Fixed mass
264+
particles.append(particle)
265+
266+
# Add deterministic initial orbital velocities
267+
for p in particles:
268+
# Vector from center to particle
269+
dx = p.position.x - center_x
270+
dy = p.position.y - center_y
271+
distance = math.sqrt(dx*dx + dy*dy)
272+
273+
if distance > 0:
274+
# Direction perpendicular to radial direction
275+
perp_x = -dy / distance
276+
perp_y = dx / distance
277+
278+
# Orbital velocity (approximating balanced centripetal force)
279+
orbital_velocity = math.sqrt(0.1 * (distance + 100)) * 0.3
280+
281+
p.velocity.x = perp_x * orbital_velocity
282+
p.velocity.y = perp_y * orbital_velocity
283+
284+
return particles
285+
286+
def calculate_system_energy(particles):
287+
"""Calculate the total energy of the system (kinetic + potential)"""
288+
energy = 0.0
289+
290+
# Calculate potential energy
291+
for i in range(len(particles)):
292+
for j in range(i + 1, len(particles)):
293+
p1 = particles[i]
294+
p2 = particles[j]
295+
296+
dx = p1.position.x - p2.position.x
297+
dy = p1.position.y - p2.position.y
298+
distance = math.sqrt(dx*dx + dy*dy)
299+
300+
# Avoid division by zero
301+
if distance < SOFTENING:
302+
distance = SOFTENING
303+
304+
# Gravitational potential energy
305+
energy -= G * p1.mass * p2.mass / distance
306+
307+
# Calculate kinetic energy
308+
for p in particles:
309+
v_squared = p.velocity.x * p.velocity.x + p.velocity.y * p.velocity.y
310+
energy += 0.5 * p.mass * v_squared
311+
312+
return energy
313+
314+
def advance_system(particles, theta, time_step, width, height):
315+
"""Advance the n-body system by one time step using the quadtree"""
316+
# Create a fresh quadtree
317+
boundary = Rectangle(width / 2, height / 2, width / 2, height / 2)
318+
quadtree = QuadTree(boundary)
319+
320+
# Insert all particles into the quadtree
321+
for particle in particles:
322+
quadtree.insert(particle)
323+
324+
# Calculate and apply forces to all particles
325+
for particle in particles:
326+
fx, fy = quadtree.calculate_force(particle, theta)
327+
particle.apply_force(fx, fy)
328+
329+
# Update all particles
330+
for particle in particles:
331+
particle.update(time_step)
332+
333+
def bench_quadtree_nbody(loops, num_particles, iterations, theta):
334+
# Initialize simulation space
335+
width = 1000
336+
height = 800
337+
338+
# Create deterministic galaxy distribution
339+
particles = create_deterministic_galaxy(num_particles, width / 2, height / 2)
340+
341+
# Calculate initial energy
342+
initial_energy = calculate_system_energy(particles)
343+
344+
range_it = range(loops)
345+
t0 = pyperf.perf_counter()
346+
347+
for _ in range_it:
348+
# Run simulation for specified iterations
349+
for _ in range(iterations):
350+
advance_system(particles, theta, TIME_STEP, width, height)
351+
352+
# Calculate final energy
353+
final_energy = calculate_system_energy(particles)
354+
355+
return pyperf.perf_counter() - t0
356+
357+
def add_cmdline_args(cmd, args):
358+
cmd.extend(("--iterations", str(args.iterations)))
359+
cmd.extend(("--particles", str(args.particles)))
360+
cmd.extend(("--theta", str(args.theta)))
361+
362+
if __name__ == '__main__':
363+
runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
364+
runner.metadata['description'] = "Quadtree N-body benchmark"
365+
366+
runner.argparser.add_argument("--iterations",
367+
type=int, default=DEFAULT_ITERATIONS,
368+
help="Number of simulation steps per benchmark loop "
369+
"(default: %s)" % DEFAULT_ITERATIONS)
370+
371+
runner.argparser.add_argument("--particles",
372+
type=int, default=DEFAULT_PARTICLES,
373+
help="Number of particles in the simulation "
374+
"(default: %s)" % DEFAULT_PARTICLES)
375+
376+
runner.argparser.add_argument("--theta",
377+
type=float, default=DEFAULT_THETA,
378+
help="Barnes-Hut approximation threshold "
379+
"(default: %s)" % DEFAULT_THETA)
380+
381+
args = runner.parse_args()
382+
runner.bench_time_func('quadtree_nbody', bench_quadtree_nbody,
383+
args.particles, args.iterations, args.theta)

0 commit comments

Comments
 (0)