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()