Update app.py

app.py

@@ -6,6 +6,7 @@ from scipy.spatial import distance
 from sklearn.cluster import KMeans
 import networkx as nx
 from collections import deque
+from scipy.signal import convolve2d
 
 # Constants
 GRID_SIZE = 200
@@ -53,7 +54,8 @@ class Ant:
         if random.random() < self.genome['exploration_rate']:  
             action = random.choice(possible_actions)
         else:
-            action = np.argmax(self.q_table[self.position[0], self.position[1], possible_actions])
+            q_values = [self.q_table[self.position[0], self.position[1], action] for action in possible_actions]
+            action = possible_actions[np.argmax(q_values)]
 
         reward = self.calculate_reward()
         self.update_q_table(action, reward)
@@ -126,7 +128,7 @@ def diffuse_pheromones(pheromone_grid):
                        [0.1,  0.4, 0.1],
                        [0.05, 0.1, 0.05]])
     for i in range(3):  # For each pheromone type
-        pheromone_grid[:,:,i] = np.convolve2d(pheromone_grid[:,:,i], kernel, mode='same', boundary='wrap')
+        pheromone_grid[:,:,i] = convolve2d(pheromone_grid[:,:,i], kernel, mode='same', boundary='wrap')
 
 # Genetic Algorithm
 def crossover(parent1, parent2):
@@ -223,10 +225,12 @@ start_simulation = st.sidebar.button("Start Simulation")
 stop_simulation = st.sidebar.button("Stop Simulation")
 reset_simulation = st.sidebar.button("Reset Simulation")
 
+# Initialize variables
+total_food_collected = 0
+iterations = 0
+
 # Main simulation loop
 if start_simulation:
-    total_food_collected = 0
-    iterations = 0
     cluster_centers = np.array([[0, 0], [0, 0], [0, 0]])
     
     progress_bar = st.progress(0)
@@ -255,34 +259,37 @@ if reset_simulation:
                  'learning_rate': learning_rate, 
                  'discount_factor': discount_factor}) 
             for _ in range(num_ants)]
+    total_food_collected = 0
+    iterations = 0
 
-# Display final statistics
-st.write("## Final Statistics")
-st.write(f"Total Food Collected: {total_food_collected}")
-st.write(f"Average Food per Iteration: {total_food_collected / iterations if iterations > 0 else 0}")
-
-# Display heatmap of pheromone concentration
-st.write("## Pheromone Concentration Heatmap")
-fig, ax = plt.subplots(figsize=(10, 10))
-heatmap = ax.imshow(np.sum(pheromone_grid, axis=2), cmap='hot', interpolation='nearest')
-plt.colorbar(heatmap)
-st.pyplot(fig)
-
-# Display ant role distribution
-roles = [ant.role for ant in ants]
-role_counts = {role: roles.count(role) for role in set(roles)}
-st.write("## Ant Role Distribution")
-st.bar_chart(role_counts)
-
-# Display network graph of ant communication
-st.write("## Ant Communication Network")
-G = nx.Graph()
-for ant in ants:
-    G.add_node(ant.position)
-    for nearby_ant in ant.nearby_ants:
-        G.add_edge(ant.position, nearby_ant.position)
-
-fig, ax = plt.subplots(figsize=(10, 10))
-pos = nx.spring_layout(G)
-nx.draw(G, pos, with_labels=False, node_size=30, node_color='skyblue', edge_color='gray', ax=ax)
-st.pyplot(fig)
+# Display final statistics only if simulation has run
+if iterations > 0:
+    st.write("## Final Statistics")
+    st.write(f"Total Food Collected: {total_food_collected}")
+    st.write(f"Average Food per Iteration: {total_food_collected / iterations}")
+
+    # Display heatmap of pheromone concentration
+    st.write("## Pheromone Concentration Heatmap")
+    fig, ax = plt.subplots(figsize=(10, 10))
+    heatmap = ax.imshow(np.sum(pheromone_grid, axis=2), cmap='hot', interpolation='nearest')
+    plt.colorbar(heatmap)
+    st.pyplot(fig)
+
+    # Display ant role distribution
+    roles = [ant.role for ant in ants]
+    role_counts = {role: roles.count(role) for role in set(roles)}
+    st.write("## Ant Role Distribution")
+    st.bar_chart(role_counts)
+
+    # Display network graph of ant communication
+    st.write("## Ant Communication Network")
+    G = nx.Graph()
+    for ant in ants:
+        G.add_node(ant.position)
+        for nearby_ant in ant.nearby_ants:
+            G.add_edge(ant.position, nearby_ant.position)
+
+    fig, ax = plt.subplots(figsize=(10, 10))
+    pos = nx.spring_layout(G)
+    nx.draw(G, pos, with_labels=False, node_size=30, node_color='skyblue', edge_color='gray', ax=ax)
+    st.pyplot(fig)