More pymunk explorations

[villares]

I think that the first time I saw pymunk was thanks to some post by @tabreturn. So much fun!

From time to time I try to put my students to use it, but there is a lot of complexity to deal with. Along the years I have always tought of ways of making it simpler to draw the simulations, but nothing really worked very well. Talking to @ramalho, he suggested a "wrapper class" approach, and I now think I have something of a proof of concept!

I even tried my hand at a tiny little bit of meta-programming to help add factory methods ... that stuff might be"hidden" if this becomes a library. Maybe using match/case on the main sketch part could be a bit excessive for beginners, I'm not sure. This is work in progress.

This depends on pymunk, shapely, trimesh and trimesh's "optional dependency" mapbox_earcut.
If python -m pip install pymunk trimesh[easy] doesn't work, then try python -m pip install trimesh shapely mapbox_earcut.

# This is a py5 "module mode" sketch
# learn about py5 modes at https://py5coding.org

"""
Press " " (space bar) to create many balls
Press "w" to drag the mouse and create a static segment (wall)
Press "b" to drag the mouse and create a box
Press "k" to drag the mouse and create a kinematic box
Press "p" to drag the mouse and create a polygon
Press <delete> or <backspace> to remove objects under the mouse
Press "c" to remove dynamic and kinematic objects
You can drag kinematic objects (and nudge a bit the dynamic ones)
You can use the mouse wheel to rotate some objects
"""

import py5
import pymunk
import shapely
from trimesh.creation import triangulate_polygon # pip install trimesh[easy]

ongoing_creation = None

def setup():
    py5.size(600, 600)
    global sim
    sim = Simulation()
    sim.add_segment(50, 500, 550, 500) # could be just Segment(...)

def draw():
    py5.background(200)
    sim.draw_and_update()  # draw and clean up
    # create many balls
    if py5.is_key_pressed and py5.key == " ":
        d = py5.random(10, 50)
        c = py5.color(d * 5, 0, 255 - d * 5)
        Ball(py5.mouse_x + py5.random(-1, 1), py5.mouse_y, d, c)
    # preview & object creation
    py5.no_fill()
    py5.stroke_weight(2)
    py5.stroke(255)
    x, y = py5.mouse_x, py5.mouse_y
    match ongoing_creation:
        case pts, "p":
            # poly creation ongoing!
            with py5.begin_shape():
                py5.vertices(pts)
                py5.vertex(x, y)
            if (len(pts) == 0) or py5.dist(x, y, *pts[-1]) > 20:
                pts.append((x, y))
        case sx, sy, "w":
            with py5.push_style(), py5.begin_shape():
                py5.line(sx, sy, x, y)
        case sx, sy, "b" | "k":
            py5.rect_mode(py5.CORNERS)
            py5.rect(sx, sy, x, y)
    # advance simulation
    sim.step(1 / 60)

def key_pressed():
    if py5.key in (py5.BACKSPACE, py5.DELETE):
        for obj in sim:
            if obj.under_mouse():
                obj.remove_from_sim()
    elif py5.key == "c":
        for obj in sim:
            if obj.body is not obj.space.static_body:
                obj.remove_from_sim() 

def mouse_pressed():
    global ongoing_creation
    if py5.key == "k" and py5.is_key_pressed:
        ongoing_creation = (py5.mouse_x, py5.mouse_y, "k")
    elif py5.key == "b" and py5.is_key_pressed:
        ongoing_creation = (py5.mouse_x, py5.mouse_y, "c")
    elif py5.key == "w" and py5.is_key_pressed:
        ongoing_creation = (py5.mouse_x, py5.mouse_y, "w")
    elif py5.key == "p" and py5.is_key_pressed:
        ongoing_creation = ([], "p")

def mouse_released():
    global ongoing_creation
    mx, my = py5.mouse_x, py5.mouse_y
    match ongoing_creation:
        case sx, sy, "w":
            Segment(sx, sy, mx, my)
        case sx, sy, box_type:
            w, h = abs(sx - mx), abs(sy - my)
            x, y = (sx + mx) / 2, (sy + my) / 2
            Box(x, y, w, h, kinematic=(box_type == "k"))
        case pts, "p":
            shapely_poly = shapely.Polygon(pts) if len(pts) >= 3 else shapely.Polygon()
            if shapely_poly.area > 100 and shapely_poly.is_simple:
                Poly(shapely_poly, fill_color=py5.color(255))
    ongoing_creation = None

def mouse_dragged():
    if not py5.is_key_pressed:
        for obj in sim:
            if obj.under_mouse():
                obj.translate(py5.mouse_x - py5.pmouse_x, py5.mouse_y - py5.pmouse_y)

