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 294909c

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 294909c
Copy full SHA for 294909c

File tree

2 files changed

+382
-0
lines changed
Filter options

2 files changed

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