Codon bouncing ball simulation

simulation.py

import sys
import time
import math
import random


class AsciiCanvas:
    width: int
    height: int
    fill_char: str
    canvas: list[list[str]]
    def __init__(self, width, height, fill_char=' '):
        self.width = width
        self.height = height
        self.fill_char = fill_char
        self.canvas = [
            [fill_char for _ in range(width)] for _ in range(height)]

    def draw_line(self, x1, y1, x2, y2, char='*'):
        # Bresenham's line algorithm
        dx = abs(x2 - x1)
        dy = abs(y2 - y1)
        x, y = x1, y1
        sx = -1 if x1 > x2 else 1
        sy = -1 if y1 > y2 else 1

        if dx > dy:
            err = dx / 2.0
            while x != x2:
                self.set_pixel(x, y, char)
                err -= dy
                if err < 0:
                    y += sy
                    err += dx
                x += sx
        else:
            err = dy / 2.0
            while y != y2:
                self.set_pixel(x, y, char)
                err -= dx
                if err < 0:
                    x += sx
                    err += dy
                y += sy

        self.set_pixel(x, y, char)

    def draw_rect(self, x, y, width, height, char='*'):
        for i in range(height):
            for j in range(width):
                if i == 0 or i == height - 1 or j == 0 or j == width - 1:
                    self.set_pixel(x + j, y + i, char)

    def set_pixel(self, x, y, char):
        if 0 <= x < self.width and 0 <= y < self.height:
            self.canvas[y][x] = char

    def __str__(self):
        return '\n'.join([''.join(row) for row in self.canvas])

    def clear(self):
        self.canvas = [[self.fill_char for _ in range(
            self.width)] for _ in range(self.height)]

    def update(self):
        sys.stdout.write('\r' + str(self) + '\n')
        sys.stdout.flush()


