Update my_pages/multiverse.py

my_pages/multiverse.py

@@ -54,7 +54,7 @@ def build_tree_and_trace_path(selected_path, spread_factor=500):
         next_nodes = []
 
         # scaling: huge spread at stage 1, tapering off deeper
-        scale = spread_factor ** (1.0 / stage_idx)
+        scale = spread_factor ** (1.0 / stage_idx**(0.5))
 
         for parent_order, parent_idx in enumerate(prev_nodes):
             px, py = node_positions[parent_idx]
@@ -162,8 +162,29 @@ def render():
             hoverinfo='none'
         ))
 
+    # --- Add stage labels at the top of each layer ---
+    stage_label_traces = []
+    for stage_idx, stage in enumerate(choices_list, start=1):
+        # all nodes in this stage
+        stage_nodes = [i for i, s in enumerate(range(len(positions))) if i < len(node_stage) and node_stage[i] == stage_idx]
+        if not stage_nodes:
+            continue
+        # find vertical max position among this stage's nodes
+        max_y = max(positions[i][1] for i in stage_nodes)
+        x = float(stage_idx)
+        y = max_y + 0.2 * spread_factor  # small offset above top node
+        stage_label_traces.append(go.Scatter(
+            x=[x], y=[y],
+            text=[stage["label"]],
+            mode="text",
+            textfont=dict(size=16, color="white"),
+            hoverinfo="none",
+            showlegend=False
+        ))
+
     # --- Render figure ---
-    fig = go.Figure(data=edge_traces + [node_trace])
+    # fig = go.Figure(data=edge_traces + [node_trace])
+    fig = go.Figure(data=edge_traces + stage_label_traces + [node_trace])
     fig.update_layout(
         showlegend=False,
         xaxis=dict(visible=False),