Update my_pages/multiverse.py

my_pages/multiverse.py

@@ -165,14 +165,14 @@ def render():
     # --- 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]
+        # find all nodes belonging to this stage (x == stage_idx)
+        stage_nodes = [i for i, (x, y) in enumerate(positions) if x == float(stage_idx)]
         if not stage_nodes:
             continue
-        # find vertical max position among this stage's nodes
+        # max y among these 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
+        y = max_y + 0.2 * spread_factor  # offset above top node
         stage_label_traces.append(go.Scatter(
             x=[x], y=[y],
             text=[stage["label"]],
@@ -182,6 +182,7 @@ def render():
             showlegend=False
         ))
 
+
     # --- Render figure ---
     # fig = go.Figure(data=edge_traces + [node_trace])
     fig = go.Figure(data=edge_traces + stage_label_traces + [node_trace])