Files
INF6B/simulations/balls/main.py
rattatwinko d0eaabdd87 some new stuff.
idk its all pretty fun! some C++ too!
2025-10-15 11:16:51 +02:00

275 lines
9.1 KiB
Python

import pygame
import math
import random
# Initialize Pygame
pygame.init()
# Screen setup
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Enhanced Physics Bouncing Balls")
# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (220, 50, 50)
GREEN = (50, 220, 50)
CYAN = (50, 220, 220)
YELLOW = (220, 220, 50)
# Clock
clock = pygame.time.Clock()
FPS = 60
# Global physics
gravity = 0.5
elasticity = 0.85
air_resistance = 0.99 # 1.0 = no resistance, lower = more resistance
friction = 0.98 # Ground friction when rolling
# Ball class
class Ball:
def __init__(self, x, y):
self.radius = 20
self.x = x
self.y = y
self.mass = self.radius / 10.0 # Mass proportional to radius
# Give random initial velocity in both x and y
self.vx = random.uniform(-5, 5)
self.vy = random.uniform(-5, 0)
self.tracer_length = 50
self.tracers = []
self.dragging = False
self.drag_start_pos = None
self.drag_last_pos = (x, y)
self.drag_history = []
def update(self):
if not self.dragging:
# Apply gravity
self.vy += gravity
# Apply air resistance (quadratic drag)
speed = math.hypot(self.vx, self.vy)
if speed > 0:
drag_force = (1 - air_resistance) * speed
drag_x = -drag_force * (self.vx / speed)
drag_y = -drag_force * (self.vy / speed)
self.vx += drag_x
self.vy += drag_y
# Update position
self.x += self.vx
self.y += self.vy
# Bounce off floor with friction
if self.y + self.radius > HEIGHT:
self.y = HEIGHT - self.radius
self.vy = -self.vy * elasticity
self.vx *= friction # Apply friction when bouncing on ground
# Stop if moving very slowly
if abs(self.vy) < 0.1:
self.vy = 0
if abs(self.vx) < 0.1:
self.vx = 0
# Bounce off ceiling
if self.y - self.radius < 0:
self.y = self.radius
self.vy = -self.vy * elasticity
# Bounce off right wall
if self.x + self.radius > WIDTH:
self.x = WIDTH - self.radius
self.vx = -self.vx * elasticity
# Bounce off left wall
if self.x - self.radius < 0:
self.x = self.radius
self.vx = -self.vx * elasticity
# Update tracers
self.tracers.append((int(self.x), int(self.y)))
if len(self.tracers) > self.tracer_length:
self.tracers.pop(0)
def start_drag(self, mx, my):
self.dragging = True
self.drag_start_pos = (mx, my)
self.drag_last_pos = (mx, my)
self.drag_history = [(mx, my)]
self.vx = 0
self.vy = 0
def update_drag(self, mx, my):
if self.dragging:
self.x, self.y = mx, my
self.drag_history.append((mx, my))
if len(self.drag_history) > 5: # Keep last 5 positions for velocity calculation
self.drag_history.pop(0)
self.drag_last_pos = (mx, my)
def end_drag(self):
if self.dragging and len(self.drag_history) >= 2:
# Calculate velocity from drag motion (throw mechanics)
dx = self.drag_history[-1][0] - self.drag_history[0][0]
dy = self.drag_history[-1][1] - self.drag_history[0][1]
frames = len(self.drag_history)
# Set velocity based on drag motion
self.vx = dx / frames * 2
self.vy = dy / frames * 2
self.dragging = False
self.drag_start_pos = None
self.drag_history = []
def draw(self, surface):
# Draw tracers with fade effect
for i, pos in enumerate(self.tracers):
alpha = int(255 * (i / len(self.tracers)))
size = int(3 + 2 * (i / len(self.tracers)))
color = (50, 220, 220)
pygame.draw.circle(surface, color, pos, size)
# Draw ball
pygame.draw.circle(surface, RED, (int(self.x), int(self.y)), self.radius)
# Draw velocity vector (scaled for visibility)
if not self.dragging:
vel_scale = 3
pygame.draw.line(surface, GREEN,
(int(self.x), int(self.y)),
(int(self.x + self.vx * vel_scale),
int(self.y + self.vy * vel_scale)), 3)
# Draw drag trajectory preview
if self.dragging and self.drag_start_pos:
pygame.draw.line(surface, YELLOW,
self.drag_start_pos,
(int(self.x), int(self.y)), 2)
# Draw arrow at the end
pygame.draw.circle(surface, YELLOW, (int(self.x), int(self.y)), 5)
# Ball-to-ball collision detection and response
def handle_ball_collisions(balls):
for i, ball1 in enumerate(balls):
for ball2 in balls[i+1:]:
dx = ball2.x - ball1.x
dy = ball2.y - ball1.y
distance = math.hypot(dx, dy)
# Check if balls are colliding
if distance < ball1.radius + ball2.radius and distance > 0:
# Normalize collision vector
nx = dx / distance
ny = dy / distance
# Separate balls
overlap = (ball1.radius + ball2.radius) - distance
ball1.x -= nx * overlap * 0.5
ball1.y -= ny * overlap * 0.5
ball2.x += nx * overlap * 0.5
ball2.y += ny * overlap * 0.5
# Calculate relative velocity
dvx = ball2.vx - ball1.vx
dvy = ball2.vy - ball1.vy
# Calculate relative velocity in collision normal direction
dot_product = dvx * nx + dvy * ny
# Only resolve if balls are moving toward each other
if dot_product < 0:
# Calculate impulse (conservation of momentum with equal mass)
impulse = 2 * dot_product / (ball1.mass + ball2.mass)
# Apply impulse to both balls
ball1.vx += impulse * ball2.mass * nx * elasticity
ball1.vy += impulse * ball2.mass * ny * elasticity
ball2.vx -= impulse * ball1.mass * nx * elasticity
ball2.vy -= impulse * ball1.mass * ny * elasticity
# List of balls
balls = [Ball(WIDTH//2, HEIGHT//2)]
running = True
while running:
mx, my = pygame.mouse.get_pos()
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
# Drag start
elif event.type == pygame.MOUSEBUTTONDOWN:
for ball in balls:
if math.hypot(mx - ball.x, my - ball.y) < ball.radius + 10:
ball.start_drag(mx, my)
break
# Drag end
elif event.type == pygame.MOUSEBUTTONUP:
for ball in balls:
if ball.dragging:
ball.end_drag()
# Spawn new ball
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_SPACE:
balls.append(Ball(random.randint(50, WIDTH-50),
random.randint(50, HEIGHT-50)))
elif event.key == pygame.K_c:
balls.clear()
balls.append(Ball(WIDTH//2, HEIGHT//2))
# Keyboard control for physics
keys = pygame.key.get_pressed()
if keys[pygame.K_UP]:
gravity = min(2.0, gravity + 0.01)
if keys[pygame.K_DOWN]:
gravity = max(0, gravity - 0.01)
if keys[pygame.K_RIGHT]:
elasticity = min(1.0, elasticity + 0.01)
if keys[pygame.K_LEFT]:
elasticity = max(0, elasticity - 0.01)
if keys[pygame.K_w]:
air_resistance = min(1.0, air_resistance + 0.001)
if keys[pygame.K_s]:
air_resistance = max(0.9, air_resistance - 0.001)
# Update dragging balls
for ball in balls:
if ball.dragging:
ball.update_drag(mx, my)
# Update balls
for ball in balls:
ball.update()
# Handle collisions between balls
handle_ball_collisions(balls)
# Drawing
screen.fill(BLACK)
for ball in balls:
ball.draw(screen)
# Physics info
font = pygame.font.SysFont(None, 22)
info = [
f"Gravity: {gravity:.2f} (↑↓) Elasticity: {elasticity:.2f} (←→)",
f"Air Resistance: {air_resistance:.3f} (WS) Balls: {len(balls)}",
f"SPACE: Add ball C: Clear Drag to throw!"
]
for i, text in enumerate(info):
surface = font.render(text, True, WHITE)
screen.blit(surface, (10, 10 + i * 25))
pygame.display.flip()
clock.tick(FPS)
pygame.quit()