Update app.py

app.py

@@ -173,13 +173,13 @@ class Environment:
                 count = len([cell for cell in self.cells if cell.modifications])
             self.population_history[cell_type].append(count)
 
-    def get_visualization_data(self):
+     def get_visualization_data(self):
         cell_data = {
-            "prokaryote": {"x": [], "y": [], "size": [], "color": "lightblue"},
-            "early_eukaryote": {"x": [], "y": [], "size": [], "color": "green"},
-            "advanced_eukaryote": {"x": [], "y": [], "size": [], "color": "red"},
-            "plant_like": {"x": [], "y": [], "size": [], "color": "darkgreen"},
-            "modified": {"x": [], "y": [], "size": [], "color": "purple"}
+            "prokaryote": {"x": [], "y": [], "size": [], "color": "lightblue", "symbol": "circle"},
+            "early_eukaryote": {"x": [], "y": [], "size": [], "color": "green", "symbol": "square"},
+            "advanced_eukaryote": {"x": [], "y": [], "size": [], "color": "red", "symbol": "diamond"},
+            "plant_like": {"x": [], "y": [], "size": [], "color": "darkgreen", "symbol": "star"},
+            "modified": {"x": [], "y": [], "size": [], "color": "purple", "symbol": "cross"}
         }
 
         for cell in self.cells:
@@ -191,26 +191,48 @@ class Environment:
         return cell_data, self.population_history
 
 def setup_figure(env):
-    fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))
+    cell_types = ["prokaryote", "early_eukaryote", "advanced_eukaryote", "plant_like", "modified"]
+    fig = make_subplots(rows=2, cols=3, 
+                        subplot_titles=("Cell Distribution", "Total Population", 
+                                        "Prokaryotes", "Early Eukaryotes", 
+                                        "Advanced Eukaryotes", "Plant-like & Modified"),
+                        vertical_spacing=0.1,
+                        horizontal_spacing=0.05)
 
     # 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"]),
+            marker=dict(color=data["color"], size=data["size"], symbol=data["symbol"]),
             name=cell_type
         ), row=1, col=1)
 
-    # Population over time
+    # 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
+    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', 
+                                     name=cell_type, line=dict(color="purple")), row=2, col=3)
+        elif cell_type == "plant_like":
+            fig.add_trace(go.Scatter(y=env.population_history[cell_type], mode='lines', 
+                                     name=cell_type, line=dict(color="darkgreen")), row=2, col=3)
+        else:
+            fig.add_trace(go.Scatter(y=env.population_history[cell_type], mode='lines', 
+                                     name=cell_type), row=2, col=i+1)
+
     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="Advanced Cell Evolution Simulation")
+    for i in range(1, 4):
+        fig.update_xaxes(title_text="Time", row=2, col=i)
+        fig.update_yaxes(title_text="Population", row=2, col=i)
+
+    fig.update_layout(height=800, width=1200, title_text="Advanced Cell Evolution Simulation")
 
     return fig
 
@@ -247,6 +269,11 @@ if st.button("Run Simulation"):
             
             for i, (cell_type, counts) in enumerate(population_history.items()):
                 fig.data[i+5].y = counts  # +5 because we have 5 cell types in the first subplot
+                if cell_type != "modified" and cell_type != "plant_like":
+                    fig.data[i+10].y = counts  # Update individual population charts
+                else:
+                    fig.data[13].y = population_history["plant_like"]
+                    fig.data[14].y = population_history["modified"]
         
         fig.layout.title.text = f"Advanced Cell Evolution Simulation (Time: {env.time})"