Pymunk Physics Engine - Joint Builder

This uses the Pymunk physics engine to simulate items with joints

Screenshot of pegboard example.
pymunk_joint_builder.py
  1"""
  2Pymunk 2
  3
  4If Python and Arcade are installed, this example can be run from the command line with:
  5python -m arcade.examples.pymunk_joint_builder
  6"""
  7import arcade
  8import pymunk
  9import timeit
 10import math
 11
 12WINDOW_WIDTH = 1200
 13WINDOW_HEIGHT = 800
 14WINDOW_TITLE = "Pymunk 2 Example"
 15
 16"""
 17Key bindings:
 18
 191 - Drag mode
 202 - Make box mode
 213 - Make PinJoint mode
 224 - Make DampedSpring mode
 23
 24S - No gravity or friction
 25L - Layout, no gravity, lots of friction
 26G - Gravity, little bit of friction
 27
 28Right-click, fire coin
 29
 30"""
 31
 32
 33class PhysicsSprite(arcade.Sprite):
 34    def __init__(self, pymunk_shape, filename):
 35        super().__init__(
 36            filename,
 37            center_x=pymunk_shape.body.position.x,
 38            center_y=pymunk_shape.body.position.y,
 39        )
 40        self.pymunk_shape = pymunk_shape
 41
 42
 43class CircleSprite(PhysicsSprite):
 44    def __init__(self, pymunk_shape, filename):
 45        super().__init__(pymunk_shape, filename)
 46        self.width = pymunk_shape.radius * 4
 47        self.height = pymunk_shape.radius * 4
 48
 49
 50class BoxSprite(PhysicsSprite):
 51    def __init__(self, pymunk_shape, filename, width, height):
 52        super().__init__(pymunk_shape, filename)
 53        self.width = width
 54        self.height = height
 55
 56
 57class GameView(arcade.View):
 58    """ Main application class. """
 59
 60    def __init__(self):
 61        super().__init__()
 62
 63        self.background_color = arcade.color.DARK_SLATE_GRAY
 64
 65        # -- Pymunk
 66        self.space = pymunk.Space()
 67        self.space.gravity = (0.0, -900.0)
 68
 69        # Lists of sprites or lines
 70        self.sprite_list: arcade.SpriteList[PhysicsSprite] = arcade.SpriteList()
 71        self.static_lines = []
 72
 73        # Used for dragging shapes around with the mouse
 74        self.shape_being_dragged = None
 75        self.last_mouse_position = 0, 0
 76
 77        self.processing_time_text = None
 78        self.draw_time_text = None
 79        self.draw_mode_text = None
 80        self.shape_1 = None
 81        self.shape_2 = None
 82        self.draw_time = 0
 83        self.processing_time = 0
 84        self.joints = []
 85
 86        self.physics = "Normal"
 87        self.mode = "Make Box"
 88
 89        # Create the floor
 90        self.floor_height = 80
 91        body = pymunk.Body(body_type=pymunk.Body.STATIC)
 92        shape = pymunk.Segment(
 93            body,
 94            (0, self.floor_height),
 95            (WINDOW_WIDTH, self.floor_height),
 96            0.0,
 97        )
 98        shape.friction = 10
 99        self.space.add(shape, body)