class Ball:
    x: float
    y: float
    char: str
    vx: float
    vy: float
    ax: float
    ay: float
    gravity: float
    friction: float
    charge: float
    mass: float
    radius: float
    angular_momentum: float
    def __init__(self, x, y, char='o', vx=1, vy=1, gravity=0.2, friction=0.98, charge=1, mass=1, radius=1, angular_momentum=0):
        self.x = x
        self.y = y
        self.char = char
        self.vx = vx
        self.vy = vy
        self.ax = 0
        self.ay = 0
        self.gravity = gravity
        self.friction = friction
        self.charge = charge
        self.radius = radius
        self.mass = mass
        self.angular_momentum = angular_momentum

    def update(self, canvas, box, charge_transfer_rate=0.01):
        next_x = self.x + self.vx
        next_y = self.y + self.vy

        if next_x <= box.x or next_x >= box.x + box.width - 1:
            self.vx = -self.vx * self.friction
            charge_transfer = charge_transfer_rate * self.charge
            self.charge -= charge_transfer
            box.charge += charge_transfer

        if next_y <= box.y or next_y >= box.y + box.height - 1:
            self.vy = -self.vy * self.friction
            charge_transfer = charge_transfer_rate * self.charge
            self.charge -= charge_transfer
            box.charge += charge_transfer

        else:
            self.vy += self.gravity

        self.x += self.vx
        self.y += self.vy
        
        self.ax = self.charge / self.mass
        self.ay = self.charge / self.mass
        


    def draw(self, canvas):
        # Calculate the magnitude of the drag force
        drag_magnitude = math.sqrt(self.vx ** 2 + self.vy ** 2) * self.mass

        # Choose a character to represent the ball based on the drag magnitude
        if drag_magnitude < 1.0:
            char = "·"
        elif drag_magnitude < 5.0:
            char = "o"
        else:
            char = "O"

        canvas.set_pixel(int(self.x), int(self.y), char)

    def check_collision(self, other):
        dx = self.x - other.x
        dy = self.y - other.y
        distance = math.sqrt(dx ** 2 + dy ** 2)

        if distance < self.radius + other.radius:
            angle = math.atan2(dy, dx)

            self_v_mag = math.sqrt(self.vx ** 2 + self.vy ** 2)
            other_v_mag = math.sqrt(other.vx ** 2 + other.vy ** 2)

            self_v_angle = math.atan2(self.vy, self.vx) - angle
            other_v_angle = math.atan2(other.vy, other.vx) - angle

            self_vx_after = ((self.mass - other.mass) * self_v_mag * math.cos(self_v_angle) +
                             2 * other.mass * other_v_mag * math.cos(other_v_angle)) / (self.mass + other.mass)
            self_vy_after = self_v_mag * math.sin(self_v_angle)

            other_vx_after = (2 * self.mass * self_v_mag * math.cos(self_v_angle) - (
                self.mass - other.mass) * other_v_mag * math.cos(other_v_angle)) / (self.mass + other.mass)
            other_vy_after = other_v_mag * math.sin(other_v_angle)

            self.vx = math.cos(angle) * self_vx_after - \
                math.sin(angle) * self_vy_after
            self.vy = math.sin(angle) * self_vx_after + \
                math.cos(angle) * self_vy_after

            other.vx = math.cos(angle) * other_vx_after - \
                math.sin(angle) * other_vy_after
            other.vy = math.sin(angle) * other_vx_after + \
                math.cos(angle) * other_vy_after

            while distance < self.radius + other.radius:
                self.x += self.vx
                self.y += self.vy
                other.x += other.vx
                other.y += other.vy

                dx = self.x - other.x
                dy = self.y - other.y
                distance = math.sqrt(dx ** 2 + dy ** 2)
                

    def apply_electrical_force(self, other, k=10000, charge_transfer_rate=0.01):
        dx = self.x - other.x
        dy = self.y - other.y
        distance = math.sqrt(dx ** 2 + dy ** 2)

        if distance < 2 * self.radius:
            force = k * self.charge * other.charge / (distance ** 2)

            fx = force * dx / distance
            fy = force * dy / distance

            self.vx += fx / self.mass
            self.vy += fy / self.mass

            other.vx -= fx / other.mass
            other.vy -= fy / other.mass

            charge_transfer = charge_transfer_rate * (self.charge - other.charge) / 2
            self.charge -= charge_transfer
            other.charge += charge_transfer
            
    def apply_magnetic_force(self, B=0.01):
        fx = self.charge * self.vy * B
        fy = -self.charge * self.vx * B

        self.vx += fx
        self.vy += fy

    def apply_magnetic_interaction(self, other, k=0.001):
        dx = self.x - other.x
        dy = self.y - other.y
        distance = math.sqrt(dx ** 2 + dy ** 2)

        if distance > 0:
            angle = math.atan2(dy, dx)
            magnetic_field = k * other.charge * other.vy / (distance ** 2)
            force = self.charge * magnetic_field

            fx = force * math.sin(angle)
            fy = -force * math.cos(angle)

            self.vx += fx / self.mass
            self.vy += fy / self.mass
            
    def apply_air_resistance(self, balls, drag_coefficient=0.001, vortex_factor=0.5, vortex_distance_factor=3):
        drag_force_x = -drag_coefficient * self.vx * abs(self.vx)
        drag_force_y = -drag_coefficient * self.vy * abs(self.vy)

        # Adjust drag force based on the proximity of other balls
        for other in balls:
            if other is not self:
                dx = self.x - other.x
                dy = self.y - other.y
                distance = math.sqrt(dx ** 2 + dy ** 2)

                if distance < self.radius * vortex_distance_factor:
                    drag_reduction = vortex_factor * (1 - distance / (self.radius * vortex_distance_factor))
                    drag_force_x *= (1 - drag_reduction)
                    drag_force_y *= (1 - drag_reduction)

        self.vx += drag_force_x / self.mass
        self.vy += drag_force_y / self.mass
    
    def check_collision_with_ramp(self, ramp, charge_transfer_rate=0.01):
        x1, y1, x2, y2 = ramp.x1, ramp.y1, ramp.x2, ramp.y2
        dx = x2 - x1
        dy = y2 - y1
        ramp_length = math.sqrt(dx ** 2 + dy ** 2)

        # Check if the ball is within the bounding box of the ramp
        if (x1 <= self.x <= x2 or x1 >= self.x >= x2) and (y1 <= self.y <= y2 or y1 >= self.y >= y2):
            # Calculate the distance between the ball and the ramp
            distance = abs(dy * self.x - dx * self.y + x2 * y1 - y2 * x1) / ramp_length

            # Check for a collision
            if distance <= self.radius:
                # Calculate the angle of the ramp
                ramp_angle = math.atan2(dy, dx)

                # Calculate the normal vector components
                nx = math.cos(ramp_angle + math.pi / 2)
                ny = math.sin(ramp_angle + math.pi / 2)

                # Calculate the relative velocity components
                rvx = self.vx - nx * (nx * self.vx + ny * self.vy)
                rvy = self.vy - ny * (nx * self.vx + ny * self.vy)

                # Update the ball's velocity with the new components
                self.vx = rvx
                self.vy = rvy

                charge_transfer = charge_transfer_rate * self.charge
                self.charge -= charge_transfer
                ramp.charge += charge_transfer
    
    def apply_strong_nuclear_force(self, other):
        distance = math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
        force_constant = 10000  # Adjust this value to control the strength of the strong nuclear force

        if distance < 1:  # Apply the strong nuclear force only for very short distances
            fx = force_constant * (self.x - other.x) / distance
            fy = force_constant * (self.y - other.y) / distance

            self.ax += fx / self.mass
            self.ay += fy / self.mass
            other.ax -= fx / other.mass
            other.ay -= fy / other.mass
        

