238 lines
8.5 KiB
Python
238 lines
8.5 KiB
Python
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()
|