100        self.static_lines.append(shape)
101
102    def on_draw(self):
103        """
104        Render the screen.
105        """
106
107        # This command has to happen before we start drawing
108        self.clear()
109
110        # Start timing how long this takes
111        draw_start_time = timeit.default_timer()
112
113        # Draw all the sprites
114        self.sprite_list.draw()
115
116        # Draw the lines that aren't sprites
117        for line in self.static_lines:
118            body = line.body
119
120            pv1 = body.position + line.a.rotated(body.angle)
121            pv2 = body.position + line.b.rotated(body.angle)
122            arcade.draw_line(pv1.x, pv1.y, pv2.x, pv2.y, arcade.color.WHITE, 2)
123
124        for joint in self.joints:
125            color = arcade.color.WHITE
126            if isinstance(joint, pymunk.DampedSpring):
127                color = arcade.color.DARK_GREEN
128            arcade.draw_line(joint.a.position.x,
129                             joint.a.position.y,
130                             joint.b.position.x,
131                             joint.b.position.y,
132                             color, 3)
133
134        # arcade.draw_text(output, 10, 20, arcade.color.WHITE, 14)
135        # Display timings
136        output = f"Processing time: {self.processing_time:.3f}"
137        arcade.draw_text(output, 20, WINDOW_HEIGHT - 20, arcade.color.WHITE)
138
139        output = f"Drawing time: {self.draw_time:.3f}"
140        arcade.draw_text(output, 20, WINDOW_HEIGHT - 40, arcade.color.WHITE)
141
142        self.draw_time = timeit.default_timer() - draw_start_time
143
144        output = f"Mode: {self.mode}"
145        arcade.draw_text(output, 20, WINDOW_HEIGHT - 60, arcade.color.WHITE)
146
147        output = f"Physics: {self.physics}"
148        arcade.draw_text(output, 20, WINDOW_HEIGHT - 80, arcade.color.WHITE)
149
150    def make_box(self, x, y):
151        size = 45
152        mass = 12.0
153        moment = pymunk.moment_for_box(mass, (size, size))
154        body = pymunk.Body(mass, moment)
155        body.position = pymunk.Vec2d(x, y)
156        shape = pymunk.Poly.create_box(body, (size, size))
157        shape.friction = 0.3
158        self.space.add(body, shape)
159
160        sprite = BoxSprite(
161            shape,
162            ":resources:images/tiles/boxCrate_double.png",
163            width=size,
164            height=size,
165        )
166        self.sprite_list.append(sprite)
167
168    def make_circle(self, x, y):
169        size = 20
170        mass = 12.0
171        moment = pymunk.moment_for_circle(mass, 0, size, (0, 0))
172        body = pymunk.Body(mass, moment)
173        body.position = pymunk.Vec2d(x, y)
174        shape = pymunk.Circle(body, size, pymunk.Vec2d(0, 0))
175        shape.friction = 0.3
176        self.space.add(body, shape)
177
178        sprite = CircleSprite(shape, ":resources:images/items/coinGold.png")
179        self.sprite_list.append(sprite)
180
181    def make_pin_joint(self, x, y):
182        shape_selected = self.get_shape(x, y)
183        if shape_selected is None:
184            return
185
186        if self.shape_1 is None:
187            print("Shape 1 Selected")
188            self.shape_1 = shape_selected
189        elif self.shape_2 is None and self.shape_1.shape != shape_selected.shape:
190            print("Shape 2 Selected")
191            self.shape_2 = shape_selected
192            joint = pymunk.PinJoint(self.shape_1.shape.body, self.shape_2.shape.body)
193            self.space.add(joint)
194            self.joints.append(joint)
195            self.shape_1 = None
196            self.shape_2 = None
197            print("Joint Made")
198        else:
199            print("Shapes Deselected")
200            self.shape_1 = self.shape_2 = None
201
202    def make_damped_spring(self, x, y):
203        shape_selected = self.get_shape(x, y)
204        if shape_selected is None:
205            return
206
207        if self.shape_1 is None:
208            print("Shape 1 Selected")
209            self.shape_1 = shape_selected
210        elif self.shape_2 is None and self.shape_1.shape != shape_selected.shape:
211            print("Shape 2 Selected")
212            self.shape_2 = shape_selected
213            joint = pymunk.DampedSpring(
214                self.shape_1.shape.body,
215                self.shape_2.shape.body,
216                (0, 0), (0, 0),
217                45, 300, 30,
218            )
219            self.space.add(joint)
220            self.joints.append(joint)
221            self.shape_1 = None
222            self.shape_2 = None
223            print("Joint Made")
224        else:
225            print("Shapes deselected")
226            self.shape_1 = self.shape_2 = None
227
228    def get_shape(self, x, y):
229        # See if we clicked on anything
230        shape_list = self.space.point_query((x, y), 1, pymunk.ShapeFilter())
231
232        # If we did, remember what we clicked on
233        if len(shape_list) > 0:
234            shape = shape_list[0]
235        else:
236            shape = None
237
238        return shape
239
240    def on_mouse_press(self, x, y, button, modifiers):
241
242        if button == 1 and self.mode == "Drag":
243            self.last_mouse_position = x, y
244            self.shape_being_dragged = self.get_shape(x, y)
245
246        elif button == 1 and self.mode == "Make Box":
247            self.make_box(x, y)
248
249        elif button == 1 and self.mode == "Make Circle":
250            self.make_circle(x, y)
251
252        elif button == 1 and self.mode == "Make PinJoint":
253            self.make_pin_joint(x, y)
254
255        elif button == 1 and self.mode == "Make DampedSpring":
256            self.make_damped_spring(x, y)
257
258        elif button == 4:
259            # With right mouse button, shoot a heavy coin fast.
260            mass = 60
261            radius = 10
262            inertia = pymunk.moment_for_circle(mass, 0, radius, (0, 0))
263            body = pymunk.Body(mass, inertia)
264            body.position = x, y
265            body.velocity = 2000, 0
266            shape = pymunk.Circle(body, radius, pymunk.Vec2d(0, 0))
267            shape.friction = 0.3
268            self.space.add(body, shape)
269
270            sprite = CircleSprite(shape, ":resources:images/items/coinGold.png")
271            self.sprite_list.append(sprite)
272
273    def on_mouse_release(self, x, y, button, modifiers):
274        if button == 1:
275            # Release the item we are holding (if any)
276            self.shape_being_dragged = None
277
278    def on_mouse_motion(self, x, y, dx, dy):
279        if self.shape_being_dragged is not None:
280            # If we are holding an object, move it with the mouse
281            self.last_mouse_position = x, y
282            self.shape_being_dragged.shape.body.position = self.last_mouse_position
283            self.shape_being_dragged.shape.body.velocity = dx * 20, dy * 20
284
285    def on_key_press(self, symbol: int, modifiers: int):
286        if symbol == arcade.key.KEY_1:
287            self.mode = "Drag"
288        elif symbol == arcade.key.KEY_2:
289            self.mode = "Make Box"
290        elif symbol == arcade.key.KEY_3:
291            self.mode = "Make Circle"
292
293        elif symbol == arcade.key.KEY_4:
294            self.mode = "Make PinJoint"
295        elif symbol == arcade.key.KEY_5:
296            self.mode = "Make DampedSpring"
297
298        elif symbol == arcade.key.S:
299            self.space.gravity = (0.0, 0.0)
300            self.space.damping = 1
301            self.physics = "Outer Space"
302        elif symbol == arcade.key.L:
303            self.space.gravity = (0.0, 0.0)
304            self.space.damping = 0
305            self.physics = "Layout"
306        elif symbol == arcade.key.G:
307            self.space.damping = 0.95
308            self.space.gravity = (0.0, -900.0)
309            self.physics = "Normal"
310
311    def on_update(self, delta_time):
312        start_time = timeit.default_timer()
313
314        # Check for balls that fall off the screen
315        for sprite in self.sprite_list:
316            if sprite.pymunk_shape.body.position.y < 0:
317                # Remove balls from physics space
318                self.space.remove(sprite.pymunk_shape, sprite.pymunk_shape.body)
319                # Remove balls from physics list
320                sprite.kill()
321
322        # Update physics
323        self.space.step(1 / 80.0)
324
325        # If we are dragging an object, make sure it stays with the mouse. Otherwise
326        # gravity will drag it down.
327        if self.shape_being_dragged is not None:
328            self.shape_being_dragged.shape.body.position = self.last_mouse_position
329            self.shape_being_dragged.shape.body.velocity = 0, 0
330
331        # Move sprites to where physics objects are
332        for sprite in self.sprite_list:
333            sprite.center_x = sprite.pymunk_shape.body.position.x
334            sprite.center_y = sprite.pymunk_shape.body.position.y
335            # Reverse angle because pymunk rotates ccw
336            sprite.angle = -math.degrees(sprite.pymunk_shape.body.angle)
337
338        # Save the time it took to do this.
339        self.processing_time = timeit.default_timer() - start_time
340
341
342def main():
343    """ Main function """
344    # Create a window class. This is what actually shows up on screen
345    window = arcade.Window(WINDOW_WIDTH, WINDOW_HEIGHT, WINDOW_TITLE)
346
347    # Create the GameView
348    game = GameView()
349
350    # Show GameView on screen
351    window.show_view(game)
352
353    # Start the arcade game loop
354    arcade.run()
355
356
357if __name__ == "__main__":
358    main()