Speedup Pymunk collision handling

Question

I want to implement a particle simulation with Python Arcade and Pymunk. While Arcade is very fast even for drawing / handling eg. 10000 particles the collision handling with Pymunk becomes very slow (14 FPS for 3000 particles). I assume that Pymunk is checking every possible particle pairs for collision.

In particle collision simulation usually there is a grid on the window and particle collision is only checked for particles in adjacent grid cells. This is much faster.

EDIT: Here is my code so far:

import arcade
import pymunk
import random

SCREEN_WIDTH = 1200
SCREEN_HEIGHT = 800
PARTICLE_COUNT = 3000
PARTICLE_RADIUS = 4
GRAVITY = (0, -981)

class Particle:
    def __init__(self, space, position, velocity, weight):
        self.body = pymunk.Body()
        self.body.position = position
        self.body.velocity = velocity
        self.shape = pymunk.Circle(self.body, radius=PARTICLE_RADIUS)
        self.shape.density = weight
        self.shape.elasticity = 0.8
        self.color = arcade.color.WHITE
        if weight > 1:
            self.color = arcade.color.RED
        elif weight < 1:
            self.color = arcade.color.BLUE      
        space.add(self.body, self.shape)

    def draw(self):
        x, y = self.body.position
        arcade.draw_circle_filled(x, y, PARTICLE_RADIUS, self.color)

class Simulation(arcade.Window):
    def __init__(self, width, height):
        super().__init__(width, height, "Particle Simulation")
        arcade.set_background_color(arcade.color.BLACK)

        self.space = pymunk.Space()
        self.space.gravity = GRAVITY
        
        wThick = 50
        dv     = 49
        walls = [
            pymunk.Segment(self.space.static_body, (-dv, 0), (-dv, height), wThick),
            pymunk.Segment(self.space.static_body, (0, height+dv), (width, height+dv), wThick),
            pymunk.Segment(self.space.static_body, (width+dv, height), (width+dv, 0), wThick),
            pymunk.Segment(self.space.static_body, (0, 0-dv), (width, 0-dv), wThick)
        ]
        for wall in walls:
            wall.elasticity = 1
            wall.friction = 0.5
            wall.color = arcade.color.RED
        self.space.add(*walls)

        walls = [
            pymunk.Segment(self.space.static_body, (0,height//2-dv), (width//2, 0-dv), wThick),
        ]
        for wall in walls:
            wall.elasticity = 1
            wall.friction = 1
            wall.color = arcade.color.RED
        self.space.add(*walls)

        self.particles = []
        for i in range(PARTICLE_COUNT):
            position = random.uniform(PARTICLE_RADIUS, width//2 - PARTICLE_RADIUS), \
                       random.uniform(height//2+PARTICLE_RADIUS, height - PARTICLE_RADIUS)
            velocity = random.uniform(-200, 200), random.uniform(-200, 200)
            weight = random.uniform(0.1, 2)
            self.particles.append(Particle(self.space, position, velocity, weight))

    def on_draw(self):
        arcade.start_render()
        for particle in self.particles:
            particle.draw()
        # Draw the FPS counter in the bottom-left corner of the window
        fps = arcade.get_fps()
        fps_text = "FPS: %.2f" % (fps)
        arcade.draw_text(fps_text, 10, 10, arcade.color.WHITE, 14)            

    def on_update(self, delta_time):
        self.space.step(1/120)
        

def main():
    sim = Simulation(SCREEN_WIDTH, SCREEN_HEIGHT)
    arcade.enable_timings()
    arcade.run()

if __name__ == '__main__':
    main()

Is it possible to realize this with Pymunk?

Answer 1

It depends on your requirements. Pymunk is not optimized for particle systems, but you can try to use the spatial hash instead of the tree based collision structure for some speedup.

Its very easy to enable, but requires some tweaking. After you create the space you can enable the spatial hash with the use_spatial_hash function.

self.space.use_spatial_hash(PARTICLE_RADIUS*2, 10000)

Both arguments can be tweaked, easiest is to try different values and measure which works best (For example for dim, the first argument, try with PARTICLE_RADIUS, PARTICLE_RADIUS2 and PARTICLE_RADIUS4 and for the second argument with 1000, 5000, 10000, 50000 and 100000).

Note that drawing the particles probably also take a fair amount of processing. To determine if its significant and worth optimizing, try to toggle the drawing off, by commenting out the drawing for loop and again measure the performance. I have done some experiments with a batch API that would help with this case, but could not fully decide of a suitable API or how useful it would be...