Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Commit 878111d

Browse filesBrowse files
committed
Add a new quadtree nbody simulation using the Barnes Hut algorithm
Signed-off-by: Pablo Galindo <pablogsal@gmail.com>
1 parent 9e29e3f commit 878111d
Copy full SHA for 878111d

File tree

2 files changed

+393
-0
lines changed
Filter options

2 files changed

+393
-0
lines changed
+10Lines 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"
+383Lines 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 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.