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