蚂蚁觅食模拟

import pygame
import math
import random
import numpy as np
from scipy.signal import convolve2d  # 修正卷积函数导入

# 初始化参数
# 调小窗口尺寸
WIDTH, HEIGHT = 800, 600
NEST_POS = (WIDTH // 2, HEIGHT // 2)
FOOD_SOURCES = [(200, 200), (WIDTH - 200, HEIGHT - 200), (WIDTH // 2, HEIGHT - 300)]  # 增加食物源
NUM_ANTS = 1000
PHEROMONE_DECAY = 0.985
FOOD_CAPACITY = 10000
MAX_SPEED = 2.2


class Ant:
    def __init__(self):
        self.pos = np.array(NEST_POS, dtype=np.float32)
        self.angle = random.uniform(0, 2 * math.pi)
        self.speed = MAX_SPEED * random.uniform(0.8, 1.2)
        self.has_food = False
        self.search_time = 0
        self.pheromone_strength = 1.0

    def sense_environment(self, grid):
        x, y = int(self.pos[0]), int(self.pos[1])
        return grid[
               max(0, x - 1):min(WIDTH, x + 2),
               max(0, y - 1):min(HEIGHT, y + 2)
               ]

    def follow_pheromone(self, sensed, homing=False):
        try:
            if sensed.size == 0:
                self.random_walk()
                return

            if homing:
                weights = sensed
            else:
                weights = 1 / (sensed + 1e-6)  # 避免除零错误

            total = np.sum(weights)
            if total > 0:
                probs = weights.flatten() / total
                choice = np.random.choice(len(probs), p=probs)
                dy, dx = np.unravel_index(choice, sensed.shape)
                target_vec = np.array([dx - 1, dy - 1])  # 转换为方向向量

                if np.linalg.norm(target_vec) > 0:
                    new_angle = math.atan2(target_vec[1], target_vec[0])
                    self.angle = new_angle + random.uniform(-0.1, 0.1)
        except:
            self.random_walk()

    def random_walk(self):
        self.angle += random.uniform(-0.5, 0.5)
        self.speed = MAX_SPEED * random.uniform(0.7, 1.1)

    def move(self, food_grid, pheromone_grid):
        # 环境感知
        sensed = self.sense_environment(pheromone_grid if self.has_food else food_grid)

        # 行为决策
        if self.has_food:
            self.follow_pheromone(sensed, homing=True)
            pheromone_grid[int(self.pos[0]), int(self.pos[1])] += 0.5  # 留下返巢信息素
        else:
            if self.search_time > 150:
                self.random_walk()
            else:
                self.follow_pheromone(sensed)

        # 位置更新
        self.pos += np.array([
            math.cos(self.angle) * self.speed,
            math.sin(self.angle) * self.speed
        ])
        self.pos = np.clip(self.pos, [10, 10], [WIDTH - 10, HEIGHT - 10])

        # 边界反弹
        if self.pos[0] <= 10 or self.pos[0] >= WIDTH - 10:
            self.angle = math.pi - self.angle
        if self.pos[1] <= 10 or self.pos[1] >= HEIGHT - 10:
            self.angle = -self.angle

        self.search_time += 1


class Simulation:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.display.set_caption("超大蚂蚁觅食模拟")
        self.clock = pygame.time.Clock()

        self.ants = [Ant() for _ in range(NUM_ANTS)]
        self.food_grid = np.zeros((WIDTH, HEIGHT))
        for pos in FOOD_SOURCES:
            self.food_grid[pos[0] - 20:pos[0] + 20, pos[1] - 20:pos[1] + 20] = FOOD_CAPACITY
        self.pheromone_grid = np.zeros((WIDTH, HEIGHT))

    def update_pheromones(self):
        self.pheromone_grid *= PHEROMONE_DECAY
        kernel = np.array([[0.05, 0.1, 0.05],
                           [0.1, 0.4, 0.1],
                           [0.05, 0.1, 0.05]])
        self.pheromone_grid = convolve2d(
            self.pheromone_grid, kernel, mode='same', boundary='symm'
        )

    def draw(self):
        self.screen.fill((25, 25, 40))

        # 绘制信息素
        pheromone_surface = pygame.surfarray.make_surface(
            np.clip(self.pheromone_grid * 50, 0, 255).astype(np.uint8).T
        )
        self.screen.blit(pheromone_surface, (0, 0))

        # 绘制蚂蚁
        for ant in self.ants:
            color = (220, 80, 80) if ant.has_food else (120, 180, 250)
            pygame.draw.circle(self.screen, color, ant.pos.astype(int), 3)

        # 绘制巢穴和食物
        pygame.draw.circle(self.screen, (250, 250, 100), NEST_POS, 20)
        for pos in FOOD_SOURCES:
            pygame.draw.circle(self.screen, (80, 220, 80), pos, 25)

        pygame.display.flip()

    def run(self):
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False

            for ant in self.ants:
                ant.move(self.food_grid, self.pheromone_grid)

                # 食物交互
                x, y = int(ant.pos[0]), int(ant.pos[1])
                if not ant.has_food and self.food_grid[x, y] > 0:
                    ant.has_food = True
                    ant.search_time = 0
                    self.food_grid[x, y] -= 1
                    self.pheromone_grid[x, y] += 1.0  # 留下觅食信息素
                elif ant.has_food and np.linalg.norm(ant.pos - NEST_POS) < 25:
                    ant.has_food = False
                    self.pheromone_grid[x, y] += 3.0

            self.update_pheromones()
            self.draw()
            self.clock.tick(60)

        pygame.quit()


if __name__ == "__main__":
    sim = Simulation()
    sim.run()