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