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