def mouse_wheel(e):
    for obj in sim:
        if obj.under_mouse():
            obj.rotate(py5.radians(e.get_count()))


class Simulation:
    _default = None
    mass_scale = 0.1
    poly_friction = 0.99
    box_friction = 0.99
    segment_friction = 0.99

    def __init__(self, name=None, gravity=(0, 900)):
        self.name = name
        self.space = pymunk.Space()
        self.space.gravity = gravity
        # object registry management
        self.living_set = set()
        self.for_removal = set()
        # if no name is provided, make this the default simulation
        if name is None:
            Simulation._default = self

    def __iter__(self):
        """To traverse living objects."""
        return iter(self.living_set)

    @classmethod
    def get_default(cls):
        """Get or create the default simulation."""
        if cls._default is None:
            cls._default = Simulation()  # Creates unnamed default
        return cls._default

    def step(self, s):
        """Advance the simulation step."""
        self.space.step(s)

    def draw_and_update(self):
        """Draw all objects and remove marked ones."""
        for obj in self.living_set:
            obj.draw()
        self.living_set.difference_update(self.for_removal)
        self.for_removal.clear()

    @classmethod
    def add_factory(cls, sobj_cls):
        """Decorate a class to add a factory method"""
        import inspect
        method_name = f"add_{sobj_cls.__name__.lower()}"
       
        def factory(self, *args, **kwargs):
            return sobj_cls(*args, **kwargs, simulation=self)
        
        sig = inspect.signature(sobj_cls.__init__)
        factory.__signature__ = inspect.Signature([
            p for name, p in sig.parameters.items() 
            if name != 'simulation'
        ])
        factory.__name__ = method_name
        factory.__doc__ = sobj_cls.__doc__
        setattr(cls, method_name, factory)
        return sobj_cls


class SObj:

    def __init__(self, simulation=None):
        self.simulation = simulation or Simulation.get_default()
        self.space = self.simulation.space

    def register_object(self):
        """Add body and shape to space, and register for drawing."""
        if not hasattr(self, "shapes"):
            self.shapes = (self.shape,)
        if self.body is not self.space.static_body:
            self.space.add(self.body, *self.shapes)
        else:
            self.space.add(*self.shapes)
        self.simulation.living_set.add(self)

    def remove_from_sim(self):
        """Remove from space and mark for removal from drawing registry."""
        if self.body is not self.space.static_body:
            self.space.remove(self.body, *self.shapes)
        else:
            self.space.remove(*self.shapes)
        self.simulation.for_removal.add(self)

    def rotate(self, rot):
        self.body.angle += rot

    def translate(self, dx, dy):
        self.body.position += pymunk.Vec2d(dx, dy)

    def under_mouse(self):
        if hasattr(self, "poly"):
            translated_mouse = shapely.Point(
                (py5.mouse_x - self.body.position.x,
                 py5.mouse_y - self.body.position.y))
            return self.poly.contains(translated_mouse)
        else:
            info = self.shape.point_query((py5.mouse_x, py5.mouse_y))
            return info.distance < 2

    def draw(self):
        raise NotImplementedError # must be defined in subclasses.


@Simulation.add_factory
class Segment(SObj):
    """Static wall-like line segment."""
    def __init__(self, xa, ya, xb, yb, radius=2, simulation=None):
        super().__init__(simulation)
        self.thickness = radius * 2
        self.body = self.space.static_body
        self.shape = pymunk.Segment(self.body, (xa, ya), (xb, yb), radius=radius)
        self.shape.friction = self.simulation.segment_friction
        self.register_object()

    def draw(self):
        py5.stroke_weight(self.thickness)
        py5.stroke(0)
        py5.line(self.shape.a.x, self.shape.a.y, self.shape.b.x, self.shape.b.y)

@Simulation.add_factory
class Ball(SObj):
    """A color-filled ball."""
    def __init__(self, x, y, diameter, fill_color=None, simulation=None):
        super().__init__(simulation)
        self.diameter = diameter
        radius = diameter / 2
        self.fill_color = fill_color or py5.color(255)
        mass = py5.PI * radius**2 * self.simulation.mass_scale
        moment = pymunk.moment_for_circle(mass, 0, radius)
        self.body = pymunk.Body(mass, moment)
        self.body.position = x, y
        self.shape = pymunk.Circle(self.body, radius, (0, 0))
        self.register_object()

    def draw(self):
        py5.no_stroke()
        py5.fill(self.fill_color)
        py5.circle(self.body.position.x, self.body.position.y, self.shape.radius * 2)
        if self.body.position.y > py5.height + self.diameter:
            self.remove_from_sim()

