cognitive/Things/Ant_Colony/visualization.py
Daniel Ari Friedman e555897efa Updates
2025-02-07 09:36:14 -08:00

374 строки
13 KiB
Python

"""
Enhanced Visualization Module for Ant Colony Simulation
This module provides advanced visualization capabilities for the ant colony simulation,
including animated pheromone trails, agent movements, and colony statistics.
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import seaborn as sns
from typing import Dict, List, Optional
from matplotlib.patches import Circle, Wedge, Path, PathPatch
from matplotlib.collections import PatchCollection
import colorsys
import networkx as nx
class ColonyVisualizer:
"""Advanced visualization for the ant colony simulation."""
def __init__(self, config: dict):
"""Initialize visualizer."""
self.config = config
self.fig = None
self.axes = {}
self.artists = {}
self.animation = None
# Color schemes
self.color_schemes = {
'terrain': plt.cm.terrain,
'pheromone': {
'food': plt.cm.Greens,
'home': plt.cm.Reds,
'alarm': plt.cm.Oranges,
'trail': plt.cm.Blues
},
'agents': {
'foraging': '#2ecc71',
'maintenance': '#3498db',
'nursing': '#9b59b6',
'defense': '#e74c3c',
'exploration': '#f1c40f'
}
}
# Setup visualization
self._setup_visualization()
def _setup_visualization(self):
"""Set up the visualization layout."""
plt.style.use('dark_background')
# Create figure with custom layout
self.fig = plt.figure(figsize=(16, 9))
gs = self.fig.add_gridspec(3, 3)
# Main simulation view (larger)
self.axes['main'] = self.fig.add_subplot(gs[0:2, 0:2])
self.axes['main'].set_title('Colony Simulation')
# Pheromone levels
self.axes['pheromone'] = self.fig.add_subplot(gs[0, 2])
self.axes['pheromone'].set_title('Pheromone Levels')
# Resource levels
self.axes['resources'] = self.fig.add_subplot(gs[1, 2])
self.axes['resources'].set_title('Resource Levels')
# Task distribution
self.axes['tasks'] = self.fig.add_subplot(gs[2, 0])
self.axes['tasks'].set_title('Task Distribution')
# Efficiency metrics
self.axes['metrics'] = self.fig.add_subplot(gs[2, 1])
self.axes['metrics'].set_title('Colony Metrics')
# Network view
self.axes['network'] = self.fig.add_subplot(gs[2, 2])
self.axes['network'].set_title('Social Network')
plt.tight_layout()
def create_animation(self, simulation, interval: int = 50) -> FuncAnimation:
"""Create animation of the colony simulation."""
def update(frame):
# Update simulation state
if not simulation.paused:
simulation.step()
# Clear all axes
for ax in self.axes.values():
ax.clear()
# Update all plots
artists = []
artists.extend(self._plot_main_view(simulation))
artists.extend(self._plot_pheromone_levels(simulation))
artists.extend(self._plot_resource_levels(simulation))
artists.extend(self._plot_task_distribution(simulation))
artists.extend(self._plot_efficiency_metrics(simulation))
artists.extend(self._plot_social_network(simulation))
return artists
self.animation = FuncAnimation(
self.fig,
update,
frames=None,
interval=interval,
blit=True
)
return self.animation
def _plot_main_view(self, simulation) -> List:
"""Plot main simulation view with enhanced graphics."""
ax = self.axes['main']
artists = []
# Plot terrain with enhanced colormap
terrain = simulation.environment.terrain.height_map
terrain_img = ax.imshow(terrain, cmap=self.color_schemes['terrain'])
artists.append(terrain_img)
# Plot pheromone trails with alpha blending
for ptype, pdata in simulation.environment.pheromones.layers.items():
pheromone_grid = pdata['grid']
if np.any(pheromone_grid > 0):
pheromone_img = ax.imshow(
pheromone_grid,
cmap=self.color_schemes['pheromone'][ptype],
alpha=0.5
)
artists.append(pheromone_img)
# Plot agents with directional markers
for agent in simulation.colony.agents:
agent_color = self.color_schemes['agents'][agent.current_task.value]
# Create ant shape
ant = self._create_ant_shape(agent.position, agent.orientation,
size=1.0, color=agent_color)
ax.add_patch(ant)
artists.append(ant)
# Add carrying indicator if needed
if agent.carrying is not None:
carry_indicator = Circle(
(agent.position[0], agent.position[1]),
radius=0.3,
color='yellow',
alpha=0.7
)
ax.add_patch(carry_indicator)
artists.append(carry_indicator)
# Plot resources
for resource in simulation.environment.resources:
if resource.type == 'food':
marker = '*'
color = 'yellow'
else:
marker = 's'
color = 'cyan'
resource_dot = ax.scatter(
resource.position.x,
resource.position.y,
c=color,
marker=marker,
s=50
)
artists.append(resource_dot)
# Plot nest with concentric circles
nest_x = simulation.colony.nest_position.x
nest_y = simulation.colony.nest_position.y
for radius in [1, 2, 3]:
nest_circle = Circle(
(nest_x, nest_y),
radius,
color='white',
fill=False,
alpha=0.5
)
ax.add_patch(nest_circle)
artists.append(nest_circle)
# Add nest center
nest_center = ax.scatter(
[nest_x], [nest_y],
c='white',
marker='s',
s=100
)
artists.append(nest_center)
# Set limits and title
ax.set_xlim(0, simulation.environment.size[0])
ax.set_ylim(0, simulation.environment.size[1])
ax.set_title(f'Colony Simulation (Step {simulation.current_step})')
return artists
def _create_ant_shape(self, position, orientation, size=1.0, color='white'):
"""Create an ant-like shape for visualization."""
# Create ant body segments
body_verts = [
(-0.5, -0.2), # Abdomen
(0.0, -0.1), # Thorax
(0.5, 0.0), # Head
(0.0, 0.1),
(-0.5, 0.2)
]
# Scale and rotate vertices
cos_theta = np.cos(orientation)
sin_theta = np.sin(orientation)
transformed_verts = []
for x, y in body_verts:
x_rot = x * cos_theta - y * sin_theta
y_rot = x * sin_theta + y * cos_theta
transformed_verts.append(
(position[0] + x_rot * size,
position[1] + y_rot * size)
)
# Create path
codes = [Path.MOVETO] + [Path.LINETO] * (len(transformed_verts) - 1)
path = Path(transformed_verts, codes)
return PathPatch(path, facecolor=color, edgecolor='none', alpha=0.7)
def _plot_pheromone_levels(self, simulation) -> List:
"""Plot pheromone concentration levels."""
ax = self.axes['pheromone']
artists = []
# Get pheromone data
pheromone_levels = []
labels = []
colors = []
for ptype, pdata in simulation.environment.pheromones.layers.items():
pheromone_levels.append(np.mean(pdata['grid']))
labels.append(ptype)
colors.append(self.color_schemes['agents'].get(ptype, '#95a5a6'))
# Create bar plot
if pheromone_levels:
bars = ax.bar(range(len(pheromone_levels)), pheromone_levels,
color=colors)
artists.extend(bars)
ax.set_xticks(range(len(labels)))
ax.set_xticklabels(labels, rotation=45)
ax.set_ylabel('Average Concentration')
return artists
def _plot_resource_levels(self, simulation) -> List:
"""Plot resource levels over time."""
ax = self.axes['resources']
artists = []
if len(simulation.data['time']) > 0:
for resource_type in simulation.colony.resources.keys():
values = [d[resource_type] for d in simulation.data['resources']]
line = ax.plot(simulation.data['time'][-100:], values[-100:],
label=resource_type)
artists.extend(line)
ax.legend()
ax.set_xlabel('Time')
ax.set_ylabel('Amount')
return artists
def _plot_task_distribution(self, simulation) -> List:
"""Plot task distribution with enhanced graphics."""
ax = self.axes['tasks']
artists = []
if len(simulation.data['task_distribution']) > 0:
latest_dist = simulation.data['task_distribution'][-1]
tasks = list(latest_dist.keys())
counts = [latest_dist[task] for task in tasks]
colors = [self.color_schemes['agents'][task.value] for task in tasks]
# Create stacked bars for current and needed agents
bars = ax.bar(range(len(tasks)), counts, color=colors)
artists.extend(bars)
# Add target lines for needed agents
for i, task in enumerate(tasks):
need = simulation.colony.task_needs[task] * len(simulation.colony.agents)
line = ax.hlines(need, i-0.4, i+0.4, colors='white', linestyles='--')
artists.append(line)
ax.set_xticks(range(len(tasks)))
ax.set_xticklabels([task.value for task in tasks], rotation=45)
return artists
def _plot_efficiency_metrics(self, simulation) -> List:
"""Plot efficiency metrics with enhanced graphics."""
ax = self.axes['metrics']
artists = []
if len(simulation.data['efficiency_metrics']) > 0:
metrics = simulation.data['efficiency_metrics'][-1]
values = list(metrics.values())
labels = list(metrics.keys())
# Create radar chart
angles = np.linspace(0, 2*np.pi, len(labels), endpoint=False)
values = np.concatenate((values, [values[0]])) # Close the polygon
angles = np.concatenate((angles, [angles[0]])) # Close the polygon
# Plot radar
line = ax.plot(angles, values)
artists.extend(line)
# Fill radar
fill = ax.fill(angles, values, alpha=0.25)
artists.extend(fill)
# Add labels
for angle, label in zip(angles[:-1], labels):
ha = 'right' if np.cos(angle) < 0 else 'left'
va = 'top' if np.sin(angle) < 0 else 'bottom'
ax.text(angle, 1.2, label,
ha=ha, va=va,
rotation=np.degrees(angle))
ax.set_ylim(0, 1)
ax.set_xticks([])
return artists
def _plot_social_network(self, simulation) -> List:
"""Plot social network with enhanced graphics."""
ax = self.axes['network']
artists = []
# Get network data
G = simulation.colony.interaction_network
if len(G.nodes) > 0:
# Calculate node positions using spring layout
pos = nx.spring_layout(G)
# Draw edges
for edge in G.edges():
x = [pos[edge[0]][0], pos[edge[1]][0]]
y = [pos[edge[0]][1], pos[edge[1]][1]]
line = ax.plot(x, y, 'w-', alpha=0.2)
artists.extend(line)
# Draw nodes
for node in G.nodes():
color = self.color_schemes['agents'][node.current_task.value]
dot = ax.scatter(pos[node][0], pos[node][1],
c=color, s=50)
artists.append(dot)
ax.set_xticks([])
ax.set_yticks([])
return artists