class Box:
    x: int
    y: int
    width: int
    height: int
    char: str
    charge: float
    def __init__(self, x, y, width, height, char='*', charge=0):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.char = char
        self.charge = charge
        
class Ramp:
    x1: int
    y1: int
    x2: int
    y2: int
    charge: float
    def __init__(self, x1, y1, x2, y2, charge=0):
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2
        self.charge = charge

def update_balls(balls, canvas, box, ramps):
    for i in range(len(balls)):
        for j in range(i + 1, len(balls)):
            balls[i].check_collision(balls[j])
            balls[i].apply_electrical_force(balls[j])
            balls[j].apply_electrical_force(balls[i])

            balls[i].apply_magnetic_interaction(balls[j])
            balls[j].apply_magnetic_interaction(balls[i])
            
            balls[i].apply_strong_nuclear_force(balls[j])
            balls[j].apply_strong_nuclear_force(balls[i])

        balls[i].apply_magnetic_force()
        balls[i].apply_air_resistance(balls)
        balls[i].angular_momentum = balls[i].mass * (balls[i].vx ** 2 + balls[i].vy ** 2) * balls[i].radius
        for ramp in ramps:
            balls[i].check_collision_with_ramp(ramp)

        balls[i].update(canvas, box)
        balls[i].draw(canvas)


def get_magnetic_field_at_point(x, y, balls, k=0.001):
    total_magnetic_field: float = 0

    for ball in balls:
        dx = x - ball.x
        dy = y - ball.y
        distance = math.sqrt(dx ** 2 + dy ** 2)

        if distance > 0:
            magnetic_field = k * ball.charge * ball.vy / (distance ** 2)
            total_magnetic_field += magnetic_field * (dy / distance)

    return abs(total_magnetic_field)


def draw_magnetic_field(balls, canvas, field_resolution=2):
    field_strength_chars = ' _-:=+*%@#'

    # Calculate the maximum magnetic field strength in the scene
    max_field_strength: float = 0
    for y in range(0, canvas.height, field_resolution):
        for x in range(0, canvas.width, field_resolution):
            field_strength = get_magnetic_field_at_point(x, y, balls)
            max_field_strength = max(max_field_strength, field_strength)

    # Draw the magnetic field using normalized field strength
    for y in range(0, canvas.height, field_resolution):
        for x in range(0, canvas.width, field_resolution):
            field_strength = get_magnetic_field_at_point(x, y, balls)

            if max_field_strength > 0:
                normalized_field_strength = field_strength / max_field_strength
            else:
                normalized_field_strength = 0

            field_strength_index = min(len(field_strength_chars) - 1, int(
                normalized_field_strength * (len(field_strength_chars) - 1)))
            canvas.set_pixel(x, y, field_strength_chars[field_strength_index])


def add_ball_to_box(balls, box, radius=1, charge_range=(-1, 1), mass_range=(1, 5)):
    x = random.uniform(box.x + radius, box.x + box.width - radius)
    y = random.uniform(box.y + radius, box.y + box.height - radius)
    vx = random.uniform(-1, 1)
    vy = random.uniform(-1, 1)
    charge = random.uniform(*charge_range)
    mass = random.uniform(*mass_range)

    new_ball = Ball(x, y, 'o', vx, vy, radius=radius, charge=charge, mass=mass)
    balls.append(new_ball)


def calculate_average_position(balls):
    total_x = sum(ball.x for ball in balls)
    total_y = sum(ball.y for ball in balls)
    return total_x / len(balls), total_y / len(balls)


def calculate_average_velocity(balls):
    total_vx = sum(ball.vx for ball in balls)
    total_vy = sum(ball.vy for ball in balls)
    return total_vx / len(balls), total_vy / len(balls)

def calculate_average_charge(balls, boxes, ramps):
    ball_charge = sum(ball.charge for ball in balls) / len(balls)
    box_charge = sum(box.charge for box in boxes) / len(boxes)
    ramp_charge = sum(ramp.charge for ramp in ramps) / len(ramps)
    return sum((ball_charge, box_charge, ramp_charge)) / 3

def calculate_average_momentum(balls):
    total_px = sum(ball.vx * ball.charge for ball in balls)
    total_py = sum(ball.vy * ball.charge for ball in balls)
    return total_px / len(balls), total_py / len(balls)

def calculate_average_angular_momentum(balls):
    return sum(ball.angular_momentum for ball in balls) / len(balls)

def draw_ramps(ramps, canvas):
    for ramp in ramps:
        canvas.draw_line(ramp.x1, ramp.y1, ramp.x2, ramp.y2)

if __name__ == '__main__':
    box = Box(0, 0, 100, 40)
    canvas = AsciiCanvas(box.width + 2, box.height + 2)

    balls: list[Ball] = []
    for _ in range(50):
        add_ball_to_box(balls, box)
        
    # Define ramps
    ramps = [
        Ramp(10, 10, 20, 20),
        Ramp(30, 15, 40, 5),
        Ramp(50, 5, 60, 15),
        Ramp(70, 20, 80, 10),
    ]

    while True:
        canvas.clear()
        draw_magnetic_field(balls, canvas)
        canvas.draw_rect(box.x, box.y, box.width, box.height, box.char)
        draw_ramps(ramps, canvas)
        update_balls(balls, canvas, box, ramps)
        canvas.update()

        # Calculate statistical averages
        avg_position = calculate_average_position(balls)
        avg_velocity = calculate_average_velocity(balls)
        avg_momentum = calculate_average_momentum(balls)
        avg_charge = calculate_average_charge(balls, [box], ramps)
        avg_angular_momentum = calculate_average_angular_momentum(balls)

        # Print statistical averages
        print("Average Position:", avg_position)
        print("Average Velocity:", avg_velocity)
        print("Average Momentum:", avg_momentum)
        print("Average Charge:", avg_charge)
        print("Average Angular Momentum:", avg_angular_momentum)

        # Add delay (in seconds) between updates for animation
        time.sleep(1/144)