@Simulation.add_factory
class Poly(SObj):
    """A simulated polygon-shaped object, from a shapely.Polygon. Kinematic or Dynamic."""
    def __init__(self, poly: shapely.Polygon, fill_color=None, kinematic=False, simulation=None):
        super().__init__(simulation)
        cx, cy = poly.centroid.x, poly.centroid.y
        self.poly = shapely.affinity.translate(poly, -cx, -cy)
        if kinematic:
            self.body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
        else:
            mass = poly.area * self.simulation.mass_scale
            moment = pymunk.moment_for_poly(mass, self.poly.exterior.coords, (0, 0))
            self.body = pymunk.Body(mass, moment)
        self.body.position = (cx, cy)
        self.shapes = []
        vs, faces = triangulate_polygon(poly)
        triangles = [(vs[a], vs[b], vs[c]) for a, b, c in faces]
        for tri in triangles:
            translated_tri = [(x - cx, y - cy) for x, y in tri]
            shape = pymunk.Poly(self.body, translated_tri)
            shape.friction = self.simulation.poly_friction
            self.shapes.append(shape)
        self.fill_color = fill_color or py5.color(255)
        self.register_object()

    def draw(self):
        py5.fill(self.fill_color)
        py5.no_stroke()
        with py5.push_matrix():
            py5.translate(self.body.position.x, self.body.position.y)
            py5.rotate(self.body.angle)
            with py5.begin_closed_shape():
                py5.vertices(self.poly.exterior.coords)
        if self.body.position.y > py5.height + 200:
            self.remove_from_sim()

@Simulation.add_factory
class Box(SObj):
    """A rectangular polygon. Kinematic or Dynamic."""
    def __init__(self, x, y, w, h, fill_color=None, kinematic=True, simulation=None):
        super().__init__(simulation)
        if kinematic:
            self.body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
        else:
            mass = w * h * self.simulation.mass_scale
            moment = pymunk.moment_for_box(mass, (w, h))
            self.body = pymunk.Body(mass, moment)
        self.body.position = x, y
        self.shape = pymunk.Poly.create_box(self.body, (w, h))
        self.shape.friction = 0.80
        self.vertices = self.shape.get_vertices()
        if fill_color is None:
            self.fill_color = (
                py5.color(100, 0, 0) if kinematic else py5.color(0, 100, 0)
            )
        else:
            self.fill_color = fill_color
        self.register_object()

    def draw(self):
        py5.fill(self.fill_color)
        py5.no_stroke()
        with py5.push_matrix():
            py5.translate(self.body.position.x, self.body.position.y)
            py5.rotate(self.body.angle)
            with py5.begin_closed_shape():
                py5.vertices(self.vertices)
        if self.body.position.y > py5.height + 200:
            self.remove_from_sim()


py5.run_sketch(block=False)

[villares]

@hx2A I have tried moving the classes to another file for importing and something weird is going on... now I'm feeling very stupid. I'll see if I can prepare a minimal example to reproduce the problem.
Update: Maybe a fluke? Suddenly it works again... 🥳
Another update: I removed the meta-programming thing for .add_ methods, as it breaks Thonny auto-complete 😕
The following file should work both with module mode and imported mode ☺️

EDIT: Removed outdated code ... see comments bellow for a more recent version. 

[villares]

I have to collect and refactor some other pymunk sketches that I have over the years:

• Experiments with "joints", some are real fun with "soft bodies" made from connecting a grid of circles with "springs".
• Code to tackle shapely.MultiPolygons an and polygons with holes (that I used to play with typographical glyphs + physics).
• Contact listeners, like the ones used on this Pinball simulator https://github.com/villares/pymunk-pinball-paulista

[hx2A]

Wow you've written a lot of code here. Perhaps this can be packaged up as a pymonk submodule? I have seen on their documentation website they have submodules for pygame and pyglet.

https://www.pymunk.org/en/latest/pymunk.html

[villares]

Yeah, I saw the pygame and pyglet utils and thought of porting one of them to py5... that could be a useful contribution too... but expanding this wrapper, for now, seems more useful: I even tried it in class Tuesday and people seemed to like it.

[hx2A]

Another update: I removed the meta-programming thing for .add_ methods, as it breaks Thonny auto-complete 😕

Hmmm, that's very strange that auto-complete would be impacted that way. I know little about meta programming. py5 does a little bit of that with the__new__() methods in classes, and that part took a lot of thought to get right.

[villares]

