Stochastic circles in Python

import pygame

import random

import math


# --- Configuration ---

WIDTH, HEIGHT = 900, 700

NUM_CIRCLES = 350

MIN_RADIUS = 4

MAX_RADIUS = 7

MAX_SPEED =3

GROUPS = 7

FORCE_SCALE = -300

ALIGNMENT_SCALE = .0025

COHESION_SCALE = 0.0004

CENTER_GRAVITY = -2  # you wanted this high — now used as a force magnitude, not multiplied by distance

TRAIL_FADE = 240

CELL_SIZE = 80  # for spatial partitioning


# --- Initialize ---

pygame.init()

screen = pygame.display.set_mode((WIDTH, HEIGHT))

clock = pygame.time.Clock()

trail_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)


# Base colors per group

BASE_COLORS = [

    (255, 80, 80),    # 

    (80, 255, 80),    # 

    (80, 80, 255),    # 

    (255, 255, 80),   # 

    (255, 80, 255),   # 

    (10, 80, 55),   # 

    (110, 10, 200),   # 


]


# --- Circle class ---

class Circle:

    def __init__(self, group):

        self.group = group

        self.radius = random.choices(

            population=[MIN_RADIUS, 8, 10, 12, 15, 20, MAX_RADIUS],

            weights=[20, 15, 10, 10, 7, 5, 3],

            k=1

        )[0]

        self.mass = self.radius ** 2

        self.x = random.uniform(self.radius, WIDTH - self.radius)

        self.y = random.uniform(self.radius, HEIGHT - self.radius)

        self.vx = random.uniform(-MAX_SPEED, MAX_SPEED)

        self.vy = random.uniform(-MAX_SPEED, MAX_SPEED)

        r, g, b = BASE_COLORS[group]

        self.color = (

            min(255, max(0, r + random.randint(-40, 40))),

            min(255, max(0, g + random.randint(-40, 40))),

            min(255, max(0, b + random.randint(-40, 40))),

        )


    def move(self):

        self.x += self.vx

        self.y += self.vy

        speed = math.hypot(self.vx, self.vy)

        if speed > MAX_SPEED:

            scale = MAX_SPEED / speed

            self.vx *= scale

            self.vy *= scale


        # Bounce off walls

        if self.x < self.radius or self.x > WIDTH - self.radius:

            self.vx *= -2

            self.x = max(self.radius, min(WIDTH - self.radius, self.x))

        if self.y < self.radius or self.y > HEIGHT - self.radius:

            self.vy *= -2

            self.y = max(self.radius, min(HEIGHT - self.radius, self.y))


    def draw(self, surf):

        glow_intensity = max(40, min(255, int(self.mass ** 0.5 * 2)))

        color = (*self.color, glow_intensity)

        pygame.draw.circle(surf, color, (int(self.x), int(self.y)), self.radius)


# --- Moving center of gravity ---

class MovingCenter:

    def __init__(self, x, y):

        self.x = x

        self.y = y

        self.vx = random.choice([-5, 5])

        self.vy = random.choice([-5, 5])


    def move(self):

        self.x += self.vx

        self.y += self.vy

        if self.x < 0 or self.x > WIDTH:

            self.vx *= -1

        if self.y < 0 or self.y > HEIGHT:

            self.vy *= -1


# --- Optimized forces using uniform grid ---

