Add a new quadtree nbody simulation using the Barnes Hut algorithm

Author: pablogsal

pyperformance/data-files/benchmarks/bm_barnes_hut/pyproject.toml

@@ -0,0 +1,10 @@
+[project]
+name = "pyperformance_bm_nbody"
+requires-python = ">=3.8"
+dependencies = ["pyperf"]
+urls = {repository = "https://github.com/python/pyperformance"}
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "nbody"
+tags = "math"

pyperformance/data-files/benchmarks/bm_barnes_hut/run_benchmark.py

@@ -0,0 +1,372 @@
+"""
+Quadtree N-body benchmark using the pyperf framework.
+
+This benchmark simulates gravitational interactions between particles using
+a quadtree for spatial partitioning and the Barnes-Hut approximation algorithm.
+No visualization, pure Python implementation without dependencies.
+"""
+
+import pyperf
+import math
+import random
+
+DEFAULT_ITERATIONS = 300
+DEFAULT_PARTICLES = 100
+DEFAULT_THETA = 0.5
+
+# Constants
+G = 6.67430e-11  # Gravitational constant
+SOFTENING = 5.0  # Softening parameter to avoid singularities
+TIME_STEP = 0.1  # Time step for simulation
+
+class Point:
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+class Particle:
+    def __init__(self, x, y, mass=1.0):
+        self.position = Point(x, y)
+        self.velocity = Point(0, 0)
+        self.acceleration = Point(0, 0)
+        self.mass = mass
+    
+    def update(self, time_step):
+        # Update velocity using current acceleration
+        self.velocity.x += self.acceleration.x * time_step
+        self.velocity.y += self.acceleration.y * time_step
+        
+        # Update position using updated velocity
+        self.position.x += self.velocity.x * time_step
+        self.position.y += self.velocity.y * time_step
+        
+        # Reset acceleration for next frame
+        self.acceleration.x = 0
+        self.acceleration.y = 0
+    
+    def apply_force(self, fx, fy):
+        # F = ma -> a = F/m
+        self.acceleration.x += fx / self.mass
+        self.acceleration.y += fy / self.mass
+
+class Rectangle:
+    def __init__(self, x, y, w, h):
+        self.x = x
+        self.y = y
+        self.w = w
+        self.h = h
+    
+    def contains(self, point):
+        return (
+            point.x >= self.x - self.w and
+            point.x < self.x + self.w and
+            point.y >= self.y - self.h and
+            point.y < self.y + self.h
+        )
+    
+    def intersects(self, range_rect):
+        return not (
+            range_rect.x - range_rect.w > self.x + self.w or
+            range_rect.x + range_rect.w < self.x - self.w or
+            range_rect.y - range_rect.h > self.y + self.h or
+            range_rect.y + range_rect.h < self.y - self.h
+        )
+
+class QuadTree:
+    def __init__(self, boundary, capacity=4):
+        self.boundary = boundary
+        self.capacity = capacity
+        self.particles = []
+        self.divided = False
+        self.center_of_mass = Point(0, 0)
+        self.total_mass = 0
+        self.northeast = None
+        self.northwest = None
+        self.southeast = None
+        self.southwest = None
+    
+    def insert(self, particle):
+        if not self.boundary.contains(particle.position):
+            return False
+        
+        if len(self.particles) < self.capacity and not self.divided:
+            self.particles.append(particle)
+            self.update_mass_distribution(particle)
+            return True
+        
+        if not self.divided:
+            self.subdivide()
+        
+        if self.northeast.insert(particle): return True
+        if self.northwest.insert(particle): return True
+        if self.southeast.insert(particle): return True
+        if self.southwest.insert(particle): return True
+        
+        # This should never happen if the boundary check is correct
+        return False
+    
+    def update_mass_distribution(self, particle):
+        # Update center of mass and total mass when adding a particle
+        total_mass_new = self.total_mass + particle.mass
+        
+        # Calculate new center of mass
+        if total_mass_new > 0:
+            self.center_of_mass.x = (self.center_of_mass.x * self.total_mass + 
+                                    particle.position.x * particle.mass) / total_mass_new
+            self.center_of_mass.y = (self.center_of_mass.y * self.total_mass + 
+                                    particle.position.y * particle.mass) / total_mass_new
+        
+        self.total_mass = total_mass_new
+    
+    def subdivide(self):
+        x = self.boundary.x
+        y = self.boundary.y
+        w = self.boundary.w / 2
+        h = self.boundary.h / 2
+        
+        ne = Rectangle(x + w, y - h, w, h)
+        self.northeast = QuadTree(ne, self.capacity)
+        
+        nw = Rectangle(x - w, y - h, w, h)
+        self.northwest = QuadTree(nw, self.capacity)
+        
+        se = Rectangle(x + w, y + h, w, h)
+        self.southeast = QuadTree(se, self.capacity)
+        
+        sw = Rectangle(x - w, y + h, w, h)
+        self.southwest = QuadTree(sw, self.capacity)
+        
+        self.divided = True
+        
+        # Redistribute particles to children
+        for particle in self.particles:
+            self.northeast.insert(particle)
+            self.northwest.insert(particle)
+            self.southeast.insert(particle)
+            self.southwest.insert(particle)
+        
+        # Clear the particles at this node
+        self.particles = []
+    
+    def calculate_force(self, particle, theta):
+        if self.total_mass == 0:
+            return 0, 0
+        
+        # If this is an external node (leaf with one particle) and it's the same particle, skip
+        if len(self.particles) == 1 and self.particles[0] == particle:
+            return 0, 0
+        
+        # Calculate distance between particle and center of mass
+        dx = self.center_of_mass.x - particle.position.x
+        dy = self.center_of_mass.y - particle.position.y
+        distance = math.sqrt(dx*dx + dy*dy)
+        
+        # If this is a leaf node or the distance is sufficient for approximation
+        if not self.divided or (self.boundary.w * 2) / distance < theta:
+            # Avoid division by zero and extreme forces at small distances
+            if distance < SOFTENING:
+                distance = SOFTENING
+                
+            # Calculate gravitational force
+            f = G * particle.mass * self.total_mass / (distance * distance)
+            
+            # Resolve force into x and y components
+            fx = f * dx / distance
+            fy = f * dy / distance
+            
+            return fx, fy
+        
+        # Otherwise, recursively calculate forces from children
+        fx, fy = 0, 0
+        
+        if self.northeast:
+            fx_ne, fy_ne = self.northeast.calculate_force(particle, theta)
+            fx += fx_ne
+            fy += fy_ne
+        
+        if self.northwest:
+            fx_nw, fy_nw = self.northwest.calculate_force(particle, theta)
+            fx += fx_nw
+            fy += fy_nw
+        
+        if self.southeast:
+            fx_se, fy_se = self.southeast.calculate_force(particle, theta)
+            fx += fx_se
+            fy += fy_se
+        
+        if self.southwest:
+            fx_sw, fy_sw = self.southwest.calculate_force(particle, theta)
+            fx += fx_sw
+            fy += fy_sw
+        
+        return fx, fy
+
+def create_galaxy_distribution(num_particles, center_x, center_y, radius=300, spiral_factor=0.1):
+    particles = []
+    
+    # Create central bulge (30% of particles)
+    bulge_count = int(num_particles * 0.3)
+    for _ in range(bulge_count):
+        # Use gaussian-like distribution for the bulge
+        # Using Box-Muller transform for gaussian approximation
+        u1 = random.random()
+        u2 = random.random()
+        r = radius / 5 * math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
+        angle = random.uniform(0, 2 * math.pi)
+        
+        x = center_x + r * math.cos(angle)
+        y = center_y + r * math.sin(angle)
+        
+        # Heavier particles in the center
+        mass = random.uniform(50, 100) * (1 - r / radius) + 1
+        
+        particle = Particle(x, y, mass)
+        particles.append(particle)
+    
+    # Create spiral arms (70% of particles)
+    spiral_count = num_particles - bulge_count
+    arms = 2  # Number of spiral arms
+    
+    for i in range(spiral_count):
+        # Choose one of the spiral arms
+        arm = i % arms
+        base_angle = 2 * math.pi * arm / arms
+        
+        # Distance from center (more particles further out)
+        distance = random.uniform(radius * 0.2, radius)
+        
+        # Spiral angle based on distance from center
+        angle = base_angle + spiral_factor * distance
+        
+        # Add some randomness to the spiral
+        angle += random.uniform(-0.2, 0.2)
+        
+        x = center_x + distance * math.cos(angle)
+        y = center_y + distance * math.sin(angle)
+        
+        # Lighter particles in the spiral arms
+        mass = random.uniform(1, 10)
+        
+        particle = Particle(x, y, mass)
+        particles.append(particle)
+    
+    # Add initial orbital velocities
+    for p in particles:
+        # Vector from center to particle
+        dx = p.position.x - center_x
+        dy = p.position.y - center_y
+        distance = math.sqrt(dx*dx + dy*dy)
+        
+        if distance > 0:
+            # Direction perpendicular to radial direction
+            perp_x = -dy / distance
+            perp_y = dx / distance
+            
+            # Orbital velocity (approximating balanced centripetal force)
+            # v = sqrt(G * M / r) where M is the mass inside the orbit
+            # We'll simplify this for visual appeal
+            orbital_velocity = math.sqrt(0.1 * (distance + 100)) * 0.3
+            
+            p.velocity.x = perp_x * orbital_velocity
+            p.velocity.y = perp_y * orbital_velocity
+    
+    return particles
+
+def calculate_system_energy(particles):
+    """Calculate the total energy of the system (kinetic + potential)"""
+    energy = 0.0
+    
+    # Calculate potential energy
+    for i in range(len(particles)):
+        for j in range(i + 1, len(particles)):
+            p1 = particles[i]
+            p2 = particles[j]
+            
+            dx = p1.position.x - p2.position.x
+            dy = p1.position.y - p2.position.y
+            distance = math.sqrt(dx*dx + dy*dy)
+            
+            # Avoid division by zero
+            if distance < SOFTENING:
+                distance = SOFTENING
+            
+            # Gravitational potential energy
+            energy -= G * p1.mass * p2.mass / distance
+    
+    # Calculate kinetic energy
+    for p in particles:
+        v_squared = p.velocity.x * p.velocity.x + p.velocity.y * p.velocity.y
+        energy += 0.5 * p.mass * v_squared
+    
+    return energy
+
+def advance_system(particles, theta, time_step, width, height):
+    """Advance the n-body system by one time step using the quadtree"""
+    # Create a fresh quadtree
+    boundary = Rectangle(width / 2, height / 2, width / 2, height / 2)
+    quadtree = QuadTree(boundary)
+    
+    # Insert all particles into the quadtree
+    for particle in particles:
+        quadtree.insert(particle)
+    
+    # Calculate and apply forces to all particles
+    for particle in particles:
+        fx, fy = quadtree.calculate_force(particle, theta)
+        particle.apply_force(fx, fy)
+    
+    # Update all particles
+    for particle in particles:
+        particle.update(time_step)
+
+def bench_quadtree_nbody(loops, num_particles, iterations, theta):
+    # Initialize simulation space
+    width = 1000
+    height = 800
+    
+    # Create galaxy distribution
+    particles = create_galaxy_distribution(num_particles, width / 2, height / 2)
+    
+    # Calculate initial energy
+    initial_energy = calculate_system_energy(particles)
+    
+    range_it = range(loops)
+    t0 = pyperf.perf_counter()
+    
+    for _ in range_it:
+        # Run simulation for specified iterations
+        for _ in range(iterations):
+            advance_system(particles, theta, TIME_STEP, width, height)
+        
+        # Calculate final energy
+        final_energy = calculate_system_energy(particles)
+    
+    return pyperf.perf_counter() - t0
+
+def add_cmdline_args(cmd, args):
+    cmd.extend(("--iterations", str(args.iterations)))
+    cmd.extend(("--particles", str(args.particles)))
+    cmd.extend(("--theta", str(args.theta)))
+
+if __name__ == '__main__':
+    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
+    runner.metadata['description'] = "Quadtree N-body benchmark"
+    
+    runner.argparser.add_argument("--iterations",
+                                  type=int, default=DEFAULT_ITERATIONS,
+                                  help="Number of simulation steps per benchmark loop "
+                                       "(default: %s)" % DEFAULT_ITERATIONS)
+    
+    runner.argparser.add_argument("--particles",
+                                  type=int, default=DEFAULT_PARTICLES,
+                                  help="Number of particles in the simulation "
+                                       "(default: %s)" % DEFAULT_PARTICLES)
+    
+    runner.argparser.add_argument("--theta",
+                                  type=float, default=DEFAULT_THETA,
+                                  help="Barnes-Hut approximation threshold "
+                                       "(default: %s)" % DEFAULT_THETA)
+    
+    args = runner.parse_args()
+    runner.bench_time_func('quadtree_nbody', bench_quadtree_nbody,
+                          args.particles, args.iterations, args.theta)