So, I'm still working on the importable pymunk_helpers.py idea:
https://github.com/villares/sketch-a-day/blob/main/2026/sketch_2026_04_11/pymunk_helpers.py

Today I struggled to add two constraints/joints, PivotJoint and PinJoint.

From https://github.com/villares/sketch-a-day/blob/main/2026/sketch_2026_04_11/ :

# This is a py5 "module mode" sketch
# learn about py5 modes at https://py5coding.org

"""
Press " " (space bar) to create many balls
Press "w" to drag the mouse and create a static segment (wall)
Press "b" to drag the mouse and create a box
Press "k" to drag the mouse and create a kinematic box
Press "p" to drag the mouse and create a polygon
Press <delete> or <backspace> to remove objects under the mouse
Press "c" to remove dynamic objects
You can drag kinematic objects (and nudge a bit the dynamic ones)
You can use the mouse wheel to rotate some objects
"""

import py5
import shapely
from pymunk_helpers import Simulation, SObj, Constraint
from pymunk_helpers import Segment, Ball, Box, Poly, PivotJoint, PinJoint

ongoing_creation = None

def setup():
    py5.size(600, 600)
    global sim
    sim = Simulation(gravity=(0,0))
#
    Segment(50, 550, 550, 550, static=True)
    Segment(50, 50, 550, 50, static=True)
    Segment(50, 50, 50, 550, fill_color=128, kinematic=True)
    Segment(550, 50, 550, 550, fill_color=128, kinematic=True)
    ba = Ball(100, 100, 50)
    bb = Ball(200, 50, 100)
    PinJoint(ba, bb, anchor_a=(100, 120))
    sa = Segment(300, 100, 400, 100, fill_color=255, radius=20)
    sb = Segment(400, 100, 500, 200, fill_color=255, radius=20)
    PivotJoint(sa, sb, (400, 100))


def draw():
    py5.background(200)
    sim.draw_and_update(step=1/60)  # draw and clean up
    # create many balls
    if py5.is_key_pressed and py5.key == " ":
        d = py5.random(10, 50)
        c = py5.color(d * 5, 0, 255 - d * 5)
        Ball(py5.mouse_x + py5.random(-1, 1), py5.mouse_y, d, c)
    # preview & object creation
    py5.no_fill()
    py5.stroke_weight(2)
    py5.stroke(255)
    x, y = py5.mouse_x, py5.mouse_y
    match ongoing_creation:
        case pts, "p":
            # poly creation ongoing!
            with py5.begin_shape():
                py5.vertices(pts)
                py5.vertex(x, y)
            if (len(pts) == 0) or py5.dist(x, y, *pts[-1]) > 20:
                pts.append((x, y))
        case sx, sy, "w":
            with py5.push_style(), py5.begin_shape():
                py5.line(sx, sy, x, y)
        case sx, sy, "b" | "k":
            py5.rect_mode(py5.CORNERS)
            py5.rect(sx, sy, x, y)

def key_pressed():
    if py5.key in (py5.BACKSPACE, py5.DELETE):
        for obj in sim:
            if obj.under_mouse():
                obj.remove_from_sim()
    elif py5.key == "c":
        for obj in sim:
            if isinstance(obj, Constraint) or obj.body.body_type not in (sim.STATIC, sim.KINEMATIC):
                obj.remove_from_sim() 
    elif py5.key == 's':
        py5.save_frame('###.png')

def mouse_pressed():
    global ongoing_creation
    if py5.key == "k" and py5.is_key_pressed:
        ongoing_creation = (py5.mouse_x, py5.mouse_y, "k")
    elif py5.key == "b" and py5.is_key_pressed:
        ongoing_creation = (py5.mouse_x, py5.mouse_y, "c")
    elif py5.key == "W" and py5.is_key_pressed:
            ongoing_creation = (py5.mouse_x, py5.mouse_y, "W")
    elif py5.key == "w" and py5.is_key_pressed:
            ongoing_creation = (py5.mouse_x, py5.mouse_y, "w")
    elif py5.key == "p" and py5.is_key_pressed:
        ongoing_creation = ([], "p")

def mouse_released():
    global ongoing_creation
    mx, my = py5.mouse_x, py5.mouse_y
    match ongoing_creation:
        case sx, sy, "w":
            sim.add_kinematic_segment(sx, sy, mx, my)
        case sx, sy, "W":
            sim.add_segment(sx, sy, mx, my, radius=20)
        case sx, sy, box_type:
            w, h = abs(sx - mx), abs(sy - my)
            x, y = (sx + mx) / 2, (sy + my) / 2
            Box(x, y, w, h, kinematic=(box_type == "k"))
        case pts, "p":
            shapely_poly = shapely.Polygon(pts) if len(pts) >= 3 else shapely.Polygon()
            if shapely_poly.area > 100 and shapely_poly.is_simple:
                Poly(shapely_poly, fill_color=py5.color(255))
    ongoing_creation = None

