""" Simulation Runner for Ant Colony Simulation This module provides the main simulation runner that coordinates the colony, environment, and visualization components. """ import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation import seaborn as sns import time import logging import yaml import h5py from pathlib import Path import pandas as pd from typing import Dict, List, Tuple, Optional, Any from .visualization import ColonyVisualizer from .colony import Colony from .environment import World from .utils.data_collection import DataCollector from agents.nestmate import TaskType class Simulation: """Main simulation runner class.""" def __init__(self, config_path: str): """Initialize simulation with configuration.""" self.config = self._load_config(config_path) self._setup_logging() # Initialize components self.environment = World(self.config['environment']) self.colony = Colony(self.config['colony'], self.environment) self.visualizer = ColonyVisualizer(self.config['visualization']) self.data_collector = DataCollector(self.config['data_collection']) # Simulation state self.current_step = 0 self.max_steps = self.config['simulation']['max_steps'] self.timestep = self.config['simulation']['timestep'] self.paused = False self.data = { 'time': [], 'resources': [], 'task_distribution': [], 'efficiency_metrics': [] } def _load_config(self, config_path: str) -> Dict[str, Any]: """Load configuration from YAML file.""" with open(config_path, 'r') as f: config = yaml.safe_load(f) return config def _setup_logging(self): """Setup logging configuration.""" log_config = self.config['debug']['logging'] logging.basicConfig( level=getattr(logging, log_config['level']), format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', filename=log_config['file'] ) self.logger = logging.getLogger(__name__) def step(self): """Execute one simulation step.""" # Update environment self.environment.update(self.timestep) # Update colony self.colony.update(self.timestep) # Collect data self._collect_data() # Increment step counter self.current_step += 1 def _collect_data(self): """Collect simulation data for analysis and visualization.""" self.data['time'].append(self.current_step * self.timestep) # Collect resource data resource_data = { rtype: sum(r.amount for r in resources) for rtype, resources in self.environment.resources.items() } self.data['resources'].append(resource_data) # Collect task distribution task_dist = self.colony.get_task_distribution() self.data['task_distribution'].append(task_dist) # Collect efficiency metrics metrics = self.colony.compute_efficiency_metrics() self.data['efficiency_metrics'].append(metrics) def run(self, headless: bool = False): """Run the simulation.""" self.logger.info("Starting simulation...") if not headless: # Create and show animation animation = self.visualizer.create_animation(self) plt.show() else: # Run without visualization try: while self.current_step < self.max_steps and not self.paused: self.step() # Save data periodically if self.current_step % self.config['data_collection']['frequency'] == 0: self.data_collector.save_data(self.data) except KeyboardInterrupt: self.logger.info("Simulation interrupted by user") finally: # Save final data self.data_collector.save_data(self.data) self.logger.info(f"Simulation completed after {self.current_step} steps") def pause(self): """Pause the simulation.""" self.paused = True self.logger.info("Simulation paused") def resume(self): """Resume the simulation.""" self.paused = False self.logger.info("Simulation resumed") def reset(self): """Reset the simulation to initial state.""" self.logger.info("Resetting simulation...") self.current_step = 0 self.environment.reset() self.colony.reset() self.data = { 'time': [], 'resources': [], 'task_distribution': [], 'efficiency_metrics': [] } def _update_visualization(self): """Update visualization plots.""" # Clear all axes for ax in [self.ax_main, self.ax_resources, self.ax_tasks, self.ax_metrics]: ax.clear() # Main simulation view self._plot_simulation_state() # Resource levels self._plot_resource_levels() # Task distribution self._plot_task_distribution() # Efficiency metrics self._plot_efficiency_metrics() plt.draw() plt.pause(0.01) def _plot_simulation_state(self): """Plot current simulation state.""" # Plot terrain terrain = self.environment.terrain.height_map self.ax_main.imshow(terrain, cmap='terrain') # Plot agents agent_positions = np.array([agent.position for agent in self.colony.agents]) agent_colors = [self._get_task_color(agent.current_task) for agent in self.colony.agents] if len(agent_positions) > 0: self.ax_main.scatter(agent_positions[:, 0], agent_positions[:, 1], c=agent_colors, alpha=0.6) # Plot resources resource_positions = np.array([[r.position.x, r.position.y] for r in self.environment.resources]) if len(resource_positions) > 0: self.ax_main.scatter(resource_positions[:, 0], resource_positions[:, 1], c='yellow', marker='*') # Plot nest self.ax_main.scatter([self.colony.nest_position.x], [self.colony.nest_position.y], c='white', marker='s', s=100) self.ax_main.set_title(f'Step {self.current_step}') def _plot_resource_levels(self): """Plot resource level history.""" if len(self.data['time']) > 0: for resource_type in self.colony.resources.keys(): values = [d[resource_type] for d in self.data['resources']] self.ax_resources.plot(self.data['time'], values, label=resource_type) self.ax_resources.legend() self.ax_resources.set_xlabel('Time') self.ax_resources.set_ylabel('Amount') def _plot_task_distribution(self): """Plot task distribution.""" if len(self.data['task_distribution']) > 0: latest_dist = self.data['task_distribution'][-1] tasks = list(latest_dist.keys()) counts = [latest_dist[task] for task in tasks] colors = [self._get_task_color(task) for task in tasks] self.ax_tasks.bar(range(len(tasks)), counts, color=colors) self.ax_tasks.set_xticks(range(len(tasks))) self.ax_tasks.set_xticklabels([task.value for task in tasks], rotation=45) def _plot_efficiency_metrics(self): """Plot efficiency metrics.""" if len(self.data['efficiency_metrics']) > 0: latest_metrics = self.data['efficiency_metrics'][-1] metrics = list(latest_metrics.keys()) values = [latest_metrics[metric] for metric in metrics] self.ax_metrics.bar(range(len(metrics)), values) self.ax_metrics.set_xticks(range(len(metrics))) self.ax_metrics.set_xticklabels(metrics, rotation=45) self.ax_metrics.set_ylim(0, 1) def _get_task_color(self, task: TaskType) -> str: """Get color for visualization based on task type.""" color_map = { TaskType.FORAGING: 'green', TaskType.MAINTENANCE: 'blue', TaskType.NURSING: 'purple', TaskType.DEFENSE: 'red', TaskType.EXPLORATION: 'orange' } return color_map.get(task, 'gray') def _save_data(self): """Save simulation data to file.""" if not self.config['data']['export']['enabled']: return # Create output directory if it doesn't exist output_dir = Path(self.config['data']['storage']['path']) output_dir.mkdir(parents=True, exist_ok=True) # Save to HDF5 if 'hdf5' in self.config['data']['storage']['format']: self._save_to_hdf5(output_dir / f'simulation_data_{self.current_step}.h5') # Save to CSV if 'csv' in self.config['data']['export']['format']: self._save_to_csv(output_dir) def _save_to_hdf5(self, filepath: Path): """Save data to HDF5 format.""" with h5py.File(filepath, 'w') as f: # Create groups sim_group = f.create_group('simulation') colony_group = f.create_group('colony') perf_group = f.create_group('performance') # Save simulation data sim_group.create_dataset('time', data=np.array(self.data['time'])) sim_group.create_dataset('step', data=self.current_step) # Save colony data colony_group.create_dataset('population', data=np.array(self.colony.stats.population)) # Create datasets for dictionary data for key in ['resources', 'task_distribution', 'efficiency_metrics']: if len(self.data[key]) > 0: group = colony_group.create_group(key) for metric_key in self.data[key][0].keys(): values = [d[metric_key] for d in self.data[key]] group.create_dataset(metric_key, data=np.array(values)) # Save performance data perf_group.create_dataset('step_times', data=np.array(self.performance_metrics['step_times'])) def _save_to_csv(self, output_dir: Path): """Save data to CSV format.""" # Save basic metrics pd.DataFrame({ 'time': self.data['time'], 'population': self.colony.stats.population }).to_csv(output_dir / 'basic_metrics.csv', index=False) # Save resource data if len(self.data['resources']) > 0: pd.DataFrame(self.data['resources']).to_csv( output_dir / 'resources.csv', index=False ) # Save task distribution if len(self.data['task_distribution']) > 0: pd.DataFrame(self.data['task_distribution']).to_csv( output_dir / 'task_distribution.csv', index=False ) # Save efficiency metrics if len(self.data['efficiency_metrics']) > 0: pd.DataFrame(self.data['efficiency_metrics']).to_csv( output_dir / 'efficiency_metrics.csv', index=False ) def _monitor_performance(self): """Monitor and log performance metrics.""" if len(self.performance_metrics['step_times']) > 0: avg_step_time = np.mean(self.performance_metrics['step_times'][-100:]) current_fps = 1.0 / avg_step_time self.performance_metrics['fps'].append(current_fps) self.logger.info(f"Current FPS: {current_fps:.2f}") # Check performance against targets if current_fps < self.config['performance']['optimization']['target_fps']: self._optimize_performance() def _optimize_performance(self): """Attempt to optimize simulation performance.""" if not self.config['performance']['optimization']['auto_tune']: return # Example optimization: Reduce visualization frequency if self.config['visualization']['realtime']['enabled']: current_freq = self.config['visualization']['realtime']['update_frequency'] self.config['visualization']['realtime']['update_frequency'] = min(current_freq * 2, 100) self.logger.info("Adjusted visualization frequency for performance optimization") def _cleanup(self): """Cleanup resources and save final state.""" # Save final data self._save_data() # Close visualization if self.config['visualization']['realtime']['enabled']: plt.close(self.fig) # Log final statistics self.logger.info(f"Simulation completed after {self.current_step} steps") self.logger.info(f"Final population: {self.colony.stats.population}") self.logger.info(f"Final resource levels: {self.colony.stats.resource_levels}") if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='Run Ant Colony Simulation') parser.add_argument('--config', type=str, required=True, help='Path to configuration file') parser.add_argument('--headless', action='store_true', help='Run without visualization') args = parser.parse_args() # Create and run simulation sim = Simulation(args.config) sim.run(headless=args.headless)