def handle_forces(circles, center):

    # spatial grid

    grid = {}

    for c in circles:

        key = (int(c.x // CELL_SIZE), int(c.y // CELL_SIZE))

        grid.setdefault(key, []).append(c)


    # compute group centers

    group_centers = [[0.0, 0.0, 0] for _ in range(GROUPS)]

    for c in circles:

        gx, gy, cnt = group_centers[c.group]

        group_centers[c.group] = [gx + c.x, gy + c.y, cnt + 1]

    group_positions = [

        (gx / cnt, gy / cnt) if cnt > 0 else (0.0, 0.0)

        for gx, gy, cnt in group_centers

    ]


    # neighbor offsets for 3x3 cells

    neighbor_offsets = [

        (-1, -1), (-1, 0), (-1, 1),

        (0, -1),  (0, 0),  (0, 1),

        (1, -1),  (1, 0),  (1, 1)

    ]


    # pairwise interactions but only among local cells

    for (cx, cy), cell_circles in list(grid.items()):

        for ox, oy in neighbor_offsets:

            neighbors = grid.get((cx + ox, cy + oy))

            if not neighbors:

                continue

            for c1 in cell_circles:

                for c2 in neighbors:

                    if c1 is c2:

                        continue

                    # avoid processing each pair twice: only process if id(c1) < id(c2)

                    if id(c1) >= id(c2):

                        continue


                    dx = c2.x - c1.x

                    dy = c2.y - c1.y

                    dist = math.hypot(dx, dy)


                    # ignore far pairs to reduce work (tuned threshold)

                    if dist > CELL_SIZE * 1.5:

                        continue


                    # protect against zero distance

                    if dist < 1e-4:

                        # tiny jitter so direction is defined

                        dx = (random.random() - 0.5) * 0.01

                        dy = (random.random() - 0.5) * 0.01

                        dist = math.hypot(dx, dy)


                    nx, ny = dx / dist, dy / dist


                    # compute force magnitude (scalar)

                    if c1.group == c2.group:

                        # gravitational-like attraction, soften denominator

                        force_mag = FORCE_SCALE * (c1.mass + c2.mass) / (dist * dist + 1e-3)

                        fx = nx * force_mag

                        fy = ny * force_mag

                    else:

                        # repulsion proportional to size, implemented as force magnitude

                        repulsion = 0.02 * (c1.radius + c2.radius)

                        fx = -nx * repulsion

                        fy = -ny * repulsion


                    # convert force -> acceleration and apply (a = F / m)

                    c1.vx += fx / c1.mass

                    c1.vy += fy / c1.mass

                    c2.vx -= fx / c2.mass

                    c2.vy -= fy / c2.mass


                    # Collision bounce with mass-aware elastic impulse

                    min_dist = c1.radius + c2.radius

                    if dist < min_dist:

                        overlap = min_dist - dist

                        # separate

                        c1.x -= nx * overlap * 0.5

                        c1.y -= ny * overlap * 0.5

                        c2.x += nx * overlap * 0.5

                        c2.y += ny * overlap * 0.5


                        # relative velocity along normal

                        rel = (c1.vx * nx + c1.vy * ny) - (c2.vx * nx + c2.vy * ny)

                        if rel < 0:  # closing

                            # compute impulse scalar for perfectly elastic collision (e=1)

                            j = -(1.0 + 1.0) * rel / (1.0 / c1.mass + 1.0 / c2.mass)

                            # apply impulse

                            c1.vx += (j * nx) / c1.mass

                            c1.vy += (j * ny) / c1.mass

                            c2.vx -= (j * nx) / c2.mass

                            c2.vy -= (j * ny) / c2.mass


    # Flocking (cohesion + alignment) and attraction to moving center

    for c in circles:

        # cohesion toward group center

        gx, gy = group_positions[c.group]

        c.vx += (gx - c.x) * COHESION_SCALE

        c.vy += (gy - c.y) * COHESION_SCALE


        # alignment using local cell neighbors (cheaper than global scan)

        cell_key = (int(c.x // CELL_SIZE), int(c.y // CELL_SIZE))

        local = grid.get(cell_key, [])

        neigh_vx = neigh_vy = 0.0

        cnt = 0

        for other in local:

            if other is c or other.group != c.group:

                continue

            dx = other.x - c.x

            dy = other.y - c.y

            if dx*dx + dy*dy < (60.0 * 60.0):

                neigh_vx += other.vx

                neigh_vy += other.vy

                cnt += 1

        if cnt > 0:

            avg_vx = neigh_vx / cnt

            avg_vy = neigh_vy / cnt

            c.vx += (avg_vx - c.vx) * ALIGNMENT_SCALE

            c.vy += (avg_vy - c.vy) * ALIGNMENT_SCALE


        # FIXED central gravity: normalized direction, force -> acceleration (divided by mass)

        dx = center.x - c.x

        dy = center.y - c.y

        dist = math.hypot(dx, dy)

        if dist > 1e-4:

            dirx, diry = dx / dist, dy / dist

            # use CENTER_GRAVITY as force magnitude, convert to accel by /mass

            # optionally cap force_mag for stability

            force_mag = min(CENTER_GRAVITY, 5.0)

            ax = dirx * force_mag / c.mass

            ay = diry * force_mag / c.mass

            c.vx += ax

            c.vy += ay


# --- Main loop ---

circles = []

for g in range(GROUPS):

    circles.extend([Circle(g) for _ in range(NUM_CIRCLES // GROUPS)])


center = MovingCenter(WIDTH / 2, HEIGHT / 2)


running = True

while running:

    for event in pygame.event.get():

        if event.type == pygame.QUIT:

            running = False


    trail_surface.fill((0, 0, 0, TRAIL_FADE))


    center.move()

    handle_forces(circles, center)

    for c in circles:

        c.move()

        c.draw(trail_surface)

    FORCE_SCALE-=1

    if FORCE_SCALE <-500:

    FORCE_SCALE =-100

    screen.blit(trail_surface, (0, 0))

    #pygame.draw.circle(screen, (255, 255, 255), (int(center.x), int(center.y)), 6)

    pygame.display.flip()

    clock.tick(60)


pygame.quit()


No comments:

Post a Comment