def mouse_dragged():
    dx, dy = py5.mouse_x - py5.pmouse_x, py5.mouse_y - py5.pmouse_y
    if not py5.is_key_pressed:
        for obj in sim:
            if obj.under_mouse():
                if obj.type == sim.KINEMATIC:
                    obj.translate(dx, dy)
                if obj.type == sim.DYNAMIC:
                    obj.vel_update(dx, dy)
                    
def mouse_wheel(e):
    for obj in sim:
        if obj.under_mouse() and isinstance(obj, SObj):
            obj.rotate(py5.radians(e.get_count()))

py5.run_sketch(block=False)

[villares]

Aaaand... I tried monkey patching Py5Color now 😄

Seems to work!

from py5 import Py5Color

def _reconstruct_py5color(color_as_int):
    return py5.color(color_as_int)

def __reduce__(self):
    return (
        Py5Color._reconstruct_py5color,
        (int(self),)
    )

Py5Color._reconstruct_py5color = _reconstruct_py5color
Py5Color.__reduce__ = __reduce__

Update: Looks like you don't even need __getstate__!
Would you consider making Py5Color picklable adding to it something similar to this?

[villares]

you don't even need __getstate__ in fact...

[hx2A]

Yes, actually, this is a good idea. I will open an issue for it.

[hx2A]

Wait, the reason why it maintains a link to the creator Sketch is because the Sketch's color mode can change, which I think will change the __repr__ and __str__ outputs. But this isn't necessary. But I want to think more about the implementation.

[hx2A]

#787

[villares]

Today I updated the Poly class so it can draw the shapely.Polygons with holes, and even take shapely.MultiPolygons :)

from https://github.com/villares/sketch-a-day/tree/main/2026/sketch_2026_04_12

class Poly(SObj):
    """
    A simulated polygonal, Kinematic or Dynamic, custom object.
    It can be created from a sequence of vertices, or shapely [Multi]Polygons.
    """
    def __init__(self, geometry:
                     Sequence[Sequence] | shapely.Polygon | shapely.MultiPolygon,
                 fill_color=None,
                 kinematic=False,
                 simulation=None):
        super().__init__(simulation)
        if isinstance(geometry, shapely.Polygon):
            cx, cy = geometry.centroid.x, geometry.centroid.y
            geometry = shapely.affinity.translate(geometry, -cx, -cy)
            polygons = (geometry,)
        elif isinstance(geometry, shapely.MultiPolygon):
            cx, cy = geometry.centroid.x, geometry.centroid.y
            geometry = shapely.affinity.translate(geometry, -cx, -cy)
            polygons = tuple(geometry.geoms)
        else:
            geometry = shapely.Polygon(geometry)
            cx, cy = geometry.centroid.x, geometry.centroid.y
            geometry = shapely.affinity.translate(geometry, -cx, -cy)
            polygons = (geometry,)
        self.fill_color = fill_color or py5.color(255)
        self.geometry = geometry
        self.py5shape = py5.convert_shape(geometry)
        self.py5shape.disable_style()
        self.shapes = []
        triangles = [] # the translated to origin triangles
        for poly in polygons:
            vs, faces = triangulate_polygon(poly)
            triangles.extend((tuple(vs[a]), tuple(vs[b]), tuple(vs[c]))
                             for a, b, c in faces)
        if kinematic:
            self.body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
        else:
            mass = geometry.area * self.simulation.mass_scale
            moment = sum(
                pymunk.moment_for_poly(mass / len(triangles), tri, (0, 0))
                for tri in triangles
            )
            self.body = pymunk.Body(mass, moment)       
        for tri in triangles:
            shape = pymunk.Poly(self.body, tri)
            shape.friction = self.friction
            shape.elasticity = self.elasticity
            self.shapes.append(shape)
        self.body.position = (cx, cy)
        self.register_object()

    def draw(self):
        s = self.simulation.current_sketch
        s.fill(self.fill_color)
        s.no_stroke()
        if self.py5shape is None:
            self.py5shape = py5.convert_shape(self.geometry)        
        with s.push_matrix():
            s.translate(self.body.position.x, self.body.position.y)
            s.rotate(self.body.angle)
            s.shape(self.py5shape)
        # trigger removal of falling objects... might be improved
        if self.body.position.y > s.height + 200:
            self.remove_from_sim()

[hx2A]

Nice, this is useful!