Update app.py

app.py

@@ -134,55 +134,31 @@ class Environment:
     def visualize(self):
         fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))
 
-        # Cell distribution
-        cell_types = set(cell.cell_type for cell in self.cells)
-        for cell_type in cell_types:
-            x = [cell.x for cell in self.cells if cell.cell_type == cell_type]
-            y = [cell.y for cell in self.cells if cell.cell_type == cell_type]
-            color = next(cell.color for cell in self.cells if cell.cell_type == cell_type)
-            size = next(cell.size * 3 for cell in self.cells if cell.cell_type == cell_type)
-            fig.add_trace(go.Scatter(x=x, y=y, mode='markers', marker=dict(color=color, size=size),
-                                     name=cell_type), row=1, col=1)
-
-        fig.update_xaxes(title_text="X", row=1, col=1)
-        fig.update_yaxes(title_text="Y", row=1, col=1)
-
-        # Population over time
-        population_counts = {
-            "prokaryote": [],
-            "early_eukaryote": [],
-            "advanced_eukaryote": [],
-            "plant_like": []
-        }
-
-        for cell_type in population_counts:
-            count = len([cell for cell in self.cells if cell.cell_type == cell_type])
-            population_counts[cell_type].append(count)
-
-        for cell_type, counts in population_counts.items():
-            fig.add_trace(go.Scatter(y=counts, mode='lines', name=cell_type), row=1, col=2)
-
-        fig.update_xaxes(title_text="Time", row=1, col=2)
-        fig.update_yaxes(title_text="Population", row=1, col=2)
-
-        fig.update_layout(height=600, width=1200,
-                          title_text=f"Cell Evolution Simulation (Time: {self.time})")
-        return fig
-
-def run_simulation(num_steps, initial_cells):
-    env = Environment(100, 100)
-    
-    # Add initial cells
-    for _ in range(initial_cells):
-        cell = Cell(random.uniform(0, env.width), random.uniform(0, env.height))
-        env.add_cell(cell)
-    
-    # Run simulation
-    for step in range(num_steps):
-        env.update()
-        if step % 10 == 0:  # Visualize every 10 steps
-            yield env.visualize()
+        def setup_figure(env):
+    fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))
+
+    # Cell distribution
+    for cell_type, data in env.get_visualization_data()[0].items():
+        fig.add_trace(go.Scatter(
+            x=data["x"], y=data["y"], mode='markers',
+            marker=dict(color=data["color"], size=data["size"]),
+            name=cell_type
+        ), row=1, col=1)
+
+    # 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)
+
+    fig.update_xaxes(title_text="X", row=1, col=1)
+    fig.update_yaxes(title_text="Y", row=1, col=1)
+    fig.update_xaxes(title_text="Time", row=1, col=2)
+    fig.update_yaxes(title_text="Population", row=1, col=2)
+
+    fig.update_layout(height=600, width=1200, title_text="Cell Evolution Simulation")
+
+    return fig
 
+# Streamlit app
 st.title("Cell Evolution Simulation")
 
 num_steps = st.slider("Number of simulation steps", 100, 1000, 500)