Update app.py

app.py

@@ -213,7 +213,9 @@ def setup_figure(env):
 
     # Total population over time
     for cell_type, counts in env.population_history.items():
-        fig.add_trace(go.Scatter(y=counts, mode='lines', name=cell_type), row=1, col=2)# Individual population charts
+        fig.add_trace(go.Scatter(y=counts, mode='lines', name=cell_type), row=1, col=2)
+
+    # Individual population charts
     for i, cell_type in enumerate(cell_types):
         if cell_type == "modified":
             fig.add_trace(go.Scatter(y=env.population_history[cell_type], mode='lines', 
@@ -282,7 +284,8 @@ def run_simulation():
                     fig.data[13].y = population_history["plant_like"]
                     fig.data[14].y = population_history["modified"]
         
-        q.put(fig)
+        if q.empty():
+            q.put(fig)
 
 # Create a function to start the simulation
 def start_simulation():
@@ -290,9 +293,14 @@ def start_simulation():
     thread.start()
 
 # Create a function to stop the simulation
+stop_button_pressed = False
+
 def stop_simulation():
+    global stop_button_pressed
+    stop_button_pressed = True
     q.put(None)
 
+# Create buttons to start and stop the simulation
 start_button = st.button("Start Simulation", on_click=start_simulation)
 stop_button = st.button("Stop Simulation", on_click=stop_simulation)
 
@@ -302,8 +310,9 @@ while True:
         if fig is not None:
             chart_placeholder.plotly_chart(fig, use_container_width=True)
         else:
-            print("Received None figure from queue. Stopping simulation.")
             break
     except queue.Empty:
         time.sleep(0.1)
+        if stop_button_pressed:
+            break
         continue