import pygame import sys import math # ------------------------------- # Pygame and Global Setup # ------------------------------- pygame.init() screen_width, screen_height = 1000, 600 screen = pygame.display.set_mode((screen_width, screen_height)) pygame.display.set_caption("Mars Rover Suspension Simulation") clock = pygame.time.Clock() # ------------------------------- # Colors and Constants # ------------------------------- WHITE = (255, 255, 255) BLACK = (0, 0, 0) GRAY = (200, 200, 200) DARKGRAY = (50, 50, 50) GREEN = (80, 200, 80) BLUE = (100, 100, 255) # Ground control point spacing and scrolling speed. SEGMENT_SPACING = 1 # Horizontal spacing between control points SCROLL_SPEED = 1 # Pixels per frame # ------------------------------- # Ground: Dynamic "Plot" # ------------------------------- # We keep a list of control points in screen coordinates. ground_points = [] # Initialize control points so the ground covers a bit more than the screen width. num_points = screen_width // SEGMENT_SPACING + 3 start_x = -SEGMENT_SPACING for i in range(num_points): # For the initial ground, set a default height (for example, 450) ground_points.append([start_x + i * SEGMENT_SPACING, 450]) def update_ground(mouse_y): """ Update ground control points: - Shift all points left by SCROLL_SPEED. - If the leftmost point is off-screen, remove it. - Append a new point at the right end with x determined by last point + SEGMENT_SPACING and y given by current mouse_y. """ # Shift every point to the left. for pt in ground_points: pt[0] -= SCROLL_SPEED # Remove points that have moved off-screen. if ground_points[0][0] < -SEGMENT_SPACING: ground_points.pop(0) # Add a new control point at the right if needed. if ground_points[-1][0] < screen_width + SEGMENT_SPACING: new_x = ground_points[-1][0] + SEGMENT_SPACING # Use the mouse's y (clamped to a reasonable range) for the new control point. new_y = max(300, min(mouse_y, screen_height - 50)) ground_points.append([new_x, new_y]) def get_ground_y(x): """ Given an x coordinate (in screen space), interpolate the ground's y value using the control points. """ # Find the segment that contains x. for i in range(len(ground_points) - 1): x0, y0 = ground_points[i] x1, y1 = ground_points[i+1] if x0 <= x <= x1: t = (x - x0) / (x1 - x0) return y0 * (1 - t) + y1 * t # Fallback: if x is beyond our control points. return ground_points[-1][1] def draw_ground(surface): """Draw the ground as a polyline connecting all control points, then fill below.""" pts = [(pt[0], pt[1]) for pt in ground_points if -SEGMENT_SPACING <= pt[0] <= screen_width + SEGMENT_SPACING] if len(pts) >= 2: pygame.draw.lines(surface, GREEN, False, pts, 4) # Fill ground below the line. fill_points = pts.copy() fill_points.append((pts[-1][0], screen_height)) fill_points.append((pts[0][0], screen_height)) pygame.draw.polygon(surface, (200,255,200), fill_points) # ------------------------------- # Rover Suspension Simulation # ------------------------------- # The rover is fixed horizontally at the center. rover_x = screen_width // 2 # Wheel horizontal offsets relative to rover center. wheel_offsets = { 'left': -50, 'center': 0, 'right': 50 } wheel_radius = 20 # Suspension parameters (vertical springs). L_rest = 60 # Resting (ideal) suspension length. k = 0.08 # Spring constant. damping = 0.1 # Damping coefficient. # For each wheel, we simulate the vertical position of the chassis attachment (the top of the suspension). # We store for each wheel its current attachment y and vertical velocity. suspension = { 'left': 450 - L_rest, 'center': 450 - L_rest, 'right': 450 - L_rest } suspension_vel = {'left': 0, 'center': 0, 'right': 0} def update_suspension(): """ Update suspension for each wheel. For a given wheel, the desired chassis attachment y is: ground_y(wheel_x) - L_rest We then update the current attachment using a spring–damper simulation. """ for key, offset in wheel_offsets.items(): # World x coordinate for the wheel. x_pos = rover_x + offset # The ground contact y at the wheel's x. ground_contact = get_ground_y(x_pos) desired_attach = ground_contact - L_rest # Current state. attach_y = suspension[key] vel = suspension_vel[key] # Compute spring force. force = k * (desired_attach - attach_y) # Acceleration (assume mass = 1 for simplicity). acc = force - damping * vel # Update velocity and position. vel += acc attach_y += vel suspension[key] = attach_y suspension_vel[key] = vel def draw_spring(surface, start, end, coils=6, amplitude=8): """ Draws a zigzag spring between start and end. """ sx, sy = start ex, ey = end total_dist = math.hypot(ex - sx, ey - sy) # Define number of segments. num_points = coils * 2 + 1 points = [] for i in range(num_points): t = i / (num_points - 1) # Linear interpolation for y. x = sx + t * (ex - sx) y = sy + t * (ey - sy) if i != 0 and i != num_points - 1: offset = amplitude if i % 2 == 0 else -amplitude # Offset perpendicular to the line (for vertical line, horizontal offset; here assume nearly vertical) x += offset points.append((x, y)) pygame.draw.lines(surface, DARKGRAY, False, points, 3) def draw_rover(surface): """ Draw the rover. - Wheels: circles drawn at their contact points on the ground. - Suspension: springs from wheel contact to chassis attachment. - Chassis: drawn as a rectangle whose top edge connects the left and right attachments. """ # Determine wheel positions. wheels = {} for key, offset in wheel_offsets.items(): world_x = rover_x + offset world_y = get_ground_y(world_x) wheels[key] = (world_x, world_y) # Draw wheel. pygame.draw.circle(surface, DARKGRAY, (int(world_x), int(world_y)), wheel_radius) pygame.draw.circle(surface, BLACK, (int(world_x), int(world_y)), wheel_radius, 2) # Update suspension so that each wheel has its chassis attachment position. update_suspension() attachments = {} for key, offset in wheel_offsets.items(): attach_x = rover_x + offset attach_y = suspension[key] attachments[key] = (attach_x, attach_y) # Draw spring from wheel contact to attachment. draw_spring(surface, wheels[key], attachments[key], coils=8, amplitude=6) # Draw the attachment point. pygame.draw.circle(surface, BLACK, (int(attach_x), int(attach_y)), 4) # Draw chassis. # For simplicity, use the left and right attachment points to define the top edge. left_attach = attachments['left'] right_attach = attachments['right'] # Compute the chassis center as the average of the three attachments. center_attach = ( (attachments['left'][0] + attachments['center'][0] + attachments['right'][0]) / 3, (attachments['left'][1] + attachments['center'][1] + attachments['right'][1]) / 3, ) # Chassis properties. chassis_height = 20 # The top edge will be the line from left_attach to right_attach. # For a rough rectangle, offset the top edge downward by chassis_height to form the body. top_edge = [left_attach, right_attach] bottom_edge = [(left_attach[0], left_attach[1] + chassis_height), (right_attach[0], right_attach[1] + chassis_height)] body_points = [left_attach, right_attach, bottom_edge[1], bottom_edge[0]] pygame.draw.polygon(surface, BLUE, body_points) pygame.draw.polygon(surface, BLACK, body_points, 2) # ------------------------------- # Main Simulation Loop # ------------------------------- running = True while running: for event in pygame.event.get(): if event.type == pygame.QUIT: running = False # The mouse vertical position controls the next ground segment's height. mouse_x, mouse_y = pygame.mouse.get_pos() # Update ground control points. update_ground(mouse_y) # Clear the screen. screen.fill(WHITE) # Draw the ground. draw_ground(screen) # Draw the rover (suspension, wheels, chassis). draw_rover(screen) pygame.display.flip() clock.tick(60) pygame.quit